local reward model

1 files changed, 382 insertions, 0 deletions

diff --git a/scripts/test_reward_comparison.py b/scripts/test_reward_comparison.py
new file mode 100644
index 0000000..baddbdf
--- /dev/null
+++ b/scripts/test_reward_comparison.py
@@ -0,0 +1,382 @@
+#!/usr/bin/env python3
+"""
+Compare GPT-4o-mini vs Llama-3.1-8B for reward classification.
+
+Tests both models on the same scenarios using the same prompt.
+"""
+import argparse
+import asyncio
+import json
+import os
+import sys
+import time
+from dataclasses import dataclass
+from typing import List, Dict, Optional, Tuple
+
+# Add src to path
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'src'))
+
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from openai import AsyncOpenAI
+
+# Same prompt as llm_reward.py
+JUDGE_SYSTEM_PROMPT = """\
+You are a feedback classifier. Given a user query (q_t), the assistant's response (a_t), \
+and the user's next message (q_{t+1}), classify the user's follow-up into exactly one label.
+
+Labels (mutually exclusive):
+- neg_constraint_restate: User reasserts constraints/preferences as correction (e.g., "as I said…", "remember…", "按我说的…").
+- neg_correction: User indicates the content is wrong or the assistant failed to answer.
+- neg_confusion: User indicates confusion or requests re-explanation.
+- pos_praise: Explicit praise or satisfaction with the response.
+- pos_progress: Constructive continuation (examples, extensions, what-if, next steps) without complaint.
+- neutral: Ambiguous or minimal feedback, not clearly positive or negative.
+- topic_shift: User switches to a new unrelated task/topic.
+
+Output a JSON object with fields: label, confidence (0-1), rationale (one short sentence)."""
+
+JUDGE_USER_TEMPLATE = """\
+q_t: {query_t}
+
+a_t: {answer_t}
+
+q_{{t+1}}: {query_t1}"""
+
+REWARD_MAP = {
+    "neg_constraint_restate": -1.0,
+    "neg_correction": -0.8,
+    "neg_confusion": -0.6,
+    "pos_praise": +0.8,
+    "pos_progress": +0.1,
+    "neutral": 0.0,
+    "topic_shift": 0.0,
+}
+
+VALID_LABELS = set(REWARD_MAP.keys())
+
+# Test cases with expected labels
+TEST_CASES = [
+    {
+        "name": "neg_constraint_restate - format preference",
+        "query_t": "Explain how sorting works in Python. Please use bullet points.",
+        "answer_t": "Sorting in Python can be done using the sorted() function or the list.sort() method. The sorted() function returns a new sorted list, while sort() modifies the list in place. Both accept a key parameter for custom sorting and a reverse parameter for descending order.",
+        "query_t1": "I asked for bullet points. Can you reformat that with bullet points please?",
+        "expected": "neg_constraint_restate",
+    },
+    {
+        "name": "neg_constraint_restate - step by step",
+        "query_t": "Solve x^2 - 5x + 6 = 0. Show step by step.",
+        "answer_t": "The solutions are x = 2 and x = 3.",
+        "query_t1": "As I said, I need to see the step-by-step solution, not just the answer.",
+        "expected": "neg_constraint_restate",
+    },
+    {
+        "name": "neg_correction - wrong answer",
+        "query_t": "What is the capital of Australia?",
+        "answer_t": "The capital of Australia is Sydney.",
+        "query_t1": "That's incorrect. Sydney is not the capital of Australia.",
+        "expected": "neg_correction",
+    },
+    {
+        "name": "neg_confusion - unclear explanation",
+        "query_t": "What is recursion in programming?",
+        "answer_t": "Recursion is when a function calls itself in a self-similar way to solve problems.",
+        "query_t1": "I'm confused. What do you mean by 'self-similar way'? Can you explain more clearly?",
+        "expected": "neg_confusion",
+    },
+    {
+        "name": "pos_praise - explicit thanks",
+        "query_t": "How do I center a div in CSS?",
+        "answer_t": "You can center a div using flexbox: set the parent to `display: flex; justify-content: center; align-items: center;`. Alternatively, use `margin: 0 auto;` for horizontal centering with a defined width.",
+        "query_t1": "Perfect, thank you! That's exactly what I needed.",
+        "expected": "pos_praise",
+    },
+    {
+        "name": "pos_praise - great explanation",
+        "query_t": "Explain how photosynthesis works.",
+        "answer_t": "Photosynthesis is the process by which plants convert sunlight, water, and CO2 into glucose and oxygen. It occurs in chloroplasts, with light-dependent reactions in the thylakoid membrane and the Calvin cycle in the stroma.",
+        "query_t1": "Great explanation! This really helped me understand the concept.",
+        "expected": "pos_praise",
+    },
+    {
+        "name": "pos_progress - follow-up question",
+        "query_t": "What is a binary search tree?",
+        "answer_t": "A binary search tree (BST) is a data structure where each node has at most two children. The left subtree contains only nodes with values less than the parent, and the right subtree only nodes with values greater than the parent.",
+        "query_t1": "Interesting! How would I implement insertion into a BST?",
+        "expected": "pos_progress",
+    },
+    {
+        "name": "pos_progress - extension",
+        "query_t": "How do I read a file in Python?",
+        "answer_t": "Use `with open('file.txt', 'r') as f: content = f.read()`. The 'with' statement ensures the file is properly closed.",
+        "query_t1": "Got it. What if I want to read it line by line instead?",
+        "expected": "pos_progress",
+    },
+    {
+        "name": "neutral - minimal response",
+        "query_t": "What's 2 + 2?",
+        "answer_t": "2 + 2 = 4",
+        "query_t1": "Ok.",
+        "expected": "neutral",
+    },
+    {
+        "name": "topic_shift - new topic",
+        "query_t": "What is the Pythagorean theorem?",
+        "answer_t": "The Pythagorean theorem states that in a right triangle, a² + b² = c², where c is the hypotenuse.",
+        "query_t1": "By the way, can you help me write a poem about nature?",
+        "expected": "topic_shift",
+    },
+    {
+        "name": "neg_constraint_restate - language preference",
+        "query_t": "Explain machine learning in simple terms.",
+        "answer_t": "Machine learning is a subset of artificial intelligence that uses statistical techniques to enable computers to learn from data. It involves training models on datasets to make predictions or decisions without being explicitly programmed for specific tasks.",
+        "query_t1": "Remember I asked for simple terms? That's too technical. Can you explain like I'm 5?",
+        "expected": "neg_constraint_restate",
+    },
+    {
+        "name": "neg_correction - incomplete answer",
+        "query_t": "List all the planets in our solar system.",
+        "answer_t": "The planets are Mercury, Venus, Earth, Mars, Jupiter, and Saturn.",
+        "query_t1": "You're missing Uranus and Neptune. There are 8 planets, not 6.",
+        "expected": "neg_correction",
+    },
+]
+
+
+def parse_json_response(text: str) -> Dict:
+    """Parse JSON from model response, handling markdown code blocks."""
+    text = text.strip()
+    # Remove markdown code blocks if present
+    if text.startswith("```"):
+        lines = text.split("\n")
+        text = "\n".join(lines[1:-1] if lines[-1].strip() == "```" else lines[1:])
+    try:
+        return json.loads(text)
+    except json.JSONDecodeError:
+        # Try to find JSON object in text
+        import re
+        match = re.search(r'\{[^}]+\}', text, re.DOTALL)
+        if match:
+            try:
+                return json.loads(match.group())
+            except:
+                pass
+        return {"label": "neutral", "confidence": 0.0, "rationale": "parse_error"}
+
+
+class LocalLLMJudge:
+    """Local LLM judge using Qwen2.5-1.5B."""
+
+    def __init__(self, model_path: str, device: str = "cuda"):
+        self.device = device
+        print(f"Loading {model_path}...")
+        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model_path,
+            torch_dtype=torch.bfloat16,
+            device_map=device,
+        )
+        self.model.eval()
+        print("Model loaded.")
+
+    def judge(self, query_t: str, answer_t: str, query_t1: str) -> Dict:
+        """Classify a single turn."""
+        user_content = JUDGE_USER_TEMPLATE.format(
+            query_t=query_t,
+            answer_t=answer_t,
+            query_t1=query_t1,
+        )
+
+        messages = [
+            {"role": "system", "content": JUDGE_SYSTEM_PROMPT},
+            {"role": "user", "content": user_content},
+        ]
+
+        text = self.tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True,
+        )
+
+        inputs = self.tokenizer(text, return_tensors="pt").to(self.device)
+
+        with torch.no_grad():
+            outputs = self.model.generate(
+                **inputs,
+                max_new_tokens=256,
+                temperature=0.1,
+                do_sample=True,
+                pad_token_id=self.tokenizer.eos_token_id,
+            )
+
+        response = self.tokenizer.decode(
+            outputs[0][inputs.input_ids.shape[1]:],
+            skip_special_tokens=True,
+        )
+
+        return parse_json_response(response), response
+
+
+class GPTJudge:
+    """GPT-5-nano judge using OpenAI API."""
+
+    def __init__(self):
+        self.client = AsyncOpenAI(
+            api_key=os.getenv("OPENAI_API_KEY"),
+        )
+
+    async def judge(self, query_t: str, answer_t: str, query_t1: str) -> Dict:
+        """Classify a single turn."""
+        user_content = JUDGE_USER_TEMPLATE.format(
+            query_t=query_t,
+            answer_t=answer_t,
+            query_t1=query_t1,
+        )
+
+        messages = [
+            {"role": "system", "content": JUDGE_SYSTEM_PROMPT},
+            {"role": "user", "content": user_content},
+        ]
+
+        response = await self.client.chat.completions.create(
+            model="gpt-4o-mini",
+            messages=messages,
+            max_tokens=256,
+            temperature=0.1,
+            response_format={"type": "json_object"},
+        )
+
+        raw = response.choices[0].message.content
+        return parse_json_response(raw), raw
+
+
+async def run_comparison(local_model_path: str, device: str = "cuda"):
+    """Run comparison between local LLM and GPT-5-nano."""
+
+    print("=" * 80)
+    print("Reward Model Comparison: Llama-3.1-8B vs GPT-4o-mini")
+    print("=" * 80)
+    print()
+
+    # Initialize models
+    local_judge = LocalLLMJudge(local_model_path, device)
+    gpt_judge = GPTJudge()
+
+    results = []
+
+    print(f"Running {len(TEST_CASES)} test cases...\n")
+
+    for i, tc in enumerate(TEST_CASES, 1):
+        print(f"--- Test {i}/{len(TEST_CASES)}: {tc['name']} ---")
+        print(f"Expected: {tc['expected']}")
+
+        # Local model
+        t0 = time.time()
+        local_result, local_raw = local_judge.judge(
+            tc["query_t"], tc["answer_t"], tc["query_t1"]
+        )
+        local_time = time.time() - t0
+
+        # GPT model
+        t0 = time.time()
+        gpt_result, gpt_raw = await gpt_judge.judge(
+            tc["query_t"], tc["answer_t"], tc["query_t1"]
+        )
+        gpt_time = time.time() - t0
+
+        local_label = local_result.get("label", "unknown")
+        local_conf = local_result.get("confidence", 0.0)
+        gpt_label = gpt_result.get("label", "unknown")
+        gpt_conf = gpt_result.get("confidence", 0.0)
+
+        # Check correctness
+        local_correct = local_label == tc["expected"]
+        gpt_correct = gpt_label == tc["expected"]
+        agreement = local_label == gpt_label
+
+        print(f"  Local (Llama): {local_label} (conf={local_conf:.2f}) [{local_time:.2f}s] {'✓' if local_correct else '✗'}")
+        print(f"  GPT-4o-mini:  {gpt_label} (conf={gpt_conf:.2f}) [{gpt_time:.2f}s] {'✓' if gpt_correct else '✗'}")
+        print(f"  Agreement: {'Yes' if agreement else 'NO'}")
+        print()
+
+        results.append({
+            "name": tc["name"],
+            "expected": tc["expected"],
+            "local_label": local_label,
+            "local_conf": local_conf,
+            "local_time": local_time,
+            "local_correct": local_correct,
+            "gpt_label": gpt_label,
+            "gpt_conf": gpt_conf,
+            "gpt_time": gpt_time,
+            "gpt_correct": gpt_correct,
+            "agreement": agreement,
+        })
+
+    # Summary
+    print("=" * 80)
+    print("SUMMARY")
+    print("=" * 80)
+
+    local_accuracy = sum(r["local_correct"] for r in results) / len(results)
+    gpt_accuracy = sum(r["gpt_correct"] for r in results) / len(results)
+    agreement_rate = sum(r["agreement"] for r in results) / len(results)
+
+    local_avg_time = sum(r["local_time"] for r in results) / len(results)
+    gpt_avg_time = sum(r["gpt_time"] for r in results) / len(results)
+
+    print(f"Local (Llama-3.1-8B) Accuracy: {local_accuracy*100:.1f}% ({sum(r['local_correct'] for r in results)}/{len(results)})")
+    print(f"GPT-4o-mini Accuracy:          {gpt_accuracy*100:.1f}% ({sum(r['gpt_correct'] for r in results)}/{len(results)})")
+    print(f"Agreement Rate:                {agreement_rate*100:.1f}%")
+    print()
+    print(f"Local Avg Time: {local_avg_time:.2f}s")
+    print(f"GPT Avg Time:   {gpt_avg_time:.2f}s")
+    print(f"Speedup:        {gpt_avg_time/local_avg_time:.1f}x faster (local)")
+
+    # Show disagreements
+    disagreements = [r for r in results if not r["agreement"]]
+    if disagreements:
+        print()
+        print(f"Disagreements ({len(disagreements)}):")
+        for r in disagreements:
+            print(f"  - {r['name']}: Local={r['local_label']}, GPT={r['gpt_label']}, Expected={r['expected']}")
+
+    # Show errors by model
+    local_errors = [r for r in results if not r["local_correct"]]
+    gpt_errors = [r for r in results if not r["gpt_correct"]]
+
+    if local_errors:
+        print()
+        print(f"Local Model Errors ({len(local_errors)}):")
+        for r in local_errors:
+            print(f"  - {r['name']}: Got {r['local_label']}, Expected {r['expected']}")
+
+    if gpt_errors:
+        print()
+        print(f"GPT Model Errors ({len(gpt_errors)}):")
+        for r in gpt_errors:
+            print(f"  - {r['name']}: Got {r['gpt_label']}, Expected {r['expected']}")
+
+    print()
+    print("=" * 80)
+
+    return results
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--local-model",
+        type=str,
+        default="/projects/bfqt/users/yurenh2/ml-projects/personalization-user-model/models/llama-3.1-8b-instruct",
+        help="Path to local model",
+    )
+    parser.add_argument("--device", type=str, default="cuda")
+    args = parser.parse_args()
+
+    asyncio.run(run_comparison(args.local_model, args.device))
+
+
+if __name__ == "__main__":
+    main()