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"]