karl
a9b6592330
Core CI / core (ubuntu-latest) (push) Failing after 1m22s
Core CI / core (macos-latest) (push) Canceled after 0s
Core CI / core (windows-latest) (push) Canceled after 0s
Core CI / desktop (ubuntu-latest) (push) Successful in 3m46s
Core CI / desktop (macos-latest) (push) Canceled after 0s
Core CI / desktop (windows-latest) (push) Canceled after 0s
136 lines
6.3 KiB
Python
136 lines
6.3 KiB
Python
"""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()
|