searchi/tools/rerank-eval/reference.py

"""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()