1069 lines
43 KiB
Python
1069 lines
43 KiB
Python
|
#!/usr/bin/env python3
|
|||
|
|
# coding=utf-8
|
|||
|
|
# Copyright (c) Ant Group. All rights reserved.
|
|||
|
|
|
|||
|
|
import itertools
|
|||
|
|
from typing import Any, Dict, List, Optional, Union
|
|||
|
|
|
|||
|
|
import torch
|
|||
|
|
from transformers import PreTrainedTokenizerFast
|
|||
|
|
from transformers.tokenization_utils_base import AddedToken, BatchEncoding
|
|||
|
|
from transformers.utils import TensorType, logging
|
|||
|
|
|
|||
|
|
logger = logging.get_logger(__name__)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def is_system(msg):
|
|||
|
|
return msg['role'].lower() == 'system'
|
|||
|
|
|
|||
|
|
|
|||
|
|
def is_user(msg):
|
|||
|
|
return msg['role'].lower() in ['human', 'user']
|
|||
|
|
|
|||
|
|
|
|||
|
|
def is_assistant(msg):
|
|||
|
|
return msg['role'].lower() == 'assistant'
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _convert_to_conversation(query, system=None):
|
|||
|
|
conversation = []
|
|||
|
|
if system:
|
|||
|
|
conversation.append({"role": "SYSTEM", "content": system})
|
|||
|
|
if isinstance(query, str):
|
|||
|
|
conversation.append({"role": "HUMAN", "content": query})
|
|||
|
|
elif isinstance(query, List):
|
|||
|
|
conversation.extend(query)
|
|||
|
|
elif isinstance(query, Dict):
|
|||
|
|
if "messages" in query:
|
|||
|
|
conversation.extend(query["messages"])
|
|||
|
|
if "system_message" in query and len(conversation) > 0 and not is_system(conversation[0]):
|
|||
|
|
conversation.insert(0, {"role": "SYSTEM", "content": query["system_message"]})
|
|||
|
|
else:
|
|||
|
|
conversation.append(query)
|
|||
|
|
return conversation
|
|||
|
|
|
|||
|
|
|
|||
|
|
class BailingTokenizer(PreTrainedTokenizerFast):
|
|||
|
|
is_bailing_tokenizer = True
|
|||
|
|
model_input_names = ["input_ids", "attention_mask"]
|
|||
|
|
slow_tokenizer_class = None
|
|||
|
|
|
|||
|
|
# add gmask_token
|
|||
|
|
SPECIAL_TOKENS_ATTRIBUTES = [
|
|||
|
|
"bos_token",
|
|||
|
|
"eos_token",
|
|||
|
|
"unk_token",
|
|||
|
|
"sep_token",
|
|||
|
|
"pad_token",
|
|||
|
|
"cls_token",
|
|||
|
|
"mask_token",
|
|||
|
|
"gmask_token",
|
|||
|
|
"additional_special_tokens",
|
|||
|
|
]
|
|||
|
|
|
|||
|
|
def __init__(
|
|||
|
|
self,
|
|||
|
|
vocab_file=None,
|
|||
|
|
merges_file=None,
|
|||
|
|
tokenizer_file=None,
|
|||
|
|
clean_up_tokenization_spaces=False,
|
|||
|
|
bos_token="<|startoftext|>",
|
|||
|
|
eos_token="<|endoftext|>",
|
|||
|
|
cls_token="[CLS]",
|
|||
|
|
pad_token="<|endoftext|>",
|
|||
|
|
gmask_token="[gMASK]",
|
|||
|
|
add_bos_token=False,
|
|||
|
|
add_eos_token=False,
|
|||
|
|
**kwargs,
|
|||
|
|
):
|
|||
|
|
self.add_bos_token = add_bos_token
|
|||
|
|
|
|||
|
|
self._gmask_token = (
|
|||
|
|
AddedToken(gmask_token, lstrip=False, rstrip=False, normalized=False)
|
|||
|
|
if isinstance(gmask_token, str)
|
|||
|
|
else gmask_token
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
self._sop_token = (
|
|||
|
|
AddedToken(bos_token, lstrip=False, rstrip=False, normalized=False)
|
|||
|
|
if isinstance(bos_token, str)
|
|||
|
|
else bos_token
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
self._eop_token = (
|
|||
|
|
AddedToken(eos_token, lstrip=False, rstrip=False, normalized=False)
|
|||
|
|
if isinstance(eos_token, str)
|
|||
|
|
else eos_token
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
super().__init__(
|
|||
|
|
vocab_file=vocab_file,
|
|||
|
|
merges_file=merges_file,
|
|||
|
|
tokenizer_file=tokenizer_file,
|
|||
|
|
clean_up_tokenization_spaces=clean_up_tokenization_spaces,
|
|||
|
|
bos_token=bos_token,
|
|||
|
|
eos_token=eos_token,
|
|||
|
|
cls_token=cls_token,
|
|||
|
|
pad_token=pad_token,
|
|||
|
|
gmask_token=gmask_token,
|
|||
|
|
add_bos_token=add_bos_token,
|
|||
|
|
add_eos_token=add_eos_token,
|
|||
|
|
**kwargs,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
self.check_special_tokens()
|
|||
|
|
|
|||
|
|
def check_special_tokens(self):
|
|||
|
|
'''
|
|||
|
|
eos_token, cls_token, mask_token
|
|||
|
|
special tokens should init, check special token is not None
|
|||
|
|
'''
|
|||
|
|
for name, special_token in zip(
|
|||
|
|
['eos', 'bos', 'cls', 'gmask'],
|
|||
|
|
[self.eos_token, self.bos_token, self.cls_token, self.gmask_token],
|
|||
|
|
):
|
|||
|
|
assert special_token is not None, f'should init special token [{name}] in tokenizer_config.json'
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def gmask_token(self) -> Optional[str]:
|
|||
|
|
if self._gmask_token is None:
|
|||
|
|
if self.verbose:
|
|||
|
|
logger.error("Using gmask_token, but it is not set yet.")
|
|||
|
|
return None
|
|||
|
|
return str(self._gmask_token)
|
|||
|
|
|
|||
|
|
@gmask_token.setter
|
|||
|
|
def gmask_token(self, value):
|
|||
|
|
if not isinstance(value, (str, AddedToken)) and value is not None:
|
|||
|
|
raise ValueError("Cannot set a non-string value as the gmask token")
|
|||
|
|
self._gmask_token = value
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def gmask_token_id(self) -> Optional[int]:
|
|||
|
|
if self._gmask_token is None:
|
|||
|
|
return None
|
|||
|
|
return self.convert_tokens_to_ids(self.gmask_token)
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def sop_token(self) -> Optional[str]:
|
|||
|
|
if self._sop_token is None:
|
|||
|
|
if self.verbose:
|
|||
|
|
logger.error("Using sop_token, but it is not set yet.")
|
|||
|
|
return None
|
|||
|
|
return str(self._sop_token)
|
|||
|
|
|
|||
|
|
@sop_token.setter
|
|||
|
|
def sop_token(self, value):
|
|||
|
|
if not isinstance(value, (str, AddedToken)) and value is not None:
|
|||
|
|
raise ValueError("Cannot set a non-string value as the sop token")
|
|||
|
|
self._sop_token = value
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def sop_token_id(self) -> Optional[int]:
|
|||
|
|
if self._sop_token is None:
|
|||
|
|
return None
|
|||
|
|
return self.convert_tokens_to_ids(self.sop_token)
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def eop_token(self) -> Optional[str]:
|
|||
|
|
if self._eop_token is None:
|
|||
|
|
if self.verbose:
|
|||
|
|
logger.error("Using eop_token, but it is not set yet.")
|
|||
|
|
return None
|
|||
|
|
return str(self._eop_token)
|
|||
|
|
|
|||
|
|
@eop_token.setter
|
|||
|
|
def eop_token(self, value):
|
|||
|
|
if not isinstance(value, (str, AddedToken)) and value is not None:
|
|||
|
|
raise ValueError("Cannot set a non-string value as the eop token")
|
|||
|
|
self._eop_token = value
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def eop_token_id(self) -> Optional[int]:
|
|||
|
|
if self._eop_token is None:
|
|||
|
|
return None
|
|||
|
|
return self.convert_tokens_to_ids(self.eop_token)
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def vocab_size(self):
|
|||
|
|
return len(self.get_vocab())
|
|||
|
|
|
|||
|
|
def _chat_from_json(self, chat, chat_format="antglm_chat", system=None):
|
|||
|
|
msgs = chat if "messages" not in chat else chat["messages"]
|
|||
|
|
_msgs = []
|
|||
|
|
sys_msg = None
|
|||
|
|
for msg in msgs:
|
|||
|
|
if is_system(msg):
|
|||
|
|
sys_msg = msg['content']
|
|||
|
|
else:
|
|||
|
|
_msgs.append(msg)
|
|||
|
|
chat = {"messages": _msgs}
|
|||
|
|
system = system or sys_msg
|
|||
|
|
if system:
|
|||
|
|
chat['system_message'] = system
|
|||
|
|
from .chat_format import Chat
|
|||
|
|
|
|||
|
|
return Chat.from_json(chat, name=chat_format)
|
|||
|
|
|
|||
|
|
def apply_chat_template(
|
|||
|
|
self,
|
|||
|
|
conversation: Union[List[Dict[str, str]], List[List[Dict[str, str]]]],
|
|||
|
|
tools: Optional[List[Dict]] = None,
|
|||
|
|
documents: Optional[List[Dict[str, str]]] = None,
|
|||
|
|
chat_template: Optional[str] = None,
|
|||
|
|
add_generation_prompt: bool = False,
|
|||
|
|
system: str = None, # only used for legacy chatml
|
|||
|
|
tokenize=False,
|
|||
|
|
padding: bool = False,
|
|||
|
|
truncation: bool = False,
|
|||
|
|
max_length: Optional[int] = None,
|
|||
|
|
return_tensors: Optional[Union[str, TensorType]] = None,
|
|||
|
|
return_dict: bool = False,
|
|||
|
|
return_assistant_tokens_mask: bool = False,
|
|||
|
|
tokenizer_kwargs: Optional[Dict[str, Any]] = None,
|
|||
|
|
**kwargs,
|
|||
|
|
):
|
|||
|
|
if hasattr(self, "chat_template") and self.chat_template:
|
|||
|
|
if isinstance(conversation, Dict) and "messages" in conversation:
|
|||
|
|
conversation = conversation["messages"]
|
|||
|
|
# use transformers built-in method
|
|||
|
|
return super().apply_chat_template(
|
|||
|
|
conversation=conversation,
|
|||
|
|
tools=tools,
|
|||
|
|
documents=documents,
|
|||
|
|
chat_template=chat_template,
|
|||
|
|
add_generation_prompt=add_generation_prompt,
|
|||
|
|
tokenize=tokenize,
|
|||
|
|
padding=padding,
|
|||
|
|
truncation=truncation,
|
|||
|
|
return_tensors=return_tensors,
|
|||
|
|
return_dict=return_dict,
|
|||
|
|
return_assistant_tokens_mask=return_assistant_tokens_mask,
|
|||
|
|
tokenizer_kwargs=tokenizer_kwargs,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# 非chat_template方式后续将不再支持。
|
|||
|
|
logger.warning("Please set chat_template in tokenizer_config.json!")
|
|||
|
|
|
|||
|
|
chat_format = kwargs.get('chat_format', 'antglm_chat')
|
|||
|
|
|
|||
|
|
is_batched = False
|
|||
|
|
|
|||
|
|
if isinstance(conversation, List) and (
|
|||
|
|
isinstance(conversation[0], (list, tuple)) or "messages" in conversation[0]
|
|||
|
|
):
|
|||
|
|
conversations = conversation
|
|||
|
|
is_batched = True
|
|||
|
|
|
|||
|
|
if not is_batched:
|
|||
|
|
conversations = [conversation]
|
|||
|
|
|
|||
|
|
rendered = []
|
|||
|
|
for chat in conversations:
|
|||
|
|
rendered_chat = self._chat_from_json(chat, chat_format=chat_format, system=system).prompt_str
|
|||
|
|
rendered.append(rendered_chat)
|
|||
|
|
|
|||
|
|
if not is_batched:
|
|||
|
|
rendered = rendered[0]
|
|||
|
|
|
|||
|
|
if tokenize:
|
|||
|
|
out = self(
|
|||
|
|
rendered,
|
|||
|
|
padding=padding,
|
|||
|
|
truncation=truncation,
|
|||
|
|
max_length=max_length,
|
|||
|
|
add_special_tokens=False,
|
|||
|
|
return_tensors=return_tensors,
|
|||
|
|
)
|
|||
|
|
if return_dict:
|
|||
|
|
return out
|
|||
|
|
else:
|
|||
|
|
return out["input_ids"]
|
|||
|
|
else:
|
|||
|
|
return rendered
|
|||
|
|
|
|||
|
|
def _build_position_ids(
|
|||
|
|
self,
|
|||
|
|
mask_pos: int,
|
|||
|
|
bos_pos: int,
|
|||
|
|
max_output_length: int,
|
|||
|
|
rotary_type: Optional[str] = "none",
|
|||
|
|
**kwargs,
|
|||
|
|
) -> List[List[int]]:
|
|||
|
|
window_size = kwargs.get("window_size", 1024) - 1
|
|||
|
|
block_position_ids = [0] * bos_pos
|
|||
|
|
|
|||
|
|
# 获得mask所在的位置，用于后面output positionid的构造
|
|||
|
|
if "1d" in rotary_type:
|
|||
|
|
position_ids = list(range(bos_pos)) + list(range(mask_pos + 1, mask_pos + max_output_length + 2))
|
|||
|
|
block_position_ids = block_position_ids + list(range(1, max_output_length + 2))
|
|||
|
|
elif "2d" in rotary_type:
|
|||
|
|
# 后面input_ids要加一个bos_id
|
|||
|
|
position_ids = list(range(bos_pos))
|
|||
|
|
position_ids = position_ids + [mask_pos] * (1 + max_output_length)
|
|||
|
|
block_position_ids = block_position_ids + list(range(1, max_output_length + 2))
|
|||
|
|
else:
|
|||
|
|
# build position ids
|
|||
|
|
position_ids = []
|
|||
|
|
repeat_times = bos_pos // window_size
|
|||
|
|
for _ in range(repeat_times):
|
|||
|
|
position_ids += list(range(window_size))
|
|||
|
|
position_ids += list(range(bos_pos - window_size * repeat_times))
|
|||
|
|
# need consider additional bos_id after input_ids
|
|||
|
|
mask_pos = position_ids[-1]
|
|||
|
|
position_ids += [mask_pos] * (max_output_length + 1)
|
|||
|
|
|
|||
|
|
block_repeat_times = max_output_length // (window_size - 1)
|
|||
|
|
additional_block_position_ids = []
|
|||
|
|
for _ in range(block_repeat_times):
|
|||
|
|
additional_block_position_ids += list(range(1, window_size))
|
|||
|
|
additional_block_position_ids += list(
|
|||
|
|
range(1, max_output_length + 2 - (window_size - 1) * block_repeat_times)
|
|||
|
|
)
|
|||
|
|
block_position_ids = block_position_ids + additional_block_position_ids
|
|||
|
|
|
|||
|
|
position_ids = [position_ids, block_position_ids]
|
|||
|
|
return position_ids
|
|||
|
|
|
|||
|
|
def _build_inputs_for_generation(
|
|||
|
|
self,
|
|||
|
|
input_ids: List[int],
|
|||
|
|
max_input_length=None,
|
|||
|
|
left_truncate=True,
|
|||
|
|
max_output_length=1024,
|
|||
|
|
rotary_type="none",
|
|||
|
|
unidirectional_attention: bool = True,
|
|||
|
|
attention_dtype=None,
|
|||
|
|
**kwargs,
|
|||
|
|
):
|
|||
|
|
if max_input_length and len(input_ids) > max_input_length:
|
|||
|
|
if left_truncate:
|
|||
|
|
input_ids = input_ids[-max_input_length:]
|
|||
|
|
else:
|
|||
|
|
input_ids = input_ids[:max_input_length]
|
|||
|
|
|
|||
|
|
is_left_padding = input_ids[0] == self.eos_token_id
|
|||
|
|
if not unidirectional_attention:
|
|||
|
|
if input_ids[0] != self.cls_token_id:
|
|||
|
|
input_ids = [self.cls_token_id] + input_ids
|
|||
|
|
|
|||
|
|
if self.gmask_token_id not in set(input_ids):
|
|||
|
|
input_ids = input_ids + [self.gmask_token_id]
|
|||
|
|
|
|||
|
|
mask_pos = input_ids.index(self.gmask_token_id)
|
|||
|
|
sep = len(input_ids)
|
|||
|
|
else:
|
|||
|
|
if self.add_bos_token:
|
|||
|
|
input_ids = input_ids + [self.bos_token_id]
|
|||
|
|
if self.eos_token_id in input_ids:
|
|||
|
|
mask_pos = input_ids.index(self.eos_token_id) - 1
|
|||
|
|
else:
|
|||
|
|
mask_pos = len(input_ids) - 1
|
|||
|
|
sep = len(input_ids) - 1
|
|||
|
|
else:
|
|||
|
|
sep = len(input_ids)
|
|||
|
|
if self.eos_token_id in input_ids:
|
|||
|
|
if is_left_padding:
|
|||
|
|
ori_input_ids = input_ids
|
|||
|
|
input_ids = input_ids[::-1]
|
|||
|
|
mask_pos = input_ids.index(self.eos_token_id) - 1
|
|||
|
|
mask_pos = max(0, mask_pos) # for empty sequence
|
|||
|
|
if is_left_padding:
|
|||
|
|
input_ids = ori_input_ids
|
|||
|
|
mask_pos = sep - 1 - mask_pos # the first non-eos token
|
|||
|
|
|
|||
|
|
else:
|
|||
|
|
mask_pos = len(input_ids) - 1
|
|||
|
|
|
|||
|
|
position_ids = self._build_position_ids(mask_pos, sep, max_output_length, rotary_type, **kwargs)
|
|||
|
|
|
|||
|
|
if is_left_padding:
|
|||
|
|
position_ids[0] = [max(0, i - mask_pos) for i in range(len(position_ids[0]))]
|
|||
|
|
|
|||
|
|
# 后面input_ids要加一个bos_id
|
|||
|
|
total_length = sep + max_output_length
|
|||
|
|
if self.add_bos_token:
|
|||
|
|
total_length += 1
|
|||
|
|
|
|||
|
|
def build_mask_matrix(seq_length, sep, mask_pos, unidirectional_attention):
|
|||
|
|
# 长序列使用bool类型节省显存
|
|||
|
|
if unidirectional_attention:
|
|||
|
|
attention_mask = torch.ones([seq_length, seq_length], dtype=attention_dtype)
|
|||
|
|
attention_mask = torch.tril(attention_mask)
|
|||
|
|
if is_left_padding:
|
|||
|
|
attention_mask[:, :mask_pos] = 0
|
|||
|
|
else:
|
|||
|
|
attention_mask[:, mask_pos + 1 : sep] = 0
|
|||
|
|
else:
|
|||
|
|
attention_mask = torch.zeros([seq_length, seq_length], dtype=attention_dtype)
|
|||
|
|
attention_mask[:, : mask_pos + 1] = 1
|
|||
|
|
for i in range(sep, total_length):
|
|||
|
|
attention_mask[i, sep : i + 1] = 1
|
|||
|
|
return attention_mask
|
|||
|
|
|
|||
|
|
if self.add_bos_token:
|
|||
|
|
attention_mask = build_mask_matrix(total_length, sep + 1, mask_pos, unidirectional_attention)
|
|||
|
|
else:
|
|||
|
|
attention_mask = build_mask_matrix(total_length, sep, mask_pos, unidirectional_attention)
|
|||
|
|
attention_mask = torch.unsqueeze(attention_mask, dim=0)
|
|||
|
|
attention_mask = torch.unsqueeze(attention_mask, dim=1)
|
|||
|
|
if attention_dtype is None:
|
|||
|
|
attention_mask = attention_mask.long()
|
|||
|
|
inputs = {
|
|||
|
|
"input_ids": torch.Tensor([input_ids]).long(),
|
|||
|
|
"position_ids": torch.Tensor([position_ids]).long(),
|
|||
|
|
"attention_mask": attention_mask,
|
|||
|
|
}
|
|||
|
|
return BatchEncoding(inputs)
|
|||
|
|
|
|||
|
|
def build_inputs_for_generation(
|
|||
|
|
self,
|
|||
|
|
input_ids: Union[List[int], List[List[int]], torch.Tensor],
|
|||
|
|
max_input_length=None,
|
|||
|
|
left_truncate=True,
|
|||
|
|
max_output_length=1024,
|
|||
|
|
rotary_type="1d",
|
|||
|
|
unidirectional_attention=True,
|
|||
|
|
attention_dtype=None,
|
|||
|
|
**kwargs,
|
|||
|
|
):
|
|||
|
|
if isinstance(input_ids, torch.Tensor):
|
|||
|
|
input_ids = input_ids.tolist()
|
|||
|
|
|
|||
|
|
if isinstance(input_ids[0], list):
|
|||
|
|
input_ids_list = []
|
|||
|
|
position_ids_list = []
|
|||
|
|
attention_mask_list = []
|
|||
|
|
for _input_ids in input_ids:
|
|||
|
|
inputs = self._build_inputs_for_generation(
|
|||
|
|
_input_ids,
|
|||
|
|
max_input_length=max_input_length,
|
|||
|
|
left_truncate=left_truncate,
|
|||
|
|
max_output_length=max_output_length,
|
|||
|
|
rotary_type=rotary_type,
|
|||
|
|
unidirectional_attention=unidirectional_attention,
|
|||
|
|
attention_dtype=attention_dtype,
|
|||
|
|
**kwargs,
|
|||
|
|
)
|
|||
|
|
input_ids_list.append(inputs['input_ids'])
|
|||
|
|
position_ids_list.append(inputs['position_ids'])
|
|||
|
|
attention_mask_list.append(inputs["attention_mask"])
|
|||
|
|
|
|||
|
|
max_ids_length = max([input.size(1) for input in input_ids_list])
|
|||
|
|
|
|||
|
|
for i in range(len(input_ids)):
|
|||
|
|
cur_ids_length = input_ids_list[i].size(1)
|
|||
|
|
if cur_ids_length < max_ids_length:
|
|||
|
|
# pad input ids
|
|||
|
|
pad_input_ids = input_ids_list[i].new_zeros((1, max_ids_length - cur_ids_length))
|
|||
|
|
input_ids_list[i] = torch.cat([pad_input_ids, input_ids_list[i]], dim=-1)
|
|||
|
|
|
|||
|
|
# pad postition ids with left pad
|
|||
|
|
# 0, 1, 2, 3, 4 ... -> 0, ..., 0, 1, 2, 3, 4, ...
|
|||
|
|
pad_position_ids = input_ids_list[i].new_zeros((1, 2, max_ids_length - cur_ids_length))
|
|||
|
|
position_ids_list[i] = torch.cat([pad_position_ids, position_ids_list[i]], dim=-1)
|
|||
|
|
|
|||
|
|
# pad generation attention mask with left and bottom pad
|
|||
|
|
new_attention_mask = input_ids_list[i].new_zeros(
|
|||
|
|
1,
|
|||
|
|
1,
|
|||
|
|
max_ids_length + max_output_length,
|
|||
|
|
max_ids_length + max_output_length,
|
|||
|
|
)
|
|||
|
|
new_attention_mask[
|
|||
|
|
:,
|
|||
|
|
:,
|
|||
|
|
max_ids_length - cur_ids_length :,
|
|||
|
|
max_ids_length - cur_ids_length :,
|
|||
|
|
] = attention_mask_list[i]
|
|||
|
|
attention_mask_list[i] = new_attention_mask.contiguous()
|
|||
|
|
|
|||
|
|
input_ids_list = torch.cat(input_ids_list, dim=0)
|
|||
|
|
position_ids_list = torch.cat(position_ids_list, dim=0)
|
|||
|
|
attention_mask_list = torch.cat(attention_mask_list, dim=0)
|
|||
|
|
|
|||
|
|
inputs = {
|
|||
|
|
"input_ids": input_ids_list,
|
|||
|
|
"position_ids": position_ids_list,
|
|||
|
|
"attention_mask": attention_mask_list,
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
return BatchEncoding(inputs)
|
|||
|
|
else:
|
|||
|
|
return self._build_inputs_for_generation(
|
|||
|
|
input_ids,
|
|||
|
|
max_input_length=max_input_length,
|
|||
|
|
left_truncate=left_truncate,
|
|||
|
|
max_output_length=max_output_length,
|
|||
|
|
rotary_type=rotary_type,
|
|||
|
|
unidirectional_attention=unidirectional_attention,
|
|||
|
|
**kwargs,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
def _build_inputs_for_train(
|
|||
|
|
self,
|
|||
|
|
inputs: Union[str, List[str]],
|
|||
|
|
outputs: Union[str, List[str]],
|
|||
|
|
new_conversation_offset: List[int] = None,
|
|||
|
|
max_length: int = 2048,
|
|||
|
|
rotary_type: str = "1d",
|
|||
|
|
left_truncate: bool = True,
|
|||
|
|
unidirectional_attention: bool = True,
|
|||
|
|
isolation_position_ids: bool = False,
|
|||
|
|
padding: bool = True,
|
|||
|
|
use_fa2: bool = True,
|
|||
|
|
use_packed: bool = True,
|
|||
|
|
use_baichuan_packed: bool = False,
|
|||
|
|
skip_truncated_turn: bool = False,
|
|||
|
|
return_attention_mask: bool = True,
|
|||
|
|
):
|
|||
|
|
r"""
|
|||
|
|
Build tensor input for model training. If inputs and outputs are list, will pack them.
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
inputs (str, List[str], List[Dict], List[List[Dict]]): the input prompts.
|
|||
|
|
outputs (str, List[str]): the output responses.
|
|||
|
|
max_length (int, Optional): the maximum length of the final input ids for training. Default: 2048
|
|||
|
|
rotary_type (str, Optional): the rotary type of position embedding. Default: 1d
|
|||
|
|
left_truncate (bool, Optional): whether truncate the inputs from left. Default: True
|
|||
|
|
use_fa2 (bool, Optional): whether to build attention mask under flash attention 2.
|
|||
|
|
new_conversation_offset (List[int], Optional): 第idx条样本是全新的对话，[0, 1]代表：inputs[0]和outputs[0]是一个对话，inputs[1]和outputs[1]是一个对话.
|
|||
|
|
"""
|
|||
|
|
if use_packed and use_baichuan_packed and unidirectional_attention:
|
|||
|
|
return self._build_baichuan_inputs_for_train(
|
|||
|
|
inputs,
|
|||
|
|
outputs,
|
|||
|
|
new_conversation_offset,
|
|||
|
|
max_length,
|
|||
|
|
rotary_type,
|
|||
|
|
left_truncate,
|
|||
|
|
skip_truncated_turn,
|
|||
|
|
use_fa2,
|
|||
|
|
padding,
|
|||
|
|
)
|
|||
|
|
if isinstance(inputs, str):
|
|||
|
|
inputs = [inputs]
|
|||
|
|
if isinstance(outputs, str):
|
|||
|
|
outputs = [outputs]
|
|||
|
|
|
|||
|
|
assert len(inputs) == len(outputs)
|
|||
|
|
|
|||
|
|
input_ids = [self(item)['input_ids'] for item in inputs]
|
|||
|
|
output_ids = [self(item)['input_ids'] for item in outputs]
|
|||
|
|
|
|||
|
|
packed_input_ids = []
|
|||
|
|
packed_output_ids = []
|
|||
|
|
if new_conversation_offset is None:
|
|||
|
|
new_conversation_offset = list(range(0, len(inputs)))
|
|||
|
|
assert 0 in new_conversation_offset, f"没有0，请检查new_conversation_offset: {new_conversation_offset}"
|
|||
|
|
current_len = 0
|
|||
|
|
|
|||
|
|
for idx, (input, output) in enumerate(zip(input_ids, output_ids)):
|
|||
|
|
num_special_tokens = 0
|
|||
|
|
if not unidirectional_attention:
|
|||
|
|
if idx in new_conversation_offset:
|
|||
|
|
# cls and gmask
|
|||
|
|
num_special_tokens += 2
|
|||
|
|
else:
|
|||
|
|
# only gmask
|
|||
|
|
num_special_tokens += 1
|
|||
|
|
else:
|
|||
|
|
# sop and eos
|
|||
|
|
if self.add_bos_token:
|
|||
|
|
num_special_tokens += 2
|
|||
|
|
else:
|
|||
|
|
num_special_tokens += 1
|
|||
|
|
|
|||
|
|
# truncate
|
|||
|
|
if len(input) + len(output) + current_len > max_length - num_special_tokens:
|
|||
|
|
if not use_packed or use_fa2 and unidirectional_attention:
|
|||
|
|
attention_mask = torch.tensor(0)
|
|||
|
|
elif use_fa2:
|
|||
|
|
attention_mask = -1 * torch.ones([2, max_length])
|
|||
|
|
else:
|
|||
|
|
attention_mask = torch.tril(torch.ones([max_length, max_length]))
|
|||
|
|
# 返回一个空的样本，该样本不参与训练
|
|||
|
|
default_return = {
|
|||
|
|
'input_ids': (torch.ones(max_length) * self.eos_token_id).long(),
|
|||
|
|
'position_ids': torch.zeros(2, max_length).long(),
|
|||
|
|
'attention_mask': (attention_mask.long()),
|
|||
|
|
'labels': (torch.ones(max_length) * -100).long(),
|
|||
|
|
}
|
|||
|
|
# 如果不截断，直接返回
|
|||
|
|
if skip_truncated_turn:
|
|||
|
|
if current_len == 0:
|
|||
|
|
return default_return
|
|||
|
|
else:
|
|||
|
|
break
|
|||
|
|
left_len = max_length - num_special_tokens - current_len
|
|||
|
|
# 如果截断，只截断prompt
|
|||
|
|
if left_len - len(output) > 0:
|
|||
|
|
if left_truncate:
|
|||
|
|
input = input[-(left_len - len(output)) :]
|
|||
|
|
else:
|
|||
|
|
input = input[: left_len - len(output)]
|
|||
|
|
else:
|
|||
|
|
# response超过left_len，直接返回
|
|||
|
|
if current_len == 0:
|
|||
|
|
return default_return
|
|||
|
|
else:
|
|||
|
|
break
|
|||
|
|
if unidirectional_attention:
|
|||
|
|
packed_input_ids.append(list(input))
|
|||
|
|
else:
|
|||
|
|
if num_special_tokens == 4:
|
|||
|
|
packed_input_ids.append([self.cls_token_id] + list(input) + [self.gmask_token_id])
|
|||
|
|
else:
|
|||
|
|
packed_input_ids.append(list(input) + [self.gmask_token_id])
|
|||
|
|
|
|||
|
|
packed_output_ids.append(list(output) + [self.eos_token_id])
|
|||
|
|
current_len += len(input) + len(output) + num_special_tokens
|
|||
|
|
|
|||
|
|
assert current_len <= max_length
|
|||
|
|
|
|||
|
|
if use_packed:
|
|||
|
|
# pack模式
|
|||
|
|
def build_mask_matrix(seq_length, sep):
|
|||
|
|
# https://github.com/pytorch/pytorch/issues/101932, fix triu/tril bf16 support
|
|||
|
|
m = torch.ones((1, seq_length, seq_length))
|
|||
|
|
mask = torch.arange(1, m.shape[-1] + 1).reshape(1, -1, 1).to(m.device)
|
|||
|
|
ids = torch.arange(1, m.shape[-1] + 1).reshape(1, 1, -1).expand(1, m.shape[-1], -1).to(m.device)
|
|||
|
|
m = (ids <= mask).type_as(m)
|
|||
|
|
|
|||
|
|
m[0, :, : int(sep)] = 1
|
|||
|
|
m = m.squeeze(0)
|
|||
|
|
return m
|
|||
|
|
|
|||
|
|
tokens = []
|
|||
|
|
attention_mask_list = []
|
|||
|
|
input_length_list = []
|
|||
|
|
position_id_list = []
|
|||
|
|
block_position_id_list = []
|
|||
|
|
for input, output in zip(packed_input_ids, packed_output_ids):
|
|||
|
|
if self.add_bos_token:
|
|||
|
|
data = input + [self.sop_token_id] + output
|
|||
|
|
mask_pos = len(input) - 1
|
|||
|
|
else:
|
|||
|
|
data = input + output
|
|||
|
|
mask_pos = len(input) - 2
|
|||
|
|
if return_attention_mask:
|
|||
|
|
if unidirectional_attention:
|
|||
|
|
attention_mask = build_mask_matrix(len(data), 0)
|
|||
|
|
else:
|
|||
|
|
attention_mask = build_mask_matrix(len(data), len(input))
|
|||
|
|
attention_mask = attention_mask.squeeze((0, 1))
|
|||
|
|
|
|||
|
|
attention_mask_list.append(attention_mask)
|
|||
|
|
input_length_list.append(len(input))
|
|||
|
|
tokens += data
|
|||
|
|
|
|||
|
|
sop_pos = mask_pos + 1
|
|||
|
|
position_ids, block_position_ids = self._build_position_ids(
|
|||
|
|
mask_pos=mask_pos, bos_pos=sop_pos, max_output_length=len(output), rotary_type=rotary_type
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
position_id_list.append(position_ids)
|
|||
|
|
block_position_id_list.append(block_position_ids)
|
|||
|
|
|
|||
|
|
labels = []
|
|||
|
|
for i in range(len(packed_input_ids)):
|
|||
|
|
if self.add_bos_token:
|
|||
|
|
labels += [-100] * len(packed_input_ids[i]) + packed_output_ids[i] + [-100]
|
|||
|
|
else:
|
|||
|
|
labels += [-100] * (len(packed_input_ids[i]) - 1) + packed_output_ids[i] + [-100]
|
|||
|
|
|
|||
|
|
total_len = 0
|
|||
|
|
if use_fa2:
|
|||
|
|
pack_attention_mask = -1 * torch.ones([2, current_len])
|
|||
|
|
else:
|
|||
|
|
pack_attention_mask = torch.tril(torch.ones([current_len, current_len]))
|
|||
|
|
|
|||
|
|
pack_position_ids = []
|
|||
|
|
pack_block_position_ids = []
|
|||
|
|
total_len = 0
|
|||
|
|
max_index = 0
|
|||
|
|
for i in range(len(position_id_list)):
|
|||
|
|
|
|||
|
|
if use_fa2:
|
|||
|
|
pack_attention_mask[0][i] = total_len
|
|||
|
|
pack_attention_mask[1][i] = total_len + input_length_list[i]
|
|||
|
|
else:
|
|||
|
|
pack_attention_mask[
|
|||
|
|
total_len : total_len + attention_mask.shape[0],
|
|||
|
|
total_len : total_len + attention_mask.shape[0],
|
|||
|
|
] = attention_mask
|
|||
|
|
position_ids = [pid + max_index for pid in position_id_list[i]]
|
|||
|
|
block_position_ids = block_position_id_list[i]
|
|||
|
|
pack_position_ids.extend(position_ids)
|
|||
|
|
pack_block_position_ids.extend(block_position_ids)
|
|||
|
|
if not isolation_position_ids:
|
|||
|
|
max_index = pack_position_ids[-1] + 1
|
|||
|
|
total_len += len(position_id_list[i])
|
|||
|
|
position_ids = [pack_position_ids, pack_block_position_ids]
|
|||
|
|
else:
|
|||
|
|
# 单输入模式
|
|||
|
|
# 真多轮下，一条样本可能会有好几轮对话，此时需要获取第一条样本的结束位置
|
|||
|
|
if len(new_conversation_offset) > 1:
|
|||
|
|
end_idx = new_conversation_offset[1]
|
|||
|
|
else:
|
|||
|
|
end_idx = 1
|
|||
|
|
input, output = list(itertools.chain(*packed_input_ids[:end_idx])), list(
|
|||
|
|
itertools.chain(*packed_output_ids[:end_idx])
|
|||
|
|
)
|
|||
|
|
if self.add_bos_token:
|
|||
|
|
tokens = input + [self.sop_token_id] + output
|
|||
|
|
else:
|
|||
|
|
tokens = input + output
|
|||
|
|
|
|||
|
|
if self.add_bos_token:
|
|||
|
|
labels = [-100] * len(input) + output + [-100]
|
|||
|
|
position_ids = self._build_position_ids(
|
|||
|
|
mask_pos=len(input) - 1, bos_pos=len(input), max_output_length=len(output), rotary_type=rotary_type
|
|||
|
|
)
|
|||
|
|
else:
|
|||
|
|
labels = [-100] * (len(input) - 1) + output + [-100]
|
|||
|
|
position_ids = self._build_position_ids(
|
|||
|
|
mask_pos=len(input) - 2,
|
|||
|
|
bos_pos=len(input) - 1,
|
|||
|
|
max_output_length=len(output),
|
|||
|
|
rotary_type=rotary_type,
|
|||
|
|
)
|
|||
|
|
attention_mask = len(input)
|
|||
|
|
assert current_len == len(tokens)
|
|||
|
|
|
|||
|
|
# 最大长度补全
|
|||
|
|
if max_length > 0 and len(tokens) < max_length and padding:
|
|||
|
|
pad_length = max_length - len(tokens)
|
|||
|
|
tokens += [self.pad_token_id] * pad_length
|
|||
|
|
labels.extend([-100] * pad_length)
|
|||
|
|
position_ids[0] += [0] * pad_length
|
|||
|
|
position_ids[1] += [0] * pad_length
|
|||
|
|
|
|||
|
|
if use_packed:
|
|||
|
|
if use_fa2:
|
|||
|
|
new_attention_mask = -1 * torch.ones([2, max_length])
|
|||
|
|
new_attention_mask[:, :current_len] = pack_attention_mask
|
|||
|
|
else:
|
|||
|
|
new_attention_mask = torch.tril(torch.ones([max_length, max_length]))
|
|||
|
|
new_attention_mask[:current_len, :current_len] = pack_attention_mask
|
|||
|
|
pack_attention_mask = new_attention_mask.contiguous()
|
|||
|
|
|
|||
|
|
assert len(tokens) == len(labels)
|
|||
|
|
|
|||
|
|
if max_length > 0 and padding:
|
|||
|
|
assert len(tokens) == max_length
|
|||
|
|
|
|||
|
|
if use_fa2 and unidirectional_attention:
|
|||
|
|
# pack_attention_mask = torch.zeros([1], dtype=torch.long)
|
|||
|
|
pack_attention_mask = torch.tensor(0)
|
|||
|
|
|
|||
|
|
if use_packed:
|
|||
|
|
if not use_fa2:
|
|||
|
|
attention_mask = pack_attention_mask.unsqueeze(0).long()
|
|||
|
|
else:
|
|||
|
|
attention_mask = pack_attention_mask
|
|||
|
|
else:
|
|||
|
|
attention_mask = torch.tensor(attention_mask).long()
|
|||
|
|
return {
|
|||
|
|
'input_ids': torch.tensor(tokens).long(),
|
|||
|
|
'position_ids': torch.tensor(position_ids).long(),
|
|||
|
|
'attention_mask': attention_mask,
|
|||
|
|
'labels': torch.tensor(labels).long(),
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
def _build_baichuan_inputs_for_train(
|
|||
|
|
self,
|
|||
|
|
inputs: Union[str, List[str]],
|
|||
|
|
outputs: Union[str, List[str]],
|
|||
|
|
new_conversation_offset: List[int] = None,
|
|||
|
|
max_length: int = 2048,
|
|||
|
|
rotary_type: str = "1d",
|
|||
|
|
left_truncate: bool = True,
|
|||
|
|
skip_truncated_turn: bool = True,
|
|||
|
|
use_fa2: bool = True,
|
|||
|
|
padding: bool = True,
|
|||
|
|
):
|
|||
|
|
'''
|
|||
|
|
input: <role> HUMAN </role> u1 <role> ASSISTANT </role> a11 a12 <role> HUMAN </role> u2 <role> ASSISTANT </role> a21 a22 <|endoftext|> <role> HUMAN </role> u1 <role> ASSISTANT </role> a11 a12 <role> HUMAN </role> u2 <role> ASSISTANT </role> a21 a22 <|endoftext|>
|
|||
|
|
output: x x x x x x a11 a12 <|endoftext|> x x x x x x a21 a22 <|endoftext|> x x x x x x x a11 a12 <|endoftext|> x x x x x x a21 a22 <|endoftext|> x
|
|||
|
|
只适用真多轮+pack数据训练单向模型，需要打开use_true_multiturn
|
|||
|
|
'''
|
|||
|
|
if isinstance(inputs, str):
|
|||
|
|
inputs = [inputs]
|
|||
|
|
if isinstance(outputs, str):
|
|||
|
|
outputs = [outputs]
|
|||
|
|
assert len(inputs) == len(outputs)
|
|||
|
|
|
|||
|
|
input_ids = [self(item)['input_ids'] for item in inputs]
|
|||
|
|
output_ids = [self(item)['input_ids'] for item in outputs]
|
|||
|
|
|
|||
|
|
packed_input_ids = []
|
|||
|
|
packed_output_ids = []
|
|||
|
|
|
|||
|
|
if new_conversation_offset is None:
|
|||
|
|
new_conversation_offset = list(range(0, len(inputs)))
|
|||
|
|
assert 0 in new_conversation_offset, f"没有0，请检查new_conversation_offset: {new_conversation_offset}"
|
|||
|
|
current_len = 0
|
|||
|
|
|
|||
|
|
for idx, (input, output) in enumerate(zip(input_ids, output_ids)):
|
|||
|
|
num_special_tokens = 0
|
|||
|
|
if idx != 0 and idx in new_conversation_offset:
|
|||
|
|
# 在input_ids加入eos，只有第0条样本不加
|
|||
|
|
num_special_tokens += 1
|
|||
|
|
|
|||
|
|
# truncate
|
|||
|
|
if len(input) + len(output) + current_len > max_length - num_special_tokens:
|
|||
|
|
if use_fa2:
|
|||
|
|
attention_mask = torch.tensor(0)
|
|||
|
|
else:
|
|||
|
|
attention_mask = torch.tril(torch.ones([max_length, max_length]))
|
|||
|
|
# 返回一个空的样本，该样本不参与训练
|
|||
|
|
default_return = {
|
|||
|
|
'input_ids': (torch.ones(max_length) * self.eos_token_id).long(),
|
|||
|
|
'position_ids': torch.zeros(2, max_length).long(),
|
|||
|
|
'attention_mask': (attention_mask.long()),
|
|||
|
|
'labels': (torch.ones(max_length) * -100).long(),
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
# 如果不截断，直接返回
|
|||
|
|
if skip_truncated_turn:
|
|||
|
|
if current_len == 0:
|
|||
|
|
return default_return
|
|||
|
|
else:
|
|||
|
|
break
|
|||
|
|
left_len = max_length - num_special_tokens - current_len
|
|||
|
|
# 如果截断，只截断prompt
|
|||
|
|
if left_len - len(output) > 0:
|
|||
|
|
if left_truncate:
|
|||
|
|
input = input[-(left_len - len(output)) :]
|
|||
|
|
else:
|
|||
|
|
input = input[: left_len - len(output)]
|
|||
|
|
else:
|
|||
|
|
# response超过left_len，直接返回
|
|||
|
|
if current_len == 0:
|
|||
|
|
return default_return
|
|||
|
|
else:
|
|||
|
|
break
|
|||
|
|
# 这里拼的是input_ids
|
|||
|
|
if num_special_tokens == 1:
|
|||
|
|
packed_input_ids.append([self.eos_token_id] + list(input))
|
|||
|
|
else:
|
|||
|
|
packed_input_ids.append(list(input))
|
|||
|
|
packed_output_ids.append(list(output))
|
|||
|
|
current_len += len(input) + len(output) + num_special_tokens
|
|||
|
|
assert current_len <= max_length
|
|||
|
|
|
|||
|
|
def build_mask_matrix(seq_length, sep):
|
|||
|
|
# https://github.com/pytorch/pytorch/issues/101932, fix triu/tril bf16 support
|
|||
|
|
m = torch.ones((1, seq_length, seq_length))
|
|||
|
|
mask = torch.arange(1, m.shape[-1] + 1).reshape(1, -1, 1).to(m.device)
|
|||
|
|
ids = torch.arange(1, m.shape[-1] + 1).reshape(1, 1, -1).expand(1, m.shape[-1], -1).to(m.device)
|
|||
|
|
m = (ids <= mask).type_as(m)
|
|||
|
|
|
|||
|
|
m[0, :, : int(sep)] = 1
|
|||
|
|
m = m.squeeze(0)
|
|||
|
|
return m
|
|||
|
|
|
|||
|
|
tokens = []
|
|||
|
|
attention_mask_list = []
|
|||
|
|
position_id_list = []
|
|||
|
|
block_position_id_list = []
|
|||
|
|
token_lens = []
|
|||
|
|
for input, output in zip(packed_input_ids, packed_output_ids):
|
|||
|
|
data = input + output
|
|||
|
|
if not use_fa2:
|
|||
|
|
attention_mask = build_mask_matrix(len(data), 0)
|
|||
|
|
attention_mask_list.append(attention_mask)
|
|||
|
|
tokens += data
|
|||
|
|
token_lens.append(len(data))
|
|||
|
|
|
|||
|
|
position_ids, block_position_ids = self._build_position_ids(
|
|||
|
|
mask_pos=len(input) - 2, bos_pos=len(input) - 1, max_output_length=len(output), rotary_type=rotary_type
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
position_id_list.append(position_ids)
|
|||
|
|
block_position_id_list.append(block_position_ids)
|
|||
|
|
|
|||
|
|
labels = []
|
|||
|
|
for i in range(len(packed_input_ids)):
|
|||
|
|
labels += [-100] * (len(packed_input_ids[i]) - 1) + packed_output_ids[i] + [self.eos_token_id]
|
|||
|
|
|
|||
|
|
total_len = 0
|
|||
|
|
if use_fa2:
|
|||
|
|
pack_attention_mask = torch.Tensor([[0], [1]])
|
|||
|
|
else:
|
|||
|
|
pack_attention_mask = torch.tril(torch.ones([max_length, max_length]))
|
|||
|
|
|
|||
|
|
pack_position_ids = []
|
|||
|
|
pack_block_position_ids = []
|
|||
|
|
total_len = 0
|
|||
|
|
max_index = 0
|
|||
|
|
for i in range(len(token_lens)):
|
|||
|
|
if not use_fa2:
|
|||
|
|
attention_mask = attention_mask_list[i]
|
|||
|
|
pack_attention_mask[
|
|||
|
|
total_len : total_len + attention_mask.shape[0], total_len : total_len + attention_mask.shape[0]
|
|||
|
|
] = attention_mask
|
|||
|
|
position_ids = [pid + max_index for pid in position_id_list[i]]
|
|||
|
|
block_position_ids = block_position_id_list[i]
|
|||
|
|
pack_position_ids.extend(position_ids)
|
|||
|
|
pack_block_position_ids.extend(block_position_ids)
|
|||
|
|
max_index = pack_position_ids[-1] + 1
|
|||
|
|
total_len += token_lens[i]
|
|||
|
|
position_ids = [pack_position_ids, pack_block_position_ids]
|
|||
|
|
|
|||
|
|
if max_length > 0 and len(tokens) < max_length and padding:
|
|||
|
|
pad_length = max_length - len(tokens)
|
|||
|
|
tokens += [self.pad_token_id] * pad_length
|
|||
|
|
labels.extend([-100] * pad_length)
|
|||
|
|
position_ids[0] += [0] * pad_length
|
|||
|
|
position_ids[1] += [0] * pad_length
|
|||
|
|
|
|||
|
|
assert len(tokens) == len(labels)
|
|||
|
|
|
|||
|
|
if not use_fa2:
|
|||
|
|
attention_mask = pack_attention_mask.unsqueeze(0).long()
|
|||
|
|
else:
|
|||
|
|
attention_mask = torch.tensor(0)
|
|||
|
|
return {
|
|||
|
|
'input_ids': torch.tensor(tokens).long(),
|
|||
|
|
'position_ids': torch.tensor(position_ids).long(),
|
|||
|
|
'attention_mask': attention_mask,
|
|||
|
|
'labels': torch.tensor(labels).long(),
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
def build_inputs_for_train(
|
|||
|
|
self,
|
|||
|
|
data: Union[Dict, List[Dict]],
|
|||
|
|
new_conversation_offset: List[int] = None,
|
|||
|
|
chat_format="antglm_chat",
|
|||
|
|
is_chat_format=True, # 如果传入的是字符串，用于说明是否已经是
|
|||
|
|
use_true_multiturn=False,
|
|||
|
|
max_length: int = 2048,
|
|||
|
|
rotary_type: str = "1d",
|
|||
|
|
left_truncate: bool = True,
|
|||
|
|
unidirectional_attention: bool = True,
|
|||
|
|
isolation_position_ids: bool = False,
|
|||
|
|
padding: bool = True,
|
|||
|
|
use_fa2: bool = True,
|
|||
|
|
use_packed: bool = True,
|
|||
|
|
use_baichuan_packed: bool = False,
|
|||
|
|
skip_truncated_turn: bool = False,
|
|||
|
|
return_attention_mask: bool = True,
|
|||
|
|
):
|
|||
|
|
r"""
|
|||
|
|
Build tensor input for model training. If inputs and outputs are list, will pack them.
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
inputs (str, List[str], List[Dict], List[List[Dict]]): the input prompts.
|
|||
|
|
outputs (str, List[str]): the output responses.
|
|||
|
|
new_conversation_offset (List[int]): the offset index of the new conversation turn.
|
|||
|
|
is_chat_format (bool): whether the input is already chatml format
|
|||
|
|
max_length (int, Optional): the maximum length of the final input ids for training. Default: 2048
|
|||
|
|
rotary_type (str, Optional): the rotary type of position embedding. Default: 1d
|
|||
|
|
left_truncate (bool, Optional): whether truncate the inputs from left. Default: True
|
|||
|
|
use_fa2 (bool, Optional): whether to build attention mask under flash attention 2.
|
|||
|
|
"""
|
|||
|
|
if isinstance(data, List):
|
|||
|
|
# chatml list
|
|||
|
|
_inputs = []
|
|||
|
|
_outputs = []
|
|||
|
|
new_conversation_offset = []
|
|||
|
|
for _input in data:
|
|||
|
|
if use_true_multiturn:
|
|||
|
|
chat = self._chat_from_json(_input, chat_format=chat_format)
|
|||
|
|
chat_data = chat.prompt_pack
|
|||
|
|
new_conversation_offset.append(len(_inputs))
|
|||
|
|
_inputs.extend(chat_data['input'])
|
|||
|
|
_outputs.extend(chat_data['output'])
|
|||
|
|
else:
|
|||
|
|
_conversation = _convert_to_conversation(_input)
|
|||
|
|
assert is_assistant(_conversation[-1])
|
|||
|
|
|
|||
|
|
_inputs.append(
|
|||
|
|
self.apply_chat_template(_conversation[:-1], tokenize=False, add_generation_prompt=True)
|
|||
|
|
)
|
|||
|
|
_outputs.append(_conversation[-1]['content'])
|
|||
|
|
|
|||
|
|
return self._build_inputs_for_train(
|
|||
|
|
inputs=_inputs,
|
|||
|
|
outputs=_outputs,
|
|||
|
|
new_conversation_offset=new_conversation_offset,
|
|||
|
|
max_length=max_length,
|
|||
|
|
rotary_type=rotary_type,
|
|||
|
|
left_truncate=left_truncate,
|
|||
|
|
unidirectional_attention=unidirectional_attention,
|
|||
|
|
isolation_position_ids=isolation_position_ids,
|
|||
|
|
padding=padding,
|
|||
|
|
use_fa2=use_fa2,
|
|||
|
|
use_packed=use_packed,
|
|||
|
|
use_baichuan_packed=use_baichuan_packed,
|
|||
|
|
skip_truncated_turn=skip_truncated_turn,
|
|||
|
|
return_attention_mask=return_attention_mask,
|
|||
|
|
)
|
|||
|
|
elif isinstance(data, Dict):
|
|||
|
|
if 'messages' in data:
|
|||
|
|
# chatml format
|
|||
|
|
if use_true_multiturn:
|
|||
|
|
chat = self._chat_from_json(data, chat_format=chat_format)
|
|||
|
|
chat_data = chat.prompt_pack
|
|||
|
|
else:
|
|||
|
|
_conversation = _convert_to_conversation(data)
|
|||
|
|
assert is_assistant(_conversation[-1])
|
|||
|
|
|
|||
|
|
chat_data = {
|
|||
|
|
"input": self.apply_chat_template(
|
|||
|
|
_conversation[:-1], tokenize=False, add_generation_prompt=True
|
|||
|
|
),
|
|||
|
|
"output": _conversation[-1]['content'],
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
return self._build_inputs_for_train(
|
|||
|
|
inputs=chat_data['input'],
|
|||
|
|
outputs=chat_data['output'],
|
|||
|
|
max_length=max_length,
|
|||
|
|
rotary_type=rotary_type,
|
|||
|
|
left_truncate=left_truncate,
|
|||
|
|
unidirectional_attention=unidirectional_attention,
|
|||
|
|
isolation_position_ids=isolation_position_ids,
|
|||
|
|
padding=padding,
|
|||
|
|
use_fa2=use_fa2,
|
|||
|
|
use_packed=use_packed,
|
|||
|
|
use_baichuan_packed=use_baichuan_packed,
|
|||
|
|
skip_truncated_turn=skip_truncated_turn,
|
|||
|
|
return_attention_mask=return_attention_mask,
|
|||
|
|
)
|
|||
|
|
else:
|
|||
|
|
inputs = data['input']
|
|||
|
|
outputs = data['output']
|
|||
|
|
|
|||
|
|
if isinstance(inputs, str):
|
|||
|
|
inputs = [inputs]
|
|||
|
|
if isinstance(outputs, str):
|
|||
|
|
outputs = [outputs]
|
|||
|
|
|
|||
|
|
if not is_chat_format and chat_format:
|
|||
|
|
inputs = [
|
|||
|
|
self.apply_chat_template(
|
|||
|
|
[{"role": "HUMAN", "content": item}], tokenize=False, chat_format=chat_format
|
|||
|
|
)
|
|||
|
|
for item in inputs
|
|||
|
|
]
|
|||
|
|
|
|||
|
|
return self._build_inputs_for_train(
|
|||
|
|
inputs=inputs,
|
|||
|
|
outputs=outputs,
|
|||
|
|
new_conversation_offset=new_conversation_offset,
|
|||
|
|
max_length=max_length,
|
|||
|
|
rotary_type=rotary_type,
|
|||
|
|
left_truncate=left_truncate,
|
|||
|
|
unidirectional_attention=unidirectional_attention,
|
|||
|
|
isolation_position_ids=isolation_position_ids,
|
|||
|
|
padding=padding,
|
|||
|
|
use_fa2=use_fa2,
|
|||
|
|
use_packed=use_packed,
|
|||
|
|
use_baichuan_packed=use_baichuan_packed,
|
|||
|
|
skip_truncated_turn=skip_truncated_turn,
|
|||
|
|
return_attention_mask=return_attention_mask,
|
|||
|
|
)
|