russell@unturf.com
87d9db15c7
Speedup (§3 plan): ShadowNLI._nli_batch batches forwards (ARBORIST_NLI_BATCH=64); device auto-detect (ARBORIST_NLI_DEVICE, else cuda-if-available); auto-prefer an ONNX export — bench/scripts/export_nli_onnx.py / make export-nli-onnx exports + int8-dynamic-quantizes the pinned checkpoint into ~/.arborist/models/nli/<ver>/onnx/ (operator state, NOT committed), _ensure_loaded loads model_quantized.onnx via optimum.onnxruntime (backend onnx-int8), falls back to torch silently. torch-cpu-batch1 ~120ms/pair → onnx-int8-cpu-batched ~32ms/pair (~4x); seconds on a 4090. optimum[onnxruntime] added to the [nli] extra; 24 tests. Gate-item-4 verdict at proper n: ARBORIST_NLI_SHADOW=1 make bench-qa BENCH_QA_N=1 → 223 cells (89 STRICT / 90 HYBRID / 44 UNGROUNDED; also surfaced + fixed a lone-surrogate bug). Shadow sweep over those: NLI-as- runtime-veto on STRICT has ~26% FP at θc 0.5, ~8% at θc 0.90, ~0% only at θc 0.99 — and θc 0.99 gives up most recombination recall (hard synthetic recombinations bottom out ~0.76). FAILS the §7 #12 gate on this design. Only untried path that might pass: a Phase-3 runtime hook running NLI on the verifier's actual matched clauses (1-3), not top-6-by-overlap. Until then: runtime NLI demotion stays off; the 2 fixtures stay permanent boundary markers; θc stays 0.5. Production verifier unchanged; falsification-hard stays 10/12.
88 lines
3.9 KiB
Python
88 lines
3.9 KiB
Python
#!/usr/bin/env python3
|
||
"""Export the pinned #000049 shadow-NLI checkpoint to ONNX (+ optional
|
||
int8 dynamic quantization) — the §3 CPU speedup path.
|
||
|
||
`onnxruntime` on a quantized cross-encoder is typically 2–4× faster than
|
||
torch on CPU and drops the torch forward path; `ShadowNLI._ensure_loaded`
|
||
will prefer the export automatically if it finds it at the conventional
|
||
location (`~/.arborist/models/nli/<model_version>/onnx/`, or wherever
|
||
`ARBORIST_NLI_ONNX_DIR` points). Run once after `make bootstrap-nli`:
|
||
|
||
python3 bench/scripts/export_nli_onnx.py # export + int8 quantize
|
||
python3 bench/scripts/export_nli_onnx.py --no-quantize
|
||
python3 bench/scripts/export_nli_onnx.py --out /some/dir
|
||
|
||
Requires `optimum[onnxruntime]` (in the `[nli]` extra). SHADOW
|
||
infrastructure — the ONNX model is the same pinned checkpoint, same
|
||
labels; nothing about audit_mode changes (cf. ticket #000049 §7).
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
REPO = Path(__file__).resolve().parents[2]
|
||
|
||
|
||
def main(argv=None) -> int:
|
||
sys.path.insert(0, str(REPO))
|
||
from arborist.qa.nli.shadow import load_manifest, ShadowNLI
|
||
|
||
manifest = load_manifest()
|
||
repo_id = manifest["hf_repo"]
|
||
rev = manifest.get("pinned_revision")
|
||
default_out = ShadowNLI(manifest)._onnx_dir()
|
||
|
||
ap = argparse.ArgumentParser()
|
||
ap.add_argument("--out", type=Path, default=default_out,
|
||
help=f"output dir (default: {default_out})")
|
||
ap.add_argument("--no-quantize", action="store_true", help="skip int8 dynamic quantization")
|
||
args = ap.parse_args(argv)
|
||
|
||
try:
|
||
from transformers import AutoTokenizer
|
||
from optimum.onnxruntime import ORTModelForSequenceClassification
|
||
if not args.no_quantize:
|
||
from optimum.onnxruntime import ORTQuantizer
|
||
from optimum.onnxruntime.configuration import AutoQuantizationConfig
|
||
except ImportError as e:
|
||
print(f"[export-nli-onnx] missing dependency: {e}\n"
|
||
f" install with: pip install 'arborist[nli]' (pulls optimum[onnxruntime])",
|
||
file=sys.stderr)
|
||
return 2
|
||
|
||
out: Path = args.out
|
||
out.mkdir(parents=True, exist_ok=True)
|
||
print(f"[export-nli-onnx] {repo_id}@{rev or 'main'} → {out} (quantize={'no' if args.no_quantize else 'int8-dynamic'})", flush=True)
|
||
|
||
tok = AutoTokenizer.from_pretrained(repo_id, revision=rev)
|
||
model = ORTModelForSequenceClassification.from_pretrained(repo_id, revision=rev, export=True)
|
||
model.save_pretrained(out)
|
||
tok.save_pretrained(out)
|
||
print(f"[export-nli-onnx] fp32 ONNX written ({sum(f.stat().st_size for f in out.glob('*.onnx'))/1e6:.1f} MB of .onnx)", flush=True)
|
||
|
||
if not args.no_quantize:
|
||
quantizer = ORTQuantizer.from_pretrained(out)
|
||
qconfig = AutoQuantizationConfig.avx2(is_static=False, per_channel=False)
|
||
quantizer.quantize(save_dir=out, quantization_config=qconfig)
|
||
print(f"[export-nli-onnx] int8 quantized; dir now {sum(f.stat().st_size for f in out.rglob('*.onnx'))/1e6:.1f} MB of .onnx", flush=True)
|
||
|
||
# sanity: load via ShadowNLI's ONNX path and run one pair
|
||
import os
|
||
os.environ["ARBORIST_NLI_ONNX_DIR"] = str(out)
|
||
n = ShadowNLI(manifest)
|
||
n._ensure_loaded()
|
||
if not n.available or not (n.backend or "").startswith("onnx"):
|
||
print(f"[export-nli-onnx] WARNING: ShadowNLI did not pick up the ONNX export "
|
||
f"(available={n.available}, backend={n.backend}, reason={n._reason})", file=sys.stderr)
|
||
return 1
|
||
pe, pn, pc = n._nli("Jupiter is the largest planet.", "Mercury is the largest planet.")
|
||
print(f"[export-nli-onnx] sanity OK — backend={n.backend} device={n.device}; "
|
||
f"NLI(Jupiter-is-largest, Mercury-is-largest) → contradiction={pc:.3f} entail={pe:.3f}")
|
||
print(f"[export-nli-onnx] done. ShadowNLI will now auto-prefer {out} (or set ARBORIST_NLI_ONNX_DIR).")
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(main())
|