sonusai 0.20.3__py3-none-any.whl → 1.0.2__py3-none-any.whl

sonusai/onnx_predict.py

@@ -14,16 +14,19 @@ The ONNX Runtime (ort) inference engine is used to execute the inference.
 
 Inputs:
     MODEL       ONNX model .onnx file of a trained model (weights are expected to be in the file).
+                The model must also include required Sonusai hyperparameters.  See theSonusai torchl_onnx command.
 
-    DATA        The input data must be one of the following:
-                * WAV
-                  Using the given model, generate feature data and run prediction. A model file must be
-                  provided. The MIXID is ignored.
-
-                * directory
-                  Using the given SonusAI mixture database directory, generate feature and truth data if not found.
-                  Run prediction. The MIXID is required.
+    DATA        A string which must be one of the following:
+                1. Path to a single file.  The prediction data is written to <filename_predict.*> in same location.
+                2. Path to a Sonusai Mixture Database directory.
+                 - Sonusai mixture database directory, prediction files will be named mixid_predict.*
+                 - MIXID will select a subset of mixture ids
+                3. Directory with audio files found recursively within.  See GLOB audio file extensions below.
+                4. Regex resolving to a list of files.
+                 - Subdirectory containing audio files with extension
+                 - Regex resolving to a list of audio files
 
+                   generate feature and truth data if not found.
 
 Note there are multiple ways to process model prediction over multiple audio data files:
 1. TSE (timestep single extension): mixture transform frames are fit into the timestep dimension and the model run as
@@ -42,33 +45,68 @@ TBD not sure below make sense, need to continue ??
 
 Outputs the following to opredict-<TIMESTAMP> directory:
     <id>
