| |
|
|
| |
|
|
| import unittest |
|
|
| from models.tt_transformers.tests.decode_test_helpers import decode_step_state, teacher_forced_decode_token |
|
|
|
|
| class DecodeStepStateTests(unittest.TestCase): |
| def test_teacher_forced_prompt_advances_before_sampling(self): |
| states = [decode_step_state(0, iteration, 5, 256) for iteration in range(6)] |
|
|
| self.assertEqual([state[0] for state in states], [1, 2, 3, 4, 5, 6]) |
| self.assertEqual([state[1] for state in states], [1, 2, 3, 4, None, None]) |
| self.assertEqual([state[2] for state in states], [1, 2, 3, 4, 5, 6]) |
|
|
| def test_decode_position_is_rejected_before_exceeding_max_sequence_length(self): |
| self.assertEqual(decode_step_state(254, 0, 5, 256), (255, 1, 1)) |
| self.assertEqual(decode_step_state(254, 1, 5, 256), (256, 2, 2)) |
| with self.assertRaisesRegex(ValueError, "decode position 256"): |
| decode_step_state(254, 2, 5, 256) |
|
|
| def test_reference_token_teacher_forces_a_shared_trajectory(self): |
| reference_token = object() |
| device_token = object() |
|
|
| selected_token = teacher_forced_decode_token( |
| reference_token=reference_token, |
| device_token=device_token, |
| ) |
|
|
| self.assertIs(selected_token, reference_token) |
|
|
| def test_device_token_is_used_without_a_reference_model(self): |
| device_token = object() |
|
|
| self.assertIs(teacher_forced_decode_token(device_token=device_token), device_token) |
|
|
| def test_decode_cannot_continue_without_a_sampled_token(self): |
| with self.assertRaisesRegex(ValueError, "reference or device token"): |
| teacher_forced_decode_token() |
|
|
|
|
| if __name__ == "__main__": |
| unittest.main() |
|
|
