wehub-resource-sync
296 lines
13 KiB
Python
296 lines
13 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn as nn
|
|
from dataclasses import dataclass, field
|
|
from typing import Any, Dict, List, Literal, Optional
|
|
|
|
from ..base import Template
|
|
from ..constant import LLMTemplateType, MLLMTemplateType
|
|
from ..register import TemplateMeta, register_template
|
|
from ..template_inputs import StdTemplateInputs
|
|
from ..utils import Context, Prompt, findall
|
|
|
|
|
|
@dataclass
|
|
class ERNIETemplateMeta(TemplateMeta):
|
|
prefix: Prompt = field(default_factory=lambda: ['<|begin_of_sentence|>'])
|
|
prompt: Prompt = field(default_factory=lambda: ['User: {{QUERY}}\nAssistant: '])
|
|
chat_sep: Optional[Prompt] = field(default_factory=lambda: ['<|end_of_sentence|>'])
|
|
suffix: Prompt = field(default_factory=lambda: ['</s>'])
|
|
system_prefix: Optional[Prompt] = field(default_factory=lambda: ['<|begin_of_sentence|>{{SYSTEM}}\n'])
|
|
|
|
|
|
register_template(ERNIETemplateMeta(LLMTemplateType.ernie))
|
|
|
|
|
|
class ErnieThinkingTemplate(Template):
|
|
|
|
def _swift_prepare_inputs(self, inputs) -> None:
|
|
super()._swift_prepare_inputs(inputs)
|
|
for message in inputs.messages:
|
|
if message['role'] == 'assistant':
|
|
if '<response>' not in message['content']:
|
|
if '</think>' in message['content']:
|
|
message['content'] = message['content'].replace('</think>', '</think>\n\n<response>\n')
|
|
message['content'] = message['content'] + '\n</response>'
|
|
if '<think>\n' not in message['content']:
|
|
message['content'] = message['content'].replace('<think>', '<think>\n')
|
|
else:
|
|
message['content'] = '<response>\n' + message['content'] + '\n</response>\n'
|
|
|
|
|
|
@dataclass
|
|
class ERNIEThinkingTemplateMeta(TemplateMeta):
|
|
prefix: Prompt = field(
|
|
default_factory=lambda:
|
|
['<|im_start|>system\n'
|
|
'<global_setting>\n'
|
|
'think_mode=True\n'
|
|
'</global_setting><|im_end|>\n\n'])
|
|
prompt: Prompt = field(
|
|
default_factory=lambda: ['<|im_start|>user\n'
|
|
'{{QUERY}}<|im_end|>\n\n'
|
|
'<|im_start|>assistant\n'])
|
|
chat_sep: Optional[Prompt] = field(default_factory=lambda: ['<|im_end|>\n\n'])
|
|
suffix: Prompt = field(default_factory=lambda: ['<|im_end|>'])
|
|
system_prefix: Optional[Prompt] = field(default_factory=lambda: [
|
|
'<|im_start|>system\n'
|
|
'<system_setting>\n'
|
|
'{{SYSTEM}}\n'
|
|
'</system_setting>\n\n'
|
|
'<global_setting>\n'
|
|
'think_mode=True\n'
|
|
'</global_setting><|im_end|>\n\n'
|
|
])
|
|
|
|
|
|
register_template(
|
|
ERNIEThinkingTemplateMeta(
|
|
LLMTemplateType.ernie_thinking,
|
|
template_cls=ErnieThinkingTemplate,
|
|
is_thinking=True,
|
|
thinking_prefix='<think>\n'))
|
|
|
|
|
|
class PaddleOCRTemplate(Template):
|
|
image_token = '<|IMAGE_PLACEHOLDER|>'
|
|
image_token_id = 100295
|
|
skip_prompt = False
|
|
version = 'v1'
|
|
|
|
def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int,
|
|
inputs: StdTemplateInputs) -> List[Context]:
|
|
assert media_type == 'image'
|
|
return ['<|IMAGE_START|><|IMAGE_PLACEHOLDER|><|IMAGE_END|>']
|
|
|
|
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)
|
|
idx_list = findall(input_ids, self.image_token_id)
|
|
processor = self.processor
|
|
images = inputs.images
|
|
if images:
|
|
processor_kwargs = {}
|
|
if self.version == 'v1_5' and inputs.chat_template_kwargs:
|
|
for key in ['shortest_edge', 'longest_edge']:
|
|
value = inputs.chat_template_kwargs.get(key, None)
|
|
if value:
|
|
processor_kwargs[key] = value
|
|
if processor_kwargs:
|
|
processor_kwargs = {'size': processor_kwargs}
|
|
image_inputs = processor.image_processor(images=images, return_tensors='pt', **processor_kwargs)
|
|
image_inputs['pixel_values'] = image_inputs['pixel_values']
|
|
image_grid_thw = image_inputs['image_grid_thw']
|
|
merge_size = processor.image_processor.merge_size**2
|
|
|
|
def _get_new_tokens(i):
|
|
img_tokens: List[int] = [self.image_token_id] * (image_grid_thw[i].prod() // merge_size)
|
|
return img_tokens
|
|
|
|
encoded['input_ids'], encoded['labels'], encoded['loss_scale'] = self._extend_tokens(
|
|
input_ids, labels, loss_scale, idx_list, _get_new_tokens)
|
|
encoded['pixel_values'] = image_inputs['pixel_values']
|
|
encoded['image_grid_thw'] = image_grid_thw
|
|
|
|
return encoded
|
|
|
|
def _post_encode(self, model: nn.Module, inputs: Dict[str, Any]) -> Dict[str, Any]:
|
|
embedding = model.get_input_embeddings()
|
|
device = embedding.weight.device
|
|
input_ids = inputs['input_ids']
|
|
inputs_embeds = embedding(input_ids).to(device=device)
|
|
pixel_values = inputs.get('pixel_values')
|
|
image_grid_thw = inputs.get('image_grid_thw')
|
|
if pixel_values is not None:
|
|
siglip_position_ids = list()
|
|
image_grid_hws = list()
|
|
sample_indices = list()
|
|
cu_seqlens = [0]
|
|
pixel_values = pixel_values.unsqueeze(0).to(device=device)
|
|
for idx, thw in enumerate(image_grid_thw):
|
|
thw_tuple = tuple(thw.detach().cpu().numpy().tolist())
|
|
numel = np.prod(thw_tuple)
|
|
image_grid_hws.append(thw_tuple)
|
|
image_position_ids = torch.arange(numel) % np.prod(thw_tuple[1:])
|
|
siglip_position_ids.append(image_position_ids)
|
|
sample_indices.append(torch.full((numel, ), idx, dtype=torch.int64))
|
|
cu_seqlens.append(cu_seqlens[-1] + numel)
|
|
|
|
siglip_position_ids = torch.concat(siglip_position_ids, dim=0).to(pixel_values.device)
|
|
cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32).to(pixel_values.device)
|
|
sample_indices = torch.concat(sample_indices, dim=0).to(pixel_values.device)
|
|
|
|
vision_outputs = model.visual(
|
|
pixel_values=pixel_values,
|
|
image_grid_thw=image_grid_hws,
|
|
position_ids=siglip_position_ids,
|
|
vision_return_embed_list=True,
|
|
interpolate_pos_encoding=True,
|
|
sample_indices=sample_indices,
|
|
cu_seqlens=cu_seqlens,
|
|
return_pooler_output=False,
|
|
use_rope=True,
|
|
window_size=-1,
|
|
)
|
|
image_embeds = vision_outputs.last_hidden_state
|
|
image_embeds = model.mlp_AR(image_embeds, image_grid_thw)
|
|
n_image_tokens = (input_ids == self.image_token_id).sum().item()
|
|
image_embeds = torch.cat(image_embeds, dim=0)
|
|
n_image_features = image_embeds.shape[0]
|
|
if n_image_tokens != n_image_features:
|
|
raise ValueError('Image features and image tokens do not match: tokens: '
|
|
f'{n_image_tokens}, features {n_image_features}')
|
|
|
|
mask = input_ids == self.image_token_id
|
|
mask_unsqueezed = mask.unsqueeze(-1)
|
|
mask_expanded = mask_unsqueezed.expand_as(inputs_embeds)
|
|
image_mask = mask_expanded.to(inputs_embeds.device)
|
|
|
|
image_embeds = image_embeds.to(inputs_embeds.device, inputs_embeds.dtype)
|
|
|
|
inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds)
|
|
|
|
return {'inputs_embeds': inputs_embeds}
|
|
|
|
|
|
register_template(ERNIETemplateMeta(MLLMTemplateType.paddle_ocr, template_cls=PaddleOCRTemplate))
|
|
|
|
|
|
class ERNIE_VLTemplate(Template):
|
|
placeholder_tokens = ['<|IMAGE_PLACEHOLDER|>']
|
|
|
|
def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int,
|
|
inputs: StdTemplateInputs) -> List[Context]:
|
|
assert media_type == 'image'
|
|
return [f'Picture {index + 1}:<|IMAGE_PLACEHOLDER|>']
|
|
|
|
def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]:
|
|
encoded = super()._encode(inputs)
|
|
input_ids = encoded['input_ids']
|
|
labels = encoded['labels']
|
|
loss_scale = encoded['loss_scale']
|
|
image_token = self._tokenize('<|IMAGE_PLACEHOLDER|>')[0]
|
|
idx_list = findall(input_ids, image_token)
|
|
if idx_list:
|
|
split_token = self._tokenize('\n')[0]
|
|
new_inputs = self.processor(
|
|
text=['\n'.join(['<|IMAGE_START|><|image@placeholder|><|IMAGE_END|>'] * len(idx_list))],
|
|
images=inputs.images,
|
|
videos=inputs.videos,
|
|
padding=True,
|
|
return_tensors='pt',
|
|
)
|
|
splited_tokens = self._split_list(new_inputs['input_ids'][0].tolist(), split_token)
|
|
# Insert image tokens into input_ids
|
|
input_ids_len = len(input_ids)
|
|
input_ids, labels, loss_scale = self._extend_tokens(input_ids, labels, loss_scale, idx_list,
|
|
lambda i: splited_tokens[i])
|
|
idx_list.append(input_ids_len)
|
|
splited_tokens.append([])
|
|
token_type_ids = []
|
|
position_ids = []
|
|
text_i, image_i, n_text_token = 0, 0, 0
|
|
for i, idx in enumerate(idx_list):
|
|
image_idx = image_i + len(splited_tokens[i])
|
|
text_len = idx - text_i
|
|
token_type_ids.append(torch.tensor([0] * (text_len))[None])
|
|
token_type_ids.append(new_inputs['token_type_ids'][:, image_i:image_idx])
|
|
text_position_ids = torch.arange(0, text_len)[None, :, None]
|
|
start_idx = 0
|
|
if position_ids:
|
|
start_idx = position_ids[-1][0, -1].max() + 1
|
|
position_ids.append(torch.concat([text_position_ids + start_idx for _ in range(3)], dim=2))
|
|
n_text_token += text_len
|
|
position_ids.append(new_inputs['position_ids'][:, image_i:image_idx] + n_text_token)
|
|
text_i = idx + 1
|
|
n_text_token -= 1 # '\n'
|
|
image_i = image_idx + 1
|
|
token_type_ids = torch.cat(token_type_ids, dim=1)
|
|
position_ids = torch.cat(position_ids, dim=1)
|
|
encoded.update(new_inputs)
|
|
encoded['token_type_ids'] = token_type_ids
|
|
encoded['position_ids'] = position_ids
|
|
encoded['input_ids'] = input_ids
|
|
encoded['labels'] = labels
|
|
encoded['loss_scale'] = loss_scale
|
|
return encoded
|
|
|
|
def _data_collator(self, batch: List[Dict[str, Any]], *, padding_to: Optional[int] = None) -> Dict[str, Any]:
|
|
res = {}
|
|
for key in ['images', 'grid_thw', 'image_type_ids']:
|
|
res[key] = self.concat_tensor(batch, key, 0)
|
|
res.update(super()._data_collator(batch, padding_to=padding_to))
|
|
return res
|
|
|
|
def generate(self, model, *args, **kwargs):
|
|
kwargs['use_cache'] = False
|
|
return super().generate(model, *args, **kwargs)
|
|
|
|
|
|
register_template(
|
|
ERNIETemplateMeta(
|
|
MLLMTemplateType.ernie_vl, template_cls=ERNIE_VLTemplate, is_thinking=True, thinking_prefix='<think>'))
|
|
|
|
ERNIE_VL_SYSTEM = ('You are a multimodal AI assistant called ERNIE developed by Baidu based on the PaddlePaddle '
|
|
'framework.')
|
|
|
|
register_template(
|
|
ERNIETemplateMeta(
|
|
MLLMTemplateType.ernie_vl_thinking,
|
|
template_cls=ERNIE_VLTemplate,
|
|
is_thinking=True,
|
|
thinking_prefix='\n<think>\n',
|
|
default_system=ERNIE_VL_SYSTEM))
|
|
|
|
|
|
class PaddleOCR1_5Template(PaddleOCRTemplate):
|
|
version = 'v1_5'
|
|
skip_prompt = True
|
|
support_padding_free = True
|
|
|
|
def _post_encode(self, model: nn.Module, inputs: Dict[str, Any]) -> Dict[str, Any]:
|
|
if not self.is_training:
|
|
return inputs
|
|
base_model = self.get_base_model(model)
|
|
input_ids = inputs['input_ids']
|
|
pixel_values = inputs.pop('pixel_values')
|
|
image_grid_thw = inputs.get('image_grid_thw')
|
|
inputs_embeds = base_model.model.language_model.embed_tokens(input_ids)
|
|
if pixel_values is not None:
|
|
image_embeds = base_model.model.get_image_features(
|
|
pixel_values, image_grid_thw, return_dict=True).pooler_output
|
|
image_embeds = image_embeds.to(inputs_embeds.device, inputs_embeds.dtype)
|
|
image_mask = base_model.model.get_placeholder_mask(
|
|
input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds)
|
|
inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds)
|
|
|
|
return {'inputs_embeds': inputs_embeds}
|
|
|
|
|
|
register_template(
|
|
ERNIETemplateMeta(
|
|
MLLMTemplateType.paddle_ocr_1_5, prompt=['User: {{QUERY}}\nAssistant:\n'], template_cls=PaddleOCR1_5Template))
|