44 lines
1.5 KiB
Python
44 lines
1.5 KiB
Python
import unittest
|
|
|
|
from trainer.reward_utils import rep_penalty
|
|
|
|
|
|
class StubTokenizer:
|
|
def __init__(self, token_ids):
|
|
self.token_ids = token_ids
|
|
self.calls = []
|
|
|
|
def encode(self, text, add_special_tokens=True):
|
|
self.calls.append((text, add_special_tokens))
|
|
return self.token_ids
|
|
|
|
|
|
class RepetitionPenaltyTests(unittest.TestCase):
|
|
def test_repeated_chinese_tokens_receive_penalty(self):
|
|
tokenizer = StubTokenizer([1, 2, 1, 2, 1, 2, 1, 2, 1, 2])
|
|
|
|
self.assertEqual(rep_penalty("你好你好你好你好你好", tokenizer), 0.5)
|
|
self.assertEqual(tokenizer.calls, [("你好你好你好你好你好", False)])
|
|
|
|
def test_repeated_english_tokens_receive_penalty_case_insensitively(self):
|
|
tokenizer = StubTokenizer([3, 3, 3, 3])
|
|
|
|
self.assertEqual(rep_penalty("Hello HELLO hello hello", tokenizer), 0.5)
|
|
self.assertEqual(tokenizer.calls, [("hello hello hello hello", False)])
|
|
|
|
def test_unique_and_short_sequences_have_no_penalty(self):
|
|
self.assertEqual(rep_penalty("unique", StubTokenizer([1, 2, 3, 4, 5])), 0.0)
|
|
self.assertEqual(rep_penalty("short", StubTokenizer([1, 2])), 0.0)
|
|
|
|
def test_invalid_configuration_is_rejected(self):
|
|
tokenizer = StubTokenizer([])
|
|
|
|
with self.assertRaisesRegex(ValueError, "n must be greater than zero"):
|
|
rep_penalty("text", tokenizer, n=0)
|
|
with self.assertRaisesRegex(ValueError, "cap must be non-negative"):
|
|
rep_penalty("text", tokenizer, cap=-0.1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|