Upload 2 files

configuration_sora.py

@@ -0,0 +1,21 @@
+from transformers import PretrainedConfig
+
+class SoraConfig(PretrainedConfig):
+    model_type = "sora_slm"
+
+    def __init__(
+        self,
+        vocab_size=2628,
+        hidden_size=512,
+        num_layers=8,
+        num_heads=8,
+        max_position_embeddings=512,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.num_layers = num_layers
+        self.num_heads = num_heads
+        self.max_position_embeddings = max_position_embeddings
+        

modeling_sora.py

@@ -0,0 +1,56 @@
+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):
+    config_class = SoraConfig
+
+    def __init__(self, config):
+        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)
+
+        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"
+            ) for _ in range(config.num_layers)
+        ])
+
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.embeddings
+
+    def prepare_inputs_for_generation(self, input_ids, **kwargs):
+        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)
+        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:
+            x = layer(x)
+
+        logits = self.lm_head(x)
+
+        loss = None
+        if labels is not None:
+            # Shift pour l'entraînement causal
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = input_ids[..., 1:].contiguous()
+            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)