scripts/build_embeddings.py

"""Offline asset build: BioLORD-embed the Hallmark descriptions → committed assets.

Outputs (committed, ~150 KB total — the only persistent runtime data):
  assets/pathway_descriptions.json   name -> short description
  assets/pathway_embeddings.npz      names: str (N,) | vectors: float32 (N, 768), L2-normalized

Run once (downloads BioLORD-2023 ~420 MB to the HF cache on first use; CPU; ~1 min):
    python scripts/build_embeddings.py
    python scripts/build_embeddings.py --model BAAI/bge-base-en-v1.5   # contrast baseline

100% free / CPU / no paid services. Network is used only to download the open model; the
resulting assets are fully offline at runtime.
"""
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

import numpy as np

# Make the repo root importable when run as `python scripts/build_embeddings.py`.
ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT))
sys.path.insert(0, str(ROOT / "scripts"))

import config  # noqa: E402
from _hallmark_descriptions import HALLMARK_DESCRIPTIONS  # noqa: E402


def build(model_id: str = config.EMBEDDER_DEFAULT) -> None:
    descriptions: dict[str, str] = dict(HALLMARK_DESCRIPTIONS)
    config.ASSETS_DIR.mkdir(parents=True, exist_ok=True)

    # 1. descriptions.json
    config.DESCRIPTIONS_JSON.write_text(json.dumps(descriptions, indent=2, sort_keys=True))
    print(f"[build] wrote {config.DESCRIPTIONS_JSON.name} ({len(descriptions)} pathways)")

    # 2. embeddings.npz
    from sentence_transformers import SentenceTransformer

    names = list(descriptions.keys())
    texts = [descriptions[n] for n in names]
    print(f"[build] loading {model_id} on {config.EMBED_DEVICE} (first run downloads the model)…")
    model = SentenceTransformer(model_id, device=config.EMBED_DEVICE)
    vectors = model.encode(
        texts, normalize_embeddings=True, batch_size=32, show_progress_bar=True
    ).astype(np.float32)

    # 3. verify invariants before committing
    assert vectors.shape == (len(names), config.EMBED_DIM), vectors.shape
    norms = np.linalg.norm(vectors, axis=1)
    assert np.allclose(norms, 1.0, atol=1e-4), f"vectors not L2-normalized: {norms.min()}..{norms.max()}"
    assert not np.isnan(vectors).any(), "NaN in embeddings"

    np.savez_compressed(
        config.EMBEDDINGS_NPZ,
        names=np.array(names, dtype=str),
        vectors=vectors,
    )
    size_kb = config.EMBEDDINGS_NPZ.stat().st_size / 1024
    print(f"[build] wrote {config.EMBEDDINGS_NPZ.name} "
          f"({len(names)} x {config.EMBED_DIM}, {size_kb:.0f} KB, model={model_id})")


if __name__ == "__main__":
    ap = argparse.ArgumentParser(description="Build pathway description + embedding assets.")
    ap.add_argument("--model", default=config.EMBEDDER_DEFAULT, help="sentence-transformers model id")
    build(ap.parse_args().model)