Add `ChatGLMTokenizerFast` and `ChatGLMTokenizerConverter`

tokenization_chatglm.py

@@ -1,13 +1,39 @@
 import json
 import os
-import torch
+import re
 from typing import List, Optional, Union, Dict
 from sentencepiece import SentencePieceProcessor
-from transformers import PreTrainedTokenizer
+from transformers import AddedToken, PreTrainedTokenizer, PreTrainedTokenizerFast
+from transformers.convert_slow_tokenizer import (
+    SLOW_TO_FAST_CONVERTERS,
+    SpmConverter,
+    decoders,
+    normalizers,
+    pre_tokenizers,
+    processors,
+)
 from transformers.utils import logging, PaddingStrategy
 from transformers.tokenization_utils_base import EncodedInput, BatchEncoding
 
 
+logger = logging.get_logger(__name__)
+
+ADDITIONAL_SPECIAL_TOKENS = [
+    "[MASK]",
+    "[gMASK]",
+    "[sMASK]",
+    "<!sop!>",
+    "<!eop!>",
+    "<|system|>",
+    "<|user|>",
+    "<|assistant|>",
+    "<|observation|>",
+]
+PREFIX_TOKENS = ["[gMASK]", "<!sop!>"]
+
+DUMMY_PREFIX_INDICATOR_FOR_FAST = "<!dummy-prefix!>"
+
+
 class SPTokenizer:
     def __init__(self, model_path: str):
         # reload tokenizer
@@ -21,17 +47,29 @@ class SPTokenizer:
         self.pad_id: int = self.sp_model.unk_id()
         assert self.sp_model.vocab_size() == self.sp_model.get_piece_size()
 
-        special_tokens = ["[MASK]", "[gMASK]", "[sMASK]", "sop", "eop", "<|system|>", "<|user|>", "<|assistant|>",
-                          "<|observation|>"]
+        special_tokens = ADDITIONAL_SPECIAL_TOKENS
         self.special_tokens = {}
         self.index_special_tokens = {}
         for token in special_tokens:
             self.special_tokens[token] = self.n_words
             self.index_special_tokens[self.n_words] = token
             self.n_words += 1
-
-    def tokenize(self, s: str):
-        return self.sp_model.EncodeAsPieces(s)
+        self.role_special_token_expression = "|".join([re.escape(token) for token in special_tokens]) # for apply_chat_template
+
+    def tokenize(self, s: str, encode_special_tokens=False):
+        if encode_special_tokens:
+            last_index = 0
+            t = []
+            for match in re.finditer(self.role_special_token_expression, s):
+                if last_index < match.start():
+                    t.extend(self.sp_model.EncodeAsPieces(s[last_index:match.start()]))
+                t.append(s[match.start():match.end()])
+                last_index = match.end()
+            if last_index < len(s):
+                t.extend(self.sp_model.EncodeAsPieces(s[last_index:]))
+            return t
+        else:
+            return self.sp_model.EncodeAsPieces(s)
 
     def encode(self, s: str, bos: bool = False, eos: bool = False) -> List[int]:
         assert type(s) is str
@@ -70,27 +108,40 @@ class SPTokenizer:
         """Converts an index (integer) in a token (str) using the vocab."""
         if index in self.index_special_tokens:
             return self.index_special_tokens[index]
-        if index in [self.eos_id, self.bos_id, self.pad_id] or index < 0:
+        if index in [self.eos_id, self.bos_id, self.pad_id] or index < 0 or index >= self.sp_model.vocab_size():
             return ""
         return self.sp_model.IdToPiece(index)
 
 
 class ChatGLMTokenizer(PreTrainedTokenizer):
-    vocab_files_names = {"vocab_file": "tokenizer.model"}
 
+    vocab_files_names = {"vocab_file": "tokenizer.model"}
     model_input_names = ["input_ids", "attention_mask", "position_ids"]
 
-    def __init__(self, vocab_file, padding_side="left", clean_up_tokenization_spaces=False, **kwargs):
+    def __init__(
+        self,
+        vocab_file,
+        padding_side="left",
+        clean_up_tokenization_spaces=False,
+        encode_special_tokens=False,
+        **kwargs
+    ):
         self.name = "GLMTokenizer"
-
         self.vocab_file = vocab_file
         self.tokenizer = SPTokenizer(vocab_file)
         self.special_tokens = {
             "<bos>": self.tokenizer.bos_id,
             "<eos>": self.tokenizer.eos_id,
+            "<unk>": self.tokenizer.pad_id,
             "<pad>": self.tokenizer.pad_id
         }
-        super().__init__(padding_side=padding_side, clean_up_tokenization_spaces=clean_up_tokenization_spaces, **kwargs)
+        self.encode_special_tokens = encode_special_tokens
+
+        super().__init__(
+            padding_side=padding_side,
+            clean_up_tokenization_spaces=clean_up_tokenization_spaces,
+            **kwargs
+        )
 
     def get_command(self, token):
         if token in self.special_tokens:
@@ -100,24 +151,40 @@ class ChatGLMTokenizer(PreTrainedTokenizer):
 
     @property
     def unk_token(self) -> str:
-        return "<unk>"
+        return self.tokenizer.sp_model.IdToPiece(self.get_command("<unk>"))
 
     @property
     def pad_token(self) -> str:
-        return "<unk>"
+        return self.tokenizer.sp_model.IdToPiece(self.get_command("<pad>"))
 
     @property
-    def pad_token_id(self):
-        return self.get_command("<pad>")
+    def eos_token(self) -> str:
+        return self.tokenizer.sp_model.IdToPiece(self.get_command("<eos>"))
 
     @property
-    def eos_token(self) -> str:
-        return "</s>"
+    def unk_token_id(self) -> int:
+        return self.get_command("<unk>")
+
+    @property
+    def pad_token_id(self) -> int:
+        return self.get_command("<pad>")
 
     @property
     def eos_token_id(self):
         return self.get_command("<eos>")
 
+    @unk_token.setter
+    def unk_token(self, value):
+        logger.warning("Setting unk_token is not supported, use the default one.")
+
+    @pad_token.setter
+    def pad_token(self, value):
+        logger.warning("Setting pad_token is not supported, use the default one.")
+
+    @eos_token.setter
+    def eos_token(self, value):
+        logger.warning("Setting eos_token is not supported, use the default one.")
+
     @property
     def vocab_size(self):
         return self.tokenizer.n_words
@@ -129,7 +196,7 @@ class ChatGLMTokenizer(PreTrainedTokenizer):
         return vocab
 
     def _tokenize(self, text, **kwargs):
-        return self.tokenizer.tokenize(text)
+        return self.tokenizer.tokenize(text, encode_special_tokens=self.encode_special_tokens)
 
     def _convert_token_to_id(self, token):
         """ Converts a token (str) in an id using the vocab. """
@@ -171,8 +238,7 @@ class ChatGLMTokenizer(PreTrainedTokenizer):
         return (vocab_file,)
 
     def get_prefix_tokens(self):
-        prefix_tokens = [self.get_command("[gMASK]"), self.get_command("sop")]
-        return prefix_tokens
+        return list(map(self.get_command, PREFIX_TOKENS))
 
     def build_single_message(self, role, metadata, message):
         assert role in ["system", "user", "assistant", "observation"], role
@@ -195,7 +261,7 @@ class ChatGLMTokenizer(PreTrainedTokenizer):
         return self.batch_encode_plus([input_ids], return_tensors="pt", is_split_into_words=True)
 
     def build_inputs_with_special_tokens(
-            self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
+        self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
     ) -> List[int]:
         """
         Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
@@ -220,12 +286,12 @@ class ChatGLMTokenizer(PreTrainedTokenizer):
         return token_ids_0
 
     def _pad(
-            self,
-            encoded_inputs: Union[Dict[str, EncodedInput], BatchEncoding],
-            max_length: Optional[int] = None,
-            padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD,
-            pad_to_multiple_of: Optional[int] = None,
-            return_attention_mask: Optional[bool] = None,
+        self,
+        encoded_inputs: Union[Dict[str, EncodedInput], BatchEncoding],
+        max_length: Optional[int] = None,
+        padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD,
+        pad_to_multiple_of: Optional[int] = None,
+        return_attention_mask: Optional[bool] = None,
     ) -> dict:
         """
         Pad encoded inputs (on left/right and up to predefined length or max length in the batch)
@@ -281,3 +347,160 @@ class ChatGLMTokenizer(PreTrainedTokenizer):
             encoded_inputs[self.model_input_names[0]] = [self.pad_token_id] * difference + required_input
 
         return encoded_inputs
