File size: 1,023 Bytes
969573c | 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 | 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()
|