sonusai 0.18.9__py3-none-any.whl → 0.19.6__py3-none-any.whl

sonusai/utils/engineering_number.py

@@ -2,90 +2,92 @@ from __future__ import annotations
 
 from decimal import Decimal
 from typing import Any
+from typing import ClassVar
 
 
 class EngineeringNumber:
     """Easy manipulation of numbers which use engineering notation"""
 
-    _suffix_lookup = {
-        'Y': 'e24',
-        'Z': 'e21',
-        'E': 'e18',
-        'P': 'e15',
-        'T': 'e12',
-        'G': 'e9',
-        'M': 'e6',
-        'k': 'e3',
-        '':  'e0',
-        'm': 'e-3',
-        'u': 'e-6',
-        'n': 'e-9',
-        'p': 'e-12',
-        'f': 'e-15',
-        'a': 'e-18',
-        'z': 'e-21',
-        'y': 'e-24',
+    _suffix_lookup: ClassVar = {
+        "Y": "e24",
+        "Z": "e21",
+        "E": "e18",
+        "P": "e15",
+        "T": "e12",
+        "G": "e9",
+        "M": "e6",
+        "k": "e3",
+        "": "e0",
+        "m": "e-3",
+        "u": "e-6",
+        "n": "e-9",
+        "p": "e-12",
+        "f": "e-15",
+        "a": "e-18",
+        "z": "e-21",
+        "y": "e-24",
     }
 
-    _exponent_lookup_scaled = {
-        '-12': 'Y',
-        '-15': 'Z',
-        '-18': 'E',
-        '-21': 'P',
-        '-24': 'T',
-        '-27': 'G',
-        '-30': 'M',
-        '-33': 'k',
-        '-36': '',
-        '-39': 'm',
-        '-42': 'u',
-        '-45': 'n',
-        '-48': 'p',
-        '-51': 'f',
-        '-54': 'a',
-        '-57': 'z',
-        '-60': 'y',
+    _exponent_lookup_scaled: ClassVar = {
+        "-12": "Y",
+        "-15": "Z",
+        "-18": "E",
+        "-21": "P",
+        "-24": "T",
+        "-27": "G",
+        "-30": "M",
+        "-33": "k",
+        "-36": "",
+        "-39": "m",
+        "-42": "u",
+        "-45": "n",
+        "-48": "p",
+        "-51": "f",
+        "-54": "a",
+        "-57": "z",
+        "-60": "y",
     }
 
-    def __init__(self,
-                 value: str | float | int | EngineeringNumber,
-                 precision: int = 2,
-                 significant: int = 0):
+    def __init__(
+        self,
+        value: str | float | int | EngineeringNumber,
+        precision: int = 2,
+        significant: int = 0,
+    ):
         """
         :param value: string, integer, or float representing the numeric value of the number
         :param precision: the precision past the decimal
         :param significant: the number of significant digits
         if given, significant takes precedence over precision
         """
-        from sonusai import SonusAIError
-
         self.precision = precision
         self.significant = significant
 
         if isinstance(value, str):
-            suffix_keys = [key for key in self._suffix_lookup.keys() if key != '']
+            suffix_keys = [key for key in self._suffix_lookup if key != ""]
 
+            str_value = str(value)
             for suffix in suffix_keys:
-                if suffix in value:
-                    value = value[:-1] + self._suffix_lookup[suffix]
+                if suffix in str_value:
+                    str_value = str_value[:-1] + self._suffix_lookup[suffix]
                     break
 
-            self.number = Decimal(value)
+            self.number = Decimal(str_value)
 
-        elif isinstance(value, int) or isinstance(value, float) or isinstance(value, EngineeringNumber):
+        elif isinstance(value, int | float | EngineeringNumber):
             self.number = Decimal(str(value))
 
         else:
-            raise SonusAIError('value has unsupported type')
+            raise TypeError("value has unsupported type")
 
     def __repr__(self):
         """Returns the string representation"""
         # The Decimal class only really converts numbers that are very small into engineering notation.
         # So we will simply make all numbers small numbers and take advantage of the Decimal class.
-        number_str = self.number * Decimal('10e-37')
+        number_str = self.number * Decimal("10e-37")
         number_str = number_str.to_eng_string().lower()
 
-        base, exponent = number_str.split('e')
+        base, exponent = number_str.split("e")
 
         if self.significant > 0:
             abs_base = abs(Decimal(base))
@@ -98,12 +100,12 @@ class EngineeringNumber:
 
         base = str(round(Decimal(base), num_digits))
 
-        if 'e' in base.lower():
+        if "e" in base.lower():
             base = str(int(Decimal(base)))
 
         # Remove trailing decimal
-        if '.' in base:
-            base = base.rstrip('.')
+        if "." in base:
+            base = base.rstrip(".")
 
         return base + self._exponent_lookup_scaled[exponent]
 
@@ -159,6 +161,6 @@ class EngineeringNumber:
         return self.number >= self._to_decimal(other)
 
     def __eq__(self, other: Any) -> bool:
-        if not isinstance(other, (str, float, int, EngineeringNumber)):
+        if not isinstance(other, str | float | int | EngineeringNumber):
             return NotImplemented
         return self.number == self._to_decimal(other)

sonusai/utils/get_label_names.py

@@ -1,22 +1,20 @@
-def get_label_names(num_labels: int, file: str = None) -> list:
+def get_label_names(num_labels: int, file: str | None = None) -> list:
     """Return label names in a list. Read from CSV file, if provided."""
     import csv
 
-    from sonusai import SonusAIError
-
     if file is None:
-        return [f'Class {val + 1}' for val in range(num_labels)]
+        return [f"Class {val + 1}" for val in range(num_labels)]
 
-    label_names = [''] * num_labels
+    label_names = [""] * num_labels
     with open(file) as f:
         reader = csv.DictReader(f)
-        if 'index' not in reader.fieldnames or 'display_name' not in reader.fieldnames:
-            raise SonusAIError(f'Missing required fields in labels CSV.')
+        if reader.fieldnames is None or "index" not in reader.fieldnames or "display_name" not in reader.fieldnames:
+            raise ValueError("Missing required fields in labels CSV.")
 
         for row in reader:
-            index = int(row['index']) - 1
+            index = int(row["index"]) - 1
             if index >= num_labels:
-                raise SonusAIError(f'The number of given label names does not match the number of labels.')
-            label_names[index] = row['display_name']
+                raise ValueError("The number of given label names does not match the number of labels.")
+            label_names[index] = row["display_name"]
 
     return label_names

sonusai/utils/human_readable_size.py

@@ -1,7 +1,7 @@
 def human_readable_size(size: float, decimal_places: int = 3) -> str:
     """Convert number into string with units"""
-    for unit in ['B', 'kB', 'MB', 'GB', 'TB']:
+    for unit in ["B", "kB", "MB", "GB", "TB"]:  # noqa: B007
         if size < 1024.0:
             break
         size /= 1024.0
-    return f'{size:.{decimal_places}f} {unit}'
+    return f"{size:.{decimal_places}f} {unit}"

sonusai/utils/model_utils.py

@@ -9,12 +9,10 @@ def import_module(name: str) -> Any:
     import sys
     from importlib import import_module
 
-    from sonusai import SonusAIError
-
     try:
         path = os.path.dirname(name)
         if len(path) < 1:
-            path = './'
+            path = "./"
 
         # Add model file location to system path
         sys.path.append(os.path.abspath(path))
@@ -23,8 +21,8 @@ def import_module(name: str) -> Any:
             root = os.path.splitext(os.path.basename(name))[0]
             model = import_module(root)
         except Exception as e:
-            raise SonusAIError(f'Error: could not import model from {name}: {e}.')
+            raise OSError(f"Error: could not import model from {name}: {e}.") from e
     except Exception as e:
-        raise SonusAIError(f'Error: could not find {name}: {e}.')
+        raise OSError(f"Error: could not find {name}: {e}.") from e
 
     return model

sonusai/utils/numeric_conversion.py

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

sonusai/utils/onnx_utils.py

@@ -1,5 +1,4 @@
-from typing import Optional
-from typing import Sequence
+from collections.abc import Sequence
 
 from onnx import ModelProto
 from onnx import ValueInfoProto
@@ -7,7 +6,14 @@ from onnxruntime import InferenceSession
 from onnxruntime import NodeArg
 from onnxruntime import SessionOptions
 
-REQUIRED_HPARAMS = ('feature', 'batch_size', 'timesteps', 'flatten', 'add1ch', 'truth_mutex')
+REQUIRED_HPARAMS = (
+    "feature",
+    "batch_size",
+    "timesteps",
+    "flatten",
+    "add1ch",
+    "truth_mutex",
+)
 
 
 def _extract_shapes(io: list[ValueInfoProto]) -> list[list[int] | str]:
@@ -18,14 +24,14 @@ def _extract_shapes(io: list[ValueInfoProto]) -> list[list[int] | str]:
         # get tensor type: 0, 1, 2, etc.
         tensor_type = item.type.tensor_type
         # check if it has a shape
-        if tensor_type.HasField('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'):
+                if d.HasField("dim_value"):
                     # known dimension, int value
                     tmp_shape.append(d.dim_value)
-                elif d.HasField('dim_param'):
+                elif d.HasField("dim_param"):
                     # dynamic dim with symbolic name of d.dim_param; set size to 0
                     tmp_shape.append(0)
                 else:
@@ -34,19 +40,21 @@ def _extract_shapes(io: list[ValueInfoProto]) -> list[list[int] | str]:
             # add as a list
             shapes.append(tmp_shape)
         else:
-            shapes.append('unknown rank')
+            shapes.append("unknown rank")
 
     return shapes
 
 
-def get_and_check_inputs(model: ModelProto) -> tuple[list[ValueInfoProto], list[list[int] | str]]:
+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')
+        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]
@@ -55,12 +63,14 @@ def get_and_check_inputs(model: ModelProto) -> tuple[list[ValueInfoProto], list[
     return inputs, shapes
 
 
-def get_and_check_outputs(model: ModelProto) -> tuple[list[ValueInfoProto], list[list[int] | str]]:
+def get_and_check_outputs(
+    model: ModelProto,
+) -> tuple[list[ValueInfoProto], list[list[int] | str]]:
     from sonusai import logger
 
-    outputs = [o for o in model.graph.output]
+    outputs = list(model.graph.output)
     if len(outputs) != 1:
-        logger.warning(f'Warning: ONNX model has {len(outputs)} outputs; expected only 1')
+        logger.warning(f"Warning: ONNX model has {len(outputs)} outputs; expected only 1")
 
     shapes = _extract_shapes(outputs)
 
@@ -85,38 +95,39 @@ def add_sonusai_metadata(model: ModelProto, hparams: dict) -> ModelProto:
     from sonusai import logger
 
     # Note hparams should be a dict (i.e., extracted from checkpoint)
-    assert eval(str(hparams)) == hparams
+    if eval(str(hparams)) != hparams:  # noqa: S307
+        raise TypeError("hparams is not a dict")
     for key in REQUIRED_HPARAMS:
-        if key not in hparams.keys():
-            logger.warning(f'Warning: SonusAI hyperparameters are missing: {key}')
+        if key not in hparams:
+            logger.warning(f"Warning: SonusAI hyperparameters are missing: {key}")
 
     meta = model.metadata_props.add()
-    meta.key = 'hparams'
+    meta.key = "hparams"
     meta.value = str(hparams)
 
     return model
 
 
-def get_sonusai_metadata(session: InferenceSession) -> Optional[dict]:
-    """Get SonusAI hyperparameter metadata from an ONNX Runtime session.
-    """
+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.keys():
+    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'])
+    hparams = eval(meta.custom_metadata_map["hparams"])  # noqa: S307
     for key in REQUIRED_HPARAMS:
-        if key not in hparams.keys():
-            logger.warning(f'Warning: ONNX model does not have required SonusAI hyperparameters: {key}')
+        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) -> tuple[
-    InferenceSession, SessionOptions, str, dict, list[NodeArg], list[NodeArg]]:
+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
@@ -127,27 +138,27 @@ def load_ort_session(model_path: str, providers: Sequence[str | tuple[str, dict]
     from sonusai import logger
 
     if providers is None:
-        providers = ['CPUExecutionProvider']
+        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}')
+        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)
+            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}')
+        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}.')
+    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]}')
+            logger.info(f"  {key:12} {hparams[key]}")
 
     inputs = session.get_inputs()
     outputs = session.get_outputs()

