embeddings.py

"""Encode text chunks into dense vectors using SentenceTransformers."""

import numpy as np
from sentence_transformers import SentenceTransformer

from config import EMBEDDING_MODEL

_model = None  # lazy singleton


def get_embedding_model() -> SentenceTransformer:
    global _model
    if _model is None:
        print(f"[embeddings] Loading embedding model: {EMBEDDING_MODEL}")
        _model = SentenceTransformer(EMBEDDING_MODEL)
    return _model


def embed_texts(texts: list[str]) -> np.ndarray:
    """Return a (N, D) float32 array of embeddings."""
    model = get_embedding_model()
    embeddings = model.encode(texts, show_progress_bar=True, convert_to_numpy=True)
    return embeddings.astype("float32")


def embed_query(query: str) -> np.ndarray:
    """Return a (1, D) float32 array for a single query."""
    model = get_embedding_model()
    vec = model.encode([query], convert_to_numpy=True)
    return vec.astype("float32")