support gemma

Former-commit-id: b9674aa2f6f1b6b09b2a37375313d8d5abfcd453
2024-02-21 23:27:36 +08:00
parent a274900188
commit 0c7d0bf172
4 changed files with 33 additions and 2 deletions
--- a/src/llmtuner/data/template.py
+++ b/src/llmtuner/data/template.py
@@ -433,6 +433,16 @@ _register_template(
 )


+_register_template(
+    name="gemma",
+    format_user=StringFormatter(slots=["<start_of_turn>user\n{{content}}<end_of_turn>\n<start_of_turn>model\n"]),
+    format_system=StringFormatter(slots=["{{content}}"]),
+    format_separator=EmptyFormatter(slots=["\n"]),
+    stop_words=["<end_of_turn>"],
+    replace_eos=True,
+)
+
+
 _register_template(
    name="intern",
    format_user=StringFormatter(slots=["<|User|>:{{content}}", {"token": "<eoh>"}, "\n<|Bot|>:"]),