higgs_audio_collator.py

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)