arborist/bench/scripts/nli_backend_ab.py

#!/usr/bin/env python3
"""NLI inference-backend A/B — #000049 §7 #28.

Loads the *same pinned shadow-NLI checkpoint* under each available
inference backend and answers two questions, in priority order:

  1. AGREEMENT (the gate): does the backend produce the same
     ``(p_entail, p_neutral, p_contra)`` as the ``torch`` reference,
     within tolerance, over an identical fixed pair set? A faster
     engine that *diverges* is not a faster shadow — it is a
     *different* shadow signal (and, under the deliberately-unwired
     Phase-3 demotion hook, a different ``nli_policy_hash``). So
     divergence fails the A/B regardless of speed.
  2. LATENCY (secondary): p50/p95 per-pair wall time.

Backends probed (only those that actually load are reported):
  • ``torch``      — the reference (`transformers` + torch fwd path)
  • ``onnx-int8``  — the §7 #22 `optimum.onnxruntime` export, if present
  • ``tinygrad``   — the §7 #28 tinygrad ONNX-frontend backend, if
                     tinygrad is importable AND its frontend covers
                     this checkpoint's ops (UNVERIFIED by design —
                     this script *measures* it, never assumes it)

Deterministic: a fixed in-repo pair set, no LLM, no verifier, no n=3
noise, no 5pp floor — the `recall_at_k` instrument discipline. SHADOW
only: writes nothing but ``bench/results/nli-backend-ab.json``.

This belongs on the GPU producer box (the #000051 "heavy passes never
in arborist's python+sqlite3 core" pattern); tinygrad is not an
arborist dependency. Install tinygrad there, then:

    make bench-nli-backends
    python3 bench/scripts/nli_backend_ab.py --repeats 5
"""
from __future__ import annotations

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

REPO = Path(__file__).resolve().parents[2]

# Fixed, deterministic (premise, hypothesis) probe set — a spread of
# entailment / contradiction / neutral over factual content. Not a
# quality benchmark; an engine-agreement fixture (same input every run).
PAIRS = [
    ("The Eiffel Tower is located in Paris, France.", "The Eiffel Tower is in Paris."),
    ("The Eiffel Tower is located in Paris, France.", "The Eiffel Tower is in Berlin."),
    ("The Eiffel Tower is located in Paris, France.", "Paris has many museums."),
    ("Water boils at 100 degrees Celsius at sea level.", "Water boils at 100 C at sea level."),
    ("Water boils at 100 degrees Celsius at sea level.", "Water boils at 50 degrees Celsius at sea level."),
    ("Water boils at 100 degrees Celsius at sea level.", "Ice is frozen water."),
    ("Mount Everest is the highest mountain above sea level.", "Mount Everest is the tallest mountain above sea level."),
    ("Mount Everest is the highest mountain above sea level.", "K2 is the highest mountain above sea level."),
    ("Mount Everest is the highest mountain above sea level.", "Many climbers attempt Everest each year."),
    ("The novel was written by George Orwell in 1949.", "George Orwell wrote the novel in 1949."),
    ("The novel was written by George Orwell in 1949.", "The novel was written by Aldous Huxley."),
    ("The novel was written by George Orwell in 1949.", "The novel has been widely translated."),
    ("Photosynthesis converts sunlight into chemical energy in plants.", "Plants use sunlight to make chemical energy."),
    ("Photosynthesis converts sunlight into chemical energy in plants.", "Photosynthesis converts chemical energy into sunlight."),
    ("Photosynthesis converts sunlight into chemical energy in plants.", "Some plants grow in shade."),
    ("The company reported a net profit of two billion dollars last year.", "The company was profitable last year."),
    ("The company reported a net profit of two billion dollars last year.", "The company reported a net loss last year."),
    ("The company reported a net profit of two billion dollars last year.", "The company has offices worldwide."),
    ("Light travels faster than sound.", "Sound travels faster than light."),
    ("Light travels faster than sound.", "Light is faster than sound."),
    ("The treaty was signed by both nations in 1815.", "Both nations signed the treaty in 1815."),
    ("The treaty was signed by both nations in 1815.", "Neither nation signed the treaty."),
    ("The species is native to the islands of Southeast Asia.", "The species comes from Southeast Asian islands."),
    ("The species is native to the islands of Southeast Asia.", "The species is native to northern Europe."),
]


def _fresh_shadow(env_overrides: dict):
    """Construct a ShadowNLI with a scoped os.environ, force-load it,
    restore the environ. Returns the (loaded) instance."""
    from arborist.qa.nli.shadow import ShadowNLI

    saved = {k: os.environ.get(k) for k in env_overrides}
    try:
        for k, v in env_overrides.items():
            if v is None:
                os.environ.pop(k, None)
            else:
                os.environ[k] = v
        nli = ShadowNLI()
        nli._ensure_loaded()
        return nli
    finally:
        for k, old in saved.items():
            if old is None:
                os.environ.pop(k, None)
            else:
                os.environ[k] = old


def _probe(nli, repeats: int):
    """Run the fixed PAIRS through _nli_batch `repeats` times. Returns
    (per_pair_probs, per_pair_seconds_list)."""
    _ = nli._nli_batch(PAIRS[:2])  # warmup (graph build / kernel JIT)
    per_iter_secs = []
    probs = None
    for _ in range(repeats):
        t0 = time.perf_counter()
        out = nli._nli_batch(PAIRS)
        per_iter_secs.append((time.perf_counter() - t0) / len(PAIRS))
        probs = out  # deterministic — last iter is representative
    return probs, per_iter_secs


