huggingface-api-haystack 1.0.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 (39) hide show
  1. {huggingface_api_haystack-1.0.0 → huggingface_api_haystack-2.0.0}/CHANGELOG.md +15 -0
  2. {huggingface_api_haystack-1.0.0 → huggingface_api_haystack-2.0.0}/PKG-INFO +4 -3
  3. {huggingface_api_haystack-1.0.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-1.0.0 → huggingface_api_haystack-2.0.0}/pyproject.toml +10 -0
  6. {huggingface_api_haystack-1.0.0 → huggingface_api_haystack-2.0.0}/src/haystack_integrations/common/huggingface_api/utils.py +25 -0
  7. {huggingface_api_haystack-1.0.0 → huggingface_api_haystack-2.0.0}/src/haystack_integrations/components/embedders/huggingface_api/__init__.py +4 -4
  8. {huggingface_api_haystack-1.0.0 → huggingface_api_haystack-2.0.0}/src/haystack_integrations/components/embedders/huggingface_api/document_embedder.py +120 -16
  9. {huggingface_api_haystack-1.0.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-1.0.0 → huggingface_api_haystack-2.0.0}/src/haystack_integrations/components/embedders/huggingface_api/text_embedder.py +70 -14
  12. {huggingface_api_haystack-1.0.0 → huggingface_api_haystack-2.0.0}/src/haystack_integrations/components/rankers/huggingface_api/ranker.py +76 -3
  13. {huggingface_api_haystack-1.0.0 → huggingface_api_haystack-2.0.0}/tests/test_chat_generator.py +2 -4
  14. {huggingface_api_haystack-1.0.0 → huggingface_api_haystack-2.0.0}/tests/test_document_embedder.py +185 -23
  15. {huggingface_api_haystack-1.0.0 → huggingface_api_haystack-2.0.0}/tests/test_ranker.py +174 -3
  16. huggingface_api_haystack-2.0.0/tests/test_sparse_document_embedder.py +621 -0
  17. huggingface_api_haystack-2.0.0/tests/test_sparse_text_embedder.py +499 -0
  18. {huggingface_api_haystack-1.0.0 → huggingface_api_haystack-2.0.0}/tests/test_text_embedder.py +143 -23
  19. {huggingface_api_haystack-1.0.0 → huggingface_api_haystack-2.0.0}/tests/test_utils.py +45 -0
  20. huggingface_api_haystack-1.0.0/docker-compose.yml +0 -14
  21. huggingface_api_haystack-1.0.0/src/haystack_integrations/components/embedders/huggingface_api/sparse_text_embedder.py +0 -116
  22. huggingface_api_haystack-1.0.0/tests/test_sparse_document_embedder.py +0 -342
  23. huggingface_api_haystack-1.0.0/tests/test_sparse_text_embedder.py +0 -263
  24. {huggingface_api_haystack-1.0.0 → huggingface_api_haystack-2.0.0}/.gitignore +0 -0
  25. {huggingface_api_haystack-1.0.0 → huggingface_api_haystack-2.0.0}/LICENSE.txt +0 -0
  26. {huggingface_api_haystack-1.0.0 → huggingface_api_haystack-2.0.0}/pydoc/config_docusaurus.yml +0 -0
  27. {huggingface_api_haystack-1.0.0 → huggingface_api_haystack-2.0.0}/src/haystack_integrations/common/huggingface_api/__init__.py +0 -0
  28. {huggingface_api_haystack-1.0.0 → huggingface_api_haystack-2.0.0}/src/haystack_integrations/common/py.typed +0 -0
  29. {huggingface_api_haystack-1.0.0 → huggingface_api_haystack-2.0.0}/src/haystack_integrations/components/embedders/huggingface_api/sparse_embedding_utils.py +0 -0
  30. {huggingface_api_haystack-1.0.0 → huggingface_api_haystack-2.0.0}/src/haystack_integrations/components/embedders/py.typed +0 -0
  31. {huggingface_api_haystack-1.0.0 → huggingface_api_haystack-2.0.0}/src/haystack_integrations/components/generators/huggingface_api/__init__.py +0 -0
  32. {huggingface_api_haystack-1.0.0 → huggingface_api_haystack-2.0.0}/src/haystack_integrations/components/generators/huggingface_api/chat/__init__.py +0 -0
  33. {huggingface_api_haystack-1.0.0 → huggingface_api_haystack-2.0.0}/src/haystack_integrations/components/generators/huggingface_api/chat/chat_generator.py +0 -0
  34. {huggingface_api_haystack-1.0.0 → huggingface_api_haystack-2.0.0}/src/haystack_integrations/components/generators/py.typed +0 -0
  35. {huggingface_api_haystack-1.0.0 → huggingface_api_haystack-2.0.0}/src/haystack_integrations/components/rankers/huggingface_api/__init__.py +0 -0
  36. {huggingface_api_haystack-1.0.0 → huggingface_api_haystack-2.0.0}/src/haystack_integrations/components/rankers/py.typed +0 -0
  37. {huggingface_api_haystack-1.0.0 → huggingface_api_haystack-2.0.0}/tests/__init__.py +0 -0
  38. {huggingface_api_haystack-1.0.0 → huggingface_api_haystack-2.0.0}/tests/conftest.py +0 -0
  39. {huggingface_api_haystack-1.0.0 → huggingface_api_haystack-2.0.0}/tests/test_files/apple.jpg +0 -0
