| """HumanEval (Python coding): complete each function, run its unit tests, pass@1 greedy.""" |
| import argparse |
| import os |
| import subprocess |
| import sys |
| import tempfile |
|
|
| from datasets import load_dataset |
|
|
| from ar.data.tokenizer import Tokenizer |
| from ar.eval.batch_gen import generate_batch |
| from ar.eval.generate import load_model |
|
|
| |
| STOP = ["\ndef ", "\nclass ", "\nif __name__", "\nprint(", "\n#"] |
|
|
|
|
| def passes(program, timeout=10): |
| |
| |
| with tempfile.NamedTemporaryFile("w", suffix=".py", delete=False) as f: |
| f.write(program) |
| path = f.name |
| try: |
| return subprocess.run([sys.executable, path], capture_output=True, |
| timeout=timeout).returncode == 0 |
| except subprocess.TimeoutExpired: |
| return False |
| finally: |
| os.remove(path) |
|
|
|
|
| if __name__ == "__main__": |
| p = argparse.ArgumentParser() |
| p.add_argument("--model", required=True) |
| p.add_argument("--batch", type=int, default=16) |
| args = p.parse_args() |
|
|
| device = "cuda" |
| tok = Tokenizer() |
| model = load_model(args.model, device) |
| ds = load_dataset("openai/openai_humaneval", split="test") |
|
|
| solved = 0 |
| for i in range(0, len(ds), args.batch): |
| chunk = ds.select(range(i, min(i + args.batch, len(ds)))) |
| outs = generate_batch(model, tok, [ex["prompt"] for ex in chunk], seq_len=1536, |
| max_new=384, stop=STOP, device=device) |
| for ex, completion in zip(chunk, outs): |
| program = (ex["prompt"] + completion + "\n\n" + ex["test"] |
| + f"\n\ncheck({ex['entry_point']})\n") |
| solved += passes(program) |
| done = min(i + args.batch, len(ds)) |
| print(f"[{done}/{len(ds)}] pass@1 so far: {100 * solved / done:.1f}%", flush=True) |
| if i == 0: |
| print("--- example completion ---\n" + chunk[0]["prompt"] + outs[0] |
| + "\n--------------------------", flush=True) |
|
|
| print(f"HumanEval pass@1 (greedy): {100 * solved / len(ds):.1f}% ({solved}/{len(ds)})") |
