Repository: deepset-ai/haystack

Issue:
Add run_async for `HuggingFaceAPIDocumentEmbedder`



Discussion:

Patch:
diff --git a/haystack/components/embedders/hugging_face_api_document_embedder.py b/haystack/components/embedders/hugging_face_api_document_embedder.py
index 20e25a48a7..da43f0f5af 100644
--- a/haystack/components/embedders/hugging_face_api_document_embedder.py
+++ b/haystack/components/embedders/hugging_face_api_document_embedder.py
@@ -2,9 +2,10 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 from tqdm import tqdm
+from tqdm.asyncio import tqdm as async_tqdm
 
 from haystack import component, default_from_dict, default_to_dict, logging
 from haystack.dataclasses import Document
@@ -14,7 +15,7 @@
 from haystack.utils.url_validation import is_valid_http_url
 
 with LazyImport(message="Run 'pip install \"huggingface_hub>=0.27.0\"'") as huggingface_hub_import:
-    from huggingface_hub import InferenceClient
+    from huggingface_hub import AsyncInferenceClient, InferenceClient
 
 logger = logging.getLogger(__name__)
 
@@ -167,6 +168,8 @@ def __init__(
             msg = f"Unknown api_type {api_type}"
             raise ValueError(msg)
 
+        client_args: Dict[str, Any] = {"model": model_or_url, "token": token.resolve_value() if token else None}
+
         self.api_type = api_type
         self.api_params = api_params
         self.token = token
@@ -178,7 +181,8 @@ def __init__(
         self.progress_bar = progress_bar
         self.meta_fields_to_embed = meta_fields_to_embed or []
         self.embedding_separator = embedding_separator
-        self._client = InferenceClient(model_or_url, token=token.resolve_value() if token else None)
+        self._client = InferenceClient(**client_args)
+        self._async_client = AsyncInferenceClient(**client_args)
 
     def to_dict(self) -> Dict[str, Any]:
         """
@@ -232,14 +236,14 @@ def _prepare_texts_to_embed(self, documents: List[Document]) -> List[str]:
             texts_to_embed.append(text_to_embed)
         return texts_to_embed
 
-    def _embed_batch(self, texts_to_embed: List[str], batch_size: int) -> List[List[float]]:
+    @staticmethod
+    def _adjust_api_parameters(
+        truncate: Optional[bool], normalize: Optional[bool], api_type: HFEmbeddingAPIType
+    ) -> Tuple[Optional[bool], Optional[bool]]:
         """
-        Embed a list of texts in batches.
+        Adjust the truncate and normalize parameters based on the API type.
         """
-        truncate = self.truncate
-        normalize = self.normalize
-
-        if self.api_type == HFEmbeddingAPIType.SERVERLESS_INFERENCE_API:
+        if api_type == HFEmbeddingAPIType.SERVERLESS_INFERENCE_API:
             if truncate is not None:
                 msg = "`truncate` parameter is not supported for Serverless Inference API. It will be ignored."
                 logger.warning(msg)
@@ -248,6 +252,13 @@ def _embed_batch(self, texts_to_embed: List[str], batch_size: int) -> List[List[
                 msg = "`normalize` parameter is not supported for Serverless Inference API. It will be ignored."
                 logger.warning(msg)
                 normalize = None
+        return truncate, normalize
+
+    def _embed_batch(self, texts_to_embed: List[str], batch_size: int) -> List[List[float]]:
+        """
+        Embed a list of texts in batches.
+        """
+        truncate, normalize = self._adjust_api_parameters(self.truncate, self.normalize, self.api_type)
 
         all_embeddings: List = []
         for i in tqdm(
@@ -269,6 +280,32 @@ def _embed_batch(self, texts_to_embed: List[str], batch_size: int) -> List[List[
 
         return all_embeddings
 
+    async def _embed_batch_async(self, texts_to_embed: List[str], batch_size: int) -> List[List[float]]:
+        """
+        Embed a list of texts in batches asynchronously.
+        """
+        truncate, normalize = self._adjust_api_parameters(self.truncate, self.normalize, self.api_type)
+
+        all_embeddings: List = []
+        for i in async_tqdm(
+            range(0, len(texts_to_embed), batch_size), disable=not self.progress_bar, desc="Calculating embeddings"
+        ):
+            batch = texts_to_embed[i : i + batch_size]
+
+            np_embeddings = await self._async_client.feature_extraction(
+                # this method does not officially support list of strings, but works as expected
+                text=batch,  # type: ignore[arg-type]
+                truncate=truncate,
+                normalize=normalize,
+            )
+
+            if np_embeddings.ndim != 2 or np_embeddings.shape[0] != len(batch):
+                raise ValueError(f"Expected embedding shape ({batch_size}, embedding_dim), got {np_embeddings.shape}")
+
+            all_embeddings.extend(np_embeddings.tolist())
+
+        return all_embeddings
+
     @component.output_types(documents=List[Document])
     def run(self, documents: List[Document]):
         """
@@ -295,3 +332,30 @@ def run(self, documents: List[Document]):
             doc.embedding = emb
 
         return {"documents": documents}
+
+    @component.output_types(documents=List[Document])
+    async def run_async(self, documents: List[Document]):
+        """
+        Embeds a list of documents asynchronously.
+
+        :param documents:
+            Documents to embed.
+
+        :returns:
+            A dictionary with the following keys:
+            - `documents`: A list of documents with embeddings.
+        """
+        if not isinstance(documents, list) or documents and not isinstance(documents[0], Document):
+            raise TypeError(
+                "HuggingFaceAPIDocumentEmbedder expects a list of Documents as input."
+                " In case you want to embed a string, please use the HuggingFaceAPITextEmbedder."
+            )
+
+        texts_to_embed = self._prepare_texts_to_embed(documents=documents)
+
+        embeddings = await self._embed_batch_async(texts_to_embed=texts_to_embed, batch_size=self.batch_size)
+
+        for doc, emb in zip(documents, embeddings):
+            doc.embedding = emb
+
+        return {"documents": documents}
diff --git a/releasenotes/notes/add-async-support-for-HuggingFaceAPIDocumentEmbedder-016bbb1f8964cf3f.yaml b/releasenotes/notes/add-async-support-for-HuggingFaceAPIDocumentEmbedder-016bbb1f8964cf3f.yaml
new file mode 100644
index 0000000000..d8c9a0245a
--- /dev/null
+++ b/releasenotes/notes/add-async-support-for-HuggingFaceAPIDocumentEmbedder-016bbb1f8964cf3f.yaml
@@ -0,0 +1,5 @@
+---
+features:
+  - |
+    Add `run_async` method to HuggingFaceAPIDocumentEmbedder. This method enriches Documents with embeddings.
+    It supports the same parameters as the `run` method. It returns a coroutine that can be awaited.


Test patch:
diff --git a/test/components/embedders/test_hugging_face_api_document_embedder.py b/test/components/embedders/test_hugging_face_api_document_embedder.py
index 3ba40b1b23..9836635082 100644
--- a/test/components/embedders/test_hugging_face_api_document_embedder.py
+++ b/test/components/embedders/test_hugging_face_api_document_embedder.py
@@ -356,6 +356,19 @@ def test_run_custom_batch_size(self, mock_check_valid_model):
             assert len(doc.embedding) == 384
             assert all(isinstance(x, float) for x in doc.embedding)
 
+    def test_adjust_api_parameters(self):
+        truncate, normalize = HuggingFaceAPIDocumentEmbedder._adjust_api_parameters(
+            True, False, HFEmbeddingAPIType.SERVERLESS_INFERENCE_API
+        )
+        assert truncate is None
+        assert normalize is None
+
+        truncate, normalize = HuggingFaceAPIDocumentEmbedder._adjust_api_parameters(
+            True, False, HFEmbeddingAPIType.TEXT_EMBEDDINGS_INFERENCE
+        )
+        assert truncate is True
+        assert normalize is False
+
     @pytest.mark.flaky(reruns=5, reruns_delay=5)
     @pytest.mark.integration
     @pytest.mark.skipif(
@@ -384,3 +397,162 @@ def test_live_run_serverless(self):
             assert isinstance(doc.embedding, list)
             assert len(doc.embedding) == 384
             assert all(isinstance(x, float) for x in doc.embedding)
+
+    @pytest.mark.asyncio
+    async def test_embed_batch_async(self, mock_check_valid_model, caplog):
+        texts = ["text 1", "text 2", "text 3", "text 4", "text 5"]
+
+        with patch("huggingface_hub.AsyncInferenceClient.feature_extraction") as mock_embedding_patch:
+            mock_embedding_patch.side_effect = mock_embedding_generation
+
+            embedder = HuggingFaceAPIDocumentEmbedder(
+                api_type=HFEmbeddingAPIType.SERVERLESS_INFERENCE_API,
+                api_params={"model": "BAAI/bge-small-en-v1.5"},
+                token=Secret.from_token("fake-api-token"),
+            )
+            embeddings = await embedder._embed_batch_async(texts_to_embed=texts, batch_size=2)
+
+            assert mock_embedding_patch.call_count == 3
+
+        assert isinstance(embeddings, list)
+        assert len(embeddings) == len(texts)
+        for embedding in embeddings:
+            assert isinstance(embedding, list)
+            assert len(embedding) == 384
+            assert all(isinstance(x, float) for x in embedding)
+
+        # Check that logger warnings about ignoring truncate and normalize are raised
+        assert len(caplog.records) == 2
+        assert "truncate" in caplog.records[0].message
+        assert "normalize" in caplog.records[1].message
+
+    @pytest.mark.asyncio
+    async def test_embed_batch_async_wrong_embedding_shape(self, mock_check_valid_model):
+        texts = ["text 1", "text 2", "text 3", "text 4", "text 5"]
+
+        # embedding ndim != 2
+        with patch("huggingface_hub.AsyncInferenceClient.feature_extraction") as mock_embedding_patch:
+            mock_embedding_patch.return_value = array([0.1, 0.2, 0.3])
+
+            embedder = HuggingFaceAPIDocumentEmbedder(
+                api_type=HFEmbeddingAPIType.SERVERLESS_INFERENCE_API,
+                api_params={"model": "BAAI/bge-small-en-v1.5"},
+                token=Secret.from_token("fake-api-token"),
+            )
+
+            with pytest.raises(ValueError):
+                await embedder._embed_batch_async(texts_to_embed=texts, batch_size=2)
+
+        # embedding ndim == 2 but shape[0] != len(batch)
+        with patch("huggingface_hub.AsyncInferenceClient.feature_extraction") as mock_embedding_patch:
+            mock_embedding_patch.return_value = array([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6], [0.7, 0.8, 0.9]])
+
+            embedder = HuggingFaceAPIDocumentEmbedder(
+                api_type=HFEmbeddingAPIType.SERVERLESS_INFERENCE_API,
+                api_params={"model": "BAAI/bge-small-en-v1.5"},
+                token=Secret.from_token("fake-api-token"),
+            )
+
+            with pytest.raises(ValueError):
+                await embedder._embed_batch_async(texts_to_embed=texts, batch_size=2)
+
+    @pytest.mark.asyncio
+    async def test_run_async_wrong_input_format(self, mock_check_valid_model):
+        embedder = HuggingFaceAPIDocumentEmbedder(
+            api_type=HFEmbeddingAPIType.SERVERLESS_INFERENCE_API, api_params={"model": "BAAI/bge-small-en-v1.5"}
+        )
+
+        list_integers_input = [1, 2, 3]
+
+        with pytest.raises(TypeError):
+            await embedder.run_async(text=list_integers_input)
+
+    @pytest.mark.asyncio
+    async def test_run_async_on_empty_list(self, mock_check_valid_model):
+        embedder = HuggingFaceAPIDocumentEmbedder(
+            api_type=HFEmbeddingAPIType.SERVERLESS_INFERENCE_API,
+            api_params={"model": "BAAI/bge-small-en-v1.5"},
+            token=Secret.from_token("fake-api-token"),
+        )
+
+        empty_list_input = []
+        result = await embedder.run_async(documents=empty_list_input)
+
+        assert result["documents"] is not None
+        assert not result["documents"]  # empty list.
+
+    @pytest.mark.asyncio
+    async def test_run_async(self, mock_check_valid_model):
+        docs = [
+            Document(content="I love cheese", meta={"topic": "Cuisine"}),
+            Document(content="A transformer is a deep learning architecture", meta={"topic": "ML"}),
+        ]
+
+        with patch("huggingface_hub.AsyncInferenceClient.feature_extraction") as mock_embedding_patch:
+            mock_embedding_patch.side_effect = mock_embedding_generation
+
+            embedder = HuggingFaceAPIDocumentEmbedder(
+                api_type=HFEmbeddingAPIType.SERVERLESS_INFERENCE_API,
+                api_params={"model": "BAAI/bge-small-en-v1.5"},
+                token=Secret.from_token("fake-api-token"),
+                prefix="prefix ",
+                suffix=" suffix",
+                meta_fields_to_embed=["topic"],
+                embedding_separator=" | ",
+            )
+
+            result = await embedder.run_async(documents=docs)
+
+            mock_embedding_patch.assert_called_once_with(
+                text=[
+                    "prefix Cuisine | I love cheese suffix",
+                    "prefix ML | A transformer is a deep learning architecture suffix",
+                ],
+                truncate=None,
+                normalize=None,
+            )
+
+        documents_with_embeddings = result["documents"]
+
+        assert isinstance(documents_with_embeddings, list)
+        assert len(documents_with_embeddings) == len(docs)
+        for doc in documents_with_embeddings:
+            assert isinstance(doc, Document)
+            assert isinstance(doc.embedding, list)
+            assert len(doc.embedding) == 384
+            assert all(isinstance(x, float) for x in doc.embedding)
+
+    @pytest.mark.asyncio
+    async def test_run_async_custom_batch_size(self, mock_check_valid_model):
+        docs = [
+            Document(content="I love cheese", meta={"topic": "Cuisine"}),
+            Document(content="A transformer is a deep learning architecture", meta={"topic": "ML"}),
+        ]
+
+        with patch("huggingface_hub.AsyncInferenceClient.feature_extraction") as mock_embedding_patch:
+            mock_embedding_patch.side_effect = mock_embedding_generation
+
+            embedder = HuggingFaceAPIDocumentEmbedder(
+                api_type=HFEmbeddingAPIType.SERVERLESS_INFERENCE_API,
+                api_params={"model": "BAAI/bge-small-en-v1.5"},
+                token=Secret.from_token("fake-api-token"),
+                prefix="prefix ",
+                suffix=" suffix",
+                meta_fields_to_embed=["topic"],
+                embedding_separator=" | ",
+                batch_size=1,
+            )
+
+            result = await embedder.run_async(documents=docs)
+
+            assert mock_embedding_patch.call_count == 2
+
+        documents_with_embeddings = result["documents"]
+
+        assert isinstance(documents_with_embeddings, list)
+        assert len(documents_with_embeddings) == len(docs)
+        for doc in documents_with_embeddings:
+            assert isinstance(doc, Document)
+            assert isinstance(doc.embedding, list)
+            assert len(doc.embedding) == 384
+            assert all(isinstance(x, float) for x in doc.embedding)