fix: NaN guard in generate_text + _keys_to_ignore_on_load_missing

modeling_quark.py

@@ -129,6 +129,8 @@ class QuarkPreTrainedModel(PreTrainedModel):
 
 
 class QuarkForCausalLM(QuarkPreTrainedModel):
+    # lm_head è tied con embed_tokens — non è missing, è intenzionale
+    _keys_to_ignore_on_load_missing = ["lm_head.weight"]
     """
     Quark autoregressive language model.
     Compatibile con AutoModelForCausalLM.from_pretrained(..., trust_remote_code=True)
@@ -192,17 +194,27 @@ class QuarkForCausalLM(QuarkPreTrainedModel):
         for _ in range(max_new_tokens):
             out    = self(ctx[:, -self.config.max_seq_len:])
             logits = out.logits[0, -1, :].float()
-            if temperature > 0:
-                logits /= temperature
-                probs   = F.softmax(logits, dim=-1)
+            # Greedy se temperature=0 o logits degeneri
+            if temperature <= 0 or torch.isnan(logits).any() or torch.isinf(logits).any():
+                token = logits.argmax().view(1, 1)
+            else:
+                logits = logits / temperature
+                # Stabilizza prima del softmax
+                logits = logits - logits.max()
+                probs  = F.softmax(logits, dim=-1)
+                probs  = torch.clamp(probs, min=0.0)
+                # Top-p nucleus
                 sorted_p, sorted_i = torch.sort(probs, descending=True)
-                cum_p = torch.cumsum(sorted_p, dim=-1)
-                remove = cum_p - sorted_p > top_p
+                cum_p  = torch.cumsum(sorted_p, dim=-1)
+                remove = (cum_p - sorted_p) > top_p
+                sorted_p = sorted_p.clone()
                 sorted_p[remove] = 0.0
-                sorted_p /= sorted_p.sum()
-                token = sorted_i[torch.multinomial(sorted_p, 1)].unsqueeze(0).unsqueeze(0)
-            else:
-                token = logits.argmax().view(1, 1)
+                total = sorted_p.sum()
+                if total <= 0:
+                    token = sorted_i[0].view(1, 1)
+                else:
+                    sorted_p /= total
+                    token = sorted_i[torch.multinomial(sorted_p, 1)].unsqueeze(0).unsqueeze(0)
             ctx = torch.cat([ctx, token], dim=1)
             if eos_token_id is not None and token.item() == eos_token_id:
                 break