104 lines
3.5 KiB
Python
104 lines
3.5 KiB
Python
"""BAAI/bge-m3 running locally — the zero-API-cost control in the benchmark.
|
|
|
|
Its job is to answer "how much is the paid model actually buying us on
|
|
Vietnamese medical prose?". Without a free baseline in the same harness, a
|
|
cloud model's recall number has nothing to be better *than*.
|
|
|
|
Two properties are taken from the published model card and have **not** been
|
|
verified on this machine (no local run has happened yet — see the coordination
|
|
handoff): the dense vector is 1024-dimensional, and bge-m3 needs no
|
|
instruction prefix on either the corpus or the query side, unlike the earlier
|
|
English bge models. Both are asserted at runtime rather than trusted: the
|
|
dimension is checked on every vector by `EmbeddingProvider._check_dimensions`,
|
|
so a wrong assumption fails on the first call instead of producing a
|
|
quietly unusable collection.
|
|
|
|
`sentence-transformers` is imported lazily and the encoder is injectable, so
|
|
this module costs nothing to import and can be tested without the ~2 GB of
|
|
model weights.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from typing import Callable, List, Optional, Sequence
|
|
|
|
from .ports import EmbeddingProvider, EmbeddingVector, text_digest
|
|
|
|
MODEL_ID = "BAAI/bge-m3"
|
|
PROVIDER_NAME = "bge-m3"
|
|
|
|
DENSE_DIMENSIONS = 1024
|
|
|
|
Encoder = Callable[[Sequence[str]], Sequence[Sequence[float]]]
|
|
|
|
|
|
class BgeM3Local(EmbeddingProvider):
|
|
def __init__(
|
|
self,
|
|
encoder: Optional[Encoder] = None,
|
|
batch_size: int = 16,
|
|
device: Optional[str] = None,
|
|
):
|
|
if batch_size < 1:
|
|
raise ValueError(f"batch_size must be >= 1, got {batch_size}")
|
|
self._encoder = encoder
|
|
self._batch_size = batch_size
|
|
self._device = device
|
|
# Only the encoder built below is known to normalize. An injected one
|
|
# is somebody else's function, so its output is recorded as unknown.
|
|
self._normalized: Optional[bool] = None if encoder is not None else True
|
|
|
|
@property
|
|
def name(self) -> str:
|
|
return PROVIDER_NAME
|
|
|
|
@property
|
|
def model_id(self) -> str:
|
|
return MODEL_ID
|
|
|
|
@property
|
|
def dimensions(self) -> int:
|
|
return DENSE_DIMENSIONS
|
|
|
|
@property
|
|
def max_batch_size(self) -> int:
|
|
return self._batch_size
|
|
|
|
def _load_encoder(self) -> Encoder:
|
|
if self._encoder is None:
|
|
from sentence_transformers import SentenceTransformer
|
|
|
|
model = SentenceTransformer(MODEL_ID, device=self._device)
|
|
|
|
def encode(texts: Sequence[str]) -> Sequence[Sequence[float]]:
|
|
return model.encode(
|
|
list(texts), normalize_embeddings=True
|
|
).tolist()
|
|
|
|
self._encoder = encode
|
|
return self._encoder
|
|
|
|
def _embed_batch(
|
|
self, texts: Sequence[str], input_kind: str
|
|
) -> List[EmbeddingVector]:
|
|
rows = self._load_encoder()(texts)
|
|
if len(rows) != len(texts):
|
|
raise ValueError(
|
|
f"{MODEL_ID} returned {len(rows)} vectors for {len(texts)} texts"
|
|
)
|
|
|
|
vectors: List[EmbeddingVector] = []
|
|
for text, values in zip(texts, rows, strict=True):
|
|
self._check_dimensions(values)
|
|
vectors.append(
|
|
EmbeddingVector(
|
|
values=list(values),
|
|
text_sha256=text_digest(text),
|
|
provider=PROVIDER_NAME,
|
|
model_id=MODEL_ID,
|
|
dimensions=DENSE_DIMENSIONS,
|
|
input_kind=input_kind,
|
|
normalized=self._normalized,
|
|
)
|
|
)
|
|
return vectors
|