Add `ChatGLMTokenizerFast` and `ChatGLMTokenizerConverter`

tokenization_chatglm.py

@@ -3,13 +3,36 @@ import os
 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):
@@ -24,8 +47,7 @@ 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
+        special_tokens = ADDITIONAL_SPECIAL_TOKENS
         self.special_tokens = {}
         self.index_special_tokens = {}
         for token in special_tokens:
@@ -86,7 +108,7 @@ 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 or index > self.sp_model.vocab_size():
+        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)
 
@@ -216,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
@@ -326,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"
+    ]
   }
 }