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

sonusai/mixture/target_class_balancing.py

@@ -5,14 +5,15 @@ from sonusai.mixture.datatypes import TargetFile
 from sonusai.mixture.datatypes import TargetFiles
 
 
-def balance_targets(augmented_targets: AugmentedTargets,
-                    targets: TargetFiles,
-                    target_augmentations: AugmentationRules,
-                    class_balancing_augmentation: AugmentationRule,
-                    num_classes: int,
-                    truth_mutex: bool,
-                    num_ir: int,
-                    mixups: list[int] = None) -> tuple[AugmentedTargets, AugmentationRules]:
+def balance_targets(
+    augmented_targets: AugmentedTargets,
+    targets: TargetFiles,
+    target_augmentations: AugmentationRules,
+    class_balancing_augmentation: AugmentationRule,
+    num_classes: int,
+    num_ir: int,
+    mixups: list[int] | None = None,
+) -> tuple[AugmentedTargets, AugmentationRules]:
     import math
 
     from .augmentation import get_mixups
@@ -34,14 +35,13 @@ def balance_targets(augmented_targets: AugmentedTargets,
             target_augmentations=target_augmentations,
             mixup=mixup,
             num_classes=num_classes,
-            truth_mutex=truth_mutex)
+        )
 
         largest = max([len(item) for item in augmented_target_indices_by_class])
         largest = math.ceil(largest / mixup) * mixup
         for at_indices in augmented_target_indices_by_class:
             additional_augmentations_needed = largest - len(at_indices)
-            target_ids = sorted(
-                list(set([augmented_targets[at_index].target_id for at_index in at_indices])))
+            target_ids = sorted({augmented_targets[at_index].target_id for at_index in at_indices})
 
             tfi_idx = 0
             for _ in range(additional_augmentations_needed):
@@ -55,50 +55,55 @@ def balance_targets(augmented_targets: AugmentedTargets,
                     target_id=target_id,
                     mixup=mixup,
                     num_ir=num_ir,
-                    first_cba_id=first_cba_id)
-                augmented_target = AugmentedTarget(target_id=target_id,
-                                                   target_augmentation_id=augmentation_index)
+                    first_cba_id=first_cba_id,
+                )
+                augmented_target = AugmentedTarget(target_id=target_id, target_augmentation_id=augmentation_index)
                 augmented_targets.append(augmented_target)
 
     return augmented_targets, target_augmentations
 
 
-def _get_unused_balancing_augmentation(augmented_targets: AugmentedTargets,
-                                       targets: TargetFiles,
-                                       target_augmentations: AugmentationRules,
-                                       class_balancing_augmentation: AugmentationRule,
-                                       target_id: int,
-                                       mixup: int,
-                                       num_ir: int,
-                                       first_cba_id: int) -> tuple[int, AugmentationRules]:
-    """Get an unused balancing augmentation for a given target file index
-    """
+def _get_unused_balancing_augmentation(
+    augmented_targets: AugmentedTargets,
+    targets: TargetFiles,
+    target_augmentations: AugmentationRules,
+    class_balancing_augmentation: AugmentationRule,
+    target_id: int,
+    mixup: int,
+    num_ir: int,
+    first_cba_id: int,
+) -> tuple[int, AugmentationRules]:
+    """Get an unused balancing augmentation for a given target file index"""
     from dataclasses import asdict
 
     from .augmentation import get_augmentation_rules
 
     balancing_augmentations = [item for item in range(len(target_augmentations)) if item >= first_cba_id]
-    used_balancing_augmentations = [at.target_augmentation_id for at in augmented_targets if
-                                    at.target_id == target_id and
-                                    at.target_augmentation_id in balancing_augmentations]
-
-    augmentation_indices = [item for item in balancing_augmentations if
-                            item not in used_balancing_augmentations and
-                            target_augmentations[item].mixup == mixup]
+    used_balancing_augmentations = [
+        at.target_augmentation_id
+        for at in augmented_targets
+        if at.target_id == target_id and at.target_augmentation_id in balancing_augmentations
+    ]
+
+    augmentation_indices = [
+        item
+        for item in balancing_augmentations
+        if item not in used_balancing_augmentations and target_augmentations[item].mixup == mixup
+    ]
     if len(augmentation_indices) > 0:
         return augmentation_indices[0], target_augmentations
 
