fastembed-haystack 0.0.3__tar.gz → 0.0.4__tar.gz

{fastembed_haystack-0.0.3 → fastembed_haystack-0.0.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: fastembed-haystack
-Version: 0.0.3
+Version: 0.0.4
 Summary: Haystack 2.x component to embed strings and Documents using fastembed embedding model
 Project-URL: Source, https://github.com/deepset-ai/haystack-core-integrations
 Project-URL: Documentation, https://github.com/deepset-ai/haystack-core-integrations/blob/main/integrations/fastembed/README.md

fastembed_haystack-0.0.4/example/example.py

@@ -0,0 +1,33 @@
+from haystack import Document, Pipeline
+from haystack.components.retrievers.in_memory import InMemoryEmbeddingRetriever
+from haystack.document_stores.in_memory import InMemoryDocumentStore
+from haystack_integrations.components.embedders.fastembed import FastembedDocumentEmbedder, FastembedTextEmbedder
+
+document_store = InMemoryDocumentStore(embedding_similarity_function="cosine")
+
+documents = [
+    Document(content="My name is Wolfgang and I live in Berlin"),
+    Document(content="I saw a black horse running"),
+    Document(content="Germany has many big cities"),
+    Document(content="fastembed is supported by and maintained by Qdrant."),
+]
+
+document_embedder = FastembedDocumentEmbedder()
+document_embedder.warm_up()
+documents_with_embeddings = document_embedder.run(documents)["documents"]
+document_store.write_documents(documents_with_embeddings)
+
+query_pipeline = Pipeline()
+query_pipeline.add_component("text_embedder", FastembedTextEmbedder())
+query_pipeline.add_component("retriever", InMemoryEmbeddingRetriever(document_store=document_store))
+query_pipeline.connect("text_embedder.embedding", "retriever.query_embedding")
+
+query = "Who supports fastembed?"
+
+result = query_pipeline.run({"text_embedder": {"text": query}})
+
+print(result["retriever"]["documents"][0])  # noqa: T201
+
+# Document(id=...,
+#  content: 'fastembed is supported by and maintained by Qdrant.',
+#  score: 0.758..)

{fastembed_haystack-0.0.3 → fastembed_haystack-0.0.4}/src/haystack_integrations/components/embedders/fastembed/fastembed_document_embedder.py

@@ -57,6 +57,7 @@ class FastembedDocumentEmbedder:
         suffix: str = "",
         batch_size: int = 256,
         progress_bar: bool = True,
+        parallel: Optional[int] = None,
         meta_fields_to_embed: Optional[List[str]] = None,
         embedding_separator: str = "\n",
     ):
@@ -69,6 +70,10 @@ class FastembedDocumentEmbedder:
         :param suffix: A string to add to the end of each text.
         :param batch_size: Number of strings to encode at once.
         :param progress_bar: If true, displays progress bar during embedding.
+        :param parallel:
+                If > 1, data-parallel encoding will be used, recommended for offline encoding of large datasets.
+                If 0, use all available cores.
+                If None, don't use data-parallel processing, use default onnxruntime threading instead.
         :param meta_fields_to_embed: List of meta fields that should be embedded along with the Document content.
         :param embedding_separator: Separator used to concatenate the meta fields to the Document content.
         """
@@ -78,6 +83,7 @@ class FastembedDocumentEmbedder:
         self.suffix = suffix
         self.batch_size = batch_size
         self.progress_bar = progress_bar
+        self.parallel = parallel
         self.meta_fields_to_embed = meta_fields_to_embed or []
         self.embedding_separator = embedding_separator
 
@@ -92,6 +98,7 @@ class FastembedDocumentEmbedder:
             suffix=self.suffix,
             batch_size=self.batch_size,
             progress_bar=self.progress_bar,
+            parallel=self.parallel,
             meta_fields_to_embed=self.meta_fields_to_embed,
             embedding_separator=self.embedding_separator,
         )
@@ -139,6 +146,7 @@ class FastembedDocumentEmbedder:
             texts_to_embed,
             batch_size=self.batch_size,
             show_progress_bar=self.progress_bar,
+            parallel=self.parallel,
         )
 
         for doc, emb in zip(documents, embeddings):

{fastembed_haystack-0.0.3 → fastembed_haystack-0.0.4}/src/haystack_integrations/components/embedders/fastembed/fastembed_text_embedder.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Optional
 
 from haystack import component, default_to_dict
 
@@ -35,6 +35,7 @@ class FastembedTextEmbedder:
         suffix: str = "",
         batch_size: int = 256,
         progress_bar: bool = True,
+        parallel: Optional[int] = None,
     ):
         """
         Create a FastembedTextEmbedder component.
@@ -44,6 +45,11 @@ class FastembedTextEmbedder:
         :param batch_size: Number of strings to encode at once.
         :param prefix: A string to add to the beginning of each text.
         :param suffix: A string to add to the end of each text.
