Peaceuai/GLM_model4

1import os2import re3from typing import List, Optional, Union, Dict4from sentencepiece import SentencePieceProcessor5from transformers import PreTrainedTokenizer6from transformers.utils import logging, PaddingStrategy7from transformers.tokenization_utils_base import EncodedInput, BatchEncoding8 9 10class SPTokenizer:11    def __init__(self, model_path: str):12        # reload tokenizer13        assert os.path.isfile(model_path), model_path14        self.sp_model = SentencePieceProcessor(model_file=model_path)15 16        # BOS / EOS token IDs17        self.n_words: int = self.sp_model.vocab_size()18        self.bos_id: int = self.sp_model.bos_id()19        self.eos_id: int = self.sp_model.eos_id()20        self.pad_id: int = self.sp_model.unk_id()21        assert self.sp_model.vocab_size() == self.sp_model.get_piece_size()22 23        special_tokens = ["[MASK]", "[gMASK]", "[sMASK]", "sop", "eop"]24        self.special_tokens = {}25        self.index_special_tokens = {}26        for token in special_tokens:27            self.special_tokens[token] = self.n_words28            self.index_special_tokens[self.n_words] = token29            self.n_words += 130        self.role_special_token_expression = "|".join([re.escape(token) for token in special_tokens]) # for apply_chat_template31 32    def tokenize(self, s: str, encode_special_tokens=False):33        if encode_special_tokens:34            last_index = 035            t = []36            for match in re.finditer(self.role_special_token_expression, s):37                if last_index < match.start():38                    t.extend(self.sp_model.EncodeAsPieces(s[last_index:match.start()]))39                t.append(s[match.start():match.end()])40                last_index = match.end()41            if last_index < len(s):42                t.extend(self.sp_model.EncodeAsPieces(s[last_index:]))43            return t44        else:45            return self.sp_model.EncodeAsPieces(s)46 47    def encode(self, s: str, bos: bool = False, eos: bool = False) -> List[int]:48        assert type(s) is str49        t = self.sp_model.encode(s)50        if bos:51            t = [self.bos_id] + t52        if eos:53            t = t + [self.eos_id]54        return t55 56    def decode(self, t: List[int]) -> str:57        text, buffer = "", []58        for token in t:59            if token in self.index_special_tokens:60                if buffer:61                    text += self.sp_model.decode(buffer)62                    buffer = []63                text += self.index_special_tokens[token]64            else:65                buffer.append(token)66        if buffer:67            text += self.sp_model.decode(buffer)68        return text69 70    def decode_tokens(self, tokens: List[str]) -> str:71        text = self.sp_model.DecodePieces(tokens)72        return text73 74    def convert_token_to_id(self, token):75        """ Converts a token (str) in an id using the vocab. """76        if token in self.special_tokens:77            return self.special_tokens[token]78        return self.sp_model.PieceToId(token)79 80    def convert_id_to_token(self, index):81        """Converts an index (integer) in a token (str) using the vocab."""82        if index in self.index_special_tokens or index in [self.eos_id, self.bos_id, self.pad_id] or index < 0:83            return ""84        return self.sp_model.IdToPiece(index)85 86 87class ChatGLMTokenizer(PreTrainedTokenizer):88    vocab_files_names = {"vocab_file": "tokenizer.model"}89 90    model_input_names = ["input_ids", "attention_mask", "position_ids"]91 92    def __init__(self, vocab_file, padding_side="left", clean_up_tokenization_spaces=False, encode_special_tokens=False, **kwargs):93        self.name = "GLMTokenizer"94 95        self.vocab_file = vocab_file96        self.tokenizer = SPTokenizer(vocab_file)97        self.special_tokens = {98            "<bos>": self.tokenizer.bos_id,99            "<eos>": self.tokenizer.eos_id,100            "<pad>": self.tokenizer.pad_id101        }102        self.encode_special_tokens = encode_special_tokens103        super().__init__(padding_side=padding_side, clean_up_tokenization_spaces=clean_up_tokenization_spaces, **kwargs)104 105    def get_command(self, token):106        if token in self.special_tokens:107            return self.special_tokens[token]108        assert token in self.tokenizer.special_tokens, f"{token} is not a special token for {self.name}"109        return self.tokenizer.special_tokens[token]110 111    @property112    def pad_token(self) -> str:113        return "<unk>"114 115    @property116    def pad_token_id(self):117        return self.get_command("<pad>")118 119    @property120    def eos_token(self) -> str:121        return "</s>"122 123    @property124    def eos_token_id(self):125        return self.get_command("<eos>")126 127    @property128    def vocab_size(self):129        return self.tokenizer.n_words130 131    def get_vocab(self):132        """ Returns vocab as a dict """133        vocab = {self._convert_id_to_token(i): i for i in range(self.vocab_size)}134        vocab.update(self.added_tokens_encoder)135        return vocab136 137    def _tokenize(self, text, **kwargs):138        return self.tokenizer.tokenize(text, encode_special_tokens=self.encode_special_tokens)139 140    def _convert_token_to_id(self, token):141        """ Converts a token (str) in an id using the vocab. """142        return self.tokenizer.convert_token_to_id(token)143 144    def _convert_id_to_token(self, index):145        """Converts an index (integer) in a token (str) using the vocab."""146        return self.tokenizer.convert_id_to_token(index)147 148    def convert_tokens_to_string(self, tokens: List[str]) -> str:149        return self.tokenizer.decode_tokens(tokens)150 151    def save_vocabulary(self, save_directory, filename_prefix=None):152        """153        Save the vocabulary and special tokens file to a directory.154 155        Args:156            save_directory (`str`):157                The directory in which to save the vocabulary.158            filename_prefix (`str`, *optional*):159                An optional prefix to add to the named of the saved files.160 161        Returns:162            `Tuple(str)`: Paths to the files saved.163        """164        if os.path.isdir(save_directory):165            vocab_file = os.path.join(166                save_directory, self.vocab_files_names["vocab_file"]167            )168        else:169            vocab_file = save_directory170 171        with open(self.vocab_file, 'rb') as fin:172            proto_str = fin.read()173 174        with open(vocab_file, "wb") as writer:175            writer.write(proto_str)176 177        return (vocab_file,)178 179    def get_prefix_tokens(self):180        prefix_tokens = [self.get_command("[gMASK]"), self.get_command("sop")]181        return prefix_tokens182 183    def build_prompt(self, query, history=None):184        if history is None:185            history = []186        prompt = ""187        for i, (old_query, response) in enumerate(history):188            prompt += "[Round {}]\n\n问：{}\n\n答：{}\n\n".format(i + 1, old_query, response)189        prompt += "[Round {}]\n\n问：{}\n\n答：".format(len(history) + 1, query)190        return prompt191 192    def build_inputs_with_special_tokens(193            self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None194    ) -> List[int]:195        """196        Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and197        adding special tokens. A BERT sequence has the following format:198 199        - single sequence: `[CLS] X [SEP]`200        - pair of sequences: `[CLS] A [SEP] B [SEP]`201 202        Args:203            token_ids_0 (`List[int]`):204                List of IDs to which the special tokens will be added.205            token_ids_1 (`List[int]`, *optional*):206                Optional second list of IDs for sequence pairs.207 208        Returns:209            `List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens.210        """211        prefix_tokens = self.get_prefix_tokens()212        token_ids_0 = prefix_tokens + token_ids_0213        if token_ids_1 is not None:214            token_ids_0 = token_ids_0 + token_ids_1 + [self.get_command("<eos>")]215        return token_ids_0216 217    def _pad(218            self,219            encoded_inputs: Union[Dict[str, EncodedInput], BatchEncoding],220            max_length: Optional[int] = None,221            padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD,222            pad_to_multiple_of: Optional[int] = None,223            return_attention_mask: Optional[bool] = None,224    ) -> dict:225        """226        Pad encoded inputs (on left/right and up to predefined length or max length in the batch)227 228        Args:229            encoded_inputs:230                Dictionary of tokenized inputs (`List[int]`) or batch of tokenized inputs (`List[List[int]]`).231            max_length: maximum length of the returned list and optionally padding length (see below).232                Will truncate by taking into account the special tokens.233            padding_strategy: PaddingStrategy to use for padding.234 235                - PaddingStrategy.LONGEST Pad to the longest sequence in the batch236                - PaddingStrategy.MAX_LENGTH: Pad to the max length (default)237                - PaddingStrategy.DO_NOT_PAD: Do not pad238                The tokenizer padding sides are defined in self.padding_side:239 240                    - 'left': pads on the left of the sequences241                    - 'right': pads on the right of the sequences242            pad_to_multiple_of: (optional) Integer if set will pad the sequence to a multiple of the provided value.243                This is especially useful to enable the use of Tensor Core on NVIDIA hardware with compute capability244                `>= 7.5` (Volta).245            return_attention_mask:246                (optional) Set to False to avoid returning attention mask (default: set to model specifics)247        """248        # Load from model defaults249        assert self.padding_side == "left"250 251        required_input = encoded_inputs[self.model_input_names[0]]252        seq_length = len(required_input)253 254        if padding_strategy == PaddingStrategy.LONGEST:255            max_length = len(required_input)256 257        if max_length is not None and pad_to_multiple_of is not None and (max_length % pad_to_multiple_of != 0):258            max_length = ((max_length // pad_to_multiple_of) + 1) * pad_to_multiple_of259 260        needs_to_be_padded = padding_strategy != PaddingStrategy.DO_NOT_PAD and len(required_input) != max_length261 262        # Initialize attention mask if not present.263        if "attention_mask" not in encoded_inputs:264            encoded_inputs["attention_mask"] = [1] * seq_length265 266        if "position_ids" not in encoded_inputs:267            encoded_inputs["position_ids"] = list(range(seq_length))268 269        if needs_to_be_padded:270            difference = max_length - len(required_input)271 272            if "attention_mask" in encoded_inputs:273                encoded_inputs["attention_mask"] = [0] * difference + encoded_inputs["attention_mask"]274            if "position_ids" in encoded_inputs:275                encoded_inputs["position_ids"] = [0] * difference + encoded_inputs["position_ids"]276            encoded_inputs[self.model_input_names[0]] = [self.pad_token_id] * difference + required_input277 278        return encoded_inputs279