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.
- {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-2.0.0}/CHANGELOG.md +35 -1
- {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-2.0.0}/PKG-INFO +6 -5
- {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-2.0.0}/README.md +1 -2
- huggingface_api_haystack-2.0.0/docker-compose.yml +50 -0
- {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-2.0.0}/pyproject.toml +11 -1
- {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-2.0.0}/src/haystack_integrations/common/huggingface_api/utils.py +66 -8
- {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-2.0.0}/src/haystack_integrations/components/embedders/huggingface_api/__init__.py +4 -4
- {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-2.0.0}/src/haystack_integrations/components/embedders/huggingface_api/document_embedder.py +167 -17
- {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
- huggingface_api_haystack-2.0.0/src/haystack_integrations/components/embedders/huggingface_api/sparse_text_embedder.py +191 -0
- {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-2.0.0}/src/haystack_integrations/components/embedders/huggingface_api/text_embedder.py +102 -9
- {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
- {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-2.0.0}/src/haystack_integrations/components/rankers/huggingface_api/ranker.py +76 -3
- {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-2.0.0}/tests/test_chat_generator.py +199 -82
- {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-2.0.0}/tests/test_document_embedder.py +319 -26
- {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-2.0.0}/tests/test_ranker.py +174 -3
- huggingface_api_haystack-2.0.0/tests/test_sparse_document_embedder.py +621 -0
- huggingface_api_haystack-2.0.0/tests/test_sparse_text_embedder.py +499 -0
- huggingface_api_haystack-2.0.0/tests/test_text_embedder.py +502 -0
- huggingface_api_haystack-2.0.0/tests/test_utils.py +128 -0
- huggingface_api_haystack-0.6.0/docker-compose.yml +0 -14
- huggingface_api_haystack-0.6.0/src/haystack_integrations/components/embedders/huggingface_api/sparse_text_embedder.py +0 -116
- huggingface_api_haystack-0.6.0/tests/test_sparse_document_embedder.py +0 -342
- huggingface_api_haystack-0.6.0/tests/test_sparse_text_embedder.py +0 -263
- huggingface_api_haystack-0.6.0/tests/test_text_embedder.py +0 -257
- huggingface_api_haystack-0.6.0/tests/test_utils.py +0 -49
- {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-2.0.0}/.gitignore +0 -0
- {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-2.0.0}/LICENSE.txt +0 -0
- {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-2.0.0}/pydoc/config_docusaurus.yml +0 -0
- {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-2.0.0}/src/haystack_integrations/common/huggingface_api/__init__.py +0 -0
- {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-2.0.0}/src/haystack_integrations/common/py.typed +0 -0
- {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
- {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-2.0.0}/src/haystack_integrations/components/embedders/py.typed +0 -0
- {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-2.0.0}/src/haystack_integrations/components/generators/huggingface_api/__init__.py +0 -0
- {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-2.0.0}/src/haystack_integrations/components/generators/huggingface_api/chat/__init__.py +0 -0
- {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-2.0.0}/src/haystack_integrations/components/generators/py.typed +0 -0
- {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-2.0.0}/src/haystack_integrations/components/rankers/huggingface_api/__init__.py +0 -0
- {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-2.0.0}/src/haystack_integrations/components/rankers/py.typed +0 -0
- {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-2.0.0}/tests/__init__.py +0 -0
- {huggingface_api_haystack-0.6.0 → huggingface_api_haystack-2.0.0}/tests/conftest.py +0 -0
- {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
|
-
## [
|
|
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.
|
|
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>=
|
|
22
|
+
Requires-Dist: haystack-ai>=3.0.0
|
|
23
23
|
Requires-Dist: httpx>=0.27.0
|
|
24
|
-
Requires-Dist: huggingface-hub[inference]>=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
|
|
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:
|
|
@@ -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>=
|
|
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
|
|
87
|
+
Check if the provided model ID corresponds to a valid model on Hugging Face Hub.
|
|
86
88
|
|
|
87
|
-
Also check if the model
|
|
89
|
+
Also check if the model supports the requested embedding or generation task.
|
|
88
90
|
|
|
89
|
-
:param model_id: A string representing the
|
|
90
|
-
:param model_type:
|
|
91
|
-
:param token: The optional authentication
|
|
92
|
-
:raises ValueError: If the model is not found or
|
|
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 =
|
|
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 =
|
|
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
|
|
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,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
|
|
122
|
-
`TEXT_EMBEDDINGS_INFERENCE`.
|
|
123
|
-
: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.
|
|
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
|
-
|
|
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.
|
|
197
|
-
self.
|
|
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
|
-
|
|
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
|
-
|
|
303
|
-
|
|
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
|
-
|
|
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
|
-
|
|
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):
|