selfcheck.py

"""Reproduction self-check: run every REPRODUCE.md row from the frozen artifacts and assert each
recomputed value against its Expected (paper) value. Exit 0 iff all load-bearing numbers match.

Run:  uv run --extra experiments --with scikit-learn --with sentence-transformers python scripts/selfcheck.py
Tiers: OFFLINE rows need no network; HF rows need the free (keyless) HuggingFace prompt fetch.
"""
from __future__ import annotations

import contextlib
import io
import json
import os
import statistics as st
import sys

# repo/bundle root on path so `experiments` + `scripts` resolve regardless of launch dir
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import numpy as np
from sklearn.metrics import roc_auc_score
from sklearn.model_selection import GroupKFold, GroupShuffleSplit

from experiments.path_b import route_eval as RE
from experiments.path_b import route_embed as EMB
from experiments.path_b import route_llm as RLLM
from experiments.path_b import endurance_capstone, endurance_xfamily, endurance_xfamily_router

os.makedirs("data/path_b", exist_ok=True)   # scripts write summaries here; self-create on a cold checkout
ROWS, FAILED = [], 0


def record(num, label, got, expected, ok, tier="OFFLINE", note="", diagnostic=False):
    """diagnostic=True rows surface a finding but never fail the load-bearing verdict/exit code."""
    global FAILED
    FAILED += (not ok and not diagnostic)
    ROWS.append((ok, num, label, got, expected, tier, note, diagnostic))
    tag = "WARN" if diagnostic else ("PASS" if ok else "FAIL")
    print(f"[{tag}] #{num:<3} {label}")
    print(f"        got      = {got}")
    print(f"        expected = {expected}   [{tier}]" + (f"\n        note     = {note}" if note else ""))


def near(a, b, tol):
    return a is not None and abs(a - b) <= tol


@contextlib.contextmanager
def quiet():
    with contextlib.redirect_stdout(io.StringIO()):
        yield


def summary(path):
    return json.load(open(path))


# ---- #1 per-feature ROC-AUC (OFFLINE; degrades to prompt-free, AUCs use frozen axes) -------------
AX = ["decomposability", "sequential_dependency", "tool_load", "ambiguity", "verifiability"]
EXPECT_AUC = {"decomposability": 0.480, "sequential_dependency": 0.527, "tool_load": 0.466,
              "ambiguity": 0.518, "verifiability": 0.465}   # the exact heights of Figure 1's bars
with quiet():
    tf = RE.load_task_features()
    rows = [r for r in RE.load_rows("speed") if r["task_id"] in tf]
    y = [0 if r["winlabel"] == "single" else 1 for r in rows]
    aucs = {a: float(roc_auc_score(y, [tf[r["task_id"]][a] for r in rows])) for a in AX}
lo, hi = min(aucs.values()), max(aucs.values())
per_feature_ok = all(near(aucs[a], EXPECT_AUC[a], 0.005) for a in AX)
record(1, "per-feature ROC-AUC (each Figure 1 bar, exact)",
       "; ".join(f"{a}={aucs[a]:.3f}" for a in AX),
       "0.480/0.527/0.466/0.518/0.465 (band 0.46-0.53)",
       per_feature_ok and 0.44 <= lo and hi <= 0.55)

# ---- #2 embedding baselines (OFFLINE; frozen vectors) -------------------------------------------
with quiet():
    loc = EMB.eval_rep("local", "speed")
    fro = EMB.eval_rep("frontier", "speed")
ok2 = (near(loc["held_out_accuracy"], 0.646, 0.01) and near(loc["macro_ovr_auc"], 0.46, 0.02)
       and near(fro["held_out_accuracy"], 0.635, 0.012) and near(fro["macro_ovr_auc"], 0.52, 0.02)
       and near(loc["majority_rate"], 0.646, 0.01))
record(2, "embedding baselines vs floor",
       f"bge acc={loc['held_out_accuracy']} AUC={loc['macro_ovr_auc']}; "
       f"oai acc={fro['held_out_accuracy']} AUC={fro['macro_ovr_auc']}; floor={loc['majority_rate']}",
       "bge 0.646/0.46, oai 0.635/0.52, floor 0.646", ok2)

