processing_robobrain_audio.py

import re
import types
from dataclasses import dataclass
from typing import List, Optional, Sequence, Union

import numpy as np
import torch
from transformers import BatchFeature
from transformers.processing_utils import ProcessorMixin
from transformers.models.whisper.feature_extraction_whisper import WhisperFeatureExtractor


@dataclass
class MelConfig:
    mel_sr: int = 16000
    mel_dim: int = 128
    mel_n_fft: int = 400
    mel_hop_length: int = 160
    mel_dtype: torch.dtype = torch.bfloat16
    use_whisper_feature_extractor: bool = True


def _normalize_mel_config(mel_config) -> dict:
    default_config = MelConfig()
    if mel_config is None:
        source = {}
    elif isinstance(mel_config, MelConfig):
        source = {key: getattr(mel_config, key) for key in MelConfig.__dataclass_fields__.keys()}
    else:
        source = dict(mel_config)

    normalized = {}
    for key in MelConfig.__dataclass_fields__.keys():
        value = source.get(key, getattr(default_config, key))
        if key == "mel_dtype":
            if isinstance(value, torch.dtype):
                value = str(value).removeprefix("torch.")
            elif isinstance(value, str) and value.startswith("torch."):
                value = value.removeprefix("torch.")
        normalized[key] = value
    return normalized


def _build_mel_config(mel_config_dict: dict) -> MelConfig:
    default_config = MelConfig()

    def _int_value(key, default):
        value = mel_config_dict.get(key, default)
        if isinstance(value, bool):
            return int(value)
        if isinstance(value, (int, str)):
            return int(value)
        return default

    def _bool_value(key, default):
        value = mel_config_dict.get(key, default)
        if isinstance(value, bool):
            return value
        if isinstance(value, str):
            return value.lower() in {"1", "true", "yes", "on"}
        if isinstance(value, int):
            return bool(value)
        return default

    mel_dtype_value = mel_config_dict.get("mel_dtype", default_config.mel_dtype)
    if isinstance(mel_dtype_value, str):
        mel_dtype = getattr(torch, mel_dtype_value.removeprefix("torch."))
    elif isinstance(mel_dtype_value, torch.dtype):
        mel_dtype = mel_dtype_value
    else:
        mel_dtype = default_config.mel_dtype

    return MelConfig(
        mel_sr=_int_value("mel_sr", default_config.mel_sr),
        mel_dim=_int_value("mel_dim", default_config.mel_dim),
        mel_n_fft=_int_value("mel_n_fft", default_config.mel_n_fft),
        mel_hop_length=_int_value("mel_hop_length", default_config.mel_hop_length),
        mel_dtype=mel_dtype,
        use_whisper_feature_extractor=_bool_value("use_whisper_feature_extractor", default_config.use_whisper_feature_extractor),
    )


