wehub-resource-sync
247 lines
9.2 KiB
Python
247 lines
9.2 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
import re
|
|
import torch
|
|
from transformers import PreTrainedTokenizerBase, StoppingCriteria
|
|
from typing import Any, Dict, List, Optional, Set, Tuple, Type, Union
|
|
|
|
from swift.utils import get_logger
|
|
|
|
logger = get_logger()
|
|
|
|
Tool = Dict[str, Union[str, Dict]]
|
|
History = List[Union[Tuple[str, str], List[str]]]
|
|
Message = Dict[str, Union[str, List[Dict[str, Any]], List[int], None]]
|
|
Messages = List[Message]
|
|
Prompt = List[Union[str, List[int], List[str]]]
|
|
Word = Union[str, List[int]]
|
|
Context = Word
|
|
|
|
|
|
class ContextType:
|
|
RESPONSE = 'response'
|
|
SUFFIX = 'suffix'
|
|
OTHER = 'other'
|
|
|
|
|
|
class StopWordsCriteria(StoppingCriteria):
|
|
"""Adding extra stop words in template to prevent unstoppable generation
|
|
Like suffixes and chat seps in the template.
|
|
"""
|
|
|
|
def __init__(self, tokenizer: PreTrainedTokenizerBase, stop_words: List[Word], **tokenizer_kwargs) -> None:
|
|
self.tokenizer = tokenizer
|
|
self.stop_words = stop_words
|
|
self.tokenizer_kwargs = tokenizer_kwargs
|
|
self.start_idx = -1
|
|
self.is_done = None
|
|
|
|
def __call__(self, input_ids: torch.Tensor, scores: torch.Tensor, **kwargs) -> torch.Tensor:
|
|
if self.start_idx == -1:
|
|
self.start_idx = len(input_ids[0]) - 1
|
|
self.is_done = torch.full((input_ids.shape[0], ), False, device=input_ids.device, dtype=torch.bool)
|
|
# [-20:]: Assuming the end tokens do not exceed 20 tokens,
|
|
# to avoid input_ids being too long and affecting efficiency.
|
|
start_idx = max(self.start_idx, input_ids.shape[1] - 20)
|
|
text_list = self.tokenizer.batch_decode(input_ids[:, start_idx:], **self.tokenizer_kwargs)
|
|
for i, text in enumerate(text_list):
|
|
if self.is_done[i]:
|
|
continue
|
|
is_finished = False
|
|
for stop_word in self.stop_words:
|
|
if isinstance(stop_word, str) and stop_word in text or isinstance(
|
|
stop_word, list) and input_ids[i][-len(stop_word):].tolist() == stop_word:
|
|
is_finished = True
|
|
break
|
|
self.is_done[i] = is_finished
|
|
return self.is_done
|
|
|
|
|
|
def fetch_one(element: Union[Tuple, List, Set, Dict, Any], item_type: Optional[Type] = None) -> Any:
|
|
if isinstance(element, (tuple, set, list)):
|
|
for ele in element:
|
|
out = fetch_one(ele)
|
|
if out and (item_type is None or isinstance(out, item_type)):
|
|
return out
|
|
elif isinstance(element, dict):
|
|
return fetch_one(list(element.values()))
|
|
else:
|
|
return element
|
|
|
|
|
|
def findall(token_list: List[int], sub_token_list: Union[int, List[int]]) -> List[int]:
|
|
"""Find the index of a token in the token_list."""
|
|
if isinstance(sub_token_list, int):
|
|
sub_token_list = [sub_token_list]
|
|
res = []
|
|
idx = -1
|
|
try:
|
|
while True:
|
|
idx = token_list.index(sub_token_list[0], idx + 1)
|
|
if len(sub_token_list) == 1 or sub_token_list == token_list[idx:idx + len(sub_token_list)]:
|
|
res.append(idx)
|
|
except ValueError:
|
|
pass
|
|
return res
|
|
|
|
|
|
def align_image_inputs(input_ids: List[int], labels: List[int], new_input_ids,
|
|
image_token: int) -> Tuple[List[int], List[int]]:
|
|
if isinstance(new_input_ids, torch.Tensor):
|
|
new_input_ids = new_input_ids.tolist()
|
|
|
|
# Find the tokens after the image_token in input_ids, and then align them.
|
|
i, j = 0, 0
|
|
while i < len(input_ids):
|
|
x = input_ids[i]
|
|
if x == image_token:
|
|
assert i + 1 < len(input_ids), f'input_ids[-10:]: {input_ids[-10:]}'
|
|
assert i - 1 >= 0, f'input_ids[:10]: {input_ids[:10]}'
|
|
# [1, 2, 3(i-1), image_token(i), 4(i+1) ,5, 6]
|
|
# [1, 2, 3(j_begin), a(j'), a, a, a, 4(j) ,5, 6]
|
|
j_begin = j - 1
|
|
for k in range(5): # Increase robustness.
|
|
if j_begin + k < len(new_input_ids) and new_input_ids[j_begin + k] == input_ids[i - 1]:
|
|
j_begin += k
|
|
break
|
|
if j_begin - k >= 0 and new_input_ids[j_begin - k] == input_ids[i - 1]:
|
|
j_begin -= k
|
|
break
|
|
else:
|
|
raise ValueError(f'new_input_ids: {new_input_ids}, input_ids: {input_ids}')
|
|
j_begin += 1
|
|
while j < len(new_input_ids) and new_input_ids[j] != input_ids[i + 1]:
|
|
j += 1
|
|
input_ids = input_ids[:i] + new_input_ids[j_begin:j] + input_ids[i + 1:]
|
|
if labels:
|
|
labels = labels[:i] + [-100] * (j - j_begin) + labels[i + 1:]
|
|
i += j - j_begin
|
|
else:
|
|
j += 1
|
|
i += 1
|
|
return input_ids, labels
|
|
|
|
|
|
def _split_str_by_regex(text: str, regex_delimiters: List[str]) -> List[str]:
|
|
combined_pattern = '|'.join(f'({pattern})' for pattern in regex_delimiters)
|
|
parts = re.split(combined_pattern, text, flags=re.DOTALL)
|
|
parts = [part for part in parts if part is not None]
|
|
if parts[0] == '':
|
|
parts.pop(0)
|
|
else:
|
|
parts.insert(0, '')
|
|
assert len(parts) % 2 == 0, f'result: {parts}'
|
|
assert ''.join(parts) == text, f'split_result: {parts}, text: {text}'
|
|
return parts
|
|
|
|
|
|
def split_str_parts_by(text: str, delimiters: List[str], regex_mode: bool = False) -> List[Dict[str, str]]:
|
|
"""Split the text field into parts.
|
|
|
|
Args:
|
|
text: A text to be split.
|
|
delimiters: The delimiters.
|
|
|
|
Returns:
|
|
The split text in list of dicts.
|
|
"""
|
|
assert isinstance(text, str), f'text: {text}'
|
|
delimiters_origin = delimiters
|
|
if not regex_mode:
|
|
delimiters = [re.escape(delimiter) for delimiter in delimiters]
|
|
parts = _split_str_by_regex(text, delimiters) if delimiters else ['', text]
|
|
res = []
|
|
if regex_mode:
|
|
parts = [part for part in parts if part]
|
|
for part in parts:
|
|
for delimiter, delimiter_origin in zip(delimiters, delimiters_origin):
|
|
if re.match(delimiter, part, re.DOTALL):
|
|
break
|
|
else:
|
|
delimiter_origin = ''
|
|
res.append({'key': delimiter_origin, 'content': part})
|
|
else:
|
|
for key, content in zip(parts[::2], parts[1::2]):
|
|
res.append({'key': key, 'content': content})
|
|
return res
|
|
|
|
|
|
def get_last_user_round(messages):
|
|
"""Get the index of the last occurrence of user role"""
|
|
for i in range(len(messages) - 1, -1, -1):
|
|
if messages[i]['role'] == 'user':
|
|
return i
|
|
return -1
|
|
|
|
|
|
def history_to_messages(history: History,
|
|
system: Optional[str] = None,
|
|
roles: Optional[List[List[str]]] = None) -> 'Messages':
|
|
"""
|
|
history: [['query1', 'response1'], ['query2', 'response2']]
|
|
or [['query1', 'response1'], ['query2', None]]
|
|
"""
|
|
messages = []
|
|
if not roles:
|
|
roles = [['user', 'assistant']] * len(history)
|
|
else:
|
|
assert len(roles) == len(history), f'len(roles): {len(roles)}, len(history): {len(history)}'
|
|
if system is not None:
|
|
messages.append({'role': 'system', 'content': system})
|
|
|
|
for role, h in zip(roles, history):
|
|
assert isinstance(h, (list, tuple))
|
|
if h[0] is not None:
|
|
messages.append({'role': role[0], 'content': h[0]})
|
|
if h[1] is not None:
|
|
messages.append({'role': role[1], 'content': h[1]})
|
|
return messages
|
|
|
|
|
|
def messages_to_history(messages: 'Messages') -> Dict[str, Any]:
|
|
system = None
|
|
messages = messages.copy()
|
|
if messages[0]['role'] == 'system':
|
|
system = messages[0]['content']
|
|
messages = messages[1::]
|
|
if len(messages) % 2 == 1:
|
|
messages.append({'role': 'assistant', 'content': None})
|
|
history = []
|
|
history_roles = []
|
|
for user_message, assistant_message in zip(messages[::2], messages[1::2]):
|
|
assert user_message['role'] in {'tool', 'user'}, f'user_message {user_message}'
|
|
assert assistant_message['role'] == 'assistant', f'assistant_message: {assistant_message}'
|
|
history.append([user_message['content'], assistant_message['content']])
|
|
history_roles.append([user_message['role'], assistant_message['role']])
|
|
query, response = history.pop() if history else (None, None)
|
|
query_role = history_roles.pop()[0] if history_roles else None
|
|
return {
|
|
'history': history,
|
|
'history_roles': history_roles,
|
|
'query': query,
|
|
'query_role': query_role,
|
|
'response': response,
|
|
'system': system,
|
|
}
|
|
|
|
|
|
def update_generation_config_eos_token(generation_config, template):
|
|
if generation_config is None:
|
|
return
|
|
stop_words = template.template_meta.stop_words
|
|
eos_token_id = generation_config.eos_token_id
|
|
if eos_token_id is None:
|
|
eos_token_id = []
|
|
elif isinstance(eos_token_id, int):
|
|
eos_token_id = [eos_token_id]
|
|
modified = False
|
|
for stop_word in stop_words:
|
|
if stop_word is None:
|
|
continue
|
|
if isinstance(stop_word, str):
|
|
stop_word = template._tokenize(stop_word)
|
|
if isinstance(stop_word, (list, tuple)) and len(stop_word) == 1 and stop_word[0] not in eos_token_id:
|
|
eos_token_id.append(stop_word[0])
|
|
modified = True
|
|
if modified:
|
|
generation_config.eos_token_id = eos_token_id
|