Blame for metadata_search/embedder.py

"""Sentence-transformer embeddings for metadata vector search."""

from __future__ import annotations

from functools import lru_cache

from metadata_search.config import EMBEDDING_MODEL


@lru_cache(maxsize=1)
def _load_model(model_name: str):
    try:
        from sentence_transformers import SentenceTransformer
    except ImportError as exc:
        raise SystemExit(
            "Missing dependency: sentence-transformers. Install with `pip install -r requirements.txt`."
        ) from exc

    try:
        return SentenceTransformer(model_name, local_files_only=True)
    except Exception:
        return SentenceTransformer(model_name)


class MetadataEmbedder:
    def __init__(self, model_name: str = EMBEDDING_MODEL, *, batch_size: int = 64):
        self.model_name = model_name
        self.batch_size = batch_size
        self.model = _load_model(model_name)

    def embed_documents(self, texts: list[str]) -> list[list[float]]:
        if not texts:
            return []
        vectors = self.model.encode(
            texts,
            batch_size=self.batch_size,
            normalize_embeddings=True,
            show_progress_bar=len(texts) > 128,
        )
        return vectors.tolist()

    def embed_query(self, query: str) -> list[float]:
        vector = self.model.encode(
            query,
            normalize_embeddings=True,
            show_progress_bar=False,
        )
        return vector.tolist()