huggingface-api-haystack 0.1.0__py3-none-any.whl

haystack_integrations/components/common/huggingface_api/__init__.py

@@ -0,0 +1,3 @@
+# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0

haystack_integrations/components/common/huggingface_api/utils.py

@@ -0,0 +1,112 @@
+# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from enum import Enum
+
+from haystack.utils import Secret
+from huggingface_hub import HfApi
+from huggingface_hub.errors import RepositoryNotFoundError
+
+
+class HFGenerationAPIType(Enum):
+    """
+    API type to use for Hugging Face API Generators.
+    """
+
+    # HF [Text Generation Inference (TGI)](https://github.com/huggingface/text-generation-inference).
+    TEXT_GENERATION_INFERENCE = "text_generation_inference"
+
+    # HF [Inference Endpoints](https://huggingface.co/inference-endpoints).
+    INFERENCE_ENDPOINTS = "inference_endpoints"
+
+    # HF [Serverless Inference API](https://huggingface.co/inference-api).
+    SERVERLESS_INFERENCE_API = "serverless_inference_api"
+
+    def __str__(self) -> str:
+        return self.value
+
+    @staticmethod
+    def from_str(string: str) -> "HFGenerationAPIType":
+        """
+        Convert a string to a HFGenerationAPIType enum.
+
+        :param string: The string to convert.
+        :return: The corresponding HFGenerationAPIType enum.
+        """
+        enum_map = {e.value: e for e in HFGenerationAPIType}
+        mode = enum_map.get(string)
+        if mode is None:
+            msg = f"Unknown Hugging Face API type '{string}'. Supported types are: {list(enum_map.keys())}"
+            raise ValueError(msg)
+        return mode
+
+
+class HFEmbeddingAPIType(Enum):
+    """
+    API type to use for Hugging Face API Embedders.
+    """
+
+    # HF [Text Embeddings Inference (TEI)](https://github.com/huggingface/text-embeddings-inference).
+    TEXT_EMBEDDINGS_INFERENCE = "text_embeddings_inference"
+
+    # HF [Inference Endpoints](https://huggingface.co/inference-endpoints).
+    INFERENCE_ENDPOINTS = "inference_endpoints"
+
+    # HF [Serverless Inference API](https://huggingface.co/inference-api).
+    SERVERLESS_INFERENCE_API = "serverless_inference_api"
+
+    def __str__(self) -> str:
+        return self.value
+
+    @staticmethod
+    def from_str(string: str) -> "HFEmbeddingAPIType":
+        """
+        Convert a string to a HFEmbeddingAPIType enum.
+
+        :param string: The string to convert.
+        :return: The corresponding HFEmbeddingAPIType enum.
+        """
+        enum_map = {e.value: e for e in HFEmbeddingAPIType}
+        mode = enum_map.get(string)
+        if mode is None:
+            msg = f"Unknown Hugging Face API type '{string}'. Supported types are: {list(enum_map.keys())}"
+            raise ValueError(msg)
+        return mode
+
+
+class HFModelType(Enum):
+    EMBEDDING = 1
+    GENERATION = 2
+
+
+def _check_valid_model(model_id: str, model_type: HFModelType, token: Secret | None) -> None:
+    """
+    Check if the provided model ID corresponds to a valid model on HuggingFace Hub.
+
+    Also check if the model is an embedding or generation model.
+
+    :param model_id: A string representing the HuggingFace model ID.
+    :param model_type: the model type, HFModelType.EMBEDDING or HFModelType.GENERATION
+    :param token: The optional authentication token.
+    :raises ValueError: If the model is not found or is not a embedding model.
+    """
+    api = HfApi()
+    try:
+        model_info = api.model_info(model_id, token=token.resolve_value() if token else None)
+    except RepositoryNotFoundError as e:
+        msg = f"Model {model_id} not found on HuggingFace Hub. Please provide a valid HuggingFace model_id."
+        raise ValueError(msg) from e
+
+    if model_type == HFModelType.EMBEDDING:
+        allowed_model = model_info.pipeline_tag in ["sentence-similarity", "feature-extraction"]
+        error_msg = f"Model {model_id} is not a embedding model. Please provide a embedding model."
+    elif model_type == HFModelType.GENERATION:
+        allowed_model = model_info.pipeline_tag in ["text-generation", "text2text-generation", "image-text-to-text"]
+        error_msg = f"Model {model_id} is not a text generation model. Please provide a text generation model."
+    else:
+        allowed_model = False
+        error_msg = f"Unknown model type for {model_id}"
+
+    if not allowed_model:
+        raise ValueError(error_msg)

haystack_integrations/components/embedders/huggingface_api/__init__.py

@@ -0,0 +1,7 @@
+# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0
+from .document_embedder import HuggingFaceAPIDocumentEmbedder
+from .text_embedder import HuggingFaceAPITextEmbedder
+
+__all__ = ["HuggingFaceAPIDocumentEmbedder", "HuggingFaceAPITextEmbedder"]

haystack_integrations/components/embedders/huggingface_api/document_embedder.py

@@ -0,0 +1,382 @@
+# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from asyncio import Semaphore, gather
+from dataclasses import replace
+from itertools import chain
+from typing import Any
+
+from haystack import component, default_from_dict, default_to_dict, logging
+from haystack.dataclasses import Document
+from haystack.utils import Secret
+from haystack.utils.url_validation import is_valid_http_url
+from huggingface_hub import AsyncInferenceClient, InferenceClient
+from tqdm import tqdm
+
+from haystack_integrations.components.common.huggingface_api.utils import (
+    HFEmbeddingAPIType,
+    HFModelType,
+    _check_valid_model,
+)
+
+logger = logging.getLogger(__name__)
+
+# batched embeddings returned by the API are expected to have shape (batch_size, embedding_dim)
+_EXPECTED_EMBEDDING_NDIM = 2
+
+
+@component
+class HuggingFaceAPIDocumentEmbedder:
+    """
+    Embeds documents using Hugging Face APIs.
+
+    Use it with the following Hugging Face APIs:
+    - [Free Serverless Inference API](https://huggingface.co/inference-api)
+    - [Paid Inference Endpoints](https://huggingface.co/inference-endpoints)
+    - [Self-hosted Text Embeddings Inference](https://github.com/huggingface/text-embeddings-inference)
+
+
+    ### Usage examples
+
+    #### With free serverless inference API
+
+    ```python
+    from haystack_integrations.components.embedders.huggingface_api import HuggingFaceAPIDocumentEmbedder
+    from haystack.utils import Secret
+    from haystack.dataclasses import Document
+
+    doc = Document(content="I love pizza!")
+
+    doc_embedder = HuggingFaceAPIDocumentEmbedder(api_type="serverless_inference_api",
+                                                  api_params={"model": "BAAI/bge-small-en-v1.5"},
+                                                  token=Secret.from_token("<your-api-key>"))
+
+    result = document_embedder.run([doc])
+    print(result["documents"][0].embedding)
+
+    # [0.017020374536514282, -0.023255806416273117, ...]
+    ```
+
+    #### With paid inference endpoints
+
+    ```python
+    from haystack_integrations.components.embedders.huggingface_api import HuggingFaceAPIDocumentEmbedder
+    from haystack.utils import Secret
+    from haystack.dataclasses import Document
+
+    doc = Document(content="I love pizza!")
+
+    doc_embedder = HuggingFaceAPIDocumentEmbedder(api_type="inference_endpoints",
+                                                  api_params={"url": "<your-inference-endpoint-url>"},
+                                                  token=Secret.from_token("<your-api-key>"))
+
+    result = document_embedder.run([doc])
+    print(result["documents"][0].embedding)
+
+    # [0.017020374536514282, -0.023255806416273117, ...]
+    ```
+
+    #### With self-hosted text embeddings inference
+
+    ```python
+    from haystack_integrations.components.embedders.huggingface_api import HuggingFaceAPIDocumentEmbedder
+    from haystack.dataclasses import Document
+
+    doc = Document(content="I love pizza!")
+
+    doc_embedder = HuggingFaceAPIDocumentEmbedder(api_type="text_embeddings_inference",
+                                                  api_params={"url": "http://localhost:8080"})
+
+    result = document_embedder.run([doc])
+    print(result["documents"][0].embedding)
+
+    # [0.017020374536514282, -0.023255806416273117, ...]
+    ```
+    """
+
+    def __init__(
+        self,
+        api_type: HFEmbeddingAPIType | str,
+        api_params: dict[str, str],
+        token: Secret | None = Secret.from_env_var(["HF_API_TOKEN", "HF_TOKEN"], strict=False),
+        prefix: str = "",
+        suffix: str = "",
+        truncate: bool | None = True,
+        normalize: bool | None = False,
+        batch_size: int = 32,
+        progress_bar: bool = True,
+        meta_fields_to_embed: list[str] | None = None,
+        embedding_separator: str = "\n",
+        concurrency_limit: int = 4,
+    ) -> None:
+        """
+        Creates a HuggingFaceAPIDocumentEmbedder component.
+
+        :param api_type:
+            The type of Hugging Face API to use.
+        :param api_params:
+            A dictionary with the following keys:
+            - `model`: Hugging Face model ID. Required when `api_type` is `SERVERLESS_INFERENCE_API`.
+            - `url`: URL of the inference endpoint. Required when `api_type` is `INFERENCE_ENDPOINTS` or
+            `TEXT_EMBEDDINGS_INFERENCE`.
+        :param token: The Hugging Face token to use as HTTP bearer authorization.
+            Check your HF token in your [account settings](https://huggingface.co/settings/tokens).
+        :param prefix:
+            A string to add at the beginning of each text.
+        :param suffix:
+            A string to add at the end of each text.
+        :param truncate:
+            Truncates the input text to the maximum length supported by the model.
+            Applicable when `api_type` is `TEXT_EMBEDDINGS_INFERENCE`, or `INFERENCE_ENDPOINTS`
+            if the backend uses Text Embeddings Inference.
+            If `api_type` is `SERVERLESS_INFERENCE_API`, this parameter is ignored.
+        :param normalize:
+            Normalizes the embeddings to unit length.
+            Applicable when `api_type` is `TEXT_EMBEDDINGS_INFERENCE`, or `INFERENCE_ENDPOINTS`
+            if the backend uses Text Embeddings Inference.
+            If `api_type` is `SERVERLESS_INFERENCE_API`, this parameter is ignored.
+        :param batch_size:
+            Number of documents to process at once.
+        :param progress_bar:
+            If `True`, shows a progress bar when running.
+        :param meta_fields_to_embed:
+            List of metadata fields to embed along with the document text.
+        :param embedding_separator:
+            Separator used to concatenate the metadata fields to the document text.
+        :param concurrency_limit:
+            The maximum number of requests that should be allowed to run concurrently.
+            This parameter is only used in the `run_async` method.
+        """
+        if isinstance(api_type, str):
+            api_type = HFEmbeddingAPIType.from_str(api_type)
+
+        api_params = api_params or {}
+
+        if api_type == HFEmbeddingAPIType.SERVERLESS_INFERENCE_API:
+            model = api_params.get("model")
+            if model is None:
+                msg = "To use the Serverless Inference API, you need to specify the `model` parameter in `api_params`."
+                raise ValueError(msg)
+            _check_valid_model(model, HFModelType.EMBEDDING, token)
+            model_or_url = model
+        elif api_type in [HFEmbeddingAPIType.INFERENCE_ENDPOINTS, HFEmbeddingAPIType.TEXT_EMBEDDINGS_INFERENCE]:
+            url = api_params.get("url")
+            if url is None:
+                msg = (
+                    "To use Text Embeddings Inference or Inference Endpoints, you need to specify the `url` "
+                    "parameter in `api_params`."
+                )
+                raise ValueError(msg)
+            if not is_valid_http_url(url):
+                msg = f"Invalid URL: {url}"
+                raise ValueError(msg)
+            model_or_url = url
+        else:
+            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
+        self.prefix = prefix
+        self.suffix = suffix
+        self.truncate = truncate
+        self.normalize = normalize
+        self.batch_size = batch_size
+        self.progress_bar = progress_bar
+        self.meta_fields_to_embed = meta_fields_to_embed or []
+        self.embedding_separator = embedding_separator
+        self.concurrency_limit = concurrency_limit
+        self._client = InferenceClient(**client_args)
+        self._async_client = AsyncInferenceClient(**client_args)
+
+    def to_dict(self) -> dict[str, Any]:
+        """
+        Serializes the component to a dictionary.
+
+        :returns:
+            Dictionary with serialized data.
+        """
+        return default_to_dict(
+            self,
+            api_type=str(self.api_type),
+            api_params=self.api_params,
+            prefix=self.prefix,
+            suffix=self.suffix,
+            token=self.token,
+            truncate=self.truncate,
+            normalize=self.normalize,
+            batch_size=self.batch_size,
+            progress_bar=self.progress_bar,
+            meta_fields_to_embed=self.meta_fields_to_embed,
+            embedding_separator=self.embedding_separator,
+            concurrency_limit=self.concurrency_limit,
+        )
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "HuggingFaceAPIDocumentEmbedder":
+        """
+        Deserializes the component from a dictionary.
+
+        :param data:
+            Dictionary to deserialize from.
+        :returns:
+            Deserialized component.
+        """
+        return default_from_dict(cls, data)
+
+    def _prepare_texts_to_embed(self, documents: list[Document]) -> list[str]:
+        """
+        Prepare the texts to embed by concatenating the Document text with the metadata fields to embed.
+        """
+        texts_to_embed = []
+        for doc in documents:
+            meta_values_to_embed = [
+                str(doc.meta[key]) for key in self.meta_fields_to_embed if key in doc.meta and doc.meta[key] is not None
+            ]
+
+            text_to_embed = (
+                self.prefix + self.embedding_separator.join([*meta_values_to_embed, doc.content or ""]) + self.suffix
+            )
+
+            texts_to_embed.append(text_to_embed)
+        return texts_to_embed
+
+    @staticmethod
+    def _adjust_api_parameters(
+        truncate: bool | None, normalize: bool | None, api_type: HFEmbeddingAPIType
+    ) -> tuple[bool | None, bool | None]:
+        """
+        Adjust the truncate and normalize parameters based on the API type.
+        """
+        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)
+                truncate = None
+            if normalize is not None:
+                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(
+            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 = self._client.feature_extraction(text=batch, truncate=truncate, normalize=normalize)
+
+            if np_embeddings.ndim != _EXPECTED_EMBEDDING_NDIM or np_embeddings.shape[0] != len(batch):
+                msg = f"Expected embedding shape ({batch_size}, embedding_dim), got {np_embeddings.shape}"
+                raise ValueError(msg)
+
+            all_embeddings.extend(np_embeddings.tolist())
+
+        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)
+        sem = Semaphore(max(1, self.concurrency_limit))
+        num_batches = (len(texts_to_embed) + batch_size - 1) // batch_size
+        pbar = tqdm(total=num_batches, disable=not self.progress_bar, desc="Calculating embeddings")
+
+        async def _runner(batch: list[str]) -> list[list[float]]:
+            async with sem:
+                np_embeddings = await self._async_client.feature_extraction(
+                    text=batch, truncate=truncate, normalize=normalize
+                )
+
+                if np_embeddings.ndim != _EXPECTED_EMBEDDING_NDIM or np_embeddings.shape[0] != len(batch):
+                    msg = f"Expected embedding shape ({batch_size}, embedding_dim), got {np_embeddings.shape}"
+                    raise ValueError(msg)
+
+                pbar.update(1)
+                return np_embeddings.tolist()
+
+        try:
+            all_embeddings = [
+                *chain(
+                    *await gather(
+                        *[
+                            _runner(texts_to_embed[i : i + batch_size])
+                            for i in range(0, len(texts_to_embed), batch_size)
+                        ]
+                    )
+                )
+            ]
+        finally:
+            pbar.close()
+
+        return all_embeddings
+
+    @component.output_types(documents=list[Document])
+    def run(self, documents: list[Document]) -> dict[str, list[Document]]:
+        """
+        Embeds a list of documents.
+
+        :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)):
+            msg = (
+                "HuggingFaceAPIDocumentEmbedder expects a list of Documents as input."
+                " In case you want to embed a string, please use the HuggingFaceAPITextEmbedder."
+            )
+            raise TypeError(msg)
+
+        texts_to_embed = self._prepare_texts_to_embed(documents=documents)
+
+        embeddings = self._embed_batch(texts_to_embed=texts_to_embed, batch_size=self.batch_size)
+
+        new_documents = []
+        for doc, emb in zip(documents, embeddings, strict=True):
+            new_documents.append(replace(doc, embedding=emb))
+
+        return {"documents": new_documents}
+
+    @component.output_types(documents=list[Document])
+    async def run_async(self, documents: list[Document]) -> dict[str, 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)):
+            msg = (
+                "HuggingFaceAPIDocumentEmbedder expects a list of Documents as input."
+                " In case you want to embed a string, please use the HuggingFaceAPITextEmbedder."
+            )
+            raise TypeError(msg)
+
+        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)
+
+        new_documents = []
+        for doc, emb in zip(documents, embeddings, strict=True):
+            new_documents.append(replace(doc, embedding=emb))
+
+        return {"documents": new_documents}