DRAFT: Add a fast tokenizer implementation and converter

tokenization_chatglm.py

@@ -1,11 +1,37 @@
 import json
 import os
-import re
-from typing import List, Optional, Union, Dict
+import warnings
+from typing import Dict, List, Optional, Tuple, Union
+
 from sentencepiece import SentencePieceProcessor
-from transformers import PreTrainedTokenizer
-from transformers.utils import logging, PaddingStrategy
-from transformers.tokenization_utils_base import EncodedInput, BatchEncoding
+from tokenizers import AddedToken, decoders, normalizers, processors
+from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast
+from transformers.convert_slow_tokenizer import SLOW_TO_FAST_CONVERTERS, SpmConverter
+from transformers.tokenization_utils_base import (
+    BatchEncoding,
+    EncodedInput,
+    PreTokenizedInput,
+    PreTokenizedInputPair,
+    TextInput,
+    TextInputPair,
+    TruncationStrategy,
+)
+from transformers.utils import PaddingStrategy
+
+ADDITIONAL_SPECIAL_TOKENS = [
+    "[MASK]",
+    "[gMASK]",
+    "[sMASK]",
+    "<!sop!>",
+    "<!eop!>",
+    "<|system|>",
+    "<|user|>",
+    "<|assistant|>",
+    "<|observation|>",
+]
+PREFIX_TOKENS = ["[gMASK]", "<!sop!>"]
+
+ENCODE_SEP_TOKEN_FOR_FAST = "<!encode-sep!>"
 
 
 class SPTokenizer:
@@ -21,30 +47,15 @@ class SPTokenizer:
         self.pad_id: int = self.sp_model.unk_id()
         assert self.sp_model.vocab_size() == self.sp_model.get_piece_size()
 
-        role_special_tokens = ["<|system|>", "<|user|>", "<|assistant|>", "<|observation|>"]
-        special_tokens = ["[MASK]", "[gMASK]", "[sMASK]", "sop", "eop"] + role_special_tokens
         self.special_tokens = {}
         self.index_special_tokens = {}
-        for token in special_tokens:
+        for token in ADDITIONAL_SPECIAL_TOKENS:
             self.special_tokens[token] = self.n_words
             self.index_special_tokens[self.n_words] = token
             self.n_words += 1
-        self.role_special_token_expression = "|".join([re.escape(token) for token in role_special_tokens])
-
-    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 tokenize(self, s: str):
+        return self.sp_model.EncodeAsPieces(s)
 
     def encode(self, s: str, bos: bool = False, eos: bool = False) -> List[int]:
         assert type(s) is str
@@ -93,8 +104,7 @@ class ChatGLMTokenizer(PreTrainedTokenizer):
 
     model_input_names = ["input_ids", "attention_mask", "position_ids"]
 
-    def __init__(self, vocab_file, padding_side="left", clean_up_tokenization_spaces=False, encode_special_tokens=False,
-                 **kwargs):
+    def __init__(self, vocab_file, padding_side="left", clean_up_tokenization_spaces=False, **kwargs):
         self.name = "GLMTokenizer"
 
         self.vocab_file = vocab_file
@@ -104,10 +114,7 @@ class ChatGLMTokenizer(PreTrainedTokenizer):
             "<eos>": self.tokenizer.eos_id,
             "<pad>": self.tokenizer.pad_id
         }
-        self.encode_special_tokens = encode_special_tokens
-        super().__init__(padding_side=padding_side, clean_up_tokenization_spaces=clean_up_tokenization_spaces,
-                         encode_special_tokens=encode_special_tokens,
-                         **kwargs)
+        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:
@@ -146,7 +153,7 @@ class ChatGLMTokenizer(PreTrainedTokenizer):
         return vocab
 
     def _tokenize(self, text, **kwargs):
-        return self.tokenizer.tokenize(text, encode_special_tokens=self.encode_special_tokens)
+        return self.tokenizer.tokenize(text)
 
     def _convert_token_to_id(self, token):
         """ Converts a token (str) in an id using the vocab. """
