sonusai 1.0.16__cp311-abi3-macosx_10_12_x86_64.whl

sonusai/utils/keyboard_interrupt.py

@@ -0,0 +1,12 @@
+def register_keyboard_interrupt() -> None:
+    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)

sonusai/utils/load_object.py

@@ -0,0 +1,21 @@
+from functools import lru_cache
+from typing import Any
+
+
+def load_object(name: str, use_cache: bool = True) -> Any:
+    """Load an object from a pickle file"""
+    if use_cache:
+        return _load_object(name)
+    return _load_object.__wrapped__(name)
+
+
+@lru_cache
+def _load_object(name: str) -> Any:
+    import pickle
+    from os.path import exists
+
+    if exists(name):
+        with open(name, "rb") as f:
+            return pickle.load(f)  # noqa: S301
+
+    raise FileNotFoundError(name)

sonusai/utils/max_text_width.py

@@ -0,0 +1,9 @@
+def max_text_width(number_of_items: int) -> int:
+    """Compute maximum text width for the indices of a sequence of items.
+
+    :param number_of_items: Total number of items in sequence
+    :return: Text width of largest item index
+    """
+    import numpy as np
+
+    return int(np.ceil(np.log10(number_of_items)))

sonusai/utils/model_utils.py

@@ -0,0 +1,28 @@
+from typing import Any
+
+
+def import_module(name: str) -> Any:
+    """Import a Python module adding the module file's directory to the Python system path so that relative package
+    imports are found correctly.
+    """
+    import os
+    import sys
+    from importlib import import_module
+
+    try:
+        path = os.path.dirname(name)
+        if len(path) < 1:
+            path = "./"
+
+        # Add model file location to system path
+        sys.path.append(os.path.abspath(path))
+
+        try:
+            root = os.path.splitext(os.path.basename(name))[0]
+            model = import_module(root)
+        except Exception as e:
+            raise OSError(f"Error: could not import model from {name}: {e}.") from e
+    except Exception as e:
+        raise OSError(f"Error: could not find {name}: {e}.") from e
+
+    return model

sonusai/utils/numeric_conversion.py

@@ -0,0 +1,11 @@
+import numpy as np
+
+
+def int16_to_float(x: np.ndarray) -> np.ndarray:
+    """Convert an int16 array to a floating point array with range +/- 1"""
+    return x.astype(np.float32) / 32768
+
+
+def float_to_int16(x: np.ndarray) -> np.ndarray:
+    """Convert a floating point array with range +/- 1 to an int16 array"""
+    return (x * 32768).astype(np.int16)

sonusai/utils/onnx_utils.py

