huggingface-api-haystack 1.0.0__tar.gz → 2.1.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.
- {huggingface_api_haystack-1.0.0 → huggingface_api_haystack-2.1.0}/CHANGELOG.md +26 -0
- {huggingface_api_haystack-1.0.0 → huggingface_api_haystack-2.1.0}/PKG-INFO +4 -3
- {huggingface_api_haystack-1.0.0 → huggingface_api_haystack-2.1.0}/README.md +1 -2
- huggingface_api_haystack-2.1.0/docker-compose.yml +50 -0
- {huggingface_api_haystack-1.0.0 → huggingface_api_haystack-2.1.0}/pyproject.toml +10 -0
- {huggingface_api_haystack-1.0.0 → huggingface_api_haystack-2.1.0}/src/haystack_integrations/common/huggingface_api/utils.py +25 -0
- {huggingface_api_haystack-1.0.0 → huggingface_api_haystack-2.1.0}/src/haystack_integrations/components/embedders/huggingface_api/__init__.py +4 -4
- {huggingface_api_haystack-1.0.0 → huggingface_api_haystack-2.1.0}/src/haystack_integrations/components/embedders/huggingface_api/document_embedder.py +120 -16
- {huggingface_api_haystack-1.0.0 → huggingface_api_haystack-2.1.0}/src/haystack_integrations/components/embedders/huggingface_api/sparse_document_embedder.py +147 -20
- huggingface_api_haystack-2.1.0/src/haystack_integrations/components/embedders/huggingface_api/sparse_text_embedder.py +191 -0
- {huggingface_api_haystack-1.0.0 → huggingface_api_haystack-2.1.0}/src/haystack_integrations/components/embedders/huggingface_api/text_embedder.py +70 -14
- {huggingface_api_haystack-1.0.0 → huggingface_api_haystack-2.1.0}/src/haystack_integrations/components/generators/huggingface_api/chat/chat_generator.py +2 -15
- {huggingface_api_haystack-1.0.0 → huggingface_api_haystack-2.1.0}/src/haystack_integrations/components/rankers/huggingface_api/ranker.py +76 -3
- {huggingface_api_haystack-1.0.0 → huggingface_api_haystack-2.1.0}/tests/test_chat_generator.py +2 -113
- {huggingface_api_haystack-1.0.0 → huggingface_api_haystack-2.1.0}/tests/test_document_embedder.py +185 -23
- {huggingface_api_haystack-1.0.0 → huggingface_api_haystack-2.1.0}/tests/test_ranker.py +174 -3
- huggingface_api_haystack-2.1.0/tests/test_sparse_document_embedder.py +621 -0
- huggingface_api_haystack-2.1.0/tests/test_sparse_text_embedder.py +499 -0
- {huggingface_api_haystack-1.0.0 → huggingface_api_haystack-2.1.0}/tests/test_text_embedder.py +143 -23
- {huggingface_api_haystack-1.0.0 → huggingface_api_haystack-2.1.0}/tests/test_utils.py +45 -0
- huggingface_api_haystack-1.0.0/docker-compose.yml +0 -14
- huggingface_api_haystack-1.0.0/src/haystack_integrations/components/embedders/huggingface_api/sparse_text_embedder.py +0 -116
- huggingface_api_haystack-1.0.0/tests/test_sparse_document_embedder.py +0 -342
- huggingface_api_haystack-1.0.0/tests/test_sparse_text_embedder.py +0 -263
- {huggingface_api_haystack-1.0.0 → huggingface_api_haystack-2.1.0}/.gitignore +0 -0
- {huggingface_api_haystack-1.0.0 → huggingface_api_haystack-2.1.0}/LICENSE.txt +0 -0
- {huggingface_api_haystack-1.0.0 → huggingface_api_haystack-2.1.0}/pydoc/config_docusaurus.yml +0 -0
- {huggingface_api_haystack-1.0.0 → huggingface_api_haystack-2.1.0}/src/haystack_integrations/common/huggingface_api/__init__.py +0 -0
- {huggingface_api_haystack-1.0.0 → huggingface_api_haystack-2.1.0}/src/haystack_integrations/common/py.typed +0 -0
- {huggingface_api_haystack-1.0.0 → huggingface_api_haystack-2.1.0}/src/haystack_integrations/components/embedders/huggingface_api/sparse_embedding_utils.py +0 -0
- {huggingface_api_haystack-1.0.0 → huggingface_api_haystack-2.1.0}/src/haystack_integrations/components/embedders/py.typed +0 -0
- {huggingface_api_haystack-1.0.0 → huggingface_api_haystack-2.1.0}/src/haystack_integrations/components/generators/huggingface_api/__init__.py +0 -0
- {huggingface_api_haystack-1.0.0 → huggingface_api_haystack-2.1.0}/src/haystack_integrations/components/generators/huggingface_api/chat/__init__.py +0 -0
- {huggingface_api_haystack-1.0.0 → huggingface_api_haystack-2.1.0}/src/haystack_integrations/components/generators/py.typed +0 -0
- {huggingface_api_haystack-1.0.0 → huggingface_api_haystack-2.1.0}/src/haystack_integrations/components/rankers/huggingface_api/__init__.py +0 -0
- {huggingface_api_haystack-1.0.0 → huggingface_api_haystack-2.1.0}/src/haystack_integrations/components/rankers/py.typed +0 -0
- {huggingface_api_haystack-1.0.0 → huggingface_api_haystack-2.1.0}/tests/__init__.py +0 -0
- {huggingface_api_haystack-1.0.0 → huggingface_api_haystack-2.1.0}/tests/conftest.py +0 -0
- {huggingface_api_haystack-1.0.0 → huggingface_api_haystack-2.1.0}/tests/test_files/apple.jpg +0 -0
|
@@ -1,5 +1,31 @@
|
|
|
1
1
|
# Changelog
|
|
2
2
|
|
|
3
|
+
## [integrations/huggingface_api-v2.0.0] - 2026-10-01
|
|
4
|
+
|
|
5
|
+
### 🚀 Features
|
|
6
|
+
|
|
7
|
+
- [**breaking**] Add grpc support for the huggingface TEI components; rename HuggingFace API Sparse Embedders -> Hugging Face TEI Sparse Embedders (#3870)
|
|
8
|
+
|
|
9
|
+
### 🧪 Testing
|
|
10
|
+
|
|
11
|
+
- HF API - use Groq for gpt-oss-20b (#4001)
|
|
12
|
+
|
|
13
|
+
|
|
14
|
+
## [integrations/huggingface_api-v1.0.0] - 2026-09-07
|
|
15
|
+
|
|
16
|
+
### 🚀 Features
|
|
17
|
+
|
|
18
|
+
- [**breaking**] HF API - add lifecycle handling (#3917)
|
|
19
|
+
|
|
20
|
+
### 🧪 Testing
|
|
21
|
+
|
|
22
|
+
- Replace retired Qwen2.5 models in `huggingface_api` live tests (#3866)
|
|
23
|
+
|
|
24
|
+
### ⚙️ CI
|
|
25
|
+
|
|
26
|
+
- Improve changelog generation; fix existing changelogs (#3883)
|
|
27
|
+
|
|
28
|
+
|
|
3
29
|
## [integrations/huggingface_api-v0.6.0] - 2026-08-27
|
|
4
30
|
|
|
5
31
|
### 🚀 Features
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.5
|
|
2
2
|
Name: huggingface-api-haystack
|
|
3
|
-
Version: 1.0
|
|
3
|
+
Version: 2.1.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
|
|
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
|
|
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
|
|
6
|
-
from .sparse_text_embedder import
|
|
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
|
|
123
|
-
`TEXT_EMBEDDINGS_INFERENCE`.
|
|
124
|
-
:param token: The Hugging Face token to use as
|
|
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
|
|
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
|
|
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
|
|
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
|
|
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
|
-
|
|
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
|
-
|
|
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):
|