-    class_balancing_augmentation = get_class_balancing_augmentation(target=targets[target_id],
-                                                                    default_cba=class_balancing_augmentation)
+    class_balancing_augmentation = get_class_balancing_augmentation(
+        target=targets[target_id], default_cba=class_balancing_augmentation
+    )
     new_augmentation = get_augmentation_rules(rules=asdict(class_balancing_augmentation), num_ir=num_ir)[0]
     new_augmentation.mixup = mixup
     target_augmentations.append(new_augmentation)
     return len(target_augmentations) - 1, target_augmentations
 
 
-def get_class_balancing_augmentation(target: TargetFile, default_cba: AugmentationRule) -> AugmentationRule | None:
-    """ Get the class balancing augmentation rule for the given target
-    """
+def get_class_balancing_augmentation(target: TargetFile, default_cba: AugmentationRule) -> AugmentationRule:
+    """Get the class balancing augmentation rule for the given target"""
     if target.class_balancing_augmentation is not None:
         return target.class_balancing_augmentation
     return default_cba

sonusai/mixture/targets.py

@@ -1,13 +1,14 @@
 from sonusai.mixture.datatypes import AugmentationRules
 from sonusai.mixture.datatypes import AugmentedTarget
 from sonusai.mixture.datatypes import AugmentedTargets
-from sonusai.mixture.datatypes import TargetFile
 from sonusai.mixture.datatypes import TargetFiles
 
 
-def get_augmented_targets(target_files: TargetFiles,
-                          target_augmentations: AugmentationRules,
-                          mixups: list[int] = None) -> AugmentedTargets:
+def get_augmented_targets(
+    target_files: TargetFiles,
+    target_augmentations: AugmentationRules,
+    mixups: list[int] | None = None,
+) -> AugmentedTargets:
     from .augmentation import get_augmentation_indices_for_mixup
     from .augmentation import get_mixups
 
