DRAFT: Add a fast tokenizer implementation and converter

tokenization_chatglm.py

@@ -1,11 +1,37 @@
 import json
 import os
-import torch
-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,11 +47,9 @@ 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|>"]
         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
@@ -171,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
@@ -281,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

@@ -1,12 +1,12 @@
 {
-  "name_or_path": "THUDM/chatglm2-6b",
+  "name_or_path": "THUDM/chatglm3-6b",
   "remove_space": false,
   "do_lower_case": false,
   "tokenizer_class": "ChatGLMTokenizer",
   "auto_map": {
     "AutoTokenizer": [
       "tokenization_chatglm.ChatGLMTokenizer",
-      null
-      ]
+      "tokenization_chatglm.ChatGLMTokenizerFast"
+    ]
   }
 }