Apply repository Python formatting

Signed-off-by: Piotr Żelasko <pzelasko@nvidia.com>

Author: pzelasko

nemo/collections/common/data/lhotse/audio_token_estimator.py

@@ -40,35 +40,22 @@ def from_config(cls, config: Mapping[str, Any]) -> ConvSubsamplingSpec:
         required = {"kernel_size", "stride", "padding"}
         missing = required - set(config)
         if missing:
-            raise ValueError(
-                f"audio_token_estimator.subsampling is missing: {sorted(missing)}"
-            )
+            raise ValueError(f"audio_token_estimator.subsampling is missing: {sorted(missing)}")
         ans = cls(
             kernel_size=int(config["kernel_size"]),
             stride=int(config["stride"]),
             padding=int(config["padding"]),
             repeat=int(config.get("repeat", 1)),
             ceil_mode=bool(config.get("ceil_mode", False)),
         )
-        if (
-            ans.kernel_size <= 0
-            or ans.stride <= 0
-            or ans.padding < 0
-            or ans.repeat <= 0
-        ):
-            raise ValueError(
-                f"Invalid audio_token_estimator.subsampling values: {config}"
-            )
+        if ans.kernel_size <= 0 or ans.stride <= 0 or ans.padding < 0 or ans.repeat <= 0:
+            raise ValueError(f"Invalid audio_token_estimator.subsampling values: {config}")
         return ans
 
     def __call__(self, length: int) -> int:
         for _ in range(self.repeat):
             numerator = length + 2 * self.padding - self.kernel_size
