diff --git a/integrations/fastembed/src/haystack_integrations/components/embedders/fastembed/embedding_backend/fastembed_backend.py b/integrations/fastembed/src/haystack_integrations/components/embedders/fastembed/embedding_backend/fastembed_backend.py index 66f797549..db2eab9fb 100644 --- a/integrations/fastembed/src/haystack_integrations/components/embedders/fastembed/embedding_backend/fastembed_backend.py +++ b/integrations/fastembed/src/haystack_integrations/components/embedders/fastembed/embedding_backend/fastembed_backend.py @@ -5,6 +5,7 @@ from fastembed import TextEmbedding from fastembed.sparse.sparse_text_embedding import SparseTextEmbedding +from fastembed.common.types import OnnxProvider class _FastembedEmbeddingBackendFactory: @@ -20,6 +21,7 @@ def get_embedding_backend( cache_dir: Optional[str] = None, threads: Optional[int] = None, local_files_only: bool = False, + onnx_providers: Optional[List[OnnxProvider]] = None, ): embedding_backend_id = f"{model_name}{cache_dir}{threads}" @@ -27,7 +29,7 @@ def get_embedding_backend( return _FastembedEmbeddingBackendFactory._instances[embedding_backend_id] embedding_backend = _FastembedEmbeddingBackend( - model_name=model_name, cache_dir=cache_dir, threads=threads, local_files_only=local_files_only + model_name=model_name, cache_dir=cache_dir, threads=threads, local_files_only=local_files_only, onnx_providers=onnx_providers ) _FastembedEmbeddingBackendFactory._instances[embedding_backend_id] = embedding_backend return embedding_backend @@ -44,9 +46,10 @@ def __init__( cache_dir: Optional[str] = None, threads: Optional[int] = None, local_files_only: bool = False, + onnx_providers: Optional[List[OnnxProvider]] = None, ): self.model = TextEmbedding( - model_name=model_name, cache_dir=cache_dir, threads=threads, local_files_only=local_files_only + model_name=model_name, cache_dir=cache_dir, threads=threads, local_files_only=local_files_only, providers=onnx_providers ) def embed(self, data: List[str], progress_bar=True, **kwargs) -> List[List[float]]: @@ -73,6 +76,7 @@ def get_embedding_backend( cache_dir: Optional[str] = None, threads: Optional[int] = None, local_files_only: bool = False, + onnx_providers: Optional[List[OnnxProvider]] = None, ): embedding_backend_id = f"{model_name}{cache_dir}{threads}" @@ -80,7 +84,7 @@ def get_embedding_backend( return _FastembedSparseEmbeddingBackendFactory._instances[embedding_backend_id] embedding_backend = _FastembedSparseEmbeddingBackend( - model_name=model_name, cache_dir=cache_dir, threads=threads, local_files_only=local_files_only + model_name=model_name, cache_dir=cache_dir, threads=threads, local_files_only=local_files_only, onnx_providers=onnx_providers ) _FastembedSparseEmbeddingBackendFactory._instances[embedding_backend_id] = embedding_backend return embedding_backend @@ -97,9 +101,10 @@ def __init__( cache_dir: Optional[str] = None, threads: Optional[int] = None, local_files_only: bool = False, + onnx_providers: Optional[List[OnnxProvider]] = None, ): self.model = SparseTextEmbedding( - model_name=model_name, cache_dir=cache_dir, threads=threads, local_files_only=local_files_only + model_name=model_name, cache_dir=cache_dir, threads=threads, local_files_only=local_files_only, providers=onnx_providers ) def embed(self, data: List[List[str]], progress_bar=True, **kwargs) -> List[SparseEmbedding]: diff --git a/integrations/fastembed/src/haystack_integrations/components/embedders/fastembed/fastembed_document_embedder.py b/integrations/fastembed/src/haystack_integrations/components/embedders/fastembed/fastembed_document_embedder.py index 8b63582c5..b2e362d1f 100644 --- a/integrations/fastembed/src/haystack_integrations/components/embedders/fastembed/fastembed_document_embedder.py +++ b/integrations/fastembed/src/haystack_integrations/components/embedders/fastembed/fastembed_document_embedder.py @@ -2,7 +2,7 @@ from haystack import Document, component, default_to_dict -from .embedding_backend.fastembed_backend import _FastembedEmbeddingBackendFactory +from .embedding_backend.fastembed_backend import _FastembedEmbeddingBackendFactory, OnnxProvider @component @@ -68,6 +68,7 @@ def __init__( local_files_only: bool = False, meta_fields_to_embed: Optional[List[str]] = None, embedding_separator: str = "\n", + onnx_providers: Optional[List[OnnxProvider]] = None, ): """ Create an FastembedDocumentEmbedder component. @@ -89,6 +90,7 @@ def __init__( :param local_files_only: If `True`, only use the model files in the `cache_dir`. :param meta_fields_to_embed: List of meta fields that should be embedded along with the Document content. :param embedding_separator: Separator used to concatenate the meta fields to the Document content. + :param onnx_providers: ONNX Runtime's Execution providers (EPs) """ self.model_name = model @@ -102,6 +104,7 @@ def __init__( self.local_files_only = local_files_only self.meta_fields_to_embed = meta_fields_to_embed or [] self.embedding_separator = embedding_separator + self.onnx_providers = onnx_providers def to_dict(self) -> Dict[str, Any]: """ @@ -122,6 +125,7 @@ def to_dict(self) -> Dict[str, Any]: local_files_only=self.local_files_only, meta_fields_to_embed=self.meta_fields_to_embed, embedding_separator=self.embedding_separator, + onnx_providers=self.onnx_providers, ) def warm_up(self): @@ -134,6 +138,7 @@ def warm_up(self): cache_dir=self.cache_dir, threads=self.threads, local_files_only=self.local_files_only, + onnx_providers=self.onnx_providers, ) def _prepare_texts_to_embed(self, documents: List[Document]) -> List[str]: diff --git a/integrations/fastembed/src/haystack_integrations/components/embedders/fastembed/fastembed_sparse_document_embedder.py b/integrations/fastembed/src/haystack_integrations/components/embedders/fastembed/fastembed_sparse_document_embedder.py index 4b72389fa..8903e9e62 100644 --- a/integrations/fastembed/src/haystack_integrations/components/embedders/fastembed/fastembed_sparse_document_embedder.py +++ b/integrations/fastembed/src/haystack_integrations/components/embedders/fastembed/fastembed_sparse_document_embedder.py @@ -2,7 +2,7 @@ from haystack import Document, component, default_to_dict -from .embedding_backend.fastembed_backend import _FastembedSparseEmbeddingBackendFactory +from .embedding_backend.fastembed_backend import _FastembedSparseEmbeddingBackendFactory, OnnxProvider @component @@ -62,6 +62,7 @@ def __init__( local_files_only: bool = False, meta_fields_to_embed: Optional[List[str]] = None, embedding_separator: str = "\n", + onnx_providers: Optional[List[OnnxProvider]] = None, ): """ Create an FastembedDocumentEmbedder component. @@ -92,6 +93,7 @@ def __init__( self.local_files_only = local_files_only self.meta_fields_to_embed = meta_fields_to_embed or [] self.embedding_separator = embedding_separator + self.onnx_providers = onnx_providers def to_dict(self) -> Dict[str, Any]: """ @@ -110,6 +112,7 @@ def to_dict(self) -> Dict[str, Any]: local_files_only=self.local_files_only, meta_fields_to_embed=self.meta_fields_to_embed, embedding_separator=self.embedding_separator, + onnx_providers=self.onnx_providers, ) def warm_up(self): @@ -122,6 +125,7 @@ def warm_up(self): cache_dir=self.cache_dir, threads=self.threads, local_files_only=self.local_files_only, + onnx_providers=self.onnx_providers ) def _prepare_texts_to_embed(self, documents: List[Document]) -> List[str]: diff --git a/integrations/fastembed/src/haystack_integrations/components/embedders/fastembed/fastembed_sparse_text_embedder.py b/integrations/fastembed/src/haystack_integrations/components/embedders/fastembed/fastembed_sparse_text_embedder.py index 67348b2bd..4e825f90e 100644 --- a/integrations/fastembed/src/haystack_integrations/components/embedders/fastembed/fastembed_sparse_text_embedder.py +++ b/integrations/fastembed/src/haystack_integrations/components/embedders/fastembed/fastembed_sparse_text_embedder.py @@ -1,9 +1,9 @@ -from typing import Any, Dict, Optional +from typing import Any, Dict, Optional, List from haystack import component, default_to_dict from haystack.dataclasses.sparse_embedding import SparseEmbedding -from .embedding_backend.fastembed_backend import _FastembedSparseEmbeddingBackendFactory +from .embedding_backend.fastembed_backend import _FastembedSparseEmbeddingBackendFactory, OnnxProvider @component @@ -35,6 +35,7 @@ def __init__( progress_bar: bool = True, parallel: Optional[int] = None, local_files_only: bool = False, + onnx_providers: Optional[List[OnnxProvider]] = None, ): """ Create a FastembedSparseTextEmbedder component. @@ -58,6 +59,7 @@ def __init__( self.progress_bar = progress_bar self.parallel = parallel self.local_files_only = local_files_only + self.onnx_providers = onnx_providers def to_dict(self) -> Dict[str, Any]: """ @@ -74,6 +76,7 @@ def to_dict(self) -> Dict[str, Any]: progress_bar=self.progress_bar, parallel=self.parallel, local_files_only=self.local_files_only, + onnx_providers=self.onnx_providers, ) def warm_up(self): @@ -86,6 +89,7 @@ def warm_up(self): cache_dir=self.cache_dir, threads=self.threads, local_files_only=self.local_files_only, + onnx_providers=self.onnx_providers, ) @component.output_types(sparse_embedding=SparseEmbedding) diff --git a/integrations/fastembed/src/haystack_integrations/components/embedders/fastembed/fastembed_text_embedder.py b/integrations/fastembed/src/haystack_integrations/components/embedders/fastembed/fastembed_text_embedder.py index a7f56ff97..e19981052 100644 --- a/integrations/fastembed/src/haystack_integrations/components/embedders/fastembed/fastembed_text_embedder.py +++ b/integrations/fastembed/src/haystack_integrations/components/embedders/fastembed/fastembed_text_embedder.py @@ -2,7 +2,7 @@ from haystack import component, default_to_dict -from .embedding_backend.fastembed_backend import _FastembedEmbeddingBackendFactory +from .embedding_backend.fastembed_backend import _FastembedEmbeddingBackendFactory, OnnxProvider @component @@ -36,6 +36,7 @@ def __init__( progress_bar: bool = True, parallel: Optional[int] = None, local_files_only: bool = False, + onnx_providers: Optional[List[OnnxProvider]] = None, ): """ Create a FastembedTextEmbedder component. @@ -63,6 +64,7 @@ def __init__( self.progress_bar = progress_bar self.parallel = parallel self.local_files_only = local_files_only + self.onnx_providers = onnx_providers def to_dict(self) -> Dict[str, Any]: """ @@ -81,6 +83,7 @@ def to_dict(self) -> Dict[str, Any]: progress_bar=self.progress_bar, parallel=self.parallel, local_files_only=self.local_files_only, + onnx_providers=self.onnx_providers, ) def warm_up(self): @@ -93,6 +96,7 @@ def warm_up(self): cache_dir=self.cache_dir, threads=self.threads, local_files_only=self.local_files_only, + onnx_providers=self.onnx_providers, ) @component.output_types(embedding=List[float])