embedder.py

"""
embedder.py
Generates dense vector embeddings using sentence-transformers.
Model: all-MiniLM-L6-v2 (384-dim, fast, high-quality for semantic search)
"""

from sentence_transformers import SentenceTransformer
from typing import List, Union
import numpy as np


MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"
EMBEDDING_DIM = 384


class Embedder:
    """Wraps SentenceTransformer to produce normalized float embeddings."""

    def __init__(self, model_name: str = MODEL_NAME):
        print(f"[Embedder] Loading model: {model_name}")
        self.model = SentenceTransformer(model_name)
        self.dimension = EMBEDDING_DIM

    def embed(self, texts: Union[str, List[str]]) -> List[List[float]]:
        """
        Embed one or more texts.
        Returns a list of float lists (one per input text).
        """
        if isinstance(texts, str):
            texts = [texts]
        embeddings = self.model.encode(
            texts,
            normalize_embeddings=True,   # cosine similarity ready
            show_progress_bar=False,
            batch_size=32,
        )
        return embeddings.tolist()

    def embed_single(self, text: str) -> List[float]:
        """Convenience method for a single string."""
        return self.embed(text)[0]