chore: import upstream snapshot with attribution

nemo/README.md

@@ -0,0 +1,26 @@
+NeMo (**Ne**ural **Mo**dules) is a toolkit for creating AI applications built around **neural modules**, conceptual blocks of neural networks that take *typed* inputs and produce *typed* outputs.
+
+## **collections/**
+* **ASR** - Collection of modules and models for building speech recognition networks.
+* **TTS** - Collection of modules and models for building speech synthesis networks.
+* **Audio** - Collection of modules and models for building audio processing networks.
+* **SpeechLM2** - Collection of modules and models for building multimodal LLM.
+
+## **core/**
+Provides fundamental APIs and utilities for NeMo modules, including:
+- **Classes** - Base classes for datasets, models, and losses.
+- **Config** - Configuration management utilities.
+- **Neural Types** - Typed inputs/outputs for module interaction.
+- **Optim** - Optimizers and learning rate schedulers.
+
+## **lightning/**
+Integration with PyTorch Lightning for training and distributed execution:
+- **Strategies & Plugins** - Custom Lightning strategies.
+- **Fabric** - Lightweight wrapper for model training.
+- **Checkpointing & Logging** - Utilities for managing model states.
+
+## **utils/**
+General utilities for debugging, distributed training, logging, and model management:
+- **callbacks/** - Hooks for training processes.
+- **loggers/** - Logging utilities for different backends.
+- **debugging & profiling** - Performance monitoring tools.

nemo/__init__.py

@@ -0,0 +1,28 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from nemo.package_info import (
+    __contact_emails__,
+    __contact_names__,
+    __description__,
+    __download_url__,
+    __homepage__,
+    __keywords__,
+    __license__,
+    __package_name__,
+    __repository_url__,
+    __shortversion__,
+    __version__,
+)

nemo/agents/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

nemo/agents/voice_agent/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

nemo/agents/voice_agent/pipecat/__init__.py

@@ -0,0 +1,18 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+try:
+    import pipecat
+except ImportError:
+    raise ImportError("pipecat is not installed. Please install it with `pip install pipecat-ai`.")

nemo/agents/voice_agent/pipecat/frames/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

nemo/agents/voice_agent/pipecat/frames/frames.py

@@ -0,0 +1,26 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from dataclasses import dataclass
+import numpy as np
+from pipecat.frames.frames import DataFrame
+
+
+@dataclass
+class DiarResultFrame(DataFrame):
+    """Diarization frame."""
+
+    diar_result: np.ndarray | int
+    stream_id: str = "default"

nemo/agents/voice_agent/pipecat/processors/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

nemo/agents/voice_agent/pipecat/processors/frameworks/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

nemo/agents/voice_agent/pipecat/processors/frameworks/rtvi.py

@@ -0,0 +1,72 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from loguru import logger
+from pipecat.frames.frames import Frame, LLMFullResponseEndFrame, LLMFullResponseStartFrame, TTSTextFrame
+from pipecat.observers.base_observer import FramePushed
+from pipecat.processors.frameworks.rtvi import (
+    RTVIBotLLMStartedMessage,
+    RTVIBotLLMStoppedMessage,
+    RTVIBotTranscriptionMessage,
+    RTVIBotTTSTextMessage,
+)
+from pipecat.processors.frameworks.rtvi import RTVIObserver as _RTVIObserver
+from pipecat.processors.frameworks.rtvi import RTVIProcessor, RTVITextMessageData
+from pipecat.transports.base_output import BaseOutputTransport
+
+
+class RTVIObserver(_RTVIObserver):
+    """
+    An observer that processes RTVI frames and pushes them to the transport.
+    """
+
+    def __init__(self, rtvi: RTVIProcessor, *args, **kwargs):
+        super().__init__(rtvi, *args, **kwargs)
+
+    async def on_push_frame(self, data: FramePushed):
+        """Process a frame being pushed through the pipeline.
+
+        Args:
+            data: Frame push event data containing source, frame, direction, and timestamp.
+        """
+        src = data.source
+        frame: Frame = data.frame
+
+        if frame.id in self._frames_seen:
+            return
+
+        if not self._params.bot_llm_enabled:
+            if isinstance(frame, LLMFullResponseStartFrame):
+                await self.send_rtvi_message(RTVIBotLLMStartedMessage())
+                self._frames_seen.add(frame.id)
+            elif isinstance(frame, LLMFullResponseEndFrame):
+                await self.send_rtvi_message(RTVIBotLLMStoppedMessage())
+                self._frames_seen.add(frame.id)
+            elif isinstance(frame, TTSTextFrame) and isinstance(src, BaseOutputTransport):
+                message = RTVIBotTTSTextMessage(data=RTVITextMessageData(text=frame.text))
+                await self.send_rtvi_message(message)
+                await self._push_bot_transcription(frame.text)
+                self._frames_seen.add(frame.id)
+            else:
+                await super().on_push_frame(data)
+        else:
+            await super().on_push_frame(data)
+
+    async def _push_bot_transcription(self, text: str):
+        """Push accumulated bot transcription as a message."""
+        if len(text.strip()) > 0:
+            message = RTVIBotTranscriptionMessage(data=RTVITextMessageData(text=text))
+            logger.debug(f"Pushing bot transcription: `{text}`")
+            await self.send_rtvi_message(message)

nemo/agents/voice_agent/pipecat/services/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

nemo/agents/voice_agent/pipecat/services/nemo/__init__.py

@@ -0,0 +1,19 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .diar import NemoDiarService
+from .llm import HuggingFaceLLMService
+from .stt import NemoSTTService
+from .tts import NeMoFastPitchHiFiGANTTSService
+from .turn_taking import NeMoTurnTakingService

nemo/agents/voice_agent/pipecat/services/nemo/audio_logger.py

@@ -0,0 +1,844 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+import threading
+import wave
+from datetime import datetime
+from pathlib import Path
+from typing import Optional, Union
+
+import librosa
+import numpy as np
+from loguru import logger
+from pipecat.frames.frames import TranscriptionFrame
+from pipecat.observers.base_observer import BaseObserver, FramePushed
+
+
+class AudioLogger:
+    """
+    Utility class for logging audio data and transcriptions during voice agent interactions.
+
+    This logger saves:
+    - Audio files in WAV format
+    - Transcriptions with metadata in JSON format
+    - Session information and metadata
+
+    File structure:
+        log_dir/
+        ├── session_YYYYMMDD_HHMMSS/
+        │   ├── user/
+        │   │   ├── 00001_HHMMSS.wav
+        │   │   ├── 00001_HHMMSS.json
+        │   │   ├── 00002_HHMMSS.wav
+        │   │   └── 00002_HHMMSS.json
+        │   ├── agent/
+        │   │   ├── 00001_HHMMSS.wav
+        │   │   ├── 00001_HHMMSS.json
+        │   └── session_metadata.json
+
+    Args:
+        log_dir: Base directory for storing logs (default: "./audio_logs")
+        session_id: Optional custom session ID. If None, auto-generated from timestamp
+        enabled: Whether logging is enabled (default: True)
+
+    # 12/19/2025 Note: Stereo conversation recording is implemented,
+    # but -0.8 seconds offset needs to be applied to make the session sound synced.
+    """
+
+    def __init__(
+        self,
+        log_dir: Union[str, Path] = "./audio_logs",
+        session_id: Optional[str] = None,
+        enabled: bool = True,
+        user_audio_sample_rate: int = 16000,
+        pre_roll_time_sec: float = 0.8,
+        round_precision: int = 2,
+    ):
+        self.enabled = enabled
+        if not self.enabled:
+            logger.info("[AudioLogger] AudioLogger is disabled")
+            return
+
+        self.log_dir = Path(log_dir)
+
+        # Generate session ID if not provided
+        self.session_start_time = datetime.now()
+        if session_id is None:
+            session_id = f"session_{self.session_start_time.strftime('%Y%m%d_%H%M%S')}"
+        self.first_audio_timestamp = None
+        self.session_id = session_id
+        self.session_dir = self.log_dir / session_id
+
+        # Create directories
+        self.user_dir = self.session_dir / "user"
+        self.agent_dir = self.session_dir / "agent"
+
+        self.user_dir.mkdir(parents=True, exist_ok=True)
+        self.agent_dir.mkdir(parents=True, exist_ok=True)
+
+        # Counters for file naming (thread-safe)
+        self._user_counter = 0
+        self._agent_counter = 0
+        self._turn_index = 0  # Turn index for conversation turns
+        self._current_speaker = None  # Track current speaker for turn transitions
+        self._agent_turn_start_time = None  # Captured when BotStartedSpeakingFrame is received
+        self._lock = threading.Lock()
+        self.staged_metadata = None
+        self._staged_audio_data = None
+        self._pre_roll_time_sec = pre_roll_time_sec
+        self._round_precision = round_precision
+
+        self.turn_audio_buffer = []
+        self.continuous_user_audio_buffer = []
+        self.turn_transcription_buffer = []
+
+        # Stereo conversation recording (left=agent, right=user)
+        self._stereo_conversation_filename = "conversation_stereo.wav"
+        self._stereo_conversation_file = self.session_dir / self._stereo_conversation_filename
+        self._stereo_sample_rate = user_audio_sample_rate  # Use user audio sample rate (downsample agent audio)
+        self._stereo_audio_buffer_left: list = []  # Agent audio (left channel)
+        self._stereo_audio_buffer_right: list = []  # User audio (right channel)
+
+        # Session metadata
+        # agent_entries is a list of lists: each sublist contains segments for one turn
+        # e.g., [[seg1, seg2, seg3], [seg4, seg5], ...]  where each [] is a turn
+        self.session_metadata = {
+            "session_id": session_id,
+            "start_time": self.session_start_time.isoformat(),
+            "user_entries": [],
+            "agent_entries": [],  # List of turns, each turn is a list of segments
+        }
+
+        logger.info(f"[AudioLogger] AudioLogger initialized: {self.session_dir}")
+
+    def append_continuous_user_audio(self, audio_data: bytes):
+        """
+        Append audio data to the continuous user audio buffer for stereo conversation.
+
+        This method should be called for EVERY audio frame received from the user,
+        regardless of VAD state, to record the complete conversation audio.
+
+        Args:
+            audio_data: Raw audio data as bytes
+        """
+        if not self.enabled:
+            return
+
+        self.continuous_user_audio_buffer.append(audio_data)
+
+    def _resample_audio(
+        self,
+        audio_data: Union[bytes, np.ndarray],
+        orig_sr: int,
+        target_sr: int,
+    ) -> np.ndarray:
+        """
+        Resample audio data to a target sample rate using librosa.
+
+        Args:
+            audio_data: Audio data as bytes (int16) or numpy array
+            orig_sr: Original sample rate
+            target_sr: Target sample rate
+
+        Returns:
+            Resampled audio as numpy array (float32)
+        """
+        # Convert bytes to numpy array if needed
+        if isinstance(audio_data, bytes):
+            audio_array = np.frombuffer(audio_data, dtype=np.int16).astype(np.float32) / 32768.0
+        elif audio_data.dtype == np.int16:
+            audio_array = audio_data.astype(np.float32) / 32768.0
+        else:
+            audio_array = audio_data.astype(np.float32)
+
+        # Resample if needed
+        if orig_sr != target_sr:
+            audio_array = librosa.resample(audio_array, orig_sr=orig_sr, target_sr=target_sr)
+
+        return audio_array
+
+    def _append_to_stereo_conversation(
+        self,
+        audio_data: Union[bytes, np.ndarray],
+        channel: str,
+        start_time: float,
+        sample_rate: int,
+    ):
+        """
+        Append audio to the stereo conversation buffer at the correct time position.
+
+        Args:
+            audio_data: Audio data as bytes or numpy array
+            channel: "left" for agent, "right" for user
+            start_time: Start time in seconds from session start
+            sample_rate: Sample rate of the input audio
+        """
+        if not self.enabled:
+            return
+
+        try:
+            # Resample to stereo sample rate if needed
+            audio_float = self._resample_audio(audio_data, sample_rate, self._stereo_sample_rate)
+
+            # Calculate the sample position for this audio
+            start_sample = int(start_time * self._stereo_sample_rate)
+
+            # Get the appropriate buffer
+            if channel == "left":
+                buffer = self._stereo_audio_buffer_left
+            else:
+                buffer = self._stereo_audio_buffer_right
+
+            # Extend buffer with zeros if needed to reach start position
+            current_length = len(buffer)
+            if start_sample > current_length:
+                buffer.extend([0.0] * (start_sample - current_length))
+
+            # Append or overwrite audio samples
+            for i, sample in enumerate(audio_float):
+                pos = start_sample + i
+                if pos < len(buffer):
+                    # Mix with existing audio (in case of overlap)
+                    buffer[pos] = np.clip(buffer[pos] + sample, -1.0, 1.0)
+                else:
+                    buffer.append(sample)
+
+            logger.debug(
+                f"[AudioLogger] Appended {len(audio_float)} samples to {channel} channel "
+                f"at position {start_sample} (buffer now {len(buffer)} samples)"
+            )
+
+        except Exception as e:
+            logger.error(f"[AudioLogger] Error appending to stereo conversation: {e}")
+
+    def save_stereo_conversation(self):
+        """
+        Save the stereo conversation buffer to a WAV file.
+        Left channel = Agent, Right channel = User.
+
+        User audio comes from continuous_user_audio_buffer (not affected by VAD).
+        """
+        if not self.enabled:
+            return
+
+        if not self._stereo_audio_buffer_left and not self.continuous_user_audio_buffer:
+            logger.warning("[AudioLogger] No stereo conversation audio to save")
+            return
+
+        try:
+            # Build right channel (user) from continuous buffer
+            # This is raw bytes at user sample rate, no resampling needed since stereo uses user sample rate
+            if self.continuous_user_audio_buffer:
+                continuous_audio_bytes = b"".join(self.continuous_user_audio_buffer)
+                right_array = np.frombuffer(continuous_audio_bytes, dtype=np.int16).astype(np.float32) / 32768.0
+            else:
+                right_array = np.array([], dtype=np.float32)
+
+            left_array = np.array(self._stereo_audio_buffer_left, dtype=np.float32)
+
+            # Pad the shorter buffer with zeros
+            max_length = max(len(left_array), len(right_array))
+
+            # Pad to same length
+            if len(left_array) < max_length:
+                left_array = np.pad(left_array, (0, max_length - len(left_array)))
+            if len(right_array) < max_length:
+                right_array = np.pad(right_array, (0, max_length - len(right_array)))
+
+            # Create stereo array (interleaved: L, R, L, R, ...)
+            stereo_array = np.column_stack((left_array, right_array))
+
+            # Convert to int16
+            stereo_int16 = (stereo_array * 32767).astype(np.int16)
+
+            # Save as WAV
+            with wave.open(str(self._stereo_conversation_file), 'wb') as wav_file:  # type: ignore[union-attr]
+                wav_file.setnchannels(2)  # Stereo
+                wav_file.setsampwidth(2)  # 16-bit
+                wav_file.setframerate(self._stereo_sample_rate)
+                wav_file.writeframes(stereo_int16.tobytes())
+
+            duration_sec = max_length / self._stereo_sample_rate
+            logger.info(
+                f"[AudioLogger] Saved stereo conversation: {self._stereo_conversation_file} "
+                f"({duration_sec:.2f} seconds, {max_length} samples)"
+            )
+
+        except Exception as e:
+            logger.error(f"[AudioLogger] Error saving stereo conversation: {e}")
+
+    def get_time_from_start_of_session(self, timestamp: datetime = None) -> float:
+        """Get the time from the start of the session to the given datetime string."""
+        # get the time difference in seconds.
+        if self.first_audio_timestamp is None:
+            raise ValueError("First audio timestamp is not set. Aborting time calculation.")
+        time_diff = (timestamp if timestamp else datetime.now()) - self.first_audio_timestamp
+        return time_diff.total_seconds()
+
+    def _get_next_counter(self, speaker: str) -> int:
+        """Get the next counter value for a speaker in a thread-safe manner."""
+        with self._lock:
+            if speaker == "user":
+                self._user_counter += 1
+                return self._user_counter
+            else:
+                self._agent_counter += 1
+                return self._agent_counter
+
+    def increment_turn_index(self, speaker: str = None) -> int:
+        """
+        Increment the turn index if the speaker has changed.
+
+        Args:
+            speaker: "user" or "agent". If provided, only increments
+                    if this is different from the current speaker.
+                    If None, always increments.
+
+        Returns:
+            The current turn index after any increment.
+        """
+        with self._lock:
+            if speaker is None:
+                # Always increment if no speaker specified
+                self._turn_index += 1
+                logger.debug(f"[AudioLogger] Turn index incremented to {self._turn_index}")
+            elif speaker != self._current_speaker:
+                # Only increment if speaker changed
+                self._current_speaker = speaker
+                self._turn_index += 1
+                # Reset agent turn start time when speaker changes
+                if speaker == "agent":
+                    self._agent_turn_start_time = None
+                logger.debug(
+                    f"[AudioLogger] Speaker changed to {speaker}, turn index incremented to {self._turn_index}"
+                )
+            # else: same speaker, no increment
+            return self._turn_index
+
+    def set_agent_turn_start_time(self):
+        """
+        Set the start time for the current agent turn.
+
+        This should be called when BotStartedSpeakingFrame is received,
+        which indicates the audio is actually starting to play (not just generated).
+        This provides more accurate timing than capturing time during TTS generation.
+        """
+        if not self.enabled:
+            return
+
+        # Only set if not already set for this turn
+        if self._agent_turn_start_time is None:
+            self._agent_turn_start_time = self.get_time_from_start_of_session()
+            logger.debug(f"[AudioLogger] Agent turn start time set to {self._agent_turn_start_time:.3f}s")
+
+    def _save_audio_wav(
+        self,
+        audio_data: Union[bytes, np.ndarray],
+        file_path: Path,
+        sample_rate: int,
+        num_channels: int = 1,
+    ):
+        """
+        Save audio data to a WAV file.
+
+        Args:
+            audio_data: Audio data as bytes or numpy array
+            file_path: Path to save the WAV file
+            sample_rate: Audio sample rate in Hz
+            num_channels: Number of audio channels (default: 1)
+        """
+        try:
+            # Convert audio data to bytes if it's a numpy array
+            if isinstance(audio_data, np.ndarray):
+                if audio_data.dtype in [np.float32, np.float64]:
+                    # Convert float [-1, 1] to int16 [-32768, 32767]
+                    audio_data = np.clip(audio_data, -1.0, 1.0)
+                    audio_data = (audio_data * 32767).astype(np.int16)
+                elif audio_data.dtype != np.int16:
+                    audio_data = audio_data.astype(np.int16)
+                audio_bytes = audio_data.tobytes()
+            else:
+                audio_bytes = audio_data
+
+            # Write WAV file
+            with wave.open(str(file_path), 'wb') as wav_file:  # type: ignore[union-attr]
+                wav_file.setnchannels(num_channels)
+                wav_file.setsampwidth(2)  # 16-bit audio
+                wav_file.setframerate(sample_rate)
+                wav_file.writeframes(audio_bytes)
+
+            logger.debug(f"[AudioLogger] Saved audio to {file_path}")
+        except Exception as e:
+            logger.error(f"[AudioLogger] Error saving audio to {file_path}: {e}")
+            raise
+
+    def _save_metadata_json(self, metadata: dict, file_path: Path):
+        """Save metadata to a JSON file."""
+        try:
+            with open(file_path, 'w', encoding='utf-8') as f:
+                json.dump(metadata, f, indent=2, ensure_ascii=False)
+            logger.debug(f"[AudioLogger] Saved metadata to {file_path}")
+        except Exception as e:
+            logger.error(f"[AudioLogger] Error saving metadata to {file_path}: {e}")
+            raise
+
+    def clear_user_audio_buffer(self):
+        """
+        Clear the user audio buffer if the user stopped speaking detected by VAD.
+        """
+        # Clear turn buffers if logging wasn't completed (e.g., no final transcription)
+        if len(self.turn_audio_buffer) > 0 or len(self.turn_transcription_buffer) > 0:
+            logger.debug(
+                "[AudioLogger] Clearing turn audio and transcription buffers due to VAD user stopped speaking"
+            )
+            self.turn_audio_buffer = []
+            self.turn_transcription_buffer = []
+
+    def stage_user_audio(
+        self,
+        timestamp_now: datetime,
+        transcription: str,
+        sample_rate: int = 16000,
+        num_channels: int = 1,
+        is_first_frame: bool = False,
+        is_backchannel: bool = False,
+        additional_metadata: Optional[dict] = None,
+    ) -> Optional[dict]:
+        """
+        Stage user audio metadata and transcription (from STT).
+        This data will be saved when the turn is complete by `save_user_audio` method.
+        Audio data is retrieved from continuous_user_audio_buffer based on timestamps.
+
+        Args:
+            timestamp_now: Timestamp when the audio was received
+            transcription: Transcribed text
+            sample_rate: Audio sample rate in Hz (default: 16000)
+            num_channels: Number of audio channels (default: 1)
+            is_first_frame: Whether this is the first frame of a turn (default: False)
+            is_backchannel: Whether this is a backchannel utterance (default: False)
+            additional_metadata: Additional metadata to include
+
+        Returns:
+            Dictionary with logged file paths, or None if logging is disabled
+        """
+        if not self.enabled:
+            return None
+
+        try:
+            # Get counter and generate filenames
+            counter = self._get_next_counter("user")
+            # timestamp_now = datetime.now()
+            base_name = f"{counter:05d}_{timestamp_now.strftime('%H%M%S')}"
+
+            audio_file = self.user_dir / f"{base_name}.wav"
+            metadata_file = self.user_dir / f"{base_name}.json"
+
+            if is_first_frame or self.staged_metadata is None or "start_time" not in self.staged_metadata:
+                raw_start_time = self.get_time_from_start_of_session(timestamp=timestamp_now)
+                # Apply pre-roll: go back pre_roll_time_sec, but don't go before the last entry's end time
+                pre_roll_start = raw_start_time - self._pre_roll_time_sec
+                if self.session_metadata["user_entries"]:
+                    last_entry_end_time = self.session_metadata["user_entries"][-1]["end_time"]
+                    _start_time = max(pre_roll_start, last_entry_end_time)
+                else:
+                    # No previous entries, just ensure we don't go negative
+                    _start_time = max(pre_roll_start, 0.0)
+            else:
+                # start_time is stored as float (seconds from session start), not ISO string
+                _start_time = self.staged_metadata["start_time"]
+
+            # Make end time into float (seconds from session start)
+            _end_time = self.get_time_from_start_of_session(timestamp=datetime.now())
+            audio_duration_sec = round(_end_time - _start_time, self._round_precision)
+
+            # Prepare metadata (initialize if None to allow update)
+            if self.staged_metadata is None:
+                self.staged_metadata = {}
+            self.staged_metadata.update(
+                {
+                    "base_name": base_name,
+                    "counter": counter,
+                    "turn_index": self._turn_index,
+                    "speaker": "user",
+                    "timestamp": timestamp_now.isoformat(),
+                    "start_time": _start_time,
+                    "end_time": _end_time,
+                    "transcription": transcription,
+                    "audio_file": audio_file.name,
+                    "sample_rate": sample_rate,
+                    "num_channels": num_channels,
+                    "audio_duration_sec": audio_duration_sec,
+                    "is_backchannel": is_backchannel,
+                }
+            )
+
+            if additional_metadata:
+                self.staged_metadata.update(additional_metadata)
+
+            return {
+                "audio_file": str(audio_file),
+                "metadata_file": str(metadata_file),
+                "counter": counter,
+            }
+
+        except Exception as e:
+            logger.error(f"Error logging user audio: {e}")
+            return None
+
+    def stage_turn_audio_and_transcription(
+        self,
+        timestamp_now: datetime,
+        is_first_frame: bool = False,
+        additional_metadata: Optional[dict] = None,
+    ):
+        """
+        Stage the complete turn audio and accumulated transcriptions.
+
+        This method is called when a final transcription is received.
+        It joins all accumulated audio and transcription chunks and stages them together.
+
+        Args:
+            timestamp_now: Timestamp when the audio was received
+            is_first_frame: Whether this is the first frame of a turn (default: False)
+            additional_metadata: Additional metadata to include (e.g., model, backend info)
+        """
+        if not self.turn_audio_buffer or not self.turn_transcription_buffer:
+            logger.debug("[AudioLogger] No audio or transcription to stage")
+            return
+
+        try:
+            complete_transcription = "".join(self.turn_transcription_buffer)
+
+            logger.debug(
+                f"[AudioLogger] Staging a turn with: {len(self.turn_audio_buffer)} audio chunks, "
+                f"{len(self.turn_transcription_buffer)} transcription chunks"
+            )
+
+            metadata = {
+                "num_transcription_chunks": len(self.turn_transcription_buffer),
+                "num_audio_chunks": len(self.turn_audio_buffer),
+            }
+            if additional_metadata:
+                metadata.update(additional_metadata)
+
+            self.stage_user_audio(
+                timestamp_now=timestamp_now,
+                transcription=complete_transcription,
+                sample_rate=self._stereo_sample_rate,
+                num_channels=1,
+                is_first_frame=is_first_frame,
+                additional_metadata=metadata,
+            )
+
+            logger.info(
+                f"[AudioLogger] Staged the audio and transcription for turn: '{complete_transcription[:50]}...'"
+            )
+
+        except Exception as e:
+            logger.warning(f"[AudioLogger] Failed to stage user audio: {e}")
+
+    def save_user_audio(self, is_backchannel: bool = False, float_divisor: float = 32768.0):
+        """Save the user audio to the disk.
+
+        Args:
+            is_backchannel: Whether this audio is a backchannel utterance (default: False)
+        """
+        # Safety check: ensure staged metadata exists and has required fields
+        if self.staged_metadata is None or "base_name" not in self.staged_metadata:
+            # This is expected - multiple TranscriptionFrames may be pushed but only one has audio staged
+            logger.debug("[AudioLogger] No staged metadata to save (this is normal for multiple frame pushes)")
+            return
+
+        try:
+            # Add backchannel metadata (only set if not already True to preserve turn-taking detection)
+            if is_backchannel or not self.staged_metadata.get("is_backchannel", False):
+                self.staged_metadata["is_backchannel"] = is_backchannel
+
+            audio_file = self.user_dir / f"{self.staged_metadata['base_name']}.wav"
+            metadata_file = self.user_dir / f"{self.staged_metadata['base_name']}.json"
+
+            # Get the audio data from continuous user audio buffer
+            stt, end = self.staged_metadata["start_time"], self.staged_metadata["end_time"]
+            continuous_audio_bytes = b"".join(self.continuous_user_audio_buffer)
+            full_audio_array = np.frombuffer(continuous_audio_bytes, dtype=np.int16).astype(np.float32) / float_divisor
+            start_idx = int(stt * self._stereo_sample_rate)
+            end_idx = int(end * self._stereo_sample_rate)
+            staged_audio_data = full_audio_array[start_idx:end_idx]
+
+            self._save_audio_wav(
+                audio_data=staged_audio_data,
+                file_path=audio_file,
+                sample_rate=self.staged_metadata["sample_rate"],
+            )
+
+            self._save_metadata_json(metadata=self.staged_metadata, file_path=metadata_file)
+            backchannel_label = " [BACKCHANNEL]" if is_backchannel else ""
+            transcription_preview = self.staged_metadata['transcription'][:50]
+            ellipsis = '...' if len(self.staged_metadata['transcription']) > 50 else ''
+            logger.info(
+                f"[AudioLogger] Saved user audio #{self.staged_metadata['counter']}"
+                f"{backchannel_label}: '{transcription_preview}{ellipsis}'"
+            )
+
+            # Note: User audio for stereo conversation is handled via continuous_user_audio_buffer
+            # which is populated in append_continuous_user_audio() (not affected by VAD)
+
+            # Update session metadata
+            with self._lock:
+                self.session_metadata["user_entries"].append(self.staged_metadata)
+                self._save_session_metadata()
+
+            self.clear_user_audio_buffer()
+
+            # Clear staged data after successful save
+            self.staged_metadata = None
+            self._staged_audio_data = None
+        except Exception as e:
+            logger.error(f"[AudioLogger] Error saving user audio: {e}")
+            raise
+
+    def log_agent_audio(
+        self,
+        audio_data: Union[bytes, np.ndarray],
+        text: str,
+        sample_rate: int = 22050,
+        num_channels: int = 1,
+        additional_metadata: Optional[dict] = None,
+        tts_generation_time: Optional[float] = None,
+    ) -> Optional[dict]:
+        """
+        Log agent audio and text (from TTS).
+
+        Args:
+            audio_data: Generated audio data as bytes or numpy array
+            text: Input text that was synthesized
+            sample_rate: Audio sample rate in Hz (default: 22050)
+            num_channels: Number of audio channels (default: 1)
+            additional_metadata: Additional metadata to include
+            tts_generation_time: Time when TTS generation started (seconds from session start).
+                                Used to calculate actual start_time for first segment of a turn.
+
+        Returns:
+            Dictionary with logged file paths, or None if logging is disabled
+        """
+        if not self.enabled:
+            return None
+
+        try:
+            # Get counter and generate filenames
+            counter = self._get_next_counter("agent")
+            timestamp_now = datetime.now()
+            base_name = f"{counter:05d}_{timestamp_now.strftime('%H%M%S')}"
+
+            audio_file = self.agent_dir / f"{base_name}.wav"
+            metadata_file = self.agent_dir / f"{base_name}.json"
+
+            # Save audio
+            self._save_audio_wav(audio_data, audio_file, sample_rate, num_channels)
+
+            # Calculate audio duration
+            audio_duration_sec = (
+                len(audio_data) / (sample_rate * num_channels * 2)
+                if isinstance(audio_data, bytes)
+                else len(audio_data) / sample_rate
+            )
+
+            # Determine start_time based on previous segment in the same turn
+            # If this is the first segment of the turn, use tts_generation_time
+            # Otherwise, use the previous segment's end_time for sequential playback
+            start_time = None
+            with self._lock:
+                agent_entries = self.session_metadata["agent_entries"]
+                # agent_entries is a list of turns, each turn is a list of segments
+                if agent_entries and agent_entries[-1]:  # If there's a current turn with segments
+                    last_segment = agent_entries[-1][-1]  # Last segment of last turn
+                    if last_segment["turn_index"] == self._turn_index:
+                        # Same turn - start after previous segment ends
+                        start_time = last_segment["end_time"]
+
+            if start_time is None:
+                # First segment of the turn - use agent_turn_start_time (from BotStartedSpeakingFrame)
+                # This is more accurate than tts_generation_time as it reflects actual playback start
+                if self._agent_turn_start_time is not None:
+                    start_time = self._agent_turn_start_time
+                elif tts_generation_time is not None:
+                    # Fallback to tts_generation_time if agent_turn_start_time not set
+                    start_time = tts_generation_time
+                else:
+                    start_time = self.get_time_from_start_of_session(timestamp=timestamp_now)
+
+            end_time = start_time + audio_duration_sec
+
+            # Prepare metadata
+            # cutoff_time is None by default (no interruption)
+            # It will be set by set_agent_cutoff_time() if TTS is interrupted
+            metadata = {
+                "base_name": base_name,
+                "counter": counter,
+                "turn_index": self._turn_index,
+                "speaker": "agent",
+                "timestamp": timestamp_now.isoformat(),
+                "start_time": round(start_time, self._round_precision),
+                "end_time": round(end_time, self._round_precision),
+                "cutoff_time": None,  # None means not interrupted; float if interrupted
+                "text": text,
+                "audio_file": audio_file.name,
+                "sample_rate": sample_rate,
+                "num_channels": num_channels,
+                "audio_duration_sec": round(audio_duration_sec, self._round_precision),
+            }
+
+            if additional_metadata:
+                metadata.update(additional_metadata)
+
+            # Save metadata
+            self._save_metadata_json(metadata, metadata_file)
+
+            # Append to stereo conversation (left channel = agent)
+            self._append_to_stereo_conversation(
+                audio_data=audio_data,
+                channel="left",
+                start_time=start_time,
+                sample_rate=sample_rate,
+            )
+
+            # Update session metadata
+            # agent_entries is a list of turns, each turn is a list of segments
+            with self._lock:
+                agent_entries = self.session_metadata["agent_entries"]
+                # Check if we need to start a new turn or append to existing turn
+                if not agent_entries or agent_entries[-1][-1]["turn_index"] != self._turn_index:
+                    # Start a new turn (new sublist)
+                    agent_entries.append([metadata])
+                else:
+                    # Append to current turn
+                    agent_entries[-1].append(metadata)
+                self._save_session_metadata()
+
+            logger.info(f"[AudioLogger] Logged agent audio #{counter}: '{text[:50]}{'...' if len(text) > 50 else ''}'")
+
+            return {
+                "audio_file": str(audio_file),
+                "metadata_file": str(metadata_file),
+                "counter": counter,
+            }
+
+        except Exception as e:
+            logger.error(f"[AudioLogger] Error logging agent audio: {e}")
+            return None
+
+    def set_agent_cutoff_time(self, cutoff_time: Optional[float] = None):
+        """
+        Set the cutoff time for the most recent agent audio entry.
+
+        This method should be called when TTS is interrupted by user speech.
+        The cutoff_time represents when the agent audio was actually cut off,
+        which may be earlier than the natural end_time.
+
+        Args:
+            cutoff_time: The cutoff time in seconds from session start.
+                        If None, uses current time from session start.
+        """
+        if not self.enabled:
+            return
+
+        if cutoff_time is None:
+            cutoff_time = self.get_time_from_start_of_session()
+
+        with self._lock:
+            agent_entries = self.session_metadata["agent_entries"]
+            if not agent_entries or not agent_entries[-1]:
+                logger.warning("[AudioLogger] No agent entries to set cutoff time")
+                return
+
+            # Get the current turn (last sublist) and update ALL segments in it
+            current_turn = agent_entries[-1]
+            turn_index = current_turn[0]["turn_index"]
+
+            # Update cutoff_time for ALL segments in the current turn
+            for segment in current_turn:
+                segment["cutoff_time"] = cutoff_time
+                # Also update individual JSON files
+                try:
+                    metadata_file = self.agent_dir / f"{segment['base_name']}.json"
+                    self._save_metadata_json(segment, metadata_file)
+                except Exception as e:
+                    logger.error(f"[AudioLogger] Error updating agent cutoff time for segment: {e}")
+
+            # Truncate the stereo buffer (left channel = agent) at the cutoff point
+            cutoff_sample = int(cutoff_time * self._stereo_sample_rate)
+            if cutoff_sample < len(self._stereo_audio_buffer_left):
+                # Zero out agent audio after cutoff point
+                for i in range(cutoff_sample, len(self._stereo_audio_buffer_left)):
+                    self._stereo_audio_buffer_left[i] = 0.0
+                logger.debug(
+                    f"[AudioLogger] Truncated agent stereo buffer at sample {cutoff_sample} "
+                    f"(cutoff_time={cutoff_time:.3f}s)"
+                )
+
+            logger.info(
+                f"[AudioLogger] Set cutoff_time={cutoff_time:.3f}s for turn {turn_index} "
+                f"({len(current_turn)} segments)"
+            )
+
+            # Save updated session metadata
+            self._save_session_metadata()
+
+    def _save_session_metadata(self):
+        """Save the session metadata to disk."""
+        if not self.enabled:
+            return
+
+        try:
+            metadata_file = self.session_dir / "session_metadata.json"
+            self.session_metadata["last_updated"] = datetime.now().isoformat()
+            self._save_metadata_json(self.session_metadata, metadata_file)
+        except Exception as e:
+            logger.error(f"[AudioLogger] Error saving session metadata: {e}")
+
+    def finalize_session(self):
+        """Finalize the session and save final metadata."""
+        if not self.enabled:
+            return
+
+        # Save stereo conversation before finalizing
+        self.save_stereo_conversation()
+
+        self.session_metadata["end_time"] = datetime.now().isoformat()
+        self.session_metadata["total_user_entries"] = self._user_counter
+        self.session_metadata["total_agent_segments"] = self._agent_counter
+        self.session_metadata["total_agent_turns"] = len(self.session_metadata["agent_entries"])
+        self._save_session_metadata()
+        logger.info(
+            f"[AudioLogger] Session finalized: {self.session_id} "
+            f"(User: {self._user_counter}, Agent: {self._agent_counter} segments in "
+            f"{len(self.session_metadata['agent_entries'])} turns)"
+        )
+
+
+class RTVIAudioLoggerObserver(BaseObserver):
+    """Observer that triggers audio logging when TranscriptionFrame is pushed."""
+
+    def __init__(self, audio_logger: AudioLogger):
+        super().__init__()
+        self._audio_logger = audio_logger
+
+    async def on_push_frame(self, data: FramePushed):
+        """Handle frame push events and save user audio on TranscriptionFrame."""
+        frame = data.frame
+        if isinstance(frame, TranscriptionFrame) and self._audio_logger:
+            self._audio_logger.save_user_audio()
+        # Call parent class's on_push_frame method
+        await super().on_push_frame(data)

nemo/agents/voice_agent/pipecat/services/nemo/diar.py

@@ -0,0 +1,360 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import asyncio
+from typing import AsyncGenerator, Optional
+
+import numpy as np
+from loguru import logger
+from pipecat.frames.frames import (
+    CancelFrame,
+    EndFrame,
+    ErrorFrame,
+    Frame,
+    StartFrame,
+    VADUserStartedSpeakingFrame,
+    VADUserStoppedSpeakingFrame,
+)
+from pipecat.processors.frame_processor import FrameDirection
+from pipecat.services.stt_service import STTService
+from pipecat.transcriptions.language import Language
+from pipecat.utils.time import time_now_iso8601
+from pipecat.utils.tracing.service_decorators import traced_stt
+from pydantic import BaseModel
+
+from nemo.agents.voice_agent.pipecat.frames.frames import DiarResultFrame
+from nemo.agents.voice_agent.pipecat.services.nemo.streaming_diar import DiarizationConfig, NeMoStreamingDiarService
+
+
+class NeMoDiarInputParams(BaseModel):
+    threshold: Optional[float] = (
+        0.4  # threshold value used to determine if a speaker exists or not, setting it to a lower value will increase the sensitivity of the diarization model
+    )
+    language: Optional[Language] = Language.EN_US
+    frame_len_in_secs: Optional[float] = 0.08  # 80ms for FastConformer model
+    config_path: Optional[str] = None  # path to the Niva ASR config file
+    raw_audio_frame_len_in_secs: Optional[float] = 0.016  # 16ms for websocket transport
+    buffer_size: Optional[int] = (
+        30  # number of audio frames to buffer, 1 frame is 16ms, streaming Sortformer was trained with 6*0.08=0.48s chunks
+    )
+
+
+class NemoDiarService(STTService):
+    def __init__(
+        self,
+        *,
+        model: Optional[str] = "",
+        device: Optional[str] = "cuda:0",
+        sample_rate: Optional[int] = 16000,
+        params: Optional[NeMoDiarInputParams] = None,
+        use_vad: bool = True,
+        audio_passthrough: bool = True,
+        backend: Optional[str] = "legacy",
+        enabled: bool = True,
+        **kwargs,
+    ):
+        super().__init__(audio_passthrough=audio_passthrough, **kwargs)
+
+        self._enabled = enabled
+        self._queue = asyncio.Queue()
+        self._response_queue = asyncio.Queue()  # Add response queue
+        self._processing_task = None  # Add processing task
+        self._response_task = None  # Add response task
+        self._device = device
+        self._sample_rate = sample_rate
+        self._audio_passthrough = audio_passthrough
+        params.buffer_size = params.frame_len_in_secs // params.raw_audio_frame_len_in_secs
+        self._params = params
+        self._model_name = model
+        self._use_vad = use_vad
+        self._backend = backend
+        if not params:
+            raise ValueError("params is required")
+
+        self._load_model()
+
+        self._vad_user_speaking = False
+        self._audio_buffer = []
+        self._current_speaker_id = None
+        self._processing_running = False
+
+        if not self._use_vad:
+            self._vad_user_speaking = True
+
+    def _load_model(self):
+        if not self._enabled or not self._model_name:
+            self._model = None
+            self._enabled = False
+            return
+
+        if self._backend == "legacy":
+            cfg = DiarizationConfig()
+            cfg.device = self._device
+            self._model = NeMoStreamingDiarService(
+                cfg, self._model_name, frame_len_in_secs=self._params.frame_len_in_secs, sample_rate=self.sample_rate
+            )
+        else:
+            raise ValueError(f"Invalid backend: {self._backend}")
+        logger.info(f"Diarization service initialized on device: {self._device}")
+
+    def can_generate_metrics(self) -> bool:
+        """Indicates whether this service can generate metrics.
+
+        Returns:
+            bool: True, as this service supports metric generation.
+        """
+        return True
+
+    async def start(self, frame: StartFrame):
+        """Handle service start."""
+        await super().start(frame)
+
+        # Initialize the model if not already done
+        if not hasattr(self, "_model"):
+            self._load_model()
+
+        # Start background processing task
+        if not self._processing_task:
+            self._processing_task = self.create_task(self._processing_task_handler())
+
+        # Start response handling task
+        if not self._response_task:
+            self._response_task = self.create_task(self._response_task_handler())
+
+    async def stop(self, frame: EndFrame):
+        """Handle service stop."""
+        await super().stop(frame)
+        await self._stop_tasks()
+
+    async def cancel(self, frame: CancelFrame):
+        """Handle service cancellation."""
+        await super().cancel(frame)
+        await self._stop_tasks()
+
+    async def _stop_tasks(self):
+        """Stop background processing tasks."""
+        await self._queue.put(None)  # Signal to stop processing
+        if self._processing_task:
+            await self.cancel_task(self._processing_task)
+            self._processing_task = None
+
+        if self._response_task:
+            await self.cancel_task(self._response_task)
+            self._response_task = None
+
+    def _diarization_processor(self):
+        """Background processor that handles diarization calls."""
+        try:
+            while self._processing_running:
+                try:
+                    # Get audio from queue - blocking call that will be interrupted by cancellation
+                    future = asyncio.run_coroutine_threadsafe(self._queue.get(), self.get_event_loop())
+                    audio = future.result()
+
+                    if audio is None:  # Stop signal
+                        logger.debug("Received stop signal in background processor")
+                        break
+
+                    # Process diarization
+                    diar_result = self._model.diarize(audio)
+
+                    # Send result back to async loop
+                    asyncio.run_coroutine_threadsafe(self._response_queue.put(diar_result), self.get_event_loop())
+
+                except Exception as e:
+                    logger.error(f"Error in background diarization processor: {e}")
+                    # Send error back to async loop
+                    asyncio.run_coroutine_threadsafe(self._response_queue.put(('error', e)), self.get_event_loop())
+
+        except Exception as e:
+            logger.error(f"Background diarization processor fatal error: {e}")
+        finally:
+            logger.debug("Background diarization processor stopped")
+
+    async def _processing_task_handler(self):
+        """Handler for background processing task."""
+        try:
+            self._processing_running = True
+            logger.debug("Starting background processing task")
+            await asyncio.to_thread(self._diarization_processor)
+        except asyncio.CancelledError:
+            logger.debug("Background processing task cancelled")
+            self._processing_running = False
+            raise
+        finally:
+            self._processing_running = False
+
+    async def _handle_diarization_result(self, diar_result):
+        """Handle diarization result from background processing."""
+        try:
+            if diar_result is None:
+                return
+            dominant_speaker_id = self._get_dominant_speaker_id(diar_result)
+            # logger.debug(f"Dominant speaker ID: {dominant_speaker_id}")
+            if dominant_speaker_id is not None and dominant_speaker_id != self._current_speaker_id:
+                self._current_speaker_id = dominant_speaker_id
+                logger.debug(f"Pushing DiarResultFrame with speaker {dominant_speaker_id}")
+                await self.push_frame(DiarResultFrame(dominant_speaker_id, stream_id="default"))
+        except Exception as e:
+            logger.error(f"Error handling diarization result: {e}")
+            await self.push_frame(
+                ErrorFrame(
+                    str(e),
+                    time_now_iso8601(),
+                )
+            )
+
+    async def _response_task_handler(self):
+        """Handler for processing diarization results."""
+        logger.debug("Response task handler started")
+        try:
+            while True:
+                try:
+                    result = await self._response_queue.get()
+
+                    if isinstance(result, tuple) and result[0] == 'error':
+                        # Handle error from background processing
+                        error = result[1]
+                        logger.error(f"Error in NeMo Diarization processing: {error}")
+                        await self.push_frame(
+                            ErrorFrame(
+                                str(error),
+                                time_now_iso8601(),
+                            )
+                        )
+                    else:
+                        # Handle successful diarization result
+                        await self._handle_diarization_result(result)
+
+                except Exception as e:
+                    logger.error(f"Error in response task handler: {e}")
+        except asyncio.CancelledError:
+            logger.debug("Response task handler cancelled")
+            raise
+
+    async def run_stt(self, audio: bytes) -> AsyncGenerator[Frame, None]:
+        """Process audio data and generate transcription frames.
+
+        Args:
+            audio: Raw audio bytes to transcribe
+
+        Yields:
+            Frame: Transcription frames containing the results
+        """
+        if self._vad_user_speaking and self._enabled:
+            self._audio_buffer.append(audio)
+            if len(self._audio_buffer) >= self._params.buffer_size:
+                await self.start_ttfb_metrics()
+                await self.start_processing_metrics()
+                audio = b"".join(self._audio_buffer)
+                self._audio_buffer = []
+                # Queue audio for background processing
+                await self._queue.put(audio)
+        yield None
+
+    @traced_stt
+    async def _handle_transcription(self, transcript: str, is_final: bool, language: Optional[str] = None):
+        """Handle a transcription result.
+
+        Args:
+            transcript: The transcribed text
+            is_final: Whether this is a final transcription
+            language: The language of the transcription
+        """
+        pass  # Base implementation - can be extended for specific handling needs
+
+    async def set_language(self, language: Language):
+        """Update the service's recognition language.
+
+        Args:
+            language: New language for recognition
+        """
+        if self._params:
+            self._params.language = language
+        else:
+            self._params = NeMoDiarInputParams(language=language)
+
+        logger.info(f"Switching STT language to: {language}")
+
+    async def set_model(self, model: str):
+        """Update the service's model.
+
+        Args:
+            model: New model name/path to use
+        """
+        await super().set_model(model)
+        self._model_name = model
+        self._load_model()
+
+    async def process_frame(self, frame: Frame, direction: FrameDirection):
+        """Process audio data and generate transcription frames.
+
+        Args:
+            audio: Raw audio bytes to transcribe
+
+        Yields:
+            Frame: Transcription frames containing the results
+        """
+        if not self._enabled:
+            # if diarization is disabled, just pass the frame through
+            await self.push_frame(frame, direction)
+            return
+
+        await super().process_frame(frame, direction)
+        if isinstance(frame, VADUserStartedSpeakingFrame):
+            self._vad_user_speaking = True
+            self._audio_buffer = []
+            logger.debug("VADUserStartedSpeakingFrame received")
+        elif isinstance(frame, VADUserStoppedSpeakingFrame):
+            self._vad_user_speaking = False
+            logger.debug("VADUserStoppedSpeakingFrame received")
+            self._current_speaker_id = None
+            self._audio_buffer = []
+
+    def reset(self):
+        """Reset the diarization service."""
+        self._current_speaker_id = None
+        self._audio_buffer = []
+        self._vad_user_speaking = False
+        self._model.reset_state()
+
+    def _get_dominant_speaker_id(self, spk_pred: np.ndarray):
+        spk_pred = (spk_pred > self._params.threshold).astype(int)
+        dominant_speaker_id = None
+        if spk_pred.sum() > 0:
+            # get the dominant speaker id
+            # Filter to only keep frames that have any speaker probability > 0.0
+            valid_frame_mask = spk_pred.sum(axis=1) > 0
+
+            # Filter diar_result to only keep valid frames
+            filtered_diar_result = spk_pred[valid_frame_mask]  # ndarray of shape [num_valid_frames, num_speakers]
+
+            # Get the primary speaker for each valid frame
+            primary_spk = np.argmax(filtered_diar_result, axis=1)  # ndarray of shape [num_valid_frames]
+            # logger.debug(f"Primary speaker for valid frames: {primary_spk}")
+
+            # count the number of different speakers in the primary speaker sequence
+            num_speakers = len(np.unique(primary_spk))
+            # logger.debug(f"Number of different speakers: {num_speakers}")
+
+            # If there are multiple speakers, get the dominant one
+            if num_speakers > 1:
+                # Count occurrences of each speaker
+                speaker_counts = np.bincount(primary_spk)
+                dominant_speaker_id = np.argmax(speaker_counts)
+            else:
+                # Only one speaker, return that speaker ID
+                dominant_speaker_id = primary_spk[0]
+        return dominant_speaker_id

nemo/agents/voice_agent/pipecat/services/nemo/llm.py

@@ -0,0 +1,760 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import asyncio
+import os
+import socket
+import subprocess
+import time
+import uuid
+from threading import Thread
+from typing import AsyncGenerator, List, Mapping, Optional
+
+import psutil
+import requests
+from jinja2.exceptions import TemplateError
+from loguru import logger
+from omegaconf import DictConfig, OmegaConf
+from openai import APITimeoutError, AsyncStream, BadRequestError
+from openai.types.chat import ChatCompletionChunk, ChatCompletionMessageParam
+from pipecat.adapters.services.open_ai_adapter import OpenAILLMInvocationParams
+from pipecat.frames.frames import (
+    CancelFrame,
+    EndFrame,
+    LLMFullResponseEndFrame,
+    LLMFullResponseStartFrame,
+    LLMTextFrame,
+)
+from pipecat.processors.aggregators.openai_llm_context import OpenAILLMContext
+from pipecat.services.openai.llm import OpenAILLMService
+from transformers import AsyncTextIteratorStreamer, AutoModelForCausalLM, AutoTokenizer
+from vllm.config import ModelConfig as vllmModelConfig
+
+DEFAULT_GENERATION_KWARGS = {
+    "max_new_tokens": 256,
+    "temperature": 0.7,
+    "top_p": 0.9,
+    "do_sample": True,
+}
+
+
+class LLMUtilsMixin:
+    """Utils for local LLM services."""
+
+    def _maybe_add_user_message(self, messages: List[ChatCompletionMessageParam]) -> List[ChatCompletionMessageParam]:
+        """
+        Some LLMs like "nvidia/Llama-3.1-Nemotron-Nano-8B-v1" requires a user turn after the system prompt,
+        this function is used to add a dummy user turn if the system prompt is followed by an assistant turn.
+        """
+        if len(messages) > 1 and messages[0]["role"] == "system" and messages[1]["role"] == "assistant":
+            message = {"role": "user", "content": "Hi"}
+            messages.insert(1, message)
+        elif len(messages) == 1 and messages[0]["role"] == "system":
+            messages.append({"role": "user", "content": "Hi"})
+        return messages
+
+    def _maybe_merge_consecutive_user_turns(
+        self, messages: List[ChatCompletionMessageParam]
+    ) -> List[ChatCompletionMessageParam]:
+        """
+        Merge consecutive user turns into a single turn,
+        since some LLMs like "nvidia/Llama-3.1-Nemotron-Nano-8B-v1" do not support consecutive user turns.
+        """
+        if not messages:
+            return messages
+
+        merged_messages = []
+
+        user_content = ""
+        for message in messages:
+            role = message["role"]
+            if role != "user":
+                # check if there's any preceeding user content, add them first
+                if user_content:
+                    merged_messages.append({"role": "user", "content": user_content})
+                    user_content = ""
+                merged_messages.append(message)
+            else:
+                if user_content:
+                    user_content += "; " + message["content"]
+                else:
+                    user_content = message["content"]
+
+        # add the last user content
+        if user_content:
+            merged_messages.append({"role": "user", "content": user_content})
+
+        return merged_messages
+
+
+class HuggingFaceLLMLocalService(LLMUtilsMixin):
+    """
+    HuggingFace LLM local service.
+    """
+
+    def __init__(
+        self,
+        model: str = "meta-llama/Meta-Llama-3-8B-Instruct",
+        device: str = "cuda:0",
+        dtype: str = "bfloat16",
+        thinking_budget: int = 0,
+        generation_kwargs: dict = None,
+        apply_chat_template_kwargs: dict = None,
+    ):
+        self.device = device
+        self.dtype = dtype
+        self.thinking_budget = thinking_budget
+        self.tokenizer = AutoTokenizer.from_pretrained(model)
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model, device_map=device, dtype=dtype, trust_remote_code=True
+        )  # type: AutoModelForCausalLM
+
+        self.generation_kwargs = generation_kwargs if generation_kwargs else DEFAULT_GENERATION_KWARGS
+        logger.debug(f"LLM generation kwargs: {self.generation_kwargs}")
+
+        self.apply_chat_template_kwargs = apply_chat_template_kwargs if apply_chat_template_kwargs else {}
+        if "tokenize" in self.apply_chat_template_kwargs:
+            if self.apply_chat_template_kwargs["tokenize"] is not False:
+                logger.warning(
+                    f"Found `tokenize=True` in apply_chat_template_kwargs={self.apply_chat_template_kwargs},"
+                    "it will be ignored and forced to `False`"
+                )
+            self.apply_chat_template_kwargs.pop("tokenize")
+
+        logger.debug(f"LLM apply_chat_template kwargs: {self.apply_chat_template_kwargs}")
+
+    def _apply_chat_template(self, messages: List[ChatCompletionMessageParam]) -> str:
+        """
+        Apply the chat template to the messages.
+        """
+        return self.tokenizer.apply_chat_template(messages, tokenize=False, **self.apply_chat_template_kwargs)
+
+    def _get_prompt_from_messages(self, messages: List[ChatCompletionMessageParam]) -> str:
+        """
+        Get the formatted prompt from the conversation history messages.
+        This function also tries to fix the messages if the LLM cannot handle consecutive turns of the same role,
+        or requires a user turn after the system prompt.
+        """
+        try:
+            prompt = self._apply_chat_template(messages)
+            return prompt
+        except TemplateError as e:
+            logger.warning(f"Got TemplateError: {e}.")
+
+        logger.debug(f"Input LLM messages: {messages}")
+        if len(messages) > 1 and messages[0]["role"] == "system" and messages[1]["role"] == "assistant":
+            logger.warning("Trying to fix by adding dummy user message after system prompt...")
+            try:
+                messages = self._maybe_add_user_message(messages)
+                logger.debug(f"LLM messages after adding dummy user message: {messages}")
+                prompt = self._apply_chat_template(messages)
+                return prompt
+            except TemplateError as e:
+                logger.warning(f"Got TemplateError: {e}. Trying to fix by merging consecutive turns if possible.")
+
+        try:
+            new_messages = self._maybe_merge_consecutive_user_turns(messages)
+            logger.debug(f"LLM messages after merging consecutive user turns: {new_messages}")
+            prompt = self._apply_chat_template(new_messages)
+            # Update the messages in place if successful
+            messages.clear()
+            messages.extend(new_messages)
+            return prompt
+        except Exception as e:
+            logger.warning(f"Got Exception: {e}, messages: {messages}")
+            raise e
+
+    async def generate_stream(
+        self, messages: List[ChatCompletionMessageParam], **kwargs
+    ) -> AsyncGenerator[ChatCompletionChunk, None]:
+        """
+        Generate a stream of chat completion chunks from the messages.
+        """
+        # Convert messages to prompt format
+        prompt = self._get_prompt_from_messages(messages)
+
+        logger.debug(f"LLM prompt: {prompt}")
+
+        inputs = self.tokenizer(prompt, add_special_tokens=False, return_tensors="pt").to(self.device)
+
+        # Generate with streaming
+        streamer = AsyncTextIteratorStreamer(self.tokenizer, skip_prompt=True, skip_special_tokens=True)
+        generation_kwargs = {
+            **inputs,
+            "streamer": streamer,
+            **self.generation_kwargs,
+        }
+
+        # Start generation in background
+        thread = Thread(
+            target=self.model.generate,
+            kwargs=generation_kwargs,
+        )
+        thread.start()
+
+        # Stream the output
+        async for text in streamer:
+            # logger.debug(f"Streamer yielded text: {text}")
+            chunk = ChatCompletionChunk(
+                id="hf-" + str(uuid.uuid4()),
+                choices=[{"delta": {"content": text}, "finish_reason": None, "index": 0}],
+                created=int(time.time()),
+                model=self.model.config._name_or_path,
+                object="chat.completion.chunk",
+            )
+            yield chunk
+
+
+class HuggingFaceLLMService(OpenAILLMService):
+    """
+    LLM service that hosts a HuggingFace model.
+    """
+
+    def __init__(
+        self,
+        *,
+        model: str = "google/gemma-7b-it",
+        device: str = "cuda",
+        dtype: str = "bfloat16",
+        thinking_budget: int = 0,
+        generation_kwargs: dict = None,
+        apply_chat_template_kwargs: dict = None,
+        **kwargs,
+    ):
+        self._model_name = model
+        self._device = device
+        self._dtype = dtype
+        self._thinking_budget = thinking_budget
+        self._generation_kwargs = generation_kwargs if generation_kwargs is not None else DEFAULT_GENERATION_KWARGS
+        self._apply_chat_template_kwargs = apply_chat_template_kwargs if apply_chat_template_kwargs is not None else {}
+        super().__init__(model=model, **kwargs)
+
+    def create_client(self, api_key=None, base_url=None, **kwargs):
+        """
+        Create a HuggingFaceLLMLocalService client.
+        """
+        return HuggingFaceLLMLocalService(
+            model=self._model_name,
+            device=self._device,
+            dtype=self._dtype,
+            thinking_budget=self._thinking_budget,
+            generation_kwargs=self._generation_kwargs,
+            apply_chat_template_kwargs=self._apply_chat_template_kwargs,
+        )
+
+    async def _process_context(self, context: OpenAILLMContext):
+        """Process a context through the LLM and push text frames.
+
+        Args:
+            context (OpenAILLMContext): The context to process, containing messages
+                and other information needed for the LLM interaction.
+        """
+        await self.push_frame(LLMFullResponseStartFrame())
+        cumulative_text = ""
+        try:
+            await self.start_ttfb_metrics()
+            messages = context.get_messages()
+            async for chunk in self._client.generate_stream(messages):
+                if chunk.choices[0].delta.content:
+                    await self.stop_ttfb_metrics()
+                    text = chunk.choices[0].delta.content
+                    cumulative_text += text
+                    frame = LLMTextFrame(text)
+                    await self.push_frame(frame)
+        except Exception as e:
+            logger.error(f"Error in _process_context: {e}", exc_info=True)
+            raise
+        finally:
+            cumulative_text = " ".join(cumulative_text.split()).strip()
+            if not cumulative_text:
+                logger.warning(f"LLM response is empty for context: {context}")
+            await self.push_frame(LLMFullResponseEndFrame())
+
+    async def get_chat_completions(
+        self, params_from_context: OpenAILLMInvocationParams
+    ) -> AsyncGenerator[ChatCompletionChunk, None]:
+        """Create a streaming chat completion using HuggingFace model.
+
+        Args:
+            context (OpenAILLMContext): The context object containing tools configuration
+                and other settings for the chat completion.
+            messages (List[ChatCompletionMessageParam]): The list of messages comprising
+                the conversation history and current request.
+
+        Returns:
+            AsyncGenerator[ChatCompletionChunk]: A streaming response of chat completion
+                chunks that can be processed asynchronously.
+        """
+        messages = params_from_context["messages"]
+
+        return self._client.generate_stream(messages)
+
+
+class VLLMService(OpenAILLMService, LLMUtilsMixin):
+    """
+    LLM service that hosts a vLLM server.
+    """
+
+    def __init__(
+        self,
+        *,
+        model: str,
+        device: str = "cuda",
+        api_key="None",
+        base_url="http://localhost:8000/v1",
+        organization="None",
+        project="None",
+        default_headers: Optional[Mapping[str, str]] = None,
+        params: Optional[OpenAILLMService.InputParams] = None,
+        thinking_budget: int = 0,
+        start_vllm_on_init: bool = False,
+        vllm_server_params: Optional[str] = None,
+        vllm_server_max_wait_time: int = 3600,  # 1 hour max wait time
+        vllm_server_check_interval: int = 5,  # check server every 5 seconds
+        **kwargs,
+    ):
+        self._device = device
+        self._vllm_server_max_wait_time = vllm_server_max_wait_time
+        self._vllm_server_check_interval = vllm_server_check_interval
+        if start_vllm_on_init:
+            base_url = self._start_vllm_server(model, vllm_server_params, base_url)
+
+        super().__init__(
+            model=model,
+            api_key=api_key,
+            base_url=base_url,
+            organization=organization,
+            project=project,
+            default_headers=default_headers,
+            params=params,
+            **kwargs,
+        )
+        self._thinking_budget = thinking_budget
+        self._vllm_server_params = vllm_server_params
+        self._start_vllm_on_init = start_vllm_on_init
+
+        # TODO: handle thinking budget
+        logger.info(
+            f"VLLMService initialized with model: {model}, api_key: {api_key}, base_url: {base_url},"
+            f"params: {params}, thinking_budget: {thinking_budget}"
+        )
+
+    def _start_vllm_server(
+        self, model: str, vllm_server_params: Optional[str] = None, base_url: Optional[str] = None
+    ) -> str:
+        """
+        Start a vllm server and return the base url.
+        """
+
+        requested_port = None
+        # If base_url is provided, extract port from it
+        if base_url:
+            try:
+                # Extract port from base_url like "http://localhost:8003/v1"
+                from urllib.parse import urlparse
+
+                parsed_url = urlparse(base_url)
+                if parsed_url.port:
+                    requested_port = parsed_url.port
+            except Exception as e:
+                logger.warning(
+                    f"Could not parse port from base_url {base_url}: {e}, using port from vllm_server_params"
+                )
+
+        # Parse port from vllm_server_params, default to 8000
+        if vllm_server_params:
+            params_list = vllm_server_params.split()
+            for i, param in enumerate(params_list):
+                if param == "--port" and i + 1 < len(params_list):
+                    try:
+                        param_port = int(params_list[i + 1])
+                        if requested_port is None:
+                            requested_port = param_port
+                        else:
+                            if param_port != requested_port:
+                                logger.warning(
+                                    f"Port {param_port} from vllm_server_params is different from base_url port"
+                                    f"{requested_port}, using new port {param_port}"
+                                )
+                                requested_port = param_port
+                        break
+                    except ValueError:
+                        logger.warning(f"Invalid port number: {params_list[i + 1]}, using default 8000")
+
+        if requested_port is None:
+            # try to use default port
+            requested_port = 8000
+
+        def find_available_port(start_port: int) -> int:
+            """Find an available port starting from start_port"""
+            for port in range(start_port, start_port + 100):  # Try up to 100 ports
+                try:
+                    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+                        s.bind(('localhost', port))
+                        return port
+                except OSError:
+                    continue
+            raise RuntimeError(f"Could not find an available port starting from {start_port}")
+
+        def get_pid_on_port(port: int) -> Optional[int]:
+            for conn in psutil.net_connections(kind="inet"):
+                if conn.laddr.port == port and conn.status == psutil.CONN_LISTEN:
+                    return conn.pid
+            return None
+
+        def check_server_model(port: int, verbose: bool = False) -> tuple[bool, str]:
+            """Check if server is running on port and return (is_running, model_name)"""
+            try:
+                response = requests.get(f"http://localhost:{port}/v1/models", timeout=5)
+                if response.status_code == 200:
+                    # get the PID for the server process
+                    pid = get_pid_on_port(port)
+                    if pid is not None and verbose:
+                        logger.warning(
+                            f"Found vLLM server process (PID: {pid}) on port {port}, you can use `lsof -i :{port}`"
+                            "to find the process and kill it if you want to start a new server."
+                        )
+                    models_data = response.json()
+                    if "data" in models_data and models_data["data"]:
+                        served_model = models_data["data"][0].get("id", "")
+                        return True, served_model
+                    return True, ""
+                return False, ""
+            except (requests.exceptions.RequestException, requests.exceptions.Timeout):
+                return False, ""
+
+        # First, check if vLLM server is already running on the requested port
+        is_running, served_model = check_server_model(requested_port, verbose=True)
+        if is_running:
+            if served_model == model:
+                final_base_url = f"http://localhost:{requested_port}/v1"
+                logger.info(f"vLLM server is already running at {final_base_url} with the correct model: {model}")
+                return final_base_url
+            else:
+                logger.warning(
+                    f"vLLM server on port {requested_port} is serving model '{served_model}' but we need '{model}'."
+                    "Finding new port..."
+                )
+
+        # Find an available port for our model
+        port = find_available_port(requested_port)
+        if port != requested_port:
+            logger.info(f"Using port {port} instead of requested port {requested_port}")
+
+        final_base_url = f"http://localhost:{port}/v1"
+
+        # Check if there's already a vLLM process running on the same port and model
+        for proc in psutil.process_iter(['pid', 'name', 'cmdline']):
+            try:
+                if proc.info['cmdline'] and any('vllm' in arg and 'serve' in arg for arg in proc.info['cmdline']):
+                    # Check if this process is using the same port and model
+                    cmdline_str = ' '.join(proc.info['cmdline'])
+                    if f"--port {port}" in cmdline_str:
+                        # Extract the model from the command line
+                        cmdline_parts = proc.info['cmdline']
+                        model_index = -1
+                        for i, arg in enumerate(cmdline_parts):
+                            if arg == "serve" and i + 1 < len(cmdline_parts):
+                                model_index = i + 1
+                                break
+
+                        if model_index != -1 and model_index < len(cmdline_parts):
+                            running_model = cmdline_parts[model_index]
+                            if running_model == model:
+                                logger.info(
+                                    f"Found existing vLLM server process (PID: {proc.info['pid']}) on port {port}"
+                                    f"serving model {model}"
+                                )
+                                # Wait a bit and check if it's responding
+                                time.sleep(2)
+                                is_running, served_model = check_server_model(port)
+                                if is_running and served_model == model:
+                                    logger.info(
+                                        f"Existing vLLM server is responding at {final_base_url} with correct model"
+                                    )
+                                    return final_base_url
+                                else:
+                                    logger.warning(
+                                        f"Existing vLLM process found on port {port} but not responding correctly,"
+                                        "will start new server"
+                                    )
+                            else:
+                                logger.info(
+                                    f"Found vLLM process on port {port} but serving different model '{running_model}'"
+                                    f"(need '{model}'). Will start new server."
+                                )
+            except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
+                continue
+
+        # Build the command with the determined port
+        cmd_parts = ["vllm", "serve", model]
+
+        # Parse and modify vllm_server_params to use the correct port
+        if vllm_server_params:
+            # parse the vllm_server_params and add the port to the command
+            params_list = vllm_server_params.split()
+            modified_params = []
+            i = 0
+            while i < len(params_list):
+                if params_list[i] == "--port" and i + 1 < len(params_list):
+                    # Replace the port with our determined port
+                    modified_params.extend(["--port", str(port)])
+                    i += 2  # Skip the original port value
+                else:
+                    modified_params.append(params_list[i])
+                    i += 1
+            cmd_parts.extend(modified_params)
+        else:
+            # Add port if vllm_server_params is not provided
+            cmd_parts.extend(["--port", str(port)])
+
+        logger.info(f"Starting vLLM server with command: {' '.join(cmd_parts)}")
+        logger.warning("It will take a while to download the model if it's not already downloaded.")
+        # Set up environment variables for device configuration
+        env = os.environ.copy()
+        if self._device and self._device != "cpu":
+            # Extract CUDA device number if it's in format "cuda:0", "cuda:1", etc.
+            if self._device.startswith("cuda:"):
+                device_id = self._device.split(":")[1]
+                env["CUDA_VISIBLE_DEVICES"] = device_id
+                logger.info(f"Setting CUDA_VISIBLE_DEVICES={device_id}")
+            elif self._device == "cuda":
+                # Use default CUDA device (don't set CUDA_VISIBLE_DEVICES)
+                logger.info("Using default CUDA device")
+            else:
+                # For other device strings, try to extract device number
+                logger.warning(f"Unknown device format: {self._device}, using as-is")
+                env["CUDA_VISIBLE_DEVICES"] = self._device
+        elif self._device == "cpu":
+            env["CUDA_VISIBLE_DEVICES"] = ""
+            logger.info("Setting CUDA_VISIBLE_DEVICES='' to use CPU")
+
+        try:
+            # Start the vLLM server process with environment variables
+            process = subprocess.Popen(
+                cmd_parts,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                text=True,
+                env=env,
+                preexec_fn=os.setsid if os.name != 'nt' else None,  # Create new process group
+            )
+
+            # Store the process for potential cleanup later
+            self._vllm_process = process
+
+            # Wait for server to start up
+            max_wait_time = self._vllm_server_max_wait_time
+            check_interval = self._vllm_server_check_interval
+            waited_time = 0
+
+            logger.info(f"Waiting for vLLM server to start on port {port}...")
+            while waited_time < max_wait_time:
+                is_running, served_model = check_server_model(port)
+                if is_running and served_model == model:
+                    logger.info(f"vLLM server started successfully at {final_base_url} serving model: {model}")
+                    return final_base_url
+                elif is_running and served_model != model:
+                    logger.warning(
+                        f"vLLM server started but serving wrong model '{served_model}' instead of '{model}'."
+                        "Continuing to wait..."
+                    )
+
+                # Check if process is still running
+                if process.poll() is not None:
+                    # Process has terminated
+                    stdout, stderr = process.communicate()
+                    logger.error(f"vLLM server process terminated unexpectedly. stdout: {stdout}, stderr: {stderr}")
+                    raise RuntimeError(f"Failed to start vLLM server: {stderr}")
+
+                time.sleep(check_interval)
+                waited_time += check_interval
+                logger.debug(f"Still waiting for vLLM server on port {port}... ({waited_time}s)")
+
+            # If we get here, server didn't start in time
+            logger.error(f"vLLM server failed to start within {max_wait_time} seconds on port {port}")
+            process.terminate()
+            raise RuntimeError(f"vLLM server failed to start within {max_wait_time} seconds on port {port}")
+
+        except FileNotFoundError:
+            logger.error("vLLM not found. Please install vLLM: pip install vllm")
+            raise RuntimeError("vLLM not found. Please install vLLM: pip install vllm")
+        except Exception as e:
+            logger.error(f"Failed to start vLLM server: {e}")
+            self._stop_vllm_server()
+            raise e
+
+    def _stop_vllm_server(self):
+        """Stop the vLLM server process if it's running."""
+        if hasattr(self, '_vllm_process') and self._vllm_process:
+            logger.info(f"Stopping vLLM server process {self._vllm_process.pid}")
+            self._vllm_process.terminate()
+
+    async def stop(self, frame: EndFrame):
+        """Stop the LLM service.
+
+        Args:
+            frame: The end frame.
+        """
+        await super().stop(frame)
+        self._stop_vllm_server()
+
+    async def cancel(self, frame: CancelFrame):
+        """Cancel the LLM service.
+
+        Args:
+            frame: The cancel frame.
+        """
+        await super().cancel(frame)
+        self._stop_vllm_server()
+
+    async def get_chat_completions(
+        self, params_from_context: OpenAILLMInvocationParams
+    ) -> AsyncStream[ChatCompletionChunk]:
+        """Get streaming chat completions from OpenAI API.
+
+        Args:
+            context: The LLM context containing tools and configuration.
+            messages: List of chat completion messages to send.
+
+        Returns:
+            Async stream of chat completion chunks.
+        """
+
+        params = self.build_chat_completion_params(params_from_context)
+        messages = params_from_context["messages"]
+        if self._retry_on_timeout:
+            try:
+                chunks = await asyncio.wait_for(
+                    self._get_response_from_client(messages, params), timeout=self._retry_timeout_secs
+                )
+                return chunks
+            except (APITimeoutError, asyncio.TimeoutError):
+                # Retry, this time without a timeout so we get a response
+                logger.debug(f"{self}: Retrying chat completion due to timeout")
+                chunks = await self._get_response_from_client(messages, params)
+                return chunks
+        else:
+            chunks = await self._get_response_from_client(messages, params)
+            return chunks
+
+    async def _get_response_from_client(
+        self, messages: List[ChatCompletionMessageParam], params: dict
+    ) -> AsyncStream[ChatCompletionChunk]:
+        """Get a response from the client."""
+
+        try:
+            chunks = await self._client.chat.completions.create(**params)
+        except BadRequestError as e:
+            logger.error(f"Error in _get_response_from_client: {e}, trying to fix...")
+            logger.debug(f"LLM messages before fixing: {messages}")
+            messages = self._maybe_add_user_message(messages)
+            messages = self._maybe_merge_consecutive_user_turns(messages)
+            logger.debug(f"LLM messages after fixing: {messages}")
+            params["messages"] = messages
+            chunks = await self._client.chat.completions.create(**params)
+
+        return chunks
+
+
+def get_llm_service_from_config(config: DictConfig) -> OpenAILLMService:
+    """Get an LLM service from the configuration."""
+    backend = config.type
+
+    logger.info(f"Initializing LLM service from config: {config}")
+
+    # If backend is "auto", try to detect the best backend
+    if backend == "auto":
+        model_name = config.get("model")
+        if not model_name:
+            raise ValueError("Model name is required for LLM")
+
+        try:
+            _ = vllmModelConfig(model_name, trust_remote_code=True)
+            backend = "vllm"
+            logger.info(f"Auto-detected vLLM as the best backend for model {model_name}")
+        except Exception as e:
+            logger.info(
+                f"The LLM doesn't seem to be supported by vLLM yet (error: {e}), using HuggingFace as the backend"
+                f"for model: {model_name}. If you are sure that the LLM is supported by vLLM, you can set `type: vllm`"
+                "in the config file to force using vLLM."
+            )
+            backend = "hf"
+
+    assert backend in [
+        "hf",
+        "vllm",
+        "auto",
+    ], f"Invalid backend: {backend}, only `hf`, `vllm`, and `auto` are supported."
+
+    if backend == "hf":
+        llm_model = config.model
+        llm_device = config.device
+        llm_dtype = config.dtype
+        llm_generation_kwargs = config.get("generation_kwargs", {})
+        if llm_generation_kwargs is not None:
+            llm_generation_kwargs = OmegaConf.to_container(llm_generation_kwargs, resolve=True)
+        llm_apply_chat_template_kwargs = config.get("apply_chat_template_kwargs", None)
+        if llm_apply_chat_template_kwargs is not None:
+            llm_apply_chat_template_kwargs = OmegaConf.to_container(llm_apply_chat_template_kwargs, resolve=True)
+        llm_thinking_budget = config.get("thinking_budget", 0)
+        return HuggingFaceLLMService(
+            model=llm_model,
+            device=llm_device,
+            dtype=llm_dtype,
+            generation_kwargs=llm_generation_kwargs,
+            apply_chat_template_kwargs=llm_apply_chat_template_kwargs,
+            thinking_budget=llm_thinking_budget,
+        )
+    elif backend == "vllm":
+        llm_model = config.get("model", "vllm_server")
+        llm_api_key = config.get("api_key", "None")
+        llm_base_url = config.get("base_url", "http://localhost:8000/v1")
+        llm_organization = config.get("organization", "None")
+        llm_project = config.get("project", "None")
+        llm_default_headers = config.get("default_headers", None)
+        llm_params = config.get("vllm_generation_params", None)
+        llm_dtype = config.dtype
+        vllm_server_params = config.get("vllm_server_params", None)
+        if vllm_server_params is not None:
+            if "dtype" not in vllm_server_params:
+                vllm_server_params = f"--dtype {llm_dtype} {vllm_server_params}"
+                logger.info(f"Adding dtype {llm_dtype} to vllm_server_params: {vllm_server_params}")
+        if llm_params is not None:
+            # cast into OpenAILLMService.InputParams object
+            llm_params = OmegaConf.to_container(llm_params, resolve=True)
+            extra = llm_params.get("extra", None)
+            # ensure extra is a dictionary
+            if extra is None:
+                llm_params["extra"] = {}
+            elif not isinstance(extra, dict):
+                raise ValueError(f"extra must be a dictionary, got {type(extra)}")
+            llm_params = OpenAILLMService.InputParams(**llm_params)
+        else:
+            llm_params = OpenAILLMService.InputParams()
+        llm_thinking_budget = config.get("thinking_budget", 0)
+        return VLLMService(
+            model=llm_model,
+            api_key=llm_api_key,
+            base_url=llm_base_url,
+            organization=llm_organization,
+            project=llm_project,
+            default_headers=llm_default_headers,
+            params=llm_params,
+            thinking_budget=llm_thinking_budget,
+            start_vllm_on_init=config.get("start_vllm_on_init", False),
+            vllm_server_params=vllm_server_params,
+        )
+    else:
+        raise ValueError(f"Invalid LLM backend: {backend}")

nemo/agents/voice_agent/pipecat/services/nemo/streaming_asr.py

@@ -0,0 +1,319 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# NOTE: This file will be deprecated in the future, as the new inference pipeline will replace it.
+
+import math
+import time
+from dataclasses import dataclass
+from typing import List, Optional
+
+import numpy as np
+import torch
+from omegaconf import open_dict
+
+import nemo.collections.asr as nemo_asr
+from nemo.agents.voice_agent.pipecat.services.nemo.utils import CacheFeatureBufferer
+from nemo.collections.asr.parts.utils.rnnt_utils import Hypothesis
+from nemo.collections.common.tokenizers.sentencepiece_tokenizer import SentencePieceTokenizer
+
+
+@dataclass
+class ASRResult:
+    text: str
+    is_final: bool
+    eou_prob: Optional[float] = None
+    eob_prob: Optional[float] = None
+    eou_latency: Optional[float] = None
+    eob_latency: Optional[float] = None
+    processing_time: Optional[float] = None
+
+
+class NemoStreamingASRService:
+    def __init__(
+        self,
+        model: str = "nvidia/parakeet_realtime_eou_120m-v1",
+        att_context_size: List[int] = [70, 1],
+        device: str = "cuda",
+        eou_string: str = "<EOU>",
+        eob_string: str = "<EOB>",
+        decoder_type: str = None,
+        chunk_size: int = -1,
+        shift_size: int = -1,
+        left_chunks: int = 2,
+        sample_rate: int = 16000,
+        frame_len_in_secs: float = 0.08,
+        use_amp: bool = False,
+        chunk_size_in_secs: float = 0.08,
+    ):
+        self.model = model
+        self.eou_string = eou_string
+        self.eob_string = eob_string
+        self.device = device
+        self.att_context_size = att_context_size
+        self.decoder_type = decoder_type
+        self.chunk_size = chunk_size
+        self.shift_size = shift_size
+        self.left_chunks = left_chunks
+        self.asr_model = self._load_model(model)
+        self.tokenizer: SentencePieceTokenizer = self.asr_model.tokenizer
+        self.use_amp = use_amp
+        self.pad_and_drop_preencoded = False
+        self.blank_id = self.get_blank_id()
+        self.chunk_size_in_secs = chunk_size_in_secs
+
+        assert len(self.att_context_size) == 2, "Att context size must be a list of two integers"
+        assert (
+            self.att_context_size[0] >= 0
+        ), f"Left att context size must be greater than 0: {self.att_context_size[0]}"
+        assert (
+            self.att_context_size[1] >= 0
+        ), f"Right att context size must be greater than 0: {self.att_context_size[1]}"
+
+        window_stride_in_secs = self.asr_model.cfg.preprocessor.window_stride
+        model_stride = self.asr_model.cfg.encoder.subsampling_factor
+        self.model_chunk_size = self.asr_model.encoder.streaming_cfg.chunk_size
+        if isinstance(self.model_chunk_size, list):
+            self.model_chunk_size = self.model_chunk_size[1]
+        self.pre_encode_cache_size = self.asr_model.encoder.streaming_cfg.pre_encode_cache_size
+        if isinstance(self.pre_encode_cache_size, list):
+            self.pre_encode_cache_size = self.pre_encode_cache_size[1]
+        self.pre_encode_cache_size_in_secs = self.pre_encode_cache_size * window_stride_in_secs
+
+        self.tokens_per_frame = math.ceil(np.trunc(self.chunk_size_in_secs / window_stride_in_secs) / model_stride)
+        # overwrite the encoder streaming params with proper shift size for cache aware streaming
+        self.asr_model.encoder.setup_streaming_params(
+            chunk_size=self.model_chunk_size // model_stride, shift_size=self.tokens_per_frame
+        )
+
+        model_chunk_size_in_secs = self.model_chunk_size * window_stride_in_secs
+
+        self.buffer_size_in_secs = self.pre_encode_cache_size_in_secs + model_chunk_size_in_secs
+
+        self._audio_buffer = CacheFeatureBufferer(
+            sample_rate=sample_rate,
+            buffer_size_in_secs=self.buffer_size_in_secs,
+            chunk_size_in_secs=self.chunk_size_in_secs,
+            preprocessor_cfg=self.asr_model.cfg.preprocessor,
+            device=self.device,
+        )
+        self._reset_cache()
+        self._previous_hypotheses = self._get_blank_hypothesis()
+        self._last_transcript_timestamp = time.time()
+        print(f"NemoStreamingASRService initialized with model `{model}` on device `{self.device}`")
+
+    def _reset_cache(self):
+        (
+            self._cache_last_channel,  # [17, B, 70, 512]
+            self._cache_last_time,  # [17, B, 512, 8]
+            self._cache_last_channel_len,  # B
+        ) = self.asr_model.encoder.get_initial_cache_state(
+            1
+        )  # batch size is 1
+
+    def _get_blank_hypothesis(self) -> List[Hypothesis]:
+        blank_hypothesis = Hypothesis(score=0.0, y_sequence=[], dec_state=None, timestamp=[], last_token=None)
+        return [blank_hypothesis]
+
+    @property
+    def drop_extra_pre_encoded(self):
+        return self.asr_model.encoder.streaming_cfg.drop_extra_pre_encoded
+
+    def get_blank_id(self):
+        return len(self.tokenizer.vocab)
+
+    def get_text_from_tokens(self, tokens: List[int]) -> str:
+        sep = "\u2581"  # '▁'
+        tokens = [int(t) for t in tokens if t != self.blank_id]
+        if tokens:
+            pieces = self.tokenizer.ids_to_tokens(tokens)
+            text = "".join([p.replace(sep, ' ') if p.startswith(sep) else p for p in pieces])
+        else:
+            text = ""
+        return text
+
+    def _load_model(self, model: str):
+        if model.endswith(".nemo"):
+            asr_model = nemo_asr.models.ASRModel.restore_from(model, map_location=torch.device(self.device))
+        else:
+            asr_model = nemo_asr.models.ASRModel.from_pretrained(model, map_location=torch.device(self.device))
+
+        if self.decoder_type is not None and hasattr(asr_model, "cur_decoder"):
+            asr_model.change_decoding_strategy(decoder_type=self.decoder_type)
+        elif isinstance(asr_model, nemo_asr.models.EncDecCTCModel):
+            self.decoder_type = "ctc"
+        elif isinstance(asr_model, nemo_asr.models.EncDecRNNTModel):
+            self.decoder_type = "rnnt"
+        else:
+            raise ValueError("Decoder type not supported for this model.")
+
+        if self.att_context_size is not None:
+            if hasattr(asr_model.encoder, "set_default_att_context_size"):
+                asr_model.encoder.set_default_att_context_size(att_context_size=self.att_context_size)
+            else:
+                raise ValueError("Model does not support multiple lookaheads.")
+        else:
+            self.att_context_size = asr_model.cfg.encoder.att_context_size
+
+        decoding_cfg = asr_model.cfg.decoding
+        with open_dict(decoding_cfg):
+            decoding_cfg.strategy = "greedy"
+            decoding_cfg.compute_timestamps = False
+            decoding_cfg.preserve_alignments = True
+            if hasattr(asr_model, 'joint'):  # if an RNNT model
+                decoding_cfg.greedy.max_symbols = 10
+                decoding_cfg.fused_batch_size = -1
+            asr_model.change_decoding_strategy(decoding_cfg)
+
+        if hasattr(asr_model.encoder, "set_default_att_context_size"):
+            asr_model.encoder.set_default_att_context_size(att_context_size=self.att_context_size)
+
+        # chunk_size is set automatically for models trained for streaming.
+        # For models trained for offline mode with full context, we need to pass the chunk_size explicitly.
+        if self.chunk_size > 0:
+            if self.shift_size < 0:
+                shift_size = self.chunk_size
+            else:
+                shift_size = self.shift_size
+            asr_model.encoder.setup_streaming_params(
+                chunk_size=self.chunk_size, left_chunks=self.left_chunks, shift_size=shift_size
+            )
+
+        asr_model.eval()
+        return asr_model
+
+    def _get_best_hypothesis(self, encoded, encoded_len, partial_hypotheses=None):
+        if self.decoder_type == "ctc":
+            best_hyp = self.asr_model.decoding.ctc_decoder_predictions_tensor(
+                encoded,
+                encoded_len,
+                return_hypotheses=True,
+            )
+        elif self.decoder_type == "rnnt":
+            best_hyp = self.asr_model.decoding.rnnt_decoder_predictions_tensor(
+                encoded, encoded_len, return_hypotheses=True, partial_hypotheses=partial_hypotheses
+            )
+        else:
+            raise ValueError("Decoder type not supported for this model.")
+        return best_hyp
+
+    def _get_tokens_and_probs_from_alignments(self, alignments):
+        tokens = []
+        probs = []
+        if self.decoder_type == "ctc":
+            all_logits = alignments[0]
+            all_tokens = alignments[1]
+            for i in range(len(all_tokens)):
+                token_id = int(all_tokens[i])
+                if token_id != self.blank_id:
+                    tokens.append(token_id)
+                    logits = all_logits[i]  # shape (vocab_size,)
+                    probs_i = torch.softmax(logits, dim=-1)[token_id].item()
+                    probs.append(probs_i)
+        elif self.decoder_type == "rnnt":
+            for t in range(len(alignments)):
+                for u in range(len(alignments[t])):
+                    logits, token_id = alignments[t][u]  # (logits, token_id)
+                    token_id = int(token_id)
+                    if token_id != self.blank_id:
+                        tokens.append(token_id)
+                        probs_i = torch.softmax(logits, dim=-1)[token_id].item()
+                        probs.append(probs_i)
+        else:
+            raise ValueError("Decoder type not supported for this model.")
+
+        return tokens, probs
+
+    def transcribe(self, audio: bytes, stream_id: str = "default") -> ASRResult:
+        start_time = time.time()
+
+        # Convert bytes to numpy array
+        audio_array = np.frombuffer(audio, dtype=np.int16).astype(np.float32) / 32768.0
+
+        self._audio_buffer.update(audio_array)
+
+        features = self._audio_buffer.get_feature_buffer()
+        feature_lengths = torch.tensor([features.shape[1]], device=self.device)
+        features = features.unsqueeze(0)  # Add batch dimension
+
+        with torch.no_grad():
+            (
+                encoded,
+                encoded_len,
+                cache_last_channel,
+                cache_last_time,
+                cache_last_channel_len,
+            ) = self.asr_model.encoder.cache_aware_stream_step(
+                processed_signal=features,
+                processed_signal_length=feature_lengths,
+                cache_last_channel=self._cache_last_channel,
+                cache_last_time=self._cache_last_time,
+                cache_last_channel_len=self._cache_last_channel_len,
+                keep_all_outputs=False,
+                drop_extra_pre_encoded=self.drop_extra_pre_encoded,
+            )
+
+        best_hyp = self._get_best_hypothesis(encoded, encoded_len, partial_hypotheses=self._previous_hypotheses)
+
+        self._previous_hypotheses = best_hyp
+        self._cache_last_channel = cache_last_channel
+        self._cache_last_time = cache_last_time
+        self._cache_last_channel_len = cache_last_channel_len
+
+        tokens, probs = self._get_tokens_and_probs_from_alignments(best_hyp[0].alignments)
+
+        text = self.get_text_from_tokens(tokens)
+
+        is_final = False
+        eou_latency = None
+        eob_latency = None
+        eou_prob = None
+        eob_prob = None
+        current_timestamp = time.time()
+        if self.eou_string in text or self.eob_string in text:
+            is_final = True
+            if self.eou_string in text:
+                eou_latency = (
+                    current_timestamp - self._last_transcript_timestamp if text.strip() == self.eou_string else 0.0
+                )
+                eou_prob = self.get_eou_probability(tokens, probs, self.eou_string)
+            if self.eob_string in text:
+                eob_latency = (
+                    current_timestamp - self._last_transcript_timestamp if text.strip() == self.eob_string else 0.0
+                )
+                eob_prob = self.get_eou_probability(tokens, probs, self.eob_string)
+            self.reset_state(stream_id=stream_id)
+        if text.strip():
+            self._last_transcript_timestamp = current_timestamp
+
+        processing_time = time.time() - start_time
+        return ASRResult(
+            text=text,
+            is_final=is_final,
+            eou_latency=eou_latency,
+            eob_latency=eob_latency,
+            eou_prob=eou_prob,
+            eob_prob=eob_prob,
+            processing_time=processing_time,
+        )
+
+    def reset_state(self, stream_id: str = "default"):
+        self._audio_buffer.reset()
+        self._reset_cache()
+        self._previous_hypotheses = self._get_blank_hypothesis()
+        self._last_transcript_timestamp = time.time()
+
+    def get_eou_probability(self, tokens: List[int], probs: List[float], eou_string: str = "<EOU>") -> float:
+        text_tokens = self.tokenizer.ids_to_tokens(tokens)
+        eou_index = text_tokens.index(eou_string)
+        return probs[eou_index]

nemo/agents/voice_agent/pipecat/services/nemo/streaming_diar.py

@@ -0,0 +1,212 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# NOTE: This file will be deprecated in the future, as the new inference pipeline will replace it.
+
+from dataclasses import dataclass
+from typing import Tuple
+
+import numpy as np
+import torch
+from torch import Tensor
+
+from nemo.agents.voice_agent.pipecat.services.nemo.utils import CacheFeatureBufferer
+from nemo.collections.asr.models import SortformerEncLabelModel
+
+from nemo.collections.asr.modules.sortformer_modules import StreamingSortformerState
+
+
+@dataclass
+class DiarizationConfig:
+    """Diarization configuration parameters for inference."""
+
+    model_path: str = "nvidia/diar_streaming_sortformer_4spk-v2"
+    device: str = "cuda"
+
+    log: bool = False  # If True, log will be printed
+    max_num_speakers: int = 4
+    spkcache_len: int = 188
+    spkcache_refresh_rate: int = 144
+    fifo_len: int = 188
+    chunk_len: int = 6
+    chunk_left_context: int = 1
+    chunk_right_context: int = 7
+
+
+class NeMoStreamingDiarService:
+    def __init__(
+        self,
+        cfg: DiarizationConfig,
+        model: str,
+        frame_len_in_secs: float = 0.08,
+        sample_rate: int = 16000,
+        left_offset: int = 8,
+        right_offset: int = 8,
+        use_amp: bool = False,
+        compute_dtype: torch.dtype = torch.float32,
+    ):
+        self.model = model
+        self.cfg = cfg
+        self.cfg.model_path = model
+        self.diarizer = self.build_diarizer()
+        self.device = cfg.device
+        self.use_amp = use_amp
+        self.compute_dtype = compute_dtype
+        self.frame_len_in_secs = frame_len_in_secs
+        self.left_offset = left_offset
+        self.right_offset = right_offset
+        self.chunk_size = self.cfg.chunk_len
+        self.buffer_size_in_secs = (
+            self.cfg.chunk_len * self.frame_len_in_secs + (self.left_offset + self.right_offset) * 0.01
+        )
+        self.max_num_speakers = self.cfg.max_num_speakers
+
+        self.feature_bufferer = CacheFeatureBufferer(
+            sample_rate=sample_rate,
+            buffer_size_in_secs=self.buffer_size_in_secs,
+            chunk_size_in_secs=self.cfg.chunk_len * self.frame_len_in_secs,
+            preprocessor_cfg=self.diarizer.cfg.preprocessor,
+            device=self.device,
+        )
+        self.streaming_state = self.init_streaming_state(batch_size=1)
+        self.total_preds = torch.zeros((1, 0, self.max_num_speakers), device=self.diarizer.device)
+
+        print(f"NeMoStreamingDiarService initialized with model `{model}` on device `{self.device}`")
+
+    def build_diarizer(self):
+        if self.cfg.model_path.endswith(".nemo"):
+            diar_model = SortformerEncLabelModel.restore_from(self.cfg.model_path, map_location=self.cfg.device)
+        else:
+            diar_model = SortformerEncLabelModel.from_pretrained(self.cfg.model_path, map_location=self.cfg.device)
+
+        # Steaming mode setup
+        diar_model.sortformer_modules.chunk_len = self.cfg.chunk_len
+        diar_model.sortformer_modules.spkcache_len = self.cfg.spkcache_len
+        diar_model.sortformer_modules.chunk_left_context = self.cfg.chunk_left_context
+        diar_model.sortformer_modules.chunk_right_context = self.cfg.chunk_right_context
+        diar_model.sortformer_modules.fifo_len = self.cfg.fifo_len
+        diar_model.sortformer_modules.log = self.cfg.log
+        diar_model.sortformer_modules.spkcache_refresh_rate = self.cfg.spkcache_refresh_rate
+        diar_model.eval()
+
+        return diar_model
+
+    def print_diar_result(self, diar_result: np.ndarray):
+        for t in range(diar_result.shape[0]):
+            spk_probs = ""
+            for s in range(diar_result.shape[1]):
+                spk_probs += f"{diar_result[t, s]:.2f} "
+            print(f"Time {t}: {spk_probs}")
+
+    def diarize(self, audio: bytes, stream_id: str = "default") -> str:
+
+        audio_array = np.frombuffer(audio, dtype=np.int16).astype(np.float32) / 32768.0
+
+        self.feature_bufferer.update(audio_array)
+
+        features = self.feature_bufferer.get_feature_buffer()
+        feature_buffers = features.unsqueeze(0)  # add batch dimension
+        feature_buffers = feature_buffers.transpose(1, 2)  # [batch, feature, time] -> [batch, time, feature]
+        feature_buffer_lens = torch.tensor([feature_buffers.shape[1]], device=self.device)
+        self.streaming_state, chunk_preds = self.stream_step(
+            processed_signal=feature_buffers,
+            processed_signal_length=feature_buffer_lens,
+            streaming_state=self.streaming_state,
+            total_preds=self.total_preds,
+            left_offset=self.left_offset,
+            right_offset=self.right_offset,
+        )
+        self.total_preds = chunk_preds
+        diar_result = chunk_preds[:, -self.chunk_size :, :].clone().cpu().numpy()
+        return diar_result[0]  # tensor of shape [6, 4]
+
+    def reset_state(self, stream_id: str = "default"):
+        self.feature_bufferer.reset()
+        self.streaming_state = self.init_streaming_state(batch_size=1)
+        self.total_preds = torch.zeros((1, 0, self.max_num_speakers), device=self.diarizer.device)
+
+    def init_streaming_state(self, batch_size: int = 1) -> StreamingSortformerState:
+        """
+        Initialize the streaming state for the diarization model.
+
+        Args:
+            batch_size: The batch size to use.
+
+        Returns:
+            SortformerStreamingState: The initialized streaming state.
+        """
+        # Use the model's init_streaming_state method but convert to SortformerStreamingState format
+        nemo_state = self.diarizer.sortformer_modules.init_streaming_state(
+            batch_size=batch_size, async_streaming=self.diarizer.async_streaming, device=self.device
+        )
+
+        return nemo_state
+
+    def stream_step(
+        self,
+        processed_signal: Tensor,
+        processed_signal_length: Tensor,
+        streaming_state: StreamingSortformerState,
+        total_preds: Tensor,
+        left_offset: int = 0,
+        right_offset: int = 0,
+    ) -> Tuple[StreamingSortformerState, Tensor]:
+        """
+        Execute a single streaming step for diarization.
+
+        Args:
+            processed_signal: The processed audio signal.
+            processed_signal_length: The length of the processed signal.
+            streaming_state: The current streaming state.
+            total_preds: The total predictions so far.
+            left_offset: The left offset for the current chunk.
+            right_offset: The right offset for the current chunk.
+
+        Returns:
+            Tuple[SortformerStreamingState, Tensor]: The updated streaming state and predictions.
+        """
+        # Move tensors to correct device
+        if processed_signal.device != self.device:
+            processed_signal = processed_signal.to(self.device)
+
+        if processed_signal_length.device != self.device:
+            processed_signal_length = processed_signal_length.to(self.device)
+
+        if total_preds is not None and total_preds.device != self.device:
+            total_preds = total_preds.to(self.device)
+
+        with (
+            torch.amp.autocast(device_type=self.device, dtype=self.compute_dtype, enabled=self.use_amp),
+            torch.inference_mode(),
+            torch.no_grad(),
+        ):
+            try:
+                # Call the model's forward_streaming_step method
+                streaming_state, diar_pred_out_stream = self.diarizer.forward_streaming_step(
+                    processed_signal=processed_signal,
+                    processed_signal_length=processed_signal_length,
+                    streaming_state=streaming_state,
+                    total_preds=total_preds,
+                    left_offset=left_offset,
+                    right_offset=right_offset,
+                )
+            except Exception as e:
+                print(f"Error in diarizer streaming step: {e}")
+                # print the stack trace
+                import traceback
+
+                traceback.print_exc()
+                # Return the existing state and preds if there's an error
+                return streaming_state, total_preds
+
+        return streaming_state, diar_pred_out_stream

nemo/agents/voice_agent/pipecat/services/nemo/stt.py

@@ -0,0 +1,316 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import asyncio
+from datetime import datetime
+from typing import AsyncGenerator, List, Optional
+
+from loguru import logger
+from pipecat.frames.frames import (
+    CancelFrame,
+    EndFrame,
+    ErrorFrame,
+    Frame,
+    InterimTranscriptionFrame,
+    StartFrame,
+    TranscriptionFrame,
+    VADUserStartedSpeakingFrame,
+    VADUserStoppedSpeakingFrame,
+)
+from pipecat.processors.frame_processor import FrameDirection
+from pipecat.services.stt_service import STTService
+from pipecat.transcriptions.language import Language
+from pipecat.utils.time import time_now_iso8601
+from pipecat.utils.tracing.service_decorators import traced_stt
+from pydantic import BaseModel
+
+from nemo.agents.voice_agent.pipecat.services.nemo.audio_logger import AudioLogger
+from nemo.agents.voice_agent.pipecat.services.nemo.streaming_asr import NemoStreamingASRService
+
+ASR_EOU_MODELS = ["nvidia/parakeet_realtime_eou_120m-v1"]
+
+try:
+    # disable nemo logging
+    from nemo.utils import logging
+
+    level = logging.getEffectiveLevel()
+    logging.setLevel(logging.CRITICAL)
+
+
+except ModuleNotFoundError as e:
+    logger.error(f"Exception: {e}")
+    logger.error('In order to use NVIDIA NeMo STT, you need to `pip install "nemo-toolkit[all]"`.')
+    raise Exception(f"Missing module: {e}")
+
+
+class NeMoSTTInputParams(BaseModel):
+    """Input parameters for NeMo STT service."""
+
+    language: Optional[Language] = Language.EN_US
+    att_context_size: Optional[List] = [70, 1]
+    frame_len_in_secs: Optional[float] = 0.08  # 80ms for FastConformer model
+    config_path: Optional[str] = None  # path to the Niva ASR config file
+    raw_audio_frame_len_in_secs: Optional[float] = 0.016  # 16ms for websocket transport
+    buffer_size: Optional[int] = 5  # number of audio frames to buffer, 1 frame is 16ms
+
+
+class NemoSTTService(STTService):
+    """NeMo Speech-to-Text service for Pipecat integration."""
+
+    def __init__(
+        self,
+        *,
+        model: Optional[str] = "nnvidia/parakeet_realtime_eou_120m-v1",
+        device: Optional[str] = "cuda:0",
+        sample_rate: Optional[int] = 16000,
+        params: Optional[NeMoSTTInputParams] = None,
+        has_turn_taking: Optional[bool] = None,  # if None, it will be set by the model name
+        backend: Optional[str] = "legacy",
+        decoder_type: Optional[str] = "rnnt",
+        audio_logger: Optional[AudioLogger] = None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self._queue = asyncio.Queue()
+        self._sample_rate = sample_rate
+        self._params = params or NeMoSTTInputParams()
+        self._model_name = model
+        if has_turn_taking is None:
+            has_turn_taking = True if model in ASR_EOU_MODELS else False
+            logger.info(f"Setting has_turn_taking to `{has_turn_taking}` based on model name: `{model}`")
+        self._has_turn_taking = has_turn_taking
+        self._backend = backend
+        self._decoder_type = decoder_type
+        self._audio_logger = audio_logger
+        self._is_vad_active = False
+        logger.info(f"NeMoSTTInputParams: {self._params}")
+
+        self._device = device
+
+        self._load_model()
+
+        self.audio_buffer = []
+        self.user_is_speaking = False
+
+    def _load_model(self):
+        if self._backend == "legacy":
+            self._model = NemoStreamingASRService(
+                self._model_name,
+                self._params.att_context_size,
+                device=self._device,
+                decoder_type=self._decoder_type,
+                frame_len_in_secs=self._params.frame_len_in_secs,
+            )
+        else:
+            raise ValueError(f"Invalid ASR backend: {self._backend}")
+
+    def can_generate_metrics(self) -> bool:
+        """Indicates whether this service can generate metrics.
+
+        Returns:
+            bool: True, as this service supports metric generation.
+        """
+        return True
+
+    async def start(self, frame: StartFrame):
+        """Handle service start.
+
+        Args:
+            frame: StartFrame containing initial configuration
+        """
+        await super().start(frame)
+
+        # Initialize the model if not already done
+        if not hasattr(self, "_model"):
+            self._load_model()
+
+    async def stop(self, frame: EndFrame):
+        """Handle service stop.
+
+        Args:
+            frame: EndFrame that triggered this method
+        """
+        await super().stop(frame)
+        # Clear any internal state if needed
+        await self._queue.put(None)  # Signal to stop processing
+
+    async def cancel(self, frame: CancelFrame):
+        """Handle service cancellation.
+
+        Args:
+            frame: CancelFrame that triggered this method
+        """
+        await super().cancel(frame)
+        # Clear any internal state
+        await self._queue.put(None)  # Signal to stop processing
+        self._queue = asyncio.Queue()  # Reset the queue
+
+    async def run_stt(self, audio: bytes) -> AsyncGenerator[Frame, None]:
+        """Process audio data and generate transcription frames.
+
+        Args:
+            audio: Raw audio bytes to transcribe
+
+        Yields:
+            Frame: Transcription frames containing the results
+        """
+        timestamp_now = datetime.now()
+        await self.start_ttfb_metrics()
+        await self.start_processing_metrics()
+        if self._audio_logger is not None and self._audio_logger.first_audio_timestamp is None:
+            self._audio_logger.first_audio_timestamp = timestamp_now
+
+        try:
+            is_final = False
+            user_has_finished = False
+            transcription = None
+            self.audio_buffer.append(audio)
+            if len(self.audio_buffer) >= self._params.buffer_size:
+                audio = b"".join(self.audio_buffer)
+                self.audio_buffer = []
+
+                # Append to continuous user audio buffer for stereo conversation recording
+                if self._audio_logger is not None:
+                    self._audio_logger.append_continuous_user_audio(audio)
+
+                asr_result = self._model.transcribe(audio)
+                transcription = asr_result.text
+                is_final = asr_result.is_final
+                if self._audio_logger is not None:
+                    if self._is_vad_active:
+                        is_first_frame = False
+                        self._audio_logger.turn_audio_buffer.append(audio)
+                        # Accumulate transcriptions for turn-based logging
+                        if transcription:
+                            self._audio_logger.turn_transcription_buffer.append(transcription)
+                            self._audio_logger.stage_turn_audio_and_transcription(
+                                timestamp_now=timestamp_now,
+                                is_first_frame=is_first_frame,
+                                additional_metadata={
+                                    "model": self._model_name,
+                                    "backend": self._backend,
+                                },
+                            )
+                eou_latency = asr_result.eou_latency
+                eob_latency = asr_result.eob_latency
+                eou_prob = asr_result.eou_prob
+                eob_prob = asr_result.eob_prob
+                if eou_latency is not None:
+                    logger.debug(
+                        f"EOU latency: {eou_latency: .4f} seconds. EOU probability: {eou_prob: .2f}."
+                        f"Processing time: {asr_result.processing_time: .4f} seconds."
+                    )
+                    user_has_finished = True
+                if eob_latency is not None:
+                    logger.debug(
+                        f"EOB latency: {eob_latency: .4f} seconds. EOB probability: {eob_prob: .2f}."
+                        f"Processing time: {asr_result.processing_time: .4f} seconds."
+                    )
+                    user_has_finished = True
+                await self.stop_ttfb_metrics()
+                await self.stop_processing_metrics()
+
+            if transcription:
+                logger.debug(f"Transcription (is_final={is_final}): `{transcription}`")
+                self.user_is_speaking = True if not user_has_finished else False
+
+                # Get the language from params or default to EN_US
+                language = self._params.language if self._params else Language.EN_US
+
+                # Create and push the transcription frame
+                if self._has_turn_taking:
+                    # if turn taking is enabled, we push interim transcription frames
+                    # and let the turn taking service handle the final transcription
+                    frame_type = InterimTranscriptionFrame
+                else:
+                    # otherwise, we use the is_final flag to determine the frame type
+                    frame_type = TranscriptionFrame if is_final else InterimTranscriptionFrame
+                await self.push_frame(
+                    frame_type(
+                        transcription,
+                        "",  # No speaker ID in this implementation
+                        time_now_iso8601(),
+                        language,
+                        result={"text": transcription},
+                    )
+                )
+
+                # Handle the transcription
+                await self._handle_transcription(
+                    transcript=transcription,
+                    is_final=is_final,
+                    language=language,
+                )
+
+            yield None
+
+        except Exception as e:
+            logger.error(f"Error in NeMo STT processing: {e}")
+            await self.push_frame(
+                ErrorFrame(
+                    str(e),
+                    time_now_iso8601(),
+                )
+            )
+            yield None
+
+    @traced_stt
+    async def _handle_transcription(self, transcript: str, is_final: bool, language: Optional[str] = None):
+        """Handle a transcription result.
+
+        Args:
+            transcript: The transcribed text
+            is_final: Whether this is a final transcription
+            language: The language of the transcription
+        """
+        pass  # Base implementation - can be extended for specific handling needs
+
+    async def set_language(self, language: Language):
+        """Update the service's recognition language.
+
+        Args:
+            language: New language for recognition
+        """
+        if self._params:
+            self._params.language = language
+        else:
+            self._params = NeMoSTTInputParams(language=language)
+
+        logger.info(f"Switching STT language to: {language}")
+
+    async def set_model(self, model: str):
+        """Update the service's model.
+
+        Args:
+            model: New model name/path to use
+        """
+        await super().set_model(model)
+        self._model_name = model
+        self._load_model()
+
+    async def process_frame(self, frame: Frame, direction: FrameDirection):
+        """Process incoming frames and handle VAD events."""
+        if isinstance(frame, VADUserStoppedSpeakingFrame) and isinstance(self._model, NemoStreamingASRService):
+            # manualy reset the state of the model when end of utterance is detected by VAD
+            logger.debug("Resetting state of the model due to VADUserStoppedSpeakingFrame")
+            if self.user_is_speaking:
+                logger.debug(
+                    "[EOU missing] STT failed to detect end of utterance before VAD detected user stopped speaking"
+                )
+            self._model.reset_state()
+            self._is_vad_active = False
+        elif isinstance(frame, VADUserStartedSpeakingFrame):
+            self._is_vad_active = True
+
+        await super().process_frame(frame, direction)

nemo/agents/voice_agent/pipecat/services/nemo/tts.py

@@ -0,0 +1,892 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import asyncio
+import inspect
+import uuid
+from collections.abc import AsyncGenerator
+from datetime import datetime
+from typing import Iterator, List, Optional
+
+import numpy as np
+import torch
+from loguru import logger
+from omegaconf import DictConfig, OmegaConf
+from pipecat.frames.frames import (
+    CancelFrame,
+    EndFrame,
+    ErrorFrame,
+    Frame,
+    LLMTextFrame,
+    StartFrame,
+    TTSAudioRawFrame,
+    TTSStartedFrame,
+    TTSStoppedFrame,
+)
+from pipecat.services.llm_service import FunctionCallParams
+from pipecat.services.tts_service import TTSService
+
+from nemo.agents.voice_agent.pipecat.services.nemo.audio_logger import AudioLogger
+from nemo.agents.voice_agent.pipecat.utils.text.simple_text_aggregator import SimpleSegmentedTextAggregator
+from nemo.agents.voice_agent.utils.tool_calling.mixins import ToolCallingMixin
+from nemo.collections.tts.models import FastPitchModel, HifiGanModel
+
+
+class BaseNemoTTSService(TTSService, ToolCallingMixin):
+    """Text-to-Speech service using Nemo TTS models.
+
+    This service works with any TTS model that exposes a generate(text) method
+    that returns audio data. The TTS generation runs in a dedicated background thread to
+    avoid blocking the main asyncio event loop, following the same pattern as NemoDiarService.
+
+    Args:
+        model: TTS model instance with a generate(text) method
+        sample_rate: Audio sample rate in Hz (defaults to 22050)
+        **kwargs: Additional arguments passed to TTSService
+    """
+
+    def __init__(
+        self,
+        *,
+        model,
+        device: str = "cuda",
+        sample_rate: int = 22050,
+        think_tokens: Optional[List[str]] = None,
+        audio_logger: Optional[AudioLogger] = None,
+        ignore_strings: Optional[List[str]] = None,
+        **kwargs,
+    ):
+        super().__init__(sample_rate=sample_rate, **kwargs)
+        logger.info(f"Initializing TTS service with model: {model} and device: {device}")
+        self._model_name = model
+        self._device = device
+        self._model = self._setup_model()
+        self._think_tokens = think_tokens
+        self._audio_logger = audio_logger
+        if think_tokens is not None:
+            assert (
+                isinstance(think_tokens, list) and len(think_tokens) == 2
+            ), f"think_tokens must be a list of two strings, but got type {type(think_tokens)}: {think_tokens}"
+        self._ignore_strings = set(ignore_strings) if ignore_strings is not None else None
+        # Background processing infrastructure - no response handler needed
+        self._tts_queue = asyncio.Queue()
+        self._processing_task = None
+        self._processing_running = False
+
+        # Track pending requests with their response queues
+        self._pending_requests = {}
+        self._have_seen_think_tokens = False
+
+    def reset(self):
+        """Reset the TTS service."""
+        self._text_aggregator.reset()
+
+    def setup_tool_calling(self):
+        """
+        Setup the tool calling mixin by registering all available tools.
+        """
+        pass  # No tools by default
+
+    def _setup_model(self):
+        raise NotImplementedError("Subclass must implement _setup_model")
+
+    def _generate_audio(self, text: str) -> Iterator[np.ndarray]:
+        raise NotImplementedError("Subclass must implement _generate_audio")
+
+    def can_generate_metrics(self) -> bool:
+        """If the TTS service can generate metrics."""
+        return True
+
+    async def start(self, frame: StartFrame):
+        """Handle service start."""
+        await super().start(frame)
+
+        # Initialize the model if not already done
+        if not hasattr(self, "_model") or self._model is None:
+            self._model = self._setup_model()
+
+        # Only start background processing task - no response handler needed
+        if not self._processing_task:
+            self._processing_task = self.create_task(self._processing_task_handler())
+
+    async def stop(self, frame: EndFrame):
+        """Handle service stop."""
+        await super().stop(frame)
+        await self._stop_tasks()
+
+    async def cancel(self, frame: CancelFrame):
+        """Handle service cancellation."""
+        await super().cancel(frame)
+        await self._stop_tasks()
+
+    async def _stop_tasks(self):
+        """Stop background processing tasks."""
+        self._processing_running = False
+        await self._tts_queue.put(None)  # Signal to stop processing
+
+        if self._processing_task:
+            await self.cancel_task(self._processing_task)
+            self._processing_task = None
+
+    def _tts_processor(self):
+        """Background processor that handles TTS generation calls."""
+        try:
+            while self._processing_running:
+                try:
+                    future = asyncio.run_coroutine_threadsafe(self._tts_queue.get(), self.get_event_loop())
+                    request = future.result()
+
+                    if request is None:  # Stop signal
+                        logger.debug("Received stop signal in TTS background processor")
+                        break
+
+                    text, request_id = request
+                    logger.debug(f"Processing TTS request for text: [{text}]")
+
+                    # Get the response queue for this request
+                    response_queue = None
+                    future = asyncio.run_coroutine_threadsafe(
+                        self._get_response_queue(request_id), self.get_event_loop()
+                    )
+                    response_queue = future.result()
+
+                    if response_queue is None:
+                        logger.warning(f"No response queue found for request {request_id}")
+                        continue
+
+                    # Process TTS generation
+                    try:
+                        audio_result = self._generate_audio(text)
+
+                        # Send result directly to the waiting request
+                        asyncio.run_coroutine_threadsafe(
+                            response_queue.put(('success', audio_result)), self.get_event_loop()
+                        )
+                    except Exception as e:
+                        logger.error(f"Error in TTS generation: {e}")
+                        # Send error directly to the waiting request
+                        asyncio.run_coroutine_threadsafe(response_queue.put(('error', e)), self.get_event_loop())
+
+                except Exception as e:
+                    logger.error(f"Error in background TTS processor: {e}")
+
+        except Exception as e:
+            logger.error(f"Background TTS processor fatal error: {e}")
+        finally:
+            logger.debug("Background TTS processor stopped")
+
+    async def _get_response_queue(self, request_id: str):
+        """Get the response queue for a specific request."""
+        return self._pending_requests.get(request_id)
+
+    async def _processing_task_handler(self):
+        """Handler for background processing task."""
+        try:
+            self._processing_running = True
+            logger.debug("Starting background TTS processing task")
+            await asyncio.to_thread(self._tts_processor)
+        except asyncio.CancelledError:
+            logger.debug("Background TTS processing task cancelled")
+            self._processing_running = False
+            raise
+        finally:
+            self._processing_running = False
+
+    def _handle_think_tokens(self, text: str) -> Optional[str]:
+        """
+        Handle the thinking tokens for TTS.
+        If the thinking tokens are not provided, return the text as it is.
+        Otherwise:
+            If both thinking tokens appear in the text, return the text after the end of thinking tokens.
+            If the LLM is thinking, return None.
+            If the LLM is done thinking, return the text after the end of thinking tokens.
+            If the LLM starts thinking, return the text before the start of thinking tokens.
+            If the LLM is not thinking, return the text as is.
+        """
+        if not self._think_tokens or not text:
+            return text
+        elif self._think_tokens[0] in text and self._think_tokens[1] in text:
+            # LLM finishes thinking in one chunk or outputs dummy thinking tokens
+            logger.debug(f"LLM finishes thinking: {text}")
+            idx = text.index(self._think_tokens[1])
+            # only return the text after the end of thinking tokens
+            text = text[idx + len(self._think_tokens[1]) :]
+            self._have_seen_think_tokens = False
+            logger.debug(f"Returning text after thinking: {text}")
+            return text
+        elif self._have_seen_think_tokens:
+            # LLM is thinking
+            if self._think_tokens[1] not in text:
+                logger.debug(f"LLM is still thinking: {text}")
+                # LLM is still thinking
+                return None
+            else:
+                # LLM is done thinking
+                logger.debug(f"LLM is done thinking: {text}")
+                idx = text.index(self._think_tokens[1])
+                # only return the text after the end of thinking tokens
+                text = text[idx + len(self._think_tokens[1]) :]
+                self._have_seen_think_tokens = False
+                logger.debug(f"Returning text after thinking: {text}")
+                return text
+        elif self._think_tokens[0] in text:
+            # LLM now starts thinking
+            logger.debug(f"LLM starts thinking: {text}")
+            self._have_seen_think_tokens = True
+            # return text before the start of thinking tokens
+            idx = text.index(self._think_tokens[0])
+            text = text[:idx]
+            logger.debug(f"Returning text before thinking: {text}")
+            return text
+        else:
+            # LLM is not thinking
+            return text
+
+    def _drop_special_tokens(self, text: str) -> Optional[str]:
+        """
+        Drop the special tokens from the text.
+        """
+        if self._ignore_strings is None:
+            return text
+        for ignore_string in self._ignore_strings:
+            if ignore_string in text:
+                logger.debug(f"Dropping string `{ignore_string}` from text: `{text}`")
+                text = text.replace(ignore_string, "")
+        return text
+
+    async def run_tts(self, text: str) -> AsyncGenerator[Frame, None]:
+        """Generate speech from text using the Nemo TTS model."""
+
+        if self._think_tokens is not None:
+            text = self._handle_think_tokens(text)
+
+        if not text:
+            yield None
+            return
+
+        if self._ignore_strings is not None:
+            text = self._drop_special_tokens(text)
+
+        logger.debug(f"{self}: Generating TTS [{text}]")
+
+        try:
+            await self.start_ttfb_metrics()
+            yield TTSStartedFrame()
+
+            # Increment turn index at the start of agent speaking (only if speaker changed)
+            if self._audio_logger is not None:
+                self._audio_logger.increment_turn_index(speaker="agent")
+
+            # Generate unique request ID
+
+            request_id = str(uuid.uuid4())
+
+            # Create response queue for this specific request
+            request_queue = asyncio.Queue()
+            self._pending_requests[request_id] = request_queue
+
+            try:
+                # Queue the TTS request for background processing
+                await self._tts_queue.put((text, request_id))
+
+                # Wait for the result directly from our request queue
+                result = await request_queue.get()
+                status, data = result
+
+                if status == 'error':
+                    logger.error(f"{self} TTS generation error: {data}")
+                    yield ErrorFrame(error=f"TTS generation error: {str(data)}")
+                    return
+
+                audio_result = data
+                if audio_result is None:
+                    logger.error(f"{self} TTS model returned None for text: [{text}]")
+                    yield ErrorFrame(error="TTS generation failed - no audio returned")
+                    return
+
+                await self.start_tts_usage_metrics(text)
+
+                # Collect all audio for logging
+                all_audio_bytes = b""
+                # Capture the start time when TTS begins (not when it ends)
+                if self._audio_logger is not None and self._audio_logger.first_audio_timestamp is None:
+                    self._audio_logger.first_audio_timestamp = datetime.now()
+
+                # Process the audio result (same as before)
+                if (
+                    inspect.isgenerator(audio_result)
+                    or hasattr(audio_result, '__iter__')
+                    and hasattr(audio_result, '__next__')
+                ):
+                    # Handle generator case
+                    first_chunk = True
+                    for audio_chunk in audio_result:
+                        if first_chunk:
+                            await self.stop_ttfb_metrics()
+                            first_chunk = False
+                            # Capture start time on first chunk
+                            if self._audio_logger is not None:
+                                tts_start_time = self._audio_logger.get_time_from_start_of_session()
+
+                        if audio_chunk is None:
+                            break
+
+                        audio_bytes = self._convert_to_bytes(audio_chunk)
+                        all_audio_bytes += audio_bytes
+                        chunk_size = self.chunk_size
+                        for i in range(0, len(audio_bytes), chunk_size):
+                            audio_chunk_bytes = audio_bytes[i : i + chunk_size]
+                            if not audio_chunk_bytes:
+                                break
+
+                            frame = TTSAudioRawFrame(
+                                audio=audio_chunk_bytes, sample_rate=self.sample_rate, num_channels=1
+                            )
+                            yield frame
+                else:
+                    # Handle single result case
+                    await self.stop_ttfb_metrics()
+                    # Capture start time for single result
+                    if self._audio_logger is not None:
+                        tts_start_time = self._audio_logger.get_time_from_start_of_session()
+                    audio_bytes = self._convert_to_bytes(audio_result)
+                    all_audio_bytes = audio_bytes
+
+                    chunk_size = self.chunk_size
+                    for i in range(0, len(audio_bytes), chunk_size):
+                        chunk = audio_bytes[i : i + chunk_size]
+                        if not chunk:
+                            break
+
+                        frame = TTSAudioRawFrame(audio=chunk, sample_rate=self.sample_rate, num_channels=1)
+                        yield frame
+
+                # Log the complete audio if logger is available
+                if self._audio_logger is not None and all_audio_bytes:
+                    try:
+                        self._audio_logger.log_agent_audio(
+                            audio_data=all_audio_bytes,
+                            text=text,
+                            sample_rate=self.sample_rate,
+                            num_channels=1,
+                            additional_metadata={
+                                "model": self._model_name,
+                            },
+                            tts_generation_time=tts_start_time,
+                        )
+                    except Exception as e:
+                        logger.warning(f"Failed to log agent audio: {e}")
+
+                yield TTSStoppedFrame()
+
+            finally:
+                # Clean up the pending request
+                if request_id in self._pending_requests:
+                    del self._pending_requests[request_id]
+
+        except Exception as e:
+            logger.exception(f"{self} error generating TTS: {e}")
+            error_message = f"TTS generation error: {str(e)}"
+            yield ErrorFrame(error=error_message)
+
+    def _convert_to_bytes(self, audio_data) -> bytes:
+        """Convert various audio data formats to bytes."""
+        if isinstance(audio_data, (bytes, bytearray)):
+            return bytes(audio_data)
+
+        if isinstance(audio_data, np.ndarray):
+            # Ensure it's in the right format (16-bit PCM)
+            if audio_data.dtype in [np.float32, np.float64]:
+                # Convert float [-1, 1] to int16 [-32768, 32767]
+                audio_data = np.clip(audio_data, -1.0, 1.0)  # Ensure values are in range
+                audio_data = (audio_data * 32767).astype(np.int16)
+            elif audio_data.dtype != np.int16:
+                # Convert other integer types to int16
+                audio_data = audio_data.astype(np.int16)
+            return audio_data.tobytes()
+        elif hasattr(audio_data, 'tobytes'):
+            return audio_data.tobytes()
+        else:
+            return bytes(audio_data)
+
+
+class NeMoFastPitchHiFiGANTTSService(BaseNemoTTSService):
+    """Text-to-Speech service using NeMo FastPitch-Hifigan model.
+
+    More info: https://huggingface.co/nvidia/tts_en_fastpitch
+
+    Args:
+        fastpitch_model: FastPitch model name
+        hifigan_model: Hifigan model name
+        device: Device to run on (default: 'cuda')
+        **kwargs: Additional arguments passed to BaseNemoTTSService
+    """
+
+    def __init__(
+        self,
+        fastpitch_model: str = "nvidia/tts_en_fastpitch",
+        hifigan_model: str = "nvidia/tts_hifigan",
+        device: str = "cuda",
+        **kwargs,
+    ):
+        model_name = f"{fastpitch_model}+{hifigan_model}"
+        self._fastpitch_model_name = fastpitch_model
+        self._hifigan_model_name = hifigan_model
+        super().__init__(model=model_name, device=device, **kwargs)
+        self.setup_tool_calling()
+
+    def _setup_model(self):
+        logger.info(
+            f"Loading FastPitch model={self._fastpitch_model_name} and HiFiGAN model={self._hifigan_model_name}"
+        )
+        self._fastpitch_model = self._setup_fastpitch_model(self._fastpitch_model_name)
+        self._hifigan_model = self._setup_hifigan_model(self._hifigan_model_name)
+        return self._fastpitch_model, self._hifigan_model
+
+    def _setup_fastpitch_model(self, model_name: str):
+        if model_name.endswith(".nemo"):
+            fastpitch_model = FastPitchModel.restore_from(model_name, map_location=torch.device(self._device))
+        else:
+            fastpitch_model = FastPitchModel.from_pretrained(model_name, map_location=torch.device(self._device))
+        fastpitch_model.eval()
+        return fastpitch_model
+
+    def _setup_hifigan_model(self, model_name: str):
+        if model_name.endswith(".nemo"):
+            hifigan_model = HifiGanModel.restore_from(model_name, map_location=torch.device(self._device))
+        else:
+            hifigan_model = HifiGanModel.from_pretrained(model_name, map_location=torch.device(self._device))
+        hifigan_model.eval()
+        return hifigan_model
+
+    def _generate_audio(self, text: str) -> Iterator[np.ndarray]:
+        with torch.no_grad():
+            parsed = self._fastpitch_model.parse(text)
+            spectrogram = self._fastpitch_model.generate_spectrogram(tokens=parsed)
+            audio = self._hifigan_model.convert_spectrogram_to_audio(spec=spectrogram)
+            audio = audio.detach().view(-1).cpu().numpy()
+            yield audio
+
+
+class KokoroTTSService(BaseNemoTTSService):
+    """Text-to-Speech service using Kokoro-82M model.
+
+    Kokoro is an open-weight TTS model with 82 million parameters.
+    More info: https://huggingface.co/hexgrad/Kokoro-82M
+
+    Args:
+        lang_code: Language code for the model (default: 'a' for American English)
+        voice: Voice to use (default: 'af_heart')
+        device: Device to run on (default: 'cuda')
+        sample_rate: Audio sample rate in Hz (default: 24000 for Kokoro)
+        download_all: Download all models for different languages (default: True)
+        cache_models: Cache models on GPU for faster switching between languages (default: True)
+        **kwargs: Additional arguments passed to BaseNemoTTSService
+    """
+
+    def __init__(
+        self,
+        model: str = "hexgrad/Kokoro-82M",
+        lang_code: str = "a",
+        voice: str = "af_heart",
+        device: str = "cuda",
+        sample_rate: int = 24000,
+        speed: float = 1.0,
+        download_all: bool = True,
+        cache_models: bool = True,
+        **kwargs,
+    ):
+        self._lang_code = lang_code
+        self._voice = voice
+        self._speed = speed
+        assert speed > 0, "Speed must be greater than 0"
+        self._original_speed = speed
+        self._original_voice = voice
+        self._gender = 'female' if voice[1] == 'f' else 'male'
+        self._original_gender = self._gender
+        self._original_lang_code = self._lang_code
+        if download_all:
+            self._model_maps = self._download_all_models(
+                lang_code=["a", "b"], device=device, repo_id=model, cache_models=cache_models
+            )
+        else:
+            self._model_maps = {}
+        super().__init__(model=model, device=device, sample_rate=sample_rate, **kwargs)
+        self.setup_tool_calling()
+
+    def _setup_model(self, lang_code: Optional[str] = None, voice: Optional[str] = None):
+        """Initialize the Kokoro pipeline."""
+        try:
+            from kokoro import KPipeline
+        except ImportError:
+            raise ImportError(
+                "kokoro package is required for KokoroTTSService. Install it with: `pip install kokoro>=0.9.2`"
+            )
+        if lang_code is None:
+            lang_code = self._lang_code
+        if voice is None:
+            voice = self._voice
+        logger.info(f"Loading Kokoro TTS model with model={self._model_name}, lang_code={lang_code}, voice={voice}")
+        if lang_code in self._model_maps:
+            pipeline = self._model_maps[lang_code]
+        else:
+            pipeline = KPipeline(lang_code=lang_code, device=self._device, repo_id=self._model_name)
+            self._model_maps[lang_code] = pipeline
+        return pipeline
+
+    def _download_all_models(
+        self, lang_code: List[str] = ['a', 'b'], device="cuda", repo_id="hexgrad/Kokoro-82M", cache_models=True
+    ):
+        """Download all models for Kokoro TTS service."""
+        logger.info(f"Downloading all models for Kokoro TTS service with lang_code={lang_code}")
+        from kokoro import KPipeline
+
+        model_maps = {}
+
+        for lang in lang_code:
+            pipeline = KPipeline(lang_code=lang, device=device, repo_id=repo_id)
+            if cache_models:
+                model_maps[lang] = pipeline
+        torch.cuda.empty_cache()
+        return model_maps
+
+    def _generate_audio(self, text: str) -> Iterator[np.ndarray]:
+        """Generate audio using the Kokoro pipeline.
+
+        Args:
+            text: Text to convert to speech
+
+        Yields:
+            Audio data as numpy arrays
+        """
+        try:
+            # Generate audio using Kokoro pipeline
+            generator = self._model(text, voice=self._voice, speed=self._speed)
+
+            # The generator yields tuples of (gs, ps, audio)
+            # We only need the audio component
+            for i, (gs, ps, audio) in enumerate(generator):
+                logger.debug(
+                    f"Kokoro generated audio chunk {i}: gs={gs}, ps={ps},"
+                    f"audio_shape={audio.shape if hasattr(audio, 'shape') else len(audio)}"
+                )
+                if isinstance(audio, torch.Tensor):
+                    audio = audio.detach().cpu().numpy()
+                # Kokoro returns audio as numpy array in float32 format [-1, 1]
+                # The base class will handle conversion to int16 bytes
+                yield audio
+
+        except Exception as e:
+            logger.error(f"Error generating audio with Kokoro: {e}")
+            raise
+
+    async def tool_tts_set_speed(self, params: FunctionCallParams, speed_lambda: float):
+        """
+        Set a specific speaking speed of the assistant's voice.
+        This tool should be called only when the user specifies the speed explicitly,
+        such as "speak twice as fast" or "speak half as slow" or "speak 1.5 times as fast".
+
+        Inform user of the result of this tool call. After calling this tool, continue the previous
+        response if it was unfinished and was interrupted by the user, otherwise start a new response
+        and ask if the user needs help on anything else. Avoid repeating previous responses.
+
+        Args:
+            speed_lambda: positive float, the relative change of the speaking speed to the original speed.
+                        E.g., 1.0 for original speed, 1.25 for 25% faster than original speed,
+                        0.8 for 20% slower than original speed.
+
+        """
+        if speed_lambda <= 0:
+            result = {
+                "success": False,
+                "message": f"Speed remains unchanged since the change is not a positive number: {speed_lambda}",
+            }
+            logger.debug(f"Speed remains unchanged since the change is not a positive number: {speed_lambda}")
+        else:
+            self._speed = speed_lambda * self._speed
+            result = {
+                "success": True,
+                "message": f"Speed set to {speed_lambda} of the previous speed",
+            }
+            logger.debug(f"Speed set to {speed_lambda} of the previous speed {self._original_speed}")
+        await params.result_callback(result)
+
+    async def tool_tts_reset_speed(self, params: FunctionCallParams):
+        """
+        Reset the speaking speed to the original speed.
+
+        Inform user of the result of this tool call. After calling this tool, continue the previous
+        response if it was unfinished and was interrupted by the user, otherwise start a new response
+        and ask if the user needs help on anything else. Avoid repeating previous responses.
+        """
+        self._speed = self._original_speed
+        result = {"success": True, "message": "Speaking speed is reset to the original one"}
+        logger.debug(f"Speaking speed is reset to the original speed {self._original_speed}")
+        await params.result_callback(result)
+
+    async def tool_tts_speak_faster(self, params: FunctionCallParams):
+        """
+        Speak faster by increasing the speaking speed 15% faster each time this function is called.
+
+        Inform user of the result of this tool call. After calling this tool, continue the previous
+        response if it was unfinished and was interrupted by the user, otherwise start a new response
+        and ask if the user needs help on anything else. Avoid repeating previous responses.
+        """
+        speed_lambda = 1.15
+        self._speed = speed_lambda * self._speed
+        result = {
+            "success": True,
+            "message": f"Speaking speed is increased to {speed_lambda} of the previous speed",
+        }
+        logger.debug(f"Speed is set to {speed_lambda} of the previous speed, new speed is {self._speed}")
+        await params.result_callback(result)
+
+    async def tool_tts_speak_slower(self, params: FunctionCallParams):
+        """
+        Speak slower by decreasing the speaking speed 15% slower each time this function is called.
+
+        Inform user of the result of this tool call. After calling this tool, continue the previous
+        response if it was unfinished and was interrupted by the user, otherwise start a new response
+        and ask if the user needs help on anything else. Avoid repeating previous responses.
+        """
+        speed_lambda = 0.85
+        self._speed = speed_lambda * self._speed
+        result = {
+            "success": True,
+            "message": f"Speaking speed is decreased to {speed_lambda} of the previous speed",
+        }
+        logger.debug(f"Speed is set to {speed_lambda} of the previous speed, new speed is {self._speed}")
+        await params.result_callback(result)
+
+    async def tool_tts_set_voice(self, params: FunctionCallParams, accent: str, gender: str):
+        """
+        Set the accent and gender of the assistant's voice.
+        This tool should be called only when the user specifies the accent and/or gender explicitly.
+
+        Inform user of the result of this tool call. After calling this tool, continue the previous
+        response if it was unfinished and was interrupted by the user, otherwise start a new response
+        and ask if the user needs help on anything else. Avoid repeating previous responses.
+
+        Args:
+            accent: Accent for the TTS model. Must be one of 'American English', 'British English'
+                    or 'current' for keeping the current accent.
+            gender: gender of the assistant's voice. Must be one of 'male', 'female',
+                    or 'current' for keeping the current gender.
+        """
+        await params.llm.push_frame(LLMTextFrame("Just a moment."))
+
+        lang_code = "a" if accent == "American English" else "b" if accent == "British English" else "current"
+        new_lang_code = self._lang_code
+        new_gender = self._gender
+        if lang_code != 'current':
+            new_lang_code = lang_code
+        if gender != 'current':
+            new_gender = gender
+
+        if new_lang_code == 'a':
+            new_voice = 'af_heart' if new_gender == 'female' else 'am_michael'
+        elif new_lang_code == 'b':
+            new_voice = 'bf_emma' if new_gender == 'female' else 'bm_george'
+        else:
+            await params.result_callback(
+                {
+                    "success": False,
+                    "message": f"Invalid language code: {new_lang_code} or gender: {new_gender}",
+                }
+            )
+            return
+
+        new_model = await asyncio.to_thread(self._setup_model, new_lang_code, new_voice)
+        self._model = new_model
+        self._lang_code = new_lang_code
+        self._gender = new_gender
+        self._voice = new_voice
+        logger.debug(f"Language and voice are set to {new_lang_code} and {new_voice}")
+        await params.result_callback({"success": True, "message": "Voice has been updated."})
+
+    async def tool_tts_reset_voice(self, params: FunctionCallParams):
+        """
+        Reset the accent and voice to the original ones.
+
+        Inform user of the result of this tool call. After calling this tool, continue the previous
+        response if it was unfinished and was interrupted by the user, otherwise start a new response
+        and ask if the user needs help on anything else. Avoid repeating previous responses.
+
+        """
+        await params.llm.push_frame(LLMTextFrame("Of course."))
+
+        new_model = await asyncio.to_thread(self._setup_model, self._original_lang_code, self._original_voice)
+        self._model = new_model
+        self._lang_code = self._original_lang_code
+        self._gender = self._original_gender
+        self._voice = self._original_voice
+        logger.debug(
+            f"Language and voice are reset to the original ones {self._original_lang_code} and {self._original_voice}"
+        )
+        await params.result_callback({"success": True, "message": "Voice has been reset to the original one."})
+
+    def setup_tool_calling(self):
+        """
+        Setup the tool calling mixin by registering all available tools.
+        """
+        self.register_direct_function("tool_tts_reset_speed", self.tool_tts_reset_speed)
+        self.register_direct_function("tool_tts_speak_faster", self.tool_tts_speak_faster)
+        self.register_direct_function("tool_tts_speak_slower", self.tool_tts_speak_slower)
+        self.register_direct_function("tool_tts_set_speed", self.tool_tts_set_speed)
+        self.register_direct_function("tool_tts_set_voice", self.tool_tts_set_voice)
+        self.register_direct_function("tool_tts_reset_voice", self.tool_tts_reset_voice)
+
+    def reset(self):
+        """
+        Reset the voice and speed to the original ones.
+        """
+        self._text_aggregator.reset()
+        self._speed = self._original_speed
+        self._model = self._setup_model(self._original_lang_code, self._original_voice)
+        self._lang_code = self._original_lang_code
+        self._gender = self._original_gender
+        self._voice = self._original_voice
+
+
+class MagpieTTSService(BaseNemoTTSService):
+    """Text-to-Speech service using Magpie TTS model.
+
+    Magpie is a multilingual TTS model with 357 million parameters.
+    More info: https://huggingface.co/nvidia/magpie_tts_multilingual_357m
+
+    Args:
+        model: Model name or path to the Magpie TTS model.
+        language: Language code for the model (default: 'en' for English)
+        speaker: Speaker to use for the model (default: 'Sofia')
+        apply_TN: Whether to apply text normalization (default: False)
+        device: Device to run on (default: 'cuda')
+        **kwargs: Additional arguments passed to BaseNemoTTSService
+    """
+
+    SPEAKER_MAP = {"John": 0, "Sofia": 1, "Aria": 2, "Jason": 3, "Leo": 4}
+
+    def __init__(
+        self,
+        model: str = "nvidia/magpie_tts_multilingual_357m",
+        language: str = "en",
+        speaker: str = "Sofia",
+        apply_TN: bool = False,
+        device: str = "cuda",
+        **kwargs,
+    ):
+        if speaker not in self.SPEAKER_MAP:
+            raise ValueError(f"Invalid speaker: {speaker}, must be one of {list(self.SPEAKER_MAP.keys())}")
+        self._language = language
+        self._current_speaker = speaker
+        self._apply_TN = apply_TN
+        super().__init__(model=model, device=device, **kwargs)
+        self.setup_tool_calling()
+
+    def _setup_model(self):
+        from nemo.collections.tts.models import MagpieTTSModel
+
+        if self._model_name.endswith(".nemo"):
+            model = MagpieTTSModel.restore_from(self._model_name, map_location=torch.device(self._device))
+        else:
+            model = MagpieTTSModel.from_pretrained(self._model_name, map_location=torch.device(self._device))
+        model.eval()
+
+        text = "Warming up the Magpie TTS model, this will help the model to respond faster for later requests."
+        with torch.no_grad():
+            _, _ = model.do_tts(
+                text,
+                language=self._language,
+                apply_TN=self._apply_TN,
+                speaker_index=self.SPEAKER_MAP[self._current_speaker],
+            )
+        torch.cuda.empty_cache()
+        return model
+
+    def _generate_audio(self, text: str) -> Iterator[np.ndarray]:
+        audio, audio_len = self._model.do_tts(
+            text,
+            language=self._language,
+            apply_TN=self._apply_TN,
+            speaker_index=self.SPEAKER_MAP[self._current_speaker],
+        )
+        audio_len = audio_len.view(-1).item()
+        audio = audio.detach().view(-1).cpu().numpy()
+        yield audio[:audio_len]
+
+    def setup_tool_calling(self):
+        """No tools for now for Magpie TTS service."""
+        pass
+
+
+def get_tts_service_from_config(config: DictConfig, audio_logger: Optional[AudioLogger] = None) -> BaseNemoTTSService:
+    """Get the TTS service from the configuration.
+
+    Args:
+        config: The DictConfig object containing the TTS configuration.
+        audio_logger: The audio logger to use for audio logging.
+    Returns:
+        The TTS service.
+    """
+    if isinstance(config, DictConfig):
+        config = OmegaConf.to_container(config, resolve=True)
+    model = config.get("model", None)
+    device = config.get("device", "cuda")
+    if config.get("type", None) != "nemo":
+        raise ValueError(f"Invalid TTS type: {config.get('type', None)}, only 'nemo' is supported")
+    if model is None:
+        raise ValueError("Model is required for Nemo TTS service")
+
+    text_aggregator = SimpleSegmentedTextAggregator(
+        punctuation_marks=config.get("extra_separator", None),
+        ignore_marks=config.get("ignore_strings", None),
+        min_sentence_length=config.get("min_sentence_length", 5),
+        use_legacy_eos_detection=config.get("use_legacy_eos_detection", False),
+    )
+
+    if model == "fastpitch-hifigan":
+        return NeMoFastPitchHiFiGANTTSService(
+            fastpitch_model=config.get("main_model_id", None),
+            hifigan_model=config.get("sub_model_id", None),
+            device=device,
+            text_aggregator=text_aggregator,
+            think_tokens=config.get("think_tokens", None),
+            audio_logger=audio_logger,
+            ignore_strings=config.get("ignore_strings", None),
+        )
+    elif model == "magpie":
+        return MagpieTTSService(
+            model=config.get("main_model_id", None),
+            language=config.get("language", "en"),
+            speaker=config.get("speaker", "Sofia"),
+            apply_TN=config.get("apply_TN", False),
+            device=device,
+            text_aggregator=text_aggregator,
+            think_tokens=config.get("think_tokens", None),
+            audio_logger=audio_logger,
+            ignore_strings=config.get("ignore_strings", None),
+        )
+    elif model == "kokoro":
+        return KokoroTTSService(
+            model=config.get("main_model_id", "hexgrad/Kokoro-82M"),
+            voice=config.get("sub_model_id", "af_heart"),
+            device=device,
+            speed=config.get("speed", 1.0),
+            text_aggregator=text_aggregator,
+            think_tokens=config.get("think_tokens", None),
+            sample_rate=24000,
+            audio_logger=audio_logger,
+            ignore_strings=config.get("ignore_strings", None),
+        )
+    else:
+        raise ValueError(f"Invalid model: {model}, only 'fastpitch-hifigan', 'magpie' and 'kokoro' are supported")

nemo/agents/voice_agent/pipecat/services/nemo/turn_taking.py

@@ -0,0 +1,441 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import time
+from datetime import datetime
+from pathlib import Path
+from typing import List, Optional, Union
+
+import yaml
+from loguru import logger
+from pipecat.frames.frames import (
+    BotStartedSpeakingFrame,
+    BotStoppedSpeakingFrame,
+    Frame,
+    InterimTranscriptionFrame,
+    StartInterruptionFrame,
+    TranscriptionFrame,
+    UserStartedSpeakingFrame,
+    UserStoppedSpeakingFrame,
+    VADUserStartedSpeakingFrame,
+    VADUserStoppedSpeakingFrame,
+)
+from pipecat.processors.frame_processor import FrameDirection, FrameProcessor
+from pipecat.transcriptions.language import Language
+from pipecat.utils.time import time_now_iso8601
+
+from nemo.agents.voice_agent.pipecat.frames.frames import DiarResultFrame
+from nemo.agents.voice_agent.pipecat.services.nemo.audio_logger import AudioLogger
+
+
+class NeMoTurnTakingService(FrameProcessor):
+    """Service for handling turn-taking in voice conversations with backchannel detection."""
+
+    def __init__(
+        self,
+        backchannel_phrases: Union[str, List[str]] = None,
+        eou_string: str = "<EOU>",
+        eob_string: str = "<EOB>",
+        language: Language = Language.EN_US,
+        use_vad: bool = True,
+        use_diar: bool = False,
+        max_buffer_size: int = 2,
+        bot_stop_delay: float = 0.5,
+        audio_logger: Optional[AudioLogger] = None,
+        can_create_user_frames: bool = True,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.eou_string = eou_string
+        self.eob_string = eob_string
+        self.language = language
+        self.use_vad = use_vad
+        self.use_diar = use_diar
+        self.max_buffer_size = max_buffer_size
+
+        self.backchannel_phrases = self._load_backchannel_phrases(backchannel_phrases)
+        self.backchannel_phrases_nopc = set([self.clean_text(phrase) for phrase in self.backchannel_phrases])
+        self.bot_stop_delay = bot_stop_delay
+        self.can_create_user_frames = can_create_user_frames
+        # internal data
+        self._current_speaker_id = None
+        self._prev_speaker_id = None
+        self._bot_stop_time = None
+        self._bot_speaking = False
+        self._vad_user_speaking = False
+        self._have_sent_user_started_speaking = False
+        self._user_speaking_buffer = ""
+        self._audio_logger = audio_logger
+        if not self.use_vad:
+            # if vad is not used, we assume the user is always speaking
+            self._vad_user_speaking = True
+
+    def _load_backchannel_phrases(self, backchannel_phrases: Optional[Union[str, List[str]]] = None):
+        if not backchannel_phrases:
+            return []
+
+        if isinstance(backchannel_phrases, str) and Path(backchannel_phrases).is_file():
+            logger.info(f"Loading backchannel phrases from file: {backchannel_phrases}")
+            if not Path(backchannel_phrases).exists():
+                raise FileNotFoundError(f"Backchannel phrases file not found: {backchannel_phrases}")
+            with open(backchannel_phrases, "r") as f:
+                backchannel_phrases = yaml.safe_load(f)
+            if not isinstance(backchannel_phrases, list):
+                raise ValueError(f"Backchannel phrases must be a list, got {type(backchannel_phrases)}")
+            logger.info(f"Loaded {len(backchannel_phrases)} backchannel phrases from file: {backchannel_phrases}")
+        elif isinstance(backchannel_phrases, list):
+            logger.info(f"Using backchannel phrases from list: {backchannel_phrases}")
+        else:
+            raise ValueError(f"Invalid backchannel phrases: {backchannel_phrases}")
+        return backchannel_phrases
+
+    def reset(self):
+        """
+        Reset the turn-taking service.
+        """
+        self._current_speaker_id = None
+        self._prev_speaker_id = None
+        self._bot_stop_time = None
+        self._bot_speaking = False
+        self._vad_user_speaking = False
+        self._have_sent_user_started_speaking = False
+        self._user_speaking_buffer = ""
+        if not self.use_vad:
+            # if vad is not used, we assume the user is always speaking
+            self._vad_user_speaking = True
+        logger.debug("TurnTaking service reset complete")
+
+    def clean_text(self, text: str) -> str:
+        """
+        Clean the text so that it can be used for backchannel detection.
+        """
+        if self.language != Language.EN_US:
+            raise ValueError(f"Language {self.language} not supported, currently only English is supported.")
+        for eou_string in [self.eou_string, self.eob_string]:
+            if text.endswith(eou_string):
+                text = text[: -len(eou_string)].strip()
+        text = text.lower()
+        valid_chars = "abcdefghijklmnopqrstuvwxyz'"
+        text = ''.join([c for c in text if c in valid_chars or c.isspace() or c == "'"])
+        return " ".join(text.split()).strip()
+
+    def is_backchannel(self, text: str) -> bool:
+        """
+        Check if the text is a backchannel phrase.
+        """
+        if not self.backchannel_phrases:
+            return False
+        if text.startswith("<speaker_"):
+            # if the text starts with a speaker tag, we remove it
+            text = text[len("<speaker_0>") :]
+        text = self.clean_text(text)
+        return text in self.backchannel_phrases_nopc
+
+    async def process_frame(self, frame: Frame, direction: FrameDirection):
+        """Process incoming frames and handle turn-taking logic."""
+        await super().process_frame(frame, direction)
+
+        if self._bot_stop_time is not None:
+            # check if the bot has stopped speaking for more than the delay
+            if time.time() - self._bot_stop_time > self.bot_stop_delay:
+                # set the _bot_speaking flag to False to actually consider the bot as stopped speaking
+                logger.debug(
+                    f"Bot stopped speaking for more than {self.bot_stop_delay} seconds, setting _bot_speaking to False"
+                )
+                self._bot_stop_time = None
+                self._bot_speaking = False
+
+        if isinstance(frame, (TranscriptionFrame, InterimTranscriptionFrame)):
+            await self._handle_transcription(frame, direction)
+        elif isinstance(frame, VADUserStartedSpeakingFrame):
+            await self._handle_vad_user_started_speaking(frame, direction)
+        elif isinstance(frame, VADUserStoppedSpeakingFrame):
+            await self._handle_vad_user_stopped_speaking(frame, direction)
+        elif isinstance(frame, BotStartedSpeakingFrame):
+            logger.debug("BotStartedSpeakingFrame received")
+            self._bot_speaking = True
+            # Capture the actual start time when audio starts playing
+            # This is more accurate than capturing during TTS generation
+            if self._audio_logger:
+                self._audio_logger.set_agent_turn_start_time()
+        elif isinstance(frame, BotStoppedSpeakingFrame):
+            logger.debug("BotStoppedSpeakingFrame received")
+            self._bot_stop_time = time.time()
+            if self.bot_stop_delay is None or self.bot_stop_delay <= 0:
+                # only set the flag if the delay is not set or is 0
+                self._bot_speaking = False
+                logger.debug("Setting _bot_speaking to False")
+        elif isinstance(frame, DiarResultFrame):
+            logger.debug("DiarResultFrame received")
+            await self._handle_diar_result(frame, direction)
+        else:
+            await self.push_frame(frame, direction)
+
+    async def _handle_backchannel_text(self, text: str):
+        # ignore the backchannel string while bot is speaking
+        # push the backchannel string upstream, not downstream
+        await self.push_frame(
+            TranscriptionFrame(
+                text=f"({text})",
+                user_id="",
+                timestamp=time_now_iso8601(),
+                language=self.language if self.language else Language.EN_US,
+                result={"text": f"Backchannel detected: {text}"},
+            ),
+            direction=FrameDirection.UPSTREAM,
+        )
+
+    async def _handle_transcription(
+        self, frame: TranscriptionFrame | InterimTranscriptionFrame, direction: FrameDirection
+    ):
+        text_segment = frame.text
+        if self._vad_user_speaking:
+            self._user_speaking_buffer += text_segment
+            has_eou = self._user_speaking_buffer.endswith(self.eou_string)
+            has_eob = self._user_speaking_buffer.endswith(self.eob_string)
+            if has_eou:
+                # EOU detected, user is done speaking - push completed text and interrupt bot
+                logger.debug(f"<EOU> Detected: `{self._user_speaking_buffer}`")
+                completed_text = self._user_speaking_buffer[: -len(self.eou_string)].strip()
+                if self._bot_speaking and self.is_backchannel(completed_text):
+                    logger.debug(f"<EOU> detected for a backchannel phrase while bot is speaking: `{completed_text}`")
+                    await self._handle_backchannel_text(completed_text)
+                    if self._audio_logger:
+                        if self._audio_logger.staged_metadata is None:
+                            self._audio_logger.staged_metadata = {"is_backchannel": True, "start_time": datetime.now()}
+                        else:
+                            self._audio_logger.staged_metadata["is_backchannel"] = True
+
+                else:
+                    await self._handle_completed_text(completed_text, direction)
+                    await self._handle_user_interruption(UserStoppedSpeakingFrame())
+                self._user_speaking_buffer = ""
+                self._have_sent_user_started_speaking = False  # user is done speaking, so we reset the flag
+            elif has_eob and self._bot_speaking:
+                logger.debug(f"<EOB> detected while bot is speaking: `{self._user_speaking_buffer}`")
+                await self._handle_backchannel_text(str(self._user_speaking_buffer))
+                if self._audio_logger:
+                    if self._audio_logger.staged_metadata is None:
+                        self._audio_logger.staged_metadata = {"is_backchannel": True, "start_time": datetime.now()}
+                    else:
+                        self._audio_logger.staged_metadata["is_backchannel"] = True
+                self._user_speaking_buffer = ""
+                self._have_sent_user_started_speaking = False  # user is done speaking, so we reset the flag
+            else:
+                # if bot is not speaking, the backchannel string is not considered a backchannel phrase
+                # user is still speaking, so we append the text segment to the buffer
+                logger.debug(f"User is speaking: `{self._user_speaking_buffer}`")
+                if has_eob:
+                    logger.debug(
+                        f"{self.eob_string} detected but ignored (bot NOT speaking): "
+                        f"`{self._user_speaking_buffer}`"
+                    )
+                    self._user_speaking_buffer = self._user_speaking_buffer[: -len(self.eob_string)].strip()
+                # assume the last word is not completed
+                completed_words = self._user_speaking_buffer.strip().split()[:-1]
+                if len(completed_words) >= self.max_buffer_size:
+                    completed_text = " ".join(completed_words)
+                    await self._handle_completed_text(completed_text, direction, is_final=False)
+
+        else:
+            # if vad is not detecting user speaking
+            logger.debug(
+                f"VAD is not detecting user speaking, but still received text segment from STT: `{text_segment}`"
+            )
+            is_backchannel = self.is_backchannel(text_segment)
+            if text_segment.endswith(self.eob_string):
+                is_backchannel = True
+                logger.debug(f"Dropping EOB token: `{text_segment}`")
+                text_segment = text_segment[: -len(self.eob_string)].strip()
+            elif text_segment.endswith(self.eou_string):
+                logger.debug(f"Dropping EOU token: `{text_segment}`")
+                text_segment = text_segment[: -len(self.eou_string)].strip()
+
+            if not text_segment.strip():
+                return
+            if is_backchannel and self._bot_speaking:
+                logger.debug(f"Backchannel detected while bot is speaking: `{text_segment}`")
+                # push the backchannel string upstream, not downstream
+                curr_text = str(self._user_speaking_buffer + text_segment)
+                self._user_speaking_buffer = ""
+                if self._audio_logger:
+                    if self._audio_logger.staged_metadata is None:
+                        self._audio_logger.staged_metadata = {"is_backchannel": True, "start_time": datetime.now()}
+                    else:
+                        self._audio_logger.staged_metadata["is_backchannel"] = True
+                await self.push_frame(
+                    TranscriptionFrame(
+                        text=f"({curr_text})",
+                        user_id="",
+                        timestamp=time_now_iso8601(),
+                        language=self.language if self.language else Language.EN_US,
+                        result={"text": f"Backchannel detected: {self._user_speaking_buffer+text_segment}"},
+                    ),
+                    direction=FrameDirection.UPSTREAM,
+                )
+
+            else:
+                # if the text segment is not empty and have non-space characters, we append it to the buffer
+                self._user_speaking_buffer += text_segment
+                if self.is_backchannel(self._user_speaking_buffer):
+                    logger.debug(f"Backchannel detected: `{self._user_speaking_buffer}`")
+                    self._user_speaking_buffer = ""
+                    self._have_sent_user_started_speaking = False
+                    return
+                logger.debug(f"Appending text segment to user speaking buffer: `{self._user_speaking_buffer}`")
+
+    async def _handle_completed_text(self, completed_text: str, direction: FrameDirection, is_final: bool = True):
+        if not self._have_sent_user_started_speaking:
+            # if we haven't sent the user started speaking frame, we send it now
+            # so that the bot can be interrupted and be ready to respond to the new user turn
+            await self._handle_user_interruption(UserStartedSpeakingFrame())
+            self._have_sent_user_started_speaking = True
+
+        completed_text = completed_text.strip()
+        completed_text = completed_text.replace(self.eou_string, "").replace(self.eob_string, "")
+
+        if self.use_diar and not completed_text.startswith("<speaker_") and self._prev_speaker_id is not None:
+            # Add the previous speaker tag to the beginning of the text
+            completed_text = f"<speaker_{self._prev_speaker_id}> {completed_text}"
+
+        frame_type = TranscriptionFrame if is_final else InterimTranscriptionFrame
+        text_frame = frame_type(
+            text=completed_text,
+            user_id="",  # No speaker ID in this implementation
+            timestamp=time_now_iso8601(),
+            language=self.language if self.language else Language.EN_US,
+            result={"text": completed_text},
+        )
+        logger.debug(f"Pushing text frame: {text_frame}")
+        await self.push_frame(text_frame, direction)
+
+    def _contains_only_speaker_tags(self, text: str) -> bool:
+        """
+        Check if the text contains only speaker tags.
+        """
+        return text.strip().startswith("<speaker_") and text.strip().endswith(">")
+
+    async def _handle_vad_user_started_speaking(self, frame: VADUserStartedSpeakingFrame, direction: FrameDirection):
+        """
+        Handle the user started speaking frame.
+
+        If there are no backchannel phrases and we haven't sent the user started speaking frame, we send it now
+        so that the bot can be interrupted and be ready to respond to the new user turn
+        """
+        self._vad_user_speaking = True
+        logger.debug("NeMoTurnTakingService: VADUserStartedSpeakingFrame")
+        await self.push_frame(frame, direction)
+        if not self.backchannel_phrases and not self._have_sent_user_started_speaking:
+            await self._handle_user_interruption(UserStartedSpeakingFrame())
+            self._have_sent_user_started_speaking = True
+
+    async def _handle_vad_user_stopped_speaking(self, frame: VADUserStoppedSpeakingFrame, direction: FrameDirection):
+        """
+        Handle the user stopped speaking frame.
+
+        If the buffer is not empty:
+            - If bot is not speaking: push completed text regardless of backchannel
+            - If bot is speaking: ignore backchannel strings
+        If the buffer is empty, do nothing.
+        """
+        if self.use_vad:
+            self._vad_user_speaking = False
+        logger.debug("NeMoTurnTakingService: VADUserStoppedSpeakingFrame")
+        await self.push_frame(frame, direction)
+
+        # if user buffer only contains speaker tags, we don't push the completed text frame
+        if self._contains_only_speaker_tags(self._user_speaking_buffer):
+            logger.debug(f"User buffer only contains speaker tags: `{self._user_speaking_buffer}`, ignoring")
+            return
+
+        is_backchannel = self.is_backchannel(self._user_speaking_buffer)
+        if not self._user_speaking_buffer:
+            return
+        if not self._bot_speaking or not is_backchannel:
+            logger.debug(f"Bot talking: {self._bot_speaking}, backchannel: {is_backchannel}")
+            logger.debug(f"Pushing completed text frame for VAD user stopped speaking: {self._user_speaking_buffer}")
+            await self._handle_completed_text(self._user_speaking_buffer, direction)
+            self._user_speaking_buffer = ""
+            if self._have_sent_user_started_speaking:
+                await self._handle_user_interruption(UserStoppedSpeakingFrame())
+                self._have_sent_user_started_speaking = False
+        elif is_backchannel:
+            logger.debug(f"Backchannel detected: `{self._user_speaking_buffer}`")
+            if self._audio_logger:
+                self._audio_logger.save_user_audio(is_backchannel=True)
+                logger.debug(
+                    f"[TurnTakingService] Saved backchannel audio (VAD stopped): {self._user_speaking_buffer}"
+                )
+
+            await self.push_frame(
+                TranscriptionFrame(
+                    text=f"({self._user_speaking_buffer})",
+                    user_id="",
+                    timestamp=time_now_iso8601(),
+                    language=self.language if self.language else Language.EN_US,
+                    result={"text": f"Backchannel detected: {self._user_speaking_buffer}"},
+                ),
+                direction=FrameDirection.UPSTREAM,
+            )
+            self._user_speaking_buffer = ""
+            self._have_sent_user_started_speaking = False
+
+    async def _handle_user_interruption(self, frame: Frame):
+        # Adapted from BaseInputTransport._handle_user_interruption
+        if isinstance(frame, UserStartedSpeakingFrame):
+            logger.debug("User started speaking")
+            if self.can_create_user_frames:
+                logger.debug("Pushing UserStartedSpeakingFrame and StartInterruptionFrame")
+                await self.push_frame(frame)
+                await self.push_frame(StartInterruptionFrame(), direction=FrameDirection.DOWNSTREAM)
+            else:
+                logger.debug(
+                    "Skipping UserStartedSpeakingFrame and StartInterruptionFrame because can_create_user_frames is False"
+                )
+            # Record cutoff time for agent audio when TTS is interrupted
+            if self._audio_logger and self._bot_speaking:
+                self._audio_logger.set_agent_cutoff_time()
+            # Increment turn index when user starts speaking (only if speaker changed)
+            if self._audio_logger:
+                self._audio_logger.increment_turn_index(speaker="user")
+        elif isinstance(frame, UserStoppedSpeakingFrame):
+            logger.debug("User stopped speaking")
+            if self.can_create_user_frames:
+                logger.debug("Pushing UserStoppedSpeakingFrame")
+                await self.push_frame(frame)
+            else:
+                logger.debug("Skipping UserStoppedSpeakingFrame because can_create_user_frames is False")
+        else:
+            logger.debug(f"Unknown frame type for _handle_user_interruption: {type(frame)}")
+
+    async def _handle_diar_result(self, frame: DiarResultFrame, direction: FrameDirection):
+        if not self.use_diar:
+            logger.debug("Diarization is disabled, skipping")
+            return
+
+        new_speaker_id = frame.diar_result  # speaker id of the dominant speaker
+
+        # logger.debug(f"Dominant speaker ID: {dominant_speaker_id}")
+        self._prev_speaker_id = self._current_speaker_id
+        last_speaker_id = self._current_speaker_id
+
+        if not self._user_speaking_buffer.startswith("<speaker_"):
+            # add speaker tag <speaker_{speaker_id}> to the beginning of the current utterance
+            self._user_speaking_buffer = f"<speaker_{new_speaker_id}> {self._user_speaking_buffer}"
+        elif last_speaker_id != new_speaker_id:
+            # change the speaker tag to the dominant speaker id
+            self._user_speaking_buffer = self._user_speaking_buffer[len("<speaker_0>") :]
+            self._user_speaking_buffer = f"<speaker_{new_speaker_id}> {self._user_speaking_buffer}"
+        logger.debug(f"Speaker changed from {last_speaker_id} to {new_speaker_id}")
+        self._current_speaker_id = new_speaker_id

nemo/agents/voice_agent/pipecat/services/nemo/utils.py

@@ -0,0 +1,197 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# NOTE: This file will be deprecated in the future, as the new inference pipeline will replace it.
+
+import math
+
+import numpy as np
+import torch
+from omegaconf import DictConfig
+
+import nemo.collections.asr as nemo_asr
+
+LOG_MEL_ZERO = -16.635
+
+
+class AudioBufferer:
+    def __init__(self, sample_rate: int, buffer_size_in_secs: float):
+        self.buffer_size = int(buffer_size_in_secs * sample_rate)
+        self.sample_buffer = torch.zeros(self.buffer_size, dtype=torch.float32)
+
+    def reset(self) -> None:
+        """
+        Reset the buffer to zero
+        """
+        self.sample_buffer.zero_()
+
+    def update(self, audio: np.ndarray) -> None:
+        """
+        Update the buffer with the new frame
+        Args:
+            frame (Frame): frame to update the buffer with
+        """
+        if not isinstance(audio, torch.Tensor):
+            audio = torch.from_numpy(audio)
+
+        audio_size = audio.shape[0]
+        if audio_size > self.buffer_size:
+            raise ValueError(f"Frame size ({audio_size}) exceeds buffer size ({self.buffer_size})")
+
+        shift = audio_size
+        self.sample_buffer[:-shift] = self.sample_buffer[shift:].clone()
+        self.sample_buffer[-shift:] = audio.clone()
+
+    def get_buffer(self) -> torch.Tensor:
+        """
+        Get the current buffer
+        Returns:
+            torch.Tensor: current state of the buffer
+        """
+        return self.sample_buffer.clone()
+
+    def is_buffer_empty(self) -> bool:
+        """
+        Check if the buffer is empty
+        Returns:
+            bool: True if the buffer is empty, False otherwise
+        """
+        return self.sample_buffer.sum() == 0
+
+
+class CacheFeatureBufferer:
+    def __init__(
+        self,
+        sample_rate: int,
+        buffer_size_in_secs: float,
+        chunk_size_in_secs: float,
+        preprocessor_cfg: DictConfig,
+        device: torch.device,
+        fill_value: float = LOG_MEL_ZERO,
+    ):
+
+        if buffer_size_in_secs < chunk_size_in_secs:
+            raise ValueError(
+                f"Buffer size ({buffer_size_in_secs}s) should be no less than chunk size ({chunk_size_in_secs}s)"
+            )
+
+        self.sample_rate = sample_rate
+        self.buffer_size_in_secs = buffer_size_in_secs
+        self.chunk_size_in_secs = chunk_size_in_secs
+        self.device = device
+
+        if hasattr(preprocessor_cfg, 'log') and preprocessor_cfg.log:
+            self.ZERO_LEVEL_SPEC_DB_VAL = LOG_MEL_ZERO  # Log-Mel spectrogram value for zero signals
+        else:
+            self.ZERO_LEVEL_SPEC_DB_VAL = fill_value
+
+        self.n_feat = preprocessor_cfg.features
+        self.timestep_duration = preprocessor_cfg.window_stride
+        self.n_chunk_look_back = int(self.timestep_duration * self.sample_rate)
+        self.chunk_size = int(self.chunk_size_in_secs * self.sample_rate)
+        self.sample_buffer = AudioBufferer(sample_rate, buffer_size_in_secs)
+
+        self.feature_buffer_len = int(buffer_size_in_secs / self.timestep_duration)
+        self.feature_chunk_len = int(chunk_size_in_secs / self.timestep_duration)
+        self.feature_buffer = torch.full(
+            [self.n_feat, self.feature_buffer_len],
+            self.ZERO_LEVEL_SPEC_DB_VAL,
+            dtype=torch.float32,
+            device=self.device,
+        )
+
+        self.preprocessor = nemo_asr.models.ASRModel.from_config_dict(preprocessor_cfg)
+        self.preprocessor.to(self.device)
+
+    def is_buffer_empty(self) -> bool:
+        """
+        Check if the buffer is empty
+        Returns:
+            bool: True if the buffer is empty, False otherwise
+        """
+        return self.sample_buffer.is_buffer_empty()
+
+    def reset(self) -> None:
+        """
+        Reset the buffer to zero
+        """
+        self.sample_buffer.reset()
+        self.feature_buffer.fill_(self.ZERO_LEVEL_SPEC_DB_VAL)
+
+    def _update_feature_buffer(self, feat_chunk: torch.Tensor) -> None:
+        """
+        Add an extracted feature to `feature_buffer`
+        """
+        self.feature_buffer[:, : -self.feature_chunk_len] = self.feature_buffer[:, self.feature_chunk_len :].clone()
+        self.feature_buffer[:, -self.feature_chunk_len :] = feat_chunk.clone()
+
+    def preprocess(self, audio_signal: torch.Tensor) -> torch.Tensor:
+        """
+        Preprocess the audio signal using the preprocessor
+        Args:
+            audio_signal (torch.Tensor): audio signal
+        Returns:
+            torch.Tensor: preprocessed features
+        """
+        audio_signal = audio_signal.unsqueeze_(0).to(self.device)
+        audio_signal_len = torch.tensor([audio_signal.shape[1]], device=self.device)
+        features, _ = self.preprocessor(
+            input_signal=audio_signal,
+            length=audio_signal_len,
+        )
+        features = features.squeeze()
+        return features
+
+    def update(self, audio: np.ndarray) -> None:
+        """
+        Update the sample anf feature buffers with the new frame
+        Args:
+            frame (Frame): frame to update the buffer with
+        """
+
+        # Update the sample buffer with the new frame
+        self.sample_buffer.update(audio)
+
+        if math.isclose(self.buffer_size_in_secs, self.chunk_size_in_secs):
+            # If the buffer size is equal to the chunk size, just take the whole buffer
+            samples = self.sample_buffer.sample_buffer.clone()
+        else:
+            # Add look_back to have context for the first feature
+            samples = self.sample_buffer.sample_buffer[-(self.n_chunk_look_back + self.chunk_size) :]
+
+        # Get the mel spectrogram
+        features = self.preprocess(samples)
+
+        # If the features are longer than supposed to be, drop the last frames
+        # Drop the last diff frames because they might be incomplete
+        if (diff := features.shape[1] - self.feature_chunk_len - 1) > 0:
+            features = features[:, :-diff]
+
+        # Update the feature buffer with the new features
+        self._update_feature_buffer(features[:, -self.feature_chunk_len :])
+
+    def get_buffer(self) -> torch.Tensor:
+        """
+        Get the current sample buffer
+        Returns:
+            torch.Tensor: current state of the buffer
+        """
+        return self.sample_buffer.get_buffer()
+
+    def get_feature_buffer(self) -> torch.Tensor:
+        """
+        Get the current feature buffer
+        Returns:
+            torch.Tensor: current state of the feature buffer
+        """
+        return self.feature_buffer.clone()

nemo/agents/voice_agent/pipecat/transports/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

nemo/agents/voice_agent/pipecat/transports/base_input.py

@@ -0,0 +1,58 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from loguru import logger
+from pipecat.audio.vad.vad_analyzer import VADState
+from pipecat.frames.frames import (
+    InputAudioRawFrame,
+    UserStartedSpeakingFrame,
+    UserStoppedSpeakingFrame,
+    VADUserStartedSpeakingFrame,
+    VADUserStoppedSpeakingFrame,
+)
+from pipecat.transports.base_input import BaseInputTransport as _BaseInputTransport
+
+
+class BaseInputTransport(_BaseInputTransport):
+    async def _handle_vad(self, audio_frame: InputAudioRawFrame, vad_state: VADState):
+        """Handle Voice Activity Detection results and generate appropriate frames."""
+        new_vad_state = await self._vad_analyze(audio_frame)
+        if new_vad_state != vad_state and new_vad_state != VADState.STARTING and new_vad_state != VADState.STOPPING:
+            frame = None
+            # If the turn analyser is enabled, this will prevent:
+            # - Creating the UserStoppedSpeakingFrame
+            # - Creating the UserStartedSpeakingFrame multiple times
+            can_create_user_frames = (
+                self._params.turn_analyzer is None or not self._params.turn_analyzer.speech_triggered
+            ) and self._params.can_create_user_frames
+
+            if new_vad_state == VADState.SPEAKING:
+                await self.push_frame(VADUserStartedSpeakingFrame())
+                if can_create_user_frames:
+                    frame = UserStartedSpeakingFrame()
+                else:
+                    logger.debug("base_input: VAD state changed to SPEAKING but can_create_user_frames is False")
+            elif new_vad_state == VADState.QUIET:
+                await self.push_frame(VADUserStoppedSpeakingFrame())
+                if can_create_user_frames:
+                    frame = UserStoppedSpeakingFrame()
+                else:
+                    logger.debug("base_input: VAD state changed to QUIET but can_create_user_frames is False")
+
+            if frame:
+                await self._handle_user_interruption(frame)
+
+            vad_state = new_vad_state
+        return vad_state

nemo/agents/voice_agent/pipecat/transports/base_transport.py

@@ -0,0 +1,20 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from pipecat.transports.base_transport import TransportParams as _TransportParams
+
+
+class TransportParams(_TransportParams):
+    can_create_user_frames: bool = True

nemo/agents/voice_agent/pipecat/transports/network/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

nemo/agents/voice_agent/pipecat/transports/network/websocket_server.py

@@ -0,0 +1,304 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import asyncio
+from typing import Optional
+
+from loguru import logger
+from pipecat.frames.frames import CancelFrame, EndFrame, InputAudioRawFrame, StartFrame
+from pipecat.serializers.base_serializer import FrameSerializer
+from pipecat.transports.base_transport import BaseTransport
+from pipecat.transports.network.websocket_server import (
+    WebsocketServerCallbacks,
+    WebsocketServerOutputTransport,
+    WebsocketServerParams,
+)
+
+from nemo.agents.voice_agent.pipecat.transports.base_input import BaseInputTransport
+from nemo.agents.voice_agent.pipecat.transports.base_transport import TransportParams
+
+try:
+    import websockets
+except ModuleNotFoundError as e:
+    logger.error(f"Exception: {e}")
+    logger.error("In order to use websockets, you need to `pip install pipecat-ai[websocket]`.")
+    raise Exception(f"Missing module: {e}")
+
+
+class WebsocketServerParams(TransportParams):
+    """Configuration parameters for WebSocket server transport.
+
+    Parameters:
+        add_wav_header: Whether to add WAV headers to audio frames.
+        serializer: Frame serializer for message encoding/decoding.
+        session_timeout: Timeout in seconds for client sessions.
+    """
+
+    add_wav_header: bool = False
+    serializer: Optional[FrameSerializer] = None
+    session_timeout: Optional[int] = None
+
+
+class WebsocketServerInputTransport(BaseInputTransport):
+    """WebSocket server input transport for receiving client data.
+
+    Handles incoming WebSocket connections, message processing, and client
+    session management including timeout monitoring and connection lifecycle.
+    """
+
+    def __init__(
+        self,
+        transport: BaseTransport,
+        host: str,
+        port: int,
+        params: WebsocketServerParams,
+        callbacks: WebsocketServerCallbacks,
+        **kwargs,
+    ):
+        """Initialize the WebSocket server input transport.
+
+        Args:
+            transport: The parent transport instance.
+            host: Host address to bind the WebSocket server to.
+            port: Port number to bind the WebSocket server to.
+            params: WebSocket server configuration parameters.
+            callbacks: Callback functions for WebSocket events.
+            **kwargs: Additional arguments passed to parent class.
+        """
+        super().__init__(params, **kwargs)
+
+        self._transport = transport
+        self._host = host
+        self._port = port
+        self._params = params
+        self._callbacks = callbacks
+
+        self._websocket: Optional[websockets.WebSocketServerProtocol] = None
+
+        self._server_task = None
+
+        # This task will monitor the websocket connection periodically.
+        self._monitor_task = None
+
+        self._stop_server_event = asyncio.Event()
+
+        # Whether we have seen a StartFrame already.
+        self._initialized = False
+
+    async def start(self, frame: StartFrame):
+        """Start the WebSocket server and initialize components.
+
+        Args:
+            frame: The start frame containing initialization parameters.
+        """
+        await super().start(frame)
+
+        if self._initialized:
+            return
+
+        self._initialized = True
+
+        if self._params.serializer:
+            await self._params.serializer.setup(frame)
+        if not self._server_task:
+            self._server_task = self.create_task(self._server_task_handler())
+        await self.set_transport_ready(frame)
+
+    async def stop(self, frame: EndFrame):
+        """Stop the WebSocket server and cleanup resources.
+
+        Args:
+            frame: The end frame signaling transport shutdown.
+        """
+        await super().stop(frame)
+        self._stop_server_event.set()
+        if self._monitor_task:
+            await self.cancel_task(self._monitor_task)
+            self._monitor_task = None
+        if self._server_task:
+            await self.wait_for_task(self._server_task)
+            self._server_task = None
+
+    async def cancel(self, frame: CancelFrame):
+        """Cancel the WebSocket server and stop all processing.
+
+        Args:
+            frame: The cancel frame signaling immediate cancellation.
+        """
+        await super().cancel(frame)
+        if self._monitor_task:
+            await self.cancel_task(self._monitor_task)
+            self._monitor_task = None
+        if self._server_task:
+            await self.cancel_task(self._server_task)
+            self._server_task = None
+
+    async def cleanup(self):
+        """Cleanup resources and parent transport."""
+        await super().cleanup()
+        await self._transport.cleanup()
+
+    async def _server_task_handler(self):
+        """Handle WebSocket server startup and client connections."""
+        logger.info(f"Starting websocket server on {self._host}:{self._port}")
+        async with websockets.serve(self._client_handler, self._host, self._port) as server:
+            await self._callbacks.on_websocket_ready()
+            await self._stop_server_event.wait()
+
+    async def _client_handler(self, websocket: websockets.WebSocketServerProtocol, path: Optional[str] = None):
+        """Handle individual client connections and message processing."""
+        logger.info(f"New client connection from {websocket.remote_address}")
+        if self._websocket:
+            await self._websocket.close()
+            logger.warning("Only one client connected, using new connection")
+
+        self._websocket = websocket
+
+        # Notify
+        await self._callbacks.on_client_connected(websocket)
+
+        # Create a task to monitor the websocket connection
+        if not self._monitor_task and self._params.session_timeout:
+            self._monitor_task = self.create_task(self._monitor_websocket(websocket, self._params.session_timeout))
+
+        # Handle incoming messages
+        try:
+            async for message in websocket:
+                if not self._params.serializer:
+                    continue
+
+                frame = await self._params.serializer.deserialize(message)
+
+                if not frame:
+                    continue
+
+                if isinstance(frame, InputAudioRawFrame):
+                    await self.push_audio_frame(frame)
+                else:
+                    await self.push_frame(frame)
+        except Exception as e:
+            logger.error(f"{self} exception receiving data: {e.__class__.__name__} ({e})")
+
+        # Notify disconnection
+        await self._callbacks.on_client_disconnected(websocket)
+
+        await self._websocket.close()
+        self._websocket = None
+
+        logger.info(f"Client {websocket.remote_address} disconnected")
+
+    async def _monitor_websocket(self, websocket: websockets.WebSocketServerProtocol, session_timeout: int):
+        """Monitor WebSocket connection for session timeout."""
+        try:
+            await asyncio.sleep(session_timeout)
+            if not websocket.closed:
+                await self._callbacks.on_session_timeout(websocket)
+        except asyncio.CancelledError:
+            logger.info(f"Monitoring task cancelled for: {websocket.remote_address}")
+            raise
+
+
+class WebsocketServerTransport(BaseTransport):
+    """WebSocket server transport for bidirectional real-time communication.
+
+    Provides a complete WebSocket server implementation with separate input and
+    output transports, client connection management, and event handling for
+    real-time audio and data streaming applications.
+    """
+
+    def __init__(
+        self,
+        params: WebsocketServerParams,
+        host: str = "localhost",
+        port: int = 8765,
+        input_name: Optional[str] = None,
+        output_name: Optional[str] = None,
+    ):
+        """Initialize the WebSocket server transport.
+
+        Args:
+            params: WebSocket server configuration parameters.
+            host: Host address to bind the server to. Defaults to "localhost".
+            port: Port number to bind the server to. Defaults to 8765.
+            input_name: Optional name for the input processor.
+            output_name: Optional name for the output processor.
+        """
+        super().__init__(input_name=input_name, output_name=output_name)
+        self._host = host
+        self._port = port
+        self._params = params
+
+        self._callbacks = WebsocketServerCallbacks(
+            on_client_connected=self._on_client_connected,
+            on_client_disconnected=self._on_client_disconnected,
+            on_session_timeout=self._on_session_timeout,
+            on_websocket_ready=self._on_websocket_ready,
+        )
+        self._input: Optional[WebsocketServerInputTransport] = None
+        self._output: Optional[WebsocketServerOutputTransport] = None
+        self._websocket: Optional[websockets.WebSocketServerProtocol] = None
+
+        # Register supported handlers. The user will only be able to register
+        # these handlers.
+        self._register_event_handler("on_client_connected")
+        self._register_event_handler("on_client_disconnected")
+        self._register_event_handler("on_session_timeout")
+        self._register_event_handler("on_websocket_ready")
+
+    def input(self) -> WebsocketServerInputTransport:
+        """Get the input transport for receiving client data.
+
+        Returns:
+            The WebSocket server input transport instance.
+        """
+        if not self._input:
+            self._input = WebsocketServerInputTransport(
+                self, self._host, self._port, self._params, self._callbacks, name=self._input_name
+            )
+        return self._input
+
+    def output(self) -> WebsocketServerOutputTransport:
+        """Get the output transport for sending data to clients.
+
+        Returns:
+            The WebSocket server output transport instance.
+        """
+        if not self._output:
+            self._output = WebsocketServerOutputTransport(self, self._params, name=self._output_name)
+        return self._output
+
+    async def _on_client_connected(self, websocket):
+        """Handle client connection events."""
+        if self._output:
+            await self._output.set_client_connection(websocket)
+            await self._call_event_handler("on_client_connected", websocket)
+        else:
+            logger.error("A WebsocketServerTransport output is missing in the pipeline")
+
+    async def _on_client_disconnected(self, websocket):
+        """Handle client disconnection events."""
+        if self._output:
+            await self._output.set_client_connection(None)
+            await self._call_event_handler("on_client_disconnected", websocket)
+        else:
+            logger.error("A WebsocketServerTransport output is missing in the pipeline")
+
+    async def _on_session_timeout(self, websocket):
+        """Handle client session timeout events."""
+        await self._call_event_handler("on_session_timeout", websocket)
+
+    async def _on_websocket_ready(self):
+        """Handle WebSocket server ready events."""
+        await self._call_event_handler("on_websocket_ready")

nemo/agents/voice_agent/pipecat/utils/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

nemo/agents/voice_agent/pipecat/utils/text/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

nemo/agents/voice_agent/pipecat/utils/text/simple_text_aggregator.py

@@ -0,0 +1,238 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import re
+from typing import AsyncIterator, Optional
+
+from loguru import logger
+from pipecat.utils.string import match_endofsentence
+from pipecat.utils.text.base_text_aggregator import Aggregation, AggregationType
+from pipecat.utils.text.simple_text_aggregator import SimpleTextAggregator
+
+
+def has_partial_decimal(text: str) -> bool:
+    """Check if the text ends with a partial decimal.
+
+    Returns True if the text ends with a number that looks like it could
+    be a partial decimal (e.g., "3.", "3.14", "($3.14)"), but NOT if it's
+    clearly a complete sentence (e.g., "It costs $3.14.") or a bullet point
+    (e.g., "1. Alpha; 2.").
+    """
+
+    # Check for bullet point pattern: ends with 1-3 digits followed by period
+    # Examples: "1.", "12.", "123.", or "text; 2."
+    # Bullet points are typically small numbers (1-999) at the end
+    bullet_match = re.search(r'(?:^|[\s;,]|[^\d])(\d{1,3})\.$', text)
+    if bullet_match:
+        # It's likely a bullet point, not a partial decimal
+        return False
+
+    # Pattern to find decimal numbers near the end, allowing for trailing
+    # non-word characters like ), ], ", ', etc.
+    # Match: digit(s) + period + optional digit(s) + optional trailing non-word chars
+    match = re.search(r'\d+\.(?:\d+)?([^\w\s]*)$', text)
+
+    if not match:
+        return False
+
+    trailing = match.group(1)  # e.g., ")" or "" or "."
+
+    # If trailing contains a period, it's sentence-ending punctuation
+    # e.g., "3.14." means complete sentence
+    if '.' in trailing:
+        return False
+
+    # Otherwise, it's a partial decimal (either incomplete like "3."
+    # or complete number but sentence not finished like "($3.14)")
+    return True
+
+
+def find_last_period_index(text: str) -> int:
+    """
+    Find the last occurrence of a period in the text,
+    but return -1 if the text doesn't seem to be a complete sentence.
+    """
+    num_periods = text.count(".")
+    if num_periods == 0:
+        return -1
+
+    if num_periods == 1:
+        if has_partial_decimal(text):
+            # if the only period in the text is part of a number, return -1
+            return -1
+        # Check if the only period is a bullet point (e.g., "1. Alpha" or incomplete "1.")
+        if re.search(r'(?:^|[\s;,]|[^\d])(\d{1,3})\.(?:\s+\w|\s*$)', text):
+            # The period is after a bullet point number, either:
+            # - followed by content (e.g., "1. Alpha")
+            # - or at the end with optional whitespace (e.g., "1." or "1. ")
+            return -1
+
+        # Check if any of the abbreviations "e.", "i." "g.", "etc." are present in the text
+        if re.search(r'\b(e\.|i\.|g\.)\b', text):
+            # The period is after a character/word that is likely to be a abbreviation, return -1
+            return -1
+
+    # otherwise, check the last occurrence of a period
+    idx = text.rfind(".")
+    if idx <= 0:
+        return idx
+    if text[idx - 1].isdigit():
+        # if the period is after a digit, it's likely a partial decimal, return -1
+        return -1
+    elif text[idx - 1].isupper():
+        # if the period is after a capital letter (e.g., "Washington, D.C."), it's likely a abbreviation, return -1
+        return -1
+    elif idx > 1 and text[idx - 2 : idx + 1].lower() in ["a.m.", "p.m."]:
+        # if the period is after a.m. or p.m., it's likely a time, return -1
+        return -1
+    elif idx > 2 and text[idx - 3 : idx + 1] in ["e.g.", "i.e.", "etc."]:
+        # The period is after a character/word that is likely to be a abbreviation, return -1
+        return -1
+    elif idx >= 2 and text[idx - 2 : idx + 1].lower() in ["st.", "mr.", "mrs.", "ms.", "dr."]:
+        # if the period is after a character/word that is likely to be a abbreviation, return -1
+        return -1
+
+    # the text seems to have a complete sentence, return the index of the last period
+    return idx
+
+
+def find_last_comma_index(text: str, min_residual_length: int = 5) -> int:
+    """
+    Find the last occurrence of a valid comma in the text,
+    ignoring the commas in the numbers (e.g., "1,234,567").
+    If the leftover text after the comma is too short, it may be an abbreviation, return -1.
+
+    Args:
+        text: The text to find the last occurrence of a valid comma.
+        min_residual_length: The minimum length of the leftover text after the rightmost comma
+                             to be considered as a valid sentence (e.g., "Santa Clara, CA, US.").
+    Returns:
+        The index of the last occurrence of a valid comma, or -1 if no valid comma is found.
+    """
+    # find the last occurrence of a comma in the text
+    idx = text.rfind(",")
+    if idx == -1:
+        return -1
+    # check if the comma is in a number
+    if re.search(r'\d+,\d+', text[: idx + 1]):
+        # the comma is in a number, return -1
+        return -1
+
+    # check if the leftover text after the comma is too short
+    if len(text[idx + 1 :]) <= min_residual_length:
+        # the leftover text is too short, it may be an abbreviation, return -1
+        return -1
+
+    # the comma is not in a number, return the index of the comma
+    return idx
+
+
+class SimpleSegmentedTextAggregator(SimpleTextAggregator):
+    """A simple text aggregator that segments the text into sentences based on punctuation marks."""
+
+    def __init__(
+        self,
+        punctuation_marks: str | list[str] = ".,!?;:\n",
+        ignore_marks: str | list[str] = "*",
+        min_sentence_length: int = 0,
+        use_legacy_eos_detection: bool = False,
+        **kwargs,
+    ):
+        """
+        Args:
+            punctuation_marks: The punctuation marks to use for sentence detection.
+            ignore_marks: The strings to ignore in the text (e.g., "*").
+            min_sentence_length: The minimum length of a sentence to be considered.
+            use_legacy_eos_detection: Whether to use the legacy EOS detection from pipecat.
+            **kwargs: Additional arguments to pass to the SimpleTextAggregator constructor.
+        """
+        super().__init__(**kwargs)
+        self._use_legacy_eos_detection = use_legacy_eos_detection
+        self._min_sentence_length = min_sentence_length
+        self._ignore_marks = set(["*"] if ignore_marks is None else set(ignore_marks))
+        if not punctuation_marks:
+            self._punctuation_marks = list()
+        else:
+            punctuation_marks = (
+                [c for c in punctuation_marks] if isinstance(punctuation_marks, str) else punctuation_marks
+            )
+            if "." in punctuation_marks:
+                punctuation_marks.remove(".")
+            # put period at the end of the list to ensure it's the last punctuation mark to be matched
+            punctuation_marks += ["."]
+            self._punctuation_marks = punctuation_marks
+
+    def _find_segment_end(self, text: str) -> Optional[int]:
+        """find the end of text segment.
+
+        Args:
+            text: The text to find the end of the segment.
+
+        Returns:
+            The index of the end of the segment, or None if the text is too short.
+        """
+        # drop leading whitespace but keep trailing whitespace to
+        # allow "\n" to trigger the end of the sentence
+        text_len = len(text)
+        text = text.lstrip()
+        offset = text_len - len(text)
+        if len(text) < self._min_sentence_length:
+            return None
+
+        for punc in self._punctuation_marks:
+            if punc == ".":
+                idx = find_last_period_index(text)
+            elif punc == ",":
+                idx = find_last_comma_index(text)
+            else:
+                idx = text.find(punc)
+            if idx != -1:
+                # add the offset to the index to account for the leading whitespace
+                return idx + 1 + offset
+        return None
+
+    async def aggregate(self, text: str) -> AsyncIterator[Aggregation]:
+        """Aggregate the input text and return the first complete sentence in the text.
+
+        Args:
+            text: The text to aggregate.
+
+        Returns:
+            The first complete sentence in the text, or None if none is found.
+        """
+        result: Optional[str] = None
+        self._text += str(text)
+
+        eos_end_index = self._find_segment_end(self._text)
+
+        if not eos_end_index and not has_partial_decimal(self._text) and self._use_legacy_eos_detection:
+            # if the text doesn't have partial decimal, and no punctuation marks,
+            # we use match_endofsentence to find the end of the sentence
+            eos_end_index = match_endofsentence(self._text)
+
+        if eos_end_index:
+            result = self._text[:eos_end_index]
+            if len(result.strip()) < self._min_sentence_length:
+                logger.debug(
+                    f"Text is too short, skipping: `{result}`, full text: `{self._text}`, input text: `{text}`"
+                )
+                result = None
+            else:
+                logger.debug(f"Text Aggregator Result: `{result}`, full text: `{self._text}`, input text: `{text}`")
+                self._text = self._text[eos_end_index:]
+
+        if result:
+            for ignore_mark in self._ignore_marks:
+                result = result.replace(ignore_mark, "")
+            yield Aggregation(text=result, type=AggregationType.SENTENCE)

nemo/agents/voice_agent/utils/__init__.py

@@ -0,0 +1,15 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from nemo.agents.voice_agent.utils.config_manager import ConfigManager

nemo/agents/voice_agent/utils/config_manager.py

@@ -0,0 +1,312 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+from typing import Any, Dict, Optional
+
+from loguru import logger
+from omegaconf import OmegaConf
+from pipecat.audio.vad.silero import VADParams
+
+from nemo.agents.voice_agent.pipecat.services.nemo.diar import NeMoDiarInputParams
+from nemo.agents.voice_agent.pipecat.services.nemo.stt import NeMoSTTInputParams
+
+
+class ConfigManager:
+    """
+    Manages configuration for the voice agent server.
+    Handles loading, merging, and providing access to all configuration parameters.
+    """
+
+    def __init__(self, server_base_path: str, server_config_path: Optional[str] = None):
+        """
+        Initialize the configuration manager.
+
+        Args:
+            config_path: Path to the main server configuration file.
+                        If None, uses default path from environment variable.
+        """
+        if not os.path.exists(server_base_path):
+            raise FileNotFoundError(f"Server base path not found at {server_base_path}")
+
+        self._server_base_path = server_base_path
+        if server_config_path is not None:
+            self._server_config_path = server_config_path
+        else:
+            self._server_config_path = f"{os.path.abspath(self._server_base_path)}/server_configs/default.yaml"
+
+        if not os.path.exists(self._server_config_path):
+            raise FileNotFoundError(f"Server configuration file not found at {self._server_config_path}")
+
+        # Load model registry
+        self.model_registry_path = f"{os.path.abspath(self._server_base_path)}/model_registry.yaml"
+        self.model_registry = self._load_model_registry()
+
+        # Load and process main configuration
+        self.server_config = self._load_server_config()
+
+        # Initialize configuration parameters
+        self._initialize_config_parameters()
+
+        self._generic_hf_llm_model_id = "hf_llm_generic"
+
+        logger.info(f"Configuration loaded from: {self._server_config_path}")
+        logger.info(f"Model registry loaded from: {self.model_registry_path}")
+
+    def _load_model_registry(self) -> Dict[str, Any]:
+        """Load model registry from YAML file."""
+        try:
+            return OmegaConf.load(self.model_registry_path)
+        except Exception as e:
+            logger.error(f"Failed to load model registry: {e}")
+            raise ValueError(f"Failed to load model registry: {e}")
+
+    def _load_server_config(self) -> OmegaConf:
+        """Load and process the main server configuration."""
+        server_config = OmegaConf.load(self._server_config_path)
+        server_config = OmegaConf.to_container(server_config, resolve=True)
+        server_config = OmegaConf.create(server_config)
+        return server_config
+
+    def _initialize_config_parameters(self):
+        """Initialize all configuration parameters from the loaded config."""
+        # Default constants
+        self.SAMPLE_RATE = 16000
+        self.RAW_AUDIO_FRAME_LEN_IN_SECS = 0.016
+        self.SYSTEM_PROMPT = " ".join(
+            [
+                "You are a helpful AI agent named Lisa.",
+                "Begin by warmly greeting the user and introducing yourself in one sentence.",
+                "Keep your answers concise and to the point.",
+            ]
+        )
+
+        # Transport configuration
+        self.TRANSPORT_AUDIO_OUT_10MS_CHUNKS = self.server_config.transport.audio_out_10ms_chunks
+
+        # VAD configuration
+        self.vad_params = VADParams(
+            confidence=self.server_config.vad.confidence,
+            start_secs=self.server_config.vad.start_secs,
+            stop_secs=self.server_config.vad.stop_secs,
+            min_volume=self.server_config.vad.min_volume,
+        )
+        # STT configuration
+        self._configure_stt()
+
+        # Diarization configuration
+        self._configure_diarization()
+
+        # Turn taking configuration
+        self._configure_turn_taking()
+
+        # LLM configuration
+        self._configure_llm()
+
+        # TTS configuration
+        self._configure_tts()
+
+    def _configure_stt(self):
+        """Configure STT parameters."""
+        self.STT_MODEL = self.server_config.stt.model
+        self.STT_DEVICE = self.server_config.stt.device
+        # Apply STT-specific configuration based on model type
+        # Try to get STT config file name from server config first
+        if self.server_config.stt.get("model_config", None) is not None:
+            yaml_file_name = os.path.basename(self.server_config.stt.model_config)
+        else:
+            # Get STT configuration from registry
+            if str(self.STT_MODEL).endswith(".nemo"):
+                model_name = os.path.splitext(os.path.basename(self.STT_MODEL))[0]
+            else:
+                model_name = self.STT_MODEL
+            if model_name in self.model_registry.stt_models:
+                yaml_file_name = self.model_registry.stt_models[model_name].yaml_id
+            else:
+                error_msg = f"STT model {model_name} is not in model registry: {self.model_registry.stt_models}."
+                logger.error(error_msg)
+                raise ValueError(error_msg)
+
+        stt_config_path = f"{os.path.abspath(self._server_base_path)}/server_configs/stt_configs/{yaml_file_name}"
+        if not os.path.exists(stt_config_path):
+            raise FileNotFoundError(f"STT config file not found at {stt_config_path}")
+        stt_config = OmegaConf.load(stt_config_path)
+
+        # merge stt config with server config
+        for key in stt_config:
+            if key in self.server_config.stt and self.server_config.stt[key] != stt_config[key]:
+                logger.info(
+                    f"STT config field `{key}` is overridden from `{self.server_config.stt[key]}` "
+                    f"to `{stt_config[key]}` by {stt_config_path}"
+                )
+            self.server_config.stt[key] = stt_config[key]
+
+        logger.info(f"Final STT config: {self.server_config.stt}")
+
+        audio_chunk_size_in_secs = self.server_config.stt.get("audio_chunk_size_in_secs", 0.08)
+        buffer_size = audio_chunk_size_in_secs // self.RAW_AUDIO_FRAME_LEN_IN_SECS
+        self.stt_params = NeMoSTTInputParams(
+            att_context_size=self.server_config.stt.att_context_size,
+            frame_len_in_secs=self.server_config.stt.frame_len_in_secs,
+            raw_audio_frame_len_in_secs=self.RAW_AUDIO_FRAME_LEN_IN_SECS,
+            buffer_size=buffer_size,
+        )
+
+    def _configure_diarization(self):
+        """
+        Configure diarization parameters.
+        Currently only NeMo End-to-End Diarization is supported.
+        """
+        self.DIAR_MODEL = self.server_config.diar.model
+        self.USE_DIAR = self.server_config.diar.enabled
+        self.diar_params = NeMoDiarInputParams(
+            frame_len_in_secs=self.server_config.diar.frame_len_in_secs,
+            threshold=self.server_config.diar.threshold,
+        )
+
+    def _configure_turn_taking(self):
+        """Configure turn taking parameters."""
+        self.TURN_TAKING_BACKCHANNEL_PHRASES_PATH = self.server_config.turn_taking.backchannel_phrases_path
+        self.TURN_TAKING_MAX_BUFFER_SIZE = self.server_config.turn_taking.max_buffer_size
+        self.TURN_TAKING_BOT_STOP_DELAY = self.server_config.turn_taking.bot_stop_delay
+
+    def _configure_llm(self):
+        """Configure LLM parameters."""
+        llm_model_id = self.server_config.llm.model
+        is_registry_model = False
+
+        # Try to get LLM config file name from server config first
+        if self.server_config.llm.get("model_config", None) is not None:
+            yaml_file_name = os.path.basename(self.server_config.llm.model_config)
+        else:
+            # Get LLM configuration from registry
+            if llm_model_id in self.model_registry.llm_models:
+                yaml_file_name = self.model_registry.llm_models[llm_model_id].yaml_id
+                is_registry_model = True
+            else:
+                logger.warning(
+                    f"LLM model {llm_model_id} is not included in the model registry. "
+                    "Using a generic HuggingFace LLM config instead."
+                )
+                yaml_file_name = self.model_registry.llm_models[self._generic_hf_llm_model_id].yaml_id
+
+        # Load and merge LLM configuration
+        llm_config_path = f"{os.path.abspath(self._server_base_path)}/server_configs/llm_configs/{yaml_file_name}"
+
+        if (
+            is_registry_model
+            and self.model_registry.llm_models[llm_model_id].get("reasoning_supported", False)
+            and self.server_config.llm.get("enable_reasoning", False)
+        ):
+            llm_config_path = llm_config_path.replace(".yaml", "_think.yaml")
+
+        if not os.path.exists(llm_config_path):
+            raise FileNotFoundError(f"LLM config file not found at {llm_config_path}")
+        logger.info(f"Loading LLM config from: {llm_config_path}")
+
+        llm_config = OmegaConf.load(llm_config_path)
+        # merge llm config with server config
+        # print the override keys
+        for key in llm_config:
+            if key in self.server_config.llm and self.server_config.llm[key] != llm_config[key]:
+                logger.info(
+                    f"LLM config field `{key}` is overridden from `{self.server_config.llm[key]}` to "
+                    f"`{llm_config[key]}` by {llm_config_path}"
+                )
+            self.server_config.llm[key] = llm_config[key]
+
+        logger.info(f"Final LLM config: {self.server_config.llm}")
+
+        # Configure system prompt
+        self.SYSTEM_ROLE = self.server_config.llm.get("system_role", "system")
+        if self.server_config.llm.get("system_prompt", None) is not None:
+            system_prompt = self.server_config.llm.system_prompt
+            if os.path.isfile(system_prompt):
+                with open(system_prompt, "r") as f:
+                    system_prompt = f.read()
+            self.SYSTEM_PROMPT = system_prompt
+        else:
+            logger.info(f"No system prompt provided, using default system prompt: {self.SYSTEM_PROMPT}")
+
+        if self.server_config.llm.get("system_prompt_suffix", None) is not None:
+            self.SYSTEM_PROMPT += "\n" + self.server_config.llm.system_prompt_suffix
+            logger.info(f"Adding system prompt suffix: {self.server_config.llm.system_prompt_suffix}")
+
+        logger.info(f"System prompt: {self.SYSTEM_PROMPT}")
+
+    def _configure_tts(self):
+        """Configure TTS parameters."""
+        tts_model_id = self.server_config.tts.model
+
+        # Try to get TTS config file name from server config first
+        if self.server_config.tts.get("model_config", None) is not None:
+            yaml_file_name = os.path.basename(self.server_config.tts.model_config)
+        else:
+            # Get TTS configuration from registry
+            if tts_model_id in self.model_registry.tts_models:
+                yaml_file_name = self.model_registry.tts_models[tts_model_id].yaml_id
+            else:
+                error_msg = f"TTS model {tts_model_id} is not in model registry: {self.model_registry.tts_models}"
+                logger.error(error_msg)
+                raise ValueError(error_msg)
+
+        tts_config_path = f"{os.path.abspath(self._server_base_path)}/server_configs/tts_configs/{yaml_file_name}"
+        if not os.path.exists(tts_config_path):
+            raise FileNotFoundError(f"Default TTS config file not found at {tts_config_path}")
+        tts_config = OmegaConf.load(tts_config_path)
+
+        # merge tts config with server config
+        for key in tts_config:
+            if key in self.server_config.tts and self.server_config.tts[key] != tts_config[key]:
+                logger.info(
+                    f"TTS config field `{key}` is overridden from `{self.server_config.tts[key]}` to "
+                    f"`{tts_config[key]}` by {tts_config_path}"
+                )
+            self.server_config.tts[key] = tts_config[key]
+
+        logger.info(f"Final TTS config: {self.server_config.tts}")
+
+        # Extract TTS parameters
+        self.TTS_MAIN_MODEL_ID = self.server_config.tts.get("main_model_id", None)
+        self.TTS_SUB_MODEL_ID = self.server_config.tts.get("sub_model_id", None)
+        self.TTS_DEVICE = self.server_config.tts.get("device", None)
+
+        # Handle optional TTS parameters
+        self.TTS_THINK_TOKENS = self.server_config.tts.get("think_tokens", None)
+        if self.TTS_THINK_TOKENS is not None:
+            self.TTS_THINK_TOKENS = OmegaConf.to_container(self.TTS_THINK_TOKENS)
+
+        self.TTS_EXTRA_SEPARATOR = self.server_config.tts.get("extra_separator", None)
+        if self.TTS_EXTRA_SEPARATOR is not None:
+            self.TTS_EXTRA_SEPARATOR = OmegaConf.to_container(self.TTS_EXTRA_SEPARATOR)
+
+    def get_server_config(self) -> OmegaConf:
+        """Get the complete server configuration."""
+        return self.server_config
+
+    def get_model_registry(self) -> Dict[str, Any]:
+        """Get the model registry configuration."""
+        return self.model_registry
+
+    def get_vad_params(self) -> VADParams:
+        """Get VAD parameters."""
+        return self.vad_params
+
+    def get_stt_params(self) -> NeMoSTTInputParams:
+        """Get STT parameters."""
+        return self.stt_params
+
+    def get_diar_params(self) -> NeMoDiarInputParams:
+        """Get diarization parameters."""
+        return self.diar_params

nemo/agents/voice_agent/utils/tool_calling/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

nemo/agents/voice_agent/utils/tool_calling/basic_tools.py

@@ -0,0 +1,72 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import asyncio
+import python_weather
+from loguru import logger
+from pipecat.frames.frames import LLMTextFrame, TTSSpeakFrame
+from pipecat.processors.frame_processor import FrameDirection
+from pipecat.services.llm_service import FunctionCallParams
+
+HTTP_REQUEST_TIMEOUT = 10.0
+
+
+async def tool_get_city_weather(params: FunctionCallParams, city_name: str):
+    """Get the current weather of a city. The result includes city name, weather description,
+    temperature, wind speed, wind direction, precipitation, humidity, visibility, and UV index.
+
+    Args:
+        city_name: The name of the city to get the weather of. For example, "London", "Beijing", "Paris".
+                Other examples are: "Paris, TX, US", "Paris, FR" and "Tokyo, JP".
+    """
+    message = f"Looking up weather data for {city_name}. Please wait a moment..."
+    # Send the message to upstream so that RTVI can log it while doesn't block the actual tool call.
+    await params.llm.push_frame(LLMTextFrame(message), direction=FrameDirection.UPSTREAM)
+    # Send the message to TTS directly so that the user can hear it immediately.
+    await params.llm.push_frame(TTSSpeakFrame(message))
+
+    # The measuring unit defaults to metric (Celsius)
+    # Use imperial for Fahrenheit: python_weather.IMPERIAL
+    async with python_weather.Client(unit=python_weather.METRIC) as client:
+        # Fetch a weather forecast from a city
+        logger.debug(f"Fetching weather forecast for `{city_name}`")
+        try:
+            weather: python_weather.Forecast = await asyncio.wait_for(
+                client.get(city_name),
+                timeout=HTTP_REQUEST_TIMEOUT,
+            )
+        except asyncio.TimeoutError:
+            error_msg = f"python_weather API request timed out after {HTTP_REQUEST_TIMEOUT} seconds for `{city_name}`"
+            logger.error(error_msg)
+            await params.result_callback({"error": error_msg})
+            return
+        except Exception as e:
+            error_msg = f"Error fetching weather forecast for `{city_name}`: {str(e)}"
+            logger.error(error_msg)
+            await params.result_callback({"error": error_msg})
+            return
+
+        results = {
+            "city": city_name,
+            "description": str(weather.description),
+            "temperature": f"{weather.temperature} degrees Celsius",
+            "wind_speed": f"{weather.wind_speed} kilometers per hour",
+            "wind_direction": str(weather.wind_direction.name),
+            "precipitation": f"{weather.precipitation} millimeters",
+            "humidity": f"{weather.humidity} percent",
+            "visibility": f"{weather.visibility} kilometers",
+            "uv_index": str(weather.ultraviolet),
+        }
+        logger.debug(f"Weather results for {city_name}: {results}")
+        await params.result_callback(results)

nemo/agents/voice_agent/utils/tool_calling/mixins.py

@@ -0,0 +1,106 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from loguru import logger
+from pipecat.adapters.schemas.direct_function import DirectFunction
+from pipecat.adapters.schemas.tools_schema import ToolsSchema
+from pipecat.processors.aggregators.openai_llm_context import OpenAILLMContext
+from pipecat.services.openai.llm import OpenAILLMService
+
+
+class ToolCallingMixin:
+    """
+    A mixin class for tool calling.
+    Subclasses must implement the `setup_tool_calling` method to register all available tools
+    using `self.register_direct_function()`. Then the `__init__` method of the subclass should
+    call the `setup_tool_calling` method to register the tools.
+    """
+
+    def setup_tool_calling(self):
+        """
+        Setup the tool calling mixin by registering all available tools using self.register_direct_function().
+        """
+        raise NotImplementedError(
+            "Subclasses must implement this method to register all available functions "
+            "using self.register_direct_function()"
+        )
+
+    def register_direct_function(self, function_name: str, function: DirectFunction):
+        """
+        Register a direct function to be called by the LLM.
+
+        Args:
+            function_name: The name of the function to register.
+            function: The direct function to register.
+        """
+        if not hasattr(self, "direct_functions"):
+            self.direct_functions = {}
+        logger.info(
+            f"[{self.__class__.__name__}] Registering direct function name {function_name} to "
+            f"{function.__module__ + '.' + function.__qualname__}"
+        )
+        self.direct_functions[function_name] = function
+
+    @property
+    def available_tools(self) -> dict[str, DirectFunction]:
+        """
+        Return a dictionary of available tools, where the key is the tool name and the value is the direct function.
+        """
+        tools = {}
+        if not hasattr(self, "direct_functions"):
+            return tools
+        for function_name, function in self.direct_functions.items():
+            tools[function_name] = function
+        return tools
+
+
+def register_direct_tools_to_llm(
+    *,
+    llm: OpenAILLMService,
+    context: OpenAILLMContext,
+    tool_mixins: list[ToolCallingMixin] = [],
+    tools: list[DirectFunction] = [],
+    cancel_on_interruption: bool = True,
+) -> None:
+    """
+    Register direct tools to the LLM.
+    Args:
+        llm: The LLM service to use.
+        context: The LLM context to use.
+        tools: The list of tools (instances of either `DirectFunction` or `ToolCallingMixin`) to use.
+    """
+    all_tools = []
+    for tool in tool_mixins:
+        if not isinstance(tool, ToolCallingMixin):
+            logger.warning(f"Tool {tool.__class__.__name__} is not a ToolCallingMixin, skipping.")
+            continue
+        for function_name, function in tool.available_tools.items():
+            logger.info(f"Registering direct function {function_name} from {tool.__class__.__name__}")
+            all_tools.append(function)
+
+    for tool in tools:
+        logger.info(f"Registering direct function: {tool.__module__ + '.' + tool.__qualname__}")
+        all_tools.append(tool)
+
+    if not all_tools:
+        logger.warning("No direct tools provided.")
+        return
+    else:
+        logger.info(f"Registering {len(all_tools)} direct tools to the LLM.")
+
+    tools_schema = ToolsSchema(standard_tools=all_tools)
+    context.set_tools(tools_schema)
+
+    for tool in all_tools:
+        llm.register_direct_function(tool, cancel_on_interruption=cancel_on_interruption)

nemo/collections/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

nemo/collections/asr/README.md

@@ -0,0 +1,37 @@
+# Automatic Speech Recognition (ASR)
+
+## Key Features
+
+* [HuggingFace Space for Audio Transcription (File, Microphone and YouTube)](https://huggingface.co/spaces/smajumdar/nemo_multilingual_language_id)
+* [Pretrained models](https://ngc.nvidia.com/catalog/collections/nvidia:nemo_asr) available in 14+ languages
+* [Automatic Speech Recognition (ASR)](https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/asr/intro.html)
+    * Supported ASR [models](https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/stable/asr/models.html):
+        * Jasper, QuartzNet, CitriNet, ContextNet
+        * Conformer-CTC, Conformer-Transducer, FastConformer-CTC, FastConformer-Transducer
+        * Squeezeformer-CTC and Squeezeformer-Transducer
+        * LSTM-Transducer (RNNT) and LSTM-CTC
+    * Supports the following decoders/losses:
+        * CTC
+        * Transducer/RNNT
+        * Hybrid Transducer/CTC
+        * NeMo Original [Multi-blank Transducers](https://arxiv.org/abs/2211.03541) and [Token-and-Duration Transducers (TDT)](https://arxiv.org/abs/2304.06795)
+    * Streaming/Buffered ASR (CTC/Transducer) - [Chunked Inference Examples](https://github.com/NVIDIA/NeMo/tree/stable/examples/asr/asr_chunked_inference)
+    * [Cache-aware Streaming Conformer](https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/stable/asr/models.html#cache-aware-streaming-conformer) with multiple lookaheads (including microphone streaming [tutorial](https://github.com/NVIDIA/NeMo/blob/main/tutorials/asr/Online_ASR_Microphone_Demo_Cache_Aware_Streaming.ipynb).
+    * Beam Search decoding
+    * [Language Modelling for ASR (CTC and RNNT)](https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/asr/asr_language_modeling.html): N-gram LM in fusion with Beam Search decoding, Neural Rescoring with Transformer
+    * [Support of long audios for Conformer with memory efficient local attention](https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/asr/results.html#inference-on-long-audio)
+* [Speech Classification, Speech Command Recognition and Language Identification](https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/asr/speech_classification/intro.html): MatchboxNet (Command Recognition), AmberNet (LangID)
+* [Voice activity Detection (VAD)](https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/stable/asr/speech_classification/models.html#marblenet-vad): MarbleNet
+    * ASR with VAD Inference - [Example](https://github.com/NVIDIA/NeMo/tree/stable/examples/asr/asr_vad)
+* [Speaker Recognition](https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/asr/speaker_recognition/intro.html): TitaNet, ECAPA_TDNN, SpeakerNet
+* [Speaker Diarization](https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/asr/speaker_diarization/intro.html)
+    * Clustering Diarizer: TitaNet, ECAPA_TDNN, SpeakerNet
+    * Neural Diarizer: Sortformer
+* [Speech Intent Detection and Slot Filling](https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/asr/speech_intent_slot/intro.html): Conformer-Transformer
+
+You can also get a high-level overview of NeMo ASR by watching the talk *NVIDIA NeMo: Toolkit for Conversational AI*, presented at PyData Yerevan 2022:
+
+
+[![NVIDIA NeMo: Toolkit for Conversational AI](https://img.youtube.com/vi/J-P6Sczmas8/maxres3.jpg
+)](https://www.youtube.com/embed/J-P6Sczmas8?mute=0&start=14&autoplay=0
+ "NeMo presentation at PyData@Yerevan 2022")

nemo/collections/asr/__init__.py

@@ -0,0 +1,25 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from nemo.collections.asr import data, losses, models, modules
+from nemo.package_info import __version__
+
+# Set collection version equal to NeMo version.
+__version = __version__
+
+# Authorship.
+__author__ = "NVIDIA Corporation"
+
+# Set collection name.
+__description__ = "Automatic Speech Recognition collection"

nemo/collections/asr/data/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

nemo/collections/asr/data/audio_to_ctm_dataset.py

@@ -0,0 +1,97 @@
+# Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+import os
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, List, Tuple
+
+from nemo.collections.asr.data.audio_to_text_dataset import ASRPredictionWriter
+from nemo.utils import logging
+
+if TYPE_CHECKING:
+    from lightning.pytorch import LightningModule
+
+
+@dataclass
+class FrameCtmUnit:
+    """A container class for one CTM unit with start and length countable in frames."""
+
+    label: str
+    start_frame: int
+    length: int
+    probability: float
+
+    def __repr__(self) -> str:
+        return f"{self.label}\t({self.probability:1.3f}): [{self.start_frame:6d}, {self.length:6d}]"
+
+    @property
+    def end_frame(self):
+        return self.start_frame + self.length
+
+    def to_ctm_str(self, time_per_frame: int) -> str:
+        """Represents the data as part of the CTM line.
+
+        The CTM line format is
+            <utterance_name> <channel> <start_time> <duration> <label_str> <probability>
+        This method prepares the last four entities."""
+        return f"{self.start_frame * time_per_frame :.3f} {self.length * time_per_frame :.3f} {self.label} {self.probability :1.3f}"
+
+
+class ASRCTMPredictionWriter(ASRPredictionWriter):
+    def __init__(self, dataset, output_file: str, output_ctm_dir: str, time_per_frame: float):
+        super().__init__(dataset, output_file)
+        self.output_ctm_dir = output_ctm_dir
+        self.time_per_frame = time_per_frame
+        os.makedirs(self.output_ctm_dir, exist_ok=True)
+
+    def write_ctm(self, name, filepath, frameCtmUnits):
+        with open(filepath, "tw", encoding="utf-8") as f:
+            for unit in frameCtmUnits:
+                f.write(f"{name} 1 {unit.to_ctm_str(self.time_per_frame)}\n")
+
+    def write_on_batch_end(
+        self,
+        trainer,
+        pl_module: 'LightningModule',
+        prediction: Tuple[int, List[FrameCtmUnit]],
+        batch_indices: List[int],
+        batch: Any,
+        batch_idx: int,
+        dataloader_idx: int,
+    ):
+        for sample_id, units in prediction:
+            sample = self.dataset.get_manifest_sample(sample_id)
+            with_ctm = True
+            if len(units) == 0:
+                logging.warning(
+                    f"""Do not producing CTM output for item `{sample.audio_file}`.
+                    Check if text is empty or if duration is too short: `{sample.text_raw}`, {sample.duration}"""
+                )
+                with_ctm = False
+            item = {}
+            item["audio_filepath"] = sample.audio_file
+            item["duration"] = sample.duration
+            item["text"] = sample.text_raw
+            if with_ctm:
+                utt_name = Path(sample.audio_file).stem
+                ctm_filepath = os.path.join(self.output_ctm_dir, utt_name) + ".ctm"
+                self.write_ctm(utt_name, ctm_filepath, units)
+                item["ctm_filepath"] = ctm_filepath
+            else:
+                item["ctm_filepath"] = ""
+            self.outf.write(json.dumps(item) + "\n")
+            self.samples_num += 1
+        return

nemo/collections/asr/data/audio_to_diar_label.py

@@ -0,0 +1,562 @@
+# Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+from typing import Dict, List, Optional, Tuple
+
+import numpy as np
+import torch
+
+from nemo.collections.asr.parts.utils.speaker_utils import convert_rttm_line, get_subsegments
+from nemo.collections.common.parts.preprocessing.collections import EndtoEndDiarizationSpeechLabel
+from nemo.core.classes import Dataset
+from nemo.core.neural_types import AudioSignal, LengthsType, NeuralType, ProbsType
+from nemo.utils import logging
+
+
+def get_subsegments_to_timestamps(
+    subsegments: List[Tuple[float, float]], feat_per_sec: int = 100, max_end_ts: float = None, decimals=2
+):
+    """
+    Convert subsegment timestamps to scale timestamps by multiplying with the feature rate (`feat_per_sec`)
+    and rounding. Segment is consisted of many subsegments and sugsegments are equivalent to `frames`
+    in end-to-end speaker diarization models.
+
+    Args:
+        subsegments (List[Tuple[float, float]]):
+            A list of tuples where each tuple contains the start and end times of a subsegment
+            (frames in end-to-end models).
+            >>> subsegments = [[t0_start, t0_duration], [t1_start, t1_duration],..., [tN_start, tN_duration]]
+        feat_per_sec (int, optional):
+            The number of feature frames per second. Defaults to 100.
+        max_end_ts (float, optional):
+            The maximum end timestamp to clip the results. If None, no clipping is applied. Defaults to None.
+        decimals (int, optional):
+            The number of decimal places to round the timestamps. Defaults to 2.
+
+    Example:
+        Segments starting from 0.0 and ending at 69.2 seconds.
+        If hop-length is 0.08 and the subsegment (frame) length is 0.16 seconds,
+        there are 864 = (69.2 - 0.16)/0.08 + 1 subsegments (frames in end-to-end models) in this segment.
+        >>> subsegments = [[[0.0, 0.16], [0.08, 0.16], ..., [69.04, 0.16], [69.12, 0.08]]
+
+    Returns:
+        ts (torch.tensor):
+            A tensor containing the scaled and rounded timestamps for each subsegment.
+    """
+    seg_ts = (torch.tensor(subsegments) * feat_per_sec).float()
+    ts_round = torch.round(seg_ts, decimals=decimals)
+    ts = ts_round.long()
+    ts[:, 1] = ts[:, 0] + ts[:, 1]
+    if max_end_ts is not None:
+        ts = np.clip(ts, 0, int(max_end_ts * feat_per_sec))
+    return ts
+
+
+def extract_frame_info_from_rttm(offset, duration, rttm_lines, round_digits=3):
+    """
+    Extracts RTTM lines containing speaker labels, start time, and end time for a given audio segment.
+
+    Args:
+        uniq_id (str): Unique identifier for the audio file and corresponding RTTM file.
+        offset (float): The starting time offset for the segment of interest.
+        duration (float): The duration of the segment of interest.
+        rttm_lines (list): List of RTTM lines in string format.
+        round_digits (int, optional): Number of decimal places to round the start and end times. Defaults to 3.
+
+    Returns:
+        rttm_mat (tuple): A tuple containing lists of start times, end times, and speaker labels.
+        sess_to_global_spkids (dict): A mapping from session-specific speaker indices to global speaker identifiers.
+    """
+    rttm_stt, rttm_end = offset, offset + duration
+    stt_list, end_list, speaker_list, speaker_set = [], [], [], []
+    sess_to_global_spkids = dict()
+
+    for rttm_line in rttm_lines:
+        start, end, speaker = convert_rttm_line(rttm_line)
+
+        # Skip invalid RTTM lines where the start time is greater than the end time.
+        if start > end:
+            continue
+
+        # Check if the RTTM segment overlaps with the specified segment of interest.
+        if (end > rttm_stt and start < rttm_end) or (start < rttm_end and end > rttm_stt):
+            # Adjust the start and end times to fit within the segment of interest.
+            start, end = max(start, rttm_stt), min(end, rttm_end)
+        else:
+            continue
+
+        # Round the start and end times to the specified number of decimal places.
+        end_list.append(round(end, round_digits))
+        stt_list.append(round(start, round_digits))
+
+        # Assign a unique index to each speaker and maintain a mapping.
+        if speaker not in speaker_set:
+            speaker_set.append(speaker)
+        speaker_list.append(speaker_set.index(speaker))
+        sess_to_global_spkids.update({speaker_set.index(speaker): speaker})
+
+    rttm_mat = (stt_list, end_list, speaker_list)
+    return rttm_mat, sess_to_global_spkids
+
+
+def get_frame_targets_from_rttm(
+    rttm_timestamps: list,
+    offset: float,
+    duration: float,
+    round_digits: int,
+    feat_per_sec: int,
+    max_spks: int,
+):
+    """
+    Create a multi-dimensional vector sequence containing speaker timestamp information in RTTM.
+    The unit-length is the frame shift length of the acoustic feature. The feature-level annotations
+    `feat_level_target` will later be converted to base-segment level diarization label.
+
+    Args:
+        rttm_timestamps (list):
+            List containing start and end time for each speaker segment label.
+            stt_list, end_list and speaker_list are contained.
+        feat_per_sec (int):
+            Number of feature frames per second.
+            This quantity is determined by window_stride variable in preprocessing module.
+        target_spks (tuple):
+            Speaker indices that are generated from combinations. If there are only one or two speakers,
+            only a single target_spks variable is generated.
+
+    Returns:
+        feat_level_target (torch.tensor):
+            Tensor containing label for each feature level frame.
+    """
+    stt_list, end_list, speaker_list = rttm_timestamps
+    sorted_speakers = sorted(list(set(speaker_list)))
+    total_fr_len = int(duration * feat_per_sec)
+    if len(sorted_speakers) > max_spks:
+        logging.warning(
+            f"Number of speakers in RTTM file {len(sorted_speakers)} exceeds the maximum number of speakers: "
+            f"{max_spks}! Only {max_spks} first speakers remain, and this will affect frame metrics!"
+        )
+    feat_level_target = torch.zeros(total_fr_len, max_spks)
+    for count, (stt, end, spk_rttm_key) in enumerate(zip(stt_list, end_list, speaker_list)):
+        if end < offset or stt > offset + duration:
+            continue
+        stt, end = max(offset, stt), min(offset + duration, end)
+        spk = spk_rttm_key
+        if spk < max_spks:
+            stt_fr, end_fr = int((stt - offset) * feat_per_sec), int((end - offset) * feat_per_sec)
+            feat_level_target[stt_fr:end_fr, spk] = 1
+    return feat_level_target
+
+
+class _AudioToSpeechE2ESpkDiarDataset(Dataset):
+    """
+    Dataset class that loads a json file containing paths to audio files,
+    RTTM files and number of speakers. This Dataset class is designed for
+    training or fine-tuning speaker embedding extractor and diarization decoder
+    at the same time.
+
+    Example:
+    {"audio_filepath": "/path/to/audio_0.wav", "num_speakers": 2,
+    "rttm_filepath": "/path/to/diar_label_0.rttm}
+    ...
+    {"audio_filepath": "/path/to/audio_n.wav", "num_speakers": 2,
+    "rttm_filepath": "/path/to/diar_label_n.rttm}
+
+    Args:
+        manifest_filepath (str):
+            Path to input manifest json files.
+        multiargs_dict (dict):
+            Dictionary containing the parameters for multiscale segmentation and clustering.
+        soft_label_thres (float):
+            Threshold that determines the label of each segment based on RTTM file information.
+        featurizer:
+            Featurizer instance for generating audio_signal from the raw waveform.
+        window_stride (float):
+            Window stride for acoustic feature. This value is used for calculating the numbers of feature-level frames.
+    """
+
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        """Returns definitions of module output ports."""
+        output_types = {
+            "audio_signal": NeuralType(('B', 'T'), AudioSignal()),
+            "audio_length": NeuralType(('B'), LengthsType()),
+            "targets": NeuralType(('B', 'T', 'C'), ProbsType()),
+            "target_len": NeuralType(('B'), LengthsType()),
+        }
+
+        return output_types
+
+    def __init__(
+        self,
+        *,
+        manifest_filepath: str,
+        soft_label_thres: float,
+        session_len_sec: float,
+        num_spks: int,
+        featurizer,
+        fb_featurizer,
+        window_stride: float,
+        min_subsegment_duration: float = 0.03,
+        global_rank: int = 0,
+        dtype=torch.float16,
+        round_digits: int = 2,
+        soft_targets: bool = False,
+        subsampling_factor: int = 8,
+        device: str = 'cpu',
+    ):
+        super().__init__()
+        self.collection = EndtoEndDiarizationSpeechLabel(
+            manifests_files=manifest_filepath.split(','),
+            round_digits=round_digits,
+        )
+        self.featurizer = featurizer
+        self.fb_featurizer = fb_featurizer
+        # STFT and subsampling factor parameters
+        self.n_fft = self.fb_featurizer.n_fft
+        self.hop_length = self.fb_featurizer.hop_length
+        self.stft_pad_amount = self.fb_featurizer.stft_pad_amount
+        self.subsampling_factor = subsampling_factor
+        # Annotation and target length parameters
+        self.round_digits = round_digits
+        self.feat_per_sec = int(1 / window_stride)
+        self.diar_frame_length = round(subsampling_factor * window_stride, round_digits)
+        self.session_len_sec = session_len_sec
+        self.soft_label_thres = soft_label_thres
+        self.max_spks = num_spks
+        self.min_subsegment_duration = min_subsegment_duration
+        self.dtype = dtype
+        self.use_asr_style_frame_count = True
+        self.soft_targets = soft_targets
+        self.round_digits = 2
+        self.floor_decimal = 10**self.round_digits
+        self.device = device
+        self.global_rank = global_rank
+
+    def __len__(self):
+        return len(self.collection)
+
+    def get_frame_count_from_time_series_length(self, seq_len):
+        """
+        This function is used to get the sequence length of the audio signal. This is required to match
+        the feature frame length with ASR (STT) models. This function is copied from
+        NeMo/nemo/collections/asr/parts/preprocessing/features.py::FilterbankFeatures::get_seq_len.
+
+        Args:
+            seq_len (int):
+                The sequence length of the time-series data.
+
+        Returns:
+            seq_len (int):
+                The sequence length of the feature frames.
+        """
+        pad_amount = self.stft_pad_amount * 2 if self.stft_pad_amount is not None else self.n_fft // 2 * 2
+        seq_len = torch.floor_divide((seq_len + pad_amount - self.n_fft), self.hop_length).to(dtype=torch.long)
+        frame_count = int(np.ceil(seq_len / self.subsampling_factor))
+        return frame_count
+
+    def get_uniq_id_with_range(self, sample, deci=3):
+        """
+        Generate unique training sample ID from unique file ID, offset and duration. The start-end time added
+        unique ID is required for identifying the sample since multiple short audio samples are generated from a single
+        audio file. The start time and end time of the audio stream uses millisecond units if `deci=3`.
+
+        Args:
+            sample:
+                `EndtoEndDiarizationSpeechLabel` instance from collections.
+
+        Returns:
+            uniq_id (str):
+                Unique sample ID which includes start and end time of the audio stream.
+                Example: abc1001_3122_6458
+        """
+        bare_uniq_id = os.path.splitext(os.path.basename(sample.rttm_file))[0]
+        offset = str(int(round(sample.offset, deci) * pow(10, deci)))
+        endtime = str(int(round(sample.offset + sample.duration, deci) * pow(10, deci)))
+        uniq_id = f"{bare_uniq_id}_{offset}_{endtime}"
+        return uniq_id
+
+    def parse_rttm_for_targets_and_lens(self, rttm_file, offset, duration, target_len):
+        """
+        Generate target tensor variable by extracting groundtruth diarization labels from an RTTM file.
+        This function converts (start, end, speaker_id) format into base-scale (the finest scale) segment level
+        diarization label in a matrix form.
+
+        Example of seg_target:
+            [[0., 1.], [0., 1.], [1., 1.], [1., 0.], [1., 0.], ..., [0., 1.]]
+        """
+        if rttm_file in [None, '']:
+            num_seg = torch.max(target_len)
+            targets = torch.zeros(num_seg, self.max_spks)
+            return targets
+
+        with open(rttm_file, 'r') as f:
+            rttm_lines = f.readlines()
+
+        rttm_timestamps, sess_to_global_spkids = extract_frame_info_from_rttm(offset, duration, rttm_lines)
+
+        fr_level_target = get_frame_targets_from_rttm(
+            rttm_timestamps=rttm_timestamps,
+            offset=offset,
+            duration=duration,
+            round_digits=self.round_digits,
+            feat_per_sec=self.feat_per_sec,
+            max_spks=self.max_spks,
+        )
+
+        soft_target_seg = self.get_soft_targets_seg(feat_level_target=fr_level_target, target_len=target_len)
+        if self.soft_targets:
+            step_target = soft_target_seg
+        else:
+            step_target = (soft_target_seg >= self.soft_label_thres).float()
+        return step_target
+
+    def get_soft_targets_seg(self, feat_level_target, target_len):
+        """
+        Generate the final targets for the actual diarization step.
+        Here, frame level means step level which is also referred to as segments.
+        We follow the original paper and refer to the step level as "frames".
+
+        Args:
+            feat_level_target (torch.tensor):
+                Tensor variable containing hard-labels of speaker activity in each feature-level segment.
+            target_len (torch.tensor):
+                Numbers of ms segments
+
+        Returns:
+            soft_target_seg (torch.tensor):
+                Tensor variable containing soft-labels of speaker activity in each step-level segment.
+        """
+        num_seg = torch.max(target_len)
+        targets = torch.zeros(num_seg, self.max_spks)
+        stride = int(self.feat_per_sec * self.diar_frame_length)
+        for index in range(num_seg):
+            if index == 0:
+                seg_stt_feat = 0
+            else:
+                seg_stt_feat = stride * index - 1 - int(stride / 2)
+            if index == num_seg - 1:
+                seg_end_feat = feat_level_target.shape[0]
+            else:
+                seg_end_feat = stride * index - 1 + int(stride / 2)
+            targets[index] = torch.mean(feat_level_target[seg_stt_feat : seg_end_feat + 1, :], axis=0)
+        return targets
+
+    def get_segment_timestamps(
+        self,
+        duration: float,
+        offset: float = 0,
+        sample_rate: int = 16000,
+    ):
+        """
+        Get start and end time of segments in each scale.
+
+        Args:
+            sample:
+                `EndtoEndDiarizationSpeechLabel` instance from preprocessing.collections
+        Returns:
+            segment_timestamps (torch.tensor):
+                Tensor containing Multiscale segment timestamps.
+            target_len (torch.tensor):
+                Number of segments for each scale. This information is used for reshaping embedding batch
+                during forward propagation.
+        """
+        subsegments = get_subsegments(
+            offset=offset,
+            window=round(self.diar_frame_length * 2, self.round_digits),
+            shift=self.diar_frame_length,
+            duration=duration,
+            min_subsegment_duration=self.min_subsegment_duration,
+            use_asr_style_frame_count=self.use_asr_style_frame_count,
+            sample_rate=sample_rate,
+            feat_per_sec=self.feat_per_sec,
+        )
+        if self.use_asr_style_frame_count:
+            effective_dur = (
+                np.ceil((1 + duration * sample_rate) / int(sample_rate / self.feat_per_sec)).astype(int)
+                / self.feat_per_sec
+            )
+        else:
+            effective_dur = duration
+        ts_tensor = get_subsegments_to_timestamps(
+            subsegments, self.feat_per_sec, decimals=2, max_end_ts=(offset + effective_dur)
+        )
+        target_len = torch.tensor([ts_tensor.shape[0]])
+        return target_len
+
+    def __getitem__(self, index):
+        sample = self.collection[index]
+        if sample.offset is None:
+            sample.offset = 0
+        offset = sample.offset
+        if self.session_len_sec < 0:
+            session_len_sec = sample.duration
+        else:
+            session_len_sec = min(sample.duration, self.session_len_sec)
+
+        audio_signal = self.featurizer.process(sample.audio_file, offset=offset, duration=session_len_sec)
+
+        # We should resolve the length mis-match from the round-off errors between these two variables:
+        # `session_len_sec` and `audio_signal.shape[0]`
+        session_len_sec = (
+            np.floor(audio_signal.shape[0] / self.featurizer.sample_rate * self.floor_decimal) / self.floor_decimal
+        )
+        audio_signal = audio_signal[: round(self.featurizer.sample_rate * session_len_sec)]
+        audio_signal_length = torch.tensor(audio_signal.shape[0]).long()
+
+        # Target length should be following the ASR feature extraction convention: Use self.get_frame_count_from_time_series_length.
+        target_len = self.get_segment_timestamps(duration=session_len_sec, sample_rate=self.featurizer.sample_rate)
+        target_len = torch.clamp(target_len, max=self.get_frame_count_from_time_series_length(audio_signal.shape[0]))
+
+        targets = self.parse_rttm_for_targets_and_lens(
+            rttm_file=sample.rttm_file, offset=offset, duration=session_len_sec, target_len=target_len
+        )
+        targets = targets[:target_len, :]
+        return audio_signal, audio_signal_length, targets, target_len
+
+
+def _eesd_train_collate_fn(self, batch):
+    """
+    Collate a batch of variables needed for training the end-to-end speaker diarization (EESD) model
+    from raw waveforms to diarization labels. The following variables are included in the training/validation batch:
+
+    Args:
+        batch (tuple):
+            A tuple containing the variables for diarization training.
+
+    Returns:
+        audio_signal (torch.Tensor):
+            A tensor containing the raw waveform samples (time series) loaded from the `audio_filepath`
+            in the input manifest file.
+        feature_length (torch.Tensor):
+            A tensor containing the lengths of the raw waveform samples.
+        targets (torch.Tensor):
+            Groundtruth speaker labels for the given input embedding sequence.
+        target_lens (torch.Tensor):
+            A tensor containing the number of segments for each sample in the batch, necessary for
+            reshaping inputs to the EESD model.
+    """
+    packed_batch = list(zip(*batch))
+    audio_signal, feature_length, targets, target_len = packed_batch
+    audio_signal_list, feature_length_list = [], []
+    target_len_list, targets_list = [], []
+
+    max_raw_feat_len = max([x.shape[0] for x in audio_signal])
+    max_target_len = max([x.shape[0] for x in targets])
+    if max([len(feat.shape) for feat in audio_signal]) > 1:
+        max_ch = max([feat.shape[1] for feat in audio_signal])
+    else:
+        max_ch = 1
+    for feat, feat_len, tgt, segment_ct in batch:
+        seq_len = tgt.shape[0]
+        if len(feat.shape) > 1:
+            pad_feat = (0, 0, 0, max_raw_feat_len - feat.shape[0])
+        else:
+            pad_feat = (0, max_raw_feat_len - feat.shape[0])
+        if feat.shape[0] < feat_len:
+            feat_len_pad = feat_len - feat.shape[0]
+            feat = torch.nn.functional.pad(feat, (0, feat_len_pad))
+        pad_tgt = (0, 0, 0, max_target_len - seq_len)
+        padded_feat = torch.nn.functional.pad(feat, pad_feat)
+        padded_tgt = torch.nn.functional.pad(tgt, pad_tgt)
+        if max_ch > 1 and padded_feat.shape[1] < max_ch:
+            feat_ch_pad = max_ch - padded_feat.shape[1]
+            padded_feat = torch.nn.functional.pad(padded_feat, (0, feat_ch_pad))
+        audio_signal_list.append(padded_feat)
+        feature_length_list.append(feat_len.clone().detach())
+        target_len_list.append(segment_ct.clone().detach())
+        targets_list.append(padded_tgt)
+        audio_signal = torch.stack(audio_signal_list)
+    feature_length = torch.stack(feature_length_list)
+    target_lens = torch.stack(target_len_list).squeeze(1)
+    targets = torch.stack(targets_list)
+    return audio_signal, feature_length, targets, target_lens
+
+
+class AudioToSpeechE2ESpkDiarDataset(_AudioToSpeechE2ESpkDiarDataset):
+    """
+    Dataset class for loading a JSON file containing paths to audio files,
+    RTTM (Rich Transcription Time Marked) files, and the number of speakers.
+    This class is designed for training or fine-tuning a speaker embedding
+    extractor and diarization decoder simultaneously.
+
+    The JSON manifest file should have entries in the following format:
+
+    Example:
+    {
+        "audio_filepath": "/path/to/audio_0.wav",
+        "num_speakers": 2,
+        "rttm_filepath": "/path/to/diar_label_0.rttm"
+    }
+    ...
+    {
+        "audio_filepath": "/path/to/audio_n.wav",
+        "num_speakers": 2,
+        "rttm_filepath": "/path/to/diar_label_n.rttm"
+    }
+
+    Args:
+        manifest_filepath (str):
+            Path to the input manifest JSON file containing paths to audio and RTTM files.
+        soft_label_thres (float):
+            Threshold for assigning soft labels to segments based on RTTM file information.
+        session_len_sec (float):
+            Duration of each session (in seconds) for training or fine-tuning.
+        num_spks (int):
+            Number of speakers in the audio files.
+        featurizer:
+            Instance of a featurizer for generating features from the raw waveform.
+        window_stride (float):
+            Window stride (in seconds) for extracting acoustic features, used to calculate
+            the number of feature frames.
+        global_rank (int):
+            Global rank of the current process (used for distributed training).
+        soft_targets (bool):
+            Whether or not to use soft targets during training.
+
+    Methods:
+        eesd_train_collate_fn(batch):
+            Collates a batch of data for end-to-end speaker diarization training.
+    """
+
+    def __init__(
+        self,
+        *,
+        manifest_filepath: str,
+        soft_label_thres: float,
+        session_len_sec: float,
+        num_spks: int,
+        featurizer,
+        fb_featurizer,
+        window_stride,
+        global_rank: int,
+        soft_targets: bool,
+        device: str,
+    ):
+        super().__init__(
+            manifest_filepath=manifest_filepath,
+            soft_label_thres=soft_label_thres,
+            session_len_sec=session_len_sec,
+            num_spks=num_spks,
+            featurizer=featurizer,
+            fb_featurizer=fb_featurizer,
+            window_stride=window_stride,
+            global_rank=global_rank,
+            soft_targets=soft_targets,
+            device=device,
+        )
+
+    def eesd_train_collate_fn(self, batch):
+        """Collate a batch of data for end-to-end speaker diarization training."""
+        return _eesd_train_collate_fn(self, batch)

nemo/collections/asr/data/audio_to_diar_label_lhotse.py

@@ -0,0 +1,114 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Dict, Optional, Tuple
+
+import torch.utils.data
+from lhotse.dataset import AudioSamples
+from lhotse.dataset.collation import collate_matrices
+
+from nemo.collections.asr.parts.utils.asr_multispeaker_utils import (
+    get_hidden_length_from_sample_length,
+    speaker_to_target,
+)
+from nemo.core.neural_types import AudioSignal, LabelsType, LengthsType, NeuralType
+from nemo.utils import logging
+
+
+class LhotseAudioToSpeechE2ESpkDiarDataset(torch.utils.data.Dataset):
+    """
+    This dataset is a Lhotse version of diarization dataset in audio_to_diar_label.py.
+    Unlike native NeMo datasets, Lhotse dataset defines only the mapping from
+    a CutSet (meta-data) to a mini-batch with PyTorch tensors.
+    Specifically, it performs tokenization, I/O, augmentation, and feature extraction (if any).
+    Managing data, sampling, de-duplication across workers/nodes etc. is all handled
+    by Lhotse samplers instead.
+    """
+
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        """Define the output types of the dataset."""
+        return {
+            'audio_signal': NeuralType(('B', 'T'), AudioSignal()),
+            'a_sig_length': NeuralType(tuple('B'), LengthsType()),
+            'targets': NeuralType(('B', 'T', 'N'), LabelsType()),
+            'target_length': NeuralType(tuple('B'), LengthsType()),
+            'sample_id': NeuralType(tuple('B'), LengthsType(), optional=True),
+        }
+
+    def __init__(self, cfg):
+        super().__init__()
+        self.load_audio = AudioSamples(fault_tolerant=True)
+        self.cfg = cfg
+        self.num_speakers = self.cfg.get('num_speakers', 4)
+        self.num_sample_per_mel_frame = int(
+            self.cfg.get('window_stride', 0.01) * self.cfg.get('sample_rate', 16000)
+        )  # 160 samples for every 1ms by default
+        self.num_mel_frame_per_target_frame = int(self.cfg.get('subsampling_factor', 8))
+
+    def __getitem__(self, cuts) -> Tuple[torch.Tensor, ...]:
+        # NOTE: This end-to-end diarization dataloader only loads the 1st ch of the audio file.
+        # Process cuts in a single loop: convert to mono and compute speaker activities
+        mono_cuts = []
+        speaker_activities = []
+        for cut in cuts:
+            if cut.num_channels is not None and cut.num_channels > 1:
+                logging.warning(
+                    "Multiple channels detected in cut '%s' (%d channels). "
+                    "Only the first channel will be used; remaining channels are ignored.",
+                    cut.id,
+                    cut.num_channels,
+                )
+            mono_cut = cut.with_channels(channels=[0])
+            mono_cuts.append(mono_cut)
+
+            speaker_activity = speaker_to_target(
+                a_cut=mono_cut,
+                num_speakers=self.num_speakers,
+                num_sample_per_mel_frame=self.num_sample_per_mel_frame,
+                num_mel_frame_per_asr_frame=self.num_mel_frame_per_target_frame,
+                boundary_segments=True,
+            )
+            # This line prevents dimension mismatch error in the collate_matrices function.
+            if speaker_activity.shape[1] > self.num_speakers:
+                logging.warning(
+                    "Number of speakers in the target %s is greater than "
+                    "the maximum number of speakers %s. Truncating extra speakers. "
+                    "Set the `num_speakers` to higher value to avoid this warning.",
+                    speaker_activity.shape[1],
+                    self.num_speakers,
+                )
+                speaker_activity = speaker_activity[:, : self.num_speakers]
+            speaker_activities.append(speaker_activity)
+
+        cuts = type(cuts).from_cuts(mono_cuts)
+        audio, audio_lens, cuts = self.load_audio(cuts)
+        targets = collate_matrices(speaker_activities).to(audio.dtype)  # (B, T, N)
+
+        if targets.shape[2] > self.num_speakers:
+            targets = targets[:, :, : self.num_speakers]
+        elif targets.shape[2] < self.num_speakers:
+            targets = torch.nn.functional.pad(
+                targets, (0, self.num_speakers - targets.shape[2]), mode='constant', value=0
+            )
+
+        target_lens_list = []
+        for audio_len in audio_lens:
+            target_fr_len = get_hidden_length_from_sample_length(
+                audio_len, self.num_sample_per_mel_frame, self.num_mel_frame_per_target_frame
+            )
+            target_lens_list.append(target_fr_len)
+        target_lens = torch.tensor(target_lens_list)
+
+        return audio, audio_lens, targets, target_lens

nemo/collections/asr/data/audio_to_eou_label_lhotse.py

@@ -0,0 +1,524 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+from dataclasses import dataclass
+from typing import Dict, List, Optional
+
+import numpy as np
+import torch.utils.data
+from lhotse.cut import Cut, CutSet, MixedCut
+from lhotse.dataset import AudioSamples
+from lhotse.dataset.collation import collate_vectors
+from omegaconf import DictConfig, OmegaConf
+
+from nemo.collections.asr.parts.preprocessing.perturb import process_augmentations
+from nemo.collections.asr.parts.preprocessing.segment import AudioSegment
+from nemo.collections.common.tokenizers.aggregate_tokenizer import TokenizerWrapper
+from nemo.collections.common.tokenizers.tokenizer_spec import TokenizerSpec
+from nemo.core.neural_types import AudioSignal, LabelsType, LengthsType, NeuralType
+from nemo.utils import logging
+
+NON_SPEECH_LABEL = 0
+SPEECH_LABEL = 1
+EOU_LABEL = 2
+EOB_LABEL = 3
+EOU_STRING = '<EOU>'
+EOB_STRING = '<EOB>'
+
+# These augmentations are not supported yet, since they will need to change the SOU/EOU timestamps
+EOU_INVALID_AUGMENTATIONS = ['random_segment', 'speed', 'time_stretch']
+
+
+@dataclass
+class AudioToTextEOUBatch:
+    """
+    Data class for ASR-EOU batch.
+    """
+
+    sample_ids: List | None = None
+    audio_filepaths: List | None = None
+    audio_signal: torch.Tensor | None = None
+    audio_lengths: torch.Tensor | None = None
+    text_tokens: torch.Tensor | None = None
+    text_token_lengths: torch.Tensor | None = None
+    eou_targets: torch.Tensor | None = None
+    eou_target_lengths: torch.Tensor | None = None
+
+
+@dataclass
+class RandomPaddingConfig:
+    prob: float = 0.9  # probability of applying padding
+    min_pad_duration: float = 0.0  # minimum duration of pre/post padding in seconds
+    max_pad_duration: float = 5.0  # maximum duration of pre/post padding in seconds
+    max_total_duration: float = 40.0  # maximum total duration of the padded audio in seconds
+    min_pre_pad_duration: float = 0.0  # minimum duration of pre-padding in seconds
+    min_post_pad_duration: float = 2.0  # minimum duration of post-padding in seconds
+    pad_distribution: str = 'uniform'  # distribution of padding duration, 'uniform' or 'normal' or 'constant'
+    normal_mean: float = 0.5  # mean of normal distribution for padding duration
+    normal_std: float = 2.0  # standard deviation of normal distribution for padding duration
+    pre_pad_duration: float = 0.2  # amount of left-padding when pad_distribution='constant'
+    post_pad_duration: float = 3.0  # amount of right-padding when pad_distribution='constant'
+
+
+class LhotseSpeechToTextBpeEOUDataset(torch.utils.data.Dataset):
+    """
+    This dataset processes the audio data and the corresponding text data to generate the ASR labels,
+    along with EOU labels for each frame. The audios used in this dataset should only contain speech with
+    NO precedding or following silence. The dataset also randomly pads non-speech frames before and after
+    the audio signal for training EOU prediction task.
+
+    To generate EOU labels, the last frame of utterance will be marked as "end of utterance" (labeled as `2`),
+    while if it's a backchannel utterance it'll be marked asd "end of backchannel" (labeled as `3`).
+    The rest of the speech frames will be marked as "speech" (labeled as `1`).
+    The padded non-speech signals will be marked as "non-speech" (labeled as 0).
+
+    Args:
+        cfg: DictConfig object container following keys, usually taken from your `model.train_ds`
+            or `model.validation_ds` config:
+        ```
+            sample_rate: # int, Sample rate of the audio signal
+            window_stride: # float, Window stride for audio encoder
+            subsampling_factor: # Subsampling factor for audio encoder
+            random_padding:  # Random padding configuration
+                prob: 0.9  # probability of applying padding
+                min_pad_duration: 0.5  # minimum duration of pre/post padding in seconds
+                max_pad_duration: 2.0 # maximum duration of pre/post padding in seconds
+                max_total_duration: 30.0  # maximum total duration of the padded audio in seconds
+                pad_distribution: 'uniform'  # distribution of padding duration, 'uniform' or 'normal' or 'constant'
+                normal_mean: 0.5  # mean of normal distribution for padding duration
+                normal_std: 2.0  # standard deviation of normal distribution for padding duration
+                pre_pad_duration: 0.2 # amount of left-padding when pad_distribution='constant'
+                post_pad_duration: 3.0  # amount of right-padding when pad_distribution='constant'
+        ```
+
+    Returns:
+        audio: torch.Tensor of audio signal
+        audio_lens: torch.Tensor of audio signal length
+        text_tokens: torch.Tensor of text text_tokens
+        text_token_lens: torch.Tensor of text token length
+        eou_targets (optional): torch.Tensor of EOU labels
+        eou_target_lens (optional): torch.Tensor of EOU label length
+
+    The input manifest should be a jsonl file where each line is a python dictionary.
+    Example manifest sample:
+    {
+        "audio_filepath": "/path/to/audio.wav",
+        "offset": 0.0,
+        "duration": 6.0,
+        "sou_time": [0.3, 4.0],
+        "eou_time": [1.3, 4.5],
+        "utterances": ["Tell me a joke", "Ah-ha"],
+        "is_backchannel": [False, True],
+    }
+
+    Padding logic:
+    0. Don't pad when `random_padding` is None or during validation/test
+    1. randomly draw a probability to decide whether to apply padding
+    2. if not padding or audio duration is longer than the maximum duration,
+        1) return the original audio and EOU labels
+    3. if apply padding,
+        1) get the max padding duration based on the maximum total duration and the audio duration
+        2) randomly draw a total padding duration based on the given distribution
+        3) randomly split the total padding duration into pre-padding and post-padding
+        4) randomly generate the non-speech signal (audio signal=0) for pre-padding and post-padding
+        5) concatenate the pre-padding, audio, and post-padding to get the padded audio signal
+        6) update the EOU labels accordingly
+
+    """
+
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        """Define the output types of the dataset."""
+        return {
+            'audio': NeuralType(('B', 'T'), AudioSignal()),
+            'audio_lens': NeuralType(tuple('B'), LengthsType()),
+            'eou_targets': NeuralType(('B', 'T'), LabelsType()),
+            'eou_target_lens': NeuralType(tuple('B'), LengthsType()),
+            'text_tokens': NeuralType(tuple('B', 'T'), LengthsType(), optional=True),
+            'text_token_lens': NeuralType(tuple('B'), LengthsType(), optional=True),
+        }
+
+    def __init__(self, cfg: DictConfig, tokenizer: TokenizerSpec, return_cuts: bool = False):
+        super().__init__()
+        self.cfg = cfg
+        self.return_cuts = return_cuts
+        self.eou_string = self.cfg.get('eou_string', EOU_STRING)
+        self.eob_string = self.cfg.get('eob_string', EOB_STRING)
+        if cfg.get('check_tokenizer', True):
+            self._check_special_tokens(tokenizer)
+
+        self.tokenizer = TokenizerWrapper(tokenizer)
+        self.load_audio = AudioSamples(fault_tolerant=True)
+        self.sample_rate = self.cfg.get('sample_rate', 16000)
+        self.window_stride = self.cfg.get('window_stride', 0.01)
+        self.num_sample_per_mel_frame = int(
+            self.window_stride * self.sample_rate
+        )  # 160 samples for every 1ms by default
+        self.num_mel_frame_per_target_frame = int(self.cfg.get('subsampling_factor', 8))
+        self.add_sep_before_eou = self.cfg.get('add_sep_before_eou', False)
+        self.add_eou_to_text = self.cfg.get('add_eou_to_text', True)
+        self.pad_eou_label_secs = self.cfg.get('pad_eou_label_secs', 0.0)
+        self.padding_cfg = self.cfg.get('random_padding', None)
+        if self.padding_cfg is not None:
+            self.padding_cfg = OmegaConf.to_container(self.padding_cfg, resolve=True)
+            self.padding_cfg = RandomPaddingConfig(**self.padding_cfg)
+        self.ignore_eob_label = self.cfg.get('ignore_eob_label', False)
+        self.augmentor = None
+        if self.cfg.get('augmentor', None) is not None:
+            augmentor = {}
+            aug_cfg = OmegaConf.to_container(self.cfg.augmentor, resolve=True)
+            for k, v in aug_cfg.items():
+                if k in EOU_INVALID_AUGMENTATIONS:
+                    logging.warning(f"EOU dataset does not support {k} augmentation yet, skipping.")
+                    continue
+                augmentor[k] = v
+
+            if len(augmentor) > 0:
+                logging.info(f"EOU dataset will apply augmentations: {augmentor}")
+                self.augmentor = process_augmentations(augmentor)
+
+    def _check_special_tokens(self, tokenizer: TokenizerSpec):
+        """
+        Check if the special tokens are in the tokenizer vocab.
+        """
+        special_tokens = set([self.eou_string, self.eob_string])
+        vocab_size = tokenizer.vocab_size
+        special_tokens_in_vocab = set([tokenizer.ids_to_text(vocab_size - 1), tokenizer.ids_to_text(vocab_size - 2)])
+        if special_tokens != special_tokens_in_vocab:
+            raise ValueError(
+                f"Input special tokens {special_tokens} don't match with the tokenizer vocab {special_tokens_in_vocab}. "
+                f"Please add them to tokenizer or change input `eou_string` and/or `eob_string` accordingly. "
+                "Special tokens should be added as the last two tokens in the new tokenizer. "
+                "Please refer to scripts/asr_end_of_utterance/tokenizers/add_special_tokens_to_sentencepiece.py for details."
+            )
+
+    def __getitem__(self, cuts: CutSet) -> AudioToTextEOUBatch:
+        audio, audio_lens, cuts = self.load_audio(cuts)
+        audio_signals = []
+        audio_lengths = []
+        eou_targets = []
+        text_tokens = []
+        sample_ids = []
+        audio_filepaths = []
+
+        for i in range(len(cuts)):
+            c = cuts[i]
+            if isinstance(c, MixedCut):
+                c = c.first_non_padding_cut
+
+            sample_ids.append(c.id)
+            audio_filepaths.append(c.recording.sources[0].source)
+
+            audio_i = audio[i]
+            audio_len_i = audio_lens[i]
+
+            # Get EOU labels and text tokens
+            eou_targets_i = self._get_frame_labels(c, audio_len_i)
+            text_tokens_i = self._get_text_tokens(c)
+
+            # Maybe apply random padding to both sides of the audio
+            audio_i, audio_len_i, eou_targets_i = self._random_pad_audio(audio_i, audio_len_i, eou_targets_i)
+
+            # Maybe apply augmentations to the audio signal after padding
+            audio_i, audio_len_i = self._maybe_augment_audio(audio_i, audio_len_i)
+
+            # Append the processed audio, EOU labels, and text tokens to the lists
+            audio_signals.append(audio_i)
+            audio_lengths.append(audio_len_i)
+            eou_targets.append(eou_targets_i)
+            text_tokens.append(text_tokens_i)
+
+        audio_signals = collate_vectors(audio_signals, padding_value=0)
+        audio_lengths = torch.tensor(audio_lengths, dtype=torch.long)
+        eou_target_lens = torch.tensor([t.size(0) for t in eou_targets], dtype=torch.long)
+        eou_targets = collate_vectors(eou_targets, padding_value=0)
+        text_token_lens = torch.tensor([t.size(0) for t in text_tokens], dtype=torch.long)
+        text_tokens = collate_vectors(text_tokens, padding_value=0)
+
+        if self.return_cuts:
+            return audio_signals, audio_lengths, cuts
+
+        return AudioToTextEOUBatch(
+            sample_ids=sample_ids,
+            audio_filepaths=audio_filepaths,
+            audio_signal=audio_signals,
+            audio_lengths=audio_lengths,
+            text_tokens=text_tokens,
+            text_token_lengths=text_token_lens,
+            eou_targets=eou_targets,
+            eou_target_lengths=eou_target_lens,
+        )
+
+    def _audio_len_to_frame_len(self, num_samples: int):
+        """
+        Convert the raw audio length to the number of frames after audio encoder.
+
+        self.num_sample_per_mel_frame = int(
+            self.cfg.get('window_stride', 0.01) * self.cfg.get('sample_rate', 16000)
+        )  # 160 samples for every 1ms by default
+        self.num_mel_frame_per_target_frame = int(self.cfg.get('subsampling_factor', 8))
+        """
+        mel_frame_count = math.ceil((num_samples + 1) / self.num_sample_per_mel_frame)
+        hidden_length = math.ceil(mel_frame_count / self.num_mel_frame_per_target_frame)
+        return hidden_length
+
+    def _repeat_eou_labels(self, eou_targets: torch.Tensor) -> torch.Tensor:
+        """
+        Repeat EOU labels according to self.pad_eou_label_secs
+        Args:
+            eou_targets: torch.Tensor of EOU labels, shape [T]
+        Returns:
+            eou_targets: torch.Tensor of padded EOU labels, shape [T]
+        """
+        if not self.pad_eou_label_secs or self.pad_eou_label_secs <= 0:
+            return eou_targets
+
+        eou_len = self._audio_len_to_frame_len(int(self.pad_eou_label_secs * self.sample_rate))
+
+        i = 0
+        while i < eou_targets.size(0):
+            if eou_targets[i] == EOU_LABEL or eou_targets[i] == EOB_LABEL:
+                # repeat the label for the next eou_len samples
+                start = i
+                end = min(i + eou_len, eou_targets.size(0))
+                j = start + 1
+                while j < end:
+                    if eou_targets[j] != NON_SPEECH_LABEL:
+                        # do not overwrite the label if it's not non-speech
+                        break
+                    j += 1
+                end = min(j, end)
+                # fill the non-speech label with the current EOU/EOB label
+                eou_targets[start:end] = eou_targets[i]
+                i = end
+            else:
+                i += 1
+        return eou_targets
+
+    def _get_frame_labels(self, cut: Cut, num_samples: int):
+        """
+        Get the frame-level EOU labels for a single audio segment.
+        Args:
+            cut: Cut object
+            num_samples: int, the number of samples in the audio segment
+        Returns:
+            eou_targets: torch.Tensor of EOU labels, shape [T]
+        """
+        hidden_length = self._audio_len_to_frame_len(num_samples)
+        if not "sou_time" in cut.custom or not "eou_time" in cut.custom:
+            # assume only single speech segment
+            text = cut.supervisions[0].text
+            if not text:
+                # skip empty utterances
+                return torch.zeros(hidden_length).long()
+            eou_targets = torch.ones(hidden_length).long()  # speech label
+            eou_targets[-1] = EOU_LABEL  # by default it's end of utterance
+            if cut.has_custom("is_backchannel") and cut.custom["is_backchannel"] and not self.ignore_eob_label:
+                eou_targets[-1] = EOB_LABEL  # end of backchannel
+            return eou_targets
+
+        sou_time = cut.custom["sou_time"]
+        eou_time = cut.custom["eou_time"]
+        if not isinstance(sou_time, list):
+            sou_time = [sou_time]
+        if not isinstance(eou_time, list):
+            eou_time = [eou_time]
+
+        assert len(sou_time) == len(
+            eou_time
+        ), f"Number of SOU time and EOU time do not match: SOU ({sou_time}) vs EOU ({eou_time})"
+
+        if cut.has_custom("is_backchannel"):
+            is_backchannel = cut.custom["is_backchannel"]
+            if not isinstance(is_backchannel, list):
+                is_backchannel = [is_backchannel]
+            assert len(sou_time) == len(
+                is_backchannel
+            ), f"Number of SOU and backchannel do not match: SOU ({len(sou_time)}) vs backchannel ({len(is_backchannel)})"
+        else:
+            is_backchannel = [False] * len(sou_time)
+
+        eou_targets = torch.zeros(hidden_length).long()
+        for i in range(len(sou_time)):
+            if sou_time[i] is None or eou_time[i] is None or sou_time[i] < 0 or eou_time[i] < 0:
+                # skip empty utterances
+                continue
+            sou_idx = self._audio_len_to_frame_len(int((sou_time[i] - cut.start) * self.sample_rate))
+            seg_len_in_secs = eou_time[i] - sou_time[i]
+            seg_len = self._audio_len_to_frame_len(int(seg_len_in_secs * self.sample_rate))
+            eou_targets[sou_idx : sou_idx + seg_len] = SPEECH_LABEL
+            last_idx = min(sou_idx + seg_len - 1, hidden_length - 1)
+            if is_backchannel[i] and not self.ignore_eob_label:
+                eou_targets[last_idx] = EOB_LABEL  # end of backchannel
+            else:
+                eou_targets[last_idx] = EOU_LABEL  # end of utterance
+
+        return eou_targets
+
+    def _get_text_tokens(self, cut: Cut):
+        """
+        Add EOU labels to the text and get the text tokens for a single audio segment.
+        Args:
+            cut: Cut object
+        Returns:
+            text_tokens: torch.Tensor of text tokens, shape [T]
+        """
+        if not cut.has_custom("sou_time") or not cut.has_custom("eou_time") or not cut.has_custom("utterances"):
+            # assume only single speech segment
+            utterances = [cut.supervisions[0].text]
+        else:
+            utterances = cut.custom["utterances"]
+
+        if not isinstance(utterances, list):
+            utterances = [utterances]
+
+        if cut.has_custom("is_backchannel"):
+            is_backchannel = cut.custom["is_backchannel"]
+            if not isinstance(is_backchannel, list):
+                is_backchannel = [is_backchannel]
+            assert len(utterances) == len(
+                is_backchannel
+            ), f"Number of utterances and backchannel do not match: utterance ({len(utterances)}) vs backchannel ({len(is_backchannel)})"
+        else:
+            is_backchannel = [False] * len(utterances)
+
+        total_text = ""
+        for i, text in enumerate(utterances):
+            if not text:
+                # skip empty utterances
+                continue
+            if self.add_eou_to_text:
+                eou_string = self.eob_string if is_backchannel[i] and not self.ignore_eob_label else self.eou_string
+                if self.add_sep_before_eou:
+                    eou_string = " " + eou_string
+            else:
+                eou_string = ""
+            total_text += text + eou_string + " "
+        total_text = total_text.strip()
+        return torch.as_tensor(self.tokenizer(total_text))
+
+    def _random_pad_audio(self, audio: torch.Tensor, audio_len: torch.Tensor, eou_targets: torch.Tensor):
+        """
+        Randomly pad the audio signal with non-speech signal before and after the audio signal.
+        Args:
+            audio: torch.Tensor of a single audio signal, shape [T]
+            audio_len: torch.Tensor of audio signal length, shape [1]
+            eou_targets: torch.Tensor of EOU labels, shape [T]
+        Returns:
+            padded_audio: torch.Tensor of padded audio signal, shape [T+padding]
+            padded_audio_len: torch.Tensor of padded audio signal length, shape [1]
+            padded_eou_targets: torch.Tensor of padded EOU labels, shape [T+padding]
+            padded_eou_targets_len: torch.Tensor of padded EOU label length, shape [1]
+        """
+        p = np.random.rand()
+        if self.padding_cfg is None or p > self.padding_cfg.prob:
+            # don't apply padding
+            eou_targets = self._repeat_eou_labels(eou_targets)
+            return audio, audio_len, eou_targets
+
+        duration = audio_len.item() / self.cfg.sample_rate
+        # if already longer than the maximum duration, return the original audio
+        if duration >= self.padding_cfg.max_total_duration:
+            return audio, audio_len, eou_targets
+
+        # apply padding
+        audio = audio[:audio_len]
+
+        self.padding_cfg.min_pre_pad_duration = max(
+            self.padding_cfg.min_pre_pad_duration, self.padding_cfg.min_pad_duration
+        )
+        self.padding_cfg.min_post_pad_duration = max(
+            self.padding_cfg.min_post_pad_duration, self.padding_cfg.min_pad_duration
+        )
+
+        max_padding_duration = max(0, self.padding_cfg.max_total_duration - duration)
+        if max_padding_duration <= self.padding_cfg.min_pre_pad_duration + self.padding_cfg.min_post_pad_duration:
+            min_padding_duration = 0
+        else:
+            min_padding_duration = self.padding_cfg.min_pre_pad_duration + self.padding_cfg.min_post_pad_duration
+
+        pre_padding_duration = None
+        post_padding_duration = None
+
+        if self.padding_cfg.pad_distribution == 'uniform':
+            total_padding_duration = np.random.uniform(min_padding_duration, max_padding_duration)
+        elif self.padding_cfg.pad_distribution == 'normal':
+            total_padding_duration = np.random.normal(self.padding_cfg.normal_mean, self.padding_cfg.normal_std)
+            total_padding_duration = max(min_padding_duration, min(max_padding_duration, total_padding_duration))
+        elif self.padding_cfg.pad_distribution == 'constant':
+            pass
+        else:
+            raise ValueError(f"Unknown padding distribution: {self.padding_cfg.pad_distribution}")
+
+        if self.padding_cfg.pad_distribution == 'constant':
+            pre_padding_duration = self.padding_cfg.pre_pad_duration
+            post_padding_duration = self.padding_cfg.post_pad_duration
+        elif min_padding_duration == 0:
+            pre_padding_duration = total_padding_duration / 2
+            post_padding_duration = total_padding_duration / 2
+        else:
+            post_padding_duration = np.random.uniform(
+                self.padding_cfg.min_post_pad_duration, total_padding_duration - self.padding_cfg.min_pre_pad_duration
+            )
+            pre_padding_duration = total_padding_duration - post_padding_duration
+
+        if self.padding_cfg.max_pad_duration is not None:
+            pre_padding_duration = min(pre_padding_duration, self.padding_cfg.max_pad_duration)
+            post_padding_duration = min(post_padding_duration, self.padding_cfg.max_pad_duration)
+
+        pre_padding_len = math.ceil(pre_padding_duration * self.cfg.sample_rate)
+        post_padding_len = math.ceil(post_padding_duration * self.cfg.sample_rate)
+
+        # pad the audio signal
+        pre_padding = torch.zeros(pre_padding_len, dtype=audio.dtype)
+        post_padding = torch.zeros(post_padding_len, dtype=audio.dtype)
+        padded_audio = torch.cat((pre_padding, audio, post_padding), dim=0)
+        padded_audio_len = audio_len + pre_padding_len + post_padding_len
+
+        # pad the EOU labels
+        pre_padding_eou_len = self._audio_len_to_frame_len(pre_padding_len)
+        post_padding_eou_len = self._audio_len_to_frame_len(post_padding_len)
+        pre_padding_eou = torch.zeros(pre_padding_eou_len, dtype=eou_targets.dtype)
+        post_padding_eou = torch.zeros(post_padding_eou_len, dtype=eou_targets.dtype)
+        padded_eou_targets = torch.cat((pre_padding_eou, eou_targets, post_padding_eou), dim=0)
+        padded_eou_targets = self._repeat_eou_labels(padded_eou_targets)
+        return padded_audio, padded_audio_len, padded_eou_targets
+
+    def _maybe_augment_audio(self, audio: torch.Tensor, audio_len: torch.Tensor):
+        """
+        Apply augmentation to the audio signal if augmentor is provided.
+        Args:
+            audio: torch.Tensor of a single audio signal, shape [T]
+            audio_len: torch.Tensor of audio signal length, shape [1]
+        Returns:
+            augmented_audio: torch.Tensor of augmented audio signal, shape [T]
+            augmented_audio_len: torch.Tensor of augmented audio signal length, shape [1]
+        """
+        if self.augmentor is None:
+            return audio, audio_len
+
+        # Cast to AudioSegment
+        audio_segment = AudioSegment(
+            samples=audio[:audio_len].numpy(),
+            sample_rate=self.sample_rate,
+            offset=0,
+            duration=audio_len.item() / self.sample_rate,
+        )
+        # Apply augmentation
+        self.augmentor.perturb(audio_segment)
+        audio = torch.from_numpy(audio_segment.samples).float()
+        audio_len = audio.size(0)
+
+        return audio, audio_len

nemo/collections/asr/data/audio_to_label_dataset.py

@@ -0,0 +1,316 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import copy
+
+from omegaconf import DictConfig
+
+from nemo.collections.asr.data import audio_to_label
+from nemo.collections.asr.data.audio_to_text_dataset import convert_to_config_list, get_chain_dataset
+from nemo.collections.asr.parts.preprocessing.perturb import process_augmentations
+from nemo.collections.common.data.dataset import ConcatDataset
+
+
+def get_classification_label_dataset(featurizer, config: dict) -> audio_to_label.AudioToClassificationLabelDataset:
+    """
+    Instantiates a Classification AudioLabelDataset.
+
+    Args:
+        config: Config of the AudioToClassificationLabelDataset.
+
+    Returns:
+        An instance of AudioToClassificationLabelDataset.
+    """
+    dataset = audio_to_label.AudioToClassificationLabelDataset(
+        manifest_filepath=config['manifest_filepath'],
+        labels=config['labels'],
+        featurizer=featurizer,
+        max_duration=config.get('max_duration', None),
+        min_duration=config.get('min_duration', None),
+        trim=config.get('trim_silence', False),
+        is_regression_task=config.get('is_regression_task', False),
+        cal_labels_occurrence=config.get('cal_labels_occurrence', False),
+    )
+    return dataset
+
+
+def get_speech_label_dataset(featurizer, config: dict) -> audio_to_label.AudioToSpeechLabelDataset:
+    """
+    Instantiates a Speech Label (e.g. VAD, speaker recognition) AudioLabelDataset.
+
+    Args:
+        config: Config of the AudioToSpeechLabelDataSet.
+
+    Returns:
+        An instance of AudioToSpeechLabelDataset.
+    """
+    dataset = audio_to_label.AudioToSpeechLabelDataset(
+        manifest_filepath=config['manifest_filepath'],
+        labels=config['labels'],
+        featurizer=featurizer,
+        max_duration=config.get('max_duration', None),
+        min_duration=config.get('min_duration', None),
+        trim=config.get('trim_silence', False),
+        window_length_in_sec=config.get('window_length_in_sec', 0.31),
+        shift_length_in_sec=config.get('shift_length_in_sec', 0.01),
+        normalize_audio=config.get('normalize_audio', False),
+        cal_labels_occurrence=config.get('cal_labels_occurrence', False),
+    )
+    return dataset
+
+
+def get_tarred_classification_label_dataset(
+    featurizer, config: dict, shuffle_n: int, global_rank: int, world_size: int
+) -> audio_to_label.TarredAudioToClassificationLabelDataset:
+    """
+    Instantiates a Classification TarredAudioLabelDataset.
+
+    Args:
+        config: Config of the TarredAudioToClassificationLabelDataset.
+        shuffle_n: How many samples to look ahead and load to be shuffled.
+            See WebDataset documentation for more details.
+        global_rank: Global rank of this device.
+        world_size: Global world size in the training method.
+
+    Returns:
+        An instance of TarredAudioToClassificationLabelDataset.
+    """
+    tarred_audio_filepaths = config['tarred_audio_filepaths']
+    manifest_filepaths = config['manifest_filepath']
+    datasets = []
+    tarred_audio_filepaths = convert_to_config_list(tarred_audio_filepaths)
+    manifest_filepaths = convert_to_config_list(manifest_filepaths)
+
+    bucketing_weights = config.get('bucketing_weights', None)  # For upsampling buckets
+    if bucketing_weights:
+        for idx, weight in enumerate(bucketing_weights):
+            if not isinstance(weight, int) or weight <= 0:
+                raise ValueError("bucket weights must be positive integers")
+
+    if len(manifest_filepaths) != len(tarred_audio_filepaths):
+        raise ValueError(
+            f"manifest_filepaths (length={len(manifest_filepaths)}) and tarred_audio_filepaths (length={len(tarred_audio_filepaths)}) need to have the same number of buckets."
+        )
+
+    for dataset_idx, (tarred_audio_filepath, manifest_filepath) in enumerate(
+        zip(tarred_audio_filepaths, manifest_filepaths)
+    ):
+        if len(tarred_audio_filepath) == 1:
+            tarred_audio_filepath = tarred_audio_filepath[0]
+        dataset = audio_to_label.TarredAudioToClassificationLabelDataset(
+            audio_tar_filepaths=tarred_audio_filepath,
+            manifest_filepath=manifest_filepath,
+            labels=config['labels'],
+            featurizer=featurizer,
+            shuffle_n=shuffle_n,
+            max_duration=config.get('max_duration', None),
+            min_duration=config.get('min_duration', None),
+            trim=config.get('trim_silence', False),
+            shard_strategy=config.get('tarred_shard_strategy', 'scatter'),
+            global_rank=global_rank,
+            world_size=world_size,
+            is_regression_task=config.get('is_regression_task', False),
+        )
+
+        if bucketing_weights:
+            [datasets.append(dataset) for _ in range(bucketing_weights[dataset_idx])]
+        else:
+            datasets.append(dataset)
+
+    return get_chain_dataset(datasets=datasets, ds_config=config, rank=global_rank)
+
+
+def get_concat_tarred_speech_label_dataset(
+    featurizer,
+    config: dict,
+    shuffle_n: int,
+    global_rank: int,
+    world_size: int,
+):
+    tarred_audio_filepaths = config['tarred_audio_filepaths']
+    manifest_filepaths = config['manifest_filepath']
+    datasets = []
+    for dataset_idx, (tarred_audio_filepath, manifest_filepath) in enumerate(
+        zip(tarred_audio_filepaths, manifest_filepaths)
+    ):
+        conf = copy.deepcopy(config)
+        conf['manifest_filepath'] = manifest_filepath
+        conf['tarred_audio_filepaths'] = tarred_audio_filepath
+        dataset = get_tarred_speech_label_dataset(
+            config=conf,
+            featurizer=featurizer,
+            shuffle_n=shuffle_n,
+            global_rank=global_rank,
+            world_size=world_size,
+        )
+        datasets.append(dataset)
+
+    dataset = ConcatDataset(
+        datasets,
+        sampling_technique=config.get('concat_sampling_technique', 'temperature'),
+        sampling_temperature=config.get('concat_sampling_temperature', 5),
+        sampling_probabilities=config.get('concat_sampling_probabilities', None),
+        global_rank=global_rank,
+        world_size=world_size,
+        shuffle=config['shuffle'],
+    )
+    return dataset
+
+
+def get_tarred_speech_label_dataset(
+    featurizer,
+    config: dict,
+    shuffle_n: int,
+    global_rank: int,
+    world_size: int,
+) -> audio_to_label.TarredAudioToSpeechLabelDataset:
+    """
+    InInstantiates a Speech Label (e.g. VAD, speaker recognition) TarredAudioLabelDataset.
+
+    Args:
+        config: Config of the TarredAudioToSpeechLabelDataset.
+        shuffle_n: How many samples to look ahead and load to be shuffled.
+            See WebDataset documentation for more details.
+        global_rank: Global rank of this device.
+        world_size: Global world size in the training method.
+
+    Returns:
+        An instance of TarredAudioToSpeechLabelDataset.
+    """
+    tarred_audio_filepaths = config['tarred_audio_filepaths']
+    manifest_filepaths = config['manifest_filepath']
+    datasets = []
+    tarred_audio_filepaths = convert_to_config_list(tarred_audio_filepaths)
+    manifest_filepaths = convert_to_config_list(manifest_filepaths)
+
+    bucketing_weights = config.get('bucketing_weights', None)  # For upsampling buckets
+    if bucketing_weights:
+        for idx, weight in enumerate(bucketing_weights):
+            if not isinstance(weight, int) or weight <= 0:
+                raise ValueError("bucket weights must be positive integers")
+
+    if len(manifest_filepaths) != len(tarred_audio_filepaths):
+        raise ValueError(
+            f"manifest_filepaths (length={len(manifest_filepaths)}) and tarred_audio_filepaths (length={len(tarred_audio_filepaths)}) need to have the same number of buckets."
+        )
+
+    for dataset_idx, (tarred_audio_filepath, manifest_filepath) in enumerate(
+        zip(tarred_audio_filepaths, manifest_filepaths)
+    ):
+        if len(tarred_audio_filepath) == 1:
+            tarred_audio_filepath = tarred_audio_filepath[0]
+        dataset = audio_to_label.TarredAudioToSpeechLabelDataset(
+            audio_tar_filepaths=tarred_audio_filepath,
+            manifest_filepath=manifest_filepath,
+            labels=config['labels'],
+            featurizer=featurizer,
+            shuffle_n=shuffle_n,
+            max_duration=config.get('max_duration', None),
+            min_duration=config.get('min_duration', None),
+            trim=config.get('trim_silence', False),
+            window_length_in_sec=config.get('window_length_in_sec', 8),
+            shift_length_in_sec=config.get('shift_length_in_sec', 0.075),
+            normalize_audio=config.get('normalize_audio', False),
+            shard_strategy=config.get('tarred_shard_strategy', 'scatter'),
+            global_rank=global_rank,
+            world_size=world_size,
+        )
+
+        if bucketing_weights:
+            [datasets.append(dataset) for _ in range(bucketing_weights[dataset_idx])]
+        else:
+            datasets.append(dataset)
+
+    return get_chain_dataset(datasets=datasets, ds_config=config, rank=global_rank)
+
+
+def get_audio_multi_label_dataset(cfg: DictConfig) -> audio_to_label.AudioToMultiLabelDataset:
+    if "augmentor" in cfg:
+        augmentor = process_augmentations(cfg.augmentor)
+    else:
+        augmentor = None
+
+    dataset = audio_to_label.AudioToMultiLabelDataset(
+        manifest_filepath=cfg.get("manifest_filepath"),
+        sample_rate=cfg.get("sample_rate"),
+        labels=cfg.get("labels", None),
+        int_values=cfg.get("int_values", False),
+        augmentor=augmentor,
+        min_duration=cfg.get("min_duration", None),
+        max_duration=cfg.get("max_duration", None),
+        trim_silence=cfg.get("trim_silence", False),
+        is_regression_task=cfg.get("is_regression_task", False),
+        cal_labels_occurrence=cfg.get("cal_labels_occurrence", False),
+        delimiter=cfg.get("delimiter", None),
+        normalize_audio_db=cfg.get("normalize_audio_db", None),
+    )
+    return dataset
+
+
+def get_tarred_audio_multi_label_dataset(
+    cfg: DictConfig, shuffle_n: int, global_rank: int, world_size: int
+) -> audio_to_label.TarredAudioToMultiLabelDataset:
+
+    if "augmentor" in cfg:
+        augmentor = process_augmentations(cfg.augmentor)
+    else:
+        augmentor = None
+
+    tarred_audio_filepaths = cfg['tarred_audio_filepaths']
+    manifest_filepaths = cfg['manifest_filepath']
+    datasets = []
+    tarred_audio_filepaths = convert_to_config_list(tarred_audio_filepaths)
+    manifest_filepaths = convert_to_config_list(manifest_filepaths)
+
+    bucketing_weights = cfg.get('bucketing_weights', None)  # For upsampling buckets
+    if bucketing_weights:
+        for idx, weight in enumerate(bucketing_weights):
+            if not isinstance(weight, int) or weight <= 0:
+                raise ValueError("bucket weights must be positive integers")
+
+    if len(manifest_filepaths) != len(tarred_audio_filepaths):
+        raise ValueError(
+            f"manifest_filepaths (length={len(manifest_filepaths)}) and tarred_audio_filepaths (length={len(tarred_audio_filepaths)}) need to have the same number of buckets."
+        )
+
+    for dataset_idx, (tarred_audio_filepath, manifest_filepath) in enumerate(
+        zip(tarred_audio_filepaths, manifest_filepaths)
+    ):
+        if len(tarred_audio_filepath) == 1:
+            tarred_audio_filepath = tarred_audio_filepath[0]
+
+        dataset = audio_to_label.TarredAudioToMultiLabelDataset(
+            audio_tar_filepaths=tarred_audio_filepath,
+            manifest_filepath=manifest_filepath,
+            sample_rate=cfg["sample_rate"],
+            labels=cfg['labels'],
+            shuffle_n=shuffle_n,
+            int_values=cfg.get("int_values", False),
+            augmentor=augmentor,
+            min_duration=cfg.get('min_duration', None),
+            max_duration=cfg.get('max_duration', None),
+            trim_silence=cfg.get('trim_silence', False),
+            is_regression_task=cfg.get('is_regression_task', False),
+            delimiter=cfg.get("delimiter", None),
+            shard_strategy=cfg.get('tarred_shard_strategy', 'scatter'),
+            global_rank=global_rank,
+            world_size=world_size,
+            normalize_audio_db=cfg.get("normalize_audio_db", None),
+        )
+
+        if bucketing_weights:
+            [datasets.append(dataset) for _ in range(bucketing_weights[dataset_idx])]
+        else:
+            datasets.append(dataset)
+
+    return get_chain_dataset(datasets=datasets, ds_config=cfg, rank=global_rank)

nemo/collections/asr/data/audio_to_text_dali.py

@@ -0,0 +1,777 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+import operator
+import os.path
+import time
+from collections.abc import Iterator
+from typing import Callable, List, Optional, Union
+
+import torch
+from omegaconf import DictConfig
+
+from nemo.collections.asr.data.audio_to_text import ASRManifestProcessor, expand_sharded_filepaths
+from nemo.collections.common.parts.preprocessing import parsers
+from nemo.utils import logging, model_utils
+
+try:
+    import nvidia.dali as dali
+    from nvidia.dali.pipeline import Pipeline
+    from nvidia.dali.plugin.pytorch import DALIGenericIterator as DALIPytorchIterator
+    from nvidia.dali.plugin.pytorch import LastBatchPolicy as LastBatchPolicy
+
+    HAVE_DALI = True
+except (ImportError, ModuleNotFoundError):
+    HAVE_DALI = False
+
+__all__ = [
+    'AudioToCharDALIDataset',
+    'AudioToBPEDALIDataset',
+]
+
+"""
+Below minimum version is required to access the "read_idxs" argument in
+dali.fn.readers.nemo_asr
+"""
+__DALI_MINIMUM_VERSION__ = "1.11"
+
+DALI_INSTALLATION_MESSAGE = (
+    "Could not import `nvidia.dali`.\n"
+    "Please install DALI by following the steps provided here - \n"
+    "https://docs.nvidia.com/deeplearning/dali/user-guide/docs/installation.html"
+)
+
+
+def is_dali_supported(min_version: str, verbose: bool = False) -> bool:
+    """
+    Checks if DALI in installed, and version is >= min_verion.
+
+    Args:
+        min_version: A semver str that is the minimum requirement.
+        verbose: Whether to log the installation instructions if DALI is not found.
+
+    Returns:
+        bool - whether DALI could be imported or not.
+    """
+    module_available, _ = model_utils.check_lib_version(
+        'nvidia.dali', checked_version=min_version, operator=operator.ge
+    )
+
+    # If DALI is not installed
+    if module_available is None:
+        if verbose:
+            logging.info(DALI_INSTALLATION_MESSAGE)
+
+        return False
+
+    return module_available
+
+
+class DALIOutputs(object):
+    def __init__(self, out_dict):
+        self._has_processed_signal = 'processed_signal' in out_dict and 'processed_signal_len' in out_dict
+        if not self._has_processed_signal:
+            assert 'audio' in out_dict and 'audio_len' in out_dict
+        assert 'transcript' in out_dict and 'transcript_len' in out_dict
+        if self._has_processed_signal:
+            self._outs = (
+                out_dict['processed_signal'],
+                out_dict['processed_signal_len'].reshape(-1),
+                out_dict['transcript'],
+                out_dict['transcript_len'].reshape(-1),
+            )
+        else:
+            self._outs = (
+                out_dict['audio'],
+                out_dict['audio_len'].reshape(-1),
+                out_dict['transcript'],
+                out_dict['transcript_len'].reshape(-1),
+            )
+
+    @property
+    def has_processed_signal(self):
+        return self._has_processed_signal
+
+    def __getitem__(self, key):
+        return self._outs[key]
+
+    def __len__(self):
+        return len(self._outs)
+
+
+class _AudioTextDALIDataset(Iterator):
+    """
+    NVIDIA DALI pipeline that loads tensors via one or more manifest files where each line containing a sample descriptor in JSON,
+    including audio files, transcripts, and durations (in seconds).
+    Here's an example:
+    {"audio_filepath": "/path/to/audio.wav", "text_filepath": "/path/to/audio.txt", "duration": 23.147}
+    ...
+    {"audio_filepath": "/path/to/audio.wav", "text": "the transcription", "offset": 301.75, "duration": 0.82, "utt":
+    "utterance_id", "ctm_utt": "en_4156", "side": "A"}
+
+    Args:
+        manifest_filepath: Path to manifest file with the format described above. Can be comma-separated paths.
+        device (str): Determines the device type to be used for preprocessing. Allowed values are: 'cpu', 'gpu'.
+        batch_size (int): Number of samples in a batch.
+        parser (str, callable): A str for an inbuilt parser, or a callable with signature f(str) -> List[int].
+        sample_rate (int): Sample rate to resample loaded audio to.
+        num_threads (int): Number of CPU processing threads to be created by the DALI pipeline.
+        max_duration (float): Determines the maximum allowed duration, in seconds, of the loaded audio files.
+        min_duration (float): Determines the minimum allowed duration, in seconds, of the loaded audio files.
+        bos_id (int): Id of beginning of sequence symbol to append if not None
+        eos_id (int): Id of end of sequence symbol to append if not None
+        pad_id (int): Id used to pad the input. Defaults to 0 if not provided.
+        trim (bool): If True, it will extract the nonsilent region of the loaded audio signal.
+        shuffle (bool): If set to True, the dataset will shuffled after loading.
+        drop_last (bool): If set to True, the last batch will be dropped if incomplete. This will be the case when the shard size is not divisible by the batch size.
+                          If set to False and the size of dataset is not divisible by the batch size, then the last batch will be smaller.
+        device_id (int): Index of the GPU to be used (local_rank). Only applicable when device == 'gpu'. Defaults to 0.
+        global_rank (int): Worker rank, used for partitioning shards. Defaults to 0.
+        world_size (int): Total number of processes, used for partitioning shards. Defaults to 1.
+        preprocessor_cfg (DictConfig): Preprocessor configuration. Supports AudioToMelSpectrogramPreprocessor and AudioToMFCCPreprocessor.
+        return_sample_id (bool): whether to return the sample_id as a part of each sample (not supported yet).
+    """
+
+    def __init__(
+        self,
+        manifest_filepath: str,
+        device: str,
+        batch_size: int,
+        parser: Union[str, Callable],
+        audio_tar_filepaths: Optional[Union[str, List[str]]] = None,
+        audio_tar_index_filepaths: Optional[Union[str, List[str]]] = None,
+        sample_rate: int = 16000,
+        num_threads: int = 4,
+        max_duration: float = 0.0,
+        min_duration: float = 0.0,
+        bos_id: Optional[int] = None,
+        eos_id: Optional[int] = None,
+        pad_id: int = 0,
+        trim: bool = False,
+        shuffle: bool = False,
+        drop_last: bool = False,
+        shard_strategy: str = "scatter",
+        device_id: int = 0,
+        global_rank: int = 0,
+        world_size: int = 1,
+        preprocessor_cfg: DictConfig = None,
+        return_sample_id: bool = False,
+    ):
+        self.drop_last = drop_last  # used by lr_scheduler
+        if return_sample_id:
+            raise ValueError(
+                "Currently DALI data layers don't support returning the sample_id and return_sample_id can not be enabled."
+            )
+        self.return_sample_id = return_sample_id
+
+        if not HAVE_DALI:
+            raise ModuleNotFoundError(
+                f"{self} requires NVIDIA DALI to be installed. "
+                f"See: https://docs.nvidia.com/deeplearning/dali/user-guide/docs/installation.html#id1"
+            )
+
+        if device not in ('cpu', 'gpu'):
+            raise ValueError(
+                f"{self} received an unexpected device argument {device}. Supported values are: 'cpu', 'gpu'"
+            )
+
+        device_id = device_id if device == 'gpu' else None
+
+        self.batch_size = batch_size  # Used by NeMo
+
+        self.device = device
+        self.device_id = device_id
+
+        if world_size > 1:
+            self.shard_id = global_rank
+            self.num_shards = world_size
+        else:
+            self.shard_id = None
+            self.num_shards = None
+
+        self.eos_id = eos_id
+        self.bos_id = bos_id
+        self.sample_rate = sample_rate
+
+        self.pipe = Pipeline(
+            batch_size=batch_size,
+            num_threads=num_threads,
+            device_id=self.device_id,
+            exec_async=True,
+            exec_pipelined=True,
+        )
+
+        has_preprocessor = preprocessor_cfg is not None
+        if has_preprocessor:
+            if preprocessor_cfg._target_ == "nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor":
+                feature_type = "mel_spectrogram"
+            elif preprocessor_cfg._target_ == "nemo.collections.asr.modules.AudioToMFCCPreprocessor":
+                feature_type = "mfcc"
+            else:
+                raise ValueError(
+                    f"{self} received an unexpected preprocessor configuration: {preprocessor_cfg._target_}."
+                    f" Supported preprocessors are: AudioToMelSpectrogramPreprocessor, AudioToMFCCPreprocessor"
+                )
+
+            # Default values taken from AudioToMelSpectrogramPreprocessor
+            params = preprocessor_cfg
+            self.dither = params['dither'] if 'dither' in params else 0.0
+            self.preemph = params['preemph'] if 'preemph' in params else 0.97
+            self.window_size_sec = params['window_size'] if 'window_size' in params else 0.02
+            self.window_stride_sec = params['window_stride'] if 'window_stride' in params else 0.01
+            self.sample_rate = params['sample_rate'] if 'sample_rate' in params else sample_rate
+            self.window_size = int(self.window_size_sec * self.sample_rate)
+            self.window_stride = int(self.window_stride_sec * self.sample_rate)
+
+            normalize = params['normalize'] if 'normalize' in params else 'per_feature'
+            if normalize == 'per_feature':  # Each freq channel independently
+                self.normalization_axes = (1,)
+            elif normalize == 'all_features':
+                self.normalization_axes = (0, 1)
+            else:
+                raise ValueError(
+                    f"{self} received {normalize} for the normalize parameter."
+                    f" It must be either 'per_feature' or 'all_features'."
+                )
+
+            self.window = None
+            window_name = params['window'] if 'window' in params else 'hann'
+            torch_windows = {
+                'hann': torch.hann_window,
+                'hamming': torch.hamming_window,
+                'blackman': torch.blackman_window,
+                'bartlett': torch.bartlett_window,
+                'none': None,
+            }
+
+            if window_name == 'ones':
+                window_tensor = torch.ones(self.window_size)
+            else:
+                try:
+                    window_fn = torch_windows.get(window_name, None)
+                except:
+                    raise ValueError(
+                        f"{self} received '{window_name}' for the window parameter."
+                        f" It must be one of: ('hann', 'ones', 'hamming', 'blackman', 'bartlett', None)."
+                        f" None is equivalent to 'hann'."
+                    )
+                window_tensor = window_fn(self.window_size, periodic=False) if window_fn else None
+            self.window = window_tensor.numpy().tolist() if window_tensor is not None else None
+
+            self.n_fft = params['n_fft'] if 'n_fft' in params else 2 ** math.ceil(math.log2(self.window_size))
+            self.n_mels = params['n_mels'] if 'n_mels' in params else 64
+            self.n_mfcc = params['n_mfcc'] if 'n_mfcc' in params else 64
+
+            features = params['features'] if 'features' in params else 0
+            if features > 0:
+                if feature_type == 'mel_spectrogram':
+                    self.n_mels = features
+                elif feature_type == 'mfcc':
+                    self.n_mfcc = features
+
+            # TODO Implement frame splicing
+            if 'frame_splicing' in params:
+                assert params['frame_splicing'] == 1, "Frame splicing is not implemented"
+
+            self.freq_low = params['lowfreq'] if 'lowfreq' in params else 0.0
+            self.freq_high = params['highfreq'] if 'highfreq' in params else self.sample_rate / 2.0
+            self.log_features = params['log'] if 'log' in params else True
+
+            # We want to avoid taking the log of zero
+            # There are two options: either adding or clamping to a small value
+
+            self.log_zero_guard_type = params['log_zero_guard_type'] if 'log_zero_guard_type' in params else 'add'
+            if self.log_zero_guard_type not in ["add", "clamp"]:
+                raise ValueError(
+                    f"{self} received {self.log_zero_guard_type} for the "
+                    f"log_zero_guard_type parameter. It must be either 'add' or "
+                    f"'clamp'."
+                )
+
+            self.log_zero_guard_value = params['log_zero_guard_value'] if 'log_zero_guard_value' in params else 2**-24
+            if isinstance(self.log_zero_guard_value, str):
+                if self.log_zero_guard_value == "tiny":
+                    self.log_zero_guard_value = torch.finfo(torch.float32).tiny
+                elif self.log_zero_guard_value == "eps":
+                    self.log_zero_guard_value = torch.finfo(torch.float32).eps
+                else:
+                    raise ValueError(
+                        f"{self} received {self.log_zero_guard_value} for the log_zero_guard_type parameter."
+                        f"It must be either a number, 'tiny', or 'eps'"
+                    )
+
+            self.mag_power = params['mag_power'] if 'mag_power' in params else 2
+            if self.mag_power != 1.0 and self.mag_power != 2.0:
+                raise ValueError(
+                    f"{self} received {self.mag_power} for the mag_power parameter." f" It must be either 1.0 or 2.0."
+                )
+
+            self.pad_to = max(params['pad_to'], 1) if 'pad_to' in params else 16
+            self.pad_value = params['pad_value'] if 'pad_value' in params else 0.0
+
+        with self.pipe:
+            if audio_tar_filepaths is None and audio_tar_index_filepaths is None:
+                audio, indices = dali.fn.readers.nemo_asr(
+                    name="Reader",
+                    manifest_filepaths=manifest_filepath.split(','),
+                    dtype=dali.types.FLOAT,
+                    downmix=True,
+                    sample_rate=float(self.sample_rate),
+                    min_duration=min_duration,
+                    max_duration=max_duration,
+                    read_sample_rate=False,
+                    read_text=False,
+                    read_idxs=True,
+                    random_shuffle=shuffle,
+                    shard_id=self.shard_id,
+                    num_shards=self.num_shards,
+                    pad_last_batch=True,
+                )
+
+                self.is_tarred_dataset = False
+
+            elif audio_tar_filepaths is not None and audio_tar_index_filepaths is not None:
+                audio_tar_filepaths = expand_sharded_filepaths(
+                    audio_tar_filepaths,
+                    shard_strategy=shard_strategy,
+                    world_size=world_size,
+                    global_rank=global_rank,
+                )
+
+                audio_tar_index_filepaths = expand_sharded_filepaths(
+                    audio_tar_index_filepaths,
+                    shard_strategy=shard_strategy,
+                    world_size=world_size,
+                    global_rank=global_rank,
+                )
+
+                if len(audio_tar_filepaths) != len(audio_tar_index_filepaths) and len(audio_tar_index_filepaths) != 0:
+                    raise ValueError(
+                        f"Number of filepaths provided for `audio_tar_filepaths` must match "
+                        f"`audio_tar_index_filepaths`. Got {len(audio_tar_filepaths)} audio_tar_filepaths and "
+                        f"{len(audio_tar_index_filepaths)} audio_tar_index_filepaths."
+                    )
+
+                tar_file = dali.fn.readers.webdataset(
+                    paths=audio_tar_filepaths,
+                    index_paths=audio_tar_index_filepaths,
+                    name="Reader",
+                    ext=["wav"],
+                    missing_component_behavior="error",
+                    random_shuffle=shuffle,
+                    shard_id=self.shard_id,
+                    num_shards=self.num_shards,
+                    pad_last_batch=True,
+                )
+                audio, _ = dali.fn.decoders.audio(
+                    tar_file,
+                    dtype=dali.types.FLOAT,
+                    downmix=True,
+                    sample_rate=float(self.sample_rate),
+                )
+                indices = dali.fn.get_property(tar_file, key="source_info")
+                indices = dali.fn.pad(indices)
+
+                self.is_tarred_dataset = True
+
+            else:
+                raise RuntimeError(
+                    "When using DALI datasets, either `audio_tar_filepaths` "
+                    "and `audio_tar_index_filepaths` should either both be None (sequential dataset)"
+                    "or provided (tarred dataset)."
+                )
+
+            # Extract nonsilent region, if necessary
+            if trim:
+                # Need to extract non-silent region before moving to the GPU
+                roi_start, roi_len = dali.fn.nonsilent_region(audio, cutoff_db=-60)
+                audio = audio.gpu() if self.device == 'gpu' else audio
+                audio = dali.fn.slice(
+                    audio, roi_start, roi_len, normalized_anchor=False, normalized_shape=False, axes=[0]
+                )
+            else:
+                audio = audio.gpu() if self.device == 'gpu' else audio
+
+            if not has_preprocessor:
+                # No preprocessing, the output is the audio signal
+                audio_len = dali.fn.shapes(dali.fn.reshape(audio, shape=[-1]))
+                audio = dali.fn.pad(audio)
+                self.pipe.set_outputs(audio, audio_len, indices)
+            else:
+                # Additive gaussian noise (dither)
+                if self.dither > 0.0:
+                    gaussian_noise = dali.fn.random.normal(audio)
+                    audio = audio + self.dither * gaussian_noise
+
+                # Preemphasis filter
+                if self.preemph > 0.0:
+                    audio = dali.fn.preemphasis_filter(audio, preemph_coeff=self.preemph, border='zero')
+
+                # Power spectrogram
+                spec = dali.fn.spectrogram(
+                    audio,
+                    nfft=self.n_fft,
+                    window_length=self.window_size,
+                    window_step=self.window_stride,
+                    window_fn=self.window,
+                )
+
+                if feature_type == 'mel_spectrogram' or feature_type == 'mfcc':
+                    # Spectrogram to Mel Spectrogram
+                    spec = dali.fn.mel_filter_bank(
+                        spec,
+                        sample_rate=self.sample_rate,
+                        nfilter=self.n_mels,
+                        normalize=True,
+                        freq_low=self.freq_low,
+                        freq_high=self.freq_high,
+                    )
+                    # Mel Spectrogram to MFCC
+                    if feature_type == 'mfcc':
+                        spec = dali.fn.mfcc(spec, n_mfcc=self.n_mfcc)
+
+                # Logarithm
+                if self.log_zero_guard_type == 'add':
+                    spec = spec + self.log_zero_guard_value
+
+                spec = dali.fn.to_decibels(
+                    spec, multiplier=math.log(10), reference=1.0, cutoff_db=math.log(self.log_zero_guard_value)
+                )
+
+                # Normalization
+                spec = dali.fn.normalize(spec, axes=self.normalization_axes, epsilon=1e-5**2, ddof=1)
+
+                # Extracting the length of the spectrogram
+                spec_len = dali.fn.slice(dali.fn.shapes(spec), 1, 1, axes=(0,))
+
+                # Pads feature dimension to be a multiple of `pad_to` and the temporal dimension to be as big as the largest sample (shape -1)
+                spec = dali.fn.pad(spec, fill_value=self.pad_value, axes=(0, 1), align=(self.pad_to, 1), shape=(1, -1))
+                self.pipe.set_outputs(spec, spec_len, indices)
+
+        x = time.time()
+        # Building DALI pipeline
+        self.pipe.build()
+        y = time.time()
+
+        logging.info(f"Time for pipe.build() : {(y - x)} seconds")
+
+        if has_preprocessor:
+            output_names = ['processed_signal', 'processed_signal_len', 'manifest_indices']
+        else:
+            output_names = ['audio', 'audio_len', 'manifest_indices']
+
+        x = time.time()
+        last_batch_policy = LastBatchPolicy.DROP if drop_last else LastBatchPolicy.PARTIAL
+        self._iter = DALIPytorchIterator(
+            [self.pipe],
+            output_map=output_names,
+            reader_name="Reader",
+            last_batch_policy=last_batch_policy,
+            dynamic_shape=True,
+            auto_reset=True,
+        )
+        y = time.time()
+        logging.info(f"Time for DALIPytorchIterator to initialize : {(y - x)} seconds")
+
+        # TODO come up with a better solution
+        class DummyDataset:
+            def __init__(self, parent):
+                self.parent = parent
+
+            def __len__(self):
+                return self.parent.size
+
+        self.dataset = DummyDataset(self)  # Used by NeMo
+
+        x = time.time()
+        self.manifest_processor = ASRManifestProcessor(
+            manifest_filepath=manifest_filepath,
+            parser=parser,
+            max_duration=max_duration,
+            min_duration=min_duration,
+            max_utts=0,
+            bos_id=bos_id,
+            eos_id=eos_id,
+            pad_id=pad_id,
+            index_by_file_id=self.is_tarred_dataset,
+        )
+        y = time.time()
+        logging.info(f"Time to build nemo manifest processor - {(y - x)} seconds")
+
+    def reset(self):
+        self._iter.reset()
+
+    def __iter__(self):
+        return self
+
+    def next(self):
+        return self.__next__()
+
+    @property
+    def size(self):
+        return self._iter.size
+
+    def __len__(self):
+        return len(self._iter)
+
+    def __next__(self):
+        outputs = self._iter.next()
+        assert len(outputs) == 1
+        dali_out = outputs[0]
+        manifest_indices = dali_out['manifest_indices'].numpy()
+
+        out = {}
+        out_names = ['processed_signal', 'processed_signal_len', 'audio', 'audio_len']
+        for out_name in out_names:
+            if out_name in dali_out:
+                out[out_name] = dali_out[out_name].detach().clone()
+
+        text_tokens = []
+        text_tokens_len = []
+        max_len = 0
+        batch_size = manifest_indices.shape[0]
+        for i, manifest_index in enumerate(manifest_indices):
+
+            if not self.is_tarred_dataset:
+                # Loose-file dataset. Index is integer based.
+                manifest_index = manifest_index[0]
+                text, text_length = self.manifest_processor.process_text_by_id(manifest_index)
+            else:
+                # Tarred-file dataset. Index is filename based.
+                resolved_manifest_indices = manifest_index.tobytes().decode().split(":")
+                resolved_manifest_index = resolved_manifest_indices[2]  # we require just the filename segment
+                resolved_manifest_index = os.path.splitext(resolved_manifest_index)[0]  # we dont need file extension
+                text, text_length = self.manifest_processor.process_text_by_file_id(resolved_manifest_index)
+
+            text_tokens_len.append(text_length)
+            text_tokens.append(text)
+            if text_length > max_len:
+                max_len = text_length
+
+        transcript_out = torch.full([batch_size, max_len], fill_value=self.manifest_processor.pad_id, dtype=torch.long)
+        for i, n in enumerate(text_tokens_len):
+            transcript_out[i, :n] = torch.tensor(text_tokens[i], dtype=torch.long)
+        transcript_len_out = torch.tensor(text_tokens_len, dtype=torch.long)
+
+        out['transcript'] = transcript_out
+        out['transcript_len'] = transcript_len_out
+        return DALIOutputs(out)
+
+
+class AudioToCharDALIDataset(_AudioTextDALIDataset):
+    """
+    Character based NVIDIA DALI pipeline that loads tensors via one or more manifest files where each line containing a
+    sample descriptor in JSON, including audio files, transcripts, and durations (in seconds).
+    Here's an example:
+    {"audio_filepath": "/path/to/audio.wav", "text_filepath": "/path/to/audio.txt", "duration": 23.147}
+    ...
+    {"audio_filepath": "/path/to/audio.wav", "text": "the transcription", "offset": 301.75, "duration": 0.82, "utt":
+    "utterance_id", "ctm_utt": "en_4156", "side": "A"}
+
+    Args:
+        manifest_filepath: Path to manifest file with the format described above. Can be comma-separated paths.
+        device (str): Determines the device type to be used for preprocessing. Allowed values are: 'cpu', 'gpu'.
+        batch_size (int): Number of samples in a batch.
+        labels (List[str]): String containing all the possible characters to map to.
+        sample_rate (int): Sample rate to resample loaded audio to.
+        num_threads (int): Number of CPU processing threads to be created by the DALI pipeline.
+        max_duration (float): Determines the maximum allowed duration, in seconds, of the loaded audio files.
+        min_duration (float): Determines the minimum allowed duration, in seconds, of the loaded audio files.
+        blank_index (int): blank character index, default = -1
+        unk_index (int): unk_character index, default = -1
+        normalize (bool): whether to normalize transcript text (default): True
+        bos_id (int): Id of beginning of sequence symbol to append if not None
+        eos_id (int): Id of end of sequence symbol to append if not None
+        pad_id (int): Id used to pad the input. Defaults to 0 if not provided.
+        trim (bool): If True, it will extract the nonsilent region of the loaded audio signal.
+        shuffle (bool): If set to True, the dataset will shuffled after loading.
+        drop_last (bool): If set to True, the last batch will be dropped if incomplete. This will be the case when the shard size is not divisible by the batch size.
+                          If set to False and the size of dataset is not divisible by the batch size, then the last batch will be smaller.
+        parser (str, callable): A str for an inbuilt parser, or a callable with signature f(str) -> List[int].
+        device_id (int): Index of the GPU to be used (local_rank). Only applicable when device == 'gpu'. Defaults to 0.
+        global_rank (int): Worker rank, used for partitioning shards. Defaults to 0.
+        world_size (int): Total number of processes, used for partitioning shards. Defaults to 1.
+        preprocessor_cfg (DictConfig): Preprocessor configuration. Supports AudioToMelSpectrogramPreprocessor and AudioToMFCCPreprocessor.
+        return_sample_id (bool): whether to return the sample_id as a part of each sample (not supported yet).
+    """
+
+    def __init__(
+        self,
+        manifest_filepath: str,
+        device: str,
+        batch_size: int,
+        labels: Union[str, List[str]],
+        sample_rate: int = 16000,
+        audio_tar_filepaths: Optional[Union[str, List[str]]] = None,
+        audio_tar_index_filepaths: Optional[Union[str, List[str]]] = None,
+        num_threads: int = 4,
+        max_duration: float = 0.0,
+        min_duration: float = 0.0,
+        blank_index: int = -1,
+        unk_index: int = -1,
+        normalize: bool = True,
+        bos_id: Optional[int] = None,
+        eos_id: Optional[int] = None,
+        pad_id: int = 0,
+        trim: bool = False,
+        shuffle: bool = False,
+        drop_last: bool = False,
+        parser: Union[str, Callable] = 'en',
+        shard_strategy: str = "scatter",
+        device_id: int = 0,
+        global_rank: int = 0,
+        world_size: int = 1,
+        preprocessor_cfg: DictConfig = None,
+        return_sample_id: bool = False,
+    ):
+        self.labels = labels
+
+        parser = parsers.make_parser(
+            labels=labels, name=parser, unk_id=unk_index, blank_id=blank_index, do_normalize=normalize
+        )
+
+        super().__init__(
+            manifest_filepath=manifest_filepath,
+            device=device,
+            batch_size=batch_size,
+            audio_tar_filepaths=audio_tar_filepaths,
+            audio_tar_index_filepaths=audio_tar_index_filepaths,
+            sample_rate=sample_rate,
+            num_threads=num_threads,
+            max_duration=max_duration,
+            min_duration=min_duration,
+            bos_id=bos_id,
+            eos_id=eos_id,
+            pad_id=pad_id,
+            trim=trim,
+            shuffle=shuffle,
+            drop_last=drop_last,
+            parser=parser,
+            shard_strategy=shard_strategy,
+            device_id=device_id,
+            global_rank=global_rank,
+            world_size=world_size,
+            preprocessor_cfg=preprocessor_cfg,
+            return_sample_id=return_sample_id,
+        )
+
+
+class AudioToBPEDALIDataset(_AudioTextDALIDataset):
+    """
+    Subword based NVIDIA DALI pipeline that loads tensors via one or more manifest files where each line containing a
+    sample descriptor in JSON, including audio files, transcripts, and durations (in seconds).
+    Here's an example:
+    {"audio_filepath": "/path/to/audio.wav", "text_filepath": "/path/to/audio.txt", "duration": 23.147}
+    ...
+    {"audio_filepath": "/path/to/audio.wav", "text": "the transcription", "offset": 301.75, "duration": 0.82, "utt":
+    "utterance_id", "ctm_utt": "en_4156", "side": "A"}
+
+    Args:
+        manifest_filepath: Path to manifest file with the format described above. Can be comma-separated paths.
+        tokenizer (TokenizerSpec): A TokenizerSpec implementation that wraps a tokenization implementation.
+        device (str): Determines the device type to be used for preprocessing. Allowed values are: 'cpu', 'gpu'.
+        batch_size (int): Number of samples in a batch.
+        sample_rate (int): Sample rate to resample loaded audio to.
+        num_threads (int): Number of CPU processing threads to be created by the DALI pipeline.
+        max_duration (float): Determines the maximum allowed duration, in seconds, of the loaded audio files.
+        min_duration (float): Determines the minimum allowed duration, in seconds, of the loaded audio files.
+        bos_id (int): Id of beginning of sequence symbol to append if not None. Injected from the tokenizer.
+        eos_id (int): Id of end of sequence symbol to append if not None. Injected from the tokenizer.
+        pad_id (int): Id used to pad the input. Defaults to 0 if not provided. Injected from the tokenizer.
+        trim (bool): If True, it will extract the nonsilent region of the loaded audio signal.
+        shuffle (bool): If set to True, the dataset will shuffled after loading.
+        drop_last (bool): If set to True, the last batch will be dropped if incomplete. This will be the case when the shard size is not divisible by the batch size.
+                          If set to False and the size of dataset is not divisible by the batch size, then the last batch will be smaller.
+
+        device_id (int): Index of the GPU to be used (local_rank). Only applicable when device == 'gpu'. Defaults to 0.
+        global_rank (int): Worker rank, used for partitioning shards. Defaults to 0.
+        world_size (int): Total number of processes, used for partitioning shards. Defaults to 1.
+        preprocessor_cfg (DictConfig): Preprocessor configuration. Supports AudioToMelSpectrogramPreprocessor and AudioToMFCCPreprocessor.
+        use_start_end_token (bool): Boolean which dictates whether to add [BOS] and [EOS] tokens to beginning and
+            ending of speech respectively.
+        return_sample_id (bool): whether to return the sample_id as a part of each sample (not supported yet).
+    """
+
+    def __init__(
+        self,
+        manifest_filepath: str,
+        tokenizer: 'nemo.collections.common.tokenizers.TokenizerSpec',
+        device: str,
+        batch_size: int,
+        sample_rate: int = 16000,
+        audio_tar_filepaths: Optional[Union[str, List[str]]] = None,
+        audio_tar_index_filepaths: Optional[Union[str, List[str]]] = None,
+        num_threads: int = 4,
+        max_duration: float = 0.0,
+        min_duration: float = 0.0,
+        trim: bool = False,
+        shuffle: bool = False,
+        drop_last: bool = False,
+        shard_strategy: str = "scatter",
+        device_id: int = 0,
+        global_rank: int = 0,
+        world_size: int = 1,
+        preprocessor_cfg: DictConfig = None,
+        use_start_end_token: bool = True,
+        return_sample_id: bool = False,
+    ):
+
+        if use_start_end_token and hasattr(tokenizer, 'bos_token'):
+            bos_id = tokenizer.bos_id
+        else:
+            bos_id = None
+
+        if use_start_end_token and hasattr(tokenizer, 'eos_token'):
+            eos_id = tokenizer.eos_id
+        else:
+            eos_id = None
+
+        if hasattr(tokenizer, 'pad_token'):
+            pad_id = tokenizer.pad_id
+        else:
+            pad_id = 0
+
+        class TokenizerWrapper:
+            def __init__(self, tokenizer):
+                self._tokenizer = tokenizer
+
+            def __call__(self, text):
+                t = self._tokenizer.text_to_ids(text)
+                return t
+
+        super().__init__(
+            manifest_filepath=manifest_filepath,
+            device=device,
+            batch_size=batch_size,
+            sample_rate=sample_rate,
+            audio_tar_filepaths=audio_tar_filepaths,
+            audio_tar_index_filepaths=audio_tar_index_filepaths,
+            num_threads=num_threads,
+            max_duration=max_duration,
+            min_duration=min_duration,
+            bos_id=bos_id,
+            eos_id=eos_id,
+            pad_id=pad_id,
+            trim=trim,
+            shuffle=shuffle,
+            drop_last=drop_last,
+            parser=TokenizerWrapper(tokenizer),
+            shard_strategy=shard_strategy,
+            device_id=device_id,
+            global_rank=global_rank,
+            world_size=world_size,
+            preprocessor_cfg=preprocessor_cfg,
+            return_sample_id=return_sample_id,
+        )

nemo/collections/asr/data/audio_to_text_lhotse.py

@@ -0,0 +1,89 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+from typing import Dict, Optional, Tuple
+
+import torch.utils.data
+from lhotse.dataset import AudioSamples
+from lhotse.dataset.collation import collate_vectors
+
+from nemo.collections.common.tokenizers.aggregate_tokenizer import TokenizerWrapper
+from nemo.collections.common.tokenizers.tokenizer_spec import TokenizerSpec
+from nemo.core.neural_types import AudioSignal, LabelsType, LengthsType, NeuralType
+
+
+class LhotseSpeechToTextBpeDataset(torch.utils.data.Dataset):
+    """
+    This dataset is based on BPE datasets from audio_to_text.py.
+    Unlike native NeMo datasets, Lhotse dataset defines only the mapping from
+    a CutSet (meta-data) to a mini-batch with PyTorch tensors.
+    Specifically, it performs tokenization, I/O, augmentation, and feature extraction (if any).
+    Managing data, sampling, de-duplication across workers/nodes etc. is all handled
+    by Lhotse samplers instead.
+
+    NOTE:
+    If the environment variable ``USE_AIS_GET_BATCH`` is set to ``true`` (case-insensitive),
+    then batch audio loading from AIStore will be enabled for this dataset. This will use the
+    AISBatchLoader to load the audio from AIStore. This can improve data loading efficiency in some setups.
+    """
+
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        return {
+            'audio_signal': NeuralType(('B', 'T'), AudioSignal()),
+            'a_sig_length': NeuralType(tuple('B'), LengthsType()),
+            'transcripts': NeuralType(('B', 'T'), LabelsType()),
+            'transcript_length': NeuralType(tuple('B'), LengthsType()),
+            'sample_id': NeuralType(tuple('B'), LengthsType(), optional=True),
+        }
+
+    def __init__(self, tokenizer: TokenizerSpec, return_cuts: bool = False):
+        super().__init__()
+        self.tokenizer = TokenizerWrapper(tokenizer)
+        self.use_ais_get_batch = os.environ.get("USE_AIS_GET_BATCH", "False").lower() == "true"
+
+        # Try to use use_batch_loader if available (Lhotse >= 1.32.0)
+        try:
+            self.load_audio = AudioSamples(fault_tolerant=True, use_batch_loader=self.use_ais_get_batch)
+        except TypeError:
+            # Lhotse < 1.32.0 doesn't support use_batch_loader
+            if self.use_ais_get_batch:
+                import logging
+
+                logging.warning(
+                    "AIS batch loading requested but not supported by this Lhotse version. "
+                    "Please upgrade to Lhotse >= 1.32.0"
+                )
+            self.load_audio = AudioSamples(fault_tolerant=True)
+
+        self.return_cuts = return_cuts
+
+    def __getitem__(self, cuts) -> Tuple[torch.Tensor, ...]:
+        audio, audio_lens, cuts = self.load_audio(cuts)
+        tokens = [
+            torch.cat(
+                [
+                    torch.as_tensor(s.tokens if hasattr(s, "tokens") else self.tokenizer(s.text or "", s.language))
+                    for s in c.supervisions
+                ],
+                dim=0,
+            )
+            for c in cuts
+        ]
+        token_lens = torch.tensor([t.size(0) for t in tokens], dtype=torch.long)
+        tokens = collate_vectors(tokens, padding_value=0)
+        if self.return_cuts:
+            return audio, audio_lens, tokens, token_lens, cuts.drop_in_memory_data()
+        return audio, audio_lens, tokens, token_lens

nemo/collections/asr/data/audio_to_text_lhotse_prompt.py

@@ -0,0 +1,177 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import math
+from typing import Dict, Optional, Tuple
+
+import numpy as np
+import torch
+import torch.utils.data
+from lhotse.dataset import AudioSamples
+from lhotse.dataset.collation import collate_matrices, collate_vectors
+
+from nemo.collections.common.tokenizers.aggregate_tokenizer import AggregateTokenizer
+from nemo.collections.common.tokenizers.tokenizer_spec import TokenizerSpec
+from nemo.core.neural_types import AudioSignal, LabelsType, LengthsType, NeuralType
+
+
+class LhotseSpeechToTextBpeDatasetWithPrompt(torch.utils.data.Dataset):
+    """
+    Dataset class for speech-to-text with prompt vectors.
+    Supports both language ID and custom prompts.
+    """
+
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        return {
+            'audio_signal': NeuralType(('B', 'T'), AudioSignal()),
+            'audio_signal_length': NeuralType(tuple('B'), LengthsType()),
+            'transcripts': NeuralType(('B', 'T'), LabelsType()),
+            'transcript_length': NeuralType(tuple('B'), LengthsType()),
+            'prompt': NeuralType(('B', 'T', 'D'), LabelsType()),
+        }
+
+    def __init__(self, tokenizer, cfg):
+        super().__init__()
+        self.tokenizer = TokenizerWrapper(tokenizer)
+        self.load_audio = AudioSamples(fault_tolerant=True)
+        self.cfg = cfg
+
+        # Calculate num_sample_per_mel_frame from config
+        sample_rate = cfg.get('sample_rate', 16000)
+        window_stride = cfg.get('window_stride', 0.01)
+        self.num_sample_per_mel_frame = int(sample_rate * window_stride)
+
+        self.subsampling_factor = cfg.get('subsampling_factor', 8)
+
+        # Load prompt dictionary from config if provided
+        self.prompt_dict = cfg.get('prompt_dictionary')
+        if self.prompt_dict:
+            # Set num_prompts based on the length of prompt_dictionary or a minimum value
+            # This ensures we have enough dimensions in our embedding space to add scale up without changing the model
+            self.num_prompts = cfg.get('num_prompts', 128)
+
+        # Field to use for prompt key (default to 'language')
+        self.prompt_field = cfg.get('prompt_field', 'language')
+
+    def _get_prompt_index(self, prompt_key: str) -> int:
+        """
+        Maps prompt keys to indices using the prompt dictionary.
+        """
+        if not self.prompt_dict:
+            raise ValueError("Prompt dictionary is empty. Please provide a valid prompt_dictionary in the config.")
+
+        if prompt_key not in self.prompt_dict:
+            available_keys = list(self.prompt_dict.keys())
+            raise ValueError(
+                f"Unknown prompt key: '{prompt_key}'. Available prompts: {available_keys[:10]}{'...' if len(available_keys) > 10 else ''}"
+            )
+
+        return self.prompt_dict[prompt_key]
+
+    def prompt_to_target(self, cut, num_prompts: int, window_stride: int, subsampling_factor: int):
+        """
+        Create prompt target tensor for the sequence.
+        """
+        # Calculate encoder output length based on subsampling factor
+        encoder_hidden_len = self.get_hidden_length_from_sample_length(cut.num_samples)
+
+        # Initialize prompt target matrix
+        mask = np.zeros((num_prompts, encoder_hidden_len))
+
+        # Get prompt index - default to language if prompt not specified
+        # revise supervisions to include prompt key
+        # prompt_key = getattr(cut.supervisions[0].custom_fields, cut.supervisions[0].language)cut.supervisions[0].custom_fields,
+        prompt_id = self._get_prompt_index(cut.supervisions[0].language)
+
+        # Set the corresponding prompt ID to 1 for all time steps
+        mask[prompt_id, :] = 1
+
+        return mask
+
+    def get_hidden_length_from_sample_length(self, num_samples: int) -> int:
+        """
+        Calculate the hidden length from the given number of samples.
+
+        Parameters:
+            num_samples (int): The total number of audio samples.
+
+        Returns:
+            hidden_length (int): The calculated hidden length in terms of the number of frames.
+        """
+        mel_frame_count = math.ceil((num_samples + 1) / self.num_sample_per_mel_frame)
+        hidden_length = math.ceil(mel_frame_count / self.subsampling_factor)
+        return int(hidden_length)
+
+    def __getitem__(self, cuts) -> Tuple[torch.Tensor, ...]:
+        audio, audio_lens, cuts = self.load_audio(cuts)
+        tokens = [torch.as_tensor(self.tokenizer(c.supervisions[0].text, c.supervisions[0].language)) for c in cuts]
+
+        # Create prompt targets
+        prompt_targets = [
+            torch.transpose(
+                torch.as_tensor(
+                    self.prompt_to_target(
+                        c,
+                        self.num_prompts,
+                        self.num_sample_per_mel_frame,
+                        self.subsampling_factor,
+                    ),
+                    dtype=torch.float32,
+                ),
+                0,
+                1,
+            )
+            for c in cuts
+        ]
+
+        # Create final tensors
+        token_lens = torch.tensor([t.size(0) for t in tokens], dtype=torch.long)
+        tokens = collate_vectors(tokens, padding_value=0)
+        prompt_targets = collate_matrices(prompt_targets)
+
+        return (
+            audio,  # Audio signal
+            audio_lens,  # Audio lengths
+            tokens,  # Text tokens
+            token_lens,  # Token lengths
+            prompt_targets,  # Prompt targets
+        )
+
+
+class TokenizerWrapper:
+    """
+    Provide a unified interface for NeMo Tokenizer, AggregateTokenizer, and (char) Parser.
+    """
+
+    def __init__(self, tokenizer):
+        self._tokenizer = tokenizer
+        if isinstance(tokenizer, AggregateTokenizer):
+            self._impl = self._call_agg_tokenizer
+        elif isinstance(tokenizer, TokenizerSpec):
+            self._impl = self._call_tokenizer
+        else:
+            self._impl = self._call_parser
+
+    def __call__(self, text: str, lang: Optional[str] = None):
+        return self._impl(text, lang)
+
+    def _call_agg_tokenizer(self, text: str, lang: Optional[str] = None):
+        assert lang is not None, "Expected 'lang' to be set for AggregateTokenizer."
+        return self._tokenizer.text_to_ids(text, lang)
+
+    def _call_tokenizer(self, text: str, lang: Optional[str] = None):
+        return self._tokenizer.text_to_ids(text)
+
+    def _call_parser(self, text: str, lang: Optional[str] = None):
+        return self._tokenizer(text)

nemo/collections/asr/data/audio_to_text_lhotse_prompt_index.py

@@ -0,0 +1,154 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Simplified Lhotse dataset that returns language ID indices instead of full prompt tensors.
+The model creates the prompt tensor using the actual encoded length.
+"""
+
+import random
+from typing import Dict, Optional, Tuple
+
+import torch
+import torch.utils.data
+from lhotse.dataset import AudioSamples
+from lhotse.dataset.collation import collate_vectors
+
+from nemo.collections.common.tokenizers.aggregate_tokenizer import TokenizerWrapper
+from nemo.collections.common.tokenizers.tokenizer_spec import TokenizerSpec
+from nemo.core.neural_types import AudioSignal, LabelsType, LengthsType, NeuralType
+from nemo.utils import logging
+
+
+class LhotseSpeechToTextBpeDatasetWithPromptIndex(torch.utils.data.Dataset):
+    """
+    Simplified dataset class for speech-to-text with prompt support.
+
+    Instead of computing full prompt tensors, this dataset returns just the
+    language ID index per sample. The model creates the prompt tensor using
+    the actual encoder output length, guaranteeing no size mismatch.
+
+    Returns:
+        audio_signal: Audio waveform [B, T]
+        audio_signal_length: Audio lengths [B]
+        transcripts: Token IDs [B, T]
+        transcript_length: Token lengths [B]
+        prompt_indices: Language ID indices [B] (NOT full tensors)
+    """
+
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        return {
+            'audio_signal': NeuralType(('B', 'T'), AudioSignal()),
+            'audio_signal_length': NeuralType(tuple('B'), LengthsType()),
+            'transcripts': NeuralType(('B', 'T'), LabelsType()),
+            'transcript_length': NeuralType(tuple('B'), LengthsType()),
+            'prompt_indices': NeuralType(tuple('B'), LabelsType()),  # Just indices, not full tensors
+        }
+
+    def __init__(self, tokenizer: TokenizerSpec, cfg: Dict) -> None:
+        super().__init__()
+        self.tokenizer = TokenizerWrapper(tokenizer)
+        self.load_audio = AudioSamples(fault_tolerant=True)
+        self.cfg = cfg
+
+        # Load prompt dictionary from config
+        self.prompt_dict = cfg.get('prompt_dictionary')
+        if not self.prompt_dict:
+            raise ValueError("prompt_dictionary is required in config")
+
+        self.num_prompts = cfg.get('num_prompts', 128)
+
+        # Per-dataset prompt mode is read from cut.custom["prompt_mode"] at runtime.
+        # Supported values:
+        #   "langID"  — always pass the real language ID
+        #   "auto"    — always pass auto (101)
+        #   "unified" — randomize: auto with probability unified_auto_ratio, else lang ID
+        # Set via lhotse input_cfg tags, e.g.  tags: { prompt_mode: langID }
+        self.prompt_mode_field = cfg.get('prompt_mode_field', 'prompt_mode')
+        self.default_prompt_mode = cfg.get('default_prompt_mode', 'unified')
+        self.unified_auto_ratio = cfg.get('unified_auto_ratio', 0.5)
+
+        # Index used for the language-agnostic / auto prompt
+        self.auto_index = self.prompt_dict.get('auto', 101)
+
+        logging.info(
+            f"LhotseSpeechToTextBpeDatasetWithPromptIndex: "
+            f"default_prompt_mode={self.default_prompt_mode}, "
+            f"unified_auto_ratio={self.unified_auto_ratio}"
+        )
+
+    def _get_prompt_index(self, prompt_key: str) -> int:
+        """Maps prompt keys to indices using the prompt dictionary."""
+        if prompt_key not in self.prompt_dict:
+            available_keys = list(self.prompt_dict.keys())
+            raise ValueError(
+                f"Unknown prompt key: '{prompt_key}'. Available: {available_keys[:10]}{'...' if len(available_keys) > 10 else ''}"
+            )
+        return self.prompt_dict[prompt_key]
+
+    def _get_prompt_mode(self, cut) -> str:
+        """Resolve the prompt_mode for a cut from its custom tags."""
+        if cut.custom is not None:
+            mode = cut.custom.get(self.prompt_mode_field)
+            if mode is not None:
+                return mode
+        return self.default_prompt_mode
+
+    def _get_prompt_index_for_cut(self, cut) -> int:
+        """
+        Determine the prompt index for a cut based on its prompt_mode tag.
+
+        Behaviour depends on prompt_mode (set per-dataset via lhotse
+        input_cfg tags, falling back to ``default_prompt_mode``):
+            "langID"  — always return the real language ID
+                        (use for inference / language-forced tasks)
+            "auto"    — always return auto index (language-agnostic)
+            "unified" — return auto with probability unified_auto_ratio,
+                        otherwise the real language ID
+        """
+        mode = self._get_prompt_mode(cut)
+
+        if mode == 'langID':
+            return self._get_prompt_index(cut.supervisions[0].language)
+        elif mode == 'auto':
+            return self.auto_index
+        elif mode == 'unified':
+            if random.random() < self.unified_auto_ratio:
+                return self.auto_index
+            return self._get_prompt_index(cut.supervisions[0].language)
+        else:
+            logging.warning(f"Unknown prompt_mode '{mode}', falling back to unified")
+            if random.random() < self.unified_auto_ratio:
+                return self.auto_index
+            return self._get_prompt_index(cut.supervisions[0].language)
+
+    def __getitem__(self, cuts) -> Tuple[torch.Tensor, ...]:
+        audio, audio_lens, cuts = self.load_audio(cuts)
+        tokens = [torch.as_tensor(self.tokenizer(c.supervisions[0].text, c.supervisions[0].language)) for c in cuts]
+
+        # Get prompt indices (just the language ID per sample, NOT full tensors)
+        prompt_indices = torch.tensor([self._get_prompt_index_for_cut(c) for c in cuts], dtype=torch.long)
+
+        # Create final tensors
+        token_lens = torch.tensor([t.size(0) for t in tokens], dtype=torch.long)
+        tokens = collate_vectors(tokens, padding_value=0)
+
+        return (
+            audio,  # Audio signal [B, T]
+            audio_lens,  # Audio lengths [B]
+            tokens,  # Text tokens [B, T]
+            token_lens,  # Token lengths [B]
+            prompt_indices,  # Language ID indices [B] - model creates full tensor
+        )

nemo/collections/asr/data/audio_to_text_lhotse_prompted.py

@@ -0,0 +1,269 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import os
+from dataclasses import dataclass
+from typing import Optional, Union
+
+import torch.utils.data
+from lhotse import CutSet
+from lhotse.cut import MixedCut
+from lhotse.dataset import AudioSamples
+from lhotse.dataset.collation import collate_vectors
+
+from nemo.collections.common.data import apply_prompt_format_fn
+from nemo.collections.common.prompts import PromptFormatter
+from nemo.collections.common.tokenizers import TokenizerSpec
+
+
+@dataclass
+class PromptedAudioToTextMiniBatch:
+    audio: torch.Tensor
+    audio_lens: torch.Tensor
+    transcript: torch.Tensor
+    transcript_lens: torch.Tensor
+    prompt: torch.Tensor
+    prompt_lens: torch.Tensor
+    prompted_transcript: torch.Tensor
+    prompted_transcript_lens: torch.Tensor
+    cuts: Optional[CutSet] = None
+
+    def get_decoder_inputs_outputs(self) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Returns the inputs and outputs of transformer decoder for training.
+        The input is ``prompted_transcript`` (minus last token),
+        and the output is ``prompted_transcript`` (minus first token).
+        """
+        return self.prompted_transcript[:, :-1], self.prompted_transcript[:, 1:]
+
+
+class PromptedAudioToTextLhotseDataset(torch.utils.data.Dataset):
+    """
+    This dataset is based on :class:`~nemo.collections.asr.data.audio_to_text_lhotse.LhotseSpeechToTextBpeDataset`.
+    It is a Lhotse-style dataset that converts a mini-batch of Cuts into tensors.
+    The main difference from ``LhotseSpeechToTextBpeDataset`` is that we introduce
+    a special prompt format for multitask encoder-decoder models.
+
+    To perform the prompt formatting, we accept a ``prompt_format_fn``.
+    It's expected to accept:
+    * a ``CutSet`` which it will internally iterate over for utterances, and
+    * a ``TokenizerWrapper`` object that will be internally used to tokenize the utterances
+
+    Tokenized utterances will be extended with special prompt tokens according to ``prompt_format_fn`` logic.
+    We support cuts with multiple supervision segments -- their tokenized texts will be concatenated before we add the prompt tokens.
+    This is useful, for example, in code-switched scenarios where each segment is spoken in a different language.
+
+    Chunking:
+    If `enable_chunking` is True, each audio sample is split into optimally sized chunks
+    (see `find_optimal_chunk_size` and `chunk_waveform`). This is useful for long audio inputs,
+    allowing the model to process them in manageable segments.
+
+    NOTE:
+    If the environment variable `USE_AIS_GET_BATCH` is set to `true` (case-insensitive),
+    then batch audio loading from AIStore will be enabled for this dataset. This will use the
+    AISBatchLoader to load the audio from AIStore. This can improve data loading efficiency in some setups.
+    """
+
+    def __init__(
+        self,
+        tokenizer: TokenizerSpec,
+        prompt: PromptFormatter,
+        enable_chunking: bool = False,
+    ):
+        super().__init__()
+        self.tokenizer = tokenizer
+        self.use_ais_get_batch = os.environ.get("USE_AIS_GET_BATCH", "False").lower() == "true"
+
+        # Try to use use_batch_loader if available (Lhotse >= 1.32.0)
+        try:
+            self.load_audio = AudioSamples(fault_tolerant=True, use_batch_loader=self.use_ais_get_batch)
+        except TypeError:
+            # Lhotse < 1.32.0 doesn't support use_batch_loader
+            if self.use_ais_get_batch:
+                import logging
+
+                logging.warning(
+                    "AIS batch loading requested but not supported by this Lhotse version. "
+                    "Please upgrade to Lhotse >= 1.32.0"
+                )
+            self.load_audio = AudioSamples(fault_tolerant=True)
+
+        self.padding_value = self.tokenizer.pad_id
+        self.prompt = prompt
+        self.enable_chunking = enable_chunking
+
+    def __getitem__(self, cuts: CutSet) -> PromptedAudioToTextMiniBatch:
+        # Load the audio's from AIS and add them to the CutSet
+        audio, audio_lens, cuts = self.load_audio(cuts)
+
+        # Will work if batch_size is set to 1.
+        if self.enable_chunking:
+            # If dynamic chunking is enabled, split each audio sample into chunks.
+            new_audio = []
+            new_audio_lens = []
+            for i in range(audio.shape[0]):
+                waveform = audio[i, : audio_lens[i]]
+                # Split the waveform into chunks and get their lengths.
+                chunks, chunk_lens = self._chunk_waveform(waveform)
+                new_audio.extend(chunks)
+                new_audio_lens.extend(chunk_lens)
+            # Stack all chunks into a batch.
+            audio = torch.stack(new_audio)
+            audio_lens = torch.tensor(new_audio_lens, dtype=torch.long)
+            # Adding this to allow gathering results of the same audio from different batches
+            if cuts[0].start != 0:
+                cuts[0].id = cuts[0].id + '_cut_segmented'
+        # Fast-path: the tokenization and prompt format ting was already done before sampling.
+        attrs = ("input_ids", "context_ids", "answer_ids")
+        pre_formatted = all(hasattr(c, a) for c in cuts for a in attrs)
+        if pre_formatted:
+            prompts_with_answers, prompts, answers = zip(*((c.input_ids, c.context_ids, c.answer_ids) for c in cuts))
+        else:
+            formatted = [apply_prompt_format_fn(cut, self.prompt) for cut in cuts]
+            prompts_with_answers = [ex["input_ids"] for ex in formatted]
+            prompts = [ex["context_ids"] for ex in formatted]
+            answers = [ex["answer_ids"] for ex in formatted]
+
+        transcript, transcript_lens = self._collate_tokens(answers)
+        prompts_with_answers, prompts_with_answers_lens = self._collate_tokens(prompts_with_answers)
+        prompts, prompt_lens = self._collate_tokens(prompts)
+
+        return PromptedAudioToTextMiniBatch(
+            audio=audio,
+            audio_lens=audio_lens,
+            transcript=transcript,
+            transcript_lens=transcript_lens,
+            prompt=prompts,
+            prompt_lens=prompt_lens,
+            prompted_transcript=prompts_with_answers,
+            prompted_transcript_lens=prompts_with_answers_lens,
+            cuts=_drop_in_memory_data(cuts),
+        )
+
+    def _collate_tokens(self, tokens: list[Union[list[int], torch.Tensor]]) -> tuple[torch.Tensor, torch.Tensor]:
+        tokens = [torch.as_tensor(t) for t in tokens]
+        token_lens = torch.tensor([t.size(0) for t in tokens], dtype=torch.long)
+        tokens = collate_vectors(tokens, padding_value=self.padding_value)
+        return tokens, token_lens
+
+    def _find_optimal_chunk_size(
+        self, total_len: int, min_sec: int = 30, max_sec: int = 40, sample_rate: int = 16000, overlap_sec: float = 1.0
+    ) -> int:
+        """
+        Find the optimal chunk size for audio processing that minimizes paddings to the last chunk.
+
+        Args:
+            total_len (int): Total length of the audio waveform in samples
+            min_sec (int, optional): Minimum chunk size in seconds. Defaults to 30.
+            max_sec (int, optional): Maximum chunk size in seconds. Defaults to 40.
+            sample_rate (int, optional): Audio sample rate in Hz. Defaults to 16000.
+            overlap_sec (float, optional): Overlap duration between consecutive chunks in seconds.
+                                         Defaults to 1.0.
+
+        Returns:
+            int: Optimal chunk size in samples that maximizes the last chunk length
+        """
+        best_chunk_size = min_sec * sample_rate
+        best_last_chunk_len = 0
+        if total_len < max_sec * sample_rate:
+            return total_len
+        # Try each possible chunk duration in the range
+        for sec in range(min_sec, max_sec + 1):
+            chunk_size = sec * sample_rate
+            overlap_size = int(overlap_sec * sample_rate)
+            step_size = chunk_size - overlap_size
+
+            if step_size <= 0:  # Invalid overlap
+                continue
+            if chunk_size > total_len:
+                continue
+
+            # Calculate how many chunks we'd need and the last chunk's length
+            n_chunks = (total_len + step_size - 1) // step_size
+            last_chunk_len = total_len - step_size * (n_chunks - 1)
+
+            if last_chunk_len > best_last_chunk_len:
+                best_last_chunk_len = last_chunk_len
+                best_chunk_size = chunk_size
+
+        return best_chunk_size
+
+    def _chunk_waveform(
+        self, waveform: torch.Tensor, chunk_size: int = None, overlap_sec: float = 1.0, sample_rate: int = 16000
+    ) -> tuple[list[torch.Tensor], list[int]]:
+        """
+        Split a waveform tensor into overlapping chunks.
+
+        Args:
+            waveform (torch.Tensor): Input audio waveform tensor of shape (time_samples,)
+            chunk_size (int, optional): Size of each chunk in samples. If None, automatically
+                                       determines optimal chunk size using find_optimal_chunk_size().
+                                       Defaults to None.
+            sample_rate (int, optional): Audio sample rate in Hz. Defaults to 16000.
+            overlap_sec (float, optional): Overlap duration between consecutive chunks in seconds.
+                                          Used to calculate step size. Defaults to 2.
+
+        Returns:
+            tuple[list[torch.Tensor], list[int]]: A tuple containing:
+                - List of chunk tensors, each of shape (chunk_size,)
+                - List of original lengths for each chunk before padding (useful for masking
+                  padded regions during processing.
+        """
+        # If chunk_size is None, find the optimal chunk size for this waveform
+        total_len = waveform.shape[0]
+        if chunk_size is None:
+            chunk_size = self._find_optimal_chunk_size(total_len, overlap_sec=overlap_sec)
+        if chunk_size >= total_len:
+            return [waveform], [total_len]
+        overlap_size = int(overlap_sec * sample_rate)
+        step_size = chunk_size - overlap_size
+        chunks = []
+        chunk_lens = []
+        start = 0
+        while start + overlap_size < total_len:
+            end = min(start + chunk_size, total_len)
+            chunk = waveform[start:end]
+            length = chunk.shape[0]
+            if length < chunk_size:
+                pad = torch.zeros(chunk_size - length, dtype=chunk.dtype, device=chunk.device)
+                chunk = torch.cat([chunk, pad], dim=0)
+            chunks.append(chunk)
+            chunk_lens.append(length)
+            start += step_size
+
+        return chunks, chunk_lens
+
+
+class ProbablyIncorrectLanguageKeyError(RuntimeError):
+    pass
+
+
+def _drop_in_memory_data(
+    cuts: CutSet,
+    _fields=frozenset(MixedCut.__dataclass_fields__.keys()),
+) -> CutSet:
+    """Workaround for an edge case in cuts.drop_in_memory_data() on MixedCut with Lhotse<1.29.0"""
+    ans = []
+    for c in cuts:
+        # Not a mixed cut or a mixed cut that wasn't assigned any extra attributes.
+        if not isinstance(c, MixedCut) or _fields.issuperset(c.__dict__.keys()):
+            ans.append(c.drop_in_memory_data())
+        else:
+            extra_attrs = {k: v for k, v in c.__dict__.items() if k not in _fields}
+            for k in extra_attrs:
+                delattr(c, k)
+            ans.append(c.drop_in_memory_data())
+            for k, v in extra_attrs.items():
+                setattr(ans[-1], k, v)
+                setattr(c, k, v)
+    return CutSet(ans)

nemo/collections/asr/data/audio_to_text_lhotse_speaker.py

@@ -0,0 +1,97 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import random
+from typing import Dict, Optional, Tuple
+
+import torch.utils.data
+from lhotse.dataset import AudioSamples
+from lhotse.dataset.collation import collate_vectors
+
+from nemo.collections.asr.data.audio_to_text_lhotse import TokenizerWrapper
+from nemo.collections.asr.parts.utils.asr_multispeaker_utils import speaker_to_target
+from nemo.collections.common.tokenizers.tokenizer_spec import TokenizerSpec
+from nemo.core.neural_types import AudioSignal, LabelsType, LengthsType, NeuralType
+
+
+class LhotseSpeechToTextSpkBpeDataset(torch.utils.data.Dataset):
+    """
+    This dataset is based on BPE datasets from audio_to_text.py. It has the same functionality of LhotseSpeechToTextBpeDataset but also yield speaker target tensor.
+    Unlike native NeMo datasets, Lhotse dataset defines only the mapping from
+    a CutSet (meta-data) to a mini-batch with PyTorch tensors.
+    Specifically, it performs tokenization, I/O, augmentation, and feature extraction (if any).
+    Managing data, sampling, de-duplication across workers/nodes etc. is all handled
+    by Lhotse samplers instead.
+    """
+
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        return {
+            'audio_signal': NeuralType(('B', 'T'), AudioSignal()),
+            'a_sig_length': NeuralType(tuple('B'), LengthsType()),
+            'transcripts': NeuralType(('B', 'T'), LabelsType()),
+            'transcript_length': NeuralType(tuple('B'), LengthsType()),
+            'spk_targets': NeuralType(('B', 'T'), LabelsType()),
+            'bg_spk_targets': NeuralType(('B', 'T'), LabelsType()),
+        }
+
+    def __init__(self, cfg, tokenizer: TokenizerSpec):
+        super().__init__()
+        self.tokenizer = TokenizerWrapper(tokenizer)
+        self.load_audio = AudioSamples(fault_tolerant=True)
+        self.cfg = cfg
+        self.num_speakers = self.cfg.get('num_speakers', 4)
+        self.num_sample_per_mel_frame = self.cfg.get('num_sample_per_mel_frame', 160)
+        self.num_mel_frame_per_asr_frame = self.cfg.get('num_mel_frame_per_asr_frame', 8)
+        self.fixed_spk_id = self.cfg.get('fixed_spk_id', None)
+        self.inference_mode = self.cfg.get('inference_mode', False)
+
+    def __getitem__(self, cuts) -> Tuple[torch.Tensor, ...]:
+
+        audio, audio_lens, cuts = self.load_audio(cuts)
+
+        tokens = []
+        spk_targets = []
+        bg_spk_targets = []
+
+        if self.inference_mode:
+            return audio, audio_lens, None, None, None, None
+
+        for idx, cut in enumerate(cuts):
+
+            speaker_targets, texts = speaker_to_target(
+                a_cut=cut,
+                num_speakers=self.num_speakers,
+                num_sample_per_mel_frame=self.num_sample_per_mel_frame,
+                num_mel_frame_per_asr_frame=self.num_mel_frame_per_asr_frame,
+                return_text=True,
+            )
+            speaker_targets = speaker_targets.transpose(0, 1)[: len(texts)]
+
+            target_speaker_id = random.choice(range(len(texts)))
+            non_target_speaker_ids = [i for i in range(len(texts)) if i != target_speaker_id]
+            text = texts[target_speaker_id]
+            speaker_target = speaker_targets[target_speaker_id]
+            bg_speaker_target = speaker_targets[non_target_speaker_ids].sum(dim=0) > 0
+
+            tokens.append(torch.as_tensor(self.tokenizer(text, cut.supervisions[0].language)))
+            spk_targets.append(speaker_target)
+            bg_spk_targets.append(bg_speaker_target)
+
+        token_lens = torch.tensor([t.size(0) for t in tokens], dtype=torch.long)
+        tokens = collate_vectors(tokens, padding_value=0)
+        spk_targets = collate_vectors(spk_targets, padding_value=0)
+        bg_spk_targets = collate_vectors(bg_spk_targets, padding_value=0)
+
+        return audio, audio_lens, tokens, token_lens, spk_targets, bg_spk_targets

nemo/collections/asr/data/feature_to_label.py

@@ -0,0 +1,507 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import TYPE_CHECKING, Dict, List, Optional
+
+import torch
+
+from nemo.collections.asr.parts.preprocessing.feature_loader import ExternalFeatureLoader
+from nemo.collections.common.parts.preprocessing import collections
+from nemo.core.classes import Dataset
+from nemo.core.neural_types import AcousticEncodedRepresentation, LabelsType, LengthsType, NeuralType
+from nemo.utils import logging
+
+if TYPE_CHECKING:
+    from nemo.collections.asr.parts.preprocessing.perturb import AudioAugmentor
+
+
+def _feature_collate_fn(batch):
+    """collate batch of feat sig, feat len, labels, labels len, assuming all features have the same shape.
+    Args:
+        batch (FloatTensor, LongTensor, LongTensor, LongTensor):  A tuple of tuples of feature, feature lengths,
+               encoded labels, and encoded labels length.
+    """
+    packed_batch = list(zip(*batch))
+    if len(packed_batch) == 5:
+        _, feat_lengths, _, labels_lengths, sample_ids = packed_batch
+    elif len(packed_batch) == 4:
+        sample_ids = None
+        _, feat_lengths, _, labels_lengths = packed_batch
+    else:
+        raise ValueError("Expects 4 or 5 tensors in the batch!")
+
+    features, labels = [], []
+    for b in batch:
+        feat_i, labels_i = b[0], b[2]
+        features.append(feat_i)
+        labels.append(labels_i)
+
+    features = torch.stack(features)
+    feat_lengths = torch.stack(feat_lengths)
+
+    labels = torch.stack(labels)
+    labels_lengths = torch.stack(labels_lengths)
+
+    if sample_ids is None:
+        return features, feat_lengths, labels, labels_lengths
+    else:
+        sample_ids = torch.tensor(sample_ids, dtype=torch.int32)
+        return features, feat_lengths, labels, labels_lengths, sample_ids
+
+
+def _audio_feature_collate_fn(batch, feat_pad_val, label_pad_id):
+    """collate batch of audio feature, audio len, labels, labels len
+    Args:
+        batch (Optional[FloatTensor], Optional[LongTensor], LongTensor,
+               LongTensor):  A tuple of tuples of feature, feature lengths,
+               labels, and label lengths.  This collate func assumes the
+               features are torch tensors of Log-Melspectrogram (i.e. [N_MEL, T]).
+    """
+    packed_batch = list(zip(*batch))
+    if len(packed_batch) == 5:
+        _, feat_lengths, _, labels_lengths, sample_ids = packed_batch
+    elif len(packed_batch) == 4:
+        sample_ids = None
+        _, feat_lengths, _, labels_lengths = packed_batch
+    else:
+        raise ValueError("Expects 4 or 5 tensors in the batch!")
+    max_feat_len = 0
+    has_feat = feat_lengths[0] is not None
+    if has_feat:
+        max_feat_len = max(feat_lengths).item()
+    max_labels_len = max(labels_lengths).item()
+
+    features, labels = [], []
+    for b in batch:
+        feat_i, feat_i_len, label_i, label_i_len = b[0], b[1], b[2], b[3]
+
+        if has_feat:
+            feat_i_len = feat_i_len.item()
+            if feat_i_len < max_feat_len:
+                pad = (0, max_feat_len - feat_i_len)
+                feat_i = torch.nn.functional.pad(feat_i, pad, value=feat_pad_val)
+            features.append(feat_i)
+
+        label_i_len = label_i_len.item()
+        if label_i_len < max_labels_len:
+            pad = (0, max_labels_len - label_i_len)
+            label_i = torch.nn.functional.pad(label_i, pad, value=label_pad_id)
+        labels.append(label_i)
+
+    if has_feat:
+        features = torch.stack(features)
+        feature_lengths = torch.stack(feat_lengths)
+    else:
+        features, feat_lengths = None, None
+    labels = torch.stack(labels)
+    labels_lengths = torch.stack(labels_lengths)
+
+    if sample_ids is None:
+        return features, feature_lengths, labels, labels_lengths
+    else:
+        sample_ids = torch.tensor(sample_ids, dtype=torch.int32)
+        return features, feature_lengths, labels, labels_lengths, sample_ids
+
+
+def _vad_feature_segment_collate_fn(batch, window_length_in_sec, shift_length_in_sec, frame_unit_in_sec):
+    """collate batch of audio features, features len, tokens, tokens len
+    Args:
+        batch (Optional[FloatTensor], Optional[LongTensor], LongTensor,
+            LongTensor):  A tuple of tuples of signal, signal lengths,
+            encoded tokens, and encoded tokens length.  This collate func
+            assumes the signals are 1d torch tensors (i.e. mono audio).
+            batch size equals to 1.
+    """
+    slice_length = int(window_length_in_sec / frame_unit_in_sec)
+    audio_features, feat_lengths, _, tokens_lengths = zip(*batch)
+
+    slice_length = int(min(slice_length, max(feat_lengths)))
+    shift = int(shift_length_in_sec / frame_unit_in_sec)
+    has_audio = feat_lengths[0] is not None
+
+    f_dim = audio_features[0].shape[0]
+    audio_features, num_slices, tokens, feat_lengths = [], [], [], []
+    append_len_start = torch.div(slice_length, 2, rounding_mode='trunc')
+    append_len_end = slice_length - torch.div(slice_length, 2, rounding_mode='trunc')
+    for feat_i, feat_i_len, tokens_i, _ in batch:
+        start = torch.zeros(f_dim, append_len_start)
+        end = torch.zeros(f_dim, append_len_end)
+        feat_i = torch.cat((start, feat_i, end), dim=1)
+        feat_i_len += slice_length
+
+        if has_audio:
+            slices = max(1, torch.div(feat_i_len - slice_length, shift, rounding_mode='trunc'))
+
+            for slice_id in range(slices):
+                start_idx = slice_id * shift
+                end_idx = start_idx + slice_length
+                feat_slice = feat_i[:, start_idx:end_idx]
+                audio_features.append(feat_slice)
+
+            num_slices.append(slices)
+            tokens.extend([tokens_i] * slices)
+            feat_lengths.extend([slice_length] * slices)
+
+    if has_audio:
+        audio_features = torch.stack(audio_features)
+        feat_lengths = torch.tensor(feat_lengths)
+    else:
+        audio_features, feat_lengths = None, None
+
+    tokens = torch.stack(tokens)
+    tokens_lengths = torch.tensor(num_slices)
+    return audio_features, feat_lengths, tokens, tokens_lengths
+
+
+class _FeatureSeqSpeakerLabelDataset(Dataset):
+    """
+    Dataset that loads tensors via a json file containing paths to feature files, sequences of labels. 
+    Each new line is a different sample. Example below:
+    and their target labels. JSON files should be of the following format:
+        {"feature_filepath": "/path/to/feature_0.p", "seq_label": speakerA speakerB SpeakerA ....} \
+        ...
+        {"feature_filepath": "/path/to/feature_n.p", "seq_label": target_seq_label_n} 
+    target_seq_label_n is the string of sequence of speaker label, separated by space.
+
+    Args:
+        manifest_filepath (str): Dataset parameter. Path to JSON containing data.
+        labels (Optional[list]): Dataset parameter. List of unique labels collected from all samples.
+        feature_loader : Dataset parameter. Feature loader to load (external) feature.       
+    """
+
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        """Returns definitions of module output ports."""
+        # TODO output type for external features
+        output_types = {
+            'external_feat': NeuralType(('B', 'D', 'T'), AcousticEncodedRepresentation()),
+            'feat_length': NeuralType(tuple('B'), LengthsType()),
+        }
+
+        if self.is_speaker_emb:
+            output_types.update(
+                {
+                    'embs': NeuralType(('B', 'D', 'T'), AcousticEncodedRepresentation()),
+                    'embs_length': NeuralType(tuple('B'), LengthsType()),
+                    'label': NeuralType(('B', 'T'), LabelsType()),
+                    'label_length': NeuralType(tuple('B'), LengthsType()),
+                }
+            )
+        else:
+            output_types.update(
+                {
+                    'label': NeuralType(('B', 'T'), LabelsType()),
+                    'label_length': NeuralType(tuple('B'), LengthsType()),
+                }
+            )
+
+        return output_types
+
+    def __init__(
+        self,
+        *,
+        manifest_filepath: str,
+        labels: List[str],
+        feature_loader,
+        is_speaker_emb: bool = False,
+    ):
+        super().__init__()
+        self.collection = collections.ASRFeatureSequenceLabel(
+            manifests_files=manifest_filepath.split(','),
+        )
+
+        self.feature_loader = feature_loader
+        self.labels = labels if labels else self.collection.uniq_labels
+        self.is_speaker_emb = is_speaker_emb
+
+        self.label2id, self.id2label = {}, {}
+        for label_id, label in enumerate(self.labels):
+            self.label2id[label] = label_id
+            self.id2label[label_id] = label
+
+        for idx in range(len(self.labels[:5])):
+            logging.debug(" label id {} and its mapped label {}".format(idx, self.id2label[idx]))
+
+    def __len__(self):
+        return len(self.collection)
+
+    def __getitem__(self, index):
+        sample = self.collection[index]
+
+        features = self.feature_loader.process(sample.feature_file)
+        f, fl = features, torch.tensor(features.shape[0]).long()
+
+        t = torch.tensor(sample.seq_label).float()
+        tl = torch.tensor(len(sample.seq_label)).long()
+
+        return f, fl, t, tl
+
+
+class FeatureToSeqSpeakerLabelDataset(_FeatureSeqSpeakerLabelDataset):
+    """
+    Dataset that loads tensors via a json file containing paths to feature
+    files and sequence of speakers. Each new line is a
+    different sample. Example below:
+    {"feature_filepath": "/path/to/feature_0.p", "seq_label": speakerA speakerB SpeakerA ....} \
+    ...
+    {"feature_filepath": "/path/to/feature_n.p", "seq_label": target_seq_label_n} 
+    target_seq_label_n is the string of sequence of speaker label, separated by space.
+
+    Args:
+        manifest_filepath (str): Path to manifest json as described above. Canbe comma-separated paths.
+        labels (Optional[list]): String containing all the possible labels to map to
+            if None then automatically picks from ASRFeatureSequenceLabel collection.
+        feature_loader, Feature load to loader (external) feature.
+    
+    """
+
+    def _collate_fn(self, batch):
+        return _feature_collate_fn(batch)
+
+
+class FeatureToLabelDataset(Dataset):
+    """
+    Dataset that loads tensors via a json file containing paths to feature files and their labels.
+    Each new line is a different sample. Example below:
+    and their target labels. JSON files should be of the following format:
+        {"feature_filepath": "/path/to/audio_feature.pt", "label": "1"}
+        ...
+        {"feature_filepath": "/path/to/audio_feature.pt", "label": "0"}
+    Args:
+        manifest_filepath (str): Path to JSON containing data.
+        labels (Optional[list]): List of unique labels collected from all samples.
+        augmentor (Optional): feature augmentation
+        window_length_in_sec (float): Window length in seconds.
+        shift_length_in_sec (float): Shift length in seconds.
+        is_regression_task (bool): if True, the labels are treated as for a regression task.
+        cal_labels_occurrence (bool): if True, the labels occurrence will be calculated.
+        zero_spec_db_val (float): Value to replace non-speech signals in log-melspectrogram.
+        min_duration (float): Minimum duration of the audio file in seconds.
+        max_duration (float): Maximum duration of the audio file in seconds.
+    """
+
+    ZERO_LEVEL_SPEC_DB_VAL = -16.635  # Log-Melspectrogram value for zero signal
+    FRAME_UNIT_TIME_SECS = 0.01
+
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        """Returns definitions of module output ports."""
+        output_types = {
+            'audio_feat': NeuralType(('B', 'D', 'T'), AcousticEncodedRepresentation()),
+            'feat_length': NeuralType(tuple('B'), LengthsType()),
+            'labels': NeuralType(('B'), LabelsType()),
+            'labels_length': NeuralType(tuple('B'), LengthsType()),
+        }
+
+        return output_types
+
+    def __init__(
+        self,
+        *,
+        manifest_filepath: str,
+        labels: List[str] = None,
+        augmentor: 'AudioAugmentor' = None,
+        window_length_in_sec: float = 0.63,
+        shift_length_in_sec: float = 0.01,
+        is_regression_task: bool = False,
+        cal_labels_occurrence: Optional[bool] = False,
+        zero_spec_db_val: float = -16.635,
+        min_duration: Optional[float] = None,
+        max_duration: Optional[float] = None,
+    ):
+        super().__init__()
+        self.window_length_in_sec = window_length_in_sec
+        self.shift_length_in_sec = shift_length_in_sec
+        self.zero_spec_db_val = zero_spec_db_val
+
+        if isinstance(manifest_filepath, str):
+            manifest_filepath = manifest_filepath.split(',')
+
+        self.collection = collections.ASRFeatureLabel(
+            manifests_files=manifest_filepath,
+            is_regression_task=is_regression_task,
+            cal_labels_occurrence=cal_labels_occurrence,
+            min_duration=min_duration,
+            max_duration=max_duration,
+        )
+
+        self.feature_loader = ExternalFeatureLoader(augmentor=augmentor)
+        self.labels = labels if labels else self.collection.uniq_labels
+
+        self.is_regression_task = is_regression_task
+
+        if not is_regression_task:
+            self.labels = labels if labels else self.collection.uniq_labels
+            self.num_classes = len(self.labels) if self.labels is not None else 1
+            self.label2id, self.id2label = {}, {}
+            self.id2occurrence, self.labels_occurrence = {}, []
+
+            for label_id, label in enumerate(self.labels):
+                self.label2id[label] = label_id
+                self.id2label[label_id] = label
+                if cal_labels_occurrence:
+                    self.id2occurrence[label_id] = self.collection.labels_occurrence[label]
+
+            if cal_labels_occurrence:
+                self.labels_occurrence = [self.id2occurrence[k] for k in sorted(self.id2occurrence)]
+
+            for idx in range(len(self.labels[:5])):
+                logging.debug(" label id {} and its mapped label {}".format(idx, self.id2label[idx]))
+        else:
+            self.labels = []
+            self.num_classes = 1
+
+    def __len__(self):
+        return len(self.collection)
+
+    def __getitem__(self, index):
+        sample = self.collection[index]
+
+        features = self.feature_loader.process(sample.feature_file)
+        f, fl = features, torch.tensor(features.shape[1]).long()
+
+        t = torch.tensor(self.label2id[sample.label])
+        tl = torch.tensor(1).long()
+
+        return f, fl, t, tl
+
+    def _collate_fn(self, batch):
+        return _audio_feature_collate_fn(batch, self.zero_spec_db_val, 0)
+
+    def _vad_segment_collate_fn(self, batch):
+        return _vad_feature_segment_collate_fn(
+            batch, self.window_length_in_sec, self.shift_length_in_sec, self.FRAME_UNIT_TIME_SECS
+        )
+
+
+class FeatureToMultiLabelDataset(Dataset):
+    """
+    Dataset that loads tensors via a json file containing paths to feature files and their labels.
+    Each new line is a different sample. Example below:
+    and their target labels. JSON files should be of the following format:
+        {"feature_filepath": "/path/to/audio_feature.pt", "label": "1 1 0 0 1"}
+        ...
+        {"feature_filepath": "/path/to/audio_feature.pt", "label": "0 1 0 0"}
+    Args:
+        manifest_filepath (str): Path to JSON containing data.
+        labels (Optional[list]): List of unique labels collected from all samples.
+        augmentor (Optional): feature augmentation
+        delimiter (str): delimiter to split the labels.
+        is_regression_task (bool): if True, the labels are treated as for a regression task.
+        cal_labels_occurrence (bool): if True, the labels occurrence will be calculated.
+        zero_spec_db_val (float): Value to replace non-speech signals in log-melspectrogram.
+        min_duration (float): Minimum duration of the audio file in seconds.
+        max_duration (float): Maximum duration of the audio file in seconds.
+    """
+
+    ZERO_LEVEL_SPEC_DB_VAL = -16.635  # Log-Melspectrogram value for zero signal
+
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        """Returns definitions of module output ports."""
+        output_types = {
+            'audio_feat': NeuralType(('B', 'D', 'T'), AcousticEncodedRepresentation()),
+            'feat_length': NeuralType(tuple('B'), LengthsType()),
+            'labels': NeuralType(('B', 'T'), LabelsType()),
+            'labels_length': NeuralType(tuple('B'), LengthsType()),
+        }
+
+        return output_types
+
+    def __init__(
+        self,
+        *,
+        manifest_filepath: str,
+        labels: List[str] = None,
+        augmentor: 'AudioAugmentor' = None,
+        delimiter: Optional[str] = None,
+        is_regression_task: bool = False,
+        cal_labels_occurrence: Optional[bool] = False,
+        zero_spec_db_val: float = -16.635,
+        min_duration: Optional[float] = None,
+        max_duration: Optional[float] = None,
+    ):
+        super().__init__()
+        self.delimiter = delimiter
+        self.zero_spec_db_val = zero_spec_db_val
+
+        if isinstance(manifest_filepath, str):
+            manifest_filepath = manifest_filepath.split(',')
+
+        self.collection = collections.ASRFeatureLabel(
+            manifests_files=manifest_filepath,
+            is_regression_task=is_regression_task,
+            cal_labels_occurrence=cal_labels_occurrence,
+            delimiter=delimiter,
+            min_duration=min_duration,
+            max_duration=max_duration,
+        )
+
+        self.is_regression_task = is_regression_task
+        self.feature_loader = ExternalFeatureLoader(augmentor=augmentor)
+        self.labels = labels if labels else self.collection.uniq_labels
+
+        self.label2id, self.id2label = {}, {}
+        if not is_regression_task:
+            self.labels = labels if labels else self._get_label_set()
+            self.num_classes = len(self.labels) if self.labels is not None else 1
+            self.label2id, self.id2label = {}, {}
+            for label_id, label in enumerate(self.labels):
+                self.label2id[label] = label_id
+                self.id2label[label_id] = label
+                if cal_labels_occurrence:
+                    self.id2occurrence[label_id] = self.collection.labels_occurrence[label]
+                    self.labels_occurrence.append(self.id2occurrence[label_id])
+
+            for idx in range(len(self.labels[:5])):
+                logging.debug(" label id {} and its mapped label {}".format(idx, self.id2label[idx]))
+        else:
+            self.labels = []
+            self.num_classes = 1
+
+    def _get_label_set(self):
+        labels = []
+        for sample in self.collection:
+            label_str = sample.label
+            if label_str:
+                label_str_list = label_str.split(self.delimiter) if self.delimiter else label_str.split()
+                labels.extend(label_str_list)
+        return sorted(set(labels))
+
+    def _label_str_to_tensor(self, label_str: str):
+        labels = label_str.split(self.delimiter) if self.delimiter else label_str.split()
+
+        if self.is_regression_task:
+            labels = [float(s) for s in labels]
+            labels = torch.tensor(labels).float()
+        else:
+            labels = [self.label2id[s] for s in labels]
+            labels = torch.tensor(labels).long()
+        return labels
+
+    def __len__(self):
+        return len(self.collection)
+
+    def __getitem__(self, index):
+        sample = self.collection[index]
+
+        features = self.feature_loader.process(sample.feature_file)
+        f, fl = features, torch.tensor(features.shape[1]).long()
+
+        t = self._label_str_to_tensor(sample.label)
+        tl = torch.tensor(t.size(0)).long()
+
+        return f, fl, t, tl
+
+    def _collate_fn(self, batch):
+        return _audio_feature_collate_fn(batch, self.zero_spec_db_val, 0)

nemo/collections/asr/data/feature_to_label_dataset.py

@@ -0,0 +1,68 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import Any, Optional
+
+from nemo.collections.asr.data import feature_to_label
+
+
+def get_feature_seq_speakerlabel_dataset(
+    feature_loader, config: dict
+) -> feature_to_label.FeatureToSeqSpeakerLabelDataset:
+    """
+    Instantiates a FeatureSeqSpeakerLabelDataset.
+    Args:
+        config: Config of the FeatureToSeqSpeakerLabelDataset.
+
+    Returns:
+        An instance of FeatureToSeqSpeakerLabelDataset.
+    """
+    dataset = feature_to_label.FeatureToSeqSpeakerLabelDataset(
+        manifest_filepath=config['manifest_filepath'],
+        labels=config['labels'],
+        feature_loader=feature_loader,
+    )
+    return dataset
+
+
+def get_feature_label_dataset(config: dict, augmentor: Optional[Any] = None) -> feature_to_label.FeatureToLabelDataset:
+    dataset = feature_to_label.FeatureToLabelDataset(
+        manifest_filepath=config['manifest_filepath'],
+        labels=config['labels'],
+        augmentor=augmentor,
+        window_length_in_sec=config.get("window_length_in_sec", 0.63),
+        shift_length_in_sec=config.get("shift_length_in_sec", 0.08),
+        is_regression_task=config.get("is_regression_task", False),
+        cal_labels_occurrence=config.get("cal_labels_occurrence", False),
+        zero_spec_db_val=config.get("zero_spec_db_val", -16.635),
+        max_duration=config.get('max_duration', None),
+        min_duration=config.get('min_duration', None),
+    )
+    return dataset
+
+
+def get_feature_multi_label_dataset(
+    config: dict, augmentor: Optional[Any] = None
+) -> feature_to_label.FeatureToMultiLabelDataset:
+    dataset = feature_to_label.FeatureToMultiLabelDataset(
+        manifest_filepath=config['manifest_filepath'],
+        labels=config['labels'],
+        augmentor=augmentor,
+        delimiter=config.get('delimiter', None),
+        is_regression_task=config.get("is_regression_task", False),
+        cal_labels_occurrence=config.get("cal_labels_occurrence", False),
+        zero_spec_db_val=config.get("zero_spec_db_val", -16.635),
+        max_duration=config.get('max_duration', None),
+        min_duration=config.get('min_duration', None),
+    )
+    return dataset

nemo/collections/asr/data/feature_to_text.py

@@ -0,0 +1,487 @@
+# Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Callable, Dict, List, Optional, Tuple, Union
+
+import torch
+
+from nemo.collections.asr.data.feature_to_label import _audio_feature_collate_fn
+from nemo.collections.asr.parts.preprocessing.feature_loader import ExternalFeatureLoader
+from nemo.collections.asr.parts.preprocessing.features import normalize_batch
+from nemo.collections.asr.parts.preprocessing.segment import ChannelSelectorType
+from nemo.collections.asr.parts.utils.vad_utils import load_speech_segments_from_rttm
+from nemo.collections.common import tokenizers
+from nemo.collections.common.parts.preprocessing import collections, parsers
+from nemo.core.classes import Dataset
+from nemo.core.neural_types import AcousticEncodedRepresentation, LabelsType, LengthsType, NeuralType
+
+
+class ASRFeatureManifestProcessor:
+    def __init__(
+        self,
+        manifest_filepath: str,
+        parser: Union[str, Callable],
+        max_duration: Optional[float] = None,
+        min_duration: Optional[float] = None,
+        max_utts: int = 0,
+        bos_id: Optional[int] = None,
+        eos_id: Optional[int] = None,
+        pad_id: int = 0,
+        index_by_file_id: bool = False,
+    ):
+        self.parser = parser
+        self.collection = collections.ASRFeatureText(
+            manifests_files=manifest_filepath,
+            parser=parser,
+            min_duration=min_duration,
+            max_duration=max_duration,
+            max_number=max_utts,
+            index_by_file_id=index_by_file_id,
+        )
+
+        self.eos_id = eos_id
+        self.bos_id = bos_id
+        self.pad_id = pad_id
+
+    def process_text_by_id(self, index: int) -> Tuple[List[int], int]:
+        sample = self.collection[index]
+        return self.process_text_by_sample(sample)
+
+    def process_text_by_file_id(self, file_id: str) -> Tuple[List[int], int]:
+        manifest_idx = self.collection.mapping[file_id][0]
+        sample = self.collection[manifest_idx]
+        return self.process_text_by_sample(sample)
+
+    def process_text_by_sample(self, sample: collections.ASRAudioText.OUTPUT_TYPE) -> Tuple[List[int], int]:
+        t, tl = sample.text_tokens, len(sample.text_tokens)
+
+        if self.bos_id is not None:
+            t = [self.bos_id] + t
+            tl += 1
+        if self.eos_id is not None:
+            t = t + [self.eos_id]
+            tl += 1
+
+        return t, tl
+
+
+class _FeatureTextDataset(Dataset):
+    """
+    Dataset that loads tensors via a json file containing paths to audio feature files, transcripts,
+    durations (in seconds) and optional RTTM files. Each new line is a different sample. Example below:
+    {"feature_filepath": "/path/to/audio_feature.pt", "text_filepath": "/path/to/audio.txt",
+    "rttm_filepath": "/path/to/audio_rttm.rttm", "duration": 23.147}
+    ...
+    {"feature_filepath": "/path/to/audio_feature.pt", "text": "the transcription", "offset": 301.75, "duration": 0.82, "utt":
+    "utterance_id", "ctm_utt": "en_4156", "side": "A"}
+    Args:
+        manifest_filepath (str): Path to manifest json as described above. Can be comma-separated paths.
+        parser: Str for a language specific preprocessor or a callable.
+        normalize (bool): whether and where to normalize feature, must be one of [None, "post_norm", "pre_norm"]
+        normalize_type (Union[str, dict]): how to normalize feature, see `nemo.collections.asr.parts.preprocessing.features.normalize_batch`
+        use_rttm (bool): whether to use RTTM files if there is any, default to False
+        rttm_mode (str): how to use RTTM files, must be one of ['mask', 'drop'], default to 'mask'
+        feat_min_len (int): minimum length of feature when rttm_mode=deop, default to 4.
+        feat_mask_val (Optional[float]): value used to mask features with RTTM files, default to None to use zero mel-spectralgram
+        frame_unit_time_secs (float): time in seconds for each frame
+        sample_rate (int): Sample rate to resample loaded audio to
+        int_values (bool): If true, load samples as 32-bit integers. Defauts to False.
+        augmentor (nemo.collections.asr.parts.perturb.AudioAugmentor): An AudioAugmentor object used to augment loaded audio
+        max_duration (float): If audio exceeds this length, do not include in dataset
+        min_duration (float): If audio is less than this length, do not include in dataset
+        max_utts (int): Limit number of utterances
+        trim (bool): whether or not to trim silence. Defaults to False
+        bos_id (int): Id of beginning of sequence symbol to append if not None
+        eos_id (int): Id of end of sequence symbol to append if not None
+        pad_id (int): Id of pad symbol. Defaults to 0
+        return_sample_id (bool): whether to return the sample_id as a part of each sample
+        channel_selector (int | Iterable[int] | str): select a single channel or a subset of channels from multi-channel audio. If set to `'average'`, it performs averaging across channels. Disabled if set to `None`. Defaults to `None`. Uses zero-based indexing.
+    """
+
+    ZERO_LEVEL_SPEC_DB_VAL = -16.635  # Log-Melspectrogram value for zero signal
+    NORM_MODES = ["pre_norm", "post_norm"]
+    RTTM_MODES = ["mask", "drop"]
+
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        """Returns definitions of module output ports."""
+        return {
+            'features': NeuralType(('B', 'D', 'T'), AcousticEncodedRepresentation()),
+            'feature_length': NeuralType(tuple('B'), LengthsType()),
+            'transcripts': NeuralType(('B', 'T'), LabelsType()),
+            'transcript_length': NeuralType(tuple('B'), LengthsType()),
+            'sample_id': NeuralType(tuple('B'), LengthsType(), optional=True),
+        }
+
+    def __init__(
+        self,
+        manifest_filepath: str,
+        parser: Union[str, Callable],
+        normalize: Optional[str] = "post_norm",
+        normalize_type: Union[str, dict] = "per_feature",
+        use_rttm: bool = False,
+        rttm_mode: str = "mask",
+        feat_min_len: int = 4,
+        feat_mask_val: Optional[float] = None,
+        frame_unit_time_secs: float = 0.01,
+        sample_rate: Optional[int] = 16000,
+        augmentor: 'nemo.collections.asr.parts.perturb.FeatureAugmentor' = None,
+        max_duration: Optional[int] = None,
+        min_duration: Optional[int] = None,
+        max_utts: int = 0,
+        trim: bool = False,
+        bos_id: Optional[int] = None,
+        eos_id: Optional[int] = None,
+        pad_id: int = 0,
+        return_sample_id: bool = False,
+        channel_selector: Optional[ChannelSelectorType] = None,
+    ):
+        if type(manifest_filepath) == str:
+            manifest_filepath = manifest_filepath.split(",")
+
+        self.sample_rate = sample_rate
+        self.normalize = normalize
+        self.normalize_type = normalize_type
+        self.use_rttm = use_rttm
+        self.rttm_mode = rttm_mode
+        if self.use_rttm and self.rttm_mode not in self.RTTM_MODES:
+            raise ValueError(f"`rttm_mode` must be one of {self.RTTM_MODES}, got `{rttm_mode}` instead")
+
+        self.feat_min_len = feat_min_len
+        if feat_mask_val is not None:
+            self.feat_mask_val = feat_mask_val
+        elif normalize == "pre_norm":
+            self.feat_mask_val = 0.0  # similar to SpectralAugmentation
+        else:
+            self.feat_mask_val = self.ZERO_LEVEL_SPEC_DB_VAL
+
+        if normalize is not None and normalize not in self.NORM_MODES:
+            raise ValueError(f"`normalize` must be one of {self.NORM_MODES}, got `{normalize}` instead")
+
+        self.frame_unit_time_secs = frame_unit_time_secs
+
+        self.manifest_processor = ASRFeatureManifestProcessor(
+            manifest_filepath=manifest_filepath,
+            parser=parser,
+            max_duration=max_duration,
+            min_duration=min_duration,
+            max_utts=max_utts,
+            bos_id=bos_id,
+            eos_id=eos_id,
+            pad_id=pad_id,
+        )
+        self.featurizer = ExternalFeatureLoader(augmentor=augmentor)
+        self.trim = trim
+        self.return_sample_id = return_sample_id
+        self.channel_selector = channel_selector
+
+    def get_manifest_sample(self, sample_id):
+        return self.manifest_processor.collection[sample_id]
+
+    def __getitem__(self, index):
+        sample = self.manifest_processor.collection[index]
+        offset = sample.offset
+
+        if offset is None:
+            offset = 0
+
+        features = self.featurizer.process(sample.feature_file)
+
+        f, fl = features, torch.tensor(features.shape[1]).long()
+
+        t, tl = self.manifest_processor.process_text_by_sample(sample=sample)
+
+        # Feature normalization
+        if self.normalize is None:
+            if self.use_rttm and sample.rttm_file:
+                f = self.process_features_with_rttm(f, offset, sample.rttm_file, self.feat_mask_val)
+        elif self.normalize == "post_norm":
+            # (Optional) Masking based on RTTM file
+            if self.use_rttm and sample.rttm_file:
+                f = self.process_features_with_rttm(f, offset, sample.rttm_file, self.feat_mask_val)
+
+            f = self.normalize_feature(f)
+        else:  # pre-norm
+            f = self.normalize_feature(f)
+            # (Optional) Masking based on RTTM file
+            if self.use_rttm and sample.rttm_file:
+                f = self.process_features_with_rttm(f, offset, sample.rttm_file, self.feat_mask_val)
+
+        if self.return_sample_id:
+            output = f, fl, torch.tensor(t).long(), torch.tensor(tl).long(), index
+        else:
+            output = f, fl, torch.tensor(t).long(), torch.tensor(tl).long()
+
+        return output
+
+    def process_features_with_rttm(self, features, offset, rttm_file, mask_val):
+        segments = load_speech_segments_from_rttm(rttm_file)
+        new_features = features.clone()
+        sid, fid = 0, 0
+        for i in range(features.size(1)):
+            t = offset + i * self.frame_unit_time_secs
+            while sid < len(segments) - 1 and segments[sid][1] < t:
+                sid += 1
+            if segments[sid][1] == 0 or t < segments[sid][0] or t > segments[sid][1]:
+                # not in speech segment
+                if self.rttm_mode == "drop":
+                    # drop the frame
+                    continue
+                else:
+                    # mask the frame with specified value
+                    new_features[:, i] = mask_val
+                    fid += 1
+            else:
+                # in speech segment
+                new_features[:, fid] = features[:, i]
+                fid += 1
+
+        if fid < self.feat_min_len and self.rttm_mode == "drop":
+            new_features[:, : self.feat_min_len] = mask_val
+            return new_features[:, : self.feat_min_len]
+        return new_features[:, :fid]
+
+    def __len__(self):
+        return len(self.manifest_processor.collection)
+
+    def _collate_fn(self, batch):
+        return _audio_feature_collate_fn(
+            batch, feat_pad_val=self.feat_mask_val, label_pad_id=self.manifest_processor.pad_id
+        )
+
+    def normalize_feature(self, feat):
+        """
+        Args:
+            feat: feature tensor of shape [M, T]
+        """
+        feat = feat.unsqueeze(0)  # add batch dim
+        feat, _, _ = normalize_batch(feat, torch.tensor([feat.size(-1)]), self.normalize_type)
+        return feat.squeeze(0)  # delete batch dim
+
+
+class FeatureToCharDataset(_FeatureTextDataset):
+    """
+    Dataset that loads tensors via a json file containing paths to audio feature
+    files, transcripts, durations (in seconds) and optional RTTM files. Each new line is a
+    different sample. Example below:
+    {"feature_filepath": "/path/to/audio_feature.pt", "text_filepath":
+    "/path/to/audio.txt", "duration": 23.147, "rttm_filepath": "/path/to/audio_rttm.rttm",}
+    ...
+    {"feature_filepath": "/path/to/audio_feature.pt", "text": "the
+    transcription", "offset": 301.75, "duration": 0.82, "utt":
+    "utterance_id", "ctm_utt": "en_4156", "side": "A"}
+
+    Args:
+        manifest_filepath (str): Path to manifest json as described above. Can
+            be comma-separated paths.
+        labels (str): String containing all the possible characters to map to
+        normalize (str): how to normalize feature, must be one of [None, "post_norm", "pre_norm"]
+        normalize_type (Union[str, dict]): how to normalize feature, see `nemo.collections.asr.parts.preprocessing.features.normalize_batch`
+        use_rttm (bool): whether to use RTTM files if there is any, default to False
+        rttm_mode (str): how to use RTTM files, must be one of ['mask', 'drop'], default to 'mask'
+        feat_min_len (int): minimum length of feature, default to 4
+        feat_mask_val (Optional[float]): value used to mask features with RTTM files, default to None to use zero mel-spectralgram
+        frame_unit_time_secs: time in seconds for each frame
+        sample_rate (int): Sample rate to resample loaded audio to
+        int_values (bool): If true, load samples as 32-bit integers. Defauts to False.
+        augmentor (nemo.collections.asr.parts.perturb.AudioAugmentor): An AudioAugmentor
+            object used to augment loaded audio
+        max_duration: If audio exceeds this length, do not include in dataset
+        min_duration: If audio is less than this length, do not include
+            in dataset
+        max_utts: Limit number of utterances
+        blank_index: blank character index, default = -1
+        unk_index: unk_character index, default = -1
+        bos_id: Id of beginning of sequence symbol to append if not None
+        eos_id: Id of end of sequence symbol to append if not None
+        return_sample_id (bool): whether to return the sample_id as a part of each sample
+        channel_selector (int | Iterable[int] | str): select a single channel or a subset of channels from multi-channel audio. If set to `'average'`, it performs averaging across channels. Disabled if set to `None`. Defaults to `None`. Uses zero-based indexing.
+    """
+
+    def __init__(
+        self,
+        manifest_filepath: str,
+        labels: Union[str, List[str]],
+        normalize: Optional[str] = "post_norm",
+        normalize_type: Union[str, dict] = "per_feature",
+        use_rttm: bool = False,
+        rttm_mode: str = "mask",
+        feat_min_len: int = 4,
+        feat_mask_val: Optional[float] = None,
+        frame_unit_time_secs: float = 0.01,
+        sample_rate: Optional[int] = 16000,
+        augmentor: 'nemo.collections.asr.parts.perturb.FeatureAugmentor' = None,
+        max_duration: Optional[int] = None,
+        min_duration: Optional[int] = None,
+        max_utts: int = 0,
+        blank_index: int = -1,
+        unk_index: int = -1,
+        trim: bool = False,
+        bos_id: Optional[int] = None,
+        eos_id: Optional[int] = None,
+        pad_id: int = 0,
+        parser: Union[str, Callable] = 'en',
+        return_sample_id: bool = False,
+        channel_selector: Optional[ChannelSelectorType] = None,
+    ):
+        self.labels = labels
+
+        parser = parsers.make_parser(
+            labels=labels, name=parser, unk_id=unk_index, blank_id=blank_index, do_normalize=normalize
+        )
+
+        super().__init__(
+            manifest_filepath=manifest_filepath,
+            parser=parser,
+            normalize=normalize,
+            normalize_type=normalize_type,
+            use_rttm=use_rttm,
+            rttm_mode=rttm_mode,
+            feat_min_len=feat_min_len,
+            feat_mask_val=feat_mask_val,
+            frame_unit_time_secs=frame_unit_time_secs,
+            sample_rate=sample_rate,
+            augmentor=augmentor,
+            max_duration=max_duration,
+            min_duration=min_duration,
+            max_utts=max_utts,
+            trim=trim,
+            bos_id=bos_id,
+            eos_id=eos_id,
+            pad_id=pad_id,
+            return_sample_id=return_sample_id,
+            channel_selector=channel_selector,
+        )
+
+
+class FeatureToBPEDataset(_FeatureTextDataset):
+    """
+    Dataset that loads tensors via a json file containing paths to audio feature
+    files, transcripts, durations (in seconds) and optional RTTM files. Each new line is a different sample.
+    Example below:
+    {"audio_filepath": "/path/to/audio.wav", "text_filepath":
+    "/path/to/audio.txt", "duration": 23.147, "rttm_filepath": "/path/to/audio_rttm.rttm",}
+    ...
+    {"audio_filepath": "/path/to/audio.wav", "text": "the
+    transcription", "offset": 301.75, "duration": 0.82, "utt":
+    "utterance_id", "ctm_utt": "en_4156", "side": "A"}
+
+    In practice, the dataset and manifest used for character encoding and byte pair encoding
+    are exactly the same. The only difference lies in how the dataset tokenizes the text in
+    the manifest.
+
+    Args:
+        manifest_filepath (str): Path to manifest json as described above. Can
+            be comma-separated paths.
+        tokenizer: A subclass of the Tokenizer wrapper found in the common collection,
+            nemo.collections.common.tokenizers.TokenizerSpec. ASR Models support a subset of
+            all available tokenizers.
+        normalize (str): how to normalize feature, must be one of [None, "post_norm", "pre_norm"]
+        normalize_type (Union[str, dict]): how to normalize feature, see `nemo.collections.asr.parts.preprocessing.features.normalize_batch`
+        use_rttm (bool): whether to use RTTM files if there is any, default to False
+        rttm_mode (str): how to use RTTM files, must be one of ['mask', 'drop'], default to 'mask'
+        feat_min_len (int): minimum length of feature, default to 4
+        feat_mask_val (Optional[float]): value used to mask features with RTTM files, default to None to use zero mel-spectralgram
+        frame_unit_time_secs: time in seconds for each frame
+        sample_rate (int): Sample rate to resample loaded audio to
+        int_values (bool): If true, load samples as 32-bit integers. Defauts to False.
+        augmentor (nemo.collections.asr.parts.perturb.AudioAugmentor): An AudioAugmentor
+            object used to augment loaded audio
+        max_duration: If audio exceeds this length, do not include in dataset
+        min_duration: If audio is less than this length, do not include
+            in dataset
+        max_utts: Limit number of utterances
+        trim: Whether to trim silence segments
+        use_start_end_token: Boolean which dictates whether to add [BOS] and [EOS]
+            tokens to beginning and ending of speech respectively.
+        return_sample_id (bool): whether to return the sample_id as a part of each sample
+        channel_selector (int | Iterable[int] | str): select a single channel or a subset of channels from multi-channel audio. If set to `'average'`, it performs averaging across channels. Disabled if set to `None`. Defaults to `None`. Uses zero-based indexing.
+    """
+
+    def __init__(
+        self,
+        manifest_filepath: str,
+        tokenizer: 'nemo.collections.common.tokenizers.TokenizerSpec',
+        normalize: Optional[str] = "post_norm",
+        normalize_type: Union[str, dict] = "per_feature",
+        use_rttm: bool = False,
+        rttm_mode: str = "mask",
+        feat_min_len: int = 4,
+        feat_mask_val: Optional[float] = None,
+        frame_unit_time_secs: float = 0.01,
+        sample_rate: Optional[int] = 16000,
+        augmentor: 'nemo.collections.asr.parts.perturb.FeatureAugmentor' = None,
+        max_duration: Optional[int] = None,
+        min_duration: Optional[int] = None,
+        max_utts: int = 0,
+        use_start_end_token: bool = True,
+        trim: bool = False,
+        return_sample_id: bool = False,
+        channel_selector: Optional[ChannelSelectorType] = None,
+    ):
+        if use_start_end_token and hasattr(tokenizer, "bos_id") and tokenizer.bos_id > 0:
+            bos_id = tokenizer.bos_id
+        else:
+            bos_id = None
+
+        if use_start_end_token and hasattr(tokenizer, "eos_id") and tokenizer.eos_id > 0:
+            eos_id = tokenizer.eos_id
+        else:
+            eos_id = None
+
+        if hasattr(tokenizer, "pad_id") and tokenizer.pad_id > 0:
+            pad_id = tokenizer.pad_id
+        else:
+            pad_id = 0
+
+        class TokenizerWrapper:
+            def __init__(self, tokenizer):
+                if isinstance(tokenizer, tokenizers.aggregate_tokenizer.AggregateTokenizer):
+                    self.is_aggregate = True
+                else:
+                    self.is_aggregate = False
+                self._tokenizer = tokenizer
+
+            def __call__(self, *args):
+                if isinstance(args[0], List) and self.is_aggregate:
+                    t = []
+                    for span in args[0]:
+                        t.extend(self._tokenizer.text_to_ids(span['str'], span['lang']))
+                    return t
+
+                t = self._tokenizer.text_to_ids(*args)
+                return t
+
+        super().__init__(
+            manifest_filepath=manifest_filepath,
+            parser=TokenizerWrapper(tokenizer),
+            normalize=normalize,
+            normalize_type=normalize_type,
+            use_rttm=use_rttm,
+            rttm_mode=rttm_mode,
+            feat_min_len=feat_min_len,
+            feat_mask_val=feat_mask_val,
+            frame_unit_time_secs=frame_unit_time_secs,
+            sample_rate=sample_rate,
+            augmentor=augmentor,
+            max_duration=max_duration,
+            min_duration=min_duration,
+            max_utts=max_utts,
+            trim=trim,
+            bos_id=bos_id,
+            eos_id=eos_id,
+            pad_id=pad_id,
+            return_sample_id=return_sample_id,
+            channel_selector=channel_selector,
+        )

nemo/collections/asr/data/feature_to_text_dataset.py

@@ -0,0 +1,94 @@
+# Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Optional
+
+from nemo.collections.asr.data.feature_to_text import FeatureToBPEDataset, FeatureToCharDataset
+from nemo.utils import logging
+
+
+def get_char_dataset(config: dict, augmentor: Optional['FeatureAugmentor'] = None) -> FeatureToCharDataset:
+    """
+    Instantiates a Character Encoding based FeatureToCharDataset.
+
+    Args:
+        config: Config of the FeatureToCharDataset.
+        augmentor: Optional AudioAugmentor object for augmentations on audio data.
+
+    Returns:
+        An instance of FeatureToCharDataset.
+    """
+    if 'labels' not in config:
+        logging.warning(f"dataset does not have explicitly defined labels")
+
+    dataset = FeatureToCharDataset(
+        manifest_filepath=config['manifest_filepath'],
+        labels=config.get('labels', None),
+        normalize=config.get('normalize', 'post_norm'),
+        normalize_type=config.get('normalize_type', 'per_feature'),
+        use_rttm=config.get('use_rttm', False),
+        rttm_mode=config.get('rttm_mode', 'mask'),
+        feat_min_len=config.get('feat_min_len', 4),
+        feat_mask_val=config.get('feat_mask_val', None),
+        frame_unit_time_secs=config.get('frame_unit_time_secs', 0.01),
+        sample_rate=config.get('sample_rate', 16000),
+        augmentor=augmentor,
+        max_duration=config.get('max_duration', None),
+        min_duration=config.get('min_duration', None),
+        max_utts=config.get('max_utts', 0),
+        blank_index=config.get('blank_index', -1),
+        unk_index=config.get('unk_index', -1),
+        trim=config.get('trim_silence', False),
+        parser=config.get('parser', 'en'),
+        return_sample_id=config.get('return_sample_id', False),
+        channel_selector=config.get('channel_selector', None),
+    )
+    return dataset
+
+
+def get_bpe_dataset(
+    config: dict, tokenizer: 'TokenizerSpec', augmentor: Optional['FeatureAugmentor'] = None
+) -> FeatureToBPEDataset:
+    """
+    Instantiates a Byte Pair Encoding / Word Piece Encoding based FeatureoToBPEDataset.
+
+    Args:
+        config: Config of the FeatureToBPEDataset.
+        tokenizer: An instance of a TokenizerSpec object.
+        augmentor: Optional FeatureAugmentor object for augmentations on audio features.
+
+    Returns:
+        An instance of FeatureToBPEDataset.
+    """
+    dataset = FeatureToBPEDataset(
+        manifest_filepath=config['manifest_filepath'],
+        tokenizer=tokenizer,
+        normalize=config.get('normalize', 'post_norm'),
+        normalize_type=config.get('normalize_type', 'per_feature'),
+        use_rttm=config.get('use_rttm', False),
+        rttm_mode=config.get('rttm_mode', 'mask'),
+        feat_min_len=config.get('feat_min_len', 4),
+        feat_mask_val=config.get('feat_mask_val', None),
+        frame_unit_time_secs=config.get('frame_unit_time_secs', 0.01),
+        sample_rate=config.get('sample_rate', 16000),
+        augmentor=augmentor,
+        max_duration=config.get('max_duration', None),
+        min_duration=config.get('min_duration', None),
+        max_utts=config.get('max_utts', 0),
+        trim=config.get('trim_silence', False),
+        use_start_end_token=config.get('use_start_end_token', True),
+        return_sample_id=config.get('return_sample_id', False),
+        channel_selector=config.get('channel_selector', None),
+    )
+    return dataset

nemo/collections/asr/data/huggingface/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2023, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

nemo/collections/asr/data/huggingface/hf_audio_to_text.py

@@ -0,0 +1,709 @@
+# Copyright (c) 2023, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from typing import Callable, Dict, List, Optional, Tuple, Union
+
+import datasets as hf_datasets
+import torch
+from datasets import concatenate_datasets
+from datasets.distributed import split_dataset_by_node
+from omegaconf import DictConfig, ListConfig, open_dict
+
+from nemo.collections.asr.data.audio_to_text import _speech_collate_fn
+from nemo.collections.asr.parts.preprocessing.perturb import AudioAugmentor
+from nemo.collections.asr.parts.preprocessing.segment import AudioSegment, ChannelSelectorType
+from nemo.collections.common import tokenizers
+from nemo.collections.common.parts.preprocessing import parsers
+from nemo.core.classes import Dataset, IterableDataset
+from nemo.core.neural_types import AudioSignal, LabelsType, LengthsType, NeuralType
+from nemo.utils import logging
+
+
+class HFTextProcessor:
+    """
+    Text processor for huggingface datasets, mimicing the behavior of
+    `nemo.collections.asr.data.audio_to_text.ASRManifestProcessor`.
+    Basic text cleaning is also supported.
+    Args:
+        parser: Str for a language specific preprocessor or a callable.
+        bos_id: BOS token id to add to the beginning of the transcript.
+        eos_id: EOS token id to add to the end of the transcript.
+        pad_id: PAD token id to pad transcripts to the same length.
+        normalize_text: If true, normalizes text in HFTextProcessor
+        symbols_to_keep: If not None, only keeps symbols in this list when normalizing text
+    """
+
+    def __init__(
+        self,
+        parser: Union[str, Callable],
+        bos_id: Optional[int] = None,
+        eos_id: Optional[int] = None,
+        pad_id: int = 0,
+        normalize_text: bool = False,
+        symbols_to_keep: Optional[str | List[str]] = None,
+    ):
+        self.parser = parser
+        self.eos_id = eos_id
+        self.bos_id = bos_id
+        self.pad_id = pad_id
+        self.normalize_text = normalize_text
+        self.symbols_to_keep = [x for x in symbols_to_keep] if symbols_to_keep is not None else []
+
+    def process_text(self, text: str, lang: Optional[str] = None) -> List[int]:
+
+        if self.normalize_text:
+            text = text.lower()
+            # only keep alphanumeric characters, spaces and symbols defined in self.symbols_to_keep
+            text = ''.join([c for c in text if c.isalnum() or c.isspace() or c in self.symbols_to_keep])
+
+        if hasattr(self.parser, "is_aggregate") and self.parser.is_aggregate and isinstance(text, str):
+            if lang is not None:
+                text_tokens = self.parser(text, lang)
+            # for future use if want to add language bypass to audio_to_text classes
+            # elif hasattr(parser, "lang") and parser.lang is not None:
+            #    text_tokens = parser(text, parser.lang)
+            else:
+                raise ValueError("lang required in manifest when using aggregate tokenizers")
+        else:
+            text_tokens = self.parser(text)
+        text_tokens_length = len(text_tokens)
+        if self.bos_id is not None:
+            text_tokens = [self.bos_id] + text_tokens
+            text_tokens_length += 1
+        if self.eos_id is not None:
+            text_tokens = text_tokens + [self.eos_id]
+            text_tokens_length += 1
+        return text_tokens, text_tokens_length
+
+
+def get_nested_dict_value(dictionary: dict, key: str):
+    """
+    the key should be a string of nested keys separated by `.`, e.g. `key1.key2.key3`,
+    then the returned value will be `dictionary[key1][key2][key3]`
+    """
+    nested_keys = key.split(".")
+    result = dictionary
+    for k in nested_keys:
+        if k not in result:
+            raise KeyError(
+                f"Key `{key}` not found in [{result.keys()}], target is {nested_keys}, input is {dictionary}"
+            )
+        result = result[k]
+    return result
+
+
+class _HFAudioTextDataset(Dataset):
+    """
+    A Dataset wrapper that loads from HuggingFace datasets and converts to NeMo compatible format.
+    Args:
+        audio_key: key to access audio data from the dataset
+        text_key: key to access text data from the dataset
+        sample_rate_key: key to access sample rate data from the dataset
+        hf_data_cfg: HuggingFace dataset config, params are passed to `hf_datasets.load_dataset` (trust_remote_code is always stripped for security)
+        parser: Str for a language specific preprocessor or a callable.
+        augmentor: An instance of `nemo.collections.asr.parts.perturb.AudioAugmentor` to apply on audio.
+        trim: If true, trims silence using `nemo.collections.asr.parts.preprocessing.segment.AudioSegment`
+        bos_id: BOS token id to add to the beginning of the transcript.
+        eos_id: EOS token id to add to the end of the transcript.
+        pad_id: PAD token id to pad transcripts to the same length.
+        return_sample_id: If true, returns sample id from the dataset.
+        channel_selector: ChannelSelectorType, which channel(s) to use for audio.
+        normalize_db: Target RMS value for audio normalization.
+        ref_channel: Reference channel for normalization.
+        id_key: key to access sample id from the dataset
+        normalize_text: If true, normalizes text in HFTextProcessor
+        symbols_to_keep: If not None, only keeps symbols in this list when normalizing text
+    """
+
+    def __init__(
+        self,
+        audio_key: str,
+        text_key: str,
+        sample_rate_key: str,
+        hf_data_cfg: Union[DictConfig, ListConfig],
+        parser: Union[str, Callable],
+        sample_rate: int,
+        augmentor: Optional[AudioAugmentor] = None,
+        trim: bool = False,
+        bos_id: Optional[int] = None,
+        eos_id: Optional[int] = None,
+        pad_id: int = 0,
+        return_sample_id: bool = False,
+        channel_selector: Optional[ChannelSelectorType] = None,
+        normalize_db: Optional[float] = None,
+        ref_channel: Optional[int] = None,
+        id_key: Optional[str] = None,
+        normalize_text: bool = False,
+        symbols_to_keep: Optional[str] = None,
+    ) -> None:
+        super().__init__()
+        self.audio_key = audio_key
+        self.text_key = text_key
+        self.sample_rate_key = sample_rate_key
+        self.id_key = id_key
+        self.sample_rate = sample_rate
+        self.augmentor = augmentor if augmentor is not None else AudioAugmentor()
+        self.trim = trim
+        self.return_sample_id = return_sample_id
+        self.channel_selector = channel_selector
+        self.normalize_db = normalize_db
+        self.ref_channel = ref_channel
+
+        self.text_processor = HFTextProcessor(parser, bos_id, eos_id, pad_id, normalize_text, symbols_to_keep)
+
+        data_config_list = [hf_data_cfg] if isinstance(hf_data_cfg, DictConfig) else hf_data_cfg
+        dataset_list = []
+        for data_cfg in data_config_list:
+            with open_dict(data_cfg):
+                if data_cfg.pop("trust_remote_code", False):
+                    logging.warning(
+                        "trust_remote_code=True was found in hf_data_cfg but has been removed for security. "
+                        "This key is always stripped from hf_data_cfg before calling load_dataset."
+                    )
+                if "streaming" in data_cfg and data_cfg.streaming:
+                    logging.warning(
+                        "streaming must be False for random access dataset, but you use streaming=True. "
+                        "Forcing streaming=False"
+                    )
+                data_cfg.streaming = False
+            logging.info(f"Loading HuggingFace Dataset with cfg: {data_cfg}")
+            dataset_list.append(hf_datasets.load_dataset(**data_cfg))
+            logging.info(f"Dataset loaded with {len(dataset_list[-1])} samples")
+        self.dataset = concatenate_datasets(dataset_list)
+
+        logging.info(f"Total number of samples loaded: {len(self.dataset)}")
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def __getitem__(self, index) -> Tuple:
+        item = self.dataset[index]
+
+        audio_array = get_nested_dict_value(item, self.audio_key)
+        origin_sr = get_nested_dict_value(item, self.sample_rate_key)
+        audio_segment = AudioSegment(
+            samples=audio_array,
+            sample_rate=origin_sr,
+            target_sr=self.sample_rate,
+            trim=self.trim,
+            channel_selector=self.channel_selector,
+            normalize_db=self.normalize_db,
+            ref_channel=self.ref_channel,
+        )
+        self.augmentor.perturb(audio_segment)
+        f = torch.tensor(audio_segment.samples, dtype=torch.float)
+        fl = torch.tensor(f.shape[0], dtype=torch.long)
+
+        text = get_nested_dict_value(item, self.text_key)
+        t, tl = self.text_processor.process_text(text)
+
+        index = get_nested_dict_value(item, self.id_key) if self.id_key else index
+        if self.return_sample_id:
+            output = f, fl, torch.tensor(t).long(), torch.tensor(tl).long(), index
+        else:
+            output = f, fl, torch.tensor(t).long(), torch.tensor(tl).long()
+
+        return output
+
+    def _collate_fn(self, batch):
+        return _speech_collate_fn(batch, pad_id=self.text_processor.pad_id)
+
+
+class HFAudioToCharDataset(_HFAudioTextDataset):
+    """
+    Wrapper class for loading HuggingFace dataset for a char-based ASR model
+    """
+
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        """Returns definitions of module output ports."""
+        return {
+            'audio_signal': NeuralType(('B', 'T'), AudioSignal()),
+            'a_sig_length': NeuralType(tuple('B'), LengthsType()),
+            'transcripts': NeuralType(('B', 'T'), LabelsType()),
+            'transcript_length': NeuralType(tuple('B'), LengthsType()),
+            'sample_id': NeuralType(tuple('B'), LengthsType(), optional=True),
+        }
+
+    def __init__(
+        self,
+        audio_key: str,
+        text_key: str,
+        sample_rate_key: str,
+        hf_data_cfg: DictConfig,
+        labels: List[str],
+        sample_rate: int,
+        augmentor: Optional[AudioAugmentor] = None,
+        trim: bool = False,
+        bos_id: Optional[int] = None,
+        eos_id: Optional[int] = None,
+        pad_id: int = 0,
+        return_sample_id: bool = False,
+        channel_selector: Optional[ChannelSelectorType] = None,
+        normalize_db: Optional[float] = None,
+        ref_channel: Optional[int] = None,
+        parser: Union[str, Callable] = 'en',
+        blank_index: int = -1,
+        unk_index: int = -1,
+        normalize: bool = True,
+        id_key: Optional[str] = None,
+        normalize_text: bool = False,
+        symbols_to_keep: Optional[str] = None,
+    ):
+        self.labels = labels
+
+        parser = parsers.make_parser(
+            labels=labels, name=parser, unk_id=unk_index, blank_id=blank_index, do_normalize=normalize
+        )
+
+        super().__init__(
+            audio_key=audio_key,
+            text_key=text_key,
+            sample_rate_key=sample_rate_key,
+            hf_data_cfg=hf_data_cfg,
+            parser=parser,
+            sample_rate=sample_rate,
+            augmentor=augmentor,
+            trim=trim,
+            bos_id=bos_id,
+            eos_id=eos_id,
+            pad_id=pad_id,
+            return_sample_id=return_sample_id,
+            channel_selector=channel_selector,
+            normalize_db=normalize_db,
+            ref_channel=ref_channel,
+            id_key=id_key,
+            normalize_text=normalize_text,
+            symbols_to_keep=symbols_to_keep,
+        )
+
+
+class HFAudioToBPEDataset(_HFAudioTextDataset):
+    """
+    Wrapper class for loading a HuggingFace dataset for a BPE-based ASR model
+    """
+
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        """Returns definitions of module output ports."""
+        return {
+            'audio_signal': NeuralType(('B', 'T'), AudioSignal()),
+            'a_sig_length': NeuralType(tuple('B'), LengthsType()),
+            'transcripts': NeuralType(('B', 'T'), LabelsType()),
+            'transcript_length': NeuralType(tuple('B'), LengthsType()),
+            'sample_id': NeuralType(tuple('B'), LengthsType(), optional=True),
+        }
+
+    def __init__(
+        self,
+        audio_key: str,
+        text_key: str,
+        sample_rate_key: str,
+        hf_data_cfg: DictConfig,
+        tokenizer: tokenizers.TokenizerSpec,
+        sample_rate: int,
+        augmentor: Optional[AudioAugmentor] = None,
+        trim: bool = False,
+        return_sample_id: bool = False,
+        channel_selector: Optional[ChannelSelectorType] = None,
+        normalize_db: Optional[float] = None,
+        ref_channel: Optional[int] = None,
+        use_start_end_token: bool = True,
+        id_key: Optional[str] = None,
+        normalize_text: bool = False,
+        symbols_to_keep: Optional[str] = None,
+    ):
+        if use_start_end_token and hasattr(tokenizer, "bos_id") and tokenizer.bos_id > 0:
+            bos_id = tokenizer.bos_id
+        else:
+            bos_id = None
+
+        if use_start_end_token and hasattr(tokenizer, "eos_id") and tokenizer.eos_id > 0:
+            eos_id = tokenizer.eos_id
+        else:
+            eos_id = None
+
+        if hasattr(tokenizer, "pad_id") and tokenizer.pad_id > 0:
+            pad_id = tokenizer.pad_id
+        else:
+            pad_id = 0
+
+        class TokenizerWrapper:
+            def __init__(self, tokenizer):
+                if isinstance(tokenizer, tokenizers.aggregate_tokenizer.AggregateTokenizer):
+                    self.is_aggregate = True
+                else:
+                    self.is_aggregate = False
+                self._tokenizer = tokenizer
+
+            def __call__(self, *args):
+                if isinstance(args[0], List) and self.is_aggregate:
+                    t = []
+                    for span in args[0]:
+                        t.extend(self._tokenizer.text_to_ids(span['str'], span['lang']))
+                    return t
+
+                t = self._tokenizer.text_to_ids(*args)
+                return t
+
+        super().__init__(
+            audio_key=audio_key,
+            text_key=text_key,
+            sample_rate_key=sample_rate_key,
+            hf_data_cfg=hf_data_cfg,
+            parser=TokenizerWrapper(tokenizer),
+            sample_rate=sample_rate,
+            augmentor=augmentor,
+            trim=trim,
+            bos_id=bos_id,
+            eos_id=eos_id,
+            pad_id=pad_id,
+            return_sample_id=return_sample_id,
+            channel_selector=channel_selector,
+            normalize_db=normalize_db,
+            ref_channel=ref_channel,
+            id_key=id_key,
+            normalize_text=normalize_text,
+            symbols_to_keep=symbols_to_keep,
+        )
+
+
+class _HFIterableAudioTextDataset(IterableDataset):
+    """
+    Wrapper class for loading HuggingFace IterableDataset and converts to NeMo compatible format.
+    Args:
+        audio_key: key to access audio data from the dataset
+        text_key: key to access text data from the dataset
+        sample_rate_key: key to access sample rate data from the dataset
+        hf_data_cfg: HuggingFace dataset config, params are passed to `hf_datasets.load_dataset` (trust_remote_code is always stripped for security)
+        parser: Str for a language specific preprocessor or a callable.
+        augmentor: An instance of `nemo.collections.asr.parts.perturb.AudioAugmentor` to apply on audio.
+        trim: If true, trims silence using `nemo.collections.asr.parts.preprocessing.segment.AudioSegment`
+        bos_id: BOS token id to add to the beginning of the transcript.
+        eos_id: EOS token id to add to the end of the transcript.
+        pad_id: PAD token id to pad transcripts to the same length.
+        return_sample_id: If true, returns sample id from the dataset.
+        channel_selector: ChannelSelectorType, which channel(s) to use for audio.
+        normalize_db: Target RMS value for audio normalization.
+        ref_channel: Reference channel for normalization.
+        id_key: key to access sample id from the dataset
+        global_rank: global rank of the current worker
+        world_size: total number of workers
+        shuffle_n: buffer size for shuffling
+        shuffle_seed: seed for shuffling
+        normalize_text: If true, normalizes text in HFTextProcessor
+        symbols_to_keep: If not None, only keeps symbols in this list when normalizing text
+    """
+
+    def __init__(
+        self,
+        audio_key: str,
+        text_key: str,
+        sample_rate_key: str,
+        hf_data_cfg: Union[DictConfig, ListConfig],
+        parser: Union[str, Callable],
+        sample_rate: int,
+        augmentor: Optional[AudioAugmentor] = None,
+        trim: bool = False,
+        bos_id: Optional[int] = None,
+        eos_id: Optional[int] = None,
+        pad_id: int = 0,
+        return_sample_id: bool = False,
+        channel_selector: Optional[ChannelSelectorType] = None,
+        normalize_db: Optional[float] = None,
+        ref_channel: Optional[int] = None,
+        id_key: Optional[str] = None,
+        global_rank: int = 0,
+        world_size: int = 0,
+        shuffle_n: int = 0,
+        shuffle_seed: Optional[int] = None,
+        normalize_text: bool = False,
+        symbols_to_keep: Optional[str] = None,
+    ) -> None:
+        super().__init__()
+
+        if return_sample_id and id_key is None:
+            raise ValueError("return_sample_id is True, but id_key is None")
+
+        self.audio_key = audio_key
+        self.text_key = text_key
+        self.sample_rate_key = sample_rate_key
+        self.id_key = id_key
+        self.sample_rate = sample_rate
+        self.augmentor = augmentor if augmentor is not None else AudioAugmentor()
+        self.trim = trim
+        self.return_sample_id = return_sample_id
+        self.channel_selector = channel_selector
+        self.normalize_db = normalize_db
+        self.ref_channel = ref_channel
+
+        self.text_processor = HFTextProcessor(parser, bos_id, eos_id, pad_id, normalize_text, symbols_to_keep)
+
+        data_config_list = [hf_data_cfg] if isinstance(hf_data_cfg, DictConfig) else hf_data_cfg
+        dataset_list = []
+        for data_cfg in data_config_list:
+            with open_dict(data_cfg):
+                if data_cfg.pop("trust_remote_code", False):
+                    logging.warning(
+                        "trust_remote_code=True was found in hf_data_cfg but has been removed for security. "
+                        "This key is always stripped from hf_data_cfg before calling load_dataset."
+                    )
+                if "streaming" in data_cfg and not data_cfg.streaming:
+                    logging.warning(
+                        "streaming must be True for streaming dataset, but you use streaming=False. "
+                        "Forcing streaming=True"
+                    )
+                # streaming must be True for iterable dataset
+                data_cfg.streaming = True
+            logging.info(f"Streaming HuggingFace IterableDataset with cfg: {data_cfg}")
+            dataset_list.append(hf_datasets.load_dataset(**data_cfg))
+
+        self.dataset = concatenate_datasets(dataset_list)
+        logging.info("Total number of samples cannot be extracted from HF streaming dataset")
+
+        if shuffle_n > 0:
+            self.dataset = self.dataset.shuffle(seed=shuffle_seed, buffer_size=shuffle_n)
+
+        self.dataset = split_dataset_by_node(self.dataset, global_rank, world_size)
+        self.dataset = self.dataset.map(self._build_sample)
+
+    def __len__(self):
+        raise NotImplementedError(
+            f"len() is not supported for {self.__class__.__name__}. "
+            "Please set `trainer.max_steps` to explicitly set the number of steps to train for."
+        )
+
+    def __iter__(self):
+        return self.dataset.__iter__()
+
+    def _collate_fn(self, batch):
+        a_signal = [b['audio_signal'] for b in batch]
+        a_sig_length = [b['a_sig_length'] for b in batch]
+        transcripts = [b['transcripts'] for b in batch]
+        transcript_length = [b['transcript_length'] for b in batch]
+        if self.return_sample_id:
+            sample_id = [b['sample_id'] for b in batch]
+            batch_list = list(zip(a_signal, a_sig_length, transcripts, transcript_length, sample_id))
+        else:
+            batch_list = list(zip(a_signal, a_sig_length, transcripts, transcript_length))
+
+        return _speech_collate_fn(batch_list, pad_id=self.text_processor.pad_id)
+
+    def _build_sample(self, sample):
+        audio_array = get_nested_dict_value(sample, self.audio_key)
+        origin_sr = get_nested_dict_value(sample, self.sample_rate_key)
+        audio_segment = AudioSegment(
+            samples=audio_array,
+            sample_rate=origin_sr,
+            target_sr=self.sample_rate,
+            trim=self.trim,
+            channel_selector=self.channel_selector,
+            normalize_db=self.normalize_db,
+            ref_channel=self.ref_channel,
+        )
+        self.augmentor.perturb(audio_segment)
+        f = torch.tensor(audio_segment.samples, dtype=torch.float)
+        fl = torch.tensor(f.shape[0], dtype=torch.long)
+
+        text = get_nested_dict_value(sample, self.text_key)
+        t, tl = self.text_processor.process_text(text)
+
+        output = {
+            'audio_signal': f,
+            'a_sig_length': fl,
+            'transcripts': torch.tensor(t).long(),
+            'transcript_length': torch.tensor(tl).long(),
+        }
+
+        if self.return_sample_id:
+            output['sample_id'] = get_nested_dict_value(sample, self.id_key)
+        return output
+
+
+class HFIterableAudioToCharDataset(_HFIterableAudioTextDataset):
+    """
+    Wrapper class for loading HuggingFace IterableDataset for a char-based ASR model
+    """
+
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        """Returns definitions of module output ports."""
+        return {
+            'audio_signal': NeuralType(('B', 'T'), AudioSignal()),
+            'a_sig_length': NeuralType(tuple('B'), LengthsType()),
+            'transcripts': NeuralType(('B', 'T'), LabelsType()),
+            'transcript_length': NeuralType(tuple('B'), LengthsType()),
+            'sample_id': NeuralType(tuple('B'), LengthsType(), optional=True),
+        }
+
+    def __init__(
+        self,
+        labels: List[str],
+        audio_key: str,
+        text_key: str,
+        sample_rate_key: str,
+        hf_data_cfg: DictConfig,
+        sample_rate: int,
+        augmentor: Optional[AudioAugmentor] = None,
+        trim: bool = False,
+        bos_id: int | None = None,
+        eos_id: int | None = None,
+        pad_id: int = 0,
+        return_sample_id: bool = False,
+        id_key: str | None = None,
+        channel_selector: ChannelSelectorType | None = None,
+        normalize_db: float | None = None,
+        ref_channel: int | None = None,
+        global_rank: int = 0,
+        world_size: int = 0,
+        shuffle_n: int = 0,
+        shuffle_seed: Optional[int] = None,
+        parser: Union[str, Callable] = 'en',
+        blank_index: int = -1,
+        unk_index: int = -1,
+        normalize: bool = True,
+        normalize_text: bool = False,
+        symbols_to_keep: Optional[str] = None,
+    ) -> None:
+        self.labels = labels
+
+        parser = parsers.make_parser(
+            labels=labels, name=parser, unk_id=unk_index, blank_id=blank_index, do_normalize=normalize
+        )
+
+        super().__init__(
+            audio_key=audio_key,
+            text_key=text_key,
+            sample_rate_key=sample_rate_key,
+            hf_data_cfg=hf_data_cfg,
+            parser=parser,
+            sample_rate=sample_rate,
+            augmentor=augmentor,
+            trim=trim,
+            bos_id=bos_id,
+            eos_id=eos_id,
+            pad_id=pad_id,
+            return_sample_id=return_sample_id,
+            id_key=id_key,
+            channel_selector=channel_selector,
+            normalize_db=normalize_db,
+            ref_channel=ref_channel,
+            global_rank=global_rank,
+            world_size=world_size,
+            shuffle_n=shuffle_n,
+            shuffle_seed=shuffle_seed,
+            normalize_text=normalize_text,
+            symbols_to_keep=symbols_to_keep,
+        )
+
+
+class HFIterableAudioToBPEDataset(_HFIterableAudioTextDataset):
+    """
+    Wrapper class for loading HuggingFace IterableDataset for a BPE-based ASR model
+    """
+
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        """Returns definitions of module output ports."""
+        return {
+            'audio_signal': NeuralType(('B', 'T'), AudioSignal()),
+            'a_sig_length': NeuralType(tuple('B'), LengthsType()),
+            'transcripts': NeuralType(('B', 'T'), LabelsType()),
+            'transcript_length': NeuralType(tuple('B'), LengthsType()),
+            'sample_id': NeuralType(tuple('B'), LengthsType(), optional=True),
+        }
+
+    def __init__(
+        self,
+        audio_key: str,
+        text_key: str,
+        sample_rate_key: str,
+        hf_data_cfg: DictConfig,
+        tokenizer: tokenizers.TokenizerSpec,
+        sample_rate: int,
+        augmentor: Optional[AudioAugmentor] = None,
+        trim: bool = False,
+        return_sample_id: bool = False,
+        id_key: str | None = None,
+        channel_selector: ChannelSelectorType | None = None,
+        normalize_db: float | None = None,
+        ref_channel: int | None = None,
+        global_rank: int = 0,
+        world_size: int = 0,
+        shuffle_n: int = 0,
+        shuffle_seed: Optional[int] = None,
+        use_start_end_token: bool = True,
+        normalize_text: bool = False,
+        symbols_to_keep: Optional[str] = None,
+    ) -> None:
+
+        if use_start_end_token and hasattr(tokenizer, "bos_id") and tokenizer.bos_id > 0:
+            bos_id = tokenizer.bos_id
+        else:
+            bos_id = None
+
+        if use_start_end_token and hasattr(tokenizer, "eos_id") and tokenizer.eos_id > 0:
+            eos_id = tokenizer.eos_id
+        else:
+            eos_id = None
+
+        if hasattr(tokenizer, "pad_id") and tokenizer.pad_id > 0:
+            pad_id = tokenizer.pad_id
+        else:
+            pad_id = 0
+
+        class TokenizerWrapper:
+            def __init__(self, tokenizer):
+                if isinstance(tokenizer, tokenizers.aggregate_tokenizer.AggregateTokenizer):
+                    self.is_aggregate = True
+                else:
+                    self.is_aggregate = False
+                self._tokenizer = tokenizer
+
+            def __call__(self, *args):
+                if isinstance(args[0], List) and self.is_aggregate:
+                    t = []
+                    for span in args[0]:
+                        t.extend(self._tokenizer.text_to_ids(span['str'], span['lang']))
+                    return t
+
+                t = self._tokenizer.text_to_ids(*args)
+                return t
+
+        super().__init__(
+            audio_key=audio_key,
+            text_key=text_key,
+            sample_rate_key=sample_rate_key,
+            hf_data_cfg=hf_data_cfg,
+            parser=TokenizerWrapper(tokenizer),
+            sample_rate=sample_rate,
+            augmentor=augmentor,
+            trim=trim,
+            bos_id=bos_id,
+            eos_id=eos_id,
+            pad_id=pad_id,
+            return_sample_id=return_sample_id,
+            id_key=id_key,
+            channel_selector=channel_selector,
+            normalize_db=normalize_db,
+            ref_channel=ref_channel,
+            global_rank=global_rank,
+            world_size=world_size,
+            shuffle_n=shuffle_n,
+            shuffle_seed=shuffle_seed,
+            normalize_text=normalize_text,
+            symbols_to_keep=symbols_to_keep,
+        )

nemo/collections/asr/data/huggingface/hf_audio_to_text_dataset.py

@@ -0,0 +1,139 @@
+# Copyright (c) 2023, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from omegaconf import DictConfig
+
+from nemo.collections.asr.data.huggingface.hf_audio_to_text import (
+    HFAudioToBPEDataset,
+    HFAudioToCharDataset,
+    HFIterableAudioToBPEDataset,
+    HFIterableAudioToCharDataset,
+)
+
+
+def get_hf_audio_to_text_bpe_dataset(
+    config: DictConfig,
+    global_rank: int,
+    world_size: int,
+    tokenizer,
+    augmentor=None,
+):
+    if "streaming" in config and config["streaming"]:
+        dataset = HFIterableAudioToBPEDataset(
+            audio_key=config.get('audio_key', 'audio.array'),
+            text_key=config["text_key"],
+            sample_rate_key=config.get('sample_rate_key', 'audio.sampling_rate'),
+            tokenizer=tokenizer,
+            hf_data_cfg=config["hf_data_cfg"],
+            sample_rate=config["sample_rate"],
+            augmentor=augmentor,
+            trim=config.get('trim_silence', False),
+            return_sample_id=config.get('return_sample_id', False),
+            id_key=config.get("id_key", None),
+            channel_selector=config.get('channel_selector', None),
+            normalize_db=config.get('normalize_db', None),
+            ref_channel=config.get('ref_channel', None),
+            global_rank=global_rank,
+            world_size=world_size,
+            shuffle_n=config.get("shuffle_n", 2048),
+            shuffle_seed=config.get("shuffle_seed", None),
+            use_start_end_token=config.get('use_start_end_token', True),
+            normalize_text=config.get('normalize_text', False),
+            symbols_to_keep=config.get('symbols_to_keep', None),
+        )
+    else:
+        dataset = HFAudioToBPEDataset(
+            audio_key=config.get('audio_key', 'audio.array'),
+            text_key=config["text_key"],
+            sample_rate_key=config.get('sample_rate_key', 'audio.sampling_rate'),
+            tokenizer=tokenizer,
+            hf_data_cfg=config["hf_data_cfg"],
+            sample_rate=config["sample_rate"],
+            augmentor=augmentor,
+            trim=config.get('trim_silence', False),
+            return_sample_id=config.get('return_sample_id', False),
+            id_key=config.get("id_key", None),
+            channel_selector=config.get('channel_selector', None),
+            normalize_db=config.get('normalize_db', None),
+            ref_channel=config.get('ref_channel', None),
+            use_start_end_token=config.get('use_start_end_token', True),
+            normalize_text=config.get('normalize_text', False),
+            symbols_to_keep=config.get('symbols_to_keep', None),
+        )
+
+    return dataset
+
+
+def get_hf_audio_to_text_char_dataset(
+    config: DictConfig,
+    global_rank: int,
+    world_size: int,
+    augmentor=None,
+):
+    if "streaming" in config and config["streaming"]:
+        dataset = HFIterableAudioToCharDataset(
+            labels=config["labels"],
+            audio_key=config.get('audio_key', 'audio.array'),
+            text_key=config["text_key"],
+            sample_rate_key=config.get('sample_rate_key', 'audio.sampling_rate'),
+            hf_data_cfg=config["hf_data_cfg"],
+            sample_rate=config["sample_rate"],
+            augmentor=augmentor,
+            trim=config.get('trim_silence', False),
+            return_sample_id=config.get('return_sample_id', False),
+            id_key=config.get("id_key", None),
+            channel_selector=config.get('channel_selector', None),
+            normalize_db=config.get('normalize_db', None),
+            ref_channel=config.get('ref_channel', None),
+            global_rank=global_rank,
+            world_size=world_size,
+            shuffle_n=config.get("shuffle_n", 2048),
+            shuffle_seed=config.get("shuffle_seed", None),
+            parser=config.get("parser", "en"),
+            blank_index=config.get("blank_index", -1),
+            unk_index=config.get("unk_index", -1),
+            normalize=config.get("normalize", False),
+            normalize_text=config.get('normalize_text', False),
+            symbols_to_keep=config.get('symbols_to_keep', None),
+            pad_id=config.get('pad_id', 0),
+            bos_id=config.get('bos_id', None),
+            eos_id=config.get('eos_id', None),
+        )
+    else:
+        dataset = HFAudioToCharDataset(
+            labels=config["labels"],
+            audio_key=config.get('audio_key', 'audio.array'),
+            text_key=config["text_key"],
+            sample_rate_key=config.get('sample_rate_key', 'audio.sampling_rate'),
+            hf_data_cfg=config["hf_data_cfg"],
+            sample_rate=config["sample_rate"],
+            augmentor=augmentor,
+            trim=config.get('trim_silence', False),
+            bos_id=config.get('bos_id', None),
+            eos_id=config.get('eos_id', None),
+            pad_id=config.get('pad_id', 0),
+            return_sample_id=config.get('return_sample_id', False),
+            id_key=config.get("id_key", None),
+            channel_selector=config.get('channel_selector', None),
+            normalize_db=config.get('normalize_db', None),
+            ref_channel=config.get('ref_channel', None),
+            parser=config.get("parser", "en"),
+            blank_index=config.get("blank_index", -1),
+            unk_index=config.get("unk_index", -1),
+            normalize=config.get("normalize", False),
+            normalize_text=config.get('normalize_text', False),
+            symbols_to_keep=config.get('symbols_to_keep', None),
+        )
+
+    return dataset

nemo/collections/asr/data/ssl_dataset.py

@@ -0,0 +1,697 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import copy
+import io
+import json
+import os
+from dataclasses import dataclass
+from math import isclose
+from typing import Any, Dict, List, Optional, Union
+
+import numpy as np
+import torch
+from lhotse.dataset import AudioSamples
+from omegaconf import DictConfig, ListConfig, open_dict
+from torch import Tensor
+
+from nemo.collections.asr.data import audio_to_text, audio_to_text_dataset
+from nemo.collections.asr.parts.preprocessing.perturb import WhiteNoisePerturbation, process_augmentations
+from nemo.collections.asr.parts.preprocessing.segment import AudioSegment
+from nemo.collections.asr.parts.utils.manifest_utils import read_manifest
+from nemo.collections.common.data.dataset import ConcatDataset
+from nemo.collections.common.parts.preprocessing.manifest import get_full_path
+from nemo.core.classes import Serialization
+from nemo.utils import logging
+
+
+@dataclass
+class AudioNoiseItem:
+    sample_id: str | None = None
+    audio: Union[Tensor, None] = None
+    audio_len: Union[Tensor, None] = None
+    noise: Union[Tensor, None] = None
+    noise_len: Union[Tensor, None] = None
+    noisy_audio: Union[Tensor, None] = None
+    noisy_audio_len: Union[Tensor, None] = None
+
+
+@dataclass
+class AudioNoiseBatch:
+    sample_id: list | None = None
+    audio: Union[Tensor, None] = None
+    audio_len: Union[Tensor, None] = None
+    noise: Union[Tensor, None] = None
+    noise_len: Union[Tensor, None] = None
+    noisy_audio: Union[Tensor, None] = None
+    noisy_audio_len: Union[Tensor, None] = None
+
+
+def _parse_manifest_item(line: str, manifest_file: str) -> Dict[str, Any]:
+    """
+    Specialized function to parse the manifest file by ignoring text,
+    such that nemo dataset can save time on tokenizing text.
+    """
+    item = json.loads(line)
+
+    # Audio file
+    if 'audio_filename' in item:
+        item['audio_file'] = item.pop('audio_filename')
+    elif 'audio_filepath' in item:
+        item['audio_file'] = item.pop('audio_filepath')
+    else:
+        raise KeyError(f"No 'audio_filename' or 'audio_filepath' in manifest item: {item}")
+
+    item['audio_file'] = get_full_path(audio_file=item['audio_file'], manifest_file=manifest_file)
+
+    # Duration.
+    if 'duration' not in item:
+        item['duration'] = None
+
+    # dummy text
+    item['text'] = ""
+
+    item = dict(
+        audio_file=item['audio_file'],
+        duration=item['duration'],
+        text=item['text'],
+        offset=item.get('offset', None),
+        speaker=item.get('speaker', None),
+        orig_sr=item.get('orig_sample_rate', None),
+        token_labels=item.get('token_labels', None),
+        lang=item.get('lang', None),
+    )
+    return item
+
+
+def _audio_noise_collate_fn(batch: List[AudioNoiseItem], batch_augmentor: Any = None) -> AudioNoiseBatch:
+    audios = [x.audio for x in batch]
+    audio_lengths = [x.audio_len for x in batch]
+    max_audio_len = max(audio_lengths).item()
+
+    noises = [x.noise for x in batch]
+    noise_lengths = [x.noise_len for x in batch]
+
+    audio_signal_list = []
+    noise_signal_list = []
+    for i, audio in enumerate(audios):
+        audio_len = audio.size(0)
+        if audio_len < max_audio_len:
+            pad = (0, max_audio_len - audio_len)
+            audio = torch.nn.functional.pad(audio, pad)
+        audio_signal_list.append(audio)
+
+        noise = noises[i]
+        noise_len = noise.size(0)
+        if noise_len < max_audio_len:
+            pad = (0, max_audio_len - noise_len)
+            noise = torch.nn.functional.pad(noise, pad)
+        noise_signal_list.append(noise[:max_audio_len])
+
+    audio_signal = torch.stack(audio_signal_list).float()
+    audio_lengths = torch.stack(audio_lengths).long()
+    noise_signal = torch.stack(noise_signal_list).float()
+    noise_lengths = torch.stack(noise_lengths).long()
+
+    output = AudioNoiseBatch(
+        audio=audio_signal,
+        audio_len=audio_lengths,
+        noise=noise_signal,
+        noise_len=noise_lengths,
+    )
+
+    if batch_augmentor is not None:
+        output = batch_augmentor(output)
+    else:
+        output.noisy_audio = output.audio + output.noise
+        output.noisy_audio_len = output.audio_len
+
+    return output
+
+
+def load_noise_manifest(noise_manifest: Union[str, ListConfig, None]):
+    """
+    load noise manifest from a single or a list of manifest files
+    """
+    if noise_manifest is None:
+        return []
+
+    if isinstance(noise_manifest, str):
+        noise_manifest = noise_manifest.split(',')
+
+    noise_data = []
+    for manifest in noise_manifest:
+        curr_data = read_manifest(manifest)
+        for i in range(len(curr_data)):
+            curr_data[i]['audio_filepath'] = get_full_path(curr_data[i]['audio_filepath'], manifest)
+        noise_data.extend(curr_data)
+    return noise_data
+
+
+def load_noise_audio(
+    sample: Dict[str, Any],
+    sample_rate: int,
+    max_audio_len: Optional[int] = None,
+    pad_to_max: bool = True,
+    min_white_noise_db: int = -90,
+    max_white_noise_db: int = -46,
+    max_trial: int = 100,
+):
+    """
+    Load noise audio from the manifest item, and apply white noise if the loaded noise audio is empty.
+    Args:
+        sample: a sample from the noise manifest
+        sample_rate: target sample rate to resample the noise audio
+        max_audio_len: the maximum audio length to load
+        pad_to_max: whether to pad the audio to max_audio_len
+        min_white_noise_db: the minimum white noise level in dB
+        max_white_noise_db: the maximum white noise level in dB
+        max_trial: the maximum number of trials to load noise audio before giving up
+    Returns:
+        noise: the loaded noise audio
+        noise_len: the length of the loaded noise audio
+    """
+    max_dur = None if max_audio_len is None else max_audio_len / sample_rate
+    duration = sample.get('duration', None)
+    offset = sample.get('offset', 0.0)
+
+    if max_dur is not None and duration is not None and duration > max_dur:
+        cnt = 0
+        while cnt < max_trial:
+            # randomly sample a segment of the noise
+            offset = np.random.uniform(0, duration - max_dur)
+
+            audio_segment = AudioSegment.from_file(
+                audio_file=sample['audio_filepath'],
+                offset=offset,
+                duration=max_dur,
+                target_sr=sample_rate,
+            )
+
+            if sum(audio_segment.samples) > 0:
+                # break if the segment is not empty
+                break
+            cnt += 1
+    else:
+        audio_segment = AudioSegment.from_file(
+            audio_file=sample['audio_filepath'],
+            offset=offset,
+            duration=duration,
+            target_sr=sample_rate,
+        )
+
+    if sum(audio_segment.samples) == 0:
+        logging.warning(
+            f"Loaded noise audio is empty: {sample}, with sampled offset={offset}, duration={max_dur}. Adding white noise."
+        )
+        WhiteNoisePerturbation(min_level=min_white_noise_db, max_level=max_white_noise_db).perturb(audio_segment)
+
+    noise = torch.tensor(audio_segment.samples, dtype=torch.float)
+    noise_len = torch.tensor(noise.size(0)).long()
+    # pad to max_audio_len if necessary
+    if max_audio_len is not None and pad_to_max:
+        if noise.size(0) < max_audio_len:
+            pad = (0, max_audio_len - noise.size(0))
+            noise = torch.nn.functional.pad(noise, pad)
+        else:
+            noise = noise[:max_audio_len]
+            noise_len = torch.tensor(max_audio_len).long()
+    return noise, noise_len
+
+
+def sample_noise(noise_data: List[Dict], sample_rate: int, max_audio_len: int | None = None, max_trial: int = 20):
+    """
+    Randomly sample noise audio from the noise manifest.
+    Args:
+        noise_data: the noise manifest data
+        sample_rate: target sample rate to resample the noise audio
+        max_audio_len: the maximum audio length to load
+        max_trial: the maximum number of trials to load noise audio before giving up
+    Returns:
+        noise_audio: the sampled noise audio
+        noise_len: the length of the sampled noise audio
+    """
+    cnt = 0
+    noise_audio = torch.zeros(max_audio_len).float()
+    noise_len = torch.tensor(max_audio_len).long()
+    while cnt < max_trial and len(noise_data) > 0:
+        try:
+            noise_sample = noise_data[np.random.randint(len(noise_data))]
+            noise_audio, noise_len = load_noise_audio(noise_sample, sample_rate, max_audio_len)
+            break
+        except Exception as e:
+            logging.warning(f"Error loading noise audio with config {noise_sample} and exception: {e}, retrying.")
+            cnt += 1
+            if cnt == max_trial:
+                logging.warning(f"Failed to load noise audio after {max_trial} attempts, returning zero noise.")
+                return torch.zeros(max_audio_len).float(), torch.tensor(max_audio_len).long()
+    return noise_audio, noise_len
+
+
+def pad_audio(audio: Tensor, min_len: int, pad_audio_mode) -> Tensor:
+    """
+    Pad audio to min_len with the specified mode
+    Args:
+        audio: the input audio tensor
+        min_len: the minimum length to pad to
+        pad_audio_mode: the padding mode, either 'repeat' or 'zero'
+    Returns:
+        audio: the padded audio tensor
+    """
+    allowed_mode = ['repeat', 'zero']
+    if audio.size(0) < min_len:
+        if pad_audio_mode == 'repeat' and audio.size(0) > 0:
+            num_repeats = int(np.ceil(min_len / audio.size(0)))
+            audio = audio.repeat(num_repeats)[:min_len]
+        elif pad_audio_mode == 'zero' or audio.size(0) == 0:
+            audio = torch.nn.functional.pad(audio, (0, min_len - audio.size(0)))
+        else:
+            raise ValueError(f"Unsupported pad_audio_mode: {pad_audio_mode}, must be one of {allowed_mode}")
+    return audio
+
+
+class AudioNoiseDataset(audio_to_text.AudioToCharDataset):
+    @property
+    def output_types(self):
+        # disable type checking for since it doesn't support dataclass
+        return None
+
+    def __init__(
+        self,
+        noise_manifest: str | None = None,
+        batch_augmentor: Any | None = None,
+        min_audio_len_secs: float = 1.0,
+        pad_audio_mode: str = 'repeat',
+        **kwargs,
+    ):
+        # add bos_id=0 to avoid empty text token
+        super().__init__(bos_id=0, manifest_parse_func=_parse_manifest_item, **kwargs)
+        self.noise_manifest = noise_manifest
+        self.batch_augmentor = batch_augmentor
+        self.noise_data = load_noise_manifest(noise_manifest)
+        self.min_audio_len_secs = min_audio_len_secs
+        self.pad_audio_mode = pad_audio_mode
+
+    def __getitem__(self, index) -> AudioNoiseItem:
+        sample = self.manifest_processor.collection[index]
+        offset = sample.offset
+
+        if offset is None:
+            offset = 0
+
+        audio = self.featurizer.process(
+            sample.audio_file,
+            offset=offset,
+            duration=sample.duration,
+            trim=self.trim,
+            orig_sr=sample.orig_sr,
+            channel_selector=self.channel_selector,
+        )
+        if audio.size(0) == 0:
+            logging.warning(f"Loaded audio has zero length: {sample}.")
+
+        min_len = int(self.min_audio_len_secs * self.featurizer.sample_rate)
+        audio = pad_audio(audio, min_len, self.pad_audio_mode)
+        audio_len = torch.tensor(audio.shape[0]).long()
+        noise, noise_len = sample_noise(self.noise_data, self.featurizer.sample_rate, audio_len.item())
+
+        item = AudioNoiseItem(
+            sample_id=str(index),
+            audio=audio,
+            audio_len=audio_len,
+            noise=noise,
+            noise_len=noise_len,
+        )
+        return item
+
+    def _collate_fn(self, batch: List[AudioNoiseItem]) -> AudioNoiseBatch:
+        return _audio_noise_collate_fn(batch, self.batch_augmentor)
+
+
+class TarredAudioNoiseDataset(audio_to_text.TarredAudioToCharDataset):
+    @property
+    def output_types(self):
+        # disable type checking for since it doesn't support dataclass
+        return None
+
+    def __init__(
+        self,
+        noise_manifest: str | None = None,
+        batch_augmentor: Any | None = None,
+        min_audio_len_secs: float = 1.0,
+        pad_audio_mode: str = 'repeat',
+        **kwargs,
+    ):
+        """
+        Args:
+            noise_manifest: the noise manifest file
+            batch_augmentor: the batch augmentor
+            min_audio_len_secs: the minimum audio length in seconds, audios shorter than this will be padded
+            pad_audio_mode: the padding mode for audios shorter than min_audio_len_secs, either 'repeat' or 'zero'
+            **kwargs: other arguments for TarredAudioToCharDataset
+
+        """
+        super().__init__(bos_id=0, manifest_parse_func=_parse_manifest_item, **kwargs)
+        self.noise_manifest = noise_manifest
+        self.batch_augmentor = batch_augmentor
+        self.noise_data = load_noise_manifest(noise_manifest)
+        self.min_audio_len_secs = min_audio_len_secs
+        self.pad_audio_mode = pad_audio_mode
+
+    def _build_sample(self, tup):
+        """Builds the training sample by combining the data from the WebDataset with the manifest info."""
+        audio_bytes, audio_filename, offset_id = tup
+
+        # Grab manifest entry from self.manifest_preprocessor.collection
+        file_id, _ = os.path.splitext(os.path.basename(audio_filename))
+        manifest_idx = self.manifest_processor.collection.mapping[file_id][offset_id]
+        manifest_entry = self.manifest_processor.collection[manifest_idx]
+
+        offset = manifest_entry.offset
+        if offset is None:
+            offset = 0
+
+        try:
+            # Convert audio bytes to IO stream for processing (for SoundFile to read)
+            audio_filestream = io.BytesIO(audio_bytes)
+            audio = self.featurizer.process(
+                audio_filestream,
+                offset=offset,
+                duration=manifest_entry.duration,
+                trim=self.trim,
+                orig_sr=manifest_entry.orig_sr,
+            )
+            audio_filestream.close()
+        except Exception as e:
+            raise RuntimeError(f"Error reading audio sample: {manifest_entry}, with exception: {e}.")
+
+        min_len = int(self.min_audio_len_secs * self.featurizer.sample_rate)
+        audio = pad_audio(audio, min_len, self.pad_audio_mode)
+        audio_len = torch.tensor(audio.shape[0]).long()
+        noise, noise_len = sample_noise(self.noise_data, self.featurizer.sample_rate, audio_len.item())
+
+        item = AudioNoiseItem(
+            sample_id=str(manifest_idx),
+            audio=audio,
+            audio_len=audio_len,
+            noise=noise,
+            noise_len=noise_len,
+        )
+        return item
+
+    def _pad_audio(self, audio: Tensor) -> Tensor:
+        min_len = int(self.min_audio_len_secs * self.featurizer.sample_rate)
+        if audio.size(0) < min_len:
+            if self.pad_audio_mode == 'repeat':
+                num_repeats = int(np.ceil(min_len / audio.size(0)))
+                audio = audio.repeat(num_repeats)[:min_len]
+            elif self.pad_audio_mode == 'zero':
+                audio = torch.nn.functional.pad(audio, (0, min_len - audio.size(0)))
+            else:
+                raise ValueError(
+                    f"Unsupported pad_audio_mode: {self.pad_audio_mode}, must be one of ['repeat', 'zero']"
+                )
+        return audio
+
+    def _collate_fn(self, batch: List[AudioNoiseItem]) -> AudioNoiseBatch:
+        return _audio_noise_collate_fn(batch, self.batch_augmentor)
+
+
+class LhotseAudioNoiseDataset(torch.utils.data.Dataset):
+    def __init__(self, noise_manifest: str | None = None, batch_augmentor_cfg: DictConfig = None):
+        super().__init__()
+
+        if batch_augmentor_cfg:
+            batch_augmentor = Serialization.from_config_dict(batch_augmentor_cfg)
+        else:
+            batch_augmentor = None
+
+        self.batch_augmentor = batch_augmentor
+        self.noise_data = load_noise_manifest(noise_manifest)
+        self.load_audio = AudioSamples(fault_tolerant=True)
+
+    def __getitem__(self, cuts):
+
+        audios, audio_lens, cuts = self.load_audio(cuts)
+        if len(self.noise_data) > 0:
+            sampled_noises = [sample_noise(self.noise_data, cut.sampling_rate, cut.num_samples) for cut in cuts]
+            sampled_noises, sampled_noises_lens = zip(*sampled_noises)
+            sampled_noises = torch.stack(sampled_noises).float()
+            sampled_noises_lens = torch.tensor(sampled_noises_lens).long()
+        else:
+            sampled_noises = torch.zeros_like(audios)
+            sampled_noises_lens = audio_lens
+
+        output = AudioNoiseBatch(
+            audio=audios,
+            audio_len=audio_lens,
+            noise=sampled_noises,
+            noise_len=sampled_noises_lens,
+        )
+
+        if self.batch_augmentor is not None:
+            output = self.batch_augmentor(output)
+        else:
+            output.noisy_audio = output.audio + output.noise
+            output.noisy_audio_len = output.audio_len
+
+        return output
+
+
+def get_audio_noise_dataset(
+    config: Dict[str, Any], augmentor: Any = None, batch_augmentor: Any = None
+) -> AudioNoiseDataset:
+    dataset = AudioNoiseDataset(
+        noise_manifest=config.get('noise_manifest', None),
+        batch_augmentor=batch_augmentor,
+        manifest_filepath=config['manifest_filepath'],
+        labels=config.get('labels', None),
+        sample_rate=config['sample_rate'],
+        int_values=config.get('int_values', False),
+        augmentor=augmentor,
+        max_duration=config.get('max_duration', None),
+        min_duration=config.get('min_duration', None),
+        trim=config.get('trim_silence', False),
+        channel_selector=config.get('channel_selector', None),
+    )
+    return dataset
+
+
+def get_concat_audio_noise_dataset(
+    config: Dict[str, Any], global_rank: int, world_size: int, augmentor: Any = None, batch_augmentor: Any = None
+) -> ConcatDataset:
+    manifest_filepaths = config['manifest_filepath']
+    datasets = []
+
+    # needed to support validation Concat Datasets that arrive here as
+    # [[dataset1,dataset2]] otherwise ModelPT would interfere
+    if len(manifest_filepaths) == 1 and not isinstance(manifest_filepaths[0], str):
+        logging.info(f"removing an extra nesting level from {manifest_filepaths}")
+        manifest_filepaths = config['manifest_filepath'][0]
+
+    for manifest_filepath in manifest_filepaths:
+        conf = copy.deepcopy(config)
+        conf['manifest_filepath'] = manifest_filepath
+
+        dataset = get_audio_noise_dataset(config=conf, augmentor=augmentor)
+        datasets.append(dataset)
+
+    dataset = ConcatDataset(
+        datasets,
+        sampling_technique=config.get('concat_sampling_technique', 'temperature'),
+        sampling_temperature=config.get('concat_sampling_temperature', 5),
+        sampling_scale=config.get('concat_sampling_scale', 1),
+        sampling_probabilities=config.get('concat_sampling_probabilities', None),
+        shuffle=config.get('concat_shuffle', True),
+        seed=config.get('concat_sampling_seed', None),
+        global_rank=global_rank,
+        world_size=world_size,
+    )
+    return dataset
+
+
+def get_tarred_audio_noise_dataset(config, shuffle_n, global_rank, world_size, augmentor, batch_augmentor: Any = None):
+    tarred_audio_filepaths = config['tarred_audio_filepaths']
+    manifest_filepaths = config['manifest_filepath']
+    datasets = []
+    tarred_audio_filepaths = audio_to_text_dataset.convert_to_config_list(tarred_audio_filepaths)
+    manifest_filepaths = audio_to_text_dataset.convert_to_config_list(manifest_filepaths)
+
+    bucketing_weights = config.get('bucketing_weights', None)  # For upsampling buckets
+    if bucketing_weights:
+        for idx, weight in enumerate(bucketing_weights):
+            if not isinstance(weight, int) or weight <= 0:
+                raise ValueError("bucket weights must be positive integers")
+
+    if len(manifest_filepaths) != len(tarred_audio_filepaths):
+        raise ValueError(
+            f"manifest_filepaths (length={len(manifest_filepaths)}) and tarred_audio_filepaths (length={len(tarred_audio_filepaths)}) need to have the same number of buckets."
+        )
+
+    for dataset_idx, (tarred_audio_filepath, manifest_filepath) in enumerate(
+        zip(tarred_audio_filepaths, manifest_filepaths)
+    ):
+        if len(tarred_audio_filepath) == 1:
+            tarred_audio_filepath = tarred_audio_filepath[0]
+        if len(manifest_filepath) == 1:
+            manifest_filepath = manifest_filepath[0]
+
+        is_sharded_manifest = True if "_OP_" in manifest_filepath and "_CL_" in manifest_filepath else False
+        logging.info(
+            f"Loading TarredAudioNoiseDataset from {tarred_audio_filepath} and {manifest_filepath}, shard={is_sharded_manifest}"
+        )
+        dataset = TarredAudioNoiseDataset(
+            noise_manifest=config.get('noise_manifest', None),
+            batch_augmentor=batch_augmentor,
+            audio_tar_filepaths=tarred_audio_filepath,
+            manifest_filepath=manifest_filepath,
+            labels=config.get('labels', None),
+            sample_rate=config['sample_rate'],
+            int_values=config.get('int_values', False),
+            augmentor=augmentor,
+            shuffle_n=shuffle_n,
+            max_duration=config.get('max_duration', None),
+            min_duration=config.get('min_duration', None),
+            trim=config.get('trim_silence', False),
+            shard_strategy=config.get('tarred_shard_strategy', 'scatter'),
+            shard_manifests=is_sharded_manifest,
+            global_rank=global_rank,
+            world_size=world_size,
+        )
+        if bucketing_weights:
+            [datasets.append(dataset) for _ in range(bucketing_weights[dataset_idx])]
+        else:
+            datasets.append(dataset)
+
+    return audio_to_text_dataset.get_chain_dataset(datasets=datasets, ds_config=config, rank=global_rank)
+
+
+def get_concat_tarred_audio_noise_dataset(
+    config, shuffle_n, global_rank, world_size, augmentor, batch_augmentor: Any = None
+):
+    tarred_audio_filepaths = config['tarred_audio_filepaths']
+    manifest_filepaths = config['manifest_filepath']
+    datasets = []
+    for dataset_idx, (tarred_audio_filepath, manifest_filepath) in enumerate(
+        zip(tarred_audio_filepaths, manifest_filepaths)
+    ):
+        conf = copy.deepcopy(config)
+        conf['manifest_filepath'] = manifest_filepath
+        conf['tarred_audio_filepaths'] = tarred_audio_filepath
+        dataset = get_tarred_audio_noise_dataset(
+            config=conf,
+            shuffle_n=shuffle_n,
+            global_rank=global_rank,
+            world_size=world_size,
+            augmentor=augmentor,
+            batch_augmentor=batch_augmentor,
+        )
+        datasets.append(dataset)
+
+    dataset = ConcatDataset(
+        datasets,
+        sampling_technique=config.get('concat_sampling_technique', 'temperature'),
+        sampling_temperature=config.get('concat_sampling_temperature', 5),
+        sampling_scale=config.get('concat_sampling_scale', 1),
+        sampling_probabilities=config.get('concat_sampling_probabilities', None),
+        shuffle=config.get('concat_shuffle', True),
+        seed=config.get('concat_sampling_seed', None),
+        global_rank=global_rank,
+        world_size=world_size,
+    )
+    return dataset
+
+
+def get_audio_noise_dataset_from_config(
+    config,
+    global_rank: int,
+    world_size: int,
+):
+    if 'augmentor' in config:
+        augmentor = process_augmentations(config['augmentor'], global_rank=global_rank, world_size=world_size)
+    else:
+        augmentor = None
+
+    if 'batch_augmentor' in config:
+        batch_augmentor = Serialization.from_config_dict(config['batch_augmentor'])
+    else:
+        batch_augmentor = None
+
+    is_concat = config.get('is_concat', False)
+    if is_concat:
+        if config.get('concat_sampling_technique', None) is None:
+            logging.warning(
+                f"Concat dataset requires `concat_sampling_technique` but it was not provided, using round-robin. Config: {config}"
+            )
+            config['concat_sampling_technique'] = 'round-robin'
+
+        if config['concat_sampling_technique'] == 'random':
+            if not 'concat_sampling_probabilities' in config:
+                logging.warning(
+                    f"Concat dataset requires `concat_sampling_probabilities` list, using uniform weights. Config: {config}"
+                )
+                with open_dict(config):
+                    config['concat_sampling_probabilities'] = [1 / len(config['manifest_filepath'])] * len(
+                        config['manifest_filepath']
+                    )
+            elif not isclose(sum(config['concat_sampling_probabilities']), 1, abs_tol=1e-6):
+                raise ValueError(
+                    f"`concat_sampling_probabilities` need to sum to 1 with 1e-6 tolerance. Config: {config}"
+                )
+
+    shuffle = config['shuffle']
+    if config.get('is_tarred', False):
+        if ('tarred_audio_filepaths' in config and config['tarred_audio_filepaths'] is None) or (
+            'manifest_filepath' in config and config['manifest_filepath'] is None
+        ):
+            logging.warning(
+                "Could not load dataset as `manifest_filepath` was None or "
+                f"`tarred_audio_filepaths` is None. Provided config : {config}"
+            )
+            return None
+
+        shuffle_n = config.get('shuffle_n', 4 * config['batch_size']) if shuffle else 0
+        if is_concat:
+            dataset = get_concat_tarred_audio_noise_dataset(
+                config=config,
+                shuffle_n=shuffle_n,
+                global_rank=global_rank,
+                world_size=world_size,
+                augmentor=augmentor,
+                batch_augmentor=batch_augmentor,
+            )
+        else:
+            dataset = get_tarred_audio_noise_dataset(
+                config=config,
+                shuffle_n=shuffle_n,
+                global_rank=global_rank,
+                world_size=world_size,
+                augmentor=augmentor,
+                batch_augmentor=batch_augmentor,
+            )
+    else:
+        if 'manifest_filepath' in config and config['manifest_filepath'] is None:
+            logging.warning(f"Could not load dataset as `manifest_filepath` was None. Provided config : {config}")
+            return None
+        if is_concat:
+            dataset = get_concat_audio_noise_dataset(
+                config=config,
+                global_rank=global_rank,
+                world_size=world_size,
+                augmentor=augmentor,
+                batch_augmentor=batch_augmentor,
+            )
+        else:
+            dataset = get_audio_noise_dataset(config=config, augmentor=augmentor, batch_augmentor=batch_augmentor)
+    return dataset

nemo/collections/asr/data/text_to_text.py

@@ -0,0 +1,484 @@
+# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import concurrent.futures
+import copy
+import gc
+import json
+import math
+import random
+from pathlib import Path
+from typing import Any, Callable, Dict, Iterable, List, NamedTuple, Optional, Set, Union
+
+import numpy as np
+import torch
+import torch.utils.data
+from torch.nn.utils.rnn import pad_sequence
+from tqdm.auto import tqdm
+
+from nemo.collections.asr.data.audio_to_text import _speech_collate_fn
+from nemo.collections.common.tokenizers import TokenizerSpec
+from nemo.core.classes import Dataset, IterableDataset
+from nemo.utils import logging
+
+try:
+    from nemo_text_processing.text_normalization.normalize import Normalizer
+except Exception as e:
+    pass  # Normalizer imported only for annotation purposes, error can be ignored
+
+AnyPath = Union[Path, str]
+
+
+class TextToTextItem(NamedTuple):
+    tts_text: torch.Tensor  # normalized and tokenized text for TTS
+    transcript: torch.Tensor  # tokenized text for ASR
+    speaker: int  # speaker id for multi-speaker TTS
+
+
+class TextToTextBatch(NamedTuple):
+    tts_texts: torch.Tensor  # tokenized texts for tts
+    tts_text_lengths: torch.Tensor
+    transcripts: torch.Tensor  # tokenized texts for ASR
+    transcript_lengths: torch.Tensor
+    speakers: torch.Tensor  # speaker ids for multi-speaker TTS
+
+    @staticmethod
+    def collate_fn(batch: List[TextToTextItem], asr_pad_id: int, tts_text_pad_id: int) -> TextToTextBatch:
+        return TextToTextBatch(
+            tts_texts=pad_sequence([item.tts_text for item in batch], batch_first=True, padding_value=tts_text_pad_id),
+            tts_text_lengths=torch.tensor([item.tts_text.shape[0] for item in batch]).long(),
+            transcripts=pad_sequence([item.transcript for item in batch], batch_first=True, padding_value=asr_pad_id),
+            transcript_lengths=torch.tensor([item.transcript.shape[0] for item in batch]).long(),
+            speakers=torch.tensor([item.speaker for item in batch]).long(),
+        )
+
+
+class TextOrAudioToTextBatch(NamedTuple):
+    audio_signals: torch.Tensor
+    audio_signal_lengths: torch.Tensor
+    tts_texts: torch.Tensor
+    tts_text_lengths: torch.Tensor
+    speakers: torch.Tensor
+    transcripts: torch.Tensor
+    transcript_lengths: torch.Tensor
+
+    @staticmethod
+    def collate_fn(
+        batch: List[Union[TextToTextItem, tuple]], tts_text_pad_id: int, asr_pad_id: int
+    ) -> Union[TextToTextBatch, TextOrAudioToTextBatch, tuple]:
+        """
+        Collate function for dataloader
+        Can accept mixed batch of text-to-text items and audio-text items (typical for ASR)
+        """
+        text_items: List[TextToTextItem] = [item for item in batch if isinstance(item, TextToTextItem)]
+        if not text_items:
+            # pure audio-text batch
+            return _speech_collate_fn(batch=batch, pad_id=asr_pad_id)
+
+        asr_items = [item for item in batch if not isinstance(item, TextToTextItem)]
+
+        if not asr_items:
+            # pure text-to-text batch
+            return TextToTextBatch.collate_fn(batch=text_items, asr_pad_id=asr_pad_id, tts_text_pad_id=tts_text_pad_id)
+
+        # mixed batch
+
+        # each asr item is a tuple:
+        # audio_signal (0), audio_length (1), transcript (2), transcript_length (3), sample_id (4, optional)
+        audio_signals = pad_sequence([item[0] for item in asr_items], batch_first=True, padding_value=0.0)
+        audio_signal_lengths = torch.tensor([item[1] for item in asr_items]).long()
+
+        tts_texts = pad_sequence(
+            [item.tts_text for item in text_items], batch_first=True, padding_value=tts_text_pad_id
+        )
+        tts_text_lengths = torch.tensor([item.tts_text.shape[0] for item in text_items]).long()
+        speakers = torch.tensor([item.speaker for item in text_items]).long()
+
+        transcripts = pad_sequence(
+            [item.transcript for item in text_items] + [item[2] for item in asr_items],
+            batch_first=True,
+            padding_value=asr_pad_id,
+        )
+        transcript_lengths = torch.tensor(
+            [item.transcript.shape[0] for item in text_items] + [item[3] for item in asr_items]
+        ).long()
+
+        return TextOrAudioToTextBatch(
+            audio_signals=audio_signals,
+            audio_signal_lengths=audio_signal_lengths,
+            tts_texts=tts_texts,
+            tts_text_lengths=tts_text_lengths,
+            speakers=speakers,
+            transcripts=transcripts,
+            transcript_lengths=transcript_lengths,
+        )
+
+
+def _asr_text_to_tokens(text: str) -> np.ndarray:
+    """
+    Helper function for asr tokenization with multiprocessing pool only.
+    Must be defined on the top level.
+    Expects asr_tokenizer_global, asr_bos_id_global, asr_eos_id_global to exist in the current pool process
+    """
+    ids = asr_tokenizer_global.text_to_ids(text)
+    if asr_bos_id_global is not None:
+        ids = [asr_bos_id_global] + ids
+    if asr_eos_id_global is not None:
+        ids.append(asr_eos_id_global)
+    return np.asarray(ids)
+
+
+def _tts_text_to_tokens(text: str) -> np.ndarray:
+    """
+    Helper function for asr tokenization with multiprocessing pool only.
+    Must be defined on the top level.
+    Expects tts_tokenizer_global to exist in the current pool process
+    """
+    return np.asarray(tts_tokenizer_global(text))
+
+
+def _iterate_manifest(filepath: AnyPath) -> Iterable[Dict[str, Any]]:
+    """
+    Helper function to iterate manifest
+    """
+    with open(filepath, "r", encoding="utf-8") as f:
+        for line in f:
+            record = json.loads(line)
+            yield record
+
+
+class TextToTextDatasetBase:
+    """
+    Base class for loading text-to-text manifests
+    Map-style and Iterable datasets should inherit this class
+    """
+
+    asr_pad_id: int
+    tts_text_pad_id: int
+    asr_bos_id: Optional[int] = None
+    asr_eos_id: Optional[int] = None
+    data: List[Dict[str, Any]]
+
+    def __init__(
+        self,
+        manifest_filepath: Union[AnyPath, List[AnyPath]],
+        speakers_filepath: Union[AnyPath, List[AnyPath]],
+        asr_tokenizer: TokenizerSpec,
+        asr_use_start_end_token: bool,
+        tts_parser: Callable,
+        tts_text_pad_id: int,
+        tts_text_normalizer: "Normalizer",
+        tts_text_normalizer_call_kwargs: Dict,
+        min_words: int = 1,
+        max_words: int = 1_000_000,
+        tokenizer_workers: int = 1,
+        num_parts: int = 1,
+        current_part_index: int = 0,
+    ):
+        super().__init__()
+        # ASR tokenizer setup
+        if asr_use_start_end_token and hasattr(asr_tokenizer, 'bos_token'):
+            self.asr_bos_id = asr_tokenizer.bos_id
+
+        if asr_use_start_end_token and hasattr(asr_tokenizer, 'eos_token'):
+            self.asr_eos_id = asr_tokenizer.eos_id
+
+        if hasattr(asr_tokenizer, 'pad_token'):
+            self.asr_pad_id = asr_tokenizer.pad_id
+        else:
+            self.asr_pad_id = 0
+
+        self.asr_tokenizer = asr_tokenizer
+
+        # TTS tokenizer setup
+        self.tts_parser = tts_parser
+        self.tts_normalizer = tts_text_normalizer
+        self.tts_normalizer_kwargs = tts_text_normalizer_call_kwargs
+        self.tts_text_pad_id = tts_text_pad_id
+
+        # Load speakers
+        if isinstance(speakers_filepath, str):
+            speakers_filepath = speakers_filepath.split(",")
+        elif isinstance(speakers_filepath, Path):
+            speakers_filepath = [speakers_filepath]
+        speakers: Set[int] = set()
+        for filepath in speakers_filepath:
+            with open(Path(filepath).expanduser(), "r") as f:
+                speakers.update(map(int, f.read().split()))
+        self.speakers = np.asarray(sorted(speakers))
+        logging.info(f"Loaded {len(self.speakers)} speakers")
+
+        # Load manifest
+        if isinstance(manifest_filepath, str):
+            manifest_filepath = manifest_filepath.split(",")
+        elif isinstance(manifest_filepath, Path):
+            manifest_filepath = [manifest_filepath]
+        self.manifest_paths = [Path(filepath) for filepath in manifest_filepath]
+
+        num_skipped_words = 0
+        num_skipped_utterances = 0
+        asr_texts = []
+        tts_texts = []
+        need_normalization = False
+
+        for manifest_path in self.manifest_paths:
+            for tmp_item in tqdm(_iterate_manifest(manifest_path)):
+                text = tmp_item["text"]
+                num_words = len(text.split())
+                # skip if number of works not in desired range
+                # TODO: maybe it would be valuable to sample sub-utterances from long utterances
+                if not (min_words <= num_words <= max_words):
+                    num_skipped_words += num_words
+                    num_skipped_utterances += 1
+                    continue
+                asr_texts.append(tmp_item["text"])
+                if "tts_text_normalized" in tmp_item:
+                    tts_texts.append(tmp_item["tts_text_normalized"])
+                else:
+                    tts_texts.append(tmp_item["tts_text"])
+                    need_normalization = True
+
+        if need_normalization:
+            logging.warning("TTS normalization is extremely slow! It is recommended to normalize TTS text")
+
+        if num_skipped_utterances:
+            logging.warning(f"Skipped {num_skipped_utterances} utterances " f"with {num_skipped_words}")
+
+        num_utterances = len(asr_texts)
+        # preprocessing is very costly, if we need only part - remove unnecessary utterances
+        if num_parts > 1:
+            # NB: floor division, full dataset can contain fewer utterances than original, like in tarred dataset
+            num_utterances_part = num_utterances // num_parts
+            start = num_utterances_part * current_part_index
+            end = start + num_utterances_part
+            logging.info(
+                f"Taking part of the dataset: {current_part_index} index, total {num_parts} from {start} to {end}"
+            )
+            asr_texts = asr_texts[start:end]
+            tts_texts = tts_texts[start:end]
+            num_utterances = num_utterances_part
+
+        self.data = [dict() for _ in range(num_utterances)]
+
+        if len(asr_texts) == 0:
+            # no data was loaded
+            logging.warning("Text-to-text dataset is empty")
+            return
+
+        if tokenizer_workers == 1:
+            logging.warning(
+                "Preprocessing large text with tokenizer_workers=1 may be slow with TTS tokenizer. "
+                "Prefer tokenizer_workers=(num_cpu_cores/num_gpus_per_node)"
+            )
+            for i, tokenized_text in enumerate(
+                tqdm((self._asr_text_to_tokens(text) for text in asr_texts), total=len(asr_texts))
+            ):
+                self.data[i]["asr_text_tokens"] = tokenized_text
+        else:
+            # Multiprocessing hack: use global variables for every process (not really global in program context)
+            def _init_asr_tokenize_process(tokenizer, bos_id, eos_id):
+                global asr_tokenizer_global, asr_bos_id_global, asr_eos_id_global  # process-global
+                # deepcopy to avoid serialization of parent models
+                asr_tokenizer_global = copy.deepcopy(tokenizer)
+                asr_bos_id_global = copy.deepcopy(bos_id)
+                asr_eos_id_global = copy.deepcopy(eos_id)
+
+            with concurrent.futures.ProcessPoolExecutor(
+                initializer=_init_asr_tokenize_process,
+                initargs=(asr_tokenizer, self.asr_bos_id, self.asr_eos_id),
+                max_workers=tokenizer_workers,
+            ) as pool:
+                # chunk size for pool map is empirically chosen as a trade-off between speed and responsiveness
+                for i, tokenized_text in enumerate(
+                    tqdm(pool.map(_asr_text_to_tokens, asr_texts, chunksize=1000), total=len(asr_texts))
+                ):
+                    self.data[i]["asr_text_tokens"] = tokenized_text
+        # force free memory
+        del asr_texts
+        gc.collect()
+
+        if tokenizer_workers == 1:
+            logging.warning(
+                "Preprocessing large text with tokenizer_workers=1 may be slow with TTS tokenizer. "
+                "Prefer tokenizer_workers=(num_cpu_cores/num_gpus_per_node)"
+            )
+            for i, tokenized_text in enumerate(
+                tqdm(
+                    (self._tts_text_to_tokens(text, normalize=need_normalization) for text in tts_texts),
+                    total=len(tts_texts),
+                )
+            ):
+                self.data[i]["tts_text_tokens"] = tokenized_text
+        else:
+            if need_normalization:
+                # TODO: implement, if we really need normalization inplace
+                raise NotImplementedError(
+                    "Normalization with tokenizer_workers > 1 is not implemented. "
+                    "It is not recommended to use normalization on the fly at all, since it's extremely slow"
+                )
+
+            def _init_tts_tokenize_process(tokenizer):
+                global tts_tokenizer_global  # process-global
+                tts_tokenizer_global = copy.deepcopy(tokenizer)
+
+            with concurrent.futures.ProcessPoolExecutor(
+                initializer=_init_tts_tokenize_process,
+                initargs=(tts_parser,),
+                max_workers=tokenizer_workers,
+            ) as pool:
+                # chunk size for pool map is empirically chosen as a trade-off between speed and responsiveness
+                for i, tokenized_text in enumerate(
+                    tqdm(pool.map(_tts_text_to_tokens, tts_texts, chunksize=1000), total=len(tts_texts))
+                ):
+                    self.data[i]["tts_text_tokens"] = tokenized_text
+        # force free memory
+        del tts_texts
+        gc.collect()
+
+    def _asr_text_to_tokens(self, text: str) -> np.ndarray:
+        ids = self.asr_tokenizer.text_to_ids(text)
+        if self.asr_bos_id is not None:
+            ids = [self.asr_bos_id] + ids
+        if self.asr_eos_id is not None:
+            ids.append(self.asr_eos_id)
+        return np.asarray(ids)
+
+    def _tts_text_to_tokens(self, text: str, normalize=True) -> np.ndarray:
+        if normalize:
+            text = self.tts_normalizer.normalize(text, **self.tts_normalizer_kwargs)
+        tokens = self.tts_parser(text)
+        return np.asarray(tokens)
+
+    def __getitem__(self, index):
+        item = self.data[index]
+        return TextToTextItem(
+            transcript=torch.from_numpy(item["asr_text_tokens"]).long(),
+            tts_text=torch.from_numpy(item["tts_text_tokens"]).long(),
+            speaker=random.choice(self.speakers),
+        )
+
+    def __len__(self):
+        return len(self.data)
+
+
+class TextToTextDataset(TextToTextDatasetBase, Dataset):
+    """Text-to-Text Map-style Dataset."""
+
+    def __init__(
+        self,
+        manifest_filepath: Union[AnyPath, List[AnyPath]],
+        speakers_filepath: Union[AnyPath, List[AnyPath]],
+        asr_tokenizer: TokenizerSpec,
+        asr_use_start_end_token: bool,
+        tts_parser: Callable,
+        tts_text_pad_id: int,
+        tts_text_normalizer: "Normalizer",
+        tts_text_normalizer_call_kwargs: Dict,
+        min_words: int = 1,
+        max_words: int = 1_000_000,
+        tokenizer_workers: int = 1,
+    ):
+        super().__init__(
+            manifest_filepath=manifest_filepath,
+            speakers_filepath=speakers_filepath,
+            asr_tokenizer=asr_tokenizer,
+            asr_use_start_end_token=asr_use_start_end_token,
+            tts_parser=tts_parser,
+            tts_text_pad_id=tts_text_pad_id,
+            tts_text_normalizer=tts_text_normalizer,
+            tts_text_normalizer_call_kwargs=tts_text_normalizer_call_kwargs,
+            min_words=min_words,
+            max_words=max_words,
+            tokenizer_workers=tokenizer_workers,
+            num_parts=1,
+        )
+
+    def collate_fn(
+        self, batch: List[Union[TextToTextItem, tuple]]
+    ) -> Union[TextToTextBatch, TextOrAudioToTextBatch, tuple]:
+        """
+        Collate function for dataloader
+        Can accept mixed batch of text-to-text items and audio-text items (typical for ASR)
+        """
+        return TextOrAudioToTextBatch.collate_fn(
+            batch=batch, asr_pad_id=self.asr_pad_id, tts_text_pad_id=self.tts_text_pad_id
+        )
+
+
+class TextToTextIterableDataset(TextToTextDatasetBase, IterableDataset):
+    """
+    Text-to-Text Iterable Dataset.
+    Only part necessary for current process should be loaded and stored.
+    """
+
+    def __init__(
+        self,
+        manifest_filepath: Union[AnyPath, List[AnyPath]],
+        speakers_filepath: Union[AnyPath, List[AnyPath]],
+        asr_tokenizer: TokenizerSpec,
+        asr_use_start_end_token: bool,
+        tts_parser: Callable,
+        tts_text_pad_id: int,
+        tts_text_normalizer: "Normalizer",
+        tts_text_normalizer_call_kwargs: Dict,
+        min_words: int = 1,
+        max_words: int = 1_000_000,
+        tokenizer_workers: int = 1,
+        num_parts: int = 1,
+        current_part_index: int = 0,
+    ):
+        super().__init__(
+            manifest_filepath=manifest_filepath,
+            speakers_filepath=speakers_filepath,
+            asr_tokenizer=asr_tokenizer,
+            asr_use_start_end_token=asr_use_start_end_token,
+            tts_parser=tts_parser,
+            tts_text_pad_id=tts_text_pad_id,
+            tts_text_normalizer=tts_text_normalizer,
+            tts_text_normalizer_call_kwargs=tts_text_normalizer_call_kwargs,
+            min_words=min_words,
+            max_words=max_words,
+            tokenizer_workers=tokenizer_workers,
+            num_parts=num_parts,
+            current_part_index=current_part_index,
+        )
+
+    def __iter__(self):
+        # Implementation based on docs: https://pytorch.org/docs/stable/data.html#torch.utils.data.IterableDataset
+        worker_info = torch.utils.data.get_worker_info()
+        if worker_info is None:  # single-process data loading, return the full iterator
+            start = 0
+            end = len(self)
+        else:  # in a worker process
+            # split workload
+            per_worker = int(math.ceil(len(self) / float(worker_info.num_workers)))
+            worker_id = worker_info.id
+            start = worker_id * per_worker
+            end = min(start + per_worker, len(self))
+        indices = np.arange(start, end)
+        np.random.shuffle(indices)
+        return map(self.__getitem__, indices)
+
+    def collate_fn(
+        self, batch: List[Union[TextToTextItem, tuple]]
+    ) -> Union[TextToTextBatch, TextOrAudioToTextBatch, tuple]:
+        """
+        Collate function for dataloader
+        Can accept mixed batch of text-to-text items and audio-text items (typical for ASR)
+        """
+        return TextOrAudioToTextBatch.collate_fn(
+            batch=batch, asr_pad_id=self.asr_pad_id, tts_text_pad_id=self.tts_text_pad_id
+        )

nemo/collections/asr/inference/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

nemo/collections/asr/inference/factory/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

nemo/collections/asr/inference/factory/base_builder.py

@@ -0,0 +1,203 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from omegaconf import OmegaConf, open_dict
+from omegaconf.dictconfig import DictConfig
+
+from nemo.collections.asr.inference.model_wrappers.cache_aware_ctc_inference_wrapper import (
+    CacheAwareCTCInferenceWrapper,
+)
+from nemo.collections.asr.inference.model_wrappers.cache_aware_rnnt_inference_wrapper import (
+    CacheAwareRNNTInferenceWrapper,
+)
+from nemo.collections.asr.inference.model_wrappers.ctc_inference_wrapper import CTCInferenceWrapper
+from nemo.collections.asr.inference.model_wrappers.rnnt_inference_wrapper import RNNTInferenceWrapper
+from nemo.collections.asr.inference.model_wrappers.salm_asr_inference_wrapper import SALMASRInferenceWrapper
+from nemo.collections.asr.inference.utils.enums import ASRDecodingType, PipelineType
+from nemo.collections.asr.parts.submodules.ctc_decoding import CTCDecodingConfig
+from nemo.collections.asr.parts.submodules.rnnt_decoding import RNNTDecodingConfig
+from nemo.utils import logging
+
+if TYPE_CHECKING:
+    from nemo.collections.asr.inference.itn.inverse_normalizer import AlignmentPreservingInverseNormalizer
+    from nemo.collections.asr.inference.nmt.llm_translator import LLMTranslator
+
+
+class BaseBuilder:
+    """
+    Base Builder class.
+    Builds the ASR/ITN components.
+    Derived classes should implement the `build` method which should include the logic of creating concrete pipeline.
+    """
+
+    @classmethod
+    def _build_nmt(cls, cfg: DictConfig) -> LLMTranslator | None:
+        """
+        Build the NMT model based on the config.
+        Args:
+            cfg: (DictConfig) Config
+        Returns:
+            (LLMTranslator | None) NMT model
+        """
+        nmt_model = None
+        if cfg.enable_nmt:
+            from nemo.collections.asr.inference.nmt.llm_translator import LLMTranslator
+
+            nmt_model = LLMTranslator(
+                model_name=cfg.nmt.model_name,
+                source_language=cfg.nmt.source_language,
+                target_language=cfg.nmt.target_language,
+                waitk=cfg.nmt.waitk,
+                device=cfg.nmt.device,
+                device_id=cfg.nmt.device_id,
+                batch_size=cfg.nmt.batch_size,
+                llm_params=cfg.nmt.llm_params,
+                sampling_params=cfg.nmt.sampling_params,
+            )
+            logging.info(f"NMT model `{cfg.nmt.model_name}` loaded")
+        return nmt_model
+
+    @staticmethod
+    def _apply_confidence_cfg(cfg: DictConfig, decoding_cfg: RNNTDecodingConfig) -> None:
+        """
+        Wire the separately-stored `confidence` block into the RNNT decoding confidence config so the
+        greedy or batched-beam decoder computes per-token confidence with the configured method. The streaming
+        pipelines only support non-blank confidence (`confidence.exclude_blank=true`).
+        Args:
+            cfg: (DictConfig) Full pipeline config (provides the top-level `confidence` block).
+            decoding_cfg: (RNNTDecodingConfig) Decoding config to update in place.
+        """
+        preserve_frame_confidence = decoding_cfg.greedy.get(
+            "preserve_frame_confidence", False
+        ) or decoding_cfg.beam.get("preserve_frame_confidence", False)
+        if not preserve_frame_confidence:
+            return
+        confidence_cfg = cfg.get("confidence", None)
+        if confidence_cfg is None:
+            return
+        if not confidence_cfg.get("exclude_blank", True):
+            raise ValueError(
+                "Streaming confidence supports only non-blank confidence (`confidence.exclude_blank=true`)."
+            )
+        decoding_cfg.confidence_cfg.preserve_frame_confidence = True
+        decoding_cfg.confidence_cfg.preserve_token_confidence = True
+        decoding_cfg.confidence_cfg.preserve_word_confidence = True
+        decoding_cfg.confidence_cfg.exclude_blank = True
+        decoding_cfg.confidence_cfg.aggregation = confidence_cfg.get("aggregation", "mean")
+        decoding_cfg.confidence_cfg.method_cfg = OmegaConf.merge(
+            decoding_cfg.confidence_cfg.method_cfg, confidence_cfg.method_cfg
+        )
+
+    @classmethod
+    def _build_asr(cls, cfg: DictConfig, decoding_cfg: CTCDecodingConfig | RNNTDecodingConfig | None) -> Any:
+        """
+        Build the ASR model based on the config.
+        Args:
+            cfg: (DictConfig) Config
+            decoding_cfg: (CTCDecodingConfig | RNNTDecodingConfig | None) Decoding config
+        Returns:
+            (Any) ASR inference model
+        """
+
+        asr_decoding_type = ASRDecodingType.from_str(cfg.asr_decoding_type)
+        pipeline_type = PipelineType.from_str(cfg.pipeline_type)
+        model_params = {
+            "model_name": cfg.asr.model_name,
+            "device": cfg.asr.device,
+            "device_id": cfg.asr.device_id,
+            "compute_dtype": cfg.asr.compute_dtype,
+            "use_amp": cfg.asr.use_amp,
+            "decoding_cfg": decoding_cfg,
+        }
+        match (asr_decoding_type, pipeline_type):
+            case (ASRDecodingType.CTC, PipelineType.BUFFERED):
+                asr_class = CTCInferenceWrapper
+            case (ASRDecodingType.RNNT, PipelineType.BUFFERED):
+                asr_class = RNNTInferenceWrapper
+            case (ASRDecodingType.SALM, PipelineType.BUFFERED):
+                asr_class = SALMASRInferenceWrapper
+                # remove decoding_cfg, SALM AED does not use decoding_cfg yet
+                model_params.pop("decoding_cfg")
+            case (ASRDecodingType.CTC, PipelineType.CACHE_AWARE):
+                asr_class = CacheAwareCTCInferenceWrapper
+            case (ASRDecodingType.RNNT, PipelineType.CACHE_AWARE):
+                asr_class = CacheAwareRNNTInferenceWrapper
+            case _:
+                raise ValueError(
+                    f"Wrong combination of ASR decoding type and pipeline type: {asr_decoding_type, pipeline_type}"
+                )
+
+        asr_model = asr_class(**model_params)
+        logging.info(f"ASR model `{cfg.asr.model_name}` loaded")
+        return asr_model
+
+    @classmethod
+    def _build_itn(cls, cfg: DictConfig, input_is_lower_cased: bool) -> AlignmentPreservingInverseNormalizer | None:
+        """
+        Build the ITN model based on the config.
+        Args:
+            cfg: (DictConfig) Config
+            input_is_lower_cased: (bool) Whether the input is lower cased
+        Returns:
+            (AlignmentPreservingInverseNormalizer | None) ITN model
+        """
+        itn_model = None
+        if cfg.enable_itn:
+            # Do not remove this import. It is used to avoid nemo_text_processing import when verbatim transcripts is enabled.
+            from nemo.collections.asr.inference.itn.inverse_normalizer import AlignmentPreservingInverseNormalizer
+
+            input_case = (
+                AlignmentPreservingInverseNormalizer.LOWER_CASED
+                if input_is_lower_cased
+                else AlignmentPreservingInverseNormalizer.UPPER_CASED
+            )
+
+            target_lang = getattr(cfg, "lang", getattr(cfg, "target_lang", None))
+            if target_lang is None:
+                raise ValueError("Language is not specified. Cannot load ITN model.")
+
+            itn_cfg = cfg.itn
+            with open_dict(itn_cfg):
+                itn_cfg.lang = target_lang
+                itn_cfg.input_case = input_case
+                itn_cfg.cache_dir = cfg.cache_dir
+
+            itn_model = AlignmentPreservingInverseNormalizer(
+                lang=itn_cfg.lang,
+                input_case=itn_cfg.input_case,
+                whitelist=itn_cfg.whitelist,
+                cache_dir=itn_cfg.cache_dir,
+                overwrite_cache=itn_cfg.overwrite_cache,
+                max_number_of_permutations_per_split=itn_cfg.max_number_of_permutations_per_split,
+            )
+            logging.info(f"Built inverse text normalizer with the input case: `{input_case}`.")
+
+        if itn_model is not None:
+            logging.info("ITN model loaded")
+        return itn_model
+
+    @classmethod
+    def build(cls, cfg: DictConfig) -> Any:
+        """
+        Build the pipeline based on the config.
+        Args:
+            cfg: (DictConfig) Config
+        Returns:
+            Returns object responsible for the inference
+        """
+        raise NotImplementedError("This method should be implemented in subclasses.")

nemo/collections/asr/inference/factory/buffered_pipeline_builder.py

@@ -0,0 +1,148 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from omegaconf import DictConfig, OmegaConf
+
+from nemo.collections.asr.inference.factory.base_builder import BaseBuilder
+from nemo.collections.asr.inference.pipelines.buffered_ctc_pipeline import BufferedCTCPipeline
+from nemo.collections.asr.inference.pipelines.buffered_rnnt_pipeline import BufferedRNNTPipeline
+from nemo.collections.asr.inference.pipelines.buffered_salm_pipeline import BufferedSALMPipeline
+from nemo.collections.asr.inference.utils.enums import ASRDecodingType
+from nemo.collections.asr.parts.submodules.ctc_decoding import CTCDecodingConfig
+from nemo.collections.asr.parts.submodules.rnnt_decoding import RNNTDecodingConfig
+from nemo.utils import logging
+
+
+class BufferedPipelineBuilder(BaseBuilder):
+    """
+    Buffered Pipeline Builder class.
+    Builds:
+        - buffered CTC/RNNT/TDT pipelines.
+        - buffered SALM pipeline.
+    """
+
+    @classmethod
+    def build(cls, cfg: DictConfig) -> BufferedRNNTPipeline | BufferedCTCPipeline | BufferedSALMPipeline:
+        """
+        Build the buffered streaming pipeline based on the config.
+        Args:
+            cfg: (DictConfig) Config
+        Returns:
+            Returns BufferedRNNTPipeline, BufferedCTCPipeline, or BufferedSALMPipeline object
+        """
+        asr_decoding_type = ASRDecodingType.from_str(cfg.asr_decoding_type)
+
+        if asr_decoding_type is ASRDecodingType.RNNT:
+            return cls.build_buffered_rnnt_pipeline(cfg)
+        elif asr_decoding_type is ASRDecodingType.CTC:
+            return cls.build_buffered_ctc_pipeline(cfg)
+        elif asr_decoding_type is ASRDecodingType.SALM:
+            return cls.build_buffered_salm_pipeline(cfg)
+
+        raise ValueError("Invalid asr decoding type for buffered streaming. Need to be one of ['CTC', 'RNNT', 'SALM']")
+
+    @classmethod
+    def get_rnnt_decoding_cfg(cls, cfg: DictConfig) -> RNNTDecodingConfig:
+        """
+        Get the decoding config for the RNNT pipeline.
+        Returns:
+            (RNNTDecodingConfig) Decoding config
+        """
+        base_cfg_structured = OmegaConf.structured(RNNTDecodingConfig)
+        base_cfg = OmegaConf.create(OmegaConf.to_container(base_cfg_structured))
+        decoding_cfg = OmegaConf.merge(base_cfg, cfg.asr.decoding)
+        cls._apply_confidence_cfg(cfg, decoding_cfg)
+        return decoding_cfg
+
+    @classmethod
+    def get_ctc_decoding_cfg(cls) -> CTCDecodingConfig:
+        """
+        Get the decoding config for the CTC pipeline.
+        Returns:
+            (CTCDecodingConfig) Decoding config
+        """
+        decoding_cfg = CTCDecodingConfig()
+        decoding_cfg.strategy = "greedy"
+        return decoding_cfg
+
+    @classmethod
+    def build_buffered_rnnt_pipeline(cls, cfg: DictConfig) -> BufferedRNNTPipeline:
+        """
+        Build the RNNT streaming pipeline based on the config.
+        Args:
+            cfg: (DictConfig) Config
+        Returns:
+            Returns BufferedRNNTPipeline object
+        """
+        # building ASR model
+        decoding_cfg = cls.get_rnnt_decoding_cfg(cfg)
+        asr_model = cls._build_asr(cfg, decoding_cfg)
+
+        # building ITN model
+        itn_model = cls._build_itn(cfg, input_is_lower_cased=True)
+
+        # building NMT model
+        nmt_model = cls._build_nmt(cfg)
+
+        # building RNNT pipeline
+        rnnt_pipeline = BufferedRNNTPipeline(cfg, asr_model, itn_model, nmt_model)
+        logging.info(f"`{type(rnnt_pipeline).__name__}` pipeline loaded")
+        return rnnt_pipeline
+
+    @classmethod
+    def build_buffered_ctc_pipeline(cls, cfg: DictConfig) -> BufferedCTCPipeline:
+        """
+        Build the CTC buffered streaming pipeline based on the config.
+        Args:
+            cfg: (DictConfig) Config
+        Returns:
+            Returns BufferedCTCPipeline object
+        """
+        # building ASR model
+        decoding_cfg = cls.get_ctc_decoding_cfg()
+        asr_model = cls._build_asr(cfg, decoding_cfg)
+
+        # building ITN model
+        itn_model = cls._build_itn(cfg, input_is_lower_cased=True)
+
+        # building NMT model
+        nmt_model = cls._build_nmt(cfg)
+
+        # building CTC pipeline
+        ctc_pipeline = BufferedCTCPipeline(cfg, asr_model, itn_model, nmt_model)
+        logging.info(f"`{type(ctc_pipeline).__name__}` pipeline loaded")
+        return ctc_pipeline
+
+    @classmethod
+    def build_buffered_salm_pipeline(cls, cfg: DictConfig) -> BufferedSALMPipeline:
+        """
+        Build the buffered SALM pipeline based on the config.
+        Args:
+            cfg: (DictConfig) Config
+        Returns:
+            Returns BufferedSALMPipeline object
+        """
+        # building ASR model
+        asr_model = cls._build_asr(cfg, decoding_cfg=None)
+
+        # building ITN model
+        itn_model = cls._build_itn(cfg, input_is_lower_cased=True)
+
+        # building NMT model
+        nmt_model = cls._build_nmt(cfg)
+
+        # building SALM pipeline
+        salm_pipeline = BufferedSALMPipeline(cfg, asr_model, itn_model, nmt_model)
+        logging.info(f"`{type(salm_pipeline).__name__}` pipeline loaded")
+        return salm_pipeline

nemo/collections/asr/inference/factory/cache_aware_pipeline_builder.py

@@ -0,0 +1,128 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from omegaconf import DictConfig, OmegaConf
+
+from nemo.collections.asr.inference.factory.base_builder import BaseBuilder
+from nemo.collections.asr.inference.pipelines.cache_aware_ctc_pipeline import CacheAwareCTCPipeline
+from nemo.collections.asr.inference.pipelines.cache_aware_rnnt_pipeline import CacheAwareRNNTPipeline
+from nemo.collections.asr.inference.utils.enums import ASRDecodingType
+from nemo.collections.asr.parts.submodules.ctc_decoding import CTCDecodingConfig
+from nemo.collections.asr.parts.submodules.rnnt_decoding import RNNTDecodingConfig
+from nemo.utils import logging
+
+
+class CacheAwarePipelineBuilder(BaseBuilder):
+    """
+    Cache Aware Pipeline Builder class.
+    Builds the cache aware CTC/RNNT pipelines.
+    """
+
+    @classmethod
+    def build(cls, cfg: DictConfig) -> CacheAwareCTCPipeline | CacheAwareRNNTPipeline:
+        """
+        Build the cache aware streaming pipeline based on the config.
+        Args:
+            cfg: (DictConfig) Config
+        Returns:
+            Returns CacheAwareCTCPipeline or CacheAwareRNNTPipeline object
+        """
+        asr_decoding_type = ASRDecodingType.from_str(cfg.asr_decoding_type)
+
+        if asr_decoding_type is ASRDecodingType.RNNT:
+            return cls.build_cache_aware_rnnt_pipeline(cfg)
+        elif asr_decoding_type is ASRDecodingType.CTC:
+            return cls.build_cache_aware_ctc_pipeline(cfg)
+
+        raise ValueError("Invalid asr decoding type for cache aware streaming. Need to be one of ['CTC', 'RNNT']")
+
+    @classmethod
+    def get_rnnt_decoding_cfg(cls, cfg: DictConfig) -> RNNTDecodingConfig:
+        """
+        Get the decoding config for the RNNT pipeline.
+        Returns:
+            (RNNTDecodingConfig) Decoding config
+        """
+        base_cfg_structured = OmegaConf.structured(RNNTDecodingConfig)
+        base_cfg = OmegaConf.create(OmegaConf.to_container(base_cfg_structured))
+        decoding_cfg = OmegaConf.merge(base_cfg, cfg.asr.decoding)
+        cls._apply_confidence_cfg(cfg, decoding_cfg)
+        return decoding_cfg
+
+    @classmethod
+    def get_ctc_decoding_cfg(cls) -> CTCDecodingConfig:
+        """
+        Get the decoding config for the CTC pipeline.
+        Returns:
+            (CTCDecodingConfig) Decoding config
+        """
+        decoding_cfg = CTCDecodingConfig()
+        decoding_cfg.strategy = "greedy"
+        decoding_cfg.preserve_alignments = False
+        return decoding_cfg
+
+    @classmethod
+    def build_cache_aware_rnnt_pipeline(cls, cfg: DictConfig) -> CacheAwareRNNTPipeline:
+        """
+        Build the cache aware RNNT streaming pipeline based on the config.
+        Args:
+            cfg: (DictConfig) Config
+        Returns:
+            Returns CacheAwareRNNTPipeline object
+        """
+        strategy = str(getattr(cfg.asr.decoding, "strategy", "greedy_batch"))
+        if strategy not in {"greedy_batch", "malsd_batch"}:
+            raise ValueError(
+                "Cache-aware RNNT streaming supports `greedy_batch` and `malsd_batch` only; "
+                f"configured decoding strategy is `{strategy}`."
+            )
+
+        # building ASR model
+        decoding_cfg = cls.get_rnnt_decoding_cfg(cfg)
+        asr_model = cls._build_asr(cfg, decoding_cfg)
+
+        # building ITN model
+        itn_model = cls._build_itn(cfg, input_is_lower_cased=True)
+
+        # building NMT model
+        nmt_model = cls._build_nmt(cfg)
+
+        # building cache aware RNNT pipeline
+        ca_rnnt_pipeline = CacheAwareRNNTPipeline(cfg, asr_model, itn_model, nmt_model)
+        logging.info(f"`{type(ca_rnnt_pipeline).__name__}` pipeline loaded")
+        return ca_rnnt_pipeline
+
+    @classmethod
+    def build_cache_aware_ctc_pipeline(cls, cfg: DictConfig) -> CacheAwareCTCPipeline:
+        """
+        Build the cache aware CTC streaming pipeline based on the config.
+        Args:
+            cfg: (DictConfig) Config
+        Returns:
+            Returns CacheAwareCTCPipeline object
+        """
+        # building ASR model
+        decoding_cfg = cls.get_ctc_decoding_cfg()
+        asr_model = cls._build_asr(cfg, decoding_cfg)
+
+        # building ITN model
+        itn_model = cls._build_itn(cfg, input_is_lower_cased=True)
+
+        # building NMT model
+        nmt_model = cls._build_nmt(cfg)
+
+        # building cache aware CTC pipeline
+        ca_ctc_pipeline = CacheAwareCTCPipeline(cfg, asr_model, itn_model, nmt_model)
+        logging.info(f"`{type(ca_ctc_pipeline).__name__}` pipeline loaded")
+        return ca_ctc_pipeline

nemo/collections/asr/inference/factory/pipeline_builder.py

@@ -0,0 +1,74 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import torch
+from omegaconf.dictconfig import DictConfig
+
+from nemo.collections.asr.inference.factory.buffered_pipeline_builder import BufferedPipelineBuilder
+from nemo.collections.asr.inference.factory.cache_aware_pipeline_builder import CacheAwarePipelineBuilder
+from nemo.collections.asr.inference.pipelines.base_pipeline import BasePipeline
+from nemo.collections.asr.inference.utils.enums import PipelineType
+from nemo.utils import logging
+
+
+class PipelineBuilder:
+    """Router for building the pipeline based on the pipeline type."""
+
+    @staticmethod
+    def set_matmul_precision(matmul_precision: str) -> None:
+        """
+        Set the matmul precision.
+        Args:
+            matmul_precision: (str) Matmul precision: highest, high, medium
+        """
+        choices = ["highest", "high", "medium"]
+        matmul_precision = matmul_precision.lower()
+        if matmul_precision not in choices:
+            raise ValueError(f"Invalid matmul precision: {matmul_precision}. Need to be one of {choices}")
+        torch.set_float32_matmul_precision(matmul_precision)
+        logging.info(f"Using matmul precision: {matmul_precision}")
+
+    @staticmethod
+    def set_log_level(log_level: int) -> None:
+        """
+        Set the logging level.
+        Args:
+            log_level: (int) Logging level: 0 (NOTSET), 10 (DEBUG), 20 (INFO), 30 (WARNING), 40 (ERROR), 50 (CRITICAL)
+        """
+        choices = [0, 10, 20, 30, 40, 50]
+        if log_level not in choices:
+            raise ValueError(f"Invalid log level: {log_level}. Need to be one of {choices}")
+        logging.setLevel(log_level)
+
+    @staticmethod
+    def build_pipeline(cfg: DictConfig) -> BasePipeline:
+        """
+        Build the pipeline based on the config.
+        Args:
+            cfg: (DictConfig) Config
+        Returns:
+            Returns Pipeline object
+        """
+        PipelineBuilder.set_log_level(cfg.log_level)
+        PipelineBuilder.set_matmul_precision(cfg.matmul_precision)
+        pipeline_type = PipelineType.from_str(cfg.pipeline_type)
+        if pipeline_type is PipelineType.BUFFERED:
+            builder = BufferedPipelineBuilder
+        elif pipeline_type is PipelineType.CACHE_AWARE:
+            builder = CacheAwarePipelineBuilder
+        else:
+            raise ValueError(f"Invalid pipeline type: {cfg.pipeline_type}")
+
+        return builder.build(cfg)

nemo/collections/asr/inference/itn/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

nemo/collections/asr/inference/itn/batch_inverse_normalizer.py

@@ -0,0 +1,185 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import itertools
+from typing import Callable
+
+from joblib import Parallel, delayed
+
+from nemo.collections.asr.inference.itn.inverse_normalizer import AlignmentPreservingInverseNormalizer
+from nemo.collections.asr.inference.utils.text_segment import Word
+
+
+def merge_punctuation_and_itn_tags(
+    input_words: list[str],
+    output_words: list[str],
+    word_alignment: list[tuple],
+    pnc_words: list[Word],
+    punct_marks: set,
+    sep: str,
+    conf_aggregate_fn: Callable,
+) -> list[Word]:
+    """
+    Merge the punctuation marks and ITN tags to the final text.
+    It will also preserve first letter capitalization, start and end time of the span.
+    Args:
+        input_words: (list[str]) List of input words
+        output_words: (list[str]) List of output words
+        word_alignment: (list[tuple[list[int], list[int]]]) Word alignment between the input and output words
+        pnc_words: (list[Word]) List of words with punctuation marks
+        punct_marks: (set) Punctuation marks
+        sep: (str) Separator
+        conf_aggregate_fn: (Callable) Confidence aggregation function
+    Returns:
+        (list[Word]) Final words after merging the punctuation marks and ITN tags
+    """
+    assert len(input_words) == len(pnc_words)
+    spans = []
+    for s_idx, t_idx, semiotic_class in word_alignment:
+        if len(t_idx) == 1 and len(s_idx) == 1 and input_words[s_idx[0]] == output_words[t_idx[0]]:
+            span = pnc_words[s_idx[0]]
+            span.semiotic_class = semiotic_class
+        else:
+            span_text = sep.join([output_words[i] for i in t_idx])
+            last_char = pnc_words[s_idx[-1]].text[-1]
+            first_char = pnc_words[s_idx[0]].text[0]
+
+            # preserve the first char capitalization
+            first_word = pnc_words[s_idx[0]].copy()
+            first_char_is_upper = first_word.text[0].isupper()
+            first_word.normalize_text_inplace(punct_marks, sep)
+            if span_text.startswith(first_word.text):
+                if first_char_is_upper:
+                    span_text = span_text[0].upper() + span_text[1:]
+
+            # preserve the last punctuation mark
+            if last_char in punct_marks:
+                span_text += last_char
+
+            # preserve the first punctuation mark
+            if first_char in punct_marks:
+                span_text = first_char + span_text
+
+            scores = [pnc_words[i].conf for i in s_idx]
+            conf = conf_aggregate_fn(scores) if len(scores) > 0 else 0.0
+            span = Word(
+                text=span_text,
+                start=pnc_words[s_idx[0]].start,
+                end=pnc_words[s_idx[-1]].end,
+                semiotic_class=semiotic_class,
+                conf=conf,
+            )
+        spans.append(span)
+    return spans
+
+
+class BatchAlignmentPreservingInverseNormalizer:
+    """
+    Batch Alignment Preserving Inverse Text Normalizer. It is used to apply ITN to a batch of texts.
+    joblib.Parallel is used to parallelize the processing.
+    """
+
+    def __init__(
+        self,
+        itn_model: AlignmentPreservingInverseNormalizer,
+        sep: str,
+        asr_supported_puncts: set[str],
+        post_word_punctuation: set[str],
+        conf_aggregate_fn: Callable,
+    ):
+        """
+        Batch Alignment Preserving Inverse Text Normalizer. It is used to apply ITN to a batch of texts.
+        Args:
+            itn_model: (AlignmentPreservingInverseNormalizer) Alignment Preserving Inverse Text Normalizer
+            sep: (str) Separator
+            asr_supported_puncts: (Set[str]) Punctuation marks supported by ASR model
+            post_word_punctuation: (Set[str]) Punctuation marks which usually appear after a word
+            conf_aggregate_fn: (Callable) Confidence aggregation function
+        """
+        self.itn_model = itn_model
+        self.sep = sep
+        self.asr_supported_puncts = asr_supported_puncts
+        self.conf_aggregate_fn = conf_aggregate_fn
+        self.punct_marks = self.asr_supported_puncts | post_word_punctuation
+
+    def apply_itn(
+        self, asr_words: list[Word], pnc_words: list[Word], return_alignment: bool = False
+    ) -> list[Word] | tuple[list[Word], list]:
+        """
+        Apply Alignment Preserving Inverse Text Normalization.
+        Args:
+            asr_words: (list[Word]) List of ASR words
+            pnc_words: (list[Word]) List of words with punctuation/capitalization
+            return_alignment: (bool) Flag to return the word alignment
+        Returns:
+            (list[Word]) List of words after applying ITN
+        """
+        input_words = []
+        for word in asr_words:
+            word.normalize_text_inplace(self.asr_supported_puncts, self.sep)
+            input_words.append(word.text)
+
+        input_words, output_words, word_alignment = self.itn_model.get_word_alignment(input_words, sep=self.sep)
+        spans = merge_punctuation_and_itn_tags(
+            input_words, output_words, word_alignment, pnc_words, self.punct_marks, self.sep, self.conf_aggregate_fn
+        )
+
+        if return_alignment:
+            # word alignment is needed for streaming inference
+            return spans, word_alignment
+        return spans
+
+    def __call__(
+        self,
+        asr_words_list: list[list[Word]],
+        pnc_words_list: list[list[Word]],
+        itn_params: dict,
+        return_alignment: bool = False,
+    ) -> list[list[Word]] | list[tuple]:
+        """
+        Batch Alignment Preserving Inverse Text Normalization.
+        Args:
+            asr_words_list: (list[list[Word]]) List of ASR words
+            pnc_words_list: (list[list[Word]]) List of words with punctuation/capitalization
+            itn_params: (dict) Parameters for the ITN model
+            return_alignment: (bool) Flag to return the word alignment
+        Returns:
+            (list[list[Word]]) List of words after applying ITN
+        """
+        if len(asr_words_list) == 0:
+            return []
+
+        batch_size = itn_params.get("batch_size", 1)
+        n_texts = len(asr_words_list)
+        batch_size = min(n_texts, batch_size)
+
+        def process_batch(batch_words, batch_words_with_pnc):
+            return [
+                self.apply_itn(words, words_with_pnc, return_alignment)
+                for words, words_with_pnc in zip(batch_words, batch_words_with_pnc)
+            ]
+
+        if n_texts <= 3 * batch_size or n_texts == 1:
+            # If the number of texts is less than 3 * batch_size, process the batch sequentially
+            # For small batch size, it is faster to process the batch sequentially
+            return process_batch(asr_words_list, pnc_words_list)
+
+        n_jobs = itn_params.get("n_jobs", 1)
+        itn_words_list = Parallel(n_jobs=n_jobs)(
+            delayed(process_batch)(asr_words_list[i : i + batch_size], pnc_words_list[i : i + batch_size])
+            for i in range(0, n_texts, batch_size)
+        )
+        itn_words_list = list(itertools.chain(*itn_words_list))
+        return itn_words_list

nemo/collections/asr/inference/itn/inverse_normalizer.py

@@ -0,0 +1,355 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import os
+import re
+from multiprocessing import Manager
+
+from nemo.collections.asr.inference.utils.itn_utils import (
+    DEFAULT_SEMIOTIC_CLASS,
+    fallback_to_trivial_alignment,
+    find_tokens,
+    get_semiotic_class,
+    split_text,
+)
+from nemo.utils import logging
+
+IN_MEM_CACHE = Manager().dict(lock=False)
+
+try:
+    import pynini
+    from nemo_text_processing.inverse_text_normalization.inverse_normalize import InverseNormalizer, Normalizer
+    from nemo_text_processing.text_normalization.en.graph_utils import INPUT_CASED, INPUT_LOWER_CASED
+
+    HAVE_NEMO_TEXT_PROCESSING = True
+except ImportError:
+    INPUT_CASED, INPUT_LOWER_CASED = "undefined", "undefined"
+    HAVE_NEMO_TEXT_PROCESSING = False
+
+try:
+    import diskcache
+
+    CACHING_FROM_DISK = True
+except ImportError:
+    logging.warning("diskcache package is not installed, caching from disk is disabled")
+    CACHING_FROM_DISK = False
+
+
+class AlignmentPreservingInverseNormalizer:
+    """
+    Inverse Text Normalizer that preserves the word alignment.
+    It is used to convert the spoken text to written text and preserve the alignment between the input and output words.
+    """
+
+    LOWER_CASED = INPUT_LOWER_CASED
+    UPPER_CASED = INPUT_CASED
+    GRAMMAR = "itn"
+
+    def __init__(
+        self,
+        input_case: str = LOWER_CASED,
+        lang: str = "en",
+        whitelist: str = None,
+        cache_dir: str = None,
+        overwrite_cache: bool = False,
+        max_number_of_permutations_per_split: int = 729,
+    ):
+        """
+        Inverse normalizer that converts text from spoken to written form.
+        Args:
+            input_case: Input text capitalization, set to 'cased' if text contains capital letters.
+                This flag affects normalization rules applied to the text. Note, `lower_cased` won't lower case input.
+            lang: language specifying the ITN
+            whitelist: path to a file with whitelist replacements. (each line of the file: written_form\tspoken_form\n),
+                e.g. nemo_text_processing/inverse_text_normalization/en/data/whitelist.tsv
+            cache_dir: path to a dir with .far grammar file. Set to None to avoid using cache.
+            overwrite_cache: set to True to overwrite .far files
+            max_number_of_permutations_per_split: a maximum number
+                of permutations which can be generated from input sequence of tokens.
+        """
+        if not HAVE_NEMO_TEXT_PROCESSING:
+            raise RuntimeError(
+                "In order to use AlignmentPreservingInverseNormalizer, "
+                "please install the nemo_text_processing and pynini packages."
+            )
+        self.itn_model = InverseNormalizer(
+            lang=lang,
+            input_case=input_case,
+            whitelist=whitelist,
+            cache_dir=cache_dir,
+            overwrite_cache=overwrite_cache,
+            max_number_of_permutations_per_split=max_number_of_permutations_per_split,
+        )
+        if cache_dir and CACHING_FROM_DISK:
+            self.DISK_TAG_CACHE = diskcache.Cache(os.path.join(cache_dir, "itn_tag_cache"))
+            self.DISK_VERB_CACHE = diskcache.Cache(os.path.join(cache_dir, "itn_verb_cache"))
+            self.caching_from_disk_enabled = True
+        else:
+            self.DISK_TAG_CACHE = None
+            self.DISK_VERB_CACHE = None
+            self.caching_from_disk_enabled = False
+
+    def inverse_normalize_list(self, texts: list[str], params: dict) -> list[str]:
+        """
+        Applies Inverse Text Normalization to the list of texts.
+        Args:
+            texts: (list[str]) list of input strings.
+            params: (dict) dictionary of runtime parameters.
+        Returns:
+            (list[str]) Returns converted list of input strings.
+        """
+        normalized_texts = self.itn_model.normalize_list(
+            texts,
+            verbose=params.get('verbose', False),
+            punct_pre_process=params.get("punct_pre_process", False),
+            punct_post_process=params.get("punct_post_process", False),
+            batch_size=params.get("batch_size", 1),
+            n_jobs=params.get("n_jobs", 1),
+        )
+        return normalized_texts
+
+    def verbalize(self, tokens: list, sep: str) -> str | None:
+        """
+        Appplies verbalization to the list of tokens.
+        Args:
+            tokens: (list) list of tokens
+            sep: (str) word separator
+        Returns:
+            (str | None) Returns verbalized text. If verbalization fails, returns None.
+        """
+        split_tokens = self.itn_model._split_tokens_to_reduce_number_of_permutations(tokens)
+        output_str = ""
+        for s in split_tokens:
+            try:
+                tags_reordered = self.itn_model.generate_permutations(s)
+                verbalizer_lattice = None
+                for tagged_text_r in tags_reordered:
+                    tagged_text_r = pynini.escape(tagged_text_r)
+
+                    verbalizer_lattice = self.itn_model.find_verbalizer(tagged_text_r)
+                    if verbalizer_lattice.num_states() != 0:
+                        break
+
+                if verbalizer_lattice is None:
+                    return None
+
+                verbalized_text = Normalizer.select_verbalizer(verbalizer_lattice)
+                output_str += sep + verbalized_text
+            except Exception as e:
+                logging.warning("Failed to verbalize tagged text: " + str(e))
+                return None
+
+        output_str = output_str.strip(sep)
+        return re.sub(r"({sep})+".format(sep=sep), sep, output_str)
+
+    def tag(self, text: str, no_cache: bool = False) -> str:
+        """
+        Tags the input text.
+        Args:
+            text: (str) input text
+            no_cache: (bool) whether to use cache
+        Returns:
+            (str) tagged text
+        """
+        if not no_cache:
+            # In-memory cache check
+            if text in IN_MEM_CACHE:
+                return IN_MEM_CACHE[text]
+
+            # Disk cache check
+            if self.caching_from_disk_enabled and text in self.DISK_TAG_CACHE:
+                x = self.DISK_TAG_CACHE[text]
+                IN_MEM_CACHE[text] = x
+                return x
+
+        text = text.strip()
+        if not text:
+            return text
+
+        text = pynini.escape(text)
+        tagged_lattice = self.itn_model.find_tags(text)
+        tagged_text = Normalizer.select_tag(tagged_lattice)
+        IN_MEM_CACHE[text] = tagged_text
+        if self.caching_from_disk_enabled:
+            self.DISK_TAG_CACHE[text] = tagged_text
+        return tagged_text
+
+    def parse_and_verbalize(self, tagged_text: str, sep: str) -> tuple[str, str]:
+        """
+        Tags and verbalizes the input text.
+        Args:
+            tagged_text: (str) tagged input text
+            sep: (str) word separator
+        Returns:
+            (str, str) Returns the verbalized text, and the semiotic class.
+        """
+
+        # In-memory cache check
+        if tagged_text in IN_MEM_CACHE:
+            return IN_MEM_CACHE[tagged_text]
+
+        # Disk cache check
+        if self.caching_from_disk_enabled and tagged_text in self.DISK_VERB_CACHE:
+            x = self.DISK_VERB_CACHE[tagged_text]
+            IN_MEM_CACHE[tagged_text] = x
+            return x
+
+        self.itn_model.parser(tagged_text)
+        tokens = self.itn_model.parser.parse()
+        span_text = self.verbalize(tokens, sep)
+        semiotic_class = DEFAULT_SEMIOTIC_CLASS if span_text is None else get_semiotic_class(tokens)
+
+        IN_MEM_CACHE[tagged_text] = (span_text, semiotic_class)
+        if self.caching_from_disk_enabled:
+            self.DISK_VERB_CACHE[tagged_text] = (span_text, semiotic_class)
+        return span_text, semiotic_class
+
+    def find_token_words(
+        self, token: str, start_idx: int, input_words: list[str], sep: str
+    ) -> tuple[list[int], bool, int]:
+        """
+        Finds the words that make up the token.
+        Args:
+            token: (str) token
+            start_idx: (int) start index
+            input_words: (list[str]) list of input words
+            sep: (str) word separator
+        Returns:
+            (tuple) Returns a tuple of indices, success, and the new start index
+        """
+        indices, tmp_text, success = [], "", False
+        length = len(input_words)
+        for i in range(start_idx, length):
+            tmp_text = tmp_text + sep + input_words[i] if tmp_text else input_words[i]
+            tmp_tagged_text = self.tag(tmp_text)
+
+            if tmp_tagged_text == token:
+                indices.append(i)
+
+                # Try to extend the token by one word
+                if i + 1 < length:
+                    extended_tmp_text = tmp_text + sep + input_words[i + 1]
+                    extended_tmp_tagged_text = self.tag(extended_tmp_text)
+                    if extended_tmp_tagged_text == token:
+                        continue
+
+                success = True
+                break
+            else:
+                indices.append(i)
+
+        return indices, success, i
+
+    def find_alignment(
+        self,
+        tokens: list[str],
+        input_words: list[str],
+        sep: str,
+        iwords: list[str],
+        owords: list[str],
+        word_alignment: list[tuple],
+    ) -> bool:
+        """
+        Finds the word alignment for the input text.
+        Args:
+            tokens: (list[str]) list of tokens
+            input_words: (list[str]) list of input words
+            sep: (str) word separator
+            iwords: (list[str]) list of input words to be updated
+            owords: (list[str]) list of output words to be updated
+            word_alignment: (list[tuple]) list of word alignments to be updated
+        Returns:
+            (bool) True if the word alignment is found, False otherwise
+        """
+        success = True
+        token_start_idx = word_start_idx = 0
+        iwords_len = owords_len = 0
+
+        while token_start_idx < len(tokens):
+            token = tokens[token_start_idx]
+            current_word = input_words[word_start_idx]
+            if token == f"tokens {{ name: \"{current_word}\" }}":
+                word_alignment.append(([iwords_len], [owords_len], DEFAULT_SEMIOTIC_CLASS))
+                iwords.append(current_word)
+                owords.append(current_word)
+                iwords_len += 1
+                owords_len += 1
+            else:
+                indices, success, word_start_idx = self.find_token_words(token, word_start_idx, input_words, sep)
+                if success:
+                    span_text, semiotic_class = self.parse_and_verbalize(token, sep)
+                    if span_text is None:
+                        logging.warning(f"Failed to verbalize the token: {token}")
+                        return False
+
+                    span_words, n_span_words = split_text(span_text, sep)
+                    word_alignment.append(
+                        (
+                            [iwords_len + i for i in range(len(indices))],
+                            [owords_len + i for i in range(n_span_words)],
+                            semiotic_class,
+                        )
+                    )
+                    owords.extend(span_words)
+                    iwords.extend([input_words[i] for i in indices])
+                    iwords_len += len(indices)
+                    owords_len += n_span_words
+                else:
+                    success = False
+                    break
+
+            token_start_idx += 1
+            word_start_idx += 1
+
+        return success
+
+    def get_word_alignment(self, input: str | list[str], sep: str) -> tuple[list[str], list[str], list[tuple]]:
+        """
+        Returns a word alignment for the input text.
+        Args:
+            input: (str | list[str]) input text or list of input words
+            sep: (str) word separator
+        Returns:
+            (tuple) Returns a tuple of input words, output words, and a word alignment between input and output words
+        """
+
+        if isinstance(input, str):
+            input_text = input
+            input_words, n_words = split_text(input_text, sep)
+        else:
+            input_words, n_words = input, len(input)
+            input_text = sep.join(input_words)
+
+        # If input_text is empty, return empty lists
+        if n_words == 0:
+            return [], [], []
+
+        # Tag the input text
+        tagged_text = self.tag(input_text, no_cache=False)
+
+        # Find the tokens in the tagged text
+        tokens = find_tokens(tagged_text)
+
+        # Find the word alignment
+        iwords, owords, word_alignment = [], [], []
+        success = self.find_alignment(
+            tokens, input_words, sep, iwords=iwords, owords=owords, word_alignment=word_alignment
+        )
+
+        # If the word alignment is not found, fallback to the trivial alignment
+        if not success:
+            return fallback_to_trivial_alignment(input_words, i_shift=0, o_shift=0)
+
+        return iwords, owords, word_alignment

nemo/collections/asr/inference/model_wrappers/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

nemo/collections/asr/inference/model_wrappers/asr_inference_wrapper.py

@@ -0,0 +1,285 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import copy
+from functools import cached_property
+from typing import Callable
+
+import torch
+from omegaconf import DictConfig, open_dict
+
+from nemo.collections.asr.inference.utils.constants import SENTENCEPIECE_UNDERSCORE
+from nemo.collections.asr.inference.utils.device_utils import setup_device
+from nemo.collections.asr.inference.utils.pipeline_utils import make_preprocessor_deterministic
+from nemo.collections.asr.models import ASRModel, EncDecHybridRNNTCTCModel
+from nemo.collections.asr.parts.submodules.ctc_decoding import CTCDecodingConfig
+from nemo.collections.asr.parts.submodules.rnnt_decoding import RNNTDecodingConfig
+from nemo.collections.asr.parts.utils.asr_confidence_utils import get_confidence_aggregation_bank
+
+SUPPORTED_CONFIDENCE_AGGREGATORS = get_confidence_aggregation_bank()
+
+
+class ASRInferenceWrapper:
+    """
+    Base class for ASR inference wrappers.
+    It provides a common interface for ASR inference wrappers.
+    Derived classes MUST implement the following methods:
+        - __post_init__: Additional post initialization steps that must be implemented in the derived classes.
+        - get_blank_id: Returns the blank id for the model.
+        - get_vocabulary: Returns the vocabulary for the model.
+        - get_subsampling_factor: Returns the subsampling factor for the model.
+    """
+
+    def __init__(
+        self,
+        model_name: str,
+        decoding_cfg: CTCDecodingConfig | RNNTDecodingConfig,
+        device: str = 'cuda',
+        device_id: int = 0,
+        compute_dtype: str = 'bfloat16',
+        use_amp: bool = True,
+    ):
+        """
+        Initialize the ASR inference wrapper.
+        Args:
+            model_name: (str) path to the model checkpoint or a model name from the NGC cloud.
+            decoding_cfg: (CTCDecodingConfig | RNNTDecodingConfig) decoding configuration.
+            device: (str) device to run the model on.
+            device_id: (int) device ID to run the model on.
+            compute_dtype: (str) compute dtype to run the model on.
+            use_amp: (bool) Use Automatic Mixed Precision
+        """
+
+        self.decoding_cfg = decoding_cfg
+        self.device_str, self.device_id, self.compute_dtype = setup_device(device.strip(), device_id, compute_dtype)
+        self.device = torch.device(self.device_str)
+        self.use_amp = use_amp
+        self.asr_model = self.load_model(model_name, self.device)
+        self.asr_model_cfg = self.asr_model._cfg
+        self.set_dither_to_zero()
+        self.tokenizer = self.asr_model.tokenizer
+
+        # post initialization steps that must be implemented in the derived classes
+        self.__post_init__()
+
+    @staticmethod
+    def load_model(model_name: str, map_location: torch.device) -> ASRModel:
+        """
+        Load the ASR model.
+        Args:
+            model_name: (str) path to the model checkpoint or a model name from the NGC cloud.
+            map_location: (torch.device) device to load the model on.
+        Returns:
+            (ASRModel) loaded ASR model.
+        """
+        try:
+            if model_name.endswith('.nemo'):
+                asr_model = ASRModel.restore_from(model_name, map_location=map_location)
+            else:
+                asr_model = ASRModel.from_pretrained(model_name, map_location=map_location)
+            asr_model.eval()
+            return asr_model
+        except Exception as e:
+            raise RuntimeError(f"Failed to load model {model_name}: {str(e)}") from e
+
+    @property
+    def word_separator(self) -> str:
+        """
+        Returns word separator.
+        Returns:
+            (str) word separator.
+        """
+        return self.decoding_cfg.word_seperator
+
+    @property
+    def confidence_aggregator(self) -> Callable:
+        """
+        Returns confidence aggregator function.
+        Returns:
+            (Callable) confidence aggregator function.
+        """
+        return SUPPORTED_CONFIDENCE_AGGREGATORS[self.decoding_cfg.confidence_cfg.aggregation]
+
+    def copy_asr_config(self) -> DictConfig:
+        """
+        Copies the ASR model config.
+        Returns:
+            (DictConfig) copy of the ASR model configuration.
+        """
+        return copy.deepcopy(self.asr_model_cfg)
+
+    def create_preprocessor(self) -> tuple[Callable, DictConfig]:
+        """
+        Creates a deterministic preprocessor from the ASR model configuration.
+        Disables normalization, dither and padding.
+        Returns:
+            (Callable, DictConfig) deterministic preprocessor and its configuration.
+        """
+        new_asr_config = self.copy_asr_config()
+        new_asr_config = make_preprocessor_deterministic(new_asr_config)
+        preprocessor_config = copy.deepcopy(new_asr_config.preprocessor)
+        preprocessor = ASRModel.from_config_dict(preprocessor_config)
+        preprocessor.to(self.device)
+        return preprocessor, preprocessor_config
+
+    def supports_capitalization(self) -> bool:
+        """
+        Checks if the ASR model supports capitalization.
+        Returns:
+            (bool) True if the ASR model supports capitalization, False otherwise.
+        """
+        return self.tokenizer.supports_capitalization
+
+    def supports_punctuation(self) -> bool:
+        """
+        Checks if the ASR model supports punctuation.
+        Returns:
+            (bool) True if the ASR model supports punctuation, False otherwise.
+        """
+        return self.supported_punctuation() != set()
+
+    def supported_punctuation(self) -> set:
+        """
+        Returns supported punctuation symbol set without single quote.
+        Returns:
+            (set) Set of supported punctuation symbols.
+        """
+        return self.tokenizer.supported_punctuation - set("'")
+
+    @cached_property
+    def punctuation_ids(self) -> set:
+        """
+        Returns ids of supported punctuation symbols.
+        Returns:
+            (set) Set of punctuation ids.
+        """
+        punctuation_ids = set()
+        if self.supports_punctuation():
+            for punctuation in self.supported_punctuation():
+                punctuation_ids.add(self.tokenizer.tokens_to_ids(punctuation)[0])
+        return punctuation_ids
+
+    @cached_property
+    def underscore_id(self) -> int:
+        """
+        Returns id of the underscore token.
+        Returns:
+            (int) underscore id for the model.
+        """
+        if getattr(self.asr_model.tokenizer, "spm_separator_id", None) is not None:
+            return self.asr_model.tokenizer.spm_separator_id
+        else:
+            return self.asr_model.tokenizer.tokens_to_ids(SENTENCEPIECE_UNDERSCORE)
+
+    @cached_property
+    def language_token_ids(self) -> set:
+        """
+        This property is used for some Riva models that have language tokens included in the vocabulary.
+        Returns:
+            (set) Set of language token ids.
+        """
+        vocab = self.get_vocabulary()
+        language_token_ids = set()
+        for token in vocab:
+            if token.startswith("<") and token.endswith(">") and token != "<unk>":
+                language_token_ids.add(self.asr_model.tokenizer.tokens_to_ids(token)[0])
+        return language_token_ids
+
+    def reset_decoding_strategy(self, decoder_type: str) -> None:
+        """
+        Reset the decoding strategy for the model.
+        Args:
+            decoder_type: (str) decoding type either 'ctc', 'rnnt'.
+        """
+        if isinstance(self.asr_model, EncDecHybridRNNTCTCModel):
+            self.asr_model.change_decoding_strategy(decoding_cfg=None, decoder_type=decoder_type)
+        else:
+            self.asr_model.change_decoding_strategy(None)
+
+    def set_decoding_strategy(self, decoder_type: str) -> None:
+        """
+        Set the decoding strategy for the model.
+        Args:
+            decoder_type: (str) decoding type either 'ctc', 'rnnt'.
+        """
+        if isinstance(self.asr_model, EncDecHybridRNNTCTCModel):
+            self.asr_model.change_decoding_strategy(decoding_cfg=self.decoding_cfg, decoder_type=decoder_type)
+        else:
+            self.asr_model.change_decoding_strategy(self.decoding_cfg)
+
+    def set_dither_to_zero(self) -> None:
+        """
+        To remove randomness from preprocessor set the dither value to zero.
+        """
+        self.asr_model.preprocessor.featurizer.dither = 0.0
+        with open_dict(self.asr_model_cfg):
+            self.asr_model_cfg.preprocessor.dither = 0.0
+
+    def get_window_stride(self) -> float:
+        """
+        Get the window stride for the model.
+        Returns:
+            (float) window stride for the model.
+        """
+        return self.asr_model_cfg.preprocessor.window_stride
+
+    def get_model_stride(self, in_secs: bool = False, in_milliseconds: bool = False) -> float:
+        """
+        Get the model stride in seconds for the model.
+        Args:
+            in_secs: (bool) Whether to return the model stride in seconds.
+            in_milliseconds: (bool) Whether to return the model stride in milliseconds.
+        Returns:
+            (float) model stride in seconds or milliseconds.
+        """
+        if in_secs and in_milliseconds:
+            raise ValueError("Cannot return both seconds and milliseconds at the same time.")
+        if in_secs:
+            return self.get_window_stride() * self.get_subsampling_factor()
+        if in_milliseconds:
+            return self.get_window_stride() * self.get_subsampling_factor() * 1000
+
+        return self.get_window_stride() * self.get_subsampling_factor()
+
+    # Methods that must be implemented in the derived classes.
+    def __post_init__(self):
+        """
+        Additional post initialization steps that must be implemented in the derived classes.
+        """
+        raise NotImplementedError()
+
+    def get_blank_id(self) -> int:
+        """
+        Returns id of the blank token.
+        Returns:
+            (int) blank id for the model.
+        """
+        raise NotImplementedError()
+
+    def get_vocabulary(self) -> list[str]:
+        """
+        Returns the list of vocabulary tokens.
+        Returns:
+            (list[str]) list of vocabulary tokens.
+        """
+        raise NotImplementedError()
+
+    def get_subsampling_factor(self) -> int:
+        """
+        Returns the subsampling factor for the model.
+        Returns:
+            (int) subsampling factor for the model.
+        """
+        raise NotImplementedError()

nemo/collections/asr/inference/model_wrappers/cache_aware_asr_inference_wrapper.py

@@ -0,0 +1,136 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from typing import Any
+
+from torch import Tensor
+
+from nemo.collections.asr.inference.model_wrappers.asr_inference_wrapper import ASRInferenceWrapper
+
+
+class CacheAwareASRInferenceWrapper(ASRInferenceWrapper):
+    """
+    Base class for Cache-Aware inference wrappers.
+    It provides a common interface for Cache-Aware models.
+    Derived classes MUST implement the following methods:
+        - stream_step: Executes a single streaming step.
+    """
+
+    def get_input_features(self) -> int:
+        """
+        Returns the number of channels in the input features.
+        Returns:
+            (int) number of channels in the input features.
+        """
+        return self.asr_model.encoder._feat_in
+
+    def get_sampling_frames(self) -> list[int] | int | None:
+        """
+        It is used for checking to make sure the audio chunk has enough frames to produce at least one output after downsampling.
+        Returns:
+            (list[int] | int | None) sampling frames for the encoder.
+        """
+        self.sampling_frames = None
+        if hasattr(self.asr_model.encoder, "pre_encode") and hasattr(
+            self.asr_model.encoder.pre_encode, "get_sampling_frames"
+        ):
+            self.sampling_frames = self.asr_model.encoder.pre_encode.get_sampling_frames()
+        return self.sampling_frames
+
+    def get_initial_cache_state(self, batch_size: int) -> tuple[Tensor, Tensor, Tensor]:
+        """
+        Returns the initial cache state for the encoder.
+        Returns:
+            (tuple[Tensor, Tensor, Tensor]) the initial cache state of the encoder.
+        """
+        return self.asr_model.encoder.get_initial_cache_state(batch_size=batch_size, dtype=self.cast_dtype)
+
+    def get_drop_extra_pre_encoded(self) -> int:
+        """
+        Returns the number of extra pre-encoded frames to drop.
+        Returns:
+            (int) drop_extra_pre_encoded.
+        """
+        return self.asr_model.encoder.streaming_cfg.drop_extra_pre_encoded
+
+    def get_chunk_size(self) -> list[int] | int:
+        """
+        Returns the chunk size for the encoder.
+        Returns:
+            (list[int] | int) the chunk size.
+        """
+        return self.asr_model.encoder.streaming_cfg.chunk_size
+
+    def get_shift_size(self) -> list[int] | int:
+        """
+        Returns the shift size for the encoder.
+        Returns:
+            (list[int] | int) the shift size.
+        """
+        return self.asr_model.encoder.streaming_cfg.shift_size
+
+    def get_pre_encode_cache_size(self) -> list[int] | int:
+        """
+        Returns the pre-encode cache size for the encoder.
+        Returns:
+            (list[int] | int) the pre_encode cache size.
+        """
+        return self.asr_model.encoder.streaming_cfg.pre_encode_cache_size
+
+    def get_subsampling_factor(self) -> int:
+        """
+        Returns the subsampling factor for the ASR encoder.
+        Returns:
+            (int) subsampling factor for the ASR encoder model.
+        """
+        return self.asr_model.encoder.subsampling_factor
+
+    def get_att_context_size(self) -> list:
+        """
+        Returns the attention context size for the encoder.
+        Returns:
+            (list) copy of the attention context size.
+        """
+        return self.asr_model.encoder.att_context_size.copy()
+
+    def set_default_att_context_size(self, att_context_size: list) -> None:
+        """
+        Set the default attention context size for the encoder.
+        The list of the supported look-ahead: [[70, 13], [70, 6], [70, 1], [70, 0]]
+        Args:
+            att_context_size: (list) the attention context size.
+        """
+        if hasattr(self.asr_model.encoder, "set_default_att_context_size"):
+            self.asr_model.encoder.set_default_att_context_size(att_context_size=att_context_size)
+        else:
+            raise ValueError("Model does not support multiple lookaheads.")
+
+    def setup_streaming_params(self, chunk_size: int, shift_size: int) -> None:
+        """
+        Setup the streaming parameters (chunk_size, shift_size) for the encoder.
+        Args:
+            chunk_size: (int) the chunk size.
+            shift_size: (int) the shift size.
+        """
+        self.asr_model.encoder.setup_streaming_params(chunk_size=chunk_size, shift_size=shift_size)
+
+    def stream_step(self, *args, **kwargs) -> Any:
+        """
+        Executes a single streaming step.
+        Each derived class must implement this method, with arguments and return types specific to that class.
+        """
+        raise NotImplementedError(
+            "`stream_step` method is not implemented. It is required for cache-aware transcribers."
+        )

nemo/collections/asr/inference/model_wrappers/cache_aware_ctc_inference_wrapper.py

@@ -0,0 +1,208 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import torch
+from torch import Tensor
+
+from nemo.collections.asr.inference.model_wrappers.cache_aware_asr_inference_wrapper import (
+    CacheAwareASRInferenceWrapper,
+)
+from nemo.collections.asr.inference.utils.context_manager import CacheAwareContext
+from nemo.collections.asr.models import EncDecCTCModel, EncDecHybridRNNTCTCModel
+from nemo.collections.asr.parts.mixins.streaming import StreamingEncoder
+
+
+class CacheAwareCTCInferenceWrapper(CacheAwareASRInferenceWrapper):
+    """
+    Provides a unified interface to work with Cache-Aware CTC models.
+    """
+
+    def __post_init__(self) -> None:
+        """
+        Additional post initialization step
+        Checks if the model is a ctc model and sets the decoding strategy to ctc.
+        """
+
+        if not isinstance(self.asr_model, (EncDecCTCModel, EncDecHybridRNNTCTCModel)):
+            raise ValueError(
+                "Provided model is not a CTC type. You are trying to use a CTC Inference with a non-CTC model."
+            )
+
+        if not isinstance(self.asr_model.encoder, StreamingEncoder):
+            raise NotImplementedError("Encoder of this model does not support streaming!")
+
+        decoder_type = 'ctc'
+        if isinstance(self.asr_model, EncDecHybridRNNTCTCModel):
+            self.asr_model.cur_decoder = decoder_type
+
+        # reset the decoding strategy
+        self.reset_decoding_strategy(decoder_type)
+        self.set_decoding_strategy(decoder_type)
+
+        # setup streaming parameters
+        if self.asr_model.encoder.streaming_cfg is None:
+            self.asr_model.encoder.setup_streaming_params()
+
+        self.drop_extra_pre_encoded = self.get_drop_extra_pre_encoded()
+
+        self.cast_dtype = torch.float32 if self.use_amp else self.compute_dtype
+        self.asr_model.to(self.cast_dtype)
+
+    def get_blank_id(self) -> int:
+        """
+        Returns id of the blank token.
+        Returns:
+            (int) blank id for the model.
+        """
+        if isinstance(self.asr_model, EncDecCTCModel):
+            blank_id = len(self.asr_model.decoder.vocabulary)
+        else:
+            blank_id = len(self.asr_model.ctc_decoder.vocabulary)
+        return blank_id
+
+    def get_vocabulary(self) -> list[str]:
+        """
+        Returns the list of vocabulary tokens.
+        Returns:
+            (list[str]) list of vocabulary tokens.
+        """
+        if isinstance(self.asr_model, EncDecCTCModel):
+            return self.asr_model.decoder.vocabulary
+        else:
+            return self.asr_model.ctc_decoder.vocabulary
+
+    def execute_step(
+        self,
+        processed_signal: Tensor,
+        processed_signal_length: Tensor,
+        context: CacheAwareContext,
+        drop_extra_pre_encoded: int | None,
+        keep_all_outputs: bool,
+        drop_left_context: int | None = None,
+        valid_out_len: int | None = None,
+        return_tail_result: bool = False,
+    ) -> tuple[Tensor, Tensor | None, CacheAwareContext]:
+        """
+        Executes a single streaming step.
+        Args:
+            processed_signal: (Tensor) input signal tensor.
+            processed_signal_length: (Tensor) input signal length tensor.
+            context: (CacheAwareContext) context object.
+            drop_extra_pre_encoded: (int | None) number of extra pre-encoded frames to drop.
+            keep_all_outputs: (bool) whether to keep all outputs or not.
+            drop_left_context: (int | None) number of left context frames to drop.
+            valid_out_len: (int | None) number of valid output frames.
+            return_tail_result: (bool) whether to return tail result or not.
+        Returns:
+            (tuple[Tensor, Tensor | None, CacheAwareContext]) log probabilities, tail log probabilities and new context.
+        """
+
+        (
+            encoded,
+            encoded_len,
+            cache_last_channel,
+            cache_last_time,
+            cache_last_channel_len,
+        ) = self.asr_model.encoder.cache_aware_stream_step(
+            processed_signal=processed_signal,
+            processed_signal_length=processed_signal_length,
+            cache_last_channel=context.cache_last_channel,
+            cache_last_time=context.cache_last_time,
+            cache_last_channel_len=context.cache_last_channel_len,
+            keep_all_outputs=keep_all_outputs,
+            drop_extra_pre_encoded=drop_extra_pre_encoded,
+        )
+
+        if drop_left_context:
+            # drop left context
+            encoded = encoded[:, :, drop_left_context:]
+
+        if isinstance(self.asr_model, EncDecHybridRNNTCTCModel):
+            all_log_probs = self.asr_model.ctc_decoder(encoder_output=encoded)
+        else:
+            all_log_probs = self.asr_model.decoder(encoder_output=encoded)
+
+        tail_log_probs = None
+        if valid_out_len and not keep_all_outputs:
+            # drop right context if any
+            log_probs = all_log_probs[:, :valid_out_len, :]
+            if return_tail_result:
+                tail_log_probs = all_log_probs[:, valid_out_len:, :]
+        else:
+            log_probs = all_log_probs
+
+        # create a new context
+        new_context = CacheAwareContext(
+            cache_last_channel=cache_last_channel,
+            cache_last_time=cache_last_time,
+            cache_last_channel_len=cache_last_channel_len,
+        )
+        return log_probs, tail_log_probs, new_context
+
+    def stream_step(
+        self,
+        processed_signal: Tensor,
+        processed_signal_length: Tensor,
+        context: CacheAwareContext = None,
+        drop_extra_pre_encoded: int | None = None,
+        keep_all_outputs: bool = False,
+        drop_left_context: int | None = None,
+        valid_out_len: int | None = None,
+        return_tail_result: bool = False,
+    ) -> tuple[Tensor, Tensor | None, CacheAwareContext]:
+        """
+        Executes a single streaming step.
+        Args:
+            processed_signal: (Tensor) input signal tensor.
+            processed_signal_length: (Tensor) input signal length tensor.
+            context: (CacheAwareContext) context object.
+            drop_extra_pre_encoded: (int | None) number of extra pre-encoded frames to drop.
+            keep_all_outputs: (bool) whether to keep all outputs or not.
+            drop_left_context: (int | None) number of left context frames to drop.
+            valid_out_len: (int | None) number of valid output frames.
+            return_tail_result: (bool) whether to return tail result or not.
+        Returns:
+            (tuple[Tensor, Tensor | None, CacheAwareContext]) log probabilities, tail log probabilities and new context.
+        """
+
+        if processed_signal.device != self.device:
+            processed_signal = processed_signal.to(self.device)
+
+        if processed_signal_length.device != self.device:
+            processed_signal_length = processed_signal_length.to(self.device)
+
+        processed_signal = processed_signal.to(self.cast_dtype)
+
+        if context is None:
+            # create a dummy context
+            context = CacheAwareContext()
+
+        with (
+            torch.amp.autocast(device_type=self.device_str, dtype=self.compute_dtype, enabled=self.use_amp),
+            torch.inference_mode(),
+            torch.no_grad(),
+        ):
+
+            log_probs, tail_log_probs, new_context = self.execute_step(
+                processed_signal,
+                processed_signal_length,
+                context,
+                drop_extra_pre_encoded,
+                keep_all_outputs,
+                drop_left_context,
+                valid_out_len,
+                return_tail_result,
+            )
+        return log_probs, tail_log_probs, new_context

nemo/collections/asr/inference/model_wrappers/cache_aware_rnnt_inference_wrapper.py

@@ -0,0 +1,314 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+from torch import Tensor
+
+from nemo.collections.asr.inference.model_wrappers.cache_aware_asr_inference_wrapper import (
+    CacheAwareASRInferenceWrapper,
+)
+from nemo.collections.asr.inference.streaming.state.cache_aware_rnnt_state import CacheAwareRNNTBeamStreamingState
+from nemo.collections.asr.inference.utils.context_manager import CacheAwareContext
+from nemo.collections.asr.inference.utils.per_stream_biasing import multi_biasing_ids_tensor_from_states
+from nemo.collections.asr.models import EncDecHybridRNNTCTCModel, EncDecRNNTModel
+from nemo.collections.asr.parts.mixins.streaming import StreamingEncoder
+from nemo.collections.asr.parts.submodules.rnnt_malsd_batched_computer import ModifiedALSDBatchedRNNTComputer
+from nemo.collections.asr.parts.utils.batched_beam_decoding_utils import export_batched_beam_hyps_to_cpu_lists
+from nemo.collections.asr.parts.utils.rnnt_utils import Hypothesis
+
+
+class CacheAwareRNNTInferenceWrapper(CacheAwareASRInferenceWrapper):
+    """
+    Provides a unified interface to work with Cache-Aware RNNT models.
+    """
+
+    def __post_init__(self) -> None:
+        """
+        Additional post initialization step
+        Checks if the model is a rnnt model and sets the decoding strategy to rnnt.
+        """
+        if not isinstance(self.asr_model, (EncDecRNNTModel, EncDecHybridRNNTCTCModel)):
+            raise ValueError(
+                "Provided model is not a RNNT type. You are trying to use a RNNT Inference with a non-RNNT model."
+            )
+
+        if not isinstance(self.asr_model.encoder, StreamingEncoder):
+            raise NotImplementedError("Encoder of this model does not support streaming!")
+
+        decoder_type = 'rnnt'
+        if isinstance(self.asr_model, EncDecHybridRNNTCTCModel):
+            self.asr_model.cur_decoder = decoder_type
+
+        # reset the decoding strategy
+        self.reset_decoding_strategy(decoder_type)
+        self.set_decoding_strategy(decoder_type)
+
+        # setup streaming parameters
+        if self.asr_model.encoder.streaming_cfg is None:
+            self.asr_model.encoder.setup_streaming_params()
+
+        self.drop_extra_pre_encoded = self.get_drop_extra_pre_encoded()
+
+        self.cast_dtype = torch.float32 if self.use_amp else self.compute_dtype
+        self.asr_model.to(self.cast_dtype)
+
+    def get_blank_id(self) -> int:
+        """
+        Returns id of the blank token.
+        Returns:
+            (int) blank id for the model.
+        """
+        blank_id = len(self.asr_model.joint.vocabulary)
+        return blank_id
+
+    def get_vocabulary(self) -> list[str]:
+        """
+        Returns the list of vocabulary tokens.
+        Returns:
+            (list[str]) list of vocabulary tokens.
+        """
+        return self.asr_model.joint.vocabulary
+
+    def encoder_step(
+        self,
+        processed_signal: Tensor,
+        processed_signal_length: Tensor,
+        context: CacheAwareContext,
+        drop_extra_pre_encoded: int | None,
+        keep_all_outputs: bool,
+        drop_left_context: int | None = None,
+        valid_out_len: int | None = None,
+    ) -> tuple[Tensor, Tensor, CacheAwareContext]:
+        """
+        Run the cache-aware encoder for one streaming chunk. Decoder is NOT invoked.
+        Args:
+            processed_signal: (Tensor) input signal tensor.
+            processed_signal_length: (Tensor) input signal length tensor.
+            context: (CacheAwareContext) context object.
+            drop_extra_pre_encoded: (int | None) number of extra pre-encoded frames to drop.
+            keep_all_outputs: (bool) whether to keep all outputs or not.
+            drop_left_context: (int | None) number of left context frames to drop.
+            valid_out_len: (int | None) number of valid output frames.
+        Returns:
+            (tuple[Tensor, Tensor, CacheAwareContext]) encoder output, encoder output lengths, and new context.
+        """
+        (
+            encoded,
+            encoded_len,
+            cache_last_channel,
+            cache_last_time,
+            cache_last_channel_len,
+        ) = self.asr_model.encoder.cache_aware_stream_step(
+            processed_signal=processed_signal,
+            processed_signal_length=processed_signal_length,
+            cache_last_channel=context.cache_last_channel,
+            cache_last_time=context.cache_last_time,
+            cache_last_channel_len=context.cache_last_channel_len,
+            keep_all_outputs=keep_all_outputs,
+            drop_extra_pre_encoded=drop_extra_pre_encoded,
+        )
+        new_context = CacheAwareContext(
+            cache_last_channel=cache_last_channel,
+            cache_last_time=cache_last_time,
+            cache_last_channel_len=cache_last_channel_len,
+        )
+
+        if drop_left_context:
+            # drop left context
+            encoded = encoded[:, :, drop_left_context:]
+            encoded_len = encoded_len - drop_left_context
+
+        if valid_out_len and not keep_all_outputs:
+            # drop right context if any
+            encoded = encoded[:, :, :valid_out_len]
+            encoded_len = torch.ones_like(encoded_len) * valid_out_len
+
+        return encoded, encoded_len, new_context
+
+    def execute_step(
+        self,
+        processed_signal: Tensor,
+        processed_signal_length: Tensor,
+        context: CacheAwareContext,
+        previous_hypotheses: list[Hypothesis] | None,
+        drop_extra_pre_encoded: int | None,
+        keep_all_outputs: bool,
+        drop_left_context: int | None = None,
+        valid_out_len: int | None = None,
+        prompt_vectors: Tensor | None = None,
+    ) -> tuple[list[Hypothesis], CacheAwareContext]:
+        """
+        Executes a single streaming step.
+        Args:
+            processed_signal: (Tensor) input signal tensor.
+            processed_signal_length: (Tensor) input signal length tensor.
+            context: (CacheAwareContext) context object.
+            previous_hypotheses: (list[Hypothesis] | None) list of previous hypotheses for RNNT decoding.
+            drop_extra_pre_encoded: (int | None) number of extra pre-encoded frames to drop.
+            keep_all_outputs: (bool) whether to keep all outputs or not.
+            drop_left_context: (int | None) number of left context frames to drop.
+            valid_out_len: (int | None) number of valid output frames.
+            prompt_vectors: (Tensor | None) Optional prompt vectors of shape [B, num_prompts].
+        Returns:
+            (tuple[list[Hypothesis], CacheAwareContext]) best hypothesis and new context.
+        """
+        encoded, encoded_len, new_context = self.encoder_step(
+            processed_signal=processed_signal,
+            processed_signal_length=processed_signal_length,
+            context=context,
+            drop_extra_pre_encoded=drop_extra_pre_encoded,
+            keep_all_outputs=keep_all_outputs,
+            drop_left_context=drop_left_context,
+            valid_out_len=valid_out_len,
+        )
+
+        best_hyp = self.asr_model.decoding.rnnt_decoder_predictions_tensor(
+            encoded, encoded_len, return_hypotheses=True, partial_hypotheses=previous_hypotheses
+        )
+        return best_hyp, new_context
+
+    def malsd_stream_step(
+        self,
+        malsd_computer: ModifiedALSDBatchedRNNTComputer,
+        states: list[CacheAwareRNNTBeamStreamingState],
+        processed_signal: Tensor,
+        processed_signal_length: Tensor,
+        context: CacheAwareContext,
+        drop_extra_pre_encoded: int | None,
+        keep_all_outputs: bool,
+        drop_left_context: int | None = None,
+        valid_out_len: int | None = None,
+    ) -> tuple[list[Hypothesis], CacheAwareContext]:
+        """Cache-aware MALSD encode/decode step for one chunk."""
+        if processed_signal.device != self.device:
+            processed_signal = processed_signal.to(self.device)
+
+        if processed_signal_length.device != self.device:
+            processed_signal_length = processed_signal_length.to(self.device)
+
+        carries = [state.hyp_decoding_state for state in states]
+        if all(c is None for c in carries):
+            batched_state = None
+        else:
+            batched_state = malsd_computer.merge_to_batched_state(carries)
+
+        multi_biasing_ids = multi_biasing_ids_tensor_from_states(
+            states,
+            self.device,
+            per_stream_biasing_enabled=malsd_computer.per_stream_biasing_enabled,
+        )
+
+        with (
+            torch.amp.autocast(device_type=self.device_str, dtype=self.compute_dtype, enabled=self.use_amp),
+            torch.inference_mode(),
+            torch.no_grad(),
+        ):
+            processed_signal = processed_signal.to(self.cast_dtype)
+            encoded, encoded_len, new_context = self.encoder_step(
+                processed_signal=processed_signal,
+                processed_signal_length=processed_signal_length,
+                context=context,
+                drop_extra_pre_encoded=drop_extra_pre_encoded,
+                keep_all_outputs=keep_all_outputs,
+                drop_left_context=drop_left_context,
+                valid_out_len=valid_out_len,
+            )
+            encs_dim_last = encoded.transpose(1, 2).contiguous()
+
+            best_batched_hyps, batched_state = malsd_computer(
+                encs_dim_last, encoded_len, batched_state, multi_biasing_ids=multi_biasing_ids
+            )
+
+            chunk_tokens, chunk_timestamps, chunk_confidences, root_ptrs = export_batched_beam_hyps_to_cpu_lists(
+                best_batched_hyps
+            )
+            beam_indices = best_batched_hyps.scores.argmax(dim=-1).detach().cpu().tolist()
+            scores_cpu = best_batched_hyps.scores.detach().cpu()
+
+            carry_items = malsd_computer.split_batched_state(batched_state)
+            for b, (state, ct, cts, rp, best_hyp_idx, carry) in enumerate(
+                zip(states, chunk_tokens, chunk_timestamps, root_ptrs, beam_indices, carry_items)
+            ):
+                state.append_chunk_beam_(
+                    ct,
+                    cts,
+                    rp,
+                    best_batched_hyps.beam_size,
+                    best_hyp_idx,
+                    chunk_confidences=(chunk_confidences[b] if chunk_confidences is not None else None),
+                )
+                state.hyp_decoding_state = carry
+
+        hyps = [state.get_hypothesis(float(scores_cpu[b, beam_indices[b]].item())) for b, state in enumerate(states)]
+        return hyps, new_context
+
+    def stream_step(
+        self,
+        processed_signal: Tensor,
+        processed_signal_length: Tensor,
+        context: CacheAwareContext = None,
+        previous_hypotheses: list[Hypothesis] | None = None,
+        drop_extra_pre_encoded: int | None = None,
+        keep_all_outputs: bool = False,
+        drop_left_context: int | None = None,
+        valid_out_len: int | None = None,
+        prompt_vectors: Tensor | None = None,
+    ) -> tuple[list[Hypothesis], CacheAwareContext]:
+        """
+        Executes a single streaming step.
+        Args:
+            processed_signal: (Tensor) input signal tensor.
+            processed_signal_length: (Tensor) input signal length tensor.
+            context: (CacheAwareContext) context object.
+            previous_hypotheses: (list[Hypothesis] | None) list of previous hypotheses for RNNT decoding.
+            drop_extra_pre_encoded: (int | None) number of extra pre-encoded frames to drop.
+            keep_all_outputs: (bool) whether to keep all outputs or not.
+            drop_left_context: (int | None) number of left context frames to drop.
+            valid_out_len: (int | None) number of valid output frames.
+            prompt_vectors: (Tensor | None) Optional prompt vectors of shape [B, num_prompts].
+        Returns:
+            (tuple[list[Hypothesis], CacheAwareContext]) best hypothesis and new context.
+        """
+
+        if processed_signal.device != self.device:
+            processed_signal = processed_signal.to(self.device)
+
+        if processed_signal_length.device != self.device:
+            processed_signal_length = processed_signal_length.to(self.device)
+
+        processed_signal = processed_signal.to(self.cast_dtype)
+
+        if context is None:
+            # create a dummy context
+            context = CacheAwareContext()
+
+        with (
+            torch.amp.autocast(device_type=self.device_str, dtype=self.compute_dtype, enabled=self.use_amp),
+            torch.inference_mode(),
+            torch.no_grad(),
+        ):
+
+            best_hyp, new_context = self.execute_step(
+                processed_signal,
+                processed_signal_length,
+                context,
+                previous_hypotheses,
+                drop_extra_pre_encoded,
+                keep_all_outputs,
+                drop_left_context,
+                valid_out_len,
+                prompt_vectors,
+            )
+
+        return best_hyp, new_context

nemo/collections/asr/inference/model_wrappers/ctc_inference_wrapper.py

@@ -0,0 +1,109 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+from torch import Tensor
+
+from nemo.collections.asr.inference.model_wrappers.asr_inference_wrapper import ASRInferenceWrapper
+from nemo.collections.asr.models import EncDecCTCModel, EncDecHybridRNNTCTCModel
+
+
+class CTCInferenceWrapper(ASRInferenceWrapper):
+    """
+    Provides a unified interface to work with CTC/Hybrid-CTC models.
+    """
+
+    def __post_init__(self) -> None:
+        """
+        Additional post initialization step
+        Checks if the model is a ctc model and sets the decoding strategy to ctc.
+        """
+        if not isinstance(self.asr_model, (EncDecCTCModel, EncDecHybridRNNTCTCModel)):
+            raise ValueError(
+                "Provided model is not a CTC type. You are trying to use a CTC transcriber with a non-CTC model."
+            )
+
+        decoder_type = 'ctc'
+        if isinstance(self.asr_model, EncDecHybridRNNTCTCModel):
+            self.asr_model.cur_decoder = decoder_type
+
+        # reset the decoding strategy
+        self.reset_decoding_strategy(decoder_type)
+        self.set_decoding_strategy(decoder_type)
+
+        self.cast_dtype = torch.float32 if self.use_amp else self.compute_dtype
+        self.asr_model.to(self.cast_dtype)
+
+    def get_blank_id(self) -> int:
+        """
+        Returns id of the blank token.
+        Returns:
+            (int) blank id for the model.
+        """
+        if isinstance(self.asr_model, EncDecCTCModel):
+            blank_id = len(self.asr_model.decoder.vocabulary)
+        else:
+            blank_id = len(self.asr_model.ctc_decoder.vocabulary)
+        return blank_id
+
+    def get_vocabulary(self) -> list[str]:
+        """
+        Returns the list of vocabulary tokens.
+        Returns:
+            (list[str]) list of vocabulary tokens.
+        """
+        if isinstance(self.asr_model, EncDecCTCModel):
+            return self.asr_model.decoder.vocabulary
+        else:
+            return self.asr_model.ctc_decoder.vocabulary
+
+    def get_subsampling_factor(self) -> int:
+        """
+        Returns the subsampling factor for the ASR encoder.
+        Returns:
+            (int) subsampling factor for the ASR encoder model.
+        """
+        return self.asr_model.encoder.subsampling_factor
+
+    def get_logprobs(self, processed_signal: Tensor, processed_signal_length: Tensor) -> Tensor:
+        """
+        Get log probabilities from the model. It is used for streaming inference.
+        Args:
+            processed_signal: (Tensor) processed signal. Shape is torch.Size([B, C, T]).
+            processed_signal_length: (Tensor) processed signal length. Shape is torch.Size([B]).
+        Returns:
+            (Tensor) log probabilities. Shape is torch.Size([B, T, V+1]).
+        """
+        if processed_signal.device != self.device:
+            processed_signal = processed_signal.to(self.device)
+
+        if processed_signal_length.device != self.device:
+            processed_signal_length = processed_signal_length.to(self.device)
+
+        with (
+            torch.amp.autocast(device_type=self.device_str, dtype=self.compute_dtype, enabled=self.use_amp),
+            torch.inference_mode(),
+            torch.no_grad(),
+        ):
+
+            forward_outs = self.asr_model(
+                processed_signal=processed_signal.to(self.cast_dtype), processed_signal_length=processed_signal_length
+            )
+
+        if isinstance(self.asr_model, EncDecHybridRNNTCTCModel):
+            encoded, encoded_len = forward_outs
+            log_probs = self.asr_model.ctc_decoder(encoder_output=encoded.clone())
+        else:
+            log_probs, encoded_len, predictions = forward_outs
+        return log_probs

nemo/collections/asr/inference/model_wrappers/rnnt_inference_wrapper.py

@@ -0,0 +1,147 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import torch
+from torch import Tensor
+
+from nemo.collections.asr.inference.model_wrappers.asr_inference_wrapper import ASRInferenceWrapper
+from nemo.collections.asr.models import EncDecHybridRNNTCTCModel, EncDecRNNTModel
+
+
+class RNNTInferenceWrapper(ASRInferenceWrapper):
+    """
+    Provides a unified interface to work with RNNT/TDT/Hybrid models.
+    """
+
+    def __post_init__(self) -> None:
+        """
+        Additional post initialization step
+        Checks if the model is a rnnt model and sets the decoding strategy to rnnt.
+        """
+        if not isinstance(self.asr_model, (EncDecRNNTModel, EncDecHybridRNNTCTCModel)):
+            raise ValueError(
+                "Provided model is not a RNNT type. You are trying to use a RNNT transcriber with a non-RNNT model."
+            )
+
+        decoder_type = 'rnnt'
+        if isinstance(self.asr_model, EncDecHybridRNNTCTCModel):
+            self.asr_model.cur_decoder = decoder_type
+
+        # reset the decoding strategy
+        self.reset_decoding_strategy(decoder_type)
+        self.set_decoding_strategy(decoder_type)
+
+        self.cast_dtype = torch.float32 if self.use_amp else self.compute_dtype
+        self.asr_model.to(self.cast_dtype)
+
+    def get_blank_id(self) -> int:
+        """
+        Returns id of the blank token.
+        Returns:
+            (int) blank id for the model.
+        """
+        blank_id = len(self.asr_model.joint.vocabulary)
+        return blank_id
+
+    def get_vocabulary(self) -> list[str]:
+        """
+        Returns the list of vocabulary tokens.
+        Returns:
+            (list[str]) list of vocabulary tokens.
+        """
+        return self.asr_model.joint.vocabulary
+
+    def get_subsampling_factor(self) -> int:
+        """
+        Returns the subsampling factor for the ASR encoder.
+        Returns:
+            (int) subsampling factor for the ASR encoder model.
+        """
+        return self.asr_model.encoder.subsampling_factor
+
+    def encode(
+        self, processed_signal: Tensor, processed_signal_length: Tensor, prompt_vectors: Tensor | None = None
+    ) -> tuple[Tensor, Tensor]:
+        """
+        Get encoder output from the model. It is used for streaming inference.
+        Args:
+            processed_signal: (Tensor) processed signal. Shape is torch.Size([B, C, T]).
+            processed_signal_length: (Tensor) processed signal length. Shape is torch.Size([B]).
+            prompt_vectors: (Tensor | None) Optional prompt vectors for multilingual models.
+                Shape can be torch.Size([B, num_prompts]) or torch.Size([B, T_enc, num_prompts]) if already expanded.
+        Returns:
+            (tuple[Tensor, Tensor]) encoder output and encoder output length of shape torch.Size([B, T, D]), torch.Size([B]).
+        """
+        if processed_signal.device != self.device:
+            processed_signal = processed_signal.to(self.device)
+
+        if processed_signal_length.device != self.device:
+            processed_signal_length = processed_signal_length.to(self.device)
+
+        with (
+            torch.amp.autocast(device_type=self.device_str, dtype=self.compute_dtype, enabled=self.use_amp),
+            torch.inference_mode(),
+            torch.no_grad(),
+        ):
+
+            # Prepare model arguments
+            model_args = {
+                'processed_signal': processed_signal.to(self.cast_dtype),
+                'processed_signal_length': processed_signal_length,
+            }
+            if prompt_vectors is not None:
+                model_args['prompt'] = prompt_vectors
+
+            forward_outs = self.asr_model(**model_args)
+
+        encoded, encoded_len = forward_outs
+        return encoded, encoded_len
+
+    def decode(self, encoded: Tensor, encoded_len: Tensor, partial_hypotheses: list) -> list:
+        """
+        Decode the encoder output using the RNNT decoder.
+        Args:
+            encoded: (Tensor) encoder output.
+            encoded_len: (Tensor) encoder output length.
+            partial_hypotheses: (list) list of partial hypotheses for stateful decoding.
+        Returns:
+            (list) list of best hypotheses.
+        """
+        best_hyp = self.asr_model.decoding.rnnt_decoder_predictions_tensor(
+            encoded.to(self.cast_dtype), encoded_len, return_hypotheses=True, partial_hypotheses=partial_hypotheses
+        )
+        return best_hyp
+
+    def encode_with_prompts(
+        self, processed_signal: Tensor, processed_signal_length: Tensor, prompt_vectors: Tensor
+    ) -> tuple[Tensor, Tensor]:
+        """
+        Convenience wrapper for prompt-enabled encoding.
+        Expands prompt vectors across the time dimension before calling encode.
+        Args:
+            processed_signal: (Tensor) processed signal. Shape is torch.Size([B, C, T]).
+            processed_signal_length: (Tensor) processed signal length. Shape is torch.Size([B]).
+            prompt_vectors: (Tensor) prompt vectors. Shape is torch.Size([B, num_prompts]).
+        Returns:
+            (tuple[Tensor, Tensor]) encoder output and encoder output length.
+        """
+        encoder_time_steps = processed_signal.shape[2] // self.get_subsampling_factor()
+        # Expand prompts: [B, num_prompts] -> [B, T_enc, num_prompts]
+        prompt_vectors = prompt_vectors.unsqueeze(1).expand(-1, encoder_time_steps, -1)
+        return self.encode(
+            processed_signal=processed_signal,
+            processed_signal_length=processed_signal_length,
+            prompt_vectors=prompt_vectors,
+        )

nemo/collections/asr/inference/model_wrappers/salm_asr_inference_wrapper.py

@@ -0,0 +1,166 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+
+from typing import TYPE_CHECKING
+import torch
+
+from nemo.collections.asr.inference.utils.device_utils import setup_device
+from nemo.collections.common.prompts import PromptFormatter
+
+if TYPE_CHECKING:
+    from nemo.collections.speechlm2.models import SALM
+
+
+class SALMASRInferenceWrapper:
+
+    def __init__(
+        self,
+        model_name: str,
+        device: str = 'cuda',
+        device_id: int = 0,
+        compute_dtype: str = 'bfloat16',
+        use_amp: bool = True,
+    ):
+        """
+        Initialize the SALM ASR inference wrapper.
+        Args:
+            model_name: (str) model name at Hugging Face or NGC cloud.
+            device: (str) device to run the model on.
+            device_id: (int) device ID to run the model on.
+            compute_dtype: (str) compute dtype to run the model on.
+            use_amp: (bool) Use Automatic Mixed Precision
+        """
+
+        self.device_str, self.device_id, self.compute_dtype = setup_device(device.strip(), device_id, compute_dtype)
+        self.use_amp = use_amp
+        self.device = torch.device(self.device_str)
+        self.salm_model = self.load_model(model_name, self.device)
+        self.audio_locator_tag = self.salm_model.audio_locator_tag
+        self.tokenizer = self.salm_model.tokenizer
+        self.set_dither_to_zero()
+
+    @property
+    def eos_token_ids(self) -> list[int]:
+        """Returns the end of sentence token ids."""
+        return [self.salm_model.text_eos_id]
+
+    @property
+    def word_separator(self) -> str:
+        """Returns word separator."""
+        return ' '
+
+    @property
+    def word_separator_ids(self) -> list[int]:
+        """Returns the word separator token ids."""
+        return self.tokenizer.text_to_ids(self.word_separator)
+
+    @staticmethod
+    def load_model(model_name: str, device: torch.device) -> SALM:
+        """
+        Load the SALM model.
+        Args:
+            model_name: (str) model name at Hugging Face or NGC cloud.
+            device: (torch.device) device to load the model on.
+        Returns:
+            (SALM) loaded SALM model.
+        """
+        try:
+            from nemo.collections.speechlm2.models import SALM
+
+            model = SALM.from_pretrained(model_name).eval()
+            model.to(device)
+            return model
+        except Exception as e:
+            raise RuntimeError(f"Failed to load model {model_name}: {str(e)}") from e
+
+    def get_window_stride(self) -> float:
+        """Returns the window stride of the model."""
+        return self.salm_model.cfg.perception.preprocessor.window_stride
+
+    def get_subsampling_factor(self) -> int:
+        """Returns the subsampling factor of the model."""
+        return self.salm_model.cfg.perception.encoder.subsampling_factor
+
+    def get_model_stride(self, in_secs: bool = False, in_milliseconds: bool = False) -> float:
+        """
+        Returns the model stride in seconds or milliseconds.
+        Args:
+            in_secs: (bool) Whether to return the model stride in seconds.
+            in_milliseconds: (bool) Whether to return the model stride in milliseconds.
+        Returns:
+            (float) model stride in seconds or milliseconds.
+        """
+        if in_secs and in_milliseconds:
+            raise ValueError("Cannot return both seconds and milliseconds at the same time.")
+        token_duration = self.salm_model.token_equivalent_duration
+        if in_secs:
+            return token_duration
+        if in_milliseconds:
+            return token_duration * 1000
+        return token_duration
+
+    def set_dither_to_zero(self) -> None:
+        """Sets the dither to zero."""
+        self.salm_model.cfg.perception.preprocessor.dither = 0.0
+        self.salm_model.perception.preprocessor.featurizer.dither = 0.0
+
+    def generate(
+        self,
+        prompts: list[list[dict[str]]] | torch.Tensor,
+        audios: torch.Tensor,
+        audio_lens: torch.Tensor,
+        max_new_tokens: int = 128,
+    ) -> torch.Tensor:
+        """
+        Generate the model output.
+        Args:
+            prompts: (list[list[dict[str]]] | torch.Tensor) List of prompts or token ids.
+            audios: (torch.Tensor) Audio tensor of shape (batch_size, num_samples).
+            audio_lens: (torch.Tensor) Audio length tensor of shape (batch_size).
+            max_new_tokens: (int) Maximum number of new tokens to generate.
+        Returns:
+            (torch.Tensor) Model output.
+        """
+        with (
+            torch.amp.autocast(device_type=self.device_str, dtype=self.compute_dtype, enabled=self.use_amp),
+            torch.inference_mode(),
+            torch.no_grad(),
+        ):
+            answer_ids = self.salm_model.generate(
+                prompts=prompts,
+                audios=audios,
+                audio_lens=audio_lens,
+                max_new_tokens=max_new_tokens,
+            )
+        return answer_ids
+
+    def preprocess_prompts(self, prompts: list[list[dict[str]]]) -> torch.Tensor:
+        """
+        Convert the prompts to token ids.
+        Args:
+            prompts: (list[list[dict[str]]]) List of prompts.
+        Returns:
+            (torch.Tensor) Token ids of size (batch_size, max_prompt_length).
+        """
+        from nemo.collections.speechlm2.data.salm_dataset import left_collate_vectors
+
+        formatter = PromptFormatter.resolve(self.salm_model.cfg.prompt_format)(self.tokenizer)
+        tokens = left_collate_vectors(
+            [formatter.encode_dialog(turns=prompt)["input_ids"] for prompt in prompts],
+            padding_value=self.salm_model.text_pad_id,
+        ).to(self.device)
+        return tokens

nemo/collections/asr/inference/nmt/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

nemo/collections/asr/inference/nmt/llm_translator.py

@@ -0,0 +1,296 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import os
+import string
+
+import torch
+from omegaconf import DictConfig, OmegaConf
+
+from nemo.collections.asr.inference.nmt.prompts import EuroLLMTranslatorPromptTemplate, PromptTemplate
+
+try:
+    from vllm import LLM, SamplingParams
+except ImportError as e:
+    raise ImportError("Failed to import vLLM.") from e
+
+from nemo.utils import logging
+
+EURO_LLM_INSTRUCT_SMALL = "utter-project/EuroLLM-1.7B-Instruct"
+EURO_LLM_INSTRUCT_LARGE = "utter-project/EuroLLM-9B-Instruct"
+SUPPORTED_TRANSLATION_MODELS = [EURO_LLM_INSTRUCT_SMALL, EURO_LLM_INSTRUCT_LARGE]
+
+
+class LLMTranslator:
+    """
+    A vLLM-based LLM translator for ASR transcripts.
+    It takes ASR transcripts and prefixes to start translation from, and returns corresponding continuations of translations.
+    """
+
+    def __init__(
+        self,
+        model_name: str,
+        source_language: str,
+        target_language: str,
+        waitk: int = -1,
+        device: str = "cuda",
+        device_id: int = 0,
+        batch_size: int = -1,
+        llm_params: dict | DictConfig | None = None,
+        sampling_params: dict | DictConfig | None = None,
+    ):
+        """
+        A model for translating ASR transcripts with LLM.
+        Args:
+            model_name: (str) path to the model name on HuggingFace.
+            source_language: (str) source language
+            target_language: (str) target language
+            waitk: (int) sets the maximum number of words the translation is allowed to lag behind the ASR transcript.
+                         If the translation falls more than waitk words behind, it automatically extends the prefix
+                         using the current translation. -1 disables this rule and relies on the longest common prefix (LCP)
+                         between current and previous translations. Larger values of waitk lead to more coherent translations,
+                         but the cost of generating the translation increases, because the model needs to generate more tokens.
+            device: (str) device to run the model on
+            device_id: (int) device ID to run the model on
+            batch_size: (int) batch size for the LLM model, in case of -1, the batch size is set to the number of ASR transcripts
+            llm_params: (dict | DictConfig | None) parameters for the LLM model
+            sampling_params: (dict | DictConfig | None) parameters for the sampling
+        """
+        self.model_name = model_name
+        if model_name not in SUPPORTED_TRANSLATION_MODELS:
+            raise ValueError(
+                f"Model {model_name} is not supported for translation. Supported models are: {SUPPORTED_TRANSLATION_MODELS}"
+            )
+
+        llm_params = self.convert_to_dict(llm_params)
+        sampling_params = self.convert_to_dict(sampling_params)
+
+        self.device_str, self.device_id = self.setup_device(device, device_id)
+
+        self.batch_size = batch_size
+        self.split_batch = self.batch_size > 0
+
+        self.nmt_model = self.load_model(llm_params)
+        self.sampling_params = SamplingParams(**sampling_params)
+
+        self.source_language = source_language
+        self.target_language = target_language
+        self.prompt_template = self.get_prompt_template(model_name)
+        self.waitk = waitk
+
+    @staticmethod
+    def convert_to_dict(params: dict | DictConfig | None) -> dict:
+        """
+        Convert DictConfig to dict.
+        Args:
+            params: (dict | DictConfig | None) parameters to convert
+        Returns:
+            dict: converted parameters
+        """
+        if params is None:
+            return dict()
+        if isinstance(params, DictConfig):
+            return OmegaConf.to_container(params)
+        return params
+
+    @staticmethod
+    def setup_device(device: str, device_id: int) -> tuple[str, int]:
+        """
+        Setup device for the LLM model.
+        Args:
+            device: (str) device to run the model on
+            device_id: (int) device ID to run the model on
+        Returns:
+            device_str: (str) device string, e.g. "cuda:1"
+            device_id: (int) device ID, e.g. 1
+        Raises:
+            ValueError: if device is not supported, or CUDA is not available
+        """
+        if device == "cpu":
+            raise ValueError("Currently, CPU is not supported for vLLM.")
+
+        if device == "cuda":
+            if not torch.cuda.is_available():
+                raise ValueError("CUDA is not available.")
+
+            if device_id >= torch.cuda.device_count():
+                logging.warning(f"Device ID {device_id} is not available. Using GPU 0 instead.")
+                device_id = 0
+
+            device_str = f"cuda:{device_id}"
+            return device_str, device_id
+
+        raise ValueError(f"Unsupported device: {device}")
+
+    @staticmethod
+    def get_prompt_template(model_name: str) -> PromptTemplate:
+        """
+        Returns prompt template for the LLM model.
+        Args:
+            model_name: (str) name of the model to get prompt template for
+        Returns:
+            PromptTemplate: prompt template for the LLM model
+        Raises:
+            ValueError: if model is not supported for translation
+        """
+        if model_name in [EURO_LLM_INSTRUCT_SMALL, EURO_LLM_INSTRUCT_LARGE]:
+            return EuroLLMTranslatorPromptTemplate
+
+        raise ValueError(
+            f"Model {model_name} is not supported for translation. Supported models are: {SUPPORTED_TRANSLATION_MODELS}"
+        )
+
+    def load_model(self, llm_params: dict) -> LLM:
+        """
+        Load NMT model in vLLM format.
+        Args:
+            llm_params: (dict) parameters for the LLM model
+        Returns:
+            Loaded LLM instance.
+        Raises:
+            RuntimeError: If model loading fails.
+        """
+        try:
+            os.environ["CUDA_VISIBLE_DEVICES"] = str(self.device_id)
+            model = LLM(model=self.model_name, **llm_params)
+            return model
+        except Exception as e:
+            raise RuntimeError(f"Model loading failed: {str(e)}") from e
+
+    def translate_batch(
+        self,
+        asr_transcripts: list[str],
+        prefixes: list[str],
+        src_langs: list[str],
+        tgt_langs: list[str],
+        src_contexts: list[str],
+        tgt_contexts: list[str],
+    ) -> list[str]:
+        """
+        Translate ASR transcripts starting from pre-defined prefixes in target language.
+        Args:
+            asr_transcripts: (list[str]) batch of ASR transcripts to be translated
+            prefixes: (list[str]) batch of prefixes to start translation from
+            src_langs: (list[str]) batch of source languages
+            tgt_langs: (list[str]) batch of target languages
+            src_contexts: (list[str]) batch of source contexts
+            tgt_contexts: (list[str]) batch of target contexts
+        Returns:
+            list[str] translations of ASR transcripts
+        """
+        input_texts = []
+        for src_lang, tgt_lang, src_prefix, tgt_prefix, src_context, tgt_context in zip(
+            src_langs, tgt_langs, asr_transcripts, prefixes, src_contexts, tgt_contexts
+        ):
+            text = self.prompt_template.format(src_lang, tgt_lang, src_prefix, tgt_prefix, src_context, tgt_context)
+            input_texts.append(text)
+
+        outputs = self.nmt_model.generate(input_texts, self.sampling_params, use_tqdm=False)
+        translations = []
+        for tgt_prefix, output in zip(prefixes, outputs):
+            output_text = output.outputs[0].text
+            output_text = self.prompt_template.extract(output_text)
+            translations.append(f"{tgt_prefix}{output_text}")
+        return translations
+
+    def translate(
+        self,
+        asr_transcripts: list[str],
+        prefixes: list[str],
+        src_langs: list[str],
+        tgt_langs: list[str],
+        src_contexts: list[str],
+        tgt_contexts: list[str],
+    ) -> list[str]:
+        """
+        Translate ASR transcript starting from pre-defined prefix in target language.
+        Args:
+            asr_transcripts: (list[str]) ASR transcripts to be translated
+            prefixes: (list[str]) prefixes to start translation from
+            src_langs: (list[str]) source languages
+            tgt_langs: (list[str]) target languages
+            src_contexts: (list[str]) source contexts
+            tgt_contexts: (list[str]) target contexts
+        Returns:
+            list[str] translations of ASR transcripts
+        """
+        all_translations = []
+        n_requests = len(asr_transcripts)
+        bs = self.batch_size if self.split_batch else n_requests
+        for i in range(0, n_requests, bs):
+            all_translations.extend(
+                self.translate_batch(
+                    asr_transcripts=asr_transcripts[i : i + bs],
+                    prefixes=prefixes[i : i + bs],
+                    src_langs=src_langs[i : i + bs],
+                    tgt_langs=tgt_langs[i : i + bs],
+                    src_contexts=src_contexts[i : i + bs],
+                    tgt_contexts=tgt_contexts[i : i + bs],
+                )
+            )
+        return all_translations
+
+    def get_prefixes(
+        self,
+        asr_transcripts: list[str],
+        translations: list[str],
+        prev_translations: list[str],
+    ) -> list[str]:
+        """
+        Generates new prefixes in target language for the next translation step.
+        Args:
+            asr_transcripts: (list[str]) current ASR transcripts to be translated
+            translations: (list[str]) translations obtained with LLM on current step
+            prev_translations: (list[str]) translations obtained with LLM on previous step
+        Returns:
+            list[str] new prefixes for LLM translation
+        """
+
+        new_prefixes = []
+        for asr, trans, prev_trans in zip(asr_transcripts, translations, prev_translations):
+
+            # Longest common prefix of translations on current and previous steps
+            lcp = os.path.commonprefix([prev_trans, trans])
+            had_leading_space = lcp.startswith(" ")
+
+            # If lcp happens mid-word, remove generated ending up to the first full word
+            if (len(lcp) > 0) and (lcp[-1] not in f"{string.punctuation} "):
+                lcp = " ".join(lcp.split()[:-1])
+
+            # Remove trailing whitespaces
+            lcp = lcp.strip()
+
+            # Remove hallucinations if ASR transcript is empty string
+            if len(asr) == 0:
+                lcp = ""
+
+            # If the LLM-generated translations disagree too much between steps,
+            # and the translation falls more than waitk words behind the ASR transcript,
+            # the algorithm forcibly advances the prefix based on the current translation.
+            n_asr_words = len(asr.split())
+            n_lcp_words = len(lcp.split())
+            if (self.waitk > 0) and (n_asr_words - n_lcp_words > self.waitk):
+                num_words_to_pick = n_asr_words - self.waitk
+                new_prefix = " ".join(trans.split()[:num_words_to_pick])
+            else:
+                new_prefix = lcp
+
+            # Preserve leading space if it was present in the previous translation
+            if len(new_prefix) > 0 and had_leading_space and not new_prefix.startswith(" "):
+                new_prefix = " " + new_prefix
+
+            new_prefixes.append(new_prefix)
+
+        return new_prefixes

nemo/collections/asr/inference/nmt/prompts.py

@@ -0,0 +1,96 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import re
+from abc import ABC, abstractmethod
+
+
+class PromptTemplate(ABC):
+    """
+    Base class for prompt templates.
+    Derived classes should implement the format and extract methods.
+        - format: format the prompt template with the given arguments
+        - extract: extract the answer from the response
+    """
+
+    @classmethod
+    @abstractmethod
+    def format(cls, **kwargs) -> str:
+        """
+        Format the prompt template with the given arguments.
+        """
+        raise NotImplementedError()
+
+    @classmethod
+    @abstractmethod
+    def extract(cls, response: str) -> str:
+        """
+        Extract the answer from the response.
+        """
+        raise NotImplementedError()
+
+
+class EuroLLMTranslatorPromptTemplate(PromptTemplate):
+    """
+    Provides a prompt template for the EuroLLM model to perform translation.
+    """
+
+    PROMPT_TEMPLATE = (
+        "<|im_start|>system\n<|im_end|>\n"
+        "<|im_start|>user\n"
+        "Translate the following {src_lang} source text to {tgt_lang}. Always output text in the {tgt_lang} language:\n"
+        "{src_lang}: {src_text}\n"
+        "{tgt_lang}: <|im_end|>\n"
+        "<|im_start|>assistant\n"
+        "{tgt_text}"
+    )
+
+    @classmethod
+    def format(
+        cls,
+        src_lang: str,
+        tgt_lang: str,
+        src_prefix: str,
+        tgt_prefix: str,
+        src_context: str = "",
+        tgt_context: str = "",
+    ) -> str:
+        """
+        Generate a translation prompt for the EuroLLM model.
+        Args:
+            src_lang (str): Source language name.
+            tgt_lang (str): Target language name.
+            src_prefix (str): Source text to translate.
+            tgt_prefix (str): Optional target prefix or placeholder for completion.
+            src_context (str): Optional source context to start translation from.
+            tgt_context (str): Optional target context to start translation from.
+        Returns:
+            str: Formatted translation prompt.
+        """
+        src_text = f"{src_context} {src_prefix}"
+        tgt_text = f"{tgt_context} {tgt_prefix}"
+        src_text = re.sub(r'\s+', ' ', src_text).strip()
+        tgt_text = re.sub(r'\s+', ' ', tgt_text).strip()
+        return cls.PROMPT_TEMPLATE.format(src_lang=src_lang, tgt_lang=tgt_lang, src_text=src_text, tgt_text=tgt_text)
+
+    @classmethod
+    def extract(cls, response: str) -> str:
+        """
+        Extract the first line of text from a model response.
+        Args:
+            response (str): The full response from the model.
+        Returns:
+            str: The text before the first newline.
+        """
+        return response.split('\n')[0]

nemo/collections/asr/inference/pipelines/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

nemo/collections/asr/inference/pipelines/base_pipeline.py

@@ -0,0 +1,647 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import json
+import os
+import re
+from abc import abstractmethod
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Iterable
+
+import torch
+from omegaconf import DictConfig
+from torch import Tensor
+
+from nemo.collections.asr.inference.model_wrappers.asr_inference_wrapper import ASRInferenceWrapper
+from nemo.collections.asr.inference.pipelines.pipeline_interface import PipelineInterface
+from nemo.collections.asr.inference.streaming.buffering.audio_bufferer import BatchedAudioBufferer
+from nemo.collections.asr.inference.streaming.buffering.cache_feature_bufferer import BatchedCacheFeatureBufferer
+from nemo.collections.asr.inference.streaming.buffering.feature_bufferer import BatchedFeatureBufferer
+from nemo.collections.asr.inference.streaming.framing.multi_stream import ContinuousBatchedRequestStreamer
+from nemo.collections.asr.inference.streaming.framing.request import FeatureBuffer, Frame, Request
+from nemo.collections.asr.inference.streaming.framing.request_options import ASRRequestOptions
+from nemo.collections.asr.inference.streaming.state.state import StreamingState
+from nemo.collections.asr.inference.streaming.text.text_processing import StreamingTextProcessor
+from nemo.collections.asr.inference.utils.bpe_decoder import BPEDecoder
+from nemo.collections.asr.inference.utils.context_manager import CacheAwareContextManager
+from nemo.collections.asr.inference.utils.enums import RequestType
+from nemo.collections.asr.inference.utils.pipeline_utils import (
+    check_existance_of_required_attributes,
+    get_leading_punctuation_regex_pattern,
+    ids_to_text_without_stripping,
+)
+from nemo.collections.asr.inference.utils.progressbar import ProgressBar
+from nemo.collections.asr.inference.utils.text_segment import TextSegment
+from nemo.collections.common.tokenizers.tokenizer_spec import TokenizerSpec
+
+if TYPE_CHECKING:
+    from nemo.collections.asr.inference.itn.inverse_normalizer import AlignmentPreservingInverseNormalizer
+    from nemo.collections.asr.inference.nmt.llm_translator import LLMTranslator
+
+
+@dataclass
+class TranscribeStepOutput:
+    """
+    Stores the output of a single transcribe step.
+    """
+
+    stream_id: int
+    # Final transcript is the transcript generated started from the previous EoU to the current EoU
+    # It is finalized transcript, optionally punctuated and ITN-normalized. It's not subject to further modifications.
+    # Final segments contains metadata for each word/segment in the final transcript.
+    final_transcript: str = ""
+    final_segments: list[TextSegment] | None = None
+    final_translation: str = ""
+    # Partial transcript is the transcript generated started from the previous EoU up to the current frame
+    # It is not finalized transcript, it may be subject to further modifications.
+    # It can also contain transcript from future frames.
+    partial_transcript: str = ""
+    partial_translation: str = ""
+    # Current step transcript/translation is the transcript/translation generated from the current frame
+    current_step_transcript: str = ""
+    current_step_translation: str = ""
+
+    @classmethod
+    def from_state(cls, state: StreamingState, request: Request, sep: str = ' ') -> 'TranscribeStepOutput':
+        """
+        Create a TranscribeStepOutput from a StreamingState
+        Args:
+            state (StreamingState): The state to create the output from.
+            request (Request): The request to create the output from.
+            sep (str): The separator for the text postprocessor.
+        Returns:
+            TranscribeStepOutput: The output for the step.
+        """
+        final_transcript = state.final_transcript.strip()
+        final_segments = [seg.copy() for seg in state.final_segments]
+        if len(final_segments) > 0:
+            final_segments[0].text = final_segments[0].text.lstrip(sep)
+            final_segments[-1].text = final_segments[-1].text.rstrip(sep)
+
+        if final_transcript:
+            separator = ''
+            if not request.is_first and state.concat_with_space:
+                separator = sep
+            final_transcript = separator + final_transcript
+            if len(final_segments) > 0:
+                final_segments[0].text = separator + final_segments[0].text
+        return cls(
+            stream_id=request.stream_id,
+            final_transcript=final_transcript,
+            final_segments=final_segments,
+            partial_transcript=state.partial_transcript,
+            current_step_transcript=state.current_step_transcript,
+        )
+
+    def __str__(self) -> str:
+        """
+        Return a string representation of the TranscribeStepOutput
+        """
+        info = {
+            "final_transcript": self.final_transcript,
+            "final_translation": self.final_translation,
+            "partial_transcript": self.partial_transcript,
+            "partial_translation": self.partial_translation,
+            "current_step_transcript": self.current_step_transcript,
+        }
+        return json.dumps(info, indent=4, ensure_ascii=False)
+
+
+class BasePipeline(PipelineInterface):
+    """
+    Base class for all pipelines.
+    """
+
+    def __init__(self):
+        """Initialize state pool to store the state for each stream"""
+        self._state_pool: dict[int, StreamingState] = {}
+
+    def get_state(self, stream_id: int) -> StreamingState:
+        """Retrieve state for a given stream ID."""
+        return self._state_pool.get(stream_id, None)
+
+    def get_states(self, stream_ids: Iterable[int]) -> list[StreamingState]:
+        """Retrieve states for a list of stream IDs."""
+        return [self.get_state(stream_id) for stream_id in stream_ids]
+
+    def delete_state(self, stream_id: int) -> None:
+        """Delete the state from the state pool."""
+        if stream_id in self._state_pool:
+            del self._state_pool[stream_id]
+
+    def delete_states(self, stream_ids: Iterable[int]) -> None:
+        """Delete states for a list of stream IDs."""
+        for stream_id in stream_ids:
+            self.delete_state(stream_id)
+
+    def init_state(self, stream_id: int, options: ASRRequestOptions) -> StreamingState:
+        """Initialize the state of the stream"""
+        if stream_id not in self._state_pool:
+            state = self.create_state(options)
+            self._state_pool[stream_id] = state
+        return self._state_pool[stream_id]
+
+    def reset_session(self) -> None:
+        """Reset the frame buffer and internal state pool"""
+        self._state_pool.clear()
+
+    def open_session(self) -> None:
+        """Start a new session by resetting the internal state pool"""
+        self.reset_session()
+
+    def close_session(self) -> None:
+        """Close the session by resetting the internal state pool"""
+        self.reset_session()
+
+    @abstractmethod
+    def transcribe_step_for_frames(self, frames: list[Frame]) -> None:
+        """Transcribe a step for frames"""
+        pass
+
+    @abstractmethod
+    def transcribe_step_for_feature_buffers(self, fbuffers: list[FeatureBuffer]) -> None:
+        """Transcribe a step for feature buffers"""
+        pass
+
+    @abstractmethod
+    def get_request_generator(self) -> ContinuousBatchedRequestStreamer:
+        """Return the request generator."""
+        pass
+
+    @abstractmethod
+    def get_sep(self) -> str:
+        """Return the separator for the text postprocessor."""
+        pass
+
+    def translate_step(self, states: list[StreamingState], step_outputs: list[TranscribeStepOutput]) -> None:
+        """
+        Translate step
+        Args:
+            states (list[StreamingState]): List of StreamingState objects.
+            step_outputs (list[TranscribeStepOutput]): List of TranscribeStepOutput objects.
+        """
+        src_langs, tgt_langs = [], []
+        asr_transcripts, current_prefixes, previous_translations = [], [], []
+        final_transcript_mask = []
+        states_to_translate = []
+
+        src_contexts, tgt_contexts = [], []
+        for state, step_output in zip(states, step_outputs):
+            if not state.options.enable_nmt:
+                continue
+
+            src_lang = state.options.source_language
+            tgt_lang = state.options.target_language
+            if not src_lang or not tgt_lang:
+                raise ValueError("Source and target languages must be set when NMT is enabled")
+
+            final = step_output.final_transcript
+            partial = step_output.partial_transcript
+            if not (final.strip() or partial.strip()):
+                continue
+
+            transcript = final or partial
+            is_final = bool(final)
+            prev_translation, prefix = state.previous_translation_info
+
+            states_to_translate.append((state, step_output))
+            src_langs.append(src_lang)
+            tgt_langs.append(tgt_lang)
+            asr_transcripts.append(transcript)
+            current_prefixes.append(prefix)
+            previous_translations.append(prev_translation)
+            final_transcript_mask.append(is_final)
+
+            src_context, tgt_context = state.previous_context
+            src_contexts.append(src_context)
+            tgt_contexts.append(tgt_context)
+
+        if len(states_to_translate) == 0:
+            return
+
+        translations = self.nmt_model.translate(
+            asr_transcripts, current_prefixes, src_langs, tgt_langs, src_contexts, tgt_contexts
+        )
+        new_prefixes = self.nmt_model.get_prefixes(asr_transcripts, translations, previous_translations)
+
+        for (state, step_output), translation, new_prefix, prev_prefix, is_final in zip(
+            states_to_translate, translations, new_prefixes, current_prefixes, final_transcript_mask
+        ):
+            if is_final:
+                step_output.final_translation = translation
+                step_output.partial_translation = ""
+                state.cleanup_translation_info_after_eou()
+                state.set_translation_context(step_output.final_transcript, translation)
+                new_prefix = translation
+            else:
+                step_output.partial_translation = translation
+                step_output.final_translation = ""
+                state.set_translation_info(translation, new_prefix)
+
+            lcp = os.path.commonprefix([prev_prefix, new_prefix])
+            step_output.current_step_translation = new_prefix[len(lcp) :]
+
+    def transcribe_step(self, requests: list[Request]) -> list[TranscribeStepOutput]:
+        """
+        Transcribe step
+        Args:
+            requests (list[Request]): List of Request objects.
+        Returns:
+            list[TranscribeStepOutput]: List of TranscribeStepOutput objects.
+        """
+
+        # Initialize the state if it is the first request for the stream
+        states = []
+        for request in requests:
+            if request.is_first:
+                self.init_state(request.stream_id, request.options)
+            states.append(self.get_state(request.stream_id))
+
+        # Perform the transcribe step for the frames or feature buffers
+        if isinstance(requests[0], Frame):
+            self.transcribe_step_for_frames(frames=requests)
+        elif isinstance(requests[0], FeatureBuffer):
+            self.transcribe_step_for_feature_buffers(fbuffers=requests)
+        else:
+            raise ValueError(f"Invalid request type: {type(requests[0])}")
+
+        # Create current step output for each request
+        outputs = []
+        sep = self.get_sep()
+        for request, state in zip(requests, states):
+            step_output = TranscribeStepOutput.from_state(state=state, request=request, sep=sep)
+            outputs.append(step_output)
+
+        # Perform the translation step
+        if self.nmt_enabled:
+            self.translate_step(states=states, step_outputs=outputs)
+
+        # Cleanup the states after the response is sent
+        # If last request, delete state from the state pool to free memory
+        for state, request in zip(states, requests):
+            state.cleanup_after_response()
+            if request.is_last:
+                self.delete_state(request.stream_id)
+        return outputs
+
+    def copy_asr_model_attributes(self, asr_model: ASRInferenceWrapper) -> None:
+        """
+        Copy the attributes from the ASR model
+        Args:
+            asr_model (ASRInferenceWrapper): ASR model to copy the attributes from.
+        """
+        self.asr_model = asr_model
+        self.tokenizer = asr_model.tokenizer
+        self.device = asr_model.device
+        self.supports_punctuation = asr_model.supports_punctuation()
+        self.asr_supported_puncts = asr_model.supported_punctuation()
+        self.leading_regex_pattern = get_leading_punctuation_regex_pattern(self.asr_supported_puncts)
+        self.blank_id = asr_model.get_blank_id()
+        self.vocabulary = asr_model.get_vocabulary()
+        self.sep = asr_model.word_separator
+        self.underscore_id = asr_model.underscore_id
+        self.punctuation_ids = asr_model.punctuation_ids
+        self.language_token_ids = asr_model.language_token_ids
+        self.preprocessor, self.preprocessor_config = asr_model.create_preprocessor()
+        self.subsampling_factor = asr_model.get_subsampling_factor()
+        self.window_stride = asr_model.get_window_stride()
+        self.model_stride_in_secs = asr_model.get_model_stride(in_secs=True)
+        self.model_stride_in_milliseconds = asr_model.get_model_stride(in_milliseconds=True)
+
+    def update_partial_transcript(
+        self, requests: list[Request], tokenizer: TokenizerSpec, leading_regex_pattern: str
+    ) -> None:
+        """
+        Update partial and current step transcripts from the state.
+        Args:
+            requests (list[Request]): List of Request objects.
+            tokenizer (TokenizerSpec): Used to convert tokens into text
+            leading_regex_pattern (str): Regex pattern for the punctuation marks.
+        """
+        word_separator = self.get_sep()
+        for request in requests:
+            state = self.get_state(request.stream_id)
+            # state tokens represent all tokens accumulated since the EOU
+            # incomplete segment tokens are the remaining tokens on the right side of the buffer after EOU
+            all_tokens = state.tokens + state.incomplete_segment_tokens
+            if len(all_tokens) > 0:
+                pt_string = ids_to_text_without_stripping(all_tokens, tokenizer, word_separator)
+                if leading_regex_pattern:
+                    pt_string = re.sub(leading_regex_pattern, r'\1', pt_string)
+                state.partial_transcript = pt_string
+            else:
+                state.partial_transcript = ""
+
+            current_step_tokens = state.current_step_tokens
+            if len(current_step_tokens) > 0:
+                step_transcript = ids_to_text_without_stripping(current_step_tokens, tokenizer, word_separator)
+                state.current_step_transcript = step_transcript
+            else:
+                state.current_step_transcript = ""
+
+    def init_bpe_decoder(self) -> None:
+        """Initialize the BPE decoder"""
+        check_existance_of_required_attributes(
+            self,
+            [
+                'vocabulary',
+                'tokenizer',
+                'confidence_aggregator',
+                'asr_supported_puncts',
+                'word_boundary_tolerance',
+                'model_stride_in_secs',
+            ],
+        )
+
+        self.bpe_decoder = BPEDecoder(
+            vocabulary=self.vocabulary,
+            tokenizer=self.tokenizer,
+            confidence_aggregator=self.confidence_aggregator,
+            asr_supported_puncts=self.asr_supported_puncts,
+            word_boundary_tolerance=self.word_boundary_tolerance,
+            token_duration_in_secs=self.model_stride_in_secs,
+        )
+
+    def init_text_processor(
+        self,
+        cfg: DictConfig,
+        itn_model: AlignmentPreservingInverseNormalizer | None,
+    ) -> None:
+        """
+        Initialize the text processor.
+        Args:
+            cfg: (DictConfig) Configuration parameters.
+            itn_model: (AlignmentPreservingInverseNormalizer | None) Inverse Text Normalization model.
+        """
+        check_existance_of_required_attributes(
+            self,
+            [
+                'asr_supported_puncts',
+                'supports_punctuation',
+                'confidence_aggregator',
+                'sep',
+            ],
+        )
+
+        self.text_processor = StreamingTextProcessor(
+            itn_cfg=cfg.itn,
+            itn_model=itn_model,
+            asr_supported_puncts=self.asr_supported_puncts,
+            asr_supports_punctuation=self.supports_punctuation,
+            confidence_aggregator=self.confidence_aggregator,
+            sep=self.sep,
+            enable_itn=cfg.enable_itn,
+        )
+
+    def init_nmt_model(self, nmt_model: LLMTranslator | None) -> None:
+        """
+        Initialize the Translation model.
+        Args:
+            nmt_model: (LLMTranslator | None) LLM based translation model.
+        """
+        self.nmt_model = nmt_model
+        self.nmt_enabled = nmt_model is not None
+
+    def init_bufferer_for_buffered_streaming(self) -> None:
+        """Initialize the bufferer."""
+        check_existance_of_required_attributes(
+            self,
+            [
+                'request_type',
+                'sample_rate',
+                'buffer_size_in_secs',
+                'preprocessor_config',
+                'device',
+            ],
+        )
+
+        if self.request_type is RequestType.FEATURE_BUFFER:
+            # Feature buffering: It will be used when the input is feature buffers
+            self.bufferer = BatchedFeatureBufferer(
+                sample_rate=self.sample_rate,
+                buffer_size_in_secs=self.buffer_size_in_secs,
+                preprocessor_cfg=self.preprocessor_config,
+                device=self.device,
+            )
+        elif self.request_type is RequestType.FRAME:
+            # Audio buffering: It will be used when the input is audio frames
+            self.bufferer = BatchedAudioBufferer(
+                sample_rate=self.sample_rate, buffer_size_in_secs=self.buffer_size_in_secs
+            )
+        else:
+            raise ValueError(f"Unknown request type: {self.request_type}")
+
+    def init_bufferer_for_cache_aware_streaming(self) -> None:
+        """Initialize the bufferer for cache-aware streaming."""
+        check_existance_of_required_attributes(
+            self,
+            [
+                'num_slots',
+                'use_feat_cache',
+                'chunk_size_in_secs',
+                'buffer_size_in_secs',
+                'sample_rate',
+                'preprocessor_config',
+                'device',
+            ],
+        )
+
+        if self.use_feat_cache:
+            # Only calculate mel-spec features for last chunk
+            chunk_size_for_feature_buffer = self.chunk_size_in_secs
+        else:
+            # Calculate mel-spec features for the whole buffer
+            chunk_size_for_feature_buffer = self.buffer_size_in_secs
+
+        self.bufferer = BatchedCacheFeatureBufferer(
+            num_slots=self.num_slots,
+            sample_rate=self.sample_rate,
+            buffer_size_in_secs=self.buffer_size_in_secs,
+            chunk_size_in_secs=chunk_size_for_feature_buffer,
+            preprocessor_cfg=self.preprocessor_config,
+            device=self.device,
+        )
+
+    def init_context_manager(self) -> None:
+        """Initialize the context manager."""
+        check_existance_of_required_attributes(self, ['asr_model', 'num_slots', 'use_cache'])
+        self.context_manager = CacheAwareContextManager(
+            cache_aware_model=self.asr_model, num_slots=self.num_slots, use_cache=self.use_cache
+        )
+
+    def init_prompt_support(self) -> None:
+        """Initialize prompt support for multilingual models."""
+        self.prompt_enabled = hasattr(self.asr_model.asr_model, 'concat') and self.asr_model.asr_model.concat
+
+        if self.prompt_enabled:
+            self._prompt_config = self._load_prompt_config()
+
+    def _load_prompt_config(self) -> dict:
+        """
+        Load prompt configuration from model.
+        Returns:
+            (dict) Prompt configuration containing num_prompts, prompt_dict, and compute_dtype.
+        """
+        cfg = self.asr_model.asr_model.cfg
+        if cfg and hasattr(cfg, 'model_defaults'):
+            model_defaults = cfg.model_defaults
+            num_prompts = model_defaults.get('num_prompts', None)
+            prompt_dict = model_defaults.get('prompt_dictionary', None)
+
+            # Validate and convert types once
+            num_prompts_int = int(num_prompts) if num_prompts is not None else 0
+
+            is_dict_like = isinstance(prompt_dict, dict) or (
+                hasattr(prompt_dict, 'get') and hasattr(prompt_dict, '__contains__')
+            )
+
+            if num_prompts_int > 0 and is_dict_like:
+                return {
+                    'num_prompts': num_prompts_int,
+                    'prompt_dict': prompt_dict,
+                    'compute_dtype': getattr(self.asr_model.asr_model, 'dtype', torch.float32),
+                }
+
+        return {}
+
+    def _resolve_prompt_index(self, language_code: str) -> int:
+        """
+        Resolve language_code to a strict prompt index; raise if invalid.
+        Args:
+            language_code: (str) Language code to resolve (e.g., "en-US", "es-ES").
+        Returns:
+            (int) Prompt index corresponding to the language code.
+        Raises:
+            RuntimeError: If prompt configuration is missing.
+            ValueError: If language_code is not found in prompt dictionary.
+        """
+        if not hasattr(self, '_prompt_config') or not self._prompt_config:
+            raise RuntimeError("Prompt configuration is missing for a prompt-enabled model.")
+        prompt_dict = self._prompt_config['prompt_dict']
+        lang_index = prompt_dict.get(language_code, None)
+        if lang_index is None:
+            raise ValueError(
+                f"Language code '{language_code}' not found in prompt dictionary. "
+                f"Available languages: {list(prompt_dict.keys())}"
+            )
+        return lang_index
+
+    def _create_one_hot_prompts(self, indices: Tensor) -> Tensor:
+        """
+        Create one-hot prompt vectors from indices.
+        Args:
+            indices: (Tensor) Prompt indices of shape [B].
+        Returns:
+            (Tensor) One-hot prompt vectors of shape [B, num_prompts].
+        """
+        num_prompts = self._prompt_config['num_prompts']
+        return torch.nn.functional.one_hot(indices, num_classes=num_prompts).to(self._prompt_config['compute_dtype'])
+
+    def _build_prompt_vectors(self, states: list) -> Tensor:
+        """
+        Build prompt vectors for a batch of states using one-hot encoding.
+        Args:
+            states: (list) List of streaming states.
+        Returns:
+            (Tensor) Prompt vectors of shape [B, num_prompts].
+        Raises:
+            ValueError: If any prompt index is out of range.
+        """
+        indices = torch.tensor([getattr(s, 'prompt_idx', 0) for s in states], device=self.device, dtype=torch.long)
+        num_prompts = self._prompt_config['num_prompts']
+        if torch.any((indices < 0) | (indices >= num_prompts)):
+            raise ValueError("Found out-of-range prompt index in batch.")
+        return self._create_one_hot_prompts(indices)
+
+    def run(
+        self,
+        audio_filepaths: list[str],
+        options: list[ASRRequestOptions] | None = None,
+        progress_bar: ProgressBar | None = None,
+    ) -> dict:
+        """
+        Orchestrates reading from audio_filepaths in a streaming manner,
+        transcribes them, and packs the results into a PipelineOutput.
+        Args:
+            audio_filepaths (list[str]): List of audio filepaths to transcribe.
+            options (list[ASRRequestOptions] | None): List of RequestOptions for each stream.
+            progress_bar (ProgressBar | None): Progress bar to show the progress. Default is None.
+        Returns:
+            dict: A dictionary containing transcriptions and segments for each stream.
+        """
+        if progress_bar is not None and not isinstance(progress_bar, ProgressBar):
+            raise ValueError("progress_bar must be an instance of ProgressBar.")
+
+        if options is None:
+            # Use default options if not provided
+            options = [ASRRequestOptions() for _ in audio_filepaths]
+
+        if len(options) != len(audio_filepaths):
+            raise ValueError("options must be the same length as audio_filepaths")
+
+        request_generator = self.get_request_generator()
+        request_generator.set_audio_filepaths(audio_filepaths, options)
+        request_generator.set_progress_bar(progress_bar)
+
+        pipeline_output = {}
+        sep = self.get_sep()
+        self.open_session()
+        for requests in request_generator:
+            step_outputs = self.transcribe_step(requests)
+            for step_output in step_outputs:
+                stream_id = step_output.stream_id
+                if stream_id not in pipeline_output:
+                    pipeline_output[stream_id] = {
+                        "text": "",
+                        "translation": "",
+                        "segments": [],
+                        "audio_filepath": request_generator.get_audio_filepath(stream_id),
+                        "translation_segments": [],
+                        "asr_segments": [],
+                    }
+
+                accumulated_text = pipeline_output[stream_id]["text"]
+                accumulated_translation = pipeline_output[stream_id]["translation"]
+                final_transcript = step_output.final_transcript
+                final_translation = step_output.final_translation
+                final_segments = step_output.final_segments
+                if not accumulated_text:
+                    final_transcript = final_transcript.lstrip(sep)
+                    if len(final_segments) > 0:
+                        first_segment = final_segments[0]
+                        first_segment.text = first_segment.text.lstrip(sep)
+
+                if not accumulated_translation:
+                    final_translation = final_translation.lstrip(sep)
+
+                accumulated_text += final_transcript
+                accumulated_translation += final_translation
+                pipeline_output[stream_id]["text"] = accumulated_text
+                pipeline_output[stream_id]["translation"] = accumulated_translation
+                pipeline_output[stream_id]["segments"].extend(final_segments)
+
+                # Record the text finalized in this step with the audio elapsed at finalization.
+                if final_transcript:
+                    delay = request_generator.get_elapsed_duration(stream_id)
+                    pipeline_output[stream_id]["asr_segments"].append((final_transcript, delay))
+
+                if self.nmt_enabled:
+                    step_translation = step_output.current_step_translation
+                    delay = request_generator.get_elapsed_duration(stream_id)
+                    pipeline_output[stream_id]["translation_segments"].append((step_translation, delay))
+
+        self.close_session()
+        return pipeline_output

nemo/collections/asr/inference/pipelines/buffered_ctc_pipeline.py

@@ -0,0 +1,443 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import math
+from typing import TYPE_CHECKING
+
+import torch
+from omegaconf import DictConfig
+from torch import Tensor
+
+from nemo.collections.asr.inference.model_wrappers.ctc_inference_wrapper import CTCInferenceWrapper
+from nemo.collections.asr.inference.pipelines.base_pipeline import BasePipeline
+from nemo.collections.asr.inference.streaming.decoders.greedy.greedy_ctc_decoder import ClippedCTCGreedyDecoder
+from nemo.collections.asr.inference.streaming.endpointing.greedy.greedy_ctc_endpointing import CTCGreedyEndpointing
+from nemo.collections.asr.inference.streaming.framing.multi_stream import ContinuousBatchedRequestStreamer
+from nemo.collections.asr.inference.streaming.framing.request import FeatureBuffer, Frame, Request
+from nemo.collections.asr.inference.streaming.framing.request_options import ASRRequestOptions
+from nemo.collections.asr.inference.streaming.state.ctc_state import CTCStreamingState
+from nemo.collections.asr.inference.utils.enums import FeatureBufferPaddingMode, RequestType
+from nemo.collections.asr.inference.utils.pipeline_utils import (
+    check_existance_of_required_attributes,
+    drop_trailing_features,
+    get_confidence_utils,
+    normalize_features,
+    normalize_log_probs,
+)
+
+if TYPE_CHECKING:
+    from nemo.collections.asr.inference.itn.inverse_normalizer import AlignmentPreservingInverseNormalizer
+    from nemo.collections.asr.inference.nmt.llm_translator import LLMTranslator
+
+
+class BufferedCTCPipeline(BasePipeline):
+    """Buffered CTC pipeline."""
+
+    def __init__(
+        self,
+        cfg: DictConfig,
+        asr_model: CTCInferenceWrapper,
+        itn_model: AlignmentPreservingInverseNormalizer | None = None,
+        nmt_model: LLMTranslator | None = None,
+    ):
+        """
+        Initialize the BufferedCTCPipeline.
+        Args:
+            cfg: (DictConfig) Configuration parameters.
+            asr_model: (CTCInferenceWrapper) ASR model.
+            itn_model: (AlignmentPreservingInverseNormalizer | None) Inverse Text Normalization model.
+            nmt_model: (LLMTranslator | None) LLM based translation model.
+        """
+        self.copy_asr_model_attributes(asr_model)
+        self.init_parameters(cfg)
+        self.init_bufferer_for_buffered_streaming()
+        self.conf_func, self.confidence_aggregator = get_confidence_utils(cfg.confidence)
+        self.init_endpointer()
+        self.init_bpe_decoder()
+        self.init_greedy_ctc_decoder()
+        self.init_text_processor(cfg, itn_model)
+        self.init_nmt_model(nmt_model)
+        super().__init__()
+
+    def init_parameters(self, cfg: DictConfig) -> None:
+        """
+        Initialize the configuration parameters.
+        Args:
+            cfg: (DictConfig) Configuration parameters.
+        """
+        self.sample_rate = cfg.streaming.sample_rate
+        self.asr_output_granularity = cfg.asr_output_granularity
+        self.batch_size = cfg.streaming.batch_size
+
+        self.chunk_size = cfg.streaming.chunk_size
+        self.left_padding_size = cfg.streaming.left_padding_size
+        self.right_padding_size = cfg.streaming.right_padding_size
+        self.buffer_size_in_secs = self.chunk_size + self.left_padding_size + self.right_padding_size
+        self.expected_feature_buffer_len = int(self.buffer_size_in_secs / self.window_stride)
+        self.tokens_per_frame_float = self.chunk_size / self.model_stride_in_secs
+        self.tokens_per_frame = math.ceil(self.tokens_per_frame_float)
+        self.initial_delay = (self.left_padding_size + self.right_padding_size) / self.model_stride_in_secs
+        self.mid_delay = math.ceil((self.chunk_size + self.right_padding_size) / self.model_stride_in_secs)
+
+        self.stop_history_eou_in_milliseconds = cfg.endpointing.stop_history_eou
+        self.residue_tokens_at_end = cfg.endpointing.residue_tokens_at_end
+        self.request_type = RequestType.from_str(cfg.streaming.request_type)
+        self.word_boundary_tolerance = cfg.streaming.word_boundary_tolerance
+        self.padding_mode = FeatureBufferPaddingMode.from_str(cfg.streaming.padding_mode)
+        self.right_padding = self.padding_mode is FeatureBufferPaddingMode.RIGHT
+        self.return_tail_result = cfg.return_tail_result
+
+        self.zero_log_probs = self.init_zero_log_probs() if self.right_padding else None
+
+    def init_endpointer(self) -> None:
+        """Initialize the endpointing."""
+        check_existance_of_required_attributes(
+            self,
+            [
+                'vocabulary',
+                'model_stride_in_milliseconds',
+                'stop_history_eou_in_milliseconds',
+                'residue_tokens_at_end',
+            ],
+        )
+
+        self.endpointer = CTCGreedyEndpointing(
+            vocabulary=self.vocabulary,
+            ms_per_timestep=self.model_stride_in_milliseconds,
+            stop_history_eou=self.stop_history_eou_in_milliseconds,
+            residue_tokens_at_end=self.residue_tokens_at_end,
+        )
+
+    def init_greedy_ctc_decoder(self) -> None:
+        """Initialize the CTC decoder."""
+        check_existance_of_required_attributes(self, ['vocabulary', 'conf_func', 'endpointer', 'tokens_per_frame'])
+        self.greedy_ctc_decoder = ClippedCTCGreedyDecoder(
+            vocabulary=self.vocabulary,
+            conf_func=self.conf_func,
+            endpointer=self.endpointer,
+            tokens_per_frame=self.tokens_per_frame,
+        )
+
+    def init_zero_log_probs(self) -> Tensor:
+        """
+        Initialize the log probabilities for the zero buffer.
+        Returns:
+            (Tensor) Log probabilities for the zero buffer.
+        """
+        check_existance_of_required_attributes(
+            self, ['asr_model', 'buffer_size_in_secs', 'sample_rate', 'device', 'expected_feature_buffer_len']
+        )
+        buffer_size_in_samples = int(self.buffer_size_in_secs * self.sample_rate)
+        zero_buffer = torch.zeros(1, buffer_size_in_samples, device=self.device)
+        zero_features, zero_features_len = self.preprocess(
+            buffers=zero_buffer,
+            buffer_lens=torch.tensor([zero_buffer.shape[1]], device=self.device),
+            expected_feature_buffer_len=self.expected_feature_buffer_len,
+        )
+        return self.asr_model.get_logprobs(processed_signal=zero_features, processed_signal_length=zero_features_len)[
+            0
+        ]
+
+    def create_state(self, options: ASRRequestOptions) -> CTCStreamingState:
+        """
+        Create new empty state.
+        Args:
+            options: (ASRRequestOptions) Request options for particular stream.
+        Returns:
+            (CTCStreamingState) New empty state.
+        """
+        state = CTCStreamingState()
+        state.set_global_offset(-self.initial_delay)
+        new_options = options.fill_defaults(
+            default_enable_itn=self.text_processor.itn_enabled,
+            default_enable_nmt=self.nmt_enabled,
+            default_source_language=self.nmt_model.source_language if self.nmt_enabled else None,
+            default_target_language=self.nmt_model.target_language if self.nmt_enabled else None,
+            default_stop_history_eou=self.stop_history_eou_in_milliseconds,
+            default_asr_output_granularity=self.asr_output_granularity,
+        )
+        state.set_options(new_options)
+        return state
+
+    def get_sep(self) -> str:
+        """Return the separator for the text processor."""
+        return self.sep
+
+    def get_cut_off_range(self, T: int, is_last: bool) -> tuple[int, int]:
+        """
+        Compute the start and end indices to clip the log probs.
+        Args:
+            T: (int) Time dimension of the log probabilities.
+            is_last: (bool) Whether the last frame is reached.
+        Returns:
+            (tuple[int, int]) Start and end indices to clip the log probs.
+        """
+        start = max(T - 1 - self.mid_delay, 0)
+        end = T if is_last else min(start + self.tokens_per_frame, T)
+        return start, end
+
+    def preprocess(
+        self, buffers: Tensor, buffer_lens: Tensor, expected_feature_buffer_len: int
+    ) -> tuple[Tensor, Tensor]:
+        """
+        Preprocess the buffered frames and extract features.
+        Args:
+            buffers: (Tensor) Audio buffers.
+            buffer_lens: (Tensor) Lengths of the audio buffers.
+            expected_feature_buffer_len: (int) Expected length of the feature buffers.
+        Returns:
+            (tuple[Tensor, Tensor]) Processed feature buffers and their lengths.
+        """
+        feature_buffers, feature_buffer_lens = self.preprocessor(input_signal=buffers, length=buffer_lens)
+        feature_buffers = drop_trailing_features(feature_buffers, expected_feature_buffer_len)
+        feature_buffers = normalize_features(feature_buffers, feature_buffer_lens)
+        feature_buffer_lens = feature_buffer_lens.clamp(max=feature_buffers.shape[2])
+        return feature_buffers, feature_buffer_lens
+
+    def get_logprobs_given_raw_signals(
+        self, frames: list[Frame], raw_signals: list[Tensor], left_paddings: list[int]
+    ) -> Tensor:
+        """
+        Get log probs from the CTC model.
+        Args:
+            frames: (list[Frame]) Frames to transcribe.
+            raw_signals: (list[Tensor]) Audio buffers.
+            left_paddings: (list[int]) Left paddings for audio buffers.
+        Returns:
+            (Tensor) Log probabilities.
+        """
+
+        if self.right_padding:
+            left_paddings = torch.tensor(left_paddings, dtype=torch.int64, device=self.device)
+
+        buffers = []
+        for i in range(len(raw_signals)):
+            buffer = raw_signals[i]
+            # Roll the buffered frames to the left by the left padding
+            # This is done to avoid the padding at the beginning of the buffered frames
+            # which can cause the performance degradation
+            if self.right_padding:
+                lpad = left_paddings[i].item()
+                if lpad > 0:
+                    buffer = buffer.roll(shifts=-lpad)
+            buffers.append(buffer.unsqueeze_(0))
+
+        # Only final frames have right padding
+        # Calculate right paddings
+        right_paddings = torch.tensor([frame.size - frame.valid_size for frame in frames], device=self.device).clamp(
+            min=0
+        )
+
+        # Create and adjust the buffer lens
+        buffer_lens = torch.tensor([buffers[0].size(1)] * len(buffers), device=self.device)
+        buffer_lens = buffer_lens - right_paddings
+        if self.right_padding:
+            buffer_lens = buffer_lens - left_paddings
+
+        # Preprocess the buffers with corresponding buffer lens
+        feature_buffers, feature_buffer_lens = self.preprocess(
+            buffers=torch.cat(buffers).to(self.device),
+            buffer_lens=buffer_lens,
+            expected_feature_buffer_len=self.expected_feature_buffer_len,
+        )
+
+        # Get the log probabilities from the ASR model
+        log_probs = self.asr_model.get_logprobs(
+            processed_signal=feature_buffers, processed_signal_length=feature_buffer_lens
+        ).clone()
+
+        # Roll back the log probabilities to the right
+        if self.right_padding:
+            for i in range(len(log_probs)):
+                lpad = left_paddings[i]
+                if lpad > 0:
+                    lpad = int(lpad / self.sample_rate / self.model_stride_in_secs)
+                    log_probs[i] = log_probs[i].roll(lpad, dims=0)
+                    log_probs[i][:lpad, :] = self.zero_log_probs[:lpad, :]
+        return log_probs
+
+    def get_logprobs_given_processed_signals(
+        self, fbuffers: list[FeatureBuffer], processed_signals: list[Tensor]
+    ) -> Tensor:
+        """
+        Get log probs from the ASR model.
+        Args:
+            fbuffers: (list[FeatureBuffer]) Feature buffers.
+            processed_signals: (list[Tensor]) Processed buffers.
+        Returns:
+            (Tensor) Log probabilities.
+        """
+        processed_signals = torch.cat([sig.unsqueeze_(0) for sig in processed_signals]).to(self.device)
+        processed_signals = drop_trailing_features(processed_signals, self.expected_feature_buffer_len)
+        processed_signal_lengths = torch.tensor([f.valid_size for f in fbuffers], device=self.device)
+        processed_signals = normalize_features(processed_signals, processed_signal_lengths)
+        processed_signal_lengths = processed_signal_lengths.clamp(max=processed_signals.shape[2])
+
+        log_probs = self.asr_model.get_logprobs(
+            processed_signal=processed_signals, processed_signal_length=processed_signal_lengths
+        ).clone()
+
+        if self.right_padding:
+            for i in range(len(log_probs)):
+                lpad = int(fbuffers[i].roll_size / self.subsampling_factor)
+                if lpad > 0:
+                    log_probs[i] = log_probs[i].roll(lpad, dims=0)
+                    log_probs[i][:lpad, :] = self.zero_log_probs[:lpad, :]
+        return log_probs
+
+    def compute_logprobs_from_frames(self, frames: list[Frame]) -> Tensor:
+        """
+        Buffer the frames and get the log probabilities.
+        Args:
+            frames: (list[Frame]) List of frames to transcribe.
+        Returns:
+            (Tensor) Log probabilities.
+        """
+        raw_signals, left_paddings = self.bufferer.update(frames)
+        log_probs = None
+        if len(raw_signals) > 0:
+            log_probs = self.get_logprobs_given_raw_signals(frames, raw_signals, left_paddings)
+        return log_probs
+
+    def compute_logprobs_from_feature_buffers(self, fbuffers: list[FeatureBuffer]) -> Tensor:
+        """
+        Buffer the feature buffers and get the log probabilities.
+        Args:
+            fbuffers: (list[FeatureBuffer]) List of feature buffers to transcribe.
+        Returns:
+            (Tensor) Log probabilities.
+        """
+        processed_signals = self.bufferer.update(fbuffers)
+        log_probs = None
+        if len(processed_signals) > 0:
+            log_probs = self.get_logprobs_given_processed_signals(fbuffers, processed_signals)
+        return log_probs
+
+    def run_greedy_decoder(
+        self, state: CTCStreamingState, request: Request, buffer_log_probs: Tensor, start: int, end: int
+    ) -> bool:
+        """
+        Run Greedy decoder, update state and trigger EOU detection.
+        Args:
+            state: (CTCStreamingState) Current state for the particular stream.
+            request: (Request) Current request for the particular stream.
+            buffer_log_probs: (Tensor) Log probabilities.
+            start: (int) Start index of the log probabilities.
+            end: (int) End index of the log probabilities.
+        Returns:
+            (bool) Whether EOU is detected.
+        """
+        clipped_output, tail_output, eou_detected, start_idx, end_idx = self.greedy_ctc_decoder(
+            buffer_log_probs,
+            start,
+            end,
+            request.is_last,
+            is_start=request.is_first,
+            return_partial_result=self.return_tail_result,
+            state_start_idx=state.decoder_start_idx,
+            state_end_idx=state.decoder_end_idx,
+            stop_history_eou=state.options.stop_history_eou,
+            compute_confidence=True,
+        )
+
+        state.update_state(clipped_output, eou_detected)
+        state.set_last_token(clipped_output["last_token"], clipped_output["last_token_idx"])
+        state.update_from_decoder_results(start_idx, end_idx)
+        state.increment_global_offset(self.tokens_per_frame_float)
+        state.set_incomplete_segment_tokens(tail_output["tokens"])
+        return eou_detected
+
+    def shared_transcribe_step(self, requests: list[Request], log_probs: Tensor) -> None:
+        """
+        Shared transcribe step for frames and feature buffers.
+        Args:
+            requests: (list[Request]) List of frames or feature buffers to transcribe.
+            log_probs: (Tensor) Log probabilities.
+        """
+        postponed_requests = [(ridx, request.stream_id) for ridx, request in enumerate(requests)]
+        next_postponed_requests = []
+
+        while len(postponed_requests) > 0:
+
+            ready_state_ids = set()
+            for ridx, stream_id in postponed_requests:
+
+                if stream_id in ready_state_ids:
+                    # Skip if the state is already ready
+                    next_postponed_requests.append((ridx, stream_id))
+                    continue
+
+                request = requests[ridx]
+                state = self.get_state(stream_id)
+                lp = log_probs[ridx].cpu()
+                start, end = self.get_cut_off_range(lp.shape[0], request.is_last)
+                eou_detected = self.run_greedy_decoder(state, request, lp, start, end)
+
+                if eou_detected:
+                    self.bpe_decoder.decode_bpe_tokens(state)
+                    state.cleanup_after_eou()
+                    ready_state_ids.add(stream_id)
+
+            if len(ready_state_ids) > 0:
+                self.text_processor.process([self.get_state(stream_id) for stream_id in ready_state_ids])
+                ready_state_ids.clear()
+
+            postponed_requests = next_postponed_requests.copy()
+            next_postponed_requests.clear()
+
+        self.update_partial_transcript(requests, self.tokenizer, self.leading_regex_pattern)
+
+    def transcribe_step_for_feature_buffers(self, fbuffers: list[FeatureBuffer]) -> None:
+        """
+        Transcribe a step for feature buffers.
+        Args:
+            fbuffers: (list[FeatureBuffer]) List of feature buffers to transcribe.
+        """
+        log_probs = self.compute_logprobs_from_feature_buffers(fbuffers)
+        if log_probs is not None:
+            log_probs = normalize_log_probs(log_probs)
+            self.shared_transcribe_step(requests=fbuffers, log_probs=log_probs)
+
+    def transcribe_step_for_frames(self, frames: list[Frame]) -> None:
+        """
+        Transcribe step for frames.
+        Args:
+            frames: (list[Frame]) List of frames to transcribe.
+        """
+        log_probs = self.compute_logprobs_from_frames(frames)
+        if log_probs is not None:
+            log_probs = normalize_log_probs(log_probs)
+            self.shared_transcribe_step(requests=frames, log_probs=log_probs)
+
+    def get_request_generator(self) -> ContinuousBatchedRequestStreamer:
+        """
+        Initialize the request generator.
+        Returns:
+            (ContinuousBatchedRequestStreamer) Request generator.
+        """
+        request_generator = ContinuousBatchedRequestStreamer(
+            n_frames_per_stream=1,
+            frame_size_in_secs=self.chunk_size,
+            sample_rate=self.sample_rate,
+            batch_size=self.batch_size,
+            request_type=self.request_type,
+            preprocessor=self.preprocessor,
+            buffer_size_in_secs=self.buffer_size_in_secs,
+            device=self.device,
+            pad_last_frame=True,
+            right_pad_features=self.right_padding,
+        )
+        return request_generator

nemo/collections/asr/inference/pipelines/buffered_rnnt_pipeline.py

@@ -0,0 +1,818 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import math
+from typing import TYPE_CHECKING
+
+import numpy as np
+import torch
+from omegaconf import DictConfig
+from torch import Tensor
+
+from nemo.collections.asr.inference.model_wrappers.rnnt_inference_wrapper import RNNTInferenceWrapper
+from nemo.collections.asr.inference.pipelines.base_pipeline import BasePipeline
+from nemo.collections.asr.inference.streaming.decoders.greedy.greedy_rnnt_decoder import ClippedRNNTGreedyDecoder
+from nemo.collections.asr.inference.streaming.endpointing.greedy.greedy_rnnt_endpointing import RNNTGreedyEndpointing
+from nemo.collections.asr.inference.streaming.framing.multi_stream import ContinuousBatchedRequestStreamer
+from nemo.collections.asr.inference.streaming.framing.request import FeatureBuffer, Frame, Request
+from nemo.collections.asr.inference.streaming.framing.request_options import ASRRequestOptions
+from nemo.collections.asr.inference.streaming.state.rnnt_state import RNNTStreamingState
+from nemo.collections.asr.inference.utils.enums import FeatureBufferPaddingMode, RequestType
+from nemo.collections.asr.inference.utils.pipeline_utils import (
+    adjust_vad_segments,
+    check_existance_of_required_attributes,
+    drop_trailing_features,
+    get_confidence_utils,
+    normalize_features,
+    update_punctuation_and_language_tokens_timestamps,
+)
+from nemo.collections.asr.parts.utils.rnnt_utils import Hypothesis as NemoHypothesis
+from nemo.collections.asr.parts.utils.rnnt_utils import batched_hyps_to_hypotheses
+from nemo.utils import logging
+
+if TYPE_CHECKING:
+    from nemo.collections.asr.inference.itn.inverse_normalizer import AlignmentPreservingInverseNormalizer
+    from nemo.collections.asr.inference.nmt.llm_translator import LLMTranslator
+
+
+class BufferedRNNTPipeline(BasePipeline):
+    """Buffered RNN-T/TDT pipeline."""
+
+    def __init__(
+        self,
+        cfg: DictConfig,
+        asr_model: RNNTInferenceWrapper,
+        itn_model: AlignmentPreservingInverseNormalizer | None = None,
+        nmt_model: LLMTranslator | None = None,
+    ):
+        """
+        Initialize the BufferedRNNTPipeline.
+        Args:
+            cfg: (DictConfig) Configuration parameters.
+            asr_model: (RNNTInferenceWrapper) ASR model.
+            itn_model: (AlignmentPreservingInverseNormalizer | None) Inverse Text Normalization model.
+            nmt_model: (LLMTranslator | None) LLM based translation model.
+        """
+
+        self.copy_asr_model_attributes(asr_model)
+        self.init_prompt_support()
+        self.init_parameters(cfg)
+        self.init_bufferer_for_buffered_streaming()
+        self.conf_func, self.confidence_aggregator = get_confidence_utils(cfg.confidence)
+        self.init_endpointer()
+        self.init_greedy_rnnt_decoder()
+        self.init_bpe_decoder()
+        self.init_decoding_computer()
+        self.init_text_processor(cfg, itn_model)
+        self.init_nmt_model(nmt_model)
+        super().__init__()
+
+    def init_parameters(self, cfg: DictConfig) -> None:
+        """
+        Initialize the configuration parameters.
+        Args:
+            cfg: (DictConfig) Configuration parameters.
+        """
+        self.asr_output_granularity = cfg.asr_output_granularity
+        self.sample_rate = cfg.streaming.sample_rate
+        self.stateful = cfg.streaming.stateful
+        self.stateless = not self.stateful
+        self.batch_size = cfg.streaming.batch_size
+
+        self.chunk_size = cfg.streaming.chunk_size
+        self.left_padding_size = cfg.streaming.left_padding_size
+        self.right_padding_size = cfg.streaming.right_padding_size
+        self.buffer_size_in_secs = self.chunk_size + self.left_padding_size + self.right_padding_size
+        self.expected_feature_buffer_len = int(self.buffer_size_in_secs / self.window_stride)
+
+        self.mid_delay = math.ceil((self.chunk_size + self.right_padding_size) / self.model_stride_in_secs)
+        self.tokens_per_frame_float = self.chunk_size / self.model_stride_in_secs
+        self.tokens_per_left_padding_float = self.left_padding_size / self.model_stride_in_secs
+        self.tokens_per_right_padding_float = self.right_padding_size / self.model_stride_in_secs
+        self.tokens_per_frame = math.ceil(self.tokens_per_frame_float)
+        self.tokens_per_left_padding = math.ceil(self.tokens_per_left_padding_float)
+        self.tokens_per_right_padding = math.ceil(self.tokens_per_right_padding_float)
+
+        if self.stateful:
+            self.initial_delay = self.right_padding_size / self.model_stride_in_secs
+        else:
+            self.initial_delay = (self.left_padding_size + self.right_padding_size) / self.model_stride_in_secs
+
+        if self.stateful and (
+            abs(self.tokens_per_frame_float - self.tokens_per_frame) > 1e-5
+            or abs(self.tokens_per_left_padding_float - self.tokens_per_left_padding) > 1e-5
+            or abs(self.tokens_per_right_padding_float - self.tokens_per_right_padding) > 1e-5
+        ):
+            self.tokens_per_frame_float = self.tokens_per_frame
+            self.tokens_per_left_padding_float = self.tokens_per_left_padding
+            self.left_padding_size = self.tokens_per_left_padding * self.model_stride_in_secs
+            self.chunk_size = self.tokens_per_frame * self.model_stride_in_secs
+            self.right_padding_size = self.tokens_per_right_padding * self.model_stride_in_secs
+            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.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
+        self.residue_tokens_at_end = cfg.endpointing.residue_tokens_at_end
+        self.word_boundary_tolerance = cfg.streaming.word_boundary_tolerance
+        self.return_tail_result = cfg.return_tail_result
+        self.tokens_to_move = self.punctuation_ids.union(self.language_token_ids)
+
+        self.zero_encoded = self.init_zero_enc() if self.right_padding else None
+
+    def init_endpointer(self) -> None:
+        """Initialize the endpointer."""
+        check_existance_of_required_attributes(
+            self,
+            [
+                'stateful',
+                'chunk_size',
+                'right_padding_size',
+                'buffer_size_in_secs',
+                'vocabulary',
+                'model_stride_in_milliseconds',
+                'stop_history_eou_in_milliseconds',
+                'residue_tokens_at_end',
+            ],
+        )
+
+        if self.stateful:
+            effective_buffer_size_in_secs = self.chunk_size + self.right_padding_size
+        else:
+            effective_buffer_size_in_secs = self.buffer_size_in_secs
+
+        self.endpointer = RNNTGreedyEndpointing(
+            vocabulary=self.vocabulary,
+            ms_per_timestep=self.model_stride_in_milliseconds,
+            effective_buffer_size_in_secs=effective_buffer_size_in_secs,
+            stop_history_eou=self.stop_history_eou_in_milliseconds,
+            residue_tokens_at_end=self.residue_tokens_at_end,
+        )
+
+    def init_greedy_rnnt_decoder(self) -> None:
+        """Initialize the greedy RNNT decoder."""
+        check_existance_of_required_attributes(self, ['vocabulary', 'conf_func', 'endpointer', 'tokens_per_frame'])
+        self.greedy_rnnt_decoder = ClippedRNNTGreedyDecoder(
+            vocabulary=self.vocabulary,
+            conf_func=self.conf_func,
+            endpointer=self.endpointer,
+            tokens_per_frame=self.tokens_per_frame,
+        )
+
+    def init_decoding_computer(self) -> None:
+        """Initialize the decoding computer."""
+        check_existance_of_required_attributes(self, ['stateful', 'asr_model'])
+        self.decoding_computer = None
+        if self.stateful:
+            self.decoding_computer = self.asr_model.asr_model.decoding.decoding.decoding_computer
+
+    def init_zero_enc(self) -> Tensor:
+        """
+        Initialize the encoder output for the zero buffer.
+        Returns:
+            (Tensor) Encoder output for the zero buffer.
+        """
+        check_existance_of_required_attributes(
+            self, ['buffer_size_in_secs', 'sample_rate', 'device', 'expected_feature_buffer_len']
+        )
+        buffer_size_in_samples = int(self.buffer_size_in_secs * self.sample_rate)
+        zero_buffer = torch.zeros(1, buffer_size_in_samples, device=self.device)
+        zero_features, zero_features_len = self.preprocess(
+            buffers=zero_buffer,
+            buffer_lens=torch.tensor([zero_buffer.shape[1]], device=self.device),
+            expected_feature_buffer_len=self.expected_feature_buffer_len,
+        )
+
+        if self.prompt_enabled:
+            # Use "en-US" as the default prompt for zero encoding
+            # This region is sliced out before decoding, so language choice doesn't matter
+            default_prompt_idx = self._resolve_prompt_index("en-US")
+            prompt_indices = torch.tensor([default_prompt_idx], device=self.device, dtype=torch.long)
+            prompt_vector = self._create_one_hot_prompts(prompt_indices)  # [1, num_prompts]
+
+            zero_encoded, _ = self.asr_model.encode_with_prompts(
+                processed_signal=zero_features,
+                processed_signal_length=zero_features_len,
+                prompt_vectors=prompt_vector,
+            )
+        else:
+            zero_encoded, _ = self.asr_model.encode(
+                processed_signal=zero_features, processed_signal_length=zero_features_len
+            )
+
+        return zero_encoded[0]
+
+    def create_state(self, options: ASRRequestOptions) -> RNNTStreamingState:
+        """
+        Create new empty state.
+        Args:
+            options: (ASRRequestOptions) Request options for particular stream.
+        Returns:
+            (RNNTStreamingState) New empty state.
+        """
+        state = RNNTStreamingState()
+        state.set_global_offset(-self.initial_delay)
+        new_options = options.fill_defaults(
+            default_enable_itn=self.text_processor.itn_enabled,
+            default_enable_nmt=self.nmt_enabled,
+            default_source_language=self.nmt_model.source_language if self.nmt_enabled else None,
+            default_target_language=self.nmt_model.target_language if self.nmt_enabled else None,
+            default_stop_history_eou=self.stop_history_eou_in_milliseconds,
+            default_asr_output_granularity=self.asr_output_granularity,
+            default_language_code="en-US" if self.prompt_enabled else None,
+        )
+        state.set_options(new_options)
+
+        # Create per-stream prompt index for prompt-enabled models
+        if self.prompt_enabled:
+            lang_code = getattr(new_options, "language_code", None)
+            if not isinstance(lang_code, str) or len(lang_code) == 0:
+                raise ValueError("Prompt-enabled model requires a valid language_code in request options.")
+            prompt_idx = self._resolve_prompt_index(lang_code)
+            state.set_prompt_index(prompt_idx)
+
+        return state
+
+    def get_sep(self) -> str:
+        """Return the separator for the text processor."""
+        return self.sep
+
+    def preprocess(
+        self, buffers: Tensor, buffer_lens: Tensor, expected_feature_buffer_len: int
+    ) -> tuple[Tensor, Tensor]:
+        """
+        Preprocess the buffered frames and extract features.
+        Args:
+            buffers: (Tensor) Audio buffers.
+            buffer_lens: (Tensor) Lengths of the audio buffers.
+            expected_feature_buffer_len: (int) Expected length of the feature buffers.
+        Returns:
+            (tuple[Tensor, Tensor]) Processed feature buffers and their lengths.
+        """
+        feature_buffers, feature_buffer_lens = self.preprocessor(input_signal=buffers, length=buffer_lens)
+        feature_buffers = drop_trailing_features(feature_buffers, expected_feature_buffer_len)
+        feature_buffers = normalize_features(feature_buffers, feature_buffer_lens)
+        feature_buffer_lens = feature_buffer_lens.clamp(max=feature_buffers.shape[2])
+        return feature_buffers, feature_buffer_lens
+
+    def get_cut_off_range(self, T: int, is_last: bool) -> tuple[int, int]:
+        """
+        Compute the start and end indices to clip.
+        Args:
+            T: (int) Time dimension of the alignment.
+            is_last: (bool) Whether the last frame is reached.
+        Returns:
+            (tuple[int, int]) Start and end indices to clip.
+        """
+        start = max(T - 1 - self.mid_delay, 0)
+        end = T if is_last else min(start + self.tokens_per_frame, T)
+        return start, end
+
+    def encode_raw_signals(
+        self, frames: list[Frame], raw_signals: list[Tensor], left_paddings: list[int]
+    ) -> tuple[Tensor, Tensor]:
+        """
+        Run Encoder part on the audio buffers.
+        Args:
+            frames: (list[Frame]) Frames to transcribe.
+            raw_signals: (list[Tensor]) Audio buffers.
+            left_paddings: (list[int]) Left paddings for audio buffers.
+        Returns:
+            (tuple[Tensor, Tensor]) Encoded signals and their lengths.
+        """
+
+        if self.right_padding:
+            left_paddings = torch.tensor(left_paddings, dtype=torch.int64, device=self.device)
+
+        buffers = []
+        for i in range(len(raw_signals)):
+            buffer = raw_signals[i]
+            if self.right_padding:
+                # Roll the buffered frames to the left by the left padding
+                # This is done to avoid the padding at the beginning of the buffered frames
+                # which can cause the performance degradation
+                lpad = left_paddings[i].item()
+                if lpad > 0:
+                    buffer = buffer.roll(shifts=-lpad)
+            buffers.append(buffer.unsqueeze_(0))
+
+        # Only final frames have right padding
+        # Calculate right paddings
+        right_paddings = torch.tensor([frame.size - frame.valid_size for frame in frames], device=self.device).clamp(
+            min=0
+        )
+
+        # Create and adjust the buffer lens
+        buffer_lens = torch.tensor([buffers[0].size(1)] * len(buffers), device=self.device)
+        buffer_lens = buffer_lens - right_paddings
+        if self.right_padding:
+            buffer_lens = buffer_lens - left_paddings
+
+        feature_buffers, feature_buffer_lens = self.preprocess(
+            buffers=torch.cat(buffers).to(self.device),
+            buffer_lens=buffer_lens,
+            expected_feature_buffer_len=self.expected_feature_buffer_len,
+        )
+
+        # Build prompt vectors if prompts are enabled
+        if self.prompt_enabled:
+            requests_states = [self.get_state(f.stream_id) for f in frames]
+            prompt_vectors = self._build_prompt_vectors(requests_states)
+
+            # Use encode_with_prompts which handles dimension expansion
+            encoded, encoded_len = self.asr_model.encode_with_prompts(
+                processed_signal=feature_buffers,
+                processed_signal_length=feature_buffer_lens,
+                prompt_vectors=prompt_vectors,
+            )
+        else:
+            encoded, encoded_len = self.asr_model.encode(
+                processed_signal=feature_buffers, processed_signal_length=feature_buffer_lens
+            )
+        encoded = encoded.clone()
+        encoded_len = encoded_len.clone()
+
+        # Roll back the encoded signals to the right
+        if self.right_padding:
+            for i in range(encoded.shape[0]):
+                lpad = left_paddings[i]
+                if lpad > 0:
+                    lpad = int(lpad / self.sample_rate / self.model_stride_in_secs)
+                    encoded[i] = encoded[i].roll(lpad, dims=1)
+                    encoded[i][:, :lpad] = self.zero_encoded[:, :lpad]
+                    encoded_len[i] = encoded_len[i] + lpad
+
+        return encoded, encoded_len
+
+    def encode_processed_signals(
+        self, fbuffers: list[FeatureBuffer], processed_signals: list[Tensor]
+    ) -> tuple[Tensor, Tensor]:
+        """
+        Run Encoder part on the feature buffers.
+        Args:
+            fbuffers: (list[FeatureBuffer]) Feature buffers.
+            processed_signals: (list[Tensor]) Processed buffers.
+        Returns:
+            (tuple[Tensor, Tensor]) Encoder output and their lengths.
+        """
+
+        processed_signals = torch.cat([sig.unsqueeze_(0) for sig in processed_signals]).to(self.device)
+        processed_signals = drop_trailing_features(processed_signals, self.expected_feature_buffer_len)
+        processed_signal_lengths = torch.tensor([f.valid_size for f in fbuffers], device=self.device)
+        processed_signals = normalize_features(processed_signals, processed_signal_lengths)
+        processed_signal_lengths = processed_signal_lengths.clamp(max=processed_signals.shape[2])
+
+        # Build prompt vectors if prompts are enabled
+        if self.prompt_enabled:
+            requests_states = [self.get_state(f.stream_id) for f in fbuffers]
+            prompt_vectors = self._build_prompt_vectors(requests_states)
+
+            # Use encode_with_prompts which handles dimension expansion
+            encoded, encoded_len = self.asr_model.encode_with_prompts(
+                processed_signal=processed_signals,
+                processed_signal_length=processed_signal_lengths,
+                prompt_vectors=prompt_vectors,
+            )
+        else:
+            encoded, encoded_len = self.asr_model.encode(
+                processed_signal=processed_signals, processed_signal_length=processed_signal_lengths
+            )
+        encoded = encoded.clone()
+        encoded_len = encoded_len.clone()
+
+        if self.right_padding:
+            for i in range(encoded.shape[0]):
+                lpad = int(fbuffers[i].roll_size / self.subsampling_factor)
+                if lpad > 0:
+                    encoded[i] = encoded[i].roll(lpad, dims=1)
+                    encoded[i][:, :lpad] = self.zero_encoded[:, :lpad]
+                    encoded_len[i] = encoded_len[i] + lpad
+        return encoded, encoded_len
+
+    def encode_frames(self, frames: list[Frame]) -> tuple[Tensor, Tensor]:
+        """
+        Encode the frames using the Encoder part of the ASR model.
+        Args:
+            frames: (list[Frame]) Frames to transcribe.
+        Returns:
+            (tuple[Tensor, Tensor]) Encoder output and their lengths.
+        """
+        raw_signals, left_paddings = self.bufferer.update(frames)
+        encs, enc_lens = None, None
+        if len(raw_signals) > 0:
+            encs, enc_lens = self.encode_raw_signals(frames, raw_signals, left_paddings)
+        return encs, enc_lens
+
+    def encode_feature_buffers(self, fbuffers: list[FeatureBuffer]) -> tuple[Tensor, Tensor]:
+        """
+        Encode the feature buffers using the Encoder part of the ASR model.
+        Args:
+            fbuffers: (list[FeatureBuffer]) Feature buffers to transcribe.
+        Returns:
+            (tuple[Tensor, Tensor]) Encoder output and their lengths.
+        """
+        processed_signals = self.bufferer.update(fbuffers)
+        encs, enc_lens = None, None
+        if len(processed_signals) > 0:
+            encs, enc_lens = self.encode_processed_signals(fbuffers, processed_signals)
+        return encs, enc_lens
+
+    def run_greedy_decoder(
+        self,
+        state: RNNTStreamingState,
+        request: Request,
+        timesteps: torch.Tensor,
+        tokens: torch.Tensor,
+        start: int,
+        end: int,
+        alignment_length: int,
+        timestamp_offset: int = 0,
+        vad_segments: torch.Tensor = None,
+        confidences: torch.Tensor | None = None,
+    ) -> bool:
+        """
+        Greedy RNN-T decoder.
+        Args:
+            state: (RNNTStreamingState) Current state for the particular stream.
+            request: (Request) Current request for the particular stream.
+            timesteps: (Tensor) Timesteps.
+            tokens: (Tensor) Tokens.
+            start: (int) Start index.
+            end: (int) End index.
+            alignment_length: (int) Length of the alignment.
+            timestamp_offset: (int) Timestamp offset.
+            vad_segments: (Tensor) VAD segments.
+            confidences: (Tensor | None) Per-token (non-blank) confidence scores aligned with `tokens`.
+        Returns:
+            (bool) Whether EOU is detected.
+        """
+        if self.stateful and vad_segments is not None:
+            vad_segments = adjust_vad_segments(vad_segments, self.left_padding_size)
+
+        clipped_output, tail_output, eou_detected, start_idx, end_idx = self.greedy_rnnt_decoder(
+            global_timesteps=timesteps,
+            tokens=tokens,
+            alignment_length=alignment_length,
+            clip_start=start,
+            clip_end=end,
+            is_last=request.is_last,
+            is_start=request.is_first,
+            return_tail_result=self.return_tail_result,
+            state_start_idx=state.decoder_start_idx,
+            state_end_idx=state.decoder_end_idx,
+            timestamp_offset=timestamp_offset,
+            vad_segments=vad_segments,
+            stop_history_eou=state.options.stop_history_eou,
+            confidences=confidences,
+        )
+        state.update_state(clipped_output, eou_detected)
+        state.update_from_decoder_results(start_idx, end_idx)
+        if self.stateless:
+            # For stateless mode, we need to set the last token, it will be used for filtering duplicate token
+            state.set_last_token(clipped_output["last_token"], clipped_output["last_token_idx"])
+            # For stateless mode, we need to increment the global offset
+            state.increment_global_offset(self.tokens_per_frame_float)
+        state.set_incomplete_segment_tokens(tail_output["tokens"])
+        return eou_detected
+
+    def stateless_transcribe_step(
+        self, requests: list[Request], encs: Tensor, enc_lens: Tensor, ready_state_ids: set
+    ) -> None:
+        """
+        Stateless transcribe step.
+        Stateless assumes that we don't keep track of partial hypotheses (partial_hypotheses=None).
+        Args:
+            requests: (list[Request]) List of requests to transcribe.
+            encs: (Tensor) Encoder output.
+            enc_lens: (Tensor) Encoder output lengths.
+            ready_state_ids: (set) Set of ready state IDs.
+        """
+        states = [self.get_state(request.stream_id) for request in requests]
+        best_hyp = self.asr_model.decode(encs, enc_lens, partial_hypotheses=None)
+        # For stateless mode, use zero timestamp offsets since we don't track timestamps
+        ready_states = self.decode_step(best_hyp, requests, states)
+        ready_state_ids.update(ready_states)
+
+    def stateful_transcribe_step(
+        self, requests: list[Request], encs: Tensor, enc_lens_chunk: Tensor, enc_lens: Tensor, ready_state_ids: set
+    ) -> None:
+        """
+        Stateful transcribe step.
+        Stateful assumes that we keep track of partial hypotheses.
+        Args:
+            requests: (list[Request]) List of requests to transcribe.
+            encs: (Tensor) Encoder output.
+            enc_lens_chunk: (Tensor) Encoder output lengths for the chunk.
+            enc_lens: (Tensor) Encoder output lengths.
+            ready_state_ids: (set) Set of ready state IDs.
+        """
+        states = [self.get_state(request.stream_id) for request in requests]
+        partial_hypotheses, rnnt_states = [], []
+        all_rnnt_states_are_none = True
+        all_multi_biasing_models_empty = True
+        multi_biasing_ids = np.full([len(states)], fill_value=-1)
+        for i, state in enumerate(states):
+            hyp_state = state.hyp_decoding_state
+            rnnt_states.append(hyp_state)
+            if hyp_state is not None:
+                all_rnnt_states_are_none = False
+            if state.has_biasing_request():
+                if state.options.biasing_cfg.multi_model_id is not None:
+                    all_multi_biasing_models_empty = False
+                    multi_biasing_ids[i] = state.options.biasing_cfg.multi_model_id
+                elif state.options.biasing_cfg.auto_manage_multi_model:
+                    state.options.biasing_cfg.add_to_multi_model(
+                        tokenizer=self.asr_model.tokenizer,
+                        biasing_multi_model=self.decoding_computer.biasing_multi_model,
+                    )
+                    multi_biasing_ids[i] = state.options.biasing_cfg.multi_model_id
+                    all_multi_biasing_models_empty = False
+                else:
+                    logging.warning("Biasing request is not empty, not auto managed and not compiled. Skipping")
+            if hyp_state is not None or state.has_biasing_request():
+                partial_hypotheses.append(
+                    NemoHypothesis(
+                        score=0.0,
+                        y_sequence=torch.zeros([0], dtype=torch.long),
+                        dec_state=hyp_state,
+                        biasing_cfg=state.options.biasing_cfg,
+                    )
+                )
+            else:
+                partial_hypotheses.append(None)
+
+        batched_rnnt_states = None
+        if not all_rnnt_states_are_none:
+            batched_rnnt_states = self.decoding_computer.merge_to_batched_state(rnnt_states)
+
+        if all_multi_biasing_models_empty:
+            multi_biasing_ids = None
+        else:
+            multi_biasing_ids = torch.from_numpy(multi_biasing_ids).to(device=enc_lens_chunk.device)
+
+        encs_dim_last = encs.transpose(1, 2)
+        # decode chunk
+        with torch.inference_mode(), torch.no_grad():
+            best_batched_hyps_chunk, batched_state = self.decoding_computer(
+                encs_dim_last,
+                enc_lens_chunk,
+                batched_rnnt_states,
+                multi_biasing_ids=multi_biasing_ids,
+            )
+        best_hyps = batched_hyps_to_hypotheses(best_batched_hyps_chunk, batch_size=enc_lens.shape[0])
+
+        # save state (after chunk)
+        for state, rnnt_state in zip(states, self.decoding_computer.split_batched_state(batched_state)):
+            state.hyp_decoding_state = rnnt_state
+
+        if self.tokens_per_right_padding > 0:
+            # decode right context
+            _, max_time, feat_dim = encs_dim_last.shape
+            device = encs.device
+            # we are indexing `encs_dim_last` with `shift_indices` to get a tensor where right context is at the start
+            # everything after right context is padded with `0` index (first encoder vector)
+            # padding will be ignored by decoder_computer since we pass the lengths
+            shift_indices = torch.arange(max_time, device=device, dtype=torch.long)[None, :] + enc_lens_chunk[:, None]
+            # pad with zeros everything beyond needed context
+            shift_indices = torch.where(shift_indices < max_time, shift_indices, torch.zeros_like(shift_indices))
+            with torch.inference_mode(), torch.no_grad():
+                best_batched_hyps_rc, _ = self.decoding_computer(
+                    torch.gather(encs_dim_last, dim=1, index=shift_indices[:, :, None].expand(-1, -1, feat_dim)),
+                    enc_lens - enc_lens_chunk,
+                    batched_state,
+                    multi_biasing_ids=multi_biasing_ids,
+                )
+                best_hyps_rc = batched_hyps_to_hypotheses(best_batched_hyps_rc, batch_size=enc_lens.shape[0])
+            # merge right context to chunk hypothesis
+            for hyp, hyp_rc in zip(best_hyps, best_hyps_rc):
+                hyp.merge_(hyp_rc)
+
+        ready_states = self.decode_step(best_hyps, requests, states)
+        for curr_state in states:
+            curr_state.timestamp_offset += self.tokens_per_frame_float
+        ready_state_ids.update(ready_states)
+
+        for request, state in zip(requests, states):
+            # only the first request contains biasing options; biasing options for the stream are stored in state
+            if request.is_last and state.has_biasing_request():
+                if state.options.biasing_cfg.auto_manage_multi_model:
+                    state.options.biasing_cfg.remove_from_multi_model(
+                        biasing_multi_model=self.decoding_computer.biasing_multi_model
+                    )
+
+    def decode_step(self, best_hyp: list, requests: list[Request], states: list[RNNTStreamingState]) -> set:
+        """
+        Perform greedy RNNT decoding to get the best hypothesis and update the state.
+        If EOU is detected, push the words to the state and cleanup the state.
+        Args:
+            best_hyp: (list) Best hypothesis.
+            requests: (list[Request]) List of requests to transcribe.
+            states: (list[RNNTStreamingState]) List of states.
+        Returns:
+            (set) Set of ready state IDs.
+        """
+        ready_state_ids = set()
+        for idx, hyp in enumerate(best_hyp):
+            state = states[idx]
+            request = requests[idx]
+            # Perform timestamp based decoding for the hypothesis
+            if self.stateful:
+                alignment_length = self.tokens_per_right_padding + self.tokens_per_frame
+            else:
+                if self.request_type is RequestType.FEATURE_BUFFER:
+                    alignment_length = math.ceil(request.size / self.subsampling_factor)
+                else:  # RequestType.FRAME
+                    alignment_length = math.ceil(self.expected_feature_buffer_len / self.subsampling_factor)
+
+            if self.stateful:
+                start, end = 0, self.tokens_per_frame
+            else:
+                # For stateless mode
+                if request.is_first and request.is_last:
+                    start, end = 0, alignment_length
+                else:
+                    start, end = self.get_cut_off_range(alignment_length, request.is_last)
+
+            timestamp = hyp.timestamp
+            tokens = hyp.y_sequence
+            timestamp = torch.tensor(timestamp) if isinstance(timestamp, list) else timestamp
+            tokens = torch.tensor(tokens) if isinstance(tokens, list) else tokens
+            timestamp = update_punctuation_and_language_tokens_timestamps(
+                tokens, timestamp, self.tokens_to_move, self.underscore_id
+            )
+            # Per-token non-blank confidence precomputed during RNN-T decoding (aligned with `tokens`).
+            # Populated when greedy or batched-beam preserve_frame_confidence is enabled; otherwise None.
+            confidences = hyp.non_blank_step_confidence_precomputed
+            if confidences is not None:
+                confidences = torch.tensor(confidences, dtype=torch.float32, device=tokens.device)
+            vad_segments = request.vad_segments
+            eou_detected = self.run_greedy_decoder(
+                state=state,
+                request=request,
+                timesteps=timestamp,
+                tokens=tokens,
+                start=start,
+                end=end,
+                alignment_length=alignment_length,
+                timestamp_offset=state.timestamp_offset,
+                vad_segments=vad_segments,
+                confidences=confidences,
+            )
+
+            if eou_detected:
+                self.bpe_decoder.decode_bpe_tokens(state)
+                state.cleanup_after_eou()
+                ready_state_ids.add(request.stream_id)
+        return ready_state_ids
+
+    def shared_transcribe_step_stateful(self, requests: list[Request], encs: Tensor, enc_lens: Tensor) -> None:
+        """
+        Stateful transcribe step.
+        After detecting EOU, it updates the state and run text processor.
+        If there are multiple streams, it waits until all states are ready to run text processor.
+        Args:
+            requests: (list[Request]) List of requests to transcribe.
+            encs: (Tensor) Encoder output.
+            enc_lens: (Tensor) Encoder output lengths.
+        """
+        tokens_per_left_padding_tensor = torch.tensor(self.tokens_per_left_padding, device=self.device)
+        tokens_per_frame_tensor = torch.tensor(self.tokens_per_frame, device=self.device)
+        postponed_requests = [(ridx, request.stream_id) for ridx, request in enumerate(requests)]
+        next_postponed_requests = []
+        ready_state_ids = set()
+        while len(postponed_requests) > 0:
+            request_ids_to_process = []
+            for ridx, stream_id in postponed_requests:
+                if stream_id in ready_state_ids:
+                    next_postponed_requests.append((ridx, stream_id))
+                    continue
+                request_ids_to_process.append(ridx)
+            if len(request_ids_to_process) > 0:
+                requests_to_process = [requests[jdx] for jdx in request_ids_to_process]
+                request_is_last = torch.tensor(
+                    [request.is_last for request in requests_to_process], dtype=torch.bool, device=self.device
+                )
+                enc_lens_dec = enc_lens - tokens_per_left_padding_tensor
+                enc_lens_dec_trimmed = torch.where(
+                    request_is_last,
+                    enc_lens_dec,
+                    torch.minimum(enc_lens_dec, tokens_per_frame_tensor.expand_as(enc_lens_dec)),
+                )
+                self.stateful_transcribe_step(
+                    requests_to_process,
+                    encs[request_ids_to_process][:, :, self.tokens_per_left_padding :],
+                    enc_lens_dec_trimmed,
+                    enc_lens_dec,
+                    ready_state_ids,
+                )
+            if len(ready_state_ids) > 0:
+                self.text_processor.process([self.get_state(stream_id) for stream_id in ready_state_ids])
+                ready_state_ids.clear()
+            postponed_requests = next_postponed_requests.copy()
+            next_postponed_requests.clear()
+
+        self.update_partial_transcript(requests, self.tokenizer, self.leading_regex_pattern)
+
+    def shared_transcribe_step(self, requests: list[Request], encs: Tensor, enc_lens: Tensor) -> None:
+        """
+        Stateless transcribe step.
+        After detecting EOU, it updates the state and run text processor.
+        If there are multiple streams, it waits until all stated are ready to run text processor.
+        Args:
+            requests: (list[Request]) List of requests to transcribe.
+            encs: (Tensor) Encoder output.
+            enc_lens: (Tensor) Encoder output lengths.
+        """
+        postponed_requests = [(ridx, request.stream_id) for ridx, request in enumerate(requests)]
+        next_postponed_requests = []
+        ready_state_ids = set()
+
+        while len(postponed_requests) > 0:
+
+            request_ids_to_process = []
+            for ridx, stream_id in postponed_requests:
+
+                if stream_id in ready_state_ids:
+                    # Skip if the state is already ready
+                    next_postponed_requests.append((ridx, stream_id))
+                    continue
+
+                request_ids_to_process.append(ridx)
+
+            if len(request_ids_to_process) > 0:
+                requests_to_process = [requests[jdx] for jdx in request_ids_to_process]
+                self.stateless_transcribe_step(
+                    requests_to_process,
+                    encs=encs[request_ids_to_process],
+                    enc_lens=enc_lens[request_ids_to_process],
+                    ready_state_ids=ready_state_ids,
+                )
+
+            if len(ready_state_ids) > 0:
+                self.text_processor.process([self.get_state(stream_id) for stream_id in ready_state_ids])
+                ready_state_ids.clear()
+
+            postponed_requests = next_postponed_requests.copy()
+            next_postponed_requests.clear()
+
+        self.update_partial_transcript(requests, self.tokenizer, self.leading_regex_pattern)
+
+    def transcribe_step_for_feature_buffers(self, fbuffers: list[FeatureBuffer]) -> None:
+        """
+        Transcribe a step for feature buffers.
+        Args:
+            fbuffers: (list[FeatureBuffer]) List of feature buffers to transcribe.
+        """
+        encs, enc_lens = self.encode_feature_buffers(fbuffers)
+        if encs is not None:
+            if self.stateful:
+                self.shared_transcribe_step_stateful(requests=fbuffers, encs=encs, enc_lens=enc_lens)
+            else:
+                self.shared_transcribe_step(requests=fbuffers, encs=encs, enc_lens=enc_lens)
+
+    def transcribe_step_for_frames(self, frames: list[Frame]) -> None:
+        """
+        Transcribe a step for frames.
+        Args:
+            frames: (list[Frame]) List of frames to transcribe.
+        """
+        encs, enc_lens = self.encode_frames(frames)
+        if encs is not None:
+            if self.stateful:
+                self.shared_transcribe_step_stateful(requests=frames, encs=encs, enc_lens=enc_lens)
+            else:
+                self.shared_transcribe_step(requests=frames, encs=encs, enc_lens=enc_lens)
+
+    def get_request_generator(self) -> ContinuousBatchedRequestStreamer:
+        """
+        Initialize the request generator.
+        Returns:
+            (ContinuousBatchedRequestStreamer) Request generator.
+        """
+        request_generator = ContinuousBatchedRequestStreamer(
+            n_frames_per_stream=1,
+            frame_size_in_secs=self.chunk_size,
+            sample_rate=self.sample_rate,
+            batch_size=self.batch_size,
+            request_type=self.request_type,
+            preprocessor=self.preprocessor,
+            buffer_size_in_secs=self.buffer_size_in_secs,
+            device=self.device,
+            pad_last_frame=True,
+            right_pad_features=self.right_padding,
+        )
+        return request_generator

nemo/collections/asr/inference/pipelines/buffered_salm_pipeline.py

@@ -0,0 +1,248 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+from omegaconf import DictConfig
+
+from nemo.collections.asr.inference.model_wrappers.salm_asr_inference_wrapper import SALMASRInferenceWrapper
+from nemo.collections.asr.inference.pipelines.base_pipeline import BasePipeline
+from nemo.collections.asr.inference.streaming.buffering.incremental_audio_bufferer import (
+    BatchedIncrementalAudioBufferer,
+)
+from nemo.collections.asr.inference.streaming.framing.multi_stream import ContinuousBatchedRequestStreamer
+from nemo.collections.asr.inference.streaming.framing.request import FeatureBuffer, Frame
+from nemo.collections.asr.inference.streaming.framing.request_options import ASRRequestOptions
+from nemo.collections.asr.inference.streaming.state.salm_state import SALMStreamingState
+from nemo.collections.asr.inference.utils.enums import ASROutputGranularity, MergingStrategy, RequestType
+from nemo.collections.asr.inference.utils.lcs_merge import lcs_merge
+from nemo.utils.decorators import experimental
+
+if TYPE_CHECKING:
+    from nemo.collections.asr.inference.itn.inverse_normalizer import AlignmentPreservingInverseNormalizer
+    from nemo.collections.asr.inference.nmt.llm_translator import LLMTranslator
+
+
+def parse_hyp(answer: torch.Tensor, eos_tokens: list[int]):
+    """
+    Parse the hypothesis. Extract the tokens before the EOS tokens.
+    Args:
+        answer: (torch.Tensor) Answer tensor.
+        eos_tokens: (list[int]) EOS tokens.
+    Returns:
+        (torch.Tensor) Parsed hypothesis.
+    """
+    end = torch.isin(answer, torch.tensor(eos_tokens)).nonzero(as_tuple=True)[0]
+    if end.numel() == 0:
+        return answer
+    end = end[0]
+    return answer[:end]
+
+
+@experimental
+class BufferedSALMPipeline(BasePipeline):
+    """Buffered SALM pipeline."""
+
+    def __init__(
+        self,
+        cfg: DictConfig,
+        asr_model: SALMASRInferenceWrapper,
+        itn_model: AlignmentPreservingInverseNormalizer | None = None,
+        nmt_model: LLMTranslator | None = None,
+    ):
+        """
+        Initialize the BufferedSALMPipeline.
+        Args:
+            cfg: (DictConfig) Configuration parameters.
+            asr_model: (SALMASRInferenceWrapper) ASR model.
+            itn_model: (AlignmentPreservingInverseNormalizer | None) Inverse Text Normalization model.
+            nmt_model: (LLMTranslator | None) LLM based translation model.
+        """
+        self.asr_model = asr_model
+        self.init_parameters(cfg)
+        self.init_nmt_model(nmt_model)
+        super().__init__()
+
+    def init_parameters(self, cfg: DictConfig) -> None:
+        """
+        Initialize the parameters.
+        Args:
+            cfg: (DictConfig) Configuration parameters.
+        """
+        self.sample_rate = cfg.streaming.sample_rate
+        self.asr_output_granularity = ASROutputGranularity.from_str(cfg.asr_output_granularity)
+        if self.asr_output_granularity is ASROutputGranularity.WORD:
+            raise ValueError("Word level output granularity is not supported for SALM AED pipeline")
+
+        self.batch_size = cfg.streaming.batch_size
+        self.max_new_tokens = cfg.streaming.max_new_tokens
+        self.device = self.asr_model.device
+        self.stop_history_eou_in_milliseconds = cfg.endpointing.stop_history_eou
+
+        self.chunk_size_in_secs = cfg.streaming.chunk_size
+        self.buffer_size_in_secs = cfg.streaming.buffer_size
+        self.overlap_size_in_secs = cfg.streaming.overlap_size
+        self.overlap_ratio = self.overlap_size_in_secs / self.buffer_size_in_secs
+        self.extra_overlap_tokens = 2  # extra tokens for better overlap detection
+        self.merging_strategy = MergingStrategy.from_str(cfg.streaming.merging_strategy)
+
+        self.audio_bufferer = BatchedIncrementalAudioBufferer(
+            self.sample_rate,
+            self.buffer_size_in_secs,
+            self.chunk_size_in_secs,
+            self.overlap_size_in_secs,
+        )
+
+        self.request_type = RequestType.from_str(cfg.streaming.request_type)
+        if self.request_type is RequestType.FEATURE_BUFFER:
+            raise ValueError("Feature buffer request type is not supported for SALM pipeline")
+
+        self.prompts = [[{"role": "user", "content": f"Transcribe the following: {self.asr_model.audio_locator_tag}"}]]
+        self.tokens = self.asr_model.preprocess_prompts(self.prompts)
+
+    def create_state(self, options: ASRRequestOptions) -> SALMStreamingState:
+        """
+        Create new empty state.
+        Args:
+            options: (ASRRequestOptions) Request options for particular stream.
+        Returns:
+            (SALMStreamingState) New empty state.
+        """
+        state = SALMStreamingState()
+        state.set_global_offset(0)
+        new_options = options.fill_defaults(
+            default_enable_itn=False,
+            default_enable_nmt=False,
+            default_source_language=None,
+            default_target_language=None,
+            default_stop_history_eou=self.stop_history_eou_in_milliseconds,
+            default_asr_output_granularity=self.asr_output_granularity,
+            default_language_code=None,
+        )
+        state.set_options(new_options)
+        return state
+
+    def get_sep(self) -> str:
+        """Return the separator for the text processor."""
+        return self.asr_model.word_separator
+
+    def lcs_merge(self, state: SALMStreamingState, data: list[int]) -> None:
+        """
+        Merge the buffer and data using the LCS algorithm.
+        Args:
+            state: (SALMStreamingState) The state of the streaming pipeline.
+            data: (list[int]) The new tokens to merge with the buffer.
+        """
+        if len(state.tokens) == 0:
+            state.tokens.extend(data)
+            return
+
+        # extra overlap tokens for better overlap detection
+        delay = int(self.overlap_ratio * len(data)) + self.extra_overlap_tokens
+        state.tokens = lcs_merge(
+            buffer=state.tokens,
+            data=data[:delay],
+            search_size=delay,
+            sep_id=self.asr_model.word_separator_ids,
+            min_lcs_length=1,
+            merging_strategy=self.merging_strategy,
+        )
+        state.tokens.extend(data[delay:])
+
+    def transcribe_step_for_frames(self, frames: list[Frame]) -> None:
+        """
+        Perform the transcribe step for frames.
+        Args:
+            frames: (list[Frame]) List of frames to transcribe.
+        """
+        buffers, paddings = self.audio_bufferer.update(frames)
+        paddings = torch.tensor(paddings, dtype=torch.int64, device=self.device)
+
+        # Right paddings for the final frames
+        # Only for last frames frame.size is greater than frame.valid_size
+        right_paddings = torch.tensor(
+            [int(frame.size - frame.valid_size) for frame in frames], dtype=torch.int64, device=self.device
+        ).clamp(min=0)
+
+        # stack buffers
+        audios = torch.cat([buffer.unsqueeze_(0) for buffer in buffers]).to(self.device)
+        audio_lens = torch.tensor([audios.size(1)] * audios.size(0), dtype=torch.int64, device=self.device)
+        audio_lens = audio_lens - paddings - right_paddings
+
+        answer_ids = self.asr_model.generate(
+            prompts=self.tokens.expand(len(audios), -1),
+            audios=audios,
+            audio_lens=audio_lens,
+            max_new_tokens=self.max_new_tokens,
+        ).cpu()
+
+        for i, frame in enumerate(frames):
+            state = self.get_state(frame.stream_id)
+            new_tokens = parse_hyp(answer_ids[i], self.asr_model.eos_token_ids).tolist()
+            state.incomplete_segment_tokens.clear()
+            if self.audio_bufferer.is_full(frame.stream_id) or frame.is_last:
+                self.lcs_merge(state, new_tokens)
+            else:
+                state.incomplete_segment_tokens.extend(new_tokens)
+
+            if frame.is_last:
+                state.final_transcript = self.asr_model.tokenizer.ids_to_text(state.tokens)
+                state.partial_transcript = ""
+            else:
+                all_tokens = state.tokens.copy()
+                if len(state.incomplete_segment_tokens) > 0:
+                    # extra overlap tokens for better overlap detection
+                    delay = int(self.overlap_ratio * len(state.incomplete_segment_tokens)) + self.extra_overlap_tokens
+                    all_tokens = lcs_merge(
+                        buffer=all_tokens,
+                        data=state.incomplete_segment_tokens[:delay],
+                        search_size=delay,
+                        sep_id=self.asr_model.word_separator_ids,
+                        min_lcs_length=1,
+                        merging_strategy=self.merging_strategy,
+                    )
+                    all_tokens.extend(state.incomplete_segment_tokens[delay:])
+
+                if len(all_tokens) > 0:
+                    state.partial_transcript = self.asr_model.tokenizer.ids_to_text(all_tokens)
+                else:
+                    state.partial_transcript = ""
+
+    def transcribe_step_for_feature_buffers(self, fbuffers: list[FeatureBuffer]) -> None:
+        """
+        Transcribe a step for feature buffers.
+        Args:
+            fbuffers: (list[FeatureBuffer]) List of feature buffers to transcribe.
+        """
+        raise NotImplementedError("Feature buffer request type is not supported for SALM pipeline")
+
+    def get_request_generator(self) -> ContinuousBatchedRequestStreamer:
+        """
+        Initialize the request generator.
+        Returns:
+            (ContinuousBatchedRequestStreamer) Request generator.
+        """
+        request_generator = ContinuousBatchedRequestStreamer(
+            n_frames_per_stream=1,
+            frame_size_in_secs=self.chunk_size_in_secs,
+            sample_rate=self.sample_rate,
+            batch_size=self.batch_size,
+            request_type=self.request_type,
+            pad_last_frame=True,
+        )
+        return request_generator

nemo/collections/asr/inference/pipelines/cache_aware_ctc_pipeline.py

@@ -0,0 +1,431 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import math
+from typing import TYPE_CHECKING
+
+import numpy as np
+import torch
+from omegaconf import DictConfig
+from torch import Tensor
+
+from nemo.collections.asr.inference.model_wrappers.cache_aware_ctc_inference_wrapper import (
+    CacheAwareCTCInferenceWrapper,
+)
+from nemo.collections.asr.inference.pipelines.base_pipeline import BasePipeline
+from nemo.collections.asr.inference.streaming.decoders.greedy.greedy_ctc_decoder import CTCGreedyDecoder
+from nemo.collections.asr.inference.streaming.endpointing.greedy.greedy_ctc_endpointing import CTCGreedyEndpointing
+from nemo.collections.asr.inference.streaming.framing.multi_stream import ContinuousBatchedRequestStreamer
+from nemo.collections.asr.inference.streaming.framing.request import FeatureBuffer, Frame, Request
+from nemo.collections.asr.inference.streaming.framing.request_options import ASRRequestOptions
+from nemo.collections.asr.inference.streaming.state.cache_aware_ctc_state import CacheAwareCTCStreamingState
+from nemo.collections.asr.inference.utils.endpointing_utils import millisecond_to_frames
+from nemo.collections.asr.inference.utils.enums import RequestType
+from nemo.collections.asr.inference.utils.pipeline_utils import (
+    check_existance_of_required_attributes,
+    drop_trailing_features,
+    get_confidence_utils,
+    normalize_log_probs,
+)
+
+if TYPE_CHECKING:
+    from nemo.collections.asr.inference.itn.inverse_normalizer import AlignmentPreservingInverseNormalizer
+    from nemo.collections.asr.inference.nmt.llm_translator import LLMTranslator
+
+
+class CacheAwareCTCPipeline(BasePipeline):
+    """Cache Aware CTC pipeline."""
+
+    def __init__(
+        self,
+        cfg: DictConfig,
+        asr_model: CacheAwareCTCInferenceWrapper,
+        itn_model: AlignmentPreservingInverseNormalizer | None = None,
+        nmt_model: LLMTranslator | None = None,
+    ):
+        """
+        Initialize the CacheAwareCTCPipeline.
+        Args:
+            cfg: (DictConfig) Configuration parameters.
+            asr_model: (CacheAwareCTCInferenceWrapper) ASR model.
+            itn_model: (AlignmentPreservingInverseNormalizer | None) Inverse Text Normalization model.
+        """
+        self.copy_asr_model_attributes(asr_model)
+        self.init_parameters(cfg)
+        self.init_context_manager()
+        self.init_bufferer_for_cache_aware_streaming()
+        self.conf_func, self.confidence_aggregator = get_confidence_utils(cfg.confidence)
+        self.init_bpe_decoder()
+        self.init_greedy_ctc_decoder()
+        self.init_endpointer()
+        self.init_text_processor(cfg, itn_model)
+        self.init_nmt_model(nmt_model)
+        super().__init__()
+
+    def init_parameters(self, cfg: DictConfig) -> None:
+        """
+        Initialize the configuration parameters.
+        Args:
+            cfg: (DictConfig) Configuration parameters.
+        """
+        if cfg.streaming.att_context_size is not None:
+            self.asr_model.set_default_att_context_size(att_context_size=cfg.streaming.att_context_size)
+        self.sample_rate = cfg.streaming.sample_rate
+        self.asr_output_granularity = cfg.asr_output_granularity
+
+        self.use_cache = cfg.streaming.use_cache
+        self.use_feat_cache = cfg.streaming.use_feat_cache
+        self.batch_size = cfg.streaming.batch_size
+        self.num_slots = cfg.streaming.num_slots
+        if self.num_slots < self.batch_size:
+            raise ValueError(
+                f"Number of slots in the context manager must be >= batch_size: {self.num_slots} < {self.batch_size}"
+            )
+        self.request_type = RequestType.from_str(cfg.streaming.request_type)
+        self.word_boundary_tolerance = cfg.streaming.word_boundary_tolerance
+        self.stop_history_eou_in_milliseconds = cfg.endpointing.stop_history_eou
+        self.residue_tokens_at_end = cfg.endpointing.residue_tokens_at_end
+        self.return_tail_result = cfg.return_tail_result
+
+        self.pre_encode_cache_size = self.asr_model.get_pre_encode_cache_size()
+        self.model_chunk_size = self.asr_model.get_chunk_size()
+        if isinstance(self.model_chunk_size, list):
+            self.model_chunk_size = self.model_chunk_size[1]
+
+        if cfg.streaming.get("chunk_size_in_secs", None) is not None:
+            self.chunk_size_in_secs = cfg.streaming.chunk_size_in_secs
+            self.tokens_per_frame = math.ceil(
+                np.trunc(self.chunk_size_in_secs / self.window_stride) / self.subsampling_factor
+            )
+            # overwrite the encoder streaming params with proper shift size for cache aware streaming
+            self.asr_model.setup_streaming_params(
+                chunk_size=self.model_chunk_size // self.subsampling_factor, shift_size=self.tokens_per_frame
+            )
+        else:
+            self.chunk_size_in_secs = self.model_chunk_size * self.window_stride
+            self.tokens_per_frame = math.ceil(self.model_chunk_size / self.subsampling_factor)
+
+        if isinstance(self.pre_encode_cache_size, list):
+            self.pre_encode_cache_size = self.pre_encode_cache_size[1]
+        self.pre_encode_cache_size_in_secs = self.pre_encode_cache_size * self.window_stride
+
+        model_chunk_size_in_secs = self.model_chunk_size * self.window_stride
+
+        if self.use_cache:
+            # if using cache, we need to pad some samples for pre_encode
+            self.buffer_size_in_secs = self.pre_encode_cache_size_in_secs + model_chunk_size_in_secs
+            self.drop_left_context = None
+            self.valid_out_len = None
+        else:
+            # if not using cache, we need to keep left context in buffer, but no extra padding in pre_encode
+            left_context_size = self.asr_model.get_att_context_size()[0]
+            if left_context_size < 0:
+                raise ValueError(f"Left context size should not be a negative value: {left_context_size}")
+            self.buffer_size_in_secs = (
+                model_chunk_size_in_secs + left_context_size * self.subsampling_factor * self.window_stride
+            )
+            self.drop_left_context = left_context_size
+            self.valid_out_len = self.tokens_per_frame
+
+        # Expected feature buffer length for trimming (safeguard for feature buffer inputs)
+        self.expected_feature_buffer_len = int(self.buffer_size_in_secs / self.window_stride)
+
+    def init_greedy_ctc_decoder(self) -> None:
+        """Initialize the CTC decoder."""
+        check_existance_of_required_attributes(self, ['vocabulary', 'conf_func'])
+        self.greedy_ctc_decoder = CTCGreedyDecoder(vocabulary=self.vocabulary, conf_func=self.conf_func)
+
+    def init_endpointer(self) -> None:
+        """Initialize the endpointer."""
+        check_existance_of_required_attributes(
+            self,
+            [
+                'vocabulary',
+                'model_stride_in_milliseconds',
+                'stop_history_eou_in_milliseconds',
+                'residue_tokens_at_end',
+            ],
+        )
+
+        self.endpointer = CTCGreedyEndpointing(
+            vocabulary=self.vocabulary,
+            ms_per_timestep=self.model_stride_in_milliseconds,
+            stop_history_eou=self.stop_history_eou_in_milliseconds,
+            residue_tokens_at_end=self.residue_tokens_at_end,
+        )
+
+    def create_state(self, options: ASRRequestOptions) -> CacheAwareCTCStreamingState:
+        """
+        Create new empty state.
+        Args:
+            options: (ASRRequestOptions) Request options for particular stream.
+        Returns:
+            (CacheAwareCTCStreamingState) New empty state.
+        """
+        state = CacheAwareCTCStreamingState()
+        state.set_global_offset(0)
+        new_options = options.fill_defaults(
+            default_enable_itn=self.text_processor.itn_enabled,
+            default_enable_nmt=self.nmt_enabled,
+            default_source_language=self.nmt_model.source_language if self.nmt_enabled else None,
+            default_target_language=self.nmt_model.target_language if self.nmt_enabled else None,
+            default_stop_history_eou=self.stop_history_eou_in_milliseconds,
+            default_asr_output_granularity=self.asr_output_granularity,
+        )
+
+        eou_label_buffer_size = 0
+        if new_options.stop_history_eou > 0:
+            eou_label_buffer_size = millisecond_to_frames(
+                new_options.stop_history_eou, math.ceil(self.model_stride_in_milliseconds)
+            )
+            eou_label_buffer_size += self.residue_tokens_at_end
+        state.setup_label_buffer(eou_label_buffer_size, self.blank_id)
+        state.set_options(new_options)
+        return state
+
+    def get_sep(self) -> str:
+        """Return the separator for the text processor."""
+        return self.sep
+
+    def preprocess(self, buffers: list[Tensor], right_paddings: list[int] | None = None) -> tuple[Tensor, Tensor]:
+        """
+        Preprocess the feature buffers by stacking them and computing the lengths
+        Args:
+            buffers: (list[Tensor]) List of feature buffers.
+            right_paddings: (list[int] | None) List of right paddings.
+        Returns:
+            (tuple[Tensor, Tensor]) Processed feature buffers and their lengths.
+        """
+        feature_buffers = [f_buffer.unsqueeze_(0) for f_buffer in buffers]
+        # Trim to expected feature buffer length (safeguard for external feature buffer inputs)
+        feature_buffers = [
+            drop_trailing_features(f_buffer, self.expected_feature_buffer_len) for f_buffer in feature_buffers
+        ]
+        feature_buffer_lens = torch.tensor([f_buffer.shape[2] for f_buffer in feature_buffers], device=self.device)
+        if right_paddings is not None:
+            right_paddings = torch.tensor(right_paddings, device=feature_buffer_lens.device)
+            feature_buffer_lens = feature_buffer_lens - right_paddings
+        feature_buffers = torch.cat(feature_buffers).to(self.device)
+        return feature_buffers, feature_buffer_lens
+
+    def run_greedy_decoder(self, state: CacheAwareCTCStreamingState, request: Request, log_probs: Tensor):
+        """
+        Run the greedy CTC decoder on the log_probs and update the state
+        Args:
+            state: (CacheAwareCTCStreamingState) The state of the stream
+            request: (Request) The current request (frame or feature buffer)
+            log_probs: (Tensor) The log probabilities of the current request
+        Returns:
+            (bool) Whether EOU is detected.
+        """
+        eou_detected = request.is_last
+        last_token = state.label_buffer[-1] if len(state.label_buffer) > 0 else self.blank_id
+        cur_output = self.greedy_ctc_decoder(log_probs, compute_confidence=True, previous=last_token)
+        state.update_label_buffer(cur_output["labels"])
+
+        if not eou_detected:
+            emissions = state.get_label_buffer()
+            pivot_point = len(emissions) - 1
+            eou_detected, _ = self.endpointer.detect_eou_near_pivot(
+                emissions, pivot_point, stop_history_eou=state.options.stop_history_eou
+            )
+
+        state.update_state(cur_output, eou_detected=eou_detected)
+        state.increment_global_offset(self.tokens_per_frame)
+        return eou_detected
+
+    def decode_log_probs(
+        self,
+        requests: list[Request],
+        log_probs: Tensor,
+        tail_log_probs: Tensor | None,
+        ready_state_ids: set,
+    ) -> None:
+        """
+        Decode the log probabilities and update the state
+        Args:
+            requests: (list[Request]) List of requests (frames or feature buffers) to transcribe.
+            log_probs: (Tensor) Log probabilities.
+            tail_log_probs: (Tensor | None) Tail log probabilities.
+            ready_state_ids: (set) Set of ready state IDs.
+        """
+
+        for idx, request in enumerate(requests):
+            state = self.get_state(request.stream_id)
+            eou_detected = self.run_greedy_decoder(state, request, log_probs[idx])
+
+            if eou_detected:
+                self.bpe_decoder.decode_bpe_tokens(state)
+                state.cleanup_after_eou()
+                ready_state_ids.add(request.stream_id)
+
+            if tail_log_probs is not None:
+                last_token = state.label_buffer[-1] if len(state.label_buffer) > 0 else self.blank_id
+                tail_output = self.greedy_ctc_decoder(
+                    tail_log_probs[idx], compute_confidence=False, previous=last_token
+                )
+                state.set_incomplete_segment_tokens(tail_output["tokens"])
+
+    def cache_aware_transcribe_step(
+        self,
+        requests: list[Request],
+        buffered_features: list[Tensor],
+        right_paddings: list[int] | None,
+        ready_state_ids: set,
+        keep_all_outputs: bool = False,
+    ) -> None:
+        """
+        Cache Aware Transcribe Step
+        It receives a list of requests (Frame or FeatureBuffer) and features and do the following:
+
+        1. Preprocess the features by stacking them and computing the lengths
+        2. Get the context and mapping from the context manager for cache aware streaming
+        3. Perform a streaming step with the ASR model
+        4. Update the cache and reset the cache slots for the streams that has ended
+        5. Decode the log probabilities and update the state
+
+        Args:
+            requests: (list[Request]) List of requests (frames or feature buffers) to transcribe.
+            buffered_features: (list[Tensor]) List of buffered features.
+            right_paddings: (list[int] | None) List of right paddings.
+            ready_state_ids: (set) Set of ready state IDs.
+            keep_all_outputs: (bool) Whether to keep all outputs or not.
+        """
+        feature_buffers, feature_buffer_lens = self.preprocess(buffered_features, right_paddings)
+
+        stream_ids = [request.stream_id for request in requests]
+        eos_flags = [request.is_last for request in requests]
+        context, mapping = self.context_manager.get_context(stream_ids)
+
+        drop_extra_pre_encoded = 0 if not self.use_cache else self.asr_model.drop_extra_pre_encoded
+        log_probs, tail_log_probs, new_context = self.asr_model.stream_step(
+            processed_signal=feature_buffers,
+            processed_signal_length=feature_buffer_lens,
+            context=context,
+            drop_extra_pre_encoded=drop_extra_pre_encoded,
+            keep_all_outputs=keep_all_outputs,
+            drop_left_context=self.drop_left_context,
+            valid_out_len=self.valid_out_len,
+            return_tail_result=self.return_tail_result,
+        )
+
+        if log_probs is not None:
+            log_probs = normalize_log_probs(log_probs)
+        self.context_manager.update_cache(stream_ids, new_context, mapping)
+        self.context_manager.reset_slots(stream_ids, eos_flags)
+        self.decode_log_probs(requests, log_probs, tail_log_probs, ready_state_ids)
+
+    def transcribe_step_for_frames(self, frames: list[Frame]) -> None:
+        """
+        Transcribes the frames in a streaming manner.
+        After detecting EOU, it updates the state and run text processor.
+        If there are multiple streams, it waits until all states are ready to run text processor.
+        Args:
+            frames: (list[Frame]) List of frames to transcribe.
+        """
+        all_fbuffers, right_paddings = self.bufferer.update(frames)
+
+        ready_state_ids = set()
+        if len(all_fbuffers) > 0:
+            nonfinal_frames, nonfinal_fbuffers = [], []
+            final_frames, final_fbuffers = [], []
+            final_right_paddings = []
+            for jdx, bfeature in enumerate(all_fbuffers):
+                frame = frames[jdx]
+                if frame.is_last:
+                    final_frames.append(frame)
+                    final_fbuffers.append(bfeature)
+                    final_right_paddings.append(right_paddings[jdx])
+                else:
+                    nonfinal_frames.append(frame)
+                    nonfinal_fbuffers.append(bfeature)
+
+            if len(nonfinal_frames) > 0:
+                self.cache_aware_transcribe_step(
+                    nonfinal_frames, nonfinal_fbuffers, None, ready_state_ids, keep_all_outputs=False
+                )
+            if len(final_frames) > 0:
+                self.cache_aware_transcribe_step(
+                    final_frames, final_fbuffers, final_right_paddings, ready_state_ids, keep_all_outputs=True
+                )
+
+        # Postprocess the ready states
+        if len(ready_state_ids) > 0:
+            self.text_processor.process([self.get_state(stream_id) for stream_id in ready_state_ids])
+            ready_state_ids.clear()
+
+        self.update_partial_transcript(frames, self.tokenizer, self.leading_regex_pattern)
+
+    def transcribe_step_for_feature_buffers(self, fbuffers: list[FeatureBuffer]) -> None:
+        """
+        Transcribes the feature buffers in a streaming manner.
+        After detecting EOU, it updates the state and run text processor.
+        If there are multiple streams, it waits until all states are ready to run text processor.
+        Args:
+            fbuffers: (list[FeatureBuffer]) List of feature buffers to transcribe.
+        """
+        ready_state_ids = set()
+
+        final_fbuffers, final_features = [], []
+        nonfinal_fbuffers, nonfinal_features = [], []
+        final_right_paddings = []
+
+        for fbuffer in fbuffers:
+            feature = fbuffer.features
+            right_padding = max(0, self.expected_feature_buffer_len - fbuffer.valid_size)
+
+            if fbuffer.is_last:
+                final_fbuffers.append(fbuffer)
+                final_features.append(feature)
+                final_right_paddings.append(right_padding)
+            else:
+                nonfinal_fbuffers.append(fbuffer)
+                nonfinal_features.append(feature)
+
+        if len(nonfinal_fbuffers) > 0:
+            self.cache_aware_transcribe_step(
+                nonfinal_fbuffers, nonfinal_features, None, ready_state_ids, keep_all_outputs=False
+            )
+
+        if len(final_fbuffers) > 0:
+            self.cache_aware_transcribe_step(
+                final_fbuffers, final_features, final_right_paddings, ready_state_ids, keep_all_outputs=True
+            )
+
+        if len(ready_state_ids) > 0:
+            self.text_processor.process([self.get_state(stream_id) for stream_id in ready_state_ids])
+            ready_state_ids.clear()
+
+        self.update_partial_transcript(fbuffers, self.tokenizer, self.leading_regex_pattern)
+
+    def get_request_generator(self) -> ContinuousBatchedRequestStreamer:
+        """
+        Initialize the request generator.
+        Returns:
+            (ContinuousBatchedRequestStreamer) Request generator.
+        """
+        request_generator = ContinuousBatchedRequestStreamer(
+            n_frames_per_stream=1,
+            frame_size_in_secs=self.chunk_size_in_secs,
+            sample_rate=self.sample_rate,
+            batch_size=self.batch_size,
+            request_type=self.request_type,
+            preprocessor=self.preprocessor,
+            buffer_size_in_secs=self.buffer_size_in_secs,
+            device=self.device,
+            pad_last_frame=True,
+        )
+        return request_generator

nemo/collections/asr/inference/pipelines/cache_aware_rnnt_pipeline.py

@@ -0,0 +1,591 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import math
+from typing import TYPE_CHECKING
+
+import numpy as np
+import torch
+from omegaconf import DictConfig
+from torch import Tensor
+
+from nemo.collections.asr.inference.model_wrappers.cache_aware_rnnt_inference_wrapper import (
+    CacheAwareRNNTInferenceWrapper,
+)
+from nemo.collections.asr.inference.pipelines.base_pipeline import BasePipeline
+from nemo.collections.asr.inference.streaming.decoders.greedy.greedy_rnnt_decoder import RNNTGreedyDecoder
+from nemo.collections.asr.inference.streaming.endpointing.greedy.greedy_rnnt_endpointing import RNNTGreedyEndpointing
+from nemo.collections.asr.inference.streaming.framing.multi_stream import ContinuousBatchedRequestStreamer
+from nemo.collections.asr.inference.streaming.framing.request import FeatureBuffer, Frame, Request
+from nemo.collections.asr.inference.streaming.framing.request_options import ASRRequestOptions
+from nemo.collections.asr.inference.streaming.state.cache_aware_rnnt_state import (
+    CacheAwareRNNTBeamStreamingState,
+    CacheAwareRNNTStreamingState,
+)
+from nemo.collections.asr.inference.utils.endpointing_utils import millisecond_to_frames
+from nemo.collections.asr.inference.utils.enums import RequestType
+from nemo.collections.asr.inference.utils.per_stream_biasing import (
+    build_multi_biasing_ids_np,
+    release_all_biasing_models,
+    release_auto_managed_stream_biasing,
+)
+from nemo.collections.asr.inference.utils.pipeline_utils import (
+    check_existance_of_required_attributes,
+    drop_trailing_features,
+    get_confidence_utils,
+)
+from nemo.collections.asr.parts.submodules.rnnt_malsd_batched_computer import ModifiedALSDBatchedRNNTComputer
+from nemo.collections.asr.parts.utils.rnnt_utils import Hypothesis
+from nemo.utils import logging
+
+if TYPE_CHECKING:
+    from nemo.collections.asr.inference.itn.inverse_normalizer import AlignmentPreservingInverseNormalizer
+    from nemo.collections.asr.inference.nmt.llm_translator import LLMTranslator
+
+
+class CacheAwareRNNTPipeline(BasePipeline):
+    """Cache Aware RNNT pipeline."""
+
+    def __init__(
+        self,
+        cfg: DictConfig,
+        asr_model: CacheAwareRNNTInferenceWrapper,
+        itn_model: AlignmentPreservingInverseNormalizer | None = None,
+        nmt_model: LLMTranslator | None = None,
+    ):
+        """
+        Initialize the CacheAwareRNNTPipeline.
+        Args:
+            cfg: (DictConfig) Configuration parameters.
+            asr_model: (CacheAwareRNNTInferenceWrapper) ASR model.
+            itn_model: (AlignmentPreservingInverseNormalizer | None) Inverse Text Normalization model.
+            nmt_model: (LLMTranslator | None) LLM based translation model.
+        """
+        self.copy_asr_model_attributes(asr_model)
+        self.init_prompt_support()
+        self.init_parameters(cfg)
+        self.init_context_manager()
+        self.init_bufferer_for_cache_aware_streaming()
+        self.conf_func, self.confidence_aggregator = get_confidence_utils(cfg.confidence)
+        self.init_bpe_decoder()
+        self.init_greedy_rnnt_decoder()
+        self.init_endpointer()
+        self.init_text_processor(cfg, itn_model)
+        self.init_nmt_model(nmt_model)
+        self.init_decoding_computer()
+        if self.beam_decoder_computer is not None and self.prompt_enabled:
+            raise ValueError("Cache-aware RNNT MALSD beam search does not yet support prompt vectors.")
+        super().__init__()
+
+    def init_decoding_computer(self) -> None:
+        """Initialize ``decoding_computer``."""
+        self.decoding_computer = None
+        asr_model = getattr(self.asr_model, "asr_model", None)
+        if asr_model is None:
+            return
+        decoding = getattr(getattr(asr_model, "decoding", None), "decoding", None)
+        if decoding is not None:
+            self.decoding_computer = getattr(decoding, "decoding_computer", None)
+
+    @property
+    def beam_decoder_computer(self) -> ModifiedALSDBatchedRNNTComputer | None:
+        """Return ``decoding_computer`` when beam-search decoding is active."""
+        if isinstance(self.decoding_computer, ModifiedALSDBatchedRNNTComputer):
+            return self.decoding_computer
+        return None
+
+    def init_parameters(self, cfg: DictConfig) -> None:
+        """
+        Initialize the parameters.
+        Args:
+            cfg: (DictConfig) Configuration parameters.
+        """
+        if cfg.streaming.att_context_size is not None:
+            self.asr_model.set_default_att_context_size(att_context_size=cfg.streaming.att_context_size)
+
+        self.sample_rate = cfg.streaming.sample_rate
+        self.asr_output_granularity = cfg.asr_output_granularity
+        self.pre_encode_cache_size = self.asr_model.get_pre_encode_cache_size()
+        self.model_chunk_size = self.asr_model.get_chunk_size()
+        if isinstance(self.model_chunk_size, list):
+            self.model_chunk_size = self.model_chunk_size[1]
+
+        self.use_cache = cfg.streaming.use_cache
+        self.use_feat_cache = cfg.streaming.use_feat_cache
+
+        if cfg.streaming.get("chunk_size_in_secs", None) is not None:
+            self.chunk_size_in_secs = cfg.streaming.chunk_size_in_secs
+            self.tokens_per_frame = math.ceil(
+                np.trunc(self.chunk_size_in_secs / self.window_stride) / self.subsampling_factor
+            )
+            # overwrite the encoder streaming params with proper shift size for cache aware streaming
+            self.asr_model.setup_streaming_params(
+                chunk_size=self.model_chunk_size // self.subsampling_factor, shift_size=self.tokens_per_frame
+            )
+        else:
+            self.chunk_size_in_secs = self.model_chunk_size * self.window_stride
+            self.tokens_per_frame = math.ceil(self.model_chunk_size / self.subsampling_factor)
+
+        if isinstance(self.pre_encode_cache_size, list):
+            self.pre_encode_cache_size = self.pre_encode_cache_size[1]
+        self.pre_encode_cache_size_in_secs = self.pre_encode_cache_size * self.window_stride
+
+        # Context Manager
+        self.batch_size = cfg.streaming.batch_size
+        self.num_slots = cfg.streaming.num_slots
+        if self.num_slots < self.batch_size:
+            raise ValueError(
+                f"Number of slots in the context manager must be >= batch_size: {self.num_slots} < {self.batch_size}"
+            )
+        model_chunk_size_in_secs = self.model_chunk_size * self.window_stride
+
+        if self.use_cache:
+            # if using cache, we need to pad some samples for pre_encode
+            self.buffer_size_in_secs = self.pre_encode_cache_size_in_secs + model_chunk_size_in_secs
+            self.drop_left_context = None
+            self.valid_out_len = None
+        else:
+            # if not using cache, we need to keep left context in buffer, but no extra padding in pre_encode
+            left_context_size = self.asr_model.get_att_context_size()[0]
+            if left_context_size < 0:
+                raise ValueError(f"Left context size should not be a negative value: {left_context_size}")
+            self.buffer_size_in_secs = (
+                model_chunk_size_in_secs + left_context_size * self.subsampling_factor * self.window_stride
+            )
+            self.drop_left_context = left_context_size
+            self.valid_out_len = self.tokens_per_frame
+
+        # Expected feature buffer length for trimming (safeguard for feature buffer inputs)
+        self.expected_feature_buffer_len = int(self.buffer_size_in_secs / self.window_stride)
+
+        self.stop_history_eou_in_milliseconds = cfg.endpointing.stop_history_eou
+        self.residue_tokens_at_end = cfg.endpointing.residue_tokens_at_end
+        self.word_boundary_tolerance = cfg.streaming.word_boundary_tolerance
+        self.return_tail_result = cfg.return_tail_result
+
+        self.request_type = RequestType.from_str(cfg.streaming.request_type)
+
+    def init_greedy_rnnt_decoder(self) -> None:
+        """Initialize the RNNT decoder."""
+        check_existance_of_required_attributes(self, ['vocabulary', 'conf_func'])
+        self.greedy_rnnt_decoder = RNNTGreedyDecoder(vocabulary=self.vocabulary, conf_func=self.conf_func)
+
+    def init_endpointer(self) -> None:
+        """Initialize the endpointer."""
+        check_existance_of_required_attributes(
+            self,
+            [
+                'vocabulary',
+                'model_stride_in_milliseconds',
+                'stop_history_eou_in_milliseconds',
+                'residue_tokens_at_end',
+            ],
+        )
+
+        self.endpointer = RNNTGreedyEndpointing(
+            vocabulary=self.vocabulary,
+            ms_per_timestep=self.model_stride_in_milliseconds,
+            stop_history_eou=self.stop_history_eou_in_milliseconds,
+            residue_tokens_at_end=self.residue_tokens_at_end,
+        )
+
+    def create_state(self, options: ASRRequestOptions) -> CacheAwareRNNTStreamingState:
+        """
+        Create new empty state.
+        Args:
+            options: (ASRRequestOptions) Request options for particular stream.
+        Returns:
+            (CacheAwareRNNTStreamingState) New empty state.
+        """
+        state = (
+            CacheAwareRNNTBeamStreamingState()
+            if self.beam_decoder_computer is not None
+            else CacheAwareRNNTStreamingState()
+        )
+        state.set_global_offset(0)
+        new_options = options.fill_defaults(
+            default_enable_itn=self.text_processor.itn_enabled,
+            default_enable_nmt=self.nmt_enabled,
+            default_source_language=self.nmt_model.source_language if self.nmt_enabled else None,
+            default_target_language=self.nmt_model.target_language if self.nmt_enabled else None,
+            default_stop_history_eou=self.stop_history_eou_in_milliseconds,
+            default_asr_output_granularity=self.asr_output_granularity,
+            default_language_code="en-US" if self.prompt_enabled else None,
+        )
+
+        eou_label_buffer_size = 0
+        if new_options.stop_history_eou > 0:
+            eou_label_buffer_size = millisecond_to_frames(
+                new_options.stop_history_eou, math.ceil(self.model_stride_in_milliseconds)
+            )
+            eou_label_buffer_size += self.residue_tokens_at_end
+        state.setup_label_buffer(eou_label_buffer_size, self.blank_id)
+        state.set_previous_hypothesis(None)
+        state.set_options(new_options)
+
+        # Create per-stream prompt index for prompt-enabled models
+        if self.prompt_enabled:
+            lang_code = getattr(new_options, "language_code", None)
+            if not isinstance(lang_code, str) or len(lang_code) == 0:
+                raise ValueError("Prompt-enabled model requires a valid language_code in request options.")
+            prompt_idx = self._resolve_prompt_index(lang_code)
+            state.set_prompt_index(prompt_idx)
+
+        return state
+
+    def close_session(self) -> None:
+        """Close the session and release per-stream biasing models held in the decoder."""
+        if self.decoding_computer is not None and self.decoding_computer.per_stream_biasing_enabled:
+            release_all_biasing_models(self.decoding_computer.biasing_multi_model, self._state_pool.values())
+        super().close_session()
+
+    def get_sep(self) -> str:
+        """Return the separator for the text processor."""
+        return self.sep
+
+    def preprocess(self, buffers: list[Tensor], right_paddings: list[int] | None = None) -> tuple[Tensor, Tensor]:
+        """
+        Preprocess the feature buffers by stacking them and computing the lengths
+        Args:
+            buffers: (list[Tensor]) List of feature buffers.
+            right_paddings: (list[int] | None) List of right paddings.
+        Returns:
+            (tuple[Tensor, Tensor]) Processed feature buffers and their lengths.
+        """
+        feature_buffers = [f_buffer.unsqueeze_(0) for f_buffer in buffers]
+        # Trim to expected feature buffer length (safeguard for external feature buffer inputs)
+        feature_buffers = [
+            drop_trailing_features(f_buffer, self.expected_feature_buffer_len) for f_buffer in feature_buffers
+        ]
+        feature_buffer_lens = torch.tensor([f_buffer.shape[2] for f_buffer in feature_buffers], device=self.device)
+        if right_paddings is not None:
+            right_paddings = torch.tensor(right_paddings, device=feature_buffer_lens.device)
+            feature_buffer_lens = feature_buffer_lens - right_paddings
+        feature_buffers = torch.cat(feature_buffers).to(self.device)
+        return feature_buffers, feature_buffer_lens
+
+    def _streaming_step(
+        self,
+        states: list[CacheAwareRNNTStreamingState],
+        feature_buffers: Tensor,
+        feature_buffer_lens: Tensor,
+        context,
+        previous_hypotheses: list[Hypothesis | None],
+        drop_extra_pre_encoded: int,
+        keep_all_outputs: bool,
+        prompt_vectors: Tensor | None,
+    ) -> tuple[list[Hypothesis], object]:
+        """
+        Run one cache-aware encode/decode step for the current chunk.
+        Returns per-stream hypotheses and the updated encoder cache context.
+        """
+        if self.beam_decoder_computer is None:
+            return self.asr_model.stream_step(
+                processed_signal=feature_buffers,
+                processed_signal_length=feature_buffer_lens,
+                context=context,
+                previous_hypotheses=previous_hypotheses,
+                drop_extra_pre_encoded=drop_extra_pre_encoded,
+                keep_all_outputs=keep_all_outputs,
+                drop_left_context=self.drop_left_context,
+                valid_out_len=self.valid_out_len,
+                prompt_vectors=prompt_vectors,
+            )
+        return self.asr_model.malsd_stream_step(
+            malsd_computer=self.beam_decoder_computer,
+            states=states,
+            processed_signal=feature_buffers,
+            processed_signal_length=feature_buffer_lens,
+            context=context,
+            drop_extra_pre_encoded=drop_extra_pre_encoded,
+            keep_all_outputs=keep_all_outputs,
+            drop_left_context=self.drop_left_context,
+            valid_out_len=self.valid_out_len,
+        )
+
+    def _prepare_per_stream_biasing(
+        self,
+        states: list[CacheAwareRNNTStreamingState],
+        previous_hypotheses: list[Hypothesis | None],
+    ) -> list[Hypothesis | None]:
+        if self.decoding_computer is None or not self.decoding_computer.per_stream_biasing_enabled:
+            if any(state.has_biasing_request() for state in states):
+                logging.warning(
+                    "Biasing request is not empty, but decoder does not support per-stream biasing. Skipping"
+                )
+            return previous_hypotheses
+
+        multi_biasing_ids_np = build_multi_biasing_ids_np(
+            states,
+            self.decoding_computer.biasing_multi_model,
+            self.asr_model.tokenizer,
+        )
+
+        if self.beam_decoder_computer is not None:
+            return previous_hypotheses
+
+        for i, (state, previous_hyp) in enumerate(zip(states, previous_hypotheses)):
+            if multi_biasing_ids_np[i] < 0:
+                continue
+            biasing_cfg = state.options.biasing_cfg
+            if previous_hyp is None:
+                previous_hypotheses[i] = Hypothesis.empty_with_biasing_cfg(biasing_cfg)
+            else:
+                previous_hyp.biasing_cfg = biasing_cfg
+        return previous_hypotheses
+
+    def _apply_beam_update_(self, state: CacheAwareRNNTBeamStreamingState, eou_detected: bool) -> None:
+        """After endpointing: refresh beam publish tokens and fold cumulative prefix on EOU."""
+        if eou_detected and state.hyp_decoding_state is not None:
+            beam_idx = state.select_best_beam_idx_(score_norm=True)
+            self.beam_decoder_computer.select_beam_in_state_item_(state.hyp_decoding_state, beam_idx)
+        state.update_(eou_detected)
+
+    def run_greedy_decoder(self, state: CacheAwareRNNTStreamingState, request: Request, hyp: Hypothesis) -> bool:
+        """
+        Run the greedy RNNT decoder on the hypothesis and update the state
+        Args:
+            state: (CacheAwareRNNTStreamingState) The state of the stream
+            request: (Request) The current request (frame or feature buffer)
+            hyp: (Hypothesis) The hypothesis of the current request
+        Returns:
+            (bool) Whether EOU is detected.
+        """
+        eou_detected = request.is_last
+        # Per-token non-blank confidence precomputed during RNN-T decoding (aligned with `hyp.y_sequence`).
+        # Populated when greedy or batched-beam preserve_frame_confidence is enabled; otherwise None.
+        cur_output, cur_labels, new_offset = self.greedy_rnnt_decoder(
+            global_timestamps=hyp.timestamp,
+            tokens=hyp.y_sequence,
+            length=self.tokens_per_frame,
+            offset=state.offset,
+            confidences=hyp.non_blank_step_confidence_precomputed,
+        )
+        state.set_offset(new_offset)
+
+        # cur labels contains blank tokens as well, it is needed for EOU detection
+        state.update_label_buffer(cur_labels)
+
+        if not eou_detected:
+            emissions = state.get_label_buffer()
+            pivot_point = len(emissions) - 1
+            eou_detected, _ = self.endpointer.detect_eou_near_pivot(
+                emissions, pivot_point, stop_history_eou=state.options.stop_history_eou
+            )
+
+        state.update_state(cur_output, eou_detected=eou_detected)
+        return eou_detected
+
+    def cache_aware_transcribe_step(
+        self,
+        requests: list[Request],
+        features: list[Tensor],
+        right_paddings: list[int],
+        ready_state_ids: set,
+        keep_all_outputs: bool = False,
+    ) -> None:
+        """
+        Cache Aware Transcribe Step
+        It receives a list of requests (Frame or FeatureBuffer) and features and do the following:
+
+        1. Preprocess the features by stacking them and computing the lengths
+        2. Collecting previous hypotheses for stateful decoding
+        3. Get the context and mapping from the context manager for cache aware streaming
+        4. Perform a streaming step with the ASR model
+        5. Update the cache and reset the cache slots for the streams that has ended
+        6. Update the previous hypothesis and reset the previous hypothesis for the streams that has ended
+        7. Perform greedy RNNT decoding to get the best hypothesis and update the states
+        8. Update the ready states to indicate that the state is ready for text post-processing
+        Args:
+            requests: (list[Request]) List of requests (frames or feature buffers) to transcribe.
+            features: (list[Tensor]) List of feature buffers.
+            right_paddings: (list[int] | None) List of right paddings.
+            ready_state_ids: (set) Set of ready state IDs.
+            keep_all_outputs: (bool) Whether to keep all outputs or not.
+        """
+
+        feature_buffers, feature_buffer_lens = self.preprocess(features, right_paddings)
+        states, stream_ids, eos_flags = [], [], []
+        for request in requests:
+            states.append(self.get_state(request.stream_id))
+            stream_ids.append(request.stream_id)
+            eos_flags.append(request.is_last)
+
+        previous_hypotheses = [state.get_previous_hypothesis() for state in states]
+
+        previous_hypotheses = self._prepare_per_stream_biasing(
+            states=states,
+            previous_hypotheses=previous_hypotheses,
+        )
+
+        context, mapping = self.context_manager.get_context(stream_ids)
+
+        prompt_vectors = None
+        if self.prompt_enabled:
+            prompt_vectors = self._build_prompt_vectors(states)
+
+        drop_extra_pre_encoded = 0 if not self.use_cache else self.asr_model.drop_extra_pre_encoded
+        best_hyp, new_context = self._streaming_step(
+            states=states,
+            feature_buffers=feature_buffers,
+            feature_buffer_lens=feature_buffer_lens,
+            context=context,
+            previous_hypotheses=previous_hypotheses,
+            drop_extra_pre_encoded=drop_extra_pre_encoded,
+            keep_all_outputs=keep_all_outputs,
+            prompt_vectors=prompt_vectors,
+        )
+
+        # update the cache and reset the cache slots for the streams that has ended
+        self.context_manager.update_cache(stream_ids, new_context, mapping)
+        self.context_manager.reset_slots(stream_ids, eos_flags)
+
+        # update the previous hypothesis and reset the previous hypothesis for the streams that has ended
+        for state, hyp, eos in zip(states, best_hyp, eos_flags):
+            if eos:
+                state.reset_previous_hypothesis()
+            else:
+                state.set_previous_hypothesis(hyp)
+
+        # run greedy decoder for each request-state-hypothesis tuple
+        for request, state, hyp in zip(requests, states, best_hyp):
+            eou_detected = self.run_greedy_decoder(state, request, hyp)
+            if self.beam_decoder_computer is not None:
+                self._apply_beam_update_(state, eou_detected)
+            if eou_detected:
+                self.bpe_decoder.decode_bpe_tokens(state)
+                state.cleanup_after_eou()
+                ready_state_ids.add(request.stream_id)
+
+        # Cleanup per-stream biasing models when stream ends
+        if self.decoding_computer is not None and self.decoding_computer.per_stream_biasing_enabled:
+            for request, state in zip(requests, states):
+                # only the first request contains biasing options; biasing options for the stream are stored in state
+                if request.is_last and state.has_biasing_request():
+                    release_auto_managed_stream_biasing(state, self.decoding_computer.biasing_multi_model)
+
+        if self.beam_decoder_computer is not None:
+            for state, eos in zip(states, eos_flags):
+                if eos:
+                    state.reset_beam_decoding_state_()
+
+    def transcribe_step_for_feature_buffers(self, fbuffers: list[FeatureBuffer]) -> None:
+        """
+        Transcribes the feature buffers in a streaming manner.
+        After detecting EOU, it updates the state and run text processor.
+        If there are multiple streams, it waits until all states are ready to run text processor.
+        Args:
+            fbuffers: (list[FeatureBuffer]) List of feature buffers to transcribe.
+        """
+        ready_state_ids = set()
+
+        final_fbuffers, final_features = [], []
+        nonfinal_fbuffers, nonfinal_features = [], []
+        final_right_paddings = []
+
+        for fbuffer in fbuffers:
+            feature = fbuffer.features
+            right_padding = max(0, self.expected_feature_buffer_len - fbuffer.valid_size)
+
+            if fbuffer.is_last:
+                final_fbuffers.append(fbuffer)
+                final_features.append(feature)
+                final_right_paddings.append(right_padding)
+            else:
+                nonfinal_fbuffers.append(fbuffer)
+                nonfinal_features.append(feature)
+
+        if len(nonfinal_fbuffers) > 0:
+            self.cache_aware_transcribe_step(
+                nonfinal_fbuffers, nonfinal_features, None, ready_state_ids, keep_all_outputs=False
+            )
+
+        if len(final_fbuffers) > 0:
+            self.cache_aware_transcribe_step(
+                final_fbuffers, final_features, final_right_paddings, ready_state_ids, keep_all_outputs=True
+            )
+
+        if len(ready_state_ids) > 0:
+            self.text_processor.process([self.get_state(stream_id) for stream_id in ready_state_ids])
+            ready_state_ids.clear()
+
+        self.update_partial_transcript(fbuffers, self.tokenizer, self.leading_regex_pattern)
+
+    def transcribe_step_for_frames(self, frames: list[Frame]) -> None:
+        """
+        Transcribes the frames in a streaming manner.
+        After detecting EOU, it updates the state and run text processor.
+        If there are multiple streams, it waits until all states are ready to run text processor.
+        Args:
+            frames: (list[Frame]) List of frames to transcribe.
+        """
+
+        all_fbuffers, right_paddings = self.bufferer.update(frames)
+        ready_state_ids = set()
+
+        # streams that contains multiple frames
+        if len(all_fbuffers) > 0:
+            final_frames, final_fbuffers = [], []
+            nonfinal_frames, nonfinal_fbuffers = [], []
+            final_right_paddings = []
+
+            for jdx, bfeature in enumerate(all_fbuffers):
+                bframe = frames[jdx]
+
+                if bframe.is_last:
+                    final_frames.append(bframe)
+                    final_fbuffers.append(bfeature)
+                    final_right_paddings.append(right_paddings[jdx])
+                else:
+                    nonfinal_frames.append(bframe)
+                    nonfinal_fbuffers.append(bfeature)
+
+            if len(nonfinal_frames) > 0:
+                self.cache_aware_transcribe_step(
+                    nonfinal_frames, nonfinal_fbuffers, None, ready_state_ids, keep_all_outputs=False
+                )
+
+            if len(final_frames) > 0:
+                self.cache_aware_transcribe_step(
+                    final_frames, final_fbuffers, final_right_paddings, ready_state_ids, keep_all_outputs=True
+                )
+
+        # post-process the ready states
+        if len(ready_state_ids) > 0:
+            self.text_processor.process([self.get_state(stream_id) for stream_id in ready_state_ids])
+            ready_state_ids.clear()
+
+        self.update_partial_transcript(frames, self.tokenizer, self.leading_regex_pattern)
+
+    def get_request_generator(self) -> ContinuousBatchedRequestStreamer:
+        """
+        Initialize the request generator.
+        Returns:
+            (ContinuousBatchedRequestStreamer) Request generator.
+        """
+        # for cache aware streaming we need to process one frame at a time -> n_frames_per_stream=1
+        request_generator = ContinuousBatchedRequestStreamer(
+            n_frames_per_stream=1,
+            frame_size_in_secs=self.chunk_size_in_secs,
+            sample_rate=self.sample_rate,
+            batch_size=self.batch_size,
+            request_type=self.request_type,
+            preprocessor=self.preprocessor,
+            buffer_size_in_secs=self.buffer_size_in_secs,
+            device=self.device,
+            pad_last_frame=True,
+        )
+        return request_generator

nemo/collections/asr/inference/pipelines/pipeline_interface.py

@@ -0,0 +1,79 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from abc import ABC, abstractmethod
+
+from nemo.collections.asr.inference.streaming.framing.request import Request
+from nemo.collections.asr.inference.streaming.framing.request_options import ASRRequestOptions
+
+
+class PipelineInterface(ABC):
+    """
+    The base interface for streaming speech pipelines
+    Base usage for all pipelines:
+        pipeline.start_session()
+        for requests in request_generator:
+            pipeline.transcribe_step(requests)
+        pipeline.close_session()
+    """
+
+    @abstractmethod
+    def open_session(self):
+        """
+        Open a new session
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def close_session(self):
+        """
+        End the current session
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_state(self, stream_id: int):
+        """
+        Get the state of the stream
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def delete_state(self, stream_id: int):
+        """
+        Delete the state of the stream
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def create_state(self, options: ASRRequestOptions):
+        """
+        Create a new empty state
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def init_state(self, stream_id: int, options: ASRRequestOptions):
+        """
+        Initialize the state of the stream
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def transcribe_step(self, requests: list[Request]):
+        """
+        Transcribe a step
+        """
+        raise NotImplementedError

nemo/collections/asr/inference/streaming/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

nemo/collections/asr/inference/streaming/buffering/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

nemo/collections/asr/inference/streaming/buffering/audio_bufferer.py

@@ -0,0 +1,131 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import torch
+from torch import Tensor
+from nemo.collections.asr.inference.streaming.framing.request import Frame
+
+
+class AudioBufferer:
+    """
+    Audio bufferer class
+    It buffers the audio chunks and maintains the buffer.
+    """
+
+    def __init__(self, sample_rate: int, buffer_size_in_secs: float):
+        """
+        Args:
+            sample_rate (int): sample rate
+            buffer_size_in_secs (float): buffer size in seconds
+        """
+        self.buffer_size = int(buffer_size_in_secs * sample_rate)
+        self.sample_buffer = torch.zeros(self.buffer_size, dtype=torch.float32)
+        self.left_padding = self.buffer_size
+
+    def reset(self) -> None:
+        """
+        Reset the buffer to zero
+        """
+        self.sample_buffer.zero_()
+        self.left_padding = self.buffer_size
+
+    def update(self, frame: Frame) -> None:
+        """
+        Update the buffer with the new frame
+        Args:
+            frame (Frame): frame to update the buffer with
+        """
+        if frame.size > self.buffer_size:
+            raise RuntimeError(f"Frame size ({frame.size}) exceeds buffer size ({self.buffer_size})")
+
+        shift = frame.size
+        self.sample_buffer = torch.roll(self.sample_buffer, -shift)
+        self.sample_buffer[-shift:].copy_(frame.samples)
+        self.left_padding = max(0, self.left_padding - shift)
+
+    def get_buffer(self) -> Tensor:
+        """
+        Get the current buffer
+        Returns:
+            Tensor: current state of the buffer
+        """
+        return self.sample_buffer.clone()
+
+    def get_left_padding(self) -> int:
+        """
+        Get the left padding
+        Returns:
+            int: left padding
+        """
+        return self.left_padding
+
+
+class BatchedAudioBufferer:
+    """
+    Batched audio bufferer class
+    It buffers the audio chunks from multiple streams and returns the buffers.
+    """
+
+    def __init__(self, sample_rate: int, buffer_size_in_secs: float):
+        """
+        Args:
+            sample_rate (int): sample rate
+            buffer_size_in_secs (float): buffer size in seconds
+        """
+        self.sample_rate = sample_rate
+        self.buffer_size_in_secs = buffer_size_in_secs
+        self.bufferers = {}
+
+    def reset(self) -> None:
+        """
+        Reset bufferers
+        """
+        self.bufferers = {}
+
+    def rm_bufferer(self, stream_id: int) -> None:
+        """
+        Remove bufferer for the given stream id
+        Args:
+            stream_id (int): stream id
+        """
+        self.bufferers.pop(stream_id, None)
+
+    def update(self, frames: list[Frame]) -> tuple[list[Tensor], list[int]]:
+        """
+        Update the bufferers with the new frames.
+        Frames can come from different streams (audios), so we need to maintain a bufferer for each stream
+        Args:
+            frames (list[Frame]): list of frames
+        Returns:
+            tuple[list[Tensor], list[int]]:
+                buffers: list of buffered audio tensors, one per input frame
+                left_paddings: list of left paddings, one per input frame
+        """
+        buffers, left_paddings = [], []
+        for frame in frames:
+            bufferer = self.bufferers.get(frame.stream_id, None)
+
+            if bufferer is None:
+                bufferer = AudioBufferer(self.sample_rate, self.buffer_size_in_secs)
+                self.bufferers[frame.stream_id] = bufferer
+
+            bufferer.update(frame)
+            buffers.append(bufferer.get_buffer())
+            left_paddings.append(bufferer.get_left_padding())
+
+            if frame.is_last:
+                self.rm_bufferer(frame.stream_id)
+
+        return buffers, left_paddings

nemo/collections/asr/inference/streaming/buffering/cache_feature_bufferer.py

@@ -0,0 +1,225 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import math
+from queue import Queue
+
+import torch
+from omegaconf import DictConfig
+from torch import Tensor
+
+from nemo.collections.asr.inference.streaming.buffering.audio_bufferer import AudioBufferer
+from nemo.collections.asr.inference.streaming.framing.request import Frame
+from nemo.collections.asr.inference.utils.constants import LOG_MEL_ZERO
+from nemo.collections.asr.models import ASRModel
+
+
+class BatchedCacheFeatureBufferer:
+    """
+    Batched cache feature bufferer class
+    Buffers feature chunks from multiple audio streams and manages their storage.
+    Maintains a tensor of shape (num_slots, n_feat, feature_buffer_len), where each slot
+    corresponds to a single audio stream. The number of slots equals the number of
+    active (or open) audio streams.
+    """
+
+    def __init__(
+        self,
+        num_slots: int,
+        sample_rate: int,
+        buffer_size_in_secs: float,
+        chunk_size_in_secs: float,
+        preprocessor_cfg: DictConfig,
+        device: torch.device,
+        fill_value: float = LOG_MEL_ZERO,
+    ):
+        """
+        Args:
+            num_slots (int): number of slots, where each slot contains feature buffer for a single audio stream
+            sample_rate (int): sample rate
+            buffer_size_in_secs (float): buffer size in seconds
+            chunk_size_in_secs (float): chunk size in seconds
+            preprocessor_cfg (DictConfig): preprocessor configuration
+            device (torch.device): device
+            fill_value (float): fill value for the feature buffer
+        """
+        if buffer_size_in_secs < chunk_size_in_secs:
+            raise ValueError(
+                f"Buffer size ({buffer_size_in_secs}s) should be no less than chunk size ({chunk_size_in_secs}s)"
+            )
+
+        self.num_slots = num_slots
+        self.sample_rate = sample_rate
+        self.buffer_size_in_secs = buffer_size_in_secs
+        self.chunk_size_in_secs = chunk_size_in_secs
+        self.preprocessor_cfg = preprocessor_cfg
+        self.device = device
+
+        self.is_buffer_size_equal_to_chunk_size = math.isclose(self.buffer_size_in_secs, self.chunk_size_in_secs)
+        self.plus_one = 0 if self.is_buffer_size_equal_to_chunk_size else 1
+
+        if hasattr(preprocessor_cfg, 'log') and preprocessor_cfg.log:
+            self.ZERO_LEVEL_SPEC_DB_VAL = LOG_MEL_ZERO  # Log-Mel spectrogram value for zero signals
+        else:
+            self.ZERO_LEVEL_SPEC_DB_VAL = fill_value  # Custom fill value for the feature buffer
+
+        self.n_feat = preprocessor_cfg.features
+        self.timestep_duration = preprocessor_cfg.window_stride
+        self.n_chunk_look_back = int(self.timestep_duration * self.sample_rate)
+        self.chunk_size = int(self.chunk_size_in_secs * self.sample_rate)
+        self.extended_chunk_size = self.n_chunk_look_back + self.chunk_size
+        self.audio_bufferers = [
+            AudioBufferer(self.sample_rate, self.buffer_size_in_secs) for _ in range(self.num_slots)
+        ]
+
+        self.feature_buffer_len = int(buffer_size_in_secs / self.timestep_duration)
+        self.feature_chunk_len = int(chunk_size_in_secs / self.timestep_duration)
+        self.feature_buffer = torch.full(
+            [self.num_slots, self.n_feat, self.feature_buffer_len],
+            self.ZERO_LEVEL_SPEC_DB_VAL,
+            dtype=torch.float32,
+            device=self.device,
+        )
+
+        self.preprocessor = ASRModel.from_config_dict(preprocessor_cfg)
+        self.preprocessor.to(self.device)
+
+        self.streamidx2slotidx, self.slotidx2streamidx = {}, {}
+        self.available_slots = Queue(self.num_slots)
+        for i in range(self.num_slots):
+            self.available_slots.put(i)
+
+    def free_slots(self, slot_ids: list[int]) -> None:
+        """
+        Free the slots for the given slot_ids
+        Args:
+            slot_ids (list[int]): list of slot ids
+        """
+        for slot_id in slot_ids:
+            if slot_id not in self.slotidx2streamidx:
+                continue
+            self.available_slots.put(slot_id)
+            stream_id = self.slotidx2streamidx[slot_id]
+            del self.slotidx2streamidx[slot_id], self.streamidx2slotidx[stream_id]
+
+    def reset_slots(self, slot_ids: list[int]) -> None:
+        """
+        Reset the slots for the given slot_ids
+        Args:
+            slot_ids (list[int]): list of slot ids
+        """
+        slot_ids_tensor = torch.tensor(slot_ids, device=self.device, dtype=torch.long)
+        self.feature_buffer.index_fill_(0, slot_ids_tensor, self.ZERO_LEVEL_SPEC_DB_VAL)
+        for slot_id in slot_ids:
+            self.audio_bufferers[slot_id].reset()
+
+    def preprocess(
+        self, audio_buffers: list[Tensor], right_paddings: Tensor, expected_feat_len: int
+    ) -> tuple[Tensor, Tensor]:
+        """
+        Preprocess the audio buffers with the given right paddings and expected feature length
+        Args:
+            audio_buffers (list[Tensor]): list of audio buffers
+            right_paddings (Tensor): right paddings: right paddings are not zero for last frames
+            expected_feat_len (int): expected feature length
+        Returns:
+            tuple[Tensor, Tensor]: features and right paddings
+        """
+        signals = torch.vstack(audio_buffers).to(self.device)  # B x T
+        signals_len = torch.tensor([signals.shape[1]] * signals.shape[0], device=self.device, dtype=torch.long)  # B
+        signals_len = signals_len - right_paddings.long()
+        features, feature_lens = self.preprocessor(input_signal=signals, length=signals_len)
+        if features.shape[2] > expected_feat_len:
+            features = features[:, :, :expected_feat_len]  # B x F x T
+            feature_lens = feature_lens.clamp(max=expected_feat_len)
+        right_padding = (features.shape[2] - feature_lens).clamp(min=0).to(torch.long)
+        return features, right_padding
+
+    def _update_feature_buffer(self, slot_ids: list[int], feat_chunk: Tensor) -> None:
+        """
+        Add an extracted feature to `feature_buffer`
+        Args:
+            slot_ids (list[int]): list of slot ids
+            feat_chunk (Tensor): feature chunk of shape (B, F, T)
+        """
+        for i, slot_id in enumerate(slot_ids):
+            chunk_len = feat_chunk[i].shape[-1]
+            if chunk_len > self.feature_buffer_len:
+                raise ValueError(f"feat_chunk ({chunk_len}) longer than buffer ({self.feature_buffer_len})")
+
+            shifted = self.feature_buffer[slot_id, :, chunk_len:].clone()
+            self.feature_buffer[slot_id, :, :-chunk_len].copy_(shifted)
+            self.feature_buffer[slot_id, :, -chunk_len:].copy_(feat_chunk[i])
+
+    def update(self, frames: list[Frame]) -> tuple[list[Tensor], list[int]]:
+        """
+        Update the feature bufferers with the new frames.
+        Args:
+            frames (list[Frame]): list of frames with length equal to batch size
+        Returns:
+            tuple[list[Tensor], list[int]]: feature buffers and right paddings
+        """
+        # if there are no frames, return empty lists
+        if len(frames) == 0:
+            return [], []
+
+        # if the stream_id is new, we need to assign a slot to it
+        slot_ids, slots_to_reset, slots_to_free = [], [], []
+        for frame in frames:
+            stream_id = frame.stream_id
+            slot_idx = self.streamidx2slotidx.get(stream_id, None)
+            if stream_id not in self.streamidx2slotidx:
+                if self.available_slots.empty():
+                    raise RuntimeError("No free slots available")
+                slot_idx = self.available_slots.get()
+                self.streamidx2slotidx[stream_id] = slot_idx
+                self.slotidx2streamidx[slot_idx] = stream_id
+                slots_to_reset.append(slot_idx)
+
+            slot_ids.append(slot_idx)
+            if frame.is_last:
+                slots_to_free.append(slot_idx)
+
+        # reset the slots for the new stream_ids
+        if len(slots_to_reset) > 0:
+            self.reset_slots(slots_to_reset)
+
+        right_paddings = torch.zeros(len(frames), dtype=torch.long, device=self.device)
+        audio_buffers = []
+        for i, frame in enumerate(frames):
+            slot_id = slot_ids[i]
+            right_paddings[i] = frame.size - frame.valid_size
+            self.audio_bufferers[slot_id].update(frame)
+
+            buffer = self.audio_bufferers[slot_id].sample_buffer
+            if not self.is_buffer_size_equal_to_chunk_size:
+                # Add look_back to have context for the first feature
+                audio_buffers.append(buffer[-(self.n_chunk_look_back + self.chunk_size) :])
+            else:
+                # If the buffer size is equal to the chunk size, just take the whole buffer
+                audio_buffers.append(buffer)
+
+        features, right_paddings = self.preprocess(
+            audio_buffers=audio_buffers,
+            right_paddings=right_paddings,
+            expected_feat_len=self.feature_chunk_len + self.plus_one,
+        )
+        self._update_feature_buffer(slot_ids=slot_ids, feat_chunk=features[:, :, -self.feature_chunk_len :])
+        fbuffers = list(self.feature_buffer[slot_ids].unbind(0))
+
+        if len(slots_to_free) > 0:
+            self.free_slots(slots_to_free)
+
+        return fbuffers, right_paddings.tolist()

nemo/collections/asr/inference/streaming/buffering/feature_bufferer.py

@@ -0,0 +1,160 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+from omegaconf import DictConfig
+
+from nemo.collections.asr.inference.streaming.framing.request import FeatureBuffer
+from nemo.collections.asr.inference.utils.constants import LOG_MEL_ZERO
+
+
+class FeatureBufferer:
+    """
+    Feature bufferer class
+    It buffers the feature chunks and maintains the buffer.
+    """
+
+    def __init__(
+        self,
+        sample_rate: int,
+        buffer_size_in_secs: float,
+        preprocessor_cfg: DictConfig,
+        device: torch.device,
+        fill_value: float = LOG_MEL_ZERO,
+    ):
+        """
+        Args:
+            sample_rate (int): sample rate
+            buffer_size_in_secs (float): buffer size in seconds
+            preprocessor_cfg (DictConfig): preprocessor config
+            device (torch.device): device
+            fill_value (float): value to fill the feature buffer with
+        """
+        self.sample_rate = sample_rate
+        self.buffer_size_in_secs = buffer_size_in_secs
+        self.device = device
+
+        if hasattr(preprocessor_cfg, 'log') and preprocessor_cfg.log:
+            self.ZERO_LEVEL_SPEC_DB_VAL = LOG_MEL_ZERO
+        else:
+            self.ZERO_LEVEL_SPEC_DB_VAL = fill_value
+
+        self.n_feat = preprocessor_cfg.features
+        self.feature_buffer_len = int(buffer_size_in_secs / preprocessor_cfg.window_stride)
+        self.feature_buffer = torch.full(
+            [self.n_feat, self.feature_buffer_len],
+            self.ZERO_LEVEL_SPEC_DB_VAL,
+            dtype=torch.float32,
+            device=self.device,
+        )
+
+    def reset(self) -> None:
+        """
+        Reset the buffer to zero
+        """
+        self.feature_buffer.fill_(self.ZERO_LEVEL_SPEC_DB_VAL)
+
+    def update(self, fbuffer: FeatureBuffer) -> None:
+        """
+        Replace feature buffer with new data
+        Args:
+            fbuffer (FeatureBuffer): feature buffer to update
+        """
+        # Resize if needed (optional)
+        if fbuffer.size != self.feature_buffer.shape[1]:
+            self.feature_buffer = torch.full(
+                [self.n_feat, fbuffer.size],
+                self.ZERO_LEVEL_SPEC_DB_VAL,
+                dtype=torch.float32,
+                device=self.device,
+            )
+
+        self.feature_buffer.copy_(fbuffer.features)
+
+    def get_feature_buffer(self) -> torch.Tensor:
+        """
+        Get the current feature buffer
+        Returns:
+            torch.Tensor: current state of the feature buffer
+        """
+        return self.feature_buffer.clone()
+
+
+class BatchedFeatureBufferer:
+    """
+    Batched feature bufferer class
+    It buffers the feature chunks from multiple streams and maintains the buffers.
+    """
+
+    def __init__(
+        self,
+        sample_rate: int,
+        buffer_size_in_secs: float,
+        preprocessor_cfg: DictConfig,
+        device: torch.device,
+    ):
+        """
+        Args:
+            sample_rate (int): sample rate
+            buffer_size_in_secs (float): buffer size in seconds
+            preprocessor_cfg (DictConfig): preprocessor config
+            device (torch.device): device
+        """
+        self.sample_rate = sample_rate
+        self.buffer_size_in_secs = buffer_size_in_secs
+        self.preprocessor_cfg = preprocessor_cfg
+        self.device = device
+        self.bufferers = {}
+
+    def reset(self) -> None:
+        """Reset bufferers"""
+        self.bufferers = {}
+
+    def rm_bufferer(self, stream_id: int) -> None:
+        """
+        Remove bufferer for the given stream id
+        Args:
+            stream_id (int): stream id
+        """
+        self.bufferers.pop(stream_id, None)
+
+    def update(self, fbuffers: list[FeatureBuffer]) -> list[torch.Tensor]:
+        """
+        Update the feature bufferers with the new feature buffers.
+        Feature buffers can come from different streams (audios), so we need to maintain a bufferer for each stream.
+        Args:
+            fbuffers (list[FeatureBuffer]): list of feature buffers
+        Returns:
+            list[torch.Tensor]: list of feature buffers, one per input frame
+        """
+        result_buffers = []
+        for fbuffer in fbuffers:
+            bufferer = self.bufferers.get(fbuffer.stream_id, None)
+
+            if bufferer is None:
+                bufferer = FeatureBufferer(
+                    self.sample_rate,
+                    self.buffer_size_in_secs,
+                    self.preprocessor_cfg,
+                    self.device,
+                )
+                self.bufferers[fbuffer.stream_id] = bufferer
+
+            bufferer.update(fbuffer)
+            result_buffers.append(bufferer.get_feature_buffer())
+
+            if fbuffer.is_last:
+                self.rm_bufferer(fbuffer.stream_id)
+
+        return result_buffers

nemo/collections/asr/inference/streaming/buffering/incremental_audio_bufferer.py

@@ -0,0 +1,173 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+from torch import Tensor
+from nemo.collections.asr.inference.streaming.framing.request import Frame
+
+
+class IncrementalAudioBufferer:
+    """
+    Incremental Audio bufferer class
+    It buffers the audio chunks and maintains the buffer.
+    """
+
+    def __init__(
+        self,
+        sample_rate: int,
+        buffer_size_in_secs: float,
+        chunk_size_in_secs: float,
+        overlap_size_in_secs: float,
+    ) -> None:
+        """
+        Args:
+            sample_rate (int): sample rate
+            buffer_size_in_secs (float): buffer size in seconds
+            chunk_size_in_secs (float): chunk size in seconds
+            overlap_size_in_secs (float): overlap size in seconds
+        """
+        self.sample_rate = sample_rate
+        self.buffer_size = int(buffer_size_in_secs * sample_rate)
+        self.chunk_size = int(chunk_size_in_secs * sample_rate)
+        self.overlap_size = int(overlap_size_in_secs * sample_rate)
+
+        # Ensure overlap is within buffer bounds to keep drop_size non-negative and meaningful.
+        if not (0 <= self.overlap_size <= self.buffer_size):
+            raise ValueError(
+                f"Overlap size in samples ({self.overlap_size}) must satisfy "
+                f"0 <= overlap_size <= buffer_size ({self.buffer_size})."
+            )
+
+        if self.buffer_size % self.chunk_size != 0:
+            raise ValueError(f"Buffer size ({self.buffer_size}) must be divisible by chunk size ({self.chunk_size})")
+
+        if self.overlap_size % self.chunk_size != 0:
+            raise ValueError(f"Overlap size ({self.overlap_size}) must be divisible by chunk size ({self.chunk_size})")
+
+        self.drop_size = self.buffer_size - self.overlap_size
+        self.sample_buffer = torch.zeros(self.buffer_size, dtype=torch.float32)
+        self.remaining_capacity = self.buffer_size
+        self.head = 0
+
+    def is_full(self) -> bool:
+        """
+        Check if the buffer is full
+        Returns:
+            bool: True if the buffer is full, False otherwise
+        """
+        return self.remaining_capacity == 0
+
+    def update(self, frame: Frame) -> None:
+        """
+        Update the buffer with the new frame
+        Args:
+            frame (Frame): frame to update the buffer with
+        """
+        if frame.size > self.buffer_size:
+            raise RuntimeError(f"Frame size ({frame.size}) exceeds buffer size ({self.buffer_size})")
+
+        if self.is_full():
+            # Drop the oldest chunk to make space for the new chunk
+            self.sample_buffer[0 : self.drop_size].zero_()
+            self.sample_buffer = torch.roll(self.sample_buffer, -self.drop_size)
+            self.head -= self.drop_size
+            self.remaining_capacity += self.drop_size
+
+        self.sample_buffer[self.head : self.head + frame.size].copy_(frame.samples)
+        self.head += frame.size
+        self.remaining_capacity = max(0, self.remaining_capacity - frame.size)
+
+
+class BatchedIncrementalAudioBufferer:
+    """
+    Batched incremental audio bufferer class
+    It buffers the audio chunks from multiple streams and returns the buffers.
+    """
+
+    def __init__(
+        self,
+        sample_rate: int,
+        buffer_size_in_secs: float,
+        chunk_size_in_secs: float,
+        overlap_size_in_secs: float,
+    ) -> None:
+        """
+        Args:
+            sample_rate (int): sample rate
+            buffer_size_in_secs (float): buffer size in seconds
+            chunk_size_in_secs (float): chunk size in seconds
+            overlap_size_in_secs (float): overlap size in seconds
+        """
+        self.sample_rate = sample_rate
+        self.buffer_size_in_secs = buffer_size_in_secs
+        self.chunk_size_in_secs = chunk_size_in_secs
+        self.overlap_size_in_secs = overlap_size_in_secs
+        self.bufferers = {}
+
+    def reset(self) -> None:
+        """
+        Reset bufferers
+        """
+        self.bufferers = {}
+
+    def rm_bufferer(self, stream_id: int) -> None:
+        """
+        Remove bufferer for the given stream id
+        Args:
+            stream_id (int): stream id
+        """
+        self.bufferers.pop(stream_id, None)
+
+    def is_full(self, stream_id: int) -> bool | None:
+        """
+        Check if the buffer is full for the given stream id
+        Returns:
+            bool | None: True if the buffer is full, False otherwise
+        """
+        if stream_id not in self.bufferers:
+            return None
+        return self.bufferers[stream_id].is_full()
+
+    def update(self, frames: list[Frame]) -> tuple[list[Tensor], list[int]]:
+        """
+        Update the bufferers with the new frames.
+        Frames can come from different streams (audios), so we need to maintain a bufferer for each stream
+        Args:
+            frames (list[Frame]): list of frames
+        Returns:
+            tuple[list[Tensor], list[int]]:
+                buffers: list of buffered audio tensors, one per input frame
+                paddings: list of paddings, one per input frame
+        """
+        buffers, paddings = [], []
+        for frame in frames:
+            bufferer = self.bufferers.get(frame.stream_id, None)
+
+            if bufferer is None:
+                bufferer = IncrementalAudioBufferer(
+                    self.sample_rate,
+                    self.buffer_size_in_secs,
+                    self.chunk_size_in_secs,
+                    self.overlap_size_in_secs,
+                )
+                self.bufferers[frame.stream_id] = bufferer
+
+            bufferer.update(frame)
+            buffers.append(bufferer.sample_buffer.clone())
+            paddings.append(bufferer.remaining_capacity)
+
+            if frame.is_last:
+                self.rm_bufferer(frame.stream_id)
+
+        return buffers, paddings

nemo/collections/asr/inference/streaming/decoders/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

nemo/collections/asr/inference/streaming/decoders/greedy/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

nemo/collections/asr/inference/streaming/decoders/greedy/greedy_ctc_decoder.py

@@ -0,0 +1,199 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from typing import Callable
+import torch
+
+from nemo.collections.asr.inference.streaming.decoders.greedy.greedy_decoder import GreedyDecoder
+
+
+class CTCGreedyDecoder(GreedyDecoder):
+    """CTC Greedy decoder class"""
+
+    def __init__(self, vocabulary: list[str], conf_func: Callable = None):
+        """
+        Initialize the CTCGreedyDecoder
+        Args:
+            vocabulary (list[str]): list of vocabulary tokens
+            conf_func (Callable): function to compute confidence
+        """
+
+        super().__init__(vocabulary, conf_func)
+
+    @staticmethod
+    def get_labels(log_probs: torch.Tensor) -> list[int]:
+        """
+        Perform greedy decoding on the log probabilities
+        Args:
+            log_probs (torch.Tensor): log probabilities
+        Returns:
+            list[int]: list of tokens
+        """
+        if log_probs.dim() != 2:
+            raise ValueError("log_probs must be 2D tensor")
+
+        labels = log_probs.argmax(dim=-1).cpu()  # T
+        return labels.tolist()
+
+    def __call__(self, log_probs: torch.Tensor, compute_confidence: bool = True, previous: int = None) -> dict:
+        """
+        Greedy decode the log probabilities
+        Args:
+            log_probs (torch.Tensor): log probabilities
+            compute_confidence (bool): compute confidence or not
+        Returns:
+            dict: output dictionary containing tokens, timesteps, and confidences
+        """
+
+        compute_confidence = compute_confidence and self.conf_func is not None
+
+        if log_probs.dim() != 2:
+            raise ValueError("log_probs must be 2D tensor")
+
+        if compute_confidence:
+            # Add batch dimension
+            log_probs = log_probs.unsqueeze(0)  # 1 x T x N
+            # Compute confidences
+            confidences = torch.zeros(log_probs.shape[0], log_probs.shape[1])  # 1 x T
+            confidences[0] = self.conf_func(log_probs[0], v=log_probs.shape[2])  # 1 x T
+            # Remove batch dimension and convert to list
+            confidences = confidences.squeeze(0).tolist()  # T
+            # Remove batch dimension
+            log_probs = log_probs.squeeze(0)  # T x N
+
+        labels = self.get_labels(log_probs)  # T
+        output = {"tokens": [], "timesteps": [], "confidences": []}
+        previous = self.blank_id if previous is None else previous
+        for i, p in enumerate(labels):
+            if p != previous and p != self.blank_id:
+                output["tokens"].append(p)
+                output["timesteps"].append(i)
+                if compute_confidence:
+                    output["confidences"].append(confidences[i])
+            previous = p
+
+        output["labels"] = labels
+        return output
+
+
+class ClippedCTCGreedyDecoder:
+    """
+    Clipped CTC Greedy decoder class
+    Decodes the tokens within a given clip range and returns the clipped tokens and timestamps.
+    """
+
+    def __init__(self, vocabulary: list[str], tokens_per_frame: int, conf_func: Callable = None, endpointer=None):
+        """
+        Initialize the ClippedCTCGreedyDecoder
+        Args:
+            vocabulary (list[str]): list of vocabulary tokens
+            tokens_per_frame (int): number of tokens per frame
+            conf_func (Callable): function to compute confidence
+            endpointer (Any): endpointer to detect EOU
+        """
+        self.greedy_decoder = CTCGreedyDecoder(vocabulary, conf_func)
+        self.endpointer = endpointer
+        self.tokens_per_frame = tokens_per_frame
+
+    def __call__(
+        self,
+        log_probs: torch.Tensor,
+        clip_start: int,
+        clip_end: int,
+        is_last: bool = False,
+        is_start: bool = True,
+        return_partial_result: bool = True,
+        state_start_idx: int = 0,
+        state_end_idx: int = 0,
+        stop_history_eou: int = None,
+        compute_confidence: bool = True,
+    ) -> tuple[dict, dict, bool, int, int]:
+        """
+        Decode the log probabilities within the clip range (clip_start, clip_end)
+        Args:
+            log_probs (torch.Tensor): log probabilities
+            clip_start (int): start index of the clip
+            clip_end (int): end index of the clip
+            is_last (bool): is the last frame or not
+            is_start (bool): is the first frame for this stream or not
+            return_partial_result (bool): return partial result left after clip_end in the buffer
+            state_start_idx (int): start index from stream state
+            state_end_idx (int): end index from stream state
+            stop_history_eou (int): stop history of EOU, if None then use the default stop history
+            compute_confidence (bool): compute confidence or not
+        Returns:
+            tuple[dict, dict, bool, int, int]:
+                clipped output, tail output, is_eou, updated start_idx, updated end_idx
+        """
+
+        is_eou = is_last
+        eou_detected_at = len(log_probs)
+        # Initialize state tracking variables from input parameters
+        start_idx, end_idx = state_start_idx, state_end_idx
+        # Update indices for next processing step
+        if end_idx > clip_start:
+            end_idx -= self.tokens_per_frame
+            start_idx = end_idx
+
+        if is_start or end_idx <= clip_start:
+            start_idx, end_idx = clip_start, clip_end
+
+        all_output = self.greedy_decoder(log_probs, compute_confidence=compute_confidence)
+
+        clipped_output = {"tokens": [], "timesteps": [], "confidences": [], "last_token": None, "last_token_idx": None}
+        tail_output = {"tokens": []}
+
+        # check if EOU is detected or is the last frame
+        if not is_eou and self.endpointer is not None:
+            is_eou, eou_detected_at = self.endpointer.detect_eou(
+                log_probs, pivot_point=start_idx, search_start_point=clip_start, stop_history_eou=stop_history_eou
+            )
+
+        # if EOU is detected, and it is after the clip end, update the end index to the EOU
+        if is_eou and eou_detected_at > end_idx:
+            end_idx = eou_detected_at
+
+        # if the end index is within the clip range, update the end index to the clip end
+        if clip_start <= end_idx < clip_end:
+            end_idx = clip_end
+            is_eou = False
+
+        # clip the output within the clip range [clip_start, clip_end)
+        timesteps = all_output["timesteps"]
+        i = 0
+        while i < len(timesteps):
+            if start_idx <= timesteps[i] < end_idx:
+                clipped_output["tokens"].append(all_output["tokens"][i])
+                clipped_output["timesteps"].append(timesteps[i])
+                if compute_confidence:
+                    clipped_output["confidences"].append(all_output["confidences"][i])
+            elif timesteps[i] >= end_idx:
+                break
+            i += 1
+
+        if end_idx - 1 < len(all_output["labels"]):
+            clipped_output["last_token"] = all_output["labels"][end_idx - 1]
+            clipped_output["last_token_idx"] = end_idx - 1
+
+        # return the partial result left after clip_end in the buffer
+        if return_partial_result:
+            while i < len(timesteps):
+                if timesteps[i] >= end_idx:
+                    tail_output["tokens"] = all_output["tokens"][i:]
+                    break
+                else:
+                    i += 1
+
+        return clipped_output, tail_output, is_eou, start_idx, end_idx

nemo/collections/asr/inference/streaming/decoders/greedy/greedy_decoder.py

@@ -0,0 +1,102 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from typing import Callable
+
+from nemo.collections.asr.inference.utils.constants import SENTENCEPIECE_UNDERSCORE
+
+
+class GreedyDecoder:
+    """Base class for the greedy decoder"""
+
+    def __init__(self, vocabulary: list[str], conf_func: Callable = None):
+        """
+        Initialize the GreedyDecoder
+        Args:
+            vocabulary (list[str]): list of vocabulary tokens
+            conf_func (Callable): function to compute confidence
+        """
+
+        self.vocabulary = vocabulary
+        self.blank_id = len(vocabulary)
+        self.conf_func = conf_func
+        self.is_start_tokens = [token.startswith(SENTENCEPIECE_UNDERSCORE) for token in vocabulary]
+
+    def count_silent_tokens(self, tokens: list[int], start: int, end: int) -> int:
+        """
+        Count how many silent tokens appear in [start, end).
+        Args:
+            tokens (list[int]): list of tokens
+            start (int): start index
+            end (int): end index
+        Returns:
+            int: number of silent tokens
+        """
+        if end <= start or start >= len(tokens):
+            return 0
+        return sum(self.is_token_silent(tokens[i]) for i in range(start, min(end, len(tokens))))
+
+    def is_token_start_of_word(self, token_id: int) -> bool:
+        """
+        Check if the token is the start of a word
+        Args:
+            token_id (int): token id
+        Returns:
+            bool: True if the token is the start of a word, False otherwise
+        """
+        return self.is_start_tokens[token_id]
+
+    def is_token_silent(self, token_id: int) -> bool:
+        """
+        Check if the token is silent
+        Args:
+            token_id (int): token id
+        Returns:
+            bool: True if the token is silent, False otherwise
+        """
+        return token_id == self.blank_id
+
+    def first_non_silent_token(self, tokens: list[int], start: int, end: int) -> int:
+        """
+        Return the index of the first non-silent token in [start, end).
+        If none found, return -1.
+        Args:
+            tokens (list[int]): list of tokens
+            start (int): start index
+            end (int): end index
+        Returns:
+            int: index of the first non-silent token
+        """
+        for i in range(start, min(end, len(tokens))):
+            if not self.is_token_silent(tokens[i]):
+                return i
+        return -1
+
+    def count_non_silent_tokens(self, tokens: list[int], start: int, end: int) -> int:
+        """
+        Count how many non-silent tokens appear in [start, end).
+        Args:
+            tokens (list[int]): list of tokens
+            start (int): start index
+            end (int): end index
+        Returns:
+            int: number of non-silent tokens
+        """
+        if end <= start or start >= len(tokens):
+            return 0
+        return sum(not self.is_token_silent(tokens[i]) for i in range(start, min(end, len(tokens))))
+
+    def __call__(self, *args, **kwds):
+        raise NotImplementedError("Subclass of GreedyDecoder should implement `__call__` method!")