askyer/chatglm3-ft1

1import json2import os3import re4from typing import List, Optional, Union, Dict5from sentencepiece import SentencePieceProcessor6from transformers import PreTrainedTokenizer7from transformers.utils import logging, PaddingStrategy8from transformers.tokenization_utils_base import EncodedInput, BatchEncoding9 10 11logger = logging.get_logger(__name__)12 13 14class SPTokenizer:15    def __init__(self, model_path: str):16        # reload tokenizer17        assert os.path.isfile(model_path), model_path18        self.sp_model = SentencePieceProcessor(model_file=model_path)19 20        # BOS / EOS token IDs21        self.n_words: int = self.sp_model.vocab_size()22        self.bos_id: int = self.sp_model.bos_id()23        self.eos_id: int = self.sp_model.eos_id()24        self.pad_id: int = self.sp_model.unk_id()25        assert self.sp_model.vocab_size() == self.sp_model.get_piece_size()26 27        role_special_tokens = ["<|system|>", "<|user|>", "<|assistant|>", "<|observation|>"]28        special_tokens = ["[MASK]", "[gMASK]", "[sMASK]", "sop", "eop"] + role_special_tokens29        self.special_tokens = {}30        self.index_special_tokens = {}31        for token in special_tokens:32            self.special_tokens[token] = self.n_words33            self.index_special_tokens[self.n_words] = token34            self.n_words += 135        self.role_special_token_expression = "|".join([re.escape(token) for token in special_tokens]) # for apply_chat_template36 37    def tokenize(self, s: str, encode_special_tokens=False):38        if encode_special_tokens:39            last_index = 040            t = []41            for match in re.finditer(self.role_special_token_expression, s):42                if last_index < match.start():43                    t.extend(self.sp_model.EncodeAsPieces(s[last_index:match.start()]))44                t.append(s[match.start():match.end()])45                last_index = match.end()46            if last_index < len(s):47                t.extend(self.sp_model.EncodeAsPieces(s[last_index:]))48            return t49        else:50            return self.sp_model.EncodeAsPieces(s)51 52    def encode(self, s: str, bos: bool = False, eos: bool = False) -> List[int]:53        assert type(s) is str54        t = self.sp_model.encode(s)55        if bos:56            t = [self.bos_id] + t57        if eos:58            t = t + [self.eos_id]59        return t60 61    def decode(self, t: List[int]) -> str:62        text, buffer = "", []63        for token in t:64            if token in self.index_special_tokens:65                if buffer:66                    text += self.sp_model.decode(buffer)67                    buffer = []68                text += self.index_special_tokens[token]69            else:70                buffer.append(token)71        if buffer:72            text += self.sp_model.decode(buffer)73        return text74 75    def decode_tokens(self, tokens: List[str]) -> str:76        text = self.sp_model.DecodePieces(tokens)77        return text78 79    def convert_token_to_id(self, token):80        """ Converts a token (str) in an id using the vocab. """81        if token in self.special_tokens:82            return self.special_tokens[token]83        return self.sp_model.PieceToId(token)84 85    def convert_id_to_token(self, index):86        """Converts an index (integer) in a token (str) using the vocab."""87        if index in self.index_special_tokens:88            return self.index_special_tokens[index]89        if index in [self.eos_id, self.bos_id, self.pad_id] or index < 0 or index > self.sp_model.vocab_size():90            return ""91        return self.sp_model.IdToPiece(index)92 93 94class ChatGLMTokenizer(PreTrainedTokenizer):95 96    vocab_files_names = {"vocab_file": "tokenizer.model"}97    model_input_names = ["input_ids", "attention_mask", "position_ids"]98 99    def __init__(100        self,101        vocab_file,102        padding_side="left",103        clean_up_tokenization_spaces=False,104        encode_special_tokens=False,105        **kwargs106    ):107        self.name = "GLMTokenizer"108        self.vocab_file = vocab_file109        self.tokenizer = SPTokenizer(vocab_file)110        self.special_tokens = {111            "<bos>": self.tokenizer.bos_id,112            "<eos>": self.tokenizer.eos_id,113            "<unk>": self.tokenizer.pad_id,114            "<pad>": self.tokenizer.pad_id115        }116        self.encode_special_tokens = encode_special_tokens117 118        super().__init__(119            padding_side=padding_side,120            clean_up_tokenization_spaces=clean_up_tokenization_spaces,121            **kwargs122        )123 124    def get_command(self, token):125        if token in self.special_tokens:126            return self.special_tokens[token]127        assert token in self.tokenizer.special_tokens, f"{token} is not a special token for {self.name}"128        return self.tokenizer.special_tokens[token]129 130    @property131    def unk_token(self) -> str:132        return self.tokenizer.sp_model.IdToPiece(self.get_command("<unk>"))133 134    @property135    def pad_token(self) -> str:136        return self.tokenizer.sp_model.IdToPiece(self.get_command("<pad>"))137 138    @property139    def eos_token(self) -> str:140        return self.tokenizer.sp_model.IdToPiece(self.get_command("<eos>"))141 142    @property143    def unk_token_id(self) -> int:144        return self.get_command("<unk>")145 146    @property147    def pad_token_id(self) -> int:148        return self.get_command("<pad>")149 150    @property151    def eos_token_id(self):152        return self.get_command("<eos>")153 154    @unk_token.setter155    def unk_token(self, value):156        logger.warning("Setting unk_token is not supported, use the default one.")157 158    @pad_token.setter159    def pad_token(self, value):160        logger.warning("Setting pad_token is not supported, use the default one.")161 162    @eos_token.setter163    def eos_token(self, value):164        logger.warning("Setting eos_token is not supported, use the default one.")165 166    @property167    def vocab_size(self):168        return self.tokenizer.n_words169 170    def get_vocab(self):171        """ Returns vocab as a dict """172        vocab = {self._convert_id_to_token(i): i for i in range(self.vocab_size)}173        vocab.update(self.added_tokens_encoder)174        return vocab175 176    def _tokenize(self, text, **kwargs):177        return self.tokenizer.tokenize(text, encode_special_tokens=self.encode_special_tokens)178 179    def _convert_token_to_id(self, token):180        """ Converts a token (str) in an id using the vocab. """181        return self.tokenizer.convert_token_to_id(token)182 183    def _convert_id_to_token(self, index):184        """Converts an index (integer) in a token (str) using the vocab."""185        return self.tokenizer.convert_id_to_token(index)186 187    def convert_tokens_to_string(self, tokens: List[str]) -> str:188        return self.tokenizer.decode_tokens(tokens)189 190    def save_vocabulary(self, save_directory, filename_prefix=None):191        """192        Save the vocabulary and special tokens file to a directory.193 194        Args:195            save_directory (`str`):196                The directory in which to save the vocabulary.197            filename_prefix (`str`, *optional*):198                An optional prefix to add to the named of the saved files.199 200        Returns:201            `Tuple(str)`: Paths to the files saved.202        """203        if os.path.isdir(save_directory):204            vocab_file = os.path.join(205                save_directory, self.vocab_files_names["vocab_file"]206            )207        else:208            vocab_file = save_directory209 210        with open(self.vocab_file, 'rb') as fin:211            proto_str = fin.read()212 213        with open(vocab_file, "wb") as writer:214            writer.write(proto_str)215 216        return (vocab_file,)217 218    def get_prefix_tokens(self):219        prefix_tokens = [self.get_command("[gMASK]"), self.get_command("sop")]220        return prefix_tokens221 222    def build_single_message(self, role, metadata, message):223        assert role in ["system", "user", "assistant", "observation"], role224        role_tokens = [self.get_command(f"<|{role}|>")] + self.tokenizer.encode(f"{metadata}\n")225        message_tokens = self.tokenizer.encode(message)226        tokens = role_tokens + message_tokens227        return tokens228 229    def build_chat_input(self, query, history=None, role="user"):230        if history is None:231            history = []232        input_ids = []233        for item in history:234            content = item["content"]235            if item["role"] == "system" and "tools" in item:236                content = content + "\n" + json.dumps(item["tools"], indent=4, ensure_ascii=False)237            input_ids.extend(self.build_single_message(item["role"], item.get("metadata", ""), content))238        input_ids.extend(self.build_single_message(role, "", query))239        input_ids.extend([self.get_command("<|assistant|>")])240        return self.batch_encode_plus([input_ids], return_tensors="pt", is_split_into_words=True)241 242    def build_inputs_with_special_tokens(243        self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None244    ) -> List[int]:245        """246        Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and247        adding special tokens. A BERT sequence has the following format:248 249        - single sequence: `[CLS] X [SEP]`250        - pair of sequences: `[CLS] A [SEP] B [SEP]`251 252        Args:253            token_ids_0 (`List[int]`):254                List of IDs to which the special tokens will be added.255            token_ids_1 (`List[int]`, *optional*):256                Optional second list of IDs for sequence pairs.257 258        Returns:259            `List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens.260        """261        prefix_tokens = self.get_prefix_tokens()262        token_ids_0 = prefix_tokens + token_ids_0263        if token_ids_1 is not None:264            token_ids_0 = token_ids_0 + token_ids_1 + [self.get_command("<eos>")]265        return token_ids_0266 267    def _pad(268        self,269        encoded_inputs: Union[Dict[str, EncodedInput], BatchEncoding],270        max_length: Optional[int] = None,271        padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD,272        pad_to_multiple_of: Optional[int] = None,273        return_attention_mask: Optional[bool] = None,274    ) -> dict:275        """276        Pad encoded inputs (on left/right and up to predefined length or max length in the batch)277 278        Args:279            encoded_inputs:280                Dictionary of tokenized inputs (`List[int]`) or batch of tokenized inputs (`List[List[int]]`).281            max_length: maximum length of the returned list and optionally padding length (see below).282                Will truncate by taking into account the special tokens.283            padding_strategy: PaddingStrategy to use for padding.284 285                - PaddingStrategy.LONGEST Pad to the longest sequence in the batch286                - PaddingStrategy.MAX_LENGTH: Pad to the max length (default)287                - PaddingStrategy.DO_NOT_PAD: Do not pad288                The tokenizer padding sides are defined in self.padding_side:289 290                    - 'left': pads on the left of the sequences291                    - 'right': pads on the right of the sequences292            pad_to_multiple_of: (optional) Integer if set will pad the sequence to a multiple of the provided value.293                This is especially useful to enable the use of Tensor Core on NVIDIA hardware with compute capability294                `>= 7.5` (Volta).295            return_attention_mask:296                (optional) Set to False to avoid returning attention mask (default: set to model specifics)297        """298        # Load from model defaults299        assert self.padding_side == "left"300 301        required_input = encoded_inputs[self.model_input_names[0]]302        seq_length = len(required_input)303 304        if padding_strategy == PaddingStrategy.LONGEST:305            max_length = len(required_input)306 307        if max_length is not None and pad_to_multiple_of is not None and (max_length % pad_to_multiple_of != 0):308            max_length = ((max_length // pad_to_multiple_of) + 1) * pad_to_multiple_of309 310        needs_to_be_padded = padding_strategy != PaddingStrategy.DO_NOT_PAD and len(required_input) != max_length311 312        # Initialize attention mask if not present.313        if "attention_mask" not in encoded_inputs:314            encoded_inputs["attention_mask"] = [1] * seq_length315 316        if "position_ids" not in encoded_inputs:317            encoded_inputs["position_ids"] = list(range(seq_length))318 319        if needs_to_be_padded:320            difference = max_length - len(required_input)321 322            if "attention_mask" in encoded_inputs:323                encoded_inputs["attention_mask"] = [0] * difference + encoded_inputs["attention_mask"]324            if "position_ids" in encoded_inputs:325                encoded_inputs["position_ids"] = [0] * difference + encoded_inputs["position_ids"]326            encoded_inputs[self.model_input_names[0]] = [self.pad_token_id] * difference + required_input327 328        return encoded_inputs329