code/scripts/code_eval.py

#!/usr/bin/env python3
"""HumanEval+ (164) and MBPP+ (378) through each model's own chat template (D-59b: the varied-data branches add code,
so the ship rule needs a code test). Thinking off, greedy, EvalPlus's instruct prompt: the problem in a code block,
then the reply opened with EvalPlus's response prefix and a ```python fence (rendered by /apply-template and appended
as plain text, as compare_quality.py does for multiple choice), generated to the closing fence.

The programs run with EvalPlus's own harness (each dataset row's `test`: base and plus inputs with expected outputs and
EvalPlus's float-tolerant assertion) inside one Docker container per model: no network, read-only root, user nobody,
memory and process limits, a time limit per problem. A problem passes if its program passes every test; the score is
pass@1 per benchmark (EvalPlus's "plus" number).

    python scripts/code_eval.py --server .../llama-server --data DIR_WITH_THE_PARQUETS --model NAME=PATH.gguf ...

Per-item results in <out>/<name>.code.jsonl (generations reused on a rerun), summary in <out>/<name>.code.summary.json.
"""
from __future__ import annotations

import argparse
import ast
import json
import re
import signal
import subprocess
import sys
import tempfile
import time
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

import pyarrow.parquet as pq

sys.path.insert(0, str(Path(__file__).resolve().parent))
import compare_quality as cq  # noqa: E402

REPO = Path(__file__).resolve().parents[1]
INSTRUCTION = "Please provide a self-contained Python script that solves the following problem in a markdown code block:"
RESPONSE_PREFIX = ("Below is a Python script with a self-contained function that solves the problem and passes "
                   "corresponding tests:\n```python\n")
IMAGE = "ghcr.io/open-webui/open-webui:main"  # an image already on gen with python3.11 and numpy (no pull needed)
GEN_TOKENS = 1024
# the closing fence, or top-level test code after the solution (smoke test: Qwen3.5-2B wrote asserts to the token cap;
# no "\nassert" stop: MBPP replies often echo the prompt's docstring, which holds an assert)
STOP = ["```", "\nif __name__", "\nprint(", "\n# Test", "\n# Example"]
# Problems whose own reference solution fails in this harness (self-test 28 Oct: 540 of 542 pass): HumanEval/32's
# check unpacks the float find_zero returns; Mbpp/255's reference needs more than the 2 GB per-program limit.
EXCLUDE = {"HumanEval/32", "Mbpp/255"}
RUNNER = r'''
import json, resource, subprocess, sys
from concurrent.futures import ThreadPoolExecutor
n, timeout = int(sys.argv[1]), float(sys.argv[2])
def limits():
    resource.setrlimit(resource.RLIMIT_AS, (2 << 30, 2 << 30))
    resource.setrlimit(resource.RLIMIT_FSIZE, (1 << 20, 1 << 20))
def run(i):
    try:
        p = subprocess.run([sys.executable, "-I", f"/work/{i}.py"], capture_output=True, timeout=timeout,
                           preexec_fn=limits, cwd="/tmp")
        return {"i": i, "status": "pass" if p.returncode == 0 else "fail", "err": p.stderr.decode(errors="replace")[-300:]}
    except subprocess.TimeoutExpired:
        return {"i": i, "status": "timeout", "err": ""}
with ThreadPoolExecutor(8) as ex:
    for r in ex.map(run, range(n)):
        print(json.dumps(r), flush=True)
'''


