257 lines
10 KiB
Python
257 lines
10 KiB
Python
"""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"])
|