knowledge/embedder.py

"""Multilingual embeddings via fastembed (ONNX-based, no torch dependency).

We use `sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2`:
  • 120MB ONNX model — small enough for the free HF Space
  • 384-dim output, well-supported in fastembed
  • Covers ~50 languages including French/Spanish/Arabic → close enough
    to Mauritian Kreol for retrieval to work (Creole shares heavy
    French-derived vocabulary)
  • Comparable retrieval quality to e5-small at similar size

To see fastembed's full supported-model list:
    from fastembed import TextEmbedding
    TextEmbedding.list_supported_models()
"""

from __future__ import annotations

import os
from functools import lru_cache
from typing import Iterable

import numpy as np

# Already MULTILINGUAL by default (~50 languages incl. French → covers most of
# Kreol's French-derived vocabulary). Swap to a stronger multilingual model
# without code changes, e.g.:
#     EMBED_MODEL=intfloat/multilingual-e5-large   (1024-d, 100 languages, heavier)
#     EMBED_MODEL=intfloat/multilingual-e5-small   (384-d, 100 languages)
# After changing it, delete data/knowledge.sqlite so the base re-embeds.
DEFAULT_MODEL = os.environ.get(
    "EMBED_MODEL", "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")
EMBED_DIM = int(os.environ.get("EMBED_DIM", "384"))


def _is_e5(name: str) -> bool:
    return "e5" in (name or "").lower()


@lru_cache(maxsize=2)
def _model(model_name: str):
    from fastembed import TextEmbedding
    print(f"[knowledge] Loading embedder: {model_name}")
    return TextEmbedding(model_name=model_name)


def embed_texts(
    texts: Iterable[str], model_name: str = DEFAULT_MODEL
) -> np.ndarray:
    """Embed a list of strings → (N, dim) float32, L2-normalised."""
    texts = list(texts)
    if not texts:
        return np.zeros((0, EMBED_DIM), dtype=np.float32)
    model = _model(model_name)
    embeddings = list(model.embed(texts))
    arr = np.array(embeddings, dtype=np.float32)
    norms = np.linalg.norm(arr, axis=1, keepdims=True)
    norms[norms == 0] = 1.0
    return arr / norms


def embed_passages(texts: Iterable[str]) -> np.ndarray:
    """Embed text chunks for storage (e5 models need the 'passage:' prefix)."""
    texts = list(texts)
    if _is_e5(DEFAULT_MODEL):
        texts = [f"passage: {t}" for t in texts]
    return embed_texts(texts)


def embed_query(text: str) -> np.ndarray:
    """Embed a search query (e5 models need the 'query:' prefix)."""
    if _is_e5(DEFAULT_MODEL):
        text = f"query: {text}"
    return embed_texts([text])[0]