fastembed-haystack 1.0.0__tar.gz → 1.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-1.0.0 → fastembed_haystack-1.2.0}/.gitignore +9 -0
- {fastembed_haystack-1.0.0 → fastembed_haystack-1.2.0}/PKG-INFO +1 -1
- {fastembed_haystack-1.0.0 → fastembed_haystack-1.2.0}/pydoc/config.yml +1 -1
- {fastembed_haystack-1.0.0 → fastembed_haystack-1.2.0}/pyproject.toml +25 -49
- {fastembed_haystack-1.0.0 → fastembed_haystack-1.2.0}/src/haystack_integrations/components/embedders/fastembed/embedding_backend/fastembed_backend.py +14 -4
- {fastembed_haystack-1.0.0 → fastembed_haystack-1.2.0}/src/haystack_integrations/components/embedders/fastembed/fastembed_document_embedder.py +9 -2
- {fastembed_haystack-1.0.0 → fastembed_haystack-1.2.0}/src/haystack_integrations/components/embedders/fastembed/fastembed_sparse_document_embedder.py +9 -2
- {fastembed_haystack-1.0.0 → fastembed_haystack-1.2.0}/src/haystack_integrations/components/embedders/fastembed/fastembed_sparse_text_embedder.py +10 -3
- {fastembed_haystack-1.0.0 → fastembed_haystack-1.2.0}/src/haystack_integrations/components/embedders/fastembed/fastembed_text_embedder.py +9 -2
- {fastembed_haystack-1.0.0 → fastembed_haystack-1.2.0}/tests/test_fastembed_backend.py +3 -1
- {fastembed_haystack-1.0.0 → fastembed_haystack-1.2.0}/tests/test_fastembed_document_embedder.py +11 -1
- {fastembed_haystack-1.0.0 → fastembed_haystack-1.2.0}/tests/test_fastembed_sparse_document_embedder.py +12 -2
- {fastembed_haystack-1.0.0 → fastembed_haystack-1.2.0}/tests/test_fastembed_sparse_text_embedder.py +4 -1
- {fastembed_haystack-1.0.0 → fastembed_haystack-1.2.0}/tests/test_fastembed_text_embedder.py +4 -1
- {fastembed_haystack-1.0.0 → fastembed_haystack-1.2.0}/LICENSE.txt +0 -0
- {fastembed_haystack-1.0.0 → fastembed_haystack-1.2.0}/README.md +0 -0
- {fastembed_haystack-1.0.0 → fastembed_haystack-1.2.0}/examples/example.py +0 -0
- {fastembed_haystack-1.0.0 → fastembed_haystack-1.2.0}/examples/sparse_example.py +0 -0
- {fastembed_haystack-1.0.0 → fastembed_haystack-1.2.0}/src/haystack_integrations/components/embedders/fastembed/__init__.py +0 -0
- {fastembed_haystack-1.0.0 → fastembed_haystack-1.2.0}/src/haystack_integrations/components/embedders/fastembed/embedding_backend/__init__.py +0 -0
- {fastembed_haystack-1.0.0 → fastembed_haystack-1.2.0}/tests/__init__.py +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.3
|
|
2
2
|
Name: fastembed-haystack
|
|
3
|
-
Version: 1.
|
|
3
|
+
Version: 1.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
|
|
@@ -18,7 +18,7 @@ processors:
|
|
|
18
18
|
- type: smart
|
|
19
19
|
- type: crossref
|
|
20
20
|
renderer:
|
|
21
|
-
type: haystack_pydoc_tools.renderers.
|
|
21
|
+
type: haystack_pydoc_tools.renderers.ReadmeIntegrationRenderer
|
|
22
22
|
excerpt: FastEmbed integration for Haystack
|
|
23
23
|
category_slug: integrations-api
|
|
24
24
|
title: FastEmbed
|
|
@@ -10,9 +10,7 @@ readme = "README.md"
|
|
|
10
10
|
requires-python = ">=3.8"
|
|
11
11
|
license = "Apache-2.0"
|
|
12
12
|
keywords = []
|
|
13
|
-
authors = [
|
|
14
|
-
{ name = "deepset GmbH", email = "info@deepset.ai" },
|
|
15
|
-
]
|
|
13
|
+
authors = [{ name = "deepset GmbH", email = "info@deepset.ai" }]
|
|
16
14
|
classifiers = [
|
|
17
15
|
"License :: OSI Approved :: Apache Software License",
|
|
18
16
|
"Development Status :: 4 - Beta",
|
|
@@ -25,10 +23,7 @@ classifiers = [
|
|
|
25
23
|
"Programming Language :: Python :: Implementation :: CPython",
|
|
26
24
|
"Programming Language :: Python :: Implementation :: PyPy",
|
|
27
25
|
]
|
|
28
|
-
dependencies = [
|
|
29
|
-
"haystack-ai>=2.0.1",
|
|
30
|
-
"fastembed>=0.2.5",
|
|
31
|
-
]
|
|
26
|
+
dependencies = ["haystack-ai>=2.0.1", "fastembed>=0.2.5"]
|
|
32
27
|
|
|
33
28
|
[project.urls]
|
|
34
29
|
Source = "https://github.com/deepset-ai/haystack-core-integrations"
|
|
@@ -50,50 +45,28 @@ git_describe_command = 'git describe --tags --match="integrations/fastembed-v[0-
|
|
|
50
45
|
dependencies = [
|
|
51
46
|
"coverage[toml]>=6.5",
|
|
52
47
|
"pytest",
|
|
48
|
+
"pytest-rerunfailures",
|
|
53
49
|
"ipython",
|
|
54
50
|
"haystack-pydoc-tools",
|
|
55
51
|
]
|
|
56
52
|
[tool.hatch.envs.default.scripts]
|
|
57
|
-
test = "pytest {args:tests}"
|
|
58
|
-
test-cov = "coverage run -m pytest {args:tests}"
|
|
59
|
-
cov-report = [
|
|
60
|
-
|
|
61
|
-
|
|
62
|
-
]
|
|
63
|
-
cov = [
|
|
64
|
-
"test-cov",
|
|
65
|
-
"cov-report",
|
|
66
|
-
]
|
|
67
|
-
docs = [
|
|
68
|
-
"pydoc-markdown pydoc/config.yml"
|
|
69
|
-
]
|
|
53
|
+
test = "pytest --reruns 3 --reruns-delay 30 -x {args:tests}"
|
|
54
|
+
test-cov = "coverage run -m pytest --reruns 3 --reruns-delay 30 -x {args:tests}"
|
|
55
|
+
cov-report = ["- coverage combine", "coverage report"]
|
|
56
|
+
cov = ["test-cov", "cov-report"]
|
|
57
|
+
docs = ["pydoc-markdown pydoc/config.yml"]
|
|
70
58
|
|
|
71
59
|
[[tool.hatch.envs.all.matrix]]
|
|
72
60
|
python = ["3.8", "3.9", "3.10", "3.11", "3.12"]
|
|
73
61
|
|
|
74
62
|
[tool.hatch.envs.lint]
|
|
75
63
|
detached = true
|
|
76
|
-
dependencies = [
|
|
77
|
-
"black>=23.1.0",
|
|
78
|
-
"mypy>=1.0.0",
|
|
79
|
-
"ruff>=0.0.243",
|
|
80
|
-
"numpy"
|
|
81
|
-
]
|
|
64
|
+
dependencies = ["black>=23.1.0", "mypy>=1.0.0", "ruff>=0.0.243", "numpy"]
|
|
82
65
|
[tool.hatch.envs.lint.scripts]
|
|
83
66
|
typing = "mypy --install-types --non-interactive --explicit-package-bases {args:src/ tests}"
|
|
84
|
-
style = [
|
|
85
|
-
|
|
86
|
-
|
|
87
|
-
]
|
|
88
|
-
fmt = [
|
|
89
|
-
"black {args:.}",
|
|
90
|
-
"ruff --fix {args:.}",
|
|
91
|
-
"style",
|
|
92
|
-
]
|
|
93
|
-
all = [
|
|
94
|
-
"style",
|
|
95
|
-
"typing",
|
|
96
|
-
]
|
|
67
|
+
style = ["ruff check {args:.}", "black --check --diff {args:.}"]
|
|
68
|
+
fmt = ["black {args:.}", "ruff --fix {args:.}", "style"]
|
|
69
|
+
all = ["style", "typing"]
|
|
97
70
|
|
|
98
71
|
[tool.black]
|
|
99
72
|
target-version = ["py38"]
|
|
@@ -135,11 +108,18 @@ ignore = [
|
|
|
135
108
|
"B027",
|
|
136
109
|
# Allow boolean positional values in function calls, like `dict.get(... True)`
|
|
137
110
|
"FBT003",
|
|
138
|
-
"FBT001",
|
|
111
|
+
"FBT001",
|
|
112
|
+
"FBT002",
|
|
139
113
|
# Ignore checks for possible passwords
|
|
140
|
-
"S105",
|
|
114
|
+
"S105",
|
|
115
|
+
"S106",
|
|
116
|
+
"S107",
|
|
141
117
|
# Ignore complexity
|
|
142
|
-
"C901",
|
|
118
|
+
"C901",
|
|
119
|
+
"PLR0911",
|
|
120
|
+
"PLR0912",
|
|
121
|
+
"PLR0913",
|
|
122
|
+
"PLR0915",
|
|
143
123
|
]
|
|
144
124
|
unfixable = [
|
|
145
125
|
# Don't touch unused imports
|
|
@@ -165,12 +145,8 @@ parallel = false
|
|
|
165
145
|
|
|
166
146
|
[tool.coverage.report]
|
|
167
147
|
omit = ["*/tests/*", "*/__init__.py"]
|
|
168
|
-
show_missing=true
|
|
169
|
-
exclude_lines = [
|
|
170
|
-
"no cov",
|
|
171
|
-
"if __name__ == .__main__.:",
|
|
172
|
-
"if TYPE_CHECKING:",
|
|
173
|
-
]
|
|
148
|
+
show_missing = true
|
|
149
|
+
exclude_lines = ["no cov", "if __name__ == .__main__.:", "if TYPE_CHECKING:"]
|
|
174
150
|
|
|
175
151
|
[[tool.mypy.overrides]]
|
|
176
152
|
module = [
|
|
@@ -178,6 +154,6 @@ module = [
|
|
|
178
154
|
"haystack_integrations.*",
|
|
179
155
|
"fastembed.*",
|
|
180
156
|
"pytest.*",
|
|
181
|
-
"numpy.*"
|
|
157
|
+
"numpy.*",
|
|
182
158
|
]
|
|
183
159
|
ignore_missing_imports = true
|
|
@@ -19,13 +19,16 @@ class _FastembedEmbeddingBackendFactory:
|
|
|
19
19
|
model_name: str,
|
|
20
20
|
cache_dir: Optional[str] = None,
|
|
21
21
|
threads: Optional[int] = None,
|
|
22
|
+
local_files_only: bool = False,
|
|
22
23
|
):
|
|
23
24
|
embedding_backend_id = f"{model_name}{cache_dir}{threads}"
|
|
24
25
|
|
|
25
26
|
if embedding_backend_id in _FastembedEmbeddingBackendFactory._instances:
|
|
26
27
|
return _FastembedEmbeddingBackendFactory._instances[embedding_backend_id]
|
|
27
28
|
|
|
28
|
-
embedding_backend = _FastembedEmbeddingBackend(
|
|
29
|
+
embedding_backend = _FastembedEmbeddingBackend(
|
|
30
|
+
model_name=model_name, cache_dir=cache_dir, threads=threads, local_files_only=local_files_only
|
|
31
|
+
)
|
|
29
32
|
_FastembedEmbeddingBackendFactory._instances[embedding_backend_id] = embedding_backend
|
|
30
33
|
return embedding_backend
|
|
31
34
|
|
|
@@ -40,8 +43,11 @@ class _FastembedEmbeddingBackend:
|
|
|
40
43
|
model_name: str,
|
|
41
44
|
cache_dir: Optional[str] = None,
|
|
42
45
|
threads: Optional[int] = None,
|
|
46
|
+
local_files_only: bool = False,
|
|
43
47
|
):
|
|
44
|
-
self.model = TextEmbedding(
|
|
48
|
+
self.model = TextEmbedding(
|
|
49
|
+
model_name=model_name, cache_dir=cache_dir, threads=threads, local_files_only=local_files_only
|
|
50
|
+
)
|
|
45
51
|
|
|
46
52
|
def embed(self, data: List[str], progress_bar=True, **kwargs) -> List[List[float]]:
|
|
47
53
|
# the embed method returns a Iterable[np.ndarray], so we convert it to a list of lists
|
|
@@ -66,6 +72,7 @@ class _FastembedSparseEmbeddingBackendFactory:
|
|
|
66
72
|
model_name: str,
|
|
67
73
|
cache_dir: Optional[str] = None,
|
|
68
74
|
threads: Optional[int] = None,
|
|
75
|
+
local_files_only: bool = False,
|
|
69
76
|
):
|
|
70
77
|
embedding_backend_id = f"{model_name}{cache_dir}{threads}"
|
|
71
78
|
|
|
@@ -73,7 +80,7 @@ class _FastembedSparseEmbeddingBackendFactory:
|
|
|
73
80
|
return _FastembedSparseEmbeddingBackendFactory._instances[embedding_backend_id]
|
|
74
81
|
|
|
75
82
|
embedding_backend = _FastembedSparseEmbeddingBackend(
|
|
76
|
-
model_name=model_name, cache_dir=cache_dir, threads=threads
|
|
83
|
+
model_name=model_name, cache_dir=cache_dir, threads=threads, local_files_only=local_files_only
|
|
77
84
|
)
|
|
78
85
|
_FastembedSparseEmbeddingBackendFactory._instances[embedding_backend_id] = embedding_backend
|
|
79
86
|
return embedding_backend
|
|
@@ -89,8 +96,11 @@ class _FastembedSparseEmbeddingBackend:
|
|
|
89
96
|
model_name: str,
|
|
90
97
|
cache_dir: Optional[str] = None,
|
|
91
98
|
threads: Optional[int] = None,
|
|
99
|
+
local_files_only: bool = False,
|
|
92
100
|
):
|
|
93
|
-
self.model = SparseTextEmbedding(
|
|
101
|
+
self.model = SparseTextEmbedding(
|
|
102
|
+
model_name=model_name, cache_dir=cache_dir, threads=threads, local_files_only=local_files_only
|
|
103
|
+
)
|
|
94
104
|
|
|
95
105
|
def embed(self, data: List[List[str]], progress_bar=True, **kwargs) -> List[SparseEmbedding]:
|
|
96
106
|
# The embed method returns a Iterable[SparseEmbedding], so we convert to Haystack SparseEmbedding type.
|
|
@@ -65,6 +65,7 @@ class FastembedDocumentEmbedder:
|
|
|
65
65
|
batch_size: int = 256,
|
|
66
66
|
progress_bar: bool = True,
|
|
67
67
|
parallel: Optional[int] = None,
|
|
68
|
+
local_files_only: bool = False,
|
|
68
69
|
meta_fields_to_embed: Optional[List[str]] = None,
|
|
69
70
|
embedding_separator: str = "\n",
|
|
70
71
|
):
|
|
@@ -80,11 +81,12 @@ class FastembedDocumentEmbedder:
|
|
|
80
81
|
:param prefix: A string to add to the beginning of each text.
|
|
81
82
|
:param suffix: A string to add to the end of each text.
|
|
82
83
|
:param batch_size: Number of strings to encode at once.
|
|
83
|
-
:param progress_bar: If
|
|
84
|
+
:param progress_bar: If `True`, displays progress bar during embedding.
|
|
84
85
|
:param parallel:
|
|
85
86
|
If > 1, data-parallel encoding will be used, recommended for offline encoding of large datasets.
|
|
86
87
|
If 0, use all available cores.
|
|
87
88
|
If None, don't use data-parallel processing, use default onnxruntime threading instead.
|
|
89
|
+
:param local_files_only: If `True`, only use the model files in the `cache_dir`.
|
|
88
90
|
:param meta_fields_to_embed: List of meta fields that should be embedded along with the Document content.
|
|
89
91
|
:param embedding_separator: Separator used to concatenate the meta fields to the Document content.
|
|
90
92
|
"""
|
|
@@ -97,6 +99,7 @@ class FastembedDocumentEmbedder:
|
|
|
97
99
|
self.batch_size = batch_size
|
|
98
100
|
self.progress_bar = progress_bar
|
|
99
101
|
self.parallel = parallel
|
|
102
|
+
self.local_files_only = local_files_only
|
|
100
103
|
self.meta_fields_to_embed = meta_fields_to_embed or []
|
|
101
104
|
self.embedding_separator = embedding_separator
|
|
102
105
|
|
|
@@ -116,6 +119,7 @@ class FastembedDocumentEmbedder:
|
|
|
116
119
|
batch_size=self.batch_size,
|
|
117
120
|
progress_bar=self.progress_bar,
|
|
118
121
|
parallel=self.parallel,
|
|
122
|
+
local_files_only=self.local_files_only,
|
|
119
123
|
meta_fields_to_embed=self.meta_fields_to_embed,
|
|
120
124
|
embedding_separator=self.embedding_separator,
|
|
121
125
|
)
|
|
@@ -126,7 +130,10 @@ class FastembedDocumentEmbedder:
|
|
|
126
130
|
"""
|
|
127
131
|
if not hasattr(self, "embedding_backend"):
|
|
128
132
|
self.embedding_backend = _FastembedEmbeddingBackendFactory.get_embedding_backend(
|
|
129
|
-
model_name=self.model_name,
|
|
133
|
+
model_name=self.model_name,
|
|
134
|
+
cache_dir=self.cache_dir,
|
|
135
|
+
threads=self.threads,
|
|
136
|
+
local_files_only=self.local_files_only,
|
|
130
137
|
)
|
|
131
138
|
|
|
132
139
|
def _prepare_texts_to_embed(self, documents: List[Document]) -> List[str]:
|
|
@@ -59,6 +59,7 @@ class FastembedSparseDocumentEmbedder:
|
|
|
59
59
|
batch_size: int = 32,
|
|
60
60
|
progress_bar: bool = True,
|
|
61
61
|
parallel: Optional[int] = None,
|
|
62
|
+
local_files_only: bool = False,
|
|
62
63
|
meta_fields_to_embed: Optional[List[str]] = None,
|
|
63
64
|
embedding_separator: str = "\n",
|
|
64
65
|
):
|
|
@@ -77,6 +78,7 @@ class FastembedSparseDocumentEmbedder:
|
|
|
77
78
|
If > 1, data-parallel encoding will be used, recommended for offline encoding of large datasets.
|
|
78
79
|
If 0, use all available cores.
|
|
79
80
|
If None, don't use data-parallel processing, use default onnxruntime threading instead.
|
|
81
|
+
:param local_files_only: If `True`, only use the model files in the `cache_dir`.
|
|
80
82
|
:param meta_fields_to_embed: List of meta fields that should be embedded along with the Document content.
|
|
81
83
|
:param embedding_separator: Separator used to concatenate the meta fields to the Document content.
|
|
82
84
|
"""
|
|
@@ -87,6 +89,7 @@ class FastembedSparseDocumentEmbedder:
|
|
|
87
89
|
self.batch_size = batch_size
|
|
88
90
|
self.progress_bar = progress_bar
|
|
89
91
|
self.parallel = parallel
|
|
92
|
+
self.local_files_only = local_files_only
|
|
90
93
|
self.meta_fields_to_embed = meta_fields_to_embed or []
|
|
91
94
|
self.embedding_separator = embedding_separator
|
|
92
95
|
|
|
@@ -104,6 +107,7 @@ class FastembedSparseDocumentEmbedder:
|
|
|
104
107
|
batch_size=self.batch_size,
|
|
105
108
|
progress_bar=self.progress_bar,
|
|
106
109
|
parallel=self.parallel,
|
|
110
|
+
local_files_only=self.local_files_only,
|
|
107
111
|
meta_fields_to_embed=self.meta_fields_to_embed,
|
|
108
112
|
embedding_separator=self.embedding_separator,
|
|
109
113
|
)
|
|
@@ -114,7 +118,10 @@ class FastembedSparseDocumentEmbedder:
|
|
|
114
118
|
"""
|
|
115
119
|
if not hasattr(self, "embedding_backend"):
|
|
116
120
|
self.embedding_backend = _FastembedSparseEmbeddingBackendFactory.get_embedding_backend(
|
|
117
|
-
model_name=self.model_name,
|
|
121
|
+
model_name=self.model_name,
|
|
122
|
+
cache_dir=self.cache_dir,
|
|
123
|
+
threads=self.threads,
|
|
124
|
+
local_files_only=self.local_files_only,
|
|
118
125
|
)
|
|
119
126
|
|
|
120
127
|
def _prepare_texts_to_embed(self, documents: List[Document]) -> List[str]:
|
|
@@ -152,7 +159,7 @@ class FastembedSparseDocumentEmbedder:
|
|
|
152
159
|
embeddings = self.embedding_backend.embed(
|
|
153
160
|
texts_to_embed,
|
|
154
161
|
batch_size=self.batch_size,
|
|
155
|
-
|
|
162
|
+
progress_bar=self.progress_bar,
|
|
156
163
|
parallel=self.parallel,
|
|
157
164
|
)
|
|
158
165
|
|
|
@@ -34,6 +34,7 @@ class FastembedSparseTextEmbedder:
|
|
|
34
34
|
threads: Optional[int] = None,
|
|
35
35
|
progress_bar: bool = True,
|
|
36
36
|
parallel: Optional[int] = None,
|
|
37
|
+
local_files_only: bool = False,
|
|
37
38
|
):
|
|
38
39
|
"""
|
|
39
40
|
Create a FastembedSparseTextEmbedder component.
|
|
@@ -43,11 +44,12 @@ class FastembedSparseTextEmbedder:
|
|
|
43
44
|
Can be set using the `FASTEMBED_CACHE_PATH` env variable.
|
|
44
45
|
Defaults to `fastembed_cache` in the system's temp directory.
|
|
45
46
|
:param threads: The number of threads single onnxruntime session can use. Defaults to None.
|
|
46
|
-
:param progress_bar: If
|
|
47
|
+
:param progress_bar: If `True`, displays progress bar during embedding.
|
|
47
48
|
:param parallel:
|
|
48
49
|
If > 1, data-parallel encoding will be used, recommended for offline encoding of large datasets.
|
|
49
50
|
If 0, use all available cores.
|
|
50
51
|
If None, don't use data-parallel processing, use default onnxruntime threading instead.
|
|
52
|
+
:param local_files_only: If `True`, only use the model files in the `cache_dir`.
|
|
51
53
|
"""
|
|
52
54
|
|
|
53
55
|
self.model_name = model
|
|
@@ -55,6 +57,7 @@ class FastembedSparseTextEmbedder:
|
|
|
55
57
|
self.threads = threads
|
|
56
58
|
self.progress_bar = progress_bar
|
|
57
59
|
self.parallel = parallel
|
|
60
|
+
self.local_files_only = local_files_only
|
|
58
61
|
|
|
59
62
|
def to_dict(self) -> Dict[str, Any]:
|
|
60
63
|
"""
|
|
@@ -70,6 +73,7 @@ class FastembedSparseTextEmbedder:
|
|
|
70
73
|
threads=self.threads,
|
|
71
74
|
progress_bar=self.progress_bar,
|
|
72
75
|
parallel=self.parallel,
|
|
76
|
+
local_files_only=self.local_files_only,
|
|
73
77
|
)
|
|
74
78
|
|
|
75
79
|
def warm_up(self):
|
|
@@ -78,7 +82,10 @@ class FastembedSparseTextEmbedder:
|
|
|
78
82
|
"""
|
|
79
83
|
if not hasattr(self, "embedding_backend"):
|
|
80
84
|
self.embedding_backend = _FastembedSparseEmbeddingBackendFactory.get_embedding_backend(
|
|
81
|
-
model_name=self.model_name,
|
|
85
|
+
model_name=self.model_name,
|
|
86
|
+
cache_dir=self.cache_dir,
|
|
87
|
+
threads=self.threads,
|
|
88
|
+
local_files_only=self.local_files_only,
|
|
82
89
|
)
|
|
83
90
|
|
|
84
91
|
@component.output_types(sparse_embedding=SparseEmbedding)
|
|
@@ -104,7 +111,7 @@ class FastembedSparseTextEmbedder:
|
|
|
104
111
|
|
|
105
112
|
embedding = self.embedding_backend.embed(
|
|
106
113
|
[text],
|
|
107
|
-
|
|
114
|
+
progress_bar=self.progress_bar,
|
|
108
115
|
parallel=self.parallel,
|
|
109
116
|
)[0]
|
|
110
117
|
return {"sparse_embedding": embedding}
|
|
@@ -35,6 +35,7 @@ class FastembedTextEmbedder:
|
|
|
35
35
|
suffix: str = "",
|
|
36
36
|
progress_bar: bool = True,
|
|
37
37
|
parallel: Optional[int] = None,
|
|
38
|
+
local_files_only: bool = False,
|
|
38
39
|
):
|
|
39
40
|
"""
|
|
40
41
|
Create a FastembedTextEmbedder component.
|
|
@@ -46,11 +47,12 @@ class FastembedTextEmbedder:
|
|
|
46
47
|
:param threads: The number of threads single onnxruntime session can use. Defaults to None.
|
|
47
48
|
:param prefix: A string to add to the beginning of each text.
|
|
48
49
|
:param suffix: A string to add to the end of each text.
|
|
49
|
-
:param progress_bar: If
|
|
50
|
+
:param progress_bar: If `True`, displays progress bar during embedding.
|
|
50
51
|
:param parallel:
|
|
51
52
|
If > 1, data-parallel encoding will be used, recommended for offline encoding of large datasets.
|
|
52
53
|
If 0, use all available cores.
|
|
53
54
|
If None, don't use data-parallel processing, use default onnxruntime threading instead.
|
|
55
|
+
:param local_files_only: If `True`, only use the model files in the `cache_dir`.
|
|
54
56
|
"""
|
|
55
57
|
|
|
56
58
|
self.model_name = model
|
|
@@ -60,6 +62,7 @@ class FastembedTextEmbedder:
|
|
|
60
62
|
self.suffix = suffix
|
|
61
63
|
self.progress_bar = progress_bar
|
|
62
64
|
self.parallel = parallel
|
|
65
|
+
self.local_files_only = local_files_only
|
|
63
66
|
|
|
64
67
|
def to_dict(self) -> Dict[str, Any]:
|
|
65
68
|
"""
|
|
@@ -77,6 +80,7 @@ class FastembedTextEmbedder:
|
|
|
77
80
|
suffix=self.suffix,
|
|
78
81
|
progress_bar=self.progress_bar,
|
|
79
82
|
parallel=self.parallel,
|
|
83
|
+
local_files_only=self.local_files_only,
|
|
80
84
|
)
|
|
81
85
|
|
|
82
86
|
def warm_up(self):
|
|
@@ -85,7 +89,10 @@ class FastembedTextEmbedder:
|
|
|
85
89
|
"""
|
|
86
90
|
if not hasattr(self, "embedding_backend"):
|
|
87
91
|
self.embedding_backend = _FastembedEmbeddingBackendFactory.get_embedding_backend(
|
|
88
|
-
model_name=self.model_name,
|
|
92
|
+
model_name=self.model_name,
|
|
93
|
+
cache_dir=self.cache_dir,
|
|
94
|
+
threads=self.threads,
|
|
95
|
+
local_files_only=self.local_files_only,
|
|
89
96
|
)
|
|
90
97
|
|
|
91
98
|
@component.output_types(embedding=List[float])
|
|
@@ -27,7 +27,9 @@ def test_model_initialization(mock_instructor):
|
|
|
27
27
|
_FastembedEmbeddingBackendFactory.get_embedding_backend(
|
|
28
28
|
model_name="BAAI/bge-small-en-v1.5",
|
|
29
29
|
)
|
|
30
|
-
mock_instructor.assert_called_once_with(
|
|
30
|
+
mock_instructor.assert_called_once_with(
|
|
31
|
+
model_name="BAAI/bge-small-en-v1.5", cache_dir=None, threads=None, local_files_only=False
|
|
32
|
+
)
|
|
31
33
|
# restore the factory state
|
|
32
34
|
_FastembedEmbeddingBackendFactory._instances = {}
|
|
33
35
|
|
{fastembed_haystack-1.0.0 → fastembed_haystack-1.2.0}/tests/test_fastembed_document_embedder.py
RENAMED
|
@@ -22,6 +22,7 @@ class TestFastembedDocumentEmbedder:
|
|
|
22
22
|
assert embedder.batch_size == 256
|
|
23
23
|
assert embedder.progress_bar is True
|
|
24
24
|
assert embedder.parallel is None
|
|
25
|
+
assert not embedder.local_files_only
|
|
25
26
|
assert embedder.meta_fields_to_embed == []
|
|
26
27
|
assert embedder.embedding_separator == "\n"
|
|
27
28
|
|
|
@@ -38,6 +39,7 @@ class TestFastembedDocumentEmbedder:
|
|
|
38
39
|
batch_size=64,
|
|
39
40
|
progress_bar=False,
|
|
40
41
|
parallel=1,
|
|
42
|
+
local_files_only=True,
|
|
41
43
|
meta_fields_to_embed=["test_field"],
|
|
42
44
|
embedding_separator=" | ",
|
|
43
45
|
)
|
|
@@ -49,6 +51,7 @@ class TestFastembedDocumentEmbedder:
|
|
|
49
51
|
assert embedder.batch_size == 64
|
|
50
52
|
assert embedder.progress_bar is False
|
|
51
53
|
assert embedder.parallel == 1
|
|
54
|
+
assert embedder.local_files_only
|
|
52
55
|
assert embedder.meta_fields_to_embed == ["test_field"]
|
|
53
56
|
assert embedder.embedding_separator == " | "
|
|
54
57
|
|
|
@@ -69,6 +72,7 @@ class TestFastembedDocumentEmbedder:
|
|
|
69
72
|
"batch_size": 256,
|
|
70
73
|
"progress_bar": True,
|
|
71
74
|
"parallel": None,
|
|
75
|
+
"local_files_only": False,
|
|
72
76
|
"embedding_separator": "\n",
|
|
73
77
|
"meta_fields_to_embed": [],
|
|
74
78
|
},
|
|
@@ -87,6 +91,7 @@ class TestFastembedDocumentEmbedder:
|
|
|
87
91
|
batch_size=64,
|
|
88
92
|
progress_bar=False,
|
|
89
93
|
parallel=1,
|
|
94
|
+
local_files_only=True,
|
|
90
95
|
meta_fields_to_embed=["test_field"],
|
|
91
96
|
embedding_separator=" | ",
|
|
92
97
|
)
|
|
@@ -102,6 +107,7 @@ class TestFastembedDocumentEmbedder:
|
|
|
102
107
|
"batch_size": 64,
|
|
103
108
|
"progress_bar": False,
|
|
104
109
|
"parallel": 1,
|
|
110
|
+
"local_files_only": True,
|
|
105
111
|
"meta_fields_to_embed": ["test_field"],
|
|
106
112
|
"embedding_separator": " | ",
|
|
107
113
|
},
|
|
@@ -122,6 +128,7 @@ class TestFastembedDocumentEmbedder:
|
|
|
122
128
|
"batch_size": 256,
|
|
123
129
|
"progress_bar": True,
|
|
124
130
|
"parallel": None,
|
|
131
|
+
"local_files_only": False,
|
|
125
132
|
"meta_fields_to_embed": [],
|
|
126
133
|
"embedding_separator": "\n",
|
|
127
134
|
},
|
|
@@ -135,6 +142,7 @@ class TestFastembedDocumentEmbedder:
|
|
|
135
142
|
assert embedder.batch_size == 256
|
|
136
143
|
assert embedder.progress_bar is True
|
|
137
144
|
assert embedder.parallel is None
|
|
145
|
+
assert not embedder.local_files_only
|
|
138
146
|
assert embedder.meta_fields_to_embed == []
|
|
139
147
|
assert embedder.embedding_separator == "\n"
|
|
140
148
|
|
|
@@ -153,6 +161,7 @@ class TestFastembedDocumentEmbedder:
|
|
|
153
161
|
"batch_size": 64,
|
|
154
162
|
"progress_bar": False,
|
|
155
163
|
"parallel": 1,
|
|
164
|
+
"local_files_only": True,
|
|
156
165
|
"meta_fields_to_embed": ["test_field"],
|
|
157
166
|
"embedding_separator": " | ",
|
|
158
167
|
},
|
|
@@ -166,6 +175,7 @@ class TestFastembedDocumentEmbedder:
|
|
|
166
175
|
assert embedder.batch_size == 64
|
|
167
176
|
assert embedder.progress_bar is False
|
|
168
177
|
assert embedder.parallel == 1
|
|
178
|
+
assert embedder.local_files_only
|
|
169
179
|
assert embedder.meta_fields_to_embed == ["test_field"]
|
|
170
180
|
assert embedder.embedding_separator == " | "
|
|
171
181
|
|
|
@@ -180,7 +190,7 @@ class TestFastembedDocumentEmbedder:
|
|
|
180
190
|
mocked_factory.get_embedding_backend.assert_not_called()
|
|
181
191
|
embedder.warm_up()
|
|
182
192
|
mocked_factory.get_embedding_backend.assert_called_once_with(
|
|
183
|
-
model_name="BAAI/bge-small-en-v1.5", cache_dir=None, threads=None
|
|
193
|
+
model_name="BAAI/bge-small-en-v1.5", cache_dir=None, threads=None, local_files_only=False
|
|
184
194
|
)
|
|
185
195
|
|
|
186
196
|
@patch(
|
|
@@ -21,6 +21,7 @@ class TestFastembedSparseDocumentEmbedderDoc:
|
|
|
21
21
|
assert embedder.batch_size == 32
|
|
22
22
|
assert embedder.progress_bar is True
|
|
23
23
|
assert embedder.parallel is None
|
|
24
|
+
assert not embedder.local_files_only
|
|
24
25
|
assert embedder.meta_fields_to_embed == []
|
|
25
26
|
assert embedder.embedding_separator == "\n"
|
|
26
27
|
|
|
@@ -35,6 +36,7 @@ class TestFastembedSparseDocumentEmbedderDoc:
|
|
|
35
36
|
batch_size=64,
|
|
36
37
|
progress_bar=False,
|
|
37
38
|
parallel=1,
|
|
39
|
+
local_files_only=True,
|
|
38
40
|
meta_fields_to_embed=["test_field"],
|
|
39
41
|
embedding_separator=" | ",
|
|
40
42
|
)
|
|
@@ -44,6 +46,7 @@ class TestFastembedSparseDocumentEmbedderDoc:
|
|
|
44
46
|
assert embedder.batch_size == 64
|
|
45
47
|
assert embedder.progress_bar is False
|
|
46
48
|
assert embedder.parallel == 1
|
|
49
|
+
assert embedder.local_files_only
|
|
47
50
|
assert embedder.meta_fields_to_embed == ["test_field"]
|
|
48
51
|
assert embedder.embedding_separator == " | "
|
|
49
52
|
|
|
@@ -62,6 +65,7 @@ class TestFastembedSparseDocumentEmbedderDoc:
|
|
|
62
65
|
"batch_size": 32,
|
|
63
66
|
"progress_bar": True,
|
|
64
67
|
"parallel": None,
|
|
68
|
+
"local_files_only": False,
|
|
65
69
|
"embedding_separator": "\n",
|
|
66
70
|
"meta_fields_to_embed": [],
|
|
67
71
|
},
|
|
@@ -78,6 +82,7 @@ class TestFastembedSparseDocumentEmbedderDoc:
|
|
|
78
82
|
batch_size=64,
|
|
79
83
|
progress_bar=False,
|
|
80
84
|
parallel=1,
|
|
85
|
+
local_files_only=True,
|
|
81
86
|
meta_fields_to_embed=["test_field"],
|
|
82
87
|
embedding_separator=" | ",
|
|
83
88
|
)
|
|
@@ -91,6 +96,7 @@ class TestFastembedSparseDocumentEmbedderDoc:
|
|
|
91
96
|
"batch_size": 64,
|
|
92
97
|
"progress_bar": False,
|
|
93
98
|
"parallel": 1,
|
|
99
|
+
"local_files_only": True,
|
|
94
100
|
"meta_fields_to_embed": ["test_field"],
|
|
95
101
|
"embedding_separator": " | ",
|
|
96
102
|
},
|
|
@@ -110,6 +116,7 @@ class TestFastembedSparseDocumentEmbedderDoc:
|
|
|
110
116
|
"batch_size": 32,
|
|
111
117
|
"progress_bar": True,
|
|
112
118
|
"parallel": None,
|
|
119
|
+
"local_files_only": False,
|
|
113
120
|
"meta_fields_to_embed": [],
|
|
114
121
|
"embedding_separator": "\n",
|
|
115
122
|
},
|
|
@@ -121,6 +128,7 @@ class TestFastembedSparseDocumentEmbedderDoc:
|
|
|
121
128
|
assert embedder.batch_size == 32
|
|
122
129
|
assert embedder.progress_bar is True
|
|
123
130
|
assert embedder.parallel is None
|
|
131
|
+
assert not embedder.local_files_only
|
|
124
132
|
assert embedder.meta_fields_to_embed == []
|
|
125
133
|
assert embedder.embedding_separator == "\n"
|
|
126
134
|
|
|
@@ -138,6 +146,7 @@ class TestFastembedSparseDocumentEmbedderDoc:
|
|
|
138
146
|
"batch_size": 64,
|
|
139
147
|
"progress_bar": False,
|
|
140
148
|
"parallel": 1,
|
|
149
|
+
"local_files_only": True,
|
|
141
150
|
"meta_fields_to_embed": ["test_field"],
|
|
142
151
|
"embedding_separator": " | ",
|
|
143
152
|
},
|
|
@@ -149,6 +158,7 @@ class TestFastembedSparseDocumentEmbedderDoc:
|
|
|
149
158
|
assert embedder.batch_size == 64
|
|
150
159
|
assert embedder.progress_bar is False
|
|
151
160
|
assert embedder.parallel == 1
|
|
161
|
+
assert embedder.local_files_only
|
|
152
162
|
assert embedder.meta_fields_to_embed == ["test_field"]
|
|
153
163
|
assert embedder.embedding_separator == " | "
|
|
154
164
|
|
|
@@ -163,7 +173,7 @@ class TestFastembedSparseDocumentEmbedderDoc:
|
|
|
163
173
|
mocked_factory.get_embedding_backend.assert_not_called()
|
|
164
174
|
embedder.warm_up()
|
|
165
175
|
mocked_factory.get_embedding_backend.assert_called_once_with(
|
|
166
|
-
model_name="prithvida/Splade_PP_en_v1", cache_dir=None, threads=None
|
|
176
|
+
model_name="prithvida/Splade_PP_en_v1", cache_dir=None, threads=None, local_files_only=False
|
|
167
177
|
)
|
|
168
178
|
|
|
169
179
|
@patch(
|
|
@@ -260,7 +270,7 @@ class TestFastembedSparseDocumentEmbedderDoc:
|
|
|
260
270
|
"meta_value 4\ndocument-number 4",
|
|
261
271
|
],
|
|
262
272
|
batch_size=32,
|
|
263
|
-
|
|
273
|
+
progress_bar=True,
|
|
264
274
|
parallel=None,
|
|
265
275
|
)
|
|
266
276
|
|
{fastembed_haystack-1.0.0 → fastembed_haystack-1.2.0}/tests/test_fastembed_sparse_text_embedder.py
RENAMED
|
@@ -52,6 +52,7 @@ class TestFastembedSparseTextEmbedder:
|
|
|
52
52
|
"threads": None,
|
|
53
53
|
"progress_bar": True,
|
|
54
54
|
"parallel": None,
|
|
55
|
+
"local_files_only": False,
|
|
55
56
|
},
|
|
56
57
|
}
|
|
57
58
|
|
|
@@ -65,6 +66,7 @@ class TestFastembedSparseTextEmbedder:
|
|
|
65
66
|
threads=2,
|
|
66
67
|
progress_bar=False,
|
|
67
68
|
parallel=1,
|
|
69
|
+
local_files_only=True,
|
|
68
70
|
)
|
|
69
71
|
embedder_dict = embedder.to_dict()
|
|
70
72
|
assert embedder_dict == {
|
|
@@ -75,6 +77,7 @@ class TestFastembedSparseTextEmbedder:
|
|
|
75
77
|
"threads": 2,
|
|
76
78
|
"progress_bar": False,
|
|
77
79
|
"parallel": 1,
|
|
80
|
+
"local_files_only": True,
|
|
78
81
|
},
|
|
79
82
|
}
|
|
80
83
|
|
|
@@ -131,7 +134,7 @@ class TestFastembedSparseTextEmbedder:
|
|
|
131
134
|
mocked_factory.get_embedding_backend.assert_not_called()
|
|
132
135
|
embedder.warm_up()
|
|
133
136
|
mocked_factory.get_embedding_backend.assert_called_once_with(
|
|
134
|
-
model_name="prithvida/Splade_PP_en_v1", cache_dir=None, threads=None
|
|
137
|
+
model_name="prithvida/Splade_PP_en_v1", cache_dir=None, threads=None, local_files_only=False
|
|
135
138
|
)
|
|
136
139
|
|
|
137
140
|
@patch(
|
|
@@ -59,6 +59,7 @@ class TestFastembedTextEmbedder:
|
|
|
59
59
|
"suffix": "",
|
|
60
60
|
"progress_bar": True,
|
|
61
61
|
"parallel": None,
|
|
62
|
+
"local_files_only": False,
|
|
62
63
|
},
|
|
63
64
|
}
|
|
64
65
|
|
|
@@ -74,6 +75,7 @@ class TestFastembedTextEmbedder:
|
|
|
74
75
|
suffix="suffix",
|
|
75
76
|
progress_bar=False,
|
|
76
77
|
parallel=1,
|
|
78
|
+
local_files_only=True,
|
|
77
79
|
)
|
|
78
80
|
embedder_dict = embedder.to_dict()
|
|
79
81
|
assert embedder_dict == {
|
|
@@ -86,6 +88,7 @@ class TestFastembedTextEmbedder:
|
|
|
86
88
|
"suffix": "suffix",
|
|
87
89
|
"progress_bar": False,
|
|
88
90
|
"parallel": 1,
|
|
91
|
+
"local_files_only": True,
|
|
89
92
|
},
|
|
90
93
|
}
|
|
91
94
|
|
|
@@ -150,7 +153,7 @@ class TestFastembedTextEmbedder:
|
|
|
150
153
|
mocked_factory.get_embedding_backend.assert_not_called()
|
|
151
154
|
embedder.warm_up()
|
|
152
155
|
mocked_factory.get_embedding_backend.assert_called_once_with(
|
|
153
|
-
model_name="BAAI/bge-small-en-v1.5", cache_dir=None, threads=None
|
|
156
|
+
model_name="BAAI/bge-small-en-v1.5", cache_dir=None, threads=None, local_files_only=False
|
|
154
157
|
)
|
|
155
158
|
|
|
156
159
|
@patch(
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|