small fix with torch.finfo

modeling_lsg_albert.py

@@ -198,7 +198,7 @@ class CausalAttentionProduct(nn.Module):
                 diagonal=-1
                 ) 
             causal_mask = causal_mask.T * torch.finfo(attention_scores.dtype).min
-            attention_scores[..., -causal_shape[0]:, -causal_shape[1]:] = causal_mask
+            attention_scores[..., -causal_shape[0]:, -causal_shape[1] + 1:] = causal_mask[:, 1:]
 
             del attention_mask
 
@@ -551,7 +551,8 @@ class LSGAttention(BaseSelfAttention):
         keys = keys.sum(dim=-2) / (mask + 1e-6)
         values = values.sum(dim=-2) / (mask + 1e-6)
 
-        mask = (1. - mask.clamp(0, 1)) * torch.finfo(mask.dtype).min
+        mask = (1. - mask.clamp(0, 1)) 
+        mask *= torch.finfo(mask.dtype).min
         return keys.reshape(n, h, -1, d), values.reshape(n, h, -1, d), mask.expand(-1, h, -1, -1).transpose(-1, -2)
 
     def get_sparse_tokens_with_stride(self, keys, values, mask):
@@ -616,7 +617,8 @@ class LSGAttention(BaseSelfAttention):
         keys /= mask + 1e-8
         values /= mask + 1e-8
 
-        mask = (1. - mask.clamp(0, 1)) * torch.finfo(mask.dtype).min
+        mask = (1. - mask.clamp(0, 1)) 
+        mask *= torch.finfo(mask.dtype).min
 
         return keys.reshape(n, h, -1, d), values.reshape(n, h, -1, d), mask.transpose(-1, -2).reshape(n, h, 1, -1)
 
@@ -769,6 +771,63 @@ class LSGAlbertTransformer(AlbertTransformer):
 
         self.albert_layer_groups = nn.ModuleList([LSGAlbertLayerGroup(config) for _ in range(config.num_hidden_groups)])
 
