fix: check surface form against special tokens in `_add_tokens`

tokenization_arcade100k.py

@@ -1,8 +1,8 @@
 # coding=utf-8
 # Copyright (c) 2023 Alibaba Cloud & Stability AI.
 #
-# This source code is licensed under the license found in the
-# LICENSE file in the root directory of this source tree.
+# Tongyi Qianwen LICENSE AGREEMENT:
+# https://github.com/QwenLM/Qwen/blob/5aa84bdfd3237b37f01bc88cd49b3279b9a71d0b/Tongyi%20Qianwen%20LICENSE%20AGREEMENT
 """Tokenization classes for Arcade100k."""
 
 import base64
@@ -29,47 +29,56 @@ def _load_tiktoken_bpe(tiktoken_bpe_file: str) -> Dict[bytes, int]:
     }
 
 
+ENDOFTEXT = "<|endoftext|>"
+FIM = [
+    "<|fim_prefix|>",
+    "<|fim_middle|>",
+    "<|fim_suffix|>",
+    "<|fim_pad|>",
+]
+# `StarCoder` Tokens
+CODE = [
+    "<gh_stars>",
+    "<filename>",
+    "<issue_start>",
+    "<issue_comment>",
+    "<issue_closed>",
+    "<jupyter_start>",
+    "<jupyter_text>",
+    "<jupyter_code>",
+    "<jupyter_output>",
+    "<empty_output>",
+    "<commit_before>",
+    "<commit_msg>",
+    "<commit_after>",
+    "<reponame>",
+]
+CHAT = [
+    "<|im_start|>",  # Chat: Input message start
+    "<|im_end|>",  # Chat: Input message end
+]
+PAUSE = "<|pause|>"  # Think before you speak (https://arxiv.org/abs/2310.02226)
+REGISTERS = [
+    f"<|reg{i}|>" for i in range(0, 8)
+]  # Register 0 sink token (https://arxiv.org/abs/2309.17453)
+ENDOFPROMPT = "<|endofprompt|>"
+SPECIAL_TOKENS_NAMES = (
+    [ENDOFTEXT]
+    + FIM
+    + CODE
+    + [ENDOFPROMPT]
+    + CHAT
+    + [PAUSE]
+    + REGISTERS
+    + ["<|extra0|>"]
+)
+START_ID = 100257
+SPECIAL_TOKENS = {t: START_ID + i for i, t in enumerate(SPECIAL_TOKENS_NAMES)}
+
+
 def _arcade100k(vocab_file: str):
     mergeable_ranks = _load_tiktoken_bpe(vocab_file)
 
-    ENDOFTEXT = "<|endoftext|>"
-    FIM = [
-        "<|fim_prefix|>",
-        "<|fim_middle|>",
-        "<|fim_suffix|>",
-        "<|fim_pad|>",
-    ]
-    # `StarCoder` Tokens
-    CODE = [
-        "<gh_stars>",
-        "<filename>",
-        "<issue_start>",
-        "<issue_comment>",
-        "<issue_closed>",
-        "<jupyter_start>",
-        "<jupyter_text>",
-        "<jupyter_code>",
-        "<jupyter_output>",
-        "<empty_output>",
-        "<commit_before>",
-        "<commit_msg>",
-        "<commit_after>",
-        "<reponame>"
-    ]
-    CHAT = [
-        "<|im_start|>", # Chat: Input message start
-        "<|im_end|>",   # Chat: Input message end
-    ]
-    PAUSE = "<|pause|>"       # Think before you speak (https://arxiv.org/abs/2310.02226)
-    REGISTERS = [f"<|reg{i}|>" for i in range(0, 8)]  # Register 0 sink token (https://arxiv.org/abs/2309.17453)
-    ENDOFPROMPT = "<|endofprompt|>"
-    SPECIAL_TOKENS_NAMES = [ENDOFTEXT] + FIM + CODE + [ENDOFPROMPT] + CHAT + [PAUSE] + REGISTERS + ["<|extra0|>"]
-    START_ID = len(mergeable_ranks) + 1
-    SPECIAL_TOKENS = {
-        t: START_ID + i
-        for i, t in enumerate(SPECIAL_TOKENS_NAMES)
-    }
-
     return {
         "name": NAME,
         "pat_str": r"""(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+""",
@@ -108,16 +117,13 @@ class Arcade100kTokenizer(PreTrainedTokenizer):
         # TODO: Remove this assertion
         assert (
             len(self.tokenizer._mergeable_ranks)
-            + len(self.tokenizer._special_tokens) + 1
+            + len(self.tokenizer._special_tokens)
+            + 1
             == self.tokenizer.n_vocab
         ), f"{len(self.tokenizer._mergeable_ranks) + len(self.tokenizer._special_tokens)} != {self.tokenizer.n_vocab} in encoding"
 
-        self.decoder = {
-            i: n for n, i in self.tokenizer._mergeable_ranks.items()
-        }
-        self.decoder.update(
-            {i: n for n, i in self.tokenizer._special_tokens.items()}
-        )
+        self.decoder = {i: n for n, i in self.tokenizer._mergeable_ranks.items()}
+        self.decoder.update({i: n for n, i in self.tokenizer._special_tokens.items()})
         self.eos_token = self.decoder[self.tokenizer.eot_token]
         self.pad_token = self.decoder[self.tokenizer.eot_token]
 
@@ -156,8 +162,7 @@ class Arcade100kTokenizer(PreTrainedTokenizer):
             raise ValueError("Adding regular tokens is not supported")
         for token in new_tokens:
             surface_form = token.content if isinstance(token, AddedToken) else token
-            # TODO: Expose `SPECIAL_TOKENS_SET` at higher level - currently only exists in `_arcade100k()`
-            if surface_form not in SPECIAL_TOKENS_SET:
+            if surface_form not in SPECIAL_TOKENS:
                 raise ValueError("Adding unknown special tokens is not supported")
         return 0