| """Audio processor for BailingMoE V3 Audio Hugging Face inference.""" |
|
|
| from pathlib import Path |
|
|
| import torch |
| import torchaudio |
| import soundfile as sf |
| import whisper |
| from transformers import AutoTokenizer, BatchFeature |
| from transformers.processing_utils import ProcessorMixin |
|
|
| from .configuration_bailing_moe_v3_audio import BailingMoeV3AudioConfig |
|
|
|
|
| class BailingMoeV3AudioProcessor(ProcessorMixin): |
| attributes = ["tokenizer"] |
| tokenizer_class = "AutoTokenizer" |
|
|
| ASR_PROMPT = ( |
| "Please recognize the language of this speech and transcribe it, " |
| "placing a `\t` between the identified language and the transcribe." |
| ) |
|
|
| def __init__(self, tokenizer, audio_config=None): |
| super().__init__(tokenizer) |
| self.audio_config = audio_config |
|
|
| @classmethod |
| def from_pretrained(cls, pretrained_model_name_or_path, **kwargs): |
| tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path, **kwargs) |
| config = BailingMoeV3AudioConfig.from_pretrained( |
| pretrained_model_name_or_path, **kwargs |
| ) |
| return cls(tokenizer, config.audio_config) |
|
|
| def _load_audio(self, audio): |
| if isinstance(audio, (str, Path)): |
| samples, sample_rate = sf.read(str(audio), dtype="float32", always_2d=True) |
| waveform = torch.from_numpy(samples.T.copy()) |
| elif isinstance(audio, tuple) and len(audio) == 2: |
| waveform, sample_rate = audio |
| waveform = torch.as_tensor(waveform) |
| else: |
| raise TypeError("audio must be a path or a (waveform, sample_rate) tuple") |
| if waveform.ndim == 1: |
| waveform = waveform.unsqueeze(0) |
| if waveform.ndim != 2: |
| raise ValueError("waveform must have shape [samples] or [channels, samples]") |
| waveform = waveform.float() |
| if sample_rate != 16000: |
| waveform = torchaudio.functional.resample(waveform, sample_rate, 16000) |
| return waveform.mean(dim=0) |
|
|
| def _num_audio_tokens(self, mel_length): |
| config = self.audio_config |
| chunk_size = getattr(config, "whisper_chunk_size", 3000) |
| kernels = config.downsampler_kernel_size_list |
| strides = config.downsampler_stride_list |
| paddings = getattr(config, "downsampler_padding_list", None) |
| paddings = paddings or [kernel // 2 for kernel in kernels] |
|
|
| total = 0 |
| for offset in range(0, mel_length, chunk_size): |
| length = min(chunk_size, mel_length - offset) |
| length = (length + 2 - 3) // 2 + 1 |
| for kernel, stride, padding in zip(kernels, strides, paddings): |
| length = (length + 2 * padding - kernel) // stride + 1 |
| total += length |
| return total |
|
|
| def __call__(self, audio, text=None, return_tensors="pt"): |
| messages = None |
| if isinstance(audio, list) and audio and isinstance(audio[0], dict): |
| messages = audio |
| audio_items = [ |
| item["audio"] |
| for message in messages |
| if isinstance(message.get("content"), list) |
| for item in message["content"] |
| if item.get("type") == "audio" |
| ] |
| if len(audio_items) != 1: |
| raise ValueError("exactly one audio item is required") |
| audio = audio_items[0] |
| waveform = self._load_audio(audio) |
| n_mels = getattr(self.audio_config, "n_mels", 128) |
| mel = whisper.log_mel_spectrogram(waveform, n_mels=n_mels) |
| mel = mel.transpose(0, 1).unsqueeze(0).contiguous() |
| mel_length = mel.shape[1] |
|
|
| audio_start = "<|audio_start|>" |
| audio_token = "<|audio_pad|>" |
| audio_end = "<|audio_end|>" |
| if messages is None: |
| prompt = ( |
| audio_start |
| + audio_token * self._num_audio_tokens(mel_length) |
| + audio_end |
| + (text or self.ASR_PROMPT) |
| ) |
| else: |
| prompt = self.tokenizer.apply_chat_template( |
| messages, tokenize=False, add_generation_prompt=True, |
| enable_thinking=False, |
| ) |
| placeholder = audio_start + audio_token + audio_end |
| if prompt.count(placeholder) != 1: |
| raise ValueError("chat prompt must contain exactly one audio placeholder") |
| prompt = prompt.replace( |
| placeholder, |
| audio_start + audio_token * self._num_audio_tokens(mel_length) + audio_end, |
| ) |
| inputs = self.tokenizer(prompt, return_tensors=return_tensors) |
| inputs.pop("token_type_ids", None) |
| inputs["audio_feats"] = mel |
| inputs["audio_feats_lengths"] = torch.tensor([[mel_length]], dtype=torch.long) |
| return BatchFeature(inputs) |
|
|
|
|
| __all__ = ["BailingMoeV3AudioProcessor"] |
|
|