# ---- #3 TF-IDF+axes predictor sits at the majority floor (HF: needs prompts) ---------------------
with quiet():
    kf = RE.run_kfold("speed")
pacc = kf["point_prediction_accuracy"]
record(3, "trained predictor does not beat floor",
       f"predictor acc={pacc}  (floor 0.646)", "~0.643-0.652, <= floor 0.646",
       near(pacc, 0.646, 0.012) and pacc <= 0.656, tier="HF")

# ---- #3b LLM-as-classifier collapses to single (OFFLINE re-analysis) -----------------------------
with quiet():
    RLLM.analyze()
llm = summary("data/path_b/route_llm_summary.json")
ok3b = (near(llm["llm_classifier_accuracy"], 0.646, 0.01) and near(llm["beats_floor_by"], 0.0, 0.005)
        and near(llm["fraction_predicted_single"], 1.0, 0.001) and llm["n_rows"] == 353)
record("3b", "LLM classifier predicts single for all 353",
       f"acc={llm['llm_classifier_accuracy']} beats_floor_by={llm['beats_floor_by']} "
       f"pred_single={llm['fraction_predicted_single']} n={llm['n_rows']}",
       "acc=floor=0.646, +0.0, all-single, n=353", ok3b)

# ---- #4 abstention lift 0.76 @ 25% coverage vs base 0.652 (HF) -----------------------------------
def speed_decisions():
    with quiet():
        rws = RE.load_rows("speed")
    groups = np.array([r["task_id"] for r in rws])
    pooled = []
    for trv, te in GroupKFold(n_splits=5).split(np.arange(len(rws)), groups=groups):
        trv_rows = [rws[i] for i in trv]
        g = np.array([r["task_id"] for r in trv_rows])
        tr_i, ca_i = next(GroupShuffleSplit(1, test_size=0.25, random_state=RE.SEED).split(trv_rows, groups=g))
        with quiet():
            out = RE._fit_fold([trv_rows[i] for i in tr_i], [trv_rows[i] for i in ca_i],
                               [rws[i] for i in te], "speed")
        pooled += [{"cal": o["cal"], "won": int(o["won"])} for o in out]
    return pooled


dec = speed_decisions()
base = np.mean([r["won"] for r in dec])
ranked = sorted(dec, key=lambda r: -r["cal"])
acc25 = float(np.mean([r["won"] for r in ranked[:max(1, int(0.25 * len(ranked)))]]))
record(4, "abstention lift at 25% coverage",
       f"acc@25%={acc25:.3f}  base={base:.3f}", "0.76 @ 25%, base 0.652",
       near(acc25, 0.76, 0.02) and near(float(base), 0.652, 0.006), tier="HF")

# ---- #5 LOFO calibration transfer ECE 0.0 -> 0.13 (HF) -------------------------------------------
with quiet():
    lf = RE.run_lofo("speed")
ece_in, ece_x = kf["ece"]["calibrated"]["point"], lf["ece"]["calibrated"]["point"]
record(5, "ECE in-distribution -> cross-family",
       f"in-dist ECE={ece_in:.3f}  cross-family ECE={ece_x:.3f}", "0.0 -> 0.13",
       near(ece_in, 0.0, 0.04) and near(ece_x, 0.13, 0.04), tier="HF")

# ---- #6 dose-response monotone 0.25 -> 1.00 (OFFLINE) + provenance check on 0.17/0.85 ------------
ext = [json.loads(x) for x in open(RE.data_file("endurance_extra.jsonl")) if x.strip()]
nc = {}
for r in ext:
    if r["arm"].startswith("n") and r["arm"][1:].isdigit():
        nc.setdefault(int(r["arm"][1:]), []).append(r["acc"])
