Ming-3.0-asr-tiny / tokenization_bailing.py
Publish Ming-3.0 ASR Tiny model package and bilingual documentation
2c656bd verified 10 days ago
45.6 kB
#!/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._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,
        )

        if not hasattr(self, "add_bos_token"):
            self.add_bos_token = add_bos_token

        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,
                max_length=max_length,
                return_tensors=return_tensors,
                return_dict=return_dict,
                return_assistant_tokens_mask=return_assistant_tokens_mask,
                tokenizer_kwargs=tokenizer_kwargs,
                **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,
                )