code/step2_sample_and_score.py

"""
Step 2: Sample N=16 solutions per problem and score with Skywork PRM.

This script:
1. Loads the filtered problems from Step 1
2. Generates N=16 solutions per problem using temperature sampling
3. Loads the Skywork-o1-Open-PRM and scores each solution (last step prediction)
4. Saves all solutions + scores for the Best-of-N computation

The Skywork PRM is loaded using its custom PRM_MODEL class, which wraps
AutoModelForCausalLM with a ValueHead (linear projection to scalar).
The model outputs a sigmoid-normalized score in [0,1] at each step boundary.

Co-authored with Claude (Anthropic). I can explain all code logic.
"""

import json
import os
import sys
import torch
import subprocess
from transformers import AutoTokenizer, AutoModelForCausalLM
from typing import Optional

# ──────────────────────────────────────────────────────────────────────────────
# Helper functions
# ──────────────────────────────────────────────────────────────────────────────
def extract_boxed_solution(text: str) -> Optional[str]:
    """Extract content of the last \\boxed{} in text."""
    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:
            return text[content_start : current_pos - 1].strip()
        return None
    except (ValueError, Exception):
        return None


# ──────────────────────────────────────────────────────────────────────────────
# Load filtered problems
# ──────────────────────────────────────────────────────────────────────────────
print("=" * 70)
print("STEP 2a: Loading problems and generating N=16 solutions per problem")
print("=" * 70)

with open("/Users/cmpatino/Projects/ml-intern/exercise/outputs/filtered_problems.json") as f:
    problems_data = json.load(f)
print(f"Loaded {len(problems_data)} problems")

# ──────────────────────────────────────────────────────────────────────────────
# Generate N=16 solutions per problem with temperature sampling
# ──────────────────────────────────────────────────────────────────────────────
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 = (
    "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."
)

N = 16  # Number of solutions per problem
TEMPERATURE = 0.7  # Sampling temperature — balances diversity vs quality

all_results = []
for i, p in enumerate(problems_data):
    print(f"\n  Problem {i+1}/{len(problems_data)}: {p['unique_id']} (Level {p['level']})")

    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)

    solutions = []
    for j in range(N):
        with torch.no_grad():
            output = model.generate(
                **inputs,
                max_new_tokens=2048,
                do_sample=True,
                temperature=TEMPERATURE,
                top_p=0.95,
            )
        generated = output[0][inputs["input_ids"].shape[1]:]
        solution_text = tokenizer.decode(generated, skip_special_tokens=True)
        solutions.append(solution_text)

        if (j + 1) % 4 == 0:
            print(f"    Generated {j+1}/{N} solutions")

    result = {**p, "sampled_solutions": solutions}
    all_results.append(result)

# Save solutions before scoring (in case PRM loading takes time or fails)
with open("/Users/cmpatino/Projects/ml-intern/exercise/outputs/sampled_solutions.json", "w") as f:
    json.dump(all_results, f, indent=2)
print(f"\nSaved {N} solutions per problem to outputs/sampled_solutions.json")

# Free LLM memory before loading PRM
del model
torch.cuda.empty_cache()
print("Freed LLM memory.")

# ──────────────────────────────────────────────────────────────────────────────
# Score solutions with Skywork PRM
# ──────────────────────────────────────────────────────────────────────────────
print("\n" + "=" * 70)
print("STEP 2b: Scoring solutions with Skywork-o1-Open-PRM")
print("=" * 70)

# Clone the Skywork PRM inference repo for the custom model class
PRM_REPO_PATH = "/Users/cmpatino/Projects/ml-intern/exercise/skywork-o1-prm-inference"
if not os.path.exists(PRM_REPO_PATH):
    print("Cloning Skywork PRM inference repo...")
    subprocess.run(
        ["git", "clone", "https://github.com/SkyworkAI/skywork-o1-prm-inference.git", PRM_REPO_PATH],
        check=True,
    )
sys.path.insert(0, PRM_REPO_PATH)

from model_utils.prm_model import PRM_MODEL
from model_utils.io_utils import prepare_input, prepare_batch_input_for_model, derive_step_rewards

PRM_MODEL_ID = "Skywork/Skywork-o1-Open-PRM-Qwen-2.5-1.5B"

prm_tokenizer = AutoTokenizer.from_pretrained(PRM_MODEL_ID, trust_remote_code=True)
prm_model = PRM_MODEL.from_pretrained(PRM_MODEL_ID, device_map="auto").eval()

print("PRM model loaded successfully.")


def score_solution(problem: str, solution: str) -> list[float]:
    """
    Score a single solution using the PRM.

    Returns a list of per-step scores (sigmoid-normalized, [0,1]).
    The last element is the 'last step prediction' — our final reward.

    The PRM splits the solution by newlines (\n), and assigns a score
    at the end of each step. These scores represent the model's estimate
    of correctness probability at each reasoning step.
    """
    input_ids, steps, reward_flags = prepare_input(problem, solution, prm_tokenizer, step_token="\n")

    # Prepare batch of size 1
    input_ids_t, attention_mask_t, reward_flags_t = prepare_batch_input_for_model(
        [input_ids], [reward_flags], prm_tokenizer.pad_token_id
    )

    # Move to model device
    device = next(prm_model.parameters()).device
    input_ids_t = input_ids_t.to(device)
    attention_mask_t = attention_mask_t.to(device)
    reward_flags_t = reward_flags_t.to(device)

    with torch.no_grad():
        # return_probs=True applies sigmoid internally
        _, _, rewards = prm_model(
            input_ids=input_ids_t,
            attention_mask=attention_mask_t,
            return_probs=True,
        )

    step_rewards = derive_step_rewards(rewards, reward_flags_t)
    return step_rewards[0]  # Return the single sample's step scores


# Score all solutions
print("\nScoring all solutions...")
for i, result in enumerate(all_results):
    print(f"\n  Scoring problem {i+1}/{len(all_results)}: {result['unique_id']}")
    scores = []
    extracted_answers = []

    for j, solution in enumerate(result["sampled_solutions"]):
        # Get PRM score
        step_scores = score_solution(result["problem"], solution)
        # Use last step prediction as the final reward (per DeepMind Appendix E)
        final_score = step_scores[-1] if step_scores else 0.0
        scores.append(final_score)

        # Extract the final answer from \boxed{}
        answer = extract_boxed_solution(solution)
        extracted_answers.append(answer)

        if (j + 1) % 4 == 0:
            print(f"    Scored {j+1}/{N} solutions (last score: {final_score:.4f})")

    result["prm_scores"] = scores
    result["extracted_answers"] = extracted_answers

# Save scored results
with open("/Users/cmpatino/Projects/ml-intern/exercise/outputs/scored_results.json", "w") as f:
    json.dump(all_results, f, indent=2)
print("\nSaved scored results to outputs/scored_results.json")
print("Ready for Step 3 (Best-of-N computation).")