diff --git a/examples/asr/conf/asr_streaming_inference/buffered_rnnt.yaml b/examples/asr/conf/asr_streaming_inference/buffered_rnnt.yaml index 24a344d40cab..8d678453b919 100644 --- a/examples/asr/conf/asr_streaming_inference/buffered_rnnt.yaml +++ b/examples/asr/conf/asr_streaming_inference/buffered_rnnt.yaml @@ -108,6 +108,7 @@ streaming: chunk_size: 4.8 # Audio chunk size in seconds word_boundary_tolerance: 4 # Tolerance for word boundaries request_type: feature_buffer # Type of request: frame or feature_buffer + flush_size_in_secs: 0.0 # Silence appended per stream so the decoder can flush its tail; 0.0 appends none stateful: true # Whether to use stateful processing padding_mode: right # Padding mode: left or right. How to pad frames to match the required buffer length diff --git a/examples/asr/conf/asr_streaming_inference/cache_aware_rnnt.yaml b/examples/asr/conf/asr_streaming_inference/cache_aware_rnnt.yaml index 823d26f52acc..9baef2dfef0d 100644 --- a/examples/asr/conf/asr_streaming_inference/cache_aware_rnnt.yaml +++ b/examples/asr/conf/asr_streaming_inference/cache_aware_rnnt.yaml @@ -129,6 +129,7 @@ streaming: use_feat_cache: true # Whether to cache mel-spec features, set false to re-calculate all mel-spec features in audio buffer chunk_size_in_secs: null # Amount of audio to load for each streaming step, e.g., 0.08s for FastConformer. Set to `null` for using default size equal to 1+lookahead frames. request_type: frame # Type of request: frame or feature_buffer + flush_size_in_secs: 0.0 # Silence appended per stream so the decoder can flush its tail; 0.0 appends none num_slots: 256 # Number of slots in the context manager: must be >= batch_size diff --git a/nemo/collections/asr/inference/pipelines/buffered_rnnt_pipeline.py b/nemo/collections/asr/inference/pipelines/buffered_rnnt_pipeline.py index 940092fcb6c0..1766e740c648 100644 --- a/nemo/collections/asr/inference/pipelines/buffered_rnnt_pipeline.py +++ b/nemo/collections/asr/inference/pipelines/buffered_rnnt_pipeline.py @@ -125,6 +125,7 @@ def init_parameters(self, cfg: DictConfig) -> None: self.buffer_size_in_secs = self.chunk_size + self.left_padding_size + self.right_padding_size self.request_type = RequestType.from_str(cfg.streaming.request_type) + self.flush_size_in_secs = cfg.streaming.get("flush_size_in_secs", 0.0) # 0.0 disables self.padding_mode = FeatureBufferPaddingMode.from_str(cfg.streaming.padding_mode) self.right_padding = self.padding_mode is FeatureBufferPaddingMode.RIGHT self.stop_history_eou_in_milliseconds = cfg.endpointing.stop_history_eou @@ -814,6 +815,7 @@ def get_request_generator(self) -> ContinuousBatchedRequestStreamer: buffer_size_in_secs=self.buffer_size_in_secs, device=self.device, pad_last_frame=True, + flush_size_in_secs=self.flush_size_in_secs, right_pad_features=self.right_padding, ) return request_generator diff --git a/nemo/collections/asr/inference/pipelines/cache_aware_rnnt_pipeline.py b/nemo/collections/asr/inference/pipelines/cache_aware_rnnt_pipeline.py index 28bafc986d2a..b4d809ed4755 100644 --- a/nemo/collections/asr/inference/pipelines/cache_aware_rnnt_pipeline.py +++ b/nemo/collections/asr/inference/pipelines/cache_aware_rnnt_pipeline.py @@ -188,6 +188,7 @@ def init_parameters(self, cfg: DictConfig) -> None: self.return_tail_result = cfg.return_tail_result self.request_type = RequestType.from_str(cfg.streaming.request_type) + self.flush_size_in_secs = cfg.streaming.get("flush_size_in_secs", 0.0) # 0.0 disables def init_greedy_rnnt_decoder(self) -> None: """Initialize the RNNT decoder.""" @@ -617,6 +618,7 @@ def get_request_generator(self) -> ContinuousBatchedRequestStreamer: buffer_size_in_secs=self.buffer_size_in_secs, device=self.device, pad_last_frame=True, + flush_size_in_secs=self.flush_size_in_secs, ) return request_generator diff --git a/nemo/collections/asr/inference/streaming/framing/mono_stream.py b/nemo/collections/asr/inference/streaming/framing/mono_stream.py index dabc86eef71e..21d9481a6924 100644 --- a/nemo/collections/asr/inference/streaming/framing/mono_stream.py +++ b/nemo/collections/asr/inference/streaming/framing/mono_stream.py @@ -26,18 +26,28 @@ class MonoStream(Stream): Iterates over the frames of the audio file """ - def __init__(self, rate: int, frame_size_in_secs: float, stream_id: int, pad_last_frame: bool = False): + def __init__( + self, + rate: int, + frame_size_in_secs: float, + stream_id: int, + pad_last_frame: bool = False, + flush_size_in_secs: float = 0.0, + ): """ Initialize the MonoStream Args: rate (int): sampling rate - frame_size_in_secs (int): frame length in seconds + frame_size_in_secs (float): frame length in seconds stream_id (int): stream id + pad_last_frame (bool): whether to pad the last frame up to frame_size + flush_size_in_secs (float): seconds of silence appended to the audio, 0.0 to append none """ self.rate = rate self.frame_size = int(frame_size_in_secs * rate) self.pad_last_frame = pad_last_frame + self.flush_size = int(flush_size_in_secs * rate) self.samples = None self.n_samples = None @@ -56,6 +66,8 @@ def load_audio(self, audio: str | torch.Tensor, options: RequestOptions | None = self.samples = read_audio(audio, target_sr=self.rate, mono=True) else: self.samples = audio + if self.flush_size > 0: # appended as signal, not padding, so it is not trimmed downstream + self.samples = torch.cat([self.samples, torch.zeros(self.flush_size, dtype=self.samples.dtype)]) self.n_samples = len(self.samples) self.frame_count = 0 # Reset frame count self.options = options diff --git a/nemo/collections/asr/inference/streaming/framing/multi_stream.py b/nemo/collections/asr/inference/streaming/framing/multi_stream.py index 65ba10746ea5..f2a02109a4d3 100644 --- a/nemo/collections/asr/inference/streaming/framing/multi_stream.py +++ b/nemo/collections/asr/inference/streaming/framing/multi_stream.py @@ -103,6 +103,7 @@ def __init__( batch_size: int, n_frames_per_stream: int, pad_last_frame: bool = False, + flush_size_in_secs: float = 0.0, ): """ Args: @@ -111,12 +112,14 @@ def __init__( batch_size (int): The batch size n_frames_per_stream (int): The number of frames per stream pad_last_frame (bool): Whether to pad the last frame + flush_size_in_secs (float): Seconds of silence appended to every stream, 0.0 for none """ self.sample_rate = sample_rate self.frame_size_in_secs = frame_size_in_secs self.batch_size = batch_size self.pad_last_frame = pad_last_frame + self.flush_size_in_secs = flush_size_in_secs self.multi_streamer = MultiStream(n_frames_per_stream=n_frames_per_stream) self.stream_id = 0 @@ -176,7 +179,11 @@ def add_stream(self) -> None: # Create a new stream stream = MonoStream( - self.sample_rate, self.frame_size_in_secs, stream_id=self.stream_id, pad_last_frame=self.pad_last_frame + self.sample_rate, + self.frame_size_in_secs, + stream_id=self.stream_id, + pad_last_frame=self.pad_last_frame, + flush_size_in_secs=self.flush_size_in_secs, ) # Load the next audio file audio_filepath = self.audio_filepaths[self.stream_id] @@ -244,6 +251,7 @@ def __init__( device: torch.device = None, pad_last_frame: bool = False, right_pad_features: bool = False, + flush_size_in_secs: float = 0.0, ): """ Args: @@ -274,6 +282,7 @@ def __init__( batch_size=batch_size, n_frames_per_stream=n_frames_per_stream, pad_last_frame=pad_last_frame, + flush_size_in_secs=flush_size_in_secs, ) if self.request_type is RequestType.FEATURE_BUFFER: