# /// script
# dependencies = [
#     "vllm",
#     "transformers>=4.40.0",
#     "datasets>=2.14",
#     "pyarrow<19.0",
#     "torch>=2.0.0",
#     "huggingface-hub",
#     "tqdm>=4.65",
# ]
# ///

import os
import re
import json
from datasets import load_dataset
from vllm import LLM, SamplingParams
from transformers import AutoTokenizer

hf_token = os.getenv("HF_TOKEN")
if hf_token:
    os.environ.setdefault("HF_HUB_TOKEN", hf_token)

MODEL = "Qwen/Qwen3.5-4B"
N_SAMPLES = 20

def extract_answer(text):
    text = text.strip()
    m = re.search(r"####\s*(-?\d+(?:\.\d+)?)", text)
    if m: return m.group(1).strip()
    m = re.search(r"ANSWER:\s*(-?\d+(?:\.\d+)?)", text, re.IGNORECASE)
    if m: return m.group(1).strip()
    nums = re.findall(r"-?\d+(?:\.\d+)?", text)
    if nums: return nums[-1].strip()
    return ""

def build_prompt(question):
    messages = [
        {"role": "system", "content": "You are a helpful assistant that thinks step by step."},
        {"role": "user", "content": question},
    ]
    return messages

if __name__ == '__main__':
    print("Loading GSM8K...")
    ds = load_dataset("gsm8k", "main", split="test").select(range(N_SAMPLES))

    print(f"Loading {MODEL}...")
    tokenizer = AutoTokenizer.from_pretrained(MODEL, trust_remote_code=True)
    llm = LLM(model=MODEL, trust_remote_code=True, dtype="half", gpu_memory_utilization=0.9)

    questions = []
    ground_truth = []
    for ex in ds:
        msgs = build_prompt(ex["question"])
        formatted = tokenizer.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
        questions.append(formatted)
        gt_match = re.search(r"####\s*(-?\d[\d,]*\.?\d*)", ex["answer"])
        gt = gt_match.group(1).replace(",", "") if gt_match else ex["answer"].strip().split()[-1]
        ground_truth.append(gt)

    sampling_params = SamplingParams(temperature=0.7, top_p=0.9, max_tokens=512, stop=["<|im_end|>", "<|endoftext|>"])
    outputs = llm.generate(questions, sampling_params)

    print("\n" + "=" * 70)
    print("BASELINE MODEL OUTPUTS (first 20 samples)")
    print("=" * 70)

    for i, out in enumerate(outputs):
        raw = out.outputs[0].text.strip()
        pred = extract_answer(raw)
        gt = ground_truth[i]
        match = pred == gt

        print(f"\n--- Question {i+1} ---")
        print(f"Q: {ds[i]['question']}")
        print(f"GT: {gt} | Pred: {pred} | Match: {match}")
        print(f"--- RAW OUTPUT ---")
        print(raw[:800] + ("..." if len(raw) > 800 else ""))
        print("-" * 70)

    correct = sum(1 for i in range(N_SAMPLES) if extract_answer(outputs[i].outputs[0].text.strip()) == ground_truth[i])
    print(f"\nScore on first {N_SAMPLES}: {correct}/{N_SAMPLES} = {correct/N_SAMPLES*100:.1f}%")