wehub-resource-sync
97 lines
4.0 KiB
Python
97 lines
4.0 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
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 MegrezTemplateMeta(TemplateMeta):
|
|
prefix: Prompt = field(default_factory=lambda: ['<|role_start|>system<|role_end|>{{SYSTEM}}<|turn_end|>'])
|
|
prompt: Prompt = field(default_factory=lambda:
|
|
['<|role_start|>user<|role_end|>{{QUERY}}<|turn_end|><|role_start|>assistant<|role_end|>'])
|
|
chat_sep: Optional[Prompt] = field(default_factory=lambda: ['<|turn_end|>'])
|
|
suffix: Prompt = field(default_factory=lambda: ['<|turn_end|>'])
|
|
default_system: str = '你是Megrez-3B-Instruct，将针对用户的问题给出详细的、积极的回答。'
|
|
|
|
|
|
register_template(MegrezTemplateMeta(LLMTemplateType.megrez))
|
|
|
|
|
|
class MegrezOmniTemplate(Template):
|
|
skip_prompt = False
|
|
placeholder_tokens = ['<|unk|>']
|
|
|
|
def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int,
|
|
inputs: StdTemplateInputs) -> List[Context]:
|
|
if media_type == 'image':
|
|
return [[-1], '\n']
|
|
elif media_type == 'audio':
|
|
return [f'Audio {index + 1}: ', [-2], '\n']
|
|
|
|
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)
|
|
|
|
for mm_key in ['images', 'audios']:
|
|
mm_data = getattr(inputs, mm_key)
|
|
if not mm_data:
|
|
continue
|
|
if mm_key == 'images':
|
|
idx_list = findall(input_ids, -1)
|
|
encoding = self.processor.process_image(
|
|
mm_data,
|
|
return_tensors='pt',
|
|
)
|
|
text = self.processor.insert_image_feature_placeholders(
|
|
'<s>'.join(['(<image>./</image>)'] * len(mm_data)), encoding)
|
|
encoded['image_encoding'] = encoding
|
|
else:
|
|
idx_list = findall(input_ids, -2)
|
|
encoding = self.processor.process_audio(
|
|
mm_data,
|
|
return_tensors='pt',
|
|
)
|
|
text = self.processor.insert_audio_feature_placeholders(
|
|
'<s>'.join(['(<audio>./</audio>)'] * len(mm_data)), encoding)
|
|
encoded['audio_encoding'] = encoding
|
|
|
|
padding = text.split('<s>')
|
|
|
|
def _get_new_tokens(i):
|
|
return self._tokenize(padding[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
|
|
return encoded
|
|
|
|
def _post_encode(self, model: nn.Module, inputs: Dict[str, Any]) -> Dict[str, Any]:
|
|
_, inputs_embeds, _ = model.compose_embeddings(inputs)
|
|
inputs.pop('position_ids', None)
|
|
return {'inputs_embeds': inputs_embeds}
|
|
|
|
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)
|
|
new_batch = []
|
|
for b in batch:
|
|
text_encodings = {'input_ids': torch.tensor(b['input_ids'])}
|
|
multimodal_inputs = {'image_encoding': b.get('image_encoding'), 'audio_encoding': b.get('audio_encoding')}
|
|
new_batch.append(self.processor.merge_encodings(text_encodings, multimodal_inputs))
|
|
res.update(self.processor.data_collator(new_batch))
|
|
return res
|
|
|
|
|
|
register_template(MegrezTemplateMeta(MLLMTemplateType.megrez_omni, template_cls=MegrezOmniTemplate))
|