new: make memory-mapping of static embeddings optional

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="<u2", mode="r", offset=data_start + start, shape=shape)
+            raw = _read(path, "<u2", data_start + start, shape, mmap)
             tensors[name] = (raw.astype(np.uint32) << 16).view(np.float32)
             continue
         if info["dtype"] not in _DTYPES:
             raise ValueError(f"Unsupported safetensors dtype {info['dtype']} for {name} in {path}")
-        tensors[name] = np.memmap(
-            path, dtype=_DTYPES[info["dtype"]], mode="r", offset=data_start + start, shape=shape
-        )
+        tensors[name] = _read(path, _DTYPES[info["dtype"]], data_start + start, shape, mmap)
     return tensors
+
+
+def _read(path: Path, dtype: str, offset: int, shape: tuple[int, ...], mmap: bool) -> 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)

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,

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(

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("<u2")
     path = tmp_path / "model.safetensors"
@@ -131,7 +132,8 @@ def test_load_safetensors(tmp_path) -> None:
             "mapping": ("I64", [3], np.array([2, 0, 1], dtype="<i8").tobytes()),
         },
     )
-    tensors = load_safetensors(path)
+    tensors = load_safetensors(path, mmap=mmap)
+    assert isinstance(tensors["f32"], np.memmap) == mmap
     assert set(tensors) == {"f32", "f16", "bf16", "mapping"}
     for name in ("f32", "f16", "bf16"):
         assert np.array_equal(np.asarray(tensors[name], dtype=np.float32), values), name
@@ -187,3 +189,14 @@ def test_single_text_matches_batch(model: TextEmbedding) -> 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)