class RoboBrainAudioProcessor(ProcessorMixin):
    attributes = ["tokenizer", "image_processor"]
    tokenizer_class = ("Qwen2Tokenizer", "Qwen2TokenizerFast")
    image_processor_class = ("Qwen2VLImageProcessorFast", "Qwen2VLImageProcessor")

    _AUDIO_SPAN_RE = re.compile(r"<\|audio_bos\|>(?:<\|AUDIO\|>)+<\|audio_eos\|>")

    def __init__(
        self,
        tokenizer=None,
        image_processor=None,
        mel_config=None,
        enable_time_marker: bool = True,
        audio_token_id: int = 151654,
        audio_start_id: int = 151669,
        audio_end_id: int = 151670,
        chat_template=None,
    ):
        super().__init__(tokenizer, image_processor, chat_template=chat_template)
        if tokenizer is None:
            raise ValueError("RoboBrainAudioProcessor requires a tokenizer.")

        self._base_tokenizer = tokenizer
        self.mel_config = _normalize_mel_config(mel_config)
        self.config = _build_mel_config(self.mel_config)
        self.enable_time_marker = bool(enable_time_marker)
        self.audio_token_id = int(audio_token_id)
        self.audio_start_id = int(audio_start_id)
        self.audio_end_id = int(audio_end_id)
        self._whisper_feature_extractor = None

        alias_map = {
            "<|AUDIO|>": self.audio_token_id,
            "<|audio_bos|>": self.audio_start_id,
            "<|audio_eos|>": self.audio_end_id,
        }
        orig_convert_tokens_to_ids = tokenizer.convert_tokens_to_ids

        def _patched_convert_tokens_to_ids(tokenizer_self, tokens):
            if isinstance(tokens, (list, tuple)):
                converted = [_patched_convert_tokens_to_ids(tokenizer_self, token) for token in tokens]
                return converted if isinstance(tokens, list) else tuple(converted)
            if isinstance(tokens, str) and tokens in alias_map:
                return alias_map[tokens]
            return orig_convert_tokens_to_ids(tokens)

        tokenizer.convert_tokens_to_ids = types.MethodType(_patched_convert_tokens_to_ids, tokenizer)

        self._digit_token_ids = {str(i): 15 + i for i in range(10)}
        self.audio_tokens_per_second = 12.5
        self.time_marker_every_seconds = 2
        self.time_marker_every_audio_tokens = int(self.audio_tokens_per_second * self.time_marker_every_seconds)

    @property
    def model_input_names(self):
        return ["input_ids", "attention_mask", "pixel_values", "image_grid_thw",
                "pixel_values_videos", "video_grid_thw", "audio_data", "audio_data_seqlens"]

    @staticmethod
    def _conv3_downsample_len(raw_mel_len: int) -> int:
        def conv_out_len(length: int) -> int:
            return (length - 1) // 2 + 1
        return int(conv_out_len(conv_out_len(conv_out_len(raw_mel_len))))

    def _get_whisper_feature_extractor(self):
        if self._whisper_feature_extractor is not None:
            return self._whisper_feature_extractor
        self._whisper_feature_extractor = WhisperFeatureExtractor(
            feature_size=int(self.config.mel_dim),
            sampling_rate=int(self.config.mel_sr),
            hop_length=int(self.config.mel_hop_length),
            n_fft=int(self.config.mel_n_fft),
        )
        return self._whisper_feature_extractor

    def _extract_mel(self, audio: Union[np.ndarray, torch.Tensor]) -> torch.Tensor:
        if isinstance(audio, np.ndarray):
            wav = torch.from_numpy(audio)
        else:
            wav = audio
        wav = wav.to(dtype=torch.float32)
        if wav.dim() == 1:
            wav = wav.unsqueeze(0)

        if bool(getattr(self.config, "use_whisper_feature_extractor", False)):
            fe = self._get_whisper_feature_extractor()
            wav_np = wav.detach().to("cpu", torch.float32).contiguous().numpy()
            if wav_np.ndim == 2:
                wav_np = wav_np[0]
            feats = fe._np_extract_fbank_features(wav_np[None, ...], device="cpu")
            mel = torch.from_numpy(feats[0])
        else:
            raise ValueError("RoboBrainAudioProcessor requires whisper feature extraction.")

        return mel.to(dtype=self.config.mel_dtype)

    def _get_time_marker_token_ids(self, second: int) -> List[int]:
        return [self._digit_token_ids[digit] for digit in str(second)]

    def _build_audio_tokens_with_time_markers(self, audio_seq_len: int) -> List[int]:
        total_duration_seconds = audio_seq_len / self.audio_tokens_per_second
        num_full_seconds = int(total_duration_seconds)

        token_ids: List[int] = []
        audio_tokens_consumed = 0
        for second in range(self.time_marker_every_seconds, num_full_seconds + 1, self.time_marker_every_seconds):
            marker_pos = (second // self.time_marker_every_seconds) * self.time_marker_every_audio_tokens
            audio_segment_len = marker_pos - audio_tokens_consumed
            if audio_segment_len > 0:
                token_ids.extend([self.audio_token_id] * audio_segment_len)
                audio_tokens_consumed += audio_segment_len
            token_ids.extend(self._get_time_marker_token_ids(second))

        remaining = audio_seq_len - audio_tokens_consumed
        if remaining > 0:
            token_ids.extend([self.audio_token_id] * remaining)
        return token_ids

    def _build_audio_placeholder_ids(self, num_audio_tokens: int) -> List[int]:
        if self.enable_time_marker:
            return self._build_audio_tokens_with_time_markers(num_audio_tokens)
        return [self.audio_token_id] * num_audio_tokens

    def apply_chat_template(self, conversation, add_generation_prompt=True, image_grid_thw=None, video_grid_thw=None):
        spatial_merge_size = 2
        prompt_parts = []
        image_idx = 0
        video_idx = 0

        for msg in conversation:
            role = msg.get("role", "user")
            prompt_parts.append(f"<|im_start|>{role}\n")
            content = msg.get("content", "")
            if isinstance(content, str):
                prompt_parts.append(content)
            else:
                for item in content:
                    item_type = item.get("type", "")
                    if item_type == "image":
                        if image_grid_thw is not None and image_idx < len(image_grid_thw):
                            num_tokens = int(image_grid_thw[image_idx].prod(-1).item() // spatial_merge_size**2)
                            image_idx += 1
                        else:
                            num_tokens = 1
                        prompt_parts.append("<|vision_start|>" + "<|image_pad|>" * num_tokens + "<|vision_end|>")
                    elif item_type == "video":
                        if video_grid_thw is not None and video_idx < len(video_grid_thw):
                            num_tokens = int(video_grid_thw[video_idx].prod(-1).item() // spatial_merge_size**2)
                            video_idx += 1
                        else:
                            num_tokens = 1
                        prompt_parts.append("<|vision_start|>" + "<|video_pad|>" * num_tokens + "<|vision_end|>")
                    elif item_type == "audio":
                        prompt_parts.append("<|audio_bos|><|AUDIO|><|audio_eos|>")
                    elif item_type == "text":
                        prompt_parts.append(item.get("text", ""))
                    elif "text" in item:
                        prompt_parts.append(item["text"])
            prompt_parts.append("<|im_end|>\n")
        if add_generation_prompt:
            prompt_parts.append("<|im_start|>assistant\n")
        return "".join(prompt_parts)

    def _build_default_prompt(self, text: str, has_audio: bool, has_image: bool) -> str:
        content = []
        if has_image:
            content.append({"type": "image"})
        if has_audio:
            content.append({"type": "audio"})
        content.append({"type": "text", "text": text})

        conversation = [
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": content},
        ]
        return self.apply_chat_template(conversation, add_generation_prompt=True)

    def _build_input_from_prompt(self, prompt: str, token_lens: List[int]) -> List[int]:
        spans = list(self._AUDIO_SPAN_RE.finditer(prompt))
        if len(spans) != len(token_lens):
            raise ValueError(
                f"Audio placeholder count mismatch: found {len(spans)} spans in text, "
                f"but got {len(token_lens)} audio inputs."
            )

        input_ids: List[int] = []
        cursor = 0
        for index, match in enumerate(spans):
            prefix = prompt[cursor:match.start()]
            if prefix:
                input_ids.extend(self._base_tokenizer.encode(prefix, add_special_tokens=False))
            input_ids.append(self.audio_start_id)
            input_ids.extend(self._build_audio_placeholder_ids(int(token_lens[index])))
            input_ids.append(self.audio_end_id)
            cursor = match.end()

        suffix = prompt[cursor:]
        if suffix:
            input_ids.extend(self._base_tokenizer.encode(suffix, add_special_tokens=False))
        return input_ids

    def __call__(
        self,
        text: Union[str, Sequence[str], None] = None,
        images=None,
        videos=None,
        audios: Optional[Sequence[Union[np.ndarray, torch.Tensor]]] = None,
        audio: Optional[Sequence[Union[np.ndarray, torch.Tensor]]] = None,
        return_tensors: str = "pt",
        **kwargs,
    ) -> BatchFeature:
        audio_list = audios if audios is not None else (audio if audio is not None else [])
        audio_list = [] if audio_list is None else list(audio_list)

        image_list = images if images is not None else []
        video_list = videos if videos is not None else []

        has_audio = len(audio_list) > 0
        has_image = len(image_list) > 0 or len(video_list) > 0

        if isinstance(text, str):
            prompt_text: Optional[str] = text
        elif isinstance(text, (list, tuple)):
            prompt_text = text[0] if len(text) == 1 else text
            if isinstance(prompt_text, (list, tuple)):
                prompt_text = None
        else:
            prompt_text = None

        image_data = None
        video_data = None
        image_grid_thw = None
        video_grid_thw = None
        processed_image_grid_thw = None

        if has_image and self.image_processor is not None:
            if len(image_list) > 0:
                image_outputs = self.image_processor(images=image_list, return_tensors=return_tensors)
                image_data = image_outputs.get("pixel_values")
                image_grid_thw = image_outputs.get("image_grid_thw")
                processed_image_grid_thw = image_grid_thw
            if len(video_list) > 0:
                video_outputs = self.image_processor(videos=video_list, return_tensors=return_tensors)
                video_data = video_outputs.get("pixel_values")
                video_grid_thw = video_outputs.get("video_grid_thw")

        mels: List[torch.Tensor] = []
        raw_lengths: List[int] = []
        token_lens: List[int] = []
        audio_data = None
        audio_data_seqlens = None

        if has_audio:
            for one_audio in audio_list:
                mel = self._extract_mel(one_audio)
                raw_len = int(mel.shape[-1])
                mels.append(mel)
                raw_lengths.append(raw_len)
                token_lens.append(self._conv3_downsample_len(raw_len))

            max_length = max(raw_lengths)
            audio_batch = torch.zeros((len(mels), self.config.mel_dim, max_length), dtype=self.config.mel_dtype)
            for index, mel in enumerate(mels):
                audio_batch[index, :, :mel.shape[-1]] = mel
            audio_data = audio_batch
            audio_data_seqlens = torch.tensor(raw_lengths, dtype=torch.long)

        if prompt_text is None:
            raise ValueError("RoboBrainAudioProcessor requires text input.")

        if self._AUDIO_SPAN_RE.search(prompt_text) is None and audio_list:
            prompt_text = self._build_default_prompt(prompt_text, has_audio=has_audio, has_image=has_image)

        if has_image and processed_image_grid_thw is not None:
            spatial_merge_size = 2
            img_tokens_per_image = [int(thw.prod(-1).item() // spatial_merge_size**2) for thw in processed_image_grid_thw]
            for num_tokens in img_tokens_per_image:
                old = "<|vision_start|><|image_pad|><|vision_end|>"
                new = "<|vision_start|>" + "<|image_pad|>" * num_tokens + "<|vision_end|>"
                prompt_text = prompt_text.replace(old, new, 1)

        if has_audio:
            input_ids_list = self._build_input_from_prompt(prompt_text, token_lens)
        else:
            input_ids_list = self._base_tokenizer.encode(prompt_text, add_special_tokens=False)

        input_ids_tensor = torch.tensor([input_ids_list], dtype=torch.long)
        attention_mask_tensor = torch.ones_like(input_ids_tensor)

        data = {
            "input_ids": input_ids_tensor,
            "attention_mask": attention_mask_tensor,
        }
        if audio_data is not None and audio_data_seqlens is not None:
            data["audio_data"] = audio_data
            data["audio_data_seqlens"] = audio_data_seqlens
        if image_data is not None:
            data["pixel_values"] = image_data
            data["image_grid_thw"] = image_grid_thw
        if video_data is not None:
            data["pixel_values_videos"] = video_data
            data["video_grid_thw"] = video_grid_thw

        return BatchFeature(data=data, tensor_type=return_tensors)

    def batch_decode(self, *args, **kwargs):
        return self._base_tokenizer.batch_decode(*args, **kwargs)

    def decode(self, *args, **kwargs):
        return self._base_tokenizer.decode(*args, **kwargs)


__all__ = ["MelConfig", "RoboBrainAudioProcessor"]