@@ -0,0 +1,155 @@
+from collections.abc import Sequence
+
+from onnx import ModelProto
+from onnx import ValueInfoProto
+from onnxruntime import InferenceSession
+from onnxruntime import NodeArg  # pyright: ignore [reportAttributeAccessIssue]
+from onnxruntime import SessionOptions  # pyright: ignore [reportAttributeAccessIssue]
+
+REQUIRED_HPARAMS = ("feature", "batch_size", "timesteps")
+
+
+def _extract_shapes(io: list[ValueInfoProto]) -> list[list[int] | str]:
+    shapes: list[list[int] | str] = []
+
+    # iterate through inputs of the graph to find shapes
+    for item in io:
+        # get tensor type: 0, 1, 2, etc.
+        tensor_type = item.type.tensor_type
+        # check if it has a shape
+        if tensor_type.HasField("shape"):
+            tmp_shape = []
+            # iterate through dimensions of the shape
+            for d in tensor_type.shape.dim:
+                if d.HasField("dim_value"):
+                    # known dimension, int value
+                    tmp_shape.append(d.dim_value)
+                elif d.HasField("dim_param"):
+                    # dynamic dim with symbolic name of d.dim_param; set size to 0
+                    tmp_shape.append(0)
+                else:
+                    # unknown dimension with no name; also set to 0
+                    tmp_shape.append(0)
+            # add as a list
+            shapes.append(tmp_shape)
+        else:
+            shapes.append("unknown rank")
+
+    return shapes
+
+
+def get_and_check_inputs(model: ModelProto) -> tuple[list[ValueInfoProto], list[list[int] | str]]:
+    from sonusai import logger
+
+    # ignore initializer inputs (only seen in older ONNX < v1.5)
+    initializer_names = [x.name for x in model.graph.initializer]
+    inputs = [i for i in model.graph.input if i.name not in initializer_names]
+    if len(inputs) != 1:
+        logger.warning(f"Warning: ONNX model has {len(inputs)} inputs; expected only 1")
+
+    # This one-liner works only if input has type and shape, returns a list
+    # shape0 = [d.dim_value for d in inputs[0].type.tensor_type.shape.dim]
+    shapes = _extract_shapes(inputs)
+
+    return inputs, shapes
+
+
+def get_and_check_outputs(model: ModelProto) -> tuple[list[ValueInfoProto], list[list[int] | str]]:
+    from sonusai import logger
+
+    outputs = list(model.graph.output)
+    if len(outputs) != 1:
+        logger.warning(f"Warning: ONNX model has {len(outputs)} outputs; expected only 1")
+
+    shapes = _extract_shapes(outputs)
+
+    return outputs, shapes
+
+
+def add_sonusai_metadata(model: ModelProto, hparams: dict) -> ModelProto:
+    """Add SonusAI hyperparameters as metadata to an ONNX model using 'hparams' key
+
+    :param model: ONNX model
+    :param hparams: dictionary of hyperparameters to add
+    :return: ONNX model
+
+    Note SonusAI conventions require models to have:
+        feature: Model feature type
+        batch_size: Model batch size
+        timesteps: Size of timestep dimension (0 for no dimension)
+    """
+    from sonusai import logger
+
+    # Note hparams should be a dict (i.e., extracted from checkpoint)
+    if eval(str(hparams)) != hparams:  # noqa: S307
+        raise TypeError("hparams is not a dict")
+    for key in REQUIRED_HPARAMS:
+        if key not in hparams:
+            logger.warning(f"Warning: SonusAI hyperparameters are missing: {key}")
+
+    meta = model.metadata_props.add()
+    meta.key = "hparams"
+    meta.value = str(hparams)
+
+    return model
+
+
+def get_sonusai_metadata(session: InferenceSession) -> dict | None:
+    """Get SonusAI hyperparameter metadata from an ONNX Runtime session."""
+    from sonusai import logger
+
+    meta = session.get_modelmeta()
+    if "hparams" not in meta.custom_metadata_map:
+        logger.warning("Warning: ONNX model metadata does not contain 'hparams'")
+        return None
+
+    hparams = eval(meta.custom_metadata_map["hparams"])  # noqa: S307
+    for key in REQUIRED_HPARAMS:
+        if key not in hparams:
+            logger.warning(f"Warning: ONNX model does not have required SonusAI hyperparameters: {key}")
+
+    return hparams
+
+
+def load_ort_session(
+    model_path: str, providers: Sequence[str | tuple[str, dict]] | None = None
+) -> tuple[InferenceSession, SessionOptions, str, dict | None, list[NodeArg], list[NodeArg]]:
+    from os.path import basename
+    from os.path import exists
+    from os.path import isfile
+    from os.path import splitext
+
+    import onnxruntime as ort
+
+    from sonusai import logger
+
+    if providers is None:
+        providers = ["CPUExecutionProvider"]
+
+    if exists(model_path) and isfile(model_path):
+        model_basename = basename(model_path)
+        model_root = splitext(model_basename)[0]
+        logger.info(f"Importing model from {model_basename}")
+        try:
+            session = ort.InferenceSession(model_path, providers=providers)
+            options = ort.SessionOptions()
+        except Exception as e:
+            logger.exception(f"Error: could not load ONNX model from {model_path}: {e}")
+            raise SystemExit(1) from e
+    else:
+        logger.exception(f"Error: model file does not exist: {model_path}")
+        raise SystemExit(1)
+
+    logger.info(f"Opened session with provider options: {session._provider_options}.")
+    hparams = get_sonusai_metadata(session)
+    if hparams is not None:
+        for key in REQUIRED_HPARAMS:
+            logger.info(f"  {key:12} {hparams[key]}")
+
+    inputs = session.get_inputs()
+    outputs = session.get_outputs()
+
+    # in_names = [n.name for n in session.get_inputs()]
+    # out_names = [n.name for n in session.get_outputs()]
+
+    return session, options, model_root, hparams, inputs, outputs

