chore: import upstream snapshot with attribution

nemo/collections/speechlm2/data/__init__.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 .datamodule import DataModule
+from .duplex_ear_tts_dataset import DuplexEARTTSDataset
+from .duplex_stt_dataset import DuplexSTTDataset
+from .s2s_dataset import DuplexS2SDataset
+from .salm_dataset import SALMDataset
+
+__all__ = [
+    'DataModule',
+    'DuplexS2SDataset',
+    'DuplexSTTDataset',
+    'DuplexEARTTSDataset',
+    'SALMDataset',
+]

nemo/collections/speechlm2/data/datamodule.py

@@ -0,0 +1,213 @@
+# 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 lightning import LightningDataModule
+from lightning.pytorch.utilities import CombinedLoader
+from omegaconf import DictConfig, OmegaConf, open_dict
+
+from nemo.collections.common.data.fallback import FallbackDataset
+from nemo.collections.common.data.lhotse import get_lhotse_dataloader_from_config
+from nemo.collections.common.data.lhotse.broadcasting import BroadcastingDataLoader, is_dp_source_rank
+from nemo.collections.common.tokenizers import TokenizerSpec
+
+
+class DataModule(LightningDataModule):
+    """
+    A Lightning DataModule specialized for Lhotse dataloading.
+    It takes care of setting up the proper DP ranks for dataloaders, and instantiating them.
+    Keep in mind the actual dataset paths and blend are defined by the YAML config, not Python code.
+
+    The typical structure of the YAML config used to initialize this module looks like the following:
+
+    .. code-block:: yaml
+
+        data:
+          train_ds:
+            input_cfg: path/to/input_cfg.yaml
+            num_workers: 2
+            batch_size: 4
+            # ... Other settings, see nemo/collections/common/data/lhotse/dataloader.py
+
+          validation_ds:
+            # The entries under 'datasets' are a list of separate dataloaders.
+            # The structure is <dataset-name>: {<dataloader-dict-config>}
+            # They inherit all settings from validation_ds, but can individually override them.
+            datasets:
+              val_set_0:  # rename to your dataset name, add more as needed
+                cuts_path: ???  # needs to be specified
+            batch_size: 4
+            # ... Other settings, see nemo/collections/common/data/lhotse/dataloader.py
+
+    See also the examples in ``examples/speechlm2/conf``.
+
+    Args:
+        cfg: a DictConfig instance, typically corresponding to `data` namespace in YAML configs.
+        tokenizer: a tokenizer instance, typically NeMo's AutoTokenizer wrapping HF's AutoTokenizer.
+        dataset: a torch.utils.data.Dataset instance, expected to define __getitem__ that accepts
+            a lhotse.CutSet. It converts metadata + raw data to a batch of PyTorch tensors.
+            The data sampling is controlled by Lhotse samplers rather than the dataset.
+    """
+
+    def __init__(self, cfg, tokenizer: TokenizerSpec, dataset: torch.utils.data.Dataset) -> None:
+        super().__init__()
+        self.cfg = cfg
+        with open_dict(self.cfg):
+            for k in ("validation_ds", "test_ds"):
+                if k in self.cfg:
+                    getattr(self.cfg, k).force_finite = True
+                    getattr(self.cfg, k).force_map_dataset = True
+        self.tokenizer = tokenizer
+        self.dataset = dataset
+
+    def train_dataloader(self):
+        if "train_ds" not in self.cfg:
+            return None
+        mesh = self._get_device_mesh()
+        if is_dp_source_rank(mesh):
+            source = get_lhotse_dataloader_from_config(
+                config=self.cfg.train_ds,
+                global_rank=self._get_dp_rank(),
+                world_size=self._get_world_size(),
+                dataset=FallbackDataset(self.dataset),
+                tokenizer=self.tokenizer,
+            )
+        else:
+            source = None
+        return BroadcastingDataLoader(source=source, device_mesh=mesh)
+
+    def val_dataloader(self):
+        if "validation_ds" not in self.cfg:
+            return None
+        cfg = self.cfg.validation_ds
+        return self._build_test_dataloader(cfg)
+
+    def test_dataloader(self):
+        if "test_ds" not in self.cfg:
+            return None
+        cfg = self.cfg.test_ds
+        return self._build_test_dataloader(cfg)
+
+    def predict_dataloader(self):
+        if "predict_ds" not in self.cfg:
+            return None
+        cfg = self.cfg.predict_ds
+
+        base_cfg = cfg.copy()
+        with open_dict(base_cfg):
+            del base_cfg.datasets
+        dloaders = {}
+        for name, item in cfg.datasets.items():
+            with open_dict(base_cfg):
+                item = OmegaConf.merge(base_cfg, item)
+            dloaders[name] = self._build_test_dataloader(item)
+        # NOTE(yifan): `trainer.predict()` only supports the `CombinedLoader(mode="sequential")` mode
+        # so we cannot reuse the `_build_test_dataloader` function here
+        return CombinedLoader(dloaders, mode="sequential")
+
+    def _build_test_dataloader(self, cfg: DictConfig) -> torch.utils.data.DataLoader | CombinedLoader:
+        # Single validation/test dataloader.
+        # This is internal-only: the config has to specify multiple dataloaders via "datasets" key,
+        # even for a single validation/test set.
+        if "datasets" not in cfg:
+            with open_dict(cfg):
+                cfg.force_finite = True
+                cfg.force_map_dataset = True
+            mesh = self._get_device_mesh()
+            if is_dp_source_rank(mesh):
+                source = get_lhotse_dataloader_from_config(
+                    config=cfg,
+                    global_rank=self._get_dp_rank(),
+                    world_size=self._get_world_size(),
+                    dataset=self.dataset,
+                    tokenizer=self.tokenizer,
+                )
+            else:
+                source = None
+            return BroadcastingDataLoader(source=source, device_mesh=mesh)
+
+        # Multiple validation/test dataloaders.
+        # Config looks like:
+        #
+        # validation_ds:
+        #   batch_size: ...
+        #   datasets:
+        #     easy_benchmark:
+        #       shar_path: ...
+        #     hard_benchmark:
+        #       shar_path: ...
+        base_cfg = cfg.copy()
+        with open_dict(base_cfg):
+            del base_cfg.datasets
+        dloaders = {}
+        for name, item in cfg.datasets.items():
+            with open_dict(base_cfg):
+                item = OmegaConf.merge(base_cfg, item)
+            dloaders[name] = self._build_test_dataloader(item)
+        return CombinedLoader(dloaders, mode="max_size")
+
+    def _get_device_mesh(self):
+        if not (torch.distributed.is_available() and torch.distributed.is_initialized()):
+            return None
+        if hasattr(self.trainer, "model") and hasattr(self.trainer.model, "device_mesh"):
+            return self.trainer.model.device_mesh
+        return None
+
+    def _get_dp_rank(self):
+        if torch.distributed.is_available() and torch.distributed.is_initialized():
+            if (
+                hasattr(self.trainer, "model")
+                and hasattr(self.trainer.model, "device_mesh")
+                and (dm := self.trainer.model.device_mesh) is not None
+            ):  # model parallelism
+                if "data_parallel" in dm.mesh_dim_names:  # Lightning's built-in ModelParallelStrategy
+                    dp_rank = dm["data_parallel"].get_local_rank()
+                elif (
+                    "dp_shard" in dm.mesh_dim_names and "dp_replicate" in dm.mesh_dim_names
+                ):  # AutomodelParallelStrategy
+                    try:
+                        dp_rank = dm["dp"].get_local_rank()
+                    except (KeyError, RuntimeError, ValueError):
+                        # Compatibility for older Automodel/PyTorch meshes without a flattened "dp" submesh.
+                        dp_rank = (
+                            dm["dp_replicate"].get_local_rank() * dm["dp_shard"].size()
+                            + dm["dp_shard"].get_local_rank()
+                        )
+                return dp_rank
+            else:
+                return torch.distributed.get_rank()  # plain ol' DDP
+        else:
+            return 0  # 1 GPU
+
+    def _get_world_size(self):
+        if torch.distributed.is_available() and torch.distributed.is_initialized():
+            if (
+                hasattr(self.trainer, "model")
+                and hasattr(self.trainer.model, "device_mesh")
+                and (dm := self.trainer.model.device_mesh) is not None
+            ):  # model parallelism
+                if "data_parallel" in dm.mesh_dim_names:  # Lightning's built-in ModelParallelStrategy
+                    dp_size = dm["data_parallel"].size()
+                elif (
+                    "dp_shard" in dm.mesh_dim_names and "dp_replicate" in dm.mesh_dim_names
+                ):  # AutomodelParallelStrategy
+                    try:
+                        dp_size = dm["dp"].size()
+                    except (KeyError, RuntimeError, ValueError):
+                        # Compatibility for older Automodel/PyTorch meshes without a flattened "dp" submesh.
+                        dp_size = dm["dp_replicate", "dp_shard"].size()
+                return dp_size
+            else:  # plain ol' DDP
+                return torch.distributed.get_world_size()
+        else:
+            return 1  # 1 GPU

