Open in Colab ▶️ Run this notebook in Colab

HumanEval - unbiased pass@k#

Track 06 - Evaluation · Notebook 03 · Runtime: ≈1 min on CPU

Prerequisites: 06_eval/01 (perplexity from scratch).

Paper: Chen et al. 2021, Evaluating Large Language Models Trained on Code (2107.03374), appendix A.


What#

If you generate n candidate completions per problem and count how many pass unit tests (c of them), what is the probability that sampling k n of them would include at least one passing sample?

The naive estimator 1 - ((n - c) / n)^k is biased. Chen et al. propose an unbiased one:

pass@k = 1 - C(n - c, k) / C(n, k)         (if n - c >= k)
       = 1.0                                (otherwise)

This notebook implements the unbiased estimator, runs tiny real candidates through a sandboxed exec against reference tests, and verifies the estimator on stubbed problem outcomes.

from llm_systems_cookbook.nb import bootstrap

import math
import multiprocessing as mp

import numpy as np

s = bootstrap("06_eval_03_humaneval_unbiased_pass_k")

The estimator#

pass_at_k(n, c, k) returns the exact unbiased pass@k given n samples out of which c passed. Use math.comb.

def pass_at_k(n: int, c: int, k: int) -> float:
    '''Unbiased pass@k estimator from Chen et al. 2021, eqn. (1).'''
    if n - c < k:
        return 1.0
    return 1.0 - math.comb(n - c, k) / math.comb(n, k)


# Sanity check: if every sample passes, pass@k = 1.
# If none pass, pass@k = 0. If exactly one of n passes, pass@k = k/n.
print(f"c=n=5, k=1:  {pass_at_k(5, 5, 1):.3f}  (should be 1.0)")
print(f"c=0, n=5, k=1: {pass_at_k(5, 0, 1):.3f}  (should be 0.0)")
print(f"c=1, n=5, k=1: {pass_at_k(5, 1, 1):.3f}  (should be 0.2)")
print(f"c=1, n=5, k=5: {pass_at_k(5, 1, 5):.3f}  (should be 1.0)")

Sandboxed execution#

Code-evaluation benchmarks run model-generated code. Running arbitrary completions in-process is a security hole; the minimum hardening is a separate process with a timeout. The sandbox below uses multiprocessing.Process with a deadline.

def _worker(code: str, tests: str, conn) -> None:
    ns: dict = {}
    try:
        exec(code, ns)
        exec(tests, ns)
        conn.send(("pass", None))
    except Exception as e:  # noqa: BLE001
        conn.send(("fail", f"{type(e).__name__}: {e}"))
    finally:
        conn.close()


def run_candidate(code: str, tests: str, timeout_s: float = 5.0) -> bool:
    parent, child = mp.Pipe()
    p = mp.get_context("fork").Process(target=_worker, args=(code, tests, child))
    p.start()
    p.join(timeout_s)
    if p.is_alive():
        p.terminate()
        p.join()
        return False
    if not parent.poll():
        return False
    status, _ = parent.recv()
    return status == "pass"


# Smoke test: correct and incorrect candidates for a trivial problem.
PROBLEM = "def add(a, b):\n    '''Return a + b.'''"
tests = 'assert add(2, 3) == 5\nassert add(-1, 1) == 0'

good = PROBLEM + '\n    return a + b'
bad  = PROBLEM + '\n    return a - b'
print(f"correct candidate passes: {run_candidate(good, tests)}")
print(f"buggy candidate passes:   {run_candidate(bad, tests)}")

A tiny five-problem benchmark#

Five trivial problems, each with a short reference test. We generate a set of candidates per problem (some correct, some buggy) and score.

PROBLEMS = [
    {
        "prompt": 'def add(a, b):\n    """Return a + b."""',
        "tests":  'assert add(2, 3) == 5\nassert add(-1, 1) == 0',
        "correct_bodies": ["    return a + b", "    return b + a"],
        "wrong_bodies":   ["    return a - b", "    return a * b"],
    },
    {
        "prompt": 'def square(x):\n    """Return x squared."""',
        "tests":  'assert square(3) == 9\nassert square(-4) == 16',
        "correct_bodies": ["    return x * x", "    return x ** 2"],
        "wrong_bodies":   ["    return 2 * x"],
    },
    {
        "prompt": 'def is_even(x):\n    """True if x is even."""',
        "tests":  'assert is_even(2) is True\nassert is_even(3) is False',
        "correct_bodies": ["    return x % 2 == 0"],
        "wrong_bodies":   ["    return x == 2", "    return x / 2 == 0"],
    },
    {
        "prompt": 'def first_letter(s):\n    """Return the first character of s."""',
        "tests":  'assert first_letter("hello") == "h"\nassert first_letter("Ai") == "A"',
        "correct_bodies": ["    return s[0]", "    return s[:1]"],
        "wrong_bodies":   ["    return s[-1]"],
    },
    {
        "prompt": 'def max_of(a, b):\n    """Return the larger argument."""',
        "tests":  'assert max_of(1, 2) == 2\nassert max_of(7, 3) == 7',
        "correct_bodies": ["    return a if a > b else b", "    return max(a, b)"],
        "wrong_bodies":   ["    return a + b"],
    },
]


