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

import os
import re
import json
from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
from tqdm import tqdm
import torch

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

BASE_MODEL = os.getenv("BASE_MODEL", "Qwen/Qwen3.5-4B")
ADAPTER_REPO = os.getenv("ADAPTER_REPO", "oExz/qwen3.5-4b-reasoning-lora")
MAX_SAMPLES = int(os.getenv("MAX_SAMPLES", "250"))

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):
    return (
        "<|im_start|>system\nYou are a helpful assistant that solves math problems. "
        "Show your reasoning step by step, then give the final answer as ANSWER: <number>.\n"
        "<|im_end|>\n"
        "<|im_start|>user\n" + question + "\n<|im_end|>\n"
        "<|im_start|>assistant\n"
    )

if __name__ == '__main__':
    print(f"Loading GSM8K dataset (max {MAX_SAMPLES} samples)...")
    ds = load_dataset("gsm8k", "main", split="test")
    if MAX_SAMPLES:
        ds = ds.select(range(min(MAX_SAMPLES, len(ds))))

    questions = [build_prompt(ex["question"]) for ex in ds]
    ground_truth = [extract_answer(ex["answer"]) for ex in ds]

    print(f"Loading base model {BASE_MODEL} + adapter {ADAPTER_REPO}...")
    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)
    
    model = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL,
        trust_remote_code=True,
        torch_dtype="auto",
        device_map="auto",
    )
    model = PeftModel.from_pretrained(model, ADAPTER_REPO)
    model = model.merge_and_unload()
    model.eval()

    print(f"Generating answers for {len(questions)} questions...")
    correct = 0
    results = []
    
    for i in tqdm(range(len(questions))):
        inputs = tokenizer(questions[i], return_tensors="pt").to(model.device)
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=512,
                do_sample=False,
                pad_token_id=tokenizer.pad_token_id,
                eos_token_id=tokenizer.eos_token_id,
            )
        pred_text = tokenizer.decode(outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)
        pred_answer = extract_answer(pred_text)
        gt = ground_truth[i]
        is_correct = pred_answer == gt
        if is_correct:
            correct += 1
        results.append({
            "question": ds[i]["question"],
            "ground_truth": gt,
            "predicted": pred_answer,
            "correct": is_correct,
            "raw_output": pred_text,
        })

    accuracy = correct / len(ds) * 100
    print(f"\n=== GSM8K TUNED EVAL RESULTS ===")
    print(f"Correct: {correct}/{len(ds)}")
    print(f"Accuracy: {accuracy:.2f}%")
    
    with open("/tmp/gsm8k_tuned_results.json", "w") as f:
        json.dump({
            "base_model": BASE_MODEL,
            "adapter_repo": ADAPTER_REPO,
            "accuracy": accuracy,
            "correct": correct,
            "total": len(ds),
            "results": results,
        }, f, indent=2)
    
    print("Results saved to /tmp/gsm8k_tuned_results.json")