sonusai/utils/parallel.py

@@ -1,32 +1,41 @@
+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 typing import Callable
-from typing import Iterable
-from typing import Optional
 
-from tqdm import tqdm
+from tqdm import TqdmExperimentalWarning
+from tqdm.rich import tqdm
 
-CONTEXT = 'fork'
+warnings.filterwarnings(action="ignore", category=TqdmExperimentalWarning)
 
+track = tqdm
 
-def pp_tqdm_imap(func: Callable,
-                 *iterables: Iterable,
-                 initializer: Optional[Callable[..., None]] = None,
-                 initargs: Optional[Iterable[Any]] = None,
-                 progress: Optional[tqdm] = None,
-                 num_cpus: Optional[int] = None,
-                 total: Optional[int] = None,
-                 no_par: bool = False) -> list[Any]:
+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,
+) -> list[Any]:
     """Performs a parallel ordered imap with tqdm progress."""
-    from os import cpu_count
-    from typing import Sized
+    from collections.abc import Sized
+
+    from psutil import cpu_count
 
     if total is None:
-        total = min(len(iterable) for iterable in iterables if isinstance(iterable, Sized))
+        _total = min(len(iterable) for iterable in iterables if isinstance(iterable, Sized))
+    else:
+        _total = int(total)
 
-    results: list[Any] = [None] * total
-    n = 0
+    results: list[Any] = [None] * _total
     if no_par or current_process().daemon:
         if initializer is not None:
             if initargs is not None:
