304 lines
16 KiB
Python
304 lines
16 KiB
Python
#!/usr/bin/env python3
|
||
"""NLI shadow GRID sweep — #000049 §7 #22 follow-up: characterize the
|
||
whole {model × candidate-clause-cap × aggregation × θc × θe} space at
|
||
once, to learn whether *any* config of the lexical-candidate approach
|
||
clears the §7 #12 gate (catch the recombinations, ~0 false-demote on
|
||
already-STRICT answers) before deciding the Phase-3 verifier-matched-
|
||
clause hook is the only path.
|
||
|
||
Inputs (both shapes from `nli_shadow_sweep.py` are accepted per record:
|
||
`{answer_text|claim, context|source}`; `_meta` lines skipped):
|
||
--should-demote records where demoting IS correct (false claims /
|
||
recombination fixtures). Default: the in-repo
|
||
`bench/fixtures/5f/falsification-hard-v1.jsonl`
|
||
(12; the 2 recombination fixtures + 10 other false
|
||
near-misses). For the richer 28-case synthetic set
|
||
point this at `arborist-nli-bench/eval/recombination.jsonl`.
|
||
--should-not-demote records where demoting is a FALSE POSITIVE — a
|
||
bench-qa JSONL; rows are filtered to `audit_mode ==
|
||
STRICT` (the confidently-grounded answers) unless
|
||
--no-strict-filter. (Add `arborist-nli-bench/eval/
|
||
legit_summary.jsonl` too if you have it.)
|
||
|
||
How it stays cheap: NLI inference runs ONCE per (model, record) over
|
||
the record's top-`KMAX` candidate clauses; the k/aggregation/θ grid is
|
||
then pure arithmetic on the cached per-clause scores. On a GPU the
|
||
whole grid is seconds. Models come from `arborist/qa/nli/manifest.json`
|
||
(`nli_model_version`+`hf_repo`+`pinned_revision`, plus `alternates`).
|
||
|
||
SHADOW ONLY — writes nothing but `bench/results/nli-shadow-grid.json`.
|
||
|
||
Usage:
|
||
python3 bench/scripts/nli_shadow_grid.py --should-not-demote bench/qa_results/<...>.jsonl
|
||
python3 bench/scripts/nli_shadow_grid.py --should-not-demote a.jsonl --should-demote ~/git/arborist-nli-bench/eval/recombination.jsonl --models all
|
||
python3 bench/scripts/nli_shadow_grid.py --should-not-demote a.jsonl --models nli-shadow-v1-minilm2-l6-h768,nli-shadow-v1-bart-large-mnli
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import json
|
||
import sys
|
||
import time
|
||
from pathlib import Path
|
||
|
||
REPO = Path(__file__).resolve().parents[2]
|
||
KMAX = 12 # NLI is run on the top-KMAX candidate clauses; the grid sub-selects k <= KMAX.
|
||
K_GRID = (1, 2, 3, 4, 6, 8, 12)
|
||
# how to combine the per-clause contradiction scores over the top-k clauses:
|
||
# max / mean — over all k; top2 / top3 — the 2nd / 3rd highest (needs agreement);
|
||
# meanTop2 / meanTop3 — mean of the 2 / 3 highest; margin — max over clauses of
|
||
# (p_contra - p_entail) of THAT clause (folds the entailment guard into the score).
|
||
AGGS = ("max", "mean", "top2", "top3", "meanTop2", "meanTop3", "margin")
|
||
# the entailment guard variant: max_entail = max p_entail over the top-k clauses;
|
||
# paired_entail = p_entail of the SAME clause that produced the aggregated contra
|
||
# (the principled version — "the contradicting clause itself doesn't entail").
|
||
GUARDS = ("max_entail", "paired_entail")
|
||
THETA_C = (0.30, 0.40, 0.50, 0.60, 0.70, 0.80, 0.85, 0.90, 0.92, 0.94, 0.95, 0.96,
|
||
0.97, 0.975, 0.98, 0.985, 0.99, 0.995, 0.998, 0.999)
|
||
THETA_E = (0.30, 0.50, 0.70, 0.90, 0.95, 1.01) # 1.01 == "no guard"
|
||
|
||
|
||
def _records(path: Path):
|
||
for ln in path.read_text().splitlines():
|
||
ln = ln.strip()
|
||
if not ln:
|
||
continue
|
||
o = json.loads(ln)
|
||
if "_meta" in o:
|
||
continue
|
||
claim = o.get("answer_text") or o.get("claim")
|
||
source = o.get("context") or o.get("source")
|
||
if not claim or not source:
|
||
continue
|
||
yield {"id": o.get("id", path.stem), "claim": claim, "source": source,
|
||
"audit_mode": o.get("audit_mode")}
|
||
|
||
|
||
def _models_from_manifest(which: str | None):
|
||
sys.path.insert(0, str(REPO))
|
||
from arborist.qa.nli.shadow import load_manifest
|
||
m = load_manifest()
|
||
entries = [{"nli_model_version": m["nli_model_version"], "hf_repo": m["hf_repo"],
|
||
"pinned_revision": m.get("pinned_revision")}]
|
||
entries += [{"nli_model_version": a["nli_model_version"], "hf_repo": a["hf_repo"],
|
||
"pinned_revision": a.get("pinned_revision")} for a in m.get("alternates", [])]
|
||
if which in (None, ""):
|
||
return entries[:1]
|
||
if which == "all":
|
||
return entries
|
||
want = {w.strip() for w in which.split(",")}
|
||
sel = [e for e in entries if e["nli_model_version"] in want or e["hf_repo"] in want]
|
||
return sel or entries[:1]
|
||
|
||
|
||
def _aggregate(top: list[tuple], how: str) -> tuple[float, int]:
|
||
"""Return (aggregated_contradiction_score, idx_of_the_clause_it_came_from)
|
||
over `top` = [(overlap, pe, pn, pc), ...] (top-k clauses, overlap-ranked).
|
||
For pooled aggs (mean/topN/meanTopN) the 'idx' is the highest-pc clause in
|
||
the pool — used for the paired-entail guard."""
|
||
if not top:
|
||
return 0.0, -1
|
||
pcs = [t[3] for t in top]
|
||
order = sorted(range(len(top)), key=lambda i: pcs[i], reverse=True)
|
||
if how == "max":
|
||
return pcs[order[0]], order[0]
|
||
if how == "mean":
|
||
return sum(pcs) / len(pcs), order[0]
|
||
if how == "top2":
|
||
return (pcs[order[1]], order[1]) if len(order) >= 2 else (0.0, order[0])
|
||
if how == "top3":
|
||
return (pcs[order[2]], order[2]) if len(order) >= 3 else (0.0, order[-1])
|
||
if how == "meanTop2":
|
||
sel = order[:2]; return sum(pcs[i] for i in sel) / len(sel), order[0]
|
||
if how == "meanTop3":
|
||
sel = order[:3]; return sum(pcs[i] for i in sel) / len(sel), order[0]
|
||
if how == "margin": # max over clauses of (pc - pe); guard folded in
|
||
margins = [t[3] - t[1] for t in top]
|
||
bi = max(range(len(top)), key=lambda i: margins[i])
|
||
return margins[bi], bi
|
||
raise ValueError(how)
|
||
|
||
|
||
def _cache_scores(nli, recs: list[dict]) -> list[dict]:
|
||
"""For each record: segment, take top-KMAX candidate clauses, NLI each.
|
||
Returns [{id, audit_mode, n_clauses, clause_scores=[(overlap,pe,pn,pc),...]}]."""
|
||
from arborist.qa.nli.shadow import clauses as _clauses, candidate_clauses as _cand
|
||
out = []
|
||
for r in recs:
|
||
n_clauses = len(_clauses(r["source"]))
|
||
cand = _cand(r["claim"], r["source"], KMAX) # list[(clause, overlap)]
|
||
triples = nli._nli_batch([(cl, r["claim"]) for cl, _ in cand]) if cand else []
|
||
clause_scores = [(ov, pe, pn, pc) for (_, ov), (pe, pn, pc) in zip(cand, triples)]
|
||
out.append({"id": r["id"], "audit_mode": r.get("audit_mode"),
|
||
"n_clauses": n_clauses, "clause_scores": clause_scores})
|
||
return out
|
||
|
||
|
||
def _rate(cached: list[dict], k: int, agg: str, guard: str, tc: float, te: float) -> float:
|
||
"""fraction of records that would_demote under (k, agg, guard, θc, θe).
|
||
For agg='margin' the guard is folded into the score (margin = pc-pe), so
|
||
only θc applies; the explicit guard loop is skipped for it (te is ignored)."""
|
||
n = len(cached)
|
||
if not n:
|
||
return 0.0
|
||
fired = 0
|
||
is_margin = (agg == "margin")
|
||
for c in cached:
|
||
top = c["clause_scores"][:k]
|
||
if not top:
|
||
continue
|
||
agg_c, src_i = _aggregate(top, agg)
|
||
if agg_c < tc:
|
||
continue
|
||
if is_margin:
|
||
fired += 1
|
||
continue
|
||
if guard == "max_entail":
|
||
guard_e = max(t[1] for t in top)
|
||
else: # paired_entail — the entail of the clause that produced agg_c
|
||
guard_e = top[src_i][1] if 0 <= src_i < len(top) else max(t[1] for t in top)
|
||
if guard_e < te:
|
||
fired += 1
|
||
return fired / n
|
||
|
||
|
||
def _pareto(pts: list[dict]) -> list[dict]:
|
||
"""Pareto frontier: sorted by fp asc, keep a point only if its catch beats
|
||
the best catch seen at a lower-or-equal fp."""
|
||
out, best = [], -1.0
|
||
for p in sorted(pts, key=lambda p: (p["fp"], -p["catch"])):
|
||
if p["catch"] > best:
|
||
out.append(p); best = p["catch"]
|
||
return out
|
||
|
||
|
||
def main(argv=None) -> int:
|
||
ap = argparse.ArgumentParser()
|
||
ap.add_argument("--should-demote", type=Path,
|
||
default=REPO / "bench" / "fixtures" / "5f" / "falsification-hard-v1.jsonl")
|
||
ap.add_argument("--should-not-demote", type=Path, required=True,
|
||
help="a bench-qa JSONL (rows filtered to audit_mode==STRICT)")
|
||
ap.add_argument("--no-strict-filter", action="store_true",
|
||
help="don't filter the should-not-demote set to STRICT rows")
|
||
ap.add_argument("--models", default=None, help="manifest model id(s) comma-sep, or 'all' (default: main only)")
|
||
ap.add_argument("--extra-models", default=None, help="comma-sep HF repo ids to also test (no manifest pin needed) — for 'no stone unturned' breadth")
|
||
ap.add_argument("--out", type=Path, default=REPO / "bench" / "results" / "nli-shadow-grid.json")
|
||
args = ap.parse_args(argv)
|
||
|
||
try:
|
||
import torch # noqa
|
||
sys.path.insert(0, str(REPO))
|
||
from arborist.qa.nli.shadow import ShadowNLI
|
||
except ImportError as e:
|
||
print(f"[nli-grid] missing dependency: {e} (pip install 'arborist[nli]')", file=sys.stderr)
|
||
return 2
|
||
|
||
sd = list(_records(args.should_demote))
|
||
snd = list(_records(args.should_not_demote))
|
||
if not args.no_strict_filter:
|
||
snd_all = len(snd)
|
||
snd = [r for r in snd if r.get("audit_mode") == "STRICT"]
|
||
print(f"[nli-grid] should-not-demote: {len(snd)}/{snd_all} rows are audit_mode==STRICT", flush=True)
|
||
print(f"[nli-grid] should-demote: {len(sd)} records ({args.should_demote.name})", flush=True)
|
||
if not sd or not snd:
|
||
print("[nli-grid] need non-empty inputs on both sides", file=sys.stderr)
|
||
return 2
|
||
|
||
models = _models_from_manifest(args.models)
|
||
if args.extra_models:
|
||
for repo in [r.strip() for r in args.extra_models.split(",") if r.strip()]:
|
||
mv = "extra:" + repo.split("/")[-1]
|
||
if not any(m["hf_repo"] == repo for m in models):
|
||
models.append({"nli_model_version": mv, "hf_repo": repo, "pinned_revision": None})
|
||
print(f"[nli-grid] models ({len(models)}): {[m['nli_model_version'] for m in models]}", flush=True)
|
||
|
||
report = {"generated_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
|
||
"should_demote_file": str(args.should_demote), "n_should_demote": len(sd),
|
||
"should_not_demote_file": str(args.should_not_demote), "n_should_not_demote": len(snd),
|
||
"k_grid": list(K_GRID), "aggs": list(AGGS), "guards": list(GUARDS),
|
||
"theta_c_grid": list(THETA_C), "theta_e_grid": list(THETA_E),
|
||
"kmax": KMAX, "models": {}}
|
||
|
||
for me in models:
|
||
mv = me["nli_model_version"]
|
||
print(f"[nli-grid] === {mv} ({me['hf_repo']}) ===", flush=True)
|
||
manifest_override = {"nli_model_version": mv, "hf_repo": me["hf_repo"],
|
||
"pinned_revision": me.get("pinned_revision"),
|
||
"thresholds": {"contradiction_veto": 0.5, "entailment_block_veto": 0.9},
|
||
"max_length": 256, "max_candidate_clauses": KMAX}
|
||
nli = ShadowNLI(manifest=manifest_override)
|
||
nli._ensure_loaded()
|
||
if not nli.available:
|
||
print(f"[nli-grid] skip — {nli._reason}", file=sys.stderr)
|
||
continue
|
||
t0 = time.time()
|
||
cache_sd = _cache_scores(nli, sd)
|
||
cache_snd = _cache_scores(nli, snd)
|
||
infer_s = time.time() - t0
|
||
n_pairs = sum(len(c["clause_scores"]) for c in cache_sd + cache_snd)
|
||
print(f"[nli-grid] backend={nli.backend} device={nli.device} · {n_pairs} NLI pairs in {infer_s:.2f}s", flush=True)
|
||
|
||
configs = [] # best fp=0 (else fp<=.05 / knee) point per (k, agg, guard)
|
||
all_pts = [] # every (model-internal) point, for the global pareto + global best
|
||
for k in K_GRID:
|
||
for agg in AGGS:
|
||
guard_list = ("max_entail",) if agg == "margin" else GUARDS
|
||
te_list = (1.01,) if agg == "margin" else THETA_E
|
||
for guard in guard_list:
|
||
pts = []
|
||
for te in te_list:
|
||
for tc in THETA_C:
|
||
catch = _rate(cache_sd, k, agg, guard, tc, te)
|
||
fp = _rate(cache_snd, k, agg, guard, tc, te)
|
||
p = {"k": k, "agg": agg, "guard": guard, "theta_c": tc,
|
||
"theta_e": te, "catch": round(catch, 4), "fp": round(fp, 4)}
|
||
pts.append(p); all_pts.append(p)
|
||
zero = [p for p in pts if p["fp"] == 0.0]
|
||
low = [p for p in pts if p["fp"] <= 0.05]
|
||
if zero:
|
||
best = max(zero, key=lambda p: p["catch"]); best_kind = "fp=0"
|
||
elif low:
|
||
best = max(low, key=lambda p: p["catch"]); best_kind = "fp<=0.05"
|
||
else:
|
||
best = max(pts, key=lambda p: p["catch"] - p["fp"]); best_kind = "knee"
|
||
configs.append({"k": k, "agg": agg, "guard": guard, "best": best, "best_kind": best_kind})
|
||
def _ck(c):
|
||
b, kind = c["best"], c["best_kind"]
|
||
return (0 if kind == "fp=0" else 1 if kind == "fp<=0.05" else 2, -b["catch"], b["fp"])
|
||
configs.sort(key=_ck)
|
||
# global best fp=0 point across ALL (k,agg,guard,θc,θe), then the pareto frontier
|
||
zero_all = [p for p in all_pts if p["fp"] == 0.0]
|
||
global_best = max(zero_all, key=lambda p: p["catch"]) if zero_all else max(all_pts, key=lambda p: p["catch"] - p["fp"])
|
||
pareto = _pareto(all_pts)
|
||
report["models"][mv] = {"hf_repo": me["hf_repo"], "backend": nli.backend, "device": nli.device,
|
||
"n_nli_pairs": n_pairs, "infer_seconds": round(infer_s, 2),
|
||
"global_best": global_best, "pareto": pareto, "configs_ranked": configs}
|
||
gb = global_best
|
||
print(f"[nli-grid] GLOBAL best fp=0: k={gb['k']} agg={gb['agg']} guard={gb['guard']} θc={gb['theta_c']} θe={gb['theta_e']} "
|
||
f"→ catch={gb['catch']} fp={gb['fp']}", flush=True)
|
||
|
||
args.out.parent.mkdir(parents=True, exist_ok=True)
|
||
args.out.write_text(json.dumps(report, indent=2))
|
||
print(f"[nli-grid] wrote {args.out}")
|
||
|
||
# summary table
|
||
print("\n" + "=" * 100)
|
||
print(f"NLI SHADOW GRID — global-best fp=0 config per model ({len(sd)} should-demote vs {len(snd)} should-not-demote)")
|
||
print("=" * 100)
|
||
print(f"{'model':<32} {'hf_repo':<44} {'k':>3} {'agg':>9} {'guard':>14} {'θc':>6} {'θe':>5} {'catch':>7} {'fp':>6}")
|
||
print("-" * 100)
|
||
ranked = sorted(report["models"].items(), key=lambda kv: (-(kv[1]["global_best"]["catch"] if kv[1]["global_best"]["fp"] == 0 else -1), kv[1]["global_best"]["fp"]))
|
||
for mv, md in ranked:
|
||
g = md["global_best"]
|
||
print(f"{mv:<32} {md['hf_repo']:<44} {g['k']:>3} {g['agg']:>9} {g['guard']:>14} {g['theta_c']:>6} {g['theta_e']:>5} {g['catch']:>7.3f} {g['fp']:>6.3f}")
|
||
print("-" * 100)
|
||
print("catch = fraction of should-demote (recombination falsehoods) the veto fires on (higher better)")
|
||
print("fp = fraction of should-not-demote (confidently-grounded STRICT answers) the veto WRONGLY fires on (≈0 clears §7 #12)")
|
||
print("'global_best' = the highest-catch point with fp==0 across the FULL {k × agg × guard × θc × θe} grid for that model.")
|
||
print("Each model's 'pareto' (in the JSON) is the catch-vs-fp frontier — read it to pick an operating point above fp=0 if a small false-demote rate is acceptable.")
|
||
print("If no model reaches a useful catch at fp=0, the lexical-candidate approach is precision-limited — Phase-3 verifier-matched-clause hook is the path.")
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(main())
|