sonusai/utils/parallel.py

@@ -0,0 +1,162 @@
+import warnings
+from collections.abc import Callable
+from collections.abc import Iterable
+from multiprocessing import current_process
+from multiprocessing import get_context
+from typing import Any
+
+from tqdm import TqdmExperimentalWarning
+from tqdm.rich import tqdm
+
+warnings.filterwarnings(action="ignore", category=TqdmExperimentalWarning)
+
+track = tqdm
+
+CONTEXT = "fork"
+
+
+def par_track(
+    func: Callable,
+    *iterables: Iterable,
+    initializer: Callable[..., None] | None = None,
+    initargs: Iterable[Any] | None = None,
+    progress: tqdm | None = None,
+    num_cpus: int | float | None = None,
+    total: int | None = None,
+    no_par: bool = False,
+    pass_index: bool = False,
+) -> list[Any]:
+    """Performs a parallel ordered imap with tqdm progress."""
+    total_items = _calculate_total_items(iterables, total)
+    results: list[Any] = [None] * total_items
+
+    if no_par or current_process().daemon:
+        _execute_sequential(func, iterables, initializer, initargs, results, progress, pass_index)
+    else:
+        cpu_count = _determine_cpu_count(num_cpus, total_items)
+        _execute_parallel(func, iterables, initializer, initargs, results, progress, pass_index, cpu_count)
+
+    if progress is not None:
+        progress.close()
+    return results
+
+
+def _calculate_total_items(iterables: tuple[Iterable, ...], total: int | None) -> int:
+    """Calculate the total number of items to process."""
+    from collections.abc import Sized
+
+    if total is None:
+        return min(len(iterable) for iterable in iterables if isinstance(iterable, Sized))
+    return int(total)
+
+
+def _cpu_count() -> int:
+    """Get the number of CPUs available."""
+    from psutil import cpu_count
+
+    count = cpu_count()
+    if count is None:
+        return 1
+    return count
+
+
+def _determine_cpu_count(num_cpus: int | float | None, total_items: int) -> int:
+    """Determine the optimal number of CPUs to use."""
+    if num_cpus is None:
+        # Reserve 2 CPUs for system, minimum 1
+        optimal_cpus = max(_cpu_count() - 2, 1)
+    elif isinstance(num_cpus, float):
+        optimal_cpus = int(round(num_cpus * _cpu_count()))
+    else:
+        optimal_cpus = int(num_cpus)
+
+    return min(optimal_cpus, total_items)
+
+
+def _create_indexed_iterables(iterables: tuple[Iterable, ...]) -> tuple[Iterable, ...]:
+    """Create iterables that include the index as the first argument."""
+    # Get the first iterable to enumerate over
+    first_iterable = iterables[0]
+    remaining_iterables = iterables[1:]
+
+    # Create an enumerated version: (index, first_item), second_item, third_item, ...
+    indexed_first = enumerate(first_iterable)
+
+    if remaining_iterables:
+        return (indexed_first,) + remaining_iterables
+    else:
+        return (indexed_first,)
+
+
+class _IndexedFunctionWrapper:
+    """Pickle-able wrapper class for functions that need an index as the first argument."""
+
+    def __init__(self, func: Callable):
+        self.func = func
+
+    def __call__(self, indexed_first_arg, *remaining_args):
+        index, first_arg = indexed_first_arg
+        return self.func(index, first_arg, *remaining_args)
+
+
+def _wrap_function_with_index(func: Callable) -> Callable:
+    """Wrap a function to handle indexed arguments."""
+    return _IndexedFunctionWrapper(func)
+
+
+def _execute_sequential(
+    func: Callable,
+    iterables: tuple[Iterable, ...],
+    initializer: Callable[..., None] | None,
+    initargs: Iterable[Any] | None,
+    results: list[Any],
+    progress: tqdm | None,
+    pass_index: bool,
+) -> None:
+    """Execute a function sequentially without using multiprocessing."""
+    if initializer is not None:
+        if initargs is not None:
+            initializer(*initargs)
+        else:
+            initializer()
+
+    if pass_index:
+        mapped_iterables = _create_indexed_iterables(iterables)
+        wrapped_func = _wrap_function_with_index(func)
+        iterator = map(wrapped_func, *mapped_iterables)
+    else:
+        iterator = map(func, *iterables)
+
+    for index, result in enumerate(iterator):
+        results[index] = result
+        if progress is not None:
+            progress.update()
+
+
+def _execute_parallel(
+    func: Callable,
+    iterables: tuple[Iterable, ...],
+    initializer: Callable[..., None] | None,
+    initargs: Iterable[Any] | None,
+    results: list[Any],
+    progress: tqdm | None,
+    pass_index: bool,
+    cpu_count: int,
+) -> None:
+    """Execute a function in parallel using multiprocessing."""
+    init_args = initargs if initargs is not None else []
+
+    if pass_index:
+        mapped_iterables = _create_indexed_iterables(iterables)
+        wrapped_func = _wrap_function_with_index(func)
+    else:
+        mapped_iterables = iterables
+        wrapped_func = func
+
+    with get_context(CONTEXT).Pool(processes=cpu_count, initializer=initializer, initargs=init_args) as pool:
+        for index, result in enumerate(pool.imap(wrapped_func, *mapped_iterables, chunksize=1)):
+            results[index] = result
+            if progress is not None:
+                progress.update()
+        pool.close()
+        pool.join()

