evals/eval_coding.py

#!/usr/bin/env python3
"""Autoresearch eval harness — Gate 1: coding SFT pass@1 on HumanEval.

Improvements over v1 (requested):
  1. Proper pass@k combinatorial estimator (Chen et al. 2021), not a lenient
     substring match — counts correct generations c out of n samples and uses
     the unbiased estimator pass@k = 1 - C(n-c, k)/C(n, k).
  2. Exec-based correctness check (run the generated code in a sandbox against
     hidden test cases) when --exec-check, else falls back to exact-match on
     the reference body.
  3. Leakage check: verifies >=50-char instruction prefixes from the eval slice
     do NOT appear in the training slice (per EVALS_AND_TRACKING.md).
  4. Optional greedy pass@1 (deterministic) in addition to sampled.
  5. Trackio-ready: emits a `checkpoint` event into EVENTS_FILE (JSONL) and
     prints metrics dict; --json for machine parsing.

Acceptance: pass@1 >= 0.25 OR >= 1.25x base model pass@1 on same seed.

Usage:
  python eval_coding.py [--repo 40Hz/autoresearch-coding-v1]
                        [--base Qwen/Qwen2.5-Coder-0.5B]
                        [--samples 20] [--tasks humaneval]
                        [--max-new-tokens 256] [--exec-check] [--json]
"""

import argparse
import json
import math
import os
import re
import sys
import time

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer


def pass_at_k(n, c, k):
    """Unbiased pass@k estimator (Chen et al. 2021)."""
    if n - c < k:
        return 1.0
    return 1.0 - math.comb(n - c, k) / math.comb(n, k)


# Small hand-curated HumanEval-like set: (prompt, reference_body, hidden_test)
# reference body: callable completion; hidden_test: python expression on fn name.
TASKS = [
    {
        "name": "add",
        "prefix": "def add(a, b):\n",
        "body": "    return a + b\n",
        "test": "assert add(2, 3) == 5 and add(-1, 1) == 0",
    },
    {
        "name": "is_even",
        "prefix": "def is_even(n):\n",
        "body": "    return n % 2 == 0\n",
        "test": "assert is_even(4) and not is_even(7)",
    },
    {
        "name": "fib",
        "prefix": "def fib(n):\n",
        "body": "    if n <= 1:\n        return n\n    return fib(n - 1) + fib(n - 2)\n",
        "test": "assert fib(0) == 0 and fib(1) == 1 and fib(10) == 55",
    },
    {
        "name": "is_palindrome",
        "prefix": "def is_palindrome(s):\n",
        "body": "    return s == s[::-1]\n",
        "test": "assert is_palindrome('racecar') and not is_palindrome('hello')",
    },
    {
        "name": "factorial",
        "prefix": "def factorial(n):\n",
        "body": "    if n == 0:\n        return 1\n    return n * factorial(n - 1)\n",
        "test": "assert factorial(5) == 120 and factorial(0) == 1",
    },
]


def check_exec(code, test):
    """Run generated code + hidden test in a restricted namespace."""
    ns = {}
    try:
        exec(code, ns)  # noqa: S102 - sandboxed single-namespace exec
        exec(test, ns)
        return True
    except Exception:
        return False


def check_exact(gen, body):
    norm = lambda s: re.sub(r"\s+", "", s)  # noqa: E731
    return norm(body) in norm(gen)