def items(data: Path) -> list[dict]:
    out = []
    for f in sorted(data.glob("*.parquet")):
        rows = pq.read_table(f).to_pylist()
        for r in rows:
            if "entry_point" in r:  # HumanEval+
                out.append({"task": "humanevalplus", "id": r["task_id"], "code": r["prompt"].strip(),
                            "prefix": r["prompt"], "entry": r["entry_point"], "imports": [], "test": r["test"]})
            else:  # MBPP+: EvalPlus's prompt is the statement plus its first assert, as a docstring
                defs = re.findall(r"^def (\w+)\(", r["code"], re.M)
                called = [d for d in defs if re.search(rf"\b{d}\(", r["test_list"][0])]
                out.append({"task": "mbppplus", "id": f"Mbpp/{r['task_id']}",
                            "code": f'"""\n{r["prompt"]}\n{r["test_list"][0]}\n"""', "prefix": "",
                            "entry": (called or defs)[-1] if (called or defs) else None,
                            "imports": list(r["test_imports"]), "test": r["test"]})
    return [it for it in out if it["id"] not in EXCLUDE]


def ask(port: int, it: dict) -> dict:
    msgs = [{"role": "user", "content": f"{INSTRUCTION}\n```\n{it['code']}\n```\n"}]
    try:
        prompt = cq.post(port, "/apply-template", {"messages": msgs,
                                                     "chat_template_kwargs": {"enable_thinking": False}})["prompt"]
        res = cq.post(port, "/completion", {"prompt": prompt + RESPONSE_PREFIX, "n_predict": GEN_TOKENS,
                                            "temperature": 0.0, "stop": STOP})
        return {"content": res["content"], "n_out": res.get("tokens_predicted"),
                "stopped": bool(res.get("stopped_word") or res.get("stopping_word"))}
    except Exception as e:  # noqa: BLE001 — scored as a failure
        return {"error": str(e)[:300]}


def sanitize(code: str) -> str:
    """EvalPlus-style: the longest prefix that parses (a reply cut at the token cap), then only its imports, classes,
    functions and constant assignments; top-level calls, asserts, prints and __main__ blocks are dropped."""
    lines = code.splitlines()
    for k in range(len(lines), 0, -1):
        try:
            tree = ast.parse("\n".join(lines[:k]))
            break
        except SyntaxError:
            continue
    else:
        return code
    keep = [n for n in tree.body if isinstance(n, (ast.Import, ast.ImportFrom, ast.FunctionDef, ast.AsyncFunctionDef,
                                                   ast.ClassDef))
            or (isinstance(n, (ast.Assign, ast.AnnAssign)) and n.value is not None
                and not any(isinstance(c, ast.Call) for c in ast.walk(n.value)))]
    return "\n\n".join(ast.unparse(n) for n in keep)


def program(it: dict, completion: str) -> str:
    code = sanitize(completion)
    if it["task"] == "humanevalplus" and not re.search(rf"^def {re.escape(it['entry'])}\(", code, re.M):
        code = it["prefix"] + code  # the reply continued the function body
    head = "\n".join(it["imports"])
    tail = f"\n\ncheck({it['entry']})\n" if it["task"] == "humanevalplus" else "\n"
    return f"{head}\n{code}\n\n{it['test']}{tail}"


def execute(progs: list[str], timeout: float) -> list[dict]:
    with tempfile.TemporaryDirectory(prefix="code_eval_") as d:
        w = Path(d)
        for i, p in enumerate(progs):
            (w / f"{i}.py").write_text(p)
        (w / "runner.py").write_text(RUNNER)
        w.chmod(0o755)
        for f in w.iterdir():
            f.chmod(0o644)
        cmd = ["docker", "run", "--rm", "--network", "none", "--read-only", "--tmpfs", "/tmp:size=256m,mode=1777",
               "--user", "65534:65534", "--memory", "8g", "--memory-swap", "8g", "--cpus", "8", "--pids-limit", "256",
               "--security-opt", "no-new-privileges", "--cap-drop", "ALL", "-v", f"{w}:/work:ro",
               "-e", "OPENBLAS_NUM_THREADS=1", "-e", "OMP_NUM_THREADS=1", "-e", "MKL_NUM_THREADS=1",
               "--entrypoint", "python3", IMAGE, "-I", "/work/runner.py", str(len(progs)), str(timeout)]
        p = subprocess.run(cmd, capture_output=True, timeout=timeout * len(progs) / 4 + 600)
        res = [json.loads(l) for l in p.stdout.decode().splitlines() if l.startswith("{")]
        if len(res) != len(progs):
            raise RuntimeError(f"sandbox returned {len(res)} of {len(progs)} results: {p.stderr.decode()[-500:]}")
        return sorted(res, key=lambda r: r["i"])


