wehub-resource-sync
681 lines
30 KiB
Python
681 lines
30 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
import inspect
|
|
import torch
|
|
from dataclasses import dataclass, field
|
|
from typing import Any, Dict, List, Literal, Optional
|
|
|
|
from swift.utils import get_env_args, get_packed_seq_params
|
|
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, Word, findall
|
|
from ..vision_utils import load_batch, load_video_cogvlm2, load_video_hf
|
|
|
|
|
|
@dataclass
|
|
class GLMTemplateMeta(TemplateMeta):
|
|
auto_add_bos: bool = True
|
|
|
|
|
|
class GLM4Template(Template):
|
|
strip_newline = True
|
|
|
|
def _swift_encode(self, inputs: StdTemplateInputs):
|
|
res_context_list, loss_scale_list, answer_len = super()._swift_encode(inputs)
|
|
if self.strip_newline:
|
|
for i, res_context in enumerate(res_context_list):
|
|
# The last round or is tool_call.
|
|
if isinstance(res_context, str) and (res_context.endswith('<|assistant|>\n')
|
|
or res_context.endswith('<think></think>\n')) and (
|
|
i + 1 >= len(res_context_list)
|
|
or '<|observation|>' in res_context_list[i + 1]):
|
|
res_context_list[i] = res_context_list[i][:-len('\n')]
|
|
return res_context_list, loss_scale_list, answer_len
|
|
|
|
def decode_generate_ids(self, *args, **kwargs):
|
|
response = super().decode_generate_ids(*args, **kwargs)
|
|
return response.lstrip('\n') if self.strip_newline else response
|
|
|
|
|
|
register_template(
|
|
GLMTemplateMeta(
|
|
LLMTemplateType.chatglm2,
|
|
prefix=['{{SYSTEM}}'],
|
|
prompt=['[Round {{ROUND1}}]\n\n问：{{QUERY}}\n\n答：'],
|
|
chat_sep=['\n\n']))
|
|
|
|
|
|
@dataclass
|
|
class ChatGLM4TemplateMeta(GLMTemplateMeta):
|
|
prefix: Prompt = field(default_factory=list)
|
|
prompt: Prompt = field(default_factory=lambda: ['<|user|>\n{{QUERY}}<|assistant|>\n'])
|
|
chat_sep: Optional[Prompt] = field(default_factory=list)
|
|
suffix: Prompt = field(default_factory=lambda: ['<|user|>'])
|
|
system_prefix: Optional[Prompt] = field(default_factory=lambda: ['<|system|>\n{{SYSTEM}}'])
|
|
|
|
agent_template: str = 'chatglm4'
|
|
stop_words: List[Word] = field(default_factory=lambda: ['<|endoftext|>', '<|user|>', '<|observation|>'])
|
|
|
|
|
|
@dataclass
|
|
class GLM4TemplateMeta(ChatGLM4TemplateMeta):
|
|
prefix: Prompt = field(default_factory=lambda: ['[gMASK]<sop>'])
|
|
system_prefix: Optional[Prompt] = field(default_factory=lambda: ['[gMASK]<sop><|system|>\n{{SYSTEM}}'])
|
|
agent_template: str = 'glm4'
|
|
|
|
|
|
@dataclass
|
|
class GLM4_5TemplateMeta(GLM4TemplateMeta):
|
|
agent_template: str = 'glm4_5'
|
|
is_thinking: bool = True
|
|
non_thinking_prefix: str = '<think></think>\n'
|
|
history_thinking_prefix: str = '<think></think>\n'
|
|
|
|
|
|
class ChatGLM4VTemplate(Template):
|
|
|
|
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 ['<|begin_of_image|><|endoftext|><|end_of_image|>']
|
|
return [[-100]]
|
|
|
|
def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]:
|
|
encoded = super()._encode(inputs)
|
|
input_ids = encoded['input_ids']
|
|
labels = encoded['labels']
|
|
idx_list = findall(input_ids, -100)
|
|
if idx_list:
|
|
idx = idx_list[0]
|
|
image = inputs.images[0]
|
|
placeholder = '<|begin_of_image|><|endoftext|><|end_of_image|>'
|
|
placeholder_id = self.processor.encode(placeholder, add_special_tokens=False)
|
|
input_ids = (input_ids[:idx] + placeholder_id + input_ids[idx + 1:])
|
|
if labels is not None:
|
|
labels = (labels[:idx] + [-100] * len(placeholder_id) + labels[idx + 1:])
|
|
messages = inputs.messages
|
|
messages[0]['image'] = image
|
|
inputs2: Dict[str, Any] = self.processor.apply_chat_template(messages, return_dict=True)
|
|
encoded['images'] = inputs2['images']
|
|
encoded['input_ids'] = input_ids
|
|
encoded['labels'] = labels
|
|
encoded['position_ids'] = list(range(len(input_ids)))
|
|
return encoded
|
|
|
|
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
|
|
|
|
|
|
class GLM4vPackingTemplateMixin:
|
|
support_padding_free = True # https://github.com/huggingface/transformers/issues/39685
|
|
use_model = True
|
|
|
|
def packing_row(self, row: List[Dict[str, Any]]) -> Dict[str, Any]:
|
|
for r in row:
|
|
r_copy = r.copy()
|
|
r_copy['input_ids'] = torch.tensor(r_copy['input_ids'])[None]
|
|
r.update(self._get_position_ids(r_copy))
|
|
packed = super().packing_row(row)
|
|
return packed
|
|
|
|
def _get_position_ids(self, inputs: Dict[str, Any]):
|
|
base_model = self.get_base_model(self._get_model())
|
|
attention_mask = inputs.get('attention_mask_2d')
|
|
if attention_mask is None:
|
|
attention_mask = inputs.get('attention_mask')
|
|
kwargs = {}
|
|
input_ids = inputs['input_ids']
|
|
get_rope_index = base_model.model.get_rope_index
|
|
if 'mm_token_type_ids' in inspect.signature(get_rope_index).parameters:
|
|
kwargs['mm_token_type_ids'] = self.create_mm_token_type_ids(input_ids)
|
|
elif not self.is_training:
|
|
return {}
|
|
position_ids, _ = get_rope_index(
|
|
input_ids,
|
|
image_grid_thw=inputs.get('image_grid_thw'),
|
|
video_grid_thw=inputs.get('video_grid_thw'),
|
|
attention_mask=attention_mask,
|
|
**kwargs)
|
|
return {'position_ids': self._concat_text_position_ids(position_ids)}
|
|
|
|
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)
|
|
if not self.padding_free:
|
|
res.update(self._get_position_ids(res))
|
|
if 'position_ids' in res and self.is_training:
|
|
position_ids = res['position_ids']
|
|
res['position_ids'] = position_ids[1:]
|
|
res['text_position_ids'] = text_position_ids = position_ids[0]
|
|
# https://github.com/huggingface/transformers/pull/40194
|
|
if text_position_ids.shape[0] == 1:
|
|
res.update(get_packed_seq_params(text_position_ids))
|
|
return res
|
|
|
|
def _patch_create_causal_mask(self, modeling_module):
|
|
create_causal_mask = modeling_module.create_causal_mask
|
|
|
|
def new_create_causal_mask(*args, **kwargs):
|
|
position_ids = kwargs.get('position_ids')
|
|
if position_ids is not None and position_ids.dim() == 3:
|
|
kwargs['position_ids'] = None
|
|
return create_causal_mask(*args, **kwargs)
|
|
|
|
modeling_module.create_causal_mask = new_create_causal_mask
|
|
|
|
|
|
register_template(
|
|
ChatGLM4TemplateMeta(MLLMTemplateType.chatglm4v, template_cls=ChatGLM4VTemplate, suffix=['<|endoftext|>']))
|
|
|
|
register_template(ChatGLM4TemplateMeta(LLMTemplateType.chatglm4, template_cls=GLM4Template))
|
|
|
|
|
|
class GLM4VTemplate(GLM4vPackingTemplateMixin, Template):
|
|
begin_of_image_token = 151339
|
|
end_of_image_token = 151340
|
|
begin_of_video_token = 151341
|
|
end_of_video_token = 151342
|
|
placeholder_tokens = ['<|image|>', '<|video|>']
|
|
|
|
def init_processor(self, processor) -> None:
|
|
if processor is None:
|
|
return
|
|
super().init_processor(processor)
|
|
if not getattr(GLM4VTemplate, '_patched', False) and self.padding_free:
|
|
GLM4VTemplate._patched = True
|
|
from transformers.models.glm4v import modeling_glm4v
|
|
self._patch_create_causal_mask(modeling_glm4v)
|
|
self.image_token = self._tokenize('<|image|>')[0]
|
|
self.video_token = self._tokenize('<|video|>')[0]
|
|
|
|
def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int,
|
|
inputs: StdTemplateInputs) -> List[Context]:
|
|
# TODO: model video infer bug
|
|
if self.mode == 'vllm':
|
|
if media_type == 'image':
|
|
return ['<|begin_of_image|><|image|><|end_of_image|>']
|
|
elif media_type == 'video':
|
|
return ['<|begin_of_video|><|video|><|end_of_video|>']
|
|
assert media_type in ['image']
|
|
if media_type == 'image':
|
|
return [[-100]]
|
|
elif media_type == 'video':
|
|
return [[-200]]
|
|
|
|
def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]:
|
|
encoded = super()._encode(inputs)
|
|
processor = self.processor
|
|
input_ids = encoded['input_ids']
|
|
labels = encoded['labels']
|
|
image_idx_list = findall(input_ids, -100)
|
|
video_idx_list = findall(input_ids, -200)
|
|
if image_idx_list:
|
|
images = inputs.images
|
|
image_inputs = processor.image_processor(images=images, return_tensors='pt')
|
|
encoded['pixel_values'] = image_inputs['pixel_values']
|
|
encoded['image_grid_thw'] = image_grid_thw = image_inputs['image_grid_thw']
|
|
merge_length = processor.image_processor.merge_size**2
|
|
added_tokens_len = 0
|
|
for i, idx in enumerate(image_idx_list):
|
|
num_image_tokens = image_grid_thw[i].prod() // merge_length
|
|
image_tokens = [self.begin_of_image_token
|
|
] + [self.image_token] * num_image_tokens + [self.end_of_image_token]
|
|
|
|
input_ids = input_ids[:added_tokens_len + idx] + image_tokens + input_ids[added_tokens_len + idx + 1:]
|
|
if labels is not None:
|
|
labels = labels[:added_tokens_len + idx] + [-100] * len(image_tokens) + labels[added_tokens_len
|
|
+ idx + 1:]
|
|
added_tokens_len += len(image_tokens) - 1
|
|
|
|
if video_idx_list:
|
|
# TODO: model video infer bug
|
|
assert len(
|
|
video_idx_list) <= 1, f'GLM4.1V model only support 1 video, but detected {len(video_idx_list)} <video> '
|
|
assert not image_idx_list, "GLM4.1V model doesn't support inputs containing both video and images"
|
|
|
|
video_fnames = inputs.videos
|
|
import numpy as np
|
|
from transformers.image_utils import load_image
|
|
from transformers.video_utils import load_video
|
|
video_metadata = []
|
|
videos = []
|
|
for fname in video_fnames:
|
|
if isinstance(fname, (list, tuple)) and isinstance(fname[0], str):
|
|
video = [np.array(load_image(image_fname)) for image_fname in fname]
|
|
# create a 4D video because `load_video` always returns a 4D array
|
|
video = np.stack(video)
|
|
metadata = None
|
|
else:
|
|
video, metadata = load_video(fname)
|
|
videos.append(video)
|
|
video_metadata.append(metadata)
|
|
videos = [videos]
|
|
video_metadata = [video_metadata]
|
|
|
|
videos_inputs = processor.video_processor(videos=videos, video_metadata=video_metadata, return_tensors='pt')
|
|
encoded['pixel_values_videos'] = videos_inputs['pixel_values_videos']
|
|
encoded['video_grid_thw'] = video_grid_thw = videos_inputs['video_grid_thw']
|
|
timestamps = videos_inputs.pop('timestamps')
|
|
num_frames = len(video_grid_thw)
|
|
video_structure = [self.begin_of_video_token]
|
|
if hasattr(timestamps, 'tolist'):
|
|
timestamps_list = timestamps.tolist()[0]
|
|
else:
|
|
timestamps_list = timestamps[0] if isinstance(timestamps[0], list) else timestamps
|
|
unique_timestamps = []
|
|
for idx in range(0, len(timestamps_list)):
|
|
unique_timestamps.append(timestamps_list[idx])
|
|
selected_timestamps = unique_timestamps[:num_frames]
|
|
while len(selected_timestamps) < num_frames:
|
|
selected_timestamps.append(selected_timestamps[-1] if selected_timestamps else 0)
|
|
merge_length = processor.video_processor.merge_size**2
|
|
added_tokens_len = 0
|
|
for frame_idx in range(num_frames):
|
|
timestamp_sec = selected_timestamps[frame_idx]
|
|
num_image_tokens = video_grid_thw[frame_idx].prod() // merge_length
|
|
timestamp_sec_token = processor.tokenizer(str(timestamp_sec))['input_ids']
|
|
frame_structure = [self.begin_of_image_token] + [self.image_token] * num_image_tokens + \
|
|
[self.end_of_image_token] + timestamp_sec_token
|
|
video_structure += frame_structure
|
|
video_structure += [self.end_of_video_token]
|
|
|
|
for i, idx in enumerate(video_idx_list):
|
|
# BUG in GLM4.1V?: All video placeholder take same tokens
|
|
# https://github.com/huggingface/transformers/blob/v4.53.0/src/transformers/models/glm4v/processing_glm4v.py#L165-L194
|
|
input_ids = input_ids[:added_tokens_len + idx] + video_structure + \
|
|
input_ids[added_tokens_len + idx + 1:]
|
|
if labels is not None:
|
|
labels = labels[:added_tokens_len + idx] + [-100] * len(video_structure) + \
|
|
labels[added_tokens_len + idx + 1:]
|
|
added_tokens_len += len(video_structure) - 1
|
|
|
|
encoded['input_ids'] = input_ids
|
|
encoded['labels'] = labels
|
|
return encoded
|
|
|
|
def _post_encode(self, model, inputs: Dict[str, Any]) -> Dict[str, Any]:
|
|
# TODO: check video
|
|
if not self.is_training:
|
|
return inputs
|
|
input_ids = inputs['input_ids']
|
|
inputs_embeds = model.get_input_embeddings()(input_ids)
|
|
inputs_embeds = self._get_inputs_embeds_hf(inputs_embeds, inputs, model.visual, self.processor, model.config)
|
|
return {'inputs_embeds': inputs_embeds}
|
|
|
|
|
|
register_template(GLM4TemplateMeta(LLMTemplateType.glm4, template_cls=GLM4Template, thinking_prefix='<think>'))
|
|
register_template(GLM4TemplateMeta(MLLMTemplateType.glm4v, template_cls=GLM4VTemplate))
|
|
|
|
|
|
class GLM4_5Template(GLM4Template):
|
|
|
|
def _jinja_encode(self, inputs: StdTemplateInputs):
|
|
for message in inputs.messages:
|
|
if message['role'] == 'assistant' and isinstance(message['content'],
|
|
str) and message['content'].endswith('<|observation|>'):
|
|
message['content'] = message['content'][:-len('<|observation|>')]
|
|
return super()._jinja_encode(inputs)
|
|
|
|
|
|
register_template(GLM4_5TemplateMeta(LLMTemplateType.glm4_5, template_cls=GLM4_5Template))
|
|
|
|
|
|
@dataclass
|
|
class GLM4_7TemplateMeta(GLM4_5TemplateMeta):
|
|
prompt: Prompt = field(default_factory=lambda: ['<|user|>{{QUERY}}<|assistant|>'])
|
|
system_prefix: Optional[Prompt] = field(default_factory=lambda: ['[gMASK]<sop><|system|>{{SYSTEM}}'])
|
|
|
|
thinking_prefix: str = '<think>'
|
|
non_thinking_prefix: str = '</think>'
|
|
history_thinking_prefix: str = '</think>'
|
|
|
|
|
|
register_template(GLM4_7TemplateMeta(
|
|
LLMTemplateType.glm4_7,
|
|
template_cls=GLM4_5Template,
|
|
agent_template='glm4_7',
|
|
))
|
|
|
|
register_template(GLM4_7TemplateMeta(
|
|
LLMTemplateType.glm5_1,
|
|
template_cls=GLM4_5Template,
|
|
agent_template='glm5_1',
|
|
))
|
|
|
|
|
|
class GLM5_2Template(GLM4_5Template):
|
|
|
|
def init_env_args(self):
|
|
super().init_env_args()
|
|
# reasoning_effort: "max" or "high"
|
|
self.reasoning_effort = get_env_args('reasoning_effort', str, 'max')
|
|
self.chat_template_kwargs['reasoning_effort'] = self.reasoning_effort
|
|
|
|
def _get_system(self, inputs):
|
|
system = super()._get_system(inputs)
|
|
reasoning_effort = inputs.chat_template_kwargs.get('reasoning_effort')
|
|
if reasoning_effort is None:
|
|
reasoning_effort = self.reasoning_effort
|
|
if self._get_enable_thinking(inputs):
|
|
effort_str = f'Reasoning Effort: {reasoning_effort.capitalize()}'
|
|
if system:
|
|
system = f'{effort_str}<|system|>{system}'
|
|
else:
|
|
system = effort_str
|
|
return system
|
|
|
|
|
|
register_template(
|
|
GLM4_7TemplateMeta(
|
|
LLMTemplateType.glm5_2,
|
|
template_cls=GLM5_2Template,
|
|
agent_template='glm5_1',
|
|
non_thinking_prefix='<think></think>',
|
|
history_thinking_prefix='<think></think>',
|
|
))
|
|
|
|
|
|
class GLM4_5VTemplate(GLM4vPackingTemplateMixin, GLM4_5Template):
|
|
placeholder_tokens = ['<|image|>', '<|video|>']
|
|
strip_newline = False
|
|
|
|
def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int,
|
|
inputs: StdTemplateInputs) -> List[Context]:
|
|
if media_type == 'image':
|
|
return ['<|begin_of_image|><|image|><|end_of_image|>']
|
|
elif media_type == 'video':
|
|
return ['<|begin_of_video|><|video|><|end_of_video|>']
|
|
|
|
def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]:
|
|
encoded = super()._encode(inputs)
|
|
input_ids = encoded['input_ids']
|
|
for mm_type in ['image', 'video']:
|
|
mm_token = f'<|{mm_type}|>'
|
|
mm_token_id = self._tokenize(mm_token)[0]
|
|
|
|
idx_list = findall(input_ids, mm_token_id)
|
|
if idx_list:
|
|
split_token = self._tokenize('\n')[0]
|
|
mm_data = getattr(inputs, f'{mm_type}s')
|
|
if mm_type == 'image':
|
|
kwargs = {'images': mm_data}
|
|
else:
|
|
videos, video_metadata = load_video_hf(mm_data)
|
|
kwargs = {'videos': [videos], 'video_metadata': [video_metadata]}
|
|
mm_inputs = self.processor(text='\n'.join([mm_token] * len(mm_data)), return_tensors='pt', **kwargs)
|
|
splited_tokens = self._split_list(mm_inputs['input_ids'][0].tolist(), split_token)
|
|
for key in ['input_ids', 'token_type_ids', 'attention_mask']:
|
|
mm_inputs.pop(key, None)
|
|
input_ids, encoded['labels'], encoded['loss_scale'] = self._extend_tokens(
|
|
input_ids, encoded['labels'], encoded['loss_scale'], idx_list, lambda i: splited_tokens[i])
|
|
encoded.update(mm_inputs)
|
|
encoded['input_ids'] = input_ids
|
|
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']
|
|
base_model = self.get_base_model(model)
|
|
inputs_embeds = base_model.model.language_model.embed_tokens(input_ids)
|
|
inputs_embeds = self._get_inputs_embeds_hf(inputs_embeds, inputs, model.visual, self.processor, model.config)
|
|
return {'inputs_embeds': inputs_embeds}
|
|
|
|
def init_processor(self, processor) -> None:
|
|
super().init_processor(processor)
|
|
if not getattr(GLM4_5VTemplate, '_patched', False) and self.padding_free:
|
|
GLM4_5VTemplate._patched = True
|
|
from transformers.models.glm4v_moe import modeling_glm4v_moe
|
|
self._patch_create_causal_mask(modeling_glm4v_moe)
|
|
|
|
|
|
register_template(GLM4_5TemplateMeta(MLLMTemplateType.glm4_5v, template_cls=GLM4_5VTemplate))
|
|
|
|
glm4z1rumination_system = (
|
|
'你是一个专业的深度研究助手，通过提供的工具与模拟浏览器交互，来帮助用户完成深度信息调研和报告撰写任务。'
|
|
'今年是 2025 年。\n\n'
|
|
'<核心要求>\n'
|
|
'- 首先分解用户请求，得到包含多个子要求的列表\n'
|
|
'- 制定初始研究计划\n'
|
|
'- 进行多轮迭代搜索和页面浏览（at least 10 function calls）：\n'
|
|
' * 根据已获得的信息调整研究计划和关键词\n'
|
|
' * 打开页面阅读，从发现的内容中识别新的关键概念/名词\n'
|
|
' * 从搜索结果中提取新的关键词继续搜索\n'
|
|
' * 访问并仔细阅读相关页面，识别新的关键概念/名词\n\n'
|
|
'<重要配置>\n'
|
|
'- 采用语言\n'
|
|
' * 搜索关键词：英文\n'
|
|
' * 思考：英文\n\n'
|
|
'<可调用的工具列表>\n'
|
|
'[{"name": "search", "description": "Execute a search query and return search results. '
|
|
'Use this function when you need to find information about a specific topic.", '
|
|
'"parameters": {"type": "object", "properties": {"query": {"type": "string", '
|
|
'"description": "Search query string, use English words unless it is a proper name in Chinese"}}, '
|
|
'"required": ["query"], "additionalProperties": false}}, '
|
|
'{"name": "click", "description": "Click a link in the search results and navigate to the corresponding page. '
|
|
'Use this function when you need to view detailed content of a specific search result.", '
|
|
'"parameters": {"type": "object", "properties": {"link_id": {"type": "integer", '
|
|
'"description": "The link ID to click (from the sequence number in search results)"}}, '
|
|
'"required": ["link_id"], "additionalProperties": false}}, '
|
|
'{"name": "open", "description": "Open a specific website. Get content from any website with its URL.", '
|
|
'"parameters": {"type": "object", "properties": {"url": {"type": "string", '
|
|
'"description": "The target website URL or domain"}}, "required": ["url"], "additionalProperties": false}}, '
|
|
'{"name": "finish", "description": "Finish the task. '
|
|
'Use this function when you have found the information you need.", '
|
|
'"parameters": {"type": "object", "properties": {}, "additionalProperties": false}}]')
|
|
|
|
register_template(
|
|
GLM4TemplateMeta(
|
|
LLMTemplateType.glm4_z1_rumination,
|
|
template_cls=GLM4Template,
|
|
default_system=glm4z1rumination_system,
|
|
is_thinking=True))
|
|
|
|
codegeex4_system = '你是一位智能编程助手，你叫CodeGeeX。你会为用户回答关于编程、代码、计算机方面的任何问题，并提供格式规范、可以执行、准确安全的代码，并在必要时提供详细的解释。'
|
|
|
|
register_template(ChatGLM4TemplateMeta(LLMTemplateType.codegeex4, default_system=codegeex4_system))
|
|
|
|
register_template(
|
|
TemplateMeta(
|
|
LLMTemplateType.longwriter_llama, ['[INST]'], ['{{QUERY}}[/INST]'], ['[INST]'], ['<|end_of_text|>'],
|
|
system_prefix=['<<SYS>>\n{{SYSTEM}}\n<</SYS>>\n\n']))
|
|
|
|
|
|
class CogTemplate(Template):
|
|
placeholder_tokens = ['<|reserved_special_token_0|>']
|
|
|
|
use_model = True
|
|
|
|
def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int,
|
|
inputs: StdTemplateInputs) -> List[Context]:
|
|
return []
|
|
|
|
def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]:
|
|
encoded = super()._encode(inputs)
|
|
model = self.model
|
|
image = inputs.images or []
|
|
history_inputs = inputs.to_history()
|
|
inputs2 = model.build_conversation_input_ids(
|
|
self.processor, query=history_inputs['query'], history=history_inputs['history'], images=image)
|
|
image_token_len = inputs2['token_type_ids'].sum().item()
|
|
input_ids = encoded['input_ids']
|
|
labels = encoded['labels']
|
|
encoded['token_type_ids'] = [0] + [1] * image_token_len + [0] * len(input_ids[1:])
|
|
encoded['input_ids'] = input_ids[:1] + [self.processor.pad_token_id] * image_token_len + input_ids[1:]
|
|
if labels is not None:
|
|
encoded['labels'] = labels[:1] + [-100] * image_token_len + labels[1:]
|
|
if len(image) > 0:
|
|
encoded['images'] = [[img.to(dtype=self.model_info.torch_dtype)] for img in inputs2['images']]
|
|
if 'cross_images' in inputs2:
|
|
# is cogagent
|
|
encoded['cross_images'] = [[cross_img.to(dtype=self.model_info.torch_dtype)]
|
|
for cross_img in inputs2['cross_images']]
|
|
return encoded
|
|
|
|
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)
|
|
keys = ['images', 'cross_images']
|
|
for key in keys:
|
|
if key in batch[0]:
|
|
res[key] = [b[key][0] for b in batch]
|
|
return res
|
|
|
|
|
|
register_template(
|
|
TemplateMeta(
|
|
MLLMTemplateType.cogagent_chat,
|
|
prefix=['<s>'],
|
|
prompt=[' [INST] {{QUERY}} [/INST] '],
|
|
chat_sep=[],
|
|
suffix=['</s>'],
|
|
template_cls=CogTemplate,
|
|
))
|
|
|
|
register_template(
|
|
TemplateMeta(
|
|
MLLMTemplateType.cogagent_vqa,
|
|
prefix=['<s>'],
|
|
prompt=['<EOI>Question: {{QUERY}} Answer:'],
|
|
chat_sep=None,
|
|
suffix=['</s>'],
|
|
template_cls=CogTemplate))
|
|
|
|
|
|
@dataclass
|
|
class CogVLMTemplateMeta(TemplateMeta):
|
|
prefix: Prompt = field(default_factory=lambda: [['bos_token_id']])
|
|
prompt: Prompt = field(default_factory=lambda: ['Question: {{QUERY}} Answer:'])
|
|
chat_sep: Optional[Prompt] = field(default_factory=lambda: ['\n'])
|
|
|
|
|
|
register_template(CogVLMTemplateMeta(MLLMTemplateType.cogvlm, template_cls=CogTemplate))
|
|
|
|
register_template(CogVLMTemplateMeta(MLLMTemplateType.cogvlm2, template_cls=CogTemplate))
|
|
|
|
|
|
class Cog2VideoTemplate(CogTemplate):
|
|
use_model = True
|
|
|
|
def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]:
|
|
model = self.model
|
|
encoded = super(CogTemplate, self)._encode(inputs)
|
|
videos_path = inputs.videos or []
|
|
video = load_batch(videos_path, load_video_cogvlm2)
|
|
history_inputs = inputs.to_history()
|
|
inputs2 = model.build_conversation_input_ids(
|
|
self.processor,
|
|
query=history_inputs['query'],
|
|
history=history_inputs['history'],
|
|
images=video,
|
|
template_version='chat')
|
|
video_token_len = inputs2['token_type_ids'].sum().item()
|
|
input_ids = encoded['input_ids']
|
|
labels = encoded['labels']
|
|
encoded['token_type_ids'] = [0] + [1] * video_token_len + [0] * len(input_ids[1:])
|
|
encoded['input_ids'] = input_ids[:1] + [self.processor.pad_token_id] * video_token_len + input_ids[1:]
|
|
if labels is not None:
|
|
encoded['labels'] = labels[:1] + [-100] * video_token_len + labels[1:]
|
|
if len(video) > 0:
|
|
dtype = model.dtype
|
|
encoded['images'] = [[img.to(dtype=dtype)] for img in inputs2['images']]
|
|
return encoded
|
|
|
|
|
|
register_template(CogVLMTemplateMeta(
|
|
MLLMTemplateType.cogvlm2_video,
|
|
template_cls=Cog2VideoTemplate,
|
|
))
|
|
|
|
|
|
class GLMEdgeVTemplate(Template):
|
|
placeholder_tokens = ['<|begin_of_image|>']
|
|
|
|
def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int,
|
|
inputs: StdTemplateInputs) -> List[Context]:
|
|
assert media_type == 'image'
|
|
return ['<|begin_of_image|>' * 578]
|
|
|
|
def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]:
|
|
encoded = super()._encode(inputs)
|
|
images = inputs.images
|
|
if images:
|
|
encoded['pixel_values'] = torch.tensor(self.processor(images).pixel_values)
|
|
return encoded
|
|
|
|
|
|
register_template(
|
|
ChatGLM4TemplateMeta(
|
|
MLLMTemplateType.glm_edge_v,
|
|
prompt=['<|user|>\\n{{QUERY}}\\n<|assistant|>\\n'],
|
|
chat_sep=['\\n'],
|
|
system_prefix=['<|system|>\\n{{SYSTEM}}\\n'],
|
|
suffix=['<|endoftext|>'],
|
|
template_cls=GLMEdgeVTemplate,
|
|
))
|
|
|
|
|
|
class GLMOCRTemplate(Template):
|
|
begin_of_image_token = 59256
|
|
end_of_image_token = 59257
|
|
placeholder_tokens = ['<|image|>']
|
|
|
|
def init_processor(self, processor) -> None:
|
|
if processor is None:
|
|
return
|
|
super().init_processor(processor)
|
|
self.image_token = self._tokenize('<|image|>')[0]
|
|
|
|
def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int,
|
|
inputs: StdTemplateInputs) -> List[Context]:
|
|
assert media_type in ['image']
|
|
if self.mode == 'vllm':
|
|
return ['<|begin_of_image|><|image|><|end_of_image|>']
|
|
return [[-100]]
|
|
|
|
def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]:
|
|
encoded = super()._encode(inputs)
|
|
processor = self.processor
|
|
input_ids = encoded['input_ids']
|
|
labels = encoded['labels']
|
|
image_idx_list = findall(input_ids, -100)
|
|
if image_idx_list:
|
|
images = inputs.images
|
|
image_inputs = processor.image_processor(images=images, return_tensors='pt')
|
|
encoded['pixel_values'] = image_inputs['pixel_values']
|
|
encoded['image_grid_thw'] = image_grid_thw = image_inputs['image_grid_thw']
|
|
merge_length = processor.image_processor.merge_size**2
|
|
added_tokens_len = 0
|
|
for i, idx in enumerate(image_idx_list):
|
|
num_image_tokens = image_grid_thw[i].prod() // merge_length
|
|
image_tokens = [self.begin_of_image_token
|
|
] + [self.image_token] * num_image_tokens + [self.end_of_image_token]
|
|
|
|
input_ids = input_ids[:added_tokens_len + idx] + image_tokens + input_ids[added_tokens_len + idx + 1:]
|
|
if labels is not None:
|
|
labels = labels[:added_tokens_len + idx] + [-100] * len(image_tokens) + labels[added_tokens_len
|
|
+ idx + 1:]
|
|
added_tokens_len += len(image_tokens) - 1
|
|
|
|
encoded['input_ids'] = input_ids
|
|
encoded['labels'] = labels
|
|
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']
|
|
inputs_embeds = model.get_input_embeddings()(input_ids)
|
|
inputs_embeds = self._get_inputs_embeds_hf(inputs_embeds, inputs, model.visual, self.processor, model.config)
|
|
return {'inputs_embeds': inputs_embeds}
|
|
|
|
|
|
register_template(GLM4TemplateMeta(
|
|
MLLMTemplateType.glm_ocr,
|
|
template_cls=GLMOCRTemplate,
|
|
))
|