feat: Add model and tokenizer kwargs to TransformersSimilarityRanker, SentenceTransformersDocumentEmbedder, SentenceTransformersTextEmbedder

haystack/components/embedders/backends/sentence_transformers_backend.py

@@ -2,7 +2,7 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-from typing import Dict, List, Optional, cast
+from typing import Any, Dict, List, Optional, cast
 
 import numpy as np
 
@@ -27,6 +27,8 @@ def get_embedding_backend(
  auth_token: Optional[Secret] = None,
  trust_remote_code: bool = False,
  truncate_dim: Optional[int] = None,
+ model_kwargs: Optional[Dict[str, Any]] = None,
+ tokenizer_kwargs: Optional[Dict[str, Any]] = None,
  ):
  embedding_backend_id = f"{model}{device}{auth_token}{truncate_dim}"
 
@@ -38,6 +40,8 @@ def get_embedding_backend(
  auth_token=auth_token,
  trust_remote_code=trust_remote_code,
  truncate_dim=truncate_dim,
+ model_kwargs=model_kwargs,
+ tokenizer_kwargs=tokenizer_kwargs,
  )
  _SentenceTransformersEmbeddingBackendFactory._instances[embedding_backend_id] = embedding_backend
  return embedding_backend
@@ -55,6 +59,8 @@ def __init__(
  auth_token: Optional[Secret] = None,
  trust_remote_code: bool = False,
  truncate_dim: Optional[int] = None,
+ model_kwargs: Optional[Dict[str, Any]] = None,
+ tokenizer_kwargs: Optional[Dict[str, Any]] = None,
  ):
  sentence_transformers_import.check()
  self.model = SentenceTransformer(
@@ -63,6 +69,8 @@ def __init__(
  use_auth_token=auth_token.resolve_value() if auth_token else None,
  trust_remote_code=trust_remote_code,
  truncate_dim=truncate_dim,
+ model_kwargs=model_kwargs,
+ tokenizer_kwargs=tokenizer_kwargs,
  )
 
  def embed(self, data: List[str], **kwargs) -> List[List[float]]:

haystack/components/embedders/sentence_transformers_document_embedder.py

@@ -9,6 +9,7 @@
  _SentenceTransformersEmbeddingBackendFactory,
 )
 from haystack.utils import ComponentDevice, Secret, deserialize_secrets_inplace
+from haystack.utils.hf import deserialize_hf_model_kwargs, serialize_hf_model_kwargs
 
 
 @component
