run_funs.py

import os
from pathlib import Path

from typeguard import typechecked

import diariz_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_aus")
Path(model_dir).mkdir(parents=True, exist_ok=True)

# TODO
# get conf_dir from input msg
conf_dir = os.path.join(os.environ["AI_FW_DIR"], "confs", "mmc_aus")


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

    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'

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

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

        message_body["programme"]["conf_dir"] = conf_dir

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