import librosa import torch import torch.nn.functional as F import math import numpy as np from typing import List, Tuple, Dict, Optional from dataclasses import dataclass from transformers.models.whisper.processing_whisper import WhisperProcessor from transformers import __version__ as _tf_version from packaging.version import parse as _V def get_tf_version(): return _V(_tf_version) from packaging.version import parse as V try: from .modeling_higgs_audio_xcodec import XCodecProcessor from .utils import build_delay_pattern_mask except ImportError: from modeling_higgs_audio_xcodec import XCodecProcessor from utils import build_delay_pattern_mask # Minimal dataclass — full chatml_dataset not needed for inference from dataclasses import dataclass, field from typing import Optional import torch @dataclass class ChatMLDatasetSample: input_ids: torch.Tensor label_ids: torch.Tensor audio_ids_concat: Optional[torch.Tensor] = None audio_ids_start: Optional[torch.Tensor] = None audio_waveforms_concat: Optional[torch.Tensor] = None audio_waveforms_start: Optional[torch.Tensor] = None audio_sample_rate: Optional[torch.Tensor] = None audio_speaker_indices: Optional[torch.Tensor] = None audio_label_ids_concat: Optional[torch.Tensor] = None audio_cuts: Optional[list] = None # Training-only fields (None for inference) reward: Optional[float] = None group_id: Optional[int] = None group_score_sample: Optional[torch.Tensor] = None max_score_sample: Optional[torch.Tensor] = None min_score_sample: Optional[torch.Tensor] = None def get_wv(self, idx): """Extract the idx-th waveform from the concatenated buffer.""" wv_start = self.audio_waveforms_start[idx] sr = self.audio_sample_rate[idx] if idx < len(self.audio_waveforms_start) - 1: wv_end = self.audio_waveforms_start[idx + 1] else: wv_end = self.audio_waveforms_concat.shape[-1] return self.audio_waveforms_concat[wv_start:wv_end], sr def get_audio_cuts(self, idx): """Get VAD cuts for the idx-th audio segment.""" if self.audio_cuts is None: return None return self.audio_cuts[idx] def num_audios(self): """Number of audio segments in this sample.""" if self.audio_waveforms_start is None: return 0 return len(self.audio_waveforms_start) def get_audio_codes(self, idx): """Get audio codes for the idx-th segment (not used for inference).""" return None def get_audio_codes_labels(self, idx): """Get audio code labels for the idx-th segment (not used for inference).""" return None RankedChatMLDatasetSampleTuple = None # not used for inference def _ceil_to_nearest(n, round_to): if round_to <= 1: return n # avoid division by zero return (n + round_to - 1) // round_to * round_to def _ceil_to_next_power_of_two(x): return 1 if x == 0 else 2 ** (x - 1).bit_length() def fix_and_round_time_dim( audio_features: torch.Tensor, # (B, C, T_feat) attention_mask: torch.Tensor | None, # (B, T_mask) or None round_to: int = 4, pad_value: float = 0.0, ) -> tuple[torch.Tensor, torch.Tensor | None]: """ Make feature & mask time dims match and pad to a multiple of `round_to`. Also resolves Whisper's 'longest' oddity where T_mask = T_feat +/- 1. Returns: (features_fixed, mask_fixed_or_none) """ assert audio_features.dim() == 3, "audio_features must be (B, C, T)" B, C, T_feat = audio_features.shape dev = audio_features.device dtype_feat = audio_features.dtype if attention_mask is not None: assert attention_mask.dim() == 2, "attention_mask must be (B, T)" assert attention_mask.size(0) == B, "batch size mismatch" T_mask = attention_mask.size(1) else: T_mask = T_feat # 1) First, equalize T between features and mask (crop/pad conservatively). T = max(T_feat, T_mask) if attention_mask is None: with torch.no_grad(): col_nonpad = (audio_features != pad_value).any(dim=1) # (B, T_feat) attention_mask = col_nonpad.to(dtype=torch.int32) # Now reconcile off-by-one cases first (common Whisper bug). if T_mask == T_feat + 1: if (attention_mask[:, -1] == 0).all(): attention_mask = attention_mask[:, :-1] T_mask -= 1 else: audio_features = F.pad(audio_features, (0, 1), value=pad_value) T_feat += 1 elif T_feat == T_mask + 1: if (audio_features[:, :, -1] == pad_value).all(): audio_features = audio_features[:, :, :-1] T_feat -= 1 else: attention_mask = F.pad(attention_mask, (0, 1), value=0) T_mask += 1 T = max(T_feat, T_mask) if T_feat < T: audio_features = F.pad(audio_features, (0, T - T_feat), value=pad_value) T_feat = T if T_mask < T: attention_mask = F.pad(attention_mask, (0, T - T_mask), value=0) T_mask = T # 2) Round up to multiple of `round_to` if round_to is not None and round_to > 1: T_rounded = _ceil_to_nearest(T, round_to) if T_rounded != T: audio_features = F.pad(audio_features, (0, T_rounded - T), value=pad_value) attention_mask = F.pad(attention_mask, (0, T_rounded - T), value=0) T = T_rounded audio_features = audio_features.to(dtype=dtype_feat, device=dev) attention_mask = attention_mask.to(dtype=torch.int32, device=dev) assert audio_features.size(-1) == attention_mask.size(-1), "time dims should now match" return audio_features, attention_mask from accelerate.logging import get_logger logger = get_logger(__name__) @dataclass class HiggsAudioBatchInput: input_ids: torch.LongTensor # shape (bsz, seq_len). attention_mask: torch.Tensor # shape (bsz, seq_len). audio_features: Optional[torch.Tensor] # shape (num_audio_in, feature_dim, max_mel_seq_len). audio_feature_attention_mask: Optional[torch.Tensor] # shape (num_audio_in, max_mel_seq_len). audio_out_ids: Optional[torch.LongTensor] # shape (num_codebooks, audio_out_total_length) audio_out_ids_start: Optional[torch.LongTensor] # shape (num_audio_out,) # The audio_out_ids_start_group_loc has the same length as audio_out_ids_start. It is used to recover group location in a batch for an audio segment # Currently, we concatenante audio segments along dim 0 to handle variadic audio segment length. However, in the alignment stage, we need the location information # For example, # audio_out_ids_start = [0, 2, 4, 8]; and the first two audio segments come from the same sample in a batch, and other two come from different samples. # This is a batch of 3 samples, then we will have the group location as: # audio_out_ids_start_group_loc = [0, 0, 1, 2] audio_out_ids_start_group_loc: Optional[torch.LongTensor] # shape (num_audio_out,), specify which a sample's group location in the batch audio_in_ids: Optional[torch.LongTensor] # shape (num_codebooks, audio_in_total_length) audio_in_ids_start: Optional[torch.LongTensor] # shape (num_audio_in,) label_ids: Optional[torch.LongTensor] # shape (bsz, seq_len) label_audio_ids: Optional[torch.LongTensor] # shape (num_codebooks, audio_out_total_length) reward: Optional[float] = None group_id: Optional[torch.LongTensor] = None audio_wv_lengths: Optional[torch.LongTensor] = None # shape (num_audio_in,) - true 24k waveform lengths per audio chunk (XCodec path) class HiggsAudioSampleCollator: """Sample collator for Higgs-Audio model. Args: whisper_processor (WhisperProcessor): The whisper processor. audio_in_token_id (int): The token id for audio-in. audio_out_token_id (int): The token id for audio-out. pad_token_id (int): The token id for padding. audio_stream_bos_id (int): The token id for audio-stream beginning of sentence. audio_stream_eos_id (int): The token id for audio-stream end of sentence. round_to (int): The round-to value. pad_left (bool): Whether to pad left. return_audio_in_tokens (bool): Whether to return audio-in tokens. use_delay_pattern (bool): Whether to use delay pattern. disable_audio_codes_transform (bool): Whether to add bos and eos tokens to audio codes. chunk_size_seconds (int): The chunk size in seconds. add_new_bos_eos_for_long_chunk (bool): Whether to add new bos and eos tokens for long chunks. mask_audio_out_token_label (bool): Whether to always mask the label associated with <|AUDIO_OUT|> token. Since we will always have `<|AUDIO_OUT|>` after `<|audio_bos|>`, we can safely mask <|AUDIO_OUT|>. """ def __init__( self, whisper_processor: WhisperProcessor | XCodecProcessor, # FIXME respect previous signature but allow XCodecProcessor audio_in_token_id, audio_out_token_id, pad_token_id, audio_stream_bos_id, audio_stream_eos_id, round_to=8, pad_left=False, encode_whisper_embed=True, # FIXME respect previous signature, use encoder_backend to determine if we use whisper or xcodec return_audio_in_tokens=True, audio_num_codebooks=None, use_delay_pattern=False, disable_audio_codes_transform=False, chunk_size_seconds=30, # Maximum duration for each chunk add_new_bos_eos_for_long_chunk=True, mask_audio_out_token_label=True, encoder_padding_method='max_length', duplicate_audio_triplet: bool = False, vad_cut: bool = False, split_vads: bool = False, ): self.encoder_processor = whisper_processor if isinstance(self.encoder_processor, WhisperProcessor): self.encoder_backend = "whisper" elif isinstance(self.encoder_processor, XCodecProcessor): self.encoder_backend = "xcodec" else: raise ValueError(f"Invalid encoder processor: {self.encoder_processor}") self.round_to = round_to self.pad_left = pad_left self.audio_in_token_id = audio_in_token_id self.audio_out_token_id = audio_out_token_id self.audio_stream_bos_id = audio_stream_bos_id self.audio_stream_eos_id = audio_stream_eos_id self.pad_token_id = pad_token_id self.encode_whisper_embed = encode_whisper_embed self.return_audio_in_tokens = return_audio_in_tokens self.audio_num_codebooks = audio_num_codebooks self.use_delay_pattern = use_delay_pattern self.encoder_padding_method = encoder_padding_method self.duplicate_audio_triplet = duplicate_audio_triplet if encode_whisper_embed: self.chunk_size_seconds = chunk_size_seconds self.chunk_size_samples = int(chunk_size_seconds * self.encoder_processor.feature_extractor.sampling_rate) else: self.chunk_size_seconds = None self.chunk_size_samples = None self.disable_audio_codes_transform = disable_audio_codes_transform self.add_new_bos_eos_for_long_chunk = add_new_bos_eos_for_long_chunk self.mask_audio_out_token_label = mask_audio_out_token_label self.vad_cut = vad_cut self.split_vads = split_vads def _duplicate_audio_tokens_helper(self, input_ids, audio_idx, labels, num_chunks) -> Tuple[torch.Tensor, torch.Tensor]: """Duplicate audio tokens helper.""" if self.duplicate_audio_triplet: # ----- OLD behavior: repeat the triplet ----- triplet = input_ids[audio_idx-1:audio_idx+2] # <|audio_bos|><|AUDIO|><|audio_eos|> duplicated = triplet.repeat(num_chunks) # length = 3 * num_chunks new_input_ids = torch.cat([input_ids[:audio_idx - 1], duplicated, input_ids[audio_idx + 2:]], dim=0) new_labels = None if labels is not None: triplet_lbl = labels[audio_idx-1:audio_idx+2] duplicated_labels = triplet_lbl.repeat(num_chunks) new_labels = torch.cat([labels[:audio_idx - 1], duplicated_labels, labels[audio_idx + 2:]], dim=0) else: # ----- NEW behavior: duplicate only <|AUDIO|> ----- audio_tok = input_ids[audio_idx: audio_idx + 1] # <|AUDIO|> duplicated = audio_tok.repeat(num_chunks) # length = num_chunks new_input_ids = torch.cat([input_ids[:audio_idx], duplicated, input_ids[audio_idx + 1:]], dim=0) new_labels = None if labels is not None: audio_lbl = labels[audio_idx: audio_idx + 1] duplicated_labels = audio_lbl.repeat(num_chunks) new_labels = torch.cat([labels[:audio_idx], duplicated_labels, labels[audio_idx + 1:]], dim=0) return new_input_ids, new_labels def _chunk_post_process_helper( self, ranges: List[Tuple[int, int]], ) -> List[Tuple[int, int]]: """ Given (start, end) ranges on the original waveform, split each so that every output range has length <= self.chunk_size_samples. """ out: List[Tuple[int, int]] = [] max_len = int(self.chunk_size_samples) for s, e in ranges: L = e - s if L == 0: continue # skip any degenerate chunk pos = s while pos < e: nxt = min(e, pos + max_len) out.append((pos, nxt)) pos = nxt return out def _process_and_duplicate_audio_tokens(self, input_ids: torch.Tensor, audio_idx: int, wv: torch.Tensor, sr: int, labels: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, torch.Tensor, int]: """Process long audio and duplicate corresponding audio tokens. Args: input_ids: Input token ids audio_idx: Index of the audio token in the sequence wv: Audio waveform sr: Sample rate labels: Optional label ids to be duplicated alongside input ids Returns: Tuple of: - New input ids with duplicated audio tokens - New label ids (if labels were provided) or None - Number of chunks created """ # Calculate number of chunks needed total_samples = len(wv) num_chunks = math.ceil(total_samples / self.chunk_size_samples) if num_chunks <= 1: return input_ids, labels, 1 new_input_ids, new_labels = self._duplicate_audio_tokens_helper(input_ids, audio_idx, labels, num_chunks) return new_input_ids, new_labels, num_chunks def _vad_process_audio_chunks(self, input_ids: torch.Tensor, audio_idx: int, wv: torch.Tensor, sr: int, labels: torch.Tensor, audio_cuts: List[Tuple[int, int]]) -> Tuple[torch.Tensor, torch.Tensor, int, List[torch.Tensor]]: """Process long audio and duplicate corresponding audio tokens. Similar to _process_and_duplicate_audio_tokens, but use VAD to cut the audio into chunks and respects the VAD boundaries. It also directly return the audio chunks while duplicating the audio tokens. If self.duplicate_audio_triplet == True (default / OLD): - repeat the full triplet <|audio_bos|><|AUDIO|><|audio_eos|> for each chunk Else (NEW): - keep one pair of <|audio_bos|> ... <|audio_eos|> and duplicate only <|AUDIO|> in between """ total_samples = len(wv) if self.split_vads: # if vad is [0,1] [2,4] [5,8] then return [0,1] [2,4] [5,8] aka respect vad completely and drop audios in between wv_chunks = audio_cuts else: # if vad is [0,1] [2,4] [5,8] then return [0,1] [1,4] [4,8] aka respect the start only and use audios in between wv_chunks = [] prev_e = 0 for idx, (start_t, end_t) in enumerate(audio_cuts): s, e = min(prev_e, start_t), end_t if idx == len(audio_cuts) - 1: # last chunk e = total_samples wv_chunks.append((s, e)) prev_e = e wv_chunks = self._chunk_post_process_helper(wv_chunks) new_input_ids, new_labels = self._duplicate_audio_tokens_helper(input_ids, audio_idx, labels, len(wv_chunks)) return new_input_ids, new_labels, len(wv_chunks), wv_chunks def __call__(self, batch: List[ChatMLDatasetSample]): """Collate the input data with support for long audio processing.""" label_ids = None label_audio_ids = None if all([ele.label_ids is None for ele in batch]): return_labels = False else: return_labels = True if self.encode_whisper_embed: # Process each sample in the batch to handle long audio # TODO(?) The implementation here can be optimized. processed_batch = [] audio_in_wv_l_in_batch = [] for i in range(len(batch)): sample = batch[i] audio_in_mask = sample.input_ids == self.audio_in_token_id audio_in_indices = torch.where(audio_in_mask)[0] audio_out_mask = sample.input_ids == self.audio_out_token_id # Process each audio token and duplicate if needed modified_input_ids = sample.input_ids modified_labels = sample.label_ids if return_labels else None modified_waveforms_concat = [] modified_waveforms_start = [] modified_sample_rate = [] offset = 0 # Track position changes from duplicating tokens curr_wv_offset = 0 # Process input audio tokens for idx, audio_idx in enumerate(audio_in_indices): # Get the audio for this token wv, sr = sample.get_wv(idx) # Use idx since we want the original audio index # Get the audio vad cuts for this token - fallback to full audio if not available (trigger default behavior) audio_cuts = sample.get_audio_cuts(idx) if self.vad_cut else None if audio_cuts is None or len(audio_cuts) == 0: audio_cuts = [(0, len(wv))] # No resample should happen in collator, it should be done in dataloader assert sr == self.encoder_processor.feature_extractor.sampling_rate, f"Sample rate mismatch: {sr} != {self.encoder_processor.feature_extractor.sampling_rate}" wv = wv.contiguous() # Process and duplicate tokens if necessary token_pos = audio_idx + offset if audio_cuts is not None: modified_input_ids, modified_labels, num_chunks, chunks_ranges = self._vad_process_audio_chunks( modified_input_ids, token_pos, wv, sr, modified_labels, audio_cuts, ) # Update audio data for chunk_idx in range(num_chunks): chunk_start, chunk_end = chunks_ranges[chunk_idx] chunk_wv = wv[chunk_start:chunk_end] audio_in_wv_l_in_batch.append(chunk_wv) modified_waveforms_concat.append(chunk_wv) modified_waveforms_start.append(curr_wv_offset) curr_wv_offset += len(chunk_wv) modified_sample_rate.append(sr) else: raise ValueError(f"Never expected to reach here") # Update offset for next iteration - each extra chunk inserts this many new tokens at the splice point per_chunk_insert = 3 if self.duplicate_audio_triplet else 1 offset += (num_chunks - 1) * per_chunk_insert # Create new sample with modified tokens and audio data processed_sample = ChatMLDatasetSample( input_ids=modified_input_ids, label_ids=modified_labels if return_labels else sample.label_ids, audio_ids_concat=sample.audio_ids_concat, audio_ids_start=sample.audio_ids_start, audio_waveforms_concat=torch.cat(modified_waveforms_concat) if modified_waveforms_concat else sample.audio_waveforms_concat, audio_waveforms_start=torch.tensor(modified_waveforms_start, dtype=torch.long) if modified_waveforms_start else sample.audio_waveforms_start, audio_sample_rate=torch.tensor(modified_sample_rate) if modified_sample_rate else sample.audio_sample_rate, audio_speaker_indices=torch.tensor([]), # FIXME(sxjscience): The logic here is not correct for audio_label_ids_concat. audio_label_ids_concat=sample.audio_label_ids_concat, ) # audio_in_chunk_len = len(torch.where(modified_input_ids == self.audio_in_token_id)[0]) # assert audio_in_chunk_len == processed_sample.num_audios(), f"Mismatch: audio_in_chunk_len={audio_in_chunk_len}, processed_sample.num_audios()={processed_sample.num_audios()}" processed_batch.append(processed_sample) else: processed_batch = batch # Get the max sequence length based on processed batch max_seq_length = _ceil_to_nearest(max([len(sample.input_ids) for sample in processed_batch]), self.round_to) # Get the ids for audio-in and audio-out for each batch audio_in_wv_l = [] audio_in_ids_l = [] audio_out_ids_l = [] audio_out_ids_group_loc_l = [] audio_in_label_ids_l = None audio_out_label_ids_l = None reward_l = [] group_id_l = [] if return_labels: audio_out_no_train_flag = [] # Whether the audio-out data should be trained on or not. # Process the audio inputs and outputs for i in range(len(processed_batch)): audio_in_mask = processed_batch[i].input_ids == self.audio_in_token_id audio_out_mask = processed_batch[i].input_ids == self.audio_out_token_id audio_ids = torch.ones_like(processed_batch[i].input_ids) audio_ids[audio_in_mask ^ audio_out_mask] = torch.cumsum(audio_ids[audio_in_mask ^ audio_out_mask], 0) - 1 audio_in_ids = audio_ids[audio_in_mask] audio_out_ids = audio_ids[audio_out_mask] if return_labels: audio_out_no_train_flag.append(processed_batch[i].label_ids[audio_out_mask] < 0) if self.mask_audio_out_token_label: processed_batch[i].label_ids[audio_out_mask] = -100 # Process audio inputs if self.return_audio_in_tokens: audio_in_ids_l.extend([processed_batch[i].get_audio_codes(idx)[:self.audio_num_codebooks, :] for idx in audio_in_ids]) if processed_batch[i].audio_label_ids_concat is not None: if audio_in_label_ids_l is None: audio_in_label_ids_l = [] audio_in_label_ids_l.extend([processed_batch[i].get_audio_codes_labels(idx)[:self.audio_num_codebooks, :] for idx in audio_in_ids]) audio_out_ids_l.extend([processed_batch[i].get_audio_codes(idx)[:self.audio_num_codebooks, :] for idx in audio_out_ids]) audio_out_ids_group_loc_l.append(i) if processed_batch[i].reward is not None: reward_l.append(processed_batch[i].reward) if processed_batch[i].group_id is not None: group_id_l.append(processed_batch[i].group_id) if processed_batch[i].audio_label_ids_concat is not None: if audio_out_label_ids_l is None: audio_out_label_ids_l = [] audio_out_label_ids_l.extend([processed_batch[i].get_audio_codes_labels(idx)[:self.audio_num_codebooks, :] for idx in audio_out_ids]) if self.encode_whisper_embed: for chunk in audio_in_wv_l_in_batch: # this is for the whole batch audio_in_wv_l.append(chunk.cpu().numpy()) if return_labels: audio_out_no_train_flag = torch.cat(audio_out_no_train_flag, dim=0) # Process all audio features audio_wv_lengths = None audio_feature_attention_mask = None if len(audio_in_wv_l) > 0: if self.encoder_backend == "whisper": feature_ret = self.encoder_processor.feature_extractor( audio_in_wv_l, sampling_rate=self.encoder_processor.feature_extractor.sampling_rate, return_attention_mask=True, padding=self.encoder_padding_method, ) audio_features = torch.from_numpy(feature_ret["input_features"]) audio_feature_attention_mask = torch.from_numpy(feature_ret["attention_mask"]) if self.encoder_padding_method == "longest": # for max_length, we don't need to fix the time dim audio_features, audio_feature_attention_mask = fix_and_round_time_dim( audio_features, audio_feature_attention_mask, round_to=16, # FIXME - hard code to 16 to avoid SP issue later pad_value=0.0, ) elif self.encoder_backend == "xcodec": feature_ret = self.encoder_processor.feature_extractor( audio_in_wv_l, sampling_rate=self.encoder_processor.feature_extractor.sampling_rate, return_attention_mask=False, padding=self.encoder_padding_method, ) audio_features = torch.from_numpy(feature_ret["input_features"]).to(dtype=torch.bfloat16) audio_wv_lengths = torch.from_numpy(feature_ret["audio_wv_lengths"]) else: raise ValueError(f"Invalid encoder backend: {self.encoder_backend}") else: if self.encode_whisper_embed: if self.encoder_backend == "whisper": audio_features = torch.zeros( (0, self.encoder_processor.feature_extractor.feature_size, self.encoder_processor.feature_extractor.nb_max_frames), dtype=torch.float32 ) audio_feature_attention_mask = torch.zeros( (0, self.encoder_processor.feature_extractor.nb_max_frames), dtype=torch.int32 ) elif self.encoder_backend == "xcodec": # xcodec processor can handle empty input feature_ret = self.encoder_processor.feature_extractor( audio_in_wv_l, sampling_rate=self.encoder_processor.feature_extractor.sampling_rate, return_attention_mask=False, padding=self.encoder_padding_method, ) audio_features = torch.from_numpy(feature_ret["input_features"]) audio_wv_lengths = torch.from_numpy(feature_ret["audio_wv_lengths"]) else: raise ValueError(f"Invalid encoder backend: {self.encoder_backend}") else: audio_features = None audio_feature_attention_mask = None # Process audio input tokens if len(audio_in_ids_l) > 0: # Append audio-stream-bos and eos tokens new_audio_in_ids_l = [] for ele in audio_in_ids_l: if self.disable_audio_codes_transform: # Do not add audio-stream-bos or eos tokens. # This may indicate that the sample comes from ConstantLengthDatasetWithBuffer. audio_codes = ele else: audio_codes = torch.cat( [ torch.full((ele.shape[0], 1), self.audio_stream_bos_id, dtype=torch.long), ele, torch.full((ele.shape[0], 1), self.audio_stream_eos_id, dtype=torch.long) ], dim=1 ) if self.use_delay_pattern: audio_codes = build_delay_pattern_mask( audio_codes.unsqueeze(0), bos_token_id=self.audio_stream_bos_id, pad_token_id=self.audio_stream_eos_id )[0].squeeze(0) new_audio_in_ids_l.append(audio_codes) audio_in_ids = torch.cat(new_audio_in_ids_l, dim=1).long() audio_in_ids_start = torch.cumsum(torch.tensor([0] + [audio_codes.shape[1] for audio_codes in new_audio_in_ids_l[:-1]]), dim=0) else: audio_in_ids = torch.zeros((0, 0), dtype=torch.long) audio_in_ids_start = torch.zeros(0, dtype=torch.long) # Process audio output tokens audio_out_ids_start_group_loc = None if len(audio_out_ids_l) > 0: new_audio_out_ids_l = [] label_audio_ids_l = [] for idx, ele in enumerate(audio_out_ids_l): if self.disable_audio_codes_transform: # Do not add audio-stream-bos or eos tokens. # This may indicate that the sample comes from ConstantLengthDatasetWithBuffer. audio_codes = ele if return_labels: label_audio_ids = audio_out_label_ids_l[idx] else: audio_codes = torch.cat( [ torch.full((ele.shape[0], 1), self.audio_stream_bos_id, dtype=torch.long), ele, torch.full((ele.shape[0], 1), self.audio_stream_eos_id, dtype=torch.long) ], dim=1 ) if return_labels: label_audio_ids = torch.cat( [ torch.full((ele.shape[0], 1), -100, dtype=torch.long), ele, torch.full((ele.shape[0], 1), self.audio_stream_eos_id, dtype=torch.long) ], dim=1 ) if self.use_delay_pattern: audio_codes = build_delay_pattern_mask( audio_codes.unsqueeze(0), bos_token_id=self.audio_stream_bos_id, pad_token_id=self.audio_stream_eos_id )[0].squeeze(0) if return_labels: label_audio_ids = build_delay_pattern_mask( label_audio_ids.unsqueeze(0), bos_token_id=-100, pad_token_id=-100, )[0].squeeze(0) new_audio_out_ids_l.append(audio_codes) if return_labels: if audio_out_no_train_flag[idx]: label_audio_ids[:] = -100 label_audio_ids_l.append(label_audio_ids) audio_out_ids = torch.cat(new_audio_out_ids_l, dim=1).long() if return_labels: label_audio_ids = torch.cat(label_audio_ids_l, dim=1).long() audio_out_ids_start = torch.cumsum( torch.tensor([0] + [audio_codes.shape[1] for audio_codes in new_audio_out_ids_l[:-1]]), dim=0 ) audio_out_ids_start_group_loc = torch.tensor(audio_out_ids_group_loc_l, dtype=torch.long) else: audio_out_ids = torch.zeros((0, 0), dtype=torch.long) audio_out_ids_start = torch.zeros(0, dtype=torch.long) if return_labels: label_audio_ids = torch.zeros((0, 0), dtype=torch.long) reward = torch.tensor(reward_l, dtype=torch.float32) group_id = torch.tensor(group_id_l, dtype=torch.long) # Handle padding for input ids and attention mask if self.pad_left: input_ids = torch.stack( [F.pad(ele.input_ids, (max_seq_length - len(ele.input_ids), 0), value=self.pad_token_id) for ele in processed_batch] ) if return_labels: label_ids = torch.stack( [F.pad(ele.label_ids, (max_seq_length - len(ele.label_ids), 0), value=-100) for ele in processed_batch] ) attention_mask = torch.stack( [F.pad(torch.ones_like(ele.input_ids), (max_seq_length - len(ele.input_ids), 0), value=0) for ele in processed_batch] ) else: input_ids = torch.stack( [F.pad(ele.input_ids, (0, max_seq_length - len(ele.input_ids)), value=self.pad_token_id) for ele in processed_batch] ) if return_labels: label_ids = torch.stack( [F.pad(ele.label_ids, (0, max_seq_length - len(ele.label_ids)), value=-100) for ele in processed_batch] ) attention_mask = torch.stack( [F.pad(torch.ones_like(ele.input_ids), (0, max_seq_length - len(ele.input_ids)), value=0) for ele in processed_batch] ) if not self.return_audio_in_tokens: audio_in_ids = None audio_in_ids_start = None # Apply audio_num_codebooks limit if specified if self.audio_num_codebooks is not None: if audio_in_ids is not None: audio_in_ids = audio_in_ids[:self.audio_num_codebooks] if audio_out_ids is not None: audio_out_ids = audio_out_ids[:self.audio_num_codebooks] if label_audio_ids is not None: label_audio_ids = label_audio_ids[:self.audio_num_codebooks] return HiggsAudioBatchInput( input_ids=input_ids, attention_mask=attention_mask, audio_features=audio_features, audio_feature_attention_mask=audio_feature_attention_mask, audio_out_ids=audio_out_ids, audio_out_ids_start=audio_out_ids_start, audio_out_ids_start_group_loc=audio_out_ids_start_group_loc, audio_in_ids=audio_in_ids, audio_in_ids_start=audio_in_ids_start, label_ids=label_ids, label_audio_ids=label_audio_ids, reward=reward, group_id=group_id, audio_wv_lengths=audio_wv_lengths, ) class HiggsAudioDPOSamplesCollator(HiggsAudioSampleCollator): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) def __call__(self, batch: List[RankedChatMLDatasetSampleTuple]) -> HiggsAudioBatchInput: # flatten ranked chatml samples chosen = [] rejected = [] for sample in batch: chosen.append(sample.max_score_sample()) rejected.append(sample.min_score_sample()) merged = chosen merged.extend(rejected) return super().__call__(batch=merged) class HiggsAudioGRPOSamplesCollator(HiggsAudioSampleCollator): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) def __call__(self, batch: List[RankedChatMLDatasetSampleTuple]) -> HiggsAudioBatchInput: merged = [] for idx, sample in enumerate(batch): merged.extend(sample.group_score_sample(idx)) return super().__call__(batch=merged)