@@ -34,21 +43,26 @@ def pp_tqdm_imap(func: Callable,
             else:
                 initializer()
 
-        for result in map(func, *iterables):
+        for n, result in enumerate(map(func, *iterables)):
             results[n] = result
-            n += 1
             if progress is not None:
                 progress.update()
     else:
         if num_cpus is None:
-            num_cpus = cpu_count()
-        elif num_cpus is float:
-            num_cpus = int(round(num_cpus * cpu_count()))
+            _num_cpus = max(cpu_count() - 2, 1)
+        elif isinstance(num_cpus, float):
+            _num_cpus = int(round(num_cpus * cpu_count()))
+        else:
+            _num_cpus = int(num_cpus)
 
-        if total < num_cpus:
-            num_cpus = total
-        with get_context(CONTEXT).Pool(processes=num_cpus, initializer=initializer, initargs=initargs) as pool:
-            for result in pool.imap(func, *iterables):  # type: ignore
+        _num_cpus = min(_num_cpus, _total)
+
+        if initargs is None:
+            initargs = []
+
+        with get_context(CONTEXT).Pool(processes=_num_cpus, initializer=initializer, initargs=initargs) as pool:
+            n = 0
+            for result in pool.imap(func, *iterables):  # type: ignore[arg-type]
                 results[n] = result
                 n += 1
                 if progress is not None:
@@ -59,6 +73,3 @@ def pp_tqdm_imap(func: Callable,
     if progress is not None:
         progress.close()
     return results
-
-
-pp_imap = pp_tqdm_imap

sonusai/utils/print_mixture_details.py

@@ -1,37 +1,38 @@
-from typing import Callable
+from collections.abc import Callable
 
 from sonusai.mixture import ClassCount
 from sonusai.mixture import MixtureDatabase
 
 
-def print_mixture_details(mixdb: MixtureDatabase,
-                          mixid: int = None,
-                          desc_len: int = 1,
-                          print_fn: Callable = print) -> None:
+def print_mixture_details(
+    mixdb: MixtureDatabase,
+    mixid: int | None = None,
+    desc_len: int = 1,
+    print_fn: Callable = print,
+) -> None:
     import numpy as np
 
-    from sonusai import SonusAIError
     from sonusai.mixture import SAMPLE_RATE
     from sonusai.utils import seconds_to_hms
 
     if mixid is not None:
         if 0 < mixid >= mixdb.num_mixtures:
-            raise SonusAIError(f'Given mixid is outside valid range of 0:{mixdb.num_mixtures - 1}.')
+            raise ValueError(f"Given mixid is outside valid range of 0:{mixdb.num_mixtures - 1}.")
 
-        print_fn(f'Mixture {mixid} details')
+        print_fn(f"Mixture {mixid} details")
         mixture = mixdb.mixture(mixid)
         target_files = [mixdb.target_files[target.file_id] for target in mixture.targets]
         target_augmentations = [target.augmentation for target in mixture.targets]
         noise_file = mixdb.noise_file(mixture.noise.file_id)
         for t_idx, target_file in enumerate(target_files):
-            print_fn(f'  Target {t_idx}')
+            print_fn(f"  Target {t_idx}")
             print_fn(f'{"    Name":{desc_len}} {target_file.name}')
             print_fn(f'{"    Duration":{desc_len}} {seconds_to_hms(target_file.duration)}')
-            for ts_idx, truth_setting in enumerate(target_file.truth_settings):
-                print_fn(f'    Truth setting {ts_idx}')
-                print_fn(f'{"      Index":{desc_len}} {truth_setting.index}')
-                print_fn(f'{"      Function":{desc_len}} {truth_setting.function}')
-                print_fn(f'{"      Config":{desc_len}} {truth_setting.config}')
+            for truth_name, truth_config in target_file.truth_configs.items():
+                print_fn(f"    Truth config: '{truth_name}'")
+                print_fn(f'{"      Function":{desc_len}} {truth_config.function}')
+                print_fn(f'{"      Stride reduction":{desc_len}} {truth_config.stride_reduction}')
+                print_fn(f'{"      Config":{desc_len}} {truth_config.config}')
             print_fn(f'{"    Augmentation":{desc_len}} {target_augmentations[t_idx]}')
         print_fn(f'{"  Samples":{desc_len}} {mixture.samples}')
         print_fn(f'{"  Feature frames":{desc_len}} {mixdb.mixture_feature_frames(mixid)}')
@@ -42,18 +43,20 @@ def print_mixture_details(mixdb: MixtureDatabase,
         print_fn(f'{"  Target gain":{desc_len}} {[target.gain for target in mixture.targets]}')
         print_fn(f'{"  Target SNR gain":{desc_len}} {mixture.target_snr_gain}')
         print_fn(f'{"  Noise SNR gain":{desc_len}} {mixture.noise_snr_gain}')
-        print_fn('')
+        print_fn("")
 
 
-def print_class_count(class_count: ClassCount,
-                      length: int,
-                      print_fn: Callable = print,
-                      all_class_counts: bool = False) -> None:
+def print_class_count(
+    class_count: ClassCount,
+    length: int,
+    print_fn: Callable = print,
+    all_class_counts: bool = False,
+) -> None:
     from sonusai.utils import max_text_width
 
-    print_fn(f'Class count:')
+    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}')
+            desc = f"  class {idx + 1:{idx_len}}"
+            print_fn(f"{desc:{length}} {count}")