arborist/bench/scripts/relevance_shadow_sweep.py

#!/usr/bin/env python3
"""Relevance shadow sweep over pooled bench-qa STRICT — #000052 §3.2.2
step 2 (the load-bearing measurement).

Take the pooled bench-qa STRICT (question, answer) pairs the LLM
actually produced and score each one with every manifest reranker.
Report per-model:
  - score distribution (min, p10, p50, p90, max) on STRICT pairs
  - at each candidate-bench θ recorded in the manifest, the
    *fraction of STRICT pairs that would fire* (= the false-demote
    rate on confidently-grounded answers, which must be ≈0 to clear
    the §7 #12 gate)
  - the smallest θ that achieves zero false-demote on this pool
    (the real-traffic-calibrated threshold)

The §000049 §7 #18→#27 lesson applies: clean candidate-bench θ
values WILL move when the denominator grows. This script measures
how much.

SHADOW ONLY — writes nothing but bench/results/relevance-shadow-sweep.json.

Usage:
  python3 bench/scripts/relevance_shadow_sweep.py
  python3 bench/scripts/relevance_shadow_sweep.py --models all
  python3 bench/scripts/relevance_shadow_sweep.py --strict-pool bench/qa_results/*.jsonl --out path.json
"""
from __future__ import annotations

import argparse
import json
import sys
import time
from pathlib import Path

REPO = Path(__file__).resolve().parents[2]
_DEFAULT_STRICT_POOL = [
    REPO / "bench" / "qa_results" / "2026-05-12T20-53-11Z.jsonl",
    REPO / "bench" / "qa_results" / "2026-05-12T21-58-58Z.jsonl",
    REPO / "bench" / "qa_results" / "2026-05-12T22-44-30Z.jsonl",
]


def _load_strict(paths: list[Path]) -> list[dict]:
    out = []
    for p in paths:
        if not p.exists():
            print(f"[relevance-sweep] skip (missing): {p}", file=sys.stderr)
            continue
        for ln in p.read_text().splitlines():
            ln = ln.strip()
            if not ln:
                continue
            o = json.loads(ln)
            if (o.get("audit_mode") == "STRICT"
                    and o.get("question") and o.get("answer_text")):
                out.append({"question": o["question"], "answer": o["answer_text"]})
    return out


def _models_from_manifest(which: str | None):
    sys.path.insert(0, str(REPO))
    from arborist.qa.relevance.shadow import load_manifest
    m = load_manifest()
    entries = [{"relevance_model_version": m["relevance_model_version"],
                "hf_repo": m["hf_repo"], "pinned_revision": m.get("pinned_revision")}]
    entries += [{"relevance_model_version": a["relevance_model_version"],
                 "hf_repo": a["hf_repo"], "pinned_revision": a.get("pinned_revision")}
                for a in m.get("alternates", [])]
    if which in (None, "", "primary"):
        return entries[:1]
    if which == "all":
        return entries
    want = {w.strip() for w in which.split(",")}
    sel = [e for e in entries if e["relevance_model_version"] in want or e["hf_repo"] in want]
    return sel or entries[:1]


def _quantiles(xs: list[float], qs=(0.0, 0.1, 0.5, 0.9, 1.0)) -> dict:
    s = sorted(xs)
    n = len(s)
    if n == 0:
        return {f"p{int(q*100)}": None for q in qs}
    return {f"p{int(q*100)}": round(s[min(n - 1, max(0, int(round(q * (n - 1)))))], 4) for q in qs}


def _fp_rate_at_theta(scores: list[float], theta: float) -> float:
    """Fraction of STRICT pairs that would FIRE the veto at θ
    (= score < θ). Want this ≈0."""
    if not scores:
        return 0.0
    fires = sum(1 for s in scores if s < theta)
    return fires / len(scores)


def _smallest_theta_for_zero_fp(scores: list[float]) -> float | None:
    """The smallest θ that yields fp=0 on this STRICT pool — i.e.
    θ ≤ min(scores). Anything below min(scores) gives 0 STRICT FP."""
    return min(scores) if scores else None