def run_problem(p: dict, n_candidates: int) -> int:
    '''Sample ``n_candidates`` bodies (rotating between correct + wrong
    with a 60/40 split) and return how many passed.'''
    rng = np.random.default_rng(abs(hash(p["prompt"])) % (2**32))
    n_correct_target = int(round(n_candidates * 0.6))
    bodies = (
        list(rng.choice(p["correct_bodies"], size=n_correct_target))
        + list(rng.choice(p["wrong_bodies"], size=n_candidates - n_correct_target))
    )
    rng.shuffle(bodies)
    return sum(run_candidate(p["prompt"] + "\n" + body, p["tests"]) for body in bodies)


N = 10
per_problem_c = [run_problem(p, N) for p in PROBLEMS]
print(f"passed count per problem (n={N}):  {per_problem_c}")

pass@k across the benchmark#

Report pass@1, pass@5, and pass@10 averaged over problems.

def benchmark_pass_at_k(per_problem_c: list[int], n: int, ks: list[int]) -> dict[int, float]:
    out: dict[int, float] = {}
    for k in ks:
        out[k] = float(np.mean([pass_at_k(n, c, k) for c in per_problem_c]))
    return out


pk = benchmark_pass_at_k(per_problem_c, N, [1, 5, 10])
for k, v in pk.items():
    print(f"pass@{k:<2}  = {v:.3f}")

# Naive (biased) comparison.
def naive_pass_at_k(n: int, c: int, k: int) -> float:
    return 1.0 - ((n - c) / n) ** k if n > 0 else 0.0

pk_naive = {k: float(np.mean([naive_pass_at_k(N, c, k) for c in per_problem_c])) for k in (1, 5, 10)}
for k in (1, 5, 10):
    print(f"naive pass@{k:<2} = {pk_naive[k]:.3f}   (vs unbiased {pk[k]:.3f})")
s.assert_close("pass_at_1_equals_c_over_n_on_average",
               actual=pk[1],
               expected=np.mean(per_problem_c) / N,
               rtol=1e-9)
s.check(
    "pass_at_k_monotone_in_k",
    lambda: pk[1] <= pk[5] <= pk[10] + 1e-9,
    msg=f"{pk}",
)
s.check(
    "unbiased_matches_naive_at_k_equals_1",
    lambda: abs(pk[1] - pk_naive[1]) < 1e-9,
    msg=f"unbiased pass@1 = {pk[1]:.4f}  naive = {pk_naive[1]:.4f}",
)
s.check(
    "naive_underestimates_for_k_gt_1",
    lambda: pk_naive[5] <= pk[5] + 1e-9,
    msg=f"naive pass@5 = {pk_naive[5]:.4f}  unbiased = {pk[5]:.4f}",
)
s.check(
    "every_problem_has_some_correct_candidates",
    lambda: all(c > 0 for c in per_problem_c),
    msg=f"per-problem c = {per_problem_c}",
)

pass@k curve - biased vs unbiased#

A single number is hard to interpret; the full pass@k curve over k = 1 n tells the real story. Plot both estimators on the same axes and the naive one consistently overshoots for k > 1. The gap is exactly the sampling-without-replacement correction Chen et al. derive in equation (1).

import matplotlib.pyplot as plt

ks = list(range(1, N + 1))
unbiased = [float(np.mean([pass_at_k(N, c, k)       for c in per_problem_c])) for k in ks]
biased   = [float(np.mean([naive_pass_at_k(N, c, k) for c in per_problem_c])) for k in ks]

fig, ax = plt.subplots(figsize=(6.8, 4.0))
ax.plot(ks, unbiased, "o-", color="tab:blue",  label="unbiased (Chen 2021 eq. 1)")
ax.plot(ks, biased,   "s--", color="tab:red",   label="naive  1 − ((n−c)/n)^k")
for k_mark in (1, 5, 10):
    ax.axvline(k_mark, color="tab:gray", alpha=0.25, linewidth=1)
ax.set_xlabel("k (samples drawn)")
ax.set_ylabel("pass@k  (mean over problems)")
ax.set_title(f"pass@k curve  (n={N} candidates/problem, {len(per_problem_c)} problems)")
ax.set_xticks([1, 2, 3, 4, 5, 7, 10])
ax.set_ylim(0, 1.05)
ax.grid(alpha=0.3)
ax.legend(loc="lower right")
fig.tight_layout()
plt.show()

Exercises#

  1. Break the sandbox on purpose - write a candidate that time.sleeps for 100 seconds and verify the timeout kills it without wedging the notebook.

  2. Compute pass@k confidence intervals via bootstrap over problems. Production HumanEval results often report 95% CI bars.

  3. Replace the stubbed candidates with real ones - generate k completions from a small model (Qwen-Coder-0.5B) and evaluate.

References#

  • Chen et al. 2021, Evaluating Large Language Models Trained on Code, Appendix A for the unbiased estimator.

  • OpenAI’s human-eval repository - the reference implementation this notebook mirrors.

s.summary()
s.save()