| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| 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}%") |
|
|
