build_embeddings.py

"""
build_embeddings.py — builds FAISS index from enriched corpus.
No parallelism, sequential GPU encoding with progress bar.
"""
import json
import logging
import sys
from pathlib import Path

import faiss
import numpy as np
import torch
from sentence_transformers import SentenceTransformer

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(message)s")
log = logging.getLogger(__name__)

CORPUS = Path("data/corpus/corpus/protocols_corpus_enriched.jsonl")
FAISS_OUT = Path("data/faiss_index.bin")
EMB_OUT = Path("data/e5_embeddings_cache.npy")
MODEL_NAME = "intfloat/multilingual-e5-large"
BATCH = 16
MAX_DIAG = 500

def main():
    # ── Load corpus ──────────────────────────────────────────────────────
    log.info("Reading corpus: %s", CORPUS)
    passages = []
    with open(CORPUS, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            obj = json.loads(line)
            title   = obj.get("title", "").strip()
            codes   = ", ".join(obj.get("icd_codes", []))
            text    = obj.get("text", "")[:MAX_DIAG]
            passage = f"passage: {title}. МКБ-10: {codes}. {text}"
            passages.append(passage)
    log.info("Loaded %d passages", len(passages))

    # ── Load model (FP16 on GPU) ─────────────────────────────────────────
    device = "cuda" if torch.cuda.is_available() else "cpu"
    log.info("Loading model %s on %s in FP16...", MODEL_NAME, device)
    model = SentenceTransformer(MODEL_NAME, device=device)
    model.half()

    # ── Encode sequentially — NO ThreadPoolExecutor ─────────────────────
    log.info("Encoding %d passages (batch=%d)...", len(passages), BATCH)
    embeddings = model.encode(
        passages,
        batch_size=BATCH,
        normalize_embeddings=True,
        convert_to_numpy=True,
        show_progress_bar=True,
    ).astype(np.float32)
    log.info("Embeddings shape: %s", embeddings.shape)

    # ── Build + save FAISS index ─────────────────────────────────────────
    dim = embeddings.shape[1]
    index = faiss.IndexFlatIP(dim)
    index.add(embeddings)
    faiss.write_index(index, str(FAISS_OUT))
    np.save(str(EMB_OUT), embeddings)
    log.info("FAISS index saved to %s (%d vectors)", FAISS_OUT, index.ntotal)
    log.info("Embeddings cache saved to %s", EMB_OUT)

if __name__ == "__main__":
    main()