-            quotient = (
-                -(-numerator // self.stride)
-                if self.ceil_mode
-                else numerator // self.stride
-            )
+            quotient = -(-numerator // self.stride) if self.ceil_mode else numerator // self.stride
             length = quotient + 1
         return length
 
@@ -85,9 +72,7 @@ def from_config(cls, config: Mapping[str, Any]) -> FeatureStackingSubsamplingSpe
             raise ValueError("audio_token_estimator.subsampling is missing: ['factor']")
         ans = cls(factor=int(config["factor"]))
         if ans.factor <= 0:
-            raise ValueError(
-                f"Invalid audio_token_estimator.subsampling values: {config}"
-            )
+            raise ValueError(f"Invalid audio_token_estimator.subsampling values: {config}")
         return ans
 
     def __call__(self, length: int) -> int:
@@ -100,17 +85,15 @@ def __call__(self, length: int) -> int:
 def _subsampling_spec_from_config(config: Mapping[str, Any]) -> SubsamplingSpec:
     if not isinstance(config, Mapping):
         raise TypeError(
-            "Each audio_token_estimator.subsampling stage must be a mapping, "
-            f"got {type(config).__name__}"
+            "Each audio_token_estimator.subsampling stage must be a mapping, " f"got {type(config).__name__}"
         )
     stage_type = config.get("type", "conv")
     if stage_type == "conv":
         return ConvSubsamplingSpec.from_config(config)
     if stage_type == "feature_stacking":
         return FeatureStackingSubsamplingSpec.from_config(config)
     raise ValueError(
-        "audio_token_estimator.subsampling.type must be 'conv' or "
-        f"'feature_stacking', got {stage_type!r}"
+        "audio_token_estimator.subsampling.type must be 'conv' or " f"'feature_stacking', got {stage_type!r}"
     )
 
 
@@ -152,30 +135,20 @@ def from_config(
         required = {"n_fft", "hop_length", "stft_pad_amount"}
         missing = required - set(preprocessor)
         if missing:
-            raise ValueError(
-                f"audio_token_estimator.preprocessor is missing: {sorted(missing)}"
-            )
+            raise ValueError(f"audio_token_estimator.preprocessor is missing: {sorted(missing)}")
 
         raw_subsampling = config.get("subsampling")
         if isinstance(raw_subsampling, Mapping):
             raw_subsampling = [raw_subsampling]
-        if not isinstance(raw_subsampling, Sequence) or isinstance(
-            raw_subsampling, (str, bytes)
-        ):
-            raise TypeError(
-                "audio_token_estimator.subsampling must be a mapping or list of mappings"
-            )
-        subsampling = tuple(
-            _subsampling_spec_from_config(stage) for stage in raw_subsampling
-        )
+        if not isinstance(raw_subsampling, Sequence) or isinstance(raw_subsampling, (str, bytes)):
+            raise TypeError("audio_token_estimator.subsampling must be a mapping or list of mappings")
+        subsampling = tuple(_subsampling_spec_from_config(stage) for stage in raw_subsampling)
 
         chunk_size_seconds = config.get("chunk_size_seconds")
         if chunk_size_seconds is not None:
             chunk_size_seconds = float(chunk_size_seconds)
             if chunk_size_seconds <= 0:
-                raise ValueError(
-                    "audio_token_estimator.chunk_size_seconds must be positive or null"
-                )
+                raise ValueError("audio_token_estimator.chunk_size_seconds must be positive or null")
 
         ans = cls(
             sample_rate=int(sample_rate),
@@ -185,15 +158,8 @@ def from_config(
             subsampling=subsampling,
             chunk_size_seconds=chunk_size_seconds,
         )
-        if (
-            ans.sample_rate <= 0
-            or ans.n_fft <= 0
-            or ans.hop_length <= 0
-            or ans.stft_pad_amount < 0
-        ):
-            raise ValueError(
-                f"Invalid audio_token_estimator.preprocessor values: {preprocessor}"
-            )
+        if ans.sample_rate <= 0 or ans.n_fft <= 0 or ans.hop_length <= 0 or ans.stft_pad_amount < 0:
+            raise ValueError(f"Invalid audio_token_estimator.preprocessor values: {preprocessor}")
         return ans
 
     def estimate_cut(self, cut: Cut) -> int:
@@ -214,20 +180,15 @@ def estimate_samples(self, num_samples: int) -> int:
         if chunk_size is None or num_samples <= chunk_size:
             return self._estimate_single_pass(num_samples)
 
-        spans = [
-            (begin, min(begin + chunk_size, num_samples))
-            for begin in range(0, num_samples, chunk_size)
-        ]
+        spans = [(begin, min(begin + chunk_size, num_samples)) for begin in range(0, num_samples, chunk_size)]
         min_chunk_size = self._min_chunk_size_samples()
         if len(spans) > 1 and spans[-1][1] - spans[-1][0] < min_chunk_size:
             spans[-2] = (spans[-2][0], spans[-1][1])
             spans.pop()
         return sum(self._estimate_single_pass(end - begin) for begin, end in spans)
 
     def _estimate_single_pass(self, num_samples: int) -> int:
-        length = (
-            num_samples + 2 * self.stft_pad_amount - self.n_fft
-        ) // self.hop_length
+        length = (num_samples + 2 * self.stft_pad_amount - self.n_fft) // self.hop_length
         for stage in self.subsampling:
             length = stage(length)
         return max(1, length)
@@ -242,9 +203,7 @@ def _min_chunk_size_samples(self) -> int:
         # find the first hop-aligned waveform producing at least two feature frames.
         samples = self.hop_length
         for _ in range(16):
-            feature_frames = (
-                samples + 2 * self.stft_pad_amount - self.n_fft
-            ) // self.hop_length
+            feature_frames = (samples + 2 * self.stft_pad_amount - self.n_fft) // self.hop_length
             if feature_frames >= 2:
                 return samples
             samples += self.hop_length

nemo/collections/common/data/lhotse/dataloader.py

@@ -43,17 +43,13 @@
 from lhotse.utils import fix_random_seed
 from omegaconf import DictConfig, OmegaConf
 
-from nemo.collections.common.data.lhotse.audio_token_estimator import (
-    AudioTokenEstimator,
-)
+from nemo.collections.common.data.lhotse.audio_token_estimator import AudioTokenEstimator
 from nemo.collections.common.data.lhotse.cutset import (
     IncompleteConfigError,
     guess_parse_cutset,
     read_cutset_from_config,
 )
-from nemo.collections.common.data.lhotse.packed_sequence_sampler import (
-    PackedSequenceDynamicCutSampler,
-)
+from nemo.collections.common.data.lhotse.packed_sequence_sampler import PackedSequenceDynamicCutSampler
 from nemo.collections.common.data.lhotse.sampling import (
     BucketingFilter,
     CERFilter,
@@ -980,11 +976,7 @@ def get_lhotse_sampler_from_config(config, global_rank, world_size, tokenizer=No
             f"Creating a Lhotse DynamicCutSampler (bucketing is disabled, "
             f"(max_batch_duration={config.batch_duration} max_batch_size={config.batch_size})"
         )
-        sampler_cls = (
-            PackedSequenceDynamicCutSampler
-            if config.use_packed_sequence_sampling
-            else DynamicCutSampler
-        )
+        sampler_cls = PackedSequenceDynamicCutSampler if config.use_packed_sequence_sampling else DynamicCutSampler
         sampler = sampler_cls(
             cuts,
             constraint=constraint,

nemo/collections/common/data/lhotse/packed_sequence_sampler.py

@@ -15,16 +15,13 @@
 from collections.abc import Sequence
 from typing import Any
 
+from lhotse import CutSet
 from lhotse.dataset import DynamicCutSampler
 from lhotse.dataset.sampling.dynamic import DurationBatcher, Filter
 from lhotse.lazy import get_graph_origin, resolve_iterator_source
 
-from lhotse import CutSet
-
 
-def _select_best_fit_indices(
-    lengths: Sequence[int], capacity: int, max_items: int | None = None
-) -> list[int]:
+def _select_best_fit_indices(lengths: Sequence[int], capacity: int, max_items: int | None = None) -> list[int]:
     """Select an exact best-fit subset, preferring earlier items on ties."""
     if capacity < 0:
         raise ValueError(f"capacity must be non-negative (got {capacity})")
@@ -91,8 +88,7 @@ def __init__(self, *args, packing_buffer_size: int, **kwargs):
         super().__init__(*args, **kwargs)
         if packing_buffer_size <= 0:
             raise ValueError(
-                "shuffle_buffer_size must be a positive packing-buffer size "
-                f"(got {packing_buffer_size})"
+                "shuffle_buffer_size must be a positive packing-buffer size " f"(got {packing_buffer_size})"
             )
         self.packing_buffer_size = packing_buffer_size
         self._source_exhausted = False
@@ -108,32 +104,20 @@ def _detuplify(examples):
 
     @staticmethod
     def _measured_example(example_or_tuple):
-        return (
-            example_or_tuple[0]
-            if isinstance(example_or_tuple, tuple)
-            else example_or_tuple
-        )
+        return example_or_tuple[0] if isinstance(example_or_tuple, tuple) else example_or_tuple
 
     def _fill_packing_buffer(self) -> None:
-        while (
-            len(self.reuse_cuts_buffer) < self.packing_buffer_size
-            and not self._source_exhausted
-        ):
+        while len(self.reuse_cuts_buffer) < self.packing_buffer_size and not self._source_exhausted:
             try:
                 self.reuse_cuts_buffer.append(next(self.cuts_iter))
             except StopIteration:
                 self._source_exhausted = True
 
     def _measure_integer_length(self, example_or_tuple) -> int:
-        length = self.constraint.measure_length(
-            self._measured_example(example_or_tuple)
-        )
+        length = self.constraint.measure_length(self._measured_example(example_or_tuple))
         integer_length = int(length)
         if integer_length != length:
-            raise ValueError(
-                "Packed sequence sampling requires integer token lengths, "
-                f"but measured {length!r}."
-            )
+            raise ValueError("Packed sequence sampling requires integer token lengths, " f"but measured {length!r}.")
         return integer_length
 
     def _limits(self) -> tuple[int, int | None]:
@@ -144,18 +128,14 @@ def _limits(self) -> tuple[int, int | None]:
             max_tokens = getattr(internal, "max_tokens", None)
             max_examples = getattr(internal, "max_examples", max_examples)
         if max_tokens is None:
-            raise ValueError(
-                "Packed sequence sampling requires batch_tokens to define the exact token cap."
-            )
+            raise ValueError("Packed sequence sampling requires batch_tokens to define the exact token cap.")
         max_tokens = int(max_tokens)
         if max_tokens <= 0:
             raise ValueError(f"batch_tokens must be positive (got {max_tokens})")
         if max_examples is not None:
             max_examples = int(max_examples)
             if max_examples <= 0:
-                raise ValueError(
-                    f"batch_size must be positive or null (got {max_examples})"
-                )
+                raise ValueError(f"batch_size must be positive or null (got {max_examples})")
         return max_tokens, max_examples
 
     def _discard(self, examples) -> None:
@@ -182,35 +162,21 @@ def _collect_batch(self):
             )
 
         remaining_items = None if max_examples is None else max_examples - 1
-        tail_indices = _select_best_fit_indices(
-            lengths[1:], max_tokens - anchor_length, max_items=remaining_items
-        )
+        tail_indices = _select_best_fit_indices(lengths[1:], max_tokens - anchor_length, max_items=remaining_items)
         selected_indices = {0, *(index + 1 for index in tail_indices)}
-        examples = [
-            example for index, example in enumerate(pool) if index in selected_indices
-        ]
-        deferred = [
-            example
-            for index, example in enumerate(pool)
-            if index not in selected_indices
-        ]
+        examples = [example for index, example in enumerate(pool) if index in selected_indices]
+        deferred = [example for index, example in enumerate(pool) if index not in selected_indices]
         self.reuse_cuts_buffer.clear()
         self.reuse_cuts_buffer.extend(deferred)
 
         self.constraint.reset()
         for example in examples:
             self.constraint.add(self._measured_example(example))
         if self.constraint.exceeded():
-            raise AssertionError(
-                "Best-fit packed batch exceeded its configured constraint."
-            )
+            raise AssertionError("Best-fit packed batch exceeded its configured constraint.")
 
         is_final_batch = self._source_exhausted and not self.reuse_cuts_buffer
-        if (
-            is_final_batch
-            and self.drop_last
-            and not self.constraint.close_to_exceeding()
-        ):
+        if is_final_batch and self.drop_last and not self.constraint.close_to_exceeding():
             self._discard(examples)
             raise StopIteration()
 
@@ -234,8 +200,7 @@ def __init__(
     ):
         if shuffle_buffer_size is None or shuffle_buffer_size <= 0:
             raise ValueError(
-                "shuffle_buffer_size must be a positive packing-buffer size "
-                f"(got {shuffle_buffer_size})"
+                "shuffle_buffer_size must be a positive packing-buffer size " f"(got {shuffle_buffer_size})"
             )
         # Consume the public `shuffle` argument for config compatibility, but
         # do not allocate DynamicCutSampler's second, reservoir-style buffer.
@@ -252,19 +217,13 @@ def __init__(
         self._inject_restored_packing_buffer = False
 
     def _uses_indexed_restore(self) -> bool:
-        return bool(self.cuts) and all(
-            getattr(source, "has_constant_time_access", False) for source in self.cuts
-        )
+        return bool(self.cuts) and all(getattr(source, "has_constant_time_access", False) for source in self.cuts)
 
     @staticmethod
     def _capture_packing_buffer_tokens(buffer) -> list[tuple[Any, ...]]:
         saved = []
         for example_or_tuple in buffer:
-            examples = (
-                example_or_tuple
-                if isinstance(example_or_tuple, tuple)
-                else (example_or_tuple,)
-            )
+            examples = example_or_tuple if isinstance(example_or_tuple, tuple) else (example_or_tuple,)
             tokens = tuple(get_graph_origin(example) for example in examples)
             if any(token is None for token in tokens):
                 raise RuntimeError(
@@ -278,13 +237,9 @@ def state_dict(self) -> dict[str, Any]:
         state = super().state_dict()
         if self._uses_indexed_restore():
             if self._batcher is not None:
-                state["packing_buffer_tokens"] = self._capture_packing_buffer_tokens(
-                    self._batcher.reuse_cuts_buffer
-                )
+                state["packing_buffer_tokens"] = self._capture_packing_buffer_tokens(self._batcher.reuse_cuts_buffer)
             else:
-                state["packing_buffer_tokens"] = list(
-                    self._restored_packing_buffer_tokens
-                )
+                state["packing_buffer_tokens"] = list(self._restored_packing_buffer_tokens)
         else:
             # Replay restoration deterministically rebuilds the post-filter pool.
             state["packing_buffer_tokens"] = None
@@ -329,25 +284,18 @@ def _restore_packing_buffer(self) -> list[tuple[Any, ...]]:
                     f"{len(tokens)} != {len(active_sources)}."
                 )
             restored.append(
-                tuple(
-                    resolve_iterator_source(source)[token]
-                    for source, token in zip(active_sources, tokens)
-                )
+                tuple(resolve_iterator_source(source)[token] for source, token in zip(active_sources, tokens))
             )
         restored.extend(self._restored_legacy_examples)
         return restored
 
     def _initialize_epoch_iterator(self, *, rebuild_sources: bool) -> None:
         if rebuild_sources or self._active_cuts is None:
             self._active_cuts = self._make_epoch_sources()
-        source_iterators = [
-            iter(resolve_iterator_source(source)) for source in self._active_cuts
-        ]
+        source_iterators = [iter(resolve_iterator_source(source)) for source in self._active_cuts]
         filtered_examples = Filter(
             iterator=zip(*source_iterators),
-            predicate=lambda examples: all(
-                self._filter_fn(example) for example in examples
-            ),
+            predicate=lambda examples: all(self._filter_fn(example) for example in examples),
             diagnostics=self.diagnostics,
         )
         self._batcher = ExactTokenBatcher(