Implemented robust type handling and tokenizer exception safety

tokenization_kimi.py

@@ -94,9 +94,16 @@ class TikTokenTokenizer(PreTrainedTokenizer):
                 "<|im_middle|>",
             ] 
         
-        special_tokens_mapping = {
-        i: added_tokens_decoder[i].content for i in added_tokens_decoder
-    }
+        special_tokens_mapping = {}
+        if added_tokens_decoder is not None:
+            for k, v in added_tokens_decoder.items():
+                if isinstance(v, dict):
+                    content = v.get("content", "")
+                elif hasattr(v, "content"):
+                    content = v.content
+                else:
+                    content = str(v)
+                special_tokens_mapping[int(k)] = content
 
         self.vocab_file = vocab_file
         mergeable_ranks = load_tiktoken_bpe(vocab_file)
@@ -108,8 +115,6 @@ class TikTokenTokenizer(PreTrainedTokenizer):
             )
         }
 
-
-
         self.model = tiktoken.Encoding(
             name=Path(vocab_file).name,
             pat_str=self.pat_str,
@@ -119,27 +124,39 @@ class TikTokenTokenizer(PreTrainedTokenizer):
         logger.info(f"Reloaded tiktoken model from {vocab_file}")
 
         self.n_words: int = self.model.n_vocab
-        # BOS / EOS token IDs
-        self.bos_id: int = self.special_tokens[str(bos_token)]
-        self.eos_id: int = self.special_tokens[str(eos_token)]
+        
+        # BOS / EOS / PAD / UNK string representation mapping
+        bos_str = bos_token.content if hasattr(bos_token, "content") else str(bos_token) if bos_token is not None else None
+        eos_str = eos_token.content if hasattr(eos_token, "content") else str(eos_token) if eos_token is not None else None
+        pad_str = pad_token.content if hasattr(pad_token, "content") else str(pad_token) if pad_token is not None else None
+        unk_str = unk_token.content if hasattr(unk_token, "content") else str(unk_token) if unk_token is not None else None
+
+        self.bos_id: Optional[int] = self.special_tokens.get(bos_str) if bos_str is not None else None
+        self.eos_id: Optional[int] = self.special_tokens.get(eos_str) if eos_str is not None else None
+        self.pad_id: Optional[int] = self.special_tokens.get(pad_str) if pad_str is not None else None
+        self.unk_id: Optional[int] = self.special_tokens.get(unk_str) if unk_str is not None else None
+
         logger.info(
             f"#words: {self.n_words} - BOS ID: {self.bos_id} - EOS ID: {self.eos_id}"
         )
 
-        self.pad_id: int = self.special_tokens[str(pad_token)]
-        self.unk_id: int = self.special_tokens[str(unk_token)]
-
         self.byte_encoder = bytes_to_unicode()
         self.byte_decoder = {v: k for k, v in self.byte_encoder.items()}
 
         self.decoder = {}
+        id_to_special = {v: k for k, v in self.special_tokens.items()}
         for i in range(self.n_words):
-            # Taken from https://gist.github.com/xenova/a452a6474428de0182b17605a98631ee
-            decoding = ''.join([
-                self.byte_encoder[ord(char)] for char in
-                self.model.decode_single_token_bytes(i).decode('latin-1')
-            ])
-            self.decoder[i] = decoding
+            if i in id_to_special:
+                self.decoder[i] = id_to_special[i]
+            else:
+                try:
+                    decoding = ''.join([
+                        self.byte_encoder[ord(char)] for char in
+                        self.model.decode_single_token_bytes(i).decode('latin-1')
+                    ])
+                    self.decoder[i] = decoding
+                except Exception:
+                    self.decoder[i] = f"<|reserved_token_{i}|>"
 
         self.encoder = {}
         for i in range(self.n_words):
@@ -180,7 +197,7 @@ class TikTokenTokenizer(PreTrainedTokenizer):
             logger.warning( f"Calling super().encode with {kwargs}" )
             return super().encode(text, **kwargs)
 
-        assert type(text) is str
+        assert isinstance(text, str), f"text must be a string, got {type(text)}"
 
         # The tiktoken tokenizer can handle <=400k chars without
         # pyo3_runtime.PanicException.
@@ -244,8 +261,13 @@ class TikTokenTokenizer(PreTrainedTokenizer):
         if len(kwargs) > 0:
             return super().decode(token_ids, **kwargs)
 
-        if type(token_ids) is int:
-            token_ids = [token_ids]
+        if hasattr(token_ids, "tolist"):
+            token_ids = token_ids.tolist()
+
+        if isinstance(token_ids, (int, float)):
+            token_ids = [int(token_ids)]
+        else:
+            token_ids = [int(x) for x in token_ids]
 
         return self.model.decode(cast(List[int], token_ids))