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, 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") + 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. + """ + + # 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 + + 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): 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 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) + 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_model(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 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) + + 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 = 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) + + 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. 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, + 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 + """ + 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.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 + + def token_count( + self, texts: str | Iterable[str], batch_size: int = 1024, **kwargs: Any + ) -> int: + """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)) + + +class StaticTextEmbeddingWorker(Worker): + 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": + 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.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/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..78a8d5c7 --- /dev/null +++ b/tests/test_static_embeddings.py @@ -0,0 +1,175 @@ +import json +import os +import warnings +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: + 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) + + +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) diff --git a/tests/test_text_onnx_embeddings.py b/tests/test_text_onnx_embeddings.py index 632bb70c..8232e23c 100644 --- a/tests/test_text_onnx_embeddings.py +++ b/tests/test_text_onnx_embeddings.py @@ -9,6 +9,7 @@ from fastembed.text.custom_text_embedding import CustomTextEmbedding from fastembed.text.last_token_normalized_embedding import LastTokenNormalizedEmbedding from fastembed.text.onnx_embedding import OnnxTextEmbedding +from fastembed.text.static_embedding import StaticTextEmbedding from fastembed.text.text_embedding import TextEmbedding from tests.utils import delete_model_cache, is_manual_run, should_test_model @@ -256,8 +257,12 @@ def test_quantized_model_reports_onnxruntime_requirement(monkeypatch) -> 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: