fastembed-haystack 2.1.0__tar.gz → 2.3.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 (31) hide show
  1. {fastembed_haystack-2.1.0 → fastembed_haystack-2.3.0}/CHANGELOG.md +30 -0
  2. {fastembed_haystack-2.1.0 → fastembed_haystack-2.3.0}/PKG-INFO +2 -2
  3. fastembed_haystack-2.3.0/pydoc/config_docusaurus.yml +18 -0
  4. {fastembed_haystack-2.1.0 → fastembed_haystack-2.3.0}/pyproject.toml +16 -4
  5. {fastembed_haystack-2.1.0 → fastembed_haystack-2.3.0}/src/haystack_integrations/components/embedders/fastembed/embedding_backend/fastembed_backend.py +2 -2
  6. {fastembed_haystack-2.1.0 → fastembed_haystack-2.3.0}/src/haystack_integrations/components/embedders/fastembed/fastembed_document_embedder.py +2 -0
  7. {fastembed_haystack-2.1.0 → fastembed_haystack-2.3.0}/src/haystack_integrations/components/embedders/fastembed/fastembed_sparse_document_embedder.py +1 -0
  8. {fastembed_haystack-2.1.0 → fastembed_haystack-2.3.0}/src/haystack_integrations/components/rankers/fastembed/__init__.py +2 -1
  9. fastembed_haystack-2.3.0/src/haystack_integrations/components/rankers/fastembed/late_interaction_ranker.py +230 -0
  10. {fastembed_haystack-2.1.0 → fastembed_haystack-2.3.0}/src/haystack_integrations/components/rankers/fastembed/ranker.py +17 -9
  11. fastembed_haystack-2.3.0/tests/test_fastembed_late_interaction_ranker.py +350 -0
  12. {fastembed_haystack-2.1.0 → fastembed_haystack-2.3.0}/tests/test_fastembed_ranker.py +24 -0
  13. fastembed_haystack-2.1.0/pydoc/config_docusaurus.yml +0 -32
  14. {fastembed_haystack-2.1.0 → fastembed_haystack-2.3.0}/.gitignore +0 -0
  15. {fastembed_haystack-2.1.0 → fastembed_haystack-2.3.0}/LICENSE.txt +0 -0
  16. {fastembed_haystack-2.1.0 → fastembed_haystack-2.3.0}/README.md +0 -0
  17. {fastembed_haystack-2.1.0 → fastembed_haystack-2.3.0}/examples/example.py +0 -0
  18. {fastembed_haystack-2.1.0 → fastembed_haystack-2.3.0}/examples/ranker_example.py +0 -0
  19. {fastembed_haystack-2.1.0 → fastembed_haystack-2.3.0}/examples/sparse_example.py +0 -0
  20. {fastembed_haystack-2.1.0 → fastembed_haystack-2.3.0}/src/haystack_integrations/components/embedders/fastembed/__init__.py +0 -0
  21. {fastembed_haystack-2.1.0 → fastembed_haystack-2.3.0}/src/haystack_integrations/components/embedders/fastembed/embedding_backend/__init__.py +0 -0
  22. {fastembed_haystack-2.1.0 → fastembed_haystack-2.3.0}/src/haystack_integrations/components/embedders/fastembed/fastembed_sparse_text_embedder.py +0 -0
  23. {fastembed_haystack-2.1.0 → fastembed_haystack-2.3.0}/src/haystack_integrations/components/embedders/fastembed/fastembed_text_embedder.py +0 -0
  24. {fastembed_haystack-2.1.0 → fastembed_haystack-2.3.0}/src/haystack_integrations/components/embedders/py.typed +0 -0
  25. {fastembed_haystack-2.1.0 → fastembed_haystack-2.3.0}/src/haystack_integrations/components/rankers/py.typed +0 -0
  26. {fastembed_haystack-2.1.0 → fastembed_haystack-2.3.0}/tests/__init__.py +0 -0
  27. {fastembed_haystack-2.1.0 → fastembed_haystack-2.3.0}/tests/test_fastembed_backend.py +0 -0
  28. {fastembed_haystack-2.1.0 → fastembed_haystack-2.3.0}/tests/test_fastembed_document_embedder.py +0 -0
  29. {fastembed_haystack-2.1.0 → fastembed_haystack-2.3.0}/tests/test_fastembed_sparse_document_embedder.py +0 -0
  30. {fastembed_haystack-2.1.0 → fastembed_haystack-2.3.0}/tests/test_fastembed_sparse_text_embedder.py +0 -0
  31. {fastembed_haystack-2.1.0 → fastembed_haystack-2.3.0}/tests/test_fastembed_text_embedder.py +0 -0
@@ -1,5 +1,35 @@
1
1
  # Changelog
2
2
 
