bench.py

"""Vue « benchmark réel » : agrégation des métriques sur T tâches + fiabilité (IC).

Données : SWE-bench Lite Samples (Large Language Monkeys). Pour chaque tâche on a
(n_i ≈ 250, c_i) -> taux estimé p_i = c_i / n_i.

Courbes agrégées (modèle i.i.d. par tâche, moyenne sur les tâches), pour k = 1..K :
  pass@k   = moyenne_i [ 1 - (1-p_i)^k ]            (best-of-k, suppose un oracle)
  pass+50% = moyenne_i [ P(Binom(k, p_i) > k/2) ]   (vote majoritaire)
  pass^k   = moyenne_i [ p_i^k ]                     (fiabilité : k réussites d'affilée)

« Le maximum » = les deux plafonds quand k -> +inf :
  couverture (oracle) = fraction de tâches résolubles = moyenne_i [ p_i > 0 ]
  plafond vote        = fraction de tâches fiables     = moyenne_i [ p_i > 0.5 ]

Fiabilité de l'estimation avec un budget (n échantillons/tâche, T tâches) : bootstrap.
On rejoue l'évaluation B fois — T tâches tirées sans remise parmi les ~300, et n
échantillons retirés sans remise des ~250 disponibles (loi hypergéométrique) — puis on
regarde la dispersion des plafonds estimés. Plus n et T grandissent, plus l'intervalle
de confiance se resserre (et le biais du plafond oracle, sous-estimé à petit n car on
rate les réussites rares, s'efface).
"""

from __future__ import annotations

import json
import os
from math import comb

import numpy as np

DATA = os.path.join(os.path.dirname(__file__), "data", "swebench_lite_samples.json")


def load_dataset(path: str = DATA) -> dict:
    with open(path) as f:
        d = json.load(f)
    N = np.array([t["n"] for t in d["tasks"]], dtype=float)
    C = np.array([t["c"] for t in d["tasks"]], dtype=float)
    return {
        "N": N, "C": C, "P": C / N,
        "task_ids": [t["task_id"] for t in d["tasks"]],
        "n_tasks": int(d["n_tasks"]),
        "n_max": int(N.max()), "n_min": int(N.min()),
        "model": d["model"], "benchmark": d["benchmark"],
        "paper": d.get("paper", ""), "source": d.get("source", ""),
    }