sonusai/utils/path_info.py

@@ -0,0 +1,7 @@
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class PathInfo:
+    abs_path: str
+    audio_filepath: str

sonusai/utils/print_mixture_details.py

@@ -0,0 +1,60 @@
+from collections.abc import Callable
+
+from ..datatypes import ClassCount
+from ..mixture.helpers import mixture_all_speech_metadata
+from ..mixture.mixdb import MixtureDatabase
+
+
+def print_mixture_details(
+    mixdb: MixtureDatabase,
+    mixid: int | None = None,
+    print_fn: Callable = print,
+) -> None:
+    from ..utils.seconds_to_hms import seconds_to_hms
+
+    if mixid is not None:
+        if 0 < mixid >= mixdb.num_mixtures:
+            raise ValueError(f"Given mixid is outside valid range of 0:{mixdb.num_mixtures - 1}.")
+
+        print_fn(f"Mixture {mixid} details")
+        mixture = mixdb.mixture(mixid)
+        speech_metadata = mixture_all_speech_metadata(mixdb, mixture)
+        for category, source in mixture.all_sources.items():
+            source_file = mixdb.source_file(source.file_id)
+            print_fn(f"  {category}")
+            print_fn(f"    name: {source_file.name}")
+            print_fn(f"    effects: {source.effects.to_dict()}")
+            print_fn(f"    pre_tempo: {source.pre_tempo}")
+            print_fn(f"    duration: {seconds_to_hms(source_file.duration)}")
+            print_fn(f"    start: {source.start}")
+            print_fn(f"    repeat: {source.loop}")
+            print_fn(f"    snr: {source.snr}")
+            print_fn(f"    random_snr: {source.snr.is_random}")
+            print_fn(f"    snr_gain: {source.snr_gain}")
+            for key in source_file.truth_configs:
+                print_fn(f"    truth '{key}' function: {source_file.truth_configs[key].function}")
+                print_fn(f"    truth '{key}' config: {source_file.truth_configs[key].config}")
+                print_fn(
+                    f"    truth '{key}' stride_reduction: {source_file.truth_configs[key].stride_reduction}"
+                )
+            for key in speech_metadata[category]:
+                print_fn(f"{category} speech {key}: {speech_metadata[category][key]}")
+        print_fn(f"  samples: {mixture.samples}")
+        print_fn(f"  feature frames: {mixdb.mixture_feature_frames(mixid)}")
+        print_fn("")
+
+
+def print_class_count(
+    class_count: ClassCount,
+    length: int,
+    print_fn: Callable = print,
+    all_class_counts: bool = False,
+) -> None:
+    from ..utils.max_text_width import max_text_width
+
+    print_fn("Class count:")
+    idx_len = max_text_width(len(class_count))
+    for idx, count in enumerate(class_count):
+        if all_class_counts or count > 0:
+            desc = f"  class {idx + 1:{idx_len}}"
+            print_fn(f"{desc:{length}} {count}")