@@ -37,7 +38,7 @@ class SentenceTransformersDocumentEmbedder:
  ```
  """
 
- def __init__(
+ def __init__( # noqa: PLR0913
  self,
  model: str = "sentence-transformers/all-mpnet-base-v2",
  device: Optional[ComponentDevice] = None,
@@ -51,6 +52,8 @@ def __init__(
  embedding_separator: str = "\n",
  trust_remote_code: bool = False,
  truncate_dim: Optional[int] = None,
+ model_kwargs: Optional[Dict[str, Any]] = None,
+ tokenizer_kwargs: Optional[Dict[str, Any]] = None,
  ):
  """
  Creates a SentenceTransformersDocumentEmbedder component.
@@ -86,6 +89,12 @@ def __init__(
  The dimension to truncate sentence embeddings to. `None` does no truncation.
  If the model wasn't trained with Matryoshka Representation Learning,
  truncating embeddings can significantly affect performance.
+ :param model_kwargs:
+ Additional keyword arguments for `AutoModelForSequenceClassification.from_pretrained`
+ when loading the model. Refer to specific model documentation for available kwargs.
+ :param tokenizer_kwargs:
+ Additional keyword arguments for `AutoTokenizer.from_pretrained` when loading the tokenizer.
+ Refer to specific model documentation for available kwargs.
  """
 
  self.model = model
@@ -100,6 +109,9 @@ def __init__(
  self.embedding_separator = embedding_separator
  self.trust_remote_code = trust_remote_code
  self.truncate_dim = truncate_dim
+ self.model_kwargs = model_kwargs
+ self.tokenizer_kwargs = tokenizer_kwargs
+ self.embedding_backend = None
 
  def _get_telemetry_data(self) -> Dict[str, Any]:
  """
@@ -114,7 +126,7 @@ def to_dict(self) -> Dict[str, Any]:
  :returns:
  Dictionary with serialized data.
  """
- return default_to_dict(
+ serialization_dict = default_to_dict(
  self,
  model=self.model,
  device=self.device.to_dict(),
@@ -128,7 +140,12 @@ def to_dict(self) -> Dict[str, Any]:
  embedding_separator=self.embedding_separator,
  trust_remote_code=self.trust_remote_code,
  truncate_dim=self.truncate_dim,
+ model_kwargs=self.model_kwargs,
+ tokenizer_kwargs=self.tokenizer_kwargs,
  )
+ if serialization_dict["init_parameters"].get("model_kwargs") is not None:
+ serialize_hf_model_kwargs(serialization_dict["init_parameters"]["model_kwargs"])
+ return serialization_dict
 
  @classmethod
  def from_dict(cls, data: Dict[str, Any]) -> "SentenceTransformersDocumentEmbedder":
@@ -144,19 +161,23 @@ def from_dict(cls, data: Dict[str, Any]) -> "SentenceTransformersDocumentEmbedde
  if init_params.get("device") is not None:
  init_params["device"] = ComponentDevice.from_dict(init_params["device"])
  deserialize_secrets_inplace(init_params, keys=["token"])
+ if init_params.get("model_kwargs") is not None:
+ deserialize_hf_model_kwargs(init_params["model_kwargs"])
  return default_from_dict(cls, data)
 
  def warm_up(self):
  """
  Initializes the component.
  """
- if not hasattr(self, "embedding_backend"):
+ if self.embedding_backend is None:
  self.embedding_backend = _SentenceTransformersEmbeddingBackendFactory.get_embedding_backend(
  model=self.model,
  device=self.device.to_torch_str(),
  auth_token=self.token,
  trust_remote_code=self.trust_remote_code,
  truncate_dim=self.truncate_dim,
+ model_kwargs=self.model_kwargs,
+ tokenizer_kwargs=self.tokenizer_kwargs,
  )
 
  @component.output_types(documents=List[Document])
@@ -176,11 +197,9 @@ def run(self, documents: List[Document]):
  "SentenceTransformersDocumentEmbedder expects a list of Documents as input."
  "In case you want to embed a list of strings, please use the SentenceTransformersTextEmbedder."
  )
- if not hasattr(self, "embedding_backend"):
+ if self.embedding_backend is None:
  raise RuntimeError("The embedding model has not been loaded. Please call warm_up() before running.")
 
- # TODO: once non textual Documents are properly supported, we should also prepare them for embedding here
-
  texts_to_embed = []
  for doc in documents:
  meta_values_to_embed = [

haystack/components/embedders/sentence_transformers_text_embedder.py

@@ -9,6 +9,7 @@
  _SentenceTransformersEmbeddingBackendFactory,
 )
 from haystack.utils import ComponentDevice, Secret, deserialize_secrets_inplace
+from haystack.utils.hf import deserialize_hf_model_kwargs, serialize_hf_model_kwargs
 
 
 @component
@@ -45,6 +46,8 @@ def __init__(
  normalize_embeddings: bool = False,
  trust_remote_code: bool = False,
  truncate_dim: Optional[int] = None,
+ model_kwargs: Optional[Dict[str, Any]] = None,
+ tokenizer_kwargs: Optional[Dict[str, Any]] = None,
  ):
  """
  Create a SentenceTransformersTextEmbedder component.
@@ -76,6 +79,12 @@ def __init__(
  The dimension to truncate sentence embeddings to. `None` does no truncation.
  If the model has not been trained with Matryoshka Representation Learning,
  truncation of embeddings can significantly affect performance.
+ :param model_kwargs:
+ Additional keyword arguments for `AutoModelForSequenceClassification.from_pretrained`
+ when loading the model. Refer to specific model documentation for available kwargs.
+ :param tokenizer_kwargs:
+ Additional keyword arguments for `AutoTokenizer.from_pretrained` when loading the tokenizer.
+ Refer to specific model documentation for available kwargs.
  """
 
  self.model = model
@@ -88,6 +97,9 @@ def __init__(
  self.normalize_embeddings = normalize_embeddings
  self.trust_remote_code = trust_remote_code
  self.truncate_dim = truncate_dim
+ self.model_kwargs = model_kwargs
+ self.tokenizer_kwargs = tokenizer_kwargs
+ self.embedding_backend = None
 
  def _get_telemetry_data(self) -> Dict[str, Any]:
  """
@@ -102,7 +114,7 @@ def to_dict(self) -> Dict[str, Any]:
  :returns:
  Dictionary with serialized data.
  """
- return default_to_dict(
+ serialization_dict = default_to_dict(
  self,
  model=self.model,
  device=self.device.to_dict(),
@@ -114,7 +126,12 @@ def to_dict(self) -> Dict[str, Any]:
  normalize_embeddings=self.normalize_embeddings,
  trust_remote_code=self.trust_remote_code,
  truncate_dim=self.truncate_dim,
+ model_kwargs=self.model_kwargs,
+ tokenizer_kwargs=self.tokenizer_kwargs,
  )
+ if serialization_dict["init_parameters"].get("model_kwargs") is not None:
+ serialize_hf_model_kwargs(serialization_dict["init_parameters"]["model_kwargs"])
+ return serialization_dict
 
  @classmethod
  def from_dict(cls, data: Dict[str, Any]) -> "SentenceTransformersTextEmbedder":
@@ -130,19 +147,23 @@ def from_dict(cls, data: Dict[str, Any]) -> "SentenceTransformersTextEmbedder":
  if init_params.get("device") is not None:
  init_params["device"] = ComponentDevice.from_dict(init_params["device"])
  deserialize_secrets_inplace(init_params, keys=["token"])
+ if init_params.get("model_kwargs") is not None:
+ deserialize_hf_model_kwargs(init_params["model_kwargs"])
  return default_from_dict(cls, data)
 
  def warm_up(self):
  """
  Initializes the component.
  """
- if not hasattr(self, "embedding_backend"):
+ if self.embedding_backend is None:
  self.embedding_backend = _SentenceTransformersEmbeddingBackendFactory.get_embedding_backend(
  model=self.model,
  device=self.device.to_torch_str(),
  auth_token=self.token,
  trust_remote_code=self.trust_remote_code,
  truncate_dim=self.truncate_dim,
+ model_kwargs=self.model_kwargs,
+ tokenizer_kwargs=self.tokenizer_kwargs,
  )
 
  @component.output_types(embedding=List[float])
@@ -162,7 +183,7 @@ def run(self, text: str):
  "SentenceTransformersTextEmbedder expects a string as input."
  "In case you want to embed a list of Documents, please use the SentenceTransformersDocumentEmbedder."
  )
- if not hasattr(self, "embedding_backend"):
+ if self.embedding_backend is None:
  raise RuntimeError("The embedding model has not been loaded. Please call warm_up() before running.")
 
  text_to_embed = self.prefix + text + self.suffix

haystack/components/rankers/transformers_similarity.py

@@ -56,6 +56,7 @@ def __init__(
  calibration_factor: Optional[float] = 1.0,
  score_threshold: Optional[float] = None,
  model_kwargs: Optional[Dict[str, Any]] = None,
+ tokenizer_kwargs: Optional[Dict[str, Any]] = None,
  ):
  """
  Creates an instance of TransformersSimilarityRanker.
@@ -89,6 +90,9 @@ def __init__(
  :param model_kwargs:
  Additional keyword arguments for `AutoModelForSequenceClassification.from_pretrained`
  when loading the model. Refer to specific model documentation for available kwargs.
+ :param tokenizer_kwargs:
+ Additional keyword arguments for `AutoTokenizer.from_pretrained` when loading the tokenizer.
+ Refer to specific model documentation for available kwargs.
 
  :raises ValueError:
  If `top_k` is not > 0.
@@ -112,6 +116,7 @@ def __init__(
 
  model_kwargs = resolve_hf_device_map(device=device, model_kwargs=model_kwargs)
  self.model_kwargs = model_kwargs
+ self.tokenizer_kwargs = tokenizer_kwargs or {}
 
  # Parameter validation
  if self.scale_score and self.calibration_factor is None:
@@ -137,7 +142,9 @@ def warm_up(self):
  self.model_name_or_path, token=self.token.resolve_value() if self.token else None, **self.model_kwargs
  )
  self.tokenizer = AutoTokenizer.from_pretrained(
- self.model_name_or_path, token=self.token.resolve_value() if self.token else None
+ self.model_name_or_path,
+ token=self.token.resolve_value() if self.token else None,
+ **self.tokenizer_kwargs,
  )
  self.device = ComponentDevice.from_multiple(device_map=DeviceMap.from_hf(self.model.hf_device_map))
 
@@ -162,6 +169,7 @@ def to_dict(self) -> Dict[str, Any]:
  calibration_factor=self.calibration_factor,
  score_threshold=self.score_threshold,
  model_kwargs=self.model_kwargs,
+ tokenizer_kwargs=self.tokenizer_kwargs,
  )
 
  serialize_hf_model_kwargs(serialization_dict["init_parameters"]["model_kwargs"])

releasenotes/notes/add-model-and-tokenizer-kwargs-4b7618806665f8ba.yaml

@@ -0,0 +1,5 @@
+---
+enhancements:
+ - |
+ Adds model_kwargs and tokenizer_kwargs to the components TransformersSimilarityRanker, SentenceTransformersDocumentEmbedder, SentenceTransformersTextEmbedder.
+ This allows passing things like model_max_length or torch_dtype for better management of model inference.

test/components/embedders/test_sentence_transformers_document_embedder.py

@@ -5,6 +5,7 @@
 
 import numpy as np
 import pytest
+import torch
 
 from haystack import Document
 from haystack.components.embedders.sentence_transformers_document_embedder import SentenceTransformersDocumentEmbedder
@@ -73,6 +74,8 @@ def test_to_dict(self):
  "meta_fields_to_embed": [],
  "trust_remote_code": False,
  "truncate_dim": None,
+ "model_kwargs": None,
+ "tokenizer_kwargs": None,
  },
  }
 
@@ -90,6 +93,8 @@ def test_to_dict_with_custom_init_parameters(self):
  embedding_separator=" - ",
  trust_remote_code=True,
  truncate_dim=256,
+ model_kwargs={"torch_dtype": torch.float32},
+ tokenizer_kwargs={"model_max_length": 512},
  )
  data = component.to_dict()
 
@@ -108,6 +113,8 @@ def test_to_dict_with_custom_init_parameters(self):
  "trust_remote_code": True,
  "meta_fields_to_embed": ["meta_field"],
  "truncate_dim": 256,
+ "model_kwargs": {"torch_dtype": "torch.float32"},
+ "tokenizer_kwargs": {"model_max_length": 512},
  },
  }
 
@@ -209,7 +216,13 @@ def test_warmup(self, mocked_factory):
  mocked_factory.get_embedding_backend.assert_not_called()
  embedder.warm_up()
  mocked_factory.get_embedding_backend.assert_called_once_with(
- model="model", device="cpu", auth_token=None, trust_remote_code=False, truncate_dim=None
+ model="model",
+ device="cpu",
+ auth_token=None,
+ trust_remote_code=False,
+ truncate_dim=None,
+ model_kwargs=None,
+ tokenizer_kwargs=None,
  )
 
  @patch(

test/components/embedders/test_sentence_transformers_embedding_backend.py

@@ -40,6 +40,8 @@ def test_model_initialization(mock_sentence_transformer):
  use_auth_token="fake-api-token",
  trust_remote_code=True,
  truncate_dim=256,
+ model_kwargs=None,
+ tokenizer_kwargs=None,
  )