From 8944de03c444728b958c05fff9a169e2ad9fe9ba Mon Sep 17 00:00:00 2001
From: Arthur Camara <camara@zeta-alpha.com>
Date: Mon, 30 Sep 2024 10:59:39 +0200
Subject: [PATCH] Remove unrelated changes

---
 sentence_transformers/SentenceTransformer.py | 2 +-
 sentence_transformers/models/Pooling.py      | 2 +-
 sentence_transformers/models/Transformer.py  | 5 ++---
 3 files changed, 4 insertions(+), 5 deletions(-)

diff --git a/sentence_transformers/SentenceTransformer.py b/sentence_transformers/SentenceTransformer.py
index 0053f4600..429abdd80 100644
--- a/sentence_transformers/SentenceTransformer.py
+++ b/sentence_transformers/SentenceTransformer.py
@@ -658,7 +658,7 @@ def encode(
 
         return all_embeddings
 
-    def forward(self, input: dict[str, Tensor], **kwargs) -> dict[str, Tensor]:
+    def forward(self, input: dict[str, torch.Tensor], **kwargs) -> dict[str, torch.Tensor]:
         if self.module_kwargs is None:
             return super().forward(input)
 
diff --git a/sentence_transformers/models/Pooling.py b/sentence_transformers/models/Pooling.py
index 125b5f573..85ad3f2a4 100644
--- a/sentence_transformers/models/Pooling.py
+++ b/sentence_transformers/models/Pooling.py
@@ -57,7 +57,7 @@ def __init__(
         pooling_mode_mean_sqrt_len_tokens: bool = False,
         pooling_mode_weightedmean_tokens: bool = False,
         pooling_mode_lasttoken: bool = False,
-        include_prompt = True,
+        include_prompt=True,
     ) -> None:
         super().__init__()
 
diff --git a/sentence_transformers/models/Transformer.py b/sentence_transformers/models/Transformer.py
index 73350fd46..2d3786e15 100644
--- a/sentence_transformers/models/Transformer.py
+++ b/sentence_transformers/models/Transformer.py
@@ -116,6 +116,7 @@ def forward(self, features: dict[str, torch.Tensor], **kwargs) -> dict[str, torc
         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"]
+
         output_states = self.auto_model(**trans_features, **kwargs, return_dict=False)
         output_tokens = output_states[0]
 
@@ -135,9 +136,7 @@ def get_word_embedding_dimension(self) -> int:
         return self.auto_model.config.hidden_size
 
     def tokenize(
-        self,
-        texts: list[str] | list[dict] | list[tuple[str, str]],
-        padding: str | bool = True
+        self, texts: list[str] | list[dict] | list[tuple[str, str]], padding: str | bool = True
     ) -> dict[str, torch.Tensor]:
         """Tokenizes a text and maps tokens to token-ids"""
         output = {}