Yuanheng Zhao
c0ccbb872f
Signed-off-by: yuanheng <jonathan.zhaoyh@gmail.com> Signed-off-by: Yuanheng Zhao <jonathan.zhaoyh@gmail.com>
431 lines
17 KiB
Python
431 lines
17 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# Copyright 2025 The vLLM-Omni team.
|
|
# Copyright 2024 ANT Group and the HuggingFace Inc. team.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
import torch
|
|
from transformers import AutoFeatureExtractor, AutoProcessor
|
|
from transformers.feature_extraction_utils import BatchFeature, FeatureExtractionMixin
|
|
from transformers.processing_utils import ProcessorMixin
|
|
from transformers.tokenization_utils_base import PreTokenizedInput, TextInput
|
|
|
|
DEFAULT_IMAGE_PATCH_TOKEN = "<imagePatch>"
|
|
DEFAULT_IM_START_TOKEN = "<image>"
|
|
DEFAULT_IM_END_TOKEN = "</image>"
|
|
DEFAULT_VID_START_TOKEN = "<video>"
|
|
DEFAULT_VID_END_TOKEN = "</video>"
|
|
DEFAULT_FRAME_PATCH_TOKEN = "<framePatch>"
|
|
|
|
DEFAULT_AUDIO_PATCH_TOKEN = "<audioPatch>"
|
|
DEFAULT_AU_START_TOKEN = "<audio>"
|
|
DEFAULT_AU_END_TOKEN = "</audio>"
|
|
|
|
# High-level placeholders used in user prompts
|
|
PLACEHOLDER_IMAGE_TOKEN_IN_TEXT = "<IMAGE>"
|
|
PLACEHOLDER_VIDEO_TOKEN_IN_TEXT = "<VIDEO>"
|
|
PLACEHOLDER_AUDIO_TOKEN_IN_TEXT = "<AUDIO>"
|
|
|
|
# Chat template constants
|
|
USER_PREFIX = "<role>HUMAN</role>"
|
|
ASSISTANT_PREFIX = "<role>ASSISTANT</role>"
|
|
SYSTEM_PROMPT_NOTHINK = "<role>SYSTEM</role>你是一个友好的AI助手。\n\ndetailed thinking off"
|
|
SYSTEM_PROMPT_THINK = "<role>SYSTEM</role>你是一个友好的AI助手。\n\ndetailed thinking on"
|
|
|
|
|
|
_NORM_FACTOR_FOR_DTYPE = {
|
|
torch.int8: 2**7,
|
|
torch.int16: 2**15,
|
|
torch.int32: 2**31,
|
|
torch.int64: 2**63,
|
|
torch.float32: 1,
|
|
torch.float64: 1,
|
|
}
|
|
|
|
|
|
def _normalize_audio_tensor(
|
|
waveform: torch.Tensor,
|
|
sample_rate: int,
|
|
target_sample_rate: int = 16000,
|
|
) -> torch.Tensor:
|
|
"""Normalize waveform to float32, mono, and optionally resample."""
|
|
norm_factor = _NORM_FACTOR_FOR_DTYPE.get(waveform.dtype, 1)
|
|
waveform = waveform.to(torch.float32) / norm_factor
|
|
|
|
# Remove channel dimension
|
|
while len(waveform.shape) > 1:
|
|
waveform = waveform[0]
|
|
|
|
# Resample if needed
|
|
if sample_rate != target_sample_rate:
|
|
import torchaudio
|
|
|
|
resampler = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=target_sample_rate)
|
|
waveform = resampler(waveform.unsqueeze(0)).squeeze(0)
|
|
|
|
return waveform
|
|
|
|
|
|
class MingWhisperFeatureExtractor(FeatureExtractionMixin):
|
|
"""Whisper log-mel feature extractor for Ming-flash-omni-2.0.
|
|
|
|
Produces audio_feats in the time-first packed format.
|
|
|
|
Adapted from Ming's WhisperAudioEncoder
|
|
https://github.com/inclusionAI/Ming/blob/070dc3c13f95d97952ab7d22030df0c9e28a5122/modeling_whisper_encoder.py
|
|
and HF transformers WhisperFeatureExtractor
|
|
https://github.com/huggingface/transformers/blob/f842abaca95a7dbf3fc6e16122e7409109bc1431/src/transformers/models/whisper/feature_extraction_whisper.py#L33
|
|
"""
|
|
|
|
model_input_names = ["audio_feats", "audio_feats_lengths"]
|
|
|
|
def __init__(self, feature_size: int = 128, sampling_rate: int = 16000, **kwargs):
|
|
# feature_size == n_mels; stored so to_dict() serialises it correctly.
|
|
self.feature_size = feature_size
|
|
self.sampling_rate = sampling_rate
|
|
super().__init__(**kwargs)
|
|
|
|
@property
|
|
def n_mels(self) -> int:
|
|
return self.feature_size
|
|
|
|
def __call__(
|
|
self,
|
|
audios: tuple | list,
|
|
return_tensors: str | None = None,
|
|
**kwargs,
|
|
) -> BatchFeature:
|
|
"""Preprocess audio(s) into Whisper log-mel spectrograms"""
|
|
import whisper
|
|
|
|
if not isinstance(audios, list):
|
|
audios = [audios]
|
|
|
|
audio_feat_list = []
|
|
for waveform, sr in audios:
|
|
if isinstance(waveform, np.ndarray):
|
|
waveform = torch.from_numpy(waveform)
|
|
waveform = _normalize_audio_tensor(waveform, sr, target_sample_rate=self.sampling_rate)
|
|
mel = whisper.log_mel_spectrogram(waveform, n_mels=self.n_mels)
|
|
audio_feat_list.append(mel.transpose(0, 1)) # [T, n_mels]
|
|
|
|
audio_feats_lengths = torch.tensor([[feat.shape[0] for feat in audio_feat_list]], dtype=torch.long)
|
|
# Two stride-2 convolutions in series:
|
|
# 1. WhisperAudioEncoder conv2: kernel=3, stride=2, padding=1
|
|
# (conv1 has stride=1 and does not change T)
|
|
# 2. AudioProjector Conv1d: kernel=3, stride=2, padding=1
|
|
# Combined: T → ((T-1)//2 + 1 - 1)//2 + 1
|
|
# See also: AudioProjector.compute_output_length()
|
|
encoder_feats_lengths = ((audio_feats_lengths - 3 + 2 * 1) // 2 + 1 - 3 + 2 * 1) // 2 + 1
|
|
audio_feats = torch.cat(audio_feat_list, dim=0).unsqueeze(0) # [1, T_total, n_mels]
|
|
|
|
data = {
|
|
# [1, T_total, n_mels], all audio clips concatenated
|
|
"audio_feats": audio_feats.numpy(),
|
|
# [1, n_audios], actual frame count
|
|
"audio_feats_lengths": audio_feats_lengths.numpy(),
|
|
# [1, n_audios]
|
|
"encoder_feats_lengths": encoder_feats_lengths,
|
|
}
|
|
return BatchFeature(data=data, tensor_type=return_tensors)
|
|
|
|
|
|
class MingFlashOmniProcessor(ProcessorMixin):
|
|
"""Top-level multimodal processor for Ming-flash-omni 2.0.
|
|
|
|
Adapted from Ming's BailingMM2Processor
|
|
https://github.com/inclusionAI/Ming/blob/3954fcb880ff5e61ff128bcf7f1ec344d46a6fe3/processing_bailingmm2.py
|
|
|
|
Subprocessors include:
|
|
- Qwen2VLImageProcessor (image/video)
|
|
- MingWhisperFeatureExtractor (modified audio processor using Whisper's log-mel spectrogram)
|
|
"""
|
|
|
|
attributes = ["image_processor", "audio_processor", "tokenizer"]
|
|
image_processor_class = "AutoImageProcessor"
|
|
audio_processor_class = "AutoFeatureExtractor"
|
|
tokenizer_class = "AutoTokenizer"
|
|
|
|
def __init__(
|
|
self,
|
|
image_processor=None,
|
|
audio_processor=None,
|
|
tokenizer=None,
|
|
merge_size: int = 2,
|
|
**kwargs,
|
|
):
|
|
# Enforce that all sub-processors exist
|
|
# Keep None defaults in the signature for HF ProcessorMixin compatibility
|
|
if image_processor is None:
|
|
raise ValueError("MingFlashOmniProcessor requires `image_processor`.")
|
|
if audio_processor is None:
|
|
raise ValueError("MingFlashOmniProcessor requires `audio_processor`.")
|
|
if tokenizer is None:
|
|
raise ValueError("MingFlashOmniProcessor requires `tokenizer`.")
|
|
|
|
self.spatial_merge_size = merge_size
|
|
self.image_token = PLACEHOLDER_IMAGE_TOKEN_IN_TEXT
|
|
self.video_token = PLACEHOLDER_VIDEO_TOKEN_IN_TEXT
|
|
self.audio_token = PLACEHOLDER_AUDIO_TOKEN_IN_TEXT
|
|
super().__init__(
|
|
image_processor=image_processor,
|
|
audio_processor=audio_processor,
|
|
tokenizer=tokenizer,
|
|
)
|
|
|
|
# Fall back to the tokenizer's own chat_template.
|
|
if self.chat_template is None:
|
|
self.chat_template = getattr(tokenizer, "chat_template", None)
|
|
|
|
def __call__(
|
|
self,
|
|
text: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput],
|
|
images: Any | None = None,
|
|
videos: Any | None = None,
|
|
audios: tuple[np.ndarray, int] | list[tuple[np.ndarray, int]] | None = None,
|
|
**kwargs,
|
|
) -> BatchFeature:
|
|
# This should always be parallel implementations that mirror
|
|
# `_get_prompt_updates` logic in Ming processor, and vice versa.
|
|
# Ensure text is a list
|
|
if isinstance(text, str):
|
|
text = [text]
|
|
elif not isinstance(text, list):
|
|
raise ValueError("text must be a string or list of strings")
|
|
|
|
data: dict[str, Any] = {}
|
|
|
|
if images is not None:
|
|
image_outputs = self.image_processor(
|
|
images=images,
|
|
videos=None,
|
|
return_tensors="pt",
|
|
**kwargs.get("images_kwargs", {}),
|
|
)
|
|
data.update(image_outputs)
|
|
if "image_grid_thw" in image_outputs:
|
|
text = self._expand_image_tokens(text, image_outputs["image_grid_thw"])
|
|
|
|
if videos is not None:
|
|
video_outputs = self.image_processor(
|
|
images=None,
|
|
videos=videos,
|
|
return_tensors="pt",
|
|
**kwargs.get("videos_kwargs", {}),
|
|
)
|
|
if "pixel_values" in video_outputs:
|
|
video_outputs["pixel_values_videos"] = video_outputs.pop("pixel_values")
|
|
if "image_grid_thw" in video_outputs:
|
|
video_outputs["video_grid_thw"] = video_outputs.pop("image_grid_thw")
|
|
data.update(video_outputs)
|
|
if "video_grid_thw" in video_outputs:
|
|
text = self._expand_video_tokens(text, video_outputs["video_grid_thw"])
|
|
|
|
if audios is not None:
|
|
audio_outputs = self.audio_processor(
|
|
audios,
|
|
return_tensors="pt",
|
|
**kwargs.get("audio_kwargs", {}),
|
|
)
|
|
data.update(audio_outputs)
|
|
if "encoder_feats_lengths" in audio_outputs:
|
|
text = self._expand_audio_tokens(text, audio_outputs["encoder_feats_lengths"])
|
|
|
|
text_outputs = self.tokenizer(
|
|
text,
|
|
return_tensors="pt",
|
|
**kwargs.get("text_kwargs", {}),
|
|
)
|
|
data.update(text_outputs)
|
|
return BatchFeature(data=data)
|
|
|
|
def _expand_image_tokens(
|
|
self,
|
|
text: list[str],
|
|
image_grid_thw: torch.Tensor,
|
|
special_token: str = PLACEHOLDER_IMAGE_TOKEN_IN_TEXT,
|
|
) -> list[str]:
|
|
merge_size = self.spatial_merge_size
|
|
num_patches_per_image = torch.prod(image_grid_thw, dim=1) // (merge_size**2)
|
|
prompt_strings = []
|
|
image_index = 0
|
|
for sample in text:
|
|
num_images = sample.count(special_token)
|
|
if num_images > 0:
|
|
for i in range(image_index, num_images + image_index):
|
|
num_patches = int(num_patches_per_image[i].item())
|
|
img_text = (
|
|
DEFAULT_IM_START_TOKEN + (DEFAULT_IMAGE_PATCH_TOKEN * num_patches) + DEFAULT_IM_END_TOKEN + "\n"
|
|
)
|
|
sample = sample.replace(special_token, img_text, 1)
|
|
image_index += num_images
|
|
prompt_strings.append(sample)
|
|
return prompt_strings
|
|
|
|
def _expand_video_tokens(
|
|
self,
|
|
text: list[str],
|
|
video_grid_thw: torch.Tensor,
|
|
special_token: str = PLACEHOLDER_VIDEO_TOKEN_IN_TEXT,
|
|
) -> list[str]:
|
|
merge_size = self.spatial_merge_size
|
|
num_patches_per_video = torch.prod(video_grid_thw, dim=1) // (merge_size**2)
|
|
prompt_strings = []
|
|
video_index = 0
|
|
for sample in text:
|
|
num_videos = sample.count(special_token)
|
|
if num_videos > 0:
|
|
for i in range(video_index, num_videos + video_index):
|
|
num_patches = int(num_patches_per_video[i].item())
|
|
video_text = (
|
|
DEFAULT_VID_START_TOKEN
|
|
+ (DEFAULT_FRAME_PATCH_TOKEN * num_patches)
|
|
+ DEFAULT_VID_END_TOKEN
|
|
+ "\n"
|
|
)
|
|
sample = sample.replace(special_token, video_text, 1)
|
|
video_index += num_videos
|
|
prompt_strings.append(sample)
|
|
return prompt_strings
|
|
|
|
def _expand_audio_tokens(
|
|
self,
|
|
text: list[str],
|
|
encoder_feats_lengths: torch.Tensor,
|
|
special_token: str = PLACEHOLDER_AUDIO_TOKEN_IN_TEXT,
|
|
) -> list[str]:
|
|
prompt_strings = []
|
|
for sample, lengths_tensor in zip(text, encoder_feats_lengths):
|
|
for length in lengths_tensor:
|
|
num_patches = int(length.item())
|
|
if num_patches == 0:
|
|
continue
|
|
audio_text = DEFAULT_AU_START_TOKEN + (DEFAULT_AUDIO_PATCH_TOKEN * num_patches) + DEFAULT_AU_END_TOKEN
|
|
if special_token in sample:
|
|
sample = sample.replace(special_token, audio_text, 1)
|
|
else:
|
|
sample = sample + audio_text + "\n"
|
|
prompt_strings.append(sample)
|
|
return prompt_strings
|
|
|
|
def apply_system_template(
|
|
self,
|
|
sys_prompt_exp: str | None = None,
|
|
use_cot_system_prompt: bool = False,
|
|
) -> str:
|
|
sys_prompt = SYSTEM_PROMPT_THINK if use_cot_system_prompt else SYSTEM_PROMPT_NOTHINK
|
|
if sys_prompt_exp is not None:
|
|
sys_prompt = sys_prompt.replace("你是一个友好的AI助手。", sys_prompt_exp)
|
|
return sys_prompt
|
|
|
|
def apply_chat_template(
|
|
self,
|
|
conversation: list[dict[str, Any]],
|
|
sys_prompt_exp: str | None = None,
|
|
use_cot_system_prompt: bool = False,
|
|
**kwargs,
|
|
) -> str:
|
|
eos = self.tokenizer.eos_token
|
|
text = self.apply_system_template(sys_prompt_exp, use_cot_system_prompt) + eos
|
|
|
|
for idx, message in enumerate(conversation):
|
|
assert message["role"] in ["HUMAN", "ASSISTANT"], (
|
|
f"Invalid role: {message['role']}. Must be 'HUMAN' or 'ASSISTANT'"
|
|
)
|
|
if idx == len(conversation) - 1:
|
|
assert message["role"] == "HUMAN", "Last message must be from HUMAN"
|
|
|
|
text += USER_PREFIX if message["role"] == "HUMAN" else ASSISTANT_PREFIX
|
|
|
|
content = message["content"]
|
|
if isinstance(content, str):
|
|
# text-only
|
|
text += content
|
|
elif isinstance(content, list):
|
|
# structured content with multimodal elements
|
|
# Count existing placeholders from text items only
|
|
image_placeholders = 0
|
|
video_placeholders = 0
|
|
audio_placeholders = 0
|
|
for content_item in content:
|
|
if content_item.get("type", "text") == "text":
|
|
t = content_item.get("text", "")
|
|
image_placeholders += t.count(PLACEHOLDER_IMAGE_TOKEN_IN_TEXT)
|
|
video_placeholders += t.count(PLACEHOLDER_VIDEO_TOKEN_IN_TEXT)
|
|
audio_placeholders += t.count(PLACEHOLDER_AUDIO_TOKEN_IN_TEXT)
|
|
|
|
if video_placeholders > 1:
|
|
raise ValueError("Video count must be at most 1 per message!")
|
|
|
|
# Insert placeholders only for media items not already covered
|
|
for content_item in content:
|
|
content_type = content_item.get("type", "text")
|
|
|
|
if content_type == "image":
|
|
image_data = content_item.get("image")
|
|
if image_data is not None:
|
|
from PIL import Image as PILImage
|
|
|
|
num_images = 1 if isinstance(image_data, (str, PILImage.Image)) else len(image_data)
|
|
for _ in range(num_images):
|
|
if image_placeholders > 0:
|
|
image_placeholders -= 1
|
|
else:
|
|
text += PLACEHOLDER_IMAGE_TOKEN_IN_TEXT
|
|
|
|
elif content_type == "video":
|
|
if video_placeholders > 0:
|
|
video_placeholders -= 1
|
|
else:
|
|
text += PLACEHOLDER_VIDEO_TOKEN_IN_TEXT
|
|
elif content_type == "audio":
|
|
audio_data = content_item.get("audio")
|
|
if audio_data is not None:
|
|
num_audios = 1 if isinstance(audio_data, str) else len(audio_data)
|
|
for _ in range(num_audios):
|
|
if audio_placeholders > 0:
|
|
audio_placeholders -= 1
|
|
else:
|
|
text += PLACEHOLDER_AUDIO_TOKEN_IN_TEXT
|
|
|
|
elif content_type == "text":
|
|
text += content_item.get("text", "")
|
|
|
|
# Add EOS token after each message except the last one
|
|
text += eos
|
|
|
|
text += ASSISTANT_PREFIX
|
|
return text
|
|
|
|
def batch_decode(self, *args, **kwargs):
|
|
return self.tokenizer.batch_decode(*args, **kwargs)
|
|
|
|
def decode(self, *args, **kwargs):
|
|
return self.tokenizer.decode(*args, **kwargs)
|
|
|
|
@property
|
|
def model_input_names(self):
|
|
names = (
|
|
self.tokenizer.model_input_names
|
|
+ self.image_processor.model_input_names
|
|
+ self.audio_processor.model_input_names
|
|
)
|
|
return list(dict.fromkeys(names))
|
|
|
|
|
|
AutoFeatureExtractor.register("MingWhisperFeatureExtractor", MingWhisperFeatureExtractor)
|
|
AutoProcessor.register("MingFlashOmniProcessor", MingFlashOmniProcessor)
|