fastembed-haystack 2.2.0__tar.gz → 2.4.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (30) hide show
  1. {fastembed_haystack-2.2.0 → fastembed_haystack-2.4.0}/CHANGELOG.md +27 -6
  2. {fastembed_haystack-2.2.0 → fastembed_haystack-2.4.0}/PKG-INFO +1 -1
  3. {fastembed_haystack-2.2.0 → fastembed_haystack-2.4.0}/pyproject.toml +1 -0
  4. {fastembed_haystack-2.2.0 → fastembed_haystack-2.4.0}/src/haystack_integrations/components/embedders/fastembed/embedding_backend/fastembed_backend.py +15 -3
  5. {fastembed_haystack-2.2.0 → fastembed_haystack-2.4.0}/src/haystack_integrations/components/embedders/fastembed/fastembed_document_embedder.py +6 -0
  6. {fastembed_haystack-2.2.0 → fastembed_haystack-2.4.0}/src/haystack_integrations/components/embedders/fastembed/fastembed_sparse_document_embedder.py +2 -2
  7. {fastembed_haystack-2.2.0 → fastembed_haystack-2.4.0}/src/haystack_integrations/components/embedders/fastembed/fastembed_sparse_text_embedder.py +2 -2
  8. {fastembed_haystack-2.2.0 → fastembed_haystack-2.4.0}/src/haystack_integrations/components/embedders/fastembed/fastembed_text_embedder.py +6 -0
  9. {fastembed_haystack-2.2.0 → fastembed_haystack-2.4.0}/src/haystack_integrations/components/rankers/fastembed/ranker.py +9 -0
  10. {fastembed_haystack-2.2.0 → fastembed_haystack-2.4.0}/tests/test_fastembed_backend.py +17 -0
  11. {fastembed_haystack-2.2.0 → fastembed_haystack-2.4.0}/tests/test_fastembed_document_embedder.py +41 -1
  12. {fastembed_haystack-2.2.0 → fastembed_haystack-2.4.0}/tests/test_fastembed_ranker.py +24 -0
  13. {fastembed_haystack-2.2.0 → fastembed_haystack-2.4.0}/tests/test_fastembed_text_embedder.py +41 -1
  14. {fastembed_haystack-2.2.0 → fastembed_haystack-2.4.0}/.gitignore +0 -0
  15. {fastembed_haystack-2.2.0 → fastembed_haystack-2.4.0}/LICENSE.txt +0 -0
  16. {fastembed_haystack-2.2.0 → fastembed_haystack-2.4.0}/README.md +0 -0
  17. {fastembed_haystack-2.2.0 → fastembed_haystack-2.4.0}/examples/example.py +0 -0
  18. {fastembed_haystack-2.2.0 → fastembed_haystack-2.4.0}/examples/ranker_example.py +0 -0
  19. {fastembed_haystack-2.2.0 → fastembed_haystack-2.4.0}/examples/sparse_example.py +0 -0
  20. {fastembed_haystack-2.2.0 → fastembed_haystack-2.4.0}/pydoc/config_docusaurus.yml +0 -0
  21. {fastembed_haystack-2.2.0 → fastembed_haystack-2.4.0}/src/haystack_integrations/components/embedders/fastembed/__init__.py +0 -0
  22. {fastembed_haystack-2.2.0 → fastembed_haystack-2.4.0}/src/haystack_integrations/components/embedders/fastembed/embedding_backend/__init__.py +0 -0
  23. {fastembed_haystack-2.2.0 → fastembed_haystack-2.4.0}/src/haystack_integrations/components/embedders/py.typed +0 -0
  24. {fastembed_haystack-2.2.0 → fastembed_haystack-2.4.0}/src/haystack_integrations/components/rankers/fastembed/__init__.py +0 -0
  25. {fastembed_haystack-2.2.0 → fastembed_haystack-2.4.0}/src/haystack_integrations/components/rankers/fastembed/late_interaction_ranker.py +0 -0
  26. {fastembed_haystack-2.2.0 → fastembed_haystack-2.4.0}/src/haystack_integrations/components/rankers/py.typed +0 -0
  27. {fastembed_haystack-2.2.0 → fastembed_haystack-2.4.0}/tests/__init__.py +0 -0
  28. {fastembed_haystack-2.2.0 → fastembed_haystack-2.4.0}/tests/test_fastembed_late_interaction_ranker.py +0 -0
  29. {fastembed_haystack-2.2.0 → fastembed_haystack-2.4.0}/tests/test_fastembed_sparse_document_embedder.py +0 -0
  30. {fastembed_haystack-2.2.0 → fastembed_haystack-2.4.0}/tests/test_fastembed_sparse_text_embedder.py +0 -0
