"""``Model2VecEmbedder``: local static embeddings, computed in this process.
A local, CPU-backed :class:`~fastrag.ingestion.steps.embedding.base.Embedder`:
loads a small static model into memory and encodes in-process, so indexing
needs no embedding API key.
"""
from __future__ import annotations
import asyncio
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Any, Literal
from fastrag.exceptions import ConfigurationError, EmbeddingError
from fastrag.ingestion.cache import CacheStore
from fastrag.ingestion.policy import Policy
from fastrag.ingestion.steps.embedding.base import Embedder
if TYPE_CHECKING:
from model2vec import StaticModel
try:
from model2vec import StaticModel as _StaticModelImpl
_MODEL2VEC_AVAILABLE = True
except ImportError: # pragma: no cover - exercised by installs without the extra
_MODEL2VEC_AVAILABLE = False
#: The model class, or ``None`` when the optional "model2vec" extra is
#: missing; :meth:`Model2VecEmbedder.__init__` refuses to run without it.
_StaticModel: Any = _StaticModelImpl if _MODEL2VEC_AVAILABLE else None
MISSING_EXTRA = (
'Model2VecEmbedder needs the "model2vec" extra: pip install "fastrag-ai[model2vec]"'
)
DEFAULT_MODEL = "minishlab/potion-base-8M"
[docs]
class Model2VecEmbedder(Embedder):
"""Embed with a static model, computed in this process.
Choose this over
:class:`~fastrag.ingestion.steps.embedding.openai.OpenAIEmbedder` when the
run must not depend on a server: a CI pipeline, a laptop with no endpoint,
a demo. The model is a lookup table of token vectors averaged per text —
no attention, no GPU, no server — which is why it is fast and why it is
weaker than a transformer encoder on paraphrase.
.. code-block:: python
from fastrag.ingestion.steps.embedding import Model2VecEmbedder
embedder = Model2VecEmbedder(model="minishlab/potion-base-8M", cache=cache)
Args:
model: HuggingFace model id or a local directory of weights.
Defaults to ``minishlab/potion-base-8M``: ~30 MB, 256 dimensions.
dimensions: Truncate the model's output to this width. ``None`` keeps
the model's full width. Changing it also changes :attr:`model_id`
(and therefore the document-cache identity), so a truncated model
never shares a cached vector with the full one.
normalize: Whether to L2-normalize vectors. ``None`` keeps the model's
own default (potion models ship normalized).
batch_size: Texts per model2vec ``encode`` call. Distinct from
``policy.batch_size`` (how many documents this step's batcher
groups into one call to :meth:`embed`) — this instead bounds
memory *inside* one such call, exactly as in the old service.
token: HuggingFace access token, for gated models.
The first construction downloads ~30 MB into ``~/.cache/huggingface``;
every later one reads the cache, and ``HF_HUB_OFFLINE=1`` forbids the
network outright. Pass a local directory as ``model`` for a fully
air-gapped install.
"""
def __init__(
self,
*,
model: str = DEFAULT_MODEL,
dimensions: int | None = None,
normalize: bool | None = None,
batch_size: int = 1024,
token: str | None = None,
cache: CacheStore | Literal[False] | None = None,
policy: Policy | None = None,
) -> None:
"""Configure and load the model, resolving its dimension immediately.
Raises:
fastrag.exceptions.ConfigurationError: The "model2vec" extra is
not installed, ``dimensions < 1``, ``batch_size < 1``, or the
model failed to load.
"""
static_model = _StaticModel
if static_model is None:
raise ConfigurationError(MISSING_EXTRA)
if dimensions is not None and dimensions < 1:
raise ConfigurationError("Model2VecEmbedder.dimensions must be at least 1")
if batch_size < 1:
raise ConfigurationError("Model2VecEmbedder.batch_size must be at least 1")
super().__init__(cache=cache, policy=policy)
self.model = model
self.dimensions = dimensions
self.normalize = normalize
self.batch_size = batch_size
self.token = token
try:
self._model: StaticModel = static_model.from_pretrained(
model,
token=token,
normalize=normalize,
dimensionality=dimensions,
# model2vec defaults this to True, which re-downloads the
# weights on every construction. The cached copy is the whole
# point: after the first load this embedder touches no network.
force_download=False,
)
except Exception as error:
raise ConfigurationError(
f"could not load model2vec model {model!r}: {error}"
) from error
self._dimension = int(self._model.dim)
[docs]
def config(self) -> Mapping[str, Any]:
"""This embedder's configuration, for cache-identity hashing."""
return {
"model": self.model,
"dimensions": self.dimensions,
"normalize": self.normalize,
}
@property
def model_id(self) -> str:
"""The model, suffixed with the pinned dimension when there is one."""
if self.dimensions is None:
return self.model
return f"{self.model}@{self.dimensions}"
@property
def dimension(self) -> int:
"""The width of every returned vector, read from the loaded model."""
return self._dimension
[docs]
async def embed(self, texts: Sequence[str]) -> list[list[float]]:
"""Embed a batch of texts.
The work happens in a worker thread: ``encode`` is synchronous and
CPU-bound, and offloading it keeps other batches' concurrent work
(cache lookups, other steps) running while this one encodes.
"""
if not texts:
return []
try:
vectors = await asyncio.to_thread(
self._model.encode,
list(texts),
batch_size=self.batch_size,
show_progress_bar=False,
# model2vec forks worker processes above 10 000 texts; forking
# from a worker thread inside an asyncio program is a bad
# trade for a batch that encodes in milliseconds.
use_multiprocessing=False,
)
except Exception as error:
# No TransientEmbeddingError from this backend: a local matrix
# lookup has no transient failure mode — if it raises, it raises
# identically next time. Marking it retryable would make a
# permanent bug take three times as long to fail.
raise EmbeddingError(f"embedding failed: {error}") from error
return vectors.tolist()