Merge pull request #84 from yusufipk/word-timestamps-for-models-without-segments

Build subtitle cues out of word times where a model marks no segments

dikte/api.py

@@ -384,8 +384,8 @@ def _transcribe_request(target, audio_path, language, prompt, response_format,
     # takes it as the initial prompt, the way OpenAI does.
     if prompt and target.provider != "openrouter":
         fields.append(("prompt", prompt))
-    if granularity:
+    for level in granularity or ():
-        fields.append(("timestamp_granularities[]", granularity))
+        fields.append(("timestamp_granularities[]", level))
     body, ctype = _multipart(fields, "file", audio_path)
     # An hour of meeting takes the local server a while, and the idle unload has
     # to count that as the model being used rather than as nobody wanting it.
@@ -443,6 +443,96 @@ def _merge_word_splits(segments):
     return merged
 
 
+# A cue built here is one a reader has time for: about two lines of subtitle,
+# and no longer on screen than a sentence takes to say. Neither is a hard rule
+# for a sentence that ends early, only the point past which one is broken.
+MAX_CUE_SECONDS = 7.0
+MAX_CUE_CHARS = 84
+# The other end of it: a cue nobody can read because it was gone before they
+# looked. A full stop this early in a cue is not the end of anything worth
+# breaking on, which is what "1." and "Dr." are, and a cue that ends up short
+# anyway is held on screen until the next one needs the space.
+MIN_CUE_SECONDS = 1.2
+# No whisper segment is longer than the window it was heard in, so a segment
+# that runs past this came from a model that is not marking segments at all.
+WHISPER_WINDOW = 30.0
+SENTENCE_END = ".!?…"
+
+
+def _too_coarse(segments):
+    """Whether these segments are too long to be cues, or are not there at all.
+
+    Not every model behind /audio/transcriptions marks segments the way whisper
+    does. Some fill the field with one entry per paragraph, or with a single one
+    covering the whole file, which turns a fourteen minute video into three
+    subtitles. Word times are what those models do give, and cues built from
+    them are better than what the segments would have been.
+    """
+    if not segments:
+        return True
+    return any(float(seg.get("end") or 0.0) - float(seg.get("start") or 0.0)
+               > WHISPER_WINDOW for seg in segments)
+
+
+def cues_from_words(words):
+    """[(start, end, text)] cut out of word times, where segments were no use.
+
+    A cue ends where a sentence does, and failing that wherever it has grown too
+    long to read or too long to leave up. Nothing is ever cut between two words:
+    the times that arrive are per word, and so are the ones that leave.
+    """
+    cues = []
+    start = end = 0.0
+    current = []
+
+    def flush():
+        nonlocal current
+        if current:
+            cues.append((start, max(end, start), " ".join(current)))
+            current = []
+
+    for word in words:
+        text = (word.get("word") or "").strip()
+        if not text:
+            continue
+        at = float(word.get("start") or 0.0)
+        until = float(word.get("end") or at)
+        if current:
+            grown = len(" ".join(current)) + 1 + len(text)
+            if grown > MAX_CUE_CHARS or until - start > MAX_CUE_SECONDS:
+                flush()
+        if not current:
+            start = at
+        current.append(text)
+        end = until
+        # A sentence can end inside the punctuation that closes a quote. What
+        # is too short to have been a sentence is a list marker or a shortened
+        # word, and the cue goes on rather than ending on it.
+        if (end - start >= MIN_CUE_SECONDS
+                and text.rstrip("\"')]»”’").endswith(tuple(SENTENCE_END))):
+            flush()
+    flush()
+    return _held(cues)
+
+
+def _held(cues):
+    """Keep a cue that is still too short on screen, without covering the next.
+
+    A one word sentence is a fifth of a second of audio and so a fifth of a
+    second of subtitle, which is a flicker. It stays up until the cue after it
+    starts, or for as long as it takes to read, whichever comes first.
+    """
+    out = []
+    for index, (start, end, text) in enumerate(cues):
+        if end - start < MIN_CUE_SECONDS:
+            room = start + MIN_CUE_SECONDS
+            if index + 1 < len(cues):
+                room = min(room, cues[index + 1][0])
+            end = max(end, room)
+        out.append((start, end, text))
+    return out
+
+
 def transcribe(target, audio_path, language="", prompt="", timeout=300, aborter=None):
     data = _transcribe_request(
         target, audio_path, language, prompt, "json", timeout=timeout, aborter=aborter
@@ -459,15 +549,34 @@ def transcribe(target, audio_path, language="", prompt="", timeout=300, aborter=
 def transcribe_segments(target, audio_path, language="", prompt="", timeout=300,
                         aborter=None):
     """[(start_seconds, end_seconds, text)] using whisper-1's verbose response."""
-    data = _transcribe_request(
+    target = target._replace(model=timestamp_model(target.provider, target.model,
-        target._replace(model=timestamp_model(target.provider, target.model,
+                                                   target.file_model))
-                                              target.file_model)),
+    ask = dict(language=language, prompt=prompt, response_format="verbose_json",
-        audio_path, language, prompt, "verbose_json",
+               timeout=timeout, aborter=aborter)
-        granularity="segment", timeout=timeout, aborter=aborter,
+    # Word times are the way out of a model that does not mark segments, and
-    )
+    # whisper.cpp is not one of those, so the local server is only ever asked
+    # for what it has always been asked for. A hosted model that refuses the
+    # field says so with a 400, and the request it used to answer is still
+    # there to fall back on rather than losing the run over a field it did not
+    # need in the first place.
+    if target.provider == "local":
+        data = _transcribe_request(target, audio_path, granularity=("segment",), **ask)
+    else:
+        try:
+            data = _transcribe_request(target, audio_path,
+                                       granularity=("segment", "word"), **ask)
+        except ApiError as exc:
+            if exc.status != 400:
+                raise
+            data = _transcribe_request(target, audio_path,
+                                       granularity=("segment",), **ask)
     segments = data.get("segments") or []
     if target.provider == "local":
         segments = _merge_word_splits(segments)
+    if _too_coarse(segments):
+        cues = cues_from_words(data.get("words") or [])
+        if cues:
+            return cues
     out = []
     for seg in segments:
         text = (seg.get("text") or "").strip()

dikte/filetranscribe.py

@@ -283,11 +283,20 @@ def to_srt(text, segments):
         hours, minutes, secs = (int(g or 0) for g in match.groups())
         cues.append([hours * 3600 + minutes * 60 + secs, None, body])
 
+    # Several cues can share a whole second, so a second holds every segment
+    # that began in it and they are handed out in the order they were spoken.
     timing = {}
     for start, end, _ in segments:
-        timing.setdefault(int(start), (start, end))
+        timing.setdefault(int(start), []).append((start, end))
     for cue in cues:
-        cue[0], cue[1] = timing.get(cue[0], (float(cue[0]), 0.0))
+        found = timing.get(cue[0])
+        if found:
+            # The last one stays, so a second with more lines than it has
+            # timings hands the last of them out again rather than falling back
+            # to the bare second, which would run backwards from the line above.
+            cue[0], cue[1] = found.pop(0) if len(found) > 1 else found[0]
+        else:
+            cue[0], cue[1] = float(cue[0]), 0.0
     for index, cue in enumerate(cues):
         following = cues[index + 1][0] if index + 1 < len(cues) else 0.0
         if following > cue[0]:

tests/test_api.py

@@ -322,7 +322,12 @@ class TranscribeSegments(DikteTest):
         fields = multipart_fields(calls[0])
         self.assertEqual(fields["model"], "whisper-1")
         self.assertEqual(fields["response_format"], "verbose_json")
-        self.assertEqual(fields["timestamp_granularities[]"], "segment")
+        # Both are asked for: whisper answers with segments, and a model that
+        # does not mark them still answers with word times.
+        body = calls[0].data.decode("utf-8", "replace")
+        for level in ("segment", "word"):
+            self.assertIn(
+                f'name="timestamp_granularities[]"\r\n\r\n{level}\r\n', body)
 
     def test_openrouter_uses_the_namespaced_id(self):
         with fake_urlopen(self.reply([{"start": 0, "end": 1, "text": "hi"}])) as calls:
@@ -363,6 +368,74 @@ class TranscribeSegments(DikteTest):
             self.assertEqual(api.transcribe_segments(OPENAI, self.wav),
                              [(5.0, 5.0, "hi")])
 
+    def test_a_long_sentence_is_broken_where_it_gets_too_long_to_read(self):
+        words = [{"word": "word", "start": i * 0.2, "end": i * 0.2 + 0.2}
+                 for i in range(60)]
+        cues = api.cues_from_words(words)
+        self.assertGreater(len(cues), 1)
+        for start, end, text in cues:
+            self.assertLessEqual(len(text), api.MAX_CUE_CHARS)
+            self.assertLessEqual(end - start, api.MAX_CUE_SECONDS + 0.2)
+
+    def test_a_pause_between_short_sentences_does_not_join_them(self):
+        cues = api.cues_from_words([
+            {"word": "Yes.", "start": 0.0, "end": 0.3},
+            {"word": "No.", "start": 9.0, "end": 9.3},
+        ])
+        self.assertEqual([(start, text) for start, _, text in cues],
+                         [(0.0, "Yes."), (9.0, "No.")])
+
+    def test_a_cue_too_short_to_read_is_held_until_the_next_one(self):
+        cues = api.cues_from_words([
+            {"word": "Yes.", "start": 0.0, "end": 0.3},
+            {"word": "No.", "start": 9.0, "end": 9.3},
+        ])
+        # The first has the room for it, the last has nothing after it to wait for.
+        self.assertEqual(cues[0][1], api.MIN_CUE_SECONDS)
+        self.assertEqual(cues[1][1], 9.0 + api.MIN_CUE_SECONDS)
+
+    def test_a_list_marker_does_not_end_a_cue_on_its_own(self):
+        cues = api.cues_from_words([
+            {"word": "1.", "start": 0.0, "end": 0.2},
+            {"word": "Antivirus.", "start": 0.4, "end": 1.6},
+        ])
+        self.assertEqual([text for _, _, text in cues], ["1. Antivirus."])
+
+    def test_a_sentence_ending_inside_a_quote_still_ends_the_cue(self):
+        cues = api.cues_from_words([
+            {"word": '"Stop', "start": 0.0, "end": 1.0},
+            {"word": 'there."', "start": 1.1, "end": 2.0},
+            {"word": "Then", "start": 2.2, "end": 2.6},
+        ])
+        self.assertEqual([text for _, _, text in cues],
+                         ['"Stop there."', "Then"])
+
+    def test_word_times_take_over_from_segments_too_long_to_read(self):
+        # What a model that does not mark segments answers with: one entry for
+        # the whole file, and the real timing in the words beside it.
+        reply = {
+            "text": "One. Two.",
+            "segments": [{"start": 0, "end": 60, "text": "One. Two."}],
+            "words": [
+                {"word": "One.", "start": 0.1, "end": 1.5},
+                {"word": "Two.", "start": 1.7, "end": 3.0},
+            ],
+        }
+        with fake_urlopen(reply):
+            self.assertEqual(api.transcribe_segments(OPENAI, self.wav),
+                             [(0.1, 1.5, "One."), (1.7, 3.0, "Two.")])
+
+    def test_whisper_segments_are_left_alone_when_words_come_too(self):
+        reply = {
+            "text": "hi there",
+            "segments": [{"start": 0, "end": 2, "text": "hi there"}],
+            "words": [{"word": "hi", "start": 0.0, "end": 0.5},
+                      {"word": "there", "start": 0.5, "end": 2.0}],
+        }
+        with fake_urlopen(reply):
+            self.assertEqual(api.transcribe_segments(OPENAI, self.wav),
+                             [(0.0, 2.0, "hi there")])
+
     def test_a_model_that_returned_no_segments_still_gives_its_text(self):
         with fake_urlopen(self.reply([], text="the whole thing")):
             self.assertEqual(api.transcribe_segments(OPENAI, self.wav),