File size: 1,217 Bytes
9040f4d | 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 | import unittest
import numpy as np
import torch
from wannabe_tokenizer.models.pipeline import chunk_ranges
from wannabe_tokenizer.quantizer import CVQHead, to_token_clock
class TokenClockTests(unittest.TestCase):
def test_offset_delays_and_repeats_edges(self):
values = np.arange(5)
np.testing.assert_array_equal(to_token_clock(values, 1), [0, 0, 1, 2, 3])
np.testing.assert_array_equal(to_token_clock(values, 0), values)
def test_context_chunks_cover_every_frame_once(self):
ranges = chunk_ranges(100, prefix_tokens=24, context=100)
self.assertEqual(ranges[0][0], 0)
self.assertEqual(ranges[-1][1], 100)
self.assertEqual(sum(end - start for start, end in ranges), 100)
for (_, previous_end), (next_start, _) in zip(ranges, ranges[1:]):
self.assertEqual(previous_end, next_start)
def test_cvq_head_output_shape(self):
torch.manual_seed(0)
head = CVQHead(width=16, dim=4, hidden=8, tau=0.07).eval()
output = head(torch.randn(2, 3, 16))
self.assertEqual(tuple(output.shape), (2, 3, 32768))
self.assertTrue(torch.isfinite(output).all())
if __name__ == "__main__":
unittest.main()
|