File size: 2,761 Bytes
02c4d82 | 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 74 75 76 77 78 79 80 81 82 83 84 | # /// 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}%")
|