fastembed-haystack 2.5.0__tar.gz → 2.6.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.
- {fastembed_haystack-2.5.0 → fastembed_haystack-2.6.0}/CHANGELOG.md +12 -0
- {fastembed_haystack-2.5.0 → fastembed_haystack-2.6.0}/PKG-INFO +2 -2
- {fastembed_haystack-2.5.0 → fastembed_haystack-2.6.0}/src/haystack_integrations/components/embedders/fastembed/fastembed_document_embedder.py +17 -1
- {fastembed_haystack-2.5.0 → fastembed_haystack-2.6.0}/src/haystack_integrations/components/embedders/fastembed/fastembed_text_embedder.py +17 -1
- {fastembed_haystack-2.5.0 → fastembed_haystack-2.6.0}/src/haystack_integrations/components/rankers/fastembed/late_interaction_ranker.py +23 -0
- {fastembed_haystack-2.5.0 → fastembed_haystack-2.6.0}/src/haystack_integrations/components/rankers/fastembed/ranker.py +23 -0
- {fastembed_haystack-2.5.0 → fastembed_haystack-2.6.0}/tests/test_fastembed_late_interaction_ranker.py +38 -0
- {fastembed_haystack-2.5.0 → fastembed_haystack-2.6.0}/tests/test_fastembed_ranker.py +36 -0
- {fastembed_haystack-2.5.0 → fastembed_haystack-2.6.0}/.gitignore +0 -0
- {fastembed_haystack-2.5.0 → fastembed_haystack-2.6.0}/LICENSE.txt +0 -0
- {fastembed_haystack-2.5.0 → fastembed_haystack-2.6.0}/README.md +0 -0
- {fastembed_haystack-2.5.0 → fastembed_haystack-2.6.0}/examples/example.py +0 -0
- {fastembed_haystack-2.5.0 → fastembed_haystack-2.6.0}/examples/ranker_example.py +0 -0
- {fastembed_haystack-2.5.0 → fastembed_haystack-2.6.0}/examples/sparse_example.py +0 -0
- {fastembed_haystack-2.5.0 → fastembed_haystack-2.6.0}/pydoc/config_docusaurus.yml +0 -0
- {fastembed_haystack-2.5.0 → fastembed_haystack-2.6.0}/pyproject.toml +0 -0
- {fastembed_haystack-2.5.0 → fastembed_haystack-2.6.0}/src/haystack_integrations/components/embedders/fastembed/__init__.py +0 -0
- {fastembed_haystack-2.5.0 → fastembed_haystack-2.6.0}/src/haystack_integrations/components/embedders/fastembed/embedding_backend/__init__.py +0 -0
- {fastembed_haystack-2.5.0 → fastembed_haystack-2.6.0}/src/haystack_integrations/components/embedders/fastembed/embedding_backend/fastembed_backend.py +0 -0
- {fastembed_haystack-2.5.0 → fastembed_haystack-2.6.0}/src/haystack_integrations/components/embedders/fastembed/fastembed_sparse_document_embedder.py +0 -0
- {fastembed_haystack-2.5.0 → fastembed_haystack-2.6.0}/src/haystack_integrations/components/embedders/fastembed/fastembed_sparse_text_embedder.py +0 -0
- {fastembed_haystack-2.5.0 → fastembed_haystack-2.6.0}/src/haystack_integrations/components/embedders/py.typed +0 -0
- {fastembed_haystack-2.5.0 → fastembed_haystack-2.6.0}/src/haystack_integrations/components/rankers/fastembed/__init__.py +0 -0
- {fastembed_haystack-2.5.0 → fastembed_haystack-2.6.0}/src/haystack_integrations/components/rankers/py.typed +0 -0
- {fastembed_haystack-2.5.0 → fastembed_haystack-2.6.0}/tests/__init__.py +0 -0
- {fastembed_haystack-2.5.0 → fastembed_haystack-2.6.0}/tests/test_fastembed_backend.py +0 -0
- {fastembed_haystack-2.5.0 → fastembed_haystack-2.6.0}/tests/test_fastembed_document_embedder.py +0 -0
- {fastembed_haystack-2.5.0 → fastembed_haystack-2.6.0}/tests/test_fastembed_sparse_document_embedder.py +0 -0
- {fastembed_haystack-2.5.0 → fastembed_haystack-2.6.0}/tests/test_fastembed_sparse_text_embedder.py +0 -0
- {fastembed_haystack-2.5.0 → fastembed_haystack-2.6.0}/tests/test_fastembed_text_embedder.py +0 -0
|
@@ -1,5 +1,17 @@
|
|
|
1
1
|
# Changelog
|
|
2
2
|
|
|
3
|
+
## [unreleased]
|
|
4
|
+
|
|
5
|
+
### 🐛 Bug Fixes
|
|
6
|
+
|
|
7
|
+
- `FastembedLateInteractionRanker` reordered the top_k fallback to happen after validation (#3746)
|
|
8
|
+
- `FastEmbedder` reordered the top_k fallback to happen after validation (#3745)
|
|
9
|
+
|
|
10
|
+
### 📚 Documentation
|
|
11
|
+
|
|
12
|
+
- Fix TypeError in FastembedSparseDocumentEmbedder usage example (#3731)
|
|
13
|
+
|
|
14
|
+
|
|
3
15
|
## [integrations/fastembed-v2.4.0] - 2026-07-30
|
|
4
16
|
|
|
5
17
|
### 🚀 Features
|
|
@@ -1,6 +1,6 @@
|
|
|
1
|
-
Metadata-Version: 2.
|
|
1
|
+
Metadata-Version: 2.5
|
|
2
2
|
Name: fastembed-haystack
|
|
3
|
-
Version: 2.
|
|
3
|
+
Version: 2.6.0
|
|
4
4
|
Summary: Haystack 2.x component to embed strings and Documents using fastembed embedding model
|
|
5
5
|
Project-URL: Source, https://github.com/deepset-ai/haystack-core-integrations
|
|
6
6
|
Project-URL: Documentation, https://github.com/deepset-ai/haystack-core-integrations/blob/main/integrations/fastembed/README.md
|
|
@@ -57,6 +57,21 @@ class FastembedDocumentEmbedder:
|
|
|
57
57
|
print(f"Document Embedding: {result['documents'][0].embedding}")
|
|
58
58
|
print(f"Embedding Dimension: {len(result['documents'][0].embedding)}")
|
|
59
59
|
```
|
|
60
|
+
|
|
61
|
+
Running on GPU:
|
|
62
|
+
```python
|
|
63
|
+
# NVIDIA GPU (requires onnxruntime-gpu)
|
|
64
|
+
doc_embedder = FastembedDocumentEmbedder(
|
|
65
|
+
model="BAAI/bge-small-en-v1.5",
|
|
66
|
+
model_kwargs={"providers": ["CUDAExecutionProvider"]},
|
|
67
|
+
)
|
|
68
|
+
|
|
69
|
+
# Intel GPU / XPU (requires onnxruntime-openvino)
|
|
70
|
+
doc_embedder = FastembedDocumentEmbedder(
|
|
71
|
+
model="BAAI/bge-small-en-v1.5",
|
|
72
|
+
model_kwargs={"providers": ["OpenVINOExecutionProvider"]},
|
|
73
|
+
)
|
|
74
|
+
```
|
|
60
75
|
"""
|
|
61
76
|
|
|
62
77
|
def __init__(
|
|
@@ -95,7 +110,8 @@ class FastembedDocumentEmbedder:
|
|
|
95
110
|
:param meta_fields_to_embed: List of meta fields that should be embedded along with the Document content.
|
|
96
111
|
:param embedding_separator: Separator used to concatenate the meta fields to the Document content.
|
|
97
112
|
:param model_kwargs: Dictionary containing additional keyword arguments to pass to the Fastembed model,
|
|
98
|
-
such as `providers` (e.g. `["CUDAExecutionProvider"]`
|
|
113
|
+
such as `providers` (e.g. `["CUDAExecutionProvider"]` for NVIDIA GPU or
|
|
114
|
+
`["OpenVINOExecutionProvider"]` for Intel GPU/XPU), `cuda`, or `device_ids`.
|
|
99
115
|
"""
|
|
100
116
|
|
|
101
117
|
self.model_name = model
|
|
@@ -27,6 +27,21 @@ class FastembedTextEmbedder:
|
|
|
27
27
|
|
|
28
28
|
embedding = text_embedder.run(text)["embedding"]
|
|
29
29
|
```
|
|
30
|
+
|
|
31
|
+
Running on GPU:
|
|
32
|
+
```python
|
|
33
|
+
# NVIDIA GPU (requires onnxruntime-gpu)
|
|
34
|
+
text_embedder = FastembedTextEmbedder(
|
|
35
|
+
model="BAAI/bge-small-en-v1.5",
|
|
36
|
+
model_kwargs={"providers": ["CUDAExecutionProvider"]},
|
|
37
|
+
)
|
|
38
|
+
|
|
39
|
+
# Intel GPU / XPU (requires onnxruntime-openvino)
|
|
40
|
+
text_embedder = FastembedTextEmbedder(
|
|
41
|
+
model="BAAI/bge-small-en-v1.5",
|
|
42
|
+
model_kwargs={"providers": ["OpenVINOExecutionProvider"]},
|
|
43
|
+
)
|
|
44
|
+
```
|
|
30
45
|
"""
|
|
31
46
|
|
|
32
47
|
def __init__(
|
|
@@ -58,7 +73,8 @@ class FastembedTextEmbedder:
|
|
|
58
73
|
If None, don't use data-parallel processing, use default onnxruntime threading instead.
|
|
59
74
|
:param local_files_only: If `True`, only use the model files in the `cache_dir`.
|
|
60
75
|
:param model_kwargs: Dictionary containing additional keyword arguments to pass to the Fastembed model,
|
|
61
|
-
such as `providers` (e.g. `["CUDAExecutionProvider"]`
|
|
76
|
+
such as `providers` (e.g. `["CUDAExecutionProvider"]` for NVIDIA GPU or
|
|
77
|
+
`["OpenVINOExecutionProvider"]` for Intel GPU/XPU), `cuda`, or `device_ids`.
|
|
62
78
|
"""
|
|
63
79
|
|
|
64
80
|
self.model_name = model
|
|
@@ -38,6 +38,21 @@ class FastembedLateInteractionRanker:
|
|
|
38
38
|
|
|
39
39
|
# Berlin
|
|
40
40
|
```
|
|
41
|
+
|
|
42
|
+
Running on GPU:
|
|
43
|
+
```python
|
|
44
|
+
# NVIDIA GPU (requires onnxruntime-gpu)
|
|
45
|
+
ranker = FastembedLateInteractionRanker(
|
|
46
|
+
model_name="colbert-ir/colbertv2.0",
|
|
47
|
+
model_kwargs={"providers": ["CUDAExecutionProvider"]},
|
|
48
|
+
)
|
|
49
|
+
|
|
50
|
+
# Intel GPU / XPU (requires onnxruntime-openvino)
|
|
51
|
+
ranker = FastembedLateInteractionRanker(
|
|
52
|
+
model_name="colbert-ir/colbertv2.0",
|
|
53
|
+
model_kwargs={"providers": ["OpenVINOExecutionProvider"]},
|
|
54
|
+
)
|
|
55
|
+
```
|
|
41
56
|
"""
|
|
42
57
|
|
|
43
58
|
def __init__(
|
|
@@ -52,6 +67,7 @@ class FastembedLateInteractionRanker:
|
|
|
52
67
|
meta_fields_to_embed: list[str] | None = None,
|
|
53
68
|
meta_data_separator: str = "\n",
|
|
54
69
|
score_threshold: float | None = None,
|
|
70
|
+
model_kwargs: dict[str, Any] | None = None,
|
|
55
71
|
) -> None:
|
|
56
72
|
"""
|
|
57
73
|
Creates an instance of the 'FastembedLateInteractionRanker'.
|
|
@@ -75,6 +91,9 @@ class FastembedLateInteractionRanker:
|
|
|
75
91
|
to the Document content.
|
|
76
92
|
:param score_threshold: If provided, only documents with a score above the threshold are returned.
|
|
77
93
|
Note that ColBERT scores are unnormalized sums and typically range from 3 to 25.
|
|
94
|
+
:param model_kwargs: Dictionary containing additional keyword arguments to pass to the Fastembed model,
|
|
95
|
+
such as `providers` (e.g. `["CUDAExecutionProvider"]` to run on an NVIDIA GPU or
|
|
96
|
+
`["OpenVINOExecutionProvider"]` to run on an Intel GPU), `cuda`, or `device_ids`.
|
|
78
97
|
"""
|
|
79
98
|
if top_k <= 0:
|
|
80
99
|
msg = f"top_k must be > 0, but got {top_k}"
|
|
@@ -90,6 +109,7 @@ class FastembedLateInteractionRanker:
|
|
|
90
109
|
self.meta_fields_to_embed = meta_fields_to_embed or []
|
|
91
110
|
self.meta_data_separator = meta_data_separator
|
|
92
111
|
self.score_threshold = score_threshold
|
|
112
|
+
self.model_kwargs = model_kwargs
|
|
93
113
|
self._model: LateInteractionTextEmbedding | None = None
|
|
94
114
|
|
|
95
115
|
def to_dict(self) -> dict[str, Any]:
|
|
@@ -111,6 +131,7 @@ class FastembedLateInteractionRanker:
|
|
|
111
131
|
meta_fields_to_embed=self.meta_fields_to_embed,
|
|
112
132
|
meta_data_separator=self.meta_data_separator,
|
|
113
133
|
score_threshold=self.score_threshold,
|
|
134
|
+
model_kwargs=self.model_kwargs,
|
|
114
135
|
)
|
|
115
136
|
|
|
116
137
|
@classmethod
|
|
@@ -130,11 +151,13 @@ class FastembedLateInteractionRanker:
|
|
|
130
151
|
Initializes the component.
|
|
131
152
|
"""
|
|
132
153
|
if self._model is None:
|
|
154
|
+
model_kwargs = self.model_kwargs or {}
|
|
133
155
|
self._model = LateInteractionTextEmbedding(
|
|
134
156
|
model_name=self.model_name,
|
|
135
157
|
cache_dir=self.cache_dir,
|
|
136
158
|
threads=self.threads,
|
|
137
159
|
local_files_only=self.local_files_only,
|
|
160
|
+
**model_kwargs,
|
|
138
161
|
)
|
|
139
162
|
|
|
140
163
|
def _prepare_fastembed_input_docs(self, documents: list[Document]) -> list[str]:
|
|
@@ -35,6 +35,21 @@ class FastembedRanker:
|
|
|
35
35
|
|
|
36
36
|
# Berlin
|
|
37
37
|
```
|
|
38
|
+
|
|
39
|
+
Running on GPU:
|
|
40
|
+
```python
|
|
41
|
+
# NVIDIA GPU (requires onnxruntime-gpu)
|
|
42
|
+
ranker = FastembedRanker(
|
|
43
|
+
model_name="Xenova/ms-marco-MiniLM-L-6-v2",
|
|
44
|
+
model_kwargs={"providers": ["CUDAExecutionProvider"]},
|
|
45
|
+
)
|
|
46
|
+
|
|
47
|
+
# Intel GPU / XPU (requires onnxruntime-openvino)
|
|
48
|
+
ranker = FastembedRanker(
|
|
49
|
+
model_name="Xenova/ms-marco-MiniLM-L-6-v2",
|
|
50
|
+
model_kwargs={"providers": ["OpenVINOExecutionProvider"]},
|
|
51
|
+
)
|
|
52
|
+
```
|
|
38
53
|
"""
|
|
39
54
|
|
|
40
55
|
def __init__(
|
|
@@ -49,6 +64,7 @@ class FastembedRanker:
|
|
|
49
64
|
meta_fields_to_embed: list[str] | None = None,
|
|
50
65
|
meta_data_separator: str = "\n",
|
|
51
66
|
score_threshold: float | None = None,
|
|
67
|
+
model_kwargs: dict[str, Any] | None = None,
|
|
52
68
|
) -> None:
|
|
53
69
|
"""
|
|
54
70
|
Creates an instance of the 'FastembedRanker'.
|
|
@@ -71,6 +87,9 @@ class FastembedRanker:
|
|
|
71
87
|
to the Document content.
|
|
72
88
|
:param score_threshold: If provided, only documents with a score above the threshold are returned.
|
|
73
89
|
Applied after `top_k`, so the output may contain fewer than `top_k` documents.
|
|
90
|
+
:param model_kwargs: Dictionary containing additional keyword arguments to pass to the Fastembed model,
|
|
91
|
+
such as `providers` (e.g. `["CUDAExecutionProvider"]` to run on an NVIDIA GPU or
|
|
92
|
+
`["OpenVINOExecutionProvider"]` to run on an Intel GPU), `cuda`, or `device_ids`.
|
|
74
93
|
"""
|
|
75
94
|
if top_k <= 0:
|
|
76
95
|
msg = f"top_k must be > 0, but got {top_k}"
|
|
@@ -86,6 +105,7 @@ class FastembedRanker:
|
|
|
86
105
|
self.meta_fields_to_embed = meta_fields_to_embed or []
|
|
87
106
|
self.meta_data_separator = meta_data_separator
|
|
88
107
|
self.score_threshold = score_threshold
|
|
108
|
+
self.model_kwargs = model_kwargs
|
|
89
109
|
self._model: TextCrossEncoder | None = None
|
|
90
110
|
|
|
91
111
|
def to_dict(self) -> dict[str, Any]:
|
|
@@ -107,6 +127,7 @@ class FastembedRanker:
|
|
|
107
127
|
meta_fields_to_embed=self.meta_fields_to_embed,
|
|
108
128
|
meta_data_separator=self.meta_data_separator,
|
|
109
129
|
score_threshold=self.score_threshold,
|
|
130
|
+
model_kwargs=self.model_kwargs,
|
|
110
131
|
)
|
|
111
132
|
|
|
112
133
|
@classmethod
|
|
@@ -126,11 +147,13 @@ class FastembedRanker:
|
|
|
126
147
|
Initializes the component.
|
|
127
148
|
"""
|
|
128
149
|
if self._model is None:
|
|
150
|
+
model_kwargs = self.model_kwargs or {}
|
|
129
151
|
self._model = TextCrossEncoder(
|
|
130
152
|
model_name=self.model_name,
|
|
131
153
|
cache_dir=self.cache_dir,
|
|
132
154
|
threads=self.threads,
|
|
133
155
|
local_files_only=self.local_files_only,
|
|
156
|
+
**model_kwargs,
|
|
134
157
|
)
|
|
135
158
|
|
|
136
159
|
def _prepare_fastembed_input_docs(self, documents: list[Document]) -> list[str]:
|
|
@@ -54,6 +54,19 @@ class TestFastembedLateInteractionRanker:
|
|
|
54
54
|
assert ranker.meta_fields_to_embed == ["test_field"]
|
|
55
55
|
assert ranker.meta_data_separator == " | "
|
|
56
56
|
|
|
57
|
+
def test_init_with_model_kwargs_parameters(self):
|
|
58
|
+
"""
|
|
59
|
+
Test initialization of FastembedLateInteractionRanker with model_kwargs parameters.
|
|
60
|
+
"""
|
|
61
|
+
model_kwargs = {"providers": ["CUDAExecutionProvider"]}
|
|
62
|
+
|
|
63
|
+
ranker = FastembedLateInteractionRanker(
|
|
64
|
+
model_name="colbert-ir/colbertv2.0",
|
|
65
|
+
model_kwargs=model_kwargs,
|
|
66
|
+
)
|
|
67
|
+
|
|
68
|
+
assert ranker.model_kwargs == model_kwargs
|
|
69
|
+
|
|
57
70
|
def test_init_with_incorrect_input(self):
|
|
58
71
|
"""
|
|
59
72
|
Test for checking incorrect input format on init.
|
|
@@ -90,6 +103,7 @@ class TestFastembedLateInteractionRanker:
|
|
|
90
103
|
"meta_fields_to_embed": [],
|
|
91
104
|
"meta_data_separator": "\n",
|
|
92
105
|
"score_threshold": None,
|
|
106
|
+
"model_kwargs": None,
|
|
93
107
|
},
|
|
94
108
|
}
|
|
95
109
|
|
|
@@ -107,6 +121,7 @@ class TestFastembedLateInteractionRanker:
|
|
|
107
121
|
local_files_only=True,
|
|
108
122
|
meta_fields_to_embed=["test_field"],
|
|
109
123
|
meta_data_separator=" | ",
|
|
124
|
+
model_kwargs={"providers": ["CUDAExecutionProvider"]},
|
|
110
125
|
)
|
|
111
126
|
ranker_dict = ranker.to_dict()
|
|
112
127
|
assert ranker_dict == {
|
|
@@ -123,6 +138,7 @@ class TestFastembedLateInteractionRanker:
|
|
|
123
138
|
"meta_fields_to_embed": ["test_field"],
|
|
124
139
|
"meta_data_separator": " | ",
|
|
125
140
|
"score_threshold": None,
|
|
141
|
+
"model_kwargs": {"providers": ["CUDAExecutionProvider"]},
|
|
126
142
|
},
|
|
127
143
|
}
|
|
128
144
|
|
|
@@ -173,6 +189,7 @@ class TestFastembedLateInteractionRanker:
|
|
|
173
189
|
"local_files_only": True,
|
|
174
190
|
"meta_fields_to_embed": ["test_field"],
|
|
175
191
|
"meta_data_separator": " | ",
|
|
192
|
+
"model_kwargs": {"providers": ["CUDAExecutionProvider"]},
|
|
176
193
|
},
|
|
177
194
|
}
|
|
178
195
|
ranker = default_from_dict(FastembedLateInteractionRanker, ranker_dict)
|
|
@@ -185,6 +202,7 @@ class TestFastembedLateInteractionRanker:
|
|
|
185
202
|
assert ranker.local_files_only
|
|
186
203
|
assert ranker.meta_fields_to_embed == ["test_field"]
|
|
187
204
|
assert ranker.meta_data_separator == " | "
|
|
205
|
+
assert ranker.model_kwargs == {"providers": ["CUDAExecutionProvider"]}
|
|
188
206
|
|
|
189
207
|
def test_run_incorrect_input_format(self):
|
|
190
208
|
"""
|
|
@@ -297,6 +315,26 @@ class TestFastembedLateInteractionRanker:
|
|
|
297
315
|
ranker.warm_up()
|
|
298
316
|
mock_cls.assert_called_once()
|
|
299
317
|
|
|
318
|
+
def test_warm_up_forwards_model_kwargs(self):
|
|
319
|
+
"""
|
|
320
|
+
Test that model_kwargs (e.g. GPU providers) are forwarded to the underlying Fastembed model.
|
|
321
|
+
"""
|
|
322
|
+
ranker = FastembedLateInteractionRanker(
|
|
323
|
+
model_name="colbert-ir/colbertv2.0",
|
|
324
|
+
model_kwargs={"providers": ["CUDAExecutionProvider"]},
|
|
325
|
+
)
|
|
326
|
+
with patch(
|
|
327
|
+
"haystack_integrations.components.rankers.fastembed.late_interaction_ranker.LateInteractionTextEmbedding"
|
|
328
|
+
) as mock_cls:
|
|
329
|
+
ranker.warm_up()
|
|
330
|
+
mock_cls.assert_called_once_with(
|
|
331
|
+
model_name="colbert-ir/colbertv2.0",
|
|
332
|
+
cache_dir=None,
|
|
333
|
+
threads=None,
|
|
334
|
+
local_files_only=False,
|
|
335
|
+
providers=["CUDAExecutionProvider"],
|
|
336
|
+
)
|
|
337
|
+
|
|
300
338
|
def test_run_calls_warm_up(self):
|
|
301
339
|
"""
|
|
302
340
|
Unit test to check that warm_up is called when run is called for the first time.
|
|
@@ -57,6 +57,19 @@ class TestFastembedRanker:
|
|
|
57
57
|
assert ranker.meta_data_separator == " | "
|
|
58
58
|
assert ranker.score_threshold == 0.5
|
|
59
59
|
|
|
60
|
+
def test_init_with_model_kwargs_parameters(self):
|
|
61
|
+
"""
|
|
62
|
+
Test initialization of FastembedRanker with model_kwargs parameters.
|
|
63
|
+
"""
|
|
64
|
+
model_kwargs = {"providers": ["CUDAExecutionProvider"]}
|
|
65
|
+
|
|
66
|
+
ranker = FastembedRanker(
|
|
67
|
+
model_name="BAAI/bge-reranker-base",
|
|
68
|
+
model_kwargs=model_kwargs,
|
|
69
|
+
)
|
|
70
|
+
|
|
71
|
+
assert ranker.model_kwargs == model_kwargs
|
|
72
|
+
|
|
60
73
|
def test_init_with_incorrect_input(self):
|
|
61
74
|
"""
|
|
62
75
|
Test for checking incorrect input format on init
|
|
@@ -92,6 +105,7 @@ class TestFastembedRanker:
|
|
|
92
105
|
"meta_fields_to_embed": [],
|
|
93
106
|
"meta_data_separator": "\n",
|
|
94
107
|
"score_threshold": None,
|
|
108
|
+
"model_kwargs": None,
|
|
95
109
|
},
|
|
96
110
|
}
|
|
97
111
|
|
|
@@ -110,6 +124,7 @@ class TestFastembedRanker:
|
|
|
110
124
|
meta_fields_to_embed=["test_field"],
|
|
111
125
|
meta_data_separator=" | ",
|
|
112
126
|
score_threshold=0.5,
|
|
127
|
+
model_kwargs={"providers": ["CUDAExecutionProvider"]},
|
|
113
128
|
)
|
|
114
129
|
ranker_dict = ranker.to_dict()
|
|
115
130
|
assert ranker_dict == {
|
|
@@ -125,6 +140,7 @@ class TestFastembedRanker:
|
|
|
125
140
|
"meta_fields_to_embed": ["test_field"],
|
|
126
141
|
"meta_data_separator": " | ",
|
|
127
142
|
"score_threshold": 0.5,
|
|
143
|
+
"model_kwargs": {"providers": ["CUDAExecutionProvider"]},
|
|
128
144
|
},
|
|
129
145
|
}
|
|
130
146
|
|
|
@@ -176,6 +192,7 @@ class TestFastembedRanker:
|
|
|
176
192
|
"meta_fields_to_embed": ["test_field"],
|
|
177
193
|
"meta_data_separator": " | ",
|
|
178
194
|
"score_threshold": 0.5,
|
|
195
|
+
"model_kwargs": {"providers": ["CUDAExecutionProvider"]},
|
|
179
196
|
},
|
|
180
197
|
}
|
|
181
198
|
ranker = default_from_dict(FastembedRanker, ranker_dict)
|
|
@@ -189,6 +206,7 @@ class TestFastembedRanker:
|
|
|
189
206
|
assert ranker.meta_fields_to_embed == ["test_field"]
|
|
190
207
|
assert ranker.meta_data_separator == " | "
|
|
191
208
|
assert ranker.score_threshold == 0.5
|
|
209
|
+
assert ranker.model_kwargs == {"providers": ["CUDAExecutionProvider"]}
|
|
192
210
|
|
|
193
211
|
def test_run_incorrect_input_format(self):
|
|
194
212
|
"""
|
|
@@ -301,6 +319,24 @@ class TestFastembedRanker:
|
|
|
301
319
|
|
|
302
320
|
mock_warm_up.assert_called_once()
|
|
303
321
|
|
|
322
|
+
@patch("haystack_integrations.components.rankers.fastembed.ranker.TextCrossEncoder")
|
|
323
|
+
def test_warm_up_forwards_model_kwargs(self, mocked_cross_encoder):
|
|
324
|
+
"""
|
|
325
|
+
Test that model_kwargs (e.g. GPU providers) are forwarded to the underlying Fastembed model.
|
|
326
|
+
"""
|
|
327
|
+
ranker = FastembedRanker(
|
|
328
|
+
model_name="BAAI/bge-reranker-base",
|
|
329
|
+
model_kwargs={"providers": ["CUDAExecutionProvider"]},
|
|
330
|
+
)
|
|
331
|
+
ranker.warm_up()
|
|
332
|
+
mocked_cross_encoder.assert_called_once_with(
|
|
333
|
+
model_name="BAAI/bge-reranker-base",
|
|
334
|
+
cache_dir=None,
|
|
335
|
+
threads=None,
|
|
336
|
+
local_files_only=False,
|
|
337
|
+
providers=["CUDAExecutionProvider"],
|
|
338
|
+
)
|
|
339
|
+
|
|
304
340
|
@pytest.mark.integration
|
|
305
341
|
def test_run(self):
|
|
306
342
|
ranker = FastembedRanker(model_name="Xenova/ms-marco-MiniLM-L-6-v2", top_k=2)
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
{fastembed_haystack-2.5.0 → fastembed_haystack-2.6.0}/tests/test_fastembed_document_embedder.py
RENAMED
|
File without changes
|
|
File without changes
|
{fastembed_haystack-2.5.0 → fastembed_haystack-2.6.0}/tests/test_fastembed_sparse_text_embedder.py
RENAMED
|
File without changes
|
|
File without changes
|