Support Atom Model

Former-commit-id: da3e76f22aca9acaf772ff821b7eb03c2a2ac869
2024-02-21 18:35:03 +08:00
parent c5eb3fba62
commit d0ea203694
2 changed files with 20 additions and 0 deletions
--- a/src/llmtuner/data/template.py
+++ b/src/llmtuner/data/template.py
@@ -644,3 +644,9 @@ _register_template(
    format_user=StringFormatter(slots=[{"token": "<human>"}, ":{{content}}\n", {"token": "<bot>"}, ":"]),
    format_separator=EmptyFormatter(slots=["\n"]),
 )
+
+_register_template(
+    name="atom",
+    format_user=StringFormatter(slots=[{"bos_token"}, "Human: {{content}}\n", {"eos_token"}, {"bos_token"}, "Assistant:"]),
+    format_assistant=StringFormatter(slots=["{{content}}\n", {"eos_token"}]),
+)