File size: 2,238 Bytes
b16c3a6 | 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 | #!/usr/bin/env python3
"""HumanEval pass@1 through vLLM (AR path of a dual-head bundle). Same prompts / truncation / checker as humaneval.py.
python3 scripts/vllm_humaneval.py --model exports/jev-judge-qwen35-9b-v0.8 --out reports/humaneval_vllm_9b.json
"""
from __future__ import annotations
import argparse
import json
import os
import sys
import time
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from humaneval import load_problems, run_check, truncate # noqa: E402
def main() -> None:
ap = argparse.ArgumentParser()
ap.add_argument("--model", required=True)
ap.add_argument("--out", required=True)
ap.add_argument("--max-new-tokens", type=int, default=384)
ap.add_argument("--gpu-mem", type=float, default=0.85)
args = ap.parse_args()
from vllm import LLM, SamplingParams
problems = load_problems()
t0 = time.time()
llm = LLM(model=args.model, dtype="bfloat16", gpu_memory_utilization=args.gpu_mem, max_model_len=4096,
enable_prefix_caching=False)
load_s = time.time() - t0
sp = SamplingParams(temperature=0.0, max_tokens=args.max_new_tokens)
t0 = time.time()
outs = llm.generate([p["prompt"] for p in problems], sp, use_tqdm=False)
gen_s = time.time() - t0
n_tok = sum(len(o.outputs[0].token_ids) for o in outs)
results, passed = [], 0
for p, o in zip(problems, outs):
comp = truncate(o.outputs[0].text)
ok, err = run_check(p["prompt"] + comp + "\n\n" + p["test"] + "\n" + f"check({p['entry_point']})\n")
passed += ok
results.append({"task_id": p["task_id"], "passed": ok, "completion": comp, "error": err})
summary = {"model": args.model, "engine": "vllm", "n": len(problems), "passed": passed, "pass@1": passed / len(problems),
"generation_s": gen_s, "generated_tokens": n_tok, "gen_tok_per_s": n_tok / gen_s, "load_s": load_s, "results": results}
with open(args.out, "w") as f:
json.dump(summary, f, indent=1)
print(f"vLLM HumanEval pass@1 = {passed/len(problems):.4f} ({passed}/{len(problems)}) | generation {gen_s:.1f}s "
f"({n_tok/gen_s:.0f} generated tok/s) | load {load_s:.0f}s -> {args.out}")
if __name__ == "__main__":
main()
|