File size: 3,354 Bytes
7ea035c | 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 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 | # /// 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
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 = os.getenv("MODEL", "Qwen/Qwen3.5-4B")
MAX_SAMPLES = int(os.getenv("MAX_SAMPLES", "250"))
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(f"Loading GSM8K dataset (max {MAX_SAMPLES} samples)...")
ds = load_dataset("gsm8k", "main", split="test")
if MAX_SAMPLES:
ds = ds.select(range(min(MAX_SAMPLES, len(ds))))
print(f"Loading model {MODEL} with vLLM...")
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)
ground_truth.append(extract_answer(ex["answer"]))
# SAME sampling params as tuned eval (temperature=0.7, do_sample)
sampling_params = SamplingParams(
temperature=0.7,
top_p=0.9,
max_tokens=512,
stop=["<|im_end|>", "<|endoftext|>"],
)
print(f"Generating answers for {len(questions)} questions...")
outputs = llm.generate(questions, sampling_params)
correct = 0
results = []
for i, out in enumerate(outputs):
pred_text = out.outputs[0].text
pred_answer = extract_answer(pred_text)
gt = ground_truth[i]
is_correct = pred_answer == gt
if is_correct:
correct += 1
results.append({
"question": ds[i]["question"],
"ground_truth": gt,
"predicted": pred_answer,
"correct": is_correct,
"raw_output": pred_text,
})
if i < 5:
print(f"Q{i+1}: GT={gt} Pred={pred_answer} OK={is_correct}")
accuracy = correct / len(ds) * 100
print(f"\n=== GSM8K BASELINE (FAIR) ===")
print(f"Correct: {correct}/{len(ds)}")
print(f"Accuracy: {accuracy:.2f}%")
import json
with open("/tmp/gsm8k_baseline_fair_results.json", "w") as f:
json.dump({
"model": MODEL,
"accuracy": accuracy,
"correct": correct,
"total": len(ds),
"results": results,
}, f, indent=2)
print("Results saved to /tmp/gsm8k_baseline_fair_results.json")
|