From 0dd632fe9e5bbf08605d4b9c6887208b7a127317 Mon Sep 17 00:00:00 2001
From: hzhaoy <hzywong@gmail.com>
Date: Wed, 29 May 2024 15:00:37 +0800
Subject: [PATCH] add TeleChat-12B/TeleChat-12B-v2 models

---
 src/llamafactory/data/template.py    | 10 ++++++++++
 src/llamafactory/extras/constants.py | 16 ++++++++++++++++
 2 files changed, 26 insertions(+)

diff --git a/src/llamafactory/data/template.py b/src/llamafactory/data/template.py
index f13b303b..1a9f08ff 100644
--- a/src/llamafactory/data/template.py
+++ b/src/llamafactory/data/template.py
@@ -918,3 +918,13 @@ _register_template(
     format_user=StringFormatter(slots=["<human>:{{content}}\n<bot>:"]),
     format_separator=EmptyFormatter(slots=["\n"]),
 )
+
+
+_register_template(
+    name="telechat",
+    format_user=StringFormatter(slots=["<_user>{{content}}<_bot>"]),
+    format_system=StringFormatter(slots=["<_system>{{content}}<_end>"]),
+    default_system="",
+    stop_words=["<_end>"],
+    replace_eos=True,
+)
diff --git a/src/llamafactory/extras/constants.py b/src/llamafactory/extras/constants.py
index d1fcd4f6..bf2388ab 100644
--- a/src/llamafactory/extras/constants.py
+++ b/src/llamafactory/extras/constants.py
@@ -1355,3 +1355,19 @@ register_model_group(
     },
     template="zephyr",
 )
+
+
+register_model_group(
+    models={
+        "TeleChat-12B-Chat": {
+            DownloadSource.DEFAULT: "Tele-AI/TeleChat-12B",
+            DownloadSource.MODELSCOPE: "TeleAI/TeleChat-12B",
+        },
+        "TeleChat-12B-v2-Chat": {
+            DownloadSource.DEFAULT: "Tele-AI/TeleChat-12B-v2",
+            DownloadSource.MODELSCOPE: "TeleAI/TeleChat-12B-v2",
+        },
+    },
+    module='query,key_value',
+    template="telechat",
+)