codex-py/codex/embed.py

"""Dense (+ optional sparse) embeddings via BGE-M3 (ADR-0002).

Uses :class:`FlagEmbedding.BGEM3FlagModel` for encoding because the
``return_dense`` / ``return_sparse`` kwargs are part of the FlagEmbedding
API — not ``sentence-transformers``.  The dense output is identical to a
vanilla ``SentenceTransformer`` load of ``BAAI/bge-m3`` (same weights,
same model); sparse output is a list of ``{token_id: weight}`` dicts per
text.

Status: only the **dense** path is wired into ingest/search today; the search
"hybrid" is dense + Postgres FTS. :meth:`Embedder.encode_sparse` / :meth:`encode`
are provided for a future sparse-retrieval layer but are not yet consumed by the
pipeline (audit C-14).

Notes for callers
-----------------
* Empty input is handled explicitly — no model call is issued.
* Dense vectors are L2-normalised row-wise so cosine similarity reduces
  to a dot product.
* :func:`get_embedder` returns a process-wide singleton so the model is
  loaded at most once per Python process.
"""

from __future__ import annotations

import numpy as np
import torch
from FlagEmbedding import BGEM3FlagModel


class Embedder:
    """Thin wrapper over :class:`BGEM3FlagModel` with a stable API.

    The wrapper hides the FlagEmbedding-specific encode-dict and exposes
    three operations relevant to the ingestion pipeline:

    * :meth:`encode_dense` — dense float32 matrix, L2-normalised.
    * :meth:`encode_sparse` — list of ``{token_id: weight}`` dicts.
    * :meth:`encode` — both in a single forward pass.
    """

    def __init__(
        self,
        model_name: str = "BAAI/bge-m3",
        dim: int = 1024,
        device: str | None = None,
        batch_size: int = 32,
    ) -> None:
        resolved_device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        self._model = BGEM3FlagModel(
            model_name,
            use_fp16=False,
            devices=[resolved_device],
        )
        self.dim = dim
        self.batch_size = batch_size

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------
    @staticmethod
    def _l2_normalise(matrix: np.ndarray) -> np.ndarray:
        """Return ``matrix`` with every row scaled to unit L2 norm.

        Zero-rows are left untouched (we divide by 1 rather than 0 to
        avoid ``nan``s — the BGE-M3 encoder never emits zero vectors in
        practice, but the guard keeps the function total).
        """
        norms = np.linalg.norm(matrix, axis=1, keepdims=True)
        norms = np.where(norms == 0.0, 1.0, norms)
        return (matrix / norms).astype(np.float32)

    @staticmethod
    def _coerce_sparse(weights: list[dict[int | str, float]]) -> list[dict[int, float]]:
        """Cast token ids to ``int`` and weights to ``float``.

        FlagEmbedding returns ids as strings in some versions and as
        ints in others; we normalise to ``int`` so downstream code can
        rely on a stable key type.
        """
        return [
            {int(token_id): float(weight) for token_id, weight in row.items()} for row in weights
        ]

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------
    def encode_dense(self, texts: list[str]) -> np.ndarray:
        """Encode ``texts`` to a dense, L2-normalised float32 matrix.

        Returns a ``(0, dim)`` zero-matrix for an empty input.
        """
        if not texts:
            return np.zeros((0, self.dim), dtype=np.float32)

        result = self._model.encode(
            texts,
            batch_size=self.batch_size,
            return_dense=True,
            return_sparse=False,
        )
        dense = np.asarray(result["dense_vecs"], dtype=np.float32)
        return self._l2_normalise(dense)

    def encode_sparse(self, texts: list[str]) -> list[dict[int, float]]:
        """Encode ``texts`` to a list of sparse ``{token_id: weight}`` dicts.

        Returns ``[]`` for an empty input.
        """
        if not texts:
            return []

        result = self._model.encode(
            texts,
            batch_size=self.batch_size,
            return_dense=False,
            return_sparse=True,
        )
        return self._coerce_sparse(result["lexical_weights"])

    def encode(self, texts: list[str]) -> tuple[np.ndarray, list[dict[int, float]]]:
        """Encode ``texts`` to dense **and** sparse in one forward pass.

        Returns ``(zeros((0, dim)), [])`` for an empty input.
        """
        if not texts:
            return np.zeros((0, self.dim), dtype=np.float32), []

        result = self._model.encode(
            texts,
            batch_size=self.batch_size,
            return_dense=True,
            return_sparse=True,
        )
        dense = self._l2_normalise(np.asarray(result["dense_vecs"], dtype=np.float32))
        sparse = self._coerce_sparse(result["lexical_weights"])
        return dense, sparse


# ---------------------------------------------------------------------------
# Process-wide singleton
# ---------------------------------------------------------------------------
_embedder: Embedder | None = None


def get_embedder() -> Embedder:
    """Return the process-wide :class:`Embedder` singleton.

    The first call constructs the embedder using values from
    :class:`codex.config.Settings`; subsequent calls return the cached
    instance.  Tests can reset the cache by setting
    ``codex.embed._embedder`` back to ``None``.
    """
    global _embedder
    if _embedder is None:
        from codex.config import Settings

        s = Settings()
        _embedder = Embedder(model_name=s.embedding_model, dim=s.embedding_dim)
    return _embedder