searchi/tools/semantic-proof/proof.py

"""Real, local ONNX text embeddings with an enforced CPU-only session."""
import argparse
import hashlib
import json
import os
from pathlib import Path
import platform
import statistics
import sys
import time

# Fail rather than silently contacting a model service or downloading an asset.
os.environ["HF_HUB_OFFLINE"] = "1"
os.environ["TRANSFORMERS_OFFLINE"] = "1"
os.environ["TOKENIZERS_PARALLELISM"] = "false"


def reject_network(event, args):
    if event in {"socket.connect", "socket.getaddrinfo", "socket.bind"}:
        raise RuntimeError("Semantic proof inference is offline; network access is forbidden")


sys.addaudithook(reject_network)

import numpy as np
import onnxruntime as ort
from tokenizers import Tokenizer

HERE = Path(__file__).resolve().parent


class CpuEmbedder:
    """Bounded reference contract for a future Rust ONNX adapter."""

    def __init__(self, model_dir, threads=1):
        if not 1 <= threads <= 8:
            raise ValueError("threads must be between 1 and 8")
        self.manifest = json.loads((HERE / "model-manifest.json").read_text())
        for spec in self.manifest["files"]:
            path = model_dir / spec["path"]
            if not path.is_file():
                raise ValueError(f"Missing local artifact {path}; explicitly run download.py first")
            with path.open("rb") as source:
                checksum = hashlib.file_digest(source, "sha256").hexdigest()
            if path.stat().st_size != spec["bytes"] or checksum != spec["sha256"]:
                raise ValueError(f"Artifact size/SHA256 mismatch: {path}")
        options = ort.SessionOptions()
        options.intra_op_num_threads = threads
        options.inter_op_num_threads = 1
        options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
        self.session = ort.InferenceSession(
            str(model_dir / "onnx/model_quantized.onnx"),
            sess_options=options,
            providers=["CPUExecutionProvider"],
        )
        if self.session.get_providers() != ["CPUExecutionProvider"]:
            raise RuntimeError("Only CPUExecutionProvider is allowed")
        self.tokenizer = Tokenizer.from_file(str(model_dir / "tokenizer.json"))
        self.tokenizer.enable_truncation(max_length=self.manifest["max_tokens"])
        self.tokenizer.enable_padding(pad_id=1, pad_token="<pad>")

    def encode(self, texts):
        if not texts or len(texts) > 8:
            raise ValueError("Batch size must be between 1 and 8")
        if any(not text.strip() or len(text.encode("utf-8")) > 65536 for text in texts):
            raise ValueError("Text must be nonempty and at most 64 KiB UTF-8")
        encoded = self.tokenizer.encode_batch(texts)
        available = {
            "input_ids": np.asarray([item.ids for item in encoded], dtype=np.int64),
            "attention_mask": np.asarray([item.attention_mask for item in encoded], dtype=np.int64),
            "token_type_ids": np.asarray([item.type_ids for item in encoded], dtype=np.int64),
        }
        inputs = {item.name: available[item.name] for item in self.session.get_inputs()}
        hidden = self.session.run(None, inputs)[0]
        mask = available["attention_mask"][..., None].astype(np.float32)
        pooled = (hidden * mask).sum(axis=1) / np.maximum(mask.sum(axis=1), 1e-9)
        vectors = pooled / np.maximum(np.linalg.norm(pooled, axis=1, keepdims=True), 1e-9)
        if vectors.shape != (len(texts), self.manifest["dimensions"]) or not np.isfinite(vectors).all():
            raise RuntimeError("Invalid model embedding shape or nonfinite values")
        return vectors.astype(np.float32)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model-dir", type=Path, default=HERE / "model")
    parser.add_argument("--fixture", type=Path, default=HERE / "fixture.json")
    parser.add_argument("--output", type=Path)
    parser.add_argument("--threads", type=int, default=1, choices=range(1, 9))
    parser.add_argument("--repeats", type=int, default=10, choices=range(1, 101))
    args = parser.parse_args()
    fixture = json.loads(args.fixture.read_text(encoding="utf-8"))
    started = time.perf_counter()
    embedder = CpuEmbedder(args.model_dir, args.threads)
    load_ms = (time.perf_counter() - started) * 1000
    started = time.perf_counter()
    vectors = embedder.encode([item["text"] for item in fixture["documents"]])
    corpus_ms = (time.perf_counter() - started) * 1000
    rankings = []
    timings = []
    for query in fixture["queries"]:
        vector = embedder.encode([query["text"]])[0]  # Warm this input shape.
        for _ in range(args.repeats):
            started = time.perf_counter()
            vector = embedder.encode([query["text"]])[0]
            scores = vectors @ vector
            order = np.argsort(-scores)
            timings.append((time.perf_counter() - started) * 1000)
        ranked = [{"id": fixture["documents"][int(i)]["id"], "cosine": float(scores[i])} for i in order]
        expected_rank = next(i + 1 for i, result in enumerate(ranked) if result["id"] == query["expected"])
        rankings.append({**query, "expected_rank": expected_rank, "top3": ranked[:3]})
    report = {
        "platform": platform.platform(), "python": platform.python_version(),
        "machine": platform.machine(), "logical_cpus": os.cpu_count(),
        "onnxruntime": ort.__version__, "numpy": np.__version__,
        "providers": embedder.session.get_providers(), "threads": args.threads,
        "model": embedder.manifest, "documents": len(fixture["documents"]),
        "queries": len(rankings), "timed_query_runs": len(timings),
        "artifact_verify_and_load_ms": load_ms, "corpus_embedding_ms": corpus_ms,
        "query_p50_ms": statistics.median(timings),
        "query_p95_ms": float(np.percentile(timings, 95)),
        "top1_accuracy": sum(item["expected_rank"] == 1 for item in rankings) / len(rankings),
        "recall_at_3": sum(item["expected_rank"] <= 3 for item in rankings) / len(rankings),
        "network": "Python socket audit guard active during import, model load, and inference",
        "rankings": rankings,
    }
    serialized = json.dumps(report, ensure_ascii=False, indent=2)
    if args.output:
        args.output.write_text(serialized + "\n", encoding="utf-8")
    print(serialized)
    if report["top1_accuracy"] < 1:
        raise SystemExit("Fixture top-1 quality gate failed")


if __name__ == "__main__":
    main()