Upload analyze_baseline_small.py with huggingface_hub

analyze_baseline_small.py

@@ -0,0 +1,72 @@
+# /// 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()