def main() -> None:
    signal.signal(signal.SIGTERM, lambda *_: sys.exit(143))
    ap = argparse.ArgumentParser()
    ap.add_argument("--server", required=True)
    ap.add_argument("--data", type=Path, default=Path("/path/to/data/evalsets/evalplus_dl/data"))
    ap.add_argument("--model", action="append", required=True, help="NAME=PATH.gguf")
    ap.add_argument("--out", type=Path, default=REPO / "reports" / "compare")
    ap.add_argument("--port", type=int, default=18190)
    ap.add_argument("--ngl", default="99")
    ap.add_argument("--workers", type=int, default=8)
    ap.add_argument("--timeout", type=float, default=60.0, help="seconds per problem (all its tests)")
    ap.add_argument("--limit", type=int, default=0, help="first N problems per benchmark (smoke test)")
    a = ap.parse_args()
    its = items(a.data)
    if a.limit:
        its = [it for t in ("humanevalplus", "mbppplus") for it in [x for x in its if x["task"] == t][:a.limit]]
    print(f"{len(its)} problems", flush=True)
    for spec in a.model:
        name, gguf = spec.split("=", 1)
        path = a.out / f"{name}.code.jsonl"
        done = {(r["task"], r["id"]): r for r in map(json.loads, path.open())} if path.exists() else {}
        todo = [it for it in its if (it["task"], it["id"]) not in done]
        if todo:
            t0 = time.time()
            p = cq.serve(a.server, gguf, a.port, a.ngl, a.out / f"{name}.code.server.log")
            try:
                with ThreadPoolExecutor(a.workers) as ex:
                    answers = list(ex.map(lambda it: ask(a.port, it), todo))
            finally:
                p.terminate()
                p.wait(60)
            print(f"{name}: generated {len(todo)} in {time.time() - t0:.0f}s", flush=True)
            progs = [program(it, ans.get("content", "")) for it, ans in zip(todo, answers)]
            runs = execute(progs, a.timeout)
            with path.open("a") as f:
                for it, ans, r in zip(todo, answers, runs):
                    row = {"task": it["task"], "id": it["id"], "correct": "error" not in ans and r["status"] == "pass",
                           "status": "error" if "error" in ans else r["status"], "n_out": ans.get("n_out"),
                           "stopped": ans.get("stopped"), "completion": ans.get("content", ans.get("error", ""))[-1500:],
                           "err": r["err"]}
                    done[(it["task"], it["id"])] = row
                    f.write(json.dumps(row, ensure_ascii=False) + "\n")
        rows = [done[(it["task"], it["id"])] for it in its]
        tasks = {}
        for t in ("humanevalplus", "mbppplus"):
            rs = [r for r in rows if r["task"] == t]
            if rs:
                tasks[t] = sum(r["correct"] for r in rs) / len(rs)
        st = {}
        for r in rows:
            st[r["status"]] = st.get(r["status"], 0) + 1
        s = {"model": name, "groups": {"HumanEval+": tasks.get("humanevalplus"), "MBPP+": tasks.get("mbppplus")},
             "tasks": tasks, "status": st, "n": len(rows), "gguf": gguf, "bytes": Path(gguf).stat().st_size}
        (a.out / f"{name}.code.summary.json").write_text(json.dumps(s, indent=1))
        print(json.dumps({"model": name, "suite": "code", **{k: round(v, 4) for k, v in tasks.items()}, "status": st}),
              flush=True)


if __name__ == "__main__":
    main()