File size: 1,784 Bytes
babe29f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 | # SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
# SPDX-License-Identifier: Apache-2.0
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()
|