Upload code/step1_filter_and_greedy.py with huggingface_hub

code/step1_filter_and_greedy.py

@@ -0,0 +1,213 @@
+"""
+Step 1: Filter MATH-500 to 20 level 1-3 problems and generate greedy (N=1) solutions.
+
+This script:
+1. Loads the MATH-500 dataset and filters to level 1-3 problems
+2. Randomly samples 20 problems (with a fixed seed for reproducibility)
+3. Generates a single greedy solution per problem using Qwen2.5-1.5B-Instruct
+4. Extracts answers from \boxed{} format and computes accuracy
+5. Saves results as JSON for the next steps
+
+Co-authored with Claude (Anthropic) — used for structuring the pipeline and
+prompt engineering. I can explain all code logic.
+"""
+
+import json
+import os
+import random
+import torch
+from datasets import load_dataset
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from typing import Optional
+
+
+# ──────────────────────────────────────────────────────────────────────────────
+# Helper: Extract answer from \boxed{...}
+# Source: https://gist.github.com/lewtun/9c2ce1937b741404090a3dc4c7c022b3
+# ──────────────────────────────────────────────────────────────────────────────
+def extract_boxed_solution(text: str) -> Optional[str]:
+    """
+    Extracts the content of the last \\boxed{} in a given LaTeX-style text.
+    Uses bracket-balanced parsing to handle nested braces correctly.
+    """
+    try:
+        start_index = text.rindex("\\boxed{")
+        content_start = start_index + 7
+        bracket_count = 1
+        current_pos = content_start
+
+        while bracket_count > 0 and current_pos < len(text):
+            if text[current_pos] == "{":
+                bracket_count += 1
+            elif text[current_pos] == "}":
+                bracket_count -= 1
+            current_pos += 1
+
+        if bracket_count == 0:
+            content = text[content_start : current_pos - 1].strip()
+            return content
+        else:
+            return None
+    except ValueError:
+        return None
+    except Exception:
+        return None
+
+
+# ──────────────────────────────────────────────────────────────────────────────
+# Step 1a: Filter dataset to level 1-3 and sample 20 problems
+# ──────────────────────────────────────────────────────────────────────────────
+print("=" * 70)
+print("STEP 1: Loading and filtering MATH-500 dataset")
+print("=" * 70)
+
+dataset = load_dataset("HuggingFaceH4/MATH-500", split="test")
+print(f"Total problems in MATH-500: {len(dataset)}")
+
+# Filter to levels 1-3 (easier problems suitable for small models)
+filtered = dataset.filter(lambda x: x["level"] in [1, 2, 3])
+print(f"Problems at levels 1-3: {len(filtered)}")
+
+# Sample 20 problems with a fixed seed for reproducibility
+random.seed(42)
+indices = random.sample(range(len(filtered)), k=20)
+problems = filtered.select(indices)
+
+# Display the selected problems
+print(f"\nSelected {len(problems)} problems:")
+for i, p in enumerate(problems):
+    print(f"  [{i+1}] Level {p['level']} | {p['subject']} | {p['unique_id']}")
+    print(f"       Answer: {p['answer']}")
+    # Show first 80 chars of problem
+    preview = p["problem"][:80].replace("\n", " ")
+    print(f"       Problem: {preview}...")
+
+# Save filtered problems for later steps
+problems_data = [
+    {
+        "idx": i,
+        "problem": p["problem"],
+        "solution": p["solution"],
+        "answer": p["answer"],
+        "subject": p["subject"],
+        "level": p["level"],
+        "unique_id": p["unique_id"],
+    }
+    for i, p in enumerate(problems)
+]
+
+os.makedirs("/Users/cmpatino/Projects/ml-intern/exercise/outputs", exist_ok=True)
+with open("/Users/cmpatino/Projects/ml-intern/exercise/outputs/filtered_problems.json", "w") as f:
+    json.dump(problems_data, f, indent=2)
+print(f"\nSaved {len(problems_data)} problems to outputs/filtered_problems.json")
+
+# ──────────────────────────────────────────────────────────────────────────────
+# Step 1b: Generate greedy (N=1) solutions
+# ──────────────────────────────────────────────────────────────────────────────
+print("\n" + "=" * 70)
+print("STEP 2: Generating greedy solutions with Qwen2.5-1.5B-Instruct")
+print("=" * 70)
+
+MODEL_ID = "Qwen/Qwen2.5-1.5B-Instruct"
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    torch_dtype=torch.bfloat16,
+    device_map="auto",
+)
+
+# System prompt encouraging chain-of-thought and \boxed{} format
+SYSTEM_PROMPT = (
+    "You are a helpful math assistant. Solve the problem step by step, "
+    "showing your reasoning clearly. Put your final answer inside "
+    "\\boxed{answer} at the end of your solution."
+)
+
+
+def generate_solutions(problems_data, model, tokenizer, n=1, temperature=None, do_sample=False):
+    """
+    Generate n solutions per problem.
+
+    Args:
+        problems_data: list of problem dicts
+        model: the language model
+        tokenizer: the tokenizer
+        n: number of solutions to generate per problem
+        temperature: sampling temperature (None for greedy)
+        do_sample: whether to sample (False = greedy)
+
+    Returns:
+        list of dicts with problem info + generated solutions
+    """
+    results = []
+
+    for i, p in enumerate(problems_data):
+        print(f"\n  Generating for problem {i+1}/{len(problems_data)}: {p['unique_id']}")
+
+        # Format the chat prompt
+        messages = [
+            {"role": "system", "content": SYSTEM_PROMPT},
+            {"role": "user", "content": p["problem"]},
+        ]
+        prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+
+        # Generation kwargs
+        gen_kwargs = {
+            "max_new_tokens": 2048,
+            "do_sample": do_sample,
+        }
+        if do_sample and temperature is not None:
+            gen_kwargs["temperature"] = temperature
+
+        solutions = []
+        for j in range(n):
+            with torch.no_grad():
+                output = model.generate(**inputs, **gen_kwargs)
+
+            # Decode only the generated tokens (exclude the prompt)
+            generated = output[0][inputs["input_ids"].shape[1]:]
+            solution_text = tokenizer.decode(generated, skip_special_tokens=True)
+            solutions.append(solution_text)
+
+            if n > 1 and (j + 1) % 4 == 0:
+                print(f"    Generated {j+1}/{n} solutions")
+
+        result = {**p, "generated_solutions": solutions}
+        results.append(result)
+
+    return results
+
+
+# Generate greedy solutions (N=1, no sampling)
+greedy_results = generate_solutions(problems_data, model, tokenizer, n=1, do_sample=False)
+
+# ──────────────────────────────────────────────────────────────────────────────
+# Step 1c: Evaluate greedy accuracy
+# ──────────────────────────────────────────────────────────────────────────────
+print("\n" + "=" * 70)
+print("STEP 3: Evaluating greedy accuracy")
+print("=" * 70)
+
+correct = 0
+for r in greedy_results:
+    extracted = extract_boxed_solution(r["generated_solutions"][0])
+    r["greedy_extracted_answer"] = extracted
+    r["greedy_correct"] = (extracted is not None) and (extracted == r["answer"])
+    if r["greedy_correct"]:
+        correct += 1
+    status = "✓" if r["greedy_correct"] else "✗"
+    print(f"  {status} [{r['unique_id']}] Expected: {r['answer']} | Got: {extracted}")
+
+greedy_accuracy = correct / len(greedy_results)
+print(f"\nGreedy accuracy: {correct}/{len(greedy_results)} = {greedy_accuracy:.1%}")
+
+# Save greedy results
+with open("/Users/cmpatino/Projects/ml-intern/exercise/outputs/greedy_results.json", "w") as f:
+    json.dump(greedy_results, f, indent=2)
+print("Saved greedy results to outputs/greedy_results.json")
+
+# Clean up model to free memory for PRM scoring
+del model
+torch.cuda.empty_cache()
+print("\nFreed LLM memory. Ready for Step 2 (sampling + PRM scoring).")