add model

rita_modeling.py

@@ -13,6 +13,7 @@ from transformers.modeling_outputs import (
     BaseModelOutputWithPastAndCrossAttentions,
     CausalLMOutputWithCrossAttentions,
     CausalLMOutputWithPast,
+    CausalLMOutput,
 )
 
 from transformers.modeling_utils import PreTrainedModel
@@ -222,18 +223,50 @@ class RITAModel(PreTrainedModel):
         self.final_norm = nn.LayerNorm(config.d_model)
         self.projector = nn.Linear(config.d_model, config.vocab_size, bias = False)
 
-    def forward(self, input_ids, attn_mask=None, padding_mask=None, return_hidden=False) -> torch.FloatTensor:
-        x = self.embedding(input_ids)  # N x L x D
-        if attn_mask == None:
-            attn_mask = (torch.triu(torch.ones(input_ids.size(1), input_ids.size(1))) == 0).transpose(0, 1).contiguous().to(input_ids.device)
+    def forward(
+        self,
+        input_ids=None,
+        past_key_values=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,
+        labels=None,
+        use_cache=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None) -> 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)
         for layer in self.layers:
-            x = layer(x, attn_mask=attn_mask, padding_mask=padding_mask)
+            x = layer(x, attn_mask=attention_mask)
         x = self.final_norm(x)  # N x L x D
 
-        if return_hidden:
-            return x
-        else:
-            return self.projector(x)
+        logits = self.projector(x)
+        loss = None
+        if labels is not None:
+            # Shift so that tokens < n predict n
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            # Flatten the tokens
+            loss_fct = CrossEntropyLoss()
+            loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
+
+        return CausalLMOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=x,
+        )
+
 
     #Some common HF functions.
     def get_input_embeddings(self):