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.
- {fastembed_haystack-2.2.0 → fastembed_haystack-2.4.0}/CHANGELOG.md +27 -6
- {fastembed_haystack-2.2.0 → fastembed_haystack-2.4.0}/PKG-INFO +1 -1
- {fastembed_haystack-2.2.0 → fastembed_haystack-2.4.0}/pyproject.toml +1 -0
- {fastembed_haystack-2.2.0 → fastembed_haystack-2.4.0}/src/haystack_integrations/components/embedders/fastembed/embedding_backend/fastembed_backend.py +15 -3
- {fastembed_haystack-2.2.0 → fastembed_haystack-2.4.0}/src/haystack_integrations/components/embedders/fastembed/fastembed_document_embedder.py +6 -0
- {fastembed_haystack-2.2.0 → fastembed_haystack-2.4.0}/src/haystack_integrations/components/embedders/fastembed/fastembed_sparse_document_embedder.py +2 -2
- {fastembed_haystack-2.2.0 → fastembed_haystack-2.4.0}/src/haystack_integrations/components/embedders/fastembed/fastembed_sparse_text_embedder.py +2 -2
- {fastembed_haystack-2.2.0 → fastembed_haystack-2.4.0}/src/haystack_integrations/components/embedders/fastembed/fastembed_text_embedder.py +6 -0
- {fastembed_haystack-2.2.0 → fastembed_haystack-2.4.0}/src/haystack_integrations/components/rankers/fastembed/ranker.py +9 -0
- {fastembed_haystack-2.2.0 → fastembed_haystack-2.4.0}/tests/test_fastembed_backend.py +17 -0
- {fastembed_haystack-2.2.0 → fastembed_haystack-2.4.0}/tests/test_fastembed_document_embedder.py +41 -1
- {fastembed_haystack-2.2.0 → fastembed_haystack-2.4.0}/tests/test_fastembed_ranker.py +24 -0
- {fastembed_haystack-2.2.0 → fastembed_haystack-2.4.0}/tests/test_fastembed_text_embedder.py +41 -1
- {fastembed_haystack-2.2.0 → fastembed_haystack-2.4.0}/.gitignore +0 -0
- {fastembed_haystack-2.2.0 → fastembed_haystack-2.4.0}/LICENSE.txt +0 -0
- {fastembed_haystack-2.2.0 → fastembed_haystack-2.4.0}/README.md +0 -0
- {fastembed_haystack-2.2.0 → fastembed_haystack-2.4.0}/examples/example.py +0 -0
- {fastembed_haystack-2.2.0 → fastembed_haystack-2.4.0}/examples/ranker_example.py +0 -0
- {fastembed_haystack-2.2.0 → fastembed_haystack-2.4.0}/examples/sparse_example.py +0 -0
- {fastembed_haystack-2.2.0 → fastembed_haystack-2.4.0}/pydoc/config_docusaurus.yml +0 -0
- {fastembed_haystack-2.2.0 → fastembed_haystack-2.4.0}/src/haystack_integrations/components/embedders/fastembed/__init__.py +0 -0
- {fastembed_haystack-2.2.0 → fastembed_haystack-2.4.0}/src/haystack_integrations/components/embedders/fastembed/embedding_backend/__init__.py +0 -0
- {fastembed_haystack-2.2.0 → fastembed_haystack-2.4.0}/src/haystack_integrations/components/embedders/py.typed +0 -0
- {fastembed_haystack-2.2.0 → fastembed_haystack-2.4.0}/src/haystack_integrations/components/rankers/fastembed/__init__.py +0 -0
- {fastembed_haystack-2.2.0 → fastembed_haystack-2.4.0}/src/haystack_integrations/components/rankers/fastembed/late_interaction_ranker.py +0 -0
- {fastembed_haystack-2.2.0 → fastembed_haystack-2.4.0}/src/haystack_integrations/components/rankers/py.typed +0 -0
- {fastembed_haystack-2.2.0 → fastembed_haystack-2.4.0}/tests/__init__.py +0 -0
- {fastembed_haystack-2.2.0 → fastembed_haystack-2.4.0}/tests/test_fastembed_late_interaction_ranker.py +0 -0
- {fastembed_haystack-2.2.0 → fastembed_haystack-2.4.0}/tests/test_fastembed_sparse_document_embedder.py +0 -0
- {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.
|
|
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.
|
|
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
|
|
@@ -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,
|
|
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,
|
|
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
|
|
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
|
|
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
|
-
- `
|
|
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
|
|
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")
|
{fastembed_haystack-2.2.0 → fastembed_haystack-2.4.0}/tests/test_fastembed_document_embedder.py
RENAMED
|
@@ -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(
|
|
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.2.0 → fastembed_haystack-2.4.0}/tests/test_fastembed_sparse_text_embedder.py
RENAMED
|
File without changes
|