improved3?

modeling_stldec.py

@@ -12,6 +12,7 @@ class STLPreTrainedModel(PreTrainedModel):
     config_class = STLDecoderConfig
     base_model_prefix = "model"
     def _init_weights(self, module):
+        """Migliorata con Xavier Uniform per evitare gradienti esplosivi nelle fasi iniziali."""
         if isinstance(module, nn.Linear):
             torch.nn.init.xavier_uniform_(module.weight)
             if module.bias is not None:
@@ -83,9 +84,10 @@ class STLDecoderBlock(nn.Module):
         return self.internal_forward(hidden_states, encoder_hidden_states, past_key_value, attention_mask)
 
     def internal_forward(self, hidden_states, encoder_hidden_states=None, past_key_value=None, attention_mask=None):
+        """Modificata in Pre-Norm per garantire la stabilità del gradiente."""
         # 1. Self-Attention
         residual = hidden_states
-        hidden_states = self.ln1(hidden_states) # LN PRIMA dell'operazione
+        hidden_states = self.ln1(hidden_states) # LN PRIMA
         hidden_states, pkv = self.self_attn(hidden_states, past_key_value=past_key_value, attention_mask=attention_mask)
         hidden_states = residual + self.dropout(hidden_states)
 
@@ -107,12 +109,14 @@ class STLDecoderModel(STLPreTrainedModel, GenerationMixin):
     def __init__(self, config):
         super().__init__(config)
         self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, config.pad_token_id)
+        # Posizionali Sinusoidali rimossi dal __init__ perché calcolati dinamicamente nel forward
         self.layers = nn.ModuleList([STLDecoderBlock(config) for _ in range(config.num_hidden_layers)])
         self.norm = nn.LayerNorm(config.hidden_size)
         self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
         self.post_init()
 
     def get_sinusoidal_embeddings(self, seq_len, d_model, device):
+        """Genera posizioni matematiche stabili, evitando errori di indice della tabella fixed."""
         inv_freq = 1.0 / (10000 ** (torch.arange(0, d_model, 2).float() / d_model)).to(device)
         pos = torch.arange(seq_len, device=device).type_as(inv_freq)
         sin_inp = torch.einsum("i,j->ij", pos, inv_freq)
@@ -140,9 +144,11 @@ class STLDecoderModel(STLPreTrainedModel, GenerationMixin):
 
         hidden_states = self.embed_tokens(input_ids)
         
+        # Sostituzione con sinusoidali (più robusti dello Script 2)
         pos_emb = self.get_sinusoidal_embeddings(seq_len, self.config.hidden_size, input_ids.device)
         hidden_states = hidden_states + pos_emb[:, :seq_len, :]
 
+        # Maschera causale ottimizzata
         causal_mask = torch.full((seq_len, seq_len + past_len), float("-inf"), device=input_ids.device, dtype=hidden_states.dtype)
         causal_mask.triu_(diagonal=past_len + 1)
         causal_mask = causal_mask[None, None, :, :]
@@ -161,11 +167,14 @@ class STLDecoderModel(STLPreTrainedModel, GenerationMixin):
             shift_logits = logits[..., :-1, :].contiguous()
             shift_labels = labels[..., 1:].contiguous()
             
-            vocab_size = logits.size(-1)
+            # --- MODIFICA DI SICUREZZA ---
+            # Prendiamo il vocab_size dai logits correnti, non dal config fisso,
+            # per evitare l'Assertion t < n_classes se hai fatto un resize_token_embeddings.
+            current_vocab_size = logits.size(-1)
             loss = F.cross_entropy(
-                shift_logits.view(-1, vocab_size), 
+                shift_logits.view(-1, current_vocab_size), 
                 shift_labels.view(-1),
-                ignore_index=self.config.pad_token_id
+                ignore_index=-100 # Standard HF per ignorare padding trasformato
             )
 
         if not return_dict: