Merge pull request #69 from pabl-o-ce/autocoder

add autocoder template
Maximilian-Winter · Jun 6, 2024 · afd6b5b · afd6b5b
2 parents f7b087d + 993c8e4
commit afd6b5b
Showing 1 changed file with 17 additions and 0 deletions.
diff --git a/src/llama_cpp_agent/messages_formatter.py b/src/llama_cpp_agent/messages_formatter.py
@@ -25,6 +25,7 @@ class MessagesFormatterType(Enum):
     LLAMA_3 = 12
     PHI_3 = 13
     OPEN_INTERPRETER = 14
+    AUTOCODER = 15
 
 @dataclass
 class PromptMarkers:
@@ -223,6 +224,12 @@ def _format_response(
     Roles.assistant: PromptMarkers("### Response:\n", "\n"),
     Roles.tool: PromptMarkers("", ""),
 }
+autocoder_chat_prompt_markers = {
+    Roles.system: PromptMarkers("", "\n"),
+    Roles.user: PromptMarkers("Human: ", "\n"),
+    Roles.assistant: PromptMarkers("Assistant: ", "<|EOT|>\n"),
+    Roles.tool: PromptMarkers("", ""),
+}
 
 """
 ### Instruction:
@@ -336,6 +343,15 @@ def _format_response(
     use_user_role_for_function_call_result=True,
 )
 
+autocoder_chat_formatter = MessagesFormatter(
+    "",
+    autocoder_chat_prompt_markers,
+    True,
+    ["<|EOT|>"],
+    bos_token="<｜begin▁of▁sentence｜>",
+    eos_token="<|EOT|>",
+)
+
 predefined_formatter = {
     MessagesFormatterType.MISTRAL: mixtral_formatter,
     MessagesFormatterType.CHATML: chatml_formatter,
@@ -351,6 +367,7 @@ def _format_response(
     MessagesFormatterType.LLAMA_3: llama_3_formatter,
     MessagesFormatterType.PHI_3: phi_3_chat_formatter,
     MessagesFormatterType.OPEN_INTERPRETER: open_interpreter_chat_formatter,
+    MessagesFormatterType.AUTOCODER: autocoder_chat_formatter,
 }