⚡️ Resume transcription tasks

bugbakery · Aug 6, 2023 · 6ddd4cf · 6ddd4cf
1 parent 2c084db
commit 6ddd4cf
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 10 deletions.
diff --git a/worker/transcribee_worker/whisper_transcribe.py b/worker/transcribee_worker/whisper_transcribe.py
@@ -40,6 +40,7 @@ def get_context(model_name: str) -> api.Context:
 def _transcription_work(
     queue: SubmissionQueue,
     data: NDArray[Any],
+    start_offset: float,
     model_name: str,
     lang_code: Optional[str],
     progress_callback: Optional[ProgressCallbackType],
@@ -112,9 +113,9 @@ def handle_new_segment(
                         text=text,
                         conf=conf,
                         # 10·ms -> seconds
-                        start=start / 100,
+                        start=(start / 100) + start_offset,
                         # 10·ms -> seconads
-                        end=end / 100,
+                        end=(end / 100) + start_offset,
                         conf_ts=conf_ts,
                     )
                 )
@@ -161,10 +162,19 @@ def handle_new_segment(
 
 
 def transcribe(
-    data: NDArray, model_name: str, lang_code="en", progress_callback=None
+    data: NDArray,
+    start_offset: float,
+    model_name: str,
+    lang_code="en",
+    progress_callback=None,
 ) -> AsyncIterator[Paragraph]:
     return async_task(
-        _transcription_work, data, model_name, lang_code, progress_callback
+        _transcription_work,
+        data,
+        start_offset,
+        model_name,
+        lang_code,
+        progress_callback,
     )
 
 
@@ -233,7 +243,11 @@ async def remove_leading_whitespace_from_paragraph(
 
 
 async def transcribe_clean(
-    data: NDArray, model_name: str, lang_code: str = "en", progress_callback=None
+    data: NDArray,
+    start_offset: float,
+    model_name: str,
+    lang_code: str = "en",
+    progress_callback=None,
 ):
     chain = (
         recombine_split_words,
@@ -243,6 +257,7 @@ async def transcribe_clean(
     iter = aiter(
         transcribe(
             data=data,
+            start_offset=start_offset,
             model_name=model_name,
             lang_code=lang_code,
             progress_callback=progress_callback,

diff --git a/worker/transcribee_worker/worker.py b/worker/transcribee_worker/worker.py
@@ -62,6 +62,16 @@ def ensure_atom_invariants(doc: EditorDocument):
         prev_atom = atom
 
 
+def get_last_atom_end(doc: EditorDocument):
+    for paragraph_idx in reversed(range(len(doc.children))):
+        for atom_idx in reversed(range(len(doc.children[paragraph_idx].children))):
+            atom = doc.children[paragraph_idx].children[atom_idx]
+            if atom.end is not None:
+                return atom.end
+
+    return 0
+
+
 class Worker:
     base_url: str
     token: str
@@ -166,14 +176,21 @@ async def transcribe(
         audio = self.load_document_audio(task.document)
 
         async with self.api_client.document(task.document.id) as doc:
+
             async with doc.transaction("Reset Document") as d:
-                d.children = []
+                if d.children is None:
+                    d.children = []
+
+                start_offset = get_last_atom_end(d)
+
+            audio = audio[int(start_offset * settings.SAMPLE_RATE) :]
 
             async for paragraph in transcribe_clean(
-                audio,
-                task.task_parameters.model,
-                task.task_parameters.lang,
-                progress_callback,
+                data=audio,
+                start_offset=start_offset,
+                model_name=task.task_parameters.model,
+                lang_code=task.task_parameters.lang,
+                progress_callback=progress_callback,
             ):
                 async with doc.transaction("Automatic Transcription") as d:
                     p = paragraph.dict()