Apply isort and black reformatting

Signed-off-by: andrusenkoau <andrusenkoau@users.noreply.github.com>

Author: andrusenkoau

examples/asr/asr_chunked_inference/rnnt/speech_to_text_streaming_infer_rnnt.py

@@ -133,7 +133,9 @@ class TranscriptionConfig:
     )
     right_context_secs: float = 2  # right context
 
-    att_context_size_as_chunk: bool = True  # whether to use the att_context_size as chunk size (importand for extra-low latency)
+    att_context_size_as_chunk: bool = (
+        True  # whether to use the att_context_size as chunk size (importand for extra-low latency)
+    )
 
     # Set `cuda` to int to define CUDA device. If 'None', will look for CUDA
     # device anyway, and do inference on CPU only if CUDA device is not found.
@@ -303,7 +305,7 @@ def main(cfg: TranscriptionConfig) -> TranscriptionConfig:
     # unified ASR model: use the att_context_size as chunk size (important for extra-low latency)
     if asr_model.cfg.encoder.att_context_style == 'chunked_limited_with_rc' and cfg.att_context_size_as_chunk:
         asr_model.encoder.set_default_att_context_size(
-            att_context_size=[context_encoder_frames.left,context_encoder_frames.chunk,context_encoder_frames.right]
+            att_context_size=[context_encoder_frames.left, context_encoder_frames.chunk, context_encoder_frames.right]
         )
 
     logging.info(

nemo/collections/asr/losses/rnnt.py

@@ -154,7 +154,7 @@ class RNNTLossConfig:
         installation_msg="Pure Pytorch implementation of TDT loss. Slow and for debugging purposes only.",
     ),
     "rnnt_triton": RNNTLossConfig(
-        loss_name="rnnt_triton", # will be added later
+        loss_name="rnnt_triton",  # will be added later
         lib_name="torch",
         min_version='0.0',
         is_available=True,
@@ -331,7 +331,7 @@ def resolve_rnnt_loss(loss_name: str, blank_idx: int, loss_kwargs: dict = None)
         loss_kwargs = _clean_kwargs(loss_name, loss_kwargs, GraphWTransducerLoss.__init__, ignore_params={"blank"})
         loss_func = GraphWTransducerLoss(blank=blank_idx, **loss_kwargs)
     elif loss_name == "rnnt_triton":
-        loss_func = None # will be added later
+        loss_func = None  # will be added later
     else:
         raise ValueError(
             f"Invalid value of `loss_name`: {loss_name}. Allowed loss names are :" f"{loss_function_names}"
@@ -343,8 +343,7 @@ def resolve_rnnt_loss(loss_name: str, blank_idx: int, loss_kwargs: dict = None)
 class RNNTLoss(Loss):
     @property
     def input_types(self):
-        """Input types definitions for CTCLoss.
-        """
+        """Input types definitions for CTCLoss."""
         return {
             "log_probs": NeuralType(('B', 'T', 'T', 'D'), LogprobsType()),
             "targets": NeuralType(('B', 'T'), LabelsType()),
@@ -405,7 +404,7 @@ def __init__(self, num_classes, reduction: str = 'mean_batch', loss_name: str =
                                  standard blank, and the standard blank is the last symbol in the vocab)
                 TDT: num_classes = V. Note, V here does not include any of the "duration outputs".
 
-            reduction: Type of reduction to perform on loss. Possible values are 
+            reduction: Type of reduction to perform on loss. Possible values are
                 `mean_batch`, 'mean_volume`, `mean`, `sum` or None.
                 `None` will return a torch vector comprising the individual loss values of the batch.
                 `mean_batch` will average the losses in the batch

nemo/collections/asr/modules/conformer_encoder.py

@@ -361,7 +361,6 @@ def __init__(
         self.use_pytorch_sdpa_backends = use_pytorch_sdpa_backends
         self.sync_max_audio_length = sync_max_audio_length
 
-
         assert conv_context_style in ["regular", "dcc", "dcc_rc"], f"Invalid conv_context_style: {conv_context_style}!"
         self.conv_context_style = conv_context_style
         self.conv_kernel_size = conv_kernel_size
@@ -370,8 +369,12 @@ def __init__(
         self.dual_mode_training = dual_mode_training
         self.unified_asr_prob = unified_asr_prob
         if att_chunk_context_size is not None:
-            assert att_context_style == "chunked_limited_with_rc", "att_chunk_context_size is only supported for chunked_limited_with_rc attention style!"
-            assert len(att_chunk_context_size) == 3, "att_chunk_context_size must have 3 elements: [left_context, chunk_size, right_context]"
+            assert (
+                att_context_style == "chunked_limited_with_rc"
+            ), "att_chunk_context_size is only supported for chunked_limited_with_rc attention style!"
+            assert (
+                len(att_chunk_context_size) == 3
+            ), "att_chunk_context_size must have 3 elements: [left_context, chunk_size, right_context]"
             self.att_chunk_context_size = att_chunk_context_size
         else:
             self.att_chunk_context_size = None
@@ -855,8 +858,10 @@ def _create_masks(self, att_context_size, padding_length, max_audio_length, offs
                     )
                     att_mask = torch.logical_and(att_mask, chunked_limited_mask.unsqueeze(0))
             elif self.att_context_style == "chunked_limited_with_rc" and sum(att_context_size) != -3:
-                assert len(att_context_size) == 3, "att_context_size must have 3 elements: [left_context, chunk_size, right_context]"
-                
+                assert (
+                    len(att_context_size) == 3
+                ), "att_context_size must have 3 elements: [left_context, chunk_size, right_context]"
+
                 left_context_frames = att_context_size[0]
                 chunk_size_frames = att_context_size[1]
                 right_context_frames = att_context_size[2]
@@ -867,22 +872,23 @@ def _create_masks(self, att_context_size, padding_length, max_audio_length, offs
                 window_start = chunk_idx * chunk_size_frames - left_context_frames
                 window_start = torch.maximum(window_start, torch.zeros_like(window_start))
                 window_end = chunk_idx * chunk_size_frames + chunk_size_frames - 1 + right_context_frames
-                
+
                 if self.training and self.skip_att_chunk_rc_prob > 0.0:
                     chunks_num = max_audio_length // chunk_size_frames
                     for chunk_step in range(chunks_num):
                         if random.random() < self.skip_att_chunk_rc_prob:
-                            window_end[chunk_step*chunk_size_frames:chunk_step*chunk_size_frames+chunk_size_frames] -= right_context_frames
-                
+                            window_end[
+                                chunk_step * chunk_size_frames : chunk_step * chunk_size_frames + chunk_size_frames
+                            ] -= right_context_frames
+
                 window_end = torch.minimum(window_end, torch.full_like(window_end, max_audio_length - 1))
                 # Create the mask: frame i can see frame j if window_start[i] <= j <= window_end[i]
                 j_indices = frame_idx.unsqueeze(0)  # [1, T]
                 window_start_expanded = window_start.unsqueeze(1)  # [T, 1]
-                window_end_expanded = window_end.unsqueeze(1)      # [T, 1]
+                window_end_expanded = window_end.unsqueeze(1)  # [T, 1]
 
                 chunked_limited_mask = torch.logical_and(
-                    j_indices >= window_start_expanded,
-                    j_indices <= window_end_expanded
+                    j_indices >= window_start_expanded, j_indices <= window_end_expanded
                 )
                 att_mask = torch.logical_and(att_mask, chunked_limited_mask.unsqueeze(0))
         else: