#!/usr/bin/env python3 """LockedIn FDE Training: inspect real local sentence embeddings. All records are fictional. `prepare` downloads a pinned, Apache-2.0 model. Other commands block Python network connections and run CPU inference locally. No chat model, vector database, API key or production identity is involved. """ from __future__ import annotations import argparse import hashlib import json import math import os from pathlib import Path import socket import sys import tempfile import unittest from urllib.request import urlopen from urllib.error import URLError from unittest.mock import patch MODEL = "sentence-transformers/all-MiniLM-L6-v2" REVISION = "1110a243fdf4706b3f48f1d95db1a4f5529b4d41" FILES = { "onnx/model.onnx": "6fd5d72fe4589f189f8ebc006442dbb529bb7ce38f8082112682524616046452", "tokenizer.json": "be50c3628f2bf5bb5e3a7f17b1f74611b2561a3a27eeab05e5aa30f411572037", } DOCUMENTS = [ {"id": "P1", "tenant": "north", "current": True, "text": "Use the account recovery page to choose a new password."}, {"id": "P2", "tenant": "north", "current": True, "text": "A locked account needs an administrator to restore access."}, {"id": "P3", "tenant": "north", "current": True, "text": "Export invoices from the billing dashboard."}, {"id": "P4", "tenant": "north", "current": True, "text": "Do not use the account recovery page to choose a new password."}, {"id": "P5", "tenant": "north", "current": False, "text": "How do I reset my password?"}, {"id": "P9", "tenant": "south", "current": True, "text": "How do I reset my password?"}, ] QUERIES = [ {"id": "exact", "label": "Same words", "text": DOCUMENTS[0]["text"]}, {"id": "paraphrase", "label": "Different wording", "text": "How do I reset my password?"}, {"id": "negation", "label": "Opposite instruction", "text": DOCUMENTS[3]["text"]}, {"id": "unrelated", "label": "Outside the collection", "text": "What is the weather forecast for tomorrow?"}, ] class SetupNeeded(ValueError): pass def cache_dir() -> Path: return Path(os.environ.get("FDE_EMBEDDING_CACHE", ".embedding-cache")) def verified_file(root: Path, name: str) -> Path: target = root / name if not target.is_file(): raise SetupNeeded("Model file missing or changed. Run: python fde-embeddings-first-search.py prepare") with target.open("rb") as source: if hashlib.file_digest(source, "sha256").hexdigest() != FILES[name]: raise SetupNeeded("Model file missing or changed. Run: python fde-embeddings-first-search.py prepare") return target def prepare() -> None: root = cache_dir() for name, expected in FILES.items(): target = root / name try: verified_file(root, name) continue except SetupNeeded: pass target.parent.mkdir(parents=True, exist_ok=True) # Atomic replacement only after checking the complete pinned download. temporary = None try: with tempfile.NamedTemporaryFile(dir=target.parent, delete=False) as out: temporary = Path(out.name) with urlopen(f"https://huggingface.co/{MODEL}/resolve/{REVISION}/{name}", timeout=60) as source: while chunk := source.read(1024 * 1024): out.write(chunk) with temporary.open("rb") as downloaded: if hashlib.file_digest(downloaded, "sha256").hexdigest() != expected: raise SetupNeeded(f"Checksum mismatch: {name}. No model file installed.") temporary.replace(target) finally: if temporary: temporary.unlink(missing_ok=True) print("Pinned model and tokenizer ready (about 91 MB). Other commands run locally.") def no_network(*_args, **_kwargs): raise RuntimeError("Network disabled during this exercise. Use prepare for model setup.") class Encoder: def __init__(self): import numpy as np import onnxruntime as ort from tokenizers import Tokenizer self.np = np self.tokenizer = Tokenizer.from_file(str(verified_file(cache_dir(), "tokenizer.json"))) # The model's sentence configuration supports 256 word pieces. # Reject longer text instead of silently dropping part of a procedure. self.tokenizer.no_truncation() self.tokenizer.enable_padding(pad_id=0, pad_token="[PAD]") options = ort.SessionOptions() options.intra_op_num_threads = 1 self.session = ort.InferenceSession(str(verified_file(cache_dir(), "onnx/model.onnx")), sess_options=options, providers=["CPUExecutionProvider"]) def encode(self, texts: list[str]) -> list[list[float]]: if not texts or any(not isinstance(text, str) or not text.strip() for text in texts): raise ValueError("Supply non-empty text.") tokens = self.tokenizer.encode_batch(texts) if any(sum(item.attention_mask) > 256 for item in tokens): raise ValueError("Text exceeds 256 word pieces. Split it deliberately; nothing was truncated.") np = self.np inputs = { "input_ids": np.asarray([t.ids for t in tokens], dtype=np.int64), "attention_mask": np.asarray([t.attention_mask for t in tokens], dtype=np.int64), "token_type_ids": np.asarray([t.type_ids for t in tokens], dtype=np.int64), } hidden = self.session.run(["last_hidden_state"], inputs)[0] mask = inputs["attention_mask"][..., None] pooled = (hidden * mask).sum(axis=1) / mask.sum(axis=1) normalized = pooled / np.linalg.norm(pooled, axis=1, keepdims=True) return normalized.tolist() def cosine(a: list[float], b: list[float]) -> float: if not a or len(a) != len(b) or any(not math.isfinite(n) for n in [*a, *b]): raise ValueError("Vectors must be finite, non-empty and have matching dimensions.") norm = math.sqrt(sum(n * n for n in a) * sum(n * n for n in b)) if norm == 0: raise ValueError("Cosine is undefined for a zero vector.") return max(-1.0, min(1.0, sum(x * y for x, y in zip(a, b)) / norm)) def eligible(doc: dict, tenant: str = "north") -> bool: # Fixed trusted fixture scope, never obtained from a query or model argument. return doc["tenant"] == tenant and doc["current"] def record(encoder: Encoder, full: bool = False, query: str | None = None) -> dict: allowed = [doc for doc in DOCUMENTS if eligible(doc)] queries = QUERIES if query is None else [{"id": "custom", "label": "Your local query", "text": query}] vectors = encoder.encode([doc["text"] for doc in allowed] + [q["text"] for q in queries]) documents = [{**doc, "preview": [round(n, 6) for n in vector[:12]]} for doc, vector in zip(allowed, vectors)] cases = [] for i, q in enumerate(queries): vector = vectors[len(allowed) + i] ranked = sorted( [{"id": doc["id"], "score": round(cosine(vector, dv), 6)} for doc, dv in zip(allowed, vectors)], key=lambda row: (-row["score"], row["id"])) cases.append({**q, "preview": [round(n, 6) for n in vector[:12]], "ranked": ranked, **({"vector": vector} if full else {})}) return { "model": MODEL, "revision": REVISION, "dimensions": 384, "pooling": "attention-masked mean, then L2 normalization", "documents": documents, "cases": cases, "excluded": [{"id": d["id"], "reason": "different organization" if d["tenant"] != "north" else "retired source"} for d in DOCUMENTS if not eligible(d)], "generation_calls": 0, **({"document_vectors": {d["id"]: v for d, v in zip(allowed, vectors)}} if full else {}), } class ExerciseTests(unittest.TestCase): @classmethod def setUpClass(cls): cls.encoder = Encoder() cls.output = record(cls.encoder, full=True) def test_shape_and_normalization(self): for vector in self.output["document_vectors"].values(): self.assertEqual(len(vector), 384) self.assertAlmostEqual(sum(x*x for x in vector), 1.0, places=6) def test_same_text_is_nearest(self): self.assertEqual(self.output["cases"][0]["ranked"][0], {"id": "P1", "score": 1.0}) def test_padding_does_not_change_meaning(self): a = self.encoder.encode(["Reset my password."])[0] b = self.encoder.encode(["Reset my password.", "A somewhat longer sentence used only to exercise padding."])[0] self.assertGreater(cosine(a, b), 0.99999) def test_negation_can_still_be_similar(self): vectors = self.output["document_vectors"] self.assertGreater(cosine(vectors["P1"], vectors["P4"]), 0.8) def test_exclusion_precedes_encoding(self): with patch.object(self.encoder, "encode", wraps=self.encoder.encode) as encode: result = record(self.encoder) self.assertEqual(len(encode.call_args.args[0]), 4 + len(QUERIES)) self.assertEqual({r["id"] for c in result["cases"] for r in c["ranked"]}, {"P1", "P2", "P3", "P4"}) self.assertEqual({r["id"] for r in result["excluded"]}, {"P5", "P9"}) def test_no_automatic_answer_or_confidence(self): self.assertEqual(self.output["generation_calls"], 0) self.assertEqual(len(self.output["cases"][-1]["ranked"]), 4) self.assertNotIn("answer", self.output) self.assertNotIn("confidence", self.output) def test_cosine_limits_and_bad_vectors(self): self.assertEqual(cosine([1, 0], [0, 1]), 0) self.assertEqual(cosine([1, 0], [-1, 0]), -1) for a, b in [([], []), ([0], [1]), ([1], [1, 2]), ([math.nan], [1])]: with self.assertRaises(ValueError): cosine(a, b) def test_empty_and_overlong_input_rejected(self): for texts in [[], [""], ["word " * 300]]: with self.assertRaises(ValueError): self.encoder.encode(texts) def main() -> int: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("command", choices=["prepare", "inspect", "test"]) parser.add_argument("--full", action="store_true", help="include all 384 values") parser.add_argument("--query", help="embed this local text instead of the four reference queries") args = parser.parse_args() try: if args.command == "prepare": prepare() return 0 with patch.object(socket.socket, "connect", no_network), \ patch.object(socket.socket, "connect_ex", no_network), \ patch.object(socket, "getaddrinfo", no_network): if args.command == "test": return 0 if unittest.TextTestRunner(verbosity=2).run(unittest.defaultTestLoader.loadTestsFromTestCase(ExerciseTests)).wasSuccessful() else 1 print(json.dumps(record(Encoder(), args.full, args.query), indent=2)) return 0 except (SetupNeeded, ValueError, ImportError, OSError, URLError) as error: print(f"Setup/input: {error}", file=sys.stderr) return 1 if __name__ == "__main__": raise SystemExit(main())