def _maj_iid_matrix(P: np.ndarray, K: int) -> np.ndarray:
    """(T, K) : P(Binom(k, p_i) > k/2) pour k = 1..K, vectorisé sur les tâches."""
    P = np.asarray(P, float)
    out = np.empty((P.shape[0], K))
    for ki in range(1, K + 1):
        js = np.arange(ki // 2 + 1, ki + 1)  # majorité stricte
        coeff = np.array([comb(ki, int(j)) for j in js], float)
        pj = P[:, None] ** js[None, :]
        qj = (1.0 - P)[:, None] ** (ki - js)[None, :]
        out[:, ki - 1] = (coeff[None, :] * pj * qj).sum(axis=1)
    return out


def aggregate_curves(P: np.ndarray, K: int) -> dict:
    """Courbes moyennes (et écart-type sur les tâches) pour k = 1..K."""
    P = np.asarray(P, float)
    ks = np.arange(1, K + 1)
    mats = {
        "at": 1.0 - (1.0 - P)[:, None] ** ks[None, :],
        "maj": _maj_iid_matrix(P, K),
        "pow": P[:, None] ** ks[None, :],
    }
    res = {"ks": ks}
    for name, M in mats.items():
        res[f"{name}_mean"] = M.mean(axis=0)
        res[f"{name}_std"] = M.std(axis=0)
    return res


def task_ci_halfwidth(std: np.ndarray, T: int, T_full: int, z: float = 1.645) -> np.ndarray:
    """Demi-largeur d'IC due à l'échantillonnage des tâches (CLT + correction pop. finie).

    Vaut z·std/sqrt(T) corrigé par sqrt(1 - T/T_full) -> 0 quand on prend toutes les tâches.
    """
    fpc = max(0.0, 1.0 - T / T_full)
    return z * np.asarray(std) * np.sqrt(fpc / max(T, 1))


def ceilings_true(P: np.ndarray) -> dict:
    """Valeurs « vraies » (toutes les données) : plafonds + score single-shot."""
    P = np.asarray(P, float)
    return {
        "coverage": float((P > 0).mean()),    # best-of-inf (oracle)
        "voting": float((P > 0.5).mean()),     # vote majoritaire à l'infini
        "pass1": float(P.mean()),              # pass@1 moyen
    }


def ceiling_bootstrap(N, C, n: int, T: int, B: int = 400, seed: int = 0) -> dict:
    """IC des deux plafonds pour un budget (n échantillons/tâche, T tâches).

    seed fixe -> résultat déterministe (la bande ne « tremble » pas en déplaçant un curseur).
    """
    N = np.asarray(N, float)
    C = np.asarray(C, float)
    T_full = N.shape[0]
    n, T, B = int(n), int(min(T, T_full)), int(B)

    rng = np.random.default_rng(seed)
    idx = np.argsort(rng.random((B, T_full)), axis=1)[:, :T]  # T tâches sans remise / rep
    Nsub, Csub = N[idx], C[idx]
    n_eff = np.minimum(n, Nsub).astype(int)                  # budget effectif <= n_i
    cprime = rng.hypergeometric(Csub.astype(int),
                                (Nsub - Csub).astype(int), n_eff)  # (B, T)

    cov = (cprime > 0).mean(axis=1)            # couverture estimée par rep
    vote = (cprime / n_eff > 0.5).mean(axis=1)  # plafond vote estimé par rep

    def stat(x):
        return {"mean": float(x.mean()),
                "lo": float(np.percentile(x, 5)),
                "hi": float(np.percentile(x, 95))}

    return {"coverage": stat(cov), "voting": stat(vote)}


if __name__ == "__main__":
    ds = load_dataset()
    P, N, C = ds["P"], ds["N"], ds["C"]
    assert ds["n_tasks"] == 300

    tr = ceilings_true(P)
    print(f"couverture (oracle) = {tr['coverage']:.3f}   "
          f"plafond vote = {tr['voting']:.3f}   pass@1 moyen = {tr['pass1']:.3f}")
    assert abs(tr["coverage"] - 168 / 300) < 1e-9
    assert abs(tr["voting"] - 38 / 300) < 1e-9

    K = ds["n_min"]
    cur = aggregate_curves(P, K)
    # En k=1, les 3 métriques valent pass@1 moyen.
    for m in ("at", "maj", "pow"):
        assert abs(cur[f"{m}_mean"][0] - tr["pass1"]) < 1e-9
    # Encadrement pass^k <= pass+50% <= pass@k, partout.
    assert np.all(cur["pow_mean"] <= cur["maj_mean"] + 1e-9)
    assert np.all(cur["maj_mean"] <= cur["at_mean"] + 1e-9)
    # Monotonie.
    assert np.all(np.diff(cur["at_mean"]) >= -1e-9)
    assert np.all(np.diff(cur["pow_mean"]) <= 1e-9)
    # pass@k ne dépasse jamais la couverture (best-of-K <= best-of-inf).
    assert cur["at_mean"][-1] <= tr["coverage"] + 1e-9

    # Budget max -> bootstrap = valeurs vraies, IC nul.
    bmax = ceiling_bootstrap(N, C, n=ds["n_max"], T=ds["n_tasks"], B=200)
    assert abs(bmax["coverage"]["mean"] - tr["coverage"]) < 1e-9
    assert bmax["coverage"]["hi"] - bmax["coverage"]["lo"] < 1e-9
    # Petit n -> couverture sous-estimée (on rate les réussites rares).
    blow = ceiling_bootstrap(N, C, n=3, T=ds["n_tasks"], B=400)
    print(f"couverture estimée à n=3 : {blow['coverage']['mean']:.3f} "
          f"[{blow['coverage']['lo']:.3f}, {blow['coverage']['hi']:.3f}]")
    assert blow["coverage"]["mean"] < tr["coverage"]
    # Peu de tâches -> IC plus large.
    w_fewT = ceiling_bootstrap(N, C, n=ds["n_max"], T=20, B=400)["coverage"]
    assert (w_fewT["hi"] - w_fewT["lo"]) > 0.02

    print("bench.py : tous les tests passent ✔")