fastembed-haystack 2.1.0__tar.gz → 2.2.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.2.0}/CHANGELOG.md +16 -0
  2. {fastembed_haystack-2.1.0 → fastembed_haystack-2.2.0}/PKG-INFO +2 -2
  3. fastembed_haystack-2.2.0/pydoc/config_docusaurus.yml +18 -0
  4. {fastembed_haystack-2.1.0 → fastembed_haystack-2.2.0}/pyproject.toml +16 -4
  5. {fastembed_haystack-2.1.0 → fastembed_haystack-2.2.0}/src/haystack_integrations/components/embedders/fastembed/embedding_backend/fastembed_backend.py +2 -2
  6. {fastembed_haystack-2.1.0 → fastembed_haystack-2.2.0}/src/haystack_integrations/components/embedders/fastembed/fastembed_document_embedder.py +2 -0
  7. {fastembed_haystack-2.1.0 → fastembed_haystack-2.2.0}/src/haystack_integrations/components/embedders/fastembed/fastembed_sparse_document_embedder.py +1 -0
  8. {fastembed_haystack-2.1.0 → fastembed_haystack-2.2.0}/src/haystack_integrations/components/rankers/fastembed/__init__.py +2 -1
  9. fastembed_haystack-2.2.0/src/haystack_integrations/components/rankers/fastembed/late_interaction_ranker.py +230 -0
  10. {fastembed_haystack-2.1.0 → fastembed_haystack-2.2.0}/src/haystack_integrations/components/rankers/fastembed/ranker.py +8 -9
  11. fastembed_haystack-2.2.0/tests/test_fastembed_late_interaction_ranker.py +350 -0
  12. fastembed_haystack-2.1.0/pydoc/config_docusaurus.yml +0 -32
  13. {fastembed_haystack-2.1.0 → fastembed_haystack-2.2.0}/.gitignore +0 -0
  14. {fastembed_haystack-2.1.0 → fastembed_haystack-2.2.0}/LICENSE.txt +0 -0
  15. {fastembed_haystack-2.1.0 → fastembed_haystack-2.2.0}/README.md +0 -0
  16. {fastembed_haystack-2.1.0 → fastembed_haystack-2.2.0}/examples/example.py +0 -0
  17. {fastembed_haystack-2.1.0 → fastembed_haystack-2.2.0}/examples/ranker_example.py +0 -0
  18. {fastembed_haystack-2.1.0 → fastembed_haystack-2.2.0}/examples/sparse_example.py +0 -0
  19. {fastembed_haystack-2.1.0 → fastembed_haystack-2.2.0}/src/haystack_integrations/components/embedders/fastembed/__init__.py +0 -0
  20. {fastembed_haystack-2.1.0 → fastembed_haystack-2.2.0}/src/haystack_integrations/components/embedders/fastembed/embedding_backend/__init__.py +0 -0
  21. {fastembed_haystack-2.1.0 → fastembed_haystack-2.2.0}/src/haystack_integrations/components/embedders/fastembed/fastembed_sparse_text_embedder.py +0 -0
  22. {fastembed_haystack-2.1.0 → fastembed_haystack-2.2.0}/src/haystack_integrations/components/embedders/fastembed/fastembed_text_embedder.py +0 -0
  23. {fastembed_haystack-2.1.0 → fastembed_haystack-2.2.0}/src/haystack_integrations/components/embedders/py.typed +0 -0
  24. {fastembed_haystack-2.1.0 → fastembed_haystack-2.2.0}/src/haystack_integrations/components/rankers/py.typed +0 -0
  25. {fastembed_haystack-2.1.0 → fastembed_haystack-2.2.0}/tests/__init__.py +0 -0
  26. {fastembed_haystack-2.1.0 → fastembed_haystack-2.2.0}/tests/test_fastembed_backend.py +0 -0
  27. {fastembed_haystack-2.1.0 → fastembed_haystack-2.2.0}/tests/test_fastembed_document_embedder.py +0 -0
  28. {fastembed_haystack-2.1.0 → fastembed_haystack-2.2.0}/tests/test_fastembed_ranker.py +0 -0
  29. {fastembed_haystack-2.1.0 → fastembed_haystack-2.2.0}/tests/test_fastembed_sparse_document_embedder.py +0 -0
  30. {fastembed_haystack-2.1.0 → fastembed_haystack-2.2.0}/tests/test_fastembed_sparse_text_embedder.py +0 -0
  31. {fastembed_haystack-2.1.0 → fastembed_haystack-2.2.0}/tests/test_fastembed_text_embedder.py +0 -0
@@ -1,5 +1,21 @@
1
1
  # Changelog
2
2
 
3
+ ## [integrations/fastembed-v2.1.0] - 2026-02-13
4
+
5
+ ### 📚 Documentation
6
+
7
+ - *(fastembed)* Remove explicit warm_up from examples (#2830)
8
+
9
+ ### 🌀 Miscellaneous
10
+
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
+ - Fix typing issue
15
+ - Fix tying
16
+ - Update docstrings
17
+ - Remove explicit warm_up calls from integration tests
18
+
3
19
  ## [integrations/fastembed-v2.0.0] - 2026-01-12
4
20
 
5
21
  ### 🧹 Chores
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: fastembed-haystack
3
- Version: 2.1.0
3
+ Version: 2.2.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,7 @@ 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
+ ) -> None:
50
52
  """
51
53
  Creates an instance of the 'FastembedRanker'.
52
54
 
@@ -114,7 +116,7 @@ class FastembedRanker:
114
116
  """
115
117
  return default_from_dict(cls, data)
116
118
 
117
- def warm_up(self):
119
+ def warm_up(self) -> None:
118
120
  """
119
121
  Initializes the component.
120
122
  """
@@ -129,6 +131,7 @@ class FastembedRanker:
129
131
  def _prepare_fastembed_input_docs(self, documents: list[Document]) -> list[str]:
130
132
  """
131
133
  Prepare the input by concatenating the document text with the metadata fields specified.
134
+
132
135
  :param documents: The list of Document objects.
133
136
 
134
137
  :return: A list of strings to be given as input to Fastembed model.
@@ -196,10 +199,6 @@ class FastembedRanker:
196
199
  # Sort the list of tuples by the score in descending order
197
200
  sorted_doc_scores = sorted(doc_scores, key=lambda x: x[1], reverse=True)
198
201
 
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)
202
+ top_k_documents = [replace(doc, score=score) for doc, score in sorted_doc_scores[:top_k]]
204
203
 
205
204
  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
@@ -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