wehub-resource-sync
418 lines
17 KiB
Python
418 lines
17 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
import itertools
|
|
import torch
|
|
from functools import partial
|
|
from typing import Any, Dict, List, Literal, Optional
|
|
|
|
from swift.utils import get_env_args, to_float_dtype
|
|
from ..base import Template
|
|
from ..constant import MLLMTemplateType
|
|
from ..register import TemplateMeta, register_template
|
|
from ..template_inputs import StdTemplateInputs
|
|
from ..utils import Context, findall
|
|
from ..vision_utils import load_batch, load_file
|
|
from .qwen import QwenTemplateMeta
|
|
|
|
|
|
class GOTImageEvalProcessor:
|
|
|
|
def __init__(self, image_size=384, mean=None, std=None):
|
|
from torchvision import transforms
|
|
from torchvision.transforms.functional import InterpolationMode
|
|
if mean is None:
|
|
mean = (0.48145466, 0.4578275, 0.40821073)
|
|
if std is None:
|
|
std = (0.26862954, 0.26130258, 0.27577711)
|
|
|
|
self.normalize = transforms.Normalize(mean, std)
|
|
|
|
self.transform = transforms.Compose([
|
|
transforms.Resize((image_size, image_size), interpolation=InterpolationMode.BICUBIC),
|
|
transforms.ToTensor(),
|
|
self.normalize,
|
|
])
|
|
|
|
def __call__(self, item):
|
|
return self.transform(item)
|
|
|
|
|
|
class GOT_OCR2Template(Template):
|
|
placeholder_tokens = ['<imgpad>']
|
|
|
|
def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int,
|
|
inputs: StdTemplateInputs) -> List[Context]:
|
|
# 'OCR: '
|
|
# 'OCR with format: '
|
|
assert media_type == 'image'
|
|
return ['<img>' + '<imgpad>' * 256 + '</img>\n']
|
|
|
|
def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]:
|
|
encoded = super()._encode(inputs)
|
|
images = inputs.images
|
|
image_processor_high = GOTImageEvalProcessor(image_size=1024)
|
|
for i, image in enumerate(images):
|
|
images[i] = image_processor_high(image)[None].to(self.model_info.torch_dtype)
|
|
if images:
|
|
encoded['images'] = images
|
|
return encoded
|
|
|
|
def _data_collator(self, batch: List[Dict[str, Any]], *, padding_to: Optional[int] = None) -> Dict[str, Any]:
|
|
res = super()._data_collator(batch, padding_to=padding_to)
|
|
images = self.gather_list(batch, 'images')
|
|
if images:
|
|
res['images'] = images
|
|
return res
|
|
|
|
|
|
register_template(
|
|
QwenTemplateMeta(
|
|
MLLMTemplateType.got_ocr2,
|
|
default_system=' You should follow the instructions carefully and explain your answers in detail.',
|
|
template_cls=GOT_OCR2Template,
|
|
agent_template=None,
|
|
))
|
|
|
|
|
|
class GOT_OCR2HfTemplate(Template):
|
|
placeholder_tokens = ['<imgpad>']
|
|
|
|
def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int,
|
|
inputs: StdTemplateInputs) -> List[Context]:
|
|
# 'OCR: '
|
|
# 'OCR with format: '
|
|
assert media_type == 'image'
|
|
return ['<img>' + '<imgpad>' * 256 + '</img>\n']
|
|
|
|
def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]: # 暂时照抄上面
|
|
encoded = super()._encode(inputs)
|
|
images = inputs.images
|
|
if images:
|
|
encoded['images'] = images
|
|
return encoded
|
|
|
|
def _data_collator(self, batch: List[Dict[str, Any]], *, padding_to: Optional[int] = None) -> Dict[str, Any]:
|
|
res = super()._data_collator(batch, padding_to=padding_to)
|
|
images = self.gather_list(batch, 'images')
|
|
_inputs = self.processor(images, return_tensors='pt')
|
|
_inputs.pop('input_ids') # this does not contain the response, so cannot be used when training
|
|
_inputs.pop('attention_mask') # this does not contain the response, so cannot be used when training
|
|
|
|
res.update(_inputs.data)
|
|
return res
|
|
|
|
|
|
register_template(
|
|
QwenTemplateMeta(
|
|
MLLMTemplateType.got_ocr2_hf,
|
|
default_system=' You should follow the instructions carefully and explain your answers in detail.',
|
|
template_cls=GOT_OCR2HfTemplate,
|
|
agent_template=None,
|
|
))
|
|
|
|
|
|
class StepAudioTemplate(Template):
|
|
use_model = True
|
|
|
|
def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int,
|
|
inputs: StdTemplateInputs) -> List[Context]:
|
|
assert media_type == 'audio', f'media_type: {media_type}'
|
|
from utils import load_audio
|
|
audio_wav, sr = load_audio(load_file(inputs.audios[index]))
|
|
audio_tokens = self.model.encoder(audio_wav, sr)
|
|
return audio_tokens
|
|
|
|
|
|
class StepAudio2MiniTemplate(Template):
|
|
use_model = True
|
|
|
|
def load_audio(self, file_path, target_rate=16000, max_length=None):
|
|
'''
|
|
Open an audio file and read as mono waveform, resampling as necessary
|
|
If max_length is provided, truncate the audio to that length
|
|
'''
|
|
import torchaudio
|
|
waveform, sample_rate = torchaudio.load(file_path)
|
|
if sample_rate != target_rate:
|
|
waveform = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=target_rate)(waveform)
|
|
audio = waveform[0] # get the first channel
|
|
|
|
# Truncate audio if it exceeds max_length
|
|
if max_length is not None and audio.shape[0] > max_length:
|
|
audio = audio[:max_length]
|
|
|
|
return audio
|
|
|
|
def _mel_filters(self, n_mels: int) -> 'torch.Tensor':
|
|
'''Load the mel filterbank matrix for projecting STFT into a Mel spectrogram.'''
|
|
import librosa
|
|
import torch
|
|
assert n_mels in {80, 128}, f'Unsupported n_mels: {n_mels}'
|
|
if n_mels == 128:
|
|
return torch.from_numpy(librosa.filters.mel(sr=16000, n_fft=400, n_mels=128))
|
|
else:
|
|
return torch.from_numpy(librosa.filters.mel(sr=16000, n_fft=400, n_mels=80))
|
|
|
|
def log_mel_spectrogram(self, audio, n_mels=128, padding=479):
|
|
'''
|
|
Compute the log-Mel spectrogram with specific padding for StepAudio
|
|
'''
|
|
import torch
|
|
import torch.nn.functional as F
|
|
if isinstance(audio, str):
|
|
audio = self.load_audio(audio)
|
|
elif not torch.is_tensor(audio):
|
|
audio = torch.from_numpy(audio)
|
|
if padding > 0:
|
|
audio = F.pad(audio, (0, padding))
|
|
window = torch.hann_window(400).to(audio.device)
|
|
stft = torch.stft(audio, 400, 160, window=window, return_complex=True)
|
|
magnitudes = stft[..., :-1].abs()**2
|
|
filters = self._mel_filters(n_mels)
|
|
mel_spec = filters @ magnitudes
|
|
|
|
log_spec = torch.clamp(mel_spec, min=1e-10).log10()
|
|
log_spec = torch.maximum(log_spec, log_spec.max() - 8.0)
|
|
log_spec = (log_spec + 4.0) / 4.0
|
|
return log_spec
|
|
|
|
def compute_token_num(self, max_feature_len):
|
|
# First, audio goes through encoder:
|
|
# 1. conv1: kernel=3, stride=1, padding=1 -> size unchanged
|
|
# 2. conv2: kernel=3, stride=2, padding=1 -> size/2
|
|
# 3. avg_pooler: kernel=2, stride=2 -> size/2
|
|
max_feature_len = max_feature_len - 2 # remove padding
|
|
encoder_output_dim = (max_feature_len + 1) // 2 // 2 # after conv2 and avg_pooler
|
|
|
|
# Then through adaptor (parameters from config file):
|
|
padding = 1
|
|
kernel_size = 3 # from config: audio_encoder_config.kernel_size
|
|
stride = 2 # from config: audio_encoder_config.adapter_stride
|
|
adapter_output_dim = (encoder_output_dim + 2 * padding - kernel_size) // stride + 1
|
|
return adapter_output_dim
|
|
|
|
def padding_mels(self, data: List['torch.Tensor']):
|
|
''' Padding the data into batch data
|
|
|
|
Parameters
|
|
----------
|
|
data: List[Tensor], shape of Tensor (128, T)
|
|
|
|
Returns:
|
|
-------
|
|
feats, feats lengths
|
|
'''
|
|
import torch
|
|
from torch.nn.utils.rnn import pad_sequence
|
|
sample = data
|
|
assert isinstance(sample, list)
|
|
feats_lengths = torch.tensor([s.size(1) - 2 for s in sample], dtype=torch.int32)
|
|
feats = [s.t() for s in sample]
|
|
padded_feats = pad_sequence(feats, batch_first=True, padding_value=0)
|
|
|
|
return padded_feats.transpose(1, 2), feats_lengths
|
|
|
|
def audio_process(self, audio):
|
|
results = []
|
|
mels = []
|
|
for i in range(0, audio.shape[0], 16000 * 25):
|
|
mel = self.log_mel_spectrogram(audio[i:i + 16000 * 25], n_mels=128, padding=479)
|
|
mels.append(mel)
|
|
audio_tokens = '<audio_patch>' * self.compute_token_num(mel.shape[1])
|
|
results.append(f'<audio_start>{audio_tokens}<audio_end>')
|
|
audio_ids = self._tokenize(''.join(results))
|
|
return audio_ids, mels
|
|
|
|
def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int,
|
|
inputs: StdTemplateInputs) -> List[Context]:
|
|
assert media_type == 'audio'
|
|
return ['<audio_patch>']
|
|
|
|
def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]:
|
|
encoded = super()._encode(inputs)
|
|
input_ids = encoded['input_ids']
|
|
labels = encoded['labels']
|
|
loss_scale = encoded.get('loss_scale', None)
|
|
sampling_rate = get_env_args('sampling_rate', int, 16000)
|
|
inputs.audios = load_batch(inputs.audios, partial(self.load_audio, target_rate=sampling_rate))
|
|
|
|
audio_token = self._tokenize('<audio_patch>')[0]
|
|
idx_list = findall(input_ids, audio_token)
|
|
|
|
if idx_list:
|
|
audio_inputs = []
|
|
mels = []
|
|
for audio in inputs.audios:
|
|
audio_input, mel = self.audio_process(audio)
|
|
audio_inputs.append(audio_input)
|
|
mels.extend(mel)
|
|
|
|
def _get_new_audio_tokens(i):
|
|
return audio_inputs[i]
|
|
|
|
input_ids, labels, loss_scale = self._extend_tokens(input_ids, labels, loss_scale, idx_list,
|
|
_get_new_audio_tokens)
|
|
encoded['input_ids'] = input_ids # Add labels to the batch
|
|
encoded['labels'] = labels # Add labels to the batch
|
|
encoded['loss_scale'] = loss_scale
|
|
encoded['mels'] = mels
|
|
wavs, wav_lens = self.padding_mels(mels)
|
|
# audio_tokens = [151688, 151690, 151689]
|
|
# for audio_token_id in audio_tokens:
|
|
# labels[labels == audio_token_id] = -100 # Mask image token IDs in labels
|
|
|
|
else:
|
|
wavs = None
|
|
wav_lens = None
|
|
|
|
encoded['wavs'] = wavs
|
|
encoded['wav_lens'] = wav_lens
|
|
|
|
return encoded
|
|
|
|
def _data_collator(self, batch: List[Dict[str, Any]], *, padding_to: Optional[int] = None) -> Dict[str, Any]:
|
|
combined_mels = list(itertools.chain.from_iterable([e.get('mels', []) for e in batch]))
|
|
batch_wavs, batch_wav_lens = self.padding_mels(combined_mels) if combined_mels else (None, None)
|
|
res = super()._data_collator(batch, padding_to=padding_to)
|
|
res['wav_lens'] = batch_wav_lens
|
|
res['wavs'] = batch_wavs
|
|
|
|
return res
|
|
|
|
|
|
register_template(
|
|
TemplateMeta(
|
|
MLLMTemplateType.step_audio2_mini,
|
|
template_cls=StepAudio2MiniTemplate,
|
|
prefix=[],
|
|
prompt=['<|BOT|>human\n{{QUERY}}<|EOT|><|BOT|>assistant\n'],
|
|
system_prefix=['<|BOT|>system\n{{SYSTEM}}<|EOT|>'],
|
|
chat_sep=['<|EOT|>'],
|
|
suffix=['<|EOT|>'],
|
|
))
|
|
|
|
register_template(
|
|
TemplateMeta(
|
|
MLLMTemplateType.step_audio,
|
|
template_cls=StepAudioTemplate,
|
|
prefix=['<s>'],
|
|
prompt=['<|BOT|>human\n{{QUERY}}<|EOT|><|BOT|>assistant\n'],
|
|
system_prefix=['<s><|BOT|>system\n{{SYSTEM}}<|EOT|>'],
|
|
chat_sep=['<|EOT|>'],
|
|
suffix=['<|EOT|>'],
|
|
))
|
|
|
|
|
|
class Step3VLTemplate(Template):
|
|
use_model = True
|
|
support_padding_free = False
|
|
image_token_id = 151679
|
|
placeholder_tokens = ['<im_patch>']
|
|
|
|
def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int,
|
|
inputs: StdTemplateInputs) -> List[Context]:
|
|
assert media_type == 'image'
|
|
return ['<im_patch>']
|
|
|
|
def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]:
|
|
encoded = super()._encode(inputs)
|
|
input_ids = encoded['input_ids']
|
|
labels = encoded['labels']
|
|
loss_scale = encoded.get('loss_scale', None)
|
|
|
|
images = inputs.images
|
|
if images:
|
|
processor = self.processor
|
|
idx_list = findall(input_ids, self.image_token_id)
|
|
splitted_images_data = processor._split_images(images)
|
|
pixel_values_lst = []
|
|
patch_pixel_values_lst = []
|
|
patch_newline_mask_lst = []
|
|
image_repl_ids_lst = []
|
|
num_patches = []
|
|
for raw_img, img_patches, patch_newline_mask in splitted_images_data:
|
|
pixel_values_lst.extend(processor._convert_images_to_pixel_values([raw_img]))
|
|
if len(img_patches) > 0:
|
|
patch_pixel_values_lst.extend(processor._convert_images_to_pixel_values(img_patches, is_patch=True))
|
|
num_patches.append(len(img_patches))
|
|
_, image_repl_ids = processor._get_image_repl_features(1, len(img_patches), patch_newline_mask)
|
|
image_repl_ids_lst.append(image_repl_ids)
|
|
if patch_newline_mask is not None:
|
|
patch_newline_mask_lst.extend(patch_newline_mask)
|
|
|
|
image_inputs = {
|
|
'pixel_values': torch.cat(pixel_values_lst),
|
|
'num_patches': num_patches,
|
|
}
|
|
if patch_pixel_values_lst:
|
|
image_inputs['patch_pixel_values'] = torch.cat(patch_pixel_values_lst)
|
|
if patch_newline_mask_lst:
|
|
image_inputs['patch_newline_mask'] = torch.tensor(patch_newline_mask_lst, dtype=torch.bool)
|
|
image_inputs = to_float_dtype(image_inputs, self.model_info.torch_dtype)
|
|
|
|
def _get_new_tokens(i):
|
|
return image_repl_ids_lst[i]
|
|
|
|
input_ids, labels, loss_scale = self._extend_tokens(input_ids, labels, loss_scale, idx_list,
|
|
_get_new_tokens)
|
|
|
|
encoded['input_ids'] = input_ids
|
|
encoded['labels'] = labels
|
|
encoded['loss_scale'] = loss_scale
|
|
encoded.update(image_inputs)
|
|
return encoded
|
|
|
|
def _post_encode(self, model, inputs: Dict[str, Any]) -> Dict[str, Any]:
|
|
if not self.is_training:
|
|
return inputs
|
|
|
|
input_ids = inputs['input_ids']
|
|
# Only one image is supported per sample. # File: modeling_step_vl.py line 319, in _process_image_input
|
|
# cur_feature.append(image_features[i].view(-1, image_features.shape[-1]))
|
|
pixel_values = inputs.get('pixel_values')
|
|
num_patches = inputs.get('num_patches')
|
|
patch_pixel_values = inputs.get('patch_pixel_values')
|
|
|
|
base_model = self.get_base_model(model)
|
|
inputs_embeds = base_model.model.language_model.embed_tokens(input_ids)
|
|
if pixel_values is not None:
|
|
img_inputs = base_model.model._parse_and_validate_image_input(
|
|
pixel_values=pixel_values, num_patches=num_patches, patch_pixel_values=patch_pixel_values)
|
|
|
|
# [image embedding or concatenation of image embedding and patch image embedding]
|
|
img_embeddings = base_model.model._process_image_input(img_inputs)
|
|
|
|
is_multimodal = input_ids == self.image_token_id
|
|
is_multimodal = is_multimodal.to(inputs_embeds.device)
|
|
bs = is_multimodal.shape[0]
|
|
for i in range(bs):
|
|
assert is_multimodal[i].sum() == img_embeddings[i].shape[0]
|
|
B, L, D = inputs_embeds.shape
|
|
flat_img_embeds = torch.cat(img_embeddings, dim=0)
|
|
flat_mask = is_multimodal.view(-1)
|
|
flat_inputs_embeds = inputs_embeds.view(-1, D)
|
|
flat_inputs_embeds[flat_mask] = flat_img_embeds
|
|
inputs_embeds = flat_inputs_embeds.view(B, L, D)
|
|
return {'inputs_embeds': inputs_embeds}
|
|
|
|
def _data_collator_mm_data(self, batch: List[Dict[str, Any]]) -> Dict[str, Any]:
|
|
res = super()._data_collator_mm_data(batch)
|
|
num_patches = self.gather_list(batch, 'num_patches')
|
|
if num_patches:
|
|
res['num_patches'] = num_patches
|
|
patch_pixel_values = [b['patch_pixel_values'] for b in batch if b.get('patch_pixel_values') is not None]
|
|
patch_newline_mask = [b['patch_newline_mask'] for b in batch if b.get('patch_newline_mask') is not None]
|
|
if patch_pixel_values:
|
|
res['patch_pixel_values'] = torch.concat(patch_pixel_values)
|
|
res['patch_newline_mask'] = torch.concat(patch_newline_mask)
|
|
return res
|
|
|
|
|
|
register_template(
|
|
QwenTemplateMeta(
|
|
MLLMTemplateType.step3_vl,
|
|
template_cls=Step3VLTemplate,
|
|
default_system=None,
|
|
is_thinking=True,
|
|
thinking_prefix='<think>\n',
|
|
))
|