From 9958675296e329310d94066ebca250e7ccdad108 Mon Sep 17 00:00:00 2001 From: George Panchuk Date: Tue, 13 Oct 2026 03:03:54 +0700 Subject: [PATCH] new: make memory-mapping of static embeddings optional --- fastembed/common/safetensors_utils.py | 24 +++++++++++++++++------- fastembed/text/static_embedding.py | 13 ++++++++++++- fastembed/text/text_embedding.py | 25 +++++++++++++++++++++++++ tests/test_static_embeddings.py | 17 +++++++++++++++-- 4 files changed, 69 insertions(+), 10 deletions(-) diff --git a/fastembed/common/safetensors_utils.py b/fastembed/common/safetensors_utils.py index a82f7dbb..bf2c9b27 100644 --- a/fastembed/common/safetensors_utils.py +++ b/fastembed/common/safetensors_utils.py @@ -19,12 +19,17 @@ } -def load_safetensors(path: str | Path) -> dict[str, NumpyArray]: +def load_safetensors(path: str | Path, mmap: bool = False) -> 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. + byte offsets of each tensor, then the raw tensor data. + + Args: + path: The safetensors file. + mmap: Memory-map the tensors read-only, so loading is instant, only the rows in use take + memory and parallel workers share the pages of one file. Otherwise the tensors are + read into memory, which avoids page faults and keeps no file open. """ path = Path(path) with open(path, "rb") as f: @@ -41,12 +46,17 @@ def load_safetensors(path: str | Path) -> dict[str, NumpyArray]: tensors[name] = np.zeros(shape, dtype=np.float32) continue if info["dtype"] == "BF16": - raw = np.memmap(path, dtype=" NumpyArray: + if mmap: + return np.memmap(path, dtype=dtype, mode="r", offset=offset, shape=shape) + count = int(np.prod(shape)) + return np.fromfile(path, dtype=dtype, count=count, offset=offset).reshape(shape) diff --git a/fastembed/text/static_embedding.py b/fastembed/text/static_embedding.py index 62699f05..853d8214 100644 --- a/fastembed/text/static_embedding.py +++ b/fastembed/text/static_embedding.py @@ -108,6 +108,7 @@ def __init__( threads: int | None = None, lazy_load: bool = False, specific_model_path: str | None = None, + mmap: bool = False, **kwargs: Any, ): """ @@ -123,6 +124,10 @@ def __init__( 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. + mmap (bool, optional): Whether to memory-map the embedding matrix instead of reading it + into memory. Memory-mapping loads instantly, takes memory only for the rows in use + and lets parallel workers share one copy; reading it into memory embeds a few + percent faster and keeps no file open. Defaults to False. """ # 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 @@ -150,6 +155,7 @@ def __init__( 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.mmap = mmap self._model_dir = self.download_model( self.model_description, self.cache_dir, @@ -166,7 +172,9 @@ def __init__( self.load_model() def load_model(self) -> None: - tensors = load_safetensors(self._model_dir / self.model_description.model_file) + tensors = load_safetensors( + self._model_dir / self.model_description.model_file, mmap=self.mmap + ) if "embeddings" not in tensors: raise ValueError(f"No `embeddings` tensor in {self._model_dir}") self.embeddings = tensors["embeddings"] @@ -276,6 +284,7 @@ def embed( parallel=parallel, local_files_only=self._local_files_only, specific_model_path=self._specific_model_path, + mmap=self.mmap, ) def _embed_documents( @@ -287,6 +296,7 @@ def _embed_documents( parallel: int | None = None, local_files_only: bool = False, specific_model_path: str | None = None, + mmap: bool = False, ) -> Iterable[NumpyArray]: is_small = False @@ -311,6 +321,7 @@ def _embed_documents( "cache_dir": cache_dir, "local_files_only": local_files_only, "specific_model_path": specific_model_path, + "mmap": mmap, } pool = ParallelWorkerPool( num_workers=parallel or 1, diff --git a/fastembed/text/text_embedding.py b/fastembed/text/text_embedding.py index 0468822c..0d4c6224 100644 --- a/fastembed/text/text_embedding.py +++ b/fastembed/text/text_embedding.py @@ -97,6 +97,31 @@ def __init__( lazy_load: bool = False, **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): The number of threads single onnxruntime session can use. + providers (Sequence[OnnxProvider], optional): The list of onnxruntime providers to use. + Mutually exclusive with the `cuda` and `device_ids` arguments. + cuda (Union[bool, Device], optional): Whether to use cuda for inference. + Mutually exclusive with `providers`. Defaults to Device.AUTO. + device_ids (list[int], optional): The list of device ids to use for data parallel + processing in workers. + lazy_load (bool, optional): Whether to load the model during class initialization or + on demand. Defaults to False. + **kwargs: Options passed on to the model class: + local_files_only (bool): Only use the files in the cache, never download them. + Defaults to False. + specific_model_path (str): A local directory with the model files to load + instead of downloading the model. + enable_cpu_mem_arena (bool): For onnx models, whether onnxruntime uses its CPU + memory arena. Defaults to onnxruntime's own default. + mmap (bool): For static models (minishlab potion), whether to memory-map the + embedding matrix instead of reading it into memory. Defaults to False. + """ super().__init__(model_name, cache_dir, threads, **kwargs) if model_name.lower() == "jinaai/jina-embeddings-v2-base-de": warnings.warn( diff --git a/tests/test_static_embeddings.py b/tests/test_static_embeddings.py index fb958c0e..7a3646b5 100644 --- a/tests/test_static_embeddings.py +++ b/tests/test_static_embeddings.py @@ -118,7 +118,8 @@ def _write_safetensors(path, tensors: dict[str, tuple[str, list[int], bytes]]) - f.write(data) -def test_load_safetensors(tmp_path) -> None: +@pytest.mark.parametrize("mmap", [True, False]) +def test_load_safetensors(tmp_path, mmap: bool) -> None: values = np.arange(6, dtype=np.float32).reshape(2, 3) / 4 bf16 = (values.view(np.uint32) >> 16).astype(" None: "mapping": ("I64", [3], np.array([2, 0, 1], dtype=" None: single = np.stack([next(iter(model.embed(doc))) for doc in docs]) assert np.allclose(single, batch, atol=1e-6) assert single.dtype == np.float32 + + +def test_mmap(model: TextEmbedding) -> None: + assert not isinstance(model.model.embeddings, np.memmap) # read into memory by default + docs = ["hello world", "", "a much longer document about vectors " * 30] * 40 + expected = np.stack(list(model.embed(docs, batch_size=16))) + mapped = TextEmbedding(CI_MODEL, mmap=True) + assert isinstance(mapped.model.embeddings, np.memmap) + assert np.array_equal(np.stack(list(mapped.embed(docs, batch_size=16))), expected) + parallel = np.stack(list(mapped.embed(docs, batch_size=16, parallel=2))) + assert np.allclose(parallel, expected, atol=1e-6)