| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| 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"])) |
|
|
| |
| 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") |
|
|