@@ -1,20 +1,41 @@
1
1
  # Changelog
2
2
 
3
- ## [integrations/fastembed-v2.1.0] - 2026-02-13
3
+ ## [integrations/fastembed-v2.3.0] - 2026-04-30
4
+
5
+ ### 🚀 Features
6
+
7
+ - Add score_threshold parameter to FastembedRanker (#3235)
8
+
9
+
10
+ ## [integrations/fastembed-v2.2.0] - 2026-04-16
11
+
12
+ ### 🚀 Features
13
+
14
+ - Update `FastembedRanker` to auto call `warm_up` on first run (#2834)
15
+ - Add FastembedLateInteractionRanker for ColBERT late-interaction reranking (#3135)
16
+
17
+ ### 🐛 Bug Fixes
18
+
19
+ - Replace in-place dataclass mutations with dataclasses.replace() (#3112)
4
20
 
5
21
  ### 📚 Documentation
6
22
 
7
23
  - *(fastembed)* Remove explicit warm_up from examples (#2830)
24
+ - Simplify pydoc configs (#2855)
25
+
26
+ ### 🧪 Testing
27
+
28
+ - Track test coverage for all integrations (#3065)
29
+
30
+ ### 🧹 Chores
31
+
32
+ - Add ANN type annotations to faiss, fastembed, firecrawl, github, google_ai (#2988)
33
+ - Enforce ruff docstring rules (D102/D103/D205/D209/D213/D417/D419) in integrations 11-20 (#3009)
8
34
 
9
35
  ### 🌀 Miscellaneous
10
36
 
11
- - Automatically call warm_up() if model is not loaded in FastembedRanker
12
- - Remove explicit call to warm_up() in ranker_example.py
13
- - Add unit tests
14
37
  - Fix typing issue
15
38
  - Fix tying
16
- - Update docstrings
17
- - Remove explicit warm_up calls from integration tests
18
39
 
19
40
  ## [integrations/fastembed-v2.0.0] - 2026-01-12
20
41
 
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: fastembed-haystack
3
- Version: 2.2.0
3
+ Version: 2.4.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
@@ -132,6 +132,7 @@ ignore = [
132
132
  "PLR0912",
133
133
  "PLR0913",
134
134
  "PLR0915",
135
+ "PLR0917",
135
136
  ]
136
137
 
137
138
  [tool.ruff.lint.isort]
@@ -20,14 +20,19 @@ class _FastembedEmbeddingBackendFactory:
20
20
  cache_dir: str | None = None,
21
21
  threads: int | None = None,
22
22
  local_files_only: bool = False,
23
+ model_kwargs: dict[str, Any] | None = None,
23
24
  ) -> "_FastembedEmbeddingBackend":
24
- embedding_backend_id = f"{model_name}{cache_dir}{threads}"
25
+ embedding_backend_id = f"{model_name}{cache_dir}{threads}{local_files_only}{model_kwargs}"
25
26
 
26
27
  if embedding_backend_id in _FastembedEmbeddingBackendFactory._instances:
27
28
  return _FastembedEmbeddingBackendFactory._instances[embedding_backend_id]
28
29
 
29
30
  embedding_backend = _FastembedEmbeddingBackend(
30
- model_name=model_name, cache_dir=cache_dir, threads=threads, local_files_only=local_files_only
31
+ model_name=model_name,
32
+ cache_dir=cache_dir,
33
+ threads=threads,
34
+ local_files_only=local_files_only,
35
+ model_kwargs=model_kwargs,
31
36
  )
32
37
  _FastembedEmbeddingBackendFactory._instances[embedding_backend_id] = embedding_backend
33
38
  return embedding_backend
@@ -44,9 +49,16 @@ class _FastembedEmbeddingBackend:
44
49
  cache_dir: str | None = None,
45
50
  threads: int | None = None,
46
51
  local_files_only: bool = False,
52
+ model_kwargs: dict[str, Any] | None = None,
47
53
  ) -> None:
54
+ model_kwargs = model_kwargs or {}
55
+
48
56
  self.model = TextEmbedding(
49
- model_name=model_name, cache_dir=cache_dir, threads=threads, local_files_only=local_files_only
57
+ model_name=model_name,
58
+ cache_dir=cache_dir,
59
+ threads=threads,
60
+ local_files_only=local_files_only,
61
+ **model_kwargs,
50
62
  )
51
63
 
52
64
  def embed(self, data: list[str], progress_bar: bool = True, **kwargs: Any) -> list[list[float]]:
@@ -72,6 +72,7 @@ class FastembedDocumentEmbedder:
72
72
  local_files_only: bool = False,
73
73
  meta_fields_to_embed: list[str] | None = None,
74
74
  embedding_separator: str = "\n",
75
+ model_kwargs: dict[str, Any] | None = None,
75
76
  ) -> None:
76
77
  """
77
78
  Create an FastembedDocumentEmbedder component.
@@ -93,6 +94,8 @@ class FastembedDocumentEmbedder:
93
94
  :param local_files_only: If `True`, only use the model files in the `cache_dir`.
94
95
  :param meta_fields_to_embed: List of meta fields that should be embedded along with the Document content.
95
96
  :param embedding_separator: Separator used to concatenate the meta fields to the Document content.
97
+ :param model_kwargs: Dictionary containing additional keyword arguments to pass to the Fastembed model,
98
+ such as `providers` (e.g. `["CUDAExecutionProvider"]` to run on GPU), `cuda`, or `device_ids`.
96
99
  """
97
100
 
98
101
  self.model_name = model
@@ -106,6 +109,7 @@ class FastembedDocumentEmbedder:
106
109
  self.local_files_only = local_files_only
107
110
  self.meta_fields_to_embed = meta_fields_to_embed or []
108
111
  self.embedding_separator = embedding_separator
112
+ self.model_kwargs = model_kwargs
109
113
  self.embedding_backend: _FastembedEmbeddingBackend | None = None
110
114
 
111
115
  def to_dict(self) -> dict[str, Any]:
@@ -128,6 +132,7 @@ class FastembedDocumentEmbedder:
128
132
  local_files_only=self.local_files_only,
129
133
  meta_fields_to_embed=self.meta_fields_to_embed,
130
134
  embedding_separator=self.embedding_separator,
135
+ model_kwargs=self.model_kwargs,
131
136
  )
132
137
 
133
138
  def warm_up(self) -> None:
@@ -140,6 +145,7 @@ class FastembedDocumentEmbedder:
140
145
  cache_dir=self.cache_dir,
141
146
  threads=self.threads,
142
147
  local_files_only=self.local_files_only,
148
+ model_kwargs=self.model_kwargs,
143
149
  )
144
150
 
145
151
  def _prepare_texts_to_embed(self, documents: list[Document]) -> list[str]:
@@ -71,7 +71,7 @@ class FastembedSparseDocumentEmbedder:
71
71
  model_kwargs: dict[str, Any] | None = None,
72
72
  ) -> None:
73
73
  """
74
- Create an FastembedDocumentEmbedder component.
74
+ Create a FastembedSparseDocumentEmbedder component.
75
75
 
76
76
  :param model: Local path or name of the model in Hugging Face's model hub,
77
77
  such as `prithivida/Splade_PP_en_v1`.
@@ -162,7 +162,7 @@ class FastembedSparseDocumentEmbedder:
162
162
  if not isinstance(documents, list) or (documents and not isinstance(documents[0], Document)):
163
163
  msg = (
164
164
  "FastembedSparseDocumentEmbedder expects a list of Documents as input. "
165
- "In case you want to embed a list of strings, please use the FastembedTextEmbedder."
165
+ "In case you want to embed a list of strings, please use the FastembedSparseTextEmbedder."
166
166
  )
167
167
  raise TypeError(msg)
168
168
 
@@ -107,13 +107,13 @@ class FastembedSparseTextEmbedder:
107
107
 
108
108
  :param text: A string to embed.
109
109
  :returns: A dictionary with the following keys:
110
- - `embedding`: A list of floats representing the embedding of the input text.
110
+ - `sparse_embedding`: The `SparseEmbedding` of the input text.
111
111
  :raises TypeError: If the input is not a string.
112
112
  """
113
113
  if not isinstance(text, str):
114
114
  msg = (
115
115
  "FastembedSparseTextEmbedder expects a string as input. "
116
- "In case you want to embed a list of Documents, please use the FastembedDocumentEmbedder."
116
+ "In case you want to embed a list of Documents, please use the FastembedSparseDocumentEmbedder."
117
117
  )
118
118
  raise TypeError(msg)
119
119
 
@@ -39,6 +39,7 @@ class FastembedTextEmbedder:
39
39
  progress_bar: bool = True,
40
40
  parallel: int | None = None,
41
41
  local_files_only: bool = False,
42
+ model_kwargs: dict[str, Any] | None = None,
42
43
  ) -> None:
43
44
  """
44
45
  Create a FastembedTextEmbedder component.
@@ -56,6 +57,8 @@ class FastembedTextEmbedder:
56
57
  If 0, use all available cores.
57
58
  If None, don't use data-parallel processing, use default onnxruntime threading instead.
58
59
  :param local_files_only: If `True`, only use the model files in the `cache_dir`.
60
+ :param model_kwargs: Dictionary containing additional keyword arguments to pass to the Fastembed model,
61
+ such as `providers` (e.g. `["CUDAExecutionProvider"]` to run on GPU), `cuda`, or `device_ids`.
59
62
  """
60
63
 
61
64
  self.model_name = model
@@ -66,6 +69,7 @@ class FastembedTextEmbedder:
66
69
  self.progress_bar = progress_bar
67
70
  self.parallel = parallel
68
71
  self.local_files_only = local_files_only
72
+ self.model_kwargs = model_kwargs
69
73
  self.embedding_backend: _FastembedEmbeddingBackend | None = None
70
74
 
71
75
  def to_dict(self) -> dict[str, Any]:
@@ -85,6 +89,7 @@ class FastembedTextEmbedder:
85
89
  progress_bar=self.progress_bar,
86
90
  parallel=self.parallel,
87
91
  local_files_only=self.local_files_only,
92
+ model_kwargs=self.model_kwargs,
88
93
  )
89
94
 
90
95
  def warm_up(self) -> None:
@@ -97,6 +102,7 @@ class FastembedTextEmbedder:
97
102
  cache_dir=self.cache_dir,
98
103
  threads=self.threads,
99
104
  local_files_only=self.local_files_only,
105
+ model_kwargs=self.model_kwargs,
100
106
  )
