Upload analyze_baseline.py with huggingface_hub

analyze_baseline.py

@@ -0,0 +1,83 @@
+# /// 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}%")