evaluate_humaneval.py

import sys
import json
import traceback
from pathlib import Path
from datasets import load_dataset
import torch
import subprocess
from unsloth import FastLanguageModel
from unsloth.chat_templates import get_chat_template
from peft import PeftModel
from pathlib import Path


# --------------------------------------
# CONFIGURATION
# --------------------------------------

# Model names
BASE_MODEL = "Qwen/Qwen2.5-Coder-7B"
FT_MODEL_DIR = str(Path("qwen25-coder-unsloth/checkpoint-376").resolve())

# Device
device = "cuda:0" if torch.cuda.is_available() else "cpu"

# Output file
OUT_FILE = Path("humaneval_results.txt")

# --------------------------------------
# Prepare dataset
# --------------------------------------

dataset = load_dataset("openai_humaneval", split="test")

# --------------------------------------
# Helper: load and prepare model
# --------------------------------------

def load_unsloth_model(model_name_or_path):
    model, tokenizer = FastLanguageModel.from_pretrained(
        model_name = model_name_or_path,
        max_seq_length = 4096,
        dtype = torch.bfloat16,
        load_in_4bit = True,
    )

    tokenizer = get_chat_template(tokenizer, chat_template="qwen25")

    return model, tokenizer

# --------------------------------------
# Generate code from prompt
# --------------------------------------

def generate_code(model, tokenizer, prompt):
    inputs = tokenizer(prompt, return_tensors="pt").to(device)
    with torch.no_grad():
        output = model.generate(**inputs, max_new_tokens=256)
    return tokenizer.decode(output[0], skip_special_tokens=True)

# --------------------------------------
# Test code for correctness
# --------------------------------------

def run_humaneval_test(generated_code, test_code):
    """
    Combines generated code with HumanEval tests and runs them.
    Returns True if tests pass (exit code 0).
    """
    tmp_file = Path("tmp_humaneval.py")
    with open(tmp_file, "w") as f:
        f.write(generated_code + "\n\n")
        f.write(test_code + "\n")

    # Run python on that file
    proc = subprocess.run(
        [sys.executable, str(tmp_file)],
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        timeout=10,
    )

    return proc.returncode == 0

# --------------------------------------
# Evaluate a model
# --------------------------------------

def evaluate(model, tokenizer, label):
    results = []
    pass1 = 0
    total = len(dataset)

    for i, ex in enumerate(dataset):
        prompt = ex["prompt"]  # HumanEval prompt
        test = ex["test"]      # HumanEval tests

        try:
            gen = generate_code(model, tokenizer, prompt)

            ok = run_humaneval_test(gen, test)

        except Exception:
            ok = False

        results.append((ex["task_id"], ok))
        if ok:
            pass1 += 1

        print(f"[{label}] {i+1}/{total} {ex['task_id']} pass1={ok}")

    return results, pass1

# --------------------------------------
# Main evaluation
# --------------------------------------

with open(OUT_FILE, "w") as out:

    out.write("=== HumanEval Results ===\n\n")

    # --------------------------------------------------
    # Evaluate fine-tuned model FIRST
    # --------------------------------------------------
    out.write("Loading fine-tuned model...\n")
    ft_model, ft_tokenizer = load_unsloth_model(BASE_MODEL)

    ft_model = PeftModel.from_pretrained(
        ft_model,
        FT_MODEL_DIR,
    )

    out.write("Evaluating fine-tuned model...\n")
    ft_results, ft_pass1 = evaluate(ft_model, ft_tokenizer, "FT")
    out.write(f"FT pass@1: {ft_pass1}/{len(dataset)}\n\n")

    # --------------------------------------------------
    # Evaluate base model SECOND
    # --------------------------------------------------
    out.write("Loading base model...\n")
    base_model, base_tokenizer = load_unsloth_model(BASE_MODEL)

    out.write("Evaluating base model...\n")
    base_results, base_pass1 = evaluate(base_model, base_tokenizer, "BASE")
    out.write(f"BASE pass@1: {base_pass1}/{len(dataset)}\n\n")

    out.write("=== Per-task results ===\n")
    out.write("task_id,ft_pass,base_pass\n")
    for (tid_f, f_ok), (tid_b, b_ok) in zip(ft_results, base_results):
        out.write(f"{tid_f},{int(f_ok)},{int(b_ok)}\n")


print("Results written to", OUT_FILE)