scripts/eval/build_local_index.py

"""Rebuild the FAISS index locally from the committed chunks.jsonl.

The committed ``data/processed/index/index.faiss`` is a git-LFS pointer (not
pulled here), so retrieval cannot run. ``chunks.jsonl`` IS real, so we re-embed
those exact chunk texts with the pinned Nemotron embedder and the same
``passage:`` prefix / L2-normalization the runtime uses, then write a usable
``index.faiss`` next to it.

This reproduces the runtime retrieval path faithfully:
  embeddings.repo / revision / query_prefix / passage_prefix / normalize
all come from configs/models.yaml.
"""
from __future__ import annotations

import json
import sys
import time
from pathlib import Path

import numpy as np

REPO_ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(REPO_ROOT / "src"))

import yaml  # noqa: E402

CHUNKS = REPO_ROOT / "data" / "processed" / "index" / "chunks.jsonl"
INDEX_OUT = REPO_ROOT / "data" / "processed" / "index" / "index.faiss"


def main() -> None:
    cfg = yaml.safe_load((REPO_ROOT / "configs" / "models.yaml").read_text())["embeddings"]
    chunks = [json.loads(l) for l in CHUNKS.read_text(encoding="utf-8").splitlines() if l.strip()]
    print(f"loaded {len(chunks)} chunks")

    from sentence_transformers import SentenceTransformer

    t = time.time()
    model = SentenceTransformer(
        cfg["repo"], revision=cfg.get("revision"), device="cpu", trust_remote_code=True
    )
    print(f"model loaded in {time.time()-t:.1f}s")

    prefix = cfg.get("passage_prefix", "")
    texts = [f"{prefix}{c['text']}" for c in chunks]
    t = time.time()
    vecs = model.encode(
        texts, batch_size=16, normalize_embeddings=cfg.get("normalize", True),
        convert_to_numpy=True, show_progress_bar=True,
    ).astype("float32")
    print(f"embedded {vecs.shape} in {time.time()-t:.1f}s")

    import faiss

    index = faiss.IndexFlatIP(vecs.shape[1])
    index.add(vecs)
    faiss.write_index(index, str(INDEX_OUT))
    print(f"wrote index: {index.ntotal} vectors, dim {vecs.shape[1]} -> {INDEX_OUT}")


if __name__ == "__main__":
    main()