models/embedder.py

"""
Lazy singleton around the multilingual sentence-embedding model.

Design constraints:
  - The torch model must never be pickled into classifier artifacts —
    classifiers store only vectors and this module reloads the model by name.
  - If the model can't load (download failure, EMBEDDINGS_ENABLED=0, low
    memory), encode() returns None and classifiers degrade to TF-IDF-only.
  - Encoding is cached by exact text so retrains re-encode only new rows.
"""
from __future__ import annotations

import threading
from typing import Optional

import numpy as np

import config
from config import get_logger

log = get_logger("embedder")

_lock = threading.Lock()
_model = None
_model_failed = False
_cache: dict[str, np.ndarray] = {}
_CACHE_MAX = 50_000


def _get_model():
    global _model, _model_failed
    if not config.EMBEDDINGS_ENABLED or _model_failed:
        return None
    if _model is not None:
        return _model
    with _lock:
        if _model is not None or _model_failed:
            return _model
        try:
            from sentence_transformers import SentenceTransformer
            log.info("loading embedding model %s (first call may download)", config.EMBEDDING_MODEL)
            _model = SentenceTransformer(config.EMBEDDING_MODEL, device="cpu")
            log.info("embedding model ready")
        except Exception:
            log.exception("embedding model unavailable — classifiers will run TF-IDF-only")
            _model_failed = True
    return _model


def available() -> bool:
    return _get_model() is not None


def encode(texts: list[str]) -> Optional[np.ndarray]:
    """L2-normalized embeddings for texts, or None if the model is unavailable."""
    model = _get_model()
    if model is None:
        return None
    todo = [t for t in dict.fromkeys(texts) if t not in _cache]
    if todo:
        vecs = model.encode(todo, batch_size=32, show_progress_bar=False,
                            normalize_embeddings=True)
        if len(_cache) + len(todo) > _CACHE_MAX:
            _cache.clear()
        _cache.update(dict(zip(todo, vecs)))
    return np.vstack([_cache[t] for t in texts])