nvidia-haystack 0.1.0__tar.gz → 0.1.2__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.
- {nvidia_haystack-0.1.0 → nvidia_haystack-0.1.2}/CHANGELOG.md +11 -1
- {nvidia_haystack-0.1.0 → nvidia_haystack-0.1.2}/PKG-INFO +3 -2
- {nvidia_haystack-0.1.0 → nvidia_haystack-0.1.2}/README.md +1 -1
- {nvidia_haystack-0.1.0 → nvidia_haystack-0.1.2}/pyproject.toml +1 -1
- nvidia_haystack-0.1.2/src/haystack_integrations/components/__init__.py +3 -0
- nvidia_haystack-0.1.2/src/haystack_integrations/components/embedders/__init__.py +3 -0
- nvidia_haystack-0.1.2/src/haystack_integrations/components/embedders/nvidia/__init__.py +9 -0
- {nvidia_haystack-0.1.0 → nvidia_haystack-0.1.2}/src/haystack_integrations/components/embedders/nvidia/document_embedder.py +25 -8
- {nvidia_haystack-0.1.0 → nvidia_haystack-0.1.2}/src/haystack_integrations/components/embedders/nvidia/text_embedder.py +29 -5
- {nvidia_haystack-0.1.0 → nvidia_haystack-0.1.2}/src/haystack_integrations/components/embedders/nvidia/truncate.py +4 -0
- nvidia_haystack-0.1.2/src/haystack_integrations/components/generators/__init__.py +3 -0
- {nvidia_haystack-0.1.0 → nvidia_haystack-0.1.2}/src/haystack_integrations/components/generators/nvidia/__init__.py +1 -0
- nvidia_haystack-0.1.2/src/haystack_integrations/components/generators/nvidia/chat/__init__.py +3 -0
- {nvidia_haystack-0.1.0 → nvidia_haystack-0.1.2}/src/haystack_integrations/components/generators/nvidia/generator.py +11 -1
- nvidia_haystack-0.1.2/src/haystack_integrations/components/rankers/__init__.py +3 -0
- nvidia_haystack-0.1.2/src/haystack_integrations/components/rankers/nvidia/__init__.py +7 -0
- {nvidia_haystack-0.1.0 → nvidia_haystack-0.1.2}/src/haystack_integrations/components/rankers/nvidia/ranker.py +91 -33
- {nvidia_haystack-0.1.0 → nvidia_haystack-0.1.2}/src/haystack_integrations/components/rankers/nvidia/truncate.py +4 -0
- nvidia_haystack-0.1.2/src/haystack_integrations/utils/__init__.py +3 -0
- nvidia_haystack-0.1.2/src/haystack_integrations/utils/nvidia/__init__.py +8 -0
- {nvidia_haystack-0.1.0 → nvidia_haystack-0.1.2}/src/haystack_integrations/utils/nvidia/nim_backend.py +29 -11
- {nvidia_haystack-0.1.0 → nvidia_haystack-0.1.2}/src/haystack_integrations/utils/nvidia/utils.py +6 -2
- {nvidia_haystack-0.1.0 → nvidia_haystack-0.1.2}/tests/__init__.py +1 -0
- {nvidia_haystack-0.1.0 → nvidia_haystack-0.1.2}/tests/conftest.py +4 -0
- {nvidia_haystack-0.1.0 → nvidia_haystack-0.1.2}/tests/test_base_url.py +4 -0
- {nvidia_haystack-0.1.0 → nvidia_haystack-0.1.2}/tests/test_document_embedder.py +47 -5
- {nvidia_haystack-0.1.0 → nvidia_haystack-0.1.2}/tests/test_embedding_truncate_mode.py +4 -0
- {nvidia_haystack-0.1.0 → nvidia_haystack-0.1.2}/tests/test_generator.py +14 -0
- {nvidia_haystack-0.1.0 → nvidia_haystack-0.1.2}/tests/test_ranker.py +112 -7
- {nvidia_haystack-0.1.0 → nvidia_haystack-0.1.2}/tests/test_text_embedder.py +37 -2
- nvidia_haystack-0.1.0/src/haystack_integrations/components/embedders/nvidia/__init__.py +0 -5
- nvidia_haystack-0.1.0/src/haystack_integrations/components/rankers/nvidia/__init__.py +0 -3
- nvidia_haystack-0.1.0/src/haystack_integrations/utils/nvidia/__init__.py +0 -4
- {nvidia_haystack-0.1.0 → nvidia_haystack-0.1.2}/.gitignore +0 -0
- {nvidia_haystack-0.1.0 → nvidia_haystack-0.1.2}/LICENSE.txt +0 -0
- {nvidia_haystack-0.1.0 → nvidia_haystack-0.1.2}/pydoc/config.yml +0 -0
- {nvidia_haystack-0.1.0/src/haystack_integrations/components/generators/nvidia/chat → nvidia_haystack-0.1.2/src/haystack_integrations}/__init__.py +0 -0
|
@@ -1,15 +1,23 @@
|
|
|
1
1
|
# Changelog
|
|
2
2
|
|
|
3
|
-
## [
|
|
3
|
+
## [integrations/nvidia-v0.1.1] - 2024-11-14
|
|
4
|
+
|
|
5
|
+
### 🐛 Bug Fixes
|
|
6
|
+
|
|
7
|
+
- Fixes to NvidiaRanker (#1191)
|
|
8
|
+
|
|
9
|
+
## [integrations/nvidia-v0.1.0] - 2024-11-13
|
|
4
10
|
|
|
5
11
|
### 🚀 Features
|
|
6
12
|
|
|
7
13
|
- Update default embedding model to nvidia/nv-embedqa-e5-v5 (#1015)
|
|
8
14
|
- Add NVIDIA NIM ranker support (#1023)
|
|
15
|
+
- Raise error when attempting to embed empty documents/strings with Nvidia embedders (#1118)
|
|
9
16
|
|
|
10
17
|
### 🐛 Bug Fixes
|
|
11
18
|
|
|
12
19
|
- Lints in `nvidia-haystack` (#993)
|
|
20
|
+
- Missing Nvidia embedding truncate mode (#1043)
|
|
13
21
|
|
|
14
22
|
### 🚜 Refactor
|
|
15
23
|
|
|
@@ -27,6 +35,8 @@
|
|
|
27
35
|
|
|
28
36
|
- Retry tests to reduce flakyness (#836)
|
|
29
37
|
- Update ruff invocation to include check parameter (#853)
|
|
38
|
+
- Update ruff linting scripts and settings (#1105)
|
|
39
|
+
- Adopt uv as installer (#1142)
|
|
30
40
|
|
|
31
41
|
### Docs
|
|
32
42
|
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.3
|
|
2
2
|
Name: nvidia-haystack
|
|
3
|
-
Version: 0.1.
|
|
3
|
+
Version: 0.1.2
|
|
4
4
|
Project-URL: Documentation, https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/nvidia#readme
|
|
5
5
|
Project-URL: Issues, https://github.com/deepset-ai/haystack-core-integrations/issues
|
|
6
6
|
Project-URL: Source, https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/nvidia
|
|
@@ -19,6 +19,7 @@ Classifier: Programming Language :: Python :: Implementation :: PyPy
|
|
|
19
19
|
Requires-Python: >=3.8
|
|
20
20
|
Requires-Dist: haystack-ai
|
|
21
21
|
Requires-Dist: requests
|
|
22
|
+
Requires-Dist: tqdm
|
|
22
23
|
Description-Content-Type: text/markdown
|
|
23
24
|
|
|
24
25
|
# nvidia-haystack
|
|
@@ -61,7 +62,7 @@ hatch run test
|
|
|
61
62
|
To only run unit tests:
|
|
62
63
|
|
|
63
64
|
```
|
|
64
|
-
hatch run test -m"not integration"
|
|
65
|
+
hatch run test -m "not integration"
|
|
65
66
|
```
|
|
66
67
|
|
|
67
68
|
To run the linters `ruff` and `mypy`:
|
|
@@ -23,7 +23,7 @@ classifiers = [
|
|
|
23
23
|
"Programming Language :: Python :: Implementation :: CPython",
|
|
24
24
|
"Programming Language :: Python :: Implementation :: PyPy",
|
|
25
25
|
]
|
|
26
|
-
dependencies = ["haystack-ai", "requests"]
|
|
26
|
+
dependencies = ["haystack-ai", "requests", "tqdm"]
|
|
27
27
|
|
|
28
28
|
[project.urls]
|
|
29
29
|
Documentation = "https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/nvidia#readme"
|
|
@@ -0,0 +1,9 @@
|
|
|
1
|
+
# SPDX-FileCopyrightText: 2024-present deepset GmbH <info@deepset.ai>
|
|
2
|
+
#
|
|
3
|
+
# SPDX-License-Identifier: Apache-2.0
|
|
4
|
+
|
|
5
|
+
from .document_embedder import NvidiaDocumentEmbedder
|
|
6
|
+
from .text_embedder import NvidiaTextEmbedder
|
|
7
|
+
from .truncate import EmbeddingTruncateMode
|
|
8
|
+
|
|
9
|
+
__all__ = ["EmbeddingTruncateMode", "NvidiaDocumentEmbedder", "NvidiaTextEmbedder"]
|
|
@@ -1,13 +1,19 @@
|
|
|
1
|
+
# SPDX-FileCopyrightText: 2024-present deepset GmbH <info@deepset.ai>
|
|
2
|
+
#
|
|
3
|
+
# SPDX-License-Identifier: Apache-2.0
|
|
4
|
+
|
|
5
|
+
import os
|
|
1
6
|
import warnings
|
|
2
7
|
from typing import Any, Dict, List, Optional, Tuple, Union
|
|
3
8
|
|
|
4
|
-
from haystack import Document, component, default_from_dict, default_to_dict
|
|
9
|
+
from haystack import Document, component, default_from_dict, default_to_dict, logging
|
|
5
10
|
from haystack.utils import Secret, deserialize_secrets_inplace
|
|
6
11
|
from tqdm import tqdm
|
|
7
12
|
|
|
13
|
+
from haystack_integrations.components.embedders.nvidia.truncate import EmbeddingTruncateMode
|
|
8
14
|
from haystack_integrations.utils.nvidia import NimBackend, is_hosted, url_validation
|
|
9
15
|
|
|
10
|
-
|
|
16
|
+
logger = logging.getLogger(__name__)
|
|
11
17
|
|
|
12
18
|
_DEFAULT_API_URL = "https://ai.api.nvidia.com/v1/retrieval/nvidia"
|
|
13
19
|
|
|
@@ -44,6 +50,7 @@ class NvidiaDocumentEmbedder:
|
|
|
44
50
|
meta_fields_to_embed: Optional[List[str]] = None,
|
|
45
51
|
embedding_separator: str = "\n",
|
|
46
52
|
truncate: Optional[Union[EmbeddingTruncateMode, str]] = None,
|
|
53
|
+
timeout: Optional[float] = None,
|
|
47
54
|
):
|
|
48
55
|
"""
|
|
49
56
|
Create a NvidiaTextEmbedder component.
|
|
@@ -71,8 +78,11 @@ class NvidiaDocumentEmbedder:
|
|
|
71
78
|
:param embedding_separator:
|
|
72
79
|
Separator used to concatenate the meta fields to the Document text.
|
|
73
80
|
:param truncate:
|
|
74
|
-
Specifies how inputs longer
|
|
81
|
+
Specifies how inputs longer than the maximum token length should be truncated.
|
|
75
82
|
If None the behavior is model-dependent, see the official documentation for more information.
|
|
83
|
+
:param timeout:
|
|
84
|
+
Timeout for request calls, if not set it is inferred from the `NVIDIA_TIMEOUT` environment variable
|
|
85
|
+
or set to 60 by default.
|
|
76
86
|
"""
|
|
77
87
|
|
|
78
88
|
self.api_key = api_key
|
|
@@ -95,6 +105,10 @@ class NvidiaDocumentEmbedder:
|
|
|
95
105
|
if is_hosted(api_url) and not self.model: # manually set default model
|
|
96
106
|
self.model = "nvidia/nv-embedqa-e5-v5"
|
|
97
107
|
|
|
108
|
+
if timeout is None:
|
|
109
|
+
timeout = float(os.environ.get("NVIDIA_TIMEOUT", 60.0))
|
|
110
|
+
self.timeout = timeout
|
|
111
|
+
|
|
98
112
|
def default_model(self):
|
|
99
113
|
"""Set default model in local NIM mode."""
|
|
100
114
|
valid_models = [
|
|
@@ -125,10 +139,11 @@ class NvidiaDocumentEmbedder:
|
|
|
125
139
|
if self.truncate is not None:
|
|
126
140
|
model_kwargs["truncate"] = str(self.truncate)
|
|
127
141
|
self.backend = NimBackend(
|
|
128
|
-
self.model,
|
|
142
|
+
model=self.model,
|
|
129
143
|
api_url=self.api_url,
|
|
130
144
|
api_key=self.api_key,
|
|
131
145
|
model_kwargs=model_kwargs,
|
|
146
|
+
timeout=self.timeout,
|
|
132
147
|
)
|
|
133
148
|
|
|
134
149
|
self._initialized = True
|
|
@@ -155,6 +170,7 @@ class NvidiaDocumentEmbedder:
|
|
|
155
170
|
meta_fields_to_embed=self.meta_fields_to_embed,
|
|
156
171
|
embedding_separator=self.embedding_separator,
|
|
157
172
|
truncate=str(self.truncate) if self.truncate is not None else None,
|
|
173
|
+
timeout=self.timeout,
|
|
158
174
|
)
|
|
159
175
|
|
|
160
176
|
@classmethod
|
|
@@ -167,7 +183,9 @@ class NvidiaDocumentEmbedder:
|
|
|
167
183
|
:returns:
|
|
168
184
|
The deserialized component.
|
|
169
185
|
"""
|
|
170
|
-
|
|
186
|
+
init_parameters = data.get("init_parameters", {})
|
|
187
|
+
if init_parameters:
|
|
188
|
+
deserialize_secrets_inplace(data["init_parameters"], keys=["api_key"])
|
|
171
189
|
return default_from_dict(cls, data)
|
|
172
190
|
|
|
173
191
|
def _prepare_texts_to_embed(self, documents: List[Document]) -> List[str]:
|
|
@@ -224,7 +242,7 @@ class NvidiaDocumentEmbedder:
|
|
|
224
242
|
if not self._initialized:
|
|
225
243
|
msg = "The embedding model has not been loaded. Please call warm_up() before running."
|
|
226
244
|
raise RuntimeError(msg)
|
|
227
|
-
elif not isinstance(documents, list) or documents and not isinstance(documents[0], Document):
|
|
245
|
+
elif not isinstance(documents, list) or (documents and not isinstance(documents[0], Document)):
|
|
228
246
|
msg = (
|
|
229
247
|
"NvidiaDocumentEmbedder expects a list of Documents as input."
|
|
230
248
|
"In case you want to embed a string, please use the NvidiaTextEmbedder."
|
|
@@ -233,8 +251,7 @@ class NvidiaDocumentEmbedder:
|
|
|
233
251
|
|
|
234
252
|
for doc in documents:
|
|
235
253
|
if not doc.content:
|
|
236
|
-
|
|
237
|
-
raise ValueError(msg)
|
|
254
|
+
logger.warning(f"Document '{doc.id}' has no content to embed.")
|
|
238
255
|
|
|
239
256
|
texts_to_embed = self._prepare_texts_to_embed(documents)
|
|
240
257
|
embeddings, metadata = self._embed_batch(texts_to_embed, self.batch_size)
|
|
@@ -1,12 +1,18 @@
|
|
|
1
|
+
# SPDX-FileCopyrightText: 2024-present deepset GmbH <info@deepset.ai>
|
|
2
|
+
#
|
|
3
|
+
# SPDX-License-Identifier: Apache-2.0
|
|
4
|
+
|
|
5
|
+
import os
|
|
1
6
|
import warnings
|
|
2
7
|
from typing import Any, Dict, List, Optional, Union
|
|
3
8
|
|
|
4
|
-
from haystack import component, default_from_dict, default_to_dict
|
|
9
|
+
from haystack import component, default_from_dict, default_to_dict, logging
|
|
5
10
|
from haystack.utils import Secret, deserialize_secrets_inplace
|
|
6
11
|
|
|
12
|
+
from haystack_integrations.components.embedders.nvidia.truncate import EmbeddingTruncateMode
|
|
7
13
|
from haystack_integrations.utils.nvidia import NimBackend, is_hosted, url_validation
|
|
8
14
|
|
|
9
|
-
|
|
15
|
+
logger = logging.getLogger(__name__)
|
|
10
16
|
|
|
11
17
|
_DEFAULT_API_URL = "https://ai.api.nvidia.com/v1/retrieval/nvidia"
|
|
12
18
|
|
|
@@ -41,6 +47,7 @@ class NvidiaTextEmbedder:
|
|
|
41
47
|
prefix: str = "",
|
|
42
48
|
suffix: str = "",
|
|
43
49
|
truncate: Optional[Union[EmbeddingTruncateMode, str]] = None,
|
|
50
|
+
timeout: Optional[float] = None,
|
|
44
51
|
):
|
|
45
52
|
"""
|
|
46
53
|
Create a NvidiaTextEmbedder component.
|
|
@@ -61,6 +68,9 @@ class NvidiaTextEmbedder:
|
|
|
61
68
|
:param truncate:
|
|
62
69
|
Specifies how inputs longer that the maximum token length should be truncated.
|
|
63
70
|
If None the behavior is model-dependent, see the official documentation for more information.
|
|
71
|
+
:param timeout:
|
|
72
|
+
Timeout for request calls, if not set it is inferred from the `NVIDIA_TIMEOUT` environment variable
|
|
73
|
+
or set to 60 by default.
|
|
64
74
|
"""
|
|
65
75
|
|
|
66
76
|
self.api_key = api_key
|
|
@@ -79,6 +89,10 @@ class NvidiaTextEmbedder:
|
|
|
79
89
|
if is_hosted(api_url) and not self.model: # manually set default model
|
|
80
90
|
self.model = "nvidia/nv-embedqa-e5-v5"
|
|
81
91
|
|
|
92
|
+
if timeout is None:
|
|
93
|
+
timeout = float(os.environ.get("NVIDIA_TIMEOUT", 60.0))
|
|
94
|
+
self.timeout = timeout
|
|
95
|
+
|
|
82
96
|
def default_model(self):
|
|
83
97
|
"""Set default model in local NIM mode."""
|
|
84
98
|
valid_models = [
|
|
@@ -86,6 +100,12 @@ class NvidiaTextEmbedder:
|
|
|
86
100
|
]
|
|
87
101
|
name = next(iter(valid_models), None)
|
|
88
102
|
if name:
|
|
103
|
+
logger.warning(
|
|
104
|
+
"Default model is set as: {model_name}. \n"
|
|
105
|
+
"Set model using model parameter. \n"
|
|
106
|
+
"To get available models use available_models property.",
|
|
107
|
+
model_name=name,
|
|
108
|
+
)
|
|
89
109
|
warnings.warn(
|
|
90
110
|
f"Default model is set as: {name}. \n"
|
|
91
111
|
"Set model using model parameter. \n"
|
|
@@ -109,10 +129,11 @@ class NvidiaTextEmbedder:
|
|
|
109
129
|
if self.truncate is not None:
|
|
110
130
|
model_kwargs["truncate"] = str(self.truncate)
|
|
111
131
|
self.backend = NimBackend(
|
|
112
|
-
self.model,
|
|
132
|
+
model=self.model,
|
|
113
133
|
api_url=self.api_url,
|
|
114
134
|
api_key=self.api_key,
|
|
115
135
|
model_kwargs=model_kwargs,
|
|
136
|
+
timeout=self.timeout,
|
|
116
137
|
)
|
|
117
138
|
|
|
118
139
|
self._initialized = True
|
|
@@ -135,6 +156,7 @@ class NvidiaTextEmbedder:
|
|
|
135
156
|
prefix=self.prefix,
|
|
136
157
|
suffix=self.suffix,
|
|
137
158
|
truncate=str(self.truncate) if self.truncate is not None else None,
|
|
159
|
+
timeout=self.timeout,
|
|
138
160
|
)
|
|
139
161
|
|
|
140
162
|
@classmethod
|
|
@@ -147,7 +169,9 @@ class NvidiaTextEmbedder:
|
|
|
147
169
|
:returns:
|
|
148
170
|
The deserialized component.
|
|
149
171
|
"""
|
|
150
|
-
|
|
172
|
+
init_parameters = data.get("init_parameters", {})
|
|
173
|
+
if init_parameters:
|
|
174
|
+
deserialize_secrets_inplace(data["init_parameters"], keys=["api_key"])
|
|
151
175
|
return default_from_dict(cls, data)
|
|
152
176
|
|
|
153
177
|
@component.output_types(embedding=List[float], meta=Dict[str, Any])
|
|
@@ -159,7 +183,7 @@ class NvidiaTextEmbedder:
|
|
|
159
183
|
The text to embed.
|
|
160
184
|
:returns:
|
|
161
185
|
A dictionary with the following keys and values:
|
|
162
|
-
- `embedding` -
|
|
186
|
+
- `embedding` - Embedding of the text.
|
|
163
187
|
- `meta` - Metadata on usage statistics, etc.
|
|
164
188
|
:raises RuntimeError:
|
|
165
189
|
If the component was not initialized.
|
|
@@ -1,6 +1,8 @@
|
|
|
1
1
|
# SPDX-FileCopyrightText: 2024-present deepset GmbH <info@deepset.ai>
|
|
2
2
|
#
|
|
3
3
|
# SPDX-License-Identifier: Apache-2.0
|
|
4
|
+
|
|
5
|
+
import os
|
|
4
6
|
import warnings
|
|
5
7
|
from typing import Any, Dict, List, Optional
|
|
6
8
|
|
|
@@ -48,6 +50,7 @@ class NvidiaGenerator:
|
|
|
48
50
|
api_url: str = _DEFAULT_API_URL,
|
|
49
51
|
api_key: Optional[Secret] = Secret.from_env_var("NVIDIA_API_KEY"),
|
|
50
52
|
model_arguments: Optional[Dict[str, Any]] = None,
|
|
53
|
+
timeout: Optional[float] = None,
|
|
51
54
|
):
|
|
52
55
|
"""
|
|
53
56
|
Create a NvidiaGenerator component.
|
|
@@ -69,6 +72,9 @@ class NvidiaGenerator:
|
|
|
69
72
|
specific to a model.
|
|
70
73
|
Search your model in the [NVIDIA NIM](https://ai.nvidia.com)
|
|
71
74
|
to find the arguments it accepts.
|
|
75
|
+
:param timeout:
|
|
76
|
+
Timeout for request calls, if not set it is inferred from the `NVIDIA_TIMEOUT` environment variable
|
|
77
|
+
or set to 60 by default.
|
|
72
78
|
"""
|
|
73
79
|
self._model = model
|
|
74
80
|
self._api_url = url_validation(api_url, _DEFAULT_API_URL, ["v1/chat/completions"])
|
|
@@ -78,6 +84,9 @@ class NvidiaGenerator:
|
|
|
78
84
|
self._backend: Optional[Any] = None
|
|
79
85
|
|
|
80
86
|
self.is_hosted = is_hosted(api_url)
|
|
87
|
+
if timeout is None:
|
|
88
|
+
timeout = float(os.environ.get("NVIDIA_TIMEOUT", 60.0))
|
|
89
|
+
self.timeout = timeout
|
|
81
90
|
|
|
82
91
|
def default_model(self):
|
|
83
92
|
"""Set default model in local NIM mode."""
|
|
@@ -109,10 +118,11 @@ class NvidiaGenerator:
|
|
|
109
118
|
msg = "API key is required for hosted NVIDIA NIMs."
|
|
110
119
|
raise ValueError(msg)
|
|
111
120
|
self._backend = NimBackend(
|
|
112
|
-
self._model,
|
|
121
|
+
model=self._model,
|
|
113
122
|
api_url=self._api_url,
|
|
114
123
|
api_key=self._api_key,
|
|
115
124
|
model_kwargs=self._model_arguments,
|
|
125
|
+
timeout=self.timeout,
|
|
116
126
|
)
|
|
117
127
|
|
|
118
128
|
if not self.is_hosted and not self._model:
|
|
@@ -1,12 +1,18 @@
|
|
|
1
|
+
# SPDX-FileCopyrightText: 2024-present deepset GmbH <info@deepset.ai>
|
|
2
|
+
#
|
|
3
|
+
# SPDX-License-Identifier: Apache-2.0
|
|
4
|
+
|
|
5
|
+
import os
|
|
1
6
|
import warnings
|
|
2
7
|
from typing import Any, Dict, List, Optional, Union
|
|
3
8
|
|
|
4
|
-
from haystack import Document, component, default_from_dict, default_to_dict
|
|
9
|
+
from haystack import Document, component, default_from_dict, default_to_dict, logging
|
|
5
10
|
from haystack.utils import Secret, deserialize_secrets_inplace
|
|
6
11
|
|
|
12
|
+
from haystack_integrations.components.rankers.nvidia.truncate import RankerTruncateMode
|
|
7
13
|
from haystack_integrations.utils.nvidia import NimBackend, url_validation
|
|
8
14
|
|
|
9
|
-
|
|
15
|
+
logger = logging.getLogger(__name__)
|
|
10
16
|
|
|
11
17
|
_DEFAULT_MODEL = "nvidia/nv-rerankqa-mistral-4b-v3"
|
|
12
18
|
|
|
@@ -51,8 +57,13 @@ class NvidiaRanker:
|
|
|
51
57
|
model: Optional[str] = None,
|
|
52
58
|
truncate: Optional[Union[RankerTruncateMode, str]] = None,
|
|
53
59
|
api_url: Optional[str] = None,
|
|
54
|
-
api_key: Optional[Secret] =
|
|
60
|
+
api_key: Optional[Secret] = Secret.from_env_var("NVIDIA_API_KEY"),
|
|
55
61
|
top_k: int = 5,
|
|
62
|
+
query_prefix: str = "",
|
|
63
|
+
document_prefix: str = "",
|
|
64
|
+
meta_fields_to_embed: Optional[List[str]] = None,
|
|
65
|
+
embedding_separator: str = "\n",
|
|
66
|
+
timeout: Optional[float] = None,
|
|
56
67
|
):
|
|
57
68
|
"""
|
|
58
69
|
Create a NvidiaRanker component.
|
|
@@ -67,6 +78,19 @@ class NvidiaRanker:
|
|
|
67
78
|
Custom API URL for the NVIDIA NIM.
|
|
68
79
|
:param top_k:
|
|
69
80
|
Number of documents to return.
|
|
81
|
+
:param query_prefix:
|
|
82
|
+
A string to add at the beginning of the query text before ranking.
|
|
83
|
+
Use it to prepend the text with an instruction, as required by reranking models like `bge`.
|
|
84
|
+
:param document_prefix:
|
|
85
|
+
A string to add at the beginning of each document before ranking. You can use it to prepend the document
|
|
86
|
+
with an instruction, as required by embedding models like `bge`.
|
|
87
|
+
:param meta_fields_to_embed:
|
|
88
|
+
List of metadata fields to embed with the document.
|
|
89
|
+
:param embedding_separator:
|
|
90
|
+
Separator to concatenate metadata fields to the document.
|
|
91
|
+
:param timeout:
|
|
92
|
+
Timeout for request calls, if not set it is inferred from the `NVIDIA_TIMEOUT` environment variable
|
|
93
|
+
or set to 60 by default.
|
|
70
94
|
"""
|
|
71
95
|
if model is not None and not isinstance(model, str):
|
|
72
96
|
msg = "Ranker expects the `model` parameter to be a string."
|
|
@@ -81,25 +105,34 @@ class NvidiaRanker:
|
|
|
81
105
|
raise TypeError(msg)
|
|
82
106
|
|
|
83
107
|
# todo: detect default in non-hosted case (when api_url is provided)
|
|
84
|
-
self.
|
|
85
|
-
self.
|
|
86
|
-
self.
|
|
108
|
+
self.model = model or _DEFAULT_MODEL
|
|
109
|
+
self.truncate = truncate
|
|
110
|
+
self.api_key = api_key
|
|
87
111
|
# if no api_url is provided, we're using a hosted model and can
|
|
88
112
|
# - assume the default url will work, because there's only one model
|
|
89
113
|
# - assume we won't call backend.models()
|
|
90
114
|
if api_url is not None:
|
|
91
|
-
self.
|
|
92
|
-
self.
|
|
115
|
+
self.api_url = url_validation(api_url, None, ["v1/ranking"])
|
|
116
|
+
self.endpoint = None # we let backend.rank() handle the endpoint
|
|
93
117
|
else:
|
|
94
|
-
if self.
|
|
118
|
+
if self.model not in _MODEL_ENDPOINT_MAP:
|
|
95
119
|
msg = f"Model '{model}' is unknown. Please provide an api_url to access it."
|
|
96
120
|
raise ValueError(msg)
|
|
97
|
-
self.
|
|
98
|
-
self.
|
|
121
|
+
self.api_url = None # we handle the endpoint
|
|
122
|
+
self.endpoint = _MODEL_ENDPOINT_MAP[self.model]
|
|
99
123
|
if api_key is None:
|
|
100
124
|
self._api_key = Secret.from_env_var("NVIDIA_API_KEY")
|
|
101
|
-
self.
|
|
125
|
+
self.top_k = top_k
|
|
102
126
|
self._initialized = False
|
|
127
|
+
self._backend: Optional[Any] = None
|
|
128
|
+
|
|
129
|
+
self.query_prefix = query_prefix
|
|
130
|
+
self.document_prefix = document_prefix
|
|
131
|
+
self.meta_fields_to_embed = meta_fields_to_embed or []
|
|
132
|
+
self.embedding_separator = embedding_separator
|
|
133
|
+
if timeout is None:
|
|
134
|
+
timeout = float(os.environ.get("NVIDIA_TIMEOUT", 60.0))
|
|
135
|
+
self.timeout = timeout
|
|
103
136
|
|
|
104
137
|
def to_dict(self) -> Dict[str, Any]:
|
|
105
138
|
"""
|
|
@@ -109,11 +142,16 @@ class NvidiaRanker:
|
|
|
109
142
|
"""
|
|
110
143
|
return default_to_dict(
|
|
111
144
|
self,
|
|
112
|
-
model=self.
|
|
113
|
-
top_k=self.
|
|
114
|
-
truncate=self.
|
|
115
|
-
api_url=self.
|
|
116
|
-
api_key=self.
|
|
145
|
+
model=self.model,
|
|
146
|
+
top_k=self.top_k,
|
|
147
|
+
truncate=self.truncate,
|
|
148
|
+
api_url=self.api_url,
|
|
149
|
+
api_key=self.api_key.to_dict() if self.api_key else None,
|
|
150
|
+
query_prefix=self.query_prefix,
|
|
151
|
+
document_prefix=self.document_prefix,
|
|
152
|
+
meta_fields_to_embed=self.meta_fields_to_embed,
|
|
153
|
+
embedding_separator=self.embedding_separator,
|
|
154
|
+
timeout=self.timeout,
|
|
117
155
|
)
|
|
118
156
|
|
|
119
157
|
@classmethod
|
|
@@ -124,7 +162,9 @@ class NvidiaRanker:
|
|
|
124
162
|
:param data: A dictionary containing the ranker's attributes.
|
|
125
163
|
:returns: The deserialized ranker.
|
|
126
164
|
"""
|
|
127
|
-
|
|
165
|
+
init_parameters = data.get("init_parameters", {})
|
|
166
|
+
if init_parameters:
|
|
167
|
+
deserialize_secrets_inplace(data["init_parameters"], keys=["api_key"])
|
|
128
168
|
return default_from_dict(cls, data)
|
|
129
169
|
|
|
130
170
|
def warm_up(self):
|
|
@@ -135,18 +175,31 @@ class NvidiaRanker:
|
|
|
135
175
|
"""
|
|
136
176
|
if not self._initialized:
|
|
137
177
|
model_kwargs = {}
|
|
138
|
-
if self.
|
|
139
|
-
model_kwargs.update(truncate=str(self.
|
|
178
|
+
if self.truncate is not None:
|
|
179
|
+
model_kwargs.update(truncate=str(self.truncate))
|
|
140
180
|
self._backend = NimBackend(
|
|
141
|
-
self.
|
|
142
|
-
api_url=self.
|
|
143
|
-
api_key=self.
|
|
181
|
+
model=self.model,
|
|
182
|
+
api_url=self.api_url,
|
|
183
|
+
api_key=self.api_key,
|
|
144
184
|
model_kwargs=model_kwargs,
|
|
185
|
+
timeout=self.timeout,
|
|
145
186
|
)
|
|
146
|
-
if not self.
|
|
147
|
-
self.
|
|
187
|
+
if not self.model:
|
|
188
|
+
self.model = _DEFAULT_MODEL
|
|
148
189
|
self._initialized = True
|
|
149
190
|
|
|
191
|
+
def _prepare_documents_to_embed(self, documents: List[Document]) -> List[str]:
|
|
192
|
+
document_texts = []
|
|
193
|
+
for doc in documents:
|
|
194
|
+
meta_values_to_embed = [
|
|
195
|
+
str(doc.meta[key])
|
|
196
|
+
for key in self.meta_fields_to_embed
|
|
197
|
+
if key in doc.meta and doc.meta[key] # noqa: RUF019
|
|
198
|
+
]
|
|
199
|
+
text_to_embed = self.embedding_separator.join([*meta_values_to_embed, doc.content or ""])
|
|
200
|
+
document_texts.append(self.document_prefix + text_to_embed)
|
|
201
|
+
return document_texts
|
|
202
|
+
|
|
150
203
|
@component.output_types(documents=List[Document])
|
|
151
204
|
def run(
|
|
152
205
|
self,
|
|
@@ -170,32 +223,37 @@ class NvidiaRanker:
|
|
|
170
223
|
msg = "The ranker has not been loaded. Please call warm_up() before running."
|
|
171
224
|
raise RuntimeError(msg)
|
|
172
225
|
if not isinstance(query, str):
|
|
173
|
-
msg = "
|
|
226
|
+
msg = "NvidiaRanker expects the `query` parameter to be a string."
|
|
174
227
|
raise TypeError(msg)
|
|
175
228
|
if not isinstance(documents, list):
|
|
176
|
-
msg = "
|
|
229
|
+
msg = "NvidiaRanker expects the `documents` parameter to be a list."
|
|
177
230
|
raise TypeError(msg)
|
|
178
231
|
if not all(isinstance(doc, Document) for doc in documents):
|
|
179
|
-
msg = "
|
|
232
|
+
msg = "NvidiaRanker expects the `documents` parameter to be a list of Document objects."
|
|
180
233
|
raise TypeError(msg)
|
|
181
234
|
if top_k is not None and not isinstance(top_k, int):
|
|
182
|
-
msg = "
|
|
235
|
+
msg = "NvidiaRanker expects the `top_k` parameter to be an integer."
|
|
183
236
|
raise TypeError(msg)
|
|
184
237
|
|
|
185
238
|
if len(documents) == 0:
|
|
186
239
|
return {"documents": []}
|
|
187
240
|
|
|
188
|
-
top_k = top_k if top_k is not None else self.
|
|
241
|
+
top_k = top_k if top_k is not None else self.top_k
|
|
189
242
|
if top_k < 1:
|
|
243
|
+
logger.warning("top_k should be at least 1, returning nothing")
|
|
190
244
|
warnings.warn("top_k should be at least 1, returning nothing", stacklevel=2)
|
|
191
245
|
return {"documents": []}
|
|
192
246
|
|
|
193
247
|
assert self._backend is not None
|
|
248
|
+
|
|
249
|
+
query_text = self.query_prefix + query
|
|
250
|
+
document_texts = self._prepare_documents_to_embed(documents=documents)
|
|
251
|
+
|
|
194
252
|
# rank result is list[{index: int, logit: float}] sorted by logit
|
|
195
253
|
sorted_indexes_and_scores = self._backend.rank(
|
|
196
|
-
|
|
197
|
-
|
|
198
|
-
endpoint=self.
|
|
254
|
+
query_text=query_text,
|
|
255
|
+
document_texts=document_texts,
|
|
256
|
+
endpoint=self.endpoint,
|
|
199
257
|
)
|
|
200
258
|
sorted_documents = []
|
|
201
259
|
for item in sorted_indexes_and_scores[:top_k]:
|
|
@@ -0,0 +1,8 @@
|
|
|
1
|
+
# SPDX-FileCopyrightText: 2024-present deepset GmbH <info@deepset.ai>
|
|
2
|
+
#
|
|
3
|
+
# SPDX-License-Identifier: Apache-2.0
|
|
4
|
+
|
|
5
|
+
from .nim_backend import Model, NimBackend
|
|
6
|
+
from .utils import is_hosted, url_validation
|
|
7
|
+
|
|
8
|
+
__all__ = ["Model", "NimBackend", "is_hosted", "url_validation"]
|
|
@@ -1,11 +1,18 @@
|
|
|
1
|
+
# SPDX-FileCopyrightText: 2024-present deepset GmbH <info@deepset.ai>
|
|
2
|
+
#
|
|
3
|
+
# SPDX-License-Identifier: Apache-2.0
|
|
4
|
+
|
|
5
|
+
import os
|
|
1
6
|
from dataclasses import dataclass, field
|
|
2
7
|
from typing import Any, Dict, List, Optional, Tuple
|
|
3
8
|
|
|
4
9
|
import requests
|
|
5
|
-
from haystack import
|
|
10
|
+
from haystack import logging
|
|
6
11
|
from haystack.utils import Secret
|
|
7
12
|
|
|
8
|
-
|
|
13
|
+
logger = logging.getLogger(__name__)
|
|
14
|
+
|
|
15
|
+
REQUEST_TIMEOUT = 60.0
|
|
9
16
|
|
|
10
17
|
|
|
11
18
|
@dataclass
|
|
@@ -31,6 +38,7 @@ class NimBackend:
|
|
|
31
38
|
api_url: str,
|
|
32
39
|
api_key: Optional[Secret] = Secret.from_env_var("NVIDIA_API_KEY"),
|
|
33
40
|
model_kwargs: Optional[Dict[str, Any]] = None,
|
|
41
|
+
timeout: Optional[float] = None,
|
|
34
42
|
):
|
|
35
43
|
headers = {
|
|
36
44
|
"Content-Type": "application/json",
|
|
@@ -46,6 +54,9 @@ class NimBackend:
|
|
|
46
54
|
self.model = model
|
|
47
55
|
self.api_url = api_url
|
|
48
56
|
self.model_kwargs = model_kwargs or {}
|
|
57
|
+
if timeout is None:
|
|
58
|
+
timeout = float(os.environ.get("NVIDIA_TIMEOUT", REQUEST_TIMEOUT))
|
|
59
|
+
self.timeout = timeout
|
|
49
60
|
|
|
50
61
|
def embed(self, texts: List[str]) -> Tuple[List[List[float]], Dict[str, Any]]:
|
|
51
62
|
url = f"{self.api_url}/embeddings"
|
|
@@ -58,10 +69,11 @@ class NimBackend:
|
|
|
58
69
|
"input": texts,
|
|
59
70
|
**self.model_kwargs,
|
|
60
71
|
},
|
|
61
|
-
timeout=
|
|
72
|
+
timeout=self.timeout,
|
|
62
73
|
)
|
|
63
74
|
res.raise_for_status()
|
|
64
75
|
except requests.HTTPError as e:
|
|
76
|
+
logger.error("Error when calling NIM embedding endpoint: Error - {error}", error=e.response.text)
|
|
65
77
|
msg = f"Failed to query embedding endpoint: Error - {e.response.text}"
|
|
66
78
|
raise ValueError(msg) from e
|
|
67
79
|
|
|
@@ -90,10 +102,11 @@ class NimBackend:
|
|
|
90
102
|
],
|
|
91
103
|
**self.model_kwargs,
|
|
92
104
|
},
|
|
93
|
-
timeout=
|
|
105
|
+
timeout=self.timeout,
|
|
94
106
|
)
|
|
95
107
|
res.raise_for_status()
|
|
96
108
|
except requests.HTTPError as e:
|
|
109
|
+
logger.error("Error when calling NIM chat completion endpoint: Error - {error}", error=e.response.text)
|
|
97
110
|
msg = f"Failed to query chat completion endpoint: Error - {e.response.text}"
|
|
98
111
|
raise ValueError(msg) from e
|
|
99
112
|
|
|
@@ -128,21 +141,22 @@ class NimBackend:
|
|
|
128
141
|
|
|
129
142
|
res = self.session.get(
|
|
130
143
|
url,
|
|
131
|
-
timeout=
|
|
144
|
+
timeout=self.timeout,
|
|
132
145
|
)
|
|
133
146
|
res.raise_for_status()
|
|
134
147
|
|
|
135
148
|
data = res.json()["data"]
|
|
136
149
|
models = [Model(element["id"]) for element in data if "id" in element]
|
|
137
150
|
if not models:
|
|
151
|
+
logger.error("No hosted model were found at URL '{u}'.", u=url)
|
|
138
152
|
msg = f"No hosted model were found at URL '{url}'."
|
|
139
153
|
raise ValueError(msg)
|
|
140
154
|
return models
|
|
141
155
|
|
|
142
156
|
def rank(
|
|
143
157
|
self,
|
|
144
|
-
|
|
145
|
-
|
|
158
|
+
query_text: str,
|
|
159
|
+
document_texts: List[str],
|
|
146
160
|
endpoint: Optional[str] = None,
|
|
147
161
|
) -> List[Dict[str, Any]]:
|
|
148
162
|
url = endpoint or f"{self.api_url}/ranking"
|
|
@@ -152,18 +166,22 @@ class NimBackend:
|
|
|
152
166
|
url,
|
|
153
167
|
json={
|
|
154
168
|
"model": self.model,
|
|
155
|
-
"query": {"text":
|
|
156
|
-
"passages": [{"text":
|
|
169
|
+
"query": {"text": query_text},
|
|
170
|
+
"passages": [{"text": text} for text in document_texts],
|
|
157
171
|
**self.model_kwargs,
|
|
158
172
|
},
|
|
159
|
-
timeout=
|
|
173
|
+
timeout=self.timeout,
|
|
160
174
|
)
|
|
161
175
|
res.raise_for_status()
|
|
162
176
|
except requests.HTTPError as e:
|
|
177
|
+
logger.error("Error when calling NIM ranking endpoint: Error - {error}", error=e.response.text)
|
|
163
178
|
msg = f"Failed to rank endpoint: Error - {e.response.text}"
|
|
164
179
|
raise ValueError(msg) from e
|
|
165
180
|
|
|
166
181
|
data = res.json()
|
|
167
|
-
|
|
182
|
+
if "rankings" not in data:
|
|
183
|
+
logger.error("Expected 'rankings' in response, got {d}", d=data)
|
|
184
|
+
msg = f"Expected 'rankings' in response, got {data}"
|
|
185
|
+
raise ValueError(msg)
|
|
168
186
|
|
|
169
187
|
return data["rankings"]
|
{nvidia_haystack-0.1.0 → nvidia_haystack-0.1.2}/src/haystack_integrations/utils/nvidia/utils.py
RENAMED
|
@@ -1,9 +1,13 @@
|
|
|
1
|
+
# SPDX-FileCopyrightText: 2024-present deepset GmbH <info@deepset.ai>
|
|
2
|
+
#
|
|
3
|
+
# SPDX-License-Identifier: Apache-2.0
|
|
4
|
+
|
|
1
5
|
import warnings
|
|
2
|
-
from typing import List
|
|
6
|
+
from typing import List, Optional
|
|
3
7
|
from urllib.parse import urlparse, urlunparse
|
|
4
8
|
|
|
5
9
|
|
|
6
|
-
def url_validation(api_url: str, default_api_url: str, allowed_paths: List[str]) -> str:
|
|
10
|
+
def url_validation(api_url: str, default_api_url: Optional[str], allowed_paths: List[str]) -> str:
|
|
7
11
|
"""
|
|
8
12
|
Validate and normalize an API URL.
|
|
9
13
|
|
|
@@ -1,3 +1,7 @@
|
|
|
1
|
+
# SPDX-FileCopyrightText: 2024-present deepset GmbH <info@deepset.ai>
|
|
2
|
+
#
|
|
3
|
+
# SPDX-License-Identifier: Apache-2.0
|
|
4
|
+
|
|
1
5
|
import os
|
|
2
6
|
|
|
3
7
|
import pytest
|
|
@@ -71,6 +75,7 @@ class TestNvidiaDocumentEmbedder:
|
|
|
71
75
|
"meta_fields_to_embed": [],
|
|
72
76
|
"embedding_separator": "\n",
|
|
73
77
|
"truncate": None,
|
|
78
|
+
"timeout": 60.0,
|
|
74
79
|
},
|
|
75
80
|
}
|
|
76
81
|
|
|
@@ -86,6 +91,7 @@ class TestNvidiaDocumentEmbedder:
|
|
|
86
91
|
meta_fields_to_embed=["test_field"],
|
|
87
92
|
embedding_separator=" | ",
|
|
88
93
|
truncate=EmbeddingTruncateMode.END,
|
|
94
|
+
timeout=45.0,
|
|
89
95
|
)
|
|
90
96
|
data = component.to_dict()
|
|
91
97
|
assert data == {
|
|
@@ -101,10 +107,11 @@ class TestNvidiaDocumentEmbedder:
|
|
|
101
107
|
"meta_fields_to_embed": ["test_field"],
|
|
102
108
|
"embedding_separator": " | ",
|
|
103
109
|
"truncate": "END",
|
|
110
|
+
"timeout": 45.0,
|
|
104
111
|
},
|
|
105
112
|
}
|
|
106
113
|
|
|
107
|
-
def
|
|
114
|
+
def test_from_dict(self, monkeypatch):
|
|
108
115
|
monkeypatch.setenv("NVIDIA_API_KEY", "fake-api-key")
|
|
109
116
|
data = {
|
|
110
117
|
"type": "haystack_integrations.components.embedders.nvidia.document_embedder.NvidiaDocumentEmbedder",
|
|
@@ -119,18 +126,38 @@ class TestNvidiaDocumentEmbedder:
|
|
|
119
126
|
"meta_fields_to_embed": ["test_field"],
|
|
120
127
|
"embedding_separator": " | ",
|
|
121
128
|
"truncate": "START",
|
|
129
|
+
"timeout": 45.0,
|
|
122
130
|
},
|
|
123
131
|
}
|
|
124
132
|
component = NvidiaDocumentEmbedder.from_dict(data)
|
|
125
|
-
assert component.model == "
|
|
133
|
+
assert component.model == "playground_nvolveqa_40k"
|
|
126
134
|
assert component.api_url == "https://example.com/v1"
|
|
127
135
|
assert component.prefix == "prefix"
|
|
128
136
|
assert component.suffix == "suffix"
|
|
137
|
+
assert component.batch_size == 10
|
|
138
|
+
assert component.progress_bar is False
|
|
139
|
+
assert component.meta_fields_to_embed == ["test_field"]
|
|
140
|
+
assert component.embedding_separator == " | "
|
|
141
|
+
assert component.truncate == EmbeddingTruncateMode.START
|
|
142
|
+
assert component.timeout == 45.0
|
|
143
|
+
|
|
144
|
+
def test_from_dict_defaults(self, monkeypatch):
|
|
145
|
+
monkeypatch.setenv("NVIDIA_API_KEY", "fake-api-key")
|
|
146
|
+
data = {
|
|
147
|
+
"type": "haystack_integrations.components.embedders.nvidia.document_embedder.NvidiaDocumentEmbedder",
|
|
148
|
+
"init_parameters": {},
|
|
149
|
+
}
|
|
150
|
+
component = NvidiaDocumentEmbedder.from_dict(data)
|
|
151
|
+
assert component.model == "nvidia/nv-embedqa-e5-v5"
|
|
152
|
+
assert component.api_url == "https://ai.api.nvidia.com/v1/retrieval/nvidia"
|
|
153
|
+
assert component.prefix == ""
|
|
154
|
+
assert component.suffix == ""
|
|
129
155
|
assert component.batch_size == 32
|
|
130
156
|
assert component.progress_bar
|
|
131
157
|
assert component.meta_fields_to_embed == []
|
|
132
158
|
assert component.embedding_separator == "\n"
|
|
133
|
-
assert component.truncate
|
|
159
|
+
assert component.truncate is None
|
|
160
|
+
assert component.timeout == 60.0
|
|
134
161
|
|
|
135
162
|
def test_prepare_texts_to_embed_w_metadata(self):
|
|
136
163
|
documents = [
|
|
@@ -326,7 +353,7 @@ class TestNvidiaDocumentEmbedder:
|
|
|
326
353
|
with pytest.raises(TypeError, match="NvidiaDocumentEmbedder expects a list of Documents as input"):
|
|
327
354
|
embedder.run(documents=list_integers_input)
|
|
328
355
|
|
|
329
|
-
def test_run_empty_document(self):
|
|
356
|
+
def test_run_empty_document(self, caplog):
|
|
330
357
|
model = "playground_nvolveqa_40k"
|
|
331
358
|
api_key = Secret.from_token("fake-api-key")
|
|
332
359
|
embedder = NvidiaDocumentEmbedder(model, api_key=api_key)
|
|
@@ -334,8 +361,10 @@ class TestNvidiaDocumentEmbedder:
|
|
|
334
361
|
embedder.warm_up()
|
|
335
362
|
embedder.backend = MockBackend(model=model, api_key=api_key)
|
|
336
363
|
|
|
337
|
-
|
|
364
|
+
# Write check using caplog that a logger.warning is raised
|
|
365
|
+
with caplog.at_level("WARNING"):
|
|
338
366
|
embedder.run(documents=[Document(content="")])
|
|
367
|
+
assert "has no content to embed." in caplog.text
|
|
339
368
|
|
|
340
369
|
def test_run_on_empty_list(self):
|
|
341
370
|
model = "playground_nvolveqa_40k"
|
|
@@ -351,6 +380,19 @@ class TestNvidiaDocumentEmbedder:
|
|
|
351
380
|
assert result["documents"] is not None
|
|
352
381
|
assert not result["documents"] # empty list
|
|
353
382
|
|
|
383
|
+
def test_setting_timeout(self, monkeypatch):
|
|
384
|
+
monkeypatch.setenv("NVIDIA_API_KEY", "fake-api-key")
|
|
385
|
+
embedder = NvidiaDocumentEmbedder(timeout=10.0)
|
|
386
|
+
embedder.warm_up()
|
|
387
|
+
assert embedder.backend.timeout == 10.0
|
|
388
|
+
|
|
389
|
+
def test_setting_timeout_env(self, monkeypatch):
|
|
390
|
+
monkeypatch.setenv("NVIDIA_API_KEY", "fake-api-key")
|
|
391
|
+
monkeypatch.setenv("NVIDIA_TIMEOUT", "45")
|
|
392
|
+
embedder = NvidiaDocumentEmbedder()
|
|
393
|
+
embedder.warm_up()
|
|
394
|
+
assert embedder.backend.timeout == 45.0
|
|
395
|
+
|
|
354
396
|
@pytest.mark.skipif(
|
|
355
397
|
not os.environ.get("NVIDIA_API_KEY", None),
|
|
356
398
|
reason="Export an env var called NVIDIA_API_KEY containing the Nvidia API key to run this test.",
|
|
@@ -1,6 +1,7 @@
|
|
|
1
1
|
# SPDX-FileCopyrightText: 2024-present deepset GmbH <info@deepset.ai>
|
|
2
2
|
#
|
|
3
3
|
# SPDX-License-Identifier: Apache-2.0
|
|
4
|
+
|
|
4
5
|
import os
|
|
5
6
|
|
|
6
7
|
import pytest
|
|
@@ -123,6 +124,19 @@ class TestNvidiaGenerator:
|
|
|
123
124
|
},
|
|
124
125
|
}
|
|
125
126
|
|
|
127
|
+
def test_setting_timeout(self, monkeypatch):
|
|
128
|
+
monkeypatch.setenv("NVIDIA_API_KEY", "fake-api-key")
|
|
129
|
+
generator = NvidiaGenerator(timeout=10.0)
|
|
130
|
+
generator.warm_up()
|
|
131
|
+
assert generator._backend.timeout == 10.0
|
|
132
|
+
|
|
133
|
+
def test_setting_timeout_env(self, monkeypatch):
|
|
134
|
+
monkeypatch.setenv("NVIDIA_API_KEY", "fake-api-key")
|
|
135
|
+
monkeypatch.setenv("NVIDIA_TIMEOUT", "45")
|
|
136
|
+
generator = NvidiaGenerator()
|
|
137
|
+
generator.warm_up()
|
|
138
|
+
assert generator._backend.timeout == 45.0
|
|
139
|
+
|
|
126
140
|
@pytest.mark.skipif(
|
|
127
141
|
not os.environ.get("NVIDIA_NIM_GENERATOR_MODEL", None) or not os.environ.get("NVIDIA_NIM_ENDPOINT_URL", None),
|
|
128
142
|
reason="Export an env var called NVIDIA_NIM_GENERATOR_MODEL containing the hosted model name and "
|
|
@@ -1,3 +1,7 @@
|
|
|
1
|
+
# SPDX-FileCopyrightText: 2024-present deepset GmbH <info@deepset.ai>
|
|
2
|
+
#
|
|
3
|
+
# SPDX-License-Identifier: Apache-2.0
|
|
4
|
+
|
|
1
5
|
import os
|
|
2
6
|
import re
|
|
3
7
|
from typing import Any, Optional, Union
|
|
@@ -15,8 +19,8 @@ class TestNvidiaRanker:
|
|
|
15
19
|
def test_init_default(self, monkeypatch):
|
|
16
20
|
monkeypatch.setenv("NVIDIA_API_KEY", "fake-api-key")
|
|
17
21
|
client = NvidiaRanker()
|
|
18
|
-
assert client.
|
|
19
|
-
assert client.
|
|
22
|
+
assert client.model == _DEFAULT_MODEL
|
|
23
|
+
assert client.api_key == Secret.from_env_var("NVIDIA_API_KEY")
|
|
20
24
|
|
|
21
25
|
def test_init_with_parameters(self):
|
|
22
26
|
client = NvidiaRanker(
|
|
@@ -25,10 +29,10 @@ class TestNvidiaRanker:
|
|
|
25
29
|
top_k=3,
|
|
26
30
|
truncate="END",
|
|
27
31
|
)
|
|
28
|
-
assert client.
|
|
29
|
-
assert client.
|
|
30
|
-
assert client.
|
|
31
|
-
assert client.
|
|
32
|
+
assert client.api_key == Secret.from_token("fake-api-key")
|
|
33
|
+
assert client.model == _DEFAULT_MODEL
|
|
34
|
+
assert client.top_k == 3
|
|
35
|
+
assert client.truncate == RankerTruncateMode.END
|
|
32
36
|
|
|
33
37
|
def test_init_fail_wo_api_key(self, monkeypatch):
|
|
34
38
|
monkeypatch.delenv("NVIDIA_API_KEY", raising=False)
|
|
@@ -39,7 +43,7 @@ class TestNvidiaRanker:
|
|
|
39
43
|
def test_init_pass_wo_api_key_w_api_url(self):
|
|
40
44
|
url = "https://url.bogus/v1"
|
|
41
45
|
client = NvidiaRanker(api_url=url)
|
|
42
|
-
assert client.
|
|
46
|
+
assert client.api_url == url
|
|
43
47
|
|
|
44
48
|
def test_warm_up_required(self):
|
|
45
49
|
client = NvidiaRanker()
|
|
@@ -256,3 +260,104 @@ class TestNvidiaRanker:
|
|
|
256
260
|
backend = client._backend
|
|
257
261
|
client.warm_up()
|
|
258
262
|
assert backend == client._backend
|
|
263
|
+
|
|
264
|
+
def test_to_dict(self) -> None:
|
|
265
|
+
client = NvidiaRanker()
|
|
266
|
+
assert client.to_dict() == {
|
|
267
|
+
"type": "haystack_integrations.components.rankers.nvidia.ranker.NvidiaRanker",
|
|
268
|
+
"init_parameters": {
|
|
269
|
+
"model": "nvidia/nv-rerankqa-mistral-4b-v3",
|
|
270
|
+
"top_k": 5,
|
|
271
|
+
"truncate": None,
|
|
272
|
+
"api_url": None,
|
|
273
|
+
"api_key": {"type": "env_var", "env_vars": ["NVIDIA_API_KEY"], "strict": True},
|
|
274
|
+
"query_prefix": "",
|
|
275
|
+
"document_prefix": "",
|
|
276
|
+
"meta_fields_to_embed": [],
|
|
277
|
+
"embedding_separator": "\n",
|
|
278
|
+
"timeout": 60.0,
|
|
279
|
+
},
|
|
280
|
+
}
|
|
281
|
+
|
|
282
|
+
def test_from_dict(self) -> None:
|
|
283
|
+
client = NvidiaRanker.from_dict(
|
|
284
|
+
{
|
|
285
|
+
"type": "haystack_integrations.components.rankers.nvidia.ranker.NvidiaRanker",
|
|
286
|
+
"init_parameters": {
|
|
287
|
+
"model": "nvidia/nv-rerankqa-mistral-4b-v3",
|
|
288
|
+
"top_k": 5,
|
|
289
|
+
"truncate": None,
|
|
290
|
+
"api_url": None,
|
|
291
|
+
"api_key": {"type": "env_var", "env_vars": ["NVIDIA_API_KEY"], "strict": True},
|
|
292
|
+
"query_prefix": "",
|
|
293
|
+
"document_prefix": "",
|
|
294
|
+
"meta_fields_to_embed": [],
|
|
295
|
+
"embedding_separator": "\n",
|
|
296
|
+
"timeout": 45.0,
|
|
297
|
+
},
|
|
298
|
+
}
|
|
299
|
+
)
|
|
300
|
+
assert client.model == "nvidia/nv-rerankqa-mistral-4b-v3"
|
|
301
|
+
assert client.top_k == 5
|
|
302
|
+
assert client.truncate is None
|
|
303
|
+
assert client.api_url is None
|
|
304
|
+
assert client.api_key == Secret.from_env_var("NVIDIA_API_KEY")
|
|
305
|
+
assert client.query_prefix == ""
|
|
306
|
+
assert client.document_prefix == ""
|
|
307
|
+
assert client.meta_fields_to_embed == []
|
|
308
|
+
assert client.embedding_separator == "\n"
|
|
309
|
+
assert client.timeout == 45.0
|
|
310
|
+
|
|
311
|
+
def test_from_dict_defaults(self) -> None:
|
|
312
|
+
client = NvidiaRanker.from_dict(
|
|
313
|
+
{
|
|
314
|
+
"type": "haystack_integrations.components.rankers.nvidia.ranker.NvidiaRanker",
|
|
315
|
+
"init_parameters": {},
|
|
316
|
+
}
|
|
317
|
+
)
|
|
318
|
+
assert client.model == "nvidia/nv-rerankqa-mistral-4b-v3"
|
|
319
|
+
assert client.top_k == 5
|
|
320
|
+
assert client.truncate is None
|
|
321
|
+
assert client.api_url is None
|
|
322
|
+
assert client.api_key == Secret.from_env_var("NVIDIA_API_KEY")
|
|
323
|
+
assert client.query_prefix == ""
|
|
324
|
+
assert client.document_prefix == ""
|
|
325
|
+
assert client.meta_fields_to_embed == []
|
|
326
|
+
assert client.embedding_separator == "\n"
|
|
327
|
+
assert client.timeout == 60.0
|
|
328
|
+
|
|
329
|
+
def test_setting_timeout(self, monkeypatch):
|
|
330
|
+
monkeypatch.setenv("NVIDIA_API_KEY", "fake-api-key")
|
|
331
|
+
client = NvidiaRanker(timeout=10.0)
|
|
332
|
+
client.warm_up()
|
|
333
|
+
assert client._backend.timeout == 10.0
|
|
334
|
+
|
|
335
|
+
def test_setting_timeout_env(self, monkeypatch):
|
|
336
|
+
monkeypatch.setenv("NVIDIA_API_KEY", "fake-api-key")
|
|
337
|
+
monkeypatch.setenv("NVIDIA_TIMEOUT", "45")
|
|
338
|
+
client = NvidiaRanker()
|
|
339
|
+
client.warm_up()
|
|
340
|
+
assert client._backend.timeout == 45.0
|
|
341
|
+
|
|
342
|
+
def test_prepare_texts_to_embed_w_metadata(self):
|
|
343
|
+
documents = [
|
|
344
|
+
Document(content=f"document number {i}:\ncontent", meta={"meta_field": f"meta_value {i}"}) for i in range(5)
|
|
345
|
+
]
|
|
346
|
+
|
|
347
|
+
ranker = NvidiaRanker(
|
|
348
|
+
model=None,
|
|
349
|
+
api_key=Secret.from_token("fake-api-key"),
|
|
350
|
+
meta_fields_to_embed=["meta_field"],
|
|
351
|
+
embedding_separator=" | ",
|
|
352
|
+
)
|
|
353
|
+
|
|
354
|
+
prepared_texts = ranker._prepare_documents_to_embed(documents)
|
|
355
|
+
|
|
356
|
+
# note that newline is replaced by space
|
|
357
|
+
assert prepared_texts == [
|
|
358
|
+
"meta_value 0 | document number 0:\ncontent",
|
|
359
|
+
"meta_value 1 | document number 1:\ncontent",
|
|
360
|
+
"meta_value 2 | document number 2:\ncontent",
|
|
361
|
+
"meta_value 3 | document number 3:\ncontent",
|
|
362
|
+
"meta_value 4 | document number 4:\ncontent",
|
|
363
|
+
]
|
|
@@ -1,3 +1,7 @@
|
|
|
1
|
+
# SPDX-FileCopyrightText: 2024-present deepset GmbH <info@deepset.ai>
|
|
2
|
+
#
|
|
3
|
+
# SPDX-License-Identifier: Apache-2.0
|
|
4
|
+
|
|
1
5
|
import os
|
|
2
6
|
|
|
3
7
|
import pytest
|
|
@@ -52,6 +56,7 @@ class TestNvidiaTextEmbedder:
|
|
|
52
56
|
"prefix": "",
|
|
53
57
|
"suffix": "",
|
|
54
58
|
"truncate": None,
|
|
59
|
+
"timeout": 60.0,
|
|
55
60
|
},
|
|
56
61
|
}
|
|
57
62
|
|
|
@@ -63,6 +68,7 @@ class TestNvidiaTextEmbedder:
|
|
|
63
68
|
prefix="prefix",
|
|
64
69
|
suffix="suffix",
|
|
65
70
|
truncate=EmbeddingTruncateMode.START,
|
|
71
|
+
timeout=10.0,
|
|
66
72
|
)
|
|
67
73
|
data = component.to_dict()
|
|
68
74
|
assert data == {
|
|
@@ -74,10 +80,11 @@ class TestNvidiaTextEmbedder:
|
|
|
74
80
|
"prefix": "prefix",
|
|
75
81
|
"suffix": "suffix",
|
|
76
82
|
"truncate": "START",
|
|
83
|
+
"timeout": 10.0,
|
|
77
84
|
},
|
|
78
85
|
}
|
|
79
86
|
|
|
80
|
-
def
|
|
87
|
+
def test_from_dict(self, monkeypatch):
|
|
81
88
|
monkeypatch.setenv("NVIDIA_API_KEY", "fake-api-key")
|
|
82
89
|
data = {
|
|
83
90
|
"type": "haystack_integrations.components.embedders.nvidia.text_embedder.NvidiaTextEmbedder",
|
|
@@ -88,6 +95,7 @@ class TestNvidiaTextEmbedder:
|
|
|
88
95
|
"prefix": "prefix",
|
|
89
96
|
"suffix": "suffix",
|
|
90
97
|
"truncate": "START",
|
|
98
|
+
"timeout": 10.0,
|
|
91
99
|
},
|
|
92
100
|
}
|
|
93
101
|
component = NvidiaTextEmbedder.from_dict(data)
|
|
@@ -95,7 +103,21 @@ class TestNvidiaTextEmbedder:
|
|
|
95
103
|
assert component.api_url == "https://example.com/v1"
|
|
96
104
|
assert component.prefix == "prefix"
|
|
97
105
|
assert component.suffix == "suffix"
|
|
98
|
-
assert component.truncate ==
|
|
106
|
+
assert component.truncate == EmbeddingTruncateMode.START
|
|
107
|
+
assert component.timeout == 10.0
|
|
108
|
+
|
|
109
|
+
def test_from_dict_defaults(self, monkeypatch):
|
|
110
|
+
monkeypatch.setenv("NVIDIA_API_KEY", "fake-api-key")
|
|
111
|
+
data = {
|
|
112
|
+
"type": "haystack_integrations.components.embedders.nvidia.text_embedder.NvidiaTextEmbedder",
|
|
113
|
+
"init_parameters": {},
|
|
114
|
+
}
|
|
115
|
+
component = NvidiaTextEmbedder.from_dict(data)
|
|
116
|
+
assert component.model == "nvidia/nv-embedqa-e5-v5"
|
|
117
|
+
assert component.api_url == "https://ai.api.nvidia.com/v1/retrieval/nvidia"
|
|
118
|
+
assert component.prefix == ""
|
|
119
|
+
assert component.suffix == ""
|
|
120
|
+
assert component.truncate is None
|
|
99
121
|
|
|
100
122
|
@pytest.mark.usefixtures("mock_local_models")
|
|
101
123
|
def test_run_default_model(self):
|
|
@@ -158,6 +180,19 @@ class TestNvidiaTextEmbedder:
|
|
|
158
180
|
with pytest.raises(ValueError, match="empty string"):
|
|
159
181
|
embedder.run(text="")
|
|
160
182
|
|
|
183
|
+
def test_setting_timeout(self, monkeypatch):
|
|
184
|
+
monkeypatch.setenv("NVIDIA_API_KEY", "fake-api-key")
|
|
185
|
+
embedder = NvidiaTextEmbedder(timeout=10.0)
|
|
186
|
+
embedder.warm_up()
|
|
187
|
+
assert embedder.backend.timeout == 10.0
|
|
188
|
+
|
|
189
|
+
def test_setting_timeout_env(self, monkeypatch):
|
|
190
|
+
monkeypatch.setenv("NVIDIA_API_KEY", "fake-api-key")
|
|
191
|
+
monkeypatch.setenv("NVIDIA_TIMEOUT", "45")
|
|
192
|
+
embedder = NvidiaTextEmbedder()
|
|
193
|
+
embedder.warm_up()
|
|
194
|
+
assert embedder.backend.timeout == 45.0
|
|
195
|
+
|
|
161
196
|
@pytest.mark.skipif(
|
|
162
197
|
not os.environ.get("NVIDIA_NIM_EMBEDDER_MODEL", None) or not os.environ.get("NVIDIA_NIM_ENDPOINT_URL", None),
|
|
163
198
|
reason="Export an env var called NVIDIA_NIM_EMBEDDER_MODEL containing the hosted model name and "
|
|
File without changes
|
|
File without changes
|
|
File without changes
|