support gemma

2024-02-21 23:27:36 +08:00
parent 1630a4cb8f
commit c99e19641a
4 changed files with 33 additions and 2 deletions
--- a/src/llmtuner/extras/constants.py
+++ b/src/llmtuner/extras/constants.py
@@ -324,6 +324,25 @@ register_model_group(
 )


+register_model_group(
+    models={
+        "Gemma-2B": {
+            DownloadSource.DEFAULT: "google/gemma-2b",
+        },
+        "Gemma-7B": {
+            DownloadSource.DEFAULT: "google/gemma-7b",
+        },
+        "Gemma-2B-Chat": {
+            DownloadSource.DEFAULT: "google/gemma-2b-it",
+        },
+        "Gemma-7B-Chat": {
+            DownloadSource.DEFAULT: "google/gemma-7b-it",
+        },
+    },
+    template="gemma",
+)
+
+
 register_model_group(
    models={
        "InternLM-7B": {