| 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() |
|
|
