Update tiktoken.py

tiktoken.py

@@ -172,26 +172,50 @@ class TiktokenTokenizerWrapper(PreTrainedTokenizer):
 
         Pinning default Chat ML template in case defaults change.
         """
+        # template = (
+        #     "{% set loop_messages = messages %}"
+        #     "{% set system_message = '' %}"
+        #     "{% if messages[0]['role'] == 'system' %}"
+        #     "{% set loop_messages = messages[1:] %}"
+        #     "{% set system_message = messages[0]['content'] %}"
+        #     "{{'<|im_start|>system\n' + system_message}}"
+        #     "{% elif USE_DEFAULT_PROMPT == true and not 'system' in messages[0]['role']%}"
+        #     "{% set loop_messages = messages %}"
+        #     "{% set system_message = 'DEFAULT_SYSTEM_PROMPT' %}"
+        #     '{% endif %}'
+        #     "{% if not (system_message == '') %}"
+        #     "{{'<|im_start|>system\n' + system_message}}"
+        #     '{% endif %}'
+        #     '{% for message in loop_messages %}'
+        #     "{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}"
+        #     '{% if (add_generation_prompt == true) %}'
+        #     "{{'<|im_start|>' + 'assistant' + '\n'}}"
+        #     '{% endif %}'
+        #     '{% endfor %}')
         template = (
-            "{% set loop_messages = messages %}"
-            "{% set system_message = '' %}"
             "{% if messages[0]['role'] == 'system' %}"
-            "{% set loop_messages = messages[1:] %}"
+            "{% set loop_messages = messages[1:] %}"  # Extract system message if it's present
             "{% set system_message = messages[0]['content'] %}"
-            "{{'<|im_start|>system\n' + system_message}}"
-            "{% elif USE_DEFAULT_PROMPT == true and not 'system' in messages[0]['role']%}"
+            "{% elif USE_DEFAULT_PROMPT == true and not 'system' in messages[0]['role'] %}"
+            "{% set loop_messages = messages %}"  # Or use the default system message if the flag is set
+            "{% set system_message = 'DEFAULT_SYSTEM_MESSAGE' %}"
+            "{% else %}"
             "{% set loop_messages = messages %}"
-            "{% set system_message = 'DEFAULT_SYSTEM_PROMPT' %}"
-            '{% endif %}'
-            "{% if not (system_message == '') %}"
-            "{{'<|im_start|>system\n' + system_message}}"
-            '{% endif %}'
-            '{% for message in loop_messages %}'
+            "{% set system_message = false %}"
+            "{% endif %}"
+            "{% for message in loop_messages %}"  # Loop over all non-system messages
+            "{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}"
+            "{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}"
+            "{% endif %}"
+            "{% if loop.index0 == 0 and system_message != false %}"  # Embed system message in first message
+            "{{ '<|im_start|>system\n' + system_message.strip() + '\n'}}"
+            "{% else %}"
             "{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}"
+            "{% endif %}"
             '{% if (add_generation_prompt == true) %}'
             "{{'<|im_start|>' + 'assistant' + '\n'}}"
-            '{% endif %}'
-            '{% endfor %}')
+            "{% else %}"
+            "{{ eos_token }}"
         template = template.replace(
             'USE_DEFAULT_PROMPT',
             'true' if self.use_default_system_prompt else 'false')