arborist/bench/scripts/nli_shadow_sweep.py

#!/usr/bin/env python3
"""NLI shadow sweep — #000049 Phase 2 / §7 #12 gate item 4 instrument.

Runs the clause-level shadow check (``arborist.qa.nli.shadow_check`` —
the §7 #5 algorithm) over a set of (answer, context) records and reports
the *would-demote* rate, bucketed by the verifier label the record
carries. SHADOW ONLY: writes nothing to any shard, touches no
``audit_mode``.

Two record shapes are accepted (auto-detected per line):
  - 5f-fixture shape: ``{"answer_text": …, "context": …,
    "expected_reason": …}`` (e.g. ``bench/fixtures/5f/*.jsonl``;
    ``_meta`` lines are skipped). ``expected_reason`` is the bucket.
  - nli-bench eval shape: ``{"claim": …, "source": …, "want": …}``
    (e.g. ``~/git/arborist-nli-bench/eval/*.jsonl``). ``want`` is the
    bucket; ``want="not_contradiction"`` rows are the false-positive
    probe — any would_demote on those is a shadow FP.

With no ``--input`` it sweeps the 5f falsification packs that have
``answer_text``/``context`` columns. Requires the ``[nli]`` extra to
produce real numbers — without it the report still renders, marked
``available: false`` (so the harness/CI never breaks).

Usage:
  python3 bench/scripts/nli_shadow_sweep.py
  python3 bench/scripts/nli_shadow_sweep.py --input path/to/records.jsonl --input more.jsonl
  python3 bench/scripts/nli_shadow_sweep.py --out bench/results/nli-shadow-sweep.json
"""
from __future__ import annotations

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

REPO = Path(__file__).resolve().parents[2]
_DEFAULT_INPUTS = [
    REPO / "bench" / "fixtures" / "5f" / "falsification-hard-v1.jsonl",
    REPO / "bench" / "fixtures" / "5f" / "falsification-live-v1.jsonl",
    REPO / "bench" / "fixtures" / "5f" / "falsification-v1.jsonl",
]


def _records(path: Path):
    for ln in path.read_text().splitlines():
        ln = ln.strip()
        if not ln:
            continue
        obj = json.loads(ln)
        if "_meta" in obj:
            continue
        claim = obj.get("answer_text") or obj.get("claim")
        source = obj.get("context") or obj.get("source")
        if not claim or not source:
            continue
        bucket = (obj.get("expected_reason") or obj.get("want")
                  or obj.get("audit_mode") or "unlabeled")
        # On real bench-qa rows there is no ground truth; treat a
        # would_demote on a STRICT row as a *potential* false positive
        # worth surfacing (the lexical verifier was confident, NLI
        # disagrees). On nli-bench eval rows `want=not_contradiction`
        # is the authoritative FP probe.
        is_fp_probe = (obj.get("want") == "not_contradiction"
                       or (obj.get("want") is None and obj.get("expected_reason") is None
                           and obj.get("audit_mode") == "STRICT"))
        yield {"id": obj.get("id", path.stem), "claim": claim, "source": source,
               "bucket": bucket, "is_fp_probe": is_fp_probe, "src_file": path.name}