def main(argv=None) -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--strict-pool", action="append", type=Path,
                    help="bench-qa JSONL file(s) to pool; defaults to the three n=1+3+5 files")
    ap.add_argument("--models", default="all",
                    help="manifest model id(s) comma-sep, or 'all' (default), or 'primary'")
    ap.add_argument("--out", type=Path,
                    default=REPO / "bench" / "results" / "relevance-shadow-sweep-pooled-strict.json")
    ap.add_argument("--no-clean", action="store_true",
                    help="skip the claim-lattice metadata cleaning before scoring (see #000052 §3.2.2 refinement: cleaning drops STRICT FP rate ~22%% relative)")
    args = ap.parse_args(argv)

    try:
        import torch  # noqa
        sys.path.insert(0, str(REPO))
        from arborist.qa.relevance.shadow import ShadowRelevance, load_manifest, clean_for_relevance
    except ImportError as e:
        print(f"[relevance-sweep] missing dependency: {e}  (pip install 'arborist[nli]')",
              file=sys.stderr)
        return 2

    pool = list(args.strict_pool or _DEFAULT_STRICT_POOL)
    strict = _load_strict(pool)
    print(f"[relevance-sweep] pooled STRICT (question, answer) pairs: {len(strict)}", flush=True)
    if not strict:
        print("[relevance-sweep] no STRICT pairs loaded", file=sys.stderr); return 2

    models = _models_from_manifest(args.models)
    print(f"[relevance-sweep] models ({len(models)}): "
          f"{[m['relevance_model_version'] for m in models]}", flush=True)

    # candidate-bench θ values (per-model, from the round-2 manifest block)
    manifest = load_manifest()
    cb = manifest.get("candidate_bench_results_2026_05_13_round2", {})
    cb_θs = {row["model"]: row["fp0_theta"] for row in cb.get("rank_by_sep_margin", [])}

    report = {"generated_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
              "strict_pool_files": [str(p) for p in pool if p.exists()],
              "n_strict_pairs": len(strict),
              "models": {}}

    for me in models:
        mv = me["relevance_model_version"]
        short = mv.replace("relevance-shadow-v1-", "")
        print(f"[relevance-sweep] === {mv}  ({me['hf_repo']}) ===", flush=True)
        manifest_override = {
            "relevance_model_version": mv, "hf_repo": me["hf_repo"],
            "pinned_revision": me.get("pinned_revision"),
            "max_length": 256, "demote_below_score": None,
        }
        rel = ShadowRelevance(manifest=manifest_override)
        rel._ensure_loaded()
        if not rel.available:
            print(f"[relevance-sweep]   skip — {rel._reason}", file=sys.stderr); continue
        t0 = time.time()
        # batched scoring: build all (q, a) pairs once, optionally
        # claim-lattice-cleaned (the §3.2.2 refinement — drops STRICT
        # FP ~22% relative on the pooled-808 sample without losing
        # true-positive deflection signal).
        if args.no_clean:
            pairs = [(p["question"], p["answer"]) for p in strict]
        else:
            pairs = [(clean_for_relevance(p["question"]),
                      clean_for_relevance(p["answer"])) for p in strict]
        scores = rel._score_batch(pairs)
        infer_s = time.time() - t0
        q = _quantiles(scores)
        cb_θ = cb_θs.get(short)
        cb_fp = _fp_rate_at_theta(scores, cb_θ) if cb_θ is not None else None
        zero_fp_θ = _smallest_theta_for_zero_fp(scores)
        print(f"[relevance-sweep]   backend={rel.backend} device={rel.device}  · "
              f"{len(pairs)} pairs in {infer_s:.2f}s", flush=True)
        print(f"[relevance-sweep]   score quantiles on STRICT: min={q['p0']}  p10={q['p10']}  "
              f"p50={q['p50']}  p90={q['p90']}  max={q['p100']}", flush=True)
        if cb_θ is not None:
            print(f"[relevance-sweep]   at candidate-bench θ={cb_θ}: "
                  f"would_fire on {cb_fp*len(scores):.0f}/{len(scores)} = {cb_fp:.4f} STRICT pairs "
                  f"(= the real-traffic FP rate of that candidate-bench θ)", flush=True)
        print(f"[relevance-sweep]   smallest θ for fp=0 on this pool: {zero_fp_θ:.4f}  "
              f"(any runtime θ must be ≤ this to never false-demote a STRICT)", flush=True)
        # what's the recommendation? a θ at min(STRICT)-ε is safe but might lose recall;
        # report a few "headroom" points
        headroom = []
        for pct in (0.0, 1.0, 5.0, 10.0):
            cut = sorted(scores)[max(0, int(round(pct / 100 * (len(scores) - 1))))]
            headroom.append({"pct_strict_below": pct, "theta_at_that_cut": round(cut, 4)})
        report["models"][mv] = {
            "hf_repo": me["hf_repo"], "backend": rel.backend, "device": rel.device,
            "n_pairs": len(pairs), "infer_seconds": round(infer_s, 2),
            "score_quantiles_on_strict": q,
            "candidate_bench_theta": cb_θ,
            "fp_rate_at_candidate_bench_theta": cb_fp,
            "smallest_theta_for_zero_fp_on_pool": round(zero_fp_θ, 4) if zero_fp_θ is not None else None,
            "headroom_table": headroom,
        }

    args.out.parent.mkdir(parents=True, exist_ok=True)
    args.out.write_text(json.dumps(report, indent=2))
    print(f"[relevance-sweep] wrote {args.out}")

    # summary table
    print("\n" + "=" * 110)
    print(f"RELEVANCE SHADOW SWEEP on {len(strict)} pooled bench-qa STRICT pairs")
    print("=" * 110)
    print(f"{'model':<40} {'cb θ':>8} {'cb-θ STRICT FP':>16} {'min(STRICT) = fp=0 θ':>24}")
    print("-" * 110)
    for mv, md in sorted(report["models"].items()):
        short = mv.replace("relevance-shadow-v1-", "")
        cb_θ = md["candidate_bench_theta"]
        cb_fp = md["fp_rate_at_candidate_bench_theta"]
        z = md["smallest_theta_for_zero_fp_on_pool"]
        cb_θ_s = f"{cb_θ:+.3f}" if cb_θ is not None else "—"
        cb_fp_s = f"{cb_fp*md['n_pairs']:.0f}/{md['n_pairs']} = {cb_fp:.3f}" if cb_fp is not None else "—"
        z_s = f"{z:+.3f}" if z is not None else "—"
        print(f"{short:<40} {cb_θ_s:>8} {cb_fp_s:>16} {z_s:>24}")
    print("-" * 110)
    print("cb θ                  = the candidate-bench fp=0 threshold from manifest §3.2.1 round-2 (set on 26 contrived pairs)")
    print("cb-θ STRICT FP        = fraction of pooled bench-qa STRICT pairs that would FIRE at the candidate-bench θ")
    print("                        (this is the load-bearing 'did the candidate-bench θ survive proper-n' number)")
    print("min(STRICT) = fp=0 θ  = the smallest θ that yields zero FP on THIS pool (= real-traffic-calibrated threshold)")
    print("→ if 'cb-θ STRICT FP' is large (≫0), the candidate-bench overestimated the safe θ and we walk back.")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())