+        :param progress_bar: If true, displays progress bar during embedding.
+        :param parallel:
+                If > 1, data-parallel encoding will be used, recommended for offline encoding of large datasets.
+                If 0, use all available cores.
+                If None, don't use data-parallel processing, use default onnxruntime threading instead.
         """
 
         # TODO add parallel
@@ -53,6 +59,7 @@ class FastembedTextEmbedder:
         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 +72,7 @@ class FastembedTextEmbedder:
             suffix=self.suffix,
             batch_size=self.batch_size,
             progress_bar=self.progress_bar,
+            parallel=self.parallel,
         )
 
     def warm_up(self):
@@ -93,6 +101,7 @@ class FastembedTextEmbedder:
                 text_to_embed,
                 batch_size=self.batch_size,
                 show_progress_bar=self.progress_bar,
+                parallel=self.parallel,
             )[0]
         )
         return {"embedding": embedding}

{fastembed_haystack-0.0.3 → fastembed_haystack-0.0.4}/tests/test_fastembed_document_embedder.py

@@ -19,6 +19,7 @@ class TestFastembedDocumentEmbedder:
         assert embedder.suffix == ""
         assert embedder.batch_size == 256
         assert embedder.progress_bar is True
+        assert embedder.parallel is None
         assert embedder.meta_fields_to_embed == []
         assert embedder.embedding_separator == "\n"
 
@@ -32,6 +33,7 @@ class TestFastembedDocumentEmbedder:
             suffix="suffix",
             batch_size=64,
             progress_bar=False,
+            parallel=1,
             meta_fields_to_embed=["test_field"],
             embedding_separator=" | ",
         )
@@ -40,6 +42,7 @@ class TestFastembedDocumentEmbedder:
         assert embedder.suffix == "suffix"
         assert embedder.batch_size == 64
         assert embedder.progress_bar is False
+        assert embedder.parallel == 1
         assert embedder.meta_fields_to_embed == ["test_field"]
         assert embedder.embedding_separator == " | "
 
@@ -57,6 +60,7 @@ class TestFastembedDocumentEmbedder:
                 "suffix": "",
                 "batch_size": 256,
                 "progress_bar": True,
+                "parallel": None,
                 "embedding_separator": "\n",
                 "meta_fields_to_embed": [],
             },
@@ -72,6 +76,7 @@ class TestFastembedDocumentEmbedder:
             suffix="suffix",
             batch_size=64,
             progress_bar=False,
+            parallel=1,
             meta_fields_to_embed=["test_field"],
             embedding_separator=" | ",
         )
@@ -84,6 +89,7 @@ class TestFastembedDocumentEmbedder:
                 "suffix": "suffix",
                 "batch_size": 64,
                 "progress_bar": False,
+                "parallel": 1,
                 "meta_fields_to_embed": ["test_field"],
                 "embedding_separator": " | ",
             },
@@ -101,6 +107,7 @@ class TestFastembedDocumentEmbedder:
                 "suffix": "",
                 "batch_size": 256,
                 "progress_bar": True,
+                "parallel": None,
                 "meta_fields_to_embed": [],
                 "embedding_separator": "\n",
             },
@@ -111,6 +118,7 @@ class TestFastembedDocumentEmbedder:
         assert embedder.suffix == ""
         assert embedder.batch_size == 256
         assert embedder.progress_bar is True
+        assert embedder.parallel is None
         assert embedder.meta_fields_to_embed == []
         assert embedder.embedding_separator == "\n"
 
@@ -126,6 +134,7 @@ class TestFastembedDocumentEmbedder:
                 "suffix": "suffix",
                 "batch_size": 64,
                 "progress_bar": False,
+                "parallel": 1,
                 "meta_fields_to_embed": ["test_field"],
                 "embedding_separator": " | ",
             },
@@ -136,6 +145,7 @@ class TestFastembedDocumentEmbedder:
         assert embedder.suffix == "suffix"
         assert embedder.batch_size == 64
         assert embedder.progress_bar is False
+        assert embedder.parallel == 1
         assert embedder.meta_fields_to_embed == ["test_field"]
         assert embedder.embedding_separator == " | "
 
@@ -232,6 +242,7 @@ class TestFastembedDocumentEmbedder:
             ],
             batch_size=256,
             show_progress_bar=True,
+            parallel=None,
         )
 
     @pytest.mark.integration

{fastembed_haystack-0.0.3 → fastembed_haystack-0.0.4}/tests/test_fastembed_text_embedder.py

@@ -19,6 +19,7 @@ class TestFastembedTextEmbedder:
         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 @@ class TestFastembedTextEmbedder:
             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 @@ class TestFastembedTextEmbedder:
                 "suffix": "",
                 "batch_size": 256,
                 "progress_bar": True,
+                "parallel": None,
             },
         }
 
@@ -64,6 +68,7 @@ class TestFastembedTextEmbedder:
             suffix="suffix",
             batch_size=64,
             progress_bar=False,
+            parallel=1,
         )
         embedder_dict = embedder.to_dict()
         assert embedder_dict == {
@@ -74,6 +79,7 @@ class TestFastembedTextEmbedder:
                 "suffix": "suffix",
                 "batch_size": 64,
                 "progress_bar": False,
+                "parallel": 1,
             },
         }
 
@@ -89,6 +95,7 @@ class TestFastembedTextEmbedder:
                 "suffix": "",
                 "batch_size": 256,
                 "progress_bar": True,
+                "parallel": None,
             },
         }
         embedder = default_from_dict(FastembedTextEmbedder, embedder_dict)
@@ -97,6 +104,7 @@ class TestFastembedTextEmbedder:
         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 @@ class TestFastembedTextEmbedder:
                 "suffix": "suffix",
                 "batch_size": 64,
                 "progress_bar": False,
+                "parallel": 1,
             },
         }
         embedder = default_from_dict(FastembedTextEmbedder, embedder_dict)
@@ -118,6 +127,7 @@ class TestFastembedTextEmbedder:
         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"