Update to follow HF naming scheme

rita_modeling.py

@@ -129,8 +129,8 @@ class SelfAttention(nn.Module):
     def forward(
         self,
         x,
-        attn_mask: Optional[torch.BoolTensor] = None,
-        padding_mask: Optional[torch.BoolTensor] = None,
+        causal_mask: Optional[torch.BoolTensor] = None,
+        attention_mask: Optional[torch.BoolTensor] = None,
     ) -> Tuple[torch.FloatTensor, torch.FloatTensor]:
 
         N, L, D = x.size()  # Batch_size, Context_size, d_model
@@ -153,14 +153,14 @@ class SelfAttention(nn.Module):
         # causal self-attention; Self-attend: (N, nh, L, hs) x (N, nh, hs, L) -> (N, nh, L, L)
         att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
         
-        if attn_mask is not None:
-            att[:,:,-L:, -L: ].masked_fill_(attn_mask.view(1, 1, L, L), float("-inf"))
+        if causal_mask is not None:
+            att[:,:,-L:, -L: ].masked_fill_(causal_mask.view(1, 1, L, L), float("-inf"))
             
         att = (
             att.transpose(0, 2)
-            .masked_fill(padding_mask.view(1, 1, N, L), float("-inf"))
+            .masked_fill(attention_mask.view(1, 1, N, L)==0, float("-inf"))
             .transpose(0, 2)
-            if padding_mask is not None
+            if attention_mask is not None
             else att
         )
         
@@ -197,11 +197,11 @@ class DecoderLayer(nn.Module):
     def forward(
         self,
         x: torch.FloatTensor,
-        attn_mask: torch.BoolTensor,
-        padding_mask: Optional[torch.BoolTensor] = None,
+        causal_mask: torch.BoolTensor,
+        attention_mask: Optional[torch.BoolTensor] = None,
     ) -> torch.FloatTensor:
         y = self.attn_norm(x)
-        y = self.self_attention(y, attn_mask=attn_mask, padding_mask=padding_mask)
+        y = self.self_attention(y, causal_mask=causal_mask, attention_mask=attention_mask)
         x = x + self.attn_dropout(y)
 
         y = self.mlp_norm(x)
@@ -228,27 +228,27 @@ class RITAModel(PreTrainedModel):
         input_ids=None,
         past_key_values=None,  # NOT USED
         attention_mask=None,
+        causal_mask=None,
         token_type_ids=None, # NOT USED
         position_ids=None, # NOT USED
         head_mask=None, # NOT USED
         inputs_embeds=None,
         encoder_hidden_states=None,  # NOT USED
-        encoder_attention_mask=None, # NOT USED
+        encoder_causal_mask=None, # NOT USED
         labels=None,
         use_cache=None, # NOT USED
         output_attentions=None, # NOT USED
         output_hidden_states=None, # NOT USED
         return_dict=None # NOT USED
         ) -> torch.FloatTensor:
-        
         if inputs_embeds == None:
             x = self.embedding(input_ids)  # N x L x D
         else:
             x = inputs_embeds
-        if attention_mask == None:
-            attention_mask = (torch.triu(torch.ones(input_ids.size(1), input_ids.size(1))) == 0).transpose(0, 1).contiguous().to(input_ids.device)
+        if causal_mask == None:
+            causal_mask = (torch.triu(torch.ones(input_ids.size(1), input_ids.size(1))) == 0).transpose(0, 1).contiguous().to(input_ids.device)
         for layer in self.layers:
-            x = layer(x, attn_mask=attention_mask)
+            x = layer(x, causal_mask=causal_mask, attention_mask=attention_mask)
         x = self.final_norm(x)  # N x L x D
 
         return BaseModelOutput(
@@ -295,23 +295,25 @@ class RITAModelForCausalLM(PreTrainedModel):
         input_ids=None,
         past_key_values=None,  # NOT USED
         attention_mask=None,
+        causal_mask=None,
         token_type_ids=None, # NOT USED
         position_ids=None, # NOT USED
         head_mask=None, # NOT USED
         inputs_embeds=None,
         encoder_hidden_states=None,  # NOT USED
-        encoder_attention_mask=None, # NOT USED
+        encoder_causal_mask=None, # NOT USED
         labels=None,
         use_cache=None, # NOT USED
         output_attentions=None, # NOT USED
         output_hidden_states=None, # NOT USED
         return_dict=None # NOT USED
         ) -> torch.FloatTensor:
-        
+
         transformer_outputs = self.transformer(
             input_ids,
             past_key_values=past_key_values,
-            attention_mask=attention_mask,
+            causal_mask=causal_mask,
+            attention_mask = attention_mask,
             token_type_ids=token_type_ids,
             position_ids=position_ids,
             head_mask=head_mask,
@@ -382,6 +384,7 @@ class RITAModelForSequenceClassification(PreTrainedModel):
         input_ids=None,
         past_key_values=None,
         attention_mask=None,
+        causal_mask=None,
         token_type_ids=None,
         position_ids=None,
         head_mask=None,
@@ -404,6 +407,7 @@ class RITAModelForSequenceClassification(PreTrainedModel):
             input_ids,
             past_key_values=past_key_values,
             attention_mask=attention_mask,
+            causal_mask=causal_mask,
             token_type_ids=token_type_ids,
             position_ids=position_ids,
             head_mask=head_mask,