| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from math import ceil |
| 5 | |
| 6 | import numpy as np |
| 7 | import torch |
| 8 | from mistral_common.tokens.tokenizers.audio import AudioEncoder |
| 9 | from transformers import ( |
| 10 | BatchFeature, |
| 11 | ProcessorMixin, |
| 12 | SequenceFeatureExtractor, |
| 13 | TensorType, |
| 14 | ) |
| 15 | from transformers.audio_utils import AudioInput |
| 16 | |
| 17 | from vllm.tokenizers.mistral import MistralTokenizer |
| 18 | |
| 19 | |
| 20 | class MistralCommonFeatureExtractor(SequenceFeatureExtractor): |
| 21 | """Provide a HF-compatible interface for |
| 22 | `mistral_common.tokens.tokenizers.multimodal.AudioEncoder`. |
| 23 | """ |
| 24 | |
| 25 | def __init__(self, audio_encoder: AudioEncoder) -> None: |
| 26 | self.audio_encoder = audio_encoder |
| 27 | |
| 28 | @property |
| 29 | def sampling_rate(self): |
| 30 | return self.audio_encoder.audio_config.sampling_rate |
| 31 | |
| 32 | @property |
| 33 | def frame_rate(self): |
| 34 | return self.audio_encoder.audio_config.frame_rate |
| 35 | |
| 36 | def __call__( |
| 37 | self, |
| 38 | audios: AudioInput, |
| 39 | return_tensors: str | TensorType | None = None, |
| 40 | **kwargs, |
| 41 | ) -> BatchFeature: |
| 42 | audios_lst = [audios] if not isinstance(audios, list) else audios |
| 43 | |
| 44 | audios_processed = list[torch.Tensor]() |
| 45 | |
| 46 | for audio in audios_lst: |
| 47 | audio = np.asarray(audio, dtype=np.float32).ravel() |
| 48 | if not self.audio_encoder.audio_config.is_streaming: |
| 49 | audio = self.audio_encoder.pad(audio, self.sampling_rate) |
| 50 | |
| 51 | audios_processed.append(torch.from_numpy(audio)) |
| 52 | |
| 53 | return BatchFeature( |
| 54 | {"audio_arrays": audios_processed}, tensor_type=return_tensors |
| 55 | ) |
| 56 | |
| 57 | def get_num_audio_tokens(self, audio_length: int) -> int: |
| 58 | return ceil(audio_length / (self.sampling_rate // self.frame_rate)) |
| 59 | |
| 60 | def fetch_audio(self, audio_url_or_urls, sampling_rate=None): |
| 61 | """HF-compatible duck-typed ``fetch_audio``. |
| 62 | |
| 63 | Mirrors :meth:`transformers.SequenceFeatureExtractor.fetch_audio` so |
| 64 | :class:`transformers.ProcessorMixin.prepare_inputs_layout` works on this |
| 65 | duck-typed feature extractor. |
| 66 | |
| 67 | Accepts the same shapes as ``SequenceFeatureExtractor.fetch_audio``: |
| 68 | |
| 69 | * ``np.ndarray`` / ``torch.Tensor`` — returned as-is. |
| 70 | * ``list[float]`` — returned as-is (a single audio sample). |
| 71 | * ``str`` URL or path — delegated to |
| 72 | :func:`transformers.audio_utils.load_audio`. |
| 73 | * ``list`` of any of the above — recursed element-wise. |
| 74 | |
| 75 | ``ProcessorMixin.prepare_inputs_layout`` always passes already-decoded |
| 76 | audio (numpy array or torch tensor), so the str / list-of-str branches |
| 77 | exist only to keep the contract identical to the upstream method. |
| 78 | """ |
| 79 | from transformers.audio_utils import is_valid_audio |
| 80 | |
| 81 | sampling_rate = sampling_rate if sampling_rate else self.sampling_rate |
| 82 | if is_valid_audio(audio_url_or_urls): |
| 83 | return audio_url_or_urls |
| 84 | if isinstance(audio_url_or_urls, (list, tuple)): |
| 85 | return [ |
| 86 | self.fetch_audio(x, sampling_rate=sampling_rate) |
| 87 | for x in audio_url_or_urls |
| 88 | ] |
| 89 | if isinstance(audio_url_or_urls, str): |
| 90 | from transformers.audio_utils import load_audio |
| 91 | |
| 92 | return load_audio(audio_url_or_urls, sampling_rate=sampling_rate) |
| 93 | raise TypeError( |
| 94 | "only a numpy array, torch tensor, str URL/path, or list of those " |
| 95 | f"is supported but got type={type(audio_url_or_urls)}" |
| 96 | ) |
| 97 | |
| 98 | |
| 99 | class MistralCommonVoxtralProcessor(ProcessorMixin): |
| 100 | attributes = ["feature_extractor", "tokenizer"] |
| 101 | |
| 102 | def __init__( |
| 103 | self, |
| 104 | tokenizer: MistralTokenizer, |
| 105 | feature_extractor: MistralCommonFeatureExtractor, |
| 106 | ) -> None: |
| 107 | self.tokenizer = tokenizer.transformers_tokenizer |
| 108 | self.feature_extractor = feature_extractor |
| 109 | |
| 110 | audio_special_ids = self.feature_extractor.audio_encoder.special_ids |
| 111 | self.audio_token_id = audio_special_ids.audio |
| 112 | self.begin_audio_token_id = audio_special_ids.begin_audio |