From b04316d9a81d1b6f3ad291ea459a0501256d0045 Mon Sep 17 00:00:00 2001
From: hiyouga <hiyouga@buaa.edu.cn>
Date: Mon, 4 Mar 2024 19:29:26 +0800
Subject: [PATCH] update readme

Former-commit-id: 24a79bd50f972008ca1a862edaff7cdd6c211cef
---
 README.md                     |  4 ++++
 README_zh.md                  |  4 ++++
 src/llmtuner/data/template.py | 26 ++++++++++++++++++++++++++
 3 files changed, 34 insertions(+)

diff --git a/README.md b/README.md
index 7f2b9068..34d53d7d 100644
--- a/README.md
+++ b/README.md
@@ -144,6 +144,8 @@ Compared to ChatGLM's [P-Tuning](https://github.com/THUDM/ChatGLM2-6B/tree/main/
 
 Please refer to [constants.py](src/llmtuner/extras/constants.py) for a full list of models we supported.
 
+You also can add a custom chat template to [template.py](src/llmtuner/data/template.py).
+
 ## Supported Training Approaches
 
 | Approach               |     Full-tuning    |    Freeze-tuning   |       LoRA         |       QLoRA        |
@@ -556,6 +558,8 @@ python src/export_model.py \
 > Merging LoRA weights into a quantized model is not supported.
 
 > [!TIP]
+> Use `--model_name_or_path path_to_export` only to use the exported model.
+> 
 > Use `--export_quantization_bit 4` and `--export_quantization_dataset data/c4_demo.json` to quantize the model after merging the LoRA weights.
 
 ### Inference with OpenAI-style API
diff --git a/README_zh.md b/README_zh.md
index e8247ed3..e1917bdb 100644
--- a/README_zh.md
+++ b/README_zh.md
@@ -144,6 +144,8 @@ https://github.com/hiyouga/LLaMA-Factory/assets/16256802/ec36a9dd-37f4-4f72-81bd
 
 项目所支持模型的完整列表请参阅 [constants.py](src/llmtuner/extras/constants.py)。
 
+您也可以在 [template.py](src/llmtuner/data/template.py) 中添加自己的对话模板。
+
 ## 训练方法
 
 | 方法                   |     全参数训练      |    部分参数训练     |       LoRA         |       QLoRA        |
@@ -555,6 +557,8 @@ python src/export_model.py \
 > 尚不支持量化模型的 LoRA 权重合并及导出。
 
 > [!TIP]
+> 仅使用 `--model_name_or_path path_to_export` 来加载导出后的模型。
+> 
 > 合并 LoRA 权重之后可再次使用 `--export_quantization_bit 4` 和 `--export_quantization_dataset data/c4_demo.json` 量化模型。
 
 ### 使用 OpenAI 风格 API 推理
diff --git a/src/llmtuner/data/template.py b/src/llmtuner/data/template.py
index 4f13f416..cb1a1811 100644
--- a/src/llmtuner/data/template.py
+++ b/src/llmtuner/data/template.py
@@ -213,6 +213,32 @@ def _register_template(
     replace_eos: Optional[bool] = False,
     force_system: Optional[bool] = False,
 ) -> None:
+    r"""
+    Registers a chat template.
+
+    To add the following chat template:
+    ```
+    [HUMAN]:
+    user prompt here
+    [AI]:
+    model response here
+
+    [HUMAN]:
+    user prompt here
+    [AI]:
+    model response here
+    ```
+
+    The corresponding code should be:
+    ```
+    _register_template(
+        name="custom",
+        format_user=StringFormatter(slots=["[HUMAN]:\n{{content}}\n[AI]:\n"]),
+        format_separator=EmptyFormatter(slots=["\n\n"]),
+        efficient_eos=True,
+    )
+    ```
+    """
     eos_slots = [] if efficient_eos else [{"eos_token"}]
     template_class = Llama2Template if name.startswith("llama2") else Template
     default_user_formatter = StringFormatter(slots=["{{content}}"])