def _max_abs_delta(a, b):
    m = 0.0
    s = 0.0
    n = 0
    for (ae, an, ac), (be, bn, bc) in zip(a, b):
        for x, y in ((ae, be), (an, bn), (ac, bc)):
            d = abs(x - y)
            m = max(m, d)
            s += d
            n += 1
    return m, (s / n if n else 0.0)


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--repeats", type=int, default=5, help="timed iterations per backend")
    ap.add_argument("--tol", type=float, default=2e-3,
                    help="max abs probability delta vs torch to PASS agreement")
    ap.add_argument("--out", default=str(REPO / "bench" / "results" / "nli-backend-ab.json"))
    args = ap.parse_args(argv)
    sys.path.insert(0, str(REPO))

    # A nonexistent dir forces the torch path even if an export exists.
    no_onnx = str(REPO / ".nonexistent-onnx-dir")
    plans = [
        ("torch", {"ARBORIST_NLI_BACKEND": None, "ARBORIST_NLI_ONNX_DIR": no_onnx}),
        ("onnx", {"ARBORIST_NLI_BACKEND": None, "ARBORIST_NLI_ONNX_DIR": None}),
        ("tinygrad", {"ARBORIST_NLI_BACKEND": "tinygrad", "ARBORIST_NLI_ONNX_DIR": None}),
    ]

    results = {}
    ref_probs = None
    for name, env in plans:
        print(f"[nli-ab] probing {name} ...", flush=True)
        try:
            nli = _fresh_shadow(env)
        except Exception as e:  # noqa: BLE001
            results[name] = {"available": False, "reason": f"{type(e).__name__}: {e}"}
            print(f"[nli-ab]   {name}: load raised — {e}", flush=True)
            continue
        if not nli.available:
            results[name] = {"available": False, "reason": nli._reason}
            print(f"[nli-ab]   {name}: unavailable — {nli._reason}", flush=True)
            continue
        # Instrument honesty: ShadowNLI *correctly* degrades to a working
        # backend in production, but for an A/B that means "requested X,
        # got Y" must be reported as X-unavailable — never relabel Y's
        # numbers as X's (false-positive-is-corruption; #000049's own
        # discipline). torch=exactly torch; onnx=either onnx variant;
        # tinygrad=exactly tinygrad.
        ok = {"torch": nli.backend == "torch",
              "onnx": nli.backend in ("onnx", "onnx-int8"),
              "tinygrad": nli.backend == "tinygrad"}.get(name, True)
        if not ok:
            reason = f"requested {name} but ShadowNLI loaded {nli.backend} (engine not usable here)"
            results[name] = {"available": False, "reason": reason}
            print(f"[nli-ab]   {name}: unavailable — {reason}", flush=True)
            continue
        probs, secs = _probe(nli, args.repeats)
        p50 = statistics.median(secs)
        p95 = sorted(secs)[max(0, int(len(secs) * 0.95) - 1)]
        rec = {
            "available": True,
            "backend": nli.backend,
            "device": nli.device,
            "p50_ms_per_pair": round(p50 * 1e3, 3),
            "p95_ms_per_pair": round(p95 * 1e3, 3),
        }
        if name == "torch":
            ref_probs = probs
            rec["agreement_vs_torch"] = "reference"
        elif ref_probs is not None and probs is not None:
            mx, mean = _max_abs_delta(ref_probs, probs)
            rec["max_abs_delta_vs_torch"] = round(mx, 6)
            rec["mean_abs_delta_vs_torch"] = round(mean, 6)
            rec["agreement_pass"] = bool(mx <= args.tol)
        else:
            rec["agreement_vs_torch"] = "no torch reference to compare against"
        results[name] = rec
        print(f"[nli-ab]   {name}: backend={nli.backend} device={nli.device} "
              f"p50={rec['p50_ms_per_pair']}ms/pair "
              f"{('Δmax=' + str(rec.get('max_abs_delta_vs_torch'))) if 'max_abs_delta_vs_torch' in rec else ''}",
              flush=True)

    out = Path(args.out)
    out.parent.mkdir(parents=True, exist_ok=True)
    payload = {
        "ticket": "#000049 §7 #28",
        "n_pairs": len(PAIRS),
        "repeats": args.repeats,
        "tolerance": args.tol,
        "shadow_only": True,
        "results": results,
    }
    out.write_text(json.dumps(payload, indent=2))
    print(f"\n[nli-ab] wrote {out}", flush=True)

    # Verdict (informational; SHADOW — never gates a build red).
    print("\n=== verdict ===", flush=True)
    print("Gate is AGREEMENT first, latency second:", flush=True)
    for name in ("onnx", "tinygrad"):
        r = results.get(name, {})
        if not r.get("available"):
            print(f"  {name}: not available — {r.get('reason','?')}", flush=True)
            continue
        if "agreement_pass" in r:
            tag = "PASS" if r["agreement_pass"] else "DIVERGES"
            print(f"  {name}: agreement {tag} (Δmax={r['max_abs_delta_vs_torch']}, "
                  f"tol={args.tol}) · {r['p50_ms_per_pair']}ms/pair p50", flush=True)
    return 0


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