wehub-resource-sync
445 lines
20 KiB
Python
445 lines
20 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from dataclasses import dataclass, field
|
|
from typing import Any, Dict, List, Literal, Optional
|
|
|
|
from swift.utils import get_logger, upper_bound
|
|
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_audio, load_vllm_video
|
|
|
|
logger = get_logger()
|
|
|
|
|
|
@dataclass
|
|
class GemmaTemplateMeta(TemplateMeta):
|
|
prefix: Prompt = field(default_factory=lambda: ['<bos>'])
|
|
prompt: Prompt = field(
|
|
default_factory=lambda: ['<start_of_turn>user\n{{QUERY}}<end_of_turn>\n<start_of_turn>model\n'])
|
|
chat_sep: Optional[Prompt] = field(default_factory=lambda: ['<end_of_turn>\n'])
|
|
suffix: Prompt = field(default_factory=lambda: ['<end_of_turn>'])
|
|
system_prefix: Optional[Prompt] = field(
|
|
default_factory=lambda: ['<bos><start_of_turn>system\n{{SYSTEM}}<end_of_turn>\n'])
|
|
|
|
|
|
register_template(GemmaTemplateMeta(LLMTemplateType.gemma))
|
|
|
|
|
|
class PaliGemmaTemplate(Template):
|
|
placeholder_tokens = ['<image>']
|
|
|
|
def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int,
|
|
inputs: StdTemplateInputs) -> List[Context]:
|
|
assert media_type == 'image'
|
|
if self.mode == 'vllm':
|
|
self.prompt = ['{{QUERY}}']
|
|
return []
|
|
else:
|
|
self.prompt = ['{{QUERY}}\n']
|
|
return ['<image>' * self.processor.image_seq_length + '<bos>']
|
|
|
|
def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]:
|
|
encoded = super()._encode(inputs)
|
|
raw_image = inputs.images
|
|
processor = self.processor
|
|
if encoded['labels'] is not None:
|
|
n = upper_bound(0, len(encoded['labels']), lambda idx: encoded['labels'][idx] == -100)
|
|
n2 = len(encoded['labels']) - n
|
|
encoded['token_type_ids'] = [0] * n + [1] * n2
|
|
else:
|
|
encoded['token_type_ids'] = [0] * len(encoded['input_ids'])
|
|
if raw_image:
|
|
model_inputs = processor(text='<image>' * len(raw_image), images=raw_image, return_tensors='pt')
|
|
encoded['pixel_values'] = model_inputs['pixel_values'].to(self.model_info.torch_dtype)
|
|
return encoded
|
|
|
|
|
|
register_template(
|
|
TemplateMeta(
|
|
MLLMTemplateType.paligemma,
|
|
prefix=[],
|
|
prompt=['{{QUERY}}\n'],
|
|
chat_sep=None,
|
|
suffix=['<eos>'],
|
|
template_cls=PaliGemmaTemplate,
|
|
))
|
|
|
|
|
|
@dataclass
|
|
class Gemma3TextTemplateMeta(TemplateMeta):
|
|
prefix: Prompt = field(default_factory=lambda: ['<bos>'])
|
|
prompt: Prompt = field(
|
|
default_factory=lambda: ['<start_of_turn>user\n{{QUERY}}<end_of_turn>\n<start_of_turn>model\n'])
|
|
chat_sep: Optional[Prompt] = field(default_factory=lambda: ['<end_of_turn>\n'])
|
|
suffix: Prompt = field(default_factory=lambda: ['<end_of_turn>'])
|
|
|
|
|
|
class Gemma3Template(Template):
|
|
|
|
def _swift_encode(self, inputs: StdTemplateInputs):
|
|
if inputs.system is not None:
|
|
system = inputs.system
|
|
inputs.system = None
|
|
inputs.messages[0]['content'] = system + '\n\n' + inputs.messages[0]['content']
|
|
for message in inputs.messages:
|
|
if message['role'] == 'assistant' and isinstance(message['content'], str):
|
|
message['content'] = message['content'].strip('\n')
|
|
return super()._swift_encode(inputs)
|
|
|
|
|
|
register_template(Gemma3TextTemplateMeta(LLMTemplateType.gemma3_text, template_cls=Gemma3Template))
|
|
|
|
|
|
class Gemma3VisionTemplate(Gemma3Template):
|
|
boi_token_id = 255999
|
|
placeholder_tokens = ['<start_of_image>']
|
|
|
|
def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int,
|
|
inputs: StdTemplateInputs) -> List[Context]:
|
|
assert media_type == 'image'
|
|
return ['<start_of_image>']
|
|
|
|
def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]:
|
|
from transformers.models.gemma3.processing_gemma3 import Gemma3ProcessorKwargs
|
|
|
|
encoded = super()._encode(inputs)
|
|
if inputs.images:
|
|
input_ids = encoded['input_ids']
|
|
labels = encoded['labels']
|
|
loss_scale = encoded.get('loss_scale', None)
|
|
idx_list = findall(input_ids, self.boi_token_id)
|
|
img_tokens = self._tokenize(self.processor.full_image_sequence)
|
|
input_ids, labels, loss_scale = self._extend_tokens(input_ids, labels, loss_scale, idx_list,
|
|
lambda _: img_tokens)
|
|
|
|
# TODO: customize
|
|
processor_kwargs = Gemma3ProcessorKwargs._defaults['images_kwargs']
|
|
image_inputs = self.processor.image_processor(inputs.images, **processor_kwargs)
|
|
image_inputs['pixel_values'] = torch.as_tensor(np.array(image_inputs['pixel_values']))
|
|
image_inputs.pop('num_crops')
|
|
|
|
array_ids = np.array(input_ids)
|
|
mm_token_type_ids = np.zeros_like(input_ids)
|
|
mm_token_type_ids[array_ids == self.processor.image_token_id] = 1
|
|
encoded['token_type_ids'] = mm_token_type_ids.tolist()
|
|
encoded['input_ids'] = input_ids
|
|
encoded['pixel_values'] = image_inputs['pixel_values']
|
|
encoded['labels'] = labels
|
|
encoded['loss_scale'] = loss_scale
|
|
return encoded
|
|
|
|
|
|
register_template(GemmaTemplateMeta(MLLMTemplateType.gemma3_vision, template_cls=Gemma3VisionTemplate))
|
|
|
|
|
|
class Gemma3nTemplate(Gemma3Template):
|
|
boi_token_id = 255999
|
|
boa_token_id = 256000
|
|
placeholder_tokens = ['<start_of_image>', '<start_of_audio>']
|
|
|
|
def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int,
|
|
inputs: StdTemplateInputs) -> List[Context]:
|
|
if media_type == 'image':
|
|
if self.mode == 'vllm':
|
|
return ['<image_soft_token>']
|
|
else:
|
|
return ['\n\n<start_of_image>']
|
|
elif media_type == 'audio':
|
|
if self.mode == 'vllm':
|
|
raise ValueError('Audio is not supported in vLLM')
|
|
inputs.audios[index] = load_audio(inputs.audios[index], self.processor.feature_extractor.sampling_rate)
|
|
return ['<start_of_audio>']
|
|
else:
|
|
raise ValueError(f'Unsupported media type: {media_type}. Supported types are: image, audio')
|
|
|
|
def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]:
|
|
from transformers.models.gemma3n.processing_gemma3n import Gemma3nProcessorKwargs
|
|
|
|
# Input validation
|
|
if not inputs.images and not inputs.audios and not inputs.messages:
|
|
raise ValueError('Provide at least one of `images`, `audios`, or `messages`.')
|
|
|
|
encoded = super()._encode(inputs)
|
|
processor = self.processor
|
|
input_ids = encoded['input_ids']
|
|
labels = encoded['labels']
|
|
loss_scale = encoded.get('loss_scale', None)
|
|
|
|
# Initialize token_type_ids and other outputs
|
|
array_ids = np.array(input_ids)
|
|
mm_token_type_ids = np.zeros_like(input_ids)
|
|
|
|
# Handle images
|
|
if inputs.images:
|
|
idx_list = findall(input_ids, self.boi_token_id)
|
|
img_tokens = self._tokenize(processor.full_image_sequence[2:])
|
|
input_ids, labels, loss_scale = self._extend_tokens(input_ids, labels, loss_scale, idx_list,
|
|
lambda _: img_tokens)
|
|
|
|
# Process images
|
|
processor_kwargs = Gemma3nProcessorKwargs._defaults.get('images_kwargs', {})
|
|
image_inputs = processor.image_processor(inputs.images, **processor_kwargs)
|
|
image_inputs['pixel_values'] = torch.as_tensor(
|
|
np.array(image_inputs['pixel_values']), dtype=self.model_info.torch_dtype)
|
|
if 'num_crops' in image_inputs:
|
|
image_inputs.pop('num_crops')
|
|
encoded.update(image_inputs)
|
|
|
|
# Handle audios
|
|
if inputs.audios:
|
|
audio_idx_list = findall(input_ids, self.boa_token_id)
|
|
if audio_idx_list:
|
|
# Get audio token sequence from processor
|
|
audio_tokens = self._tokenize(processor.full_audio_sequence)
|
|
input_ids, labels, loss_scale = self._extend_tokens(input_ids, labels, loss_scale, audio_idx_list,
|
|
lambda _: audio_tokens)
|
|
|
|
# Process audios
|
|
processor_kwargs = Gemma3nProcessorKwargs._defaults.get('audio_kwargs', {})
|
|
audio_inputs = processor.feature_extractor(inputs.audios, **processor_kwargs)
|
|
|
|
if 'input_features' in audio_inputs:
|
|
audio_inputs['input_features'] = torch.tensor(audio_inputs['input_features']).to(
|
|
self.model_info.torch_dtype)
|
|
if 'input_features_mask' in audio_inputs:
|
|
audio_inputs['input_features_mask'] = torch.tensor(audio_inputs['input_features_mask'])
|
|
encoded.update(audio_inputs)
|
|
|
|
# Update array_ids after token extension
|
|
array_ids = np.array(input_ids)
|
|
mm_token_type_ids = np.zeros_like(input_ids)
|
|
|
|
if hasattr(processor, 'image_token_id') and processor.image_token_id is not None:
|
|
mm_token_type_ids[array_ids == processor.image_token_id] = 1
|
|
|
|
if hasattr(processor, 'audio_token_id') and processor.audio_token_id is not None:
|
|
mm_token_type_ids[array_ids == processor.audio_token_id] = 3
|
|
|
|
encoded['token_type_ids'] = mm_token_type_ids.tolist()
|
|
encoded['input_ids'] = input_ids
|
|
encoded['labels'] = labels
|
|
encoded['loss_scale'] = loss_scale
|
|
return encoded
|
|
|
|
def _data_collator_mm_data(self, batch: List[Dict[str, Any]]) -> Dict[str, Any]:
|
|
"""Handle multimodal data collation for Gemma3n, including audio features"""
|
|
res = super()._data_collator_mm_data(batch)
|
|
|
|
# Handle audio features like other templates do
|
|
input_features = [b['input_features'] for b in batch if b.get('input_features') is not None]
|
|
input_features_mask = [b['input_features_mask'] for b in batch if b.get('input_features_mask') is not None]
|
|
|
|
if input_features:
|
|
res['input_features'] = torch.concat(input_features)
|
|
if input_features_mask:
|
|
res['input_features_mask'] = torch.concat(input_features_mask)
|
|
|
|
return res
|
|
|
|
|
|
register_template(GemmaTemplateMeta(MLLMTemplateType.gemma3n, template_cls=Gemma3nTemplate))
|
|
|
|
|
|
class Gemma4Template(Template):
|
|
placeholder_tokens = ['<|image|>', '<|audio|>', '<|video|>']
|
|
non_thinking_prefix_only_after_user = True
|
|
|
|
def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int,
|
|
inputs: StdTemplateInputs) -> List[Context]:
|
|
if media_type == 'image':
|
|
return ['<|image|>']
|
|
elif media_type == 'audio':
|
|
if self.mode != 'vllm':
|
|
inputs.audios[index] = load_audio(inputs.audios[index], self.processor.feature_extractor.sampling_rate)
|
|
return ['<|audio|>']
|
|
elif media_type == 'video':
|
|
if self.mode == 'vllm':
|
|
num_frames = self.processor.video_processor.num_frames
|
|
video_data, video_metadatas = load_vllm_video(inputs.videos[index], num_frames)
|
|
inputs.videos[index] = [(video_data, video_metadatas)]
|
|
return ['<|video|>']
|
|
|
|
def _get_system(self, inputs: StdTemplateInputs) -> Optional[str]:
|
|
system = super()._get_system(inputs)
|
|
if self._get_enable_thinking(inputs):
|
|
system = '<|think|>\n' + (system or '')
|
|
return system
|
|
|
|
def _add_non_thinking_prefix(self, inputs: StdTemplateInputs, thinking_prefix: str = '<|channel>thought'):
|
|
return super()._add_non_thinking_prefix(inputs, thinking_prefix=thinking_prefix)
|
|
|
|
def _remove_thinking_content(self, content: str, thinking_suffix: str = '<channel|>') -> str:
|
|
return super()._remove_thinking_content(content, thinking_suffix=thinking_suffix)
|
|
|
|
def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]:
|
|
encoded = super()._encode(inputs)
|
|
split_token = self._tokenize('\n')
|
|
media_inputs = self.processor(
|
|
text='\n'.join(['<|image|>'] * len(inputs.images) + ['<|video|>'] * len(inputs.videos)
|
|
+ ['<|audio|>'] * len(inputs.audios)),
|
|
audio=inputs.audios or None,
|
|
images=inputs.images or None,
|
|
videos=inputs.videos or None,
|
|
return_tensors='pt',
|
|
add_special_tokens=False,
|
|
)
|
|
splited_tokens = self._split_list(media_inputs['input_ids'][0].tolist(), split_token)
|
|
media_inputs.pop('input_ids')
|
|
media_inputs.pop('attention_mask')
|
|
input_ids = encoded['input_ids']
|
|
labels = encoded['labels']
|
|
loss_scale = encoded.get('loss_scale', None)
|
|
mm_mask = [False] * len(input_ids)
|
|
|
|
idx_list = []
|
|
for key in ['image', 'video', 'audio']:
|
|
token_id = getattr(self.config, f'{key}_token_id', None)
|
|
if token_id is None:
|
|
continue
|
|
idx_list += findall(input_ids, token_id)
|
|
sorted_order = sorted(range(len(idx_list)), key=lambda i: idx_list[i])
|
|
idx_list = [idx_list[i] for i in sorted_order]
|
|
splited_tokens = [splited_tokens[i] for i in sorted_order]
|
|
|
|
def _get_new_tokens(i):
|
|
return splited_tokens[i]
|
|
|
|
if idx_list:
|
|
input_ids, labels, loss_scale, mm_mask = self._extend_tokens(
|
|
input_ids, labels, loss_scale, idx_list, _get_new_tokens, mm_mask=mm_mask)
|
|
for key in [
|
|
'pixel_values', 'image_position_ids', 'pixel_values_videos', 'video_position_ids', 'input_features',
|
|
'input_features_mask'
|
|
]:
|
|
if key in media_inputs:
|
|
encoded[key] = media_inputs[key]
|
|
# unpad input_features
|
|
# https://github.com/vllm-project/vllm/blob/v0.23.0/vllm/model_executor/models/gemma4_mm.py#L747-L758
|
|
if 'input_features' in encoded and 'input_features_mask' in encoded:
|
|
masks = encoded['input_features_mask']
|
|
features = encoded['input_features']
|
|
if isinstance(masks, torch.Tensor) and masks.ndim >= 2:
|
|
bool_masks = masks.bool()
|
|
encoded['input_features'] = torch.stack([f[m] for f, m in zip(features, bool_masks)])
|
|
encoded['input_features_mask'] = torch.stack([m[m] for m in bool_masks])
|
|
encoded['input_ids'] = input_ids
|
|
encoded['labels'] = labels
|
|
encoded['loss_scale'] = loss_scale
|
|
encoded['mm_token_type_ids'] = self.create_mm_token_type_ids(input_ids, mm_mask)
|
|
return encoded
|
|
|
|
def _data_collator_mm_data(self, batch: List[Dict[str, Any]]) -> Dict[str, Any]:
|
|
res = super()._data_collator_mm_data(batch)
|
|
for key in ['image_position_ids', 'video_position_ids']:
|
|
value = [b[key] for b in batch if b.get(key) is not None]
|
|
if value:
|
|
res[key] = torch.concat(value)
|
|
input_features = [b['input_features'] for b in batch if b.get('input_features') is not None]
|
|
if input_features:
|
|
input_features_mask = [b['input_features_mask'] for b in batch if b.get('input_features_mask') is not None]
|
|
max_len = max([x.shape[1] for x in input_features_mask])
|
|
res['input_features'] = torch.concat([F.pad(x, (0, 0, 0, max_len - x.shape[1])) for x in input_features])
|
|
res['input_features_mask'] = torch.concat(
|
|
[F.pad(x, (0, max_len - x.shape[1])) for x in input_features_mask])
|
|
return res
|
|
|
|
|
|
@dataclass
|
|
class Gemma4TemplateMeta(TemplateMeta):
|
|
prefix: Prompt = field(default_factory=lambda: ['<bos>'])
|
|
prompt: Prompt = field(default_factory=lambda: ['<|turn>user\n{{QUERY}}<turn|>\n<|turn>model\n'])
|
|
chat_sep: Optional[Prompt] = field(default_factory=lambda: ['<turn|>\n'])
|
|
suffix: Prompt = field(default_factory=lambda: ['<turn|>\n'])
|
|
system_prefix: Optional[Prompt] = field(default_factory=lambda: ['<bos><|turn>system\n{{SYSTEM}}<turn|>\n'])
|
|
stop_words: List[Word] = field(default_factory=lambda: ['<eos>', '<turn|>', '<|tool_response>'])
|
|
|
|
|
|
register_template(
|
|
Gemma4TemplateMeta(MLLMTemplateType.gemma4_nothinking, template_cls=Gemma4Template, agent_template='gemma4'))
|
|
|
|
register_template(
|
|
Gemma4TemplateMeta(
|
|
MLLMTemplateType.gemma4,
|
|
template_cls=Gemma4Template,
|
|
agent_template='gemma4',
|
|
is_thinking=True,
|
|
non_thinking_prefix='<|channel>thought\n<channel|>'))
|
|
|
|
|
|
class DiffusionGemmaTemplate(Gemma4Template):
|
|
is_encoder_decoder = True
|
|
skip_prompt = True
|
|
|
|
@property
|
|
def loss_scale(self):
|
|
loss_scale = super().loss_scale
|
|
if self.is_training and loss_scale.base_strategy != 'last_round':
|
|
logger.warning_once('DiffusionGemmaTemplate only supports the `last_round` base strategy for loss scaling. '
|
|
'Setting loss_scale.base_strategy to `last_round`.')
|
|
loss_scale.base_strategy = 'last_round'
|
|
return loss_scale
|
|
|
|
def _data_collator(self, batch: List[Dict[str, Any]], *, padding_to: Optional[int] = None) -> Dict[str, Any]:
|
|
inputs = super()._data_collator(batch, padding_to=padding_to)
|
|
if self.is_training:
|
|
inputs = self._update_inputs(inputs)
|
|
return inputs
|
|
|
|
# Code reference: https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/DiffusionGemma_(26B-A4B)-Sudoku.ipynb # noqa
|
|
def _update_inputs(self, inputs):
|
|
canvas_length = self.config.canvas_length
|
|
if inputs['labels'].shape[0] > 1:
|
|
raise ValueError('per_device_train_batch_size must be 1 for diffusion gemma')
|
|
first_idx = (inputs['labels'] != -100).int().argmax().item()
|
|
prompt_ids = inputs['input_ids'][:, :first_idx]
|
|
# reserve one slot at the end of the canvas for the explicit eos token expected by
|
|
# the diffusion sampler as the termination signal.
|
|
response_length = inputs['input_ids'].shape[1] - first_idx
|
|
if response_length > canvas_length - 1:
|
|
raise ValueError(f'response length ({response_length}) exceeds canvas_length-1 ({canvas_length - 1}); '
|
|
'please use a shorter response or increase canvas_length.')
|
|
canvas_content = inputs['input_ids'][:, first_idx:first_idx + canvas_length - 1]
|
|
# x0: clean canvas padded to canvas_length; loss is only computed on response + eos.
|
|
device = prompt_ids.device
|
|
eos_token_id = self.tokenizer.eos_token_id
|
|
pad_token_id = self.tokenizer.pad_token_id
|
|
x0 = torch.full((prompt_ids.shape[0], canvas_length), pad_token_id, dtype=torch.long, device=device)
|
|
n = canvas_content.shape[1]
|
|
x0[:, :n] = canvas_content
|
|
# explicitly append eos as the canvas-end signal expected by the diffusion sampler.
|
|
# without it, sampler keeps denoising the trailing positions during inference and emits garbage.
|
|
x0[:, n] = eos_token_id
|
|
labels = x0.clone()
|
|
labels[:, n + 1:] = -100
|
|
|
|
# forward diffusion: per-sample noise level t ∈ [min, max], replace tokens with random vocab ids
|
|
t = torch.empty((), device=device).uniform_(0.1, 1.)
|
|
noise_mask = torch.rand(canvas_length, device=device) < t
|
|
random_tokens = torch.randint(0, self.config.text_config.vocab_size, (canvas_length, ), device=device)
|
|
decoder_input_ids = torch.where(noise_mask, random_tokens, x0)
|
|
return {'input_ids': prompt_ids, 'decoder_input_ids': decoder_input_ids, 'labels': labels}
|
|
|
|
def compute_sft_loss(self, model, inputs: Dict[str, Any], num_items_in_batch: Optional[int] = None, trainer=None):
|
|
if trainer.args.gradient_checkpointing:
|
|
raise ValueError('Gradient checkpointing is not supported for diffusion gemma')
|
|
outputs = model(**inputs)
|
|
logits = outputs.logits.view(-1, outputs.logits.shape[-1])
|
|
labels = inputs['labels'].view(-1)
|
|
outputs.loss = F.cross_entropy(logits, labels, reduction='sum')
|
|
outputs.loss = outputs.loss / num_items_in_batch
|
|
return outputs
|
|
|
|
|
|
register_template(
|
|
Gemma4TemplateMeta(
|
|
MLLMTemplateType.diffusion_gemma,
|
|
template_cls=DiffusionGemmaTemplate,
|
|
agent_template='gemma4',
|
|
is_thinking=True,
|
|
non_thinking_prefix='<|channel>thought\n<channel|>'))
|