Update custom_tokenizer.py

custom_tokenizer.py

@@ -5,9 +5,17 @@ class CustomGPT2Tokenizer(GPT2Tokenizer):
         super().__init__(*args, **kwargs)
         self.chat_template = "{% for message in messages %}{% if loop.first and messages[0]['role'] != 'system' %}{{ '<|im_start|>system\nYou are a helpful AI assistant named Securitron, trained by Aquilax.<|im_end|>\n' }}{% endif %}{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}"
     
-    def apply_chat_template(self, messages, add_generation_prompt=True):
+    def apply_chat_template(self, messages, add_system_prompt=True, add_generation_prompt=True):
         """
-        Applies the chat template to the provided messages.
+        Applies the chat template to the provided messages. Optionally adds the system prompt.
+        
+        Args:
+            messages (list): List of message dictionaries with 'role' and 'content'.
+            add_system_prompt (bool): If True, adds the system prompt at the beginning.
+            add_generation_prompt (bool): If True, adds a prompt for generation.
+        
+        Returns:
+            str: Formatted text with the chat template applied.
         """
         # Handle template processing here
         formatted_messages = []
@@ -17,12 +25,11 @@ class CustomGPT2Tokenizer(GPT2Tokenizer):
             if role and content:
                 formatted_messages.append(f"<|im_start|>{role}\n{content}<|im_end|>\n")
         
-        # If the first message is not from 'system', insert system message
-        if messages[0]['role'] != 'system':
+        # If the first message is not from 'system' and we want to add the system prompt, do so
+        if add_system_prompt and messages[0]['role'] != 'system':
             formatted_messages.insert(0, "<|im_start|>system\nYou are a helpful AI assistant named Securitron, trained by Aquilax.<|im_end|>\n")
         
         formatted_text = "".join(formatted_messages)
-
         
         return formatted_text