From 5f2943dc84d7b84ce76be05200831424d2d5f6d8 Mon Sep 17 00:00:00 2001
From: hiyouga <hiyouga@buaa.edu.cn>
Date: Thu, 23 Nov 2023 19:31:49 +0800
Subject: [PATCH] support Yi-34B-Chat models

Former-commit-id: ff1c289229ee382d3e76578bbb6a5e299b969ded
---
 src/llmtuner/data/template.py    | 16 ++++++++++++++++
 src/llmtuner/extras/constants.py |  8 ++++++--
 2 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/src/llmtuner/data/template.py b/src/llmtuner/data/template.py
index 2e1f31c7..ac275c50 100644
--- a/src/llmtuner/data/template.py
+++ b/src/llmtuner/data/template.py
@@ -682,6 +682,22 @@ register_template(
 )
 
 
+register_template(
+    name="yi",
+    prefix=[
+        "{{system}}"
+    ],
+    prompt=[
+        "<|im_start|>user\n{{query}}<|im_end|>\n<|im_start|>assistant\n"
+    ],
+    system="",
+    sep=[
+        "<|im_end|>\n"
+    ],
+    efficient_eos=True
+)
+
+
 register_template(
     name="zephyr",
     prefix=[
diff --git a/src/llmtuner/extras/constants.py b/src/llmtuner/extras/constants.py
index cc9ac290..81040d64 100644
--- 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"
 )