src/embedder.py

import numpy as np
from sentence_transformers import SentenceTransformer

from src.config import Config


class Embedder:
    def __init__(self, config: Config):
        self.config = config
        self._model: SentenceTransformer | None = None
        self._failed = False

    def _get_model(self) -> SentenceTransformer | None:
        if self._failed:
            return None
        if self._model is None:
            try:
                self._model = SentenceTransformer(
                    self.config.embedding_model,
                    trust_remote_code=True,
                )
            except Exception:
                try:
                    self._model = SentenceTransformer(
                        self.config.embedding_model_fallback,
                        trust_remote_code=True,
                    )
                except Exception:
                    self._failed = True
                    return None
        return self._model

    def encode(self, texts: list[str], batch_size: int = 32) -> np.ndarray | None:
        model = self._get_model()
        if model is None:
            return None
        embeddings = model.encode(
            texts,
            batch_size=batch_size,
            show_progress_bar=True,
            normalize_embeddings=True,
        )
        return np.array(embeddings, dtype=np.float32)

    def encode_query(self, query: str) -> np.ndarray | None:
        model = self._get_model()
        if model is None:
            return None
        embedding = model.encode(
            query,
            normalize_embeddings=True,
        )
        return np.array([embedding], dtype=np.float32)