101
107
 
102
108
  @component.output_types(embedding=list[float])
@@ -48,6 +48,7 @@ class FastembedRanker:
48
48
  local_files_only: bool = False,
49
49
  meta_fields_to_embed: list[str] | None = None,
50
50
  meta_data_separator: str = "\n",
51
+ score_threshold: float | None = None,
51
52
  ) -> None:
52
53
  """
53
54
  Creates an instance of the 'FastembedRanker'.
@@ -68,6 +69,8 @@ class FastembedRanker:
68
69
  with the document content for reranking.
69
70
  :param meta_data_separator: Separator used to concatenate the meta fields
70
71
  to the Document content.
72
+ :param score_threshold: If provided, only documents with a score above the threshold are returned.
73
+ Applied after `top_k`, so the output may contain fewer than `top_k` documents.
71
74
  """
72
75
  if top_k <= 0:
73
76
  msg = f"top_k must be > 0, but got {top_k}"
@@ -82,6 +85,7 @@ class FastembedRanker:
82
85
  self.local_files_only = local_files_only
83
86
  self.meta_fields_to_embed = meta_fields_to_embed or []
84
87
  self.meta_data_separator = meta_data_separator
88
+ self.score_threshold = score_threshold
85
89
  self._model: TextCrossEncoder | None = None
