fix: create final list of special tokens

tokenization_arcade100k.py

@@ -41,42 +41,52 @@ def _load_tiktoken_bpe(tiktoken_bpe_file: str) -> Dict[bytes, int]:
 def _arcade100k(vocab_file: str):
     mergeable_ranks = _load_tiktoken_bpe(vocab_file)
 
-    # Special Tokens
     ENDOFTEXT = "<|endoftext|>"
-    FIM_PREFIX = "<|fim_prefix|>"
-    FIM_MIDDLE = "<|fim_middle|>"
-    FIM_SUFFIX = "<|fim_suffix|>"
+    # StarCoder special tokens (https://huggingface.co/bigcode/starcoder/blob/main/tokenizer_config.json)
+    CODE = [
+        "<fim_prefix>",
+        "<fim_middle>",
+        "<fim_suffix>",
+        "<fim_pad>",
+        "<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|>"
-
-    # Custom Special Tokens
-    IM_START = "<|im_start|>"  # Chat: Input message start
-    IM_END = "<|im_end|>"  # Chat: Input message end
-    PAUSE = "<|pause|>"  # Think before you speak (https://arxiv.org/abs/2310.02226)
-    # Register/sink tokens (https://arxiv.org/abs/2309.17453)
-    REGISTERS = [f"<|reg{i}|>" for i in range(0, 8)]
-    custom_special_tokens = {
-        t: 100261 + i for i, t in enumerate([IM_START, IM_END, PAUSE, *REGISTERS])
-    }
-    ENDOFPROMPT_ID = 100276
-
-    # Fill-out extra tokens
-    for i in range(100261 + len(custom_special_tokens), ENDOFPROMPT_ID + 1):
-        custom_special_tokens[f"<|extra{i}|>"] = i
-
-    special_tokens = {
-        ENDOFTEXT: 100257,
-        FIM_PREFIX: 100258,
-        FIM_MIDDLE: 100259,
-        FIM_SUFFIX: 100260,
-        **custom_special_tokens,
-        ENDOFPROMPT: 100276,
+    
+    SPECIAL_TOKENS_NAMES = [ENDOFTEXT] + CODE + [ENDOFPROMPT] + CHAT + [PAUSE] + REGISTERS
+    START_ID = len(mergeable_ranks) + 1
+    SPECIAL_TOKENS = {
+        t: START_ID + i
+        for i, t in enumerate(SPECIAL_TOKENS_NAMES)
     }
 
+    print(len(mergeable_ranks))
+    print(len(SPECIAL_TOKENS))
+    print(len(mergeable_ranks) + len(SPECIAL_TOKENS))
+
     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+""",
         "mergeable_ranks": mergeable_ranks,
-        "special_tokens": special_tokens,
+        "special_tokens": SPECIAL_TOKENS,
     }
 
 
@@ -108,41 +118,44 @@ class Arcade100kTokenizer(PreTrainedTokenizer):
         self.tokenizer = tiktoken.Encoding(**self._tiktoken_config)
         # TODO: Remove this assertion
         assert (
-            len(self._tiktoken_config["mergeable_ranks"])
-            + len(self._tiktoken_config["special_tokens"])
+            len(self.tokenizer._mergeable_ranks)
+            + len(self.tokenizer._special_tokens) + 1
             == self.tokenizer.n_vocab
-        ), f"{len(self._tiktoken_config['mergeable_ranks']) + len(self._tiktoken_config['special_tokens'])} != {self.tokenizer.n_vocab} in encoding"
+        ), 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._tiktoken_config["mergeable_ranks"].items()
+            i: n for n, i in self.tokenizer._mergeable_ranks.items()
         }
         self.decoder.update(
-            {i: n for n, i in self._tiktoken_config["special_tokens"].items()}
+            {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]
 
+    def __len__(self):
+        return self.tokenizer.n_vocab
+
     @property
     def vocab_size(self):
         return self.tokenizer.n_vocab
 
     def get_vocab(self) -> Dict[bytes, int]:
-        return self._tiktoken_config["mergeable_ranks"]
+        return self.tokenizer._mergeable_ranks
 
     def convert_tokens_to_ids(
         self, tokens: Union[bytes, str, List[Union[bytes, str]]]
     ) -> List[int]:
         ids = []
         if isinstance(tokens, (str, bytes)):
-            if tokens in self._tiktoken_config["special_tokens"]:
-                return self._tiktoken_config["special_tokens"][tokens]
+            if tokens in self.tokenizer._special_tokens:
+                return self.tokenizer._special_tokens[tokens]
             else:
-                return self._tiktoken_config["mergeable_ranks"].get(tokens)
+                return self.tokenizer._mergeable_ranks.get(tokens)
         for token in tokens:
-            if token in self._tiktoken_config["special_tokens"]:
-                ids.append(self._tiktoken_config["special_tokens"][token])
+            if token in self.tokenizer._special_tokens:
+                ids.append(self.tokenizer._special_tokens[token])
             else:
-                ids.append(self._tiktoken_config["mergeable_ranks"].get(token))
+                ids.append(self.tokenizer._mergeable_ranks.get(token))
         return ids
 
     def _add_tokens(
@@ -167,7 +180,7 @@ class Arcade100kTokenizer(PreTrainedTokenizer):
         """
         file_path = os.path.join(save_directory, "qwen.tiktoken")
         with open(file_path, "w", encoding="utf8") as w:
-            for k, v in self._tiktoken_config["mergeable_ranks"].items():
+            for k, v in self.tokenizer._mergeable_ranks.items():
                 line = base64.b64encode(k).decode("utf8") + " " + str(v) + "\n"
                 w.write(line)
         return (file_path,)
@@ -236,10 +249,10 @@ class Arcade100kTokenizer(PreTrainedTokenizer):
 
     def _convert_token_to_id(self, token: Union[bytes, str]) -> int:
         """Converts a token to an id using the vocab, special tokens included"""
-        if token in self._tiktoken_config["special_tokens"]:
-            return self._tiktoken_config["special_tokens"][token]
-        if token in self._tiktoken_config["mergeable_ranks"]:
-            return self._tiktoken_config["mergeable_ranks"][token]
+        if token in self.tokenizer._special_tokens:
+            return self.tokenizer._special_tokens[token]
+        if token in self.tokenizer._mergeable_ranks:
+            return self.tokenizer._mergeable_ranks[token]
         raise ValueError("unknown token")
 
     def _tokenize(self, text: str, **kwargs):
@@ -262,4 +275,4 @@ class Arcade100kTokenizer(PreTrainedTokenizer):
             token_ids = [token_ids]
         if skip_special_tokens:
             token_ids = [i for i in token_ids if i < self.tokenizer.eot_token]
-        return self.tokenizer.decode(token_ids, errors=errors or self.errors)
+        return self.tokenizer.decode(token_ids)