| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| import os |
| import re |
| import json |
| from datasets import load_dataset |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
| from peft import PeftModel |
| from tqdm import tqdm |
| import torch |
|
|
| hf_token = os.getenv("HF_TOKEN") |
| if hf_token: |
| os.environ.setdefault("HF_HUB_TOKEN", hf_token) |
|
|
| BASE_MODEL = os.getenv("BASE_MODEL", "Qwen/Qwen3.5-4B") |
| ADAPTER_REPO = os.getenv("ADAPTER_REPO", "oExz/qwen3.5-4b-reasoning-lora") |
| 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): |
| return ( |
| "<|im_start|>system\nYou are a helpful assistant that solves math problems. " |
| "Show your reasoning step by step, then give the final answer as ANSWER: <number>.\n" |
| "<|im_end|>\n" |
| "<|im_start|>user\n" + question + "\n<|im_end|>\n" |
| "<|im_start|>assistant\n" |
| ) |
|
|
| 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)))) |
|
|
| questions = [build_prompt(ex["question"]) for ex in ds] |
| ground_truth = [extract_answer(ex["answer"]) for ex in ds] |
|
|
| print(f"Loading base model {BASE_MODEL} + adapter {ADAPTER_REPO}...") |
| tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True) |
| |
| model = AutoModelForCausalLM.from_pretrained( |
| BASE_MODEL, |
| trust_remote_code=True, |
| torch_dtype="auto", |
| device_map="auto", |
| ) |
| model = PeftModel.from_pretrained(model, ADAPTER_REPO) |
| model = model.merge_and_unload() |
| model.eval() |
|
|
| print(f"Generating answers for {len(questions)} questions...") |
| correct = 0 |
| results = [] |
| |
| for i in tqdm(range(len(questions))): |
| inputs = tokenizer(questions[i], return_tensors="pt").to(model.device) |
| with torch.no_grad(): |
| outputs = model.generate( |
| **inputs, |
| max_new_tokens=512, |
| do_sample=False, |
| pad_token_id=tokenizer.pad_token_id, |
| eos_token_id=tokenizer.eos_token_id, |
| ) |
| pred_text = tokenizer.decode(outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True) |
| 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, |
| }) |
|
|
| accuracy = correct / len(ds) * 100 |
| print(f"\n=== GSM8K TUNED EVAL RESULTS ===") |
| print(f"Correct: {correct}/{len(ds)}") |
| print(f"Accuracy: {accuracy:.2f}%") |
| |
| with open("/tmp/gsm8k_tuned_results.json", "w") as f: |
| json.dump({ |
| "base_model": BASE_MODEL, |
| "adapter_repo": ADAPTER_REPO, |
| "accuracy": accuracy, |
| "correct": correct, |
| "total": len(ds), |
| "results": results, |
| }, f, indent=2) |
| |
| print("Results saved to /tmp/gsm8k_tuned_results.json") |
|
|
