fastembed-haystack 2.0.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.
- {fastembed_haystack-2.0.0 → fastembed_haystack-2.2.0}/CHANGELOG.md +23 -0
- {fastembed_haystack-2.0.0 → fastembed_haystack-2.2.0}/PKG-INFO +2 -2
- {fastembed_haystack-2.0.0 → fastembed_haystack-2.2.0}/examples/example.py +0 -1
- {fastembed_haystack-2.0.0 → fastembed_haystack-2.2.0}/examples/ranker_example.py +0 -1
- {fastembed_haystack-2.0.0 → fastembed_haystack-2.2.0}/examples/sparse_example.py +0 -1
- fastembed_haystack-2.2.0/pydoc/config_docusaurus.yml +18 -0
- {fastembed_haystack-2.0.0 → fastembed_haystack-2.2.0}/pyproject.toml +16 -4
- {fastembed_haystack-2.0.0 → fastembed_haystack-2.2.0}/src/haystack_integrations/components/embedders/fastembed/embedding_backend/fastembed_backend.py +2 -2
- {fastembed_haystack-2.0.0 → fastembed_haystack-2.2.0}/src/haystack_integrations/components/embedders/fastembed/fastembed_document_embedder.py +2 -2
- {fastembed_haystack-2.0.0 → fastembed_haystack-2.2.0}/src/haystack_integrations/components/embedders/fastembed/fastembed_sparse_document_embedder.py +1 -2
- {fastembed_haystack-2.0.0 → fastembed_haystack-2.2.0}/src/haystack_integrations/components/embedders/fastembed/fastembed_sparse_text_embedder.py +0 -1
- {fastembed_haystack-2.0.0 → fastembed_haystack-2.2.0}/src/haystack_integrations/components/embedders/fastembed/fastembed_text_embedder.py +0 -1
- {fastembed_haystack-2.0.0 → fastembed_haystack-2.2.0}/src/haystack_integrations/components/rankers/fastembed/__init__.py +2 -1
- fastembed_haystack-2.2.0/src/haystack_integrations/components/rankers/fastembed/late_interaction_ranker.py +230 -0
- {fastembed_haystack-2.0.0 → fastembed_haystack-2.2.0}/src/haystack_integrations/components/rankers/fastembed/ranker.py +10 -12
- {fastembed_haystack-2.0.0 → fastembed_haystack-2.2.0}/tests/test_fastembed_document_embedder.py +14 -1
- fastembed_haystack-2.2.0/tests/test_fastembed_late_interaction_ranker.py +350 -0
- {fastembed_haystack-2.0.0 → fastembed_haystack-2.2.0}/tests/test_fastembed_ranker.py +15 -2
- {fastembed_haystack-2.0.0 → fastembed_haystack-2.2.0}/tests/test_fastembed_sparse_document_embedder.py +13 -2
- {fastembed_haystack-2.0.0 → fastembed_haystack-2.2.0}/tests/test_fastembed_sparse_text_embedder.py +14 -1
- {fastembed_haystack-2.0.0 → fastembed_haystack-2.2.0}/tests/test_fastembed_text_embedder.py +14 -1
- fastembed_haystack-2.0.0/pydoc/config_docusaurus.yml +0 -32
- {fastembed_haystack-2.0.0 → fastembed_haystack-2.2.0}/.gitignore +0 -0
- {fastembed_haystack-2.0.0 → fastembed_haystack-2.2.0}/LICENSE.txt +0 -0
- {fastembed_haystack-2.0.0 → fastembed_haystack-2.2.0}/README.md +0 -0
- {fastembed_haystack-2.0.0 → fastembed_haystack-2.2.0}/src/haystack_integrations/components/embedders/fastembed/__init__.py +0 -0
- {fastembed_haystack-2.0.0 → fastembed_haystack-2.2.0}/src/haystack_integrations/components/embedders/fastembed/embedding_backend/__init__.py +0 -0
- {fastembed_haystack-2.0.0 → fastembed_haystack-2.2.0}/src/haystack_integrations/components/embedders/py.typed +0 -0
- {fastembed_haystack-2.0.0 → fastembed_haystack-2.2.0}/src/haystack_integrations/components/rankers/py.typed +0 -0
- {fastembed_haystack-2.0.0 → fastembed_haystack-2.2.0}/tests/__init__.py +0 -0
- {fastembed_haystack-2.0.0 → fastembed_haystack-2.2.0}/tests/test_fastembed_backend.py +0 -0
|
@@ -1,5 +1,28 @@
|
|
|
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
|
+
|
|
19
|
+
## [integrations/fastembed-v2.0.0] - 2026-01-12
|
|
20
|
+
|
|
21
|
+
### 🧹 Chores
|
|
22
|
+
|
|
23
|
+
- [**breaking**] Fastembed - drop Python 3.9 and use X|Y typing (#2698)
|
|
24
|
+
|
|
25
|
+
|
|
3
26
|
## [integrations/fastembed-v1.6.0] - 2026-01-09
|
|
4
27
|
|
|
5
28
|
### 🚀 Features
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: fastembed-haystack
|
|
3
|
-
Version: 2.
|
|
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
|
+
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.
|
|
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
|
|
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:
|
|
@@ -29,8 +30,6 @@ class FastembedDocumentEmbedder:
|
|
|
29
30
|
batch_size=256,
|
|
30
31
|
)
|
|
31
32
|
|
|
32
|
-
doc_embedder.warm_up()
|
|
33
|
-
|
|
34
33
|
# Text taken from PubMed QA Dataset (https://huggingface.co/datasets/pubmed_qa)
|
|
35
34
|
document_list = [
|
|
36
35
|
Document(
|
|
@@ -112,6 +111,7 @@ class FastembedDocumentEmbedder:
|
|
|
112
111
|
def to_dict(self) -> dict[str, Any]:
|
|
113
112
|
"""
|
|
114
113
|
Serializes the component to a dictionary.
|
|
114
|
+
|
|
115
115
|
:returns:
|
|
116
116
|
Dictionary with serialized data.
|
|
117
117
|
"""
|
|
@@ -28,8 +28,6 @@ class FastembedSparseDocumentEmbedder:
|
|
|
28
28
|
batch_size=32,
|
|
29
29
|
)
|
|
30
30
|
|
|
31
|
-
sparse_doc_embedder.warm_up()
|
|
32
|
-
|
|
33
31
|
# Text taken from PubMed QA Dataset (https://huggingface.co/datasets/pubmed_qa)
|
|
34
32
|
document_list = [
|
|
35
33
|
Document(
|
|
@@ -108,6 +106,7 @@ class FastembedSparseDocumentEmbedder:
|
|
|
108
106
|
def to_dict(self) -> dict[str, Any]:
|
|
109
107
|
"""
|
|
110
108
|
Serializes the component to a dictionary.
|
|
109
|
+
|
|
111
110
|
:returns:
|
|
112
111
|
Dictionary with serialized data.
|
|
113
112
|
"""
|
|
@@ -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
|
-
|
|
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.
|
|
@@ -177,13 +180,12 @@ class FastembedRanker:
|
|
|
177
180
|
raise ValueError(msg)
|
|
178
181
|
|
|
179
182
|
if self._model is None:
|
|
180
|
-
|
|
181
|
-
raise RuntimeError(msg)
|
|
183
|
+
self.warm_up()
|
|
182
184
|
|
|
183
185
|
fastembed_input_docs = self._prepare_fastembed_input_docs(documents)
|
|
184
186
|
|
|
185
187
|
scores = list(
|
|
186
|
-
self._model.rerank(
|
|
188
|
+
self._model.rerank( # type: ignore[union-attr]
|
|
187
189
|
query=query,
|
|
188
190
|
documents=fastembed_input_docs,
|
|
189
191
|
batch_size=self.batch_size,
|
|
@@ -197,10 +199,6 @@ class FastembedRanker:
|
|
|
197
199
|
# Sort the list of tuples by the score in descending order
|
|
198
200
|
sorted_doc_scores = sorted(doc_scores, key=lambda x: x[1], reverse=True)
|
|
199
201
|
|
|
200
|
-
|
|
201
|
-
top_k_documents = []
|
|
202
|
-
for doc, score in sorted_doc_scores[:top_k]:
|
|
203
|
-
doc.score = score
|
|
204
|
-
top_k_documents.append(doc)
|
|
202
|
+
top_k_documents = [replace(doc, score=score) for doc, score in sorted_doc_scores[:top_k]]
|
|
205
203
|
|
|
206
204
|
return {"documents": top_k_documents}
|
{fastembed_haystack-2.0.0 → fastembed_haystack-2.2.0}/tests/test_fastembed_document_embedder.py
RENAMED
|
@@ -281,12 +281,25 @@ class TestFastembedDocumentEmbedder:
|
|
|
281
281
|
parallel=None,
|
|
282
282
|
)
|
|
283
283
|
|
|
284
|
+
def test_run_calls_warm_up(self):
|
|
285
|
+
embedder = FastembedDocumentEmbedder()
|
|
286
|
+
assert embedder.embedding_backend is None
|
|
287
|
+
|
|
288
|
+
mock_backend = MagicMock()
|
|
289
|
+
mock_backend.embed.return_value = [[0.1, 0.2, 0.3]]
|
|
290
|
+
|
|
291
|
+
with patch.object(
|
|
292
|
+
embedder, "warm_up", side_effect=lambda: setattr(embedder, "embedding_backend", mock_backend)
|
|
293
|
+
) as mock_warm_up:
|
|
294
|
+
embedder.run(documents=[Document(content="test document")])
|
|
295
|
+
|
|
296
|
+
mock_warm_up.assert_called_once()
|
|
297
|
+
|
|
284
298
|
@pytest.mark.integration
|
|
285
299
|
def test_run(self):
|
|
286
300
|
embedder = FastembedDocumentEmbedder(
|
|
287
301
|
model="BAAI/bge-small-en-v1.5",
|
|
288
302
|
)
|
|
289
|
-
embedder.warm_up()
|
|
290
303
|
|
|
291
304
|
doc = Document(content="Parton energy loss in QCD matter")
|
|
292
305
|
|
|
@@ -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
|
|
@@ -2,7 +2,7 @@
|
|
|
2
2
|
#
|
|
3
3
|
# SPDX-License-Identifier: Apache-2.0
|
|
4
4
|
|
|
5
|
-
from unittest.mock import MagicMock
|
|
5
|
+
from unittest.mock import MagicMock, patch
|
|
6
6
|
|
|
7
7
|
import numpy as np
|
|
8
8
|
import pytest
|
|
@@ -257,10 +257,23 @@ class TestFastembedRanker:
|
|
|
257
257
|
parallel=None,
|
|
258
258
|
)
|
|
259
259
|
|
|
260
|
+
def test_run_calls_warm_up(self):
|
|
261
|
+
"""
|
|
262
|
+
Unit test to check that warm_up is called when run is called for the first time.
|
|
263
|
+
"""
|
|
264
|
+
ranker = FastembedRanker()
|
|
265
|
+
|
|
266
|
+
mock_model = MagicMock()
|
|
267
|
+
mock_model.rerank.return_value = [0.5]
|
|
268
|
+
|
|
269
|
+
with patch.object(ranker, "warm_up", side_effect=lambda: setattr(ranker, "_model", mock_model)) as mock_warm_up:
|
|
270
|
+
ranker.run(query="test query", documents=[Document(content="test document")])
|
|
271
|
+
|
|
272
|
+
mock_warm_up.assert_called_once()
|
|
273
|
+
|
|
260
274
|
@pytest.mark.integration
|
|
261
275
|
def test_run(self):
|
|
262
276
|
ranker = FastembedRanker(model_name="Xenova/ms-marco-MiniLM-L-6-v2", top_k=2)
|
|
263
|
-
ranker.warm_up()
|
|
264
277
|
|
|
265
278
|
query = "Who is maintaining Qdrant?"
|
|
266
279
|
documents = [
|
|
@@ -304,6 +304,19 @@ class TestFastembedSparseDocumentEmbedderDoc:
|
|
|
304
304
|
|
|
305
305
|
assert embedder.model_kwargs == bm25_config
|
|
306
306
|
|
|
307
|
+
def test_run_calls_warm_up(self):
|
|
308
|
+
embedder = FastembedSparseDocumentEmbedder()
|
|
309
|
+
|
|
310
|
+
mock_backend = MagicMock()
|
|
311
|
+
mock_backend.embed.return_value = [{"indices": [0], "values": [0.5]}]
|
|
312
|
+
|
|
313
|
+
with patch.object(
|
|
314
|
+
embedder, "warm_up", side_effect=lambda: setattr(embedder, "embedding_backend", mock_backend)
|
|
315
|
+
) as mock_warm_up:
|
|
316
|
+
embedder.run(documents=[Document(content="test document")])
|
|
317
|
+
|
|
318
|
+
mock_warm_up.assert_called_once()
|
|
319
|
+
|
|
307
320
|
@pytest.mark.integration
|
|
308
321
|
def test_run_with_model_kwargs(self):
|
|
309
322
|
"""
|
|
@@ -317,7 +330,6 @@ class TestFastembedSparseDocumentEmbedderDoc:
|
|
|
317
330
|
model="Qdrant/bm42-all-minilm-l6-v2-attentions",
|
|
318
331
|
model_kwargs=bm42_config,
|
|
319
332
|
)
|
|
320
|
-
embedder.warm_up()
|
|
321
333
|
|
|
322
334
|
doc = Document(content="Example content using BM42")
|
|
323
335
|
|
|
@@ -336,7 +348,6 @@ class TestFastembedSparseDocumentEmbedderDoc:
|
|
|
336
348
|
embedder = FastembedSparseDocumentEmbedder(
|
|
337
349
|
model="prithivida/Splade_PP_en_v1",
|
|
338
350
|
)
|
|
339
|
-
embedder.warm_up()
|
|
340
351
|
|
|
341
352
|
doc = Document(content="Parton energy loss in QCD matter")
|
|
342
353
|
|
{fastembed_haystack-2.0.0 → fastembed_haystack-2.2.0}/tests/test_fastembed_sparse_text_embedder.py
RENAMED
|
@@ -224,6 +224,20 @@ class TestFastembedSparseTextEmbedder:
|
|
|
224
224
|
|
|
225
225
|
assert embedder.model_kwargs == bm25_config
|
|
226
226
|
|
|
227
|
+
def test_run_calls_warm_up(self):
|
|
228
|
+
embedder = FastembedSparseTextEmbedder()
|
|
229
|
+
assert embedder.embedding_backend is None
|
|
230
|
+
|
|
231
|
+
mock_backend = MagicMock()
|
|
232
|
+
mock_backend.embed.return_value = [[0.1, 0.2, 0.3]]
|
|
233
|
+
|
|
234
|
+
with patch.object(
|
|
235
|
+
embedder, "warm_up", side_effect=lambda: setattr(embedder, "embedding_backend", mock_backend)
|
|
236
|
+
) as mock_warm_up:
|
|
237
|
+
embedder.run(text="test text")
|
|
238
|
+
|
|
239
|
+
mock_warm_up.assert_called_once()
|
|
240
|
+
|
|
227
241
|
@pytest.mark.integration
|
|
228
242
|
def test_run_with_model_kwargs(self):
|
|
229
243
|
"""
|
|
@@ -258,7 +272,6 @@ class TestFastembedSparseTextEmbedder:
|
|
|
258
272
|
embedder = FastembedSparseTextEmbedder(
|
|
259
273
|
model="prithivida/Splade_PP_en_v1",
|
|
260
274
|
)
|
|
261
|
-
embedder.warm_up()
|
|
262
275
|
|
|
263
276
|
text = "Parton energy loss in QCD matter"
|
|
264
277
|
|
|
@@ -202,12 +202,25 @@ class TestFastembedTextEmbedder:
|
|
|
202
202
|
with pytest.raises(TypeError, match="FastembedTextEmbedder expects a string as input"):
|
|
203
203
|
embedder.run(text=list_integers_input)
|
|
204
204
|
|
|
205
|
+
def test_run_calls_warm_up(self):
|
|
206
|
+
embedder = FastembedTextEmbedder()
|
|
207
|
+
assert embedder.embedding_backend is None
|
|
208
|
+
|
|
209
|
+
mock_backend = MagicMock()
|
|
210
|
+
mock_backend.embed.return_value = [[0.1, 0.2, 0.3]]
|
|
211
|
+
|
|
212
|
+
with patch.object(
|
|
213
|
+
embedder, "warm_up", side_effect=lambda: setattr(embedder, "embedding_backend", mock_backend)
|
|
214
|
+
) as mock_warm_up:
|
|
215
|
+
embedder.run(text="test text")
|
|
216
|
+
|
|
217
|
+
mock_warm_up.assert_called_once()
|
|
218
|
+
|
|
205
219
|
@pytest.mark.integration
|
|
206
220
|
def test_run(self):
|
|
207
221
|
embedder = FastembedTextEmbedder(
|
|
208
222
|
model="BAAI/bge-small-en-v1.5",
|
|
209
223
|
)
|
|
210
|
-
embedder.warm_up()
|
|
211
224
|
|
|
212
225
|
text = "Parton energy loss in QCD matter"
|
|
213
226
|
|
|
@@ -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
|
|
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
|