def main(argv=None) -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--input", action="append", type=Path, help="JSONL record file(s); repeatable")
    ap.add_argument("--out", type=Path, default=REPO / "bench" / "results" / "nli-shadow-sweep.json")
    args = ap.parse_args(argv)

    inputs = args.input or [p for p in _DEFAULT_INPUTS if p.exists()]
    recs = []
    for p in inputs:
        if not p.exists():
            print(f"[nli-shadow] skip (missing): {p}", file=sys.stderr)
            continue
        n0 = len(recs)
        recs.extend(_records(p))
        print(f"[nli-shadow] {p.name}: {len(recs) - n0} records", flush=True)
    if not recs:
        print("[nli-shadow] no usable records", file=sys.stderr)
        return 2

    from arborist.qa.nli import ShadowNLI
    nli = ShadowNLI()

    by_bucket: dict[str, dict] = defaultdict(lambda: {"n": 0, "would_demote": 0})
    by_recomb_risk: dict[str, dict] = {"risk": {"n": 0, "would_demote": 0},
                                       "no_risk": {"n": 0, "would_demote": 0}}
    fp_probe = {"n": 0, "would_demote": 0}
    rows = []
    n_available = 0
    n_candidate_clauses_total = 0
    t0 = time.time()
    for r in recs:
        res = nli.check(r["claim"], r["source"])
        if res.available:
            n_available += 1
        n_candidate_clauses_total += res.n_candidate_clauses
        b = by_bucket[r["bucket"]]
        b["n"] += 1
        rr = by_recomb_risk["risk" if res.recombination_risk else "no_risk"]
        rr["n"] += 1
        if res.would_demote:
            b["would_demote"] += 1
            rr["would_demote"] += 1
        if r["is_fp_probe"]:
            fp_probe["n"] += 1
            if res.would_demote:
                fp_probe["would_demote"] += 1
        rows.append({**{k: r[k] for k in ("id", "bucket", "is_fp_probe", "src_file")},
                     **res.as_dict()})
    elapsed = time.time() - t0

    available = n_available > 0
    reason = "ok" if available else (rows[0]["reason"] if rows else "no rows")
    total_demote = sum(b["would_demote"] for b in by_bucket.values())
    report = {
        "generated_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "available": available,
        "reason": reason,
        "model_version": nli.model_version,
        "theta_contra": nli.theta_contra,
        "theta_entail": nli.theta_entail,
        "n_records": len(recs),
        "n_available": n_available,
        "elapsed_seconds": round(elapsed, 1),
        "would_demote_total": total_demote,
        "would_demote_rate": round(total_demote / len(recs), 4) if recs else 0.0,
        "mean_candidate_clauses": round(n_candidate_clauses_total / len(recs), 2) if recs else 0.0,
        "max_candidate_clauses_cap": nli.max_candidate_clauses,
        "by_bucket": {k: {**v, "rate": round(v["would_demote"] / v["n"], 4) if v["n"] else 0.0}
                      for k, v in sorted(by_bucket.items())},
        "by_recombination_risk": {k: {**v, "rate": round(v["would_demote"] / v["n"], 4) if v["n"] else 0.0}
                                  for k, v in by_recomb_risk.items()},
        "false_positive_probe": {**fp_probe,
                                 "rate": round(fp_probe["would_demote"] / fp_probe["n"], 4) if fp_probe["n"] else None,
                                 "note": "would_demote on records labeled want=not_contradiction — these are shadow FALSE POSITIVES; this is §7 #12 gate item 4 when the input is a real legit-answer sample"},
        "rows": rows,
    }
    args.out.parent.mkdir(parents=True, exist_ok=True)
    args.out.write_text(json.dumps(report, indent=2))
    print(f"[nli-shadow] wrote {args.out}")

    print("\n" + "=" * 72)
    print("NLI SHADOW SWEEP  —  #000049 Phase 2  (would-demote rate, shadow only)")
    print("=" * 72)
    if not available:
        print(f"  model UNAVAILABLE — {reason}")
        print("  (install the runtime to get real numbers:  pip install 'arborist[nli]')")
        print(f"  swept {len(recs)} records structurally; would-demote rate not measured.")
        return 0
    print(f"  model {nli.model_version}  ·  θc={nli.theta_contra} θe={nli.theta_entail}  ·  cand-clause cap {nli.max_candidate_clauses}  ·  {len(recs)} records  ·  {elapsed:.1f}s")
    print(f"  mean candidate clauses NLI'd / record: {report['mean_candidate_clauses']}  (vs the whole context — the §7 #20 restriction)")
    print(f"  would_demote overall: {total_demote}/{len(recs)} = {report['would_demote_rate']:.3f}")
    rr = report["by_recombination_risk"]
    print(f"  by recombination-risk:  risk {rr['risk']['would_demote']}/{rr['risk']['n']}={rr['risk']['rate']:.3f}   "
          f"no_risk {rr['no_risk']['would_demote']}/{rr['no_risk']['n']}={rr['no_risk']['rate']:.3f}   "
          f"(if no_risk-rate is high too, the risk-gate alone isn't enough)")
    print(f"  {'bucket':<28} {'n':>5} {'would_demote':>13} {'rate':>7}")
    print("  " + "-" * 56)
    for k, v in report["by_bucket"].items():
        print(f"  {k:<28} {v['n']:>5} {v['would_demote']:>13} {v['rate']:>7.3f}")
    print("  " + "-" * 56)
    fpp = report["false_positive_probe"]
    if fpp["n"]:
        print(f"  false-positive probe (want=not_contradiction): {fpp['would_demote']}/{fpp['n']} = {fpp['rate']:.3f}  ← lower is better; 0 clears §7 #12 item 2")
    else:
        print("  (no want=not_contradiction rows in this input — point --input at a legit-answer sample for the real gate-item-4 number)")
    print("  NB: on falsification-* packs every record is a FALSE claim, so a high would_demote rate there is the model working as intended — not a false positive.")
    return 0


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