Update tiktoken.py

tiktoken.py

@@ -179,10 +179,17 @@ class TiktokenTokenizerWrapper(PreTrainedTokenizer):
             '{% endif %}'
             '{% for message in messages %}'
             "{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}"
+            '{% if ADD_GENERATION_PROMPT == true %}'
+            "{{ '<|im_start|>' + 'assistant' + '\n' }}"
             '{% endfor %}')
         template = template.replace(
             'USE_DEFAULT_PROMPT',
             'true' if self.use_default_system_prompt else 'false')
+        template = temp.replace(
+            'ADD_GENERATION_PROMPT',
+            'true' if self.add_generation_prompt else 'false'
+        )
+        print("GENERATION PROMPT", self.add_generation_prompt)
         template = template.replace('DEFAULT_SYSTEM_PROMPT',
                                     DEFAULT_SYSTEM_PROMPT)
         return template