refactor(asr-worker): use doc IDs as ASR workflow inputs instead of paths

asr-worker/asr_worker/models.py

@@ -1,5 +1,4 @@
 import math
-from pathlib import Path
 from typing import Annotated, Any, Self
 
 from caul.asr_pipeline import ASRPipelineConfig
@@ -32,11 +31,12 @@
 
 
 DocumentSearchQuery = dict[str, Any]
+DocId = str
 
 
 class ASRArgs(DatashareModel):
     project: str
-    docs: list[Path] | DocumentSearchQuery
+    docs: list[DocId] | DocumentSearchQuery
     config: ASRPipelineConfig = Field(default=_DEFAULT_PIPELINE_CONFIG)
     batch_size: int
 

asr-worker/asr_worker/workflows.py

@@ -4,6 +4,7 @@
 from itertools import repeat
 
 from datashare_python.utils import WorkflowWithProgress, execute_activity
+from icij_common.es import has_id
 from pydantic import TypeAdapter
 from temporalio import workflow
 
@@ -30,24 +31,16 @@ async def run(self, args: ASRArgs) -> ASRResponse:
         logger = workflow.logger
         config = args.config
         batch_size = args.batch_size
-        docs = args.docs
-        if isinstance(docs, dict):
-            args = [args.project, docs, batch_size]
-            batched_input_paths = workflow.execute_activity(
-                ASRActivities.search_audios,
-                args=args,
-                start_to_close_timeout=timedelta(seconds=TEN_MINUTES),
-                task_queue=TaskQueues.IO,
-            )
-        else:
-            batched_input_paths = [
-                docs[batch_start : batch_start + batch_size]
-                for batch_start in range(0, len(docs), batch_size)
-            ]
-        # Preprocessing
-        preprocess_args = zip(
-            batched_input_paths, repeat(config.preprocessing), strict=False
+        doc_query = has_id(args.docs) if isinstance(args.docs, list) else args.docs
+        search_args = [args.project, doc_query, batch_size]
+        batch_paths = await workflow.execute_activity(
+            ASRActivities.search_audio_paths,
+            args=search_args,
+            start_to_close_timeout=timedelta(seconds=TEN_MINUTES),
+            task_queue=TaskQueues.IO,
         )
+        # Preprocessing
+        preprocess_args = zip(batch_paths, repeat(config.preprocessing), strict=False)
         preprocessing_acts = (
             execute_activity(
                 ASRActivities.preprocess,
@@ -81,7 +74,7 @@ async def run(self, args: ASRArgs) -> ASRResponse:
         postprocessing_ins = list(
             zip(
                 inference_results,
-                batched_input_paths,
+                batch_paths,
                 repeat(config.postprocessing),
                 repeat(args.project),
                 strict=False,
@@ -97,9 +90,10 @@ async def run(self, args: ASRArgs) -> ASRResponse:
             for i in postprocessing_ins
         ]
         logger.info("running postprocessing...")
-        await gather(*postprocessing_acts)
+        n_transcribed = await gather(*postprocessing_acts)
+        n_transcribed = sum(n_transcribed)
         logger.info("postprocessing complete !")
-        return ASRResponse(n_transcribed=len(args.docs))
+        return ASRResponse(n_transcribed=n_transcribed)
 
 
 REGISTRY = [ASRWorkflow]

asr-worker/tests/conftest.py

@@ -1,3 +1,5 @@
+from pathlib import Path
+
 import pytest
 from _pytest.tmpdir import TempPathFactory
 from asr_worker.config import ASRWorkerConfig
@@ -48,7 +50,11 @@ def test_worker_config(tmp_path_factory: TempPathFactory) -> ASRWorkerConfig:  #
 @pytest.fixture(scope="session")
 def doc_0() -> Document:
     return Document(
-        id="doc-0", root_document="root-0", language="ENGLISH", content_type="audio/wav"
+        id="doc-0",
+        root_document="root-0",
+        language="ENGLISH",
+        content_type="audio/wav",
+        path=Path("doc-0.wav"),
     )
 
 
@@ -59,13 +65,18 @@ def doc_1() -> Document:
         root_document="root-1",
         language="ENGLISH",
         content_type="application/json",
+        path=Path("doc-1.json"),
     )
 
 
 @pytest.fixture(scope="session")
 def doc_2() -> Document:
     return Document(
-        id="doc-2", root_document="root-2", language="FRENCH", content_type="audio/mpeg"
+        id="doc-2",
+        root_document="root-2",
+        language="FRENCH",
+        content_type="audio/mpeg",
+        path=Path("doc-2.mp3"),
     )