Suport for embedded representation (#3156)

* inputs_embeds * leave feature validation to transformers --------- Co-authored-by: rchivereanu <[email protected]>
UKPLab · Jan 10, 2025 · a7e3707 · a7e3707
1 parent a41aada
commit a7e3707
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/sentence_transformers/models/Transformer.py b/sentence_transformers/models/Transformer.py
@@ -432,9 +432,11 @@ def __repr__(self) -> str:
 
     def forward(self, features: dict[str, torch.Tensor], **kwargs) -> dict[str, torch.Tensor]:
         """Returns token_embeddings, cls_token"""
-        trans_features = {"input_ids": features["input_ids"], "attention_mask": features["attention_mask"]}
-        if "token_type_ids" in features:
-            trans_features["token_type_ids"] = features["token_type_ids"]
+        trans_features = {
+            key: value
+            for key, value in features.items()
+            if key in ["input_ids", "attention_mask", "token_type_ids", "inputs_embeds"]
+        }
 
         output_states = self.auto_model(**trans_features, **kwargs, return_dict=False)
         output_tokens = output_states[0]