-        predict.pkl
+        predict.h5
     onnx_predict.log
 
 """
 
-import signal
-
-
-def signal_handler(_sig, _frame):
-    import sys
-
-    from sonusai import logger
 
-    logger.info("Canceled due to keyboard interrupt")
-    sys.exit(1)
+def process_path(path, ext_list: list[str] | None = None):
+    """
+    Check path which can be a single file, a subdirectory, or a regex
+    return:
+      - a list of files with matching extensions to any in ext_list provided (i.e. ['.wav', '.mp3', '.acc'])
+      - the basedir of the path, if
+    """
+    import glob
+    from os.path import abspath
+    from os.path import commonprefix
+    from os.path import dirname
+    from os.path import isdir
+    from os.path import isfile
+    from os.path import join
 
+    from sonusai.utils import braced_iglob
 
-signal.signal(signal.SIGINT, signal_handler)
+    if ext_list is None:
+        ext_list = [".wav", ".WAV", ".flac", ".FLAC", ".mp3", ".aac"]
+
+    # Check if the path is a single file, and return it as a list with the dirname
+    if isfile(path):
+        if any(path.endswith(ext) for ext in ext_list):
+            basedir = dirname(path)  # base directory
+            if not basedir:
+                basedir = "./"
+            return [path], basedir
+        else:
+            return [], []
+
+    # Check if the path is a dir, recursively find all files any of the specified extensions, return file list and dir
+    if isdir(path):
+        matching_files = []
+        for ext in ext_list:
+            matching_files.extend(glob.glob(join(path, "**/*" + ext), recursive=True))
+        return matching_files, path
+
+    # Process as a regex, return list of filenames and basedir
+    apath = abspath(path)  # join(abspath(path), "**", "*.{wav,flac,WAV}")
+    matching_files = []
+    for file in braced_iglob(pathname=apath, recursive=True):
+        matching_files.append(file)
+    if matching_files:
+        basedir = commonprefix(matching_files)  # Find basedir
+        return matching_files, basedir
+    else:
+        return [], []
 
 
 def main() -> None:
     from docopt import docopt
 
-    import sonusai
+    from sonusai import __version__ as sai_version
     from sonusai.utils import trim_docstring
 
-    args = docopt(trim_docstring(__doc__), version=sonusai.__version__, options_first=True)
+    args = docopt(trim_docstring(__doc__), version=sai_version, options_first=True)
 
     verbose = args["--verbose"]
     wav = args["--write-wav"]
@@ -77,18 +115,23 @@ def main() -> None:
     model_path = args["MODEL"]
     data_paths = args["DATA"]
 
+    # Quick check of CPU and GPU devices
+    import re
+    import subprocess
+    import time
     from os import makedirs
-    from os.path import abspath
     from os.path import basename
+    from os.path import exists
     from os.path import isdir
+    from os.path import isfile
     from os.path import join
     from os.path import normpath
-    from os.path import realpath
     from os.path import splitext
 
     import h5py
     import numpy as np
     import onnxruntime as ort
+    import psutil
 
     from sonusai import create_file_handler
     from sonusai import initial_log_messages
@@ -96,66 +139,122 @@ def main() -> None:
     from sonusai import update_console_handler
     from sonusai.mixture import MixtureDatabase
     from sonusai.mixture import get_audio_from_feature
-    from sonusai.utils import PathInfo
-    from sonusai.utils import braced_iglob
     from sonusai.utils import create_ts_name
     from sonusai.utils import load_ort_session
-    from sonusai.utils import reshape_inputs
+    from sonusai.utils import seconds_to_hms
     from sonusai.utils import write_audio
 
-    mixdb_path = None
-    mixdb: MixtureDatabase | None = None
-    p_mixids: list[int] = []
-    entries: list[PathInfo] = []
-
-    if len(data_paths) == 1 and isdir(data_paths[0]):
-        # Assume it's a single path to SonusAI mixdb subdir
-        in_basename = basename(normpath(data_paths[0]))
-        mixdb_path = data_paths[0]
+    num_cpu = psutil.cpu_count()
+    cpu_percent = psutil.cpu_percent(interval=1)
+    print(f"#CPUs: {num_cpu}, current CPU utilization: {cpu_percent}%")
+    print(f"Memory utilization: {psutil.virtual_memory().percent}%")
+
+    vga_devices = [
+        line.split(" ", 3)[-1]
+        for line in subprocess.check_output("lspci | grep -i vga", shell=True).decode().splitlines()
+    ]
+    nv_devs = list(filter(lambda x: "nvidia" in x.lower(), vga_devices))
+    nv_mods = [re.search(r"\[.*?\]", device).group(0) if re.search(r"\[.*?\]", device) else None for device in nv_devs]
+    if len(nv_mods) > 0:
+        print(f"{len(nv_mods)} Nvidia devices present: {nv_mods}")  # prints model names
     else:
-        # search all data paths for .wav, .flac (or whatever is specified in include)
-        in_basename = ""
-
-    output_dir = create_ts_name("opredict-" + in_basename)
-    makedirs(output_dir, exist_ok=True)
-
-    # Setup logging file
-    create_file_handler(join(output_dir, "onnx-predict.log"))
-    update_console_handler(verbose)
-    initial_log_messages("onnx_predict")
+        print("No cuda devices present, using cpu")
+
+    avail_providers = ort.get_available_providers()
+    print(f"Loaded ONNX Runtime, available providers: {avail_providers}.")
+    if len(nv_mods) > 0:
+        print(
+            "If GPU is desired, need to replace onnxruntime with onnxruntime-gpu i.e. using pip:"
+            "> pip uninstall onnxruntime"
+            "> pip install onnxruntime-gpu\n\n"
+        )
+
+    # Quick check that model is valid
+    if exists(model_path) and isfile(model_path):
+        try:
+            session = ort.InferenceSession(model_path)
+            options = ort.SessionOptions()
+        except Exception as e:
+            print(f"Error: could not load ONNX model from {model_path}: {e}")
+            raise SystemExit(1) from e
+    else:
+        print(f"Error: model file path is not valid: {model_path}")
+        raise SystemExit(1)
 
-    providers = ort.get_available_providers()
-    logger.info(f"Loaded ONNX Runtime, available providers: {providers}.")
+    # Check datapath is valid
+    if len(data_paths) == 1 and isdir(data_paths[0]):  # Try opening as mixdb subdir
+        mixdb_path = data_paths[0]
+        try:
+            mixdb = MixtureDatabase(mixdb_path)
+        except Exception:
+            mixdb_path = None
+        in_basename = basename(normpath(data_paths[0]))
+        output_dir = create_ts_name("opredict-" + in_basename)
+        num_featparams = mixdb.feature_parameters
+        print(f"Loaded SonusAI mixdb with {mixdb.num_mixtures} mixtures and {num_featparams} classes")
+        p_mixids = mixdb.mixids_to_list(mixids)
+        feature_mode = mixdb.feature
 
+    if mixdb_path is None:
+        if verbose:
+            print(f"Checking {len(data_paths)} locations ... ")
+        # Check location, default ext are ['.wav', '.WAV', '.flac', '.FLAC', '.mp3', '.aac']
+        pfiles, basedir = process_path(data_paths)
+        if pfiles is None or len(pfiles) < 1:
+            print(f"No audio files or Sonusai mixture database found in {data_paths}, exiting ...")
+            raise SystemExit(1)
+        else:
+            pfiles = sorted(pfiles, key=basename)
+            output_dir = basedir
+
+    if mixdb_path is not None or len(pfiles) > 1:  # log file only if mixdb or more than one file
+        makedirs(output_dir, exist_ok=True)
+        # Setup logging file
+        create_file_handler(join(output_dir, "onnx-predict.log"))
+        update_console_handler(verbose)
+        initial_log_messages("onnx_predict")
+        # print some previous messages
+        logger.info(f"Loaded ONNX Runtime, available providers: {avail_providers}.")
+        if mixdb_path:
+            logger.debug(f"Loaded SonusAI mixdb with {mixdb.num_mixtures} mixtures and {num_featparams} classes")
+            if len(p_mixids) != mixdb.num_mixtures:
+                logger.info(f"Processing a subset of {len(p_mixids)} from available mixtures.")
+
+    # Reload model/session and do more thorough checking
     session, options, model_root, hparams, sess_inputs, sess_outputs = load_ort_session(model_path)
+    if "CUDAExecutionProvider" in avail_providers:
+        session.set_providers(["CUDAExecutionProvider"])
     if hparams is None:
         logger.error("Error: ONNX model does not have required SonusAI hyperparameters, cannot proceed.")
         raise SystemExit(1)
-    if len(sess_inputs) != 1:
+
+    if len(sess_inputs) != 1:  # TBD update to support state_in and state_out
         logger.error(f"Error: ONNX model does not have 1 input, but {len(sess_inputs)}. Exit due to unknown input.")
 
     in0name = sess_inputs[0].name
     in0type = sess_inputs[0].type
-    out_names = [n.name for n in session.get_outputs()]
-
-    logger.info(f"Read and compiled ONNX model from {model_path}.")
+    in0shape = sess_inputs[0].shape  # a list
+    # Check for 2 cases of model feature input shape:  batch x timesteps x fparams or batch x fparams
+    if not isinstance(in0shape[0], str):
+        model_batchsz = int(in0shape[0])
+        logger.debug(f"Onnx model has fixed batch_size: {model_batchsz}.")
+    else:
+        model_batchsz = -1
+        logger.debug("Onnx model has a dynamic batch_size.")
 
-    if mixdb_path is not None:
-        # Assume it's a single path to SonusAI mixdb subdir
-        logger.debug(f"Attempting to load mixture database from {mixdb_path}")
-        mixdb = MixtureDatabase(mixdb_path)
-        logger.info(f"SonusAI mixdb: found {mixdb.num_mixtures} mixtures with {mixdb.num_classes} classes")
-        p_mixids = mixdb.mixids_to_list(mixids)
-        if len(p_mixids) != mixdb.num_mixtures:
-            logger.info(f"Processing a subset of {p_mixids} from available mixtures.")
+    if len(in0shape) < 3:
+        model_tsteps = 0
+        model_featparams = int(in0shape[1])
     else:
-        for p in data_paths:
-            location = join(realpath(abspath(p)), "**", include)
-            logger.debug(f"Processing {location}")
-            for file in braced_iglob(pathname=location, recursive=True):
-                name = file
-                entries.append(PathInfo(abs_path=file, audio_filepath=name))
-        logger.info(f"{len(data_paths)} data paths specified, found {len(entries)} audio files.")
+        model_featparams = int(in0shape[2])
+        if not isinstance(in0shape[1], str):
+            model_tsteps = int(in0shape[1])
+            logger.debug(f"Onnx model has fixed timesteps: {model_tsteps}.")
+        else:
+            model_tsteps = -1
+            logger.debug("Onnx model has dynamic timesteps dimension size.")
+
+    out_names = [n.name for n in session.get_outputs()]
 
     if in0type.find("float16") != -1:
         model_is_fp16 = True
@@ -163,38 +262,40 @@ def main() -> None:
     else:
         model_is_fp16 = False
 
+    logger.info(f"Read and compiled ONNX model from {model_path}.")
+
+    start_time = time.monotonic()
+
     if mixdb is not None and hparams["batch_size"] == 1:
-        # mixdb input
-        # Assume (of course) that mixdb feature, etc. is what model expects
-        if hparams["feature"] != mixdb.feature:
+        if hparams["feature"] != feature_mode:  # warn on mis-match, but TBD could be sov-mode
             logger.warning("Mixture feature does not match model feature, this inference run may fail.")
-        # no choice, can't use hparams.feature since it's different from the mixdb
-        feature_mode = mixdb.feature
+        logger.info(f"Processing {len(p_mixids)} mixtures from SonusAI mixdb ...")
+        logger.info(f"Using OnnxRT provider {session.get_providers()} ...")
 
         for mixid in p_mixids:
-            # frames x stride x feature_params
-            feature, _ = mixdb.mixture_ft(mixid)
+            # feature data is now always fp32 and frames x stride x feature_params
+            feat_dat, _ = mixdb.mixture_ft(mixid)
+            if feat_dat.shape[1] > 1:  # stride mode num_frames overrides batch dim, no reshape
+                stride_mode = 1
+                batch_size = feat_dat.shape[0]  # num_frames in stride mode becomes batch size
             if hparams["timesteps"] == 0:
-                # no timestep dimension, reshape will handle
+                # no timestep dimension, remove the dimension
                 timesteps = 0
+                feat_dat = np.reshape(feat_dat, [batch_size, num_featparams])
             else:
-                # fit frames into timestep dimension (TSE mode)
-                timesteps = feature.shape[0]
-
-            feature, _ = reshape_inputs(
-                feature=feature,
-                batch_size=1,
-                timesteps=timesteps,
-                flatten=hparams["flatten"],
-                add1ch=hparams["add1ch"],
-            )
+                # fit frames into timestep dimension (TSE mode) and knowing batch_size = 1
+                timesteps = feat_dat.shape[0]
+                feat_dat = np.transpose(feat_dat, (1, 0, 2))  # transpose to 1 x frames=tsteps x feat_params
+
             if model_is_fp16:
-                feature = np.float16(feature)  # type: ignore[assignment]
+                feat_dat = np.float16(feat_dat)  # type: ignore[assignment]
+
             # run inference, ort session wants i.e. batch x timesteps x feat_params, outputs numpy BxTxFP or BxFP
-            predict = session.run(out_names, {in0name: feature})[0]
+            predict = session.run(out_names, {in0name: feat_dat})[0]
             # predict, _ = reshape_outputs(predict=predict[0], timesteps=frames)  # frames x feat_params
+
             output_fname = join(output_dir, mixdb.mixture(mixid).name)
-            with h5py.File(output_fname, "a") as f:
+            with h5py.File(output_fname + ".h5", "a") as f:
                 if "predict" in f:
                     del f["predict"]
                 f.create_dataset("predict", data=predict)
@@ -206,6 +307,57 @@ def main() -> None:
                 owav_name = splitext(output_fname)[0] + "_predict.wav"
                 write_audio(owav_name, predict_audio)
 
+    else:  # TBD add support
+        logger.info("Mixture database does not exist or batch_size is not equal to one, exiting ...")
+
+    end_time = time.monotonic()
+    logger.info(f"Completed in {seconds_to_hms(seconds=end_time - start_time)}")
+    logger.info("")
+
 
 if __name__ == "__main__":
-    main()
+    from sonusai import exception_handler
+    from sonusai.utils import register_keyboard_interrupt
+
+    register_keyboard_interrupt()
+    try:
+        main()
+    except Exception as e:
+        exception_handler(e)
+
+    # mixdb_path = None
+    # mixdb: MixtureDatabase | None = None
+    # p_mixids: list[int] = []
+    # entries: list[PathInfo] = []
+    #
+    # if len(data_paths) == 1 and isdir(data_paths[0]):
+    #     # Assume it's a single path to SonusAI mixdb subdir
+    #     in_basename = basename(normpath(data_paths[0]))
+    #     mixdb_path = data_paths[0]
+    # else:
+    #     # search all data paths for .wav, .flac (or whatever is specified in include)
+    #     in_basename = ""
+
+    # if mixdb_path is not None:  # a mixdb is found and loaded
+    #     # Assume it's a single path to SonusAI mixdb subdir
+    #     num_featparams = mixdb.feature_parameters
+    #     logger.debug(f"SonusAI mixdb: found {mixdb.num_mixtures} mixtures with {num_featparams} classes")
+    #     p_mixids = mixdb.mixids_to_list(mixids)
+    #     if len(p_mixids) != mixdb.num_mixtures:
+    #         logger.info(f"Processing a subset of {p_mixids} from available mixtures.")
+    # else:
+    # for p in data_paths:
+    #     location = join(realpath(abspath(p)), "**", include)
+    #     logger.debug(f"Processing files in {location}")
+    #     for file in braced_iglob(pathname=location, recursive=True):
+    #         name = file
+    #         entries.append(PathInfo(abs_path=file, audio_filepath=name))
+    # logger.info(f"{len(data_paths)} data paths specified, found {len(pfile)} audio files.")
+
+    # feature, _ = reshape_inputs(
+    #     feature=feature,
+    #     batch_size=1,
+    #     timesteps=timesteps,
+    #     flatten=hparams["flatten"],
+    #     add1ch=hparams["add1ch"],
+    # )

sonusai/queries/__init__.py

@@ -1,8 +1,8 @@
 # SonusAI query utilities
 # ruff: noqa: F401
 
+from .queries import get_mixids_from_class_indices
 from .queries import get_mixids_from_noise
 from .queries import get_mixids_from_snr
-from .queries import get_mixids_from_target
+from .queries import get_mixids_from_source
 from .queries import get_mixids_from_truth_function
-from .queries import get_mixids_from_class_indices

sonusai/queries/queries.py

@@ -1,8 +1,8 @@
 from collections.abc import Callable
 from typing import Any
 
-from sonusai.mixture.datatypes import GeneralizedIDs
-from sonusai.mixture.mixdb import MixtureDatabase
+from ..datatypes import GeneralizedIDs
+from ..mixture.mixdb import MixtureDatabase
 
 
 def _true_predicate(_: Any) -> bool:
@@ -29,8 +29,8 @@ def get_mixids_from_mixture_field_predicate(
     criteria_set = set()
     for m_id in mixid_out:
         value = getattr(mixdb.mixture(m_id), field)
-        if isinstance(value, list):
-            for v in value:
+        if isinstance(value, dict):
+            for v in value.values():
                 if predicate(v):
                     criteria_set.add(v)
         elif predicate(value):
@@ -42,8 +42,8 @@ def get_mixids_from_mixture_field_predicate(
         result[criterion] = []
         for m_id in mixid_out:
             value = getattr(mixdb.mixture(m_id), field)
-            if isinstance(value, list):
-                for v in value:
+            if isinstance(value, dict):
+                for v in value.values():
                     if v == criterion:
                         result[criterion].append(m_id)
             elif value == criterion:
@@ -64,7 +64,7 @@ def get_mixids_from_truth_configs_field_predicate(
         - keys are the matching field values
         - values are lists of the mixids that match the criteria
     """