seq = [st.mean(nc[k]) for k in sorted(nc)]
mono = all(b >= a - 0.02 for a, b in zip(seq, seq[1:]))
record(6, "dose-response monotone (chunk-count sweep)",
       f"n1..n16 = {[round(s,2) for s in seq]}", "0.25 -> 1.00, monotone",
       near(seq[0], 0.25, 0.03) and near(seq[-1], 1.0, 0.001) and mono)

# #6b §6 single-vs-chunk @16k on decomposable probes (the §6-opening collapse-vs-hold numbers)
with quiet():
    endurance_capstone.analyze()
cap = summary("data/path_b/endurance_capstone_summary.json")
ms, mc = cap["matters_single"], cap["matters_chunk"]
record("6b", "§6 single collapses / chunk holds @16k (decomposable, n=90)",
       f"matters_single={ms:.3f} matters_chunk={mc:.3f}",
       "single 0.18 vs chunk 0.82 @16k (n=90)",
       near(ms, 0.18, 0.012) and near(mc, 0.82, 0.012))

# ---- #7 {single,chunk} router (OFFLINE; capstone analyze already run above) ----------------------
ok7 = all(near(cap[k], v, 0.012) for k, v in
          [("always_single", 0.74), ("always_chunk", 0.48), ("judge_router", 0.75), ("true_router", 0.92)])
record(7, "single,chunk router (oracle/deployable/baselines)",
       f"always_single={cap['always_single']:.3f} always_chunk={cap['always_chunk']:.3f} "
       f"deployable={cap['judge_router']:.3f} oracle={cap['true_router']:.3f}",
       "0.74 / 0.48 / 0.75 / 0.92", ok7)

# ---- #8 cross-family replication (OFFLINE) -------------------------------------------------------
with quiet():
    endurance_xfamily.analyze()
xf = summary("data/path_b/endurance_xfamily_summary.json")["by_length"]
ok8 = (near(xf["1000"]["single"], 0.37, 0.02) and near(xf["16000"]["single"], 0.0, 0.02)
       and near(xf["1000"]["chunk"], 0.30, 0.03) and near(xf["16000"]["chunk"], 0.40, 0.05))
record(8, "Qwen single collapses, chunk holds",
       f"single 1k={xf['1000']['single']:.2f}->16k={xf['16000']['single']:.2f}; "
       f"chunk 1k={xf['1000']['chunk']:.2f}->16k={xf['16000']['chunk']:.2f}",
       "single 0.37->0.00, chunk 0.30->0.40", ok8)

# ---- #9 second-family router gap (OFFLINE) -------------------------------------------------------
with quiet():
    endurance_xfamily_router.analyze()
rt = summary("data/path_b/endurance_xfamily_router_summary.json")
ok9 = (near(rt["oracle_deployable_gap"], 0.22, 0.012) and near(rt["true_router"], 0.58, 0.012)
       and near(rt["judge_router"], 0.36, 0.012) and near(rt["always_single"], 0.53, 0.012)
       and rt["judge_accuracy"] == "3/6")
record(9, "Qwen router gap widens to 0.22",
       f"gap={rt['oracle_deployable_gap']:.3f} oracle={rt['true_router']:.3f} "
       f"deployable={rt['judge_router']:.3f} single={rt['always_single']:.3f} judge={rt['judge_accuracy']}",
       "gap 0.22, oracle 0.58, deployable 0.36, single 0.53, judge 3/6", ok9)

# ---- verdict ------------------------------------------------------------------------------------
load_bearing = [r for r in ROWS if not r[7]]
diag = [r for r in ROWS if r[7]]
n = len(load_bearing)
print("\n" + "=" * 78)
print(f"SELF-CHECK: {n - FAILED}/{n} load-bearing rows PASS, {FAILED} FAIL"
      + (f"  (+{len(diag)} diagnostic warning(s))" if diag else ""))
for ok, num, label, *_ in (r for r in ROWS if not r[0]):
    print(f"  {'WARN' if [x for x in ROWS if x[1]==num][0][7] else 'FAIL'} #{num}: {label}")
raise SystemExit(1 if FAILED else 0)