@@ -19,85 +20,82 @@ def get_augmented_targets(target_files: TargetFiles,
         augmentation_indices = get_augmentation_indices_for_mixup(target_augmentations, mixup)
         for target_index in range(len(target_files)):
             for augmentation_index in augmentation_indices:
-                augmented_targets.append(AugmentedTarget(target_id=target_index,
-                                                         target_augmentation_id=augmentation_index))
+                augmented_targets.append(
+                    AugmentedTarget(
+                        target_id=target_index,
+                        target_augmentation_id=augmentation_index,
+                    )
+                )
 
     return augmented_targets
 
 
-def get_truth_indices_for_target(target: TargetFile) -> list[int]:
-    """Get a list of truth indices for a given target."""
-    index = [truth_setting.index for truth_setting in target.truth_settings]
-
-    # flatten, uniquify, and sort
-    return sorted(list(set([item for sublist in index for item in sublist])))
-
-
-def get_truth_indices_for_augmented_target(augmented_target: AugmentedTarget, targets: TargetFiles) -> list[int]:
-    return get_truth_indices_for_target(targets[augmented_target.target_id])
+def get_class_index_for_augmented_target(augmented_target: AugmentedTarget, targets: TargetFiles) -> list[int]:
+    return targets[augmented_target.target_id].class_indices
 
 
 def get_mixup_for_augmented_target(augmented_target: AugmentedTarget, augmentations: AugmentationRules) -> int:
     return augmentations[augmented_target.target_augmentation_id].mixup
 
 
-def get_target_ids_for_truth_index(targets: TargetFiles,
-                                   truth_index: int,
-                                   allow_multiple: bool = False) -> list[int]:
-    """Get a list of target indices containing the given truth index.
+def get_target_ids_for_class_index(targets: TargetFiles, class_index: int, allow_multiple: bool = False) -> list[int]:
+    """Get a list of target indices containing the given class index.
 
-    If allow_multiple is True, then include targets that contain multiple truth indices.
+    If allow_multiple is True, then include targets that contain multiple class indices.
     """
     target_indices = set()
     for target_index, target in enumerate(targets):
-        indices = get_truth_indices_for_target(target)
+        indices = target.class_indices
         if len(indices) == 1 or allow_multiple:
             for index in indices:
-                if index == truth_index + 1:
+                if index == class_index + 1:
                     target_indices.add(target_index)
 
-    return sorted(list(target_indices))
+    return sorted(target_indices)
 
 
-def get_augmented_target_ids_for_truth_index(augmented_targets: AugmentedTargets,
-                                             targets: TargetFiles,
-                                             augmentations: AugmentationRules,
-                                             truth_index: int,
-                                             mixup: int,
-                                             allow_multiple: bool = False) -> list[int]:
-    """Get a list of augmented target indices containing the given truth index.
+def get_augmented_target_ids_for_class_index(
+    augmented_targets: AugmentedTargets,
+    targets: TargetFiles,
+    augmentations: AugmentationRules,
+    class_index: int,
+    mixup: int,
+    allow_multiple: bool = False,
+) -> list[int]:
+    """Get a list of augmented target indices containing the given class index.
 
-    If allow_multiple is True, then include targets that contain multiple truth indices.
+    If allow_multiple is True, then include targets that contain multiple class indices.
     """
     augmented_target_ids = set()
     for augmented_target_id, augmented_target in enumerate(augmented_targets):
         if get_mixup_for_augmented_target(augmented_target=augmented_target, augmentations=augmentations) == mixup:
-            indices = get_truth_indices_for_augmented_target(augmented_target=augmented_target, targets=targets)
+            indices = get_class_index_for_augmented_target(augmented_target=augmented_target, targets=targets)
             if len(indices) == 1 or allow_multiple:
                 for index in indices:
-                    if index == truth_index + 1:
+                    if index == class_index + 1:
                         augmented_target_ids.add(augmented_target_id)
 
-    return sorted(list(augmented_target_ids))
+    return sorted(augmented_target_ids)
 
 
-def get_augmented_target_ids_by_class(augmented_targets: AugmentedTargets,
-                                      targets: TargetFiles,
-                                      target_augmentations: AugmentationRules,
-                                      mixup: int,
-                                      num_classes: int,
-                                      truth_mutex: bool) -> list[list[int]]:
-    if truth_mutex:
-        num_classes -= 1
-
+def get_augmented_target_ids_by_class(
+    augmented_targets: AugmentedTargets,
+    targets: TargetFiles,
+    target_augmentations: AugmentationRules,
+    mixup: int,
+    num_classes: int,
+) -> list[list[int]]:
     indices = []
     for idx in range(num_classes):
         indices.append(
-            get_augmented_target_ids_for_truth_index(augmented_targets=augmented_targets,
-                                                     targets=targets,
-                                                     augmentations=target_augmentations,
-                                                     truth_index=idx,
-                                                     mixup=mixup))
+            get_augmented_target_ids_for_class_index(
+                augmented_targets=augmented_targets,
+                targets=targets,
+                augmentations=target_augmentations,
+                class_index=idx,
+                mixup=mixup,
+            )
+        )
     return indices
 
 
