fix: repetition_penalty boosts negative-logit tokens instead of suppressing them

The generate() method divides all previously-seen token logits by
repetition_penalty uniformly. For tokens with negative logits (which
represent less-likely tokens), dividing by penalty > 1.0 moves the
logit toward zero, which *increases* the token's softmax probability.
This is the exact opposite of the intended repetition suppression.

Example with penalty=1.5:
  logit -3.0 → -3.0/1.5 = -2.0 → probability INCREASES (bug)
  logit -3.0 → -3.0*1.5 = -4.5 → probability decreases (correct)

Fix: apply the penalty conditionally based on logit sign:
  - positive logits: divide by penalty (push toward zero)
  - negative logits: multiply by penalty (push away from zero)

Both directions reduce probability for previously-seen tokens, matching
the behavior of Hugging Face Transformers' RepetitionPenaltyLogitsProcessor.

Added 14 tests covering: direction correctness for positive/negative logits,
penalty=1.0 no-op, batch independence, duplicate token handling, large
penalties, edge cases, and a regression test proving the old code boosted
negative-logit token probabilities.

model/model_minimind.py

@@ -259,7 +259,13 @@ class MiniMindForCausalLM(PreTrainedModel, GenerationMixin):
             attention_mask = torch.cat([attention_mask, attention_mask.new_ones(attention_mask.shape[0], 1)], -1) if attention_mask is not None else None
             logits = outputs.logits[:, -1, :] / temperature
             if repetition_penalty != 1.0:
-                for i in range(input_ids.shape[0]): logits[i, torch.unique(input_ids[i])] /= repetition_penalty
+                for i in range(input_ids.shape[0]):
+                    prev_tokens = torch.unique(input_ids[i])
+                    score = logits[i, prev_tokens]
+                    # For tokens with negative logits, multiply by penalty (push further negative);
+                    # for tokens with positive logits, divide by penalty (push toward zero).
+                    # Simply dividing all logits would *boost* negative-logit tokens toward zero.
+                    logits[i, prev_tokens] = torch.where(score < 0, score * repetition_penalty, score / repetition_penalty)
             if top_k > 0: 
                 logits[logits < torch.topk(logits, top_k)[0][..., -1, None]] = -float('inf')
             if top_p < 1.0:

tests/test_repetition_penalty.py

