# 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 from copy import deepcopy import torch import torch.nn.functional as F 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.collections.speechlm2.parts.precision import fp32_precision from nemo.collections.tts.parts.utils.helpers import get_mask_from_lengths from nemo.utils import logging class DuplexEARTTSDataset(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"]. p_drop_description (float, optional): Probability of dropping text descriptions. Default: `0.0`. add_text_bos_and_eos_in_each_turn (bool, optional): If True, each conversational turn from any speaker is explicitly delimited with BOS and EOS tokens in the text stream. Default: `True`. add_audio_prompt (bool, optional): If True, an optional audio/speaker prompt is appended after the description. Default: `True`. audio_prompt_duration (float, optional): Duration (in seconds) of the audio prompt appended when `add_audio_prompt=True`. Default: `3.0`. num_delay_speech_tokens (int, optional): Number of PAD tokens to insert before speech tokens to artificially delay the start of speech. Default: `0`. Returns: A dictionary with the following keys: - sample_id: List of sample IDs for each cut in the batch [B] - non_prompt_mask: Bool tensor [B, T] marking positions that are not part of the prompt - prompt_lens: Tensor of description + audio prompt lengths [B] - aligned_attention_mask: Bool tensor [B, T] used by alignment-aware transformer models - aligned_position_ids: Tensor of position indices aligned to audio frames [B, T] - 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_text_tokens: Tensor of frame-aligned input text tokens [B, T], including BOS/EOS/PAD when enabled - target_token_lens: Tensor of target token sequence lengths [B] - source_tokens: Tensor of frame-aligned source text tokens [B, T], including BOS/EOS/PAD - source_token_lens: Tensor of source token sequence lengths [B] - target_texts: List of full target texts joined from output_roles supervisions [B] - audio_prompt: Tensor of optional speaker reference waveform samples [B, T] - audio_prompt_lens: Tensor of speaker reference audio lengths [B] - task: List indicating the task to use for each cut (default "s2s_duplex") [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, frame_length: Seconds, source_sample_rate: int, target_sample_rate: int, input_roles: list[str] = None, output_roles: list[str] = None, p_drop_description: float = 0.0, add_text_bos_and_eos_in_each_turn: bool = True, add_audio_prompt: bool = True, audio_prompt_duration: float = 3.0, num_delay_speech_tokens: int = 0, add_system_prompt: bool = False, ignore_data_system_prompt: bool = True, ): 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"])) self.p_drop_description = p_drop_description self.add_text_bos_and_eos_in_each_turn = add_text_bos_and_eos_in_each_turn self.add_audio_prompt = add_audio_prompt self.audio_prompt_duration = audio_prompt_duration self.num_delay_speech_tokens = num_delay_speech_tokens self.add_system_prompt = add_system_prompt self.ignore_data_system_prompt = ignore_data_system_prompt # compute source and target samples_per_frame self.source_samples_per_frame = int(self.source_sample_rate * self.frame_length) self.target_samples_per_frame = int(self.target_sample_rate * self.frame_length) 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) # ensures fp32 audio load to avoid issues of duration mistakes on fp16 training with fp32_precision(): 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_text_tokens, target_token_lens = collate_token_channel( cuts, self.tokenizer, self.frame_length, roles=self.output_roles, add_text_bos_and_eos_in_each_turn=self.add_text_bos_and_eos_in_each_turn, ) source_tokens, source_token_lens = collate_token_channel( cuts, self.tokenizer, self.frame_length, roles=self.input_roles, add_text_bos_and_eos_in_each_turn=self.add_text_bos_and_eos_in_each_turn, ) with fp32_precision(): audio_prompt, audio_prompt_lens = get_audio_prompt( cuts, self.target_sample_rate, roles=self.output_roles, recording_field="target_audio" ) # add speech channel delay if needed if self.num_delay_speech_tokens: source_audio, source_audio_lens, target_audio, target_audio_lens = add_speech_delay( source_audio, source_audio_lens, target_audio, target_audio_lens, self.num_delay_speech_tokens, self.target_samples_per_frame, self.source_samples_per_frame, ) if self.add_system_prompt: with fp32_precision(): system_prompts, system_prompts_lens, system_prompts_raw = collate_system_prompt( cuts, self.tokenizer, ignore_data_system_prompt=self.ignore_data_system_prompt ) else: system_prompts = None system_prompts_lens = None system_prompts_raw = None # get dataset name/type dataset_type = [getattr(c, "type", "") for c in cuts] # add audio prompt if needed ( target_text_tokens, target_token_lens, source_tokens, source_token_lens, source_audio, source_audio_lens, target_audio, target_audio_lens, prompt_lens, ) = self.maybe_add_audio_prompt( target_text_tokens, target_token_lens, source_tokens, source_token_lens, target_audio, target_audio_lens, source_audio, source_audio_lens, audio_prompt, audio_prompt_lens, system_prompts, system_prompts_lens, ) # create non_prompt_mask that should mask desc plus audio prompt if used non_prompt_mask = get_mask_from_lengths(target_token_lens) for i, frame in enumerate(prompt_lens): non_prompt_mask[i, : frame - 1] = 0.0 max_len = max(target_token_lens) # Segment IDs per sequence (padded) aligned_segment_ids = torch.stack( [ torch.nn.functional.pad(torch.full((seq_len,), i), (0, max_len - seq_len), value=-1) # -1 for padding for i, seq_len in enumerate(target_token_lens) ], dim=0, ) # [B, max_len] # Attention mask: same-segment & causal aligned_attention_mask = ( aligned_segment_ids.unsqueeze(-2) == aligned_segment_ids.unsqueeze(-1) ) & ( # [B, max_len, max_len] torch.arange(max_len).unsqueeze(0).unsqueeze(1) <= torch.arange(max_len).unsqueeze(0).unsqueeze(-1) ) # causal tril aligned_attention_mask = aligned_attention_mask.unsqueeze(1) # [B, 1, max_len, max_len] # create position IDs from the aligned length aligned_position_ids = torch.stack( [ torch.nn.functional.pad( torch.arange(seq_len), (0, max(target_token_lens) - seq_len), value=0 ) # value=0 is safe for padding for seq_len in target_token_lens ], dim=0, ) return { "sample_id": [str(cut.id) for cut in cuts], "non_prompt_mask": non_prompt_mask.bool(), "prompt_lens": prompt_lens, "aligned_attention_mask": aligned_attention_mask.bool(), "aligned_position_ids": aligned_position_ids, "source_audio": source_audio, "source_audio_lens": source_audio_lens, "target_audio": target_audio, "target_audio_lens": target_audio_lens, "target_text_tokens": target_text_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 ], "audio_prompt": audio_prompt, "audio_prompt_lens": audio_prompt_lens, "system_prompts_raw": system_prompts_raw, "dataset_type": dataset_type, "task": [getattr(cut, "task", "s2s_duplex") for cut in cuts], } def maybe_add_audio_prompt( self, target_text_tokens: torch.Tensor, target_token_lens: torch.Tensor, source_tokens: torch.Tensor, source_token_lens: torch.Tensor, target_audio: torch.Tensor, target_audio_lens: torch.Tensor, source_audio: torch.Tensor, source_audio_lens: torch.Tensor, audio_prompt: torch.Tensor, audio_prompt_lens: torch.Tensor, system_prompts: torch.Tensor = None, system_prompts_lens: torch.Tensor = None, ): """ Prepend an audio-based speaker prompt and aligned text tokens to the Duplex S2S inputs. This method optionally injects a speaker-reference audio prompt at the beginning of each sample in the batch. The prompt is inserted in the target-audio channel and aligned text padding is inserted into the text-token streams (input text tokens and source tokens). Args: target_text_tokens (torch.Tensor): Tensor of input text tokens with shape [B, T_text]. dtype: torch.long. target_token_lens (torch.Tensor): Lengths of target_text_tokens per batch element (before padding). shape [B]. source_tokens (torch.Tensor): Source-side text tokens, shape [B, T_src_text], dtype torch.long. source_token_lens (torch.Tensor): Source text token lengths per batch element, shape [B]. target_audio (torch.Tensor): Target-side audio waveforms, shape [B, T_audio], dtype torch.float32/float16. target_audio_lens (torch.Tensor): Target audio lengths per batch element, shape [B]. source_audio (torch.Tensor): Source-side audio waveforms, shape [B, T_audio], dtype torch.float32/float16. source_audio_lens (torch.Tensor): Source audio lengths per batch element, shape [B]. audio_prompt (torch.Tensor): Audio prompt waveforms to sample from, shape [B, T_prompt_audio]. audio_prompt_lens (torch.Tensor): Valid lengths for audio_prompt per batch element. Returns: Tuple containing: target_text_tokens (torch.Tensor): Updated text tokens with prepended prompt-aligned tokens. Shape [B, T']. target_token_lens (torch.Tensor): Updated token lengths per batch element. source_tokens (torch.Tensor): Updated source text tokens with prompt padding included. Shape [B, T']. source_token_lens (torch.Tensor): Updated source token lengths per batch element. source_audio (torch.Tensor): Updated source audio with silence padding. Shape [B, T_audio']. source_audio_lens (torch.Tensor): Updated source audio lengths. target_audio (torch.Tensor): Updated target audio with prompt audio and silence padding. Shape [B, T_audio']. target_audio_lens (torch.Tensor): Updated target audio lengths. prompt_lens (list[int]): Length (in text-token units) of the prompt region per batch item. """ text_pad_id = get_pad_id(self.tokenizer) target_text_tokens_ = [] source_tokens_ = [] source_audio_ = [] target_audio_ = [] prompt_lens = [] for i in range(target_text_tokens.size(0)): if system_prompts is not None: text_prompt = system_prompts[i][: system_prompts_lens[i]] else: text_prompt = torch.tensor( [self.tokenizer.eos], dtype=torch.long, device=target_text_tokens.device, ) if self.add_audio_prompt: # Compute audio prompt duration in samples (rounded to frame boundaries) prompt_audio_size = int( ((self.audio_prompt_duration * self.target_sample_rate) // self.target_samples_per_frame) * self.target_samples_per_frame ) prompt_audio = sample_audio_segments_repeat( audio_prompt, audio_prompt_lens, prompt_audio_size, sample=True ) # add silence at the end of the prompt prompt_audio[:, -int(self.target_samples_per_frame * 2) :] = 0 # Number of text tokens to insert to align with prompt_audio frames prompt_audio_text_pad_size = prompt_audio_size // self.target_samples_per_frame prompt_audio_text_pad = ( torch.ones( prompt_audio_text_pad_size, device=target_text_tokens.device, dtype=target_text_tokens.dtype, ) * text_pad_id ) prompt_audio_text_pad[-1] = self.tokenizer.eos new_target_text_tokens = torch.cat( [ text_prompt.to(target_text_tokens.dtype), prompt_audio_text_pad, target_text_tokens[i], ] ) target_text_tokens_.append(new_target_text_tokens) target_token_lens[i] += len(text_prompt) + prompt_audio_text_pad_size new_source_tokens = torch.cat([text_prompt, prompt_audio_text_pad, source_tokens[i]]) source_tokens_.append(new_source_tokens) source_token_lens[i] += len(text_prompt) + prompt_audio_text_pad_size # Silence in source audio during prompt processing pad_size_src = (len(text_prompt) * self.source_samples_per_frame) + prompt_audio.size(1) pad_audio_src = torch.zeros( pad_size_src, device=source_audio.device, dtype=source_audio.dtype, ) source_audio_.append(torch.cat([pad_audio_src, source_audio[i]])) source_audio_lens[i] += pad_size_src # Insert prompt audio in the target channel pad_size_tgt = len(text_prompt) * self.target_samples_per_frame pad_audio_tgt = torch.zeros( pad_size_tgt, device=target_audio.device, dtype=target_audio.dtype, ) target_audio_.append(torch.cat([pad_audio_tgt, prompt_audio[i], target_audio[i]])) target_audio_lens[i] += pad_size_tgt + prompt_audio.size(1) prompt_lens.append(len(text_prompt) + prompt_audio_text_pad_size - 1) else: # Add only a single text-frame (EOS) as prompt target_text_tokens_.append(torch.cat([text_prompt, target_text_tokens[i]])) target_token_lens[i] += len(text_prompt) source_tokens_.append(torch.cat([text_prompt, source_tokens[i]])) source_token_lens[i] += len(text_prompt) pad_size_src = len(text_prompt) * self.source_samples_per_frame pad_audio_src = torch.zeros( pad_size_src, device=source_audio.device, dtype=source_audio.dtype, ) source_audio_.append(torch.cat([pad_audio_src, source_audio[i]])) source_audio_lens[i] += pad_size_src pad_size_tgt = len(text_prompt) * self.target_samples_per_frame pad_audio_tgt = torch.zeros( pad_size_tgt, device=target_audio.device, dtype=target_audio.dtype, ) target_audio_.append(torch.cat([pad_audio_tgt, target_audio[i]])) target_audio_lens[i] += pad_size_tgt prompt_lens.append(len(text_prompt)) target_text_tokens = collate_vectors(target_text_tokens_, padding_value=text_pad_id) source_tokens = collate_vectors(source_tokens_, padding_value=text_pad_id) source_audio = collate_vectors(source_audio_, padding_value=0) target_audio = collate_vectors(target_audio_, padding_value=0) return ( target_text_tokens, target_token_lens, source_tokens, source_token_lens, source_audio, source_audio_lens, target_audio, target_audio_lens, prompt_lens, ) def add_speech_delay( source_audio: torch.Tensor, source_audio_lens: torch.Tensor, target_audio: torch.Tensor, target_audio_lens: torch.Tensor, num_delay_speech_tokens: int, target_samples_per_frame: int, source_samples_per_frame: int, ): """ Apply a speech delay by padding audio waveforms based on the number of delay speech tokens. Behavior: - Target audio is *left padded* to force the model to predict initial silence. - Source audio is *right padded* to maintain size consistency for attention alignment. Args: source_audio (FloatTensor [B, T_src]): Source/input audio waveforms. source_audio_lens (LongTensor [B]): Lengths of source audio in samples. target_audio (FloatTensor [B, T_tgt]): Target/output audio waveforms. target_audio_lens (LongTensor [B]): Lengths of target audio in samples. num_delay_speech_tokens (int): Number of delay tokens inserted on the text side. target_samples_per_frame (int): Number of audio samples per frame for target audio. source_samples_per_frame (int): Number of audio samples per frame for source audio. Returns: Tuple containing: - source_audio (FloatTensor [B, T_src + pad]) - source_audio_lens (LongTensor [B]) - target_audio (FloatTensor [B, T_tgt + pad]) - target_audio_lens (LongTensor [B]) """ # Compute target-side left padding for the delay extra_target_samples = int(num_delay_speech_tokens * target_samples_per_frame) # Left-pad target audio: forces model to generate silence initially target_audio = F.pad(target_audio, (extra_target_samples, 0)) target_audio_lens = target_audio_lens + extra_target_samples # Compute source-side right padding to maintain alignment extra_source_samples = int(num_delay_speech_tokens * source_samples_per_frame) # Right-pad source audio: avoids mismatch when aligning source/target frames source_audio = F.pad(source_audio, (0, extra_source_samples)) source_audio_lens = source_audio_lens + extra_source_samples return source_audio, source_audio_lens, target_audio, target_audio_lens def collate_system_prompt( cuts: CutSet, tokenizer: TokenizerSpec, ignore_data_system_prompt: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: """ Collate system prompts from cuts. System prompts should be stored in cut.custom['system_prompt']. """ pad_id = get_pad_id(tokenizer) tokens = [] system_prompts_raw = [] for c in cuts: # Check if system prompt exists in custom field if c.custom and c.custom.get("system_prompt", None) and not ignore_data_system_prompt: prompt_text = c.custom["system_prompt"] tokens.append( torch.as_tensor( [tokenizer.bos] + tokenizer.text_to_ids(prompt_text) + [tokenizer.eos], dtype=torch.long ) ) system_prompts_raw.append(prompt_text) else: # use dataset type name (defined on the config) if getattr(c, "type", None): prompt_text = c.type tokens.append( torch.as_tensor( [tokenizer.bos] + tokenizer.text_to_ids(prompt_text) + [tokenizer.eos], dtype=torch.long ) ) system_prompts_raw.append(prompt_text) else: logging.warning( "No system prompt or dataset type defined on the config! Using a eos token as system prompt!" ) # No system prompt for this cut, add just a eos that indicates that the prompt is finished tokens.append(torch.as_tensor([tokenizer.eos], dtype=torch.long)) system_prompts_raw.append("") token_lens = torch.tensor([len(tt) for tt in tokens]) tokens = collate_vectors(tokens, padding_value=pad_id) return tokens, token_lens, system_prompts_raw def get_audio_prompt( cuts: CutSet, target_sample_rate: int, roles: set[str], recording_field: str = "target_audio", ) -> tuple[torch.Tensor, torch.Tensor]: """ Retrieve an audio prompt for speaker conditioning. This function returns: - context audio if available (per-cut), - otherwise a random turn belonging to the desired speaker roles. Behavior: 1. If `cut.context_audio` exists, use it as the reference. 2. Otherwise, select a random audio turn from the same target-role speakers. Args: cuts (CutSet): Batch of cuts from which to extract reference audio. target_sample_rate (int): Sample rate to which reference audio is resampled. roles (set[str]): Set of speaker roles to sample from when selecting random turns. recording_field (str, optional): Name of the audio field in the cut ("recording", "target_audio", etc.). Used when sampling random reference turns. Returns: Tuple containing: - audio_prompt (FloatTensor [B, T]): Padded batch of reference waveforms. - audio_prompt_lens (LongTensor [B]): Lengths of each reference waveform before padding. """ # use provided context audio directly if hasattr(cuts[0], "context_audio"): audio_prompt = [] audio_prompt_lens = [] for cut in cuts: ref_audio = cut.context_audio.resample(target_sample_rate).load_audio() ref_audio = torch.tensor(ref_audio).float() # shape: [1, T] ref_audio_len = ref_audio.shape[1] audio_prompt.append(ref_audio.squeeze(0)) # [T] audio_prompt_lens.append(ref_audio_len) audio_prompt = collate_vectors(audio_prompt, padding_value=0).float() audio_prompt_lens = torch.tensor(audio_prompt_lens).long() else: # Sanitize cuts to remove out-of-bounds supervisions # this prevents crashes when sampling from truncated audio. cuts = sanitize_cuts(cuts) # sample a reference turn from the target-role speakers audio_prompt, audio_prompt_lens = collate_random_turn_audio( cuts.resample(target_sample_rate, recording_field=recording_field), roles=roles, recording_field=recording_field, ) return audio_prompt, audio_prompt_lens def sanitize_cuts(cuts: CutSet) -> CutSet: """ Adjusts supervisions to fit within the cut's truncated duration. - If a supervision extends beyond the cut end, it is TRUNCATED (using deepcopy). - If a supervision starts after the cut end, it is DROPPED. Args: cuts (CutSet): The batch of cuts to sanitize. Returns: CutSet: A new CutSet with valid, potentially truncated supervisions. """ sanitized_list = [] for cut in cuts: valid_supervisions = [] for sup in cut.supervisions: # Case 1: Supervision starts after the audio ends -> Drop if sup.start >= cut.duration: continue # Case 2: Supervision starts inside but ends after the audio -> Truncate if sup.end > cut.duration: # Calculate the remaining duration available in the audio new_duration = cut.duration - sup.start if new_duration <= 0: continue # Create a deepcopy and update duration new_sup = deepcopy(sup) new_sup.duration = new_duration valid_supervisions.append(new_sup) # Case 3: Supervision is fully inside -> Keep else: valid_supervisions.append(sup) # Update the cut with the cleaned list cut.supervisions = valid_supervisions sanitized_list.append(cut) return cuts.from_cuts(sanitized_list) def collate_random_turn_audio( cuts: CutSet, roles: set[str], recording_field: str = "target_audio", ) -> tuple[torch.Tensor, torch.Tensor]: """ Sample and collate reference audio from random speaker turns. For each cut in the batch, this function: - selects a random supervision belonging to one of the specified roles, - extracts the corresponding audio segment, - collates all segments into a padded batch. Args: cuts (CutSet): Batch of cuts to sample from. roles (set[str]): Set of speaker roles to consider when selecting random turns. recording_field (str, optional): Name of the audio field to load from the cut (e.g., "recording", "target_audio"). Returns: Tuple containing: - audio (FloatTensor [B, T]): Padded batch of sampled reference waveforms. - audio_lens (LongTensor [B]): Lengths of each waveform before padding. """ selected_turn_audios = [] selected_turn_audios_lens = [] for cut in cuts: # Filter supervisions matching roles matching_supervisions = [s for s in cut.supervisions if s.speaker in roles] # if there is no target speaker sample, prompt with a silence audio if len(matching_supervisions) == 0: # Create 5 seconds of silence target_duration = 5.0 num_samples = int(target_duration * cut.sampling_rate) # Create a zero tensor of shape [T] (assuming mono audio) silence_tensor = torch.zeros(num_samples, dtype=torch.float32) selected_turn_audios.append(silence_tensor) selected_turn_audios_lens.append(num_samples) logging.warning( "There is no target speaker supervision available on this sample! Using a silence audio as audio prompt!" ) else: # Randomly select one supervision selected_supervision = random.choice(matching_supervisions) # Truncate audio according to supervision truncated_audio = cut.truncate( offset=max(0, selected_supervision.start), duration=selected_supervision.duration ).load_custom(recording_field) selected_turn_audios.append(truncated_audio.squeeze(0)) selected_turn_audios_lens.append(truncated_audio.shape[-1]) return collate_vectors(selected_turn_audios, padding_value=0), torch.tensor(selected_turn_audios_lens) def collate_token_channel( cuts: CutSet, tokenizer: TokenizerSpec, frame_length: Seconds, roles: set[str], add_text_bos_and_eos_in_each_turn: bool = True, ) -> tuple[torch.Tensor, torch.Tensor]: """ Build and collate token channels aligned to the audio frame grid. This function converts text supervisions into frame-level token representations for each cut, pads them to a uniform length, and returns both the padded tokens and their true lengths. Args: cuts (CutSet): Batch of cuts from which to extract token channels. tokenizer (TokenizerSpec): Tokenizer used to convert text into token IDs. frame_length (Seconds): Duration of a single audio frame, used to align text tokens to audio frames. roles (set[str]): Speaker roles whose text will be included in the token channel. add_text_bos_and_eos_in_each_turn (bool, optional): Whether to insert BOS at the beginning and EOS at the end of each speaking turn. Returns: Tuple containing: - tokens (LongTensor [B, T]): Padded batch of frame-aligned token sequences. - token_lens (LongTensor [B]): Length of each sequence before padding. """ pad_id = get_pad_id(tokenizer) tokens = [ build_token_channel( c, tokenizer=tokenizer, frame_length=frame_length, roles=roles, pad_id=pad_id, add_text_bos_and_eos_in_each_turn=add_text_bos_and_eos_in_each_turn, ) 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, add_text_bos_and_eos_in_each_turn: bool = True, ) -> torch.Tensor: """ Build a frame-aligned token sequence for a single cut. This function maps speaking turns into a token channel aligned to the audio frame grid. Tokens are inserted at frame positions corresponding to supervision start times, with optional BOS and EOS insertion. Any region not covered by text is filled with `pad_id`. Args: cut (Cut): Input cut containing audio and supervisions. tokenizer (TokenizerSpec): Tokenizer used to encode text into tokens. frame_length (Seconds): Duration of one frame, used to align text to the audio grid. roles (set[str]): Speaker roles whose text should be included. pad_id (int, optional): Token ID used for padding empty frames. add_text_bos_and_eos_in_each_turn (bool, optional): Whether to insert BOS before and EOS after each supervision. Returns: tokens (LongTensor [T]): Frame-aligned token sequence for the cut. """ 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 = supervision.text if add_text_bos_and_eos_in_each_turn: text_ids = torch.as_tensor([tokenizer.bos] + tokenizer.text_to_ids(text)) else: text_ids = torch.as_tensor(tokenizer.text_to_ids(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. if add_text_bos_and_eos_in_each_turn: 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 def sample_audio_segments_repeat( prompt_audio: torch.Tensor, prompt_audio_lens: torch.Tensor, n_sample: int, sample: bool = True, ) -> torch.Tensor: """ Extract audio segments of length n_sample. If sample=True: randomly sample segments (repeating if shorter). If sample=False: always take from the beginning (repeating if shorter). Args: prompt_audio: Tensor [B, T] prompt_audio_lens: Tensor [B] with valid lengths n_sample: int, target length per segment sample: bool, whether to randomly sample (True) or take first seconds (False) Returns: Tensor [B, n_sample] """ B, T = prompt_audio.shape device = prompt_audio.device out = torch.zeros(B, n_sample, device=device, dtype=prompt_audio.dtype) for b in range(B): length = min(prompt_audio_lens[b].item(), T) # Case: empty audio if length <= 0: continue if length >= n_sample: if sample: # Random start (safe bounds) max_start = max(1, length - n_sample + 1) start = torch.randint(0, max_start, (1,), device=device).item() else: # Deterministic: take from start start = 0 out[b] = prompt_audio[b, start : start + n_sample] else: # Audio shorter than target → repeat start = 0 segment = prompt_audio[b, start:length] repeat_times = (n_sample + (length - start) - 1) // (length - start) repeated = segment.repeat(repeat_times)[:n_sample] out[b] = repeated return out