From e25f7bae16b7ea41a4a1fd1e8db1b961e55d0c5b Mon Sep 17 00:00:00 2001
From: hiyouga <hiyouga@buaa.edu.cn>
Date: Wed, 6 Dec 2023 13:33:18 +0800
Subject: [PATCH] add models

---
 src/llmtuner/data/template.py    | 3 +++
 src/llmtuner/extras/constants.py | 8 ++++++++
 2 files changed, 11 insertions(+)

diff --git a/src/llmtuner/data/template.py b/src/llmtuner/data/template.py
index d9e83326..4f8d6a51 100644
--- a/src/llmtuner/data/template.py
+++ b/src/llmtuner/data/template.py
@@ -724,6 +724,9 @@ register_template(
     sep=[
         "<|im_end|>\n"
     ],
+    stop_words=[
+        "<|im_end|>"
+    ],
     efficient_eos=True
 )
 
diff --git a/src/llmtuner/extras/constants.py b/src/llmtuner/extras/constants.py
index 56d7ebb8..ee506b65 100644
--- a/src/llmtuner/extras/constants.py
+++ b/src/llmtuner/extras/constants.py
@@ -578,10 +578,18 @@ register_model_group(
             DownloadSource.DEFAULT: "01-ai/Yi-34B",
             DownloadSource.MODELSCOPE: "01ai/Yi-34B"
         },
+        "Yi-6B-Chat": {
+            DownloadSource.DEFAULT: "01-ai/Yi-6B-Chat",
+            DownloadSource.MODELSCOPE: "01ai/Yi-6B-Chat"
+        },
         "Yi-34B-Chat": {
             DownloadSource.DEFAULT: "01-ai/Yi-34B-Chat",
             DownloadSource.MODELSCOPE: "01ai/Yi-34B-Chat"
         },
+        "Yi-6B-int8-Chat": {
+            DownloadSource.DEFAULT: "01-ai/Yi-6B-Chat-8bits",
+            DownloadSource.MODELSCOPE: "01ai/Yi-6B-Chat-8bits"
+        },
         "Yi-34B-int8-Chat": {
             DownloadSource.DEFAULT: "01-ai/Yi-34B-Chat-8bits",
             DownloadSource.MODELSCOPE: "01ai/Yi-34B-Chat-8bits"