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.
Files changed (21) hide show
  1. {fastembed_haystack-0.1.0 → fastembed_haystack-1.1.0}/PKG-INFO +2 -1
  2. {fastembed_haystack-0.1.0 → fastembed_haystack-1.1.0}/pydoc/config.yml +1 -1
  3. {fastembed_haystack-0.1.0 → fastembed_haystack-1.1.0}/pyproject.toml +1 -0
  4. {fastembed_haystack-0.1.0 → fastembed_haystack-1.1.0}/src/haystack_integrations/components/embedders/fastembed/embedding_backend/fastembed_backend.py +14 -4
  5. {fastembed_haystack-0.1.0 → fastembed_haystack-1.1.0}/src/haystack_integrations/components/embedders/fastembed/fastembed_document_embedder.py +16 -5
  6. {fastembed_haystack-0.1.0 → fastembed_haystack-1.1.0}/src/haystack_integrations/components/embedders/fastembed/fastembed_sparse_document_embedder.py +20 -12
  7. {fastembed_haystack-0.1.0 → fastembed_haystack-1.1.0}/src/haystack_integrations/components/embedders/fastembed/fastembed_sparse_text_embedder.py +15 -15
  8. {fastembed_haystack-0.1.0 → fastembed_haystack-1.1.0}/src/haystack_integrations/components/embedders/fastembed/fastembed_text_embedder.py +12 -7
  9. {fastembed_haystack-0.1.0 → fastembed_haystack-1.1.0}/tests/test_fastembed_backend.py +3 -1
  10. {fastembed_haystack-0.1.0 → fastembed_haystack-1.1.0}/tests/test_fastembed_document_embedder.py +11 -1
  11. {fastembed_haystack-0.1.0 → fastembed_haystack-1.1.0}/tests/test_fastembed_sparse_document_embedder.py +11 -1
  12. {fastembed_haystack-0.1.0 → fastembed_haystack-1.1.0}/tests/test_fastembed_sparse_text_embedder.py +4 -11
  13. {fastembed_haystack-0.1.0 → fastembed_haystack-1.1.0}/tests/test_fastembed_text_embedder.py +4 -1
  14. {fastembed_haystack-0.1.0 → fastembed_haystack-1.1.0}/.gitignore +0 -0
  15. {fastembed_haystack-0.1.0 → fastembed_haystack-1.1.0}/LICENSE.txt +0 -0
  16. {fastembed_haystack-0.1.0 → fastembed_haystack-1.1.0}/README.md +0 -0
  17. {fastembed_haystack-0.1.0 → fastembed_haystack-1.1.0}/examples/example.py +0 -0
  18. {fastembed_haystack-0.1.0 → fastembed_haystack-1.1.0}/examples/sparse_example.py +0 -0
  19. {fastembed_haystack-0.1.0 → fastembed_haystack-1.1.0}/src/haystack_integrations/components/embedders/fastembed/__init__.py +0 -0
  20. {fastembed_haystack-0.1.0 → fastembed_haystack-1.1.0}/src/haystack_integrations/components/embedders/fastembed/embedding_backend/__init__.py +0 -0
  21. {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: 0.1.0
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.ReadmePreviewRenderer
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(model_name=model_name, cache_dir=cache_dir, threads=threads)
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(model_name=model_name, cache_dir=cache_dir, threads=threads)
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(model_name=model_name, cache_dir=cache_dir, threads=threads)
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 destruction. Radical species with oxidative activity, including reactive nitrogen species, represent mediators of inflammation and cartilage damage.",
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 islet hormones, such as insulin and glucagon, have been extensively investigated, PP secretion and actions are still poorly understood.",
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
- """ # noqa: E501
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 true, displays progress bar during embedding.
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, cache_dir=self.cache_dir, threads=self.threads
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
- doc_embedder = FastembedSparseDocumentEmbedder(
18
+ sparse_doc_embedder = FastembedSparseDocumentEmbedder(
22
19
  model="prithvida/Splade_PP_en_v1",
23
20
  batch_size=32,
24
21
  )
25
22
 
26
- doc_embedder.warm_up()
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 destruction. Radical species with oxidative activity, including reactive nitrogen species, represent mediators of inflammation and cartilage damage.",
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 islet hormones, such as insulin and glucagon, have been extensively investigated, PP secretion and actions are still poorly understood.",
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 = doc_embedder.run(document_list)
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
- """ # noqa: E501
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, cache_dir=self.cache_dir, threads=self.threads
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. The disk comes and it does not, only Windows. Do Not order this if you have a Mac!!"
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
- text_embedder = FastembedSparseTextEmbedder(
21
+ sparse_text_embedder = FastembedSparseTextEmbedder(
24
22
  model="prithvida/Splade_PP_en_v1"
25
23
  )
26
- text_embedder.warm_up()
24
+ sparse_text_embedder.warm_up()
27
25
 
28
- embedding = text_embedder.run(text)["embedding"]
26
+ sparse_embedding = sparse_text_embedder.run(text)["sparse_embedding"]
29
27
  ```
30
- """ # noqa: E501
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 batch_size: Number of strings to encode at once.
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, cache_dir=self.cache_dir, threads=self.threads
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. The disk comes and it does not, only Windows. Do Not order this if you have a Mac!!"
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
- """ # noqa: E501
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 true, displays progress bar during embedding.
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, cache_dir=self.cache_dir, threads=self.threads
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(model_name="BAAI/bge-small-en-v1.5", cache_dir=None, threads=None)
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
 
@@ -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(
@@ -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(