fastembed-haystack 0.1.0__tar.gz → 1.1.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-0.1.0 → fastembed_haystack-1.1.0}/PKG-INFO +2 -1
- {fastembed_haystack-0.1.0 → fastembed_haystack-1.1.0}/pydoc/config.yml +1 -1
- {fastembed_haystack-0.1.0 → fastembed_haystack-1.1.0}/pyproject.toml +1 -0
- {fastembed_haystack-0.1.0 → fastembed_haystack-1.1.0}/src/haystack_integrations/components/embedders/fastembed/embedding_backend/fastembed_backend.py +14 -4
- {fastembed_haystack-0.1.0 → fastembed_haystack-1.1.0}/src/haystack_integrations/components/embedders/fastembed/fastembed_document_embedder.py +16 -5
- {fastembed_haystack-0.1.0 → fastembed_haystack-1.1.0}/src/haystack_integrations/components/embedders/fastembed/fastembed_sparse_document_embedder.py +20 -12
- {fastembed_haystack-0.1.0 → fastembed_haystack-1.1.0}/src/haystack_integrations/components/embedders/fastembed/fastembed_sparse_text_embedder.py +15 -15
- {fastembed_haystack-0.1.0 → fastembed_haystack-1.1.0}/src/haystack_integrations/components/embedders/fastembed/fastembed_text_embedder.py +12 -7
- {fastembed_haystack-0.1.0 → fastembed_haystack-1.1.0}/tests/test_fastembed_backend.py +3 -1
- {fastembed_haystack-0.1.0 → fastembed_haystack-1.1.0}/tests/test_fastembed_document_embedder.py +11 -1
- {fastembed_haystack-0.1.0 → fastembed_haystack-1.1.0}/tests/test_fastembed_sparse_document_embedder.py +11 -1
- {fastembed_haystack-0.1.0 → fastembed_haystack-1.1.0}/tests/test_fastembed_sparse_text_embedder.py +4 -11
- {fastembed_haystack-0.1.0 → fastembed_haystack-1.1.0}/tests/test_fastembed_text_embedder.py +4 -1
- {fastembed_haystack-0.1.0 → fastembed_haystack-1.1.0}/.gitignore +0 -0
- {fastembed_haystack-0.1.0 → fastembed_haystack-1.1.0}/LICENSE.txt +0 -0
- {fastembed_haystack-0.1.0 → fastembed_haystack-1.1.0}/README.md +0 -0
- {fastembed_haystack-0.1.0 → fastembed_haystack-1.1.0}/examples/example.py +0 -0
- {fastembed_haystack-0.1.0 → fastembed_haystack-1.1.0}/examples/sparse_example.py +0 -0
- {fastembed_haystack-0.1.0 → fastembed_haystack-1.1.0}/src/haystack_integrations/components/embedders/fastembed/__init__.py +0 -0
- {fastembed_haystack-0.1.0 → fastembed_haystack-1.1.0}/src/haystack_integrations/components/embedders/fastembed/embedding_backend/__init__.py +0 -0
- {fastembed_haystack-0.1.0 → fastembed_haystack-1.1.0}/tests/__init__.py +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.3
|
|
2
2
|
Name: fastembed-haystack
|
|
3
|
-
Version:
|
|
3
|
+
Version: 1.1.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
|
|
@@ -9,6 +9,7 @@ Author-email: deepset GmbH <info@deepset.ai>
|
|
|
9
9
|
License-Expression: Apache-2.0
|
|
10
10
|
License-File: LICENSE.txt
|
|
11
11
|
Classifier: Development Status :: 4 - Beta
|
|
12
|
+
Classifier: License :: OSI Approved :: Apache Software License
|
|
12
13
|
Classifier: Programming Language :: Python
|
|
13
14
|
Classifier: Programming Language :: Python :: 3.8
|
|
14
15
|
Classifier: Programming Language :: Python :: 3.9
|
|
@@ -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
|
|
@@ -14,6 +14,7 @@ authors = [
|
|
|
14
14
|
{ name = "deepset GmbH", email = "info@deepset.ai" },
|
|
15
15
|
]
|
|
16
16
|
classifiers = [
|
|
17
|
+
"License :: OSI Approved :: Apache Software License",
|
|
17
18
|
"Development Status :: 4 - Beta",
|
|
18
19
|
"Programming Language :: Python",
|
|
19
20
|
"Programming Language :: Python :: 3.8",
|
|
@@ -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.
|
|
@@ -29,14 +29,18 @@ class FastembedDocumentEmbedder:
|
|
|
29
29
|
# Text taken from PubMed QA Dataset (https://huggingface.co/datasets/pubmed_qa)
|
|
30
30
|
document_list = [
|
|
31
31
|
Document(
|
|
32
|
-
content="Oxidative stress generated within inflammatory joints can produce autoimmune phenomena and joint
|
|
32
|
+
content=("Oxidative stress generated within inflammatory joints can produce autoimmune phenomena and joint "
|
|
33
|
+
"destruction. Radical species with oxidative activity, including reactive nitrogen species, "
|
|
34
|
+
"represent mediators of inflammation and cartilage damage."),
|
|
33
35
|
meta={
|
|
34
36
|
"pubid": "25,445,628",
|
|
35
37
|
"long_answer": "yes",
|
|
36
38
|
},
|
|
37
39
|
),
|
|
38
40
|
Document(
|
|
39
|
-
content="Plasma levels of pancreatic polypeptide (PP) rise upon food intake. Although other pancreatic
|
|
41
|
+
content=("Plasma levels of pancreatic polypeptide (PP) rise upon food intake. Although other pancreatic "
|
|
42
|
+
"islet hormones, such as insulin and glucagon, have been extensively investigated, PP secretion "
|
|
43
|
+
"and actions are still poorly understood."),
|
|
40
44
|
meta={
|
|
41
45
|
"pubid": "25,445,712",
|
|
42
46
|
"long_answer": "yes",
|
|
@@ -49,7 +53,7 @@ class FastembedDocumentEmbedder:
|
|
|
49
53
|
print(f"Document Embedding: {result['documents'][0].embedding}")
|
|
50
54
|
print(f"Embedding Dimension: {len(result['documents'][0].embedding)}")
|
|
51
55
|
```
|
|
52
|
-
"""
|
|
56
|
+
"""
|
|
53
57
|
|
|
54
58
|
def __init__(
|
|
55
59
|
self,
|
|
@@ -61,6 +65,7 @@ class FastembedDocumentEmbedder:
|
|
|
61
65
|
batch_size: int = 256,
|
|
62
66
|
progress_bar: bool = True,
|
|
63
67
|
parallel: Optional[int] = None,
|
|
68
|
+
local_files_only: bool = False,
|
|
64
69
|
meta_fields_to_embed: Optional[List[str]] = None,
|
|
65
70
|
embedding_separator: str = "\n",
|
|
66
71
|
):
|
|
@@ -76,11 +81,12 @@ class FastembedDocumentEmbedder:
|
|
|
76
81
|
:param prefix: A string to add to the beginning of each text.
|
|
77
82
|
:param suffix: A string to add to the end of each text.
|
|
78
83
|
:param batch_size: Number of strings to encode at once.
|
|
79
|
-
:param progress_bar: If
|
|
84
|
+
:param progress_bar: If `True`, displays progress bar during embedding.
|
|
80
85
|
:param parallel:
|
|
81
86
|
If > 1, data-parallel encoding will be used, recommended for offline encoding of large datasets.
|
|
82
87
|
If 0, use all available cores.
|
|
83
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`.
|
|
84
90
|
:param meta_fields_to_embed: List of meta fields that should be embedded along with the Document content.
|
|
85
91
|
:param embedding_separator: Separator used to concatenate the meta fields to the Document content.
|
|
86
92
|
"""
|
|
@@ -93,6 +99,7 @@ class FastembedDocumentEmbedder:
|
|
|
93
99
|
self.batch_size = batch_size
|
|
94
100
|
self.progress_bar = progress_bar
|
|
95
101
|
self.parallel = parallel
|
|
102
|
+
self.local_files_only = local_files_only
|
|
96
103
|
self.meta_fields_to_embed = meta_fields_to_embed or []
|
|
97
104
|
self.embedding_separator = embedding_separator
|
|
98
105
|
|
|
@@ -112,6 +119,7 @@ class FastembedDocumentEmbedder:
|
|
|
112
119
|
batch_size=self.batch_size,
|
|
113
120
|
progress_bar=self.progress_bar,
|
|
114
121
|
parallel=self.parallel,
|
|
122
|
+
local_files_only=self.local_files_only,
|
|
115
123
|
meta_fields_to_embed=self.meta_fields_to_embed,
|
|
116
124
|
embedding_separator=self.embedding_separator,
|
|
117
125
|
)
|
|
@@ -122,7 +130,10 @@ class FastembedDocumentEmbedder:
|
|
|
122
130
|
"""
|
|
123
131
|
if not hasattr(self, "embedding_backend"):
|
|
124
132
|
self.embedding_backend = _FastembedEmbeddingBackendFactory.get_embedding_backend(
|
|
125
|
-
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,
|
|
126
137
|
)
|
|
127
138
|
|
|
128
139
|
def _prepare_texts_to_embed(self, documents: List[Document]) -> List[str]:
|
|
@@ -12,30 +12,31 @@ class FastembedSparseDocumentEmbedder:
|
|
|
12
12
|
|
|
13
13
|
Usage example:
|
|
14
14
|
```python
|
|
15
|
-
# To use this component, install the "fastembed-haystack" package.
|
|
16
|
-
# pip install fastembed-haystack
|
|
17
|
-
|
|
18
15
|
from haystack_integrations.components.embedders.fastembed import FastembedSparseDocumentEmbedder
|
|
19
16
|
from haystack.dataclasses import Document
|
|
20
17
|
|
|
21
|
-
|
|
18
|
+
sparse_doc_embedder = FastembedSparseDocumentEmbedder(
|
|
22
19
|
model="prithvida/Splade_PP_en_v1",
|
|
23
20
|
batch_size=32,
|
|
24
21
|
)
|
|
25
22
|
|
|
26
|
-
|
|
23
|
+
sparse_doc_embedder.warm_up()
|
|
27
24
|
|
|
28
25
|
# Text taken from PubMed QA Dataset (https://huggingface.co/datasets/pubmed_qa)
|
|
29
26
|
document_list = [
|
|
30
27
|
Document(
|
|
31
|
-
content="Oxidative stress generated within inflammatory joints can produce autoimmune phenomena and joint
|
|
28
|
+
content=("Oxidative stress generated within inflammatory joints can produce autoimmune phenomena and joint "
|
|
29
|
+
"destruction. Radical species with oxidative activity, including reactive nitrogen species, "
|
|
30
|
+
"represent mediators of inflammation and cartilage damage."),
|
|
32
31
|
meta={
|
|
33
32
|
"pubid": "25,445,628",
|
|
34
33
|
"long_answer": "yes",
|
|
35
34
|
},
|
|
36
35
|
),
|
|
37
36
|
Document(
|
|
38
|
-
content="Plasma levels of pancreatic polypeptide (PP) rise upon food intake. Although other pancreatic
|
|
37
|
+
content=("Plasma levels of pancreatic polypeptide (PP) rise upon food intake. Although other pancreatic "
|
|
38
|
+
"islet hormones, such as insulin and glucagon, have been extensively investigated, PP secretion "
|
|
39
|
+
"and actions are still poorly understood."),
|
|
39
40
|
meta={
|
|
40
41
|
"pubid": "25,445,712",
|
|
41
42
|
"long_answer": "yes",
|
|
@@ -43,12 +44,12 @@ class FastembedSparseDocumentEmbedder:
|
|
|
43
44
|
),
|
|
44
45
|
]
|
|
45
46
|
|
|
46
|
-
result =
|
|
47
|
+
result = sparse_doc_embedder.run(document_list)
|
|
47
48
|
print(f"Document Text: {result['documents'][0].content}")
|
|
48
|
-
print(f"Document Embedding: {result['documents'][0].sparse_embedding}")
|
|
49
|
-
print(f"Embedding Dimension: {len(result['documents'][0].sparse_embedding)}")
|
|
49
|
+
print(f"Document Sparse Embedding: {result['documents'][0].sparse_embedding}")
|
|
50
|
+
print(f"Sparse Embedding Dimension: {len(result['documents'][0].sparse_embedding)}")
|
|
50
51
|
```
|
|
51
|
-
"""
|
|
52
|
+
"""
|
|
52
53
|
|
|
53
54
|
def __init__(
|
|
54
55
|
self,
|
|
@@ -58,6 +59,7 @@ class FastembedSparseDocumentEmbedder:
|
|
|
58
59
|
batch_size: int = 32,
|
|
59
60
|
progress_bar: bool = True,
|
|
60
61
|
parallel: Optional[int] = None,
|
|
62
|
+
local_files_only: bool = False,
|
|
61
63
|
meta_fields_to_embed: Optional[List[str]] = None,
|
|
62
64
|
embedding_separator: str = "\n",
|
|
63
65
|
):
|
|
@@ -76,6 +78,7 @@ class FastembedSparseDocumentEmbedder:
|
|
|
76
78
|
If > 1, data-parallel encoding will be used, recommended for offline encoding of large datasets.
|
|
77
79
|
If 0, use all available cores.
|
|
78
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`.
|
|
79
82
|
:param meta_fields_to_embed: List of meta fields that should be embedded along with the Document content.
|
|
80
83
|
:param embedding_separator: Separator used to concatenate the meta fields to the Document content.
|
|
81
84
|
"""
|
|
@@ -86,6 +89,7 @@ class FastembedSparseDocumentEmbedder:
|
|
|
86
89
|
self.batch_size = batch_size
|
|
87
90
|
self.progress_bar = progress_bar
|
|
88
91
|
self.parallel = parallel
|
|
92
|
+
self.local_files_only = local_files_only
|
|
89
93
|
self.meta_fields_to_embed = meta_fields_to_embed or []
|
|
90
94
|
self.embedding_separator = embedding_separator
|
|
91
95
|
|
|
@@ -103,6 +107,7 @@ class FastembedSparseDocumentEmbedder:
|
|
|
103
107
|
batch_size=self.batch_size,
|
|
104
108
|
progress_bar=self.progress_bar,
|
|
105
109
|
parallel=self.parallel,
|
|
110
|
+
local_files_only=self.local_files_only,
|
|
106
111
|
meta_fields_to_embed=self.meta_fields_to_embed,
|
|
107
112
|
embedding_separator=self.embedding_separator,
|
|
108
113
|
)
|
|
@@ -113,7 +118,10 @@ class FastembedSparseDocumentEmbedder:
|
|
|
113
118
|
"""
|
|
114
119
|
if not hasattr(self, "embedding_backend"):
|
|
115
120
|
self.embedding_backend = _FastembedSparseEmbeddingBackendFactory.get_embedding_backend(
|
|
116
|
-
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,
|
|
117
125
|
)
|
|
118
126
|
|
|
119
127
|
def _prepare_texts_to_embed(self, documents: List[Document]) -> List[str]:
|
|
@@ -13,30 +13,28 @@ class FastembedSparseTextEmbedder:
|
|
|
13
13
|
|
|
14
14
|
Usage example:
|
|
15
15
|
```python
|
|
16
|
-
# To use this component, install the "fastembed-haystack" package.
|
|
17
|
-
# pip install fastembed-haystack
|
|
18
|
-
|
|
19
16
|
from haystack_integrations.components.embedders.fastembed import FastembedSparseTextEmbedder
|
|
20
17
|
|
|
21
|
-
text = "It clearly says online this will work on a Mac OS system.
|
|
18
|
+
text = ("It clearly says online this will work on a Mac OS system. "
|
|
19
|
+
"The disk comes and it does not, only Windows. Do Not order this if you have a Mac!!")
|
|
22
20
|
|
|
23
|
-
|
|
21
|
+
sparse_text_embedder = FastembedSparseTextEmbedder(
|
|
24
22
|
model="prithvida/Splade_PP_en_v1"
|
|
25
23
|
)
|
|
26
|
-
|
|
24
|
+
sparse_text_embedder.warm_up()
|
|
27
25
|
|
|
28
|
-
|
|
26
|
+
sparse_embedding = sparse_text_embedder.run(text)["sparse_embedding"]
|
|
29
27
|
```
|
|
30
|
-
"""
|
|
28
|
+
"""
|
|
31
29
|
|
|
32
30
|
def __init__(
|
|
33
31
|
self,
|
|
34
32
|
model: str = "prithvida/Splade_PP_en_v1",
|
|
35
33
|
cache_dir: Optional[str] = None,
|
|
36
34
|
threads: Optional[int] = None,
|
|
37
|
-
batch_size: int = 32,
|
|
38
35
|
progress_bar: bool = True,
|
|
39
36
|
parallel: Optional[int] = None,
|
|
37
|
+
local_files_only: bool = False,
|
|
40
38
|
):
|
|
41
39
|
"""
|
|
42
40
|
Create a FastembedSparseTextEmbedder component.
|
|
@@ -46,20 +44,20 @@ class FastembedSparseTextEmbedder:
|
|
|
46
44
|
Can be set using the `FASTEMBED_CACHE_PATH` env variable.
|
|
47
45
|
Defaults to `fastembed_cache` in the system's temp directory.
|
|
48
46
|
:param threads: The number of threads single onnxruntime session can use. Defaults to None.
|
|
49
|
-
:param
|
|
50
|
-
:param progress_bar: If true, displays progress bar during embedding.
|
|
47
|
+
:param progress_bar: If `True`, displays progress bar during embedding.
|
|
51
48
|
:param parallel:
|
|
52
49
|
If > 1, data-parallel encoding will be used, recommended for offline encoding of large datasets.
|
|
53
50
|
If 0, use all available cores.
|
|
54
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`.
|
|
55
53
|
"""
|
|
56
54
|
|
|
57
55
|
self.model_name = model
|
|
58
56
|
self.cache_dir = cache_dir
|
|
59
57
|
self.threads = threads
|
|
60
|
-
self.batch_size = batch_size
|
|
61
58
|
self.progress_bar = progress_bar
|
|
62
59
|
self.parallel = parallel
|
|
60
|
+
self.local_files_only = local_files_only
|
|
63
61
|
|
|
64
62
|
def to_dict(self) -> Dict[str, Any]:
|
|
65
63
|
"""
|
|
@@ -73,9 +71,9 @@ class FastembedSparseTextEmbedder:
|
|
|
73
71
|
model=self.model_name,
|
|
74
72
|
cache_dir=self.cache_dir,
|
|
75
73
|
threads=self.threads,
|
|
76
|
-
batch_size=self.batch_size,
|
|
77
74
|
progress_bar=self.progress_bar,
|
|
78
75
|
parallel=self.parallel,
|
|
76
|
+
local_files_only=self.local_files_only,
|
|
79
77
|
)
|
|
80
78
|
|
|
81
79
|
def warm_up(self):
|
|
@@ -84,7 +82,10 @@ class FastembedSparseTextEmbedder:
|
|
|
84
82
|
"""
|
|
85
83
|
if not hasattr(self, "embedding_backend"):
|
|
86
84
|
self.embedding_backend = _FastembedSparseEmbeddingBackendFactory.get_embedding_backend(
|
|
87
|
-
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,
|
|
88
89
|
)
|
|
89
90
|
|
|
90
91
|
@component.output_types(sparse_embedding=SparseEmbedding)
|
|
@@ -110,7 +111,6 @@ class FastembedSparseTextEmbedder:
|
|
|
110
111
|
|
|
111
112
|
embedding = self.embedding_backend.embed(
|
|
112
113
|
[text],
|
|
113
|
-
batch_size=self.batch_size,
|
|
114
114
|
show_progress_bar=self.progress_bar,
|
|
115
115
|
parallel=self.parallel,
|
|
116
116
|
)[0]
|
|
@@ -12,12 +12,10 @@ class FastembedTextEmbedder:
|
|
|
12
12
|
|
|
13
13
|
Usage example:
|
|
14
14
|
```python
|
|
15
|
-
# To use this component, install the "fastembed-haystack" package.
|
|
16
|
-
# pip install fastembed-haystack
|
|
17
|
-
|
|
18
15
|
from haystack_integrations.components.embedders.fastembed import FastembedTextEmbedder
|
|
19
16
|
|
|
20
|
-
text = "It clearly says online this will work on a Mac OS system.
|
|
17
|
+
text = ("It clearly says online this will work on a Mac OS system. "
|
|
18
|
+
"The disk comes and it does not, only Windows. Do Not order this if you have a Mac!!")
|
|
21
19
|
|
|
22
20
|
text_embedder = FastembedTextEmbedder(
|
|
23
21
|
model="BAAI/bge-small-en-v1.5"
|
|
@@ -26,7 +24,7 @@ class FastembedTextEmbedder:
|
|
|
26
24
|
|
|
27
25
|
embedding = text_embedder.run(text)["embedding"]
|
|
28
26
|
```
|
|
29
|
-
"""
|
|
27
|
+
"""
|
|
30
28
|
|
|
31
29
|
def __init__(
|
|
32
30
|
self,
|
|
@@ -37,6 +35,7 @@ class FastembedTextEmbedder:
|
|
|
37
35
|
suffix: str = "",
|
|
38
36
|
progress_bar: bool = True,
|
|
39
37
|
parallel: Optional[int] = None,
|
|
38
|
+
local_files_only: bool = False,
|
|
40
39
|
):
|
|
41
40
|
"""
|
|
42
41
|
Create a FastembedTextEmbedder component.
|
|
@@ -48,11 +47,12 @@ class FastembedTextEmbedder:
|
|
|
48
47
|
:param threads: The number of threads single onnxruntime session can use. Defaults to None.
|
|
49
48
|
:param prefix: A string to add to the beginning of each text.
|
|
50
49
|
:param suffix: A string to add to the end of each text.
|
|
51
|
-
:param progress_bar: If
|
|
50
|
+
:param progress_bar: If `True`, displays progress bar during embedding.
|
|
52
51
|
:param parallel:
|
|
53
52
|
If > 1, data-parallel encoding will be used, recommended for offline encoding of large datasets.
|
|
54
53
|
If 0, use all available cores.
|
|
55
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`.
|
|
56
56
|
"""
|
|
57
57
|
|
|
58
58
|
self.model_name = model
|
|
@@ -62,6 +62,7 @@ class FastembedTextEmbedder:
|
|
|
62
62
|
self.suffix = suffix
|
|
63
63
|
self.progress_bar = progress_bar
|
|
64
64
|
self.parallel = parallel
|
|
65
|
+
self.local_files_only = local_files_only
|
|
65
66
|
|
|
66
67
|
def to_dict(self) -> Dict[str, Any]:
|
|
67
68
|
"""
|
|
@@ -79,6 +80,7 @@ class FastembedTextEmbedder:
|
|
|
79
80
|
suffix=self.suffix,
|
|
80
81
|
progress_bar=self.progress_bar,
|
|
81
82
|
parallel=self.parallel,
|
|
83
|
+
local_files_only=self.local_files_only,
|
|
82
84
|
)
|
|
83
85
|
|
|
84
86
|
def warm_up(self):
|
|
@@ -87,7 +89,10 @@ class FastembedTextEmbedder:
|
|
|
87
89
|
"""
|
|
88
90
|
if not hasattr(self, "embedding_backend"):
|
|
89
91
|
self.embedding_backend = _FastembedEmbeddingBackendFactory.get_embedding_backend(
|
|
90
|
-
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,
|
|
91
96
|
)
|
|
92
97
|
|
|
93
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-0.1.0 → fastembed_haystack-1.1.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(
|
{fastembed_haystack-0.1.0 → fastembed_haystack-1.1.0}/tests/test_fastembed_sparse_text_embedder.py
RENAMED
|
@@ -18,7 +18,6 @@ class TestFastembedSparseTextEmbedder:
|
|
|
18
18
|
assert embedder.model_name == "prithvida/Splade_PP_en_v1"
|
|
19
19
|
assert embedder.cache_dir is None
|
|
20
20
|
assert embedder.threads is None
|
|
21
|
-
assert embedder.batch_size == 32
|
|
22
21
|
assert embedder.progress_bar is True
|
|
23
22
|
assert embedder.parallel is None
|
|
24
23
|
|
|
@@ -30,14 +29,12 @@ class TestFastembedSparseTextEmbedder:
|
|
|
30
29
|
model="prithvida/Splade_PP_en_v1",
|
|
31
30
|
cache_dir="fake_dir",
|
|
32
31
|
threads=2,
|
|
33
|
-
batch_size=64,
|
|
34
32
|
progress_bar=False,
|
|
35
33
|
parallel=1,
|
|
36
34
|
)
|
|
37
35
|
assert embedder.model_name == "prithvida/Splade_PP_en_v1"
|
|
38
36
|
assert embedder.cache_dir == "fake_dir"
|
|
39
37
|
assert embedder.threads == 2
|
|
40
|
-
assert embedder.batch_size == 64
|
|
41
38
|
assert embedder.progress_bar is False
|
|
42
39
|
assert embedder.parallel == 1
|
|
43
40
|
|
|
@@ -53,9 +50,9 @@ class TestFastembedSparseTextEmbedder:
|
|
|
53
50
|
"model": "prithvida/Splade_PP_en_v1",
|
|
54
51
|
"cache_dir": None,
|
|
55
52
|
"threads": None,
|
|
56
|
-
"batch_size": 32,
|
|
57
53
|
"progress_bar": True,
|
|
58
54
|
"parallel": None,
|
|
55
|
+
"local_files_only": False,
|
|
59
56
|
},
|
|
60
57
|
}
|
|
61
58
|
|
|
@@ -67,9 +64,9 @@ class TestFastembedSparseTextEmbedder:
|
|
|
67
64
|
model="prithvida/Splade_PP_en_v1",
|
|
68
65
|
cache_dir="fake_dir",
|
|
69
66
|
threads=2,
|
|
70
|
-
batch_size=64,
|
|
71
67
|
progress_bar=False,
|
|
72
68
|
parallel=1,
|
|
69
|
+
local_files_only=True,
|
|
73
70
|
)
|
|
74
71
|
embedder_dict = embedder.to_dict()
|
|
75
72
|
assert embedder_dict == {
|
|
@@ -78,9 +75,9 @@ class TestFastembedSparseTextEmbedder:
|
|
|
78
75
|
"model": "prithvida/Splade_PP_en_v1",
|
|
79
76
|
"cache_dir": "fake_dir",
|
|
80
77
|
"threads": 2,
|
|
81
|
-
"batch_size": 64,
|
|
82
78
|
"progress_bar": False,
|
|
83
79
|
"parallel": 1,
|
|
80
|
+
"local_files_only": True,
|
|
84
81
|
},
|
|
85
82
|
}
|
|
86
83
|
|
|
@@ -94,7 +91,6 @@ class TestFastembedSparseTextEmbedder:
|
|
|
94
91
|
"model": "prithvida/Splade_PP_en_v1",
|
|
95
92
|
"cache_dir": None,
|
|
96
93
|
"threads": None,
|
|
97
|
-
"batch_size": 32,
|
|
98
94
|
"progress_bar": True,
|
|
99
95
|
"parallel": None,
|
|
100
96
|
},
|
|
@@ -103,7 +99,6 @@ class TestFastembedSparseTextEmbedder:
|
|
|
103
99
|
assert embedder.model_name == "prithvida/Splade_PP_en_v1"
|
|
104
100
|
assert embedder.cache_dir is None
|
|
105
101
|
assert embedder.threads is None
|
|
106
|
-
assert embedder.batch_size == 32
|
|
107
102
|
assert embedder.progress_bar is True
|
|
108
103
|
assert embedder.parallel is None
|
|
109
104
|
|
|
@@ -117,7 +112,6 @@ class TestFastembedSparseTextEmbedder:
|
|
|
117
112
|
"model": "prithvida/Splade_PP_en_v1",
|
|
118
113
|
"cache_dir": "fake_dir",
|
|
119
114
|
"threads": 2,
|
|
120
|
-
"batch_size": 64,
|
|
121
115
|
"progress_bar": False,
|
|
122
116
|
"parallel": 1,
|
|
123
117
|
},
|
|
@@ -126,7 +120,6 @@ class TestFastembedSparseTextEmbedder:
|
|
|
126
120
|
assert embedder.model_name == "prithvida/Splade_PP_en_v1"
|
|
127
121
|
assert embedder.cache_dir == "fake_dir"
|
|
128
122
|
assert embedder.threads == 2
|
|
129
|
-
assert embedder.batch_size == 64
|
|
130
123
|
assert embedder.progress_bar is False
|
|
131
124
|
assert embedder.parallel == 1
|
|
132
125
|
|
|
@@ -141,7 +134,7 @@ class TestFastembedSparseTextEmbedder:
|
|
|
141
134
|
mocked_factory.get_embedding_backend.assert_not_called()
|
|
142
135
|
embedder.warm_up()
|
|
143
136
|
mocked_factory.get_embedding_backend.assert_called_once_with(
|
|
144
|
-
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
|
|
145
138
|
)
|
|
146
139
|
|
|
147
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
|
|
File without changes
|