86
90
 
87
91
  def to_dict(self) -> dict[str, Any]:
@@ -102,6 +106,7 @@ class FastembedRanker:
102
106
  local_files_only=self.local_files_only,
103
107
  meta_fields_to_embed=self.meta_fields_to_embed,
104
108
  meta_data_separator=self.meta_data_separator,
109
+ score_threshold=self.score_threshold,
105
110
  )
106
111
 
107
112
  @classmethod
@@ -201,4 +206,8 @@ class FastembedRanker:
201
206
 
202
207
  top_k_documents = [replace(doc, score=score) for doc, score in sorted_doc_scores[:top_k]]
203
208
 
209
+ score_threshold = self.score_threshold
210
+ if score_threshold is not None:
211
+ top_k_documents = [doc for doc in top_k_documents if doc.score is not None and doc.score >= score_threshold]
212
+
204
213
  return {"documents": top_k_documents}
@@ -39,6 +39,23 @@ def test_model_initialization(mock_instructor):
39
39
  _FastembedEmbeddingBackendFactory._instances = {}
40
40
 
41
41
 
42
+ @patch("haystack_integrations.components.embedders.fastembed.embedding_backend.fastembed_backend.TextEmbedding")
43
+ def test_dense_model_kwargs_initialization(mock_instructor):
44
+ model_kwargs = {"providers": ["CUDAExecutionProvider"]}
45
+
46
+ _FastembedEmbeddingBackendFactory.get_embedding_backend(
47
+ model_name="BAAI/bge-small-en-v1.5",
48
+ model_kwargs=model_kwargs,
49
+ )
50
+
51
+ mock_instructor.assert_called_once_with(
52
+ model_name="BAAI/bge-small-en-v1.5", cache_dir=None, threads=None, local_files_only=False, **model_kwargs
53
+ )
54
+
55
+ # restore the factory state
56
+ _FastembedEmbeddingBackendFactory._instances = {}
57
+
58
+
42
59
  @patch("haystack_integrations.components.embedders.fastembed.embedding_backend.fastembed_backend.TextEmbedding")
