fix: fix registering custom models in child workers (#714)

* fix: fix registering custom models in child workers

Co-authored-by: S0rryHorizon <151612757+S0rryHorizon@users.noreply.github.com>

* fix: fix mypy

---------

Co-authored-by: S0rryHorizon <151612757+S0rryHorizon@users.noreply.github.com>

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.
 

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,
+        )

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:

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,
+        )

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:

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