src/embed.py

"""Local sentence-transformers embeddings (free, no API calls)."""
from __future__ import annotations

import os
from functools import lru_cache

from .config import CONFIG

# Quiet the HF/transformers/tokenizers progress bars and warnings so CLI and Streamlit
# output stays readable. Set before the libraries are imported.
os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1")
os.environ.setdefault("TRANSFORMERS_VERBOSITY", "error")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")


@lru_cache(maxsize=1)
def _model():
    from sentence_transformers import SentenceTransformer

    return SentenceTransformer(CONFIG.embed_model)


def embed_texts(texts: list[str]) -> list[list[float]]:
    model = _model()
    vecs = model.encode(texts, normalize_embeddings=True, show_progress_bar=len(texts) > 64)
    return [v.tolist() for v in vecs]


def embed_query(text: str) -> list[float]:
    return embed_texts([text])[0]