Release v0.1.6

src/llmtuner/extras/constants.py

@@ -37,7 +37,9 @@ SUPPORTED_MODELS = {
     "InternLM-7B": "internlm/internlm-7b",
     "InternLM-7B-Chat": "internlm/internlm-chat-7b",
     "Qwen-7B": "Qwen/Qwen-7B",
-    "Qwen-7B-Chat": "Qwen/Qwen-7B-Chat"
+    "Qwen-7B-Chat": "Qwen/Qwen-7B-Chat",
+    "XVERSE-13B": "xverse/XVERSE-13B",
+    "ChatGLM2-6B": "THUDM/chatglm2-6b"
 }
 
 DEFAULT_MODULE = {
@@ -48,5 +50,7 @@ DEFAULT_MODULE = {
     "Falcon": "query_key_value",
     "Baichuan": "W_pack",
     "InternLM": "q_proj,v_proj",
-    "Qwen": "c_attn"
+    "Qwen": "c_attn",
+    "XVERSE": "q_proj,v_proj",
+    "ChatGLM2": "query_key_value"
 }

src/llmtuner/extras/template.py

@@ -178,7 +178,7 @@ def register_template(
     stop_words: List[str],
     use_history: bool
 ) -> None:
-    template_class = Llama2Template if name == "llama2" else Template
+    template_class = Llama2Template if "llama2" in name else Template
     templates[name] = template_class(
         prefix=prefix,
         prompt=prompt,
@@ -272,6 +272,23 @@ register_template(
 )
 
 
+r"""
+Supports: https://github.com/ymcui/Chinese-LLaMA-Alpaca-2
+"""
+register_template(
+    name="llama2_zh",
+    prefix=[
+        "<<SYS>>\nYou are a helpful assistant. 你是一个乐于助人的助手。\n<</SYS>>\n\n"
+    ],
+    prompt=[
+        "[INST] {{query}} [/INST] "
+    ],
+    sep=[],
+    stop_words=[],
+    use_history=True
+)
+
+
 r"""
 Supports: https://huggingface.co/tatsu-lab/alpaca-7b-wdiff
           https://github.com/ymcui/Chinese-LLaMA-Alpaca