fix: prepare sparse post-processing for lazy parallel embedding

fastembed/sparse/bm42.py

@@ -8,15 +8,16 @@
 from py_rust_stemmers import SnowballStemmer
 
 from fastembed.common import OnnxProvider
+from fastembed.common.model_description import ModelSource, SparseModelDescription
 from fastembed.common.onnx_model import OnnxOutputContext
+from fastembed.common.preprocessor_utils import load_tokenizer
 from fastembed.common.types import Device
 from fastembed.common.utils import define_cache_dir
 from fastembed.sparse.sparse_embedding_base import (
     SparseEmbedding,
     SparseTextEmbeddingBase,
 )
 from fastembed.text.onnx_text_model import OnnxTextModel, TextEmbeddingWorker
-from fastembed.common.model_description import SparseModelDescription, ModelSource
 
 supported_bm42_models: list[SparseModelDescription] = [
     SparseModelDescription(
@@ -152,6 +153,12 @@ def load_onnx_model(self) -> None:
             extra_session_options=self._extra_session_options,
         )
 
+        self._prepare_post_processing()
+
+    def _prepare_post_processing(self) -> None:
+        if getattr(self, "tokenizer", None) is None:
+            self.tokenizer, self.special_token_to_id = load_tokenizer(model_dir=self._model_dir)
+
         for token, idx in self.tokenizer.get_vocab().items():  # type: ignore[union-attr]
             self.invert_vocab[idx] = token
         self.special_tokens = set(self.special_token_to_id.keys())

fastembed/sparse/minicoil.py

@@ -1,15 +1,15 @@
 from pathlib import Path
-
-from typing import Any, Sequence, Iterable, Type
+from typing import Any, Iterable, Sequence, Type
 
 import numpy as np
 from numpy.typing import NDArray
 from py_rust_stemmers import SnowballStemmer
 from tokenizers import Tokenizer
 
-from fastembed.common.model_description import SparseModelDescription, ModelSource
-from fastembed.common.onnx_model import OnnxOutputContext
 from fastembed.common import OnnxProvider
+from fastembed.common.model_description import ModelSource, SparseModelDescription
+from fastembed.common.onnx_model import OnnxOutputContext
+from fastembed.common.preprocessor_utils import load_tokenizer
 from fastembed.common.types import Device
 from fastembed.common.utils import define_cache_dir
 from fastembed.sparse.sparse_embedding_base import (
@@ -21,7 +21,6 @@
 from fastembed.sparse.utils.vocab_resolver import VocabResolver, VocabTokenizer
 from fastembed.text.onnx_text_model import OnnxTextModel, TextEmbeddingWorker
 
-
 MINICOIL_MODEL_FILE = "minicoil.triplet.model.npy"
 MINICOIL_VOCAB_FILE = "minicoil.triplet.model.vocab"
 STOPWORDS_FILE = "stopwords.txt"
@@ -160,6 +159,14 @@ def load_onnx_model(self) -> None:
             extra_session_options=self._extra_session_options,
         )
 
+        self._prepare_post_processing()
+
+    def _prepare_post_processing(self) -> None:
+        if self.vocab_resolver is not None:
+            return
+        if self.tokenizer is None:
+            self.tokenizer, self.special_token_to_id = load_tokenizer(model_dir=self._model_dir)
+
         assert self.tokenizer is not None
 
         for token, idx in self.tokenizer.get_vocab().items():  # type: ignore[union-attr]

fastembed/text/onnx_text_model.py

@@ -7,9 +7,9 @@
 from numpy.typing import NDArray
 from tokenizers import Encoding, Tokenizer
 
-from fastembed.common.types import NumpyArray, OnnxProvider, Device
 from fastembed.common.onnx_model import EmbeddingWorker, OnnxModel, OnnxOutputContext, T
 from fastembed.common.preprocessor_utils import load_tokenizer
+from fastembed.common.types import Device, NumpyArray, OnnxProvider
 from fastembed.common.utils import iter_batch
 from fastembed.parallel_processor import ParallelWorkerPool
 
@@ -70,6 +70,9 @@ def _load_onnx_model(
     def load_onnx_model(self) -> None:
         raise NotImplementedError("Subclasses must implement this method")
 
+    def _prepare_post_processing(self) -> None:
+        """Load model-specific assets needed to process worker outputs in the parent."""
+
     def tokenize(self, documents: list[str], **kwargs: Any) -> list[Encoding]:
         return self.tokenizer.encode_batch(documents)  # type: ignore[union-attr]
 
@@ -140,6 +143,7 @@ def _embed_documents(
                     self.onnx_embed(batch, **kwargs), **kwargs
                 )
         else:
+            self._prepare_post_processing()
             if parallel == 0:
                 parallel = os.cpu_count()
 

tests/test_sparse_embeddings.py

@@ -1,14 +1,13 @@
 import os
 from contextlib import contextmanager
 
-import pytest
 import numpy as np
+import pytest
 
 from fastembed.sparse.bm25 import Bm25
 from fastembed.sparse.sparse_text_embedding import SparseTextEmbedding
 from tests.utils import delete_model_cache, should_test_model
 
-
 CANONICAL_COLUMN_VALUES = {
     "prithivida/Splade_PP_en_v1": {
         "indices": [
@@ -243,6 +242,30 @@ def test_parallel_processing(model_cache, model_name: str) -> None:
             # assert np.allclose(sparse_embedding.values, sparse_embedding_all.values, atol=1e-3)
 
 
+@pytest.mark.parametrize(
+    "model_name",
+    ["Qdrant/bm42-all-minilm-l6-v2-attentions", "Qdrant/minicoil-v1"],
+)
+def test_lazy_parallel_sparse_post_processing(model_cache, model_name: str) -> None:
+    with model_cache(model_name) as eager_model:
+        documents = ["hello world", "flag embedding"] * 2
+        expected = list(eager_model.embed(documents, batch_size=2))
+        lazy_model = SparseTextEmbedding(
+            model_name=model_name,
+            specific_model_path=str(eager_model.model._model_dir),
+            lazy_load=True,
+        )
+        assert not hasattr(lazy_model.model, "model")
+
+        actual = list(lazy_model.embed(documents, batch_size=2, parallel=2))
+
+        assert not hasattr(lazy_model.model, "model")
+        assert len(actual) == len(expected)
+        for expected_embedding, actual_embedding in zip(expected, actual):
+            assert actual_embedding.indices.tolist() == expected_embedding.indices.tolist()
+            assert np.allclose(actual_embedding.values, expected_embedding.values, atol=1e-3)
+
+
 def test_stem_with_stopwords_and_punctuation(model_cache) -> None:
     with model_cache("Qdrant/bm25") as model:
         bm25_instance = model.model