Wrapped output into SparseEmbedding dataclass + fix tests

deepset-ai · anakin87 · Sep 19, 2025 · Jul 3, 2025 · Jul 3, 2025 · Aug 3, 2025
commit 74c222e8e55efe57cba568545a3031d1fdddd189
@@ -4,6 +4,7 @@
 
 from typing import Any, Dict, List, Literal, Optional
 
+from haystack.dataclasses.sparse_embedding import SparseEmbedding
 from haystack.lazy_imports import LazyImport
 from haystack.utils.auth import Secret
 
@@ -162,6 +163,19 @@ def __init__(  # pylint: disable=too-many-positional-arguments
             backend=backend,
         )
 
-    def embed(self, data: List[str], **kwargs) -> List[List[float]]:
-        embeddings = self.model.encode(data, **kwargs).tolist()
-        return embeddings
+    def embed(self, data: List[str], **kwargs) -> List[SparseEmbedding]:
+        embeddings = self.model.encode(data, **kwargs)
+
+        sparse_embeddings = []
+
+        if isinstance(embeddings, list):
+            for embedding in embeddings:
+                sparse_embeddings.append(
+                    SparseEmbedding(indices=embedding.indices.tolist(), values=embedding.values.tolist())
+                )
+        else:
+            sparse_embeddings.append(
+                SparseEmbedding(indices=embeddings.indices.tolist(), values=embeddings.values.tolist())
+            )
+
+        return sparse_embeddings
@@ -179,7 +179,6 @@ def warm_up(self):
                 tokenizer_kwargs=self.tokenizer_kwargs,
                 config_kwargs=self.config_kwargs,
                 backend=self.backend,
-                sparse=True,
             )
             if self.tokenizer_kwargs and self.tokenizer_kwargs.get("model_max_length"):
                 self.embedding_backend.model.max_seq_length = self.tokenizer_kwargs["model_max_length"]

@@ -189,7 +189,7 @@ def test_from_dict_none_device(self):
         assert component.local_files_only is False
 
     @patch(
-        "haystack.components.embedders.sentence_transformers_text_embedder._SentenceTransformersEmbeddingBackendFactory"
+        "haystack.components.embedders.sentence_transformers_sparse_text_embedder._SentenceTransformersSparseEmbeddingBackendFactory"
     )
     def test_warmup(self, mocked_factory):
         embedder = SentenceTransformersSparseTextEmbedder(
@@ -214,7 +214,7 @@ def test_warmup(self, mocked_factory):
         )
 
     @patch(
-        "haystack.components.embedders.sentence_transformers_text_embedder._SentenceTransformersEmbeddingBackendFactory"
+        "haystack.components.embedders.sentence_transformers_sparse_text_embedder._SentenceTransformersSparseEmbeddingBackendFactory"
     )
     def test_warmup_doesnt_reload(self, mocked_factory):
         embedder = SentenceTransformersSparseTextEmbedder(model="model")
@@ -248,7 +248,7 @@ def test_run_wrong_input_format(self):
             embedder.run(text=list_integers_input)
 
     @patch(
-        "haystack.components.embedders.sentence_transformers_text_embedder._SentenceTransformersEmbeddingBackendFactory"
+        "haystack.components.embedders.sentence_transformers_sparse_text_embedder._SentenceTransformersSparseEmbeddingBackendFactory"
     )
     def test_model_onnx_backend(self, mocked_factory):
         onnx_embedder = SentenceTransformersSparseTextEmbedder(
@@ -275,7 +275,7 @@ def test_model_onnx_backend(self, mocked_factory):
         )
 
     @patch(
-        "haystack.components.embedders.sentence_transformers_text_embedder._SentenceTransformersEmbeddingBackendFactory"
+        "haystack.components.embedders.sentence_transformers_sparse_text_embedder._SentenceTransformersSparseEmbeddingBackendFactory"
     )
     def test_model_openvino_backend(self, mocked_factory):
         openvino_embedder = SentenceTransformersSparseTextEmbedder(
@@ -302,7 +302,7 @@ def test_model_openvino_backend(self, mocked_factory):
         )
 
     @patch(
-        "haystack.components.embedders.sentence_transformers_text_embedder._SentenceTransformersEmbeddingBackendFactory"
+        "haystack.components.embedders.sentence_transformers_sparse_text_embedder._SentenceTransformersSparseEmbeddingBackendFactory"
     )
     @pytest.mark.parametrize("model_kwargs", [{"torch_dtype": "bfloat16"}, {"torch_dtype": "float16"}])
     def test_dtype_on_gpu(self, mocked_factory, model_kwargs):