diff --git a/fastembed/sparse/bm42.py b/fastembed/sparse/bm42.py index 2b090f74..14c93653 100644 --- a/fastembed/sparse/bm42.py +++ b/fastembed/sparse/bm42.py @@ -8,7 +8,9 @@ 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 ( @@ -16,7 +18,6 @@ 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()) diff --git a/fastembed/sparse/minicoil.py b/fastembed/sparse/minicoil.py index 6f29c60f..ccbb3b57 100644 --- a/fastembed/sparse/minicoil.py +++ b/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] diff --git a/fastembed/text/onnx_text_model.py b/fastembed/text/onnx_text_model.py index feeb70d5..e610a545 100644 --- a/fastembed/text/onnx_text_model.py +++ b/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() diff --git a/tests/test_sparse_embeddings.py b/tests/test_sparse_embeddings.py index 9fdbb84d..f2f57f4f 100644 --- a/tests/test_sparse_embeddings.py +++ b/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