43
60
  def test_embedding_function_with_kwargs(mock_instructor): # noqa: ARG001
44
61
  embedding_backend = _FastembedEmbeddingBackendFactory.get_embedding_backend(model_name="BAAI/bge-small-en-v1.5")
@@ -80,6 +80,7 @@ class TestFastembedDocumentEmbedder:
80
80
  "local_files_only": False,
81
81
  "embedding_separator": "\n",
82
82
  "meta_fields_to_embed": [],
83
+ "model_kwargs": None,
83
84
  },
84
85
  }
85
86
 
@@ -99,6 +100,7 @@ class TestFastembedDocumentEmbedder:
99
100
  local_files_only=True,
100
101
  meta_fields_to_embed=["test_field"],
101
102
  embedding_separator=" | ",
103
+ model_kwargs={"providers": ["CUDAExecutionProvider"]},
102
104
  )
103
105
  embedder_dict = embedder.to_dict()
104
106
  assert embedder_dict == {
@@ -115,6 +117,7 @@ class TestFastembedDocumentEmbedder:
115
117
  "local_files_only": True,
116
118
  "meta_fields_to_embed": ["test_field"],
117
119
  "embedding_separator": " | ",
120
+ "model_kwargs": {"providers": ["CUDAExecutionProvider"]},
118
121
  },
