diff --git a/integrations/fastembed/src/haystack_integrations/components/embedders/fastembed/fastembed_text_embedder.py b/integrations/fastembed/src/haystack_integrations/components/embedders/fastembed/fastembed_text_embedder.py index 455a1f94b..3f01752cd 100644 --- a/integrations/fastembed/src/haystack_integrations/components/embedders/fastembed/fastembed_text_embedder.py +++ b/integrations/fastembed/src/haystack_integrations/components/embedders/fastembed/fastembed_text_embedder.py @@ -35,6 +35,7 @@ def __init__( suffix: str = "", batch_size: int = 256, progress_bar: bool = True, + parallel: int = None, ): """ Create a FastembedTextEmbedder component. @@ -53,6 +54,7 @@ def __init__( self.suffix = suffix self.batch_size = batch_size self.progress_bar = progress_bar + self.parallel = parallel def to_dict(self) -> Dict[str, Any]: """ @@ -65,6 +67,7 @@ def to_dict(self) -> Dict[str, Any]: suffix=self.suffix, batch_size=self.batch_size, progress_bar=self.progress_bar, + parallel=self.parallel, ) def warm_up(self): @@ -93,6 +96,7 @@ def run(self, text: str): text_to_embed, batch_size=self.batch_size, show_progress_bar=self.progress_bar, + parallel=self.parallel, )[0] ) return {"embedding": embedding} diff --git a/integrations/fastembed/tests/test_fastembed_text_embedder.py b/integrations/fastembed/tests/test_fastembed_text_embedder.py index 3a7588263..42134a60e 100644 --- a/integrations/fastembed/tests/test_fastembed_text_embedder.py +++ b/integrations/fastembed/tests/test_fastembed_text_embedder.py @@ -19,6 +19,7 @@ def test_init_default(self): assert embedder.suffix == "" assert embedder.batch_size == 256 assert embedder.progress_bar is True + assert embedder.parallel is None def test_init_with_parameters(self): """ @@ -30,12 +31,14 @@ def test_init_with_parameters(self): suffix="suffix", batch_size=64, progress_bar=False, + parallel=1, ) assert embedder.model_name == "BAAI/bge-small-en-v1.5" assert embedder.prefix == "prefix" assert embedder.suffix == "suffix" assert embedder.batch_size == 64 assert embedder.progress_bar is False + assert embedder.parallel == 1 def test_to_dict(self): """ @@ -51,6 +54,7 @@ def test_to_dict(self): "suffix": "", "batch_size": 256, "progress_bar": True, + "parallel": None, }, } @@ -64,6 +68,7 @@ def test_to_dict_with_custom_init_parameters(self): suffix="suffix", batch_size=64, progress_bar=False, + parallel=1, ) embedder_dict = embedder.to_dict() assert embedder_dict == { @@ -74,6 +79,7 @@ def test_to_dict_with_custom_init_parameters(self): "suffix": "suffix", "batch_size": 64, "progress_bar": False, + "parallel": 1, }, } @@ -89,6 +95,7 @@ def test_from_dict(self): "suffix": "", "batch_size": 256, "progress_bar": True, + "parallel": None, }, } embedder = default_from_dict(FastembedTextEmbedder, embedder_dict) @@ -97,6 +104,7 @@ def test_from_dict(self): assert embedder.suffix == "" assert embedder.batch_size == 256 assert embedder.progress_bar is True + assert embedder.parallel is None def test_from_dict_with_custom_init_parameters(self): """ @@ -110,6 +118,7 @@ def test_from_dict_with_custom_init_parameters(self): "suffix": "suffix", "batch_size": 64, "progress_bar": False, + "parallel": 1, }, } embedder = default_from_dict(FastembedTextEmbedder, embedder_dict) @@ -118,6 +127,7 @@ def test_from_dict_with_custom_init_parameters(self): assert embedder.suffix == "suffix" assert embedder.batch_size == 64 assert embedder.progress_bar is False + assert embedder.parallel == 1 @patch( "haystack_integrations.components.embedders.fastembed.fastembed_text_embedder._FastembedEmbeddingBackendFactory"