fastembed-haystack 1.3.0__tar.gz → 1.4.1__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-1.3.0 → fastembed_haystack-1.4.1}/CHANGELOG.md +16 -1
- {fastembed_haystack-1.3.0 → fastembed_haystack-1.4.1}/PKG-INFO +35 -9
- {fastembed_haystack-1.3.0 → fastembed_haystack-1.4.1}/README.md +32 -5
- fastembed_haystack-1.4.1/examples/ranker_example.py +22 -0
- {fastembed_haystack-1.3.0 → fastembed_haystack-1.4.1}/pydoc/config.yml +2 -1
- {fastembed_haystack-1.3.0 → fastembed_haystack-1.4.1}/pyproject.toml +8 -2
- {fastembed_haystack-1.3.0 → fastembed_haystack-1.4.1}/src/haystack_integrations/components/embedders/fastembed/fastembed_sparse_document_embedder.py +3 -3
- {fastembed_haystack-1.3.0 → fastembed_haystack-1.4.1}/src/haystack_integrations/components/embedders/fastembed/fastembed_sparse_text_embedder.py +3 -3
- fastembed_haystack-1.4.1/src/haystack_integrations/components/rankers/fastembed/__init__.py +3 -0
- fastembed_haystack-1.4.1/src/haystack_integrations/components/rankers/fastembed/ranker.py +202 -0
- fastembed_haystack-1.4.1/tests/test_fastembed_ranker.py +292 -0
- {fastembed_haystack-1.3.0 → fastembed_haystack-1.4.1}/tests/test_fastembed_sparse_document_embedder.py +18 -18
- {fastembed_haystack-1.3.0 → fastembed_haystack-1.4.1}/tests/test_fastembed_sparse_text_embedder.py +16 -16
- {fastembed_haystack-1.3.0 → fastembed_haystack-1.4.1}/.gitignore +0 -0
- {fastembed_haystack-1.3.0 → fastembed_haystack-1.4.1}/LICENSE.txt +0 -0
- {fastembed_haystack-1.3.0 → fastembed_haystack-1.4.1}/examples/example.py +0 -0
- {fastembed_haystack-1.3.0 → fastembed_haystack-1.4.1}/examples/sparse_example.py +0 -0
- {fastembed_haystack-1.3.0 → fastembed_haystack-1.4.1}/src/haystack_integrations/components/embedders/fastembed/__init__.py +0 -0
- {fastembed_haystack-1.3.0 → fastembed_haystack-1.4.1}/src/haystack_integrations/components/embedders/fastembed/embedding_backend/__init__.py +0 -0
- {fastembed_haystack-1.3.0 → fastembed_haystack-1.4.1}/src/haystack_integrations/components/embedders/fastembed/embedding_backend/fastembed_backend.py +0 -0
- {fastembed_haystack-1.3.0 → fastembed_haystack-1.4.1}/src/haystack_integrations/components/embedders/fastembed/fastembed_document_embedder.py +0 -0
- {fastembed_haystack-1.3.0 → fastembed_haystack-1.4.1}/src/haystack_integrations/components/embedders/fastembed/fastembed_text_embedder.py +0 -0
- {fastembed_haystack-1.3.0 → fastembed_haystack-1.4.1}/tests/__init__.py +0 -0
- {fastembed_haystack-1.3.0 → fastembed_haystack-1.4.1}/tests/test_fastembed_backend.py +0 -0
- {fastembed_haystack-1.3.0 → fastembed_haystack-1.4.1}/tests/test_fastembed_document_embedder.py +0 -0
- {fastembed_haystack-1.3.0 → fastembed_haystack-1.4.1}/tests/test_fastembed_text_embedder.py +0 -0
|
@@ -1,11 +1,26 @@
|
|
|
1
1
|
# Changelog
|
|
2
2
|
|
|
3
|
-
## [
|
|
3
|
+
## [integrations/fastembed-v1.4.0] - 2024-11-13
|
|
4
|
+
|
|
5
|
+
### ⚙️ Miscellaneous Tasks
|
|
6
|
+
|
|
7
|
+
- Adopt uv as installer (#1142)
|
|
8
|
+
|
|
9
|
+
## [integrations/fastembed-v1.3.0] - 2024-10-07
|
|
10
|
+
|
|
11
|
+
### 🚀 Features
|
|
12
|
+
|
|
13
|
+
- Introduce `model_kwargs` in Sparse Embedders (can be used for BM25 parameters) (#1126)
|
|
14
|
+
|
|
15
|
+
### 🧪 Testing
|
|
16
|
+
|
|
17
|
+
- Do not retry tests in `hatch run test` command (#954)
|
|
4
18
|
|
|
5
19
|
### ⚙️ Miscellaneous Tasks
|
|
6
20
|
|
|
7
21
|
- Retry tests to reduce flakyness (#836)
|
|
8
22
|
- Update ruff invocation to include check parameter (#853)
|
|
23
|
+
- Update ruff linting scripts and settings (#1105)
|
|
9
24
|
|
|
10
25
|
### Fix
|
|
11
26
|
|
|
@@ -1,13 +1,12 @@
|
|
|
1
1
|
Metadata-Version: 2.3
|
|
2
2
|
Name: fastembed-haystack
|
|
3
|
-
Version: 1.
|
|
3
|
+
Version: 1.4.1
|
|
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
|
|
7
7
|
Project-URL: Issues, https://github.com/deepset-ai/haystack-core-integrations/issues
|
|
8
8
|
Author-email: deepset GmbH <info@deepset.ai>
|
|
9
|
-
License
|
|
10
|
-
License-File: LICENSE.txt
|
|
9
|
+
License: Apache-2.0
|
|
11
10
|
Classifier: Development Status :: 4 - Beta
|
|
12
11
|
Classifier: License :: OSI Approved :: Apache Software License
|
|
13
12
|
Classifier: Programming Language :: Python
|
|
@@ -19,7 +18,7 @@ Classifier: Programming Language :: Python :: 3.12
|
|
|
19
18
|
Classifier: Programming Language :: Python :: Implementation :: CPython
|
|
20
19
|
Classifier: Programming Language :: Python :: Implementation :: PyPy
|
|
21
20
|
Requires-Python: >=3.8
|
|
22
|
-
Requires-Dist: fastembed>=0.2
|
|
21
|
+
Requires-Dist: fastembed>=0.4.2
|
|
23
22
|
Requires-Dist: haystack-ai>=2.0.1
|
|
24
23
|
Description-Content-Type: text/markdown
|
|
25
24
|
|
|
@@ -33,6 +32,7 @@ Description-Content-Type: text/markdown
|
|
|
33
32
|
**Table of Contents**
|
|
34
33
|
|
|
35
34
|
- [Installation](#installation)
|
|
35
|
+
- [Usage](#Usage)
|
|
36
36
|
- [License](#license)
|
|
37
37
|
|
|
38
38
|
## Installation
|
|
@@ -58,7 +58,7 @@ embedding = text_embedder.run(text)["embedding"]
|
|
|
58
58
|
|
|
59
59
|
```python
|
|
60
60
|
from haystack_integrations.components.embedders.fastembed import FastembedDocumentEmbedder
|
|
61
|
-
from haystack
|
|
61
|
+
from haystack import Document
|
|
62
62
|
|
|
63
63
|
embedder = FastembedDocumentEmbedder(
|
|
64
64
|
model="BAAI/bge-small-en-v1.5",
|
|
@@ -75,24 +75,50 @@ from haystack_integrations.components.embedders.fastembed import FastembedSparse
|
|
|
75
75
|
|
|
76
76
|
text = "fastembed is supported by and maintained by Qdrant."
|
|
77
77
|
text_embedder = FastembedSparseTextEmbedder(
|
|
78
|
-
model="
|
|
78
|
+
model="prithivida/Splade_PP_en_v1"
|
|
79
79
|
)
|
|
80
80
|
text_embedder.warm_up()
|
|
81
|
-
embedding = text_embedder.run(text)["
|
|
81
|
+
embedding = text_embedder.run(text)["sparse_embedding"]
|
|
82
82
|
```
|
|
83
83
|
|
|
84
84
|
```python
|
|
85
85
|
from haystack_integrations.components.embedders.fastembed import FastembedSparseDocumentEmbedder
|
|
86
|
-
from haystack
|
|
86
|
+
from haystack import Document
|
|
87
87
|
|
|
88
88
|
embedder = FastembedSparseDocumentEmbedder(
|
|
89
|
-
model="
|
|
89
|
+
model="prithivida/Splade_PP_en_v1",
|
|
90
90
|
)
|
|
91
91
|
embedder.warm_up()
|
|
92
92
|
doc = Document(content="fastembed is supported by and maintained by Qdrant.", meta={"long_answer": "no",})
|
|
93
93
|
result = embedder.run(documents=[doc])
|
|
94
94
|
```
|
|
95
95
|
|
|
96
|
+
You can use `FastembedRanker` by importing as:
|
|
97
|
+
|
|
98
|
+
```python
|
|
99
|
+
from haystack import Document
|
|
100
|
+
|
|
101
|
+
from haystack_integrations.components.rankers.fastembed import FastembedRanker
|
|
102
|
+
|
|
103
|
+
query = "Who is maintaining Qdrant?"
|
|
104
|
+
documents = [
|
|
105
|
+
Document(
|
|
106
|
+
content="This is built to be faster and lighter than other embedding libraries e.g. Transformers, Sentence-Transformers, etc."
|
|
107
|
+
),
|
|
108
|
+
Document(content="fastembed is supported by and maintained by Qdrant."),
|
|
109
|
+
]
|
|
110
|
+
|
|
111
|
+
ranker = FastembedRanker(model_name="Xenova/ms-marco-MiniLM-L-6-v2")
|
|
112
|
+
ranker.warm_up()
|
|
113
|
+
reranked_documents = ranker.run(query=query, documents=documents)["documents"]
|
|
114
|
+
|
|
115
|
+
print(reranked_documents[0])
|
|
116
|
+
|
|
117
|
+
# Document(id=...,
|
|
118
|
+
# content: 'fastembed is supported by and maintained by Qdrant.',
|
|
119
|
+
# score: 5.472434997558594..)
|
|
120
|
+
```
|
|
121
|
+
|
|
96
122
|
## License
|
|
97
123
|
|
|
98
124
|
`fastembed-haystack` is distributed under the terms of the [Apache-2.0](https://spdx.org/licenses/Apache-2.0.html) license.
|
|
@@ -8,6 +8,7 @@
|
|
|
8
8
|
**Table of Contents**
|
|
9
9
|
|
|
10
10
|
- [Installation](#installation)
|
|
11
|
+
- [Usage](#Usage)
|
|
11
12
|
- [License](#license)
|
|
12
13
|
|
|
13
14
|
## Installation
|
|
@@ -33,7 +34,7 @@ embedding = text_embedder.run(text)["embedding"]
|
|
|
33
34
|
|
|
34
35
|
```python
|
|
35
36
|
from haystack_integrations.components.embedders.fastembed import FastembedDocumentEmbedder
|
|
36
|
-
from haystack
|
|
37
|
+
from haystack import Document
|
|
37
38
|
|
|
38
39
|
embedder = FastembedDocumentEmbedder(
|
|
39
40
|
model="BAAI/bge-small-en-v1.5",
|
|
@@ -50,24 +51,50 @@ from haystack_integrations.components.embedders.fastembed import FastembedSparse
|
|
|
50
51
|
|
|
51
52
|
text = "fastembed is supported by and maintained by Qdrant."
|
|
52
53
|
text_embedder = FastembedSparseTextEmbedder(
|
|
53
|
-
model="
|
|
54
|
+
model="prithivida/Splade_PP_en_v1"
|
|
54
55
|
)
|
|
55
56
|
text_embedder.warm_up()
|
|
56
|
-
embedding = text_embedder.run(text)["
|
|
57
|
+
embedding = text_embedder.run(text)["sparse_embedding"]
|
|
57
58
|
```
|
|
58
59
|
|
|
59
60
|
```python
|
|
60
61
|
from haystack_integrations.components.embedders.fastembed import FastembedSparseDocumentEmbedder
|
|
61
|
-
from haystack
|
|
62
|
+
from haystack import Document
|
|
62
63
|
|
|
63
64
|
embedder = FastembedSparseDocumentEmbedder(
|
|
64
|
-
model="
|
|
65
|
+
model="prithivida/Splade_PP_en_v1",
|
|
65
66
|
)
|
|
66
67
|
embedder.warm_up()
|
|
67
68
|
doc = Document(content="fastembed is supported by and maintained by Qdrant.", meta={"long_answer": "no",})
|
|
68
69
|
result = embedder.run(documents=[doc])
|
|
69
70
|
```
|
|
70
71
|
|
|
72
|
+
You can use `FastembedRanker` by importing as:
|
|
73
|
+
|
|
74
|
+
```python
|
|
75
|
+
from haystack import Document
|
|
76
|
+
|
|
77
|
+
from haystack_integrations.components.rankers.fastembed import FastembedRanker
|
|
78
|
+
|
|
79
|
+
query = "Who is maintaining Qdrant?"
|
|
80
|
+
documents = [
|
|
81
|
+
Document(
|
|
82
|
+
content="This is built to be faster and lighter than other embedding libraries e.g. Transformers, Sentence-Transformers, etc."
|
|
83
|
+
),
|
|
84
|
+
Document(content="fastembed is supported by and maintained by Qdrant."),
|
|
85
|
+
]
|
|
86
|
+
|
|
87
|
+
ranker = FastembedRanker(model_name="Xenova/ms-marco-MiniLM-L-6-v2")
|
|
88
|
+
ranker.warm_up()
|
|
89
|
+
reranked_documents = ranker.run(query=query, documents=documents)["documents"]
|
|
90
|
+
|
|
91
|
+
print(reranked_documents[0])
|
|
92
|
+
|
|
93
|
+
# Document(id=...,
|
|
94
|
+
# content: 'fastembed is supported by and maintained by Qdrant.',
|
|
95
|
+
# score: 5.472434997558594..)
|
|
96
|
+
```
|
|
97
|
+
|
|
71
98
|
## License
|
|
72
99
|
|
|
73
100
|
`fastembed-haystack` is distributed under the terms of the [Apache-2.0](https://spdx.org/licenses/Apache-2.0.html) license.
|
|
@@ -0,0 +1,22 @@
|
|
|
1
|
+
from haystack import Document
|
|
2
|
+
|
|
3
|
+
from haystack_integrations.components.rankers.fastembed import FastembedRanker
|
|
4
|
+
|
|
5
|
+
query = "Who is maintaining Qdrant?"
|
|
6
|
+
documents = [
|
|
7
|
+
Document(
|
|
8
|
+
content="This is built to be faster and lighter than other embedding libraries e.g. Transformers, Sentence-Transformers, etc."
|
|
9
|
+
),
|
|
10
|
+
Document(content="fastembed is supported by and maintained by Qdrant."),
|
|
11
|
+
]
|
|
12
|
+
|
|
13
|
+
ranker = FastembedRanker(model_name="Xenova/ms-marco-MiniLM-L-6-v2")
|
|
14
|
+
ranker.warm_up()
|
|
15
|
+
reranked_documents = ranker.run(query=query, documents=documents)["documents"]
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
print(reranked_documents[0])
|
|
19
|
+
|
|
20
|
+
# Document(id=...,
|
|
21
|
+
# content: 'fastembed is supported by and maintained by Qdrant.',
|
|
22
|
+
# score: 5.472434997558594..)
|
|
@@ -6,7 +6,8 @@ loaders:
|
|
|
6
6
|
"haystack_integrations.components.embedders.fastembed.fastembed_document_embedder",
|
|
7
7
|
"haystack_integrations.components.embedders.fastembed.fastembed_text_embedder",
|
|
8
8
|
"haystack_integrations.components.embedders.fastembed.fastembed_sparse_document_embedder",
|
|
9
|
-
"haystack_integrations.components.embedders.fastembed.fastembed_sparse_text_embedder"
|
|
9
|
+
"haystack_integrations.components.embedders.fastembed.fastembed_sparse_text_embedder",
|
|
10
|
+
"haystack_integrations.components.rankers.fastembed.ranker"
|
|
10
11
|
]
|
|
11
12
|
ignore_when_discovered: ["__init__"]
|
|
12
13
|
processors:
|
|
@@ -23,7 +23,7 @@ classifiers = [
|
|
|
23
23
|
"Programming Language :: Python :: Implementation :: CPython",
|
|
24
24
|
"Programming Language :: Python :: Implementation :: PyPy",
|
|
25
25
|
]
|
|
26
|
-
dependencies = ["haystack-ai>=2.0.1", "fastembed>=0.2
|
|
26
|
+
dependencies = ["haystack-ai>=2.0.1", "fastembed>=0.4.2"]
|
|
27
27
|
|
|
28
28
|
[project.urls]
|
|
29
29
|
Source = "https://github.com/deepset-ai/haystack-core-integrations"
|
|
@@ -42,6 +42,7 @@ root = "../.."
|
|
|
42
42
|
git_describe_command = 'git describe --tags --match="integrations/fastembed-v[0-9]*"'
|
|
43
43
|
|
|
44
44
|
[tool.hatch.envs.default]
|
|
45
|
+
installer = "uv"
|
|
45
46
|
dependencies = [
|
|
46
47
|
"coverage[toml]>=6.5",
|
|
47
48
|
"pytest",
|
|
@@ -62,8 +63,9 @@ docs = ["pydoc-markdown pydoc/config.yml"]
|
|
|
62
63
|
python = ["3.8", "3.9", "3.10", "3.11", "3.12"]
|
|
63
64
|
|
|
64
65
|
[tool.hatch.envs.lint]
|
|
66
|
+
installer = "uv"
|
|
65
67
|
detached = true
|
|
66
|
-
dependencies = ["black>=23.1.0", "mypy>=1.0.0", "ruff>=0.0.243", "numpy"]
|
|
68
|
+
dependencies = ["pip", "black>=23.1.0", "mypy>=1.0.0", "ruff>=0.0.243", "numpy"]
|
|
67
69
|
[tool.hatch.envs.lint.scripts]
|
|
68
70
|
typing = "mypy --install-types --non-interactive --explicit-package-bases {args:src/ tests}"
|
|
69
71
|
style = ["ruff check {args:.}", "black --check --diff {args:.}"]
|
|
@@ -152,6 +154,10 @@ omit = ["*/tests/*", "*/__init__.py"]
|
|
|
152
154
|
show_missing = true
|
|
153
155
|
exclude_lines = ["no cov", "if __name__ == .__main__.:", "if TYPE_CHECKING:"]
|
|
154
156
|
|
|
157
|
+
[tool.pytest.ini_options]
|
|
158
|
+
minversion = "6.0"
|
|
159
|
+
markers = ["unit: unit tests", "integration: integration tests"]
|
|
160
|
+
|
|
155
161
|
[[tool.mypy.overrides]]
|
|
156
162
|
module = [
|
|
157
163
|
"haystack.*",
|
|
@@ -16,7 +16,7 @@ class FastembedSparseDocumentEmbedder:
|
|
|
16
16
|
from haystack.dataclasses import Document
|
|
17
17
|
|
|
18
18
|
sparse_doc_embedder = FastembedSparseDocumentEmbedder(
|
|
19
|
-
model="
|
|
19
|
+
model="prithivida/Splade_PP_en_v1",
|
|
20
20
|
batch_size=32,
|
|
21
21
|
)
|
|
22
22
|
|
|
@@ -53,7 +53,7 @@ class FastembedSparseDocumentEmbedder:
|
|
|
53
53
|
|
|
54
54
|
def __init__(
|
|
55
55
|
self,
|
|
56
|
-
model: str = "
|
|
56
|
+
model: str = "prithivida/Splade_PP_en_v1",
|
|
57
57
|
cache_dir: Optional[str] = None,
|
|
58
58
|
threads: Optional[int] = None,
|
|
59
59
|
batch_size: int = 32,
|
|
@@ -68,7 +68,7 @@ class FastembedSparseDocumentEmbedder:
|
|
|
68
68
|
Create an FastembedDocumentEmbedder component.
|
|
69
69
|
|
|
70
70
|
:param model: Local path or name of the model in Hugging Face's model hub,
|
|
71
|
-
such as `
|
|
71
|
+
such as `prithivida/Splade_PP_en_v1`.
|
|
72
72
|
:param cache_dir: The path to the cache directory.
|
|
73
73
|
Can be set using the `FASTEMBED_CACHE_PATH` env variable.
|
|
74
74
|
Defaults to `fastembed_cache` in the system's temp directory.
|
|
@@ -19,7 +19,7 @@ class FastembedSparseTextEmbedder:
|
|
|
19
19
|
"The disk comes and it does not, only Windows. Do Not order this if you have a Mac!!")
|
|
20
20
|
|
|
21
21
|
sparse_text_embedder = FastembedSparseTextEmbedder(
|
|
22
|
-
model="
|
|
22
|
+
model="prithivida/Splade_PP_en_v1"
|
|
23
23
|
)
|
|
24
24
|
sparse_text_embedder.warm_up()
|
|
25
25
|
|
|
@@ -29,7 +29,7 @@ class FastembedSparseTextEmbedder:
|
|
|
29
29
|
|
|
30
30
|
def __init__(
|
|
31
31
|
self,
|
|
32
|
-
model: str = "
|
|
32
|
+
model: str = "prithivida/Splade_PP_en_v1",
|
|
33
33
|
cache_dir: Optional[str] = None,
|
|
34
34
|
threads: Optional[int] = None,
|
|
35
35
|
progress_bar: bool = True,
|
|
@@ -40,7 +40,7 @@ class FastembedSparseTextEmbedder:
|
|
|
40
40
|
"""
|
|
41
41
|
Create a FastembedSparseTextEmbedder component.
|
|
42
42
|
|
|
43
|
-
:param model: Local path or name of the model in Fastembed's model hub, such as `
|
|
43
|
+
:param model: Local path or name of the model in Fastembed's model hub, such as `prithivida/Splade_PP_en_v1`
|
|
44
44
|
:param cache_dir: The path to the cache directory.
|
|
45
45
|
Can be set using the `FASTEMBED_CACHE_PATH` env variable.
|
|
46
46
|
Defaults to `fastembed_cache` in the system's temp directory.
|
|
@@ -0,0 +1,202 @@
|
|
|
1
|
+
from typing import Any, Dict, List, Optional
|
|
2
|
+
|
|
3
|
+
from haystack import Document, component, default_from_dict, default_to_dict, logging
|
|
4
|
+
|
|
5
|
+
from fastembed.rerank.cross_encoder import TextCrossEncoder
|
|
6
|
+
|
|
7
|
+
logger = logging.getLogger(__name__)
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
@component
|
|
11
|
+
class FastembedRanker:
|
|
12
|
+
"""
|
|
13
|
+
Ranks Documents based on their similarity to the query using
|
|
14
|
+
[Fastembed models](https://qdrant.github.io/fastembed/examples/Supported_Models/).
|
|
15
|
+
|
|
16
|
+
Documents are indexed from most to least semantically relevant to the query.
|
|
17
|
+
|
|
18
|
+
Usage example:
|
|
19
|
+
```python
|
|
20
|
+
from haystack import Document
|
|
21
|
+
from haystack_integrations.components.rankers.fastembed import FastembedRanker
|
|
22
|
+
|
|
23
|
+
ranker = FastembedRanker(model_name="Xenova/ms-marco-MiniLM-L-6-v2", top_k=2)
|
|
24
|
+
|
|
25
|
+
docs = [Document(content="Paris"), Document(content="Berlin")]
|
|
26
|
+
query = "What is the capital of germany?"
|
|
27
|
+
output = ranker.run(query=query, documents=docs)
|
|
28
|
+
print(output["documents"][0].content)
|
|
29
|
+
|
|
30
|
+
# Berlin
|
|
31
|
+
```
|
|
32
|
+
"""
|
|
33
|
+
|
|
34
|
+
def __init__(
|
|
35
|
+
self,
|
|
36
|
+
model_name: str = "Xenova/ms-marco-MiniLM-L-6-v2",
|
|
37
|
+
top_k: int = 10,
|
|
38
|
+
cache_dir: Optional[str] = None,
|
|
39
|
+
threads: Optional[int] = None,
|
|
40
|
+
batch_size: int = 64,
|
|
41
|
+
parallel: Optional[int] = None,
|
|
42
|
+
local_files_only: bool = False,
|
|
43
|
+
meta_fields_to_embed: Optional[List[str]] = None,
|
|
44
|
+
meta_data_separator: str = "\n",
|
|
45
|
+
):
|
|
46
|
+
"""
|
|
47
|
+
Creates an instance of the 'FastembedRanker'.
|
|
48
|
+
|
|
49
|
+
:param model_name: Fastembed model name. Check the list of supported models in the [Fastembed documentation](https://qdrant.github.io/fastembed/examples/Supported_Models/).
|
|
50
|
+
:param top_k: The maximum number of documents to return.
|
|
51
|
+
:param cache_dir: The path to the cache directory.
|
|
52
|
+
Can be set using the `FASTEMBED_CACHE_PATH` env variable.
|
|
53
|
+
Defaults to `fastembed_cache` in the system's temp directory.
|
|
54
|
+
:param threads: The number of threads single onnxruntime session can use. Defaults to None.
|
|
55
|
+
:param batch_size: Number of strings to encode at once.
|
|
56
|
+
:param parallel:
|
|
57
|
+
If > 1, data-parallel encoding will be used, recommended for offline encoding of large datasets.
|
|
58
|
+
If 0, use all available cores.
|
|
59
|
+
If None, don't use data-parallel processing, use default onnxruntime threading instead.
|
|
60
|
+
:param local_files_only: If `True`, only use the model files in the `cache_dir`.
|
|
61
|
+
:param meta_fields_to_embed: List of meta fields that should be concatenated
|
|
62
|
+
with the document content for reranking.
|
|
63
|
+
:param meta_data_separator: Separator used to concatenate the meta fields
|
|
64
|
+
to the Document content.
|
|
65
|
+
"""
|
|
66
|
+
if top_k <= 0:
|
|
67
|
+
msg = f"top_k must be > 0, but got {top_k}"
|
|
68
|
+
raise ValueError(msg)
|
|
69
|
+
|
|
70
|
+
self.model_name = model_name
|
|
71
|
+
self.top_k = top_k
|
|
72
|
+
self.cache_dir = cache_dir
|
|
73
|
+
self.threads = threads
|
|
74
|
+
self.batch_size = batch_size
|
|
75
|
+
self.parallel = parallel
|
|
76
|
+
self.local_files_only = local_files_only
|
|
77
|
+
self.meta_fields_to_embed = meta_fields_to_embed or []
|
|
78
|
+
self.meta_data_separator = meta_data_separator
|
|
79
|
+
self._model = None
|
|
80
|
+
|
|
81
|
+
def to_dict(self) -> Dict[str, Any]:
|
|
82
|
+
"""
|
|
83
|
+
Serializes the component to a dictionary.
|
|
84
|
+
|
|
85
|
+
:returns:
|
|
86
|
+
Dictionary with serialized data.
|
|
87
|
+
"""
|
|
88
|
+
return default_to_dict(
|
|
89
|
+
self,
|
|
90
|
+
model_name=self.model_name,
|
|
91
|
+
top_k=self.top_k,
|
|
92
|
+
cache_dir=self.cache_dir,
|
|
93
|
+
threads=self.threads,
|
|
94
|
+
batch_size=self.batch_size,
|
|
95
|
+
parallel=self.parallel,
|
|
96
|
+
local_files_only=self.local_files_only,
|
|
97
|
+
meta_fields_to_embed=self.meta_fields_to_embed,
|
|
98
|
+
meta_data_separator=self.meta_data_separator,
|
|
99
|
+
)
|
|
100
|
+
|
|
101
|
+
@classmethod
|
|
102
|
+
def from_dict(cls, data: Dict[str, Any]) -> "FastembedRanker":
|
|
103
|
+
"""
|
|
104
|
+
Deserializes the component from a dictionary.
|
|
105
|
+
|
|
106
|
+
:param data:
|
|
107
|
+
The dictionary to deserialize from.
|
|
108
|
+
:returns:
|
|
109
|
+
The deserialized component.
|
|
110
|
+
"""
|
|
111
|
+
return default_from_dict(cls, data)
|
|
112
|
+
|
|
113
|
+
def warm_up(self):
|
|
114
|
+
"""
|
|
115
|
+
Initializes the component.
|
|
116
|
+
"""
|
|
117
|
+
if self._model is None:
|
|
118
|
+
self._model = TextCrossEncoder(
|
|
119
|
+
model_name=self.model_name,
|
|
120
|
+
cache_dir=self.cache_dir,
|
|
121
|
+
threads=self.threads,
|
|
122
|
+
local_files_only=self.local_files_only,
|
|
123
|
+
)
|
|
124
|
+
|
|
125
|
+
def _prepare_fastembed_input_docs(self, documents: List[Document]) -> List[str]:
|
|
126
|
+
"""
|
|
127
|
+
Prepare the input by concatenating the document text with the metadata fields specified.
|
|
128
|
+
:param documents: The list of Document objects.
|
|
129
|
+
|
|
130
|
+
:return: A list of strings to be given as input to Fastembed model.
|
|
131
|
+
"""
|
|
132
|
+
concatenated_input_list = []
|
|
133
|
+
for doc in documents:
|
|
134
|
+
meta_values_to_embed = [
|
|
135
|
+
str(doc.meta[key]) for key in self.meta_fields_to_embed if key in doc.meta and doc.meta.get(key)
|
|
136
|
+
]
|
|
137
|
+
concatenated_input = self.meta_data_separator.join([*meta_values_to_embed, doc.content or ""])
|
|
138
|
+
concatenated_input_list.append(concatenated_input)
|
|
139
|
+
|
|
140
|
+
return concatenated_input_list
|
|
141
|
+
|
|
142
|
+
@component.output_types(documents=List[Document])
|
|
143
|
+
def run(self, query: str, documents: List[Document], top_k: Optional[int] = None):
|
|
144
|
+
"""
|
|
145
|
+
Returns a list of documents ranked by their similarity to the given query, using FastEmbed.
|
|
146
|
+
|
|
147
|
+
:param query:
|
|
148
|
+
The input query to compare the documents to.
|
|
149
|
+
:param documents:
|
|
150
|
+
A list of documents to be ranked.
|
|
151
|
+
:param top_k:
|
|
152
|
+
The maximum number of documents to return.
|
|
153
|
+
|
|
154
|
+
:returns:
|
|
155
|
+
A dictionary with the following keys:
|
|
156
|
+
- `documents`: A list of documents closest to the query, sorted from most similar to least similar.
|
|
157
|
+
|
|
158
|
+
:raises ValueError: If `top_k` is not > 0.
|
|
159
|
+
"""
|
|
160
|
+
if not isinstance(documents, list) or documents and not isinstance(documents[0], Document):
|
|
161
|
+
msg = "FastembedRanker expects a list of Documents as input. "
|
|
162
|
+
raise TypeError(msg)
|
|
163
|
+
if query == "":
|
|
164
|
+
msg = "No query provided"
|
|
165
|
+
raise ValueError(msg)
|
|
166
|
+
|
|
167
|
+
if not documents:
|
|
168
|
+
return {"documents": []}
|
|
169
|
+
|
|
170
|
+
top_k = top_k or self.top_k
|
|
171
|
+
if top_k <= 0:
|
|
172
|
+
msg = f"top_k must be > 0, but got {top_k}"
|
|
173
|
+
raise ValueError(msg)
|
|
174
|
+
|
|
175
|
+
if self._model is None:
|
|
176
|
+
msg = "The ranker model has not been loaded. Please call warm_up() before running."
|
|
177
|
+
raise RuntimeError(msg)
|
|
178
|
+
|
|
179
|
+
fastembed_input_docs = self._prepare_fastembed_input_docs(documents)
|
|
180
|
+
|
|
181
|
+
scores = list(
|
|
182
|
+
self._model.rerank(
|
|
183
|
+
query=query,
|
|
184
|
+
documents=fastembed_input_docs,
|
|
185
|
+
batch_size=self.batch_size,
|
|
186
|
+
parallel=self.parallel,
|
|
187
|
+
)
|
|
188
|
+
)
|
|
189
|
+
|
|
190
|
+
# Combine the two lists into a single list of tuples
|
|
191
|
+
doc_scores = list(zip(documents, scores))
|
|
192
|
+
|
|
193
|
+
# Sort the list of tuples by the score in descending order
|
|
194
|
+
sorted_doc_scores = sorted(doc_scores, key=lambda x: x[1], reverse=True)
|
|
195
|
+
|
|
196
|
+
# Get the top_k documents
|
|
197
|
+
top_k_documents = []
|
|
198
|
+
for doc, score in sorted_doc_scores[:top_k]:
|
|
199
|
+
doc.score = score
|
|
200
|
+
top_k_documents.append(doc)
|
|
201
|
+
|
|
202
|
+
return {"documents": top_k_documents}
|
|
@@ -0,0 +1,292 @@
|
|
|
1
|
+
from unittest.mock import MagicMock
|
|
2
|
+
|
|
3
|
+
import pytest
|
|
4
|
+
from haystack import Document, default_from_dict
|
|
5
|
+
|
|
6
|
+
from haystack_integrations.components.rankers.fastembed.ranker import (
|
|
7
|
+
FastembedRanker,
|
|
8
|
+
)
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
class TestFastembedRanker:
|
|
12
|
+
def test_init_default(self):
|
|
13
|
+
"""
|
|
14
|
+
Test default initialization parameters for FastembedRanker.
|
|
15
|
+
"""
|
|
16
|
+
ranker = FastembedRanker(model_name="BAAI/bge-reranker-base")
|
|
17
|
+
assert ranker.model_name == "BAAI/bge-reranker-base"
|
|
18
|
+
assert ranker.top_k == 10
|
|
19
|
+
assert ranker.cache_dir is None
|
|
20
|
+
assert ranker.threads is None
|
|
21
|
+
assert ranker.batch_size == 64
|
|
22
|
+
assert ranker.parallel is None
|
|
23
|
+
assert not ranker.local_files_only
|
|
24
|
+
assert ranker.meta_fields_to_embed == []
|
|
25
|
+
assert ranker.meta_data_separator == "\n"
|
|
26
|
+
|
|
27
|
+
def test_init_with_parameters(self):
|
|
28
|
+
"""
|
|
29
|
+
Test custom initialization parameters for FastembedRanker.
|
|
30
|
+
"""
|
|
31
|
+
ranker = FastembedRanker(
|
|
32
|
+
model_name="BAAI/bge-reranker-base",
|
|
33
|
+
top_k=64,
|
|
34
|
+
cache_dir="fake_dir",
|
|
35
|
+
threads=2,
|
|
36
|
+
batch_size=50,
|
|
37
|
+
parallel=1,
|
|
38
|
+
local_files_only=True,
|
|
39
|
+
meta_fields_to_embed=["test_field"],
|
|
40
|
+
meta_data_separator=" | ",
|
|
41
|
+
)
|
|
42
|
+
assert ranker.model_name == "BAAI/bge-reranker-base"
|
|
43
|
+
assert ranker.top_k == 64
|
|
44
|
+
assert ranker.cache_dir == "fake_dir"
|
|
45
|
+
assert ranker.threads == 2
|
|
46
|
+
assert ranker.batch_size == 50
|
|
47
|
+
assert ranker.parallel == 1
|
|
48
|
+
assert ranker.local_files_only
|
|
49
|
+
assert ranker.meta_fields_to_embed == ["test_field"]
|
|
50
|
+
assert ranker.meta_data_separator == " | "
|
|
51
|
+
|
|
52
|
+
def test_init_with_incorrect_input(self):
|
|
53
|
+
"""
|
|
54
|
+
Test for checking incorrect input format on init
|
|
55
|
+
"""
|
|
56
|
+
with pytest.raises(
|
|
57
|
+
ValueError,
|
|
58
|
+
match="top_k must be > 0, but got 0",
|
|
59
|
+
):
|
|
60
|
+
FastembedRanker(model_name="Xenova/ms-marco-MiniLM-L-12-v2", top_k=0)
|
|
61
|
+
|
|
62
|
+
with pytest.raises(
|
|
63
|
+
ValueError,
|
|
64
|
+
match="top_k must be > 0, but got -3",
|
|
65
|
+
):
|
|
66
|
+
FastembedRanker(model_name="Xenova/ms-marco-MiniLM-L-12-v2", top_k=-3)
|
|
67
|
+
|
|
68
|
+
def test_to_dict(self):
|
|
69
|
+
"""
|
|
70
|
+
Test serialization of FastembedRanker to a dictionary, using default initialization parameters.
|
|
71
|
+
"""
|
|
72
|
+
ranker = FastembedRanker(model_name="BAAI/bge-reranker-base")
|
|
73
|
+
ranker_dict = ranker.to_dict()
|
|
74
|
+
assert ranker_dict == {
|
|
75
|
+
"type": "haystack_integrations.components.rankers.fastembed.ranker.FastembedRanker",
|
|
76
|
+
"init_parameters": {
|
|
77
|
+
"model_name": "BAAI/bge-reranker-base",
|
|
78
|
+
"top_k": 10,
|
|
79
|
+
"cache_dir": None,
|
|
80
|
+
"threads": None,
|
|
81
|
+
"batch_size": 64,
|
|
82
|
+
"parallel": None,
|
|
83
|
+
"local_files_only": False,
|
|
84
|
+
"meta_fields_to_embed": [],
|
|
85
|
+
"meta_data_separator": "\n",
|
|
86
|
+
},
|
|
87
|
+
}
|
|
88
|
+
|
|
89
|
+
def test_to_dict_with_custom_init_parameters(self):
|
|
90
|
+
"""
|
|
91
|
+
Test serialization of FastembedRanker to a dictionary, using custom initialization parameters.
|
|
92
|
+
"""
|
|
93
|
+
ranker = FastembedRanker(
|
|
94
|
+
model_name="BAAI/bge-reranker-base",
|
|
95
|
+
cache_dir="fake_dir",
|
|
96
|
+
threads=2,
|
|
97
|
+
top_k=5,
|
|
98
|
+
batch_size=50,
|
|
99
|
+
parallel=1,
|
|
100
|
+
local_files_only=True,
|
|
101
|
+
meta_fields_to_embed=["test_field"],
|
|
102
|
+
meta_data_separator=" | ",
|
|
103
|
+
)
|
|
104
|
+
ranker_dict = ranker.to_dict()
|
|
105
|
+
assert ranker_dict == {
|
|
106
|
+
"type": "haystack_integrations.components.rankers.fastembed.ranker.FastembedRanker",
|
|
107
|
+
"init_parameters": {
|
|
108
|
+
"model_name": "BAAI/bge-reranker-base",
|
|
109
|
+
"cache_dir": "fake_dir",
|
|
110
|
+
"threads": 2,
|
|
111
|
+
"top_k": 5,
|
|
112
|
+
"batch_size": 50,
|
|
113
|
+
"parallel": 1,
|
|
114
|
+
"local_files_only": True,
|
|
115
|
+
"meta_fields_to_embed": ["test_field"],
|
|
116
|
+
"meta_data_separator": " | ",
|
|
117
|
+
},
|
|
118
|
+
}
|
|
119
|
+
|
|
120
|
+
def test_from_dict(self):
|
|
121
|
+
"""
|
|
122
|
+
Test deserialization of FastembedRanker from a dictionary, using default initialization parameters.
|
|
123
|
+
"""
|
|
124
|
+
ranker_dict = {
|
|
125
|
+
"type": "haystack_integrations.components.rankers.fastembed.ranker.FastembedRanker",
|
|
126
|
+
"init_parameters": {
|
|
127
|
+
"model_name": "BAAI/bge-reranker-base",
|
|
128
|
+
"cache_dir": None,
|
|
129
|
+
"threads": None,
|
|
130
|
+
"top_k": 5,
|
|
131
|
+
"batch_size": 50,
|
|
132
|
+
"parallel": None,
|
|
133
|
+
"local_files_only": False,
|
|
134
|
+
"meta_fields_to_embed": [],
|
|
135
|
+
"meta_data_separator": "\n",
|
|
136
|
+
},
|
|
137
|
+
}
|
|
138
|
+
ranker = default_from_dict(FastembedRanker, ranker_dict)
|
|
139
|
+
assert ranker.model_name == "BAAI/bge-reranker-base"
|
|
140
|
+
assert ranker.cache_dir is None
|
|
141
|
+
assert ranker.threads is None
|
|
142
|
+
assert ranker.top_k == 5
|
|
143
|
+
assert ranker.batch_size == 50
|
|
144
|
+
assert ranker.parallel is None
|
|
145
|
+
assert not ranker.local_files_only
|
|
146
|
+
assert ranker.meta_fields_to_embed == []
|
|
147
|
+
assert ranker.meta_data_separator == "\n"
|
|
148
|
+
|
|
149
|
+
def test_from_dict_with_custom_init_parameters(self):
|
|
150
|
+
"""
|
|
151
|
+
Test deserialization of FastembedRanker from a dictionary, using custom initialization parameters.
|
|
152
|
+
"""
|
|
153
|
+
ranker_dict = {
|
|
154
|
+
"type": "haystack_integrations.components.rankers.fastembed.ranker.FastembedRanker",
|
|
155
|
+
"init_parameters": {
|
|
156
|
+
"model_name": "BAAI/bge-reranker-base",
|
|
157
|
+
"cache_dir": "fake_dir",
|
|
158
|
+
"threads": 2,
|
|
159
|
+
"top_k": 5,
|
|
160
|
+
"batch_size": 50,
|
|
161
|
+
"parallel": 1,
|
|
162
|
+
"local_files_only": True,
|
|
163
|
+
"meta_fields_to_embed": ["test_field"],
|
|
164
|
+
"meta_data_separator": " | ",
|
|
165
|
+
},
|
|
166
|
+
}
|
|
167
|
+
ranker = default_from_dict(FastembedRanker, ranker_dict)
|
|
168
|
+
assert ranker.model_name == "BAAI/bge-reranker-base"
|
|
169
|
+
assert ranker.cache_dir == "fake_dir"
|
|
170
|
+
assert ranker.threads == 2
|
|
171
|
+
assert ranker.top_k == 5
|
|
172
|
+
assert ranker.batch_size == 50
|
|
173
|
+
assert ranker.parallel == 1
|
|
174
|
+
assert ranker.local_files_only
|
|
175
|
+
assert ranker.meta_fields_to_embed == ["test_field"]
|
|
176
|
+
assert ranker.meta_data_separator == " | "
|
|
177
|
+
|
|
178
|
+
def test_run_incorrect_input_format(self):
|
|
179
|
+
"""
|
|
180
|
+
Test for checking incorrect input format.
|
|
181
|
+
"""
|
|
182
|
+
ranker = FastembedRanker(model_name="Xenova/ms-marco-MiniLM-L-12-v2")
|
|
183
|
+
ranker._model = "mock_model"
|
|
184
|
+
|
|
185
|
+
query = "query"
|
|
186
|
+
string_input = "text"
|
|
187
|
+
list_integers_input = [1, 2, 3]
|
|
188
|
+
list_document = [Document("Document 1")]
|
|
189
|
+
|
|
190
|
+
with pytest.raises(
|
|
191
|
+
TypeError,
|
|
192
|
+
match="FastembedRanker expects a list of Documents as input.",
|
|
193
|
+
):
|
|
194
|
+
ranker.run(query=query, documents=string_input)
|
|
195
|
+
|
|
196
|
+
with pytest.raises(
|
|
197
|
+
TypeError,
|
|
198
|
+
match="FastembedRanker expects a list of Documents as input.",
|
|
199
|
+
):
|
|
200
|
+
ranker.run(query=query, documents=list_integers_input)
|
|
201
|
+
|
|
202
|
+
with pytest.raises(
|
|
203
|
+
ValueError,
|
|
204
|
+
match="No query provided",
|
|
205
|
+
):
|
|
206
|
+
ranker.run(query="", documents=list_document)
|
|
207
|
+
|
|
208
|
+
with pytest.raises(
|
|
209
|
+
ValueError,
|
|
210
|
+
match="top_k must be > 0, but got -3",
|
|
211
|
+
):
|
|
212
|
+
ranker.run(query=query, documents=list_document, top_k=-3)
|
|
213
|
+
|
|
214
|
+
def test_run_no_warmup(self):
|
|
215
|
+
"""
|
|
216
|
+
Test for checking error when calling without a warmup.
|
|
217
|
+
"""
|
|
218
|
+
ranker = FastembedRanker(model_name="Xenova/ms-marco-MiniLM-L-12-v2")
|
|
219
|
+
|
|
220
|
+
query = "query"
|
|
221
|
+
list_document = [Document("Document 1")]
|
|
222
|
+
|
|
223
|
+
with pytest.raises(
|
|
224
|
+
RuntimeError,
|
|
225
|
+
):
|
|
226
|
+
ranker.run(query=query, documents=list_document)
|
|
227
|
+
|
|
228
|
+
def test_run_empty_document_list(self):
|
|
229
|
+
"""
|
|
230
|
+
Test for no error when sending no documents.
|
|
231
|
+
"""
|
|
232
|
+
ranker = FastembedRanker(model_name="Xenova/ms-marco-MiniLM-L-12-v2")
|
|
233
|
+
ranker._model = "mock_model"
|
|
234
|
+
|
|
235
|
+
query = "query"
|
|
236
|
+
list_document = []
|
|
237
|
+
|
|
238
|
+
result = ranker.run(query=query, documents=list_document)
|
|
239
|
+
assert len(result["documents"]) == 0
|
|
240
|
+
|
|
241
|
+
def test_embed_metadata(self):
|
|
242
|
+
"""
|
|
243
|
+
Tests the embedding of metadata fields in document content for ranking.
|
|
244
|
+
"""
|
|
245
|
+
ranker = FastembedRanker(
|
|
246
|
+
model_name="model_name",
|
|
247
|
+
meta_fields_to_embed=["meta_field"],
|
|
248
|
+
)
|
|
249
|
+
ranker._model = MagicMock()
|
|
250
|
+
|
|
251
|
+
documents = [Document(content=f"document-number {i}", meta={"meta_field": f"meta_value {i}"}) for i in range(5)]
|
|
252
|
+
query = "test"
|
|
253
|
+
ranker.run(query=query, documents=documents)
|
|
254
|
+
|
|
255
|
+
ranker._model.rerank.assert_called_once_with(
|
|
256
|
+
query=query,
|
|
257
|
+
documents=[
|
|
258
|
+
"meta_value 0\ndocument-number 0",
|
|
259
|
+
"meta_value 1\ndocument-number 1",
|
|
260
|
+
"meta_value 2\ndocument-number 2",
|
|
261
|
+
"meta_value 3\ndocument-number 3",
|
|
262
|
+
"meta_value 4\ndocument-number 4",
|
|
263
|
+
],
|
|
264
|
+
batch_size=64,
|
|
265
|
+
parallel=None,
|
|
266
|
+
)
|
|
267
|
+
|
|
268
|
+
@pytest.mark.integration
|
|
269
|
+
def test_run(self):
|
|
270
|
+
ranker = FastembedRanker(model_name="Xenova/ms-marco-MiniLM-L-6-v2", top_k=2)
|
|
271
|
+
ranker.warm_up()
|
|
272
|
+
|
|
273
|
+
query = "Who is maintaining Qdrant?"
|
|
274
|
+
documents = [
|
|
275
|
+
Document(
|
|
276
|
+
content="This is built to be faster and lighter than other embedding \
|
|
277
|
+
libraries e.g. Transformers, Sentence-Transformers, etc."
|
|
278
|
+
),
|
|
279
|
+
Document(content="This is some random input"),
|
|
280
|
+
Document(content="fastembed is supported by and maintained by Qdrant."),
|
|
281
|
+
]
|
|
282
|
+
|
|
283
|
+
result = ranker.run(query=query, documents=documents)
|
|
284
|
+
|
|
285
|
+
assert len(result["documents"]) == 2
|
|
286
|
+
first_document = result["documents"][0]
|
|
287
|
+
second_document = result["documents"][1]
|
|
288
|
+
|
|
289
|
+
assert isinstance(first_document, Document)
|
|
290
|
+
assert isinstance(second_document, Document)
|
|
291
|
+
assert first_document.content == "fastembed is supported by and maintained by Qdrant."
|
|
292
|
+
assert first_document.score > second_document.score
|
|
@@ -15,8 +15,8 @@ class TestFastembedSparseDocumentEmbedderDoc:
|
|
|
15
15
|
"""
|
|
16
16
|
Test default initialization parameters for FastembedSparseDocumentEmbedder.
|
|
17
17
|
"""
|
|
18
|
-
embedder = FastembedSparseDocumentEmbedder(model="
|
|
19
|
-
assert embedder.model_name == "
|
|
18
|
+
embedder = FastembedSparseDocumentEmbedder(model="prithivida/Splade_PP_en_v1")
|
|
19
|
+
assert embedder.model_name == "prithivida/Splade_PP_en_v1"
|
|
20
20
|
assert embedder.cache_dir is None
|
|
21
21
|
assert embedder.threads is None
|
|
22
22
|
assert embedder.batch_size == 32
|
|
@@ -31,7 +31,7 @@ class TestFastembedSparseDocumentEmbedderDoc:
|
|
|
31
31
|
Test custom initialization parameters for FastembedSparseDocumentEmbedder.
|
|
32
32
|
"""
|
|
33
33
|
embedder = FastembedSparseDocumentEmbedder(
|
|
34
|
-
model="
|
|
34
|
+
model="prithivida/Splade_PP_en_v1",
|
|
35
35
|
cache_dir="fake_dir",
|
|
36
36
|
threads=2,
|
|
37
37
|
batch_size=64,
|
|
@@ -41,7 +41,7 @@ class TestFastembedSparseDocumentEmbedderDoc:
|
|
|
41
41
|
meta_fields_to_embed=["test_field"],
|
|
42
42
|
embedding_separator=" | ",
|
|
43
43
|
)
|
|
44
|
-
assert embedder.model_name == "
|
|
44
|
+
assert embedder.model_name == "prithivida/Splade_PP_en_v1"
|
|
45
45
|
assert embedder.cache_dir == "fake_dir"
|
|
46
46
|
assert embedder.threads == 2
|
|
47
47
|
assert embedder.batch_size == 64
|
|
@@ -55,12 +55,12 @@ class TestFastembedSparseDocumentEmbedderDoc:
|
|
|
55
55
|
"""
|
|
56
56
|
Test serialization of FastembedSparseDocumentEmbedder to a dictionary, using default initialization parameters.
|
|
57
57
|
"""
|
|
58
|
-
embedder = FastembedSparseDocumentEmbedder(model="
|
|
58
|
+
embedder = FastembedSparseDocumentEmbedder(model="prithivida/Splade_PP_en_v1")
|
|
59
59
|
embedder_dict = embedder.to_dict()
|
|
60
60
|
assert embedder_dict == {
|
|
61
61
|
"type": "haystack_integrations.components.embedders.fastembed.fastembed_sparse_document_embedder.FastembedSparseDocumentEmbedder", # noqa
|
|
62
62
|
"init_parameters": {
|
|
63
|
-
"model": "
|
|
63
|
+
"model": "prithivida/Splade_PP_en_v1",
|
|
64
64
|
"cache_dir": None,
|
|
65
65
|
"threads": None,
|
|
66
66
|
"batch_size": 32,
|
|
@@ -78,7 +78,7 @@ class TestFastembedSparseDocumentEmbedderDoc:
|
|
|
78
78
|
Test serialization of FastembedSparseDocumentEmbedder to a dictionary, using custom initialization parameters.
|
|
79
79
|
"""
|
|
80
80
|
embedder = FastembedSparseDocumentEmbedder(
|
|
81
|
-
model="
|
|
81
|
+
model="prithivida/Splade_PP_en_v1",
|
|
82
82
|
cache_dir="fake_dir",
|
|
83
83
|
threads=2,
|
|
84
84
|
batch_size=64,
|
|
@@ -92,7 +92,7 @@ class TestFastembedSparseDocumentEmbedderDoc:
|
|
|
92
92
|
assert embedder_dict == {
|
|
93
93
|
"type": "haystack_integrations.components.embedders.fastembed.fastembed_sparse_document_embedder.FastembedSparseDocumentEmbedder", # noqa
|
|
94
94
|
"init_parameters": {
|
|
95
|
-
"model": "
|
|
95
|
+
"model": "prithivida/Splade_PP_en_v1",
|
|
96
96
|
"cache_dir": "fake_dir",
|
|
97
97
|
"threads": 2,
|
|
98
98
|
"batch_size": 64,
|
|
@@ -113,7 +113,7 @@ class TestFastembedSparseDocumentEmbedderDoc:
|
|
|
113
113
|
embedder_dict = {
|
|
114
114
|
"type": "haystack_integrations.components.embedders.fastembed.fastembed_sparse_document_embedder.FastembedSparseDocumentEmbedder", # noqa
|
|
115
115
|
"init_parameters": {
|
|
116
|
-
"model": "
|
|
116
|
+
"model": "prithivida/Splade_PP_en_v1",
|
|
117
117
|
"cache_dir": None,
|
|
118
118
|
"threads": None,
|
|
119
119
|
"batch_size": 32,
|
|
@@ -125,7 +125,7 @@ class TestFastembedSparseDocumentEmbedderDoc:
|
|
|
125
125
|
},
|
|
126
126
|
}
|
|
127
127
|
embedder = default_from_dict(FastembedSparseDocumentEmbedder, embedder_dict)
|
|
128
|
-
assert embedder.model_name == "
|
|
128
|
+
assert embedder.model_name == "prithivida/Splade_PP_en_v1"
|
|
129
129
|
assert embedder.cache_dir is None
|
|
130
130
|
assert embedder.threads is None
|
|
131
131
|
assert embedder.batch_size == 32
|
|
@@ -143,7 +143,7 @@ class TestFastembedSparseDocumentEmbedderDoc:
|
|
|
143
143
|
embedder_dict = {
|
|
144
144
|
"type": "haystack_integrations.components.embedders.fastembed.fastembed_sparse_document_embedder.FastembedSparseDocumentEmbedder", # noqa
|
|
145
145
|
"init_parameters": {
|
|
146
|
-
"model": "
|
|
146
|
+
"model": "prithivida/Splade_PP_en_v1",
|
|
147
147
|
"cache_dir": "fake_dir",
|
|
148
148
|
"threads": 2,
|
|
149
149
|
"batch_size": 64,
|
|
@@ -155,7 +155,7 @@ class TestFastembedSparseDocumentEmbedderDoc:
|
|
|
155
155
|
},
|
|
156
156
|
}
|
|
157
157
|
embedder = default_from_dict(FastembedSparseDocumentEmbedder, embedder_dict)
|
|
158
|
-
assert embedder.model_name == "
|
|
158
|
+
assert embedder.model_name == "prithivida/Splade_PP_en_v1"
|
|
159
159
|
assert embedder.cache_dir == "fake_dir"
|
|
160
160
|
assert embedder.threads == 2
|
|
161
161
|
assert embedder.batch_size == 64
|
|
@@ -172,11 +172,11 @@ class TestFastembedSparseDocumentEmbedderDoc:
|
|
|
172
172
|
"""
|
|
173
173
|
Test for checking embedder instances after warm-up.
|
|
174
174
|
"""
|
|
175
|
-
embedder = FastembedSparseDocumentEmbedder(model="
|
|
175
|
+
embedder = FastembedSparseDocumentEmbedder(model="prithivida/Splade_PP_en_v1")
|
|
176
176
|
mocked_factory.get_embedding_backend.assert_not_called()
|
|
177
177
|
embedder.warm_up()
|
|
178
178
|
mocked_factory.get_embedding_backend.assert_called_once_with(
|
|
179
|
-
model_name="
|
|
179
|
+
model_name="prithivida/Splade_PP_en_v1",
|
|
180
180
|
cache_dir=None,
|
|
181
181
|
threads=None,
|
|
182
182
|
local_files_only=False,
|
|
@@ -190,7 +190,7 @@ class TestFastembedSparseDocumentEmbedderDoc:
|
|
|
190
190
|
"""
|
|
191
191
|
Test for checking backend instances after multiple warm-ups.
|
|
192
192
|
"""
|
|
193
|
-
embedder = FastembedSparseDocumentEmbedder(model="
|
|
193
|
+
embedder = FastembedSparseDocumentEmbedder(model="prithivida/Splade_PP_en_v1")
|
|
194
194
|
mocked_factory.get_embedding_backend.assert_not_called()
|
|
195
195
|
embedder.warm_up()
|
|
196
196
|
embedder.warm_up()
|
|
@@ -211,7 +211,7 @@ class TestFastembedSparseDocumentEmbedderDoc:
|
|
|
211
211
|
"""
|
|
212
212
|
Test for checking output dimensions and embedding dimensions.
|
|
213
213
|
"""
|
|
214
|
-
embedder = FastembedSparseDocumentEmbedder(model="
|
|
214
|
+
embedder = FastembedSparseDocumentEmbedder(model="prithivida/Splade_PP_en_v1")
|
|
215
215
|
embedder.embedding_backend = MagicMock()
|
|
216
216
|
embedder.embedding_backend.embed = lambda x, **kwargs: self._generate_mocked_sparse_embedding( # noqa: ARG005
|
|
217
217
|
len(x)
|
|
@@ -235,7 +235,7 @@ class TestFastembedSparseDocumentEmbedderDoc:
|
|
|
235
235
|
"""
|
|
236
236
|
Test for checking incorrect input format when creating embedding.
|
|
237
237
|
"""
|
|
238
|
-
embedder = FastembedSparseDocumentEmbedder(model="
|
|
238
|
+
embedder = FastembedSparseDocumentEmbedder(model="prithivida/Splade_PP_en_v1")
|
|
239
239
|
|
|
240
240
|
string_input = "text"
|
|
241
241
|
list_integers_input = [1, 2, 3]
|
|
@@ -330,7 +330,7 @@ class TestFastembedSparseDocumentEmbedderDoc:
|
|
|
330
330
|
@pytest.mark.integration
|
|
331
331
|
def test_run(self):
|
|
332
332
|
embedder = FastembedSparseDocumentEmbedder(
|
|
333
|
-
model="
|
|
333
|
+
model="prithivida/Splade_PP_en_v1",
|
|
334
334
|
)
|
|
335
335
|
embedder.warm_up()
|
|
336
336
|
|
{fastembed_haystack-1.3.0 → fastembed_haystack-1.4.1}/tests/test_fastembed_sparse_text_embedder.py
RENAMED
|
@@ -15,8 +15,8 @@ class TestFastembedSparseTextEmbedder:
|
|
|
15
15
|
"""
|
|
16
16
|
Test default initialization parameters for FastembedSparseTextEmbedder.
|
|
17
17
|
"""
|
|
18
|
-
embedder = FastembedSparseTextEmbedder(model="
|
|
19
|
-
assert embedder.model_name == "
|
|
18
|
+
embedder = FastembedSparseTextEmbedder(model="prithivida/Splade_PP_en_v1")
|
|
19
|
+
assert embedder.model_name == "prithivida/Splade_PP_en_v1"
|
|
20
20
|
assert embedder.cache_dir is None
|
|
21
21
|
assert embedder.threads is None
|
|
22
22
|
assert embedder.progress_bar is True
|
|
@@ -27,13 +27,13 @@ class TestFastembedSparseTextEmbedder:
|
|
|
27
27
|
Test custom initialization parameters for FastembedSparseTextEmbedder.
|
|
28
28
|
"""
|
|
29
29
|
embedder = FastembedSparseTextEmbedder(
|
|
30
|
-
model="
|
|
30
|
+
model="prithivida/Splade_PP_en_v1",
|
|
31
31
|
cache_dir="fake_dir",
|
|
32
32
|
threads=2,
|
|
33
33
|
progress_bar=False,
|
|
34
34
|
parallel=1,
|
|
35
35
|
)
|
|
36
|
-
assert embedder.model_name == "
|
|
36
|
+
assert embedder.model_name == "prithivida/Splade_PP_en_v1"
|
|
37
37
|
assert embedder.cache_dir == "fake_dir"
|
|
38
38
|
assert embedder.threads == 2
|
|
39
39
|
assert embedder.progress_bar is False
|
|
@@ -43,12 +43,12 @@ class TestFastembedSparseTextEmbedder:
|
|
|
43
43
|
"""
|
|
44
44
|
Test serialization of FastembedSparseTextEmbedder to a dictionary, using default initialization parameters.
|
|
45
45
|
"""
|
|
46
|
-
embedder = FastembedSparseTextEmbedder(model="
|
|
46
|
+
embedder = FastembedSparseTextEmbedder(model="prithivida/Splade_PP_en_v1")
|
|
47
47
|
embedder_dict = embedder.to_dict()
|
|
48
48
|
assert embedder_dict == {
|
|
49
49
|
"type": "haystack_integrations.components.embedders.fastembed.fastembed_sparse_text_embedder.FastembedSparseTextEmbedder", # noqa
|
|
50
50
|
"init_parameters": {
|
|
51
|
-
"model": "
|
|
51
|
+
"model": "prithivida/Splade_PP_en_v1",
|
|
52
52
|
"cache_dir": None,
|
|
53
53
|
"threads": None,
|
|
54
54
|
"progress_bar": True,
|
|
@@ -63,7 +63,7 @@ class TestFastembedSparseTextEmbedder:
|
|
|
63
63
|
Test serialization of FastembedSparseTextEmbedder to a dictionary, using custom initialization parameters.
|
|
64
64
|
"""
|
|
65
65
|
embedder = FastembedSparseTextEmbedder(
|
|
66
|
-
model="
|
|
66
|
+
model="prithivida/Splade_PP_en_v1",
|
|
67
67
|
cache_dir="fake_dir",
|
|
68
68
|
threads=2,
|
|
69
69
|
progress_bar=False,
|
|
@@ -74,7 +74,7 @@ class TestFastembedSparseTextEmbedder:
|
|
|
74
74
|
assert embedder_dict == {
|
|
75
75
|
"type": "haystack_integrations.components.embedders.fastembed.fastembed_sparse_text_embedder.FastembedSparseTextEmbedder", # noqa
|
|
76
76
|
"init_parameters": {
|
|
77
|
-
"model": "
|
|
77
|
+
"model": "prithivida/Splade_PP_en_v1",
|
|
78
78
|
"cache_dir": "fake_dir",
|
|
79
79
|
"threads": 2,
|
|
80
80
|
"progress_bar": False,
|
|
@@ -91,7 +91,7 @@ class TestFastembedSparseTextEmbedder:
|
|
|
91
91
|
embedder_dict = {
|
|
92
92
|
"type": "haystack_integrations.components.embedders.fastembed.fastembed_sparse_text_embedder.FastembedSparseTextEmbedder", # noqa
|
|
93
93
|
"init_parameters": {
|
|
94
|
-
"model": "
|
|
94
|
+
"model": "prithivida/Splade_PP_en_v1",
|
|
95
95
|
"cache_dir": None,
|
|
96
96
|
"threads": None,
|
|
97
97
|
"progress_bar": True,
|
|
@@ -99,7 +99,7 @@ class TestFastembedSparseTextEmbedder:
|
|
|
99
99
|
},
|
|
100
100
|
}
|
|
101
101
|
embedder = default_from_dict(FastembedSparseTextEmbedder, embedder_dict)
|
|
102
|
-
assert embedder.model_name == "
|
|
102
|
+
assert embedder.model_name == "prithivida/Splade_PP_en_v1"
|
|
103
103
|
assert embedder.cache_dir is None
|
|
104
104
|
assert embedder.threads is None
|
|
105
105
|
assert embedder.progress_bar is True
|
|
@@ -112,7 +112,7 @@ class TestFastembedSparseTextEmbedder:
|
|
|
112
112
|
embedder_dict = {
|
|
113
113
|
"type": "haystack_integrations.components.embedders.fastembed.fastembed_sparse_text_embedder.FastembedSparseTextEmbedder", # noqa
|
|
114
114
|
"init_parameters": {
|
|
115
|
-
"model": "
|
|
115
|
+
"model": "prithivida/Splade_PP_en_v1",
|
|
116
116
|
"cache_dir": "fake_dir",
|
|
117
117
|
"threads": 2,
|
|
118
118
|
"progress_bar": False,
|
|
@@ -120,7 +120,7 @@ class TestFastembedSparseTextEmbedder:
|
|
|
120
120
|
},
|
|
121
121
|
}
|
|
122
122
|
embedder = default_from_dict(FastembedSparseTextEmbedder, embedder_dict)
|
|
123
|
-
assert embedder.model_name == "
|
|
123
|
+
assert embedder.model_name == "prithivida/Splade_PP_en_v1"
|
|
124
124
|
assert embedder.cache_dir == "fake_dir"
|
|
125
125
|
assert embedder.threads == 2
|
|
126
126
|
assert embedder.progress_bar is False
|
|
@@ -133,11 +133,11 @@ class TestFastembedSparseTextEmbedder:
|
|
|
133
133
|
"""
|
|
134
134
|
Test for checking embedder instances after warm-up.
|
|
135
135
|
"""
|
|
136
|
-
embedder = FastembedSparseTextEmbedder(model="
|
|
136
|
+
embedder = FastembedSparseTextEmbedder(model="prithivida/Splade_PP_en_v1")
|
|
137
137
|
mocked_factory.get_embedding_backend.assert_not_called()
|
|
138
138
|
embedder.warm_up()
|
|
139
139
|
mocked_factory.get_embedding_backend.assert_called_once_with(
|
|
140
|
-
model_name="
|
|
140
|
+
model_name="prithivida/Splade_PP_en_v1",
|
|
141
141
|
cache_dir=None,
|
|
142
142
|
threads=None,
|
|
143
143
|
local_files_only=False,
|
|
@@ -151,7 +151,7 @@ class TestFastembedSparseTextEmbedder:
|
|
|
151
151
|
"""
|
|
152
152
|
Test for checking backend instances after multiple warm-ups.
|
|
153
153
|
"""
|
|
154
|
-
embedder = FastembedSparseTextEmbedder(model="
|
|
154
|
+
embedder = FastembedSparseTextEmbedder(model="prithivida/Splade_PP_en_v1")
|
|
155
155
|
mocked_factory.get_embedding_backend.assert_not_called()
|
|
156
156
|
embedder.warm_up()
|
|
157
157
|
embedder.warm_up()
|
|
@@ -252,7 +252,7 @@ class TestFastembedSparseTextEmbedder:
|
|
|
252
252
|
@pytest.mark.integration
|
|
253
253
|
def test_run(self):
|
|
254
254
|
embedder = FastembedSparseTextEmbedder(
|
|
255
|
-
model="
|
|
255
|
+
model="prithivida/Splade_PP_en_v1",
|
|
256
256
|
)
|
|
257
257
|
embedder.warm_up()
|
|
258
258
|
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
{fastembed_haystack-1.3.0 → fastembed_haystack-1.4.1}/tests/test_fastembed_document_embedder.py
RENAMED
|
File without changes
|
|
File without changes
|