| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| 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() |
|
|