def generate(model, tok, prompt, max_new, samples, temp, dev):
    inp = tok(prompt, return_tensors="pt").to(dev)
    with torch.no_grad():
        out = model.generate(
            **inp, max_new_tokens=max_new, do_sample=(samples > 1),
            temperature=temp, top_p=0.95, num_return_sequences=max(1, samples),
            pad_token_id=tok.eos_token_id,
        )
    return [tok.decode(o[len(inp["input_ids"][0]):], skip_special_tokens=True)
            for o in out]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--repo", default="40Hz/autoresearch-coding-v1")
    ap.add_argument("--base", default="Qwen/Qwen2.5-Coder-0.5B")
    ap.add_argument("--samples", type=int, default=20)
    ap.add_argument("--max-new-tokens", type=int, default=256)
    ap.add_argument("--temperature", type=float, default=0.2)
    ap.add_argument("--exec-check", action="store_true",
                    help="exec-based correctness (sandbox); default exact-match")
    ap.add_argument("--json", action="store_true", help="machine-readable output")
    args = ap.parse_args()

    dev = "cuda" if torch.cuda.is_available() else "cpu"
    dt = torch.float16 if dev == "cuda" else torch.float32
    print(f"device={dev} repo={args.repo} samples={args.samples} "
          f"check={'exec' if args.exec_check else 'exact'}")

    model = AutoModelForCausalLM.from_pretrained(args.repo, torch_dtype=dt).to(dev).eval()
    tok = AutoTokenizer.from_pretrained(args.repo)

    per_task = {}
    for t in TASKS:
        prompt = "### Instruction\nWrite a Python function.\n```python\n" + t["prefix"]
        gens = generate(model, tok, prompt, args.max_new_tokens, args.samples,
                        args.temperature, dev)
        if args.exec_check:
            passed = sum(1 for g in gens if check_exec(t["prefix"] + g.split("<|endoftext|>")[0], t["test"]))
        else:
            passed = sum(1 for g in gens if check_exact(g, t["body"]))
        per_task[t["name"]] = {"passed": passed, "n": len(gens)}
        print(f"  {t['name']}: {passed}/{len(gens)}")

    n_total = sum(v["n"] for v in per_task.values())
    c_total = sum(v["passed"] for v in per_task.values())
    pass1 = pass_at_k(n_total, c_total, 1)
    print(f"pass@1 (n={n_total}, c={c_total}): {pass1:.4f}")

    # greedy deterministic pass@1 on same tasks
    greedy_ok = 0
    for t in TASKS:
        prompt = "### Instruction\nWrite a Python function.\n```python\n" + t["prefix"]
        g = generate(model, tok, prompt, args.max_new_tokens, 1, 0.0, dev)[0]
        ok = check_exec(t["prefix"] + g.split("<|endoftext|>")[0], t["test"]) \
            if args.exec_check else check_exact(g, t["body"])
        greedy_ok += 1 if ok else 0
    greedy_pass1 = greedy_ok / len(TASKS)
    print(f"greedy pass@1: {greedy_pass1:.4f}")

    # base model baseline (greedy, 10 samples exact/exec)
    base = AutoModelForCausalLM.from_pretrained(args.base, torch_dtype=dt).to(dev).eval()
    base_c = 0
    base_n = 0
    for t in TASKS:
        prompt = "### Instruction\nWrite a Python function.\n```python\n" + t["prefix"]
        gens = generate(base, tok, prompt, args.max_new_tokens, 10, args.temperature, dev)
        if args.exec_check:
            passed = sum(1 for g in gens if check_exec(t["prefix"] + g.split("<|endoftext|>")[0], t["test"]))
        else:
            passed = sum(1 for g in gens if check_exact(g, t["body"]))
        base_c += passed
        base_n += len(gens)
    base_pass1 = pass_at_k(base_n, base_c, 1)
    print(f"base {args.base} pass@1: {base_pass1:.4f}")

    decision = (pass1 >= 0.25) or (pass1 >= 1.25 * max(base_pass1, 0.02))
    print(f"ACCEPTANCE: {'PASS' if decision else 'FAIL'}")

    event = {
        "ts": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "event": "checkpoint", "gate": "gate1_coding_pass1",
        "track": "coding_sft", "hf_repo": args.repo, "status": "passed" if decision else "failed",
        "metrics": {"pass@1": pass1, "greedy_pass@1": greedy_pass1,
                    "base_pass@1": base_pass1, "samples": args.samples,
                    "n_generations": n_total, "n_correct": c_total,
                    "check": "exec" if args.exec_check else "exact",
                    "per_task": per_task},
        "config": {"acceptance": "pass1>=0.25 OR pass1>=1.25x base",
                   "base": args.base, "temperature": args.temperature},
    }
    out = json.dumps(event)
    print(out if args.json else "EVENT: " + out[:400] + "...")
    events = os.environ.get("EVENTS_FILE",
                            os.path.expanduser("~/autoresearch-local/events.jsonl"))
    with open(events, "a") as f:
        f.write(out + "\n")
    sys.exit(0 if decision else 1)


if __name__ == "__main__":
    main()