Upload eval_gsm8k_tuned.py with huggingface_hub

eval_gsm8k_tuned.py

@@ -0,0 +1,137 @@
+# /// script
+# dependencies = [
+#     "vllm",
+#     "transformers>=4.40.0",
+#     "peft>=0.7.0",
+#     "datasets",
+#     "torch>=2.0.0",
+#     "huggingface-hub",
+#     "accelerate",
+# ]
+# ///
+
+import os
+import re
+import json
+from datasets import load_dataset
+from vllm import LLM, SamplingParams
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from peft import PeftModel
+
+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()
+    # Look for #### X format
+    m = re.search(r"####\s*(-?\d+(?:\.\d+)?)", text)
+    if m:
+        return m.group(1).strip()
+    # Look for ANSWER: X format
+    m = re.search(r"ANSWER:\s*(-?\d+(?:\.\d+)?)", text, re.IGNORECASE)
+    if m:
+        return m.group(1).strip()
+    # Last number fallback
+    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}...")
+    
+    # Merge adapter with base model and save to temp dir
+    print("Merging LoRA adapter with base model...")
+    base = AutoModelForCausalLM.from_pretrained(
+        BASE_MODEL,
+        trust_remote_code=True,
+        torch_dtype="auto",
+        device_map="auto",
+    )
+    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)
+    model = PeftModel.from_pretrained(base, ADAPTER_REPO)
+    model = model.merge_and_unload()
+    
+    temp_dir = "/tmp/merged-qwen-tuned"
+    print(f"Saving merged model to {temp_dir}...")
+    model.save_pretrained(temp_dir)
+    tokenizer.save_pretrained(temp_dir)
+    
+    # Free memory before vLLM loads
+    del base, model
+    import torch
+    torch.cuda.empty_cache()
+    
+    print("Loading merged model with vLLM...")
+    llm = LLM(
+        model=temp_dir,
+        trust_remote_code=True,
+        dtype="half",
+        gpu_memory_utilization=0.9,
+    )
+    
+    sampling_params = SamplingParams(
+        temperature=0,
+        max_tokens=512,
+        stop=["<|im_end|>", "<|endoftext|>"],
+    )
+
+    print(f"Generating answers for {len(questions)} questions...")
+    outputs = llm.generate(questions, sampling_params)
+
+    correct = 0
+    results = []
+    for i, out in enumerate(outputs):
+        pred_text = out.outputs[0].text
+        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}%")
+    
+    # Save results
+    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")