File size: 2,218 Bytes
97fea6f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 | # /// 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()
|