improve chunking of texts

app.py

@@ -1,6 +1,8 @@
 import base64
 import os
 import time
+from dataclasses import dataclass
+from datetime import timedelta
 
 import banana_dev as banana
 import gradio as gr
@@ -11,6 +13,57 @@ api_key = os.environ["BANANA_API_KEY"]
 model_key = os.environ["BANANA_MODEL_KEY"]
 password = os.environ["PASSWORD"]
 
+SECONDS_IN_HOUR = 3600
+SECONDS_IN_MINUTE = 60
+HOURS_IN_DAY = 24
+MICROSECONDS_IN_MILLISECOND = 1000
+
+
+def timedelta_to_srt_timestamp(timedelta_timestamp):
+    r"""
+    Convert a :py:class:`~datetime.timedelta` to an SRT timestamp.
+    .. doctest::
+        >>> import datetime
+        >>> delta = datetime.timedelta(hours=1, minutes=23, seconds=4)
+        >>> timedelta_to_srt_timestamp(delta)
+        '01:23:04,000'
+    :param datetime.timedelta timedelta_timestamp: A datetime to convert to an SRT timestamp
+    :returns: The timestamp in SRT format
+    :rtype: str
+    """
+
+    hrs, secs_remainder = divmod(timedelta_timestamp.seconds, SECONDS_IN_HOUR)
+    hrs += timedelta_timestamp.days * HOURS_IN_DAY
+    mins, secs = divmod(secs_remainder, SECONDS_IN_MINUTE)
+    msecs = timedelta_timestamp.microseconds // MICROSECONDS_IN_MILLISECOND
+    return "%02d:%02d:%02d,%03d" % (hrs, mins, secs, msecs)
+
+
+def timedelta_to_otr_timestamp(timedelta_timestamp):
+    output = timedelta_to_srt_timestamp(timedelta_timestamp)
+    if output.startswith("00:"):
+        output = output[3:]
+    return output[:-4]
+
+
+@dataclass
+class Segment:
+    text: str
+    start: float
+    end: float
+
+    @property
+    def start_ts(self) -> str:
+        return timedelta_to_otr_timestamp(timedelta(seconds=self.start))
+
+    @property
+    def end_ts(self):
+        return timedelta_to_otr_timestamp(timedelta(seconds=self.end))
+
+    def __str__(self):
+        sep = " "
+        return f'<p><span class="timestamp" data-timestamp="{self.start}">{self.start_ts}</span>{sep}{self.text}</p>'
+
 
 def transcribe(audio=None, url=None):
     if audio:
@@ -47,8 +100,18 @@ def run_demo(password, microphone, file_upload):
 
     start = time.time()
     transcription = transcribe(AudioSegment.from_file(file)[:60_000])
+
+    segments = []
+
+    for seg in transcription["segments"]:
+        text = seg["text"].strip()
+        if not segments or segments[-1].text[-1] in ".:?!":
+            segments.append(Segment(text, seg["start"], seg["end"]))
+        else:
+            segments[-1].text += " " + text
+
     logger.info(f"transcription took {time.time()-start:.3f}s")
-    return "\n\n".join([seg["text"].strip() for seg in transcription["segments"]])
+    return "\n\n".join(str(s) for s in segments)
 
 
 demo = gr.Interface(