wehub-resource-sync
224 lines
9.1 KiB
Python
224 lines
9.1 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
import torch
|
|
from dataclasses import dataclass, field
|
|
from functools import partial
|
|
from torch import nn
|
|
from typing import Any, Dict, List, Literal, Optional
|
|
|
|
from swift.utils import get_env_args
|
|
from ..base import Template
|
|
from ..constant import MLLMTemplateType
|
|
from ..register import TemplateMeta, register_template
|
|
from ..template_inputs import StdTemplateInputs
|
|
from ..utils import Context, Prompt, findall
|
|
from ..vision_utils import load_video_minicpmv_mplug_owl3
|
|
from .qwen import QwenTemplateMeta
|
|
|
|
|
|
class mPlugOwl2Template(Template):
|
|
|
|
def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int,
|
|
inputs: StdTemplateInputs) -> List[Context]:
|
|
assert media_type == 'image'
|
|
return [[-200]]
|
|
|
|
def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]:
|
|
from mplug_owl2.mm_utils import process_images
|
|
processor = self.processor
|
|
images = inputs.images
|
|
for i, image in enumerate(images):
|
|
# ref: https://modelscope.cn/models/iic/mPLUG-Owl2.1
|
|
max_edge = max(image.size)
|
|
image = image.resize((max_edge, max_edge))
|
|
images[i] = image
|
|
encoded = super()._encode(inputs)
|
|
input_ids = encoded['input_ids']
|
|
labels = encoded['labels']
|
|
res = {'input_ids': input_ids, 'labels': labels}
|
|
if images:
|
|
images = process_images(images, processor)
|
|
images = images.to(self.model_info.torch_dtype)
|
|
res['images'] = images
|
|
return res
|
|
|
|
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 = [b['images'] for b in batch if 'images' in b]
|
|
if images:
|
|
res['images'] = torch.concat(images)
|
|
return res
|
|
|
|
|
|
register_template(
|
|
TemplateMeta(
|
|
MLLMTemplateType.mplug_owl2,
|
|
template_cls=mPlugOwl2Template,
|
|
prefix=['{{SYSTEM}}'],
|
|
prompt=['USER: {{QUERY}}ASSISTANT:'],
|
|
chat_sep=['</s>'],
|
|
suffix=[['eos_token_id']],
|
|
stop_words=['<|endoftext|>', '</s>']))
|
|
|
|
|
|
class mPlugOwl3Template(Template):
|
|
version = None
|
|
|
|
def init_env_args(self):
|
|
super().init_env_args()
|
|
self.max_num_frames = get_env_args('max_num_frames', int, 16)
|
|
|
|
def _get_image_token_list(self, cut_shape):
|
|
text = self.processor.image_processor.cut_prompt_template(img_token='<|image|>', h=cut_shape[0], w=cut_shape[1])
|
|
text_list = text.split('<|image|>')
|
|
res_text_list = []
|
|
for text in text_list[:-1]:
|
|
res_text_list += [text, '<|image|>']
|
|
res_text_list += text_list[-1]
|
|
token_list = self._encode_context_list(res_text_list)[0]
|
|
return token_list
|
|
|
|
def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int,
|
|
inputs: StdTemplateInputs) -> List[Context]:
|
|
assert media_type in {'image', 'video'}
|
|
load_video = partial(load_video_minicpmv_mplug_owl3, max_num_frames=self.max_num_frames)
|
|
if media_type == 'image':
|
|
return [[-100], '\n']
|
|
elif media_type == 'video':
|
|
return self.replace_video2image(load_video, inputs, lambda i: [[-100]]) + ['\n']
|
|
|
|
def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]:
|
|
encoded = super()._encode(inputs)
|
|
images = inputs.images
|
|
videos = inputs.videos
|
|
cut_enable = not videos
|
|
input_ids = encoded['input_ids']
|
|
labels = encoded['labels']
|
|
loss_scale = encoded.get('loss_scale', None)
|
|
idx_list = findall(input_ids, -100)
|
|
processor = self.processor
|
|
encoded = {}
|
|
if images:
|
|
image_inputs = processor.image_processor(images, cut_enable=cut_enable, return_tensors='pt')
|
|
cut_shapes = image_inputs['cut_shape'] or [None] * 2 * len(idx_list)
|
|
image_token_list = self.processor.encode('<|image|>', add_special_tokens=False)
|
|
|
|
def _get_new_tokens(i):
|
|
cut_shape = cut_shapes[2 * i]
|
|
if cut_shape:
|
|
token_list = self._get_image_token_list(cut_shape)
|
|
else:
|
|
token_list = image_token_list
|
|
return token_list
|
|
|
|
input_ids, labels, loss_scale = self._extend_tokens(input_ids, labels, loss_scale, idx_list,
|
|
_get_new_tokens)
|
|
|
|
image_token_idx = torch.tensor(findall(input_ids, image_token_list))
|
|
if self.version == '241101':
|
|
media_offset = image_token_idx
|
|
else:
|
|
_range = torch.arange(len(input_ids))[:, None]
|
|
matrix = (_range > image_token_idx[None]).sum(dim=1)
|
|
media_offset = torch.stack([torch.zeros(matrix.shape[0], dtype=torch.long), matrix], dim=-1)[None]
|
|
encoded.update({
|
|
'pixel_values': image_inputs['pixel_values'],
|
|
'media_offset': media_offset,
|
|
})
|
|
encoded['input_ids'] = input_ids
|
|
encoded['labels'] = labels
|
|
encoded['loss_scale'] = loss_scale
|
|
return encoded
|
|
|
|
def _post_encode(self, model: nn.Module, inputs: Dict[str, Any]) -> Dict[str, Any]:
|
|
if 'media_offset' in inputs:
|
|
media_offset = []
|
|
cusum_offset = 0
|
|
image_embeds = []
|
|
pixel_values = inputs.pop('pixel_values')
|
|
max_sequence_length = inputs['input_ids'].shape[1]
|
|
for i, curr_media_offset in enumerate(inputs['media_offset']):
|
|
if curr_media_offset is None:
|
|
continue
|
|
if curr_media_offset.shape[1] < max_sequence_length:
|
|
padding = curr_media_offset[:, -1:, :].expand(curr_media_offset.shape[0],
|
|
max_sequence_length - curr_media_offset.shape[1],
|
|
curr_media_offset.shape[2])
|
|
curr_media_offset = torch.concat([curr_media_offset, padding], dim=1)
|
|
media_offset.append(curr_media_offset + cusum_offset)
|
|
image_embeds.append(model.forward_image(pixel_values[i]))
|
|
cusum_offset += image_embeds[-1].shape[0]
|
|
inputs['media_offset'] = torch.concat(media_offset)
|
|
inputs['image_embeds'] = torch.concat(image_embeds)
|
|
return inputs
|
|
|
|
def _data_collator(self, batch: List[Dict[str, Any]], *, padding_to: Optional[int] = None) -> Dict[str, Any]:
|
|
res = self.fetch_inputs(batch, ['media_offset', 'pixel_values'])
|
|
for b in batch:
|
|
b.pop('pixel_values', None)
|
|
res.update(super()._data_collator(batch, padding_to=padding_to))
|
|
return res
|
|
|
|
|
|
class mPlugOwl3_241101Template(mPlugOwl3Template):
|
|
version = '241101'
|
|
|
|
def _post_encode(self, model: nn.Module, inputs: Dict[str, Any]) -> Dict[str, Any]:
|
|
if 'pixel_values' in inputs:
|
|
pixel_values = inputs.pop('pixel_values')
|
|
inputs['image_embeds'] = torch.concat([model.forward_image(pv) for pv in pixel_values])
|
|
else:
|
|
inputs['media_offset'] = [None] * inputs['input_ids'].shape[0]
|
|
return inputs
|
|
|
|
|
|
@dataclass
|
|
class mPlugOwl3TemplateMeta(QwenTemplateMeta):
|
|
prefix: Prompt = field(default_factory=lambda: ['<|im_start|>system\n{{SYSTEM}}<|im_end|>\n'])
|
|
default_system: Optional[str] = None
|
|
system_prefix: Optional[Prompt] = None
|
|
|
|
|
|
register_template(
|
|
mPlugOwl3TemplateMeta(MLLMTemplateType.mplug_owl3, template_cls=mPlugOwl3Template, agent_template=None))
|
|
|
|
register_template(
|
|
mPlugOwl3TemplateMeta(
|
|
MLLMTemplateType.mplug_owl3_241101, template_cls=mPlugOwl3_241101Template, agent_template=None))
|
|
|
|
|
|
class DocOwl2Template(Template):
|
|
|
|
def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int,
|
|
inputs: StdTemplateInputs) -> List[Context]:
|
|
if media_type == 'image':
|
|
return [f'<img {index + 1}>', [-200]]
|
|
|
|
def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]:
|
|
encoded = super()._encode(inputs)
|
|
if inputs.images:
|
|
image_tensor, patch_positions, _ = self.processor._process_image(inputs.images)
|
|
image_tensor = image_tensor.to(self.model_info.torch_dtype)
|
|
encoded.update({'images': image_tensor, 'patch_positions': patch_positions})
|
|
return encoded
|
|
|
|
def _data_collator(self, batch: List[Dict[str, Any]], *, padding_to: Optional[int] = None) -> Dict[str, Any]:
|
|
keys = ['images', 'patch_positions']
|
|
res = self.fetch_inputs(batch, keys)
|
|
for key in keys:
|
|
val = res.get(key)
|
|
if val:
|
|
res[key] = torch.concat([v for v in val if v is not None])
|
|
res.update(super()._data_collator(batch, padding_to=padding_to))
|
|
return res
|
|
|
|
|
|
register_template(
|
|
TemplateMeta(
|
|
MLLMTemplateType.doc_owl2,
|
|
prefix=['<s>'],
|
|
prompt=[' USER: {{QUERY}} ASSISTANT:'],
|
|
chat_sep=['</s>'],
|
|
suffix=['</s>'],
|
|
template_cls=DocOwl2Template,
|
|
))
|