This commit is contained in:
@@ -0,0 +1,219 @@
|
||||
# Copyright (c) ModelScope Contributors. All rights reserved.
|
||||
import json
|
||||
from copy import deepcopy
|
||||
from dataclasses import dataclass, field, fields
|
||||
from PIL import Image
|
||||
from typing import Any, Dict, List, Optional, Union
|
||||
|
||||
from swift.utils import get_logger
|
||||
from .utils import Messages, Tool, get_last_user_round, messages_to_history
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
|
||||
@dataclass
|
||||
class StdTemplateInputs:
|
||||
# only user/tool/assistant
|
||||
messages: List[Dict[str, str]]
|
||||
# None: use default system; '': not use system
|
||||
system: Optional[str] = None
|
||||
tools: Optional[List[Tool]] = None
|
||||
|
||||
label: Optional[int] = None
|
||||
channel: Optional[str] = None
|
||||
|
||||
images: List[Union[str, Image.Image]] = field(default_factory=list)
|
||||
videos: List[str] = field(default_factory=list)
|
||||
audios: List[str] = field(default_factory=list)
|
||||
objects: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
margin: Optional[float] = None # for reward modeling
|
||||
chat_template_kwargs: Dict[str, Any] = field(default_factory=dict) # from dataset
|
||||
extra_kwargs: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
mm_processor_kwargs: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
def __post_init__(self):
|
||||
self.image_idx = 0
|
||||
self.audio_idx = 0
|
||||
self.video_idx = 0
|
||||
self.ref_idx = 0
|
||||
self.bbox_idx = 0
|
||||
if self.images and not isinstance(self.images, (list, tuple)):
|
||||
self.images = [self.images]
|
||||
if self.videos and not isinstance(self.videos, (list, tuple)):
|
||||
self.videos = [self.videos]
|
||||
if self.audios and not isinstance(self.audios, (list, tuple)):
|
||||
self.audios = [self.audios]
|
||||
|
||||
def to_history(self):
|
||||
if not self.messages:
|
||||
return None
|
||||
return messages_to_history(self.messages)
|
||||
|
||||
@property
|
||||
def is_multimodal(self):
|
||||
return bool(self.images or self.audios or self.videos or self.objects)
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, inputs: Dict[str, Any]) -> 'StdTemplateInputs':
|
||||
inputs = deepcopy(inputs)
|
||||
kwargs = {}
|
||||
for key in ['label', 'channel', 'margin', 'rejected_response']:
|
||||
if key in inputs:
|
||||
kwargs[key] = inputs[key]
|
||||
messages = inputs['messages']
|
||||
tools = inputs.get('tools')
|
||||
objects = inputs.get('objects') or {}
|
||||
chat_template_kwargs = inputs.get('chat_template_kwargs') or {}
|
||||
|
||||
if messages and messages[0]['role'] == 'system':
|
||||
message = messages.pop(0)
|
||||
system = message['content']
|
||||
else:
|
||||
system = None
|
||||
|
||||
for message in messages:
|
||||
if message['role'] == 'tool_response':
|
||||
message['role'] = 'tool'
|
||||
if message['role'] in {'tool_call', 'tool'} and not isinstance(message['content'], str):
|
||||
message['content'] = json.dumps(message['content'], ensure_ascii=False)
|
||||
|
||||
media_kwargs = StdTemplateInputs.remove_messages_media(messages)
|
||||
for k in list(media_kwargs.keys()):
|
||||
mm_data = media_kwargs[k]
|
||||
|
||||
inputs_mm_data = inputs.get(k)
|
||||
if isinstance(inputs_mm_data, str):
|
||||
inputs_mm_data = [inputs_mm_data]
|
||||
inputs_mm_data = (inputs_mm_data or []).copy()
|
||||
if mm_data:
|
||||
assert not inputs_mm_data, f'self.{k}: {inputs_mm_data}'
|
||||
else:
|
||||
media_kwargs[k] = inputs_mm_data
|
||||
|
||||
all_keys = set(f.name for f in fields(StdTemplateInputs))
|
||||
extra_kwargs = {k: v for k, v in inputs.items() if k not in all_keys}
|
||||
return cls(
|
||||
messages=messages,
|
||||
system=system,
|
||||
tools=tools,
|
||||
objects=objects,
|
||||
chat_template_kwargs=chat_template_kwargs,
|
||||
extra_kwargs=extra_kwargs,
|
||||
**kwargs,
|
||||
**media_kwargs)
|
||||
|
||||
@staticmethod
|
||||
def remove_messages_media(messages: Messages) -> Dict[str, Any]:
|
||||
res = {'images': [], 'audios': [], 'videos': []}
|
||||
for message in messages:
|
||||
content = message['content']
|
||||
if isinstance(content, str):
|
||||
continue
|
||||
elif (isinstance(content, list) and content
|
||||
and isinstance(content[0], int)) or (isinstance(content, dict) and 'token_ids' in content):
|
||||
continue
|
||||
# List[Dict[str, Any]]
|
||||
new_content = ''
|
||||
for item in content:
|
||||
key: str = item['type']
|
||||
value = item.get(key)
|
||||
if key == 'text':
|
||||
new_content += value
|
||||
continue
|
||||
# image/audio/video
|
||||
# image_url/audio_url/video_url
|
||||
if key.endswith('_url'):
|
||||
key = key[:-len('_url')]
|
||||
new_content += f'<{key}>'
|
||||
if isinstance(value, dict):
|
||||
value = value['url']
|
||||
if value:
|
||||
res[f'{key}s'].append(value)
|
||||
message['content'] = new_content
|
||||
return res
|
||||
|
||||
|
||||
@dataclass
|
||||
class TemplateInputs:
|
||||
chosen: StdTemplateInputs # or Dict[str, Any]
|
||||
rejected: Optional[StdTemplateInputs] = None
|
||||
positive: List[StdTemplateInputs] = field(default_factory=list) # or Dict[str, Any]
|
||||
negative: List[StdTemplateInputs] = field(default_factory=list)
|
||||
|
||||
def __post_init__(self):
|
||||
all_keys = set(f.name for f in fields(StdTemplateInputs))
|
||||
for key in ['chosen', 'rejected', 'positive', 'negative']:
|
||||
value_dict = getattr(self, key, None)
|
||||
if not isinstance(value_dict, dict):
|
||||
continue
|
||||
if key in {'chosen', 'rejected'}:
|
||||
setattr(self, key, StdTemplateInputs.from_dict(value_dict))
|
||||
else:
|
||||
res = []
|
||||
for i in range(len(value_dict['messages'])):
|
||||
kwargs = {}
|
||||
for k in all_keys:
|
||||
val = value_dict.get(k)
|
||||
if val is None:
|
||||
continue
|
||||
kwargs[k] = val[i]
|
||||
res.append(StdTemplateInputs.from_dict(kwargs))
|
||||
setattr(self, key, res)
|
||||
|
||||
@staticmethod
|
||||
def _compat_rejected_response(inputs: Dict[str, Any]):
|
||||
if 'rejected_response' not in inputs:
|
||||
return
|
||||
messages = inputs['messages']
|
||||
assert len(messages) > 0, f'messages: {messages}'
|
||||
idx = get_last_user_round(messages) + 1
|
||||
|
||||
rejected_response = inputs.pop('rejected_response')
|
||||
if isinstance(rejected_response, str):
|
||||
rejected_responses = [{'role': 'assistant', 'content': rejected_response}]
|
||||
elif isinstance(rejected_response, list):
|
||||
rejected_responses = rejected_response
|
||||
for message in rejected_responses:
|
||||
if message['role'] == 'user':
|
||||
raise ValueError(
|
||||
f"The 'user' role is not allowed in 'rejected_response' messages. Found: {message}")
|
||||
else:
|
||||
raise ValueError(f'rejected_response must be a str or list. rejected_response: {rejected_response}')
|
||||
# Check that the response is different from the rejected_response.
|
||||
if len(messages[idx:]) == 1 and len(rejected_responses) == 1:
|
||||
response = messages[idx]['content']
|
||||
rejected_response = rejected_responses[0]['content']
|
||||
assert rejected_response != response, f'rejected_response: {rejected_response}, response: {response}'
|
||||
inputs['rejected_messages'] = deepcopy(messages[:idx]) + rejected_responses
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, inputs: Dict[str, Any]) -> 'TemplateInputs':
|
||||
inputs = deepcopy(inputs)
|
||||
|
||||
has_rejected_messages = inputs.get('rejected_messages') is not None
|
||||
cls._compat_rejected_response(inputs)
|
||||
kwargs = {}
|
||||
non_chosen_keys = ['rejected', 'positive', 'negative']
|
||||
for prefix in ['chosen'] + non_chosen_keys:
|
||||
if prefix == 'chosen':
|
||||
std_inputs = {
|
||||
k: v
|
||||
for k, v in inputs.items() if not any(k.startswith(f'{p}_') for p in non_chosen_keys)
|
||||
}
|
||||
else:
|
||||
std_inputs = {k[len(f'{prefix}_'):]: v for k, v in inputs.items() if k.startswith(f'{prefix}_')}
|
||||
if std_inputs:
|
||||
kwargs[prefix] = std_inputs
|
||||
|
||||
if not has_rejected_messages and kwargs.get('rejected') is not None:
|
||||
chosen = kwargs['chosen']
|
||||
rejected = kwargs['rejected']
|
||||
# Supplement additional key-value pairs
|
||||
for k, chosen_v in chosen.items():
|
||||
rejected_v = rejected.get(k)
|
||||
if chosen_v is not None and rejected_v is None:
|
||||
rejected[k] = chosen_v
|
||||
|
||||
return cls(**kwargs)
|
||||
Reference in New Issue
Block a user