+        assert hasattr(config, "num_global_tokens")
+        self.num_global_tokens = config.num_global_tokens
+        self.pad_idx = config.pad_token_id
+
+        assert hasattr(config, "block_size") and hasattr(config, "adaptive")
+        self.block_size = config.block_size
+        self.adaptive = config.adaptive
+        self.mask_first_token = config.mask_first_token
+        self.pool_with_global = config.pool_with_global
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        head_mask: Optional[torch.FloatTensor] = None,
+        output_attentions: bool = False,
+        output_hidden_states: bool = False,
+        return_dict: bool = True,
+    ) -> Union[BaseModelOutput, Tuple]:
+
+        mask_value = torch.finfo(attention_mask.dtype).min
+        n, _, __, t = attention_mask.size()
+        
+        if self.mask_first_token:
+            attention_mask[..., 0] = mask_value
+        
+        b = self.block_size * 2
+        pad = t % self.block_size
+        
+        # Check if t is multiple of block_size and pad
+        if self.adaptive and t > b and pad > 0:
+            pad_length = self.block_size - pad
+            hidden_states = torch.nn.functional.pad(hidden_states.transpose(-1, -2), (0, pad_length), value=0.).transpose(-1, -2)
+            attention_mask = torch.nn.functional.pad(attention_mask, (0, pad_length), value=mask_value)
+
+        encoder_outputs = super().forward(
+            hidden_states=hidden_states, 
+            attention_mask=attention_mask, 
+            head_mask=head_mask, 
+            output_attentions=output_attentions, 
+            output_hidden_states=output_hidden_states, 
+            return_dict=return_dict
+            )
+
+        sequence_output = encoder_outputs[0]
+        if self.pool_with_global:
+            sequence_output[:, self.num_global_tokens] = sequence_output[:, 0]
+
+        # Adapt sequence to initial shape
+        sequence_output = sequence_output[..., self.num_global_tokens: t + self.num_global_tokens, :]
+
+        if not return_dict:
+            return (sequence_output, ) + encoder_outputs[1:]
+        
+        encoder_outputs.last_hidden_state = sequence_output 
+        return encoder_outputs
+
 
 class LSGAlbertPreTrainedModel(PreTrainedModel):
     """
@@ -806,16 +865,6 @@ class LSGAlbertModel(LSGAlbertPreTrainedModel, AlbertModel):
     def __init__(self, config, add_pooling_layer=True):
         AlbertPreTrainedModel.__init__(self, config)
 
-        assert hasattr(config, "num_global_tokens")
-        self.num_global_tokens = config.num_global_tokens
-        self.pad_idx = config.pad_token_id
-
-        assert hasattr(config, "block_size") and hasattr(config, "adaptive")
-        self.block_size = config.block_size
-        self.adaptive = config.adaptive
-        self.mask_first_token = config.mask_first_token
-        self.pool_with_global = config.pool_with_global
-
         self.config = config
         self.embeddings = LSGAlbertEmbeddings(config)
         self.encoder = LSGAlbertTransformer(config)
@@ -829,87 +878,7 @@ class LSGAlbertModel(LSGAlbertPreTrainedModel, AlbertModel):
         # 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,
-        output_attentions=None,
-        output_hidden_states=None,
-        return_dict=None,
-        ):
-
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_hidden_states = (
-            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
-        )
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-        
-        inputs_ = input_ids if input_ids is not None else inputs_embeds
-        n, t = inputs_.size()[:2]
-
-        if attention_mask is None:
-            attention_mask = torch.ones(n, t, device=inputs_.device, dtype=inputs_.dtype)
-        if self.mask_first_token:
-            attention_mask[:,0] = 0
-            
-        b = self.block_size * 2
-        pad = t % self.block_size
-        
-        # Check if t is multiple of block_size and pad
-        if self.adaptive and t > b and pad > 0:
-            pad_length = self.block_size - pad
-            if input_ids is not None:
-                input_ids = torch.nn.functional.pad(input_ids, (0, pad_length), value=self.pad_idx)
-            else:
-                inputs_embeds = torch.nn.functional.pad(inputs_embeds.transpose(-1, -2), (0, pad_length), value=0.).transpose(-1, -2)
-
-            attention_mask = torch.nn.functional.pad(attention_mask, (0, pad_length), value=0)
-
-            if token_type_ids is not None:
-                token_type_ids = torch.nn.functional.pad(token_type_ids, (0, pad_length), value=0)
-            if position_ids is not None:
-                position_ids = torch.nn.functional.pad(position_ids, (0, pad_length), value=0)
-        
-        n, t_ = attention_mask.size()
-
-        encoder_outputs = super().forward(
-            input_ids=input_ids, 
-            attention_mask=attention_mask, 
-            token_type_ids=token_type_ids, 
-            position_ids=position_ids, 
-            head_mask=head_mask, 
-            inputs_embeds=inputs_embeds, 
-            output_attentions=output_attentions, 
-            output_hidden_states=output_hidden_states, 
-            return_dict=return_dict
-            )
-
-        sequence_output = encoder_outputs[0]
-        if self.pool_with_global:
-            sequence_output[:, self.num_global_tokens] = sequence_output[:, 0]
-        
-        diff = t - t_
-        n, _, d = sequence_output.size()
-        sequence_output = sequence_output[..., self.num_global_tokens:, :]
-
-        # Adapt sequence to initial shape
-        if diff < 0:
-            sequence_output = sequence_output[:, :t]
-        
-        pooled_output = self.pooler(sequence_output) if self.pooler is not None else None
-
-        if not return_dict:
-            return (sequence_output, pooled_output) + encoder_outputs[1:]
-        
-        encoder_outputs.last_hidden_state = sequence_output 
-        encoder_outputs.pooler_output = pooled_output
-        return encoder_outputs
-
-
+    
 class LSGAlbertForPreTraining(LSGAlbertPreTrainedModel, AlbertForPreTraining):
 
     def __init__(self, config):