support Yi-34B-Chat models

Former-commit-id: 1751a79c27e7fc13e76a731a061dc0c10d828cda
2023-11-23 19:31:49 +08:00
parent 7f54008d3c
commit 953a562ec1
2 changed files with 22 additions and 2 deletions
--- a/src/llmtuner/extras/constants.py
+++ b/src/llmtuner/extras/constants.py
@@ -264,8 +264,12 @@ register_model_group(
 register_model_group(
    models={
        "Yi-6B": "01-ai/Yi-6B",
-        "Yi-34B": "01-ai/Yi-34B"
-    }
+        "Yi-34B": "01-ai/Yi-34B",
+        "Yi-34B-Chat": "01-ai/Yi-34B-Chat",
+        "Yi-34B-int8-Chat": "01-ai/Yi-34B-Chat-8bits",
+        "Yi-34B-int4-Chat": "01-ai/Yi-34B-Chat-4bits"
+    },
+    template="yi"
 )