119
122
  }
120
123
 
@@ -169,6 +172,7 @@ class TestFastembedDocumentEmbedder:
169
172
  "local_files_only": True,
170
173
  "meta_fields_to_embed": ["test_field"],
171
174
  "embedding_separator": " | ",
175
+ "model_kwargs": {"providers": ["CUDAExecutionProvider"]},
172
176
  },
173
177
  }
174
178
  embedder = default_from_dict(FastembedDocumentEmbedder, embedder_dict)
@@ -183,6 +187,20 @@ class TestFastembedDocumentEmbedder:
183
187
  assert embedder.local_files_only
184
188
  assert embedder.meta_fields_to_embed == ["test_field"]
185
189
  assert embedder.embedding_separator == " | "
190
+ assert embedder.model_kwargs == {"providers": ["CUDAExecutionProvider"]}
191
+
192
+ def test_init_with_model_kwargs_parameters(self):
193
+ """
194
+ Test initialization of FastembedDocumentEmbedder with model_kwargs parameters.
195
+ """
196
+ model_kwargs = {"providers": ["CUDAExecutionProvider"]}
197
+
198
+ embedder = FastembedDocumentEmbedder(
199
+ model="BAAI/bge-small-en-v1.5",
200
+ model_kwargs=model_kwargs,
201
+ )
202
+
203
+ assert embedder.model_kwargs == model_kwargs
186
204
 
187
205
  @patch(
188
206
  "haystack_integrations.components.embedders.fastembed.fastembed_document_embedder._FastembedEmbeddingBackendFactory"
@@ -195,7 +213,7 @@ class TestFastembedDocumentEmbedder:
195
213
  mocked_factory.get_embedding_backend.assert_not_called()
196
214
  embedder.warm_up()
197
215
  mocked_factory.get_embedding_backend.assert_called_once_with(
198
- model_name="BAAI/bge-small-en-v1.5", cache_dir=None, threads=None, local_files_only=False
216
+ model_name="BAAI/bge-small-en-v1.5", cache_dir=None, threads=None, local_files_only=False, model_kwargs=None
199
217
  )
200
218
 
201
219
  @patch(
@@ -295,6 +313,28 @@ class TestFastembedDocumentEmbedder:
295
313
 
296
314
  mock_warm_up.assert_called_once()
297
315
 
316
+ @pytest.mark.integration
317
+ def test_run_with_model_kwargs(self):
318
+ """
319
+ Integration test to check the embedding with model_kwargs parameters.
320
+ """
321
+ model_kwargs = {"providers": ["CPUExecutionProvider"], "lazy_load": True}
322
+
323
+ embedder = FastembedDocumentEmbedder(
324
+ model="BAAI/bge-small-en-v1.5",
325
+ model_kwargs=model_kwargs,
326
+ )
327
+ embedder.warm_up()
328
+
329
+ doc = Document(content="Parton energy loss in QCD matter")
330
+
331
+ result = embedder.run(documents=[doc])
332
+ embedding = result["documents"][0].embedding
333
+
334
+ assert isinstance(embedding, list)
335
+ assert len(embedding) == 384
336
+ assert all(isinstance(emb, float) for emb in embedding)
337
+
298
338
  @pytest.mark.integration
299
339
  def test_run(self):
300
340
  embedder = FastembedDocumentEmbedder(
@@ -28,6 +28,7 @@ class TestFastembedRanker:
28
28
  assert not ranker.local_files_only
29
29
  assert ranker.meta_fields_to_embed == []
30
30
  assert ranker.meta_data_separator == "\n"
31
+ assert ranker.score_threshold is None
31
32
 
32
33
  def test_init_with_parameters(self):
33
34
  """
@@ -43,6 +44,7 @@ class TestFastembedRanker:
43
44
  local_files_only=True,
44
45
  meta_fields_to_embed=["test_field"],
45
46
  meta_data_separator=" | ",
47
+ score_threshold=0.5,
46
48
  )
47
49
  assert ranker.model_name == "BAAI/bge-reranker-base"
48
50
  assert ranker.top_k == 64
@@ -53,6 +55,7 @@ class TestFastembedRanker:
53
55
  assert ranker.local_files_only
54
56
  assert ranker.meta_fields_to_embed == ["test_field"]
55
57
  assert ranker.meta_data_separator == " | "
58
+ assert ranker.score_threshold == 0.5
56
59
 
57
60
  def test_init_with_incorrect_input(self):
58
61
  """
@@ -88,6 +91,7 @@ class TestFastembedRanker:
88
91
  "local_files_only": False,
89
92
  "meta_fields_to_embed": [],
90
93
  "meta_data_separator": "\n",
94
+ "score_threshold": None,
91
95
  },
92
96
  }
93
97
 
@@ -105,6 +109,7 @@ class TestFastembedRanker:
105
109
  local_files_only=True,
106
110
  meta_fields_to_embed=["test_field"],
107
111
  meta_data_separator=" | ",
112
+ score_threshold=0.5,
108
113
  )
