add chatml

axolotl-ai-cloud · Aug 6, 2023 · b6fd675 · b6fd675
1 parent e0c8a8f
commit b6fd675
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/src/axolotl/prompters.py b/src/axolotl/prompters.py
@@ -46,6 +46,12 @@ def match_prompt_style(self):
             self.turn_format = "USER: {instruction}\n{input}\nASSISTANT:"
             self.turn_no_input_format = "USER: {instruction}\nASSISTANT:"
             self.system_format = "SYSTEM:{system}\n"
+        if self.prompt_style == PromptStyle.CHATML.value:
+            self.turn_format = "<|im_start|>user\n{instruction}\n{input}<|im_end|>\n<|im_start|>assistant\n"
+            self.turn_no_input_format = (
+                "<|im_start|>user\n{instruction}<|im_end|>\n<|im_start|>assistant\n"
+            )
+            self.system_format = "<|im_start|>system\n{system}<|im_end|>\n"
 
     def build_prompt(
         self,