@@ -1,5 +1,20 @@
1
1
  # Changelog
2
2
 
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
+
3
18
  ## [integrations/huggingface_api-v0.6.0] - 2026-08-27
4
19
 
5
20
  ### 🚀 Features
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.5
2
2
  Name: huggingface-api-haystack
3
- Version: 1.0.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
@@ -22,6 +22,8 @@ Requires-Python: >=3.10
22
22
  Requires-Dist: haystack-ai>=3.0.0
23
23
  Requires-Dist: httpx>=0.27.0
24
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:
@@ -25,6 +25,11 @@ classifiers = [
25
25
  ]
26
26
  dependencies = ["haystack-ai>=3.0.0", "httpx>=0.27.0", "huggingface_hub[inference]>=1.0.0"]
27
27
 
28
+ [project.optional-dependencies]
29
+ grpc = [
30
+ "grpc-requests>=0.1.21",
31
+ ]
32
+
28
33
  [project.urls]
29
34
  Documentation = "https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/huggingface_api#readme"
30
35
  Issues = "https://github.com/deepset-ai/haystack-core-integrations/issues"
@@ -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,6 +4,7 @@
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
@@ -143,3 +144,27 @@ def _validate_model_type(model_id: str, model_type: HFModelType, pipeline_tag: s
143
144
 
144
145
  if not allowed_model:
145
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,10 +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,
21
24
  _check_valid_model_async,
25
+ _grpc_metadata,
22
26
  )
23
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
+
24
33
  logger = logging.getLogger(__name__)
25
34
 
26
35
  # batched embeddings returned by the API are expected to have shape (batch_size, embedding_dim)
@@ -110,6 +119,7 @@ class HuggingFaceAPIDocumentEmbedder:
110
119
  meta_fields_to_embed: list[str] | None = None,
111
120
  embedding_separator: str = "\n",
112
121
  concurrency_limit: int = 4,
122
+ use_grpc: bool = False,
113
123
  ) -> None:
114
124
  """
115
125
  Creates a HuggingFaceAPIDocumentEmbedder component.
@@ -119,9 +129,9 @@ class HuggingFaceAPIDocumentEmbedder:
119
129
  :param api_params:
120
130
  A dictionary with the following keys:
121
131
  - `model`: Hugging Face model ID. Required when `api_type` is `SERVERLESS_INFERENCE_API`.
122
- - `url`: URL of the inference endpoint. Required when `api_type` is `INFERENCE_ENDPOINTS` or
123
- `TEXT_EMBEDDINGS_INFERENCE`.
124
- :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.
125
135
  Check your HF token in your [account settings](https://huggingface.co/settings/tokens).
126
136
  :param prefix:
127
137
  A string to add at the beginning of each text.
@@ -138,7 +148,7 @@ class HuggingFaceAPIDocumentEmbedder:
138
148
  if the backend uses Text Embeddings Inference.
139
149
  If `api_type` is `SERVERLESS_INFERENCE_API`, this parameter is ignored.
140
150
  :param batch_size:
141
- Number of documents to process at once.
151
+ Number of documents to process at once. Only used with HTTP.
142
152
  :param progress_bar:
143
153
  If `True`, shows a progress bar when running.
144
154
  :param meta_fields_to_embed:
@@ -146,15 +156,21 @@ class HuggingFaceAPIDocumentEmbedder:
146
156
  :param embedding_separator:
147
157
  Separator used to concatenate the metadata fields to the document text.
148
158
  :param concurrency_limit:
149
- 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.
150
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]`.
151
163
  :raises ValueError:
152
- If the required `model` or `url` is missing from `api_params`, the `url` is invalid,
153
- 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`.
154
166
  """
155
167
  if isinstance(api_type, str):
156
168
  api_type = HFEmbeddingAPIType.from_str(api_type)
157
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
+
158
174
  api_params = api_params or {}
159
175
 
160
176
  if api_type == HFEmbeddingAPIType.SERVERLESS_INFERENCE_API:
@@ -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,6 +195,9 @@ class HuggingFaceAPIDocumentEmbedder:
179
195
  msg = f"Unknown api_type {api_type}"
180
196
  raise ValueError(msg)
181
197
 
198
+ if use_grpc:
199
+ grpc_import.check()
200
+
182
201
  self.api_type = api_type
183
202
  self.api_params = api_params
184
203
  self.token = token
@@ -186,6 +205,7 @@ class HuggingFaceAPIDocumentEmbedder:
186
205
  self.suffix = suffix
187
206
  self.truncate = truncate
188
207
  self.normalize = normalize
