fix: dual-mode imports + complete ChatMLDatasetSample + transcribe output parsing

higgs_audio_collator.py

@@ -0,0 +1,722 @@
+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)