Upload modelling.py

modelling.py

@@ -5,8 +5,7 @@ from torch import nn
 from torch.nn import CrossEntropyLoss
 from transformers import ModernBertModel, ModernBertPreTrainedModel, ModernBertConfig
 from transformers.modeling_outputs import QuestionAnsweringModelOutput
-from transformers.models.modernbert.modeling_modernbert import _pad_modernbert_output, _unpad_modernbert_input, \
-    ModernBertPredictionHead
+from transformers.models.modernbert.modeling_modernbert import ModernBertPredictionHead
 
 
 class ModernBertForQuestionAnswering(ModernBertPreTrainedModel):
@@ -26,10 +25,6 @@ class ModernBertForQuestionAnswering(ModernBertPreTrainedModel):
         # Initialize weights and apply final processing
         self.post_init()
 
-    @torch.compile(dynamic=True)
-    def compiled_head(self, output: torch.Tensor) -> torch.Tensor:
-        return self.head(output)
-
     def forward(
             self,
             input_ids: Optional[torch.Tensor],
@@ -46,6 +41,7 @@ class ModernBertForQuestionAnswering(ModernBertPreTrainedModel):
             output_attentions: Optional[bool] = None,
             output_hidden_states: Optional[bool] = None,
             return_dict: Optional[bool] = None,
+            **kwargs,
     ) -> Union[Tuple[torch.Tensor], QuestionAnsweringModelOutput]:
         r"""
         start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
@@ -60,20 +56,6 @@ class ModernBertForQuestionAnswering(ModernBertPreTrainedModel):
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
         self._maybe_set_compile()
 
-        # Get sequence length and batch size if not provided
-        # if batch_size is None or seq_len is None:
-        #     batch_size, seq_len = input_ids.shape[:2]
-
-        # # Handle Flash Attention 2 unpadding
-        # if self.config._attn_implementation == "flash_attention_2":
-        #     if indices is None and cu_seqlens is None and max_seqlen is None:
-        #         if attention_mask is None:
-        #             attention_mask = torch.ones((batch_size, seq_len), device=input_ids.device, dtype=torch.bool)
-        #         with torch.no_grad():
-        #             input_ids, indices, cu_seqlens, max_seqlen, position_ids, _ = _unpad_modernbert_input(
-        #                 inputs=input_ids, attention_mask=attention_mask, position_ids=position_ids
-        #             )
-
         outputs = self.model(
             input_ids,
             attention_mask=attention_mask,
@@ -90,24 +72,12 @@ class ModernBertForQuestionAnswering(ModernBertPreTrainedModel):
         )
 
         sequence_output = outputs[0]
-        sequence_output = (
-            self.drop(self.compiled_head(sequence_output))
-            if self.config.reference_compile
-            else self.drop(self.head(sequence_output))
-        )
-        # sequence_output = self.drop(self.head(sequence_output))
+        sequence_output = self.drop(self.head(sequence_output))
 
         logits = self.qa_outputs(sequence_output)
         start_logits, end_logits = logits.split(1, dim=-1)
-        start_logits = start_logits.squeeze(-1)
-        end_logits = end_logits.squeeze(-1)
-
-        # # Handle Flash Attention 2 padding
-        # if self.config._attn_implementation == "flash_attention_2":
-        #     start_logits = _pad_modernbert_output(inputs=start_logits, indices=indices, batch=batch_size,
-        #                                           seqlen=seq_len)
-        #     end_logits = _pad_modernbert_output(inputs=end_logits, indices=indices, batch=batch_size,
-        #                                         seqlen=seq_len)
+        start_logits = start_logits.squeeze(-1).contiguous()
+        end_logits = end_logits.squeeze(-1).contiguous()
 
         total_loss = None
         if start_positions is not None and end_positions is not None: