updated saf tensors so they match

modeling_lsg_bert.py

@@ -23,7 +23,11 @@ from transformers.models.bert.modeling_bert import (
 )
 import torch
 import torch.nn as nn
-from transformers.modeling_outputs import BaseModelOutputWithPastAndCrossAttentions
+from transformers.modeling_outputs import ( 
+    BaseModelOutputWithPastAndCrossAttentions, 
+    BaseModelOutputWithPoolingAndCrossAttentions
+)
+
 from transformers.models.bert.configuration_bert import BertConfig
 import math
 import sys
@@ -817,7 +821,7 @@ class LSGSelfAttention(BaseSelfAttention):
         n, h, t, d = query_layer.size()
 
         # Cat global mask
-        attention_mask = torch.nn.functional.pad(attention_mask, (self.num_global_tokens, 0), value=0)
+        # attention_mask = torch.nn.functional.pad(attention_mask, (self.num_global_tokens, 0), value=0)
         
         # Use normal attention if local attention covers every tokens
         if t <= 2 * self.block_size + self.num_global_tokens:
@@ -1023,9 +1027,9 @@ class LSGBertModel(LSGBertPreTrainedModel, BertModel):
 
     def __init__(self, config, add_pooling_layer=True):
         
-        # ensure your LSGBertModel inherits all the necessary fields introduced in the latest Transformers.
+        # 1) Initialize the standard BertModel
         BertModel.__init__(self, config)
-        
+        # 2) Initialize our LSG PreTrained
         LSGBertPreTrainedModel.__init__(self, config)
         
         self.config = config
@@ -1041,6 +1045,78 @@ class LSGBertModel(LSGBertPreTrainedModel, BertModel):
         
         # Initialize weights and apply final processing
         self.post_init()
+        
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        past_key_values=None,
+        use_cache=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        # ----------------------------
+        # 1) Use LSG embeddings
+        embedding_output = self.embeddings(
+            input_ids=input_ids,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            inputs_embeds=inputs_embeds,
+            past_key_values_length=past_key_values[0][0].size(2) if past_key_values else 0
+        ) if (input_ids is not None or inputs_embeds is not None) else None
+
+        # 2) If we have an attention mask and some global tokens, pad the mask
+        #    by `config.num_global_tokens` so it matches embedding_output.size(1).
+        if attention_mask is not None and self.config.num_global_tokens > 0:
+            # Original shape: (batch_size, seq_len)
+            bsz, seq_len = attention_mask.shape
+
+            new_shape = (bsz, seq_len + self.config.num_global_tokens)
+            extended_mask = torch.zeros(new_shape, dtype=attention_mask.dtype, device=attention_mask.device)
+            # Fill from index `num_global_tokens` onward
+            extended_mask[:, self.config.num_global_tokens:] = attention_mask
+            attention_mask = extended_mask
+
+        # 3) Now call self.encoder with the updated mask
+        encoder_outputs = self.encoder(
+            hidden_states=embedding_output,
+            attention_mask=attention_mask.unsqueeze(1).unsqueeze(2) if attention_mask is not None else None,
+            head_mask=head_mask,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_attention_mask=encoder_attention_mask,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict
+        )
+        
+        # 4) Grab the last hidden state
+        sequence_output = encoder_outputs[0]
+        
+        # 5) Optionally apply the pooler
+        pooled_output = self.pooler(sequence_output) if self.pooler is not None else None
+
+        # Return
+        if not return_dict:
+            return (sequence_output, pooled_output) + encoder_outputs[1:]
+
+        
+        return BaseModelOutputWithPoolingAndCrossAttentions(
+            last_hidden_state=sequence_output,
+            pooler_output=pooled_output,
+            past_key_values=encoder_outputs.past_key_values,
+            hidden_states=encoder_outputs.hidden_states,
+            attentions=encoder_outputs.attentions,
+            cross_attentions=encoder_outputs.cross_attentions,
+        )
 
     def get_extended_attention_mask(self, attention_mask, input_shape, device=None):