Upload eval_gsm8k_tuned.py with huggingface_hub

eval_gsm8k_tuned.py

@@ -1,6 +1,5 @@
 # /// script
 # dependencies = [
-#     "vllm",
 #     "transformers>=4.40.0",
 #     "peft>=0.7.0",
 #     "datasets>=2.14",
@@ -16,9 +15,9 @@ import os
 import re
 import json
 from datasets import load_dataset
-from vllm import LLM, SamplingParams
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from peft import PeftModel
+from tqdm import tqdm
 
 hf_token = os.getenv("HF_TOKEN")
 if hf_token:
@@ -30,15 +29,12 @@ MAX_SAMPLES = int(os.getenv("MAX_SAMPLES", "250"))
 
 def extract_answer(text):
     text = text.strip()
-    # Look for #### X format
     m = re.search(r"####\s*(-?\d+(?:\.\d+)?)", text)
     if m:
         return m.group(1).strip()
-    # Look for ANSWER: X format
     m = re.search(r"ANSWER:\s*(-?\d+(?:\.\d+)?)", text, re.IGNORECASE)
     if m:
         return m.group(1).strip()
-    # Last number fallback
     nums = re.findall(r"-?\d+(?:\.\d+)?", text)
     if nums:
         return nums[-1].strip()
@@ -63,50 +59,33 @@ if __name__ == '__main__':
     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)
     
-    # Merge adapter with base model and save to temp dir
-    print("Merging LoRA adapter with base model...")
-    base = AutoModelForCausalLM.from_pretrained(
+    model = AutoModelForCausalLM.from_pretrained(
         BASE_MODEL,
         trust_remote_code=True,
         torch_dtype="auto",
         device_map="auto",
     )
-    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)
-    model = PeftModel.from_pretrained(base, ADAPTER_REPO)
+    model = PeftModel.from_pretrained(model, ADAPTER_REPO)
     model = model.merge_and_unload()
-    
-    temp_dir = "/tmp/merged-qwen-tuned"
-    print(f"Saving merged model to {temp_dir}...")
-    model.save_pretrained(temp_dir)
-    tokenizer.save_pretrained(temp_dir)
-    
-    # Free memory before vLLM loads
-    del base, model
-    import torch
-    torch.cuda.empty_cache()
-    
-    print("Loading merged model with vLLM...")
-    llm = LLM(
-        model=temp_dir,
-        trust_remote_code=True,
-        dtype="half",
-        gpu_memory_utilization=0.9,
-    )
-    
-    sampling_params = SamplingParams(
-        temperature=0,
-        max_tokens=512,
-        stop=["<|im_end|>", "<|endoftext|>"],
-    )
+    model.eval()
 
     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
+    
+    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
@@ -125,7 +104,6 @@ if __name__ == '__main__':
     print(f"Correct: {correct}/{len(ds)}")
     print(f"Accuracy: {accuracy:.2f}%")
     
-    # Save results
     with open("/tmp/gsm8k_tuned_results.json", "w") as f:
         json.dump({
             "base_model": BASE_MODEL,