+
+
+class ChatGLMTokenizerFast(PreTrainedTokenizerFast):
+    # multiple breaking changes, no backward-compatibility
+    slow_tokenizer_class = ChatGLMTokenizer
+    vocab_files_names = {
+        **ChatGLMTokenizer.vocab_files_names,
+        **PreTrainedTokenizerFast.vocab_files_names,
+    }
+
+    def __init__(self, **kwargs):
+        kwargs.setdefault("clean_up_tokenization_spaces", False)
+        kwargs.setdefault("bos_token", "<s>")
+        kwargs.setdefault("eos_token", "</s>")
+        kwargs.setdefault("unk_token", "<unk>")
+        kwargs.setdefault("pad_token", "<unk>")
+        super().__init__(**kwargs)
+
+    @property
+    def dummy_prefix_indicator(self):
+        return DUMMY_PREFIX_INDICATOR_FOR_FAST
+
+    @property
+    def can_save_slow_tokenizer(self) -> bool:
+        # multiple breaking changes
+        return False
+
+    def save_pretrained(self, *args, **kwargs):
+        if not self.can_save_slow_tokenizer:
+            logger.warning(
+                f"{type(self).__name__} does not support saving slow tokenizer. "
+                "Saving it at the same directory may break the original tokenizer. "
+                "Please keep a backup beforehand."
+            )
+
+        return super().save_pretrained(*args, **kwargs)
+
+    def build_single_message_prompt(self, role, metadata, message):
+        assert role in ["system", "user", "assistant", "observation"], role
+        return (
+            f"<|{role}|>"
+            f"{self.dummy_prefix_indicator}{metadata}\n"
+            f"{self.dummy_prefix_indicator}{message}"
+        )
+
+    def build_chat_prompt(self, query, history=None, role="user", metadata=""):
+        inputs = []
+
+        for item in history or []:
+            content = item["content"]
+
+            if item["role"] == "system" and "tools" in item:
+                content += "\n" + json.dumps(
+                    item["tools"], indent=4, ensure_ascii=False
+                )
+
+            inputs.append(
+                self.build_single_message_prompt(
+                    item["role"], item.get("metadata", ""), content
+                )
+            )
+
+        inputs.append(self.build_single_message_prompt(role, metadata, query))
+        inputs.append("<|assistant|>")
+
+        return "".join(inputs)
+
+    def build_chat_input(self, *args, **kwargs):
+        return self.batch_encode_plus(
+            [self.build_chat_prompt(*args, **kwargs)],
+            return_tensors="pt",
+        )
+
+
+ChatGLMTokenizer.register_for_auto_class()
+ChatGLMTokenizerFast.register_for_auto_class()
+
+
+class ChatGLMTokenizerConverter(SpmConverter):
+    handle_byte_fallback = True
+
+    def normalizer(self, proto):
+        return normalizers.Sequence(
+            [
+                normalizers.Replace(
+                    pattern=DUMMY_PREFIX_INDICATOR_FOR_FAST, content="▁"
+                ),
+                normalizers.Replace(pattern=" ", content="▁"),
+            ]
+        )
+
+    def pre_tokenizer(self, replacement, add_prefix_space):
+        # NOTE: don't use Metaspace, it won't merge spaces into one token
+        # without Metaspace: "  " => ["▁▁"]
+        # with Metaspace: "  " => ["▁", "▁"]
+        return pre_tokenizers.Split(DUMMY_PREFIX_INDICATOR_FOR_FAST, "merged_with_next")
+
+    def decoder(self, replacement, add_prefix_space):
+        return decoders.Sequence(
+            [
+                decoders.ByteFallback(),
+                decoders.Metaspace(replacement="▁", add_prefix_space=True),
+            ]
+        )
+
+    def tokenizer(self, proto):
+        tokenizer = super().tokenizer(proto)
+
+        tokenizer.model.byte_fallback = True
+
+        assert tokenizer.token_to_id("<unk>") == 0
+        assert tokenizer.token_to_id("<s>") == 1
+        assert tokenizer.token_to_id("</s>") == 2
+        special_tokens = [
+            "<unk>",
+            "<s>",
+            "</s>",
+            *ADDITIONAL_SPECIAL_TOKENS,
+        ]
+
+        tokenizer.add_special_tokens(
+            [AddedToken(token, special=True) for token in special_tokens]
+        )
+
+        return tokenizer
+
+    def converted(self):
+        tokenizer = super().converted()
+
+        # Post processors
+        prefix_token_ids = list(map(tokenizer.token_to_id, PREFIX_TOKENS))
+        assert all(i is not None for i in prefix_token_ids)
+        prefix_template = " ".join(PREFIX_TOKENS)
+
+        template_special_tokens = list(frozenset(zip(PREFIX_TOKENS, prefix_token_ids)))
+
+        if "</s>" not in PREFIX_TOKENS:
+            eos_token_id = tokenizer.token_to_id("</s>")
+            assert eos_token_id is not None
+            template_special_tokens.append(("</s>", eos_token_id))
+
+        post = processors.TemplateProcessing(
+            single=f"{prefix_template} $A",
+            pair=f"{prefix_template} $A $B:1 </s>:1",
+            special_tokens=template_special_tokens,
+        )
+        if tokenizer.post_processor is None:
+            tokenizer.post_processor = post
+        else:
+            tokenizer.post_processor = processors.Sequence(
+                [tokenizer.post_processor, post]
+            )
+
+        return tokenizer
+
+
+SLOW_TO_FAST_CONVERTERS[ChatGLMTokenizer.__name__] = ChatGLMTokenizerConverter

tokenizer_config.json

@@ -7,7 +7,7 @@
   "auto_map": {
     "AutoTokenizer": [
       "tokenization_chatglm.ChatGLMTokenizer",
-      null
-      ]
+      "tokenization_chatglm.ChatGLMTokenizerFast"
+    ]
   }
 }