@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 108 lines 3.45 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from vllm.model_executor.models.moss_transcribe_diarize import (
5 MossTranscribeDiarizeForConditionalGeneration,
6 )
7
8
9 def test_parse_diarized_transcript_preserves_moss_segments():
10 segments = MossTranscribeDiarizeForConditionalGeneration.parse_diarized_transcript(
11 "[0.48][S01]Welcome[1.66][12.26][S02]Ready[13.81]"
12 )
13
14 assert [
15 (segment.start, segment.end, segment.speaker, segment.text)
16 for segment in segments
17 ] == [
18 (0.48, 1.66, "S01", "Welcome"),
19 (12.26, 13.81, "S02", "Ready"),
20 ]
21
22
23 def test_parse_diarized_transcript_preserves_overlapping_segments():
24 segments = MossTranscribeDiarizeForConditionalGeneration.parse_diarized_transcript(
25 "[0][S01]First speaker[2][1][S02]Second speaker[3]"
26 )
27
28 assert [
29 (segment.start, segment.end, segment.speaker, segment.text)
30 for segment in segments
31 ] == [
32 (0.0, 2.0, "S01", "First speaker"),
33 (1.0, 3.0, "S02", "Second speaker"),
34 ]
35
36
37 def test_parse_diarized_transcript_preserves_numeric_text_markers():
38 segments = MossTranscribeDiarizeForConditionalGeneration.parse_diarized_transcript(
39 "[0][S01]The [2024] report is ready.[4]"
40 )
41
42 assert [segment.text for segment in segments] == ["The [2024] report is ready."]
43
44
45 def test_parse_diarized_transcript_ignores_whitespace_between_segments():
46 segments = MossTranscribeDiarizeForConditionalGeneration.parse_diarized_transcript(
47 "[0][S01]Hello[1]\n [2][S02]Hi[3]"
48 )
49
50 assert [(segment.start, segment.end, segment.text) for segment in segments] == [
51 (0.0, 1.0, "Hello"),
52 (2.0, 3.0, "Hi"),
53 ]
54
55
56 def test_parse_diarized_transcript_ignores_noise_before_a_segment():
57 segments = MossTranscribeDiarizeForConditionalGeneration.parse_diarized_transcript(
58 "noise [bad][0.1][S01]Hello[0.9]"
59 )
60
61 assert [(segment.start, segment.end, segment.text) for segment in segments] == [
62 (0.1, 0.9, "Hello"),
63 ]
64
65
66 def test_parse_diarized_transcript_preserves_timestamps_before_the_end():
67 segments = MossTranscribeDiarizeForConditionalGeneration.parse_diarized_transcript(
68 "[2][S01]The earlier timestamp is [1] not the end[3]"
69 )
70
71 assert [segment.text for segment in segments] == [
72 "The earlier timestamp is [1] not the end",
73 ]
74
75
76 def test_parse_diarized_transcript_skips_empty_segments():
77 segments = MossTranscribeDiarizeForConditionalGeneration.parse_diarized_transcript(
78 "[0][S01][1][2][S02]Complete[3]"
79 )
80
81 assert [(segment.speaker, segment.text) for segment in segments] == [
82 ("S02", "Complete"),
83 ]
84
85
86 def test_parse_diarized_transcript_fails_closed_for_incomplete_output():
87 segments = MossTranscribeDiarizeForConditionalGeneration.parse_diarized_transcript(
88 "[0][S01]Complete[1][2][S02]Incomplete"
89 )
90
91 assert segments == []
92
93
94 def test_parse_diarized_transcript_fails_closed_for_trailing_text():
95 segments = MossTranscribeDiarizeForConditionalGeneration.parse_diarized_transcript(
96 "[0][S01]Complete[1] trailing text"
97 )
98
99 assert segments == []
100
101
102 def test_parse_diarized_transcript_preserves_overlong_timestamp_markers():
103 marker = f"[{'1' * 33}]"
104 segments = MossTranscribeDiarizeForConditionalGeneration.parse_diarized_transcript(
105 f"[0][S01]Value {marker}[1]"
106 )
107
108 assert [segment.text for segment in segments] == [f"Value {marker}"]