@@ -111,36 +109,40 @@ def get_target_augmentations_for_mixup(target_augmentations: AugmentationRules,
     return [target_augmentation for target_augmentation in target_augmentations if target_augmentation.mixup == mixup]
 
 
-def get_augmented_target_ids_for_mixup(augmented_targets: AugmentedTargets,
-                                       targets: TargetFiles,
-                                       target_augmentations: AugmentationRules,
-                                       mixup: int,
-                                       num_classes: int,
-                                       truth_mutex: bool) -> list[list[int]]:
+def get_augmented_target_ids_for_mixup(
+    augmented_targets: AugmentedTargets,
+    targets: TargetFiles,
+    target_augmentations: AugmentationRules,
+    mixup: int,
+    num_classes: int,
+) -> list[list[int]]:
     from collections import deque
     from random import shuffle
 
-    from sonusai import SonusAIError
-
     mixup_indices = []
 
     if mixup == 1:
         for index, augmented_target in enumerate(augmented_targets):
-            if get_mixup_for_augmented_target(augmented_target=augmented_target,
-                                              augmentations=target_augmentations) == 1:
+            if (
+                get_mixup_for_augmented_target(
+                    augmented_target=augmented_target,
+                    augmentations=target_augmentations,
+                )
+                == 1
+            ):
                 mixup_indices.append([index])
         return mixup_indices
 
-    augmented_target_ids_by_class = get_augmented_target_ids_by_class(augmented_targets=augmented_targets,
-                                                                      targets=targets,
-                                                                      target_augmentations=target_augmentations,
-                                                                      mixup=mixup,
-                                                                      num_classes=num_classes,
-                                                                      truth_mutex=truth_mutex)
+    augmented_target_ids_by_class = get_augmented_target_ids_by_class(
+        augmented_targets=augmented_targets,
+        targets=targets,
+        target_augmentations=target_augmentations,
+        mixup=mixup,
+        num_classes=num_classes,
+    )
 
     if mixup > num_classes:
-        raise SonusAIError(
-            f'Specified mixup, {mixup}, is greater than the number of classes, {num_classes}')
+        raise ValueError(f"Specified mixup, {mixup}, is greater than the number of classes, {num_classes}")
 
     de: deque[int] = deque()
 

sonusai/mixture/tokenized_shell_vars.py

@@ -23,10 +23,10 @@ def tokenized_expand(name: str | bytes | Path) -> tuple[str, dict[str, str]]:
 
     from .constants import DEFAULT_NOISE
 
-    os.environ['default_noise'] = str(DEFAULT_NOISE)
+    os.environ["default_noise"] = str(DEFAULT_NOISE)  # noqa: SIM112
 
     if isinstance(name, bytes):
-        name = name.decode('utf-8')
+        name = name.decode("utf-8")
 
     if isinstance(name, Path):
         name = name.as_posix()
@@ -34,37 +34,37 @@ def tokenized_expand(name: str | bytes | Path) -> tuple[str, dict[str, str]]:
     name = os.fspath(name)
     token_map: dict = {}
 
-    if '$' not in name and '%' not in name:
+    if "$" not in name and "%" not in name:
         return name, token_map
 
-    var_chars = string.ascii_letters + string.digits + '_-'
-    quote = '\''
-    percent = '%'
-    brace = '{'
-    rbrace = '}'
-    dollar = '$'
+    var_chars = string.ascii_letters + string.digits + "_-"
+    quote = "'"
+    percent = "%"
+    brace = "{"
+    rbrace = "}"
+    dollar = "$"
     environ = os.environ
 
     result = name[:0]
     index = 0
     path_len = len(name)
     while index < path_len:
-        c = name[index:index + 1]
+        c = name[index : index + 1]
         if c == quote:  # no expansion within single quotes
-            name = name[index + 1:]
+            name = name[index + 1 :]
             path_len = len(name)
             try:
                 index = name.index(c)
-                result += c + name[:index + 1]
+                result += c + name[: index + 1]
             except ValueError:
                 result += c + name
                 index = path_len - 1
         elif c == percent:  # variable or '%'
-            if name[index + 1:index + 2] == percent:
+            if name[index + 1 : index + 2] == percent:
                 result += c
                 index += 1
             else:
-                name = name[index + 1:]
+                name = name[index + 1 :]
                 path_len = len(name)
                 try:
                     index = name.index(percent)
@@ -75,7 +75,7 @@ def tokenized_expand(name: str | bytes | Path) -> tuple[str, dict[str, str]]:
                     var = name[:index]
                     try:
                         if environ is None:
-                            value = os.fsencode(os.environ[os.fsdecode(var)]).decode('utf-8')
+                            value = os.fsencode(os.environ[os.fsdecode(var)]).decode("utf-8")  # type: ignore[unreachable]
                         else:
                             value = environ[var]
                         token_map[var] = value
@@ -83,11 +83,11 @@ def tokenized_expand(name: str | bytes | Path) -> tuple[str, dict[str, str]]:
                         value = percent + var + percent
                     result += value
         elif c == dollar:  # variable or '$$'
-            if name[index + 1:index + 2] == dollar:
+            if name[index + 1 : index + 2] == dollar:
                 result += c
                 index += 1
-            elif name[index + 1:index + 2] == brace:
-                name = name[index + 2:]
+            elif name[index + 1 : index + 2] == brace:
+                name = name[index + 2 :]
                 path_len = len(name)
                 try:
                     index = name.index(rbrace)
@@ -98,7 +98,7 @@ def tokenized_expand(name: str | bytes | Path) -> tuple[str, dict[str, str]]:
                     var = name[:index]
                     try:
                         if environ is None:
-                            value = os.fsencode(os.environ[os.fsdecode(var)]).decode('utf-8')
+                            value = os.fsencode(os.environ[os.fsdecode(var)]).decode("utf-8")  # type: ignore[unreachable]
                         else:
                             value = environ[var]
                         token_map[var] = value
@@ -108,14 +108,14 @@ def tokenized_expand(name: str | bytes | Path) -> tuple[str, dict[str, str]]:
             else:
                 var = name[:0]
                 index += 1
-                c = name[index:index + 1]
+                c = name[index : index + 1]
                 while c and c in var_chars:
                     var += c
                     index += 1
-                    c = name[index:index + 1]
+                    c = name[index : index + 1]
                 try:
                     if environ is None:
-                        value = os.fsencode(os.environ[os.fsdecode(var)]).decode('utf-8')
+                        value = os.fsencode(os.environ[os.fsdecode(var)]).decode("utf-8")  # type: ignore[unreachable]
                     else:
                         value = environ[var]
                     token_map[var] = value
@@ -139,5 +139,5 @@ def tokenized_replace(name: str, tokens: dict[str, str]) -> str:
     :return: replaced string
     """
     for key, value in tokens.items():
-        name = name.replace(value, f'${key}')
+        name = name.replace(value, f"${key}")
     return name

sonusai/mixture/torchaudio_audio.py

@@ -14,19 +14,18 @@ def read_impulse_response(name: str | Path) -> ImpulseResponseData:
     import torch
     import torchaudio
 
-    from sonusai import SonusAIError
     from .tokenized_shell_vars import tokenized_expand
 
     expanded_name, _ = tokenized_expand(name)
 
     # Read impulse response data from audio file
     try:
-        raw, sample_rate = torchaudio.load(expanded_name, backend='soundfile')
+        raw, sample_rate = torchaudio.load(expanded_name, backend="soundfile")
     except Exception as e:
         if name != expanded_name:
-            raise SonusAIError(f'Error reading {name} (expanded: {expanded_name}): {e}')
+            raise OSError(f"Error reading {name} (expanded: {expanded_name}): {e}") from e
         else:
-            raise SonusAIError(f'Error reading {name}: {e}')
+            raise OSError(f"Error reading {name}: {e}") from e
 
     raw = torch.squeeze(raw[0, :])
     offset = torch.argmax(raw)
@@ -49,7 +48,6 @@ def get_sample_rate(name: str | Path) -> int:
     """
     import torchaudio
 
-    from sonusai import SonusAIError
     from .tokenized_shell_vars import tokenized_expand
 
     expanded_name, _ = tokenized_expand(name)
@@ -58,9 +56,9 @@ def get_sample_rate(name: str | Path) -> int:
         return torchaudio.info(expanded_name).sample_rate
     except Exception as e:
         if name != expanded_name:
-            raise SonusAIError(f'Error reading {name} (expanded: {expanded_name}):\n{e}')
+            raise OSError(f"Error reading {name} (expanded: {expanded_name}):\n{e}") from e
         else:
-            raise SonusAIError(f'Error reading {name}:\n{e}')
+            raise OSError(f"Error reading {name}:\n{e}") from e
 
 
 def read_audio(name: str | Path) -> AudioT:
@@ -73,24 +71,25 @@ def read_audio(name: str | Path) -> AudioT:
     import torch
     import torchaudio
 
-    from sonusai import SonusAIError
     from .constants import SAMPLE_RATE
     from .tokenized_shell_vars import tokenized_expand
 
     expanded_name, _ = tokenized_expand(name)
 
     try:
-        out, samplerate = torchaudio.load(expanded_name, backend='soundfile')
+        out, samplerate = torchaudio.load(expanded_name, backend="soundfile")
         out = torch.reshape(out[0, :], (1, out.size()[1]))
-        out = torchaudio.functional.resample(out,
-                                             orig_freq=samplerate,
-                                             new_freq=SAMPLE_RATE,
-                                             resampling_method='sinc_interp_hann')
+        out = torchaudio.functional.resample(
+            out,
+            orig_freq=samplerate,
+            new_freq=SAMPLE_RATE,
+            resampling_method="sinc_interp_hann",
+        )
     except Exception as e:
         if name != expanded_name:
-            raise SonusAIError(f'Error reading {name} (expanded: {expanded_name}):\n{e}')
+            raise OSError(f"Error reading {name} (expanded: {expanded_name}):\n{e}") from e
         else:
-            raise SonusAIError(f'Error reading {name}:\n{e}')
+            raise OSError(f"Error reading {name}:\n{e}") from e
 
     result = np.squeeze(np.array(out))
     return result

sonusai/mixture/torchaudio_augmentation.py

@@ -3,9 +3,7 @@ from sonusai.mixture.datatypes import Augmentation
 from sonusai.mixture.datatypes import ImpulseResponseData
 
 
-def apply_augmentation(audio: AudioT,
-                       augmentation: Augmentation,
-                       frame_length: int = 1) -> AudioT:
+def apply_augmentation(audio: AudioT, augmentation: Augmentation, frame_length: int = 1) -> AudioT:
     """Apply augmentations to audio data using torchaudio.sox_effects
 
     :param audio: Audio
@@ -17,7 +15,6 @@ def apply_augmentation(audio: AudioT,
     import torch
     import torchaudio
 
-    from sonusai import SonusAIError
     from .augmentation import pad_audio_to_frame
     from .constants import SAMPLE_RATE
 
@@ -28,29 +25,29 @@ def apply_augmentation(audio: AudioT,
     #  Normalize to globally set level (should this be a global config parameter,
     #  or hard-coded into the script?)
     if augmentation.normalize is not None:
-        effects.append(['norm', str(augmentation.normalize)])
+        effects.append(["norm", str(augmentation.normalize)])
 
     if augmentation.gain is not None:
-        effects.append(['gain', str(augmentation.gain)])
+        effects.append(["gain", str(augmentation.gain)])
 
     if augmentation.pitch is not None:
-        effects.append(['pitch', str(augmentation.pitch)])
-        effects.append(['rate', str(SAMPLE_RATE)])
+        effects.append(["pitch", str(augmentation.pitch)])
+        effects.append(["rate", str(SAMPLE_RATE)])
 
     if augmentation.tempo is not None:
-        effects.append(['tempo', '-s', str(augmentation.tempo)])
+        effects.append(["tempo", "-s", str(augmentation.tempo)])
 
     if augmentation.eq1 is not None:
-        effects.append(['equalizer', *[str(item) for item in augmentation.eq1]])
+        effects.append(["equalizer", *[str(item) for item in augmentation.eq1]])
 
     if augmentation.eq2 is not None:
-        effects.append(['equalizer', *[str(item) for item in augmentation.eq2]])
+        effects.append(["equalizer", *[str(item) for item in augmentation.eq2]])
 
     if augmentation.eq3 is not None:
-        effects.append(['equalizer', *[str(item) for item in augmentation.eq3]])
+        effects.append(["equalizer", *[str(item) for item in augmentation.eq3]])
 
     if augmentation.lpf is not None:
-        effects.append(['lowpass', '-2', str(augmentation.lpf), '0.707'])
+        effects.append(["lowpass", "-2", str(augmentation.lpf), "0.707"])
 
     if effects:
         if audio.ndim == 1:
@@ -58,11 +55,9 @@ def apply_augmentation(audio: AudioT,
         out = torch.tensor(audio)
 
         try:
-            out, _ = torchaudio.sox_effects.apply_effects_tensor(out,
-                                                                 sample_rate=SAMPLE_RATE,
-                                                                 effects=effects)
+            out, _ = torchaudio.sox_effects.apply_effects_tensor(out, sample_rate=SAMPLE_RATE, effects=effects)
         except Exception as e:
-            raise SonusAIError(f'Error applying {augmentation}: {e}')
+            raise RuntimeError(f"Error applying {augmentation}: {e}") from e
 
         audio_out = np.squeeze(np.array(out))
     else:
@@ -84,6 +79,7 @@ def apply_impulse_response(audio: AudioT, ir: ImpulseResponseData) -> AudioT:
     import torchaudio
 
     from sonusai.utils import linear_to_db
+
     from .constants import SAMPLE_RATE
 
     # Early exit if no ir or if all audio is zero
@@ -95,20 +91,20 @@ def apply_impulse_response(audio: AudioT, ir: ImpulseResponseData) -> AudioT:
 
     # Convert audio to IR sample rate
     audio_in = torch.reshape(torch.tensor(audio), (1, len(audio)))
-    audio_out, sr = torchaudio.sox_effects.apply_effects_tensor(audio_in,
-                                                                sample_rate=SAMPLE_RATE,
-                                                                effects=[['rate', str(ir.sample_rate)]])
+    audio_out, sr = torchaudio.sox_effects.apply_effects_tensor(
+        audio_in, sample_rate=SAMPLE_RATE, effects=[["rate", str(ir.sample_rate)]]
+    )
 
     # Apply IR and convert back to global sample rate
     rir = torch.reshape(torch.tensor(ir.data), (1, len(ir.data)))
     audio_out = torchaudio.functional.fftconvolve(audio_out, rir)
-    audio_out, sr = torchaudio.sox_effects.apply_effects_tensor(audio_out,
-                                                                sample_rate=ir.sample_rate,
-                                                                effects=[['rate', str(SAMPLE_RATE)]])
+    audio_out, sr = torchaudio.sox_effects.apply_effects_tensor(
+        audio_out, sample_rate=ir.sample_rate, effects=[["rate", str(SAMPLE_RATE)]]
+    )
 
     # Reset level to previous max value
-    audio_out, sr = torchaudio.sox_effects.apply_effects_tensor(audio_out,
-                                                                sample_rate=SAMPLE_RATE,
-                                                                effects=[['norm', str(max_db)]])
+    audio_out, sr = torchaudio.sox_effects.apply_effects_tensor(
+        audio_out, sample_rate=SAMPLE_RATE, effects=[["norm", str(max_db)]]
+    )
 
-    return np.squeeze(np.array(audio_out[:, :len(audio)]))
+    return np.squeeze(np.array(audio_out[:, : len(audio)]))