wehub-resource-sync
2296 lines
105 KiB
Python
2296 lines
105 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
import hashlib
|
|
import inspect
|
|
import math
|
|
import numpy as np
|
|
import os
|
|
import random
|
|
import re
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from collections import defaultdict
|
|
from contextlib import contextmanager, nullcontext
|
|
from copy import deepcopy
|
|
from dataclasses import asdict
|
|
from functools import partial, wraps
|
|
from modelscope.hub.utils.utils import get_cache_dir
|
|
from peft import PeftModel
|
|
from PIL import Image
|
|
from torch.nn.utils.rnn import pad_sequence
|
|
from transformers import StoppingCriteriaList
|
|
from transformers.integrations import is_deepspeed_zero3_enabled
|
|
from transformers.utils import strtobool
|
|
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal, Optional, Tuple, Union
|
|
|
|
from swift.utils import Processor, ProcessorMixin, get_env_args, get_logger, remove_response, retry_decorator, to_device
|
|
from .template_inputs import StdTemplateInputs, TemplateInputs
|
|
from .utils import Context, ContextType, StopWordsCriteria, fetch_one, findall, get_last_user_round, split_str_parts_by
|
|
from .vision_utils import _check_path, load_audio, load_batch, load_image, rescale_image
|
|
|
|
logger = get_logger()
|
|
if TYPE_CHECKING:
|
|
from swift.infer_engine import InferRequest
|
|
from .template_meta import TemplateMeta
|
|
|
|
|
|
class MaxLengthError(ValueError):
|
|
pass
|
|
|
|
|
|
class Template(ProcessorMixin):
|
|
"""Base template class for formatting and processing model inputs/outputs.
|
|
|
|
This class serves as the foundation for all template implementations in the Swift framework.
|
|
It handles the conversion between conversation formats and token sequences, manages multimodal
|
|
inputs (images, videos, audio), supports various training modes (standard, RLHF, KTO), and
|
|
provides utilities for tokenization, padding, and data collation.
|
|
|
|
The Template class is designed to be flexible and extensible, supporting:
|
|
- Multiple chat formats (user/assistant conversations, system prompts, tool calls)
|
|
- Multimodal data processing (images, videos, audio, bounding boxes)
|
|
- Different training strategies (causal language modeling, sequence classification, embedding, etc.)
|
|
- Various inference engines (Transformers, vLLM, LMDeploy, SGLang)
|
|
- Advanced features like padding-free training, sequence parallelism, and loss scaling
|
|
"""
|
|
special_tokens = ['<image>', '<video>', '<audio>', '<bbox>', '<ref-object>', '<cot-process>', '<start-image>']
|
|
special_keys = ['images', 'videos', 'audios', 'objects']
|
|
|
|
image_placeholder = ['<image>']
|
|
video_placeholder = ['<video>']
|
|
audio_placeholder = ['<audio>']
|
|
cot_process_placeholder = ['ки']
|
|
placeholder_tokens = [] # For clearer printing
|
|
load_images = True
|
|
skip_prompt = True
|
|
use_model = False
|
|
norm_bbox = 'norm1000'
|
|
# For pure text models, the default is True; for multimodal models, the default is False.
|
|
support_padding_free = None
|
|
jinja_enable_thinking_key = 'enable_thinking'
|
|
# If True, only inject non_thinking_prefix when the previous turn is a 'user' turn.
|
|
# Set this in subclasses where thinking / tool_call / tool_response / follow-up
|
|
# assistant text live in the same logical turn (e.g. Gemma4, DeepSeekV3.1), so
|
|
# the assistant turn after a tool_response should NOT open a new thinking block.
|
|
non_thinking_prefix_only_after_user: bool = False
|
|
model_accepts_loss_kwargs: Optional[bool] = None
|
|
|
|
is_encoder_decoder = False
|
|
|
|
def __init__(
|
|
self,
|
|
processor: Processor,
|
|
template_meta: 'TemplateMeta',
|
|
default_system: Optional[str] = None,
|
|
max_length: Optional[int] = None,
|
|
*,
|
|
truncation_strategy: Literal['raise', 'left', 'right', 'split'] = 'raise',
|
|
max_pixels: Optional[int] = None,
|
|
agent_template: Optional[str] = None,
|
|
norm_bbox: Literal['norm1000', 'none', None] = None,
|
|
use_chat_template: bool = True,
|
|
remove_unused_columns: bool = True,
|
|
padding_side: Literal['left', 'right'] = 'right',
|
|
# only for train
|
|
padding_free: bool = False,
|
|
loss_scale: str = 'default',
|
|
is_binary_loss_scale: Optional[bool] = None,
|
|
sequence_parallel_size: int = 1,
|
|
# infer/deploy
|
|
template_backend: Literal['swift', 'jinja'] = 'swift',
|
|
# thinking
|
|
response_prefix: Optional[str] = None,
|
|
enable_thinking: Optional[bool] = None,
|
|
preserve_thinking: Optional[bool] = None,
|
|
add_non_thinking_prefix: bool = True,
|
|
) -> None:
|
|
"""
|
|
default_system: Override the default_system in the template.
|
|
max_length: Max length of the sequence
|
|
truncation_strategy: The truncation strategy
|
|
max_pixels: Rescale image to reduce memory usage, default `None` means no limitation.
|
|
e.g. 512 * 512 (H*W)
|
|
padding_side: The padding_side when the training batch_size >= 2
|
|
loss_scale: The loss scale function to use
|
|
"""
|
|
self._processor_inited = False
|
|
self._version = 'v6' # Avoid compatibility issues caused by load_from_cache_file caching.
|
|
self.max_length = max_length
|
|
self.model = None
|
|
self.dummy_model = None
|
|
|
|
if not use_chat_template:
|
|
template_meta = template_meta.to_generate_template_meta()
|
|
else:
|
|
template_meta = deepcopy(template_meta)
|
|
# if default_system is None. not change self.default_system
|
|
template_meta.check_system(default_system)
|
|
if default_system is not None:
|
|
template_meta.default_system = default_system
|
|
if enable_thinking is None:
|
|
enable_thinking = template_meta.is_thinking and not template_meta.non_thinking_prefix
|
|
self.response_prefix = response_prefix
|
|
self.template_meta: 'TemplateMeta' = template_meta
|
|
self.use_chat_template = use_chat_template
|
|
self.enable_thinking = enable_thinking
|
|
self.preserve_thinking = preserve_thinking
|
|
self.add_non_thinking_prefix = add_non_thinking_prefix
|
|
self.chat_template_kwargs = {}
|
|
self.remove_unused_columns = remove_unused_columns
|
|
self.template_backend = template_backend
|
|
self.max_length = max_length
|
|
self.truncation_strategy = truncation_strategy
|
|
self._loss_scale_cache = {}
|
|
self._agent_template_cache = {}
|
|
self._loss_scale = loss_scale
|
|
self.is_binary_loss_scale = is_binary_loss_scale
|
|
self.max_pixels = max_pixels
|
|
self.padding_side = padding_side
|
|
self.sequence_parallel_size = sequence_parallel_size
|
|
self.padding_free = padding_free # padding_free/packing
|
|
self.packing = False
|
|
agent_template = agent_template or template_meta.agent_template
|
|
self._agent_template = agent_template
|
|
self.norm_bbox = norm_bbox or self.norm_bbox
|
|
self.mode: Literal['transformers', 'vllm', 'lmdeploy', 'sglang', 'train', 'rlhf', 'kto'] = 'transformers'
|
|
self.task_type: Literal['causal_lm', 'seq_cls', 'embedding', 'prm', 'reranker',
|
|
'generative_reranker'] = 'causal_lm'
|
|
self.use_megatron = False
|
|
self._handles = []
|
|
self._deepspeed_initialize = None
|
|
|
|
if processor is not None:
|
|
self.init_processor(processor)
|
|
|
|
def _get_enable_thinking(self, inputs=None):
|
|
enable_thinking = None if inputs is None else inputs.chat_template_kwargs.get('enable_thinking')
|
|
if enable_thinking is None:
|
|
enable_thinking = self.enable_thinking
|
|
return enable_thinking
|
|
|
|
def _get_preserve_thinking(self, inputs=None):
|
|
preserve_thinking = None if inputs is None else inputs.chat_template_kwargs.get('preserve_thinking')
|
|
if preserve_thinking is None:
|
|
preserve_thinking = self.preserve_thinking
|
|
if preserve_thinking is None:
|
|
enable_thinking = self._get_enable_thinking(inputs)
|
|
if self.template_meta.is_thinking or enable_thinking:
|
|
if self.is_training and self.loss_scale.base_strategy != 'last_round':
|
|
preserve_thinking = True
|
|
else:
|
|
preserve_thinking = False
|
|
else:
|
|
preserve_thinking = True
|
|
return preserve_thinking
|
|
|
|
def _get_response_prefix(self, inputs=None):
|
|
response_prefix = None if inputs is None else inputs.chat_template_kwargs.get('response_prefix')
|
|
if response_prefix is None:
|
|
response_prefix = self.response_prefix
|
|
if response_prefix is not None:
|
|
return response_prefix
|
|
elif not self.use_chat_template:
|
|
return ''
|
|
enable_thinking = self._get_enable_thinking(inputs)
|
|
if enable_thinking:
|
|
return self.template_meta.thinking_prefix
|
|
else:
|
|
return self.template_meta.non_thinking_prefix
|
|
|
|
@property
|
|
def loss_scale(self):
|
|
from swift.loss_scale import get_loss_scale
|
|
if self._loss_scale not in self._loss_scale_cache:
|
|
self._loss_scale_cache[self._loss_scale] = get_loss_scale(self._loss_scale)
|
|
return self._loss_scale_cache[self._loss_scale]
|
|
|
|
@property
|
|
def agent_template(self):
|
|
from swift.agent_template import agent_template_map
|
|
if self._agent_template is None:
|
|
raise ValueError(
|
|
f'Failed to automatically match an agent_template for template "{self.template_meta.template_type}". '
|
|
f'Please specify it manually via `--agent_template`. '
|
|
f'Available options: {list(agent_template_map.keys())}.')
|
|
if self._agent_template not in self._agent_template_cache:
|
|
self._agent_template_cache[self._agent_template] = agent_template_map[self._agent_template]()
|
|
return self._agent_template_cache[self._agent_template]
|
|
|
|
def init_env_args(self):
|
|
if self.model_meta.is_multimodal:
|
|
self.root_image_dir = get_env_args('ROOT_IMAGE_DIR', str, None)
|
|
else:
|
|
self.root_image_dir = None
|
|
|
|
def init_processor(self, processor: Processor) -> None:
|
|
if processor is None or self._processor_inited:
|
|
return
|
|
self._processor_inited = True
|
|
self.processor = processor
|
|
self.model_info = processor.model_info
|
|
self.config = self.model_info.config
|
|
self.task_type = self.model_info.task_type
|
|
|
|
self.model_meta = processor.model_meta
|
|
if self.max_length is None:
|
|
self.max_length = self.model_info.max_model_len
|
|
logger.info(f'default_system: {repr(self.template_meta.default_system)}')
|
|
logger.info(f'max_length: {self.max_length}')
|
|
logger.info(f'response_prefix: {repr(self.response_prefix)}')
|
|
logger.info(f'agent_template: {self._agent_template}')
|
|
if self.model_meta.is_multimodal:
|
|
logger.info(f'norm_bbox: {self.norm_bbox}')
|
|
self._init_placeholder_tokens()
|
|
self.template_meta.init(self.tokenizer)
|
|
self.init_env_args()
|
|
|
|
def _init_placeholder_tokens(self):
|
|
for mm_type in ['image', 'video', 'audio']:
|
|
mm_token = getattr(self.processor, f'{mm_type}_token', None)
|
|
mm_token_id = getattr(self.processor, f'{mm_type}_token_id', None)
|
|
if mm_token_id is not None and mm_token_id not in self.placeholder_tokens:
|
|
self.placeholder_tokens.append(mm_token_id)
|
|
elif mm_token is not None and mm_token not in self.placeholder_tokens:
|
|
self.placeholder_tokens.append(mm_token)
|
|
for i, token in enumerate(self.placeholder_tokens):
|
|
if isinstance(token, str):
|
|
self.placeholder_tokens[i] = self.tokenizer.convert_tokens_to_ids(token)
|
|
|
|
def _get_model(self):
|
|
if self.model is not None:
|
|
return self.model
|
|
if self.dummy_model is None:
|
|
from swift.model import get_model_processor
|
|
with torch.device('meta'):
|
|
self.dummy_model = get_model_processor(self.model_info.model_dir, return_dummy_model=True)[0]
|
|
return self.dummy_model
|
|
|
|
@staticmethod
|
|
def _load_image(image, load_images: bool):
|
|
if load_images:
|
|
if isinstance(image, dict) and 'bytes' in image:
|
|
image = image['bytes'] or image['path']
|
|
image = load_image(image)
|
|
else:
|
|
if isinstance(image, dict):
|
|
path = image['path']
|
|
if path and (path.startswith('http') or os.path.exists(path)):
|
|
image = path
|
|
else:
|
|
image = load_image(image['bytes'])
|
|
elif not isinstance(image, str):
|
|
image = load_image(image)
|
|
return image
|
|
|
|
@staticmethod
|
|
def _get_height_width(inputs: StdTemplateInputs) -> None:
|
|
width = []
|
|
height = []
|
|
for image in inputs.images:
|
|
width.append(image.width)
|
|
height.append(image.height)
|
|
inputs.objects['width'] = width
|
|
inputs.objects['height'] = height
|
|
|
|
def normalize_bbox(self, inputs: StdTemplateInputs) -> None:
|
|
objects = inputs.objects
|
|
bbox_list = objects['bbox']
|
|
width_list = objects['width']
|
|
height_list = objects['height']
|
|
bbox_type = objects.pop('bbox_type', None) or 'real'
|
|
image_id_list = objects.pop('image_id', None) or []
|
|
image_id_list += [0] * (len(bbox_list) - len(image_id_list))
|
|
for bbox, image_id in zip(bbox_list, image_id_list):
|
|
if bbox_type == 'norm1':
|
|
width, height = 1, 1
|
|
else:
|
|
width, height = width_list[image_id], height_list[image_id]
|
|
for i, (x, y) in enumerate(zip(bbox[::2], bbox[1::2])):
|
|
if self.norm_bbox == 'norm1000':
|
|
norm_width, norm_height = 1000, 1000
|
|
elif self.norm_bbox == 'none':
|
|
image = inputs.images[image_id]
|
|
norm_width, norm_height = image.width, image.height
|
|
bbox[2 * i] = int(round(x / width * norm_width))
|
|
bbox[2 * i + 1] = int(round(y / height * norm_height))
|
|
|
|
def _preprocess_tools(self, inputs: StdTemplateInputs) -> None:
|
|
if inputs.tools:
|
|
agent_template = self.agent_template
|
|
agent_template.template_meta = self.template_meta # for hermes
|
|
if isinstance(inputs.tools, str):
|
|
inputs.tools = agent_template._parse_json(inputs.tools)
|
|
if not isinstance(inputs.tools, (list, tuple)):
|
|
inputs.tools = [inputs.tools]
|
|
elif isinstance(inputs.tools, (list, tuple)):
|
|
inputs.tools = [agent_template._parse_json(tool) for tool in inputs.tools]
|
|
else:
|
|
raise ValueError(f'inputs.tools: {inputs.tools}')
|
|
for i, tool in enumerate(inputs.tools):
|
|
inputs.tools[i] = agent_template.wrap_tool(tool)
|
|
|
|
def _preprocess_tool_call(self, inputs: StdTemplateInputs) -> None:
|
|
i = 0
|
|
messages = inputs.messages
|
|
while i < len(messages):
|
|
if messages[i]['role'] == 'tool_call':
|
|
agent_template = self.agent_template
|
|
agent_template.template_meta = self.template_meta # for hermes
|
|
i_start = i
|
|
while i + 1 < len(messages) and messages[i + 1]['role'] == 'tool_call':
|
|
i += 1
|
|
tool_call_msgs = messages[i_start:i + 1]
|
|
tool_content = agent_template._format_tool_calls(tool_call_msgs)
|
|
pre_message = messages[i_start - 1] if i_start > 0 else None
|
|
tool_content = agent_template._add_tool_call_prefix(tool_content, pre_message)
|
|
merged_message = {'role': 'assistant', 'content': tool_content}
|
|
# Preserve loss/loss_scale fields from the first tool_call message.
|
|
for msg in tool_call_msgs:
|
|
for key in ['loss', 'loss_scale']:
|
|
if key in msg and key not in merged_message:
|
|
merged_message[key] = msg[key]
|
|
messages[i_start:i + 1] = [merged_message]
|
|
i = i_start + 1
|
|
else:
|
|
i += 1
|
|
|
|
def prepare_engine_kwargs(self) -> Dict[str, Any]:
|
|
return {}
|
|
|
|
def _get_max_pixels(self, inputs=None):
|
|
max_pixels = None if inputs is None else inputs.chat_template_kwargs.get('max_pixels')
|
|
if max_pixels is None:
|
|
max_pixels = self.max_pixels
|
|
return max_pixels
|
|
|
|
def _preprocess_inputs(
|
|
self,
|
|
inputs: StdTemplateInputs,
|
|
) -> None:
|
|
self._preprocess_tools(inputs)
|
|
if self.model_meta.is_multimodal:
|
|
self._replace_image_tags(inputs)
|
|
self._replace_start_image_tags(inputs)
|
|
|
|
images = inputs.images
|
|
load_images = self.load_images or self.mode in {'vllm', 'lmdeploy'}
|
|
load_images_origin = load_images
|
|
max_pixels = self._get_max_pixels(inputs)
|
|
if max_pixels is not None or inputs.objects:
|
|
load_images = True
|
|
if images:
|
|
for i, image in enumerate(images):
|
|
images[i] = self._load_image(images[i], load_images)
|
|
if inputs.objects:
|
|
self._get_height_width(inputs)
|
|
if max_pixels is not None:
|
|
# Scale the image proportionally without affecting the scaled objects.
|
|
images = [rescale_image(img, max_pixels) for img in images]
|
|
if images and not load_images_origin: # fix pt & qwen-vl
|
|
for i, image in enumerate(images):
|
|
if isinstance(image, Image.Image):
|
|
images[i] = self._save_pil_image(image)
|
|
inputs.images = images
|
|
|
|
# Resolve video/audio paths with ROOT_IMAGE_DIR.
|
|
# Image paths are resolved by _load_image above, but video/audio paths are
|
|
# passed as raw strings to model-specific templates. Templates that delegate
|
|
# media loading to HF processors (e.g. Gemma4) need resolved absolute paths.
|
|
if self.root_image_dir:
|
|
for media_list in (inputs.videos, inputs.audios):
|
|
for i, media_file in enumerate(media_list):
|
|
if isinstance(media_file, str) and not media_file.startswith('http'):
|
|
media_list[i] = _check_path(media_file) or media_file
|
|
|
|
if self.mode == 'vllm' and inputs.audios:
|
|
sampling_rate = get_env_args('sampling_rate', int, None)
|
|
inputs.audios = load_batch(
|
|
inputs.audios, load_func=partial(load_audio, sampling_rate=sampling_rate, return_sr=True))
|
|
if inputs.is_multimodal:
|
|
self._add_default_tags(inputs)
|
|
|
|
@staticmethod
|
|
def _replace_image_tags(inputs: StdTemplateInputs):
|
|
# compat
|
|
if inputs.images:
|
|
return
|
|
images = []
|
|
pattern = r'<img>(.+?)</img>'
|
|
for message in inputs.messages:
|
|
content = message['content']
|
|
if not isinstance(content, str):
|
|
continue
|
|
for image in re.findall(pattern, content):
|
|
# only support local_path
|
|
if os.path.isfile(image):
|
|
images.append(image)
|
|
else:
|
|
logger.warning_once(f'Failed to parse image path: `{content}`.', hash_id='<img></img>')
|
|
message['content'] = re.sub(pattern, '<image>', content)
|
|
inputs.images = images
|
|
|
|
@staticmethod
|
|
def _replace_start_image_tags(inputs: StdTemplateInputs):
|
|
# compat
|
|
generate_mode = False
|
|
message = inputs.messages[-1]
|
|
content = message['content']
|
|
if message['role'] == 'user' and content.endswith('<start-image>'):
|
|
generate_mode = True
|
|
message['content'] = message['content'][:-len('<start-image>')] # remove the <start-image>
|
|
inputs.generate_mode = generate_mode
|
|
|
|
@staticmethod
|
|
def _extend_tokens(input_ids: List[int],
|
|
labels: Optional[List[int]],
|
|
loss_scale: Optional[List[float]],
|
|
replace_idx_list: List[int],
|
|
get_new_tokens: Callable[[int], List[int]],
|
|
mm_mask: Optional[List[bool]] = None):
|
|
added_tokens_len = 0
|
|
for i, idx in enumerate(replace_idx_list):
|
|
try:
|
|
new_tokens = get_new_tokens(i)
|
|
except IndexError as e:
|
|
logger.warning(f'IndexError occurs in the _extend_tokens function: {e}.')
|
|
continue
|
|
token_len = len(new_tokens)
|
|
input_ids = input_ids[:idx + added_tokens_len] + new_tokens + input_ids[added_tokens_len + idx + 1:]
|
|
if labels:
|
|
labels = labels[:idx + added_tokens_len] + [-100] * token_len + labels[added_tokens_len + idx + 1:]
|
|
if loss_scale:
|
|
scale_idx = loss_scale[idx + added_tokens_len]
|
|
loss_scale = loss_scale[:idx + added_tokens_len] + [scale_idx] * token_len + loss_scale[added_tokens_len
|
|
+ idx + 1:]
|
|
if mm_mask:
|
|
mm_mask = mm_mask[:idx + added_tokens_len] + [True] * token_len + mm_mask[added_tokens_len + idx + 1:]
|
|
added_tokens_len += token_len - 1
|
|
if mm_mask is not None:
|
|
return input_ids, labels, loss_scale, mm_mask
|
|
return input_ids, labels, loss_scale
|
|
|
|
def forward_context(self, model, inputs):
|
|
# This function is only used to handle scenarios where the model needs
|
|
# to be patched during the forward pass.
|
|
return nullcontext()
|
|
|
|
@staticmethod
|
|
def get_base_model(model):
|
|
if isinstance(model, PeftModel):
|
|
return model.model
|
|
else:
|
|
return model
|
|
|
|
def _rlhf_encode(self, inputs: TemplateInputs, check_rejected=True) -> Dict[str, Any]:
|
|
chosen = inputs.chosen
|
|
margin = chosen.margin
|
|
chosen_encoded = self._encode_truncated(chosen)
|
|
if inputs.rejected is None:
|
|
if check_rejected:
|
|
raise ValueError('inputs.rejected is None')
|
|
rejected_encoded = {}
|
|
else:
|
|
rejected_encoded = self._encode_truncated(inputs.rejected)
|
|
|
|
encoded = {}
|
|
for prefix in ['chosen', 'rejected']:
|
|
data = locals()[f'{prefix}_encoded']
|
|
for k, v in data.items():
|
|
encoded[f'{prefix}_{k}'] = v
|
|
if margin is not None:
|
|
encoded['margin'] = float(margin)
|
|
return encoded
|
|
|
|
def _kto_encode(self, inputs: TemplateInputs) -> Dict[str, Any]:
|
|
encoded = self._rlhf_encode(inputs, check_rejected=False)
|
|
encoded['label'] = bool(inputs.chosen.label)
|
|
return encoded
|
|
|
|
def _embedding_encode(self, inputs: TemplateInputs) -> Dict[str, Any]:
|
|
_encoded = {}
|
|
labels = []
|
|
|
|
if self.is_training:
|
|
anchor = inputs.chosen
|
|
anchor_encoded = self._encode_truncated(anchor)
|
|
for key in anchor_encoded:
|
|
_encoded[f'anchor_{key}'] = anchor_encoded[key]
|
|
positive = inputs.positive
|
|
if isinstance(positive, list):
|
|
positive = positive[0]
|
|
positive_encoded = self._encode_truncated(positive)
|
|
for key in positive_encoded:
|
|
_encoded[f'positive_{key}'] = positive_encoded[key]
|
|
labels.append(float(inputs.chosen.label) if inputs.chosen.label is not None else 1.0)
|
|
|
|
_all_negative_keys = set()
|
|
for idx, negative in enumerate(inputs.negative):
|
|
_tmp_negative_keys = set() # used to fill in missing keys
|
|
negative_encoded = self._encode_truncated(negative)
|
|
for key in negative_encoded:
|
|
negative_key = f'negative_{key}'
|
|
_all_negative_keys.add(negative_key)
|
|
_tmp_negative_keys.add(negative_key)
|
|
if negative_key not in _encoded:
|
|
_encoded[negative_key] = [None] * idx
|
|
_encoded[negative_key].append(negative_encoded[key])
|
|
for miss_key in (_all_negative_keys - _tmp_negative_keys):
|
|
_encoded[miss_key].append(None)
|
|
labels.append(0.0)
|
|
|
|
_encoded['labels'] = labels
|
|
else:
|
|
anchor = inputs.chosen
|
|
_encoded = self._encode_truncated(anchor)
|
|
_encoded.pop('labels', None)
|
|
return _encoded
|
|
|
|
def _reranker_encode(self, inputs: TemplateInputs) -> Dict[str, Any]:
|
|
if self.is_training:
|
|
chosen = inputs.chosen
|
|
instruction = chosen.system
|
|
|
|
_encoded = defaultdict(list)
|
|
labels = []
|
|
|
|
for positive in inputs.positive:
|
|
if instruction is not None and positive.system is None:
|
|
positive.system = instruction
|
|
positive.messages = chosen.messages + positive.messages
|
|
positive.images = chosen.images + positive.images
|
|
positive.audios = chosen.audios + positive.audios
|
|
positive.videos = chosen.videos + positive.videos
|
|
positive_encoded = self._encode_truncated(positive)
|
|
labels.append(1)
|
|
for key in positive_encoded:
|
|
_encoded[key].append(positive_encoded[key])
|
|
|
|
for negative in inputs.negative:
|
|
if instruction is not None and negative.system is None:
|
|
negative.system = instruction
|
|
negative.messages = chosen.messages + negative.messages
|
|
negative.images = chosen.images + negative.images
|
|
negative.audios = chosen.audios + negative.audios
|
|
negative.videos = chosen.videos + negative.videos
|
|
negative_encoded = self._encode_truncated(negative)
|
|
labels.append(0)
|
|
for key in negative_encoded:
|
|
_encoded[key].append(negative_encoded[key])
|
|
|
|
_encoded['labels'] = labels
|
|
else:
|
|
anchor = inputs.chosen
|
|
_encoded = self._encode_truncated(anchor)
|
|
_encoded.pop('labels', None)
|
|
return _encoded
|
|
|
|
def _seq_cls_encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]:
|
|
encoded = self._encode_truncated(inputs)
|
|
if inputs.label is not None:
|
|
labels = inputs.label
|
|
problem_type = self.config.problem_type
|
|
if problem_type == 'single_label_classification':
|
|
labels = int(labels)
|
|
encoded['labels'] = labels
|
|
return encoded
|
|
|
|
@torch.inference_mode()
|
|
@retry_decorator(3)
|
|
def encode(self,
|
|
inputs: Union[TemplateInputs, Dict[str, Any], 'InferRequest'],
|
|
return_template_inputs: bool = False,
|
|
return_length: bool = False) -> Dict[str, Any]:
|
|
"""The entrance method of Template!
|
|
|
|
Returns:
|
|
return {'input_ids': List[int], 'labels': Optional[List[int]], ...}
|
|
"""
|
|
from swift.infer_engine import InferRequest
|
|
assert self._processor_inited, ('Please initialize the processor before calling the template.encode method: '
|
|
'template.init_processor(processor).')
|
|
if isinstance(inputs, InferRequest):
|
|
inputs = asdict(inputs)
|
|
|
|
if isinstance(inputs, dict):
|
|
inputs = TemplateInputs.from_dict(inputs)
|
|
elif isinstance(inputs, TemplateInputs):
|
|
inputs = deepcopy(inputs)
|
|
assert isinstance(inputs, TemplateInputs)
|
|
|
|
chosen = inputs.chosen
|
|
if self.task_type == 'causal_lm':
|
|
if self.mode in {'train', 'transformers', 'vllm', 'lmdeploy', 'sglang'}:
|
|
encoded = self._encode_truncated(chosen)
|
|
elif self.mode == 'rlhf':
|
|
encoded = self._rlhf_encode(inputs)
|
|
elif self.mode == 'kto':
|
|
encoded = self._kto_encode(inputs)
|
|
elif self.task_type == 'seq_cls':
|
|
if self.mode == 'rlhf':
|
|
encoded = self._rlhf_encode(inputs)
|
|
for prefix in ['chosen', 'rejected']: # rm
|
|
encoded.pop(f'{prefix}_labels', None)
|
|
encoded.pop(f'{prefix}_loss_scale', None)
|
|
else:
|
|
encoded = self._seq_cls_encode(chosen)
|
|
elif self.task_type == 'prm':
|
|
encoded = self._encode_truncated(chosen)
|
|
elif self.task_type == 'embedding':
|
|
encoded = self._embedding_encode(inputs)
|
|
elif self.task_type in {'reranker', 'generative_reranker'}:
|
|
encoded = self._reranker_encode(inputs)
|
|
else:
|
|
raise ValueError(f'task_type: {self.task_type} is not supported.')
|
|
|
|
# compatible with `--truncation_strategy split`
|
|
batched = encoded
|
|
if not isinstance(batched, (list, tuple)):
|
|
batched = [batched]
|
|
for encoded in batched:
|
|
if chosen.channel is not None:
|
|
encoded['channel'] = chosen.channel
|
|
|
|
lengths = []
|
|
for key in list(encoded.keys()):
|
|
if encoded[key] is None:
|
|
encoded.pop(key)
|
|
elif key.endswith('length'):
|
|
value = encoded[key]
|
|
if isinstance(value, int):
|
|
lengths.append(value)
|
|
elif isinstance(value, (tuple, list)):
|
|
lengths += value
|
|
if return_length:
|
|
if not lengths:
|
|
raise ValueError(f'lengths should not be empty. batched: {batched}')
|
|
encoded['lengths'] = lengths
|
|
else:
|
|
encoded.pop('length', None)
|
|
if return_template_inputs:
|
|
encoded['template_inputs'] = chosen
|
|
if not self.remove_unused_columns:
|
|
encoded['_extra_kwargs'] = chosen.extra_kwargs
|
|
return batched[0] if len(batched) == 1 else batched
|
|
|
|
def packing_row(self, row: List[Dict[str, Any]]) -> Dict[str, Any]:
|
|
packed = {}
|
|
keys = set()
|
|
length = []
|
|
is_3d_position_ids = False
|
|
for r in row:
|
|
if isinstance(r.get('position_ids'), torch.Tensor) and r['position_ids'].dim() == 3:
|
|
is_3d_position_ids = True
|
|
keys.update(r.keys())
|
|
length.append(r['length'])
|
|
for key in keys:
|
|
if key == 'position_ids' and is_3d_position_ids or key in {'mm_token_type_ids'}:
|
|
packed[key] = torch.cat([x.get(key) for x in row], dim=-1)
|
|
elif key in {'input_ids', 'labels', 'loss_scale', 'position_ids', 'token_type_ids'}:
|
|
packed[key] = sum((x.get(key) or [] for x in row), start=[])
|
|
elif key == 'channel':
|
|
packed[key] = [x.get(key) for x in row]
|
|
if 'position_ids' not in packed:
|
|
packed['position_ids'] = sum((list(range(x)) for x in length), start=[])
|
|
|
|
packed.update(self._data_collator_mm_data(row))
|
|
return packed
|
|
|
|
def _post_encode(self, model: nn.Module, inputs: Dict[str, Any]) -> Dict[str, Any]:
|
|
return inputs
|
|
|
|
@staticmethod
|
|
def _get_seq_cls_logprobs(pred: int, logprobs: torch.Tensor, top_logprobs: int):
|
|
idxs = logprobs.argsort(descending=True, dim=-1)[:top_logprobs].tolist()
|
|
logprobs = logprobs.tolist()
|
|
return {
|
|
'content': [{
|
|
'index': pred,
|
|
'logprobs': [logprobs[p] for p in pred] if isinstance(pred, (list, tuple)) else logprobs[pred],
|
|
'top_logprobs': [{
|
|
'index': idx,
|
|
'logprob': logprobs[idx]
|
|
} for idx in idxs]
|
|
}]
|
|
}
|
|
|
|
def decode_seq_cls(self, logits: torch.Tensor, top_logprobs: int):
|
|
assert isinstance(logits, torch.Tensor)
|
|
problem_type = self.config.problem_type
|
|
if problem_type == 'regression':
|
|
preds = logits.squeeze(dim=-1).tolist()
|
|
logprobs = [None] * len(preds)
|
|
else:
|
|
if problem_type == 'single_label_classification':
|
|
preds = torch.argmax(logits, dim=-1).tolist()
|
|
logprobs = torch.log_softmax(logits, -1)
|
|
else:
|
|
preds = [(logprob >= 0.5).nonzero(as_tuple=True)[0].tolist() for logprob in torch.sigmoid(logits)]
|
|
logprobs = F.logsigmoid(logits)
|
|
logprobs = [self._get_seq_cls_logprobs(pred, logprobs[i], top_logprobs) for i, pred in enumerate(preds)]
|
|
return preds, logprobs
|
|
|
|
def decode_generate_ids(self,
|
|
generate_ids: List[int],
|
|
*,
|
|
is_finished: bool = True,
|
|
first_token=True,
|
|
template_inputs=None,
|
|
**kwargs) -> Any:
|
|
if kwargs.get('spaces_between_special_tokens') is None:
|
|
kwargs['spaces_between_special_tokens'] = False
|
|
generate_ids = self.skip_stop_tokens(generate_ids, is_finished)
|
|
response = self.tokenizer.decode(generate_ids, **kwargs)
|
|
response_prefix = self._get_response_prefix(template_inputs)
|
|
if first_token and response_prefix:
|
|
response = response_prefix + response
|
|
return response
|
|
|
|
def decode_prm(self, input_ids: torch.Tensor, logits: torch.Tensor) -> Any:
|
|
raise NotImplementedError
|
|
|
|
@contextmanager
|
|
def generate_context(self):
|
|
origin_mode = self.mode
|
|
if self.mode in {'train', 'rlhf', 'kto'}:
|
|
self.set_mode('transformers')
|
|
is_multimodal = self.model_meta.is_multimodal
|
|
if is_multimodal:
|
|
models = self.remove_post_encode_hook()
|
|
try:
|
|
yield
|
|
finally:
|
|
if is_multimodal:
|
|
self.register_post_encode_hook(models)
|
|
self.set_mode(origin_mode)
|
|
|
|
def generate(self, model, *args, **kwargs):
|
|
base_model = self.get_base_model(model)
|
|
signature = inspect.signature(base_model.generate)
|
|
if 'use_model_defaults' in signature.parameters and 'use_model_defaults' not in kwargs:
|
|
kwargs['use_model_defaults'] = False
|
|
return model.generate(*args, **kwargs)
|
|
|
|
def compute_sft_loss(self, model, inputs: Dict[str, Any], num_items_in_batch: Optional[int] = None, trainer=None):
|
|
# Default SFT Loss Calculation Method
|
|
outputs = model(**inputs)
|
|
if 'labels' in inputs:
|
|
labels = inputs['labels']
|
|
outputs.loss = outputs.loss.to(labels.device)
|
|
# fix https://github.com/huggingface/transformers/issues/34263
|
|
if num_items_in_batch is not None:
|
|
outputs.loss = outputs.loss * ((labels[:, 1:] != -100).sum() / num_items_in_batch)
|
|
return outputs
|
|
|
|
def skip_stop_tokens(self, generate_ids: List[int], is_finished: bool = True) -> List[int]:
|
|
# Do not print template_meta.suffix_stop and eos_token.
|
|
# However, other stop_words will be printed.
|
|
tokenizer = self.tokenizer
|
|
|
|
if len(generate_ids) > 0 and generate_ids[-1] == tokenizer.eos_token_id:
|
|
generate_ids = generate_ids[:-1]
|
|
# skip suffix and eos_token
|
|
template_suffix = self.template_meta.suffix_stop
|
|
if isinstance(template_suffix, str):
|
|
# [-1:]: fix OpenGVLab/Mini-InternVL-Chat-4B-V1-5
|
|
template_suffix = tokenizer.encode(template_suffix, add_special_tokens=False)[-1:]
|
|
|
|
len_tokens = len(template_suffix)
|
|
if is_finished and generate_ids[-len_tokens:] == template_suffix:
|
|
generate_ids = generate_ids[:-len_tokens]
|
|
elif not is_finished:
|
|
for i in range(len_tokens, 0, -1):
|
|
if generate_ids[-i:] == template_suffix[:i]:
|
|
generate_ids = generate_ids[:-i]
|
|
break
|
|
return generate_ids
|
|
|
|
def prepare_generate_kwargs(self, generate_kwargs: Dict[str, Any], *, model=None) -> Dict[str, Any]:
|
|
generation_config = generate_kwargs['generation_config']
|
|
stop_words = getattr(generation_config, 'stop_words', None) or self.template_meta.stop_words
|
|
generate_kwargs['stopping_criteria'] = StoppingCriteriaList([StopWordsCriteria(self.tokenizer, stop_words)])
|
|
return generate_kwargs
|
|
|
|
@staticmethod
|
|
def _save_pil_image(image: Image.Image) -> str:
|
|
img_bytes = image.tobytes()
|
|
img_hash = hashlib.sha256(img_bytes).hexdigest()
|
|
tmp_dir = os.path.join(get_cache_dir(), 'tmp', 'images')
|
|
logger.info_once(f'create tmp_dir: {tmp_dir}')
|
|
os.makedirs(tmp_dir, exist_ok=True)
|
|
img_path = os.path.join(tmp_dir, f'{img_hash}.png')
|
|
if not os.path.exists(img_path):
|
|
image.save(img_path)
|
|
return img_path
|
|
|
|
@staticmethod
|
|
def _concat_context_list(
|
|
context_list: List[Context],
|
|
res_context_list: List[Context], # inplace
|
|
res_context_type: List[ContextType], # inplace
|
|
system: Optional[str] = None,
|
|
query: Optional[str] = None,
|
|
response: Optional[str] = None,
|
|
round0: Optional[int] = None) -> None:
|
|
"""Concat context list and replace placeholder"""
|
|
round1 = None
|
|
if round0 is not None:
|
|
round1 = str(round0 + 1)
|
|
round0 = str(round0)
|
|
for context in context_list:
|
|
if isinstance(context, str):
|
|
if '{{RESPONSE}}' == context:
|
|
assert response is not None
|
|
res_context_list.append(response)
|
|
res_context_type.append(ContextType.RESPONSE)
|
|
continue
|
|
old_str_list = ['{{SYSTEM}}', '{{QUERY}}', '{{ROUND0}}', '{{ROUND1}}']
|
|
new_str_list = [system, query, round0, round1]
|
|
for (old_str, new_str) in zip(old_str_list, new_str_list):
|
|
if new_str is not None and old_str in context:
|
|
assert isinstance(new_str, str), f'new_str: {new_str}'
|
|
context = context.replace(old_str, new_str)
|
|
if len(context) == 0:
|
|
continue
|
|
res_context_list.append(context)
|
|
res_context_type.append(ContextType.OTHER)
|
|
|
|
def _simplify_context_list(self, context_list: List[Context], loss_scale_list: List[float],
|
|
inputs: StdTemplateInputs) -> Tuple[List[Context], List[float]]:
|
|
"""Merge anything in the context to simplify the inputs"""
|
|
context_list, loss_scale_list = self._split_special_tokens(context_list, loss_scale_list)
|
|
context_list, loss_scale_list = self._pre_tokenize(context_list, loss_scale_list, inputs)
|
|
|
|
res: List[Context] = [] # result of context_list
|
|
res_loss_scale: List[float] = [] # result of loss_scale_list
|
|
temp: List[str] = []
|
|
temp_loss_scale = 0.
|
|
for i, (context, loss_scale) in enumerate(zip(context_list, loss_scale_list)):
|
|
if isinstance(context, str) and (loss_scale == temp_loss_scale):
|
|
temp.append(context)
|
|
else:
|
|
if len(temp) > 0:
|
|
res.append(''.join(temp))
|
|
res_loss_scale.append(temp_loss_scale)
|
|
temp.clear()
|
|
if isinstance(context, str): # loss_scale diff
|
|
temp.append(context)
|
|
else:
|
|
res.append(context)
|
|
res_loss_scale.append(loss_scale)
|
|
temp_loss_scale = loss_scale
|
|
if len(temp) > 0:
|
|
res.append(''.join(temp))
|
|
res_loss_scale.append(temp_loss_scale)
|
|
|
|
return res, res_loss_scale
|
|
|
|
@staticmethod
|
|
def _split_special_tokens(context_list: List[Context],
|
|
loss_scale_list: List[float]) -> Tuple[List[Context], List[float]]:
|
|
"""Split special tokens, for example `<image>`, `<video>`, this will help the replace_tag operation"""
|
|
res: List[Context] = []
|
|
loss_scale_res: List[float] = []
|
|
for context, loss_scale in zip(context_list, loss_scale_list):
|
|
contexts = []
|
|
if isinstance(fetch_one(context), str):
|
|
for d in split_str_parts_by(context, Template.special_tokens):
|
|
contexts.extend([d['key'], d['content']])
|
|
contexts = [c for c in contexts if c]
|
|
res.extend(contexts)
|
|
loss_scale_res.extend([loss_scale] * len(contexts))
|
|
else:
|
|
res.append(context)
|
|
loss_scale_res.append(loss_scale)
|
|
return res, loss_scale_res
|
|
|
|
def _tokenize(self, context, **kwargs):
|
|
return self.tokenizer(context, return_attention_mask=False, add_special_tokens=False, **kwargs)['input_ids']
|
|
|
|
def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int,
|
|
inputs: StdTemplateInputs) -> List[Context]:
|
|
"""Override this function to do your own replace operation.
|
|
|
|
This method is used to replace standard tags like `<image>` to some tokens that the model needs.
|
|
|
|
Args:
|
|
media_type: The modal.
|
|
index: The index of the medias, for index 0 represents the first elements in `images`
|
|
inputs: The inputs
|
|
|
|
Returns:
|
|
The content or input_ids after replacement.
|
|
"""
|
|
if media_type == 'image':
|
|
if self.mode == 'lmdeploy':
|
|
return [[-100]]
|
|
return self.image_placeholder
|
|
elif media_type == 'video':
|
|
if self.mode == 'vllm':
|
|
from ..vision_utils import load_vllm_video
|
|
num_frames = get_env_args('vllm_num_frames', int, 16)
|
|
video_data, video_metadatas = load_vllm_video(inputs.videos[index], num_frames)
|
|
inputs.videos[index] = [(video_data, video_metadatas)]
|
|
return self.video_placeholder
|
|
else:
|
|
return self.video_placeholder
|
|
elif media_type == 'audio':
|
|
return self.audio_placeholder
|
|
|
|
def replace_ref(self, ref: str, index: int, inputs: StdTemplateInputs) -> List[Context]:
|
|
"""Replace objects referenced by the bbox to contents or input_ids. This is useful in the grounding task.
|
|
Override this function to do your own replace operation.
|
|
|
|
Args:
|
|
ref: Description of the bbox
|
|
index: The index in the `objects` key
|
|
inputs: The inputs
|
|
|
|
Returns:
|
|
The contents or input_ids replaced
|
|
"""
|
|
return [ref]
|
|
|
|
def replace_cot_process(self, inputs: StdTemplateInputs) -> List[Context]:
|
|
"""Replace the cot process label for PRM training or inference.
|
|
Override this function to do your own replace operation.
|
|
|
|
Args:
|
|
inputs: The inputs
|
|
|
|
Returns:
|
|
The contents or input_ids replaced
|
|
"""
|
|
return [self.cot_process_placeholder]
|
|
|
|
@staticmethod
|
|
def _get_bbox_str(bbox: List[int]) -> str:
|
|
point = []
|
|
for x, y in zip(bbox[::2], bbox[1::2]):
|
|
point.append(f'({x},{y})')
|
|
return ','.join(point)
|
|
|
|
def replace_bbox(self, bbox: List[int], index: int, inputs: StdTemplateInputs) -> List[Context]:
|
|
"""Replace bbox pointing to the objects to contents or input_ids. This is useful in the grounding task.
|
|
Override this function to do your own replace operation.
|
|
|
|
Args:
|
|
bbox: [x, y] or [x1, y1, x2, y2]
|
|
index: The index in the `objects` key
|
|
inputs: The inputs
|
|
|
|
Returns:
|
|
The contents or input_ids replaced
|
|
"""
|
|
return [f'[{self._get_bbox_str(bbox)}]']
|
|
|
|
def _pre_tokenize_images(self, context_list: List[Context], loss_scale_list: List[float],
|
|
inputs: StdTemplateInputs) -> Tuple[List[Context], List[float]]:
|
|
# https://github.com/modelscope/ms-swift/issues/3407
|
|
# Fix the bounding box position offset issue in the Qwen2.5-VL grounding task.
|
|
res: List[Context] = []
|
|
res_loss_scale: List[float] = []
|
|
inputs.image_idx = 0
|
|
|
|
for context, loss_scale in zip(context_list, loss_scale_list):
|
|
if context == '<image>' and inputs.is_multimodal and inputs.image_idx < len(inputs.images):
|
|
c_list = self.replace_tag('image', inputs.image_idx, inputs)
|
|
inputs.image_idx += 1
|
|
loss_scale = 0. if self.template_backend == 'swift' else 1.
|
|
else:
|
|
c_list = [context]
|
|
res += c_list
|
|
res_loss_scale += [loss_scale] * len(c_list)
|
|
return res, res_loss_scale
|
|
|
|
def _pre_tokenize(self, context_list: List[Context], loss_scale_list: List[float],
|
|
inputs: StdTemplateInputs) -> Tuple[List[Context], List[float]]:
|
|
"""This method happens before tokenization, replace standard tags to the contents or input_ids needed by
|
|
the model.
|
|
|
|
Args:
|
|
context_list: The content list
|
|
loss_scale_list: The loss scale list
|
|
Returns:
|
|
The context_list and loss_scale_list after replacement.
|
|
"""
|
|
context_list, loss_scale_list = self._pre_tokenize_images(context_list, loss_scale_list, inputs)
|
|
if inputs.images and inputs.objects:
|
|
self.normalize_bbox(inputs)
|
|
# replace tag/object/box
|
|
res: List[Context] = [] # result of context_list
|
|
res_loss_scale: List[float] = [] # result of loss_scale_list
|
|
|
|
# reset
|
|
for k in ['video', 'audio', 'object', 'box']:
|
|
setattr(inputs, f'{k}_idx', 0)
|
|
|
|
for context, loss_scale in zip(context_list, loss_scale_list):
|
|
for k in ['video', 'audio']:
|
|
if context == f'<{k}>' and inputs.is_multimodal and getattr(inputs, f'{k}_idx') < len(
|
|
getattr(inputs, f'{k}s')):
|
|
c_list = self.replace_tag(k, getattr(inputs, f'{k}_idx'), inputs)
|
|
setattr(inputs, f'{k}_idx', getattr(inputs, f'{k}_idx') + 1)
|
|
loss_scale = 0.
|
|
break
|
|
else:
|
|
ref = inputs.objects.get('ref') or []
|
|
bbox = inputs.objects.get('bbox') or []
|
|
if context == '<ref-object>' and inputs.ref_idx < len(ref):
|
|
idx = inputs.ref_idx
|
|
c_list = self.replace_ref(ref[idx], idx, inputs)
|
|
inputs.ref_idx += 1
|
|
elif context == '<bbox>' and inputs.bbox_idx < len(bbox):
|
|
idx = inputs.bbox_idx
|
|
c_list = self.replace_bbox(bbox[idx], idx, inputs)
|
|
inputs.bbox_idx += 1
|
|
elif context == '<cot-process>' and self.task_type == 'prm':
|
|
c_list = self.replace_cot_process(inputs)
|
|
else:
|
|
c_list = [context]
|
|
res += c_list
|
|
res_loss_scale += [loss_scale] * len(c_list)
|
|
return res, res_loss_scale
|
|
|
|
@staticmethod
|
|
def _add_default_tags(inputs: StdTemplateInputs):
|
|
total_content = []
|
|
for message in inputs.messages:
|
|
content = message['content'] or ''
|
|
if not isinstance(content, str):
|
|
if message['role'] == 'user':
|
|
# Give up adding the default tag
|
|
return
|
|
elif message['role'] == 'assistant':
|
|
continue
|
|
total_content.append(content)
|
|
total_content = '\n'.join(total_content)
|
|
if inputs.system:
|
|
total_content = f'{inputs.system}\n{total_content}'
|
|
for media_type in ['image', 'audio', 'video']:
|
|
media_key, media_tag = f'{media_type}s', f'<{media_type}>'
|
|
medias = getattr(inputs, media_key)
|
|
if not isinstance(medias, list):
|
|
medias = [medias]
|
|
if medias:
|
|
num_media_tags = len(re.findall(media_tag, total_content))
|
|
num_media = len(medias)
|
|
num_new_tags = num_media - num_media_tags
|
|
if num_new_tags > 0:
|
|
inputs.messages[0]['content'] = media_tag * num_new_tags + inputs.messages[0]['content']
|
|
elif num_new_tags < 0:
|
|
logger.warning(
|
|
f'num_media: {num_media}, num_media_tags: {num_media_tags}, total_content: {total_content}. '
|
|
'We will only replace the frontmost media_tags while keeping the subsequent media_tags.')
|
|
|
|
def _encode_context_list(self,
|
|
context_list: List[Context],
|
|
loss_scale_list: Optional[List[float]] = None) -> Tuple[List[int], List[int], List[float]]:
|
|
is_binary_loss_scale = self.is_binary_loss_scale
|
|
if is_binary_loss_scale is None:
|
|
is_binary_loss_scale = self.loss_scale.is_binary_loss_scale
|
|
input_ids: List[int] = []
|
|
labels: List[int] = []
|
|
loss_scale: List[float] = []
|
|
if loss_scale_list is None:
|
|
loss_scale_list = [0.] * len(context_list)
|
|
for i, (context, loss_weight) in enumerate(zip(context_list, loss_scale_list)):
|
|
if isinstance(context, str):
|
|
token_list = self._tokenize(context)
|
|
else:
|
|
token_list = context
|
|
input_ids += token_list
|
|
if loss_scale_list[i] > 0.0:
|
|
labels += token_list
|
|
else:
|
|
labels += [-100] * len(token_list)
|
|
if not is_binary_loss_scale:
|
|
loss_scale.extend([loss_weight] * len(token_list))
|
|
if is_binary_loss_scale:
|
|
loss_scale = None
|
|
return input_ids, labels, loss_scale
|
|
|
|
@staticmethod
|
|
def _add_dynamic_eos(input_ids: List[int], labels: List[int], loss_scale: Optional[List[int]],
|
|
suffix_tokens_id: List[int]) -> None:
|
|
suffix_len = len(suffix_tokens_id)
|
|
start = 0
|
|
for i in range(1, len(labels) + 1):
|
|
if labels[i - 1] >= 0 and i < len(labels) and labels[i] == -100:
|
|
start = i
|
|
elif start > 0 and labels[i - 1] == -100 and (i == len(labels) or labels[i] >= 0):
|
|
# [0, 1, 2, -100(start), -100, 3(i), 4]
|
|
length = i - start
|
|
if length >= suffix_len and input_ids[start:start + suffix_len] == suffix_tokens_id:
|
|
labels[start:start + suffix_len] = suffix_tokens_id
|
|
if loss_scale and loss_scale[start:start + suffix_len] == [0] * suffix_len:
|
|
loss_scale[start:start + suffix_len] = [1] * suffix_len
|
|
|
|
@staticmethod
|
|
def _get_std_messages(messages):
|
|
if messages and messages[0]['role'] == 'assistant':
|
|
messages.insert(0, {'role': 'user', 'content': ''}) # pretrain
|
|
if len(messages) % 2 == 1:
|
|
messages.append({'role': 'assistant', 'content': None}) # inference
|
|
|
|
def _jinja_encode(self, inputs: StdTemplateInputs):
|
|
messages = inputs.messages.copy()
|
|
if inputs.system is None:
|
|
# Fix default_system passed from command line being ignored.
|
|
inputs.system = self.template_meta.default_system
|
|
if inputs.system is not None:
|
|
messages.insert(0, {'role': 'system', 'content': inputs.system})
|
|
if messages[-1]['content'] is None:
|
|
messages.pop()
|
|
add_generation_prompt = messages[-1]['role'] != 'assistant'
|
|
kwargs = {}
|
|
if inputs.tools:
|
|
kwargs['tools'] = inputs.tools
|
|
enable_thinking = self._get_enable_thinking(inputs)
|
|
if self.template_meta.is_thinking or enable_thinking:
|
|
kwargs[self.jinja_enable_thinking_key] = enable_thinking
|
|
kwargs['preserve_thinking'] = self._get_preserve_thinking(inputs)
|
|
kwargs.update(self.chat_template_kwargs)
|
|
kwargs.update(inputs.chat_template_kwargs)
|
|
text = self.tokenizer.apply_chat_template(
|
|
messages, tokenize=False, add_generation_prompt=add_generation_prompt, **kwargs)
|
|
answer_len = 1 if self.is_training else 0
|
|
return [text], [1.], answer_len
|
|
|
|
def _get_system(self, inputs: StdTemplateInputs) -> Optional[str]:
|
|
template_meta = self.template_meta
|
|
system = inputs.system
|
|
tools = inputs.tools
|
|
template_meta.check_system(system)
|
|
if system is None:
|
|
system = template_meta.default_system
|
|
|
|
if tools:
|
|
system = self.agent_template._format_tools(tools, system, inputs.messages[0])
|
|
return system
|
|
|
|
def _is_add_non_thinking_round(self, messages, i: int, start_idx: int):
|
|
message = messages[i]
|
|
if not (i >= start_idx and message['role'] == 'assistant'):
|
|
return False
|
|
if self.non_thinking_prefix_only_after_user and not (i > 0 and messages[i - 1]['role'] == 'user'):
|
|
return False
|
|
return True
|
|
|
|
def _add_non_thinking_prefix(self, inputs, thinking_prefix='<think>') -> None:
|
|
messages = inputs.messages
|
|
non_thinking_prefix = self.template_meta.non_thinking_prefix
|
|
if non_thinking_prefix:
|
|
# Determine the starting index for processing messages
|
|
# During inference or when using 'last_round' strategy, only process the last round
|
|
# Otherwise, process all messages (start_idx = -1 means start from the beginning)
|
|
if not self.is_training or self.loss_scale.base_strategy == 'last_round':
|
|
start_idx = get_last_user_round(messages)
|
|
else:
|
|
start_idx = -1
|
|
for i, message in enumerate(messages):
|
|
if not self._is_add_non_thinking_round(messages, i, start_idx):
|
|
continue
|
|
content = message['content']
|
|
# After merge, content may be a list; only process the first element.
|
|
if isinstance(content, list):
|
|
_add_prefix = content and isinstance(content[0], str) and not content[0].startswith(
|
|
(thinking_prefix, non_thinking_prefix))
|
|
if _add_prefix:
|
|
content[0] = non_thinking_prefix + content[0]
|
|
elif isinstance(content, str):
|
|
if not content.startswith((thinking_prefix, non_thinking_prefix)):
|
|
message['content'] = non_thinking_prefix + content
|
|
|
|
def _remove_thinking_content(self, content: str, thinking_suffix='</think>') -> str:
|
|
content = content.split(thinking_suffix)[-1].strip()
|
|
return self.template_meta.history_thinking_prefix + content
|
|
|
|
def _remove_history_thinking(self, inputs) -> None:
|
|
messages = inputs.messages
|
|
# Delete the previous 'think' entries from the messages.
|
|
last_user_round = get_last_user_round(messages)
|
|
for i, message in enumerate(messages):
|
|
# Delete the content before '</think>' in all assistant turns except the last round.
|
|
if message['role'] == 'assistant' and i < last_user_round:
|
|
content = message['content']
|
|
# After merge, content may be a list; only process the first element.
|
|
if isinstance(content, list) and content and isinstance(content[0], str):
|
|
content[0] = self._remove_thinking_content(content[0])
|
|
elif isinstance(content, str):
|
|
message['content'] = self._remove_thinking_content(content)
|
|
|
|
def _swift_prepare_inputs(self, inputs: StdTemplateInputs):
|
|
"""
|
|
Preprocesses the list of messages in the input by merging and formatting consecutive messages
|
|
according to their roles.
|
|
|
|
Specifically, this method:
|
|
- Merges consecutive messages from the same role ('assistant' or 'user') to prevent downstream errors.
|
|
- Detects consecutive tool-related messages following an assistant message, then formats and
|
|
combines them using `agent_template._format_tool_responses` for structured output.
|
|
- Updates the messages list in-place for further processing.
|
|
|
|
Args:
|
|
inputs: An StdTemplateInputs object which contains a 'messages' attribute, which is a list of dictionaries.
|
|
Each message dictionary should have at least the keys 'role' and 'content'.
|
|
|
|
Returns:
|
|
None. The input messages list is updated in-place.
|
|
"""
|
|
self._preprocess_tool_call(inputs)
|
|
messages = inputs.messages
|
|
if len(messages) < 2:
|
|
return
|
|
i = 1
|
|
while i < len(messages):
|
|
pre_message, message = messages[i - 1], messages[i]
|
|
pre_role, pre_content = pre_message['role'], pre_message['content']
|
|
role, content = message['role'], message['content']
|
|
if pre_role == 'assistant' and role == 'tool' and self.template_backend == 'swift':
|
|
i_start = i
|
|
while i + 1 < len(messages) and messages[i + 1]['role'] == 'tool':
|
|
i += 1
|
|
pre_message['content'], tool_content = self.agent_template._format_tool_responses(
|
|
pre_content, messages[i_start:i + 1])
|
|
# where tool_content is a List.
|
|
messages[i_start:i + 1] = [{'role': 'tool', 'content': tool_content}]
|
|
i = i_start + 1
|
|
elif pre_role == 'assistant' and role == 'assistant' or pre_role == 'user' and role == 'user':
|
|
# Consecutive messages from the assistant/user role need to be merged to prevent errors.
|
|
if self.template_backend == 'swift' and pre_role == 'assistant':
|
|
for key in ['content', 'loss', 'loss_scale']:
|
|
pre_val = pre_message.get(key)
|
|
cur_val = message.get(key)
|
|
pre_message[key] = (pre_val if isinstance(pre_val, list) else [pre_val]) + \
|
|
(cur_val if isinstance(cur_val, list) else [cur_val])
|
|
else:
|
|
pre_message['content'] = pre_content + content
|
|
messages.pop(i)
|
|
else:
|
|
i += 1
|
|
|
|
def _swift_encode(self, inputs: StdTemplateInputs):
|
|
template_meta = self.template_meta
|
|
if self.use_chat_template:
|
|
if self.add_non_thinking_prefix:
|
|
self._add_non_thinking_prefix(inputs)
|
|
preserve_thinking = self._get_preserve_thinking(inputs)
|
|
if not preserve_thinking:
|
|
self._remove_history_thinking(inputs)
|
|
system = self._get_system(inputs)
|
|
else:
|
|
system = None
|
|
self._get_std_messages(inputs.messages)
|
|
n_round = len(inputs.messages) // 2
|
|
if n_round > 1 and not self.template_meta.support_multi_round:
|
|
logger.warning_once(
|
|
'The template does not support multi-round chat. Only use the last round of the conversation.')
|
|
# TODO: Multimodal models may encounter image mismatch issues.
|
|
inputs.messages = inputs.messages[-2:]
|
|
|
|
res_context_list: List[Context] = []
|
|
res_context_types: List[ContextType] = []
|
|
sep_token = None
|
|
if template_meta.auto_add_bos:
|
|
all_tokens = self.tokenizer.encode('a')
|
|
single_token = self.tokenizer.encode('a', add_special_tokens=False)
|
|
assert len(single_token) == 1
|
|
idx = all_tokens.index(single_token[0])
|
|
bos_token = all_tokens[:idx]
|
|
sep_token = all_tokens[idx + 1:]
|
|
if bos_token:
|
|
res_context_list.append(bos_token)
|
|
res_context_types.append(ContextType.OTHER)
|
|
|
|
if self.template_meta.is_post_system or not system:
|
|
prefix = template_meta.prefix
|
|
else:
|
|
prefix = template_meta.system_prefix
|
|
self._concat_context_list(prefix, res_context_list, res_context_types, system=system)
|
|
|
|
assert len(inputs.messages) > 0, f'inputs.messages: {inputs.messages}'
|
|
n_round = len(inputs.messages) // 2
|
|
for i, (query_message, response_message) in enumerate(zip(inputs.messages[::2], inputs.messages[1::2])):
|
|
query_role, query = query_message['role'], query_message['content']
|
|
response_role, response = response_message['role'], response_message['content']
|
|
# TODO: Optimize the Template mechanism.
|
|
assert query_role in {'user', 'tool'}, f'query_role: "{query_role}"'
|
|
assert response_role in {'assistant'}, f'response_role: "{response_role}"'
|
|
if query_role == 'tool':
|
|
prompt = query
|
|
query = ''
|
|
elif template_meta.is_post_system and i == n_round - 1:
|
|
prompt = template_meta.system_prompt
|
|
else:
|
|
prompt = template_meta.prompt
|
|
|
|
context_list = prompt.copy()
|
|
extra_context_list = []
|
|
extra_context_type = None
|
|
response_prefix = self._get_response_prefix(inputs)
|
|
if i < n_round - 1:
|
|
# Not the last round.
|
|
context_list.append('{{RESPONSE}}')
|
|
if inputs.messages[2 * (i + 1)]['role'] != 'tool':
|
|
extra_context_list = template_meta.chat_sep
|
|
extra_context_type = ContextType.OTHER
|
|
elif response is not None:
|
|
# It is the final round, and the response exists (during training).
|
|
context_list.append('{{RESPONSE}}')
|
|
# The GLM-4.5 assistant part (tool call) may end with <|observation|>,
|
|
# and here we avoid adding <|user|>.
|
|
response_content = response
|
|
if not isinstance(response_content, str):
|
|
if isinstance(response_content, list) and response_content and isinstance(
|
|
response_content[-1], str):
|
|
response_content = response_content[-1]
|
|
else:
|
|
token_ids = response_content if isinstance(response_content,
|
|
list) else response_content['token_ids']
|
|
response_content = self.tokenizer.decode(token_ids[-20:])
|
|
endswith_stop_words = any(
|
|
response_content.endswith(stop_word) for stop_word in template_meta.stop_words
|
|
if isinstance(stop_word, str))
|
|
# self.is_training needed because we may want to continue generation from
|
|
# the current response
|
|
add_eos = inputs.extra_kwargs.get('add_eos')
|
|
if add_eos is None:
|
|
add_eos = (self.is_training
|
|
or self.task_type != 'causal_lm') and not sep_token and not endswith_stop_words
|
|
if add_eos:
|
|
extra_context_list = template_meta.suffix
|
|
extra_context_type = ContextType.SUFFIX
|
|
elif response_prefix:
|
|
# final round and during inference.
|
|
context_list.append(response_prefix)
|
|
|
|
self._concat_context_list(
|
|
context_list,
|
|
res_context_list,
|
|
res_context_types,
|
|
query=query,
|
|
response=response,
|
|
system=system,
|
|
round0=i)
|
|
res_context_list += extra_context_list
|
|
res_context_types += [extra_context_type] * len(extra_context_list)
|
|
if template_meta.auto_add_bos and sep_token:
|
|
res_context_list.append(sep_token)
|
|
res_context_types.append(ContextType.SUFFIX)
|
|
res_context_list, loss_scale_list = self.loss_scale(res_context_list, res_context_types, inputs.messages,
|
|
**inputs.extra_kwargs)
|
|
if self.is_training:
|
|
answer_len = len(extra_context_list) + bool(response is not None)
|
|
else:
|
|
answer_len = 0
|
|
return res_context_list, loss_scale_list, answer_len
|
|
|
|
def _truncate(self, input_ids: List[int], labels: Optional[List[int]], encoded,
|
|
truncation_strategy: Literal['left', 'right']):
|
|
placeholder_tokens = torch.tensor(self.placeholder_tokens)
|
|
input_ids_tensor = torch.tensor(input_ids)
|
|
protected = (input_ids_tensor[:, None] == placeholder_tokens).any(dim=-1)
|
|
n_protected = protected.sum().item()
|
|
if n_protected < self.max_length:
|
|
non_protected = (~protected).nonzero(as_tuple=True)[0]
|
|
if truncation_strategy == 'left':
|
|
idx = non_protected[-(self.max_length - n_protected):]
|
|
else:
|
|
idx = non_protected[:self.max_length - n_protected]
|
|
protected[idx] = True
|
|
input_ids = input_ids_tensor[protected].tolist()
|
|
if labels is not None:
|
|
labels = torch.tensor(labels)[protected].tolist()
|
|
labels[0] = -100
|
|
loss_scale = encoded.get('loss_scale')
|
|
if loss_scale is not None:
|
|
loss_scale = torch.tensor(loss_scale)[protected].tolist()
|
|
loss_scale[0] = 0
|
|
encoded['loss_scale'] = loss_scale
|
|
mm_token_type_ids = encoded.get('mm_token_type_ids')
|
|
if mm_token_type_ids is not None:
|
|
encoded['mm_token_type_ids'] = mm_token_type_ids[protected]
|
|
return input_ids, labels
|
|
|
|
@staticmethod
|
|
def _get_length(input_ids, labels):
|
|
# input_ids might be a tensor.
|
|
lengths = [0]
|
|
if input_ids is not None:
|
|
lengths.append(len(input_ids))
|
|
if labels is not None:
|
|
lengths.append(len(labels))
|
|
length = max(lengths)
|
|
return length
|
|
|
|
def _encode_truncated(self, inputs: StdTemplateInputs):
|
|
self._preprocess_inputs(inputs)
|
|
if self.mode in {'vllm', 'lmdeploy', 'sglang'}:
|
|
# For multi-modal models, images do not need to be pre processed here
|
|
# vllm/lmdeploy/sglang will handle the logic
|
|
encoded = Template._encode(self, inputs)
|
|
keys = ['images', 'audios', 'videos']
|
|
if self.mode == 'vllm':
|
|
keys.append('mm_processor_kwargs')
|
|
for key in keys:
|
|
value = getattr(inputs, key)
|
|
if value:
|
|
encoded[key] = value
|
|
else:
|
|
encoded = self._encode(inputs)
|
|
input_ids = encoded.get('input_ids')
|
|
labels = encoded.get('labels')
|
|
length = self._get_length(input_ids, labels)
|
|
if self.max_length is not None and length > self.max_length:
|
|
if self.truncation_strategy in {'right', 'left'}:
|
|
input_ids, labels = self._truncate(
|
|
input_ids, labels, encoded, truncation_strategy=self.truncation_strategy)
|
|
length = self._get_length(input_ids, labels)
|
|
elif self.truncation_strategy == 'raise':
|
|
raise MaxLengthError(f'Current length of row({length}) is larger'
|
|
f' than the max_length({self.max_length}).')
|
|
elif self.truncation_strategy == 'split':
|
|
i = 0
|
|
batched = []
|
|
while i < length:
|
|
splited = {}
|
|
for key in ['input_ids', 'labels', 'loss_scale']:
|
|
value = encoded.get(key)
|
|
if value is not None:
|
|
value = value[i:i + self.max_length]
|
|
if key == 'labels' and len(value) > 0:
|
|
value[0] = -100
|
|
elif key == 'loss_scale' and len(value) > 0:
|
|
value[0] = 0
|
|
splited[key] = value
|
|
splited['length'] = self._get_length(splited.get('input_ids'), splited.get('labels'))
|
|
batched.append(splited)
|
|
i += self.max_length
|
|
return batched
|
|
else:
|
|
raise ValueError(f'Invalid truncation_strategy: {self.truncation_strategy}')
|
|
encoded['length'] = length
|
|
encoded['input_ids'] = input_ids
|
|
if self.task_type in {'seq_cls', 'embedding', 'reranker', 'generative_reranker'}:
|
|
encoded.pop('labels', None)
|
|
encoded.pop('loss_scale', None)
|
|
else:
|
|
encoded['labels'] = labels
|
|
return encoded
|
|
|
|
def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]:
|
|
inputs.messages = deepcopy(inputs.messages)
|
|
template_backend = self.template_backend
|
|
if (self.template_meta.template_type == 'dummy' and self.use_chat_template and not self.is_training
|
|
and self.task_type == 'causal_lm'):
|
|
template_backend = 'jinja'
|
|
logger.info_once(f'Setting template_backend: {template_backend}')
|
|
self._swift_prepare_inputs(inputs)
|
|
res_context_list, loss_scale_list, answer_len = (
|
|
self._swift_encode(inputs) if template_backend == 'swift' else self._jinja_encode(inputs))
|
|
encoded = {}
|
|
if self.is_encoder_decoder:
|
|
total_len = len(res_context_list)
|
|
for key, _slice in zip(['prompt', 'answer'],
|
|
[slice(0, total_len - answer_len),
|
|
slice(total_len - answer_len, total_len)]):
|
|
context_list, loss_scale = self._simplify_context_list(res_context_list[_slice],
|
|
loss_scale_list[_slice], inputs)
|
|
input_ids, labels, loss_scale = self._encode_context_list(context_list, loss_scale)
|
|
encoded[f'{key}_input_ids'] = input_ids
|
|
encoded[f'{key}_labels'] = labels
|
|
encoded[f'{key}_loss_scale'] = loss_scale
|
|
input_ids = encoded['prompt_input_ids'] + encoded['answer_input_ids']
|
|
labels = encoded['prompt_labels'] + encoded['answer_labels']
|
|
loss_scale = None
|
|
if isinstance(encoded['prompt_loss_scale'], list):
|
|
loss_scale = encoded['prompt_loss_scale'] + encoded['answer_loss_scale']
|
|
else:
|
|
res_context_list, loss_scale_list = self._simplify_context_list(res_context_list, loss_scale_list, inputs)
|
|
input_ids, labels, loss_scale = self._encode_context_list(res_context_list, loss_scale_list)
|
|
self._add_dynamic_eos(input_ids, labels, loss_scale, self._encode_context_list(self.template_meta.suffix)[0])
|
|
|
|
encoded['input_ids'] = input_ids
|
|
encoded['labels'] = labels
|
|
encoded['loss_scale'] = loss_scale
|
|
if encoded.get('labels') is not None:
|
|
encoded['labels'][0] = -100
|
|
if encoded.get('loss_scale') is not None:
|
|
encoded['loss_scale'][0] = 0
|
|
if not self.is_training:
|
|
for k in list(encoded.keys()):
|
|
if k.endswith('labels') or k.endswith('loss_scale'):
|
|
encoded[k] = None
|
|
return encoded
|
|
|
|
def _get_megatron_cp_length(self, length) -> int:
|
|
cp_size = self.sequence_parallel_size
|
|
if not self.use_megatron or cp_size == 1:
|
|
return length
|
|
return math.ceil(length / (cp_size * 2)) * (cp_size * 2)
|
|
|
|
def _handle_megatron_cp(self, batch: List[Dict[str, Any]]) -> None:
|
|
cp_size = self.sequence_parallel_size
|
|
if not self.use_megatron or cp_size == 1:
|
|
return
|
|
for encoded in batch:
|
|
input_ids = encoded['input_ids']
|
|
padding_len = math.ceil(len(input_ids) / (cp_size * 2)) * (cp_size * 2) - len(input_ids)
|
|
input_ids += [self.tokenizer.pad_token_id] * padding_len
|
|
if encoded.get('labels') is not None:
|
|
encoded['labels'] += [-100] * padding_len
|
|
if encoded.get('loss_scale') is not None:
|
|
encoded['loss_scale'] += [0] * padding_len
|
|
if encoded.get('length') is not None:
|
|
encoded['length'] += padding_len
|
|
if encoded.get('mm_token_type_ids') is not None:
|
|
encoded['mm_token_type_ids'] = F.pad(encoded['mm_token_type_ids'], (0, padding_len), value=0)
|
|
|
|
def debug_logger(self, inputs):
|
|
if not strtobool(os.getenv('SWIFT_DEBUG', 'false')):
|
|
return
|
|
if 'input_ids' in inputs:
|
|
k = 'input_ids'
|
|
val = inputs['input_ids']
|
|
elif 'generate_ids' in inputs:
|
|
k = 'generate_ids'
|
|
val = inputs['generate_ids']
|
|
for v in val:
|
|
self.print_inputs({k: v.tolist()})
|
|
|
|
@staticmethod
|
|
def _split_list(inputs: List[int], x: int) -> List[List[int]]:
|
|
idxs = findall(inputs, x)
|
|
idxs.append(len(inputs))
|
|
res = []
|
|
lo = 0
|
|
for idx in idxs:
|
|
res.append(inputs[lo:idx])
|
|
lo = idx + 1
|
|
return res
|
|
|
|
def replace_video2image(self, load_video_func, inputs, replace_tag: Callable) -> List[Context]:
|
|
context_list = []
|
|
if self.mode in {'vllm', 'lmdeploy'}:
|
|
video = inputs.videos.pop(inputs.video_idx)
|
|
inputs.video_idx -= 1
|
|
else:
|
|
video = inputs.videos[inputs.video_idx]
|
|
images = inputs.images
|
|
new_images = load_video_func(video)
|
|
inputs.images = images[:inputs.image_idx] + new_images + images[inputs.image_idx:]
|
|
for i in range(len(new_images)):
|
|
context_list += replace_tag(i)
|
|
inputs.image_idx += len(new_images)
|
|
return context_list
|
|
|
|
def get_generate_ids(self, generate_ids: Union[torch.Tensor, List[int]],
|
|
num_prompt_tokens: int) -> Union[torch.Tensor, List[int]]:
|
|
if self.skip_prompt:
|
|
generate_ids = generate_ids[..., num_prompt_tokens:]
|
|
return generate_ids
|
|
|
|
def post_process_generate_response(self, response: str, inputs: StdTemplateInputs) -> str:
|
|
return response
|
|
|
|
def pre_forward_hook(self, model: nn.Module, args, kwargs):
|
|
old_kwargs = to_device(kwargs, model.device)
|
|
kwargs = to_device(self._post_encode(model, old_kwargs), model.device)
|
|
for k, v in old_kwargs.items():
|
|
if k in {
|
|
'input_ids', 'attention_mask', 'labels', 'position_ids', 'output_hidden_states', 'logits_to_keep',
|
|
'max_length_q', 'max_length_k', 'cu_seq_lens_q', 'cu_seq_lens_k', 'mm_token_type_ids'
|
|
} and k not in kwargs:
|
|
kwargs[k] = v
|
|
if 'inputs_embeds' in kwargs:
|
|
kwargs.pop('input_ids', None)
|
|
|
|
base_model = self.get_base_model(model)
|
|
parameters = inspect.signature(base_model.forward).parameters
|
|
if 'position_ids' not in parameters:
|
|
kwargs.pop('position_ids', None)
|
|
return args, kwargs
|
|
|
|
@property
|
|
def is_training(self):
|
|
return self.mode not in {'transformers', 'vllm', 'lmdeploy', 'sglang'}
|
|
|
|
def set_mode(self, mode: Literal['transformers', 'vllm', 'lmdeploy', 'sglang', 'train', 'rlhf', 'kto']) -> None:
|
|
if mode == 'pt':
|
|
mode = 'transformers'
|
|
logger.warning("The mode 'pt' is deprecated, please use 'transformers'.")
|
|
self.mode = mode
|
|
|
|
def register_post_encode_hook(self, models: List[nn.Module]) -> None:
|
|
"""This function is important for multi-modal training, as it registers the post_encode method
|
|
as a forward hook, converting input_ids into inputs_embeds.
|
|
"""
|
|
if self._handles:
|
|
return
|
|
|
|
for model in models:
|
|
# please use torch>=2.0
|
|
handle = model.register_forward_pre_hook(self.pre_forward_hook, with_kwargs=True)
|
|
self._handles.append((model, handle))
|
|
|
|
if is_deepspeed_zero3_enabled():
|
|
import deepspeed
|
|
self._deepspeed_initialize = deepspeed.initialize
|
|
|
|
@wraps(self._deepspeed_initialize)
|
|
def _initialize(*args, **kwargs):
|
|
res = self._deepspeed_initialize(*args, **kwargs)
|
|
for model, handle in self._handles:
|
|
model._forward_pre_hooks.move_to_end(handle.id)
|
|
return res
|
|
|
|
deepspeed.initialize = _initialize
|
|
|
|
def remove_post_encode_hook(self):
|
|
models = []
|
|
for model, handle in self._handles:
|
|
models.append(model)
|
|
handle.remove()
|
|
self._handles = []
|
|
|
|
if self._deepspeed_initialize is not None:
|
|
import deepspeed
|
|
deepspeed.initialize = self._deepspeed_initialize
|
|
self._deepspeed_initialize = None
|
|
return models
|
|
|
|
def data_collator(self, batch: List[Dict[str, Any]], *, padding_to: Optional[int] = None) -> Dict[str, Any]:
|
|
from swift.dataset import RowPreprocessor
|
|
if self.packing and isinstance(batch[0], list):
|
|
batch = sum(batch, start=[])
|
|
if self.task_type == 'causal_lm':
|
|
if self.mode in {'transformers', 'train'}:
|
|
res = self._data_collator(batch, padding_to=padding_to)
|
|
elif self.mode == 'rlhf':
|
|
res = self._rlhf_data_collator(batch, padding_to=padding_to)
|
|
elif self.mode == 'kto':
|
|
res = self._kto_data_collator(batch, padding_to=padding_to)
|
|
elif self.task_type == 'prm':
|
|
res = self._data_collator(batch, padding_to=padding_to)
|
|
elif self.task_type == 'seq_cls':
|
|
if self.mode == 'rlhf':
|
|
res = self._rlhf_data_collator(batch, padding_to=padding_to)
|
|
else:
|
|
res = self._seq_cls_data_collator(batch, padding_to=padding_to)
|
|
elif self.task_type == 'embedding':
|
|
res = self._embedding_data_collator(batch, padding_to=padding_to)
|
|
elif self.task_type in {'reranker', 'generative_reranker'}:
|
|
res = self._reranker_data_collator(batch, padding_to=padding_to)
|
|
else:
|
|
raise ValueError(f'task_type: {self.task_type} is not supported.')
|
|
if not self.remove_unused_columns:
|
|
extra_kwargs = [b['_extra_kwargs'] for b in batch if b.get('_extra_kwargs') is not None]
|
|
extra_kwargs = RowPreprocessor.rows_to_batched(extra_kwargs)
|
|
res.update({k: v for k, v in extra_kwargs.items() if k not in res})
|
|
return res
|
|
|
|
@staticmethod
|
|
def _fetch_inputs_startswith(batch: List[Dict[str, Any]], prefix: str) -> List[Dict[str, Any]]:
|
|
new_batch = []
|
|
for inputs in batch:
|
|
new_inputs = {}
|
|
for k, v in inputs.items():
|
|
if k.startswith(prefix):
|
|
new_inputs[k[len(prefix):]] = v
|
|
new_batch.append(new_inputs)
|
|
return new_batch
|
|
|
|
@staticmethod
|
|
def fetch_inputs(batch: List[Dict[str, Any]], keys: Optional[List[str]] = None) -> Dict[str, Any]:
|
|
from swift.dataset import RowPreprocessor
|
|
keys = keys or []
|
|
rows = RowPreprocessor.rows_to_batched(batch)
|
|
return {k: rows[k] for k in keys if rows.get(k) is not None}
|
|
|
|
@staticmethod
|
|
def gather_list(batch: List[Dict[str, Any]], attr_name: str) -> Optional[List[Any]]:
|
|
# List[Tensor] -> List[Tensor]
|
|
res = []
|
|
for b in batch:
|
|
if b.get(attr_name) is not None:
|
|
res += b.pop(attr_name)
|
|
return res
|
|
|
|
@staticmethod
|
|
def concat_tensor(batch: List[Dict[str, Any]], attr_name: str, dim: int) -> Optional[torch.Tensor]:
|
|
res = []
|
|
for b in batch:
|
|
if b.get(attr_name) is not None:
|
|
res.append(b.pop(attr_name))
|
|
return torch.concat(res, dim=dim) if res else None
|
|
|
|
def _rlhf_data_collator(self,
|
|
batch: List[Dict[str, Any]],
|
|
*,
|
|
chosen_prefix: str = 'chosen_',
|
|
rejected_prefix: str = 'rejected_',
|
|
padding_to: Optional[int] = None) -> Dict[str, Any]:
|
|
new_batch = []
|
|
for prefix in [chosen_prefix, rejected_prefix]:
|
|
new_batch += self._fetch_inputs_startswith(batch, prefix)
|
|
res = self._data_collator(new_batch, padding_to=padding_to)
|
|
|
|
# reward modeling
|
|
margin = [b['margin'] for b in batch if b.get('margin') is not None]
|
|
if margin:
|
|
res['margin'] = torch.tensor(margin, dtype=torch.float)
|
|
|
|
return res
|
|
|
|
def _kto_data_collator(self, batch: List[Dict[str, Any]], *, padding_to: Optional[int] = None) -> Dict[str, Any]:
|
|
new_batch = self._fetch_inputs_startswith(batch, 'chosen_')
|
|
kl_batch = self._fetch_inputs_startswith(batch, 'rejected_')
|
|
|
|
res = self._data_collator(new_batch, padding_to=padding_to)
|
|
if any(kl_batch):
|
|
kl_res = self._data_collator(kl_batch, padding_to=padding_to)
|
|
else:
|
|
kl_res = {}
|
|
res = {
|
|
**{
|
|
f'completion_{k}': v
|
|
for k, v in res.items()
|
|
},
|
|
**{
|
|
f'KL_completion_{k}': v
|
|
for k, v in kl_res.items()
|
|
},
|
|
}
|
|
label = [b['label'] for b in batch if b.get('label') is not None]
|
|
if label:
|
|
res['label'] = label
|
|
return res
|
|
|
|
def _embedding_data_collator(self,
|
|
batch: List[Dict[str, Any]],
|
|
*,
|
|
padding_to: Optional[int] = None) -> Dict[str, Any]:
|
|
labels = []
|
|
new_batch = []
|
|
for b in batch:
|
|
if 'input_ids' in b:
|
|
new_batch += [b]
|
|
else:
|
|
keys = [key for key in b.keys() if 'negative' in key]
|
|
max_neg = None # number of negative samples
|
|
for key in keys:
|
|
value_list = b[key]
|
|
suffix = key[len('negative_'):]
|
|
max_neg = len(value_list)
|
|
for i, value in enumerate(value_list):
|
|
b[f'negative{i}_{suffix}'] = value
|
|
b.pop(key)
|
|
|
|
indexes = ['anchor_', 'positive_']
|
|
if max_neg is not None:
|
|
for i in range(0, max_neg):
|
|
indexes.append(f'negative{i}_')
|
|
for prefix in indexes:
|
|
new_batch += self._fetch_inputs_startswith([b], prefix)
|
|
labels.extend(b.get('labels', []))
|
|
res = self._data_collator(new_batch, padding_to=padding_to)
|
|
if labels:
|
|
res['labels'] = torch.tensor(labels, dtype=torch.float32)
|
|
return res
|
|
|
|
def _reranker_data_collator(self,
|
|
batch: List[Dict[str, Any]],
|
|
*,
|
|
padding_to: Optional[int] = None) -> Dict[str, Any]:
|
|
if self.is_training:
|
|
if not hasattr(self, 'random_state'):
|
|
# TODO: Move to `__init__`; kept here to avoid cache invalidation caused by template hash changes.
|
|
self.random_state = np.random.RandomState(42)
|
|
max_positive_samples = int(os.environ.get('MAX_POSITIVE_SAMPLES', 1))
|
|
max_negative_samples = int(os.environ.get('MAX_NEGATIVE_SAMPLES', 7))
|
|
labels_list = []
|
|
new_batch = []
|
|
for b in batch:
|
|
labels = b.pop('labels', None)
|
|
positive_num = sum(labels)
|
|
negative_num = len(labels) - positive_num
|
|
max_positive = min(positive_num, max_positive_samples)
|
|
max_negative = min(negative_num, max_negative_samples)
|
|
for i in self.random_state.choice(positive_num, max_positive, replace=False):
|
|
new_batch.append(
|
|
{key: b[key][i]
|
|
for key in b.keys() if isinstance(b[key], list) and b[key][i] is not None})
|
|
labels_list.append(1)
|
|
for j in self.random_state.choice(negative_num, max_negative, replace=False):
|
|
new_batch.append({
|
|
key: b[key][j + positive_num]
|
|
for key in b.keys() if isinstance(b[key], list) and b[key][j + positive_num] is not None
|
|
})
|
|
labels_list.append(0)
|
|
res = self._data_collator(new_batch, padding_to=padding_to)
|
|
if labels_list:
|
|
res['labels'] = torch.tensor(labels_list, dtype=torch.long)
|
|
else:
|
|
res = self._data_collator(batch, padding_to=padding_to)
|
|
return res
|
|
|
|
def _seq_cls_data_collator(self,
|
|
batch: List[Dict[str, Any]],
|
|
*,
|
|
padding_to: Optional[int] = None) -> Dict[str, Any]:
|
|
labels = [b.pop('labels') for b in batch if b.get('labels') is not None]
|
|
res = self._data_collator(batch, padding_to=padding_to)
|
|
if labels:
|
|
problem_type = self.config.problem_type
|
|
if problem_type == 'regression':
|
|
labels = torch.tensor(labels, dtype=torch.float32)
|
|
elif problem_type == 'multi_label_classification':
|
|
one_hot_labels = torch.zeros((len(labels), self.config.num_labels), dtype=torch.float32)
|
|
for i, label in enumerate(labels):
|
|
one_hot_labels[i, label] = 1
|
|
labels = one_hot_labels
|
|
else:
|
|
labels = torch.tensor(labels, dtype=torch.long)
|
|
res['labels'] = labels
|
|
return res
|
|
|
|
def _data_collator(self, batch: List[Dict[str, Any]], *, padding_to: Optional[int] = None) -> Dict[str, Any]:
|
|
"""
|
|
Args:
|
|
batch(`List[Dict[str, Any]]`): The input data in batch
|
|
padding_to(`int`, optional): Whether padding the batch to a fixed length, if none, the batch
|
|
will be padded to the `longest`
|
|
"""
|
|
assert self.tokenizer.pad_token_id is not None
|
|
padding_side = self.padding_side if self.is_training else 'left'
|
|
padding_right = padding_side == 'right'
|
|
real_seq_lens = [len(b['input_ids']) for b in batch] if self.use_megatron else None
|
|
self._handle_megatron_cp(batch)
|
|
if self.padding_free:
|
|
batch[:] = [self.packing_row(batch)]
|
|
assert 'position_ids' in batch[0], f'batch[0]: {batch[0]}'
|
|
elif self.use_megatron or self.sequence_parallel_size > 1:
|
|
assert padding_side == 'right', (
|
|
f'padding_side must be "right" when use_megatron or sequence_parallel_size > 1, got {padding_side!r}.')
|
|
for encoded in batch:
|
|
val = encoded['input_ids'] if encoded.get('labels') is None else encoded['labels']
|
|
encoded['position_ids'] = list(range(len(val)))
|
|
|
|
res = {}
|
|
gather_keys = ['labels', 'loss_scale', 'position_ids', 'token_type_ids', 'mm_token_type_ids']
|
|
if self.padding_free:
|
|
assert len(batch) == 1, f'batch: {batch}'
|
|
for k in ['input_ids', 'channel'] + gather_keys:
|
|
v = batch[0].get(k)
|
|
if v is not None:
|
|
res[k] = v if k == 'channel' else [v]
|
|
else:
|
|
inputs_embeds = [b['inputs_embeds'] for b in batch if b.get('inputs_embeds') is not None]
|
|
input_ids = [b['input_ids'] for b in batch if b.get('input_ids') is not None]
|
|
channel = [b.get('channel') for b in batch]
|
|
|
|
if inputs_embeds:
|
|
res['inputs_embeds'] = inputs_embeds
|
|
if input_ids:
|
|
res['input_ids'] = input_ids
|
|
if any(channel):
|
|
res['channel'] = channel
|
|
|
|
for key in gather_keys:
|
|
val = [b[key] for b in batch if b.get(key) is not None]
|
|
if val:
|
|
res[key] = val
|
|
|
|
pad_keys = [
|
|
'input_ids',
|
|
'inputs_embeds',
|
|
'attention_mask',
|
|
'attention_mask_2d',
|
|
] + gather_keys
|
|
pad_values = [self.tokenizer.pad_token_id, 0., 0, 0] + [-100, 0., 0, 0, 0]
|
|
# Convert to tensor and remove unnecessary dimensions.
|
|
seq_lens = None
|
|
for key in pad_keys:
|
|
if key not in res:
|
|
continue
|
|
for i, val in enumerate(res[key]):
|
|
if isinstance(val, (list, tuple)):
|
|
val = torch.tensor(val)
|
|
elif key == 'inputs_embeds' and val.ndim == 3 or key != 'inputs_embeds' and val.ndim == 2:
|
|
val = val[0]
|
|
res[key][i] = val
|
|
if not seq_lens:
|
|
seq_lens = [seq.shape[0] for seq in res[key]]
|
|
if not self.padding_free and seq_lens and ('input_ids' in res or 'inputs_embeds' in res):
|
|
attention_mask_key = 'attention_mask_2d' if self.use_megatron else 'attention_mask'
|
|
res[attention_mask_key] = [torch.ones(seq_len, dtype=torch.int64) for seq_len in seq_lens]
|
|
if self.is_training and self.padding_side == 'left':
|
|
res['position_ids'] = [torch.arange(seq_len, dtype=torch.int64) for seq_len in seq_lens]
|
|
|
|
if self.use_megatron:
|
|
if padding_to is not None:
|
|
padding_to = math.ceil(max(seq_lens) / padding_to) * padding_to
|
|
if self.padding_free:
|
|
cp_size = self.sequence_parallel_size
|
|
if cp_size > 1:
|
|
padding_len = padding_to - seq_lens[0]
|
|
position_ids = res['position_ids'][0]
|
|
extended_position_ids = torch.arange(cp_size * 2).repeat(padding_len // (cp_size * 2))
|
|
if position_ids.ndim == 3: # compat mrope
|
|
extended_position_ids = extended_position_ids[None,
|
|
None, :].expand(position_ids.shape[0], 1, -1)
|
|
res['position_ids'] = [torch.concat([position_ids, extended_position_ids], dim=-1)]
|
|
else:
|
|
seq_len = max(seq_lens) if padding_to is None else padding_to
|
|
res['attention_mask'] = torch.tril(torch.ones(
|
|
(len(seq_lens), seq_len, seq_len), dtype=torch.bool)).view(len(seq_lens), 1, seq_len, seq_len)
|
|
assert res['attention_mask'].dtype is torch.bool, f'attention_mask.dtype: {res["attention_mask"].dtype}'
|
|
for i, seq_len in enumerate(real_seq_lens):
|
|
res['attention_mask'][i, :, :, seq_len:] = 0
|
|
res['attention_mask'] = ~res['attention_mask']
|
|
|
|
for key, pad_value in zip(pad_keys, pad_values):
|
|
if key not in res:
|
|
continue
|
|
if self.use_megatron and not self.padding_free and key == 'attention_mask':
|
|
continue
|
|
if padding_to is not None and not (self.padding_free and key == 'position_ids'
|
|
and self.sequence_parallel_size > 1):
|
|
padding_len = padding_to - seq_lens[0]
|
|
if padding_len > 0:
|
|
res[key][0] = F.pad(res[key][0], (0, padding_len) if padding_right else (padding_len, 0),
|
|
'constant', pad_value)
|
|
if key == 'position_ids' and res[key][0].ndim == 3:
|
|
res[key] = self._pad_3d_position_ids(res[key], pad_value)
|
|
else:
|
|
res[key] = self._pad_sequence(res[key], pad_value)
|
|
|
|
# multimodal
|
|
res.update(self._data_collator_mm_data(batch))
|
|
if self.use_megatron:
|
|
res['seq_lens'] = real_seq_lens # CP locates the last token.
|
|
return res
|
|
|
|
def _pad_3d_position_ids(self,
|
|
position_ids: List[torch.Tensor],
|
|
padding_value: float = 0.,
|
|
batch_dim: int = 1) -> torch.Tensor:
|
|
padding_side = self.padding_side if self.is_training else 'left'
|
|
padding_right = padding_side == 'right'
|
|
# position_ids
|
|
# batch_dim 0: [1, 4, 379], [1, 4, 300] -> [2, 4, 379]
|
|
# batch_dim 1: [3/4, 1, 379], [3/4, 1, 300] -> [3/4, 2, 379]
|
|
max_len = max(pos.shape[-1] for pos in position_ids)
|
|
|
|
padded_position_ids = []
|
|
for pos in position_ids:
|
|
current_len = pos.shape[-1]
|
|
pad_len = max_len - current_len
|
|
|
|
if pad_len > 0:
|
|
pad_shape = (pos.shape[0], pos.shape[1], pad_len)
|
|
padding = pos.new_full(pad_shape, padding_value)
|
|
|
|
if padding_right:
|
|
padded_pos = torch.cat([pos, padding], dim=-1)
|
|
else:
|
|
padded_pos = torch.cat([padding, pos], dim=-1)
|
|
else:
|
|
padded_pos = pos
|
|
|
|
padded_position_ids.append(padded_pos)
|
|
|
|
result = torch.cat(padded_position_ids, dim=batch_dim)
|
|
|
|
return result
|
|
|
|
def create_mm_token_type_ids(self, input_ids: List[int], mm_mask: Optional[List[bool]] = None) -> torch.Tensor:
|
|
processor = self.processor
|
|
if not isinstance(input_ids, torch.Tensor):
|
|
input_ids = torch.tensor(input_ids)
|
|
if mm_mask is None:
|
|
mm_mask = True
|
|
elif not isinstance(mm_mask, torch.Tensor):
|
|
mm_mask = torch.tensor(mm_mask, dtype=torch.bool)
|
|
mm_token_type_ids = torch.zeros_like(input_ids)
|
|
for key, mm_token_id in zip(['image', 'video', 'audio'], [1, 2, 3]):
|
|
media_token_id = getattr(processor, f'{key}_token_id', None)
|
|
if media_token_id is None:
|
|
continue
|
|
mm_token_type_ids[(input_ids == media_token_id) & mm_mask] = mm_token_id
|
|
return mm_token_type_ids
|
|
|
|
def _data_collator_mm_data(self, batch: List[Dict[str, Any]]) -> Dict[str, Any]:
|
|
# multimodal
|
|
res = {}
|
|
pixel_values = [b['pixel_values'] for b in batch if b.get('pixel_values') is not None]
|
|
if len(pixel_values) > 0:
|
|
res['pixel_values'] = torch.concat(pixel_values)
|
|
|
|
image_sizes = [b['image_sizes'] for b in batch if b.get('image_sizes') is not None]
|
|
if len(image_sizes) > 0:
|
|
res['image_sizes'] = torch.concat(image_sizes)
|
|
|
|
pixel_values_videos = [b['pixel_values_videos'] for b in batch if b.get('pixel_values_videos') is not None]
|
|
if len(pixel_values_videos) > 0:
|
|
res['pixel_values_videos'] = torch.concat(pixel_values_videos)
|
|
|
|
for media_type in ['image', 'video']:
|
|
grid_thw = self.concat_tensor(batch, f'{media_type}_grid_thw', 0)
|
|
if grid_thw is not None:
|
|
res[f'{media_type}_grid_thw'] = grid_thw
|
|
return res
|
|
|
|
def print_inputs(self, inputs: Dict[str, Any]) -> None:
|
|
# Base keys to check
|
|
tokenizer_kwargs = inputs.pop('tokenizer_kwargs', None) or {}
|
|
base_keys = [
|
|
'input', 'labels', 'generate', 'chosen_input', 'chosen_labels', 'rejected_input', 'rejected_labels'
|
|
]
|
|
|
|
# For reranker/embedding modes, also check prefixed keys
|
|
if self.task_type in {'reranker', 'generative_reranker', 'embedding'}:
|
|
prefixes = []
|
|
if self.task_type in {'reranker', 'generative_reranker'}:
|
|
prefixes = ['positive_', 'negative_']
|
|
elif self.task_type == 'embedding':
|
|
prefixes = ['anchor_', 'positive_', 'negative_']
|
|
|
|
# Add prefixed keys for reranker/embedding modes
|
|
extended_keys = base_keys.copy()
|
|
for prefix in prefixes:
|
|
for base_key in ['input', 'labels']:
|
|
extended_keys.append(f'{prefix}{base_key}')
|
|
|
|
# Also check for numbered negative keys (negative0_, negative1_, etc.)
|
|
input_keys = list(inputs.keys())
|
|
for key in input_keys:
|
|
if any(key.startswith(f'{prefix}') for prefix in prefixes):
|
|
# Extract the base key after removing prefix
|
|
for prefix in prefixes:
|
|
if key.startswith(prefix):
|
|
base_key = key[len(prefix):]
|
|
if base_key in ['input_ids', 'labels'
|
|
] or base_key.rstrip('0123456789_') in ['input', 'labels']:
|
|
extended_keys.append(key.replace('_ids', ''))
|
|
break
|
|
|
|
keys_to_check = list(set(extended_keys))
|
|
else:
|
|
keys_to_check = base_keys
|
|
|
|
for key in keys_to_check:
|
|
# Skip labels completely for certain modes
|
|
if key.endswith('labels') and self.task_type in {'reranker', 'generative_reranker'}:
|
|
continue
|
|
|
|
val = inputs.get(key) # fix val is a tensor
|
|
if val is None:
|
|
val = inputs.get(f'{key}_ids')
|
|
if val is not None:
|
|
key_upper = key.upper()
|
|
logger.info(f'[{key_upper}_IDS] {val}')
|
|
if key.endswith('labels') and self.task_type in {'seq_cls', 'embedding'}:
|
|
continue
|
|
if isinstance(val, (list, tuple, torch.Tensor)):
|
|
# Handle nested lists (e.g., for reranker negative samples)
|
|
if isinstance(val, (list, tuple)) and len(val) > 0 and isinstance(val[0], (list, tuple)):
|
|
val_str = [self.safe_decode(sub_val, **tokenizer_kwargs) for sub_val in val]
|
|
else:
|
|
val_str = self.safe_decode(val, **tokenizer_kwargs)
|
|
logger.info(f'[{key_upper}] {val_str}')
|
|
if inputs.get('loss_scale') is not None:
|
|
val = inputs['loss_scale']
|
|
logger.info(f'[LOSS_SCALE] {val}')
|
|
|
|
async def prepare_lmdeploy_pytorch_inputs(self, inputs) -> None:
|
|
images = inputs.pop('images', None) or []
|
|
if len(images) == 0:
|
|
return
|
|
input_ids = inputs['input_ids']
|
|
idx_list = findall(input_ids, -100)
|
|
assert len(idx_list) == len(images), f'len(idx_list): {len(idx_list)}, len(images): {len(images)}'
|
|
idx_list.insert(0, -1)
|
|
new_input_ids = []
|
|
for i in range(len(idx_list) - 1):
|
|
new_input_ids += input_ids[idx_list[i] + 1:idx_list[i + 1]]
|
|
images[i]['offset'] = len(new_input_ids)
|
|
new_input_ids += [images[i]['image_token_id']] * images[i]['image_tokens']
|
|
new_input_ids += input_ids[idx_list[-1] + 1:]
|
|
inputs['input_ids'] = new_input_ids
|
|
inputs['multimodal'] = images
|
|
|
|
async def prepare_lmdeploy_turbomind_inputs(self, inputs: Dict[str, Any]) -> None:
|
|
images = inputs.pop('images', None) or []
|
|
if len(images) == 0:
|
|
return
|
|
from lmdeploy.vl.constants import IMAGE_DUMMY_TOKEN_INDEX
|
|
input_ids = inputs['input_ids']
|
|
idx_list = findall(input_ids, -100)
|
|
assert len(idx_list) == len(images), f'len(idx_list): {len(idx_list)}, len(images): {len(images)}'
|
|
idx_list.insert(0, -1)
|
|
new_input_ids = []
|
|
ranges = []
|
|
for i in range(len(idx_list) - 1):
|
|
_range = []
|
|
new_input_ids += input_ids[idx_list[i] + 1:idx_list[i + 1]]
|
|
_range.append(len(new_input_ids))
|
|
new_input_ids += [IMAGE_DUMMY_TOKEN_INDEX] * images[i].shape[0]
|
|
_range.append(len(new_input_ids))
|
|
ranges.append(_range)
|
|
new_input_ids += input_ids[idx_list[-1] + 1:]
|
|
inputs['input_embeddings'] = [image.to('cpu') for image in images]
|
|
inputs['input_embedding_ranges'] = ranges
|
|
inputs['input_ids'] = new_input_ids
|
|
|
|
def _pad_sequence(self, sequences: List[torch.Tensor], padding_value: float = 0.) -> torch.Tensor:
|
|
"""Pad sequence by some side
|
|
|
|
Args:
|
|
sequences: The input sequences in tensor.
|
|
padding_value: The padding value
|
|
|
|
Returns:
|
|
A tensor after padding
|
|
"""
|
|
padding_side = self.padding_side if self.is_training else 'left'
|
|
padding_right = padding_side == 'right'
|
|
if padding_right:
|
|
return pad_sequence(sequences, batch_first=True, padding_value=padding_value)
|
|
|
|
max_len = max([s.shape[0] for s in sequences])
|
|
|
|
padded_sequences = []
|
|
for seq in sequences:
|
|
pad_length = max_len - seq.shape[0]
|
|
pad_tuple = [0] * ((seq.dim() - 1) * 2) + [pad_length, 0]
|
|
padded_seq = F.pad(seq, tuple(pad_tuple), 'constant', padding_value)
|
|
padded_sequences.append(padded_seq)
|
|
|
|
return torch.stack(padded_sequences)
|
|
|
|
def safe_decode(self, input_ids: List[int], **kwargs) -> str:
|
|
if isinstance(self, Template):
|
|
tokenizer = self.tokenizer
|
|
placeholder_tokens = self.placeholder_tokens
|
|
else:
|
|
tokenizer = self
|
|
placeholder_tokens = []
|
|
|
|
def _is_special(token: int) -> bool:
|
|
if isinstance(token, float) or token < 0:
|
|
return True
|
|
return token in placeholder_tokens
|
|
|
|
if isinstance(input_ids, torch.Tensor):
|
|
input_ids = input_ids.tolist()
|
|
if len(input_ids) == 0:
|
|
return ''
|
|
result_str = ''
|
|
for i in range(len(input_ids)):
|
|
if i == 0:
|
|
if _is_special(input_ids[i]):
|
|
s = 0
|
|
else:
|
|
e = 0
|
|
continue
|
|
if _is_special(input_ids[i]) and not _is_special(input_ids[i - 1]):
|
|
s = i
|
|
result_str += tokenizer.decode(input_ids[e:s], **kwargs)
|
|
if not _is_special(input_ids[i]) and _is_special(input_ids[i - 1]):
|
|
e = i
|
|
result_str += f'[{input_ids[i - 1]} * {e - s}]'
|
|
if _is_special(input_ids[i]):
|
|
result_str += f'[{input_ids[i]} * {len(input_ids) - s}]'
|
|
else:
|
|
result_str += tokenizer.decode(input_ids[e:], **kwargs)
|
|
return result_str
|
|
|
|
@staticmethod
|
|
@contextmanager
|
|
def _patch_flash_attention_forward(modeling_module, position_ids, use_new_func: bool = False):
|
|
_origin_flash_attention_forward = modeling_module._flash_attention_forward
|
|
|
|
def _flash_attention_forward(*args, **kwargs):
|
|
if use_new_func:
|
|
from transformers.modeling_flash_attention_utils import \
|
|
_flash_attention_forward as flash_attention_forward
|
|
if args and isinstance(args[0], nn.Module):
|
|
args = args[1:]
|
|
if 'is_causal' not in kwargs:
|
|
kwargs['is_causal'] = True
|
|
else:
|
|
flash_attention_forward = _origin_flash_attention_forward
|
|
kwargs['position_ids'] = position_ids
|
|
if args and isinstance(args[0], torch.Tensor):
|
|
kwargs['position_ids'] = kwargs['position_ids'].to(args[0].device)
|
|
return flash_attention_forward(*args, **kwargs)
|
|
|
|
modeling_module._flash_attention_forward = _flash_attention_forward
|
|
try:
|
|
yield
|
|
finally:
|
|
modeling_module._flash_attention_forward = _origin_flash_attention_forward
|
|
|
|
@staticmethod
|
|
def _get_inputs_embeds_hf(inputs_embeds, inputs, visual, processor, config):
|
|
input_ids = inputs['input_ids']
|
|
pixel_values = inputs.get('pixel_values')
|
|
pixel_values_videos = inputs.get('pixel_values_videos')
|
|
image_grid_thw = inputs.get('image_grid_thw')
|
|
video_grid_thw = inputs.get('video_grid_thw')
|
|
dtype = visual.dtype
|
|
if pixel_values is None and pixel_values_videos is None: # plain-text
|
|
images = [Image.new('RGB', (32, 32), (0, 0, 0))]
|
|
media_inputs = processor.image_processor(images=images, return_tensors='pt')
|
|
media_inputs = to_device(media_inputs, input_ids.device)
|
|
pixel_values = media_inputs['pixel_values'].type(dtype)
|
|
image_embeds = visual(pixel_values, grid_thw=media_inputs['image_grid_thw'])
|
|
if hasattr(image_embeds, 'pooler_output'):
|
|
image_embeds = image_embeds.pooler_output
|
|
inputs_embeds = inputs_embeds + image_embeds.mean().to(device=inputs_embeds.device) * 0.
|
|
else:
|
|
if pixel_values is None:
|
|
pixel_values_mixed = pixel_values_videos
|
|
grid_thw = video_grid_thw
|
|
elif pixel_values_videos is None:
|
|
pixel_values_mixed = pixel_values
|
|
grid_thw = image_grid_thw
|
|
else:
|
|
pixel_values_mixed = torch.concat([pixel_values, pixel_values_videos], dim=0)
|
|
grid_thw = torch.concat([image_grid_thw, video_grid_thw], dim=0)
|
|
pixel_values_mixed = pixel_values_mixed.type(dtype)
|
|
mixed_embeds = visual(pixel_values_mixed, grid_thw=grid_thw)
|
|
if hasattr(mixed_embeds, 'pooler_output'):
|
|
mixed_embeds = mixed_embeds.pooler_output
|
|
if pixel_values is None:
|
|
image_embeds = None
|
|
video_embeds = mixed_embeds
|
|
elif pixel_values_videos is None:
|
|
image_embeds = mixed_embeds
|
|
video_embeds = None
|
|
else:
|
|
merge_length = processor.image_processor.merge_size**2
|
|
image_tokens = (image_grid_thw.prod(dim=-1) // merge_length).sum()
|
|
image_embeds = mixed_embeds[:image_tokens]
|
|
video_embeds = mixed_embeds[image_tokens:]
|
|
|
|
if image_embeds is not None:
|
|
image_mask = (input_ids == config.image_token_id).unsqueeze(-1).expand_as(inputs_embeds)
|
|
image_embeds = image_embeds.to(inputs_embeds.device, inputs_embeds.dtype)
|
|
image_mask = image_mask.to(inputs_embeds.device)
|
|
inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds)
|
|
|
|
if video_embeds is not None:
|
|
video_mask = (input_ids == config.video_token_id).unsqueeze(-1).expand_as(inputs_embeds)
|
|
video_embeds = video_embeds.to(inputs_embeds.device, inputs_embeds.dtype)
|
|
video_mask = video_mask.to(inputs_embeds.device)
|
|
inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds)
|
|
return inputs_embeds
|
|
|
|
@staticmethod
|
|
def _concat_text_position_ids(position_ids):
|
|
seq_len = position_ids.shape[-1]
|
|
text_position_ids = torch.arange(seq_len, device=position_ids.device).expand(1, *position_ids.shape[1:])
|
|
return torch.concat([text_position_ids, position_ids], dim=0)
|