Clean up

Signed-off-by: Vladimir Bataev <vbataev@nvidia.com>

Author: artbataev

examples/asr/asr_chunked_inference/rnnt/speech_to_text_streaming_infer_rnnt.py

@@ -84,7 +84,7 @@
 from nemo.collections.asr.parts.utils.streaming_utils import (
     AudioBatch,
     ContextSize,
-    DynamicTensor,
+    DynamicLengthTensor,
     SimpleAudioDataset,
     StreamingBatchedAudioBuffer,
 )
@@ -421,7 +421,7 @@ def main(cfg: TranscriptionConfig) -> TranscriptionConfig:
                 device=device,
             )
             rest_audio_lengths = audio_batch_lengths.clone()
-            encoder_output_aggregated: DynamicTensor | None = None
+            encoder_output_aggregated: DynamicLengthTensor | None = None
 
             # iterate over audio samples
             while left_sample < audio_batch.shape[1]:
@@ -461,7 +461,7 @@ def main(cfg: TranscriptionConfig) -> TranscriptionConfig:
                 if use_simulated_decoding:
                     # store encoder output
                     if encoder_output_aggregated is None:
-                        encoder_output_aggregated = DynamicTensor(
+                        encoder_output_aggregated = DynamicLengthTensor(
                             batch_size=batch_size,
                             init_length=encoder_output.shape[1],
                             dim_shape=encoder_output.shape[2],

nemo/collections/asr/parts/utils/streaming_utils.py

@@ -2401,7 +2401,7 @@ def __len__(self):
         return len(self.audio_filenames)
 
 
-class DynamicTensor:
+class DynamicLengthTensor:
     def __init__(
         self,
         batch_size: int,
@@ -2441,19 +2441,23 @@ def _allocate_more(self, min_add_length: int | None = None):
         self.data = torch.cat((self.data, self.data.new_zeros(add_shape)), dim=1)
         self._max_length += add_len
 
-    def to_device(self, device: str | torch.device):
+    def to_device(self, device: str | torch.device) -> "DynamicLengthTensor":
+        """Move storage to device"""
         self.device = device
         self.data.to(device=device)
         self.lengths.to(device=device)
+        return self
 
     def append_(self, data: torch.Tensor, lengths: torch.Tensor | None = None):
+        """Append new data along length dimension"""
         cur_len = self.lengths.max().item()
         other_len = data.shape[1] if lengths is None else lengths.max().item()
         if cur_len + other_len >= self._max_length:
             self._allocate_more(min_add_length=cur_len + other_len - self._max_length + 1)
         self.append_no_checks_(data=data[:, :other_len], lengths=lengths)
 
     def append_no_checks_(self, data: torch.Tensor, lengths: torch.Tensor | None = None):
+        """Append new data along length dimension without checks"""
         other_len = data.shape[1]
         indices = torch.arange(other_len, device=self.device)
         shifted_indices = self.lengths[:, None] + indices[None, :]
@@ -2463,9 +2467,9 @@ def append_no_checks_(self, data: torch.Tensor, lengths: torch.Tensor | None = N
         else:
             self.lengths += lengths
 
-    def clone(self) -> "DynamicTensor":
+    def clone(self) -> "DynamicLengthTensor":
         """Return a copy of self"""
-        new_dynamic_tensor = DynamicTensor(
+        new_dynamic_tensor = DynamicLengthTensor(
             batch_size=self.batch_size,
             init_length=self._max_length,
             device=self.device,
@@ -2475,13 +2479,13 @@ def clone(self) -> "DynamicTensor":
         new_dynamic_tensor.data.copy_(self.lengths)
         return new_dynamic_tensor
 
-    def merge_(self, other: "DynamicTensor") -> "DynamicTensor":
+    def merge_(self, other: "DynamicLengthTensor") -> "DynamicLengthTensor":
         """
         Merge two dynamic tensors
         NB: this will reallocate memory
 
         Args:
-            other: DynamicTensor
+            other: DynamicLengthTensor
         """
         self.append_(data=other.data, lengths=other.lengths)
         return self