huggingface-api-haystack 0.6.0__tar.gz → 2.0.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (41) hide show
  1. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-2.0.0}/CHANGELOG.md +35 -1
  2. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-2.0.0}/PKG-INFO +6 -5
  3. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-2.0.0}/README.md +1 -2
  4. huggingface_api_haystack-2.0.0/docker-compose.yml +50 -0
  5. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-2.0.0}/pyproject.toml +11 -1
  6. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-2.0.0}/src/haystack_integrations/common/huggingface_api/utils.py +66 -8
  7. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-2.0.0}/src/haystack_integrations/components/embedders/huggingface_api/__init__.py +4 -4
  8. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-2.0.0}/src/haystack_integrations/components/embedders/huggingface_api/document_embedder.py +167 -17
  9. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-2.0.0}/src/haystack_integrations/components/embedders/huggingface_api/sparse_document_embedder.py +147 -20
  10. huggingface_api_haystack-2.0.0/src/haystack_integrations/components/embedders/huggingface_api/sparse_text_embedder.py +191 -0
  11. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-2.0.0}/src/haystack_integrations/components/embedders/huggingface_api/text_embedder.py +102 -9
  12. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-2.0.0}/src/haystack_integrations/components/generators/huggingface_api/chat/chat_generator.py +55 -22
  13. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-2.0.0}/src/haystack_integrations/components/rankers/huggingface_api/ranker.py +76 -3
  14. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-2.0.0}/tests/test_chat_generator.py +199 -82
  15. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-2.0.0}/tests/test_document_embedder.py +319 -26
  16. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-2.0.0}/tests/test_ranker.py +174 -3
  17. huggingface_api_haystack-2.0.0/tests/test_sparse_document_embedder.py +621 -0
  18. huggingface_api_haystack-2.0.0/tests/test_sparse_text_embedder.py +499 -0
  19. huggingface_api_haystack-2.0.0/tests/test_text_embedder.py +502 -0
  20. huggingface_api_haystack-2.0.0/tests/test_utils.py +128 -0
  21. huggingface_api_haystack-0.6.0/docker-compose.yml +0 -14
  22. huggingface_api_haystack-0.6.0/src/haystack_integrations/components/embedders/huggingface_api/sparse_text_embedder.py +0 -116
  23. huggingface_api_haystack-0.6.0/tests/test_sparse_document_embedder.py +0 -342
  24. huggingface_api_haystack-0.6.0/tests/test_sparse_text_embedder.py +0 -263
  25. huggingface_api_haystack-0.6.0/tests/test_text_embedder.py +0 -257
  26. huggingface_api_haystack-0.6.0/tests/test_utils.py +0 -49
  27. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-2.0.0}/.gitignore +0 -0
  28. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-2.0.0}/LICENSE.txt +0 -0
  29. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-2.0.0}/pydoc/config_docusaurus.yml +0 -0
  30. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-2.0.0}/src/haystack_integrations/common/huggingface_api/__init__.py +0 -0
  31. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-2.0.0}/src/haystack_integrations/common/py.typed +0 -0
  32. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-2.0.0}/src/haystack_integrations/components/embedders/huggingface_api/sparse_embedding_utils.py +0 -0
  33. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-2.0.0}/src/haystack_integrations/components/embedders/py.typed +0 -0
  34. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-2.0.0}/src/haystack_integrations/components/generators/huggingface_api/__init__.py +0 -0
  35. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-2.0.0}/src/haystack_integrations/components/generators/huggingface_api/chat/__init__.py +0 -0
  36. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-2.0.0}/src/haystack_integrations/components/generators/py.typed +0 -0
  37. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-2.0.0}/src/haystack_integrations/components/rankers/huggingface_api/__init__.py +0 -0
  38. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-2.0.0}/src/haystack_integrations/components/rankers/py.typed +0 -0
  39. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-2.0.0}/tests/__init__.py +0 -0
  40. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-2.0.0}/tests/conftest.py +0 -0
  41. {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-2.0.0}/tests/test_files/apple.jpg +0 -0
@@ -1,6 +1,40 @@
1
1
  # Changelog
2
2
 
