feature(asr-worker): support search query as input of ASR workflow

asr-worker/asr_worker/activities.py

@@ -1,18 +1,30 @@
 import uuid
+from collections.abc import AsyncGenerator, AsyncIterable, Callable, Generator, Iterable
+from contextlib import contextmanager
 from pathlib import Path
+from typing import Any
 
-from caul.configs import ParakeetConfig
 from caul.model_handlers.objects import ASRModelHandlerResult
 from caul.tasks.preprocessing.objects import PreprocessedInput
 from datashare_python.types_ import ProgressRateHandler
 from datashare_python.utils import (
     ActivityWithProgress,
+    activity_contextual_id,
     activity_defn,
     debuggable_name,
     safe_dir,
     to_raw_progress,
     write_artifact,
 )
+from icij_common.es import (
+    DOC_CONTENT_TYPE,
+    ES_DOCUMENT_TYPE,
+    HITS,
+    ESClient,
+    and_query,
+    has_type,
+)
+from icij_common.iter_utils import async_batches
 from pydantic import TypeAdapter
 from temporalio import activity
 
@@ -21,38 +33,64 @@
     POSTPROCESS_ACTIVITY,
     PREPROCESS_ACTIVITY,
     RUN_INFERENCE_ACTIVITY,
+    SEARCH_AUDIOS_ACTIVITY,
+    SUPPORTED_CONTENT_TYPES,
     TRANSCRIPTION_METADATA_KEY,
     TRANSCRIPTION_METADATA_VALUE,
 )
+from .dependencies import lifespan_es_client
 from .models import Transcription
 
 _BASE_WEIGHT = 1.0
+_SEARCH_AUDIOS_WEIGHT = _BASE_WEIGHT * 2
 _PREPROCESS_WEIGHT = 5 * _BASE_WEIGHT
 _INFERENCE_WEIGHT = 10 * _PREPROCESS_WEIGHT
 
 _LIST_OF_PATH_ADAPTER = TypeAdapter(list[Path])
 
 
 class ASRActivities(ActivityWithProgress):
-    # TODO: pass this at runtime
-    _handler_config = ParakeetConfig(return_tensors=False)
+    @activity_defn(name=SEARCH_AUDIOS_ACTIVITY, progress_weight=_SEARCH_AUDIOS_WEIGHT)
+    async def search_audios(
+        self, project: str, query: dict[str, Any], batch_size: int
+    ) -> list[Path]:
+        es_client = lifespan_es_client()
+        worker_config = ASRWorkerConfig()
+        batch_dir_name = activity_contextual_id()
+        workdir = worker_config.workdir
+        batch_root = workdir / batch_dir_name
+        batch_root.mkdir(parents=True, exist_ok=True)
+        # TODO: supported content types should be args
+        query = search_audios(
+            es_client, project, query, supported_content_types=SUPPORTED_CONTENT_TYPES
+        )
+        batch_paths = [
+            p.relative_to(workdir)
+            async for p in write_audio_search_results(query, batch_root, batch_size)
+        ]
+        return batch_paths
 
     @activity_defn(name=PREPROCESS_ACTIVITY, progress_weight=_PREPROCESS_WEIGHT)
-    def preprocess(self, paths: list[Path], config: PreprocessorConfig) -> list[Path]:
+    def preprocess(
+        self, paths: list[Path] | Path, config: PreprocessorConfig
+    ) -> list[Path]:
         # TODO: this shouldn't be necessary, fix this bug
-        paths = _LIST_OF_PATH_ADAPTER.validate_python(paths)
         worker_config = ASRWorkerConfig()
         audio_root = worker_config.audios_root
         workdir = worker_config.workdir
         # TODO: load from config passed at runtime with caching
         preprocessor = Preprocessor.from_config(config)
-        with preprocessor:
+        if isinstance(paths, Path):
+            audio_cm = _read_audio_ids(paths)
+        else:
+            audio_cm = _read_audios_cm(paths)
+        with audio_cm as audios, preprocessor:
             # TODO: implement a caching strategy here, we could avoid processing files
             #  which have already been preprocessed
-            to_process = [str(audio_root / p) for p in paths]
+            to_process = (str(audio_root / p) for p in audios)
             batches = []
             # TODO: handle progress here
-            for batch in preprocessor.process(str(to_process), output_dir=workdir):
+            for batch in preprocessor.process(to_process, output_dir=workdir):
                 for preprocessed_input in batch:
                     uuid_name = uuid.uuid4().hex[:20]
                     segment_dir = safe_dir(uuid_name)
@@ -69,7 +107,7 @@
     def infer(
         self,
         preprocessed_inputs: list[Path],
         config: InferenceRunnerConfig,
         *,
         progress: ProgressRateHandler | None = None,
     ) -> list[Path]:
@@ -77,7 +115,7 @@
         workdir = worker_config.workdir
         preprocessed_inputs = _LIST_OF_PATH_ADAPTER.validate_python(preprocessed_inputs)
         # TODO: load from config passed at runtime with caching
         inference_runner = InferenceRunner.from_config(config)
         with inference_runner:
             # TODO: extract this into a function to improve testability
             paths = []
@@ -104,7 +142,7 @@
         self,
         inference_results: list[Path],
         input_paths: list[Path],
         config: PostprocessorConfig,
         project: str,
         *,
         progress: ProgressRateHandler | None = None,
@@ -114,11 +152,11 @@
         worker_config = ASRWorkerConfig()
         artifacts_root = worker_config.artifacts_root
         # TODO: load from config passed at runtime with caching
         postprocessor = Postprocessor.from_config(config)
         with postprocessor:
             if progress is not None:
                 progress = to_raw_progress(progress, max_progress=len(input_paths))
             transcriptions = post_processor.process(inference_results)
             # Strict is important here !
             for i, (original, asr_result) in enumerate(
                 zip(input_paths, transcriptions, strict=True)
@@ -156,4 +194,68 @@
     return rel_path
 
 
-REGISTRY = [ASRActivities.preprocess, ASRActivities.infer, ASRActivities.postprocess]
+async def write_audio_search_results(
+    results: AsyncIterable[dict], root: Path, batch_size: int
+) -> AsyncIterable[Path]:
+    batch_id = 0
+    async for batch in async_batches(results, batch_size):
+        batch_path = root / f"{batch_id}.txt"
+        with batch_path.open("w") as f:
+            for audio_file in batch:
+                f.write(f"{audio_file}\n")
+        yield batch_path
+        batch_id += 1
+
+
+_DOC_TYPE_QUERY = has_type(type_field="type", type_value=ES_DOCUMENT_TYPE)
+
+
+async def search_audios(
+    es_client: ESClient,
+    project: str,
+    query: dict[str, Any],
+    supported_content_types: set[str],
+) -> AsyncGenerator[str, None]:
+    with_audio_content_body = _with_audio_content(query, supported_content_types)
+    async for page in es_client.poll_search_pages(
+        index=project, body=with_audio_content_body, sort="_doc:asc", _source=False
+    ):
+        for hit in page[HITS][HITS]:
+            yield hit
+
+
+def _content_type_query(supported_content_types: set[str]) -> dict[str, Any]:
+    content_type_query = {"terms": {DOC_CONTENT_TYPE: sorted(supported_content_types)}}
+    doc_type = has_type(type_field="type", type_value=ES_DOCUMENT_TYPE)
+    return and_query(content_type_query, doc_type)
+
+
+def _with_audio_content(
+    query: dict[str, Any], supported_content_types: set[str]
+) -> dict[str, Any]:
+    return and_query(query, _content_type_query(supported_content_types))
+
+
+def _read_audio_ids(path: Path) -> Callable:
+    @contextmanager
+    def cm() -> Generator[Iterable[Path], None, None]:
+        with open(path) as f:
+            yield (Path(line.strip()) for line in f)
+
+    return cm
+
+
+def _read_audios_cm(paths: list[Path]) -> Callable:
+    @contextmanager
+    def cm() -> Generator[Iterable[Path], None, None]:
+        yield iter(paths)
+
+    return cm
+
+
+REGISTRY = [
+    ASRActivities.search_audios,
+    ASRActivities.preprocess,
+    ASRActivities.infer,
+    ASRActivities.postprocess,
+]

asr-worker/asr_worker/constants.py

@@ -20,5 +20,21 @@
 ASR_WORKFLOW = "asr.transcription"
 GET_CONFIG_ACTIVITY = "asr.transcription.config"
 PREPROCESS_ACTIVITY = "asr.transcription.preprocess"
+SEARCH_AUDIOS_ACTIVITY = "asr.transcription.search_audios"
 RUN_INFERENCE_ACTIVITY = "asr.transcription.infer"
 POSTPROCESS_ACTIVITY = "asr.transcription.postprocess"
+
+SUPPORTED_CONTENT_TYPES = {
+    "audio/aac",
+    "audio/aiff",
+    "audio/mp4",
+    "audio/ogg",
+    "audio/vnd.wave",
+    "audio/wav",
+    "audio/wave",
+    "audio/x-wav",
+    "audio/x-pn-wav",
+    "video/mp4",
+    "video/mpeg",
+    "video/mov",
+}

asr-worker/asr_worker/dependencies.py

@@ -0,0 +1,8 @@
+from datashare_python.dependencies import (  # noqa: F401
+    lifespan_es_client,
+    set_es_client,
+)
+
+DEPENDENCIES = {
+    "base": [set_es_client],
+}

asr-worker/asr_worker/models.py

@@ -1,6 +1,6 @@
 import math
 from pathlib import Path
-from typing import Self
+from typing import Any, Self
 
 from caul.model_handlers.objects import ASRModelHandlerResult
 from datashare_python.objects import DatashareModel
@@ -29,10 +29,10 @@
 
 class ASRPipelineConfig(DatashareModel):
     batch_size: int = 32
     preprocessing: ParakeetPreprocessorConfig = Field(
         default_factory=ParakeetPreprocessorConfig
     )
     inference: InferenceRunnerConfig = Field(
         default_factory=ParakeetInferenceRunnerConfig
     )
     postprocessing: PostProcessorConfig = Field(
@@ -40,9 +40,12 @@
     )
 
 
+DocumentSearchQuery = dict[str, Any]
+
+
 class ASRInputs(DatashareModel):
     project: str
-    paths: list[Path]
+    docs: list[Path] | DocumentSearchQuery
     config: ASRPipelineConfig
 
 

asr-worker/asr_worker/workflows.py

@@ -27,14 +27,24 @@ class TaskQueues(StrEnum):
 class ASRWorkflow(WorkflowWithProgress):
     @workflow.run
     async def run(self, inputs: ASRInputs) -> ASRResponse:
-        # Preprocessing
         logger = workflow.logger
         config = inputs.config
         batch_size = config.batch_size
-        batched_input_paths = (
-            inputs.paths[batch_start : batch_start + batch_size]
-            for batch_start in range(0, len(inputs.paths), batch_size)
-        )
+        if isinstance(inputs.docs, dict):
+            args = [inputs.project, inputs.docs, batch_size]
+            audio_paths = workflow.execute_activity(
+                ASRActivities.search_audios,
+                args=args,
+                start_to_close_timeout=timedelta(seconds=TEN_MINUTES),
+                task_queue=TaskQueues.IO,
+            )
+        else:
+            audio_paths = inputs.docs
+            batched_input_paths = (
+                inputs.docs[batch_start : batch_start + batch_size]
+                for batch_start in range(0, len(inputs.docs), batch_size)
+            )
+        # Preprocessing
         preprocess_args = zip(
             batched_input_paths, repeat(config.preprocessing), strict=False
         )
@@ -87,7 +97,7 @@ async def run(self, inputs: ASRInputs) -> ASRResponse:
         logger.info("running postprocessing...")
         await gather(*postprocessing_acts)
         logger.info("postprocessing complete !")
-        return ASRResponse(n_transcribed=len(inputs.paths))
+        return ASRResponse(n_transcribed=len(inputs.docs))
 
 
 REGISTRY = [ASRWorkflow]

asr-worker/tests/conftest.py

@@ -4,13 +4,18 @@
 import pytest
 from asr_worker.config import ASRWorkerConfig
 from datashare_python.conftest import (  # noqa: F401
+    TEST_PROJECT,
     event_loop,
+    index_docs,
+    populate_es,
     test_deps,
     test_temporal_client_session,
     test_worker_config,
     worker_lifetime_deps,
 )
+from icij_common.es import ESClient
 from icij_common.test_utils import reset_env  # noqa: F401
+from objects import Document
 
 
 @pytest.fixture
@@ -19,3 +24,41 @@ def mocked_worker_config_in_env(reset_env, tmp_path: Path) -> ASRWorkerConfig:
     os.environ["DS_WORKER_ARTIFACT_ROOT"] = str(tmp_path / "artifacts")
     os.environ["DS_WORKER_WORKDIR"] = str(tmp_path / "workdir")
     return ASRWorkerConfig()
+
+
+@pytest.fixture(scope="session")
+def doc_0() -> Document:
+    return Document(
+        id="doc-0", root_document="root-0", language="ENGLISH", content_type="audio/wav"
+    )
+
+
+@pytest.fixture(scope="session")
+def doc_1() -> Document:
+    return Document(
+        id="doc-1",
+        root_document="root-1",
+        language="ENGLISH",
+        content_type="application/json",
+    )
+
+
+@pytest.fixture(scope="session")
+def doc_2() -> Document:
+    return Document(
+        id="doc-2", root_document="root-2", language="FRENCH", content_type="audio/mpeg"
+    )
+
+
+@pytest.fixture
+async def populate_es_with_audio(
+    test_es_client: ESClient,
+    doc_0: Document,
+    doc_1: Document,
+    doc_2: Document,
+    doc_3: Document,
+) -> list[Document]:
+    docs = [doc_0, doc_1, doc_2, doc_3]
+    async for _ in index_docs(test_es_client, docs=docs, index_name=TEST_PROJECT):
+        pass
+    return docs