156 lines
5.4 KiB
Python
156 lines
5.4 KiB
Python
"""Hybrid dense + 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.
|
|
|
|
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
|