backend/scripts/precompute_embeddings.py

#!/usr/bin/env python3
"""
precompute_embeddings.py — build demo_data/embeddings.npz.

Collects every item-name string the matcher could feed to the semantic step
(applying the matcher's own _norm, plus the "UNKNOWN" blank-fallback) from every
served demo document across all five models, encodes them once with
all-MiniLM-L6-v2 (normalize_embeddings=True), and saves an .npz the runtime
loads as a torch-free cache. Run offline with system python3 (needs
sentence-transformers in the user site-packages).
"""
import glob
import json
import os
import sys

import numpy as np

HERE = os.path.dirname(os.path.abspath(__file__))
BACKEND = os.path.dirname(HERE)
TX_DIR = os.path.join(BACKEND, "demo_data", "transactions")
OUT = os.path.join(BACKEND, "demo_data", "embeddings.npz")

# Use the vendored matcher's exact normaliser so cached keys match runtime lookups.
sys.path.insert(0, BACKEND)
from vendor import matcher as M  # noqa: E402


def collect_names():
    names = set()
    for f in glob.glob(os.path.join(TX_DIR, "**", "*.json"), recursive=True):
        doc = json.load(open(f, encoding="utf-8"))
        for it in (doc.get("line_items") or []):
            norm = M._norm(it.get("item_name") or "")
            names.add(norm if norm else "UNKNOWN")
    names.add("UNKNOWN")
    return sorted(n for n in names if n)


def main():
    names = collect_names()
    if not names:
        print("No item names found — did you run build_demo_data.py first?")
        sys.exit(1)
    from sentence_transformers import SentenceTransformer
    model = SentenceTransformer("all-MiniLM-L6-v2")
    vectors = model.encode(names, normalize_embeddings=True).astype("float32")
    np.savez_compressed(OUT, names=np.array(names, dtype=object), vectors=vectors)
    print(f"Wrote {len(names)} embeddings (dim={vectors.shape[1]}) to {OUT}")
    print("Sample names:", names[:5])


if __name__ == "__main__":
    main()