From 3b8d169170dbee8348d5cd955bc97ffe7ad1b146 Mon Sep 17 00:00:00 2001 From: George Panchuk Date: Mon, 12 Oct 2026 20:58:01 +0700 Subject: [PATCH 1/5] new: run minishlab potion models with numpy instead of onnx --- fastembed/common/safetensors_utils.py | 52 +++++ fastembed/text/onnx_embedding.py | 36 --- fastembed/text/static_embedding.py | 302 ++++++++++++++++++++++++++ fastembed/text/text_embedding.py | 2 + tests/test_static_embeddings.py | 137 ++++++++++++ tests/test_text_onnx_embeddings.py | 9 +- 6 files changed, 500 insertions(+), 38 deletions(-) create mode 100644 fastembed/common/safetensors_utils.py create mode 100644 fastembed/text/static_embedding.py create mode 100644 tests/test_static_embeddings.py diff --git a/fastembed/common/safetensors_utils.py b/fastembed/common/safetensors_utils.py new file mode 100644 index 00000000..a82f7dbb --- /dev/null +++ b/fastembed/common/safetensors_utils.py @@ -0,0 +1,52 @@ +import json +from pathlib import Path + +import numpy as np + +from fastembed.common.types import NumpyArray + +# safetensors dtype -> numpy dtype. BF16 has no numpy dtype and is widened to float32 on load. +_DTYPES: dict[str, str] = { + "F64": " dict[str, NumpyArray]: + """Read every tensor of a safetensors file with numpy, without the `safetensors` package. + + The format is an 8 byte little-endian header size, a JSON header with the dtype, shape and + byte offsets of each tensor, then the raw tensor data. Tensors are memory-mapped read-only, so + loading is instant and parallel workers share the pages of one file. + """ + path = Path(path) + with open(path, "rb") as f: + header_size = int.from_bytes(f.read(8), "little") + header = json.loads(f.read(header_size)) + header.pop("__metadata__", None) + data_start = 8 + header_size + + tensors: dict[str, NumpyArray] = {} + for name, info in header.items(): + start, end = info["data_offsets"] + shape = tuple(info["shape"]) + if end == start: + tensors[name] = np.zeros(shape, dtype=np.float32) + continue + if info["dtype"] == "BF16": + raw = np.memmap(path, dtype=" tuple[Tokenizer, int | None]: + """Load the tokenizer of a static model: truncated to MAX_LENGTH, without padding. + + Returns the tokenizer and the id of its unknown token, if it has one. Unknown tokens carry no + meaning in a static model, so they are left out of the mean, as model2vec does. + """ + tokenizer_path = model_dir / "tokenizer.json" + if not tokenizer_path.exists(): + raise ValueError(f"Could not find tokenizer.json in {model_dir}") + + tokenizer = Tokenizer.from_file(str(tokenizer_path)) + tokenizer.enable_truncation(max_length=MAX_LENGTH) + tokenizer.no_padding() + + with open(tokenizer_path) as f: + tokenizer_model = json.load(f)["model"] + if tokenizer_model.get("type") == "Unigram": + unk_token_id = tokenizer_model.get("unk_id") + else: + unk_token = tokenizer_model.get("unk_token") + vocab = tokenizer_model.get("vocab") + unk_token_id = vocab.get(unk_token) if unk_token and isinstance(vocab, dict) else None + return tokenizer, unk_token_id + + +class StaticTextEmbedding(TextEmbeddingBase): + """Static (model2vec) embeddings: the mean of the token vectors, computed with numpy. + + There is no network to run, so tokenization dominates the cost. The vectors are read straight + from `model.safetensors` and pooled without the padding an onnx graph needs. + """ + + @classmethod + def _list_supported_models(cls) -> list[DenseModelDescription]: + return supported_static_models + + def __init__( + self, + model_name: str, + cache_dir: str | None = None, + threads: int | None = None, + lazy_load: bool = False, + specific_model_path: str | None = None, + **kwargs: Any, + ): + """ + Args: + model_name (str): The name of the model to use. + cache_dir (str, optional): The path to the cache directory. + Can be set using the `FASTEMBED_CACHE_PATH` env variable. + Defaults to `fastembed_cache` in the system's temp directory. + threads (int, optional): Unused. Tokenization runs on the thread pool of the tokenizers + library, which `TOKENIZERS_PARALLELISM=false` turns off. + lazy_load (bool, optional): Whether to load the model during class initialization or + on demand. Defaults to False. + specific_model_path (str, optional): A local directory with `model.safetensors` and + `tokenizer.json` to load instead of downloading the model. + """ + # TextEmbedding passes onnx runtime options to every model class; static models run on + # the CPU with numpy and ignore them. + providers = kwargs.pop("providers", None) + cuda = kwargs.pop("cuda", Device.AUTO) + for option in ("device_ids", "device_id"): + kwargs.pop(option, None) + if providers or cuda in (True, Device.CUDA): + logger.warning(f"{model_name} is a static model, it runs on the CPU only") + super().__init__(model_name, cache_dir, threads, **kwargs) + self.lazy_load = lazy_load + self.model_description = self._get_model_description(model_name) + self.cache_dir = str(define_cache_dir(cache_dir)) + self._specific_model_path = specific_model_path + self._model_dir = self.download_model( + self.model_description, + self.cache_dir, + local_files_only=self._local_files_only, + specific_model_path=self._specific_model_path, + ) + + self.embeddings: NumpyArray | None = None + self.weights: NumpyArray | None = None + self.token_mapping: NumpyArray | None = None + self.tokenizer: Tokenizer | None = None + self.unk_token_id: int | None = None + if not self.lazy_load: + self.load_model() + + def load_model(self) -> None: + tensors = load_safetensors(self._model_dir / self.model_description.model_file) + if "embeddings" not in tensors: + raise ValueError(f"No `embeddings` tensor in {self._model_dir}") + self.embeddings = tensors["embeddings"] + # Vocabulary-quantized model2vec models map token ids to shared rows and reweight them. + self.weights = tensors.get("weights") + self.token_mapping = tensors.get("mapping") + self.tokenizer, self.unk_token_id = load_static_tokenizer(self._model_dir) + + def _ensure_loaded(self) -> None: + if self.embeddings is None or self.tokenizer is None: + self.load_model() + + def _token_ids(self, documents: list[str]) -> tuple[NDArray[np.int64], NDArray[np.int64]]: + """Tokenize a batch into one flat id array plus the number of ids of each document.""" + assert self.tokenizer is not None + # encode_batch_fast skips the character offsets, which are not needed here. + encodings = self.tokenizer.encode_batch_fast(documents, add_special_tokens=False) + id_lists = [encoding.ids for encoding in encodings] + lengths = np.fromiter(map(len, id_lists), dtype=np.int64, count=len(documents)) + ids = np.fromiter( + itertools.chain.from_iterable(id_lists), dtype=np.int64, count=int(lengths.sum()) + ) + return ids, lengths + + def _embed_batch(self, documents: list[str]) -> NumpyArray: + self._ensure_loaded() + assert self.embeddings is not None + ids, lengths = self._token_ids(documents) + + if self.unk_token_id is not None: + known = ids != self.unk_token_id + if not known.all(): + doc_index = np.repeat(np.arange(len(documents)), lengths) + lengths = np.bincount(doc_index[known], minlength=len(documents)) + ids = ids[known] + + sums = np.zeros((len(documents), self.embeddings.shape[1]), dtype=np.float32) + ends = np.cumsum(lengths) + for start in range(0, len(documents), POOLING_CHUNK): + stop = min(start + POOLING_CHUNK, len(documents)) + first, last = ends[start] - lengths[start], ends[stop - 1] + sums[start:stop] = self._sum_rows(ids[first:last], lengths[start:stop]) + # The mean and the sum point the same way, so normalizing the sum gives the mean's + # direction. Empty documents stay zero vectors. + return normalize(sums) + + def _sum_rows(self, ids: NDArray[np.int64], lengths: NDArray[np.int64]) -> NumpyArray: + """Sum the vectors of each document's tokens. + + The sums are a product of a (documents x distinct tokens) count matrix with the rows of + those tokens: one matrix multiplication, several times faster than segment sums with + np.add.reduceat or a loop over the documents. + """ + assert self.embeddings is not None + if not ids.size: + return np.zeros((len(lengths), self.embeddings.shape[1]), dtype=np.float32) + rows = ids if self.token_mapping is None else self.token_mapping[ids] + distinct_rows, row_index = np.unique(rows, return_inverse=True) + doc_index = np.repeat(np.arange(len(lengths)), lengths) + counts = np.bincount( + doc_index * len(distinct_rows) + row_index, + weights=None if self.weights is None else self.weights[ids], + minlength=len(lengths) * len(distinct_rows), + ).astype(np.float32) + return counts.reshape(len(lengths), len(distinct_rows)) @ self.embeddings[distinct_rows] + + def embed( + self, + documents: str | Iterable[str], + batch_size: int = 256, + parallel: int | None = None, + **kwargs: Any, + ) -> Iterable[NumpyArray]: + """ + Encode a list of documents into list of embeddings. + + Args: + documents: Iterator of documents or single document to embed + batch_size: Batch size for encoding + parallel: + If > 1, data-parallel encoding will be used, recommended for offline encoding of large datasets. + If 0, use all available cores. + If None, don't use data-parallel processing. + + Returns: + List of embeddings, one per document + """ + is_small = isinstance(documents, str) or ( + isinstance(documents, list) and len(documents) < batch_size + ) + if isinstance(documents, str): + documents = [documents] + + if parallel is None or is_small: + for batch in iter_batch(documents, batch_size): + yield from self._embed_batch(batch) + return + + if parallel == 0: + parallel = os.cpu_count() + start_method = "forkserver" if "forkserver" in get_all_start_methods() else "spawn" + params = { + "model_name": self.model_name, + "cache_dir": self.cache_dir, + "local_files_only": self._local_files_only, + "specific_model_path": self._specific_model_path, + } + pool = ParallelWorkerPool( + num_workers=parallel or 1, worker=self._get_worker_class(), start_method=start_method + ) + for batch in pool.ordered_map(iter_batch(documents, batch_size), **params): + yield from batch + + @classmethod + def _get_worker_class(cls) -> Type["StaticTextEmbeddingWorker"]: + return StaticTextEmbeddingWorker + + @classmethod + def get_embedding_size(cls, model_name: str) -> int: + dim = cls._get_model_description(model_name).dim + if dim is None: + raise ValueError(f"Embedding size of {model_name} is unknown") + return dim + + @property + def embedding_size(self) -> int: + if self._embedding_size is None: + self._embedding_size = self.get_embedding_size(self.model_name) + return self._embedding_size + + def token_count( + self, texts: str | Iterable[str], batch_size: int = 1024, **kwargs: Any + ) -> int: + """Returns the number of tokens that are averaged, after truncation.""" + self._ensure_loaded() + texts = [texts] if isinstance(texts, str) else texts + return sum(int(self._token_ids(batch)[1].sum()) for batch in iter_batch(texts, batch_size)) + + +class StaticTextEmbeddingWorker(Worker): + def __init__(self, model_name: str, cache_dir: str, **kwargs: Any): + self.model = StaticTextEmbedding(model_name=model_name, cache_dir=cache_dir, **kwargs) + + @classmethod + def start(cls, model_name: str, cache_dir: str, **kwargs: Any) -> "StaticTextEmbeddingWorker": + return cls(model_name=model_name, cache_dir=cache_dir, **kwargs) + + def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, NumpyArray]]: + for idx, batch in items: + yield idx, self.model._embed_batch(batch) diff --git a/fastembed/text/text_embedding.py b/fastembed/text/text_embedding.py index 14966383..0468822c 100644 --- a/fastembed/text/text_embedding.py +++ b/fastembed/text/text_embedding.py @@ -12,6 +12,7 @@ from fastembed.text.last_token_normalized_embedding import LastTokenNormalizedEmbedding from fastembed.text.onnx_embedding import OnnxTextEmbedding from fastembed.text.siglip_embedding import SiglipOnnxTextEmbedding +from fastembed.text.static_embedding import StaticTextEmbedding from fastembed.text.text_embedding_base import TextEmbeddingBase from fastembed.common.model_description import DenseModelDescription, ModelSource, PoolingType @@ -26,6 +27,7 @@ class TextEmbedding(TextEmbeddingBase): JinaEmbeddingV3, BuiltinSentenceEmbedding, LastTokenNormalizedEmbedding, + StaticTextEmbedding, CustomTextEmbedding, ] diff --git a/tests/test_static_embeddings.py b/tests/test_static_embeddings.py new file mode 100644 index 00000000..6d65af7b --- /dev/null +++ b/tests/test_static_embeddings.py @@ -0,0 +1,137 @@ +import json +import os +from typing import Iterator + +import numpy as np +import pytest + +from fastembed.common.safetensors_utils import load_safetensors +from fastembed.text.static_embedding import StaticTextEmbedding +from fastembed.text.text_embedding import TextEmbedding +from tests.test_text_onnx_embeddings import CANONICAL_VECTOR_VALUES +from tests.utils import delete_model_cache, is_manual_run + +STATIC_MODELS = [model.model for model in StaticTextEmbedding._list_supported_models()] + +CI_MODEL = "minishlab/potion-base-8M" + + +def _models_to_test() -> list[str]: + if os.getenv("CI") and not is_manual_run(): + return [CI_MODEL] + return STATIC_MODELS + + +@pytest.fixture(scope="module") +def model() -> Iterator[TextEmbedding]: + model = TextEmbedding(CI_MODEL) + yield model + if os.getenv("CI"): + delete_model_cache(model.model._model_dir) + + +@pytest.mark.parametrize("model_name", _models_to_test()) +def test_canonical_vectors(model_name: str) -> None: + model = TextEmbedding(model_name) + assert isinstance(model.model, StaticTextEmbedding) + embeddings = np.stack(list(model.embed(["hello world", "flag embedding"]))) + assert embeddings.shape == (2, model.embedding_size) + canonical = CANONICAL_VECTOR_VALUES[model_name] + assert np.allclose(embeddings[0, : canonical.shape[0]], canonical, atol=1e-3), model_name + assert np.allclose(np.linalg.norm(embeddings, axis=1), 1.0, atol=1e-5) + if os.getenv("CI") and model_name != CI_MODEL: + delete_model_cache(model.model._model_dir) + + +def test_empty_and_unknown_input_give_zero_vectors(model: TextEmbedding) -> None: + # Emoji are outside the WordPiece vocabulary, so the second text is only unknown tokens. + embeddings = np.stack(list(model.embed(["", "🤗🤗", "hello world"]))) + assert not np.isnan(embeddings).any() + assert np.all(embeddings[0] == 0) and np.all(embeddings[1] == 0) + assert np.isclose(np.linalg.norm(embeddings[2]), 1.0, atol=1e-5) + + +def test_unknown_tokens_do_not_change_the_mean(model: TextEmbedding) -> None: + plain, with_unknown = model.embed(["hello world", "hello 🤗 world"]) + assert np.allclose(plain, with_unknown, atol=1e-6) + + +def test_batches_match_single_documents(model: TextEmbedding) -> None: + docs = ["", "a short text", "hello world " * 40, "", "flag embedding", "🤗"] * 7 + batched = np.stack(list(model.embed(docs, batch_size=4))) + single = np.stack([next(iter(model.embed(doc))) for doc in docs]) + assert np.allclose(batched, single, atol=1e-6) + + +def test_truncation_to_512_tokens(model: TextEmbedding) -> None: + long_text = " ".join(f"word{i}" for i in range(2000)) + tokenizer = model.model.tokenizer + ids = tokenizer.encode(long_text, add_special_tokens=False).ids + assert len(ids) == 512 + prefix = tokenizer.decode(ids) + assert np.allclose( + next(iter(model.embed(long_text))), next(iter(model.embed(prefix))), atol=1e-6 + ) + + +def test_parallel_matches_sequential(model: TextEmbedding) -> None: + docs = ["hello world", "flag embedding", "", "a much longer document about vectors " * 5] * 100 + sequential = np.stack(list(model.embed(docs, batch_size=10))) + parallel = np.stack(list(model.embed(docs, batch_size=10, parallel=2))) + assert np.allclose(sequential, parallel, atol=1e-6) + + +def test_lazy_load() -> None: + model = TextEmbedding(CI_MODEL, lazy_load=True) + assert model.model.embeddings is None + assert len(list(model.embed(["hello world"]))) == 1 + assert model.model.embeddings is not None + + +def test_token_count(model: TextEmbedding) -> None: + assert model.model.token_count("hello world") == 2 + assert model.model.token_count(["hello world", "", "flag embedding"]) == 6 + assert model.model.token_count(" ".join(["word"] * 1000)) == 512 + + +def test_embedding_size() -> None: + assert TextEmbedding.get_embedding_size("minishlab/potion-retrieval-32M") == 512 + assert TextEmbedding.get_embedding_size("minishlab/potion-base-8M") == 256 + + +def _write_safetensors(path, tensors: dict[str, tuple[str, list[int], bytes]]) -> None: + header, data, offset = {}, b"", 0 + for name, (dtype, shape, raw) in tensors.items(): + header[name] = { + "dtype": dtype, + "shape": shape, + "data_offsets": [offset, offset + len(raw)], + } + data += raw + offset += len(raw) + header["__metadata__"] = {"format": "pt"} + encoded = json.dumps(header).encode() + with open(path, "wb") as f: + f.write(len(encoded).to_bytes(8, "little")) + f.write(encoded) + f.write(data) + + +def test_load_safetensors(tmp_path) -> None: + values = np.arange(6, dtype=np.float32).reshape(2, 3) / 4 + bf16 = (values.view(np.uint32) >> 16).astype(" None: @pytest.mark.parametrize( "embedding_class", - # custom models are covered in test_custom_models.py - [cls for cls in TextEmbedding.EMBEDDINGS_REGISTRY if cls is not CustomTextEmbedding], + # custom models are covered in test_custom_models.py, static models have no onnx output + [ + cls + for cls in TextEmbedding.EMBEDDINGS_REGISTRY + if cls not in (CustomTextEmbedding, StaticTextEmbedding) + ], ids=lambda cls: cls.__name__, ) def test_post_processing_keeps_model_dtype(embedding_class) -> None: From d1d316340d08e69a36e1d17fa6ab3122299eb623 Mon Sep 17 00:00:00 2001 From: George Panchuk Date: Tue, 13 Oct 2026 00:04:55 +0700 Subject: [PATCH 2/5] fix: read tokenizer.json as utf-8 in static embeddings --- fastembed/text/static_embedding.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/fastembed/text/static_embedding.py b/fastembed/text/static_embedding.py index 174ab094..d6aa30e9 100644 --- a/fastembed/text/static_embedding.py +++ b/fastembed/text/static_embedding.py @@ -76,7 +76,7 @@ def load_static_tokenizer(model_dir: Path) -> tuple[Tokenizer, int | None]: tokenizer.enable_truncation(max_length=MAX_LENGTH) tokenizer.no_padding() - with open(tokenizer_path) as f: + with open(tokenizer_path, encoding="utf-8") as f: tokenizer_model = json.load(f)["model"] if tokenizer_model.get("type") == "Unigram": unk_token_id = tokenizer_model.get("unk_id") From c8d51bb4f770fa90f232655caaebaa552504c6fb Mon Sep 17 00:00:00 2001 From: George Panchuk Date: Tue, 13 Oct 2026 00:40:37 +0700 Subject: [PATCH 3/5] refactor: restructure static embeddings like other models --- fastembed/text/static_embedding.py | 137 ++++++++++++++++++----------- tests/test_static_embeddings.py | 23 +++++ 2 files changed, 110 insertions(+), 50 deletions(-) diff --git a/fastembed/text/static_embedding.py b/fastembed/text/static_embedding.py index d6aa30e9..16d87795 100644 --- a/fastembed/text/static_embedding.py +++ b/fastembed/text/static_embedding.py @@ -1,12 +1,12 @@ import itertools import json import os +import warnings from multiprocessing import get_all_start_methods from pathlib import Path from typing import Any, Iterable, Type import numpy as np -from loguru import logger from numpy.typing import NDArray from tokenizers import Tokenizer @@ -113,21 +113,35 @@ def __init__( cache_dir (str, optional): The path to the cache directory. Can be set using the `FASTEMBED_CACHE_PATH` env variable. Defaults to `fastembed_cache` in the system's temp directory. - threads (int, optional): Unused. Tokenization runs on the thread pool of the tokenizers - library, which `TOKENIZERS_PARALLELISM=false` turns off. + threads (int, optional): Not supported, a warning is raised if set. Tokenization runs on + the thread pool of the tokenizers library, which `TOKENIZERS_PARALLELISM=false` + turns off. lazy_load (bool, optional): Whether to load the model during class initialization or on demand. Defaults to False. specific_model_path (str, optional): A local directory with `model.safetensors` and `tokenizer.json` to load instead of downloading the model. """ - # TextEmbedding passes onnx runtime options to every model class; static models run on - # the CPU with numpy and ignore them. - providers = kwargs.pop("providers", None) - cuda = kwargs.pop("cuda", Device.AUTO) - for option in ("device_ids", "device_id"): - kwargs.pop(option, None) - if providers or cuda in (True, Device.CUDA): - logger.warning(f"{model_name} is a static model, it runs on the CPU only") + # TextEmbedding passes its onnx runtime options to every model class. Static models run + # on the CPU with numpy and ignore them, so only options set to something else than their + # default are reported. + ignored = {"threads": threads} if threads is not None else {} + for option, default in ( + ("providers", None), + ("cuda", Device.AUTO), + ("device_ids", None), + ("enable_cpu_mem_arena", None), + ): + value = kwargs.pop(option, default) + if value != default: + ignored[option] = value + kwargs.pop("device_id", None) # set by parallel workers, not by users + if ignored: + warnings.warn( + f"{model_name} is a static model computed with numpy on the CPU, " + f"so these options have no effect: {', '.join(sorted(ignored))}", + UserWarning, + stacklevel=3, + ) super().__init__(model_name, cache_dir, threads, **kwargs) self.lazy_load = lazy_load self.model_description = self._get_model_description(model_name) @@ -158,7 +172,7 @@ def load_model(self) -> None: self.token_mapping = tensors.get("mapping") self.tokenizer, self.unk_token_id = load_static_tokenizer(self._model_dir) - def _ensure_loaded(self) -> None: + def _ensure_model(self) -> None: if self.embeddings is None or self.tokenizer is None: self.load_model() @@ -174,8 +188,9 @@ def _token_ids(self, documents: list[str]) -> tuple[NDArray[np.int64], NDArray[n ) return ids, lengths - def _embed_batch(self, documents: list[str]) -> NumpyArray: - self._ensure_loaded() + def raw_embed(self, documents: list[str]) -> NumpyArray: + """Embed one batch of documents, without batching or parallelism.""" + self._ensure_model() assert self.embeddings is not None ids, lengths = self._token_ids(documents) @@ -237,61 +252,79 @@ def embed( Returns: List of embeddings, one per document """ - is_small = isinstance(documents, str) or ( - isinstance(documents, list) and len(documents) < batch_size + yield from self._embed_documents( + model_name=self.model_name, + cache_dir=str(self.cache_dir), + documents=documents, + batch_size=batch_size, + parallel=parallel, + local_files_only=self._local_files_only, + specific_model_path=self._specific_model_path, ) + + def _embed_documents( + self, + model_name: str, + cache_dir: str, + documents: str | Iterable[str], + batch_size: int = 256, + parallel: int | None = None, + local_files_only: bool = False, + specific_model_path: str | None = None, + ) -> Iterable[NumpyArray]: + is_small = False + if isinstance(documents, str): documents = [documents] + is_small = True + + if isinstance(documents, list): + if len(documents) < batch_size: + is_small = True if parallel is None or is_small: for batch in iter_batch(documents, batch_size): - yield from self._embed_batch(batch) - return - - if parallel == 0: - parallel = os.cpu_count() - start_method = "forkserver" if "forkserver" in get_all_start_methods() else "spawn" - params = { - "model_name": self.model_name, - "cache_dir": self.cache_dir, - "local_files_only": self._local_files_only, - "specific_model_path": self._specific_model_path, - } - pool = ParallelWorkerPool( - num_workers=parallel or 1, worker=self._get_worker_class(), start_method=start_method - ) - for batch in pool.ordered_map(iter_batch(documents, batch_size), **params): - yield from batch + yield from self.raw_embed(batch) + else: + if parallel == 0: + parallel = os.cpu_count() + + start_method = "forkserver" if "forkserver" in get_all_start_methods() else "spawn" + params = { + "model_name": model_name, + "cache_dir": cache_dir, + "local_files_only": local_files_only, + "specific_model_path": specific_model_path, + } + pool = ParallelWorkerPool( + num_workers=parallel or 1, + worker=self._get_worker_class(), + start_method=start_method, + ) + for batch in pool.ordered_map(iter_batch(documents, batch_size), **params): + yield from batch @classmethod def _get_worker_class(cls) -> Type["StaticTextEmbeddingWorker"]: return StaticTextEmbeddingWorker - @classmethod - def get_embedding_size(cls, model_name: str) -> int: - dim = cls._get_model_description(model_name).dim - if dim is None: - raise ValueError(f"Embedding size of {model_name} is unknown") - return dim - - @property - def embedding_size(self) -> int: - if self._embedding_size is None: - self._embedding_size = self.get_embedding_size(self.model_name) - return self._embedding_size - def token_count( self, texts: str | Iterable[str], batch_size: int = 1024, **kwargs: Any ) -> int: """Returns the number of tokens that are averaged, after truncation.""" - self._ensure_loaded() + self._ensure_model() texts = [texts] if isinstance(texts, str) else texts return sum(int(self._token_ids(batch)[1].sum()) for batch in iter_batch(texts, batch_size)) class StaticTextEmbeddingWorker(Worker): - def __init__(self, model_name: str, cache_dir: str, **kwargs: Any): - self.model = StaticTextEmbedding(model_name=model_name, cache_dir=cache_dir, **kwargs) + def __init__( + self, + model_name: str, + cache_dir: str, + **kwargs: Any, + ): + self.model = self.init_embedding(model_name, cache_dir, **kwargs) @classmethod def start(cls, model_name: str, cache_dir: str, **kwargs: Any) -> "StaticTextEmbeddingWorker": @@ -299,4 +332,8 @@ def start(cls, model_name: str, cache_dir: str, **kwargs: Any) -> "StaticTextEmb def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, NumpyArray]]: for idx, batch in items: - yield idx, self.model._embed_batch(batch) + yield idx, self.model.raw_embed(batch) + + @staticmethod + def init_embedding(model_name: str, cache_dir: str, **kwargs: Any) -> StaticTextEmbedding: + return StaticTextEmbedding(model_name=model_name, cache_dir=cache_dir, **kwargs) diff --git a/tests/test_static_embeddings.py b/tests/test_static_embeddings.py index 6d65af7b..235830e1 100644 --- a/tests/test_static_embeddings.py +++ b/tests/test_static_embeddings.py @@ -1,5 +1,6 @@ import json import os +import warnings from typing import Iterator import numpy as np @@ -135,3 +136,25 @@ def test_load_safetensors(tmp_path) -> None: for name in ("f32", "f16", "bf16"): assert np.array_equal(np.asarray(tensors[name], dtype=np.float32), values), name assert tensors["mapping"].tolist() == [2, 0, 1] + + +@pytest.mark.parametrize( + "option", + [ + {"threads": 2}, + {"providers": ["CPUExecutionProvider"]}, + {"cuda": True}, + {"device_ids": [0]}, + {"enable_cpu_mem_arena": False}, + ], + ids=lambda option: next(iter(option)), +) +def test_unsupported_options_warn(option: dict) -> None: + with pytest.warns(UserWarning, match=next(iter(option))): + TextEmbedding(CI_MODEL, lazy_load=True, **option) + + +def test_default_options_do_not_warn() -> None: + with warnings.catch_warnings(): + warnings.simplefilter("error") + TextEmbedding(CI_MODEL, lazy_load=True) From f1f82be04a34051ec0d39a24bef0d4e1b3cb02a4 Mon Sep 17 00:00:00 2001 From: George Panchuk Date: Tue, 13 Oct 2026 01:01:45 +0700 Subject: [PATCH 4/5] fix: cap temporary memory of static embedding pooling --- fastembed/text/static_embedding.py | 26 +++++++++++++++++--------- tests/test_static_embeddings.py | 15 +++++++++++++++ 2 files changed, 32 insertions(+), 9 deletions(-) diff --git a/fastembed/text/static_embedding.py b/fastembed/text/static_embedding.py index 16d87795..11d093e1 100644 --- a/fastembed/text/static_embedding.py +++ b/fastembed/text/static_embedding.py @@ -19,8 +19,6 @@ # Same cap as the onnx exports of these models and model2vec's default `max_length`. MAX_LENGTH = 512 -# Documents pooled per matrix product, which bounds the count matrix whatever the batch size. -POOLING_CHUNK = 256 supported_static_models: list[DenseModelDescription] = [ DenseModelDescription( @@ -94,6 +92,11 @@ class StaticTextEmbedding(TextEmbeddingBase): from `model.safetensors` and pooled without the padding an onnx graph needs. """ + # Cells of the (documents x distinct tokens) count matrix pooled per matrix product. It bounds + # the temporary memory of pooling to about 60 MiB, however long, diverse or many the documents + # are. + POOLING_CELLS = 4_000_000 + @classmethod def _list_supported_models(cls) -> list[DenseModelDescription]: return supported_static_models @@ -201,12 +204,7 @@ def raw_embed(self, documents: list[str]) -> NumpyArray: lengths = np.bincount(doc_index[known], minlength=len(documents)) ids = ids[known] - sums = np.zeros((len(documents), self.embeddings.shape[1]), dtype=np.float32) - ends = np.cumsum(lengths) - for start in range(0, len(documents), POOLING_CHUNK): - stop = min(start + POOLING_CHUNK, len(documents)) - first, last = ends[start] - lengths[start], ends[stop - 1] - sums[start:stop] = self._sum_rows(ids[first:last], lengths[start:stop]) + sums = self._sum_rows(ids, lengths) # The mean and the sum point the same way, so normalizing the sum gives the mean's # direction. Empty documents stay zero vectors. return normalize(sums) @@ -216,13 +214,23 @@ def _sum_rows(self, ids: NDArray[np.int64], lengths: NDArray[np.int64]) -> Numpy The sums are a product of a (documents x distinct tokens) count matrix with the rows of those tokens: one matrix multiplication, several times faster than segment sums with - np.add.reduceat or a loop over the documents. + np.add.reduceat or a loop over the documents. Documents whose count matrix would exceed + POOLING_CELLS are halved until it fits, which bounds the temporary memory. """ assert self.embeddings is not None if not ids.size: return np.zeros((len(lengths), self.embeddings.shape[1]), dtype=np.float32) rows = ids if self.token_mapping is None else self.token_mapping[ids] distinct_rows, row_index = np.unique(rows, return_inverse=True) + if len(lengths) > 1 and len(lengths) * len(distinct_rows) > self.POOLING_CELLS: + half = len(lengths) // 2 + split = int(lengths[:half].sum()) + return np.concatenate( + [ + self._sum_rows(ids[:split], lengths[:half]), + self._sum_rows(ids[split:], lengths[half:]), + ] + ) doc_index = np.repeat(np.arange(len(lengths)), lengths) counts = np.bincount( doc_index * len(distinct_rows) + row_index, diff --git a/tests/test_static_embeddings.py b/tests/test_static_embeddings.py index 235830e1..78a8d5c7 100644 --- a/tests/test_static_embeddings.py +++ b/tests/test_static_embeddings.py @@ -158,3 +158,18 @@ def test_default_options_do_not_warn() -> None: with warnings.catch_warnings(): warnings.simplefilter("error") TextEmbedding(CI_MODEL, lazy_load=True) + + +def test_pooling_budget_splits_without_changing_embeddings(model: TextEmbedding) -> None: + docs = [ + "", + "hello world", + "a much longer document about vectors " * 30, + "🤗", + "flag embedding", + ] * 9 + expected = np.stack(list(model.embed(docs))) + # a budget below one document's tokens halves every batch down to single documents + tiny_budget = TextEmbedding(CI_MODEL) + tiny_budget.model.POOLING_CELLS = 5 + assert np.allclose(np.stack(list(tiny_budget.embed(docs))), expected, atol=1e-6) From c855e543017e777f9abfe20cf86156ffdb72ccc0 Mon Sep 17 00:00:00 2001 From: George Panchuk Date: Tue, 13 Oct 2026 01:03:58 +0700 Subject: [PATCH 5/5] fix: Update token_count docstring in static embedding --- fastembed/text/static_embedding.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/fastembed/text/static_embedding.py b/fastembed/text/static_embedding.py index 11d093e1..241886aa 100644 --- a/fastembed/text/static_embedding.py +++ b/fastembed/text/static_embedding.py @@ -319,7 +319,7 @@ def _get_worker_class(cls) -> Type["StaticTextEmbeddingWorker"]: def token_count( self, texts: str | Iterable[str], batch_size: int = 1024, **kwargs: Any ) -> int: - """Returns the number of tokens that are averaged, after truncation.""" + """Returns the number of tokens after truncation, unknown tokens included.""" self._ensure_model() texts = [texts] if isinstance(texts, str) else texts return sum(int(self._token_ids(batch)[1].sum()) for batch in iter_batch(texts, batch_size))