@@ -188,8 +195,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
@@ -298,3 +304,220 @@ 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 more 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 encode_sep_token(self):
+        return ENCODE_SEP_TOKEN_FOR_FAST
+
+    def _batch_encode_plus(
+        self,
+        batch_text_or_text_pairs: Union[
+            List[TextInput],
+            List[TextInputPair],
+            List[PreTokenizedInput],
+            List[PreTokenizedInputPair],
+        ],
+        add_special_tokens: bool = True,
+        padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD,
+        truncation_strategy: TruncationStrategy = TruncationStrategy.DO_NOT_TRUNCATE,
+        max_length: Optional[int] = None,
+        stride: int = 0,
+        is_split_into_words: bool = False,
+        pad_to_multiple_of: Optional[int] = None,
+        return_tensors: Optional[str] = None,
+        return_token_type_ids: Optional[bool] = None,
+        return_attention_mask: Optional[bool] = None,
+        return_overflowing_tokens: bool = False,
+        return_special_tokens_mask: bool = False,
+        return_offsets_mapping: bool = False,
+        return_length: bool = False,
+        verbose: bool = True,
+    ) -> BatchEncoding:
+        def split_sep(t: Union[TextInput, PreTokenizedInput]) -> PreTokenizedInput:
+            if isinstance(t, str):
+                return t.split(self.encode_sep_token)
+
+            return [w for word in t for w in split_sep(word)]
+
+        def split_maybe_tupled(
+            t: Union[TextInput, TextInputPair, PreTokenizedInput, PreTokenizedInputPair]
+        ) -> Union[PreTokenizedInputPair, PreTokenizedInput]:
+            if isinstance(t, tuple):
+                return split_sep(t[0]), split_sep(t[1])
+
+            return split_sep(t)
+
+        return super()._batch_encode_plus(
+            list(map(split_maybe_tupled, batch_text_or_text_pairs)),  # pyright: ignore
+            add_special_tokens,
+            padding_strategy,
+            truncation_strategy,
+            max_length,
+            stride,
+            True,
+            pad_to_multiple_of,
+            return_tensors,
+            return_token_type_ids,
+            return_attention_mask,
+            return_overflowing_tokens,
+            return_special_tokens_mask,
+            return_offsets_mapping,
+            return_length,
+            verbose,
+        )
+
+    @property
+    def can_save_slow_tokenizer(self) -> bool:
+        # multiple breaking changes
+        return False
+
+    def save_pretrained(
+        self,
+        save_directory: Union[str, os.PathLike],
+        legacy_format: Optional[bool] = None,
+        filename_prefix: Optional[str] = None,
+        push_to_hub: bool = False,
+        **kwargs,
+    ) -> Tuple[str]:
+        warnings.warn(
+            f"{type(self)} does not support saving slow tokenizer. "
+            "Saving it at the same directory may break the slow tokenizer. "
+            "Please keep a backup of the original tokenizer beforehand."
+        )
+        return super().save_pretrained(
+            save_directory, legacy_format, filename_prefix, push_to_hub, **kwargs
+        )
+
+    def build_single_message(self, role, metadata, message):
+        assert role in ["system", "user", "assistant", "observation"], role
+        return f"<|{role}|>{self.encode_sep_token}{metadata}\n{self.encode_sep_token}{message}"
+
+    def build_chat_text(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(
+                    item["role"], item.get("metadata", ""), content
+                )
+            )
+
+        inputs.append(self.build_single_message(role, metadata, query))
+        inputs.append("<|assistant|>")
+
+        return "".join(inputs)
+
+    def build_chat_input(self, *args, **kwargs):
+        return self.batch_encode_plus(
+            [self.build_chat_text(*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.Prepend(prepend="▁"),
+                normalizers.Replace(pattern=" ", content="▁"),
+            ]
+        )
+
+    def pre_tokenizer(self, replacement, add_prefix_space):
+        # don't use Metaspace, it will skip merging spaces into one token
+
+        # give up to split `encode_sep_token` here, buggy
+        # return pre_tokenizers.Split(ENCODE_SEP_TOKEN_FOR_FAST, "removed")
+
+        return None
+
+    def decoder(self, replacement, add_prefix_space):
+        return decoders.Sequence(
+            [
+                decoders.ByteFallback(),
+                super().decoder(replacement, add_prefix_space),
+            ]
+        )
+
+    def tokenizer(self, proto):
+        tokenizer = super().tokenizer(proto)
+
+        tokenizer.model.byte_fallback = True
+
+        special_tokens = [
+            "<unk>",
+            "<s>",
+            "</s>",
+            *ADDITIONAL_SPECIAL_TOKENS,
+        ]
+
+        tokenizer.add_special_tokens(
+            [
+                AddedToken(token, special=True, normalized=False)
+                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

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