Training in progress, step 23750

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:a6bd3ebd29b873e4e06d10a6b58a955c0e341adb2cbfb4fb8256a6ead9f81b3c
+oid sha256:44b55f92fa7fc67b6a482ca1c47995642b88f2f2032baa4be347f71af6fb62c0
 size 611079728

modeling_stlenc.py

@@ -10,40 +10,43 @@ class STLEncoderModel(PreTrainedModel):
         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)
-
+        
         encoder_layer = nn.TransformerEncoderLayer(
             d_model=config.hidden_size,
             nhead=config.num_attention_heads,
             dim_feedforward=config.intermediate_size,
-            activation="gelu",
+            activation="gelu", # GELU è standard per i Transformer moderni
             batch_first=True
         )
         self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=config.num_hidden_layers)
-
-        # 1024 -> 512 (Bottleneck) -> Target (1024 o 512)
+        
+        # --- POTENZIAMENTO ARCHITETTURALE ---
+        # Creiamo una testa di proiezione profonda (MLP)
         self.projector = nn.Sequential(
+            nn.Linear(config.hidden_size, config.hidden_size),
+            nn.GELU(),
+            nn.LayerNorm(config.hidden_size),
+            nn.Dropout(0.1), # Aiuta a non overfittare sulle costanti numeriche
             nn.Linear(config.hidden_size, config.hidden_size // 2),
             nn.GELU(),
-            nn.LayerNorm(config.hidden_size // 2),
             nn.Linear(config.hidden_size // 2, config.embedding_dim_target)
         )
+        # ------------------------------------
+        
         self.post_init()
 
     def forward(self, input_ids, attention_mask=None, **kwargs):
         batch_size, seq_length = input_ids.size()
-        position_ids = torch.arange(seq_length, device=input_ids.device).unsqueeze(0).expand(batch_size, seq_length)
-        
+        position_ids = torch.arange(seq_length, dtype=torch.long, device=input_ids.device)
+        position_ids = position_ids.unsqueeze(0).expand(batch_size, seq_length)
         x = self.embeddings(input_ids) + self.position_embeddings(position_ids)
+        # Maschera per il padding (TransformerEncoder si aspetta True dove NON deve guardare)
         padding_mask = (attention_mask == 0) if attention_mask is not None else None
-        
+        # Encoding delle sequenze
         sequence_output = self.encoder(x, src_key_padding_mask=padding_mask)
-
-        if attention_mask is not None:
-            mask = attention_mask.unsqueeze(-1).expand(sequence_output.size()).float()
-            sum_embeddings = torch.sum(sequence_output * mask, 1)
-            sum_mask = torch.clamp(mask.sum(1), min=1e-9)
-            emb = sum_embeddings / sum_mask
-        else:
-            emb = sequence_output.mean(dim=1)
-
-        return self.projector(emb)
+        # Prendiamo il CLS (indice 0)
+        cls_token = sequence_output[:, 0, :]
+        # Passiamo per la testa di proiezione non-lineare
+        # Rimuoviamo la Tanh finale per lasciare che il kernel scalare respiri
+        pooled_output = self.projector(cls_token)
+        return pooled_output