91 lines
3.1 KiB
Python
91 lines
3.1 KiB
Python
"""BGE-M3 dense embedder with batching and CPU/GPU support.
|
|
|
|
We prefer FlagEmbedding's ``BGEM3FlagModel`` because it is the canonical
|
|
implementation and supports dense + sparse output. We fall back to
|
|
``sentence-transformers`` for portability.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Sequence
|
|
from functools import lru_cache
|
|
|
|
import numpy as np
|
|
|
|
from app.config import settings
|
|
from app.logging_config import get_logger
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
|
|
class Embedder:
|
|
def __init__(self, model_name: str, device: str, normalize: bool, batch_size: int) -> None:
|
|
self.model_name = model_name
|
|
self.device = device
|
|
self.normalize = normalize
|
|
self.batch_size = batch_size
|
|
self._impl = "flagembedding"
|
|
self._model = None
|
|
self._st_model = None
|
|
self._load()
|
|
|
|
def _load(self) -> None:
|
|
try:
|
|
from FlagEmbedding import BGEM3FlagModel # type: ignore
|
|
use_fp16 = self.device != "cpu"
|
|
self._model = BGEM3FlagModel(self.model_name, use_fp16=use_fp16, devices=self.device)
|
|
self._impl = "flagembedding"
|
|
logger.info("embedder.loaded", impl="flagembedding", model=self.model_name, device=self.device)
|
|
return
|
|
except Exception as exc: # noqa: BLE001
|
|
logger.warning("embedder.flagembedding_failed", error=str(exc))
|
|
|
|
from sentence_transformers import SentenceTransformer
|
|
self._st_model = SentenceTransformer(self.model_name, device=self.device)
|
|
self._impl = "sentence-transformers"
|
|
logger.info("embedder.loaded", impl="sentence-transformers", model=self.model_name, device=self.device)
|
|
|
|
def encode(self, texts: Sequence[str]) -> list[list[float]]:
|
|
if not texts:
|
|
return []
|
|
if self._impl == "flagembedding":
|
|
out = self._model.encode( # type: ignore[union-attr]
|
|
list(texts),
|
|
batch_size=self.batch_size,
|
|
max_length=8192,
|
|
return_dense=True,
|
|
return_sparse=False,
|
|
return_colbert_vecs=False,
|
|
)
|
|
dense = out["dense_vecs"] if isinstance(out, dict) else out
|
|
arr = np.asarray(dense, dtype=np.float32)
|
|
else:
|
|
arr = self._st_model.encode( # type: ignore[union-attr]
|
|
list(texts),
|
|
batch_size=self.batch_size,
|
|
normalize_embeddings=self.normalize,
|
|
convert_to_numpy=True,
|
|
show_progress_bar=False,
|
|
)
|
|
arr = arr.astype(np.float32)
|
|
|
|
if self.normalize and self._impl == "flagembedding":
|
|
norms = np.linalg.norm(arr, axis=1, keepdims=True)
|
|
norms[norms == 0] = 1.0
|
|
arr = arr / norms
|
|
|
|
return arr.tolist()
|
|
|
|
def encode_one(self, text: str) -> list[float]:
|
|
return self.encode([text])[0]
|
|
|
|
|
|
@lru_cache(maxsize=1)
|
|
def get_embedder() -> Embedder:
|
|
return Embedder(
|
|
model_name=settings.embedding_model,
|
|
device=settings.embedding_device,
|
|
normalize=settings.embedding_normalize,
|
|
batch_size=settings.embedding_batch_size,
|
|
)
|