From 6fc2d5cc0375a183ec95f003e7a27567b2a71514 Mon Sep 17 00:00:00 2001
From: hiyouga <hiyouga@buaa.edu.cn>
Date: Mon, 22 Jan 2024 21:26:53 +0800
Subject: [PATCH] add orion models

---
 src/llmtuner/data/template.py    |  8 ++++++++
 src/llmtuner/extras/constants.py | 29 ++++++++++++++++++++++++++++-
 2 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/src/llmtuner/data/template.py b/src/llmtuner/data/template.py
index 85516f98..2d89c3fa 100644
--- a/src/llmtuner/data/template.py
+++ b/src/llmtuner/data/template.py
@@ -462,6 +462,14 @@ register_template(
 )
 
 
+register_template(
+    name="orion",
+    format_user=StringFormatter(slots=["Human: {{content}}\n\nAssistant: </s>"]),
+    format_system=StringFormatter(slots=[{"bos_token"}, "{{content}}"]),
+    force_system=True,
+)
+
+
 register_template(
     name="qwen",
     format_user=StringFormatter(slots=["<|im_start|>user\n{{content}}<|im_end|>\n<|im_start|>assistant\n"]),
diff --git a/src/llmtuner/extras/constants.py b/src/llmtuner/extras/constants.py
index 0b4f5a07..a9818e58 100644
--- a/src/llmtuner/extras/constants.py
+++ b/src/llmtuner/extras/constants.py
@@ -455,7 +455,7 @@ register_model_group(
 register_model_group(
     models={
         "OpenChat3.5-7B-Chat": {
-            DownloadSource.DEFAULT: "openchat/openchat_3.5",
+            DownloadSource.DEFAULT: "openchat/openchat-3.5-0106",
             DownloadSource.MODELSCOPE: "myxiongmodel/openchat_3.5",
         }
     },
@@ -463,6 +463,33 @@ register_model_group(
 )
 
 
+register_model_group(
+    models={
+        "Orion-14B-Base": {
+            DownloadSource.DEFAULT: "OrionStarAI/Orion-14B-Base",
+            DownloadSource.MODELSCOPE: "OrionStarAI/Orion-14B-Base",
+        },
+        "Orion-14B-Chat": {
+            DownloadSource.DEFAULT: "OrionStarAI/Orion-14B-Chat",
+            DownloadSource.MODELSCOPE: "OrionStarAI/Orion-14B-Chat",
+        },
+        "Orion-14B-Long-Chat": {
+            DownloadSource.DEFAULT: "OrionStarAI/Orion-14B-LongChat",
+            DownloadSource.MODELSCOPE: "OrionStarAI/Orion-14B-LongChat",
+        },
+        "Orion-14B-RAG-Chat": {
+            DownloadSource.DEFAULT: "OrionStarAI/Orion-14B-Chat-RAG",
+            DownloadSource.MODELSCOPE: "OrionStarAI/Orion-14B-Chat-RAG",
+        },
+        "Orion-14B-Plugin-Chat": {
+            DownloadSource.DEFAULT: "OrionStarAI/Orion-14B-Chat-Plugin",
+            DownloadSource.MODELSCOPE: "OrionStarAI/Orion-14B-Chat-Plugin",
+        },
+    },
+    template="orion",
+)
+
+
 register_model_group(
     models={
         "Phi-1.5-1.3B": {DownloadSource.DEFAULT: "microsoft/phi-1_5", DownloadSource.MODELSCOPE: "allspace/PHI_1-5"},