Update modeling_sora.py

modeling_sora.py

@@ -11,20 +11,17 @@ class SoraForSLM(PreTrainedModel, GenerationMixin):
         super().__init__(config)
         self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size)
         self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
-        
-        # Utilisation de TransformerEncoderLayer, mais nous allons appliquer un masque causal manuellement
+
         self.layers = nn.ModuleList([
             nn.TransformerEncoderLayer(
                 d_model=config.hidden_size,
                 nhead=config.num_heads,
                 dim_feedforward=config.hidden_size * 4,
                 batch_first=True,
-                activation="gelu",
-                norm_first=True
+                activation="gelu"
             ) for _ in range(config.num_layers)
         ])
-        
-        self.ln_f = nn.LayerNorm(config.hidden_size)
+
         self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
         self.post_init()
 
@@ -35,18 +32,17 @@ class SoraForSLM(PreTrainedModel, GenerationMixin):
         return {"input_ids": input_ids}
 
     def forward(self, input_ids=None, attention_mask=None, labels=None, **kwargs):
+        # Calcul des positions
         seq_length = input_ids.size(1)
-        # Création du masque causal (triangulaire)
-        causal_mask = torch.triu(torch.ones(seq_length, seq_length, device=input_ids.device), diagonal=1).bool()
-        
         positions = torch.arange(seq_length, device=input_ids.device).unsqueeze(0)
+        
+        # Embeddings
         x = self.embeddings(input_ids) + self.position_embeddings(positions)
 
+        # Passage dans les couches (sans masque pour éviter tout conflit)
         for layer in self.layers:
-            # Application du masque pour éviter de voir le futur
-            x = layer(x, src_mask=causal_mask)
-            
-        x = self.ln_f(x)
+            x = layer(x)
+
         logits = self.lm_head(x)
 
         loss = None
@@ -57,5 +53,4 @@ class SoraForSLM(PreTrainedModel, GenerationMixin):
             loss_fct = nn.CrossEntropyLoss()
             loss = loss_fct(shift_logits.view(-1, self.config.vocab_size), shift_labels.view(-1))
             
-        return CausalLMOutput(loss=loss, logits=logits)
-        
+        return CausalLMOutput(loss=loss, logits=logits)