-    from sonusai.mixture import REQUIRED_TRUTH_CONFIGS
+    from ..mixture.constants import REQUIRED_TRUTH_CONFIGS
 
     mixid_out = mixdb.mixids_to_list(mixids)
 
@@ -79,23 +79,24 @@ def get_mixids_from_truth_configs_field_predicate(
 
     result = {}
     for value in values:
-        # Get a list of targets for each field value
+        # Get a list of sources for each field value
         indices = []
-        for t_id in mixdb.target_file_ids:
-            target = mixdb.target_file(t_id)
-            for truth_config in target.truth_configs.values():
-                if field in REQUIRED_TRUTH_CONFIGS:
-                    if value in getattr(truth_config, field):
-                        indices.append(t_id)
-                else:
-                    if value in getattr(truth_config.config, field):
-                        indices.append(t_id)
+        for s_ids in mixdb.source_file_ids.values():
+            for s_id in s_ids:
+                source = mixdb.source_file(s_id)
+                for truth_config in source.truth_configs.values():
+                    if field in REQUIRED_TRUTH_CONFIGS:
+                        if value in getattr(truth_config, field):
+                            indices.append(s_id)
+                    else:
+                        if value in getattr(truth_config.config, field):
+                            indices.append(s_id)
         indices = sorted(set(indices))
 
         mixids = []
         for index in indices:
             for m_id in mixid_out:
-                if index in [target.file_id for target in mixdb.mixture(m_id).targets]:
+                if index in [source.file_id for source in mixdb.mixture(m_id).all_sources.values()]:
                     mixids.append(m_id)
 
         mixids = sorted(set(mixids))
@@ -109,18 +110,19 @@ def get_all_truth_configs_values_from_field(mixdb: MixtureDatabase, field: str)
     """
     Generate a list of all values corresponding to the given field in truth_configs
     """
-    from sonusai.mixture import REQUIRED_TRUTH_CONFIGS
+    from ..mixture.constants import REQUIRED_TRUTH_CONFIGS
 
     result = []
-    for target in mixdb.target_files:
-        for truth_config in target.truth_configs.values():
-            if field in REQUIRED_TRUTH_CONFIGS:
-                value = getattr(truth_config, field)
-            else:
-                value = getattr(truth_config.config, field, None)
-            if not isinstance(value, list):
-                value = [value]
-            result.extend(value)
+    for sources in mixdb.source_files.values():
+        for source in sources:
+            for truth_config in source.truth_configs.values():
+                if field in REQUIRED_TRUTH_CONFIGS:
+                    value = getattr(truth_config, field)
+                else:
+                    value = getattr(truth_config.config, field, None)
+                if not isinstance(value, list):
+                    value = [value]
+                result.extend(value)
 
     return sorted(set(result))
 
@@ -139,18 +141,18 @@ def get_mixids_from_noise(
     return get_mixids_from_mixture_field_predicate(mixdb=mixdb, mixids=mixids, field="noise_id", predicate=predicate)
 
 
-def get_mixids_from_target(
+def get_mixids_from_source(
     mixdb: MixtureDatabase,
     mixids: GeneralizedIDs = "*",
     predicate: Callable[[Any], bool] | None = None,
 ) -> dict[int, list[int]]:
     """
-    Generate mixids based on a target index predicate
+    Generate mixids based on a source index predicate
     Return a dictionary where:
-        - keys are the target indices
-        - values are lists of the mixids that match the target index
+        - keys are the source indices
+        - values are lists of the mixids that match the source index
     """
-    return get_mixids_from_mixture_field_predicate(mixdb=mixdb, mixids=mixids, field="target_ids", predicate=predicate)
+    return get_mixids_from_mixture_field_predicate(mixdb=mixdb, mixids=mixids, field="source_ids", predicate=predicate)
 
 
 def get_mixids_from_snr(
@@ -178,7 +180,9 @@ def get_mixids_from_snr(
     result: dict[float, list[int]] = {}
     for snr in snrs:
         # Get a list of mixids for each SNR
-        result[snr] = sorted([i for i, mixture in enumerate(mixdb.mixtures()) if mixture.snr == snr and i in mixid_out])
+        result[snr] = sorted(
+            [i for i, mixture in enumerate(mixdb.mixtures) if mixture.noise.snr == snr and i in mixid_out]
+        )
 
     return result
 

sonusai/speech/librispeech.py

@@ -30,7 +30,7 @@ def load_text(audio: str | os.PathLike[str]) -> TimeAlignedType | None:
     """
     import string
 
-    from sonusai.mixture import get_sample_rate
+    from ..mixture.audio import get_sample_rate
 
     path = Path(audio)
     name = path.stem

sonusai/speech/mcgill.py

@@ -12,7 +12,7 @@ def load_text(audio: str | os.PathLike[str]) -> TimeAlignedType | None:
     import string
     import struct
 
-    from sonusai.mixture import get_sample_rate
+    from ..mixture.audio import get_sample_rate
 
     if not os.path.exists(audio):
         return None

sonusai/speech/timit.py

@@ -12,7 +12,7 @@ def load_text(audio: str | os.PathLike[str]) -> TimeAlignedType | None:
     """
     import string
 
-    from sonusai.mixture import get_sample_rate
+    from ..mixture.audio import get_sample_rate
 
     file = Path(audio).with_suffix(".TXT")
     if not os.path.exists(file):
@@ -52,7 +52,7 @@ def load_phonemes(audio: str | os.PathLike[str]) -> list[TimeAlignedType] | None
 
 
 def _load_ta(audio: str | os.PathLike[str], tier: str) -> list[TimeAlignedType] | None:
-    from sonusai.mixture import get_sample_rate
+    from ..mixture.audio import get_sample_rate
 
     if tier == "words":
         file = Path(audio).with_suffix(".WRD")

sonusai/summarize_metric_spenh.py

@@ -14,20 +14,6 @@ Inputs:
 
 """
 
-import signal
-
-
-def signal_handler(_sig, _frame):
-    import sys
-
-    from sonusai import logger
-
-    logger.info("Canceled due to keyboard interrupt")
-    sys.exit(1)
-
-
-signal.signal(signal.SIGINT, signal_handler)
-
 
 def summarize_metric_spenh(location: str, by: str = "MIXID", reverse: bool = False) -> str:
     import glob
@@ -56,10 +42,10 @@ def summarize_metric_spenh(location: str, by: str = "MIXID", reverse: bool = Fal
 def main():
     from docopt import docopt
 
-    import sonusai
+    from sonusai import __version__ as sai_version
     from sonusai.utils import trim_docstring
 
-    args = docopt(trim_docstring(__doc__), version=sonusai.__version__, options_first=True)
+    args = docopt(trim_docstring(__doc__), version=sai_version, options_first=True)
 
     by = args["--sort"]
     reverse = args["--reverse"]
@@ -69,4 +55,11 @@ def main():
 
 
 if __name__ == "__main__":
-    main()
+    from sonusai import exception_handler
+    from sonusai.utils import register_keyboard_interrupt
+
+    register_keyboard_interrupt()
+    try:
+        main()
+    except Exception as e:
+        exception_handler(e)