support gemma

Former-commit-id: b9674aa2f6f1b6b09b2a37375313d8d5abfcd453
2025-12-18 12:50:38 +08:00 · 2024-02-21 23:27:36 +08:00
parent 5ebd605149
commit 3b6a5951ce
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": {