Update tiktoken.py

tiktoken.py

@@ -127,7 +127,7 @@ class TiktokenTokenizerWrapper(PreTrainedTokenizer):
         self.byte_encoder = bytes_to_unicode()
         self.byte_decoder = {v: k for k, v in self.byte_encoder.items()}
 
-        self.decoder = {}
+        self.decoder: Dict[int, str] = {}
         for i in range(self.encoding.n_vocab):
             try:
                 self.encoding.decode_single_token_bytes(i)
@@ -141,7 +141,7 @@ class TiktokenTokenizerWrapper(PreTrainedTokenizer):
             ])
             self.decoder[i] = decoding
 
-        self.encoder = {}
+        self.encoder: Dict[str, int] = {}
         for i in range(self.encoding.n_vocab):
             if i in self.decoder:
                 self.encoder[self.decoder[i]] = i
@@ -173,12 +173,30 @@ class TiktokenTokenizerWrapper(PreTrainedTokenizer):
         Pinning default Chat ML template in case defaults change.
         """
         template = (
-            "{% set system_message = '' %}"
-            '{% if USE_DEFAULT_PROMPT == true %}'
-            "{{'<|im_start|>system\n' + 'DEFAULT_SYSTEM_PROMPT'}}"
+            "{% if messages[0]['role'] == 'system' %}"
+            '{% set loop_messages = messages[1:] %}'
+            "{% set system_message = messages[0]['content'] %}"
+            "{% elif USE_DEFAULT_PROMPT == true and not 'system' in messages[0]['role'] %}"
+            '{% set loop_messages = messages %}'
+            "{% set system_message = 'DEFAULT_SYSTEM_PROMPT' %}"
+            '{% else %}'
+            '{% set loop_messages = messages %}'
+            '{% set system_message = false %}'
+            '{% endif %}'
+            '{% for message in loop_messages %}'
+            '{% if loop.index0 == 0 %}'
+            '{% if system_message != false %}'
+            "{{ '<|im_start|>system\n' + system_message.strip() + '\n'}}"
+            '{% endif %}'
+            "{{ '<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' }}"
+            '{% else %}'
+            "{{ '\n' + '<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' }}"
+            '{% endif %}'
+            '{% if (add_generation_prompt == true) %}'
+            "{{ '\n' + '<|im_start|>' + 'assistant' + '\n' }}"
+            "{% elif (message['role'] == 'assistant') %}"
+            '{{ eos_token }}'
             '{% endif %}'
-            '{% for message in messages %}'
-            "{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}"
             '{% endfor %}')
         template = template.replace(
             'USE_DEFAULT_PROMPT',
@@ -188,11 +206,7 @@ class TiktokenTokenizerWrapper(PreTrainedTokenizer):
         return template
 
     def get_vocab(self) -> Dict[str, int]:
-        """Returns vocab as a dict.
-
-        Note: This function does not work properly due to difference in assumptions between tiktoken and Hugging Face tokenizers.
-        Most uses do not need to use get_vocab, so this is not a priority to fix.
-        """
+        """Returns vocab as a dict."""
         # As far as I can tell, we don't require get_vocab to completely work,
         # but when using additional_special_tokens, Hugging Face determines the next
         # token index to add with len(self.get_vocab()) so we need the _size_ of this dictionary to be correct.
@@ -227,15 +241,15 @@ class TiktokenTokenizerWrapper(PreTrainedTokenizer):
 
         return tokens
 
-    def _convert_token_to_id(self, token: str):
+    def _convert_token_to_id(self, token: str) -> Optional[int]:
         """Converts a token (str) in an id using the vocab."""
         return self.encoder.get(token, self.encoder.get(self.unk_token))
 
-    def _convert_id_to_token(self, index: int):
+    def _convert_id_to_token(self, index: int) -> Optional[str]:
         """Converts an index (integer) in a token (str) using the vocab."""
         return self.decoder.get(index)
 
-    def convert_tokens_to_string(self, tokens: List[str]):
+    def convert_tokens_to_string(self, tokens: List[str]) -> str:
         """Converts a sequence of tokens (string) in a single string."""
         text = ''.join(tokens)
         text = bytearray([self.byte_decoder[c] for c in text]).decode('utf-8')