109
114
  ranker_dict = ranker.to_dict()
110
115
  assert ranker_dict == {
@@ -119,6 +124,7 @@ class TestFastembedRanker:
119
124
  "local_files_only": True,
120
125
  "meta_fields_to_embed": ["test_field"],
121
126
  "meta_data_separator": " | ",
127
+ "score_threshold": 0.5,
122
128
  },
123
129
  }
124
130
 
@@ -138,6 +144,7 @@ class TestFastembedRanker:
138
144
  "local_files_only": False,
139
145
  "meta_fields_to_embed": [],
140
146
  "meta_data_separator": "\n",
147
+ "score_threshold": None,
141
148
  },
142
149
  }
143
150
  ranker = default_from_dict(FastembedRanker, ranker_dict)
@@ -150,6 +157,7 @@ class TestFastembedRanker:
150
157
  assert not ranker.local_files_only
151
158
  assert ranker.meta_fields_to_embed == []
152
159
  assert ranker.meta_data_separator == "\n"
160
+ assert ranker.score_threshold is None
153
161
 
154
162
  def test_from_dict_with_custom_init_parameters(self):
155
163
  """
@@ -167,6 +175,7 @@ class TestFastembedRanker:
167
175
  "local_files_only": True,
168
176
  "meta_fields_to_embed": ["test_field"],
169
177
  "meta_data_separator": " | ",
178
+ "score_threshold": 0.5,
170
179
  },
171
180
  }
172
181
  ranker = default_from_dict(FastembedRanker, ranker_dict)
@@ -179,6 +188,7 @@ class TestFastembedRanker:
179
188
  assert ranker.local_files_only
180
189
  assert ranker.meta_fields_to_embed == ["test_field"]
181
190
  assert ranker.meta_data_separator == " | "
191
+ assert ranker.score_threshold == 0.5
182
192
 
183
193
  def test_run_incorrect_input_format(self):
184
194
  """
@@ -257,6 +267,20 @@ class TestFastembedRanker:
257
267
  parallel=None,
258
268
  )
259
269
 
270
+ def test_run_with_score_threshold(self):
271
+ """
272
+ Tests that documents below score_threshold are filtered out.
273
+ """
274
+ ranker = FastembedRanker(model_name="model_name", score_threshold=0.5)
275
+ ranker._model = MagicMock()
276
+ ranker._model.rerank.return_value = [0.9, 0.3, 0.7]
277
+
278
+ documents = [Document(content=f"doc {i}") for i in range(3)]
279
+ result = ranker.run(query="test", documents=documents)
280
+
281
+ assert len(result["documents"]) == 2
282
+ assert all(doc.score >= 0.5 for doc in result["documents"])
283
+
260
284
  def test_run_calls_warm_up(self):
