feat: SentenceTransformersDocumentEmbedder and SentenceTransformersTextEmbedder can accept and pass any arguments to SentenceTransformer.encode

haystack/components/embedders/sentence_transformers_document_embedder.py

@@ -55,6 +55,7 @@ def __init__(  # noqa: PLR0913 # pylint: disable=too-many-positional-arguments
         model_kwargs: Optional[Dict[str, Any]] = None,
         tokenizer_kwargs: Optional[Dict[str, Any]] = None,
         config_kwargs: Optional[Dict[str, Any]] = None,
+        encode_kwargs: Optional[Dict[str, Any]] = None,
         precision: Literal["float32", "int8", "uint8", "binary", "ubinary"] = "float32",
     ):
         """
@@ -99,6 +100,8 @@ def __init__(  # noqa: PLR0913 # pylint: disable=too-many-positional-arguments
             Refer to specific model documentation for available kwargs.
         :param config_kwargs:
             Additional keyword arguments for `AutoConfig.from_pretrained` when loading the model configuration.
+        :param encode_kwargs:
+            Additional keyword arguments for `SentenceTransformer.encode` when embedding documents.
         :param precision:
             The precision to use for the embeddings.
             All non-float32 precisions are quantized embeddings.
@@ -121,6 +124,7 @@ def __init__(  # noqa: PLR0913 # pylint: disable=too-many-positional-arguments
         self.model_kwargs = model_kwargs
         self.tokenizer_kwargs = tokenizer_kwargs
         self.config_kwargs = config_kwargs
+        self.encode_kwargs = encode_kwargs
         self.embedding_backend = None
         self.precision = precision
 
@@ -154,6 +158,7 @@ def to_dict(self) -> Dict[str, Any]:
             model_kwargs=self.model_kwargs,
             tokenizer_kwargs=self.tokenizer_kwargs,
             config_kwargs=self.config_kwargs,
+            encode_kwargs=self.encode_kwargs,
             precision=self.precision,
         )
         if serialization_dict["init_parameters"].get("model_kwargs") is not None:
@@ -232,6 +237,7 @@ def run(self, documents: List[Document]):
             show_progress_bar=self.progress_bar,
             normalize_embeddings=self.normalize_embeddings,
             precision=self.precision,
+            **(self.encode_kwargs if self.encode_kwargs else {}),
         )
 
         for doc, emb in zip(documents, embeddings):

haystack/components/embedders/sentence_transformers_text_embedder.py

@@ -49,6 +49,7 @@ def __init__(  # noqa: PLR0913 # pylint: disable=too-many-positional-arguments
         model_kwargs: Optional[Dict[str, Any]] = None,
         tokenizer_kwargs: Optional[Dict[str, Any]] = None,
         config_kwargs: Optional[Dict[str, Any]] = None,
+        encode_kwargs: Optional[Dict[str, Any]] = None,
         precision: Literal["float32", "int8", "uint8", "binary", "ubinary"] = "float32",
     ):
         """
@@ -89,6 +90,8 @@ def __init__(  # noqa: PLR0913 # pylint: disable=too-many-positional-arguments
             Refer to specific model documentation for available kwargs.
         :param config_kwargs:
             Additional keyword arguments for `AutoConfig.from_pretrained` when loading the model configuration.
+        :param encode_kwargs:
+            Additional keyword arguments for `SentenceTransformer.encode` when embedding texts.
         :param precision:
             The precision to use for the embeddings.
             All non-float32 precisions are quantized embeddings.
@@ -109,6 +112,7 @@ def __init__(  # noqa: PLR0913 # pylint: disable=too-many-positional-arguments
         self.model_kwargs = model_kwargs
         self.tokenizer_kwargs = tokenizer_kwargs
         self.config_kwargs = config_kwargs
+        self.encode_kwargs = encode_kwargs
         self.embedding_backend = None
         self.precision = precision
 
@@ -140,6 +144,7 @@ def to_dict(self) -> Dict[str, Any]:
             model_kwargs=self.model_kwargs,
             tokenizer_kwargs=self.tokenizer_kwargs,
             config_kwargs=self.config_kwargs,
+            encode_kwargs=self.encode_kwargs,
             precision=self.precision,
         )
         if serialization_dict["init_parameters"].get("model_kwargs") is not None:
@@ -209,5 +214,6 @@ def run(self, text: str):
             show_progress_bar=self.progress_bar,
             normalize_embeddings=self.normalize_embeddings,
             precision=self.precision,
+            **(self.encode_kwargs if self.encode_kwargs else {}),
         )[0]
         return {"embedding": embedding}

releasenotes/notes/add-encode-kwargs-sentence-transformers-f4d885f6c5b1706f.yaml

@@ -0,0 +1,6 @@
+---
+enhancements:
+  - |
+    Enhanced `SentenceTransformersDocumentEmbedder` and `SentenceTransformersTextEmbedder` to accept
+    an additional parameter, which is passed directly to the underlying `SentenceTransformer.encode` method
+    for greater flexibility in embedding customization.

test/components/embedders/test_sentence_transformers_document_embedder.py

@@ -79,6 +79,7 @@ def test_to_dict(self):
                 "truncate_dim": None,
                 "model_kwargs": None,
                 "tokenizer_kwargs": None,
+                "encode_kwargs": None,
                 "config_kwargs": None,
                 "precision": "float32",
             },
@@ -101,6 +102,7 @@ def test_to_dict_with_custom_init_parameters(self):
             model_kwargs={"torch_dtype": torch.float32},
             tokenizer_kwargs={"model_max_length": 512},
             config_kwargs={"use_memory_efficient_attention": True},
+            encode_kwargs={"task": "clustering"},
             precision="int8",
         )
         data = component.to_dict()
@@ -123,6 +125,7 @@ def test_to_dict_with_custom_init_parameters(self):
                 "model_kwargs": {"torch_dtype": "torch.float32"},
                 "tokenizer_kwargs": {"model_max_length": 512},
                 "config_kwargs": {"use_memory_efficient_attention": True},
+                "encode_kwargs": {"task": "clustering"},
                 "precision": "int8",
             },
         }

test/components/embedders/test_sentence_transformers_text_embedder.py

@@ -70,6 +70,7 @@ def test_to_dict(self):
                 "truncate_dim": None,
                 "model_kwargs": None,
                 "tokenizer_kwargs": None,
+                "encode_kwargs": None,
                 "config_kwargs": None,
                 "precision": "float32",
             },
@@ -90,6 +91,7 @@ def test_to_dict_with_custom_init_parameters(self):
             model_kwargs={"torch_dtype": torch.float32},
             tokenizer_kwargs={"model_max_length": 512},
             config_kwargs={"use_memory_efficient_attention": False},
+            encode_kwargs={"task": "clustering"},
             precision="int8",
         )
         data = component.to_dict()
@@ -109,6 +111,7 @@ def test_to_dict_with_custom_init_parameters(self):
                 "model_kwargs": {"torch_dtype": "torch.float32"},
                 "tokenizer_kwargs": {"model_max_length": 512},
                 "config_kwargs": {"use_memory_efficient_attention": False},
+                "encode_kwargs": {"task": "clustering"},
                 "precision": "int8",
             },
         }