208
+ self.use_grpc = use_grpc
189
209
  self.batch_size = batch_size
190
210
  self.progress_bar = progress_bar
191
211
  self.meta_fields_to_embed = meta_fields_to_embed or []
@@ -194,36 +214,54 @@ class HuggingFaceAPIDocumentEmbedder:
194
214
  self._model_or_url = model_or_url
195
215
  self._client: InferenceClient | None = None
196
216
  self._async_client: AsyncInferenceClient | None = None
217
+ self._grpc_client: Client | None = None
218
+ self._async_grpc_client: AsyncClient | None = None
197
219
 
198
220
  def _client_kwargs(self) -> dict[str, Any]:
199
221
  """Build the keyword arguments used to create Hugging Face clients."""
200
222
  return {"model": self._model_or_url, "token": self.token.resolve_value() if self.token else None}
201
223
 
202
224
  def warm_up(self) -> None:
203
- """Create the synchronous Hugging Face client."""
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
+
204
231
  if self._client is None:
205
232
  if self.api_type == HFEmbeddingAPIType.SERVERLESS_INFERENCE_API:
206
233
  _check_valid_model(self._model_or_url, HFModelType.EMBEDDING, self.token)
207
234
  self._client = InferenceClient(**self._client_kwargs())
208
235
 
209
236
  async def warm_up_async(self) -> None:
210
- """Create the asynchronous Hugging Face client."""
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
+
211
243
  if self._async_client is None:
212
244
  if self.api_type == HFEmbeddingAPIType.SERVERLESS_INFERENCE_API:
213
245
  await _check_valid_model_async(self._model_or_url, HFModelType.EMBEDDING, self.token)
214
246
  self._async_client = AsyncInferenceClient(**self._client_kwargs())
215
247
 
216
248
  def close(self) -> None:
217
- """Close the synchronous Hugging Face client."""
249
+ """Close the synchronous client."""
218
250
  if self._client is not None:
219
251
  self._client.close()
220
252
  self._client = None
253
+ if self._grpc_client is not None:
254
+ self._grpc_client.channel.close()
255
+ self._grpc_client = None
221
256
 
222
257
  async def close_async(self) -> None:
223
- """Close the asynchronous Hugging Face client."""
258
+ """Close the asynchronous client."""
224
259
  if self._async_client is not None:
225
260
  await self._async_client.close()
226
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
227
265
 
228
266
  def to_dict(self) -> dict[str, Any]:
229
267
  """
@@ -241,6 +279,7 @@ class HuggingFaceAPIDocumentEmbedder:
241
279
  token=self.token,
242
280
  truncate=self.truncate,
243
281
  normalize=self.normalize,
282
+ use_grpc=self.use_grpc,
244
283
  batch_size=self.batch_size,
245
284
  progress_bar=self.progress_bar,
246
285
  meta_fields_to_embed=self.meta_fields_to_embed,
@@ -295,6 +334,32 @@ class HuggingFaceAPIDocumentEmbedder:
295
334
  normalize = None
296
335
  return truncate, normalize
297
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
+
298
363
  def _embed_batch(self, texts_to_embed: list[str], batch_size: int) -> list[list[float]]:
299
364
  """
300
365
  Embed a list of texts in batches.
@@ -323,6 +388,41 @@ class HuggingFaceAPIDocumentEmbedder:
323
388
 
324
389
  return all_embeddings
325
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
+
326
426
  async def _embed_batch_async(self, texts_to_embed: list[str], batch_size: int) -> list[list[float]]:
327
427
  """
328
428
  Embed a list of texts in batches asynchronously.
@@ -383,7 +483,6 @@ class HuggingFaceAPIDocumentEmbedder:
383
483
  - `documents`: A list of documents with embeddings.
384
484
  """
385
485
  self.warm_up()
386
- assert self._client is not None # noqa: S101
387
486
 
388
487
  if not isinstance(documents, list) or (documents and not isinstance(documents[0], Document)):
389
488
  msg = (
@@ -394,7 +493,10 @@ class HuggingFaceAPIDocumentEmbedder:
394
493
 
395
494
  texts_to_embed = self._prepare_texts_to_embed(documents=documents)
396
495
 
397
- 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)
398
500
 
399
501
  new_documents = []
400
502
  for doc, emb in zip(documents, embeddings, strict=True):
@@ -419,7 +521,6 @@ class HuggingFaceAPIDocumentEmbedder:
419
521
  - `documents`: A list of documents with embeddings.
420
522
  """
421
523
  await self.warm_up_async()
422
- assert self._async_client is not None # noqa: S101
423
524
 
424
525
  if not isinstance(documents, list) or (documents and not isinstance(documents[0], Document)):
425
526
  msg = (
@@ -430,7 +531,10 @@ class HuggingFaceAPIDocumentEmbedder:
430
531
 
431
532
  texts_to_embed = self._prepare_texts_to_embed(documents=documents)
432
533
 
433
- 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)
434
538
 
435
539
  new_documents = []
436
540
  for doc, emb in zip(documents, embeddings, strict=True):