Upload eval_gsm8k_baseline_fair.py with huggingface_hub

eval_gsm8k_baseline_fair.py

@@ -0,0 +1,113 @@
+# /// 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
+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 = os.getenv("MODEL", "Qwen/Qwen3.5-4B")
+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):
+    messages = [
+        {"role": "system", "content": "You are a helpful assistant that thinks step by step."},
+        {"role": "user", "content": question},
+    ]
+    return messages
+
+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))))
+
+    print(f"Loading model {MODEL} with vLLM...")
+    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)
+        ground_truth.append(extract_answer(ex["answer"]))
+
+    # SAME sampling params as tuned eval (temperature=0.7, do_sample)
+    sampling_params = SamplingParams(
+        temperature=0.7,
+        top_p=0.9,
+        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,
+        })
+        if i < 5:
+            print(f"Q{i+1}: GT={gt} Pred={pred_answer} OK={is_correct}")
+
+    accuracy = correct / len(ds) * 100
+    print(f"\n=== GSM8K BASELINE (FAIR) ===")
+    print(f"Correct: {correct}/{len(ds)}")
+    print(f"Accuracy: {accuracy:.2f}%")
+
+    import json
+    with open("/tmp/gsm8k_baseline_fair_results.json", "w") as f:
+        json.dump({
+            "model": MODEL,
+            "accuracy": accuracy,
+            "correct": correct,
+            "total": len(ds),
+            "results": results,
+        }, f, indent=2)
+    print("Results saved to /tmp/gsm8k_baseline_fair_results.json")