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 d0a32a88c..dbc4fdb64 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 @@ -1,4 +1,4 @@ -from typing import ClassVar, Dict, List, Optional +from typing import ClassVar, Dict, List, Optional, Union import numpy as np @@ -81,7 +81,7 @@ def __init__( ): self.model = SparseTextEmbedding(model_name=model_name, cache_dir=cache_dir, threads=threads) - def embed(self, data: List[List[str]], **kwargs) -> List[Dict[str, np.ndarray]]: + def embed(self, data: List[List[str]], **kwargs) -> List[Dict[str, Union[List[int], List[float]]]]: # The embed method returns a Iterable[SparseEmbedding], so we convert it to a list of dictionaries. # Each dict contains an `indices` key containing a list of int and an `values` key containing a list of floats. sparse_embeddings = [sparse_embedding.as_object() for sparse_embedding in self.model.embed(data, **kwargs)] diff --git a/integrations/fastembed/src/haystack_integrations/components/embedders/fastembed/fastembed_text_SPLADE_embedder.py b/integrations/fastembed/src/haystack_integrations/components/embedders/fastembed/fastembed_text_SPLADE_embedder.py index ef8c73803..33700ac12 100644 --- a/integrations/fastembed/src/haystack_integrations/components/embedders/fastembed/fastembed_text_SPLADE_embedder.py +++ b/integrations/fastembed/src/haystack_integrations/components/embedders/fastembed/fastembed_text_SPLADE_embedder.py @@ -1,4 +1,4 @@ -from typing import Any, Dict, List, Optional +from typing import Any, Dict, List, Optional, Union import numpy as np from haystack import component, default_to_dict @@ -94,7 +94,7 @@ def warm_up(self): model_name=self.model_name, cache_dir=self.cache_dir, threads=self.threads ) - @component.output_types(embedding=List[Dict[str, np.ndarray]]) + @component.output_types(embedding=List[Dict[str, Union[List[int], List[float]]]]) def run(self, text: str): """ Embeds text using the Fastembed model.