nemo/collections/speechlm2/data/duplex_stt_dataset.py

@@ -0,0 +1,733 @@
+# 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 random
+import re
+
+import torch
+import torch.utils.data
+from lhotse import CutSet, MonoCut, Recording, Seconds, SupervisionSegment, compute_num_frames
+from lhotse.cut import Cut
+from lhotse.dataset.collation import collate_audio, collate_vectors
+from lhotse.utils import ifnone
+
+from nemo.collections.common.data.lhotse.text_adapters import Formattable
+from nemo.collections.common.tokenizers import TokenizerSpec
+from nemo.collections.speechlm2.data.force_align import ForceAligner
+from nemo.collections.speechlm2.data.s2s_dataset import _strip_timestamps
+from nemo.collections.speechlm2.data.utils import get_pad_id
+from nemo.collections.speechlm2.parts.augmentation import AudioAugmenter
+from nemo.utils import logging
+
+MCQ_VAL_PROMPT = "Answer the following multiple choice question with an explanation for the answer."
+
+
+class DuplexSTTDataset(torch.utils.data.Dataset):
+    """
+    A dataset for duplex speech-to-text models.
+
+    Unlike DuplexS2SDataset, this dataset does not require target audio and is suitable
+    for training on standard ASR, AST, and SpeechQA datasets in addition to duplex data.
+
+    Args:
+        tokenizer (TokenizerSpec):
+            Tokenizer for converting text to token IDs. Must support BOS and EOS tokens.
+
+        frame_length (Seconds):
+            Duration of a single frame in seconds.
+
+        source_sample_rate (int):
+            Sample rate for source audio (e.g., 16000 Hz).
+
+        input_roles (list[str], optional):
+            Speaker roles to treat as inputs. Defaults to ["user"].
+
+        output_roles (list[str], optional):
+            Speaker roles to treat as outputs. Defaults to ["agent"].
+
+        aug_by_swap_role (bool, optional):
+            Whether to augment data by swapping user/agent roles. Defaults to False.
+            Note: Enabling this requires agent audio to be available in cut.custom['target_audio'].
+
+        cfg (dict, optional):
+            Dataset configuration (e.g., word_align_position).
+
+        model_cfg (dict, optional):
+            Model configuration (e.g., predict_user_text, force_align_user_text).
+    """
+
+    def __init__(
+        self,
+        tokenizer: TokenizerSpec,
+        frame_length: Seconds,
+        source_sample_rate: int,
+        input_roles: list[str] = None,
+        output_roles: list[str] = None,
+        aug_by_swap_role: bool = False,
+        cfg: dict = None,
+        model_cfg: dict = None,
+    ):
+        self.tokenizer = tokenizer
+        self.frame_length = frame_length
+        self.source_sample_rate = source_sample_rate
+        self.input_roles = set(ifnone(input_roles, ["user"]))
+        self.output_roles = set(ifnone(output_roles, ["agent"]))
+        self.aug_by_swap_role = aug_by_swap_role
+
+        self.word_align_position = cfg.get("word_align_position", "left") if cfg is not None else "left"
+        self.predict_user_text = model_cfg.get("predict_user_text", False) if model_cfg is not None else False
+        self.force_align_user_text = model_cfg.get("force_align_user_text", False) if model_cfg is not None else None
+        self.force_align_device = model_cfg.get("force_align_device", "cuda") if model_cfg is not None else "cuda"
+
+        self.prepend_word_space = cfg.get("prepend_word_space", True) if cfg is not None else True
+        self.early_interruption_prob = cfg.get("early_interruption_prob", 0.0) if cfg is not None else 0.0
+        self.add_mcq_val_prompt = cfg.get("add_mcq_val_prompt", False) if cfg is not None else False
+
+        self.cfg = cfg
+        self.model_cfg = model_cfg
+
+        self.force_aligner = None
+        if self.force_align_user_text:
+            self.force_aligner = ForceAligner(device=self.force_align_device, frame_length=self.frame_length)
+
+        self.audio_augmenter = None
+        if cfg is not None and (
+            cfg.get('use_noise_aug', None)
+            or cfg.get('use_room_ir_aug', None)
+            or cfg.get('use_mic_ir_aug', None)
+            or cfg.get('use_codec_aug', None)
+        ):
+            self.audio_augmenter = AudioAugmenter(sample_rate=source_sample_rate)
+
+        assert tokenizer.bos is not None, "BOS support in the tokenizer is required."
+        assert tokenizer.eos is not None, "EOS support in the tokenizer is required."
+
+    def _is_augmentation_task(self, task: str) -> bool:
+        if self.cfg is not None and self.cfg.get('force_use_noise_augmentation', False):
+            return True
+        return task not in ('s2s_duplex_overlap_as_s2s_duplex', 'asr')
+
+    def _apply_early_interruption_augmentation(
+        self,
+        target_tokens: torch.Tensor,
+        source_tokens: torch.Tensor,
+        source_audio: torch.Tensor,
+        source_audio_lens: torch.Tensor,
+        batch_idx: int,
+    ) -> None:
+        """Simulate early interruption by randomly truncating an agent turn with overlap."""
+        target_seq = target_tokens[batch_idx]
+        bos_id = self.tokenizer.bos
+        eos_id = self.tokenizer.eos
+        pad_id = get_pad_id(self.tokenizer)
+
+        overlap_tokens = self.cfg.get("early_interruption_overlap_tokens", 13) if self.cfg is not None else 13
+
+        bos_positions = (target_seq == bos_id).nonzero(as_tuple=True)[0]
+        eos_positions = (target_seq == eos_id).nonzero(as_tuple=True)[0]
+
+        if len(bos_positions) == 0 or len(eos_positions) == 0:
+            return
+
+        turns = []
+        for bos_pos in bos_positions:
+            matching_eos = eos_positions[eos_positions > bos_pos]
+            if len(matching_eos) > 0:
+                eos_pos = matching_eos[0]
+                turn_tokens = target_seq[bos_pos + 1 : eos_pos]
+                non_pad_mask = turn_tokens != pad_id
+                all_non_pad_positions = (bos_pos + 1 + non_pad_mask.nonzero(as_tuple=True)[0]).tolist()
+                non_pad_positions = [pos for pos in all_non_pad_positions if (eos_pos - pos) > overlap_tokens]
+
+                if len(non_pad_positions) > 0:
+                    turns.append(
+                        {'bos_pos': bos_pos.item(), 'eos_pos': eos_pos.item(), 'non_pad_positions': non_pad_positions}
+                    )
+
+        if len(turns) == 0:
+            return
+
+        selected_turn = random.choice(turns)
+        cutoff_pos = random.choice(selected_turn['non_pad_positions'])
+        original_eos_pos = selected_turn['eos_pos']
+
+        new_eos_pos = min(cutoff_pos + overlap_tokens, original_eos_pos)
+        frames_to_remove = original_eos_pos - cutoff_pos
+        if frames_to_remove <= 0:
+            return
+
+        # Update target_tokens: place eos at new_eos_pos, shift tail, pad at end
+        target_tokens[batch_idx, new_eos_pos] = eos_id
+        seq_len = target_tokens.shape[1]
+        cont_start_pos = original_eos_pos + overlap_tokens
+        tail_length = seq_len - (cont_start_pos + 1)
+        if tail_length > 0:
+            target_tokens[batch_idx, new_eos_pos + 1 : new_eos_pos + 1 + tail_length] = target_tokens[
+                batch_idx, cont_start_pos + 1 : cont_start_pos + 1 + tail_length
+            ].clone()
+        target_tokens[batch_idx, -frames_to_remove:] = pad_id
+
+        # Update source_tokens: shift tail (from cutoff_pos)
+        src_frames_to_remove = original_eos_pos - cutoff_pos
+        source_seq_len = source_tokens.shape[1]
+        source_tail_length = source_seq_len - (original_eos_pos + 1)
+        if source_tail_length > 0:
+            source_tokens[batch_idx, cutoff_pos + 1 : cutoff_pos + 1 + source_tail_length] = source_tokens[
+                batch_idx, original_eos_pos + 1 : original_eos_pos + 1 + source_tail_length
+            ].clone()
+        source_tokens[batch_idx, -src_frames_to_remove:] = pad_id
+
+        # Update source_audio: shift and pad with silence
+        old_source_len = source_audio_lens[batch_idx].item()
+        new_bos_source_sample = min(int(cutoff_pos * self.frame_length * self.source_sample_rate), old_source_len)
+        original_eos_source_sample = min(
+            int(original_eos_pos * self.frame_length * self.source_sample_rate), old_source_len
+        )
+
+        source_tail_audio_length = old_source_len - original_eos_source_sample
+        if source_tail_audio_length > 0:
+            source_audio[batch_idx, new_bos_source_sample : new_bos_source_sample + source_tail_audio_length] = (
+                source_audio[batch_idx, original_eos_source_sample:old_source_len].clone()
+            )
+
+        source_samples_to_remove = original_eos_source_sample - new_bos_source_sample
+        if new_bos_source_sample + source_tail_audio_length < source_audio.shape[1]:
+            source_audio[
+                batch_idx,
+                new_bos_source_sample
+                + source_tail_audio_length : new_bos_source_sample
+                + source_tail_audio_length
+                + source_samples_to_remove,
+            ] = 0
+
+    def _create_minimal_batch(self) -> dict:
+        """Create a minimal valid batch when all cuts are filtered out."""
+        return {
+            "sample_id": ["empty_batch"],
+            "source_audio": torch.zeros((1, 1000), dtype=torch.float32),
+            "source_audio_lens": torch.tensor([1000], dtype=torch.long),
+            "target_tokens": torch.full((1, 50), get_pad_id(self.tokenizer), dtype=torch.long),
+            "target_token_lens": torch.tensor([1], dtype=torch.long),
+            "source_tokens": torch.full((1, 50), get_pad_id(self.tokenizer), dtype=torch.long),
+            "source_token_lens": torch.tensor([1], dtype=torch.long),
+            "source_texts": [""],
+            "target_texts": [""],
+            "task": ["s2s_duplex"],
+        }
+
+    def __getitem__(self, all_cuts: CutSet) -> dict:
+        cuts = all_cuts.filter(lambda c: isinstance(c, Cut))
+        audio_data = None
+
+        if cuts and getattr(cuts[0], 'task', None) == 'asr':
+            filtered_cuts = []
+            skipped_cuts = []
+            for cut in cuts:
+                if self._has_valid_input(cut):
+                    filtered_cuts.append(cut)
+                else:
+                    skipped_cuts.append(cut.id)
+            if skipped_cuts:
+                logging.info(
+                    f"Skipped {len(skipped_cuts)} cuts with empty input text. Skipped cut ids: {', '.join(skipped_cuts)}"
+                )
+            if not filtered_cuts:
+                logging.warning(
+                    f"All cuts were filtered out! Original batch size: {len(cuts)}. Returning minimal valid batch."
+                )
+                return self._create_minimal_batch()
+            cuts = CutSet.from_cuts(filtered_cuts)
+
+        if cuts:
+            swapped_cuts = []
+
+            if self.aug_by_swap_role:
+                for cut in cuts:
+                    total_turns = cut.custom.get('total_turns', len(cut.supervisions))
+
+                    if total_turns > 4 and total_turns % 2 == 0:
+                        swapped_cut = self._create_role_swapped_cut(cut)
+                        if swapped_cut:
+                            swapped_cuts.append(swapped_cut)
+
+            if swapped_cuts:
+                all_cuts_combined = CutSet.from_cuts(list(cuts) + swapped_cuts)
+            else:
+                all_cuts_combined = cuts
+
+            prompt_tokens, prompt_token_lens = collate_system_prompt(
+                all_cuts_combined, self.tokenizer, add_mcq_val_prompt=self.add_mcq_val_prompt
+            )
+            source_audio, source_audio_lens = collate_audio(all_cuts_combined.resample(self.source_sample_rate))
+
+            target_tokens, target_token_lens = collate_token_channel(
+                all_cuts_combined,
+                self.tokenizer,
+                self.frame_length,
+                roles=self.output_roles,
+                bos_id=self.tokenizer.bos,
+                eos_id=self.tokenizer.eos,
+                remove_timestamps=True,
+            )
+
+            # Force align user text (runs in dataloader worker, overlapped with training)
+            if self.force_align_user_text and torch.is_grad_enabled():
+                all_cuts_combined = self.force_aligner.batch_force_align_user_audio(
+                    all_cuts_combined, source_sample_rate=self.source_sample_rate
+                )
+
+            source_tokens, source_token_lens = collate_token_channel(
+                all_cuts_combined,
+                self.tokenizer,
+                self.frame_length,
+                roles=self.input_roles,
+                bos_id=self.tokenizer.bos,
+                eos_id=self.tokenizer.eos,
+                word_align_position=self.word_align_position,
+                remove_timestamps=not self.predict_user_text,
+                prepend_word_space=self.prepend_word_space,
+            )
+
+            # Audio augmentation (runs in dataloader workers for performance)
+            if (
+                self.audio_augmenter is not None
+                and torch.is_grad_enabled()
+                and self._is_augmentation_task(getattr(all_cuts_combined[0], 'task', 's2s_duplex'))
+            ):
+                source_audio = self.audio_augmenter.augment_batch(self.cfg, source_audio, source_audio_lens)
+
+            # Early interruption augmentation
+            if self.early_interruption_prob > 0 and torch.is_grad_enabled():
+                for batch_idx in range(target_tokens.shape[0]):
+                    if random.random() < self.early_interruption_prob:
+                        self._apply_early_interruption_augmentation(
+                            target_tokens,
+                            source_tokens,
+                            source_audio,
+                            source_audio_lens,
+                            batch_idx,
+                        )
+
+            audio_data = {
+                "sample_id": [str(cut.id) for cut in all_cuts_combined],
+                "source_audio": source_audio,
+                "source_audio_lens": source_audio_lens,
+                "target_tokens": target_tokens,
+                "target_token_lens": target_token_lens,
+                "source_tokens": source_tokens,
+                "source_token_lens": source_token_lens,
+                "source_texts": [
+                    " ".join(_strip_timestamps(s.text) for s in cut.supervisions if s.speaker in self.input_roles)
+                    for cut in all_cuts_combined
+                ],
+                "target_texts": [
+                    " ".join(s.text for s in cut.supervisions if s.speaker in self.output_roles)
+                    for cut in all_cuts_combined
+                ],
+                "task": [getattr(cut, "task", "s2s_duplex") for cut in all_cuts_combined],
+            }
+
+            if torch.sum(prompt_token_lens) > 0:
+                audio_data['prompt_tokens'] = prompt_tokens
+                audio_data['prompt_token_lens'] = prompt_token_lens
+
+        text_cuts = all_cuts.filter(lambda c: isinstance(c, Formattable))
+        text_data = None
+        if text_cuts:
+            text_tokens = []
+            text_token_lens = []
+            for c in text_cuts:
+                text_ids = c.input_ids
+                text_tokens.append(text_ids)
+                text_token_lens.append(text_ids.shape[0])
+
+            text_tokens = collate_vectors(text_tokens, padding_value=get_pad_id(self.tokenizer))
+            text_token_lens = torch.tensor(text_token_lens, dtype=torch.long)
+            text_data = {
+                "text_tokens": text_tokens,
+                "text_token_lens": text_token_lens,
+            }
+
+        return {
+            "audio_data": audio_data,
+            "text_data": text_data,
+        }
+
+    def _create_role_swapped_cut(self, cut):
+        from io import BytesIO
+
+        import numpy as np
+        import soundfile as sf
+        from lhotse import AudioSource
+
+        assert (
+            'target_audio' in cut.custom
+        ), f"Role swapping requires target_audio in cut.custom, but cut {cut.id} does not have it. Disable aug_by_swap_role or ensure your data includes target audio."
+
+        swapped_supervisions = []
+        for sup in cut.supervisions:
+            if sup.speaker == 'User':
+                new_speaker = 'Assistant'
+            elif sup.speaker == 'Assistant':
+                new_speaker = 'User'
+            else:
+                continue
+
+            swapped_sup = SupervisionSegment(
+                id=sup.id + "_swapped",
+                recording_id=sup.recording_id,
+                start=sup.start,
+                duration=sup.duration,
+                channel=sup.channel,
+                text=sup.text,
+                language=sup.language,
+                speaker=new_speaker,
+                gender=sup.gender,
+                custom=sup.custom,
+                alignment=sup.alignment,
+            )
+            swapped_supervisions.append(swapped_sup)
+
+        swapped_supervisions = sorted(swapped_supervisions, key=lambda s: s.start)
+
+        first_agent_idx = None
+        last_user_idx = None
+
+        for i, sup in enumerate(swapped_supervisions):
+            if sup.speaker == 'Assistant' and first_agent_idx is None:
+                first_agent_idx = i
+            if sup.speaker == 'User':
+                last_user_idx = i
+
+        filtered_supervisions = []
+        for i, sup in enumerate(swapped_supervisions):
+            if i != first_agent_idx and i != last_user_idx:
+                filtered_supervisions.append(sup)
+
+        if not filtered_supervisions:
+            return None
+
+        first_remaining_start = filtered_supervisions[0].start
+        adjusted_supervisions = []
+        for sup in filtered_supervisions:
+            adjusted_sup = SupervisionSegment(
+                id=sup.id,
+                recording_id=sup.recording_id,
+                start=sup.start - first_remaining_start,
+                duration=sup.duration,
+                channel=sup.channel,
+                text=sup.text,
+                language=sup.language,
+                speaker=sup.speaker,
+                gender=sup.gender,
+                custom=sup.custom,
+                alignment=sup.alignment,
+            )
+            adjusted_supervisions.append(adjusted_sup)
+
+        total_duration = max(s.start + s.duration for s in adjusted_supervisions)
+        total_samples = int(total_duration * cut.sampling_rate)
+
+        new_source_audio = np.zeros(total_samples, dtype=np.float32)
+
+        for sup in adjusted_supervisions:
+            start_sample = int(sup.start * cut.sampling_rate)
+            end_sample = int((sup.start + sup.duration) * cut.sampling_rate)
+
+            if sup.speaker == 'User':
+                # New User was originally Assistant — audio lives in target_audio channel
+                original_start = sup.start + first_remaining_start
+                agent_audio = (
+                    cut.custom['target_audio']
+                    .to_cut()
+                    .truncate(offset=original_start, duration=sup.duration)
+                    .load_audio()
+                )
+                if len(agent_audio.shape) > 1:
+                    agent_audio = agent_audio.squeeze()
+                actual_end = min(end_sample, start_sample + len(agent_audio))
+                new_source_audio[start_sample:actual_end] = agent_audio[: actual_end - start_sample]
+
+        source_buffer = BytesIO()
+        sf.write(source_buffer, new_source_audio, cut.sampling_rate, format='wav')
+        source_buffer.seek(0)
+
+        new_source_recording = Recording(
+            id=f"{cut.id}_swapped_source",
+            sampling_rate=cut.sampling_rate,
+            num_samples=len(new_source_audio),
+            duration=total_duration,
+            sources=[AudioSource(type="memory", channels=[0], source=source_buffer.getvalue())],
+        )
+
+        swapped_cut = MonoCut(
+            id=f"{cut.id}_swapped",
+            start=0,
+            duration=total_duration,
+            channel=0,
+            supervisions=adjusted_supervisions,
+            recording=new_source_recording,
+            custom={
+                **cut.custom,
+                'total_turns': len(adjusted_supervisions),
+                'role_swapped': True,
+            },
+        )
+
+        return swapped_cut
+
+    def _has_valid_input(self, cut: Cut) -> bool:
+        return any(s.text.strip() for s in cut.supervisions if s.speaker in self.input_roles)
+
+
+def collate_token_channel(
+    cuts: CutSet,
+    tokenizer: TokenizerSpec,
+    frame_length: Seconds,
+    roles: set[str],
+    bos_id: int = None,
+    eos_id: int = None,
+    word_align_position: str = 'left',
+    remove_timestamps: bool = False,
+    prepend_word_space: bool = True,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    pad_id = get_pad_id(tokenizer)
+    tokens = [
+        _build_token_channel(
+            c,
+            tokenizer=tokenizer,
+            frame_length=frame_length,
+            roles=roles,
+            pad_id=pad_id,
+            bos_id=bos_id,
+            eos_id=eos_id,
+            word_align_position=word_align_position,
+            remove_timestamps=remove_timestamps,
+            prepend_word_space=prepend_word_space,
+        )
+        for c in cuts
+    ]
+    token_lens = torch.tensor([len(tt) for tt in tokens])
+    tokens = collate_vectors(tokens, padding_value=pad_id)
+    return tokens, token_lens
+
+
+def collate_system_prompt(
+    cuts: CutSet,
+    tokenizer: TokenizerSpec,
+    add_mcq_val_prompt: bool = False,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Collate system prompts from cuts stored in cut.custom['system_prompt']."""
+    pad_id = get_pad_id(tokenizer)
+    tokens = []
+    for c in cuts:
+        if c.custom and c.custom.get("system_prompt", None):
+            prompt_text = c.custom["system_prompt"]
+        elif add_mcq_val_prompt:
+            prompt_text = MCQ_VAL_PROMPT
+        else:
+            prompt_text = None
+
+        if prompt_text:
+            tokens.append(
+                torch.as_tensor(
+                    [tokenizer.bos] + tokenizer.text_to_ids(prompt_text) + [tokenizer.eos], dtype=torch.long
+                )
+            )
+        else:
+            tokens.append(torch.as_tensor([], dtype=torch.long))
+
+    token_lens = torch.tensor([len(tt) for tt in tokens])
+    tokens = collate_vectors(tokens, padding_value=pad_id)
+    return tokens, token_lens
+
+
+def _build_token_channel(
+    cut: Cut,
+    tokenizer: TokenizerSpec,
+    frame_length: Seconds,
+    roles: set[str],
+    pad_id: int = -1,
+    bos_id: int = None,
+    eos_id: int = None,
+    word_align_position: str = 'left',
+    remove_timestamps: bool = False,
+    prepend_word_space: bool = True,
+) -> torch.Tensor:
+    diagnostic = f"Extra info: {cut.id=}"
+    if getattr(cut, "shard_origin", None) is not None:
+        diagnostic = f"{diagnostic} {cut.shard_origin=}"
+
+    total = compute_num_frames(cut.duration, frame_length, cut.sampling_rate)
+    tokens = torch.ones(total, dtype=torch.long) * pad_id
+    for supervision in cut.supervisions:
+        if supervision.speaker in roles:
+            pos = compute_num_frames(supervision.start, frame_length, cut.sampling_rate)
+            if pos >= len(tokens):
+                logging.warning(
+                    f"Ill-constructed example: the beginning offset of a supervision {pos} is larger than or equal to the example's length {len(tokens)}. {diagnostic}"
+                )
+                continue
+            eospos = compute_num_frames(supervision.end, frame_length, cut.sampling_rate)
+            available_frames_for_text = eospos - pos
+
+            text = supervision.text
+
+            text_ids = torch.as_tensor(
+                [bos_id]
+                + _text_to_ids(
+                    text,
+                    tokenizer,
+                    available_frames_for_text=available_frames_for_text,
+                    word_align_position=word_align_position,
+                    remove_timestamps=remove_timestamps,
+                    prepend_word_space=prepend_word_space,
+                )
+            )
+
+            if available_frames_for_text > 0 and len(text_ids) > available_frames_for_text:
+                text_ids = text_ids[:available_frames_for_text]
+            elif available_frames_for_text <= 0:
+                text_ids = torch.tensor([], dtype=torch.long)
+
+            endpos = pos + len(text_ids)
+            if endpos > len(tokens):
+                trunc_len = len(tokens) - pos
+                logging.warning(
+                    f"Truncating training example's text_ids of length {len(text_ids)} by {trunc_len} because {endpos=} > {len(tokens)=}. {diagnostic}"
+                )
+                text_ids = text_ids[:trunc_len]
+                endpos = pos + len(text_ids)
+
+            try:
+                tokens[pos:endpos] = text_ids
+            except Exception as e:
+                raise RuntimeError(f"{tokens.shape=} {pos=} {endpos=} {text_ids.shape=} {diagnostic}") from e
+
+            if eospos < len(tokens) and eos_id is not None:
+                tokens[eospos] = eos_id
+
+    return tokens
+
+
+def _text_to_ids(
+    text: str,
+    tokenizer: TokenizerSpec,
+    _TIMESTAMP_PATTERN_STR=r"<\|(\d+)\|>",
+    available_frames_for_text=None,
+    word_align_position='left',
+    remove_timestamps=False,
+    prepend_word_space=True,
+):
+    if not remove_timestamps and re.compile(_TIMESTAMP_PATTERN_STR).search(text):
+        text_ids = _text_with_timestamps_to_ids(
+            text,
+            tokenizer,
+            _TIMESTAMP_PATTERN_STR,
+            available_frames_for_text,
+            word_align_position,
+            prepend_word_space=prepend_word_space,
+        )
+    else:
+        _TIMESTAMP_PATTERN = re.compile(_TIMESTAMP_PATTERN_STR)
+        text = _TIMESTAMP_PATTERN.sub("", text)
+        text = " ".join(text.strip().split())
+        text_ids = tokenizer.text_to_ids(text)
+    return text_ids
+
+
+def _text_with_timestamps_to_ids(
+    text: str,
+    tokenizer: TokenizerSpec,
+    _TIMESTAMP_PATTERN_STR=r"<\|(\d+)\|>",
+    available_frames_for_text=None,
+    word_align_position='left',
+    prepend_word_space=True,
+) -> list[int]:
+    text_ids, start_times, end_times, word_lens = _extract_text_and_time_tokens(
+        text,
+        tokenizer,
+        _TIMESTAMP_PATTERN_STR,
+        prepend_word_space=prepend_word_space,
+    )
+    text_ids_with_timestamps = _expand_text_with_timestamps_and_word_lengths(
+        text_ids,
+        word_lens,
+        start_times,
+        end_times,
+        available_frames_for_text,
+        frame_rate=0.08,
+        pad_id=get_pad_id(tokenizer),
+        word_align_position=word_align_position,
+    )
+    return text_ids_with_timestamps
+
+
+def _extract_text_and_time_tokens(
+    text, tokenizer: TokenizerSpec, _TIMESTAMP_PATTERN_STR=r"<\|(\d+)\|>", prepend_word_space=True
+):
+    time_tokens = re.findall(_TIMESTAMP_PATTERN_STR, text)
+    start_time = [int(time_tokens[i]) for i in range(0, len(time_tokens), 2)]
+    end_time = [int(time_tokens[i]) for i in range(1, len(time_tokens), 2)]
+    words = re.sub(_TIMESTAMP_PATTERN_STR, '', text).split()
+    text_ids = []
+    word_lens = []
+    for i, word in enumerate(words):
+        word_with_space = word if i == 0 or not prepend_word_space else ' ' + word
+        word_ids = tokenizer.text_to_ids(word_with_space)
+        word_len = len(word_ids)
+        text_ids.extend(word_ids)
+        word_lens.append(word_len)
+    return text_ids, start_time, end_time, word_lens
+
+
+def _expand_text_with_timestamps_and_word_lengths(
+    text_ids,
+    word_lens,
+    start_time,
+    end_time,
+    available_frames_for_text,
+    frame_rate=0.08,
+    pad_id=None,
+    word_align_position='left',
+):
+    def discretize_time(start_token, speech_frame_rate=0.08, timestamp_frame_rate=0.08):
+        return int(start_token * timestamp_frame_rate / speech_frame_rate)
+
+    if pad_id is None:
+        raise ValueError("pad_id must be provided.")
+
+    max_length = available_frames_for_text
+    text_ids_with_timestamps = [pad_id] * max_length
+
+    cur_word_idx = 0
+    for word_idx, word_len in enumerate(word_lens):
+        start_idx = discretize_time(start_time[word_idx], speech_frame_rate=frame_rate)
+        end_idx = discretize_time(end_time[word_idx], speech_frame_rate=frame_rate)
+        if word_align_position == 'left':
+            end_idx = min(start_idx + word_len, end_idx)
+        elif word_align_position == 'right':
+            start_idx = max(start_idx, end_idx - word_len)
+        else:
+            raise ValueError(f"Unknown word_align_position: {word_align_position}")
+
+        word_ids = text_ids[cur_word_idx : cur_word_idx + word_len]
+
+        for i in range(start_idx, end_idx + 1):
+            if i - start_idx < len(word_ids) and i < max_length:
+                token_id = word_ids[i - start_idx]
+                text_ids_with_timestamps[i] = token_id
+
+        cur_word_idx += word_len
+
+    return text_ids_with_timestamps

nemo/collections/speechlm2/data/force_align.py

@@ -0,0 +1,343 @@
+# 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 logging
+import multiprocessing as mp
+import re
+from typing import Any, Dict, List, Optional
+
+import numpy as np
+import torch
+from lhotse import CutSet
+
+# Use NeMo's force alignment utilities instead of torchaudio
+from nemo.collections.asr.models.asr_model import ASRModel
+from nemo.collections.asr.parts.utils.aligner_utils import (
+    add_t_start_end_to_utt_obj,
+    get_batch_variables,
+    viterbi_decoding,
+)
+
+
+class ForceAligner:
+    """
+    Force alignment utility using NeMo CTC-based ASR models for speech-to-text alignment.
+    """
+
+    def __init__(
+        self,
+        asr_model: Optional[ASRModel] = None,
+        device: str = None,
+        frame_length: float = 0.02,
+        asr_model_name: str = "stt_en_fastconformer_ctc_large",
+    ):
+        """
+        Initialize the ForceAligner.
+
+        Args:
+            asr_model: NeMo ASR model instance for alignment. If None, will load from asr_model_name
+            device: Device to run alignment on (default: auto-detect)
+            frame_length: Frame length in seconds for timestamp conversion
+            asr_model_name: Name of the NeMo ASR model to load if asr_model is None
+        """
+        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+        self.frame_length = frame_length
+        self.asr_model_name = asr_model_name
+
+        self.asr_model = asr_model
+        self.output_timestep_duration = None
+        self._model_loaded = False
+
+    def _load_asr_model(self):
+        """Load the NeMo ASR model."""
+        try:
+            if self.device == 'cuda' and mp.get_start_method(allow_none=True) == 'fork':
+                logging.warning(
+                    "Detected 'fork' multiprocessing start method with CUDA device. "
+                    "To avoid CUDA re-initialization errors in worker processes, "
+                    "falling back to CPU for force alignment. "
+                    "To use CUDA, set mp.set_start_method('spawn', force=True) in your main training script "
+                    "before creating the DataLoader."
+                )
+                self.device = 'cpu'
+
+            device = torch.device(self.device)
+            logging.info(f"Loading NeMo ASR model '{self.asr_model_name}' for force alignment on device {device}")
+
+            if self.asr_model is None:
+                # Load ASR model from pretrained
+                self.asr_model = ASRModel.from_pretrained(self.asr_model_name, map_location=device)
+            else:
+                self.asr_model = self.asr_model.to(device)
+
+            self.asr_model.eval()
+
+            # Calculate output timestep duration
+            try:
+                self.output_timestep_duration = (
+                    self.asr_model.cfg['preprocessor']['window_stride'] * self.asr_model.encoder.subsampling_factor
+                )
+            except Exception as e:
+                # Default fallback based on typical FastConformer settings
+                self.output_timestep_duration = 0.04
+                logging.warning(
+                    f"Could not calculate output_timestep_duration from model config: {e}. "
+                    f"Using default {self.output_timestep_duration}s"
+                )
+
+            logging.info(
+                f"NeMo ASR model loaded successfully for force alignment. "
+                f"Output timestep duration: {self.output_timestep_duration}s"
+            )
+        except Exception as e:
+            logging.error(f"Failed to load NeMo ASR model for force alignment: {e}")
+            self.asr_model = None
+            raise
+
+    def batch_force_align_user_audio(self, cuts: CutSet, source_sample_rate: int = 16000) -> CutSet:
+        """
+        Perform batched force alignment on all user audio segments.
+
+        Collects all user segments, writes temp files, runs a single batched
+        get_batch_variables + viterbi_decoding call, then maps results back.
+
+        Args:
+            cuts: CutSet containing all cuts to process
+            source_sample_rate: Source sample rate of the audio
+
+        Returns:
+            CutSet with updated supervision texts (timestamped where alignment succeeded)
+        """
+        if not self._model_loaded:
+            self._load_asr_model()
+            self._model_loaded = True
+
+        if self.asr_model is None:
+            logging.warning("ASR model not available for force alignment, returning empty cutset")
+            return CutSet.from_cuts([])
+
+        # Collect all user supervisions
+        user_supervisions = []
+        user_cuts = []
+        for cut in cuts:
+            for supervision in cut.supervisions:
+                if supervision.speaker.lower() == "user":
+                    user_supervisions.append(supervision)
+                    user_cuts.append(cut)
+
+        if not user_supervisions:
+            logging.info("No user supervisions found for force alignment")
+            return cuts
+
+        logging.info(f"Performing batched force alignment on {len(user_supervisions)} user audio segments")
+
+        # Prepare all audio arrays and texts for batched processing
+        audio_arrays = []
+        normalized_texts = []
+        valid_indices = []  # track which supervisions have valid audio/text
+        target_sample_rate = 16000
+
+        for i, (supervision, cut) in enumerate(zip(user_supervisions, user_cuts)):
+            try:
+                text = self._strip_timestamps(supervision.text)
+                normalized_text = self._normalize_transcript(text)
+                if not normalized_text.strip():
+                    logging.warning(f"Text became empty after normalization: {supervision.text}")
+                    continue
+
+                user_cut = cut.truncate(offset=supervision.start, duration=supervision.duration)
+                audio = user_cut.load_audio()
+                if audio.ndim > 1:
+                    audio = audio.mean(axis=0)
+
+                if source_sample_rate != target_sample_rate:
+                    from scipy import signal
+
+                    num_samples = int(len(audio) * target_sample_rate / source_sample_rate)
+                    audio = signal.resample(audio, num_samples)
+
+                # Add silence padding for better alignment at the end
+                silence_samples = int(0.64 * target_sample_rate)
+                audio = np.concatenate([audio, np.zeros(silence_samples)])
+
+                audio_arrays.append(audio)
+                normalized_texts.append(normalized_text)
+                valid_indices.append(i)
+            except Exception as e:
+                logging.error(f"Failed to prepare segment {i} for alignment: {e}")
+
+        if not audio_arrays:
+            logging.warning("No valid segments to align")
+            return cuts
+
+        # Batched ASR inference + Viterbi decoding
+        success_count = 0
+        failed_count = 0
+        try:
+            (
+                log_probs_batch,
+                y_batch,
+                T_batch,
+                U_batch,
+                utt_obj_batch,
+                output_timestep_duration,
+            ) = get_batch_variables(
+                audio=audio_arrays,
+                model=self.asr_model,
+                gt_text_batch=normalized_texts,
+                align_using_pred_text=False,
+                output_timestep_duration=self.output_timestep_duration,
+            )
+
+            alignments_batch = viterbi_decoding(
+                log_probs_batch=log_probs_batch,
+                y_batch=y_batch,
+                T_batch=T_batch,
+                U_batch=U_batch,
+                viterbi_device=torch.device(self.device),
+            )
+
+            # Map results back to supervisions
+            for batch_idx, orig_idx in enumerate(valid_indices):
+                try:
+                    if batch_idx >= len(alignments_batch) or batch_idx >= len(utt_obj_batch):
+                        failed_count += 1
+                        continue
+
+                    utt_obj = utt_obj_batch[batch_idx]
+                    if not utt_obj.token_ids_with_blanks:
+                        failed_count += 1
+                        continue
+
+                    alignment = alignments_batch[batch_idx]
+                    utt_obj = add_t_start_end_to_utt_obj(utt_obj, alignment, output_timestep_duration)
+                    word_segments = self._extract_word_timestamps(utt_obj)
+
+                    if word_segments:
+                        timestamped_text = self._convert_alignment_to_timestamped_text(
+                            word_segments, user_supervisions[orig_idx].text
+                        )
+                        user_supervisions[orig_idx].text = timestamped_text
+                        success_count += 1
+                    else:
+                        failed_count += 1
+                except Exception as e:
+                    logging.error(f"Failed to process alignment for segment {orig_idx}: {e}")
+                    failed_count += 1
+
+        except Exception as e:
+            logging.error(f"Batched force alignment failed: {e}")
+            failed_count = len(valid_indices)
+        finally:
+            if self.device == 'cuda' and torch.cuda.is_available():
+                torch.cuda.empty_cache()
+
+        if failed_count > 0:
+            logging.warning(
+                f"Force alignment failed for {failed_count}/{len(user_supervisions)} user segments. "
+                f"Keeping original text for failed alignments."
+            )
+        else:
+            logging.info(f"Force alignment succeeded for all {success_count} user segments.")
+
+        return cuts
+
+    def _extract_word_timestamps(self, utt_obj) -> List[Dict[str, Any]]:
+        """
+        Extract word-level timestamps from the utterance object returned by NeMo aligner.
+
+        Args:
+            utt_obj: Utterance object with timing information
+
+        Returns:
+            List of word segments with timing information
+        """
+        word_segments = []
+
+        for segment_or_token in utt_obj.segments_and_tokens:
+            # Check if this is a Segment object (has words_and_tokens attribute)
+            if hasattr(segment_or_token, 'words_and_tokens'):
+                segment = segment_or_token
+                for word_or_token in segment.words_and_tokens:
+                    # Check if this is a Word object (has 'text' and timing attributes)
+                    if hasattr(word_or_token, 'text') and hasattr(word_or_token, 't_start'):
+                        word = word_or_token
+                        # Skip CTC blank tokens and include only words with valid timing
+                        if (
+                            word.text not in ('<b>', '')
+                            and word.t_start is not None
+                            and word.t_end is not None
+                            and word.t_start >= 0
+                            and word.t_end >= 0
+                        ):
+                            word_segments.append(
+                                {
+                                    'word': word.text,
+                                    'start': word.t_start,
+                                    'end': word.t_end,
+                                    'score': 1.0,  # NeMo CTC alignment doesn't provide confidence scores
+                                }
+                            )
+
+        return word_segments
+
+    def _normalize_transcript(self, transcript: str) -> str:
+        """
+        Normalize transcript for the ASR model's tokenizer.
+        Keeps it simple to match common ASR preprocessing.
+        """
+        text = transcript.lower()
+        # Remove special characters except apostrophes and spaces
+        text = re.sub(r"[^a-z' ]", " ", text)
+        # Collapse multiple spaces
+        text = re.sub(r' +', ' ', text)
+        return text.strip()
+
+    def _convert_alignment_to_timestamped_text(
+        self, alignment_result: List[Dict[str, Any]], original_text: str
+    ) -> str:
+        """
+        Convert alignment results to timestamped text format.
+
+        Args:
+            alignment_result: List of word segments with timing information
+            original_text: Original text without timestamps
+
+        Returns:
+            Text with timestamp tokens in the format <|start_frame|>word<|end_frame|>
+        """
+        timestamped_words = []
+
+        for word_seg in alignment_result:
+            word = word_seg["word"]
+            start_frame = int(word_seg["start"] / self.frame_length)
+            end_frame = int(word_seg["end"] / self.frame_length)
+            timestamped_words.append(f"<|{start_frame}|> {word} <|{end_frame}|>")
+
+        return " ".join(timestamped_words)
+
+    def _strip_timestamps(self, text: str) -> str:
+        """
+        Strip timestamp tokens from text.
+
+        Args:
+            text: Text that may contain timestamp tokens
+
+        Returns:
+            Text with timestamp tokens removed
+        """
+        text = re.sub(r'<\|[0-9]+\|>', '', text)
+        text = re.sub(r' +', ' ', text)
+
+        return text.strip()

nemo/collections/speechlm2/data/s2s_dataset.py

@@ -0,0 +1,204 @@
+# 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
+
+import torch
+import torch.utils.data
+from lhotse import CutSet, Seconds, compute_num_frames
+from lhotse.cut import Cut
+from lhotse.dataset.collation import collate_audio, collate_vectors
+from lhotse.utils import ifnone
+
+from nemo.collections.common.tokenizers import TokenizerSpec
+from nemo.collections.speechlm2.data.utils import get_pad_id
+from nemo.utils import logging
+
+
+class DuplexS2SDataset(torch.utils.data.Dataset):
+    """
+    A dataset for duplex speech-to-speech models that handles bidirectional conversations.
+
+    This dataset processes Lhotse CutSet objects containing recordings with supervision segments
+    from different speakers (roles). It creates aligned representations of audio and text for
+    both source (input) and target (output) channels, preserving temporal alignment between
+    audio frames and text tokens.
+
+    Args:
+        tokenizer (TokenizerSpec):
+            Tokenizer for converting text to token IDs and vice versa. Must support BOS and EOS tokens.
+            It's expected to support PAD token as well, otherwise we will use 0 as the pad token
+            and emit a warning.
+
+        frame_length (Seconds):
+            Duration of a single frame in seconds. Used to calculate frame positions for token alignment.
+
+        source_sample_rate (int):
+            Sample rate for source audio (e.g., 16000 Hz).
+
+        target_sample_rate (int):
+            Sample rate for target audio (e.g., 22050 Hz).
+
+        input_roles (list[str], optional):
+            List of speaker roles (cut.supervisions[:].speaker) to consider as inputs. Defaults to ["user"].
+
+        output_roles (list[str], optional):
+            List of speaker roles (cut.supervisions[:].speaker) to consider as outputs. Defaults to ["agent"].
+
+    Returns:
+        A dictionary with the following keys:
+            - source_audio: Tensor of source waveform samples [B, T]
+            - source_audio_lens: Tensor of source audio lengths [B]
+            - target_audio: Tensor of target waveform samples [B, T]
+            - target_audio_lens: Tensor of target audio lengths [B]
+            - target_tokens: Tensor of target text tokens [B, T], with special tokens (BOS/EOS/PAD)
+                at positions aligned with audio frames
+            - target_token_lens: Tensor of target token sequence lengths [B]
+            - source_tokens: Tensor of source text tokens [B, T], with special tokens (BOS/EOS/PAD)
+                at positions aligned with audio frames
+            - source_token_lens: Tensor of source token sequence lengths [B]
+            - target_texts: List of full target texts joined from output_roles supervisions [B]
+
+    Notes:
+        - The dataset ensures frame-level alignment between audio and text by inserting tokens at
+          specific frame positions based on the timing of supervision segments.
+        - PAD tokens (typically 0) are used to fill gaps where there's no text.
+        - BOS tokens mark the beginning of each speech segment.
+        - EOS tokens mark the end of each speech segment.
+        - Text tokens from each speaker are placed at frame positions corresponding to their
+          timestamp in the original recording, preserving the temporal relationship.
+          This is a segment-level alignment only, not word-level alignment.
+    """
+
+    def __init__(
+        self,
+        tokenizer: TokenizerSpec,
+        frame_length: Seconds,
+        source_sample_rate: int,
+        target_sample_rate: int,
+        input_roles: list[str] = None,
+        output_roles: list[str] = None,
+    ):
+        self.tokenizer = tokenizer
+        self.frame_length = frame_length
+        self.source_sample_rate = source_sample_rate
+        self.target_sample_rate = target_sample_rate
+        self.input_roles = set(ifnone(input_roles, ["user"]))
+        self.output_roles = set(ifnone(output_roles, ["agent"]))
+
+        assert tokenizer.bos is not None, "BOS support in the tokenizer is required for S2S models."
+        assert tokenizer.eos is not None, "EOS support in the tokenizer is required for S2S models."
+
+    def __getitem__(self, cuts: CutSet) -> dict:
+        cuts = cuts.transform_text(_strip_timestamps)
+        source_audio, source_audio_lens = collate_audio(cuts.resample(self.source_sample_rate))
+        target_audio, target_audio_lens = collate_audio(
+            cuts.resample(self.target_sample_rate, recording_field="target_audio"), recording_field="target_audio"
+        )
+        target_tokens, target_token_lens = collate_token_channel(
+            cuts, self.tokenizer, self.frame_length, roles=self.output_roles
+        )
+        source_tokens, source_token_lens = collate_token_channel(
+            cuts, self.tokenizer, self.frame_length, roles=self.input_roles
+        )
+        return {
+            "source_audio": source_audio,
+            "source_audio_lens": source_audio_lens,
+            "target_audio": target_audio,
+            "target_audio_lens": target_audio_lens,
+            "target_tokens": target_tokens,
+            "target_token_lens": target_token_lens,
+            "source_tokens": source_tokens,
+            "source_token_lens": source_token_lens,
+            "target_texts": [
+                " ".join(s.text for s in cut.supervisions if s.speaker in self.output_roles) for cut in cuts
+            ],
+        }
+
+
+def collate_token_channel(
+    cuts: CutSet,
+    tokenizer: TokenizerSpec,
+    frame_length: Seconds,
+    roles: set[str],
+) -> tuple[torch.Tensor, torch.Tensor]:
+    pad_id = get_pad_id(tokenizer)
+    tokens = [
+        build_token_channel(c, tokenizer=tokenizer, frame_length=frame_length, roles=roles, pad_id=pad_id)
+        for c in cuts
+    ]
+    token_lens = torch.tensor([len(tt) for tt in tokens])
+    tokens = collate_vectors(tokens, padding_value=pad_id)
+    return tokens, token_lens
+
+
+def build_token_channel(
+    cut: Cut,
+    tokenizer: TokenizerSpec,
+    frame_length: Seconds,
+    roles: set[str],
+    pad_id: int = -1,
+) -> torch.Tensor:
+    diagnostic = f"Extra info: {cut.id=}"
+    if getattr(cut, "shard_origin", None) is not None:
+        diagnostic = f"{diagnostic} {cut.shard_origin=}"
+
+    total = compute_num_frames(cut.duration, frame_length, cut.sampling_rate)
+    tokens = torch.ones(total, dtype=torch.long) * pad_id
+    for supervision in cut.supervisions:
+        if supervision.speaker in roles:
+            text_ids = torch.as_tensor([tokenizer.bos] + tokenizer.text_to_ids(supervision.text))
+
+            # Determine the frame offset for the start of the supervision to insert the text tokens.
+            pos = compute_num_frames(supervision.start, frame_length, cut.sampling_rate)
+            if pos > len(tokens):
+                logging.warning(
+                    f"Ill-constructed example: the beginning offset of a supervision {pos} is larger than the example's length {len(tokens)}. {diagnostic}"
+                )
+                continue
+
+            # Determine the frame offset for the last non-EOS text token to form a valid range for insertion;
+            # Note that EOS will be placed possibly much later, at the frame that coincides with end of speech,
+            # rather than end of text. The gap between last non-EOS token and EOS token will be filled with `pad_id`.
+            endpos = pos + len(text_ids)
+            if endpos > len(tokens):
+                trunc_len = len(tokens) - pos
+                logging.warning(
+                    f"Truncating training example's text_ids of length {len(text_ids)} by {trunc_len} because {endpos=} > {len(tokens)=}. {diagnostic}"
+                )
+                text_ids = text_ids[:trunc_len]
+            try:
+                tokens[pos:endpos] = text_ids
+            except Exception as e:
+                raise RuntimeError(f"{tokens.shape=} {pos=} {endpos=} {text_ids.shape=} {diagnostic}") from e
+
+            # Insert EOS at the end of the supervision segment.
+            eospos = compute_num_frames(supervision.end, frame_length, cut.sampling_rate)
+            if eospos < len(tokens):  # skip otherwise - unfinished turn
+                tokens[eospos] = tokenizer.eos
+
+    return tokens
+
+
+def _strip_timestamps(
+    text: str, _TIMESTAMP_PATTERN=re.compile(r"<\|\d+\|>"), _SPACE_PATTERN=re.compile(r"\s+")
+) -> str:
+    """
+    Strips timestamp tokens from text, e.g. turns:
+      '<|0|> Hey <|3|> <|3|> how <|5|> <|7|> are <|8|> <|8|> <|10|> you? <|12|>'
+      into:
+      'Hey how are you?'
+    """
+    # Regexp pattern args are cached compiled patterns (micro-optimization).
+    text = _TIMESTAMP_PATTERN.sub("", text)  # strip timestamp tokens if present
+    return _SPACE_PATTERN.sub(" ", text).strip()  # strip multi-whitespaces

nemo/collections/speechlm2/data/salm_dataset.py

@@ -0,0 +1,279 @@
+# 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 logging
+import os
+from dataclasses import dataclass
+from itertools import groupby
+from typing import Iterable, Union
+
+import numpy as np
+import torch
+import torch.utils.data
+from lhotse import CutSet, fastcopy
+from lhotse.cut import MixedCut, MultiCut
+from lhotse.dataset import AudioSamples
+from torch.nn import CrossEntropyLoss
+from torch.nn.utils.rnn import pad_sequence
+
+from nemo.collections.asr.parts.utils.sot_speaker_alignment import (
+    collate_speaker_activity_targets,
+    ensure_single_speaker_sot,
+    fix_speaker_activity,
+    speaker_activity_from_cut,
+)
+from nemo.collections.common.data.lhotse import NeMoMultimodalConversation
+from nemo.collections.common.data.lhotse.text_adapters import (
+    AudioTurn,
+    TextTurn,
+    collate_conversation_audio_fault_tolerant,
+)
+from nemo.collections.common.data.prompt_fn import registered_prompt_format_fn
+from nemo.collections.common.prompts import Llama2PromptFormatter
+from nemo.collections.common.tokenizers import AutoTokenizer
+from nemo.collections.speechlm2.data.utils import get_pad_id
+
+
+class SALMDataset(torch.utils.data.Dataset):
+    """
+    A dataset for Speech-Augmented Language Models (SALM) that processes multimodal conversations
+    containing both text and audio turns.
+
+    This dataset handles NeMoMultimodalConversation objects which combine text messages
+    and audio segments in a conversational format. It uses audio_locator_tag in the text,
+    where each such placeholder corresponds to an entire audio segment.
+
+    Args:
+        tokenizer (AutoTokenizer):
+            Tokenizer for converting text to token IDs and vice versa. Must have a special
+            audio_locator_tag token that will be replaced with audio embeddings during model's
+            training step.
+        multispeaker_cfg (dict | None):
+            Optional Serialized Output Training (SOT) speaker-activity settings.
+            When provided, each batch additionally includes RTTM-derived
+            ``spk_targets`` / ``spk_target_length``.
+
+            [ SOT Example for overlapping speakers ]
+            Speaker-parallel transcription as a timeline:
+                <spk:0>: Well, we should focus on the most important issues first.
+                <spk:1>:                     Let me finish. John, let me finish.
+            Serialized Output Training (SOT) transcription:
+                <spk:0> Well, we should focus on <spk:1> Let me <spk:0> the most
+                <spk:1> finish, John, <spk:0> important issues <spk:1> let me finish.
+                <spk:0> first.
+
+    Returns:
+        A dictionary with the following keys:
+            - audios: Tensor of audio waveform samples [B_audio, T_samples]
+            - audio_lens: Tensor of audio lengths [B_audio]
+            - input_ids: Tensor of text token IDs [B, T_tokens], including audio_locator_tag tokens
+            - loss_mask: Boolean tensor [B, T_tokens] indicating which tokens are part of the
+                assistant's responses (True) and should be used for computing loss
+
+    Notes:
+        - Each audio_locator_tag token in input_ids corresponds to an audio segment in audios
+        - The SALM model later replaces these audio_locator_tag tokens with encoded audio embeddings
+        - The loss_mask identifies which tokens are part of the target sequences (assistant responses)
+          and which are part of the source sequences (user prompts)
+        - The input_ids and loss_mask will be expanded during model forward pass to account for
+          the variable-length audio segments that replace each audio_locator_tag token
+        - Serialized Output Training (SOT) speaker tags ``<spk:N>`` stay regular text tokens here;
+          normalization and aliasing happen upstream. Auxiliary SOT mode (off by default) is opt-in via
+          ``multispeaker_cfg`` and does not affect the default single-speaker behavior.
+    """
+
+    def __init__(self, tokenizer: AutoTokenizer, multispeaker_cfg: dict | None = None) -> None:
+        self.tokenizer = tokenizer
+        self.pad_id = get_pad_id(tokenizer)
+        # Setting USE_AIS_GET_BATCH=true makes the loader issue a single AIStore GetBatch
+        # call per minibatch, paired with URL-backed cuts produced by the multimodal
+        # conversation adapters (NeMoMultimodalConversation{Jsonl,ShareGPTJsonl}Adapter).
+        self.load_audio = AudioSamples(
+            fault_tolerant=True,
+            use_batch_loader=os.environ.get("USE_AIS_GET_BATCH", "False").lower() == "true",
+            mono_downmix=True,
+        )
+        self.multispeaker_cfg = MultiSpeakerConfig.from_dict(multispeaker_cfg)
+        self.multispeaker_processor = (
+            SALMMultiSpeakerProcessor(self.multispeaker_cfg) if self.multispeaker_cfg is not None else None
+        )
+
+    def __getitem__(self, conversations: CutSet) -> dict | None:
+        # Note: the function call below may filter out some or all conversations due to audio loading issues.
+        # If all conversations are filtered out, we'll return None, and expect users to wrap this dataset
+        # in ``nemo.collections.common.data.fallback.FallbackDataset`` to use the previous mini-batch instead.
+        try:
+            audios, audio_lens, conversations = collate_conversation_audio_fault_tolerant(
+                conversations, self.load_audio
+            )
+        except Exception as e:
+            logging.warning(f"Error collating conversations: {e}")
+            return None
+        if not conversations:
+            return None
+        batch = {
+            "audios": audios,
+            "audio_lens": audio_lens,
+            "input_ids": left_collate_vectors([c.input_ids for c in conversations], padding_value=self.pad_id),
+            "loss_mask": left_collate_vectors(
+                [getattr(c, "mask", torch.empty(0)) for c in conversations], padding_value=0
+            ).to(torch.bool),
+            "conversations": drop_in_memory_data(conversations),
+        }
+        if self.multispeaker_processor is not None:
+            self.multispeaker_processor(batch)
+        return batch
+
+
+def left_collate_vectors(
+    tensors: Iterable[Union[torch.Tensor, np.ndarray]],
+    padding_value: Union[int, float] = CrossEntropyLoss().ignore_index,
+) -> torch.Tensor:
+    tensors = [torch.as_tensor(t) for t in tensors]
+    assert all(len(t.shape) == 1 for t in tensors), "Expected only 1-D input tensors."
+    return pad_sequence(tensors, batch_first=True, padding_value=padding_value, padding_side="left")
+
+
+def drop_in_memory_data(conversations: CutSet) -> CutSet:
+    def _drop(conversation: NeMoMultimodalConversation) -> NeMoMultimodalConversation:
+        turns = []
+        for t in conversation.turns:
+            if isinstance(t, AudioTurn):
+                t = fastcopy(t, cut=t.cut.drop_in_memory_data())
+            turns.append(t)
+        return fastcopy(conversation, turns=turns)
+
+    return conversations.map(_drop, apply_fn=None)
+
+
+@registered_prompt_format_fn(NeMoMultimodalConversation, Llama2PromptFormatter)
+def default_multimodal_conversation_prompt_format_fn(
+    example: NeMoMultimodalConversation, prompt: Llama2PromptFormatter, **prompt_kwargs
+):
+    # Collapse consecutive same-role turns into single turn for proper prompt formatting.
+    turns = groupby(
+        [
+            {
+                "role": turn.role,
+                "slots": {"message": turn.value if isinstance(turn, TextTurn) else turn.audio_locator_tag},
+            }
+            for turn in example.turns
+        ],
+        key=lambda turn: turn["role"],
+    )
+    turns = [
+        {"role": role, "slots": {"message": " ".join(t["slots"]["message"] for t in turn_grp)}}
+        for role, turn_grp in turns
+    ]
+    if hasattr(example, "system_prompt"):
+        turns[0]["role"] = "system_and_user"
+        turns[0]["slots"]["system"] = example.system_prompt
+    return prompt.encode_dialog(turns, **prompt_kwargs)
+
+
+@dataclass(frozen=True)
+class MultiSpeakerConfig:
+    """Configuration for auxiliary multi-speaker SOT targets."""
+
+    num_speakers: int = 4
+    no_rttm_to_ones: bool = True
+    num_sample_per_mel_frame: int = 160
+    num_mel_frame_per_target_frame: int = 8
+
+    @staticmethod
+    def from_dict(cfg: dict | None) -> "MultiSpeakerConfig | None":
+        """Build a config from a raw settings dict, or return ``None`` when no SOT settings are given."""
+        if cfg is None:
+            return None
+        return MultiSpeakerConfig(
+            num_speakers=int(cfg.get('num_speakers', 4)),
+            no_rttm_to_ones=cfg.get('no_rttm_to_ones', True),
+            num_sample_per_mel_frame=int(cfg.get('window_stride', 0.01) * cfg.get('sample_rate', 16000)),
+            num_mel_frame_per_target_frame=int(cfg.get('subsampling_factor', 8)),
+        )
+
+
+class SALMMultiSpeakerProcessor:
+    """Adds auxiliary SOT speaker-activity targets to an otherwise prepared SALM batch."""
+
+    def __init__(self, cfg: MultiSpeakerConfig) -> None:
+        self.cfg = cfg
+
+    def __call__(self, batch: dict) -> None:
+        """Attach RTTM-derived ``spk_targets`` / ``spk_target_length`` to ``batch`` in place."""
+        cfg = self.cfg
+        speaker_activities = self._build_speaker_activities(batch["conversations"])
+        if not speaker_activities:
+            return
+        targets, target_length = collate_speaker_activity_targets(
+            speaker_activities,
+            batch["audio_lens"],
+            num_speakers=cfg.num_speakers,
+            num_sample_per_mel_frame=cfg.num_sample_per_mel_frame,
+            num_mel_frame_per_target_frame=cfg.num_mel_frame_per_target_frame,
+            dtype=batch["audios"].dtype,
+        )
+        batch["spk_targets"] = targets
+        batch["spk_target_length"] = target_length
+
+    def _build_speaker_activities(self, conversations: CutSet) -> list[torch.Tensor]:
+        cfg = self.cfg
+        speaker_activities = []
+        for conversation in conversations:
+            for turn in conversation.turns:
+                if not isinstance(turn, AudioTurn):
+                    continue
+
+                cut = self._prepare_audio_turn_cut(turn)
+                speaker_activity = speaker_activity_from_cut(
+                    cut,
+                    num_speakers=cfg.num_speakers,
+                    num_sample_per_mel_frame=cfg.num_sample_per_mel_frame,
+                    num_mel_frame_per_target_frame=cfg.num_mel_frame_per_target_frame,
+                    no_rttm_to_ones=cfg.no_rttm_to_ones,
+                )
+
+                text = self._audio_turn_text(turn, cut)
+                new_text, _, _ = ensure_single_speaker_sot(text)
+
+                speaker_activity = fix_speaker_activity(new_text, speaker_activity, cfg.num_speakers)
+                speaker_activities.append(speaker_activity)
+        return speaker_activities
+
+    @staticmethod
+    def _prepare_audio_turn_cut(turn: AudioTurn):
+        cut = turn.cut
+        if isinstance(cut, MultiCut):
+            cut = cut.to_mono(mono_downmix=True)
+        elif isinstance(cut, MixedCut):
+            pass
+        elif 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 for speaker targets; remaining channels are ignored.",
+                cut.id,
+                cut.num_channels,
+            )
+            cut = cut.with_channels(channels=[0])
+
+        if getattr(cut, "custom", None) is None:
+            cut = fastcopy(cut, custom={})
+
+        return cut
+
+    @staticmethod
+    def _audio_turn_text(turn: AudioTurn, cut) -> str:
+        text = turn.text or getattr(cut, "text", None)
+        if text:
+            return text
+        return " ".join(s.text for s in getattr(cut, "supervisions", []) if s.text)

nemo/collections/speechlm2/data/utils.py

@@ -0,0 +1,27 @@
+# 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 warnings
+
+
+def get_pad_id(tokenizer) -> int:
+    pad_id = tokenizer.pad
+    if pad_id is not None:
+        return pad_id
+    pad_id = tokenizer.unk_id
+    if pad_id is not None:
+        return pad_id
+    warnings.warn(
+        "The text tokenizer has no <pad> or <unk> tokens available, using ID 0 for padding (this may lead to silent bugs)."
+    )
+    return 0