wehub-resource-sync
117 lines
5.7 KiB
Python
117 lines
5.7 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
||
import torch
|
||
from dataclasses import dataclass, field
|
||
from typing import Any, Dict, List, Literal, Optional
|
||
|
||
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
|
||
|
||
|
||
@dataclass
|
||
class HunYuanVLTemplateMeta(TemplateMeta):
|
||
prefix: Prompt = field(default_factory=lambda: ['<｜hy_begin▁of▁sentence｜>'])
|
||
prompt: Prompt = field(default_factory=lambda: ['{{QUERY}}<｜hy_User｜>'])
|
||
chat_sep: Optional[Prompt] = field(default_factory=lambda: ['<｜hy_Assistant｜><｜hy_begin▁of▁sentence｜>'])
|
||
suffix: Prompt = field(default_factory=lambda: ['<｜hy_Assistant｜>'])
|
||
system_prefix: Optional[Prompt] = field(
|
||
default_factory=lambda: ['<｜hy_begin▁of▁sentence｜>{{SYSTEM}}<｜hy_place▁holder▁no▁3｜>'])
|
||
|
||
|
||
class HunYuanVLTemplate(Template):
|
||
image_token_id = 120120
|
||
placeholder_tokens = ['<｜hy_place▁holder▁no▁102｜>']
|
||
|
||
# support_padding_free = True # position_ids with batch_dim of 0 does not support padding_free
|
||
|
||
def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int,
|
||
inputs: StdTemplateInputs) -> List[Context]:
|
||
assert media_type == 'image'
|
||
if self.mode == 'vllm':
|
||
return ['<｜hy_place▁holder▁no▁100｜><｜hy_place▁holder▁no▁102｜><｜hy_place▁holder▁no▁101｜>']
|
||
return [[-100]]
|
||
|
||
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, -100)
|
||
processor = self.processor
|
||
images = inputs.images
|
||
if images:
|
||
image_inputs = processor.image_processor(images=images, return_tensors='pt')
|
||
image_grid_thw = image_inputs['image_grid_thw']
|
||
merge_size = processor.image_processor.merge_size
|
||
|
||
def _get_new_tokens(i):
|
||
grid_h, grid_w = image_grid_thw[i][-2:]
|
||
patch_h = grid_h // merge_size
|
||
patch_w = grid_w // merge_size
|
||
img_tokens: List[int] = [self.image_token_id] * (patch_h * (patch_w + 1) + 2)
|
||
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
|
||
|
||
input_ids = encoded['input_ids']
|
||
position_ids = torch.arange(len(input_ids))
|
||
position_ids_w = torch.arange(len(input_ids))
|
||
position_ids_h = torch.arange(len(input_ids))
|
||
position_ids_t = torch.arange(len(input_ids))
|
||
image_tokens_cumsum = [0]
|
||
for i in range(len(image_grid_thw)):
|
||
grid_h, grid_w = image_grid_thw[i][-2:]
|
||
patch_h = grid_h // merge_size
|
||
patch_w = grid_w // merge_size
|
||
num_image_tokens = patch_h * (patch_w + 1) + 2
|
||
image_tokens_cumsum.append(image_tokens_cumsum[-1] + int(num_image_tokens))
|
||
image_token_pos_indices = torch.where(torch.tensor(input_ids) == self.image_token_id)
|
||
start_pos = image_token_pos_indices[0][image_tokens_cumsum[i]] + 1
|
||
replace_num = (patch_w + 1) * patch_h
|
||
position_ids_w[start_pos:start_pos + replace_num] = torch.tensor(
|
||
list(range(patch_w + 1)) * patch_h, dtype=torch.int64)
|
||
patch_h_list = []
|
||
for h in range(patch_h):
|
||
patch_h_list += [h] * (patch_w + 1)
|
||
position_ids_h[start_pos:start_pos + replace_num] = torch.tensor(patch_h_list, dtype=torch.int64)
|
||
position_ids_t[start_pos:start_pos + replace_num] = 0
|
||
position_ids = torch.stack([position_ids, position_ids_w, position_ids_h, position_ids_t]).unsqueeze(0)
|
||
encoded['position_ids'] = position_ids
|
||
attention_mask = torch.tensor(input_ids).ne(processor.pad_id)
|
||
encoded['attention_mask'] = attention_mask
|
||
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']
|
||
pixel_values = inputs.get('pixel_values')
|
||
image_grid_thw = inputs.get('image_grid_thw')
|
||
base_model = self.get_base_model(model)
|
||
inputs_embeds = base_model.model.embed_tokens(input_ids)
|
||
|
||
if pixel_values is not None:
|
||
pixel_values = pixel_values.to(base_model.dtype)
|
||
image_embeds = base_model.vit(pixel_values, image_grid_thw)
|
||
image_embeds = image_embeds.to(input_ids.device, non_blocking=True)
|
||
image_mask, _ = base_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}
|
||
|
||
def _pad_3d_position_ids(self,
|
||
position_ids: List[torch.Tensor],
|
||
padding_value: float = 0.,
|
||
batch_dim: int = 1) -> torch.Tensor:
|
||
batch_dim = 0
|
||
return super()._pad_3d_position_ids(position_ids, padding_value, batch_dim)
|
||
|
||
|
||
register_template(HunYuanVLTemplateMeta(MLLMTemplateType.hunyuan_ocr, template_cls=HunYuanVLTemplate))
|