3
- ## [unreleased]
3
+ ## [integrations/huggingface_api-v1.0.0] - 2026-09-07
4
+
5
+ ### 🚀 Features
6
+
7
+ - [**breaking**] HF API - add lifecycle handling (#3917)
8
+
9
+ ### 🧪 Testing
10
+
11
+ - Replace retired Qwen2.5 models in `huggingface_api` live tests (#3866)
12
+
13
+ ### ⚙️ CI
14
+
15
+ - Improve changelog generation; fix existing changelogs (#3883)
16
+
17
+
18
+ ## [integrations/huggingface_api-v0.6.0] - 2026-08-27
19
+
20
+ ### 🚀 Features
21
+
22
+ - *(huggingface_api)* Add SPLADE sparse embedder components (#3835)
23
+
24
+ ## [integrations/huggingface_api-v0.5.0] - 2026-08-24
25
+
26
+ ### 🐛 Bug Fixes
27
+
28
+ - Fix new issues raised by ruff 0.16.0 (#3670)
29
+ - Serialize missing init params in to_dict (transformers, huggingface_api, ragas) (#3808)
30
+
31
+ ### 🧹 Chores
32
+
33
+ - Re-standardize READMEs (#3768)
34
+ - Clarify how generation_kwargs passed in run are handled (#3805)
35
+
36
+
37
+ ## [integrations/huggingface_api-v0.4.0] - 2026-07-09
4
38
 
5
39
  ### 🧪 Testing
6
40
 
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.5
2
2
  Name: huggingface-api-haystack
3
- Version: 0.6.0
3
+ Version: 2.0.0
4
4
  Summary: Haystack integration for Hugging Face API
5
5
  Project-URL: Documentation, https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/huggingface_api#readme
6
6
  Project-URL: Issues, https://github.com/deepset-ai/haystack-core-integrations/issues
@@ -19,9 +19,11 @@ Classifier: Programming Language :: Python :: 3.14
19
19
  Classifier: Programming Language :: Python :: Implementation :: CPython
20
20
  Classifier: Programming Language :: Python :: Implementation :: PyPy
21
21
  Requires-Python: >=3.10
22
- Requires-Dist: haystack-ai>=2.30.0
22
+ Requires-Dist: haystack-ai>=3.0.0
23
23
  Requires-Dist: httpx>=0.27.0
24
- Requires-Dist: huggingface-hub[inference]>=0.35.0
24
+ Requires-Dist: huggingface-hub[inference]>=1.0.0
25
+ Provides-Extra: grpc
26
+ Requires-Dist: grpc-requests>=0.1.21; extra == 'grpc'
25
27
  Description-Content-Type: text/markdown
26
28
 
27
29
  # huggingface-api-haystack
@@ -38,7 +40,6 @@ Description-Content-Type: text/markdown
38
40
 
39
41
  Refer to the general [Contribution Guidelines](https://github.com/deepset-ai/haystack-core-integrations/blob/main/CONTRIBUTING.md).
40
42
 
41
- To run integration tests locally, you need a Docker container running a Text Embeddings Inference (TEI) server with
42
- a sparse embedding model. Use the provided `docker-compose.yml` file to start the container: `docker compose up -d`.
43
+ To run integration tests locally, you need Docker containers running a few Text Embeddings Inference (TEI) servers. Use the provided `docker-compose.yml` file to start the containers: `docker compose up -d`.
43
44
 
44
45
  The tests that call the Hugging Face Inference API are skipped unless `HF_TOKEN` is set as an evironment variable.
@@ -12,7 +12,6 @@
12
12
 
13
13
  Refer to the general [Contribution Guidelines](https://github.com/deepset-ai/haystack-core-integrations/blob/main/CONTRIBUTING.md).
14
14
 
15
- To run integration tests locally, you need a Docker container running a Text Embeddings Inference (TEI) server with
16
- a sparse embedding model. Use the provided `docker-compose.yml` file to start the container: `docker compose up -d`.
15
+ To run integration tests locally, you need Docker containers running a few Text Embeddings Inference (TEI) servers. Use the provided `docker-compose.yml` file to start the containers: `docker compose up -d`.
17
16
 
18
17
  The tests that call the Hugging Face Inference API are skipped unless `HF_TOKEN` is set as an evironment variable.
@@ -0,0 +1,50 @@
1
+ services:
2
+ tei-sparse:
3
+ image: ghcr.io/huggingface/text-embeddings-inference:cpu-1.9.3
4
+ command: ["--model-id", "opensearch-project/opensearch-neural-sparse-encoding-doc-v2-mini", "--pooling", "splade"]
5
+ ports:
6
+ - "8080:80"
7
+ environment:
8
+ - HF_TOKEN
9
+ volumes:
10
+ - tei-data:/data
11
+ tei-grpc:
12
+ image: ghcr.io/huggingface/text-embeddings-inference:cpu-1.9.3-grpc
13
+ command: ["--model-id", "sentence-transformers/all-MiniLM-L6-v2"]
14
+ ports:
15
+ - "8081:80"
16
+ environment:
17
+ - HF_TOKEN
18
+ volumes:
19
+ - tei-data:/data
20
+ tei-sparse-grpc:
21
+ image: ghcr.io/huggingface/text-embeddings-inference:cpu-1.9.3-grpc
22
+ command: ["--model-id", "opensearch-project/opensearch-neural-sparse-encoding-doc-v2-mini", "--pooling", "splade"]
23
+ ports:
24
+ - "8082:80"
25
+ environment:
26
+ - HF_TOKEN
27
+ volumes:
28
+ - tei-data:/data
29
+
30
+ tei-ranker:
31
+ image: ghcr.io/huggingface/text-embeddings-inference:cpu-1.9.3
32
+ command: ["--model-id", "cross-encoder/ms-marco-TinyBERT-L2-v2", "--api-key", "tei-test-key"]
33
+ ports:
34
+ - "8083:80"
35
+ environment:
36
+ - HF_TOKEN
37
+ volumes:
38
+ - tei-data:/data
39
+ tei-ranker-grpc:
40
+ image: ghcr.io/huggingface/text-embeddings-inference:cpu-1.9.3-grpc
41
+ command: ["--model-id", "cross-encoder/ms-marco-TinyBERT-L2-v2", "--api-key", "tei-test-key"]
42
+ ports:
43
+ - "8084:80"
44
+ environment:
45
+ - HF_TOKEN
46
+ volumes:
47
+ - tei-data:/data
48
+
49
+ volumes:
50
+ tei-data:
@@ -23,7 +23,12 @@ classifiers = [
23
23
  "Programming Language :: Python :: Implementation :: CPython",
24
24
  "Programming Language :: Python :: Implementation :: PyPy",
25
25
  ]
26
- dependencies = ["haystack-ai>=2.30.0", "httpx>=0.27.0", "huggingface_hub[inference]>=0.35.0"]
26
+ dependencies = ["haystack-ai>=3.0.0", "httpx>=0.27.0", "huggingface_hub[inference]>=1.0.0"]
27
+
28
+ [project.optional-dependencies]
29
+ grpc = [
30
+ "grpc-requests>=0.1.21",
31
+ ]
27
32
 
28
33
  [project.urls]
29
34
  Documentation = "https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/huggingface_api#readme"
@@ -58,6 +63,7 @@ dependencies = [
58
63
  "pytest-rerunfailures",
59
64
  "mypy",
60
65
  "pip",
66
+ "grpc-requests>=0.1.21",
61
67
  "numpy",
62
68
  "pillow", # image resizing in multimodal tests
63
69
  ]
@@ -79,6 +85,10 @@ non_interactive = true
79
85
  check_untyped_defs = true
80
86
  disallow_incomplete_defs = true
81
87
 
88
+ [[tool.mypy.overrides]]
89
+ module = ["grpc_requests", "grpc_requests.*"]
90
+ ignore_missing_imports = true
91
+
82
92
  [tool.ruff]
83
93
  line-length = 120
84
94
 
@@ -4,9 +4,11 @@
4
4
 
5
5
  from enum import Enum
6
6
 
7
+ from haystack.dataclasses import SparseEmbedding
7
8
  from haystack.utils import Secret
8
9
  from huggingface_hub import HfApi
9
10
  from huggingface_hub.errors import RepositoryNotFoundError
11
+ from huggingface_hub.utils import build_hf_headers, get_async_session, hf_raise_for_status
10
12
 
11
13
 
12
14
  class HFGenerationAPIType(Enum):
@@ -82,14 +84,15 @@ class HFModelType(Enum):
82
84
 
83
85
  def _check_valid_model(model_id: str, model_type: HFModelType, token: Secret | None) -> None:
84
86
  """
85
- Check if the provided model ID corresponds to a valid model on HuggingFace Hub.
87
+ Check if the provided model ID corresponds to a valid model on Hugging Face Hub.
86
88
 
87
- Also check if the model is an embedding or generation model.
89
+ Also check if the model supports the requested embedding or generation task.
88
90
 
89
- :param model_id: A string representing the HuggingFace model ID.
90
- :param model_type: the model type, HFModelType.EMBEDDING or HFModelType.GENERATION
91
- :param token: The optional authentication token.
92
- :raises ValueError: If the model is not found or is not a embedding model.
91
+ :param model_id: A string representing the Hugging Face model ID.
92
+ :param model_type: The model type, `HFModelType.EMBEDDING` or `HFModelType.GENERATION`.
93
+ :param token: The optional authentication secret.
94
+ :raises ValueError: If the secret cannot be resolved, the model is not found, or it does not match the
95
+ requested type.
93
96
  """
94
97
  api = HfApi()
95
98
  try:
@@ -98,11 +101,42 @@ def _check_valid_model(model_id: str, model_type: HFModelType, token: Secret | N
98
101
  msg = f"Model {model_id} not found on HuggingFace Hub. Please provide a valid HuggingFace model_id."
99
102
  raise ValueError(msg) from e
100
103
 
104
+ _validate_model_type(model_id, model_type, model_info.pipeline_tag)
105
+
106
+
107
+ async def _check_valid_model_async(model_id: str, model_type: HFModelType, token: Secret | None) -> None:
108
+ """
109
+ Asynchronously check if the provided model ID corresponds to a valid model on Hugging Face Hub.
110
+
111
+ Also check if the model supports the requested embedding or generation task.
112
+
113
+ :param model_id: A string representing the Hugging Face model ID.
114
+ :param model_type: The model type, `HFModelType.EMBEDDING` or `HFModelType.GENERATION`.
115
+ :param token: The optional authentication secret.
116
+ :raises ValueError: If the secret cannot be resolved, the model is not found, or it does not match the
117
+ requested type.
118
+ """
119
+ api = HfApi()
120
+ try:
121
+ async with get_async_session() as client:
122
+ response = await client.get(
123
+ f"{api.endpoint}/api/models/{model_id}",
124
+ headers=build_hf_headers(token=token.resolve_value() if token else None),
125
+ )
126
+ hf_raise_for_status(response)
127
+ except RepositoryNotFoundError as e:
128
+ msg = f"Model {model_id} not found on HuggingFace Hub. Please provide a valid HuggingFace model_id."
129
+ raise ValueError(msg) from e
130
+
131
+ _validate_model_type(model_id, model_type, response.json().get("pipeline_tag"))
132
+
133
+
134
+ def _validate_model_type(model_id: str, model_type: HFModelType, pipeline_tag: str | None) -> None:
101
135
  if model_type == HFModelType.EMBEDDING:
102
- allowed_model = model_info.pipeline_tag in ["sentence-similarity", "feature-extraction"]
136
+ allowed_model = pipeline_tag in ["sentence-similarity", "feature-extraction"]
103
137
  error_msg = f"Model {model_id} is not a embedding model. Please provide a embedding model."
104
138
  elif model_type == HFModelType.GENERATION:
105
- allowed_model = model_info.pipeline_tag in ["text-generation", "text2text-generation", "image-text-to-text"]
139
+ allowed_model = pipeline_tag in ["text-generation", "text2text-generation", "image-text-to-text"]
106
140
  error_msg = f"Model {model_id} is not a text generation model. Please provide a text generation model."
107
141
  else:
108
142
  allowed_model = False
@@ -110,3 +144,27 @@ def _check_valid_model(model_id: str, model_type: HFModelType, token: Secret | N
110
144
 
111
145
  if not allowed_model:
112
146
  raise ValueError(error_msg)
147
+
148
+
149
+ def _parse_grpc_sparse_embedding(items: list[dict[str, int | float]]) -> SparseEmbedding:
150
+ """Convert sparse gRPC values, restoring defaults omitted by protobuf."""
151
+ return SparseEmbedding(
152
+ indices=[int(item.get("index", 0)) for item in items],
153
+ values=[float(item.get("value", 0.0)) for item in items],
154
+ )
155
+
156
+
157
+ def _build_grpc_embedding_request(text: str, truncate: bool | None, normalize: bool | None) -> dict[str, str | bool]:
158
+ """Build a gRPC embedding request."""
159
+ request: dict[str, str | bool] = {"inputs": text}
160
+ if truncate is not None:
161
+ request["truncate"] = truncate
162
+ if normalize is not None:
163
+ request["normalize"] = normalize
164
+ return request
165
+
166
+
167
+ def _grpc_metadata(token: Secret | None) -> tuple[tuple[str, str], ...]:
168
+ """Build gRPC authorization metadata."""
169
+ value = token.resolve_value() if token else None
170
+ return (("authorization", f"Bearer {value}"),) if value else ()
@@ -2,13 +2,13 @@
2
2
  #
3
3
  # SPDX-License-Identifier: Apache-2.0
4
4
  from .document_embedder import HuggingFaceAPIDocumentEmbedder
5
- from .sparse_document_embedder import HuggingFaceAPISparseDocumentEmbedder
6
- from .sparse_text_embedder import HuggingFaceAPISparseTextEmbedder
5
+ from .sparse_document_embedder import HuggingFaceTEISparseDocumentEmbedder
6
+ from .sparse_text_embedder import HuggingFaceTEISparseTextEmbedder
7
7
  from .text_embedder import HuggingFaceAPITextEmbedder
8
8
 
9
9
  __all__ = [
10
10
  "HuggingFaceAPIDocumentEmbedder",
11
- "HuggingFaceAPISparseDocumentEmbedder",
12
- "HuggingFaceAPISparseTextEmbedder",
13
11
  "HuggingFaceAPITextEmbedder",
12
+ "HuggingFaceTEISparseDocumentEmbedder",
13
+ "HuggingFaceTEISparseTextEmbedder",
14
14
  ]
@@ -3,12 +3,14 @@
3
3
  # SPDX-License-Identifier: Apache-2.0
4
4
 
5
5
  from asyncio import Semaphore, gather
6
+ from collections.abc import AsyncIterator
6
7
  from dataclasses import replace
7
8
  from itertools import chain
8
9
  from typing import Any
9
10
 
10
11
  from haystack import component, default_from_dict, default_to_dict, logging
11
12
  from haystack.dataclasses import Document
13
+ from haystack.lazy_imports import LazyImport
12
14
  from haystack.utils import Secret
13
15
  from haystack.utils.url_validation import is_valid_http_url
14
16
  from huggingface_hub import AsyncInferenceClient, InferenceClient
@@ -17,9 +19,17 @@ from tqdm import tqdm
17
19
  from haystack_integrations.common.huggingface_api.utils import (
18
20
  HFEmbeddingAPIType,
19
21
  HFModelType,
22
+ _build_grpc_embedding_request,
20
23
  _check_valid_model,
24
+ _check_valid_model_async,
25
+ _grpc_metadata,
21
26
  )
22
27
 
28
+ with LazyImport("Run 'pip install \"huggingface-api-haystack[grpc]\"' for grpc support.") as grpc_import:
29
+ from grpc_requests import Client
30
+ from grpc_requests.aio import AsyncClient
31
+
32
+
23
33
  logger = logging.getLogger(__name__)
24
34
 
25
35
  # batched embeddings returned by the API are expected to have shape (batch_size, embedding_dim)
@@ -109,6 +119,7 @@ class HuggingFaceAPIDocumentEmbedder:
109
119
  meta_fields_to_embed: list[str] | None = None,
110
120
  embedding_separator: str = "\n",
111
121
  concurrency_limit: int = 4,
122
+ use_grpc: bool = False,
112
123
  ) -> None:
113
124
  """
114
125
  Creates a HuggingFaceAPIDocumentEmbedder component.
@@ -118,9 +129,9 @@ class HuggingFaceAPIDocumentEmbedder:
118
129
  :param api_params:
119
130
  A dictionary with the following keys:
120
131
  - `model`: Hugging Face model ID. Required when `api_type` is `SERVERLESS_INFERENCE_API`.
121
- - `url`: URL of the inference endpoint. Required when `api_type` is `INFERENCE_ENDPOINTS` or
122
- `TEXT_EMBEDDINGS_INFERENCE`.
123
- :param token: The Hugging Face token to use as HTTP bearer authorization.
132
+ - `url`: URL of the inference endpoint, or gRPC target. Required when `api_type` is
133
+ `INFERENCE_ENDPOINTS` or `TEXT_EMBEDDINGS_INFERENCE`.
134
+ :param token: The Hugging Face token to use as bearer authorization.
124
135
  Check your HF token in your [account settings](https://huggingface.co/settings/tokens).
125
136
  :param prefix:
126
137
  A string to add at the beginning of each text.
@@ -137,7 +148,7 @@ class HuggingFaceAPIDocumentEmbedder:
137
148
  if the backend uses Text Embeddings Inference.
138
149
  If `api_type` is `SERVERLESS_INFERENCE_API`, this parameter is ignored.
139
150
  :param batch_size:
140
- Number of documents to process at once.
151
+ Number of documents to process at once. Only used with HTTP.
141
152
  :param progress_bar:
142
153
  If `True`, shows a progress bar when running.
143
154
  :param meta_fields_to_embed:
@@ -145,15 +156,21 @@ class HuggingFaceAPIDocumentEmbedder:
145
156
  :param embedding_separator:
146
157
  Separator used to concatenate the metadata fields to the document text.
147
158
  :param concurrency_limit:
148
- The maximum number of requests that should be allowed to run concurrently.
159
+ The maximum number of HTTP requests or gRPC streams that should be allowed to run concurrently.
149
160
  This parameter is only used in the `run_async` method.
161
+ :param use_grpc:
162
+ Use gRPC instead of HTTP. Requires `huggingface-api-haystack[grpc]`.
150
163
  :raises ValueError:
151
- If the required `model` or `url` is missing from `api_params`, the `url` is invalid,
152
- or the `api_type` is unknown.
164
+ If the required `model` or `url` is missing from `api_params`, the HTTP `url` is invalid,
165
+ or the `api_type` is unknown or is `SERVERLESS_INFERENCE_API` with `use_grpc=True`.
153
166
  """
154
167
  if isinstance(api_type, str):
155
168
  api_type = HFEmbeddingAPIType.from_str(api_type)
156
169
 
170
+ if use_grpc and api_type == HFEmbeddingAPIType.SERVERLESS_INFERENCE_API:
171
+ msg = "gRPC is not supported by the Serverless Inference API."
172
+ raise ValueError(msg)
173
+
157
174
  api_params = api_params or {}
158
175
 
159
176
  if api_type == HFEmbeddingAPIType.SERVERLESS_INFERENCE_API:
@@ -161,7 +178,6 @@ class HuggingFaceAPIDocumentEmbedder:
161
178
  if model is None:
162
179
  msg = "To use the Serverless Inference API, you need to specify the `model` parameter in `api_params`."
163
180
  raise ValueError(msg)
164
- _check_valid_model(model, HFModelType.EMBEDDING, token)
165
181
  model_or_url = model
166
182
  elif api_type in [HFEmbeddingAPIType.INFERENCE_ENDPOINTS, HFEmbeddingAPIType.TEXT_EMBEDDINGS_INFERENCE]:
167
183
  url = api_params.get("url")
@@ -171,7 +187,7 @@ class HuggingFaceAPIDocumentEmbedder:
171
187
  "parameter in `api_params`."
172
188
  )
173
189
  raise ValueError(msg)
174
- if not is_valid_http_url(url):
190
+ if not use_grpc and not is_valid_http_url(url):
175
191
  msg = f"Invalid URL: {url}"
176
192
  raise ValueError(msg)
177
193
  model_or_url = url
@@ -179,7 +195,8 @@ class HuggingFaceAPIDocumentEmbedder:
179
195
  msg = f"Unknown api_type {api_type}"
180
196
  raise ValueError(msg)
181
197
 
182
- client_args: dict[str, Any] = {"model": model_or_url, "token": token.resolve_value() if token else None}
198
+ if use_grpc:
199
+ grpc_import.check()
183
200
 
184
201
  self.api_type = api_type
185
202
  self.api_params = api_params
@@ -188,13 +205,63 @@ class HuggingFaceAPIDocumentEmbedder:
188
205
  self.suffix = suffix
189
206
  self.truncate = truncate
190
207
  self.normalize = normalize
208
+ self.use_grpc = use_grpc
191
209
  self.batch_size = batch_size
192
210
  self.progress_bar = progress_bar
193
211
  self.meta_fields_to_embed = meta_fields_to_embed or []
194
212
  self.embedding_separator = embedding_separator
195
213
  self.concurrency_limit = concurrency_limit
196
- self._client = InferenceClient(**client_args)
197
- self._async_client = AsyncInferenceClient(**client_args)
214
+ self._model_or_url = model_or_url
215
+ self._client: InferenceClient | None = None
216
+ self._async_client: AsyncInferenceClient | None = None
217
+ self._grpc_client: Client | None = None
218
+ self._async_grpc_client: AsyncClient | None = None
219
+
220
+ def _client_kwargs(self) -> dict[str, Any]:
221
+ """Build the keyword arguments used to create Hugging Face clients."""
222
+ return {"model": self._model_or_url, "token": self.token.resolve_value() if self.token else None}
223
+
224
+ def warm_up(self) -> None:
225
+ """Create the synchronous client."""
226
+ if self.use_grpc:
227
+ if self._grpc_client is None:
228
+ self._grpc_client = Client(self._model_or_url)
229
+ return
230
+
231
+ if self._client is None:
232
+ if self.api_type == HFEmbeddingAPIType.SERVERLESS_INFERENCE_API:
233
+ _check_valid_model(self._model_or_url, HFModelType.EMBEDDING, self.token)
234
+ self._client = InferenceClient(**self._client_kwargs())
235
+
236
+ async def warm_up_async(self) -> None:
237
+ """Create the asynchronous client."""
238
+ if self.use_grpc:
239
+ if self._async_grpc_client is None:
240
+ self._async_grpc_client = await AsyncClient.create(self._model_or_url)
241
+ return
242
+
243
+ if self._async_client is None:
244
+ if self.api_type == HFEmbeddingAPIType.SERVERLESS_INFERENCE_API:
245
+ await _check_valid_model_async(self._model_or_url, HFModelType.EMBEDDING, self.token)
246
+ self._async_client = AsyncInferenceClient(**self._client_kwargs())
247
+
248
+ def close(self) -> None:
249
+ """Close the synchronous client."""
250
+ if self._client is not None:
251
+ self._client.close()
252
+ self._client = None
253
+ if self._grpc_client is not None:
254
+ self._grpc_client.channel.close()
255
+ self._grpc_client = None
256
+
257
+ async def close_async(self) -> None:
258
+ """Close the asynchronous client."""
259
+ if self._async_client is not None:
260
+ await self._async_client.close()
261
+ self._async_client = None
262
+ if self._async_grpc_client is not None:
263
+ await self._async_grpc_client.channel.close()
264
+ self._async_grpc_client = None
198
265
 
199
266
  def to_dict(self) -> dict[str, Any]:
200
267
  """
@@ -212,6 +279,7 @@ class HuggingFaceAPIDocumentEmbedder:
212
279
  token=self.token,
213
280
  truncate=self.truncate,
214
281
  normalize=self.normalize,
282
+ use_grpc=self.use_grpc,
215
283
  batch_size=self.batch_size,
216
284
  progress_bar=self.progress_bar,
217
285
  meta_fields_to_embed=self.meta_fields_to_embed,
@@ -266,10 +334,37 @@ class HuggingFaceAPIDocumentEmbedder:
266
334
  normalize = None
267
335
  return truncate, normalize
268
336
 
337
+ def _embed_batch_grpc(self, texts_to_embed: list[str]) -> list[list[float]]:
338
+ """Embed all texts through a single gRPC stream."""
339
+ if not texts_to_embed:
340
+ return []
341
+
342
+ assert self._grpc_client is not None # noqa: S101
343
+ responses = self._grpc_client.stream_stream(
344
+ "tei.v1.Embed",
345
+ "EmbedStream",
346
+ (_build_grpc_embedding_request(text, self.truncate, self.normalize) for text in texts_to_embed),
347
+ metadata=_grpc_metadata(self.token),
348
+ )
349
+ embeddings = [
350
+ response["embeddings"]
351
+ for response in tqdm(
352
+ responses,
353
+ total=len(texts_to_embed),
354
+ disable=not self.progress_bar,
355
+ desc="Calculating embeddings",
356
+ )
357
+ ]
358
+ if len(embeddings) != len(texts_to_embed):
359
+ msg = f"Expected {len(texts_to_embed)} embeddings, got {len(embeddings)}"
360
+ raise ValueError(msg)
361
+ return embeddings
362
+
269
363
  def _embed_batch(self, texts_to_embed: list[str], batch_size: int) -> list[list[float]]:
270
364
  """
271
365
  Embed a list of texts in batches.
272
366
  """
367
+ assert self._client is not None # noqa: S101
273
368
  truncate, normalize = self._adjust_api_parameters(self.truncate, self.normalize, self.api_type)
274
369
 
275
370
  all_embeddings: list = []
@@ -278,7 +373,12 @@ class HuggingFaceAPIDocumentEmbedder:
278
373
  ):
279
374
  batch = texts_to_embed[i : i + batch_size]
280
375
 
281
- np_embeddings = self._client.feature_extraction(text=batch, truncate=truncate, normalize=normalize)
376
+ # a batch is accepted but text is typed as str
377
+ np_embeddings = self._client.feature_extraction(
378
+ text=batch, # type: ignore[arg-type]
379
+ truncate=truncate,
380
+ normalize=normalize,
381
+ )
282
382
 
283
383
  if np_embeddings.ndim != _EXPECTED_EMBEDDING_NDIM or np_embeddings.shape[0] != len(batch):
284
384
  msg = f"Expected embedding shape ({batch_size}, embedding_dim), got {np_embeddings.shape}"
@@ -288,10 +388,47 @@ class HuggingFaceAPIDocumentEmbedder:
288
388
 
289
389
  return all_embeddings
290
390
 
391
+ async def _embed_batch_grpc_async(self, texts_to_embed: list[str]) -> list[list[float]]:
392
+ """Embed texts through concurrent gRPC streams."""
393
+ if not texts_to_embed:
394
+ return []
395
+
396
+ # Split texts evenly into ordered streams, up to the concurrency limit.
397
+ stream_count = min(max(1, self.concurrency_limit), len(texts_to_embed))
398
+ streams = [
399
+ texts_to_embed[len(texts_to_embed) * i // stream_count : len(texts_to_embed) * (i + 1) // stream_count]
400
+ for i in range(stream_count)
401
+ ]
402
+
403
+ pbar = tqdm(total=stream_count, disable=not self.progress_bar, desc="Calculating embeddings")
404
+
405
+ async def _runner(texts: list[str]) -> list[list[float]]:
406
+ async def _requests() -> AsyncIterator[dict[str, Any]]:
407
+ for text in texts:
408
+ yield _build_grpc_embedding_request(text, self.truncate, self.normalize)
409
+
410
+ assert self._async_grpc_client is not None # noqa: S101
411
+ responses = await self._async_grpc_client.stream_stream(
412
+ "tei.v1.Embed", "EmbedStream", _requests(), metadata=_grpc_metadata(self.token)
413
+ )
414
+ embeddings = [response["embeddings"] async for response in responses]
415
+ if len(embeddings) != len(texts):
416
+ msg = f"Expected {len(texts)} embeddings, got {len(embeddings)}"
417
+ raise ValueError(msg)
418
+ pbar.update(1)
419
+ return embeddings
420
+
421
+ try:
422
+ return [*chain(*await gather(*(_runner(texts) for texts in streams)))]
423
+ finally:
424
+ pbar.close()
425
+
291
426
  async def _embed_batch_async(self, texts_to_embed: list[str], batch_size: int) -> list[list[float]]:
292
427
  """
293
428
  Embed a list of texts in batches asynchronously.
294
429
  """
430
+ assert self._async_client is not None # noqa: S101
431
+ async_client = self._async_client
295
432
  truncate, normalize = self._adjust_api_parameters(self.truncate, self.normalize, self.api_type)
296
433
  sem = Semaphore(max(1, self.concurrency_limit))
297
434
  num_batches = (len(texts_to_embed) + batch_size - 1) // batch_size
@@ -299,8 +436,11 @@ class HuggingFaceAPIDocumentEmbedder:
299
436
 
300
437
  async def _runner(batch: list[str]) -> list[list[float]]:
301
438
  async with sem:
302
- np_embeddings = await self._async_client.feature_extraction(
303
- text=batch, truncate=truncate, normalize=normalize
439
+ # a batch is accepted but text is typed as str
440
+ np_embeddings = await async_client.feature_extraction(
441
+ text=batch, # type: ignore[arg-type]
442
+ truncate=truncate,
443
+ normalize=normalize,
304
444
  )
305
445
 
306
446
  if np_embeddings.ndim != _EXPECTED_EMBEDDING_NDIM or np_embeddings.shape[0] != len(batch):
@@ -342,6 +482,8 @@ class HuggingFaceAPIDocumentEmbedder:
342
482
  A dictionary with the following keys:
343
483
  - `documents`: A list of documents with embeddings.
344
484
  """
485
+ self.warm_up()
486
+
345
487
  if not isinstance(documents, list) or (documents and not isinstance(documents[0], Document)):
346
488
  msg = (
347
489
  "HuggingFaceAPIDocumentEmbedder expects a list of Documents as input."
@@ -351,7 +493,10 @@ class HuggingFaceAPIDocumentEmbedder:
351
493
 
352
494
  texts_to_embed = self._prepare_texts_to_embed(documents=documents)
353
495
 
354
- embeddings = self._embed_batch(texts_to_embed=texts_to_embed, batch_size=self.batch_size)
496
+ if self.use_grpc:
497
+ embeddings = self._embed_batch_grpc(texts_to_embed)
498
+ else:
499
+ embeddings = self._embed_batch(texts_to_embed=texts_to_embed, batch_size=self.batch_size)
355
500
 
356
501
  new_documents = []
357
502
  for doc, emb in zip(documents, embeddings, strict=True):
@@ -375,6 +520,8 @@ class HuggingFaceAPIDocumentEmbedder:
375
520
  A dictionary with the following keys:
376
521
  - `documents`: A list of documents with embeddings.
377
522
  """
523
+ await self.warm_up_async()
524
+
378
525
  if not isinstance(documents, list) or (documents and not isinstance(documents[0], Document)):
379
526
  msg = (
380
527
  "HuggingFaceAPIDocumentEmbedder expects a list of Documents as input."
@@ -384,7 +531,10 @@ class HuggingFaceAPIDocumentEmbedder:
384
531
 
385
532
  texts_to_embed = self._prepare_texts_to_embed(documents=documents)
386
533
 
387
- embeddings = await self._embed_batch_async(texts_to_embed=texts_to_embed, batch_size=self.batch_size)
534
+ if self.use_grpc:
535
+ embeddings = await self._embed_batch_grpc_async(texts_to_embed)
536
+ else:
537
+ embeddings = await self._embed_batch_async(texts_to_embed=texts_to_embed, batch_size=self.batch_size)
388
538
 
389
539
  new_documents = []
390
540
  for doc, emb in zip(documents, embeddings, strict=True):