3
+ ## [integrations/fastembed-v2.2.0] - 2026-04-16
4
+
5
+ ### 🚀 Features
6
+
7
+ - Update `FastembedRanker` to auto call `warm_up` on first run (#2834)
8
+ - Add FastembedLateInteractionRanker for ColBERT late-interaction reranking (#3135)
9
+
10
+ ### 🐛 Bug Fixes
11
+
12
+ - Replace in-place dataclass mutations with dataclasses.replace() (#3112)
13
+
14
+ ### 📚 Documentation
15
+
16
+ - *(fastembed)* Remove explicit warm_up from examples (#2830)
17
+ - Simplify pydoc configs (#2855)
18
+
19
+ ### 🧪 Testing
20
+
21
+ - Track test coverage for all integrations (#3065)
22
+
23
+ ### 🧹 Chores
24
+
25
+ - Add ANN type annotations to faiss, fastembed, firecrawl, github, google_ai (#2988)
26
+ - Enforce ruff docstring rules (D102/D103/D205/D209/D213/D417/D419) in integrations 11-20 (#3009)
27
+
28
+ ### 🌀 Miscellaneous
29
+
30
+ - Fix typing issue
31
+ - Fix tying
32
+
3
33
  ## [integrations/fastembed-v2.0.0] - 2026-01-12
4
34
 
5
35
  ### 🧹 Chores
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: fastembed-haystack
3
- Version: 2.1.0
3
+ Version: 2.3.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
@@ -19,7 +19,7 @@ Classifier: Programming Language :: Python :: Implementation :: CPython
19
19
  Classifier: Programming Language :: Python :: Implementation :: PyPy
20
20
  Requires-Python: >=3.10
21
21
  Requires-Dist: fastembed>=0.4.2
22
- Requires-Dist: haystack-ai>=2.22.0
22
+ Requires-Dist: haystack-ai>=2.24.1
23
23
  Description-Content-Type: text/markdown
24
24
 
25
25
  # fastembed-haystack
@@ -0,0 +1,18 @@
1
+ loaders:
2
+ - modules:
3
+ - haystack_integrations.components.embedders.fastembed.fastembed_document_embedder
4
+ - haystack_integrations.components.embedders.fastembed.fastembed_text_embedder
5
+ - haystack_integrations.components.embedders.fastembed.fastembed_sparse_document_embedder
6
+ - haystack_integrations.components.embedders.fastembed.fastembed_sparse_text_embedder
7
+ - haystack_integrations.components.rankers.fastembed.ranker
8
+ - haystack_integrations.components.rankers.fastembed.late_interaction_ranker
9
+ search_path: [../src]
10
+ processors:
11
+ - type: filter
12
+ documented_only: true
13
+ skip_empty_modules: true
14
+ renderer:
15
+ description: FastEmbed integration for Haystack
16
+ id: fastembed-embedders
17
+ filename: fastembed.md
18
+ title: FastEmbed
@@ -22,7 +22,7 @@ classifiers = [
22
22
  "Programming Language :: Python :: Implementation :: CPython",
23
23
  "Programming Language :: Python :: Implementation :: PyPy",
24
24
  ]
25
- dependencies = ["haystack-ai>=2.22.0", "fastembed>=0.4.2"]
25
+ dependencies = ["haystack-ai>=2.24.1", "fastembed>=0.4.2"]
26
26
 
27
27
  [project.urls]
28
28
  Source = "https://github.com/deepset-ai/haystack-core-integrations"
@@ -44,7 +44,7 @@ git_describe_command = 'git describe --tags --match="integrations/fastembed-v[0-
44
44
  installer = "uv"
45
45
  dependencies = ["haystack-pydoc-tools", "ruff"]
46
46
  [tool.hatch.envs.default.scripts]
47
- docs = ["pydoc-markdown pydoc/config_docusaurus.yml"]
47
+ docs = ["haystack-pydoc pydoc/config_docusaurus.yml"]
48
48
  fmt = "ruff check --fix {args}; ruff format {args}"
49
49
  fmt-check = "ruff check {args} && ruff format --check {args}"
50
50
 
@@ -63,7 +63,8 @@ dependencies = [
63
63
  unit = 'pytest -m "not integration" {args:tests}'
64
64
  integration = 'pytest -m "integration" {args:tests}'
65
65
  all = 'pytest {args:tests}'
66
- cov-retry = 'pytest --cov=haystack_integrations --reruns 3 --reruns-delay 30 -x {args:tests}'
66
+ unit-cov-retry = 'pytest --cov=haystack_integrations --reruns 3 --reruns-delay 30 -x -m "not integration" {args:tests}'
67
+ integration-cov-append-retry = 'pytest --cov=haystack_integrations --cov-append --reruns 3 --reruns-delay 30 -x -m "integration" {args:tests}'
67
68
  types = """mypy -p haystack_integrations.components.embedders.fastembed \
68
69
  -p haystack_integrations.components.rankers.fastembed {args}"""
69
70
 
@@ -79,9 +80,17 @@ line-length = 120
79
80
  [tool.ruff.lint]
80
81
  select = [
81
82
  "A",
83
+ "ANN",
82
84
  "ARG",
83
85
  "B",
84
86
  "C",
87
+ "D102", # Missing docstring in public method
88
+ "D103", # Missing docstring in public function
89
+ "D205", # 1 blank line required between summary line and description
90
+ "D209", # Closing triple quotes go to new line
91
+ "D213", # summary lines must be positioned on the second physical line of the docstring
92
+ "D417", # Missing argument descriptions in the docstring
93
+ "D419", # Docstring is empty
85
94
  "DTZ",
86
95
  "E",
87
96
  "EM",
@@ -115,6 +124,8 @@ ignore = [
115
124
  "S105",
116
125
  "S106",
117
126
  "S107",
127
+ # Allow `Any` - used legitimately for dynamic types and SDK boundaries
128
+ "ANN401",
118
129
  # Ignore complexity
119
130
  "C901",
120
131
  "PLR0911",
@@ -131,13 +142,14 @@ ban-relative-imports = "parents"
131
142
 
132
143
  [tool.ruff.lint.per-file-ignores]
133
144
  # Tests can use magic values, assertions, and relative imports
134
- "tests/**/*" = ["PLR2004", "S101", "TID252"]
145
+ "tests/**/*" = ["D", "PLR2004", "S101", "TID252", "ANN"]
135
146
  # examples can contain "print" commands
136
147
  "examples/**/*" = ["T201", "E501"]
137
148
 
138
149
  [tool.coverage.run]
139
150
  source = ["haystack_integrations"]
140
151
  branch = true
152
+ relative_files = true
141
153
  parallel = false
142
154
 
143
155
  [tool.coverage.report]
@@ -44,7 +44,7 @@ class _FastembedEmbeddingBackend:
44
44
  cache_dir: str | None = None,
45
45
  threads: int | None = None,
46
46
  local_files_only: bool = False,
47
- ):
47
+ ) -> None:
48
48
  self.model = TextEmbedding(
49
49
  model_name=model_name, cache_dir=cache_dir, threads=threads, local_files_only=local_files_only
50
50
  )
@@ -103,7 +103,7 @@ class _FastembedSparseEmbeddingBackend:
103
103
  threads: int | None = None,
104
104
  local_files_only: bool = False,
105
105
  model_kwargs: dict[str, Any] | None = None,
106
- ):
106
+ ) -> None:
107
107
  model_kwargs = model_kwargs or {}
108
108
 
109
109
  self.model = SparseTextEmbedding(
@@ -14,6 +14,7 @@ from .embedding_backend.fastembed_backend import _FastembedEmbeddingBackend, _Fa
14
14
  class FastembedDocumentEmbedder:
15
15
  """
16
16
  FastembedDocumentEmbedder computes Document embeddings using Fastembed embedding models.
17
+
17
18
  The embedding of each Document is stored in the `embedding` field of the Document.
18
19
 
19
20
  Usage example:
@@ -110,6 +111,7 @@ class FastembedDocumentEmbedder:
110
111
  def to_dict(self) -> dict[str, Any]:
111
112
  """
112
113
  Serializes the component to a dictionary.
114
+
113
115
  :returns:
114
116
  Dictionary with serialized data.
115
117
  """
@@ -106,6 +106,7 @@ class FastembedSparseDocumentEmbedder:
106
106
  def to_dict(self) -> dict[str, Any]:
107
107
  """
108
108
  Serializes the component to a dictionary.
109
+
109
110
  :returns:
110
111
  Dictionary with serialized data.
111
112
  """
@@ -2,6 +2,7 @@
2
2
  #
3
3
  # SPDX-License-Identifier: Apache-2.0
4
4
 
5
+ from .late_interaction_ranker import FastembedLateInteractionRanker
5
6
  from .ranker import FastembedRanker
6
7
 
7
- __all__ = ["FastembedRanker"]
8
+ __all__ = ["FastembedLateInteractionRanker", "FastembedRanker"]
@@ -0,0 +1,230 @@
1
+ # SPDX-FileCopyrightText: 2024-present deepset GmbH <info@deepset.ai>
2
+ #
3
+ # SPDX-License-Identifier: Apache-2.0
4
+
5
+ from dataclasses import replace
6
+ from typing import Any
7
+
8
+ import numpy as np
9
+ from haystack import Document, component, default_from_dict, default_to_dict, logging
10
+ from haystack.utils.misc import _deduplicate_documents
11
+
12
+ from fastembed import LateInteractionTextEmbedding
13
+
14
+ logger = logging.getLogger(__name__)
15
+
16
+
17
+ @component
18
+ class FastembedLateInteractionRanker:
19
+ """
20
+ Ranks Documents based on their similarity to the query using ColBERT models via Fastembed.
21
+
22
+ Uses late interaction (MaxSim) scoring to compute token-level similarity between
23
+ query and document embeddings, then ranks documents accordingly.
24
+
25
+ See https://qdrant.github.io/fastembed/examples/Supported_Models/ for supported models.
26
+
27
+ Usage example:
28
+ ```python
29
+ from haystack import Document
30
+ from haystack_integrations.components.rankers.fastembed import FastembedLateInteractionRanker
31
+
32
+ ranker = FastembedLateInteractionRanker(model_name="colbert-ir/colbertv2.0", top_k=2)
33
+
34
+ docs = [Document(content="Paris"), Document(content="Berlin")]
35
+ query = "What is the capital of germany?"
36
+ output = ranker.run(query=query, documents=docs)
37
+ print(output["documents"][0].content)
38
+
39
+ # Berlin
40
+ ```
41
+ """
42
+
43
+ def __init__(
44
+ self,
45
+ model_name: str = "colbert-ir/colbertv2.0",
46
+ top_k: int = 10,
47
+ cache_dir: str | None = None,
48
+ threads: int | None = None,
49
+ batch_size: int = 64,
50
+ parallel: int | None = None,
51
+ local_files_only: bool = False,
52
+ meta_fields_to_embed: list[str] | None = None,
53
+ meta_data_separator: str = "\n",
54
+ score_threshold: float | None = None,
55
+ ) -> None:
56
+ """
57
+ Creates an instance of the 'FastembedLateInteractionRanker'.
58
+
59
+ :param model_name: Fastembed ColBERT model name. Check the list of supported models in the
60
+ [Fastembed documentation](https://qdrant.github.io/fastembed/examples/Supported_Models/).
61
+ :param top_k: The maximum number of documents to return.
62
+ :param cache_dir: The path to the cache directory.
63
+ Can be set using the `FASTEMBED_CACHE_PATH` env variable.
64
+ Defaults to `fastembed_cache` in the system's temp directory.
65
+ :param threads: The number of threads single onnxruntime session can use. Defaults to None.
66
+ :param batch_size: Number of strings to encode at once.
67
+ :param parallel:
68
+ If > 1, data-parallel encoding will be used, recommended for offline encoding of large datasets.
69
+ If 0, use all available cores.
70
+ If None, don't use data-parallel processing, use default onnxruntime threading instead.
71
+ :param local_files_only: If `True`, only use the model files in the `cache_dir`.
72
+ :param meta_fields_to_embed: List of meta fields that should be concatenated
73
+ with the document content for reranking.
74
+ :param meta_data_separator: Separator used to concatenate the meta fields
75
+ to the Document content.
76
+ :param score_threshold: If provided, only documents with a score above the threshold are returned.
77
+ Note that ColBERT scores are unnormalized sums and typically range from 3 to 25.
78
+ """
79
+ if top_k <= 0:
80
+ msg = f"top_k must be > 0, but got {top_k}"
81
+ raise ValueError(msg)
82
+
83
+ self.model_name = model_name
84
+ self.top_k = top_k
85
+ self.cache_dir = cache_dir
86
+ self.threads = threads
87
+ self.batch_size = batch_size
88
+ self.parallel = parallel
89
+ self.local_files_only = local_files_only
90
+ self.meta_fields_to_embed = meta_fields_to_embed or []
91
+ self.meta_data_separator = meta_data_separator
92
+ self.score_threshold = score_threshold
93
+ self._model: LateInteractionTextEmbedding | None = None
94
+
95
+ def to_dict(self) -> dict[str, Any]:
96
+ """
97
+ Serializes the component to a dictionary.
98
+
99
+ :returns:
100
+ Dictionary with serialized data.
101
+ """
102
+ return default_to_dict(
103
+ self,
104
+ model_name=self.model_name,
105
+ top_k=self.top_k,
106
+ cache_dir=self.cache_dir,
107
+ threads=self.threads,
108
+ batch_size=self.batch_size,
109
+ parallel=self.parallel,
110
+ local_files_only=self.local_files_only,
111
+ meta_fields_to_embed=self.meta_fields_to_embed,
112
+ meta_data_separator=self.meta_data_separator,
113
+ score_threshold=self.score_threshold,
114
+ )
115
+
116
+ @classmethod
117
+ def from_dict(cls, data: dict[str, Any]) -> "FastembedLateInteractionRanker":
118
+ """
119
+ Deserializes the component from a dictionary.
120
+
121
+ :param data:
122
+ The dictionary to deserialize from.
123
+ :returns:
124
+ The deserialized component.
125
+ """
126
+ return default_from_dict(cls, data)
127
+
128
+ def warm_up(self) -> None:
129
+ """
130
+ Initializes the component.
131
+ """
132
+ if self._model is None:
133
+ self._model = LateInteractionTextEmbedding(
134
+ model_name=self.model_name,
135
+ cache_dir=self.cache_dir,
136
+ threads=self.threads,
137
+ local_files_only=self.local_files_only,
138
+ )
139
+
140
+ def _prepare_fastembed_input_docs(self, documents: list[Document]) -> list[str]:
141
+ """
142
+ Prepare the input by concatenating the document text with the metadata fields specified.
143
+
144
+ :param documents: The list of Document objects.
145
+
146
+ :return: A list of strings to be given as input to Fastembed model.
147
+ """
148
+ concatenated_input_list = []
149
+ for doc in documents:
150
+ meta_values_to_embed = [
151
+ str(doc.meta[key]) for key in self.meta_fields_to_embed if key in doc.meta and doc.meta.get(key)
152
+ ]
153
+ concatenated_input = self.meta_data_separator.join([*meta_values_to_embed, doc.content or ""])
154
+ concatenated_input_list.append(concatenated_input)
155
+
156
+ return concatenated_input_list
157
+
158
+ def _get_telemetry_data(self) -> dict[str, Any]:
159
+ """
160
+ Data that is sent to Posthog for usage analytics.
161
+ """
162
+ return {"model": self.model_name}
163
+
164
+ @component.output_types(documents=list[Document])
165
+ def run(self, query: str, documents: list[Document], top_k: int | None = None) -> dict[str, list[Document]]:
166
+ """
167
+ Returns a list of documents ranked by their similarity to the given query using ColBERT MaxSim scoring.
168
+
169
+ :param query:
170
+ The input query to compare the documents to.
171
+ :param documents:
172
+ A list of documents to be ranked.
173
+ :param top_k:
174
+ The maximum number of documents to return.
175
+
176
+ :returns:
177
+ A dictionary with the following keys:
178
+ - `documents`: A list of documents closest to the query, sorted from most similar to least similar.
179
+
180
+ :raises ValueError: If `top_k` is not > 0.
181
+ """
182
+ if not isinstance(documents, list) or (documents and not isinstance(documents[0], Document)):
183
+ msg = "FastembedLateInteractionRanker expects a list of Documents as input. "
184
+ raise TypeError(msg)
185
+ if query == "":
186
+ msg = "No query provided"
187
+ raise ValueError(msg)
188
+
189
+ if not documents:
190
+ return {"documents": []}
191
+
192
+ top_k = top_k or self.top_k
193
+ if top_k <= 0:
194
+ msg = f"top_k must be > 0, but got {top_k}"
195
+ raise ValueError(msg)
196
+
197
+ if self._model is None:
198
+ self.warm_up()
199
+
200
+ documents = _deduplicate_documents(documents)
201
+ fastembed_input_docs = self._prepare_fastembed_input_docs(documents)
202
+
203
+ assert self._model is not None # noqa: S101
204
+ query_embedding = next(iter(self._model.query_embed(query)))
205
+ document_embeddings = list(
206
+ self._model.embed(
207
+ fastembed_input_docs,
208
+ batch_size=self.batch_size,
209
+ parallel=self.parallel,
210
+ )
211
+ )
212
+
213
+ query_emb = np.array(query_embedding)
214
+ scores = []
215
+ for raw_doc_emb in document_embeddings:
216
+ doc_emb = np.array(raw_doc_emb)
217
+ sim = np.matmul(query_emb, doc_emb.T)
218
+ max_sim = np.max(sim, axis=1)
219
+ scores.append(float(np.sum(max_sim)))
220
+
221
+ doc_scores = list(zip(documents, scores, strict=True))
222
+ sorted_doc_scores = sorted(doc_scores, key=lambda x: x[1], reverse=True)
223
+
224
+ top_k_documents = [replace(doc, score=score) for doc, score in sorted_doc_scores[:top_k]]
225
+
226
+ score_threshold = self.score_threshold
227
+ if score_threshold is not None:
228
+ top_k_documents = [doc for doc in top_k_documents if doc.score is not None and doc.score >= score_threshold]
229
+
230
+ return {"documents": top_k_documents}
@@ -2,6 +2,7 @@
2
2
  #
3
3
  # SPDX-License-Identifier: Apache-2.0
4
4
 
5
+ from dataclasses import replace
5
6
  from typing import Any
6
7
 
7
8
  from haystack import Document, component, default_from_dict, default_to_dict, logging
@@ -14,8 +15,9 @@ logger = logging.getLogger(__name__)
14
15
  @component
15
16
  class FastembedRanker:
16
17
  """
17
- Ranks Documents based on their similarity to the query using
18
- [Fastembed models](https://qdrant.github.io/fastembed/examples/Supported_Models/).
18
+ Ranks Documents based on their similarity to the query using Fastembed models.
19
+
20
+ See https://qdrant.github.io/fastembed/examples/Supported_Models/ for supported models.
19
21
 
20
22
  Documents are indexed from most to least semantically relevant to the query.
21
23
 
@@ -46,7 +48,8 @@ class FastembedRanker:
46
48
  local_files_only: bool = False,
47
49
  meta_fields_to_embed: list[str] | None = None,
48
50
  meta_data_separator: str = "\n",
49
- ):
51
+ score_threshold: float | None = None,
52
+ ) -> None:
50
53
  """
51
54
  Creates an instance of the 'FastembedRanker'.
52
55
 
@@ -66,6 +69,8 @@ class FastembedRanker:
66
69
  with the document content for reranking.
67
70
  :param meta_data_separator: Separator used to concatenate the meta fields
68
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.
69
74
  """
70
75
  if top_k <= 0:
71
76
  msg = f"top_k must be > 0, but got {top_k}"
@@ -80,6 +85,7 @@ class FastembedRanker:
80
85
  self.local_files_only = local_files_only
81
86
  self.meta_fields_to_embed = meta_fields_to_embed or []
82
87
  self.meta_data_separator = meta_data_separator
88
+ self.score_threshold = score_threshold
83
89
  self._model: TextCrossEncoder | None = None
84
90
 
85
91
  def to_dict(self) -> dict[str, Any]:
@@ -100,6 +106,7 @@ class FastembedRanker:
100
106
  local_files_only=self.local_files_only,
101
107
  meta_fields_to_embed=self.meta_fields_to_embed,
102
108
  meta_data_separator=self.meta_data_separator,
109
+ score_threshold=self.score_threshold,
103
110
  )
104
111
 
105
112
  @classmethod
@@ -114,7 +121,7 @@ class FastembedRanker:
114
121
  """
115
122
  return default_from_dict(cls, data)
116
123
 
117
- def warm_up(self):
124
+ def warm_up(self) -> None:
118
125
  """
119
126
  Initializes the component.
120
127
  """
@@ -129,6 +136,7 @@ class FastembedRanker:
129
136
  def _prepare_fastembed_input_docs(self, documents: list[Document]) -> list[str]:
130
137
  """
131
138
  Prepare the input by concatenating the document text with the metadata fields specified.
139
+
132
140
  :param documents: The list of Document objects.
133
141
 
134
142
  :return: A list of strings to be given as input to Fastembed model.
@@ -196,10 +204,10 @@ class FastembedRanker:
196
204
  # Sort the list of tuples by the score in descending order
197
205
  sorted_doc_scores = sorted(doc_scores, key=lambda x: x[1], reverse=True)
198
206
 
199
- # Get the top_k documents
200
- top_k_documents = []
201
- for doc, score in sorted_doc_scores[:top_k]:
202
- doc.score = score
203
- top_k_documents.append(doc)
207
+ top_k_documents = [replace(doc, score=score) for doc, score in sorted_doc_scores[:top_k]]
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]
204
212
 
205
213
  return {"documents": top_k_documents}
@@ -0,0 +1,350 @@
1
+ # SPDX-FileCopyrightText: 2024-present deepset GmbH <info@deepset.ai>
2
+ #
3
+ # SPDX-License-Identifier: Apache-2.0
4
+
5
+ from unittest.mock import MagicMock, patch
6
+
7
+ import numpy as np
8
+ import pytest
9
+ from haystack import Document, default_from_dict
10
+
11
+ from haystack_integrations.components.rankers.fastembed.late_interaction_ranker import (
12
+ FastembedLateInteractionRanker,
13
+ )
14
+
15
+
16
+ class TestFastembedLateInteractionRanker:
17
+ def test_init_default(self):
18
+ """
19
+ Test default initialization parameters for FastembedLateInteractionRanker.
20
+ """
21
+ ranker = FastembedLateInteractionRanker(model_name="colbert-ir/colbertv2.0")
22
+ assert ranker.model_name == "colbert-ir/colbertv2.0"
23
+ assert ranker.top_k == 10
24
+ assert ranker.cache_dir is None
25
+ assert ranker.threads is None
26
+ assert ranker.batch_size == 64
27
+ assert ranker.parallel is None
28
+ assert not ranker.local_files_only
29
+ assert ranker.meta_fields_to_embed == []
30
+ assert ranker.meta_data_separator == "\n"
31
+
32
+ def test_init_with_parameters(self):
33
+ """
34
+ Test custom initialization parameters for FastembedLateInteractionRanker.
35
+ """
36
+ ranker = FastembedLateInteractionRanker(
37
+ model_name="colbert-ir/colbertv2.0",
38
+ top_k=64,
39
+ cache_dir="fake_dir",
40
+ threads=2,
41
+ batch_size=50,
42
+ parallel=1,
43
+ local_files_only=True,
44
+ meta_fields_to_embed=["test_field"],
45
+ meta_data_separator=" | ",
46
+ )
47
+ assert ranker.model_name == "colbert-ir/colbertv2.0"
48
+ assert ranker.top_k == 64
49
+ assert ranker.cache_dir == "fake_dir"
50
+ assert ranker.threads == 2
51
+ assert ranker.batch_size == 50
52
+ assert ranker.parallel == 1
53
+ assert ranker.local_files_only
54
+ assert ranker.meta_fields_to_embed == ["test_field"]
55
+ assert ranker.meta_data_separator == " | "
56
+
57
+ def test_init_with_incorrect_input(self):
58
+ """
59
+ Test for checking incorrect input format on init.
60
+ """
61
+ with pytest.raises(
62
+ ValueError,
63
+ match="top_k must be > 0, but got 0",
64
+ ):
65
+ FastembedLateInteractionRanker(model_name="colbert-ir/colbertv2.0", top_k=0)
66
+
67
+ with pytest.raises(
68
+ ValueError,
69
+ match="top_k must be > 0, but got -3",
70
+ ):
71
+ FastembedLateInteractionRanker(model_name="colbert-ir/colbertv2.0", top_k=-3)
72
+
73
+ def test_to_dict(self):
74
+ """
75
+ Test serialization of FastembedLateInteractionRanker to a dictionary, using default initialization parameters.
76
+ """
77
+ ranker = FastembedLateInteractionRanker(model_name="colbert-ir/colbertv2.0")
78
+ ranker_dict = ranker.to_dict()
79
+ assert ranker_dict == {
80
+ "type": "haystack_integrations.components.rankers.fastembed"
81
+ ".late_interaction_ranker.FastembedLateInteractionRanker",
82
+ "init_parameters": {
83
+ "model_name": "colbert-ir/colbertv2.0",
84
+ "top_k": 10,
85
+ "cache_dir": None,
86
+ "threads": None,
87
+ "batch_size": 64,
88
+ "parallel": None,
89
+ "local_files_only": False,
90
+ "meta_fields_to_embed": [],
91
+ "meta_data_separator": "\n",
92
+ "score_threshold": None,
93
+ },
94
+ }
95
+
96
+ def test_to_dict_with_custom_init_parameters(self):
97
+ """
98
+ Test serialization of FastembedLateInteractionRanker to a dictionary, using custom initialization parameters.
99
+ """
100
+ ranker = FastembedLateInteractionRanker(
101
+ model_name="colbert-ir/colbertv2.0",
102
+ cache_dir="fake_dir",
103
+ threads=2,
104
+ top_k=5,
105
+ batch_size=50,
106
+ parallel=1,
107
+ local_files_only=True,
108
+ meta_fields_to_embed=["test_field"],
109
+ meta_data_separator=" | ",
110
+ )
111
+ ranker_dict = ranker.to_dict()
112
+ assert ranker_dict == {
113
+ "type": "haystack_integrations.components.rankers.fastembed"
114
+ ".late_interaction_ranker.FastembedLateInteractionRanker",
115
+ "init_parameters": {
116
+ "model_name": "colbert-ir/colbertv2.0",
117
+ "cache_dir": "fake_dir",
118
+ "threads": 2,
119
+ "top_k": 5,
120
+ "batch_size": 50,
121
+ "parallel": 1,
122
+ "local_files_only": True,
123
+ "meta_fields_to_embed": ["test_field"],
124
+ "meta_data_separator": " | ",
125
+ "score_threshold": None,
126
+ },
127
+ }
128
+
129
+ def test_from_dict(self):
130
+ """
131
+ Test deserialization of FastembedLateInteractionRanker from a dictionary, using default init parameters.
132
+ """
133
+ ranker_dict = {
134
+ "type": "haystack_integrations.components.rankers.fastembed"
135
+ ".late_interaction_ranker.FastembedLateInteractionRanker",
136
+ "init_parameters": {
137
+ "model_name": "colbert-ir/colbertv2.0",
138
+ "cache_dir": None,
139
+ "threads": None,
140
+ "top_k": 5,
141
+ "batch_size": 50,
142
+ "parallel": None,
143
+ "local_files_only": False,
144
+ "meta_fields_to_embed": [],
145
+ "meta_data_separator": "\n",
146
+ },
147
+ }
148
+ ranker = default_from_dict(FastembedLateInteractionRanker, ranker_dict)
149
+ assert ranker.model_name == "colbert-ir/colbertv2.0"
150
+ assert ranker.cache_dir is None
151
+ assert ranker.threads is None
152
+ assert ranker.top_k == 5
153
+ assert ranker.batch_size == 50
154
+ assert ranker.parallel is None
155
+ assert not ranker.local_files_only
156
+ assert ranker.meta_fields_to_embed == []
157
+ assert ranker.meta_data_separator == "\n"
158
+
159
+ def test_from_dict_with_custom_init_parameters(self):
160
+ """
161
+ Test deserialization of FastembedLateInteractionRanker from a dictionary, using custom init parameters.
162
+ """
163
+ ranker_dict = {
164
+ "type": "haystack_integrations.components.rankers.fastembed"
165
+ ".late_interaction_ranker.FastembedLateInteractionRanker",
166
+ "init_parameters": {
167
+ "model_name": "colbert-ir/colbertv2.0",
168
+ "cache_dir": "fake_dir",
169
+ "threads": 2,
170
+ "top_k": 5,
171
+ "batch_size": 50,
172
+ "parallel": 1,
173
+ "local_files_only": True,
174
+ "meta_fields_to_embed": ["test_field"],
175
+ "meta_data_separator": " | ",
176
+ },
177
+ }
178
+ ranker = default_from_dict(FastembedLateInteractionRanker, ranker_dict)
179
+ assert ranker.model_name == "colbert-ir/colbertv2.0"
180
+ assert ranker.cache_dir == "fake_dir"
181
+ assert ranker.threads == 2
182
+ assert ranker.top_k == 5
183
+ assert ranker.batch_size == 50
184
+ assert ranker.parallel == 1
185
+ assert ranker.local_files_only
186
+ assert ranker.meta_fields_to_embed == ["test_field"]
187
+ assert ranker.meta_data_separator == " | "
188
+
189
+ def test_run_incorrect_input_format(self):
190
+ """
191
+ Test for checking incorrect input format.
192
+ """
193
+ ranker = FastembedLateInteractionRanker(model_name="colbert-ir/colbertv2.0")
194
+ ranker._model = "mock_model"
195
+
196
+ query = "query"
197
+ string_input = "text"
198
+ list_integers_input = [1, 2, 3]
199
+ list_document = [Document("Document 1")]
200
+
201
+ with pytest.raises(
202
+ TypeError,
203
+ match=r"FastembedLateInteractionRanker expects a list of Documents as input\.",
204
+ ):
205
+ ranker.run(query=query, documents=string_input)
206
+
207
+ with pytest.raises(
208
+ TypeError,
209
+ match=r"FastembedLateInteractionRanker expects a list of Documents as input\.",
210
+ ):
211
+ ranker.run(query=query, documents=list_integers_input)
212
+
213
+ with pytest.raises(
214
+ ValueError,
215
+ match="No query provided",
216
+ ):
217
+ ranker.run(query="", documents=list_document)
218
+
219
+ with pytest.raises(
220
+ ValueError,
221
+ match="top_k must be > 0, but got -3",
222
+ ):
223
+ ranker.run(query=query, documents=list_document, top_k=-3)
224
+
225
+ def test_run_empty_document_list(self):
226
+ """
227
+ Test for no error when sending no documents.
228
+ """
229
+ ranker = FastembedLateInteractionRanker(model_name="colbert-ir/colbertv2.0")
230
+ ranker._model = "mock_model"
231
+
232
+ query = "query"
233
+ list_document = []
234
+
235
+ result = ranker.run(query=query, documents=list_document)
236
+ assert len(result["documents"]) == 0
237
+
238
+ def test_embed_metadata(self):
239
+ """
240
+ Tests the embedding of metadata fields in document content for ranking.
241
+ """
242
+ ranker = FastembedLateInteractionRanker(
243
+ model_name="colbert-ir/colbertv2.0",
244
+ meta_fields_to_embed=["meta_field"],
245
+ )
246
+ mock_model = MagicMock()
247
+ # query_embed returns one embedding with shape (num_query_tokens, embedding_dim)
248
+ mock_model.query_embed.return_value = iter([np.random.rand(32, 128)])
249
+ # embed returns one embedding per document with shape (num_doc_tokens, embedding_dim)
250
+ mock_model.embed.return_value = iter([np.random.rand(64, 128) for _ in range(5)])
251
+ ranker._model = mock_model
252
+
253
+ documents = [Document(content=f"document-number {i}", meta={"meta_field": f"meta_value {i}"}) for i in range(5)]
254
+ query = "test"
255
+ ranker.run(query=query, documents=documents)
256
+
257
+ mock_model.embed.assert_called_once_with(
258
+ [
259
+ "meta_value 0\ndocument-number 0",
260
+ "meta_value 1\ndocument-number 1",
261
+ "meta_value 2\ndocument-number 2",
262
+ "meta_value 3\ndocument-number 3",
263
+ "meta_value 4\ndocument-number 4",
264
+ ],
265
+ batch_size=64,
266
+ parallel=None,
267
+ )
268
+
269
+ def test_warm_up_called_once(self):
270
+ """
271
+ Test that calling warm_up() twice only initializes the model once.
272
+ """
273
+ ranker = FastembedLateInteractionRanker()
274
+ with patch(
275
+ "haystack_integrations.components.rankers.fastembed.late_interaction_ranker.LateInteractionTextEmbedding"
276
+ ) as mock_cls:
277
+ ranker.warm_up()
278
+ ranker.warm_up()
279
+ mock_cls.assert_called_once()
280
+
281
+ def test_run_calls_warm_up(self):
282
+ """
283
+ Unit test to check that warm_up is called when run is called for the first time.
284
+ """
285
+ ranker = FastembedLateInteractionRanker()
286
+
287
+ mock_model = MagicMock()
288
+ mock_model.query_embed.return_value = iter([np.random.rand(32, 128)])
289
+ mock_model.embed.return_value = iter([np.random.rand(64, 128)])
290
+
291
+ with patch.object(ranker, "warm_up", side_effect=lambda: setattr(ranker, "_model", mock_model)) as mock_warm_up:
292
+ ranker.run(query="test query", documents=[Document(content="test document")])
293
+
294
+ mock_warm_up.assert_called_once()
295
+
296
+ def test_run_with_mock(self):
297
+ """
298
+ Test that MaxSim scoring produces correct ranking order.
299
+ """
300
+ ranker = FastembedLateInteractionRanker(model_name="colbert-ir/colbertv2.0", top_k=2)
301
+
302
+ mock_model = MagicMock()
303
+ # Query embedding: 2 tokens, 4 dims
304
+ mock_model.query_embed.return_value = iter([np.array([[1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0]])])
305
+ # Doc embeddings: 2 docs, each with 3 tokens, 4 dims
306
+ # Doc 0: low similarity to query
307
+ # Doc 1: high similarity to query
308
+ mock_model.embed.return_value = iter(
309
+ [
310
+ np.array([[0.1, 0.1, 0.9, 0.0], [0.1, 0.0, 0.0, 0.9], [0.0, 0.1, 0.9, 0.0]]),
311
+ np.array([[0.9, 0.0, 0.0, 0.1], [0.0, 0.9, 0.1, 0.0], [0.1, 0.1, 0.8, 0.0]]),
312
+ ]
313
+ )
314
+ ranker._model = mock_model
315
+
316
+ documents = [
317
+ Document(content="low similarity doc"),
318
+ Document(content="high similarity doc"),
319
+ ]
320
+
321
+ result = ranker.run(query="test", documents=documents)
322
+
323
+ assert len(result["documents"]) == 2
324
+ assert result["documents"][0].content == "high similarity doc"
325
+ assert result["documents"][0].score > result["documents"][1].score
326
+
327
+ @pytest.mark.integration
328
+ def test_run(self):
329
+ ranker = FastembedLateInteractionRanker(model_name="colbert-ir/colbertv2.0", top_k=2)
330
+
331
+ query = "Who is maintaining Qdrant?"
332
+ documents = [
333
+ Document(
334
+ content="This is built to be faster and lighter than other embedding "
335
+ "libraries e.g. Transformers, Sentence-Transformers, etc."
336
+ ),
337
+ Document(content="This is some random input"),
338
+ Document(content="fastembed is supported by and maintained by Qdrant."),
339
+ ]
340
+
341
+ result = ranker.run(query=query, documents=documents)
342
+
343
+ assert len(result["documents"]) == 2
344
+ first_document = result["documents"][0]
345
+ second_document = result["documents"][1]
346
+
347
+ assert isinstance(first_document, Document)
348
+ assert isinstance(second_document, Document)
349
+ assert first_document.content == "fastembed is supported by and maintained by Qdrant."
350
+ assert first_document.score > second_document.score
@@ -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.
@@ -1,32 +0,0 @@
1
- loaders:
2
- - ignore_when_discovered:
3
- - __init__
4
- modules:
5
- - haystack_integrations.components.embedders.fastembed.fastembed_document_embedder
6
- - haystack_integrations.components.embedders.fastembed.fastembed_text_embedder
7
- - haystack_integrations.components.embedders.fastembed.fastembed_sparse_document_embedder
8
- - haystack_integrations.components.embedders.fastembed.fastembed_sparse_text_embedder
9
- - haystack_integrations.components.rankers.fastembed.ranker
10
- search_path:
11
- - ../src
12
- type: haystack_pydoc_tools.loaders.CustomPythonLoader
13
- processors:
14
- - do_not_filter_modules: false
15
- documented_only: true
16
- expression: null
17
- skip_empty_modules: true
18
- type: filter
19
- - type: smart
20
- - type: crossref
21
- renderer:
22
- description: FastEmbed integration for Haystack
23
- id: fastembed-embedders
24
- markdown:
25
- add_member_class_prefix: false
26
- add_method_class_prefix: true
27
- classdef_code_block: false
28
- descriptive_class_title: false
29
- descriptive_module_title: true
30
- filename: fastembed.md
31
- title: FastEmbed
32
- type: haystack_pydoc_tools.renderers.DocusaurusRenderer