feat: repetition penalty in generate_text

modeling_quark.py

@@ -176,23 +176,37 @@ class QuarkForCausalLM(QuarkPreTrainedModel):
 
     @torch.no_grad()
     def generate_text(self, input_ids, max_new_tokens=200, temperature=0.7,
-                      top_p=0.9, eos_token_id=None):
-        ctx = input_ids.clone()
+                      top_p=0.9, rep_penalty=1.0, eos_token_id=None):
+        ctx       = input_ids.clone()
+        generated = []
         for _ in range(max_new_tokens):
             out    = self(ctx[:, -self.config.max_seq_len:])
             logits = out.logits[0, -1, :].float()
+
+            # Repetition penalty — penalizza token già visti nel contesto+generati
+            if rep_penalty != 1.0:
+                seen = set(ctx[0].tolist() + generated)
+                for tid in seen:
+                    if logits[tid] > 0:
+                        logits[tid] /= rep_penalty
+                    else:
+                        logits[tid] *= rep_penalty
+
             if temperature <= 0 or logits.isnan().any():
-                token = logits.argmax().view(1, 1)
+                token_id = logits.argmax().item()
             else:
-                logits -= logits.max()
-                logits /= temperature
-                probs   = F.softmax(logits, dim=-1)
+                logits = logits - logits.max()
+                logits = logits / temperature
+                probs  = F.softmax(logits, dim=-1)
                 sorted_p, sorted_i = torch.sort(probs, descending=True)
-                cum_p = torch.cumsum(sorted_p, dim=-1)
+                cum_p  = torch.cumsum(sorted_p, dim=-1)
                 sorted_p[(cum_p - sorted_p) > top_p] = 0.0
-                total = sorted_p.sum()
-                token = sorted_i[torch.multinomial(sorted_p / (total if total > 0 else 1), 1)].view(1, 1)
-            ctx = torch.cat([ctx, token], dim=1)
-            if eos_token_id is not None and token.item() == eos_token_id:
+                total  = sorted_p.sum()
+                token_id = sorted_i[torch.multinomial(sorted_p / (total if total > 0 else 1), 1)].item()
+
+            generated.append(token_id)
+            token = torch.tensor([[token_id]], device=ctx.device)
+            ctx   = torch.cat([ctx, token], dim=1)
+            if eos_token_id is not None and token_id == eos_token_id:
                 break
         return ctx