Return correct SparseEmbedding, imports and tests

deepset-ai · anakin87 · Sep 19, 2025 · Jul 3, 2025 · Jul 3, 2025 · Aug 3, 2025
commit 4ddde7844e9ed93082d9df92e0101bded6f29262
@@ -16,6 +16,8 @@
     "openai_text_embedder": ["OpenAITextEmbedder"],
     "sentence_transformers_document_embedder": ["SentenceTransformersDocumentEmbedder"],
     "sentence_transformers_text_embedder": ["SentenceTransformersTextEmbedder"],
+    "sentence_transformers_sparse_document_embedder": ["SentenceTransformersSparseDocumentEmbedder"],
+    "sentence_transformers_sparse_text_embedder": ["SentenceTransformersSparseTextEmbedder"],
 }
 
 if TYPE_CHECKING:
@@ -28,6 +30,12 @@
     from .sentence_transformers_document_embedder import (
         SentenceTransformersDocumentEmbedder as SentenceTransformersDocumentEmbedder,
     )
+    from .sentence_transformers_sparse_document_embedder import (
+        SentenceTransformersSparseDocumentEmbedder as SentenceTransformersSparseDocumentEmbedder,
+    )
+    from .sentence_transformers_sparse_text_embedder import (
+        SentenceTransformersSparseTextEmbedder as SentenceTransformersSparseTextEmbedder,
+    )
     from .sentence_transformers_text_embedder import (
         SentenceTransformersTextEmbedder as SentenceTransformersTextEmbedder,
     )

@@ -164,18 +164,17 @@ def __init__(  # pylint: disable=too-many-positional-arguments
         )
 
     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())
-            )
+        embeddings = self.model.encode(data, **kwargs).coalesce()
+
+        rows, columns = embeddings.indices()
+        values = embeddings.values()
+        batch_size = embeddings.size(0)
+
+        sparse_embeddings: List[SparseEmbedding] = []
+        for embedding in range(batch_size):
+            mask = rows == embedding
+            embedding_columns = columns[mask].tolist()
+            embedding_values = values[mask].tolist()
+            sparse_embeddings.append(SparseEmbedding(indices=embedding_columns, values=embedding_values))
 
         return sparse_embeddings
@@ -186,7 +186,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"]

@@ -5,6 +5,7 @@
 from unittest.mock import patch
 
 import pytest
+import torch
 
 from haystack.components.embedders.backends.sentence_transformers_backend import (
     _SentenceTransformersEmbeddingBackendFactory,
@@ -101,6 +102,10 @@ def test_embedding_function_with_kwargs(mock_sentence_transformer):
 
 @patch("haystack.components.embedders.backends.sentence_transformers_backend.SparseEncoder")
 def test_sparse_embedding_function_with_kwargs(mock_sparse_encoder):
+    indices = torch.tensor([[0, 1], [1, 3]])
+    values = torch.tensor([0.5, 0.7])
+    mock_sparse_encoder.return_value.encode.return_value = torch.sparse_coo_tensor(indices, values, (2, 5))
+
     embedding_backend = _SentenceTransformersSparseEmbeddingBackendFactory.get_embedding_backend(model="model")
 
     data = ["sentence1", "sentence2"]

@@ -14,6 +14,11 @@
 )
 from haystack.utils import ComponentDevice, Secret
 
+from haystack.components.embedders.backends.sentence_transformers_backend import (
+    _SentenceTransformersSparseEmbeddingBackendFactory,
+    _SentenceTransformersSparseEncoderEmbeddingBackend,
+)
+
 
 class TestSentenceTransformersDocumentEmbedder:
     def test_init_default(self):
@@ -210,7 +215,7 @@ def test_from_dict_none_device(self):
         assert component.meta_fields_to_embed == ["meta_field"]
 
     @patch(
-        "haystack.components.embedders.sentence_transformers_document_embedder._SentenceTransformersEmbeddingBackendFactory"
+        "haystack.components.embedders.sentence_transformers_sparse_document_embedder._SentenceTransformersSparseEmbeddingBackendFactory"
     )
     def test_warmup(self, mocked_factory):
         embedder = SentenceTransformersSparseDocumentEmbedder(
@@ -236,7 +241,7 @@ def test_warmup(self, mocked_factory):
         )
 
     @patch(
-        "haystack.components.embedders.sentence_transformers_document_embedder._SentenceTransformersEmbeddingBackendFactory"
+        "haystack.components.embedders.sentence_transformers_sparse_document_embedder._SentenceTransformersSparseEmbeddingBackendFactory"
     )
     def test_warmup_doesnt_reload(self, mocked_factory):
         embedder = SentenceTransformersSparseDocumentEmbedder(model="model")
@@ -322,7 +327,7 @@ def test_prefix_suffix(self):
         )
 
     @patch(
-        "haystack.components.embedders.sentence_transformers_document_embedder._SentenceTransformersEmbeddingBackendFactory"
+        "haystack.components.embedders.sentence_transformers_sparse_document_embedder._SentenceTransformersSparseEmbeddingBackendFactory"
     )
     def test_model_onnx_backend(self, mocked_factory):
         onnx_embedder = SentenceTransformersSparseDocumentEmbedder(
@@ -349,7 +354,7 @@ def test_model_onnx_backend(self, mocked_factory):
         )
 
     @patch(
-        "haystack.components.embedders.sentence_transformers_document_embedder._SentenceTransformersEmbeddingBackendFactory"
+        "haystack.components.embedders.sentence_transformers_sparse_document_embedder._SentenceTransformersSparseEmbeddingBackendFactory"
     )
     def test_model_openvino_backend(self, mocked_factory):
         openvino_embedder = SentenceTransformersSparseDocumentEmbedder(
@@ -376,7 +381,7 @@ def test_model_openvino_backend(self, mocked_factory):
         )
 
     @patch(
-        "haystack.components.embedders.sentence_transformers_document_embedder._SentenceTransformersEmbeddingBackendFactory"
+        "haystack.components.embedders.sentence_transformers_sparse_document_embedder._SentenceTransformersSparseEmbeddingBackendFactory"
     )
     @pytest.mark.parametrize("model_kwargs", [{"torch_dtype": "bfloat16"}, {"torch_dtype": "float16"}])
     def test_dtype_on_gpu(self, mocked_factory, model_kwargs):