pop kwargs

illuin-tech · Aug 21, 2024 · 1497772 · 1497772
1 parent ae13d5a
commit 1497772
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/colpali_engine/models/paligemma_colbert_architecture.py b/colpali_engine/models/paligemma_colbert_architecture.py
@@ -22,6 +22,8 @@ def forward(self, *args, **kwargs):
         Returns:
         - torch.Tensor: Embeddings of shape (batch_size, num_tokens, dim)
         """
+        # delete output_hidden_states from kwargs
+        kwargs.pop("output_hidden_states", None)
         outputs = self.model(*args, output_hidden_states=True, **kwargs)
         last_hidden_states = outputs.hidden_states[-1]  # (batch_size, sequence_length, hidden_size)
         # pooling - last token
@@ -49,6 +51,8 @@ def forward(self, *args, **kwargs):
         Returns:
         - torch.Tensor: Embeddings of shape (batch_size, num_tokens, dim)
         """
+        # delete output_hidden_states from kwargs
+        kwargs.pop("output_hidden_states", None)
         outputs = self.model(*args, output_hidden_states=True, **kwargs)
         last_hidden_states = outputs.hidden_states[-1]  # (batch_size, sequence_length, hidden_size)
         # pooling -mean on attention mask==1
@@ -117,6 +121,8 @@ def forward(self, *args, **kwargs):
             # print(f"Doc proj: {proj.shape}")
             proj = proj / proj.norm(dim=-1, keepdim=True)
         else:
+            # delete output_hidden_states from kwargs
+            kwargs.pop("output_hidden_states", None)
             outputs = self.model(*args, output_hidden_states=True, **kwargs)
             last_hidden_states = outputs.hidden_states[-1]  # (batch_size, sequence_length, hidden_size)
             # print(f"Query: {last_hidden_states.shape}")
@@ -168,6 +174,8 @@ def forward(self, *args, **kwargs):
             norm = proj.norm(dim=-1, keepdim=True)
             proj = proj / norm
         else:
+            # delete output_hidden_states from kwargs
+            kwargs.pop("output_hidden_states", None)
             outputs = self.model(*args, output_hidden_states=True, **kwargs)
             last_hidden_states = outputs.hidden_states[-1]  # (batch_size, sequence_length, hidden_size)
             # pooling -mean on attention mask==1