Update modeling_sora.py

modeling_sora.py

@@ -1,6 +1,7 @@
 import torch
 import torch.nn as nn
 from transformers import PreTrainedModel, GenerationMixin
+from transformers.modeling_outputs import CausalLMOutput
 from .configuration_sora import SoraConfig
 
 class SoraForSLM(PreTrainedModel, GenerationMixin):
@@ -25,29 +26,39 @@ class SoraForSLM(PreTrainedModel, GenerationMixin):
         self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
         self.post_init()
 
-    # La méthode get_input_embeddings est requise par GenerationMixin
+    # Requis pour que GenerationMixin fonctionne
     def get_input_embeddings(self):
         return self.embeddings
 
-    # La méthode set_input_embeddings est requise par GenerationMixin
     def set_input_embeddings(self, value):
         self.embeddings = value
 
+    def prepare_inputs_for_generation(self, input_ids, **kwargs):
+        # Nécessaire pour la génération en mode auto-régressif
+        return {"input_ids": input_ids}
+
     def forward(self, input_ids=None, labels=None, **kwargs):
         seq_length = input_ids.size(1)
         positions = torch.arange(seq_length, device=input_ids.device).unsqueeze(0)
 
+        # Embedding + Position
         x = self.embeddings(input_ids) + self.position_embeddings(positions)
 
+        # Passage dans les couches
         for layer in self.layers:
             x = layer(x)
 
         logits = self.lm_head(x)
 
+        # Calcul de la perte si les labels sont fournis
         loss = None
         if labels is not None:
             loss_fct = nn.CrossEntropyLoss()
             loss = loss_fct(logits.view(-1, self.config.vocab_size), labels.view(-1))
             
-        # GenerationMixin attend un objet de type CausalLMOutput
-        return {"loss": loss, "logits": logits} if loss is not None else {"logits": logits}
+        # 
+        # Retourner un objet CausalLMOutput au lieu d'un simple dictionnaire
+        return CausalLMOutput(
+            loss=loss,
+            logits=logits
+        )