fixes

modeling_minitransformer.py

@@ -7,6 +7,7 @@ import torch.nn.functional as F
 
 from transformers import PreTrainedModel, PretrainedConfig
 from .configuration_minitransformer import MiniTransformerConfig
+from transformers.modeling_outputs import CausalLMOutput
 try:
     from flash_attn import flash_attn_func
 except ImportError as e:
@@ -220,16 +221,38 @@ class MiniTransformer(PreTrainedModel):
         self.apply(self._init_weights)
         print("Model Parameter Count: %.2fM\n" % (self._get_num_params() / 1e6,))
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        tok_emb = self.tok_emb(x)
-        x = self.dropout(tok_emb)
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        labels: torch.Tensor = None,
+        **kwargs
+    ) -> CausalLMOutput:
+        # Compute embeddings
+        tok_emb = self.tok_emb(input_ids)
 
         for layer in self.layers:
-            x = layer(x, self.freqs_cis)
-
-        y_hat = self.lm_head(self.norm(x))
-
-        return y_hat
+            tok_emb = layer(tok_emb, self.freqs_cis)
+
+        # Normalize and project to vocabulary
+        tok_emb = self.norm(tok_emb)
+        logits = self.lm_head(tok_emb)
+
+        loss = None
+        if labels is not None:
+            # Shift so that tokens predict the next token
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            loss_fct = nn.CrossEntropyLoss()
+            loss = loss_fct(
+                shift_logits.view(-1, shift_logits.size(-1)),
+                shift_labels.view(-1)
+            )
+
+        return CausalLMOutput(
+            loss=loss,
+            logits=logits,
+        )
+    
 
     def _get_num_params(self):
         n_params = sum(p.numel() for p in self.parameters())