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

import os
import re
from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

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

MODEL = "Qwen/Qwen3.5-4B"
N = 10

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 ""

print("Loading dataset...")
ds = load_dataset("gsm8k", "main", split="test").select(range(N))

print(f"Loading {MODEL}...")
tokenizer = AutoTokenizer.from_pretrained(MODEL, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    MODEL, trust_remote_code=True, torch_dtype=torch.bfloat16, device_map="auto",
)
model.eval()

print("\n" + "=" * 70)
print("BASELINE MODEL — FULL OUTPUTS")
print("=" * 70)

for i, ex in enumerate(ds):
    msgs = [
        {"role": "system", "content": "You are a helpful assistant that thinks step by step."},
        {"role": "user", "content": ex["question"]},
    ]
    text = tokenizer.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer(text, return_tensors="pt").to(model.device)

    with torch.no_grad():
        outputs = model.generate(**inputs, max_new_tokens=512, do_sample=True, temperature=0.7, top_p=0.9)

    raw = tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True).strip()
    pred = extract_answer(raw)
    gt_match = re.search(r"####\s*(-?\d[\d,]*\.?\d*)", ex["answer"])
    gt = gt_match.group(1).replace(",", "") if gt_match else ""

    print(f"\n{'─' * 70}")
    print(f"QUESTION {i+1} | GT: {gt} | Pred: {pred} | Match: {pred == gt}")
    print(f"{'─' * 70}")
    print(raw[:1000])
    if len(raw) > 1000:
        print("... [truncated]")
    print()