144 lines
7.5 KiB
Python
144 lines
7.5 KiB
Python
"""Owned CPU-only cross-encoder compatibility/cost diagnostic; never alters search."""
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import math
|
|
import os
|
|
from pathlib import Path
|
|
import platform
|
|
import statistics
|
|
import tempfile
|
|
import time
|
|
|
|
import numpy as np
|
|
import onnxruntime as ort
|
|
from tokenizers import Tokenizer
|
|
from download import DESTINATION, digest
|
|
|
|
EXPECTED_PLAN = "174867d4212b222e5302b1b16832b371c6bd38ec991f4145c952ed35446341f1"
|
|
EXPECTED_MANIFEST = "3c252dfc10e03ebba59a7ff0e33e6e2183a26c71862c52f0700f27ba659142be"
|
|
|
|
|
|
def checked_json(path, expected):
|
|
data = path.read_bytes()
|
|
if hashlib.sha256(data).hexdigest() != expected:
|
|
raise ValueError("Frozen diagnostic inputs changed")
|
|
return json.loads(data)
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--output", type=Path, required=True)
|
|
args = parser.parse_args()
|
|
owned_root = DESTINATION.parents[1].resolve()
|
|
output = args.output.resolve()
|
|
if not output.is_relative_to(owned_root) or output.exists():
|
|
raise ValueError("Output must be new and under owned .searchi data")
|
|
if os.environ.get("SEARCHI_FORCE_CPU") != "1":
|
|
raise ValueError("Windows VM diagnostic requires explicit SEARCHI_FORCE_CPU=1")
|
|
plan = checked_json(Path(__file__).with_name("reference-plan.json"), EXPECTED_PLAN)
|
|
manifest = checked_json(Path(__file__).with_name("reference-manifest.json"), EXPECTED_MANIFEST)
|
|
start = time.perf_counter()
|
|
for item in manifest["files"]:
|
|
path = DESTINATION / item["path"]
|
|
if path.stat().st_size != item["bytes"] or digest(path) != item["sha256"]:
|
|
raise ValueError("Pinned diagnostic asset mismatch")
|
|
verification_ms = (time.perf_counter() - start) * 1000
|
|
start = time.perf_counter()
|
|
tokenizer = Tokenizer.from_file(str(DESTINATION / "tokenizer.json"))
|
|
special = tokenizer.num_special_tokens_to_add(True)
|
|
encodings = []
|
|
for query, passage in plan["pairs"]:
|
|
query_encoding = tokenizer.encode(query, add_special_tokens=False)
|
|
query_encoding.truncate(plan["max_query_tokens"])
|
|
passage_encoding = tokenizer.encode(passage, add_special_tokens=False)
|
|
passage_encoding.truncate(plan["max_pair_tokens"] - len(query_encoding.ids) - special)
|
|
encoding = tokenizer.post_process(query_encoding, passage_encoding, add_special_tokens=True)
|
|
if len(encoding.ids) > plan["max_pair_tokens"]:
|
|
raise ValueError("Pair token budget exceeded")
|
|
encodings.append(encoding)
|
|
tokenization_ms = (time.perf_counter() - start) * 1000
|
|
options = ort.SessionOptions()
|
|
options.intra_op_num_threads = 1
|
|
options.inter_op_num_threads = 1
|
|
options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
|
|
start = time.perf_counter()
|
|
session = ort.InferenceSession(str(DESTINATION / plan["reference_model"]), sess_options=options, providers=["CPUExecutionProvider"])
|
|
startup_ms = (time.perf_counter() - start) * 1000
|
|
if session.get_providers() != ["CPUExecutionProvider"]:
|
|
raise ValueError("Unexpected execution provider")
|
|
names = [item.name for item in session.get_inputs()]
|
|
if not set(names) <= {"input_ids", "attention_mask", "token_type_ids"}:
|
|
raise ValueError("Unexpected model input contract")
|
|
def score(items):
|
|
width = max(len(item.ids) for item in items)
|
|
ids = np.full((len(items), width), 1, dtype=np.int64)
|
|
mask = np.zeros_like(ids)
|
|
types = np.zeros_like(ids)
|
|
for index, item in enumerate(items):
|
|
count = len(item.ids)
|
|
ids[index, :count] = item.ids
|
|
mask[index, :count] = item.attention_mask
|
|
types[index, :count] = item.type_ids
|
|
values = {"input_ids": ids, "attention_mask": mask, "token_type_ids": types}
|
|
logits = session.run(None, {name: values[name] for name in names})[0]
|
|
if logits.shape != (len(items), 1) or not np.isfinite(logits).all():
|
|
raise ValueError("Invalid classifier output")
|
|
return logits[:, 0]
|
|
start = time.perf_counter()
|
|
serial = np.asarray([float(score([item])[0]) for item in encodings])
|
|
first_serial_ms = (time.perf_counter() - start) * 1000
|
|
start = time.perf_counter()
|
|
batch = score(encodings)
|
|
first_batch_ms = (time.perf_counter() - start) * 1000
|
|
parity = bool(np.allclose(serial, batch, atol=plan["parity_absolute_tolerance"], rtol=plan["parity_relative_tolerance"]))
|
|
serial_timings = []
|
|
batch_timings = []
|
|
for _ in range(plan["warm_iterations"]):
|
|
start = time.perf_counter()
|
|
for item in encodings:
|
|
score([item])
|
|
serial_timings.append((time.perf_counter() - start) * 1000)
|
|
start = time.perf_counter()
|
|
score(encodings)
|
|
batch_timings.append((time.perf_counter() - start) * 1000)
|
|
reference_session = session
|
|
session = ort.InferenceSession(str(DESTINATION / plan["comparison_model"]), sess_options=options, providers=["CPUExecutionProvider"])
|
|
if session.get_providers() != ["CPUExecutionProvider"]:
|
|
raise ValueError("Unexpected comparison provider")
|
|
quantized_serial = np.asarray([float(score([item])[0]) for item in encodings])
|
|
quantized_batch = score(encodings)
|
|
session = reference_session
|
|
def timing(values):
|
|
return {"median_ms": statistics.median(values), "p95_ms": sorted(values)[math.ceil(len(values) * .95) - 1], "samples_ms": values}
|
|
report = {"complete": True, "plan_sha256": EXPECTED_PLAN, "manifest_sha256": EXPECTED_MANIFEST,
|
|
"platform": platform.platform(), "python": platform.python_version(), "onnxruntime": ort.__version__,
|
|
"numpy": np.__version__, "providers": session.get_providers(), "cpu_threads": 1,
|
|
"asset_verification_ms": verification_ms, "tokenization_ms": tokenization_ms, "session_startup_ms": startup_ms,
|
|
"pair_tokens": [len(item.ids) for item in encodings], "first_serial_ms": first_serial_ms, "first_batch_ms": first_batch_ms,
|
|
"reference_model": plan["reference_model"], "comparison_model": plan["comparison_model"],
|
|
"serial_logits": serial.tolist(), "batch_logits": batch.tolist(), "parity_passed": parity,
|
|
"quantized_serial_logits": quantized_serial.tolist(), "quantized_batch_logits": quantized_batch.tolist(),
|
|
"quantized_vs_reference_serial_max_abs": float(np.max(np.abs(serial-quantized_serial))),
|
|
"quantized_batch_max_abs": float(np.max(np.abs(quantized_serial-quantized_batch))),
|
|
"five_short_pairs_ordering_agrees": bool(np.array_equal(np.argsort(serial[:5]),np.argsort(quantized_serial[:5]))),
|
|
"max_absolute_difference": float(np.max(np.abs(serial - batch))),
|
|
"warm_serial": timing(serial_timings), "warm_batch": timing(batch_timings),
|
|
"limitations": "Five synthetic short pairs plus one 256-token cost pair; classifier compatibility/cost only. No search pipeline or absence calibration. No confidence probability or suppression."}
|
|
output.parent.mkdir(parents=True, exist_ok=True)
|
|
with tempfile.NamedTemporaryFile(mode="w", encoding="utf-8", dir=output.parent, delete=False) as handle:
|
|
temporary = Path(handle.name)
|
|
json.dump(report, handle, indent=2)
|
|
handle.write("\n")
|
|
try:
|
|
os.link(temporary, output)
|
|
finally:
|
|
temporary.unlink(missing_ok=True)
|
|
print(json.dumps({"output": str(output), "parity_passed": parity, "max_absolute_difference": report["max_absolute_difference"], "warm_serial": report["warm_serial"]["p95_ms"], "warm_batch": report["warm_batch"]["p95_ms"]}), flush=True)
|
|
if not parity:
|
|
raise SystemExit("Serial/batch parity failed; do not deploy")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|