| import unittest |
|
|
| import torch |
|
|
| from moss_transcribe_diarize.app.model_runner import generation_progress |
| from moss_transcribe_diarize.inference_utils import ProgressStreamer |
|
|
|
|
| class ProgressStreamerTest(unittest.TestCase): |
| def test_counts_generated_tokens_after_prompt_prefill(self): |
| counts = [] |
| streamer = ProgressStreamer(counts.append) |
|
|
| streamer.put(torch.tensor([[1, 2, 3]])) |
| streamer.put(torch.tensor([4])) |
| streamer.put(torch.tensor([[5, 6]])) |
| streamer.end() |
|
|
| self.assertEqual(counts, [1, 3]) |
| self.assertEqual(streamer.generated_tokens, 3) |
|
|
| def test_generation_progress_mapping(self): |
| self.assertEqual(generation_progress(0, 100), 0.25) |
| self.assertAlmostEqual(generation_progress(50, 100), 0.55) |
| self.assertEqual(generation_progress(100, 100), 0.85) |
| self.assertEqual(generation_progress(150, 100), 0.85) |
| self.assertEqual(generation_progress(10, None), 0.25) |
|
|
|
|
| if __name__ == "__main__": |
| unittest.main() |
|
|
