diff --git a/fastembed/common/onnx_model.py b/fastembed/common/onnx_model.py index d357f2c..efde77f 100644 --- a/fastembed/common/onnx_model.py +++ b/fastembed/common/onnx_model.py @@ -31,6 +31,18 @@ class OnnxModel(Generic[T]): def _get_worker_class(cls) -> Type["EmbeddingWorker[T]"]: raise NotImplementedError("Subclasses must implement this method") + def _get_worker_init_kwargs(self) -> dict[str, Any]: + """Additional kwargs a worker process needs to reconstruct this model. + + Workers are started with `spawn`/`forkserver`, hence they don't inherit class-level state + which has been set up in runtime, e.g. models registered via `add_custom_model`. + Such state has to be shipped to the workers explicitly. + + Returns: + dict[str, Any]: kwargs to pass to `_get_worker_class().init_embedding`. + """ + return {} + def _post_process_onnx_output(self, output: OnnxOutputContext, **kwargs: Any) -> Iterable[T]: """Post-process the ONNX model output to convert it into a usable format. diff --git a/fastembed/rerank/cross_encoder/custom_text_cross_encoder.py b/fastembed/rerank/cross_encoder/custom_text_cross_encoder.py index fc1f6e9..c6e5263 100644 --- a/fastembed/rerank/cross_encoder/custom_text_cross_encoder.py +++ b/fastembed/rerank/cross_encoder/custom_text_cross_encoder.py @@ -1,9 +1,10 @@ -from typing import Sequence, Any +from typing import Sequence, Any, Type from fastembed.common import OnnxProvider from fastembed.common.model_description import BaseModelDescription from fastembed.common.types import Device from fastembed.rerank.cross_encoder.onnx_text_cross_encoder import OnnxTextCrossEncoder +from fastembed.rerank.cross_encoder.onnx_text_model import TextRerankerWorker class CustomTextCrossEncoder(OnnxTextCrossEncoder): @@ -39,9 +40,39 @@ class CustomTextCrossEncoder(OnnxTextCrossEncoder): def _list_supported_models(cls) -> list[BaseModelDescription]: return cls.SUPPORTED_MODELS + @classmethod + def _get_worker_class(cls) -> Type[TextRerankerWorker]: + return CustomTextCrossEncoderWorker + + def _get_worker_init_kwargs(self) -> dict[str, Any]: + return {"model_description": self.model_description} + @classmethod def add_model( cls, model_description: BaseModelDescription, ) -> None: cls.SUPPORTED_MODELS.append(model_description) + + +class CustomTextCrossEncoderWorker(TextRerankerWorker): + def init_embedding( + self, + model_name: str, + cache_dir: str, + model_description: BaseModelDescription | None = None, + **kwargs: Any, + ) -> CustomTextCrossEncoder: + if model_description is None: + raise ValueError( + "`model_description` is required to initialize a custom model in a worker " + "process, it is provided by `CustomTextCrossEncoder._get_worker_init_kwargs`" + ) + # custom models live in a class-level registry, which spawned workers don't inherit + CustomTextCrossEncoder.add_model(model_description) + return CustomTextCrossEncoder( + model_name=model_name, + cache_dir=cache_dir, + threads=1, + **kwargs, + ) diff --git a/fastembed/rerank/cross_encoder/onnx_text_model.py b/fastembed/rerank/cross_encoder/onnx_text_model.py index 55f3ea8..e32064d 100644 --- a/fastembed/rerank/cross_encoder/onnx_text_model.py +++ b/fastembed/rerank/cross_encoder/onnx_text_model.py @@ -128,6 +128,7 @@ class OnnxCrossEncoderModel(OnnxModel[float]): "local_files_only": local_files_only, "specific_model_path": specific_model_path, **kwargs, + **self._get_worker_init_kwargs(), } if extra_session_options is not None: diff --git a/fastembed/text/custom_text_embedding.py b/fastembed/text/custom_text_embedding.py index 512b914..c3d5cf5 100644 --- a/fastembed/text/custom_text_embedding.py +++ b/fastembed/text/custom_text_embedding.py @@ -1,4 +1,4 @@ -from typing import Sequence, Any, Iterable +from typing import Sequence, Any, Iterable, Type from dataclasses import dataclass import numpy as np @@ -13,6 +13,7 @@ from fastembed.common.onnx_model import OnnxOutputContext from fastembed.common.types import NumpyArray, Device from fastembed.common.utils import normalize, mean_pooling, last_token_pooling from fastembed.text.onnx_embedding import OnnxTextEmbedding +from fastembed.text.onnx_text_model import TextEmbeddingWorker @dataclass(frozen=True) @@ -58,6 +59,16 @@ class CustomTextEmbedding(OnnxTextEmbedding): def _list_supported_models(cls) -> list[DenseModelDescription]: return cls.SUPPORTED_MODELS + @classmethod + def _get_worker_class(cls) -> Type["TextEmbeddingWorker[NumpyArray]"]: + return CustomTextEmbeddingWorker + + def _get_worker_init_kwargs(self) -> dict[str, Any]: + return { + "model_description": self.model_description, + "postprocessing_config": self.POSTPROCESSING_MAPPING[self.model_description.model], + } + def _post_process_onnx_output( self, output: OnnxOutputContext, **kwargs: Any ) -> Iterable[NumpyArray]: @@ -102,3 +113,32 @@ class CustomTextEmbedding(OnnxTextEmbedding): cls.POSTPROCESSING_MAPPING[model_description.model] = PostprocessingConfig( pooling=pooling, normalization=normalization ) + + +class CustomTextEmbeddingWorker(TextEmbeddingWorker[NumpyArray]): + def init_embedding( + self, + model_name: str, + cache_dir: str, + model_description: DenseModelDescription | None = None, + postprocessing_config: PostprocessingConfig | None = None, + **kwargs: Any, + ) -> CustomTextEmbedding: + if model_description is None or postprocessing_config is None: + raise ValueError( + "`model_description` and `postprocessing_config` are required to initialize a " + "custom model in a worker process, they are provided by " + "`CustomTextEmbedding._get_worker_init_kwargs`" + ) + # custom models live in a class-level registry, which spawned workers don't inherit + CustomTextEmbedding.add_model( + model_description, + pooling=postprocessing_config.pooling, + normalization=postprocessing_config.normalization, + ) + return CustomTextEmbedding( + model_name=model_name, + cache_dir=cache_dir, + threads=1, + **kwargs, + ) diff --git a/fastembed/text/onnx_text_model.py b/fastembed/text/onnx_text_model.py index 10a4aa1..feeb70d 100644 --- a/fastembed/text/onnx_text_model.py +++ b/fastembed/text/onnx_text_model.py @@ -151,6 +151,7 @@ class OnnxTextModel(OnnxModel[T]): "local_files_only": local_files_only, "specific_model_path": specific_model_path, **kwargs, + **self._get_worker_init_kwargs(), } if extra_session_options is not None: diff --git a/tests/test_custom_models.py b/tests/test_custom_models.py index ab067ef..1cf47f7 100644 --- a/tests/test_custom_models.py +++ b/tests/test_custom_models.py @@ -77,6 +77,30 @@ def test_text_custom_model(): delete_model_cache(model.model._model_dir) +def test_text_custom_model_parallel_processing(): + is_ci = os.getenv("CI") + custom_model_name = "intfloat/multilingual-e5-small" + dim = 384 + + TextEmbedding.add_custom_model( + custom_model_name, + pooling=PoolingType.MEAN, + normalization=True, + sources=ModelSource(hf=custom_model_name), + dim=dim, + size_in_gb=0.47, + ) + + model = TextEmbedding(custom_model_name) + docs = ["hello world", "flag embedding"] * 50 + embeddings = np.stack(list(model.embed(docs, batch_size=10, parallel=2)), axis=0) + + assert embeddings.shape == (len(docs), dim) + + if is_ci: + delete_model_cache(model.model._model_dir) + + def test_cross_encoder_custom_model(): is_ci = os.getenv("CI") custom_model_name = "Xenova/ms-marco-MiniLM-L-4-v2" @@ -114,6 +138,27 @@ def test_cross_encoder_custom_model(): delete_model_cache(model.model._model_dir) +def test_cross_encoder_custom_model_parallel_processing(): + is_ci = os.getenv("CI") + custom_model_name = "Xenova/ms-marco-MiniLM-L-4-v2" + + TextCrossEncoder.add_custom_model( + custom_model_name, + model_file="onnx/model.onnx", + sources=ModelSource(hf=custom_model_name), + size_in_gb=0.08, + ) + + model = TextCrossEncoder(custom_model_name) + pairs = [("What is AI?", "Artificial intelligence is ...")] * 50 + scores = np.stack(list(model.rerank_pairs(pairs, batch_size=10, parallel=2)), axis=0) + + assert scores.shape == (len(pairs),) + + if is_ci: + delete_model_cache(model.model._model_dir) + + def test_mock_add_custom_models(): dim = 5 size_in_gb = 0.1