File size: 4,113 Bytes
448ef61 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 | #!/usr/bin/env python3
"""HumanEval pass@1 (greedy, completion-style prompt) for text-only qwen3_5 checkpoints / our bundles.
python3 scripts/humaneval.py --model /root/models/Qwen3.5-9B --out reports/humaneval_base.json
python3 scripts/humaneval.py --model exports/jev-judge-qwen35-9b-v0.8 --out reports/humaneval_v08.json
"""
from __future__ import annotations
import argparse
import json
import os
import subprocess
import sys
import tempfile
import time
import torch
sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "src"))
from jev_judge.model import load_text_causal_lm # noqa: E402
STOPS = ["\nclass ", "\ndef ", "\n#", "\nif __name__", "\nprint(", "\nassert "]
def load_problems() -> list[dict]:
from datasets import load_dataset
ds = load_dataset("openai/openai_humaneval", split="test")
return [dict(r) for r in ds]
def truncate(completion: str) -> str:
cut = len(completion)
for s in STOPS:
i = completion.find(s)
if i != -1:
cut = min(cut, i)
return completion[:cut]
def run_check(program: str, timeout: float = 15.0) -> tuple[bool, str]:
with tempfile.TemporaryDirectory() as d:
path = os.path.join(d, "prog.py")
with open(path, "w") as f:
f.write(program)
try:
r = subprocess.run([sys.executable, path], cwd=d, capture_output=True, text=True, timeout=timeout)
return r.returncode == 0, (r.stderr[-300:] if r.returncode else "")
except subprocess.TimeoutExpired:
return False, "timeout"
@torch.no_grad()
def generate_all(model, tok, prompts: list[str], batch_size: int, max_new_tokens: int) -> list[str]:
tok.padding_side = "left"
outs: list[str] = []
order = sorted(range(len(prompts)), key=lambda i: len(prompts[i]))
result = [""] * len(prompts)
for s in range(0, len(order), batch_size):
idx = order[s : s + batch_size]
batch = [prompts[i] for i in idx]
enc = tok(batch, return_tensors="pt", padding=True, add_special_tokens=False).to("cuda")
with torch.autocast("cuda", dtype=torch.bfloat16):
gen = model.generate(**enc, max_new_tokens=max_new_tokens, do_sample=False, pad_token_id=tok.pad_token_id)
for j, i in enumerate(idx):
result[i] = tok.decode(gen[j, enc["input_ids"].shape[1]:], skip_special_tokens=True)
print(f" generated {min(s + batch_size, len(order))}/{len(order)}", flush=True)
return result
def main() -> None:
ap = argparse.ArgumentParser()
ap.add_argument("--model", required=True)
ap.add_argument("--out", required=True)
ap.add_argument("--batch-size", type=int, default=16)
ap.add_argument("--max-new-tokens", type=int, default=384)
args = ap.parse_args()
from transformers import AutoTokenizer
problems = load_problems()
tok = AutoTokenizer.from_pretrained(args.model)
if tok.pad_token_id is None:
tok.pad_token = tok.eos_token
model, _ = load_text_causal_lm(args.model)
t0 = time.time()
raw = generate_all(model, tok, [p["prompt"] for p in problems], args.batch_size, args.max_new_tokens)
gen_s = time.time() - t0
results = []
passed = 0
for p, r in zip(problems, raw):
comp = truncate(r)
program = p["prompt"] + comp + "\n\n" + p["test"] + "\n" + f"check({p['entry_point']})\n"
ok, err = run_check(program)
passed += ok
results.append({"task_id": p["task_id"], "passed": ok, "completion": comp, "error": err})
score = passed / len(problems)
summary = {"model": args.model, "n": len(problems), "passed": passed, "pass@1": score, "generation_s": gen_s,
"protocol": "greedy, completion-style prompt, stop at " + repr(STOPS), "results": results}
os.makedirs(os.path.dirname(args.out) or ".", exist_ok=True)
with open(args.out, "w") as f:
json.dump(summary, f, indent=1)
print(f"HumanEval pass@1 = {score:.4f} ({passed}/{len(problems)}) | generation {gen_s:.0f}s | -> {args.out}")
if __name__ == "__main__":
main()
|