precompute.py

"""
Precompute corpus embeddings and store in ChromaDB.

Run once locally:
    python precompute.py

Produces:  data/chroma_db/  (persistent ChromaDB directory)
"""

import os
import chromadb
import pandas as pd
from datasets import load_dataset
from sentence_transformers import SentenceTransformer

MODEL_NAME = "intfloat/multilingual-e5-small"
CHROMA_PATH = os.path.join(os.path.dirname(__file__), "data", "chroma_db")
COLLECTION_NAME = "scifact"
BATCH_SIZE = 64


def load_scifact() -> pd.DataFrame:
    corpus = load_dataset("mteb/scifact", "corpus", split="corpus").to_pandas()
    corpus["title"] = corpus["title"].fillna("").astype(str)
    corpus["text"] = corpus["text"].fillna("").astype(str)
    corpus["full_text"] = (
        corpus["title"].str.strip() + ". " + corpus["text"].str.strip()
    ).str.strip(" .")
    corpus = corpus.rename(columns={"_id": "doc_id"})[
        ["doc_id", "title", "text", "full_text"]
    ]
    return corpus.reset_index(drop=True)


def main():
    print("Loading SciFact corpus...")
    corpus_df = load_scifact()
    print(f"  {len(corpus_df)} documents loaded.")

    print(f"Loading model: {MODEL_NAME}")
    model = SentenceTransformer(MODEL_NAME)

    passages = [f"passage: {t.strip()}" for t in corpus_df["full_text"].tolist()]
    print(f"Encoding {len(passages)} passages...")
    embeddings = model.encode(
        passages,
        batch_size=BATCH_SIZE,
        show_progress_bar=True,
        normalize_embeddings=True,
    )
    print(f"  Embedding shape: {embeddings.shape}")

    os.makedirs(CHROMA_PATH, exist_ok=True)
    client = chromadb.PersistentClient(path=CHROMA_PATH)

    # Delete existing collection if present
    try:
        client.delete_collection(COLLECTION_NAME)
    except Exception:
        pass

    collection = client.create_collection(
        name=COLLECTION_NAME,
        metadata={"hnsw:space": "cosine"},
    )

    ids = [str(i) for i in range(len(corpus_df))]
    titles = corpus_df["title"].tolist()
    texts = [t[:500] for t in corpus_df["text"].tolist()]
    emb_lists = embeddings.tolist()

    # Add in batches
    ADD_BATCH = 500
    for start in range(0, len(ids), ADD_BATCH):
        end = min(start + ADD_BATCH, len(ids))
        collection.add(
            ids=ids[start:end],
            embeddings=emb_lists[start:end],
            metadatas=[
                {"title": titles[i], "text": texts[i]} for i in range(start, end)
            ],
            documents=corpus_df["full_text"].tolist()[start:end],
        )
        print(f"  Added {end}/{len(ids)} documents to ChromaDB.")

    print(f"\nChromaDB persisted to: {CHROMA_PATH}")
    print(f"Collection '{COLLECTION_NAME}': {collection.count()} documents")


if __name__ == "__main__":
    main()