sonusai/utils/rand.py

@@ -0,0 +1,13 @@
+import contextlib
+
+
+@contextlib.contextmanager
+def seed_context(seed):
+    import numpy as np
+
+    state = np.random.get_state()
+    np.random.seed(seed)
+    try:
+        yield
+    finally:
+        np.random.set_state(state)

sonusai/utils/ranges.py

@@ -0,0 +1,43 @@
+def expand_range(s: str, sort: bool = True) -> list[int]:
+    """Returns a list of integers from a string input representing a range."""
+    import re
+
+    clean_s = s.replace(":", "-")
+    clean_s = clean_s.replace(";", ",")
+    clean_s = re.sub(r" +", ",", clean_s)
+    clean_s = re.sub(r",+", ",", clean_s)
+
+    r: list[int] = []
+    for i in clean_s.split(","):
+        if "-" not in i:
+            r.append(int(i))
+        else:
+            lo, hi = map(int, i.split("-"))
+            r += range(lo, hi + 1)
+
+    if sort:
+        r = sorted(r)
+
+    return r
+
+
+def consolidate_range(r: list[int]) -> str:
+    """Returns a string representing a range from an input list of integers."""
+    from collections.abc import Generator
+
+    def ranges(i: list[int]) -> Generator[tuple[int, int], None, None]:
+        import itertools
+
+        for _, b in itertools.groupby(enumerate(i), lambda pair: pair[1] - pair[0]):
+            b_list = list(b)
+            yield b_list[0][1], b_list[-1][1]
+
+    ls: list[tuple[int, int]] = list(ranges(r))
+    result: list[str] = []
+    for val in ls:
+        entry = str(val[0])
+        if val[0] != val[1]:
+            entry += f"-{val[1]}"
+        result.append(entry)
+
+    return ", ".join(result)

sonusai/utils/read_predict_data.py

@@ -0,0 +1,32 @@
+import numpy as np
+
+from ..datatypes import Predict
+
+
+def read_predict_data(filename: str) -> Predict:
+    """Read predict data from given HDF5 file and return it."""
+    import h5py
+
+    from .. import logger
+
+    logger.debug(f"Reading prediction data from {filename}")
+    with h5py.File(filename, "r") as f:
+        # prediction data is either [frames, num_classes], or [frames, timesteps, num_classes]
+        predict = np.array(f["predict"])
+
+        if predict.ndim == 2:
+            return predict
+
+        if predict.ndim == 3:
+            frames, timesteps, num_classes = predict.shape
+
+            logger.debug(
+                f"Reshaping prediction data in {filename} "
+                f""
+                f"from [{frames}, {timesteps}, {num_classes}] "
+                f"to [{frames * timesteps}, {num_classes}]"
+            )
+            predict = np.reshape(predict, [frames * timesteps, num_classes], order="F")
+            return predict
+
+        raise RuntimeError(f"Invalid prediction data dimensions in {filename}")