diff --git a/tests/test_subtitle_pauses.py b/tests/test_subtitle_pauses.py new file mode 100644 index 000000000..155268fa4 --- /dev/null +++ b/tests/test_subtitle_pauses.py @@ -0,0 +1,23 @@ +import pytest + +from whisper.utils import WriteVTT + + +@pytest.mark.parametrize( + "next_start, expected_text", + [ + (3.5, ["hello again"]), # 2.5 seconds of silence + (4.0, ["hello again"]), # exactly the three-second threshold + (4.1, ["hello", "again"]), # more than three seconds of silence + ], +) +def test_subtitle_break_uses_gap_between_words(next_start, expected_text): + words = [ + {"start": 0.0, "end": 1.0, "word": " hello"}, + {"start": next_start, "end": next_start + 0.5, "word": " again"}, + ] + result = {"segments": [{"start": 0.0, "end": next_start + 0.5, "words": words}]} + + cues = WriteVTT(".").iterate_result(result, max_line_width=40, max_line_count=2) + + assert [text for _, _, text in cues] == expected_text diff --git a/whisper/utils.py b/whisper/utils.py index 79658fc75..bb64a9ca5 100644 --- a/whisper/utils.py +++ b/whisper/utils.py @@ -144,7 +144,7 @@ def iterate_subtitles(): line_count = 1 # the next subtitle to yield (a list of word timings with whitespace) subtitle: List[dict] = [] - last: float = get_start(result["segments"]) or 0.0 + previous_end: Optional[float] = None for segment in result["segments"]: chunk_index = 0 words_count = max_words_per_line @@ -157,7 +157,9 @@ def iterate_subtitles(): ): timing = original_timing.copy() long_pause = ( - not preserve_segments and timing["start"] - last > 3.0 + not preserve_segments + and previous_end is not None + and timing["start"] - previous_end > 3.0 ) has_room = line_len + len(timing["word"]) <= max_line_width seg_break = i == 0 and len(subtitle) > 0 and preserve_segments @@ -188,7 +190,7 @@ def iterate_subtitles(): timing["word"] = "\n" + timing["word"] line_len = len(timing["word"].strip()) subtitle.append(timing) - last = timing["start"] + previous_end = timing["end"] chunk_index += max_words_per_line if len(subtitle) > 0: yield subtitle