261
285
  """
262
286
  Unit test to check that warm_up is called when run is called for the first time.
@@ -65,6 +65,7 @@ class TestFastembedTextEmbedder:
65
65
  "progress_bar": True,
66
66
  "parallel": None,
67
67
  "local_files_only": False,
68
+ "model_kwargs": None,
68
69
  },
69
70
  }
70
71
 
@@ -81,6 +82,7 @@ class TestFastembedTextEmbedder:
81
82
  progress_bar=False,
82
83
  parallel=1,
83
84
  local_files_only=True,
85
+ model_kwargs={"providers": ["CUDAExecutionProvider"]},
84
86
  )
85
87
  embedder_dict = embedder.to_dict()
86
88
  assert embedder_dict == {
@@ -94,6 +96,7 @@ class TestFastembedTextEmbedder:
94
96
  "progress_bar": False,
95
97
  "parallel": 1,
96
98
  "local_files_only": True,
99
+ "model_kwargs": {"providers": ["CUDAExecutionProvider"]},
97
100
  },
98
101
  }
99
102
 
@@ -136,6 +139,7 @@ class TestFastembedTextEmbedder:
136
139
  "suffix": "suffix",
137
140
  "progress_bar": False,
138
141
  "parallel": 1,
142
+ "model_kwargs": {"providers": ["CUDAExecutionProvider"]},
139
143
  },
140
144
  }
141
145
  embedder = default_from_dict(FastembedTextEmbedder, embedder_dict)
@@ -146,6 +150,20 @@ class TestFastembedTextEmbedder:
146
150
  assert embedder.suffix == "suffix"
147
151
  assert embedder.progress_bar is False
148
152
  assert embedder.parallel == 1
153
+ assert embedder.model_kwargs == {"providers": ["CUDAExecutionProvider"]}
154
+
155
+ def test_init_with_model_kwargs_parameters(self):
156
+ """
157
+ Test initialization of FastembedTextEmbedder with model_kwargs parameters.
158
+ """
159
+ model_kwargs = {"providers": ["CUDAExecutionProvider"]}
160
+
161
+ embedder = FastembedTextEmbedder(
162
+ model="BAAI/bge-small-en-v1.5",
163
+ model_kwargs=model_kwargs,
164
+ )
165
+
166
+ assert embedder.model_kwargs == model_kwargs
149
167
 
150
168
  @patch(
151
169
  "haystack_integrations.components.embedders.fastembed.fastembed_text_embedder._FastembedEmbeddingBackendFactory"
@@ -158,7 +176,7 @@ class TestFastembedTextEmbedder:
158
176
  mocked_factory.get_embedding_backend.assert_not_called()
159
177
  embedder.warm_up()
160
178
  mocked_factory.get_embedding_backend.assert_called_once_with(
161
- model_name="BAAI/bge-small-en-v1.5", cache_dir=None, threads=None, local_files_only=False
179
+ model_name="BAAI/bge-small-en-v1.5", cache_dir=None, threads=None, local_files_only=False, model_kwargs=None
162
180
  )
163
181
 
164
182
  @patch(
@@ -216,6 +234,28 @@ class TestFastembedTextEmbedder:
216
234
 
217
235
  mock_warm_up.assert_called_once()
218
236
 
237
+ @pytest.mark.integration
238
+ def test_run_with_model_kwargs(self):
239
+ """
240
+ Integration test to check the embedding with model_kwargs parameters.
241
+ """
242
+ model_kwargs = {"providers": ["CPUExecutionProvider"], "lazy_load": True}
243
+
244
+ embedder = FastembedTextEmbedder(
245
+ model="BAAI/bge-small-en-v1.5",
246
+ model_kwargs=model_kwargs,
247
+ )
248
+ embedder.warm_up()
249
+
250
+ text = "Parton energy loss in QCD matter"
251
+
252
+ result = embedder.run(text=text)
253
+ embedding = result["embedding"]
254
+
255
+ assert isinstance(embedding, list)
256
+ assert len(embedding) == 384
257
+ assert all(isinstance(emb, float) for emb in embedding)
258
+
219
259
  @pytest.mark.integration
220
260
  def test_run(self):
221
261
  embedder = FastembedTextEmbedder(