retrieval/embedder.py

from sentence_transformers import SentenceTransformer
import numpy as np

_model: SentenceTransformer | None = None


def get_model() -> SentenceTransformer:
    global _model
    if _model is None:
        # 130MB, 384-dim, fastest accurate model on CPU
        _model = SentenceTransformer("BAAI/bge-small-en-v1.5")
    return _model


def embed_chunks(chunks: list[str]) -> list[list[float]]:
    model = get_model()
    vecs = model.encode(chunks, normalize_embeddings=True, batch_size=32)
    return vecs.tolist()


def embed_query(query: str) -> list[float]:
    model = get_model()
    # BGE needs this prefix for queries
    prefixed = f"Represent this sentence for searching: {query}"
    vec = model.encode(prefixed, normalize_embeddings=True)
    return vec.tolist()