run_funs.py

import os
from pathlib import Path

from typeguard import typechecked

import asr_funs

try:
    from common_utils import msg_builder, rabbitmq
except ModuleNotFoundError:
    from common_module.common_utils import msg_builder, rabbitmq

# TODO
# get vid_dir from input msg
base_dir = os.path.join(os.environ["AI_FW_DIR"], "vids")

# TODO
# get model_dir from input msg
model_dir = os.path.join(os.environ["AI_FW_DIR"], "models", "mmc_asr")
Path(model_dir).mkdir(parents=True, exist_ok=True)


@typechecked
def run(message_body: dict, worker: rabbitmq.Worker) -> bool:
    defs = {
        # module name in msg
        "mod_name": "mmc_asr",
        # metadata key in output msg
        "metadata_key": "transcription",
        # metadata type in output msg
        "metadata_type": "transcription",
        # main key in output JSON
        "out_json_key": "transcription",
        # error msg if output JSON is not found
        "not_found_msg": "cannot transcribe!",
        # error msg if input msg is invalid
        "invalid_msg": "External ID/UID/Application/Diar Required!",
    }

    extras = {"programme": {"module": defs["mod_name"]}}

    if "programme" in message_body:
        if "external_id" in message_body["programme"]:
            # name of output JSON
            defs["out_json"] = f'{message_body["programme"]["external_id"]}.json'

        for k in msg_builder.handed_over_keys():
            if k in message_body["programme"]:
                extras["programme"][k] = message_body["programme"][k]

    return msg_builder.build_msg(
        message_body,
        worker,
        "mmc_asr",
        asr_funs.dl_trs_save,
        msg_builder.validate_message,
        ["external_id", "application", "uid", "mmc_aus"],
        base_dir,
        model_dir,
        defs,
        extras,
    )