wehub-resource-sync
199 lines
9.1 KiB
Python
199 lines
9.1 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
|
|
import math
|
|
import torch
|
|
from dataclasses import dataclass, field
|
|
from PIL import Image
|
|
from torch import nn as nn
|
|
from typing import Any, Dict, List, Literal, Optional
|
|
|
|
from swift.utils import is_deepspeed_enabled, to_device
|
|
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 MoonlightTemplateMeta(TemplateMeta):
|
|
prefix: Prompt = field(default_factory=list)
|
|
prompt: Prompt = field(default_factory=lambda:
|
|
['<|im_user|>user<|im_middle|>{{QUERY}}<|im_end|><|im_assistant|>assistant<|im_middle|>'])
|
|
chat_sep: Optional[Prompt] = field(default_factory=lambda: ['<|im_end|>'])
|
|
suffix: Prompt = field(default_factory=lambda: ['<|im_end|>'])
|
|
system_prefix: Optional[Prompt] = field(
|
|
default_factory=lambda: ['<|im_system|>system<|im_middle|>{{SYSTEM}}<|im_end|>'])
|
|
default_system: Optional[str] = 'You are a helpful assistant'
|
|
|
|
|
|
register_template(MoonlightTemplateMeta(LLMTemplateType.moonlight))
|
|
|
|
register_template(
|
|
MoonlightTemplateMeta(
|
|
LLMTemplateType.kimi_k2, default_system='You are Kimi, an AI assistant created by Moonshot AI.'))
|
|
|
|
|
|
class KimiVLTemplate(Template):
|
|
placeholder_tokens = ['<|media_pad|>']
|
|
support_padding_free = True
|
|
skip_prompt = False
|
|
|
|
def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int,
|
|
inputs: StdTemplateInputs) -> List[Context]:
|
|
if media_type == 'image':
|
|
return ['<|media_start|>image<|media_content|><|media_pad|><|media_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)
|
|
media_token = self._tokenize('<|media_pad|>')[0]
|
|
idx_list = findall(input_ids, media_token)
|
|
if inputs.images:
|
|
image_processor = self.processor.image_processor
|
|
image_inputs = image_processor(inputs.images, return_tensors='pt')
|
|
image_grid_hws = image_inputs['image_grid_hws']
|
|
merge_length = image_processor.merge_kernel_size[0] * image_processor.merge_kernel_size[1]
|
|
|
|
def _get_new_tokens(i):
|
|
token_len = (image_grid_hws[i].prod() // merge_length)
|
|
return [media_token] * token_len
|
|
|
|
input_ids, labels, loss_scale = self._extend_tokens(input_ids, labels, loss_scale, idx_list,
|
|
_get_new_tokens)
|
|
|
|
encoded['loss_scale'] = loss_scale
|
|
encoded['input_ids'] = input_ids
|
|
encoded['labels'] = labels
|
|
encoded.update(image_inputs)
|
|
return encoded
|
|
|
|
def _data_collator_mm_data(self, batch: List[Dict[str, Any]]) -> Dict[str, Any]:
|
|
res = super()._data_collator_mm_data(batch)
|
|
image_grid_hws = self.concat_tensor(batch, 'image_grid_hws', 0)
|
|
if image_grid_hws is not None:
|
|
res['image_grid_hws'] = image_grid_hws
|
|
return res
|
|
|
|
def _post_encode(self, model: nn.Module, inputs: Dict[str, Any]) -> Dict[str, Any]:
|
|
input_ids = inputs['input_ids']
|
|
pixel_values = inputs.get('pixel_values')
|
|
inputs_embeds = model.get_input_embeddings()(input_ids)
|
|
|
|
if pixel_values is not None and pixel_values.size(0) > 0:
|
|
pixel_values = pixel_values.to(model.vision_tower.dtype)
|
|
image_features: torch.Tensor = model._extract_image_features(pixel_values, inputs['image_grid_hws'])
|
|
inputs_embeds = inputs_embeds.to(image_features[0].dtype).clone()
|
|
inputs_embeds = model._merge_with_image_features(inputs_embeds, input_ids, image_features)
|
|
elif is_deepspeed_enabled():
|
|
image_processor = self.processor.image_processor
|
|
dummy_image = Image.new('RGB', (32, 32), (0, 0, 0))
|
|
image_inputs = image_processor([dummy_image], return_tensors='pt')
|
|
pixel_values = image_inputs['pixel_values'].to(model.vision_tower.dtype)
|
|
image_features: torch.Tensor = model._extract_image_features(pixel_values, image_inputs['image_grid_hws'])
|
|
inputs_embeds = inputs_embeds + image_features.mean() * 0.
|
|
return {'inputs_embeds': inputs_embeds}
|
|
|
|
|
|
register_template(MoonlightTemplateMeta(MLLMTemplateType.kimi_vl, template_cls=KimiVLTemplate))
|
|
|
|
|
|
class KimiK25Template(Template):
|
|
placeholder_tokens = ['<|media_pad|>', '<|kimi_k25_video_placeholder|>']
|
|
jinja_enable_thinking_key = 'thinking'
|
|
support_padding_free = True
|
|
skip_prompt = False
|
|
|
|
def _get_system(self, inputs: StdTemplateInputs) -> Optional[str]:
|
|
system = super()._get_system(inputs)
|
|
if system is not None and '<|im_middle|>' not in system: # compat agent
|
|
system = f'system<|im_middle|>{system}'
|
|
return system
|
|
|
|
def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int,
|
|
inputs: StdTemplateInputs) -> List[Context]:
|
|
if media_type == 'image':
|
|
return ['<|media_begin|>image<|media_content|><|media_pad|><|media_end|>\n']
|
|
raise ValueError(f'KimiK25Template does not currently support {media_type}. '
|
|
'Please open an issue to request support.')
|
|
|
|
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)
|
|
media_token = self._tokenize('<|media_pad|>')[0]
|
|
idx_list = findall(input_ids, media_token)
|
|
if inputs.images:
|
|
image_processor = self.processor.image_processor
|
|
image_inputs = image_processor([{
|
|
'type': 'image',
|
|
'image': image
|
|
} for image in inputs.images],
|
|
return_tensors='pt')
|
|
grid_thws = image_inputs['grid_thws']
|
|
merge_length = math.prod(self.config.vision_config.merge_kernel_size)
|
|
|
|
def _get_new_tokens(i):
|
|
token_len = (grid_thws[i].prod() // merge_length)
|
|
return [media_token] * token_len
|
|
|
|
input_ids, labels, loss_scale = self._extend_tokens(input_ids, labels, loss_scale, idx_list,
|
|
_get_new_tokens)
|
|
|
|
encoded['loss_scale'] = loss_scale
|
|
encoded['input_ids'] = input_ids
|
|
encoded['labels'] = labels
|
|
encoded.update(image_inputs)
|
|
return encoded
|
|
|
|
def _data_collator_mm_data(self, batch: List[Dict[str, Any]]) -> Dict[str, Any]:
|
|
res = super()._data_collator_mm_data(batch)
|
|
grid_thws = self.concat_tensor(batch, 'grid_thws', 0)
|
|
if grid_thws is not None:
|
|
res['grid_thws'] = grid_thws
|
|
return res
|
|
|
|
def _post_encode(self, model: nn.Module, inputs: Dict[str, Any]) -> Dict[str, Any]:
|
|
input_ids = inputs['input_ids']
|
|
pixel_values = inputs.get('pixel_values')
|
|
inputs_embeds = model.get_input_embeddings()(input_ids)
|
|
|
|
if pixel_values is not None and pixel_values.size(0) > 0:
|
|
pixel_values = pixel_values.to(model.vision_tower.dtype)
|
|
image_features: torch.Tensor = model._extract_image_features(pixel_values, inputs['grid_thws'])
|
|
if model.mm_projector:
|
|
image_features = model.mm_projector(image_features)
|
|
image_features = torch.cat(image_features, dim=0)
|
|
inputs_embeds = inputs_embeds.to(image_features.dtype)
|
|
image_mask = (input_ids == self.config.media_placeholder_token_id).unsqueeze(-1).expand_as(inputs_embeds)
|
|
inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_features)
|
|
elif is_deepspeed_enabled():
|
|
image_processor = self.processor.image_processor
|
|
dummy_image = Image.new('RGB', (32, 32), (0, 0, 0))
|
|
image_inputs = image_processor([{'type': 'image', 'image': dummy_image}], return_tensors='pt')
|
|
image_inputs = to_device(image_inputs, inputs_embeds.device)
|
|
pixel_values = image_inputs['pixel_values'].to(model.vision_tower.dtype)
|
|
image_features: torch.Tensor = model._extract_image_features(pixel_values, image_inputs['grid_thws'])
|
|
if model.mm_projector:
|
|
image_features = model.mm_projector(image_features)
|
|
image_features = torch.cat(image_features, dim=0)
|
|
inputs_embeds = inputs_embeds + image_features.mean() * 0.
|
|
return {'inputs_embeds': inputs_embeds}
|
|
|
|
|
|
register_template(
|
|
MoonlightTemplateMeta(
|
|
MLLMTemplateType.kimi_k25,
|
|
template_cls=KimiK25Template,
|
|
system_prefix=['<|im_system|>{{SYSTEM}}<|im_end|>'],
|
|
default_system=None,
|
|
is_thinking=True,
|
|
thinking_prefix='<think>',
|
|
non_thinking_prefix='<think></think>',
|
|
history_thinking_prefix='<think></think>',
|
|
agent_template='kimi_k25',
|
|
))
|