@@ -0,0 +1,256 @@
+"""Tests for the repetition_penalty fix in MiniMindForCausalLM.generate().
+
+The original implementation divides all logits of previously-seen tokens by the
+penalty factor. For tokens with *negative* logits, division by penalty > 1.0
+moves the logit *toward* zero, which **increases** the token's probability —
+the opposite of the intended effect.
+
+The fix applies the penalty conditionally:
+  - positive logits → divide by penalty (push toward zero, lower probability)
+  - negative logits → multiply by penalty (push away from zero, lower probability)
+
+This matches the widely-adopted convention from Hugging Face Transformers:
+  https://github.com/huggingface/transformers/blob/main/src/transformers/generation/logits_process.py
+"""
+
+import math
+import pytest
+import torch
+import torch.nn.functional as F
+
+# ---------------------------------------------------------------------------
+# Helper: apply the FIXED penalty logic to a logits tensor
+# ---------------------------------------------------------------------------
+def apply_penalty_fixed(logits: torch.Tensor, token_ids: torch.Tensor, penalty: float) -> torch.Tensor:
+    """Reference implementation of the correct repetition penalty."""
+    out = logits.clone()
+    prev = torch.unique(token_ids)
+    score = out[prev]
+    out[prev] = torch.where(score < 0, score * penalty, score / penalty)
+    return out
+
+
+def apply_penalty_buggy(logits: torch.Tensor, token_ids: torch.Tensor, penalty: float) -> torch.Tensor:
+    """The old (buggy) implementation for comparison."""
+    out = logits.clone()
+    out[torch.unique(token_ids)] /= penalty
+    return out
+
+
+# ---------------------------------------------------------------------------
+# Tests
+# ---------------------------------------------------------------------------
+
+class TestRepetitionPenaltyDirection:
+    """Core tests: the penalty must ALWAYS reduce probability of seen tokens."""
+
+    def test_positive_logit_reduced(self):
+        """A previously-seen token with positive logit should have LOWER probability."""
+        logits = torch.tensor([3.0, 1.0, 0.0, -1.0, -2.0])
+        seen = torch.tensor([0])  # token 0 has logit 3.0
+        penalty = 1.5
+
+        orig_prob = F.softmax(logits, dim=-1)[0].item()
+        fixed = apply_penalty_fixed(logits, seen, penalty)
+        new_prob = F.softmax(fixed, dim=-1)[0].item()
+
+        assert new_prob < orig_prob, f"Positive logit token prob should decrease: {orig_prob:.4f} → {new_prob:.4f}"
+
+    def test_negative_logit_reduced(self):
+        """A previously-seen token with negative logit should have LOWER probability."""
+        logits = torch.tensor([3.0, 1.0, 0.0, -1.0, -2.0])
+        seen = torch.tensor([4])  # token 4 has logit -2.0
+        penalty = 1.5
+
+        orig_prob = F.softmax(logits, dim=-1)[4].item()
+        fixed = apply_penalty_fixed(logits, seen, penalty)
+        new_prob = F.softmax(fixed, dim=-1)[4].item()
+
+        assert new_prob < orig_prob, f"Negative logit token prob should decrease: {orig_prob:.4f} → {new_prob:.4f}"
+
+    def test_buggy_boosts_negative_logit(self):
+        """Demonstrate that the OLD code incorrectly BOOSTS negative-logit tokens."""
+        logits = torch.tensor([3.0, 1.0, 0.0, -1.0, -2.0])
+        seen = torch.tensor([4])
+        penalty = 1.5
+
+        orig_prob = F.softmax(logits, dim=-1)[4].item()
+        buggy = apply_penalty_buggy(logits, seen, penalty)
+        buggy_prob = F.softmax(buggy, dim=-1)[4].item()
+
+        # The buggy version INCREASES probability — this is the bug
+        assert buggy_prob > orig_prob, f"Buggy code should boost negative logit prob: {orig_prob:.4f} → {buggy_prob:.4f}"
+
+    def test_zero_logit_unchanged(self):
+        """A token with logit exactly 0.0 should remain at 0.0 regardless of penalty."""
+        logits = torch.tensor([2.0, 0.0, -1.0])
+        seen = torch.tensor([1])  # token 1 has logit 0.0
+        penalty = 2.0
+
+        fixed = apply_penalty_fixed(logits, seen, penalty)
+        # 0.0 / penalty = 0.0, and 0.0 * penalty = 0.0, either branch gives 0.0
+        assert fixed[1].item() == 0.0
+
+
+class TestRepetitionPenaltyNoOp:
+    """Penalty = 1.0 should be a no-op."""
+
+    def test_penalty_one_is_noop(self):
+        logits = torch.tensor([2.5, -1.3, 0.0, 4.1, -3.0])
+        seen = torch.tensor([0, 1, 2, 3, 4])
+        fixed = apply_penalty_fixed(logits, seen, 1.0)
+        assert torch.allclose(logits, fixed)
+
+
+class TestRepetitionPenaltyMultipleTokens:
+    """Penalty applied to multiple seen tokens at once."""
+
+    def test_all_seen_token_logits_move_away_from_zero(self):
+        """Every seen token's logit magnitude should increase (move away from zero).
+        
+        Note: softmax probability may not always decrease for every penalized token
+        when multiple tokens are penalized simultaneously (redistribution effect).
+        The correct invariant is that logits move in the right direction."""
+        logits = torch.tensor([5.0, 2.0, -1.0, -3.0, 0.5])
+        seen = torch.tensor([0, 1, 2, 3])  # 4 is unseen
+        penalty = 1.3
+
+        fixed = apply_penalty_fixed(logits, seen, penalty)
+
+        for idx in [0, 1]:  # positive logits should decrease
+            assert fixed[idx] < logits[idx], (
+                f"Token {idx} (logit={logits[idx]:.1f}) should decrease: "
+                f"{logits[idx]:.4f} → {fixed[idx]:.4f}"
+            )
+        for idx in [2, 3]:  # negative logits should become more negative
+            assert fixed[idx] < logits[idx], (
+                f"Token {idx} (logit={logits[idx]:.1f}) should become more negative: "
+                f"{logits[idx]:.4f} → {fixed[idx]:.4f}"
+            )
+
+    def test_unseen_token_prob_increases(self):
+        """Unseen tokens should gain probability share when seen tokens are penalized."""
+        logits = torch.tensor([5.0, 2.0, -1.0, -3.0, 0.5])
+        seen = torch.tensor([0, 1, 2, 3])
+        penalty = 1.3
+
+        orig_prob_unseen = F.softmax(logits, dim=-1)[4].item()
+        fixed = apply_penalty_fixed(logits, seen, penalty)
+        new_prob_unseen = F.softmax(fixed, dim=-1)[4].item()
+
+        assert new_prob_unseen > orig_prob_unseen, (
+            f"Unseen token prob should increase: {orig_prob_unseen:.4f} → {new_prob_unseen:.4f}"
+        )
+
+
+class TestRepetitionPenaltyDuplicateTokens:
+    """Duplicate token IDs in input should not cause issues."""
+
+    def test_duplicates_in_input(self):
+        logits = torch.tensor([1.0, -2.0, 3.0])
+        seen = torch.tensor([1, 1, 1, 0, 0])  # duplicates
+        penalty = 1.5
+
+        fixed = apply_penalty_fixed(logits, seen, penalty)
+        expected = apply_penalty_fixed(logits, torch.tensor([0, 1]), penalty)
+        assert torch.allclose(fixed, expected)
+
+
+class TestRepetitionPenaltyStrength:
+    """Higher penalty should produce a stronger effect."""
+
+    def test_higher_penalty_stronger_reduction(self):
+        logits = torch.tensor([4.0, -2.0, 1.0])
+        seen = torch.tensor([0, 1])
+
+        probs_15 = F.softmax(apply_penalty_fixed(logits, seen, 1.5), dim=-1)
+        probs_30 = F.softmax(apply_penalty_fixed(logits, seen, 3.0), dim=-1)
+
+        # Higher penalty → lower prob for token 0 (positive logit)
+        assert probs_30[0] < probs_15[0]
+        # Higher penalty → lower prob for token 1 (negative logit)
+        assert probs_30[1] < probs_15[1]
+
+
+class TestRepetitionPenaltyBatch:
+    """Test batch behavior matching the generate() loop."""
+
+    def test_batch_independent(self):
+        """Each batch element should be penalized independently."""
+        logits = torch.tensor([
+            [2.0, -1.0, 0.5],
+            [-2.0, 3.0, 1.0],
+        ])
+        input_ids = torch.tensor([
+            [0, 0, 1],
+            [2, 1, 1],
+        ])
+        penalty = 1.5
+
+        result = logits.clone()
+        for i in range(input_ids.shape[0]):
+            prev_tokens = torch.unique(input_ids[i])
+            score = result[i, prev_tokens]
+            result[i, prev_tokens] = torch.where(score < 0, score * penalty, score / penalty)
+
+        # Batch 0: tokens 0,1 seen → logit[0,0]=2.0/1.5, logit[0,1]=-1.0*1.5, logit[0,2] unchanged
+        assert abs(result[0, 0].item() - 2.0 / 1.5) < 1e-5
+        assert abs(result[0, 1].item() - (-1.0 * 1.5)) < 1e-5
+        assert abs(result[0, 2].item() - 0.5) < 1e-5
+
+        # Batch 1: tokens 1,2 seen → logit[1,1]=3.0/1.5, logit[1,2]=1.0/1.5, logit[1,0] unchanged
+        assert abs(result[1, 0].item() - (-2.0)) < 1e-5
+        assert abs(result[1, 1].item() - 3.0 / 1.5) < 1e-5
+        assert abs(result[1, 2].item() - 1.0 / 1.5) < 1e-5
+
+
+class TestRepetitionPenaltyEdgeCases:
+    """Edge cases and boundary conditions."""
+
+    def test_all_negative_logits(self):
+        """When all logits are negative, all seen tokens should still be pushed lower."""
+        logits = torch.tensor([-0.5, -1.0, -2.0, -3.0])
+        seen = torch.tensor([0, 1, 2, 3])
+        penalty = 1.2
+
+        orig_probs = F.softmax(logits, dim=-1)
+        # When ALL tokens are penalized equally, softmax probs change if any logit magnitudes differ
+        fixed = apply_penalty_fixed(logits, seen, penalty)
+
+        # All logits are negative → all get multiplied by penalty → all become more negative
+        for i in range(4):
+            assert fixed[i] < logits[i], f"Token {i} logit should decrease (more negative)"
+
+    def test_single_token_vocabulary(self):
+        """Single-token vocabulary: penalty doesn't change softmax (always 1.0)."""
+        logits = torch.tensor([5.0])
+        seen = torch.tensor([0])
+        penalty = 2.0
+
+        fixed = apply_penalty_fixed(logits, seen, penalty)
+        assert abs(F.softmax(fixed, dim=-1)[0].item() - 1.0) < 1e-6
+
+    def test_large_penalty(self):
+        """Even with very large penalty, probabilities should remain valid."""
+        logits = torch.tensor([3.0, -3.0, 1.0, -1.0, 0.0])
+        seen = torch.arange(5)
+        penalty = 100.0
+
+        fixed = apply_penalty_fixed(logits, seen, penalty)
+        probs = F.softmax(fixed, dim=-1)
+
+        assert torch.all(probs >= 0)
+        assert abs(probs.sum().item() - 1.0) < 1e-5
+
+    def test_empty_seen_tokens(self):
+        """No seen tokens: logits should be unchanged."""
+        logits = torch.tensor([2.0, -1.0, 0.5])
+        seen = torch.tensor([], dtype=torch.long)
+        penalty = 1.5
+
+        fixed = apply_penalty_fixed(logits, seen, penalty)
+        assert torch.allclose(logits, fixed)
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])