release v0.1.2

Former-commit-id: 04aad91b71cc3a1acaf1bcec4304ce6b2098f7dc
fix api
2023-07-20 22:33:59 +08:00 · 2023-07-20 22:14:54 +08:00 · 2023-07-20 22:12:07 +08:00 · 2023-07-20 22:09:06 +08:00 · 2023-07-20 21:49:26 +08:00 · 2023-07-20 17:23:16 +08:00
45 changed files with 735 additions and 429 deletions
--- a/README.md
+++ b/README.md
@@ -10,6 +10,8 @@

 ## Changelog

+[23/07/19] Now we support training the **LLaMA-2** models in this repo. Try `--model_name_or_path meta-llama/Llama-2-7b-hf` argument to use the LLaMA-2 model. Remember to use `--prompt_template llama2` argument when you are using the LLaMA-2-chat model.
+
 [23/07/18] Now we develop an all-in-one Web UI for training, evaluation and inference. Try `train_web.py` to fine-tune models in your Web browser. Thank [@KanadeSiina](https://github.com/KanadeSiina) and [@codemayq](https://github.com/codemayq) for their efforts in the development.

 [23/07/11] Now we support training the **Baichuan-13B** model in this repo. Please replace the Baichuan-13B model file with `tests/modeling_baichuan.py` and try `--model_name_or_path path_to_baichuan_model` and `--lora_target W_pack` arguments to train the Baichuan-13B model. Remember to use `--prompt_template baichuan` argument when you are using the Baichuan-13B-Chat model.
@@ -24,7 +26,7 @@

 [23/06/22] Now we align the [demo API](src/api_demo.py) with the [OpenAI's](https://platform.openai.com/docs/api-reference/chat) format where you can insert the fine-tuned model in **arbitrary ChatGPT-based applications**.

-[23/06/15] Now we support training the **Baichuan-7B** model in this repo. Try `--model_name_or_path baichuan-inc/Baichuan-7B` and `--lora_target W_pack` arguments to use the Baichuan-7B model. If you want to train with RTX3090, use `git checkout baichuan-7b-rtx3090` to switch to the `baichuan-7b-rtx3090` branch and try the `--baichuan_rtx_gpu true` argument. (Other RTX series GPUs can also be tried)
+[23/06/15] Now we support training the **Baichuan-7B** model in this repo. Try `--model_name_or_path baichuan-inc/Baichuan-7B` and `--lora_target W_pack` arguments to use the Baichuan-7B model.

 [23/06/03] Now we support quantized training and inference (aka **[QLoRA](https://github.com/artidoro/qlora)**). Try `--quantization_bit 4/8` argument to work with quantized model. (experimental feature)

@@ -33,6 +35,7 @@
 ## Supported Models

 - [LLaMA](https://github.com/facebookresearch/llama) (7B/13B/33B/65B)
+- [LLaMA-2](https://huggingface.co/meta-llama) (7B/13B/70B)
 - [BLOOM](https://huggingface.co/bigscience/bloom) & [BLOOMZ](https://huggingface.co/bigscience/bloomz) (560M/1.1B/1.7B/3B/7.1B/176B)
 - [Falcon](https://huggingface.co/tiiuae/falcon-7b) (7B/40B)
 - [Baichuan](https://huggingface.co/baichuan-inc/baichuan-7B) (7B/13B)
@@ -57,32 +60,32 @@
 ## Provided Datasets

 - For pre-training:
-  - [Wiki Demo](data/wiki_demo.txt)
+  - [Wiki Demo (en)](data/wiki_demo.txt)
 - For supervised fine-tuning:
-  - [Stanford Alpaca](https://github.com/tatsu-lab/stanford_alpaca)
-  - [Stanford Alpaca (Chinese)](https://github.com/ymcui/Chinese-LLaMA-Alpaca)
-  - [GPT-4 Generated Data](https://github.com/Instruction-Tuning-with-GPT-4/GPT-4-LLM)
-  - [BELLE 2M](https://huggingface.co/datasets/BelleGroup/train_2M_CN)
-  - [BELLE 1M](https://huggingface.co/datasets/BelleGroup/train_1M_CN)
-  - [BELLE 0.5M](https://huggingface.co/datasets/BelleGroup/train_0.5M_CN)
-  - [BELLE Dialogue 0.4M](https://huggingface.co/datasets/BelleGroup/generated_chat_0.4M)
-  - [BELLE School Math 0.25M](https://huggingface.co/datasets/BelleGroup/school_math_0.25M)
-  - [BELLE Multiturn Chat 0.8M](https://huggingface.co/datasets/BelleGroup/multiturn_chat_0.8M)
-  - [Guanaco Dataset](https://huggingface.co/datasets/JosephusCheung/GuanacoDataset)
-  - [Firefly 1.1M](https://huggingface.co/datasets/YeungNLP/firefly-train-1.1M)
-  - [CodeAlpaca 20k](https://huggingface.co/datasets/sahil2801/CodeAlpaca-20k)
-  - [Alpaca CoT](https://huggingface.co/datasets/QingyiSi/Alpaca-CoT)
-  - [Web QA (Chinese)](https://huggingface.co/datasets/suolyer/webqa)
-  - [UltraChat](https://github.com/thunlp/UltraChat)
-  - [Open Assistant](https://huggingface.co/datasets/OpenAssistant/oasst1)
-  - [Open Assistant (Chinese)](https://huggingface.co/datasets/OpenAssistant/oasst1)
-  - [WebNovel (Chinese)](https://huggingface.co/datasets/zxbsmk/webnovel_cn)
- For reward model training:
-  - [HH-RLHF](https://huggingface.co/datasets/Anthropic/hh-rlhf)
-  - [Open Assistant](https://huggingface.co/datasets/OpenAssistant/oasst1)
-  - [Open Assistant (Chinese)](https://huggingface.co/datasets/OpenAssistant/oasst1)
-  - [GPT-4 Generated Data](https://github.com/Instruction-Tuning-with-GPT-4/GPT-4-LLM)
-  - [GPT-4 Generated Data (Chinese)](https://github.com/Instruction-Tuning-with-GPT-4/GPT-4-LLM)
+  - [Stanford Alpaca (en)](https://github.com/tatsu-lab/stanford_alpaca)
+  - [Stanford Alpaca (zh)](https://github.com/ymcui/Chinese-LLaMA-Alpaca)
+  - [GPT-4 Generated Data (en&zh)](https://github.com/Instruction-Tuning-with-GPT-4/GPT-4-LLM)
+  - [Open Assistant (multilingual)](https://huggingface.co/datasets/OpenAssistant/oasst1)
+  - [Self-cognition (zh)](data/self_cognition.json)
+  - [ShareGPT (zh)](https://huggingface.co/datasets/QingyiSi/Alpaca-CoT/tree/main/Chinese-instruction-collection)
+  - [RefGPT (zh)](https://github.com/sufengniu/RefGPT)
+  - [Guanaco Dataset (multilingual)](https://huggingface.co/datasets/JosephusCheung/GuanacoDataset)
+  - [BELLE 2M (zh)](https://huggingface.co/datasets/BelleGroup/train_2M_CN)
+  - [BELLE 1M (zh)](https://huggingface.co/datasets/BelleGroup/train_1M_CN)
+  - [BELLE 0.5M (zh)](https://huggingface.co/datasets/BelleGroup/train_0.5M_CN)
+  - [BELLE Dialogue 0.4M (zh)](https://huggingface.co/datasets/BelleGroup/generated_chat_0.4M)
+  - [BELLE School Math 0.25M (zh)](https://huggingface.co/datasets/BelleGroup/school_math_0.25M)
+  - [BELLE Multiturn Chat 0.8M (zh)](https://huggingface.co/datasets/BelleGroup/multiturn_chat_0.8M)
+  - [Firefly 1.1M (zh)](https://huggingface.co/datasets/YeungNLP/firefly-train-1.1M)
+  - [CodeAlpaca 20k (en)](https://huggingface.co/datasets/sahil2801/CodeAlpaca-20k)
+  - [Alpaca CoT (multilingual)](https://huggingface.co/datasets/QingyiSi/Alpaca-CoT)
+  - [Web QA (zh)](https://huggingface.co/datasets/suolyer/webqa)
+  - [UltraChat (en)](https://github.com/thunlp/UltraChat)
+  - [WebNovel (zh)](https://huggingface.co/datasets/zxbsmk/webnovel_cn)
+- For reward modelling:
+  - [HH-RLHF (en)](https://huggingface.co/datasets/Anthropic/hh-rlhf)
+  - [Open Assistant (multilingual)](https://huggingface.co/datasets/OpenAssistant/oasst1)
+  - [GPT-4 Generated Data (en&zh)](https://github.com/Instruction-Tuning-with-GPT-4/GPT-4-LLM)

 Please refer to [data/README.md](data/README.md) for details.

@@ -260,34 +263,64 @@ use_cpu: false

 ```bash
 CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
-    --stage pt \
+    --stage sft \
    --model_name_or_path path_to_your_model \
    --do_eval \
    --dataset alpaca_gpt4_en \
+    --finetuning_type lora \
    --checkpoint_dir path_to_checkpoint \
    --output_dir path_to_eval_result \
    --per_device_eval_batch_size 8 \
-    --max_samples 50 \
+    --max_samples 100 \
    --predict_with_generate
 ```

 We recommend using `--per_device_eval_batch_size=1` and `--max_target_length 128` at 4/8-bit evaluation.

+### Predict
+
+```bash
+CUDA_VISIBLE_DEVICES=0 python src/train_bash.py \
+    --stage sft \
+    --model_name_or_path path_to_your_model \
+    --do_predict \
+    --dataset alpaca_gpt4_en \
+    --finetuning_type lora \
+    --checkpoint_dir path_to_checkpoint \
+    --output_dir path_to_predict_result \
+    --per_device_eval_batch_size 8 \
+    --max_samples 100 \
+    --predict_with_generate
+```
+
+If you want to predict the samples with empty responses, please kindly fill the `response` column with **dummy tokens** to ensure the sample will not be discarded throughout the preprocessing phase.
+
 ### API Demo

 ```bash
 python src/api_demo.py \
    --model_name_or_path path_to_your_model \
+    --finetuning_type lora \
    --checkpoint_dir path_to_checkpoint
 ```

-See `http://localhost:8000/docs` for API documentation.
+Visit `http://localhost:8000/docs` for API documentation.

 ### CLI Demo

 ```bash
 python src/cli_demo.py \
    --model_name_or_path path_to_your_model \
+    --finetuning_type lora \
+    --checkpoint_dir path_to_checkpoint
+```
+
+### Web Demo
+
+```bash
+python src/web_demo.py \
+    --model_name_or_path path_to_your_model \
+    --finetuning_type lora \
    --checkpoint_dir path_to_checkpoint
 ```

@@ -296,6 +329,7 @@ python src/cli_demo.py \
 ```bash
 python src/export_model.py \
    --model_name_or_path path_to_your_model \
+    --finetuning_type lora \
    --checkpoint_dir path_to_checkpoint \
    --output_dir path_to_export
 ```
@@ -307,6 +341,7 @@ This repository is licensed under the [Apache-2.0 License](LICENSE).
 Please follow the model licenses to use the corresponding model weights:

 - [LLaMA](https://github.com/facebookresearch/llama/blob/main/MODEL_CARD.md)
+- [LLaMA-2](https://ai.meta.com/llama/license/)
 - [BLOOM](https://huggingface.co/spaces/bigscience/license)
 - [Falcon](LICENSE)
 - [baichuan](https://huggingface.co/baichuan-inc/baichuan-7B/resolve/main/baichuan-7B%20%E6%A8%A1%E5%9E%8B%E8%AE%B8%E5%8F%AF%E5%8D%8F%E8%AE%AE.pdf)
--- a/data/README.md
+++ b/data/README.md
@@ -1,4 +1,5 @@
-Data format in `dataset_info.json`:
+If you are using a custom dataset, please provide your dataset definition in the following format in `dataset_info.json`.
+
 ```json
 "dataset_name": {
    "hf_hub_url": "the name of the dataset repository on the HuggingFace hub. (if specified, ignore below 3 arguments)",
@@ -14,40 +15,4 @@ Data format in `dataset_info.json`:
 }
 ```

-`dataset_info.json` 中的数据集定义格式：
-```json
-"数据集名称": {
-    "hf_hub_url": "HuggingFace上的项目地址（若指定，则忽略下列三个参数）",
-    "script_url": "包含数据加载脚本的本地文件夹名称（若指定，则忽略下列两个参数）",
-    "file_name": "该目录下数据集文件的名称（若上述参数未指定，则此项必需）",
-    "file_sha1": "数据集文件的SHA-1哈希值（可选）",
-    "columns": {
-        "prompt": "数据集代表提示词的表头名称（默认：instruction）",
-        "query": "数据集代表请求的表头名称（默认：input）",
-        "response": "数据集代表回答的表头名称（默认：output）",
-        "history": "数据集代表历史对话的表头名称（默认：None）"
-    }
-}
-```
-
-部分预置数据集简介：
-
-| 数据集名称 | 规模 | 描述 |
-| --- | --- | --- |
-| [Stanford Alpaca](https://github.com/tatsu-lab/stanford_alpaca) | 52k | 斯坦福大学开源的 Alpaca 数据集，训练了 Alpaca 这类早期基于 LLaMA 的模型 |
-| [Stanford Alpaca (Chinese)](https://github.com/ymcui/Chinese-LLaMA-Alpaca) | 51k | 使用 ChatGPT 翻译的 Alpaca 数据集 |
-| [GPT-4 Generated Data](https://github.com/Instruction-Tuning-with-GPT-4/GPT-4-LLM) | 100k+ | 基于 GPT-4 的 self-instruction 数据集 |
-| [BELLE 2M](https://huggingface.co/datasets/BelleGroup/train_2M_CN) | 2m | 包含约 200 万条由 [BELLE](https://github.com/LianjiaTech/BELLE) 项目生成的中文指令数据 |
-| [BELLE 1M](https://huggingface.co/datasets/BelleGroup/train_1M_CN) | 1m | 包含约 100 万条由 [BELLE](https://github.com/LianjiaTech/BELLE) 项目生成的中文指令数据 |
-| [BELLE 0.5M](https://huggingface.co/datasets/BelleGroup/train_0.5M_CN) | 500k  | 包含约 50 万条由 [BELLE](https://github.com/LianjiaTech/BELLE) 项目生成的中文指令数据 |
-| [BELLE Dialogue 0.4M](https://huggingface.co/datasets/BelleGroup/generated_chat_0.4M) | 400k | 包含约 40 万条由 [BELLE](https://github.com/LianjiaTech/BELLE) 项目生成的个性化角色对话数据，包含角色介绍 |
-| [BELLE School Math 0.25M](https://huggingface.co/datasets/BelleGroup/school_math_0.25M) | 250k  | 包含约 25 万条由 [BELLE](https://github.com/LianjiaTech/BELLE) 项目生成的中文数学题数据，包含解题过程 |
-| [BELLE Multiturn Chat 0.8M](https://huggingface.co/datasets/BelleGroup/multiturn_chat_0.8M) | 800k | 包含约 80 万条由 [BELLE](https://github.com/LianjiaTech/BELLE) 项目生成的用户与助手的多轮对话 |
-| [Guanaco Dataset](https://huggingface.co/datasets/JosephusCheung/GuanacoDataset) | 100k+ | 包含日文、简繁体中文、英文等多类数据，数据集原用于 Guanaco 模型训练 |
-| [Firefly 1.1M](https://huggingface.co/datasets/YeungNLP/firefly-train-1.1M) | 1.1M  | 中文对话大模型 firefly（流萤）的中文数据集，包含多个 NLP 任务 |
-| [CodeAlpaca 20k](https://huggingface.co/datasets/sahil2801/CodeAlpaca-20k) | 20k | 英文代码生成任务数据集 |
-| [Alpaca CoT](https://huggingface.co/datasets/QingyiSi/Alpaca-CoT) | 6M | 用于微调的指令数据集集合 |
-| [Web QA](https://huggingface.co/datasets/suolyer/webqa) | 36k | 百度知道汇集的中文问答数据集 |
-| [UltraChat](https://github.com/thunlp/UltraChat) | 1.57M | 清华 NLP 发布的大规模多轮对话数据集 |
-
-注：BELLE 数据集是由 ChatGPT 产生的数据集，不保证数据准确性，所有类 GPT 模型产生的 self-instruction 数据集均不能保证其准确性。
+where the `prompt` and `response` columns should contain non-empty values. The `query` column will be concatenated with the `prompt` column and used as input for the model. The `history` column should contain a list where each element is a string tuple representing a query-response pair.
--- a/data/pretrain_data/wiki_long.jsonl
+++ b/data/pretrain_data/wiki_long.jsonl
@@ -1,2 +0,0 @@
-{"id": 0,"title": "大卫·亨利","content": "大卫·亨利\n\n大卫·克莱顿·亨利（David Clayton Henrie，），美国演员。近来在迪士尼频道原创电视影集《少年魔法师》（Wizards of Waverly Place）当中演出贾斯汀·鲁索（Justin Russo）一角。\n\n大卫·亨利出生在加州Mission Viejo，在凤凰城长大。他的胞弟劳伦斯·亨利（Lorenzo Henrie）也是演员。大卫·亨利就读夏安传统学校。家中是信奉罗马天主教。 \n\n大卫在2007年拍摄少年魔法师期间认识女演员露西·海尔（Lucy Hale），之后与其交往，于2009年分手。\n\n10岁时，大卫·亨利和SAG在凤凰城签订了合约，并开始走出去试镜。 9岁的时候，在沙加缅度进行商业拍摄，SAG董事建议大卫·亨利搬到洛杉矶。在10岁那年夏天，他和他的家人搬到了好莱坞。他预定他的前2支商业试镜，扮演主要角色为汉堡王和桂格燕麦。他初演电视节目为Providence。 \n\n到了13岁，大卫有了他的第一次重大突破，在福克斯公司的喜剧The Pitts饰演 Petey Pitt一角。大卫下出作品为的Hallmark movie为Monster Maker，和琳达布莱儿、乔治甘迺迪共同演出，并要求回来Hallmark movie公司。 \n\n在18岁时，大卫得到了迪士尼频道原创系列演出机会，该节目2007年10月12日首播。大卫2008年参加了迪士尼频道的游戏节目。他是绿色团队的队长，隔年，为旋风队队长。他在迪士尼原创电影《少年魔法师》之后在《酷爸的疯狂假期》中有饰演一角。\n"}
-{"id": 1,"title": "大卫·亨利","content": "大卫·亨利\n\n大卫·克莱顿·亨利（David Clayton Henrie，），美国演员。近来在迪士尼频道原创电视影集《少年魔法师》（Wizards of Waverly Place）当中演出贾斯汀·鲁索（Justin Russo）一角。\n\n大卫·亨利出生在加州Mission Viejo，在凤凰城长大。他的胞弟劳伦斯·亨利（Lorenzo Henrie）也是演员。大卫·亨利就读夏安传统学校。家中是信奉罗马天主教。 \n\n大卫在2007年拍摄少年魔法师期间认识女演员露西·海尔（Lucy Hale），之后与其交往，于2009年分手。\n\n10岁时，大卫·亨利和SAG在凤凰城签订了合约，并开始走出去试镜。 9岁的时候，在沙加缅度进行商业拍摄，SAG董事建议大卫·亨利搬到洛杉矶。在10岁那年夏天，他和他的家人搬到了好莱坞。他预定他的前2支商业试镜，扮演主要角色为汉堡王和桂格燕麦。他初演电视节目为Providence。 \n\n到了13岁，大卫有了他的第一次重大突破，在福克斯公司的喜剧The Pitts饰演 Petey Pitt一角。大卫下出作品为的Hallmark movie为Monster Maker，和琳达布莱儿、乔治甘迺迪共同演出，并要求回来Hallmark movie公司。 \n\n在18岁时，大卫得到了迪士尼频道原创系列演出机会，该节目2007年10月12日首播。大卫2008年参加了迪士尼频道的游戏节目。他是绿色团队的队长，隔年，为旋风队队长。他在迪士尼原创电影《少年魔法师》之后在《酷爸的疯狂假期》中有饰演一角。\n"}
--- a/data/refgpt_zh_50k_p1.json.REMOVED.git-id
+++ b/data/refgpt_zh_50k_p1.json.REMOVED.git-id
@@ -0,0 +1 @@
+56405bb8f52727e52e99693739494b9b7b0d7ba6
--- a/data/refgpt_zh_50k_p2.json.REMOVED.git-id
+++ b/data/refgpt_zh_50k_p2.json.REMOVED.git-id
@@ -0,0 +1 @@
+fa935248a5d40d2bdd5649af99a72a754d40ae7a
--- a/data/sharegpt_zh_27k.json.REMOVED.git-id
+++ b/data/sharegpt_zh_27k.json.REMOVED.git-id
@@ -0,0 +1 @@
+38c89869c6aeca2a3af9ea1e09afe460f9b46810
--- a/requirements.txt
+++ b/requirements.txt
@@ -10,7 +10,7 @@ rouge-chinese
 nltk
 gradio>=3.36.0
 uvicorn
-pydantic
-fastapi
+pydantic==1.10.11
+fastapi==0.95.1
 sse-starlette
 matplotlib
--- a/src/api_demo.py
+++ b/src/api_demo.py
@@ -5,9 +5,16 @@

 import uvicorn

-from llmtuner import create_app
+from llmtuner import ChatModel
+from llmtuner.api.app import create_app
+from llmtuner.tuner import get_infer_args
+
+
+def main():
+    chat_model = ChatModel(*get_infer_args())
+    app = create_app(chat_model)
+    uvicorn.run(app, host="0.0.0.0", port=8000, workers=1)


 if __name__ == "__main__":
-    app = create_app()
-    uvicorn.run(app, host="0.0.0.0", port=8000, workers=1)
+    main()
--- a/src/cli_demo.py
+++ b/src/cli_demo.py
@@ -2,7 +2,8 @@
 # Implements stream chat in command line for fine-tuned models.
 # Usage: python cli_demo.py --model_name_or_path path_to_model --checkpoint_dir path_to_checkpoint

-from llmtuner import ChatModel, get_infer_args
+from llmtuner import ChatModel
+from llmtuner.tuner import get_infer_args


 def main():
--- a/src/export_model.py
+++ b/src/export_model.py
@@ -2,7 +2,7 @@
 # Exports the fine-tuned model.
 # Usage: python export_model.py --checkpoint_dir path_to_checkpoint --output_dir path_to_save_model

-from llmtuner import get_train_args, load_model_and_tokenizer
+from llmtuner.tuner import get_train_args, load_model_and_tokenizer


 def main():
--- a/src/llmtuner/init.py
+++ b/src/llmtuner/init.py
@@ -1,7 +1,4 @@
-from llmtuner.api import create_app
 from llmtuner.chat import ChatModel
-from llmtuner.tuner import get_train_args, get_infer_args, load_model_and_tokenizer, run_pt, run_sft, run_rm, run_ppo
-from llmtuner.webui import create_ui


-__version__ = "0.1.0"
+__version__ = "0.1.2"
--- a/src/llmtuner/api/init.py
+++ b/src/llmtuner/api/init.py
@@ -1 +0,0 @@
-from llmtuner.api.app import create_app
--- a/src/llmtuner/api/app.py
+++ b/src/llmtuner/api/app.py
@@ -1,4 +1,3 @@
-import json
 import uvicorn
 from fastapi import FastAPI, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
@@ -10,6 +9,8 @@ from llmtuner.tuner import get_infer_args
 from llmtuner.extras.misc import torch_gc
 from llmtuner.chat.stream_chat import ChatModel
 from llmtuner.api.protocol import (
+    Role,
+    Finish,
    ModelCard,
    ModelList,
    ChatMessage,
@@ -29,9 +30,7 @@ async def lifespan(app: FastAPI): # collects GPU memory
    torch_gc()


-def create_app():
-    chat_model = ChatModel(*get_infer_args())
-
+def create_app(chat_model: ChatModel) -> FastAPI:
    app = FastAPI(lifespan=lifespan)

    app.add_middleware(
@@ -49,12 +48,12 @@ def create_app():

    @app.post("/v1/chat/completions", response_model=ChatCompletionResponse)
    async def create_chat_completion(request: ChatCompletionRequest):
-        if request.messages[-1].role != "user":
+        if request.messages[-1].role != Role.USER:
            raise HTTPException(status_code=400, detail="Invalid request")
        query = request.messages[-1].content

        prev_messages = request.messages[:-1]
-        if len(prev_messages) > 0 and prev_messages[0].role == "system":
+        if len(prev_messages) > 0 and prev_messages[0].role == Role.SYSTEM:
            prefix = prev_messages.pop(0).content
        else:
            prefix = None
@@ -62,7 +61,7 @@ def create_app():
        history = []
        if len(prev_messages) % 2 == 0:
            for i in range(0, len(prev_messages), 2):
-                if prev_messages[i].role == "user" and prev_messages[i+1].role == "assistant":
+                if prev_messages[i].role == Role.USER and prev_messages[i+1].role == Role.ASSISTANT:
                    history.append([prev_messages[i].content, prev_messages[i+1].content])

        if request.stream:
@@ -81,20 +80,20 @@ def create_app():

        choice_data = ChatCompletionResponseChoice(
            index=0,
-            message=ChatMessage(role="assistant", content=response),
-            finish_reason="stop"
+            message=ChatMessage(role=Role.ASSISTANT, content=response),
+            finish_reason=Finish.STOP
        )

-        return ChatCompletionResponse(model=request.model, choices=[choice_data], usage=usage, object="chat.completion")
+        return ChatCompletionResponse(model=request.model, choices=[choice_data], usage=usage)

    async def predict(query: str, history: List[Tuple[str, str]], prefix: str, request: ChatCompletionRequest):
        choice_data = ChatCompletionResponseStreamChoice(
            index=0,
-            delta=DeltaMessage(role="assistant"),
+            delta=DeltaMessage(role=Role.ASSISTANT),
            finish_reason=None
        )
-        chunk = ChatCompletionStreamResponse(model=request.model, choices=[choice_data], object="chat.completion.chunk")
-        yield json.dumps(chunk, ensure_ascii=False)
+        chunk = ChatCompletionStreamResponse(model=request.model, choices=[choice_data])
+        yield chunk.json(exclude_unset=True, ensure_ascii=False)

        for new_text in chat_model.stream_chat(
            query, history, prefix, temperature=request.temperature, top_p=request.top_p, max_new_tokens=request.max_tokens
@@ -107,21 +106,22 @@ def create_app():
                delta=DeltaMessage(content=new_text),
                finish_reason=None
            )
-            chunk = ChatCompletionStreamResponse(model=request.model, choices=[choice_data], object="chat.completion.chunk")
-            yield json.dumps(chunk, ensure_ascii=False)
+            chunk = ChatCompletionStreamResponse(model=request.model, choices=[choice_data])
+            yield chunk.json(exclude_unset=True, ensure_ascii=False)

        choice_data = ChatCompletionResponseStreamChoice(
            index=0,
            delta=DeltaMessage(),
-            finish_reason="stop"
+            finish_reason=Finish.STOP
        )
-        chunk = ChatCompletionStreamResponse(model=request.model, choices=[choice_data], object="chat.completion.chunk")
-        yield json.dumps(chunk, ensure_ascii=False)
+        chunk = ChatCompletionStreamResponse(model=request.model, choices=[choice_data])
+        yield chunk.json(exclude_unset=True, ensure_ascii=False)
        yield "[DONE]"

    return app


 if __name__ == "__main__":
-    app = create_app()
+    chat_model = ChatModel(*get_infer_args())
+    app = create_app(chat_model)
    uvicorn.run(app, host="0.0.0.0", port=8000, workers=1)
--- a/src/llmtuner/api/protocol.py
+++ b/src/llmtuner/api/protocol.py
@@ -1,6 +1,18 @@
 import time
+from enum import Enum
 from pydantic import BaseModel, Field
-from typing import List, Literal, Optional
+from typing import List, Optional
+
+
+class Role(str, Enum):
+    USER = "user"
+    ASSISTANT = "assistant"
+    SYSTEM = "system"
+
+
+class Finish(str, Enum):
+    STOP = "stop"
+    LENGTH = "length"


 class ModelCard(BaseModel):
@@ -19,12 +31,12 @@ class ModelList(BaseModel):


 class ChatMessage(BaseModel):
-    role: Literal["user", "assistant", "system"]
+    role: Role
    content: str


 class DeltaMessage(BaseModel):
-    role: Optional[Literal["user", "assistant", "system"]] = None
+    role: Optional[Role] = None
    content: Optional[str] = None


@@ -41,13 +53,13 @@ class ChatCompletionRequest(BaseModel):
 class ChatCompletionResponseChoice(BaseModel):
    index: int
    message: ChatMessage
-    finish_reason: Literal["stop", "length"]
+    finish_reason: Finish


 class ChatCompletionResponseStreamChoice(BaseModel):
    index: int
    delta: DeltaMessage
-    finish_reason: Optional[Literal["stop", "length"]] = None
+    finish_reason: Optional[Finish] = None


 class ChatCompletionResponseUsage(BaseModel):
@@ -58,7 +70,7 @@ class ChatCompletionResponseUsage(BaseModel):

 class ChatCompletionResponse(BaseModel):
    id: Optional[str] = "chatcmpl-default"
-    object: Literal["chat.completion"]
+    object: Optional[str] = "chat.completion"
    created: Optional[int] = Field(default_factory=lambda: int(time.time()))
    model: str
    choices: List[ChatCompletionResponseChoice]
@@ -67,7 +79,7 @@ class ChatCompletionResponse(BaseModel):

 class ChatCompletionStreamResponse(BaseModel):
    id: Optional[str] = "chatcmpl-default"
-    object: Literal["chat.completion.chunk"]
+    object: Optional[str] = "chat.completion.chunk"
    created: Optional[int] = Field(default_factory=lambda: int(time.time()))
    model: str
    choices: List[ChatCompletionResponseStreamChoice]
--- a/src/llmtuner/chat/stream_chat.py
+++ b/src/llmtuner/chat/stream_chat.py
@@ -4,7 +4,7 @@ from threading import Thread
 from transformers import TextIteratorStreamer

 from llmtuner.extras.misc import get_logits_processor
-from llmtuner.extras.template import Template
+from llmtuner.extras.template import get_template
 from llmtuner.hparams import ModelArguments, DataArguments, FinetuningArguments, GeneratingArguments
 from llmtuner.tuner import load_model_and_tokenizer

@@ -19,14 +19,14 @@ class ChatModel:
        generating_args: GeneratingArguments
    ) -> None:
        self.model, self.tokenizer = load_model_and_tokenizer(model_args, finetuning_args)
-        self.template = Template(data_args.prompt_template)
-        self.source_prefix = data_args.source_prefix if data_args.source_prefix else ""
+        self.template = get_template(data_args.prompt_template)
+        self.source_prefix = data_args.source_prefix or ""
        self.generating_args = generating_args

    def process_args(
        self, query: str, history: Optional[List[Tuple[str, str]]] = None, prefix: Optional[str] = None, **input_kwargs
    ) -> Tuple[Dict[str, Any], int]:
-        prefix = prefix if prefix else self.source_prefix
+        prefix = prefix or self.source_prefix

        inputs = self.tokenizer([self.template.get_prompt(query, history, prefix)], return_tensors="pt")
        inputs = inputs.to(self.model.device)
@@ -81,5 +81,4 @@ class ChatModel:
        thread = Thread(target=self.model.generate, kwargs=gen_kwargs)
        thread.start()

-        for new_text in streamer:
-            yield new_text
+        yield from streamer
--- a/src/llmtuner/dsets/init.py
+++ b/src/llmtuner/dsets/init.py
@@ -1,2 +1,3 @@
 from llmtuner.dsets.loader import get_dataset
 from llmtuner.dsets.preprocess import preprocess_dataset
+from llmtuner.dsets.utils import split_dataset
--- a/src/llmtuner/dsets/callbacks.py
+++ b/src/llmtuner/dsets/callbacks.py
@@ -1,63 +0,0 @@
-import os
-import json
-import time
-from datetime import timedelta
-
-from transformers import (
-    TrainerCallback,
-    TrainerControl,
-    TrainerState,
-    TrainingArguments
-)
-
-
-class LogCallback(TrainerCallback):
-
-    def __init__(self, runner=None):
-        self.runner = runner
-        self.start_time = time.time()
-        self.tracker = {}
-
-    def on_step_begin(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
-        r"""
-        Event called at the beginning of a training step. If using gradient accumulation, one training step
-        might take several inputs.
-        """
-        if self.runner is not None and self.runner.aborted:
-            control.should_epoch_stop = True
-            control.should_training_stop = True
-
-    def on_substep_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
-        r"""
-        Event called at the end of an substep during gradient accumulation.
-        """
-        if self.runner is not None and self.runner.aborted:
-            control.should_epoch_stop = True
-            control.should_training_stop = True
-
-    def on_log(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs) -> None:
-        r"""
-        Event called after logging the last logs.
-        """
-        if "loss" not in state.log_history[-1]:
-            return
-        cur_time = time.time()
-        cur_steps = state.log_history[-1].get("step")
-        elapsed_time = cur_time - self.start_time
-        avg_time_per_step = elapsed_time / cur_steps if cur_steps != 0 else 0
-        remaining_steps = state.max_steps - cur_steps
-        remaining_time = remaining_steps * avg_time_per_step
-        self.tracker = {
-            "current_steps": cur_steps,
-            "total_steps": state.max_steps,
-            "loss": state.log_history[-1].get("loss", None),
-            "reward": state.log_history[-1].get("reward", None),
-            "learning_rate": state.log_history[-1].get("learning_rate", None),
-            "epoch": state.log_history[-1].get("epoch", None),
-            "percentage": round(cur_steps / state.max_steps * 100, 2) if state.max_steps != 0 else 100,
-            "elapsed_time": str(timedelta(seconds=int(elapsed_time))),
-            "remaining_time": str(timedelta(seconds=int(remaining_time)))
-        }
-        os.makedirs(args.output_dir, exist_ok=True)
-        with open(os.path.join(args.output_dir, "trainer_log.jsonl"), "a", encoding="utf-8") as f:
-            f.write(json.dumps(self.tracker) + "\n")
--- a/src/llmtuner/dsets/preprocess.py
+++ b/src/llmtuner/dsets/preprocess.py
@@ -6,7 +6,7 @@ from transformers.tokenization_utils import PreTrainedTokenizer
 from datasets import Dataset

 from llmtuner.extras.constants import IGNORE_INDEX
-from llmtuner.extras.template import Template
+from llmtuner.extras.template import get_template
 from llmtuner.hparams import DataArguments


@@ -19,7 +19,7 @@ def preprocess_dataset(
 ) -> Dataset:

    column_names = list(dataset.column_names)
-    prompt_template = Template(data_args.prompt_template)
+    prompt_template = get_template(data_args.prompt_template)

    # support question with a single answer or multiple answers
    def get_dialog(examples):
--- a/src/llmtuner/dsets/utils.py
+++ b/src/llmtuner/dsets/utils.py
@@ -0,0 +1,16 @@
+from typing import Dict
+from datasets import Dataset
+
+
+def split_dataset(
+    dataset: Dataset, dev_ratio: float, do_train: bool
+) -> Dict[str, Dataset]:
+    # Split the dataset
+    if do_train:
+        if dev_ratio > 1e-6:
+            dataset = dataset.train_test_split(test_size=dev_ratio)
+            return {"train_dataset": dataset["train"], "eval_dataset": dataset["test"]}
+        else:
+            return {"train_dataset": dataset}
+    else: # do_eval or do_predict
+        return {"eval_dataset": dataset}
--- a/src/llmtuner/extras/callbacks.py
+++ b/src/llmtuner/extras/callbacks.py
@@ -47,6 +47,9 @@ class LogCallback(TrainerCallback):
        r"""
        Event called after logging the last logs.
        """
+        if not state.is_world_process_zero:
+            return
+
        cur_time = time.time()
        cur_steps = state.log_history[-1].get("step")
        elapsed_time = cur_time - self.start_time
--- a/src/llmtuner/extras/constants.py
+++ b/src/llmtuner/extras/constants.py
@@ -13,6 +13,12 @@ SUPPORTED_MODELS = {
    "LLaMA-13B": "huggyllama/llama-13b",
    "LLaMA-30B": "huggyllama/llama-30b",
    "LLaMA-65B": "huggyllama/llama-65b",
+    "LLaMA2-7B": "meta-llama/Llama-2-7b-hf",
+    "LLaMA2-13B": "meta-llama/Llama-2-13b-hf",
+    "LLaMA2-70B": "meta-llama/Llama-2-70b-hf",
+    "LLaMA2-7B-Chat": "meta-llama/Llama-2-7b-chat-hf",
+    "LLaMA2-13B-Chat": "meta-llama/Llama-2-13b-chat-hf",
+    "LLaMA2-70B-Chat": "meta-llama/Llama-2-70b-chat-hf",
    "BLOOM-560M": "bigscience/bloom-560m",
    "BLOOM-3B": "bigscience/bloom-3b",
    "BLOOM-7B1": "bigscience/bloom-7b1",
@@ -30,8 +36,9 @@ SUPPORTED_MODELS = {
    "InternLM-7B-Chat": "internlm/internlm-chat-7b"
 }

-DEFAULT_MODULE = { # will be deprecated
+DEFAULT_MODULE = {
    "LLaMA": "q_proj,v_proj",
+    "LLaMA2": "q_proj,v_proj",
    "BLOOM": "query_key_value",
    "BLOOMZ": "query_key_value",
    "Falcon": "query_key_value",
--- a/src/llmtuner/extras/template.py
+++ b/src/llmtuner/extras/template.py
@@ -3,30 +3,13 @@ from dataclasses import dataclass


@dataclass
-class Format:
+class Template:
+
    prefix: str
    prompt: str
    sep: str
    use_history: bool

-
-templates: Dict[str, Format] = {}
-
-
-@dataclass
-class Template:
-
-    name: str
-
-    def __post_init__(self):
-        if self.name in templates:
-            self.prefix = templates[self.name].prefix
-            self.prompt = templates[self.name].prompt
-            self.sep = templates[self.name].sep
-            self.use_history = templates[self.name].use_history
-        else:
-            raise ValueError("Template {} does not exist.".format(self.name))
-
    def get_prompt(
        self, query: str, history: Optional[List[Tuple[str, str]]] = None, prefix: Optional[str] = ""
    ) -> str:
@@ -46,7 +29,7 @@ class Template:
    def _format_example(
        self, query: str, history: Optional[List[Tuple[str, str]]] = None, prefix: Optional[str] = ""
    ) -> List[str]:
-        prefix = prefix if prefix else self.prefix # use prefix if provided
+        prefix = prefix or self.prefix # use prefix if provided
        prefix = prefix + self.sep if prefix else "" # add separator for non-empty prefix
        history = history if (history and self.use_history) else []
        history = history + [(query, "<dummy>")]
@@ -61,8 +44,11 @@ class Template:
        return convs[:-1] # drop last


+templates: Dict[str, Template] = {}
+
+
 def register_template(name: str, prefix: str, prompt: str, sep: str, use_history: bool) -> None:
-    templates[name] = Format(
+    templates[name] = Template(
        prefix=prefix,
        prompt=prompt,
        sep=sep,
@@ -70,6 +56,12 @@ def register_template(name: str, prefix: str, prompt: str, sep: str, use_history
    )


+def get_template(name: str) -> Template:
+    template = templates.get(name, None)
+    assert template is not None, "Template {} does not exist.".format(name)
+    return template
+
+
 r"""
 Supports language model inference without histories.
 """
@@ -95,6 +87,27 @@ register_template(
 )


+r"""
+Supports: https://huggingface.co/meta-llama/Llama-2-7b-chat-hf
+          https://huggingface.co/meta-llama/Llama-2-13b-chat-hf
+          https://huggingface.co/meta-llama/Llama-2-70b-chat-hf
+"""
+register_template(
+    name="llama2",
+    prefix="<<SYS>>\nYou are a helpful, respectful and honest assistant. "
+           "Always answer as helpfully as possible, while being safe.  "
+           "Your answers should not include any harmful, unethical, "
+           "racist, sexist, toxic, dangerous, or illegal content. "
+           "Please ensure that your responses are socially unbiased and positive in nature.\n"
+           "If a question does not make any sense, or is not factually coherent, "
+           "explain why instead of answering something not correct. "
+           "If you don't know the answer to a question, please don't share false information.\n<</SYS>>\n\n",
+    prompt=" [INST] {query} [/INST] ",
+    sep="</s>",
+    use_history=True
+)
+
+
 r"""
 Supports: https://huggingface.co/tatsu-lab/alpaca-7b-wdiff
          https://github.com/ymcui/Chinese-LLaMA-Alpaca
@@ -203,6 +216,19 @@ register_template(
    name="baichuan",
    prefix="",
    prompt="<reserved_102>{query}<reserved_103>",
-    sep="",
+    sep="</s>",
+    use_history=True
+)
+
+
+r"""
+Supports: https://huggingface.co/HuggingFaceH4/starchat-alpha
+          https://huggingface.co/HuggingFaceH4/starchat-beta
+"""
+register_template(
+    name="starchat",
+    prefix="<|system|>\n",
+    prompt="<|user|>\n{query}<|end|>\n<|assistant|>\n",
+    sep="<|end|>\n",
    use_history=True
 )
--- a/src/llmtuner/hparams/finetuning_args.py
+++ b/src/llmtuner/hparams/finetuning_args.py
@@ -16,9 +16,10 @@ class FinetuningArguments:
        default=32,
        metadata={"help": "Number of decoder blocks in the model. \
                  LLaMA choices: [\"32\", \"40\", \"60\", \"80\"], \
+                  LLaMA-2 choices: [\"32\", \"40\", \"80\"], \
                  BLOOM choices: [\"24\", \"30\", \"70\"], \
                  Falcon choices: [\"32\", \"60\"], \
-                  Baichuan choices: [\"32\"]"}
+                  Baichuan choices: [\"32\", \"40\"]"}
    )
    num_layer_trainable: Optional[int] = field(
        default=3,
@@ -27,7 +28,7 @@ class FinetuningArguments:
    name_module_trainable: Optional[Literal["mlp", "self_attn", "self_attention"]] = field(
        default="mlp",
        metadata={"help": "Name of trainable modules for Freeze fine-tuning. \
-                  LLaMA choices: [\"mlp\", \"self_attn\"], \
+                  LLaMA & LLaMA-2 choices: [\"mlp\", \"self_attn\"], \
                  BLOOM & Falcon choices: [\"mlp\", \"self_attention\"], \
                  Baichuan choices: [\"mlp\", \"self_attn\"]"}
    )
@@ -46,7 +47,7 @@ class FinetuningArguments:
    lora_target: Optional[str] = field(
        default="q_proj,v_proj",
        metadata={"help": "Name(s) of target modules to apply LoRA. Use commas to separate multiple modules. \
-                  LLaMA choices: [\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\", \"gate_proj\", \"up_proj\", \"down_proj\"], \
+                  LLaMA & LLaMA-2 choices: [\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\", \"gate_proj\", \"up_proj\", \"down_proj\"], \
                  BLOOM & Falcon choices: [\"query_key_value\", \"self_attention.dense\", \"mlp.dense\"], \
                  Baichuan choices: [\"W_pack\", \"o_proj\", \"gate_proj\", \"up_proj\", \"down_proj\"]"}
    )
--- a/src/llmtuner/tuner/core/loader.py
+++ b/src/llmtuner/tuner/core/loader.py
@@ -10,6 +10,7 @@ from transformers import (
 )
 from transformers.utils import check_min_version
 from transformers.utils.versions import require_version
+from transformers.deepspeed import is_deepspeed_zero3_enabled
 from transformers.modeling_utils import PretrainedConfig, PreTrainedModel
 from transformers.tokenization_utils import PreTrainedTokenizerBase
 from trl import AutoModelForCausalLMWithValueHead
@@ -108,7 +109,7 @@ def load_model_and_tokenizer(
        model_to_load,
        config=config,
        torch_dtype=torch.bfloat16 if model_args.compute_dtype == torch.bfloat16 else torch.float16,
-        low_cpu_mem_usage=True,
+        low_cpu_mem_usage=(not is_deepspeed_zero3_enabled()),
        **config_kwargs
    )

--- a/src/llmtuner/tuner/ppo/trainer.py
+++ b/src/llmtuner/tuner/ppo/trainer.py
@@ -107,8 +107,12 @@ class PPOPeftTrainer(PPOTrainer, PeftTrainer):
            # Compute rewards
            replace_model(unwrapped_model, target="reward")
            with torch.no_grad():
-                _, _, values = self.model(**self.prepare_model_inputs(queries, responses))
-            rewards = [reward for reward in values[-1].to(torch.float32)] # use float32 type
+                _, _, values = self.model(
+                    **self.prepare_model_inputs(queries, responses),
+                    output_hidden_states=True,
+                    return_dict=True
+                )
+            rewards = [reward for reward in values[:, -1].to(torch.float32)] # use float32 type
            replace_model(unwrapped_model, target="default")

            # Run PPO step
--- a/src/llmtuner/tuner/pt/workflow.py
+++ b/src/llmtuner/tuner/pt/workflow.py
@@ -4,7 +4,7 @@ import math
 from typing import Optional, List
 from transformers import Seq2SeqTrainingArguments, DataCollatorForSeq2Seq, TrainerCallback

-from llmtuner.dsets import get_dataset, preprocess_dataset
+from llmtuner.dsets import get_dataset, preprocess_dataset, split_dataset
 from llmtuner.extras.callbacks import LogCallback
 from llmtuner.extras.constants import IGNORE_INDEX
 from llmtuner.extras.ploting import plot_loss
@@ -28,16 +28,6 @@ def run_pt(
        label_pad_token_id=IGNORE_INDEX if data_args.ignore_pad_token_for_loss else tokenizer.pad_token_id
    )

-    # Split the dataset
-    if training_args.do_train:
-        if data_args.dev_ratio > 1e-6:
-            dataset = dataset.train_test_split(test_size=data_args.dev_ratio)
-            trainer_kwargs = {"train_dataset": dataset["train"], "eval_dataset": dataset["test"]}
-        else:
-            trainer_kwargs = {"train_dataset": dataset}
-    else: # do_eval or do_predict
-        trainer_kwargs = {"eval_dataset": dataset}
-
    # Initialize our Trainer
    trainer = PeftTrainer(
        finetuning_args=finetuning_args,
@@ -46,7 +36,7 @@ def run_pt(
        tokenizer=tokenizer,
        data_collator=data_collator,
        callbacks=callbacks,
-        **trainer_kwargs
+        **split_dataset(dataset, data_args.dev_ratio, training_args.do_train)
    )

    # Training
--- a/src/llmtuner/tuner/rm/trainer.py
+++ b/src/llmtuner/tuner/rm/trainer.py
@@ -32,7 +32,7 @@ class PairwisePeftTrainer(PeftTrainer):
        See: https://github.com/huggingface/transformers/blob/v4.30.2/src/transformers/trainer.py#L3509
        """
        batch_size = inputs["input_ids"].size(0) // 2
-        _, _, values = model(**inputs)
+        _, _, values = model(**inputs, output_hidden_states=True, return_dict=True)
        r_accept, r_reject = values[:, -1].split(batch_size, dim=0)
        loss = -torch.log(torch.sigmoid(r_accept - r_reject)).mean()
        return (loss, [loss, r_accept, r_reject]) if return_outputs else loss
--- a/src/llmtuner/tuner/rm/workflow.py
+++ b/src/llmtuner/tuner/rm/workflow.py
@@ -5,7 +5,7 @@
 from typing import Optional, List
 from transformers import Seq2SeqTrainingArguments, TrainerCallback

-from llmtuner.dsets import get_dataset, preprocess_dataset
+from llmtuner.dsets import get_dataset, preprocess_dataset, split_dataset
 from llmtuner.extras.callbacks import LogCallback
 from llmtuner.extras.ploting import plot_loss
 from llmtuner.hparams import ModelArguments, DataArguments, FinetuningArguments
@@ -29,16 +29,6 @@ def run_rm(

    training_args.remove_unused_columns = False # important for pairwise dataset

-    # Split the dataset
-    if training_args.do_train:
-        if data_args.dev_ratio > 1e-6:
-            dataset = dataset.train_test_split(test_size=data_args.dev_ratio)
-            trainer_kwargs = {"train_dataset": dataset["train"], "eval_dataset": dataset["test"]}
-        else:
-            trainer_kwargs = {"train_dataset": dataset}
-    else: # do_eval or do_predict
-        trainer_kwargs = {"eval_dataset": dataset}
-
    # Initialize our Trainer
    trainer = PairwisePeftTrainer(
        finetuning_args=finetuning_args,
@@ -48,7 +38,7 @@ def run_rm(
        data_collator=data_collator,
        callbacks=callbacks,
        compute_metrics=compute_accuracy,
-        **trainer_kwargs
+        **split_dataset(dataset, data_args.dev_ratio, training_args.do_train)
    )

    # Training
--- a/src/llmtuner/tuner/sft/metric.py
+++ b/src/llmtuner/tuner/sft/metric.py
@@ -23,7 +23,7 @@ class ComputeMetrics:
        Uses the model predictions to compute metrics.
        """
        preds, labels = eval_preds
-        score_dict = {"rouge-1": [], "rouge-2": [], "rouge-l": [], "bleu-4": []}
+        score_dict = {"accuracy": [], "rouge-1": [], "rouge-2": [], "rouge-l": [], "bleu-4": []}

        preds = np.where(preds != IGNORE_INDEX, preds, self.tokenizer.pad_token_id)
        labels = np.where(labels != IGNORE_INDEX, labels, self.tokenizer.pad_token_id)
@@ -47,5 +47,6 @@ class ComputeMetrics:

            bleu_score = sentence_bleu([list(label)], list(pred), smoothing_function=SmoothingFunction().method3)
            score_dict["bleu-4"].append(round(bleu_score * 100, 4))
+            score_dict["accuracy"].append(float(len(label) != 0 and pred[:len(label)] == label))

        return {k: float(np.mean(v)) for k, v in score_dict.items()}
--- a/src/llmtuner/tuner/sft/trainer.py
+++ b/src/llmtuner/tuner/sft/trainer.py
@@ -36,31 +36,44 @@ class Seq2SeqPeftTrainer(PeftTrainer):
            inputs["labels"] = self._pad_tensors_to_target_len(inputs["labels"], inputs["input_ids"])
        if label_len > prompt_len:
            inputs["input_ids"] = self._pad_tensors_to_target_len(inputs["input_ids"], inputs["labels"])
+            if "attention_mask" in inputs:
+                inputs["attention_mask"] = self._pad_tensors_to_target_len(
+                    inputs["attention_mask"], inputs["labels"], pad_token_id=0
+                )
+            if "position_ids" in inputs:
+                inputs["position_ids"] = self._pad_tensors_to_target_len(
+                    inputs["position_ids"], inputs["labels"], pad_token_id=0
+                )

        loss, generated_tokens, labels = super().prediction_step(
            model, inputs, prediction_loss_only=prediction_loss_only, ignore_keys=ignore_keys
        )
-        generated_tokens = generated_tokens[:, max(prompt_len, label_len):] if generated_tokens is not None else None
+        generated_tokens = (
+            generated_tokens[:, max(prompt_len, label_len):] if generated_tokens is not None else None
+        )

        return (loss, generated_tokens, labels)

-    def _pad_tensors_to_target_len(self, src_tensor: torch.Tensor, tgt_tensor: torch.Tensor) -> torch.Tensor:
+    def _pad_tensors_to_target_len(
+        self,
+        src_tensor: torch.Tensor,
+        tgt_tensor: torch.Tensor,
+        pad_token_id: Optional[int] = None
+    ) -> torch.Tensor:
        r"""
        Pads the tensor to the same length as the target tensor.

        Should only be called when predict_with_generate=True.
        """
+        if pad_token_id is None:
            if self.tokenizer is not None and hasattr(self.tokenizer, "pad_token_id"):
                assert self.tokenizer.padding_side == "left", "This method only accepts left-padded tensor."
-            # If PAD token is not defined at least EOS token has to be defined
-            pad_token_id = (
-                self.tokenizer.pad_token_id if self.tokenizer.pad_token_id is not None else self.tokenizer.eos_token_id
-            )
+                pad_token_id = self.tokenizer.pad_token_id
            else:
                if self.model.config.pad_token_id is not None:
                    pad_token_id = self.model.config.pad_token_id
                else:
-                raise ValueError("Pad_token_id must be set in the configuration of the model, in order to pad tensors")
+                    raise ValueError("Pad_token_id must be set in the configuration of the model.")

        padded_tensor = pad_token_id * torch.ones_like(tgt_tensor)
        padded_tensor[:, -src_tensor.shape[-1]:] = src_tensor # adopt left-padding
--- a/src/llmtuner/tuner/sft/workflow.py
+++ b/src/llmtuner/tuner/sft/workflow.py
@@ -3,7 +3,7 @@
 from typing import Optional, List
 from transformers import Seq2SeqTrainingArguments, DataCollatorForSeq2Seq, TrainerCallback

-from llmtuner.dsets import get_dataset, preprocess_dataset
+from llmtuner.dsets import get_dataset, preprocess_dataset, split_dataset
 from llmtuner.extras.callbacks import LogCallback
 from llmtuner.extras.constants import IGNORE_INDEX
 from llmtuner.extras.misc import get_logits_processor
@@ -35,16 +35,6 @@ def run_sft(
    training_args.generation_num_beams = data_args.eval_num_beams if \
                data_args.eval_num_beams is not None else training_args.generation_num_beams

-    # Split the dataset
-    if training_args.do_train:
-        if data_args.dev_ratio > 1e-6:
-            dataset = dataset.train_test_split(test_size=data_args.dev_ratio)
-            trainer_kwargs = {"train_dataset": dataset["train"], "eval_dataset": dataset["test"]}
-        else:
-            trainer_kwargs = {"train_dataset": dataset}
-    else: # do_eval or do_predict
-        trainer_kwargs = {"eval_dataset": dataset}
-
    # Initialize our Trainer
    trainer = Seq2SeqPeftTrainer(
        finetuning_args=finetuning_args,
@@ -54,7 +44,7 @@ def run_sft(
        data_collator=data_collator,
        callbacks=callbacks,
        compute_metrics=ComputeMetrics(tokenizer) if training_args.predict_with_generate else None,
-        **trainer_kwargs
+        **split_dataset(dataset, data_args.dev_ratio, training_args.do_train)
    )

    # Keyword arguments for `model.generate`
--- a/src/llmtuner/webui/init.py
+++ b/src/llmtuner/webui/init.py
@@ -1 +0,0 @@
-from llmtuner.webui.interface import create_ui
--- a/src/llmtuner/webui/chat.py
+++ b/src/llmtuner/webui/chat.py
@@ -11,14 +11,22 @@ from llmtuner.webui.locales import ALERTS

 class WebChatModel(ChatModel):

-    def __init__(self):
+    def __init__(self, *args):
        self.model = None
        self.tokenizer = None
        self.generating_args = GeneratingArguments()
+        if len(args) != 0:
+            super().__init__(*args)

    def load_model(
-        self, lang: str, model_name: str, checkpoints: list,
-        finetuning_type: str, template: str, quantization_bit: str
+        self,
+        lang: str,
+        model_name: str,
+        checkpoints: List[str],
+        finetuning_type: str,
+        quantization_bit: str,
+        template: str,
+        source_prefix: str
    ):
        if self.model is not None:
            yield ALERTS["err_exists"][lang]
@@ -43,10 +51,11 @@ class WebChatModel(ChatModel):
        yield ALERTS["info_loading"][lang]
        args = dict(
            model_name_or_path=model_name_or_path,
-            finetuning_type=finetuning_type,
-            prompt_template=template,
            checkpoint_dir=checkpoint_dir,
-            quantization_bit=int(quantization_bit) if quantization_bit else None
+            finetuning_type=finetuning_type,
+            quantization_bit=int(quantization_bit) if quantization_bit else None,
+            prompt_template=template,
+            source_prefix=source_prefix
        )
        super().__init__(*get_infer_args(args))

@@ -64,6 +73,7 @@ class WebChatModel(ChatModel):
        chatbot: List[Tuple[str, str]],
        query: str,
        history: List[Tuple[str, str]],
+        prefix: str,
        max_new_tokens: int,
        top_p: float,
        temperature: float
@@ -71,7 +81,7 @@ class WebChatModel(ChatModel):
        chatbot.append([query, ""])
        response = ""
        for new_text in self.stream_chat(
-            query, history, max_new_tokens=max_new_tokens, top_p=top_p, temperature=temperature
+            query, history, prefix, max_new_tokens=max_new_tokens, top_p=top_p, temperature=temperature
        ):
            response += new_text
            new_history = history + [(query, response)]
--- a/src/llmtuner/webui/components/chatbot.py
+++ b/src/llmtuner/webui/components/chatbot.py
@@ -1,4 +1,4 @@
-from typing import Dict, Tuple
+from typing import Dict, Optional, Tuple

 import gradio as gr
 from gradio.blocks import Block
@@ -8,17 +8,16 @@ from llmtuner.webui.chat import WebChatModel


 def create_chat_box(
-    chat_model: WebChatModel
+    chat_model: WebChatModel,
+    visible: Optional[bool] = False
 ) -> Tuple[Block, Component, Component, Dict[str, Component]]:
-    with gr.Box(visible=False) as chat_box:
+    with gr.Box(visible=visible) as chat_box:
        chatbot = gr.Chatbot()

        with gr.Row():
            with gr.Column(scale=4):
-                with gr.Column(scale=12):
+                prefix = gr.Textbox(show_label=False)
                query = gr.Textbox(show_label=False, lines=8)
-
-                with gr.Column(min_width=32, scale=1):
                submit_btn = gr.Button(variant="primary")

            with gr.Column(scale=1):
@@ -35,7 +34,7 @@ def create_chat_box(

    submit_btn.click(
        chat_model.predict,
-        [chatbot, query, history, max_new_tokens, top_p, temperature],
+        [chatbot, query, history, prefix, max_new_tokens, top_p, temperature],
        [chatbot, history],
        show_progress=True
    ).then(
@@ -45,6 +44,7 @@ def create_chat_box(
    clear_btn.click(lambda: ([], []), outputs=[chatbot, history], show_progress=True)

    return chat_box, chatbot, history, dict(
+        prefix=prefix,
        query=query,
        submit_btn=submit_btn,
        clear_btn=clear_btn,
--- a/src/llmtuner/webui/components/eval.py
+++ b/src/llmtuner/webui/components/eval.py
@@ -10,8 +10,8 @@ from llmtuner.webui.utils import can_preview, get_preview

 def create_eval_tab(top_elems: Dict[str, Component], runner: Runner) -> Dict[str, Component]:
    with gr.Row():
-        dataset_dir = gr.Textbox(value=DEFAULT_DATA_DIR, interactive=True, scale=2)
-        dataset = gr.Dropdown(multiselect=True, interactive=True, scale=4)
+        dataset_dir = gr.Textbox(value=DEFAULT_DATA_DIR, scale=2)
+        dataset = gr.Dropdown(multiselect=True, scale=4)
        preview_btn = gr.Button(interactive=False, scale=1)

    preview_box, preview_count, preview_samples, close_btn = create_preview_box()
@@ -21,23 +21,36 @@ def create_eval_tab(top_elems: Dict[str, Component], runner: Runner) -> Dict[str
    preview_btn.click(get_preview, [dataset_dir, dataset], [preview_count, preview_samples, preview_box])

    with gr.Row():
-        max_samples = gr.Textbox(value="100000", interactive=True)
-        batch_size = gr.Slider(value=8, minimum=1, maximum=128, step=1, interactive=True)
-        quantization_bit = gr.Dropdown([8, 4])
+        max_source_length = gr.Slider(value=512, minimum=4, maximum=4096, step=1)
+        max_target_length = gr.Slider(value=512, minimum=4, maximum=4096, step=1)
+        max_samples = gr.Textbox(value="100000")
+        batch_size = gr.Slider(value=8, minimum=1, maximum=512, step=1)
        predict = gr.Checkbox(value=True)

    with gr.Row():
        start_btn = gr.Button()
        stop_btn = gr.Button()

+    with gr.Box():
        output_box = gr.Markdown()

    start_btn.click(
        runner.run_eval,
        [
-            top_elems["lang"], top_elems["model_name"], top_elems["checkpoints"],
-            top_elems["finetuning_type"], top_elems["template"],
-            dataset, dataset_dir, max_samples, batch_size, quantization_bit, predict
+            top_elems["lang"],
+            top_elems["model_name"],
+            top_elems["checkpoints"],
+            top_elems["finetuning_type"],
+            top_elems["quantization_bit"],
+            top_elems["template"],
+            top_elems["source_prefix"],
+            dataset_dir,
+            dataset,
+            max_source_length,
+            max_target_length,
+            max_samples,
+            batch_size,
+            predict
        ],
        [output_box]
    )
@@ -50,9 +63,10 @@ def create_eval_tab(top_elems: Dict[str, Component], runner: Runner) -> Dict[str
        preview_count=preview_count,
        preview_samples=preview_samples,
        close_btn=close_btn,
+        max_source_length=max_source_length,
+        max_target_length=max_target_length,
        max_samples=max_samples,
        batch_size=batch_size,
-        quantization_bit=quantization_bit,
        predict=predict,
        start_btn=start_btn,
        stop_btn=stop_btn,
--- a/src/llmtuner/webui/components/infer.py
+++ b/src/llmtuner/webui/components/infer.py
@@ -11,9 +11,8 @@ def create_infer_tab(top_elems: Dict[str, Component]) -> Dict[str, Component]:
    with gr.Row():
        load_btn = gr.Button()
        unload_btn = gr.Button()
-        quantization_bit = gr.Dropdown([8, 4])

-    info_box = gr.Markdown()
+    info_box = gr.Textbox(show_label=False, interactive=False)

    chat_model = WebChatModel()
    chat_box, chatbot, history, chat_elems = create_chat_box(chat_model)
@@ -21,9 +20,13 @@ def create_infer_tab(top_elems: Dict[str, Component]) -> Dict[str, Component]:
    load_btn.click(
        chat_model.load_model,
        [
-            top_elems["lang"], top_elems["model_name"], top_elems["checkpoints"],
-            top_elems["finetuning_type"], top_elems["template"],
-            quantization_bit
+            top_elems["lang"],
+            top_elems["model_name"],
+            top_elems["checkpoints"],
+            top_elems["finetuning_type"],
+            top_elems["quantization_bit"],
+            top_elems["template"],
+            top_elems["source_prefix"]
        ],
        [info_box]
    ).then(
@@ -39,7 +42,6 @@ def create_infer_tab(top_elems: Dict[str, Component]) -> Dict[str, Component]:
    )

    return dict(
-        quantization_bit=quantization_bit,
        info_box=info_box,
        load_btn=load_btn,
        unload_btn=unload_btn,
--- a/src/llmtuner/webui/components/sft.py
+++ b/src/llmtuner/webui/components/sft.py
@@ -12,8 +12,8 @@ from llmtuner.webui.utils import can_preview, get_preview, gen_plot

 def create_sft_tab(top_elems: Dict[str, Component], runner: Runner) -> Dict[str, Component]:
    with gr.Row():
-        dataset_dir = gr.Textbox(value=DEFAULT_DATA_DIR, interactive=True, scale=1)
-        dataset = gr.Dropdown(multiselect=True, interactive=True, scale=4)
+        dataset_dir = gr.Textbox(value=DEFAULT_DATA_DIR, scale=2)
+        dataset = gr.Dropdown(multiselect=True, scale=4)
        preview_btn = gr.Button(interactive=False, scale=1)

    preview_box, preview_count, preview_samples, close_btn = create_preview_box()
@@ -23,22 +23,33 @@ def create_sft_tab(top_elems: Dict[str, Component], runner: Runner) -> Dict[str,
    preview_btn.click(get_preview, [dataset_dir, dataset], [preview_count, preview_samples, preview_box])

    with gr.Row():
-        learning_rate = gr.Textbox(value="5e-5", interactive=True)
-        num_train_epochs = gr.Textbox(value="3.0", interactive=True)
-        max_samples = gr.Textbox(value="100000", interactive=True)
-        quantization_bit = gr.Dropdown([8, 4])
+        max_source_length = gr.Slider(value=512, minimum=4, maximum=4096, step=1)
+        max_target_length = gr.Slider(value=512, minimum=4, maximum=4096, step=1)
+        learning_rate = gr.Textbox(value="5e-5")
+        num_train_epochs = gr.Textbox(value="3.0")
+        max_samples = gr.Textbox(value="100000")

    with gr.Row():
-        batch_size = gr.Slider(value=4, minimum=1, maximum=128, step=1, interactive=True)
-        gradient_accumulation_steps = gr.Slider(value=4, minimum=1, maximum=32, step=1, interactive=True)
+        batch_size = gr.Slider(value=4, minimum=1, maximum=512, step=1)
+        gradient_accumulation_steps = gr.Slider(value=4, minimum=1, maximum=512, step=1)
        lr_scheduler_type = gr.Dropdown(
-            value="cosine", choices=[scheduler.value for scheduler in SchedulerType], interactive=True
+            value="cosine", choices=[scheduler.value for scheduler in SchedulerType]
        )
-        fp16 = gr.Checkbox(value=True)
+        max_grad_norm = gr.Textbox(value="1.0")
+        dev_ratio = gr.Slider(value=0, minimum=0, maximum=1, step=0.001)

+    with gr.Accordion(label="Advanced config", open=False) as advanced_tab:
        with gr.Row():
-        logging_steps = gr.Slider(value=5, minimum=5, maximum=1000, step=5, interactive=True)
-        save_steps = gr.Slider(value=100, minimum=10, maximum=2000, step=10, interactive=True)
+            logging_steps = gr.Slider(value=5, minimum=5, maximum=1000, step=5)
+            save_steps = gr.Slider(value=100, minimum=10, maximum=5000, step=10)
+            warmup_steps = gr.Slider(value=0, minimum=0, maximum=5000, step=1)
+            compute_type = gr.Radio(choices=["fp16", "bf16"], value="fp16")
+
+    with gr.Accordion(label="LoRA config", open=False) as lora_tab:
+        with gr.Row():
+            lora_rank = gr.Slider(value=8, minimum=1, maximum=1024, step=1, scale=1)
+            lora_dropout = gr.Slider(value=0, minimum=0, maximum=1, step=0.01, scale=1)
+            lora_target = gr.Textbox(scale=2)

    with gr.Row():
        start_btn = gr.Button()
@@ -47,6 +58,8 @@ def create_sft_tab(top_elems: Dict[str, Component], runner: Runner) -> Dict[str,
    with gr.Row():
        with gr.Column(scale=4):
            output_dir = gr.Textbox(interactive=True)
+
+            with gr.Box():
                output_box = gr.Markdown()

        with gr.Column(scale=1):
@@ -55,11 +68,33 @@ def create_sft_tab(top_elems: Dict[str, Component], runner: Runner) -> Dict[str,
    start_btn.click(
        runner.run_train,
        [
-            top_elems["lang"], top_elems["model_name"], top_elems["checkpoints"],
-            top_elems["finetuning_type"], top_elems["template"],
-            dataset, dataset_dir, learning_rate, num_train_epochs, max_samples,
-            fp16, quantization_bit, batch_size, gradient_accumulation_steps,
-            lr_scheduler_type, logging_steps, save_steps, output_dir
+            top_elems["lang"],
+            top_elems["model_name"],
+            top_elems["checkpoints"],
+            top_elems["finetuning_type"],
+            top_elems["quantization_bit"],
+            top_elems["template"],
+            top_elems["source_prefix"],
+            dataset_dir,
+            dataset,
+            max_source_length,
+            max_target_length,
+            learning_rate,
+            num_train_epochs,
+            max_samples,
+            batch_size,
+            gradient_accumulation_steps,
+            lr_scheduler_type,
+            max_grad_norm,
+            dev_ratio,
+            logging_steps,
+            save_steps,
+            warmup_steps,
+            compute_type,
+            lora_rank,
+            lora_dropout,
+            lora_target,
+            output_dir
        ],
        [output_box]
    )
@@ -76,16 +111,25 @@ def create_sft_tab(top_elems: Dict[str, Component], runner: Runner) -> Dict[str,
        preview_count=preview_count,
        preview_samples=preview_samples,
        close_btn=close_btn,
+        max_source_length=max_source_length,
+        max_target_length=max_target_length,
        learning_rate=learning_rate,
        num_train_epochs=num_train_epochs,
        max_samples=max_samples,
-        quantization_bit=quantization_bit,
        batch_size=batch_size,
        gradient_accumulation_steps=gradient_accumulation_steps,
        lr_scheduler_type=lr_scheduler_type,
-        fp16=fp16,
+        max_grad_norm=max_grad_norm,
+        dev_ratio=dev_ratio,
+        advanced_tab=advanced_tab,
        logging_steps=logging_steps,
        save_steps=save_steps,
+        warmup_steps=warmup_steps,
+        compute_type=compute_type,
+        lora_tab=lora_tab,
+        lora_rank=lora_rank,
+        lora_dropout=lora_dropout,
+        lora_target=lora_target,
        start_btn=start_btn,
        stop_btn=stop_btn,
        output_dir=output_dir,
--- a/src/llmtuner/webui/components/top.py
+++ b/src/llmtuner/webui/components/top.py
@@ -6,29 +6,41 @@ from gradio.components import Component
 from llmtuner.extras.constants import METHODS, SUPPORTED_MODELS
 from llmtuner.extras.template import templates
 from llmtuner.webui.common import list_checkpoint, get_model_path, save_config
+from llmtuner.webui.utils import can_quantize


 def create_top() -> Dict[str, Component]:
    available_models = list(SUPPORTED_MODELS.keys()) + ["Custom"]

    with gr.Row():
-        lang = gr.Dropdown(choices=["en", "zh"], value="en", interactive=True, scale=1)
+        lang = gr.Dropdown(choices=["en", "zh"], value="en", scale=1)
        model_name = gr.Dropdown(choices=available_models, scale=3)
        model_path = gr.Textbox(scale=3)

    with gr.Row():
-        finetuning_type = gr.Dropdown(value="lora", choices=METHODS, interactive=True, scale=1)
-        template = gr.Dropdown(value="default", choices=list(templates.keys()), interactive=True, scale=1)
-        checkpoints = gr.Dropdown(multiselect=True, interactive=True, scale=4)
+        finetuning_type = gr.Dropdown(value="lora", choices=METHODS, scale=1)
+        checkpoints = gr.Dropdown(multiselect=True, scale=5)
        refresh_btn = gr.Button(scale=1)

+    with gr.Accordion(label="Advanced config", open=False) as advanced_tab:
+        with gr.Row():
+            quantization_bit = gr.Dropdown([8, 4], scale=1)
+            template = gr.Dropdown(value="default", choices=list(templates.keys()), scale=1)
+            source_prefix = gr.Textbox(scale=2)
+
    model_name.change(
        list_checkpoint, [model_name, finetuning_type], [checkpoints]
    ).then(
        get_model_path, [model_name], [model_path]
    ) # do not save config since the below line will save
    model_path.change(save_config, [model_name, model_path])
-    finetuning_type.change(list_checkpoint, [model_name, finetuning_type], [checkpoints])
+
+    finetuning_type.change(
+        list_checkpoint, [model_name, finetuning_type], [checkpoints]
+    ).then(
+        can_quantize, [finetuning_type], [quantization_bit]
+    )
+
    refresh_btn.click(list_checkpoint, [model_name, finetuning_type], [checkpoints])

    return dict(
@@ -36,7 +48,10 @@ def create_top() -> Dict[str, Component]:
        model_name=model_name,
        model_path=model_path,
        finetuning_type=finetuning_type,
-        template=template,
        checkpoints=checkpoints,
-        refresh_btn=refresh_btn
+        refresh_btn=refresh_btn,
+        advanced_tab=advanced_tab,
+        quantization_bit=quantization_bit,
+        template=template,
+        source_prefix=source_prefix
    )
--- a/src/llmtuner/webui/interface.py
+++ b/src/llmtuner/webui/interface.py
@@ -27,7 +27,7 @@ def create_ui() -> gr.Blocks:
        with gr.Tab("Evaluate"):
            eval_elems = create_eval_tab(top_elems, runner)

-        with gr.Tab("Inference"):
+        with gr.Tab("Chat"):
            infer_elems = create_infer_tab(top_elems)

        elem_list = [top_elems, sft_elems, eval_elems, infer_elems]
--- a/src/llmtuner/webui/locales.py
+++ b/src/llmtuner/webui/locales.py
@@ -25,6 +25,14 @@ LOCALES = {
            "info": "本地模型的文件路径或 Hugging Face 的模型标识符。"
        }
    },
+    "finetuning_type": {
+        "en": {
+            "label": "Finetuning method"
+        },
+        "zh": {
+            "label": "微调方法"
+        }
+    },
    "checkpoints": {
        "en": {
            "label": "Checkpoints"
@@ -33,14 +41,6 @@ LOCALES = {
            "label": "模型断点"
        }
    },
-    "template": {
-        "en": {
-            "label": "Prompt template"
-        },
-        "zh": {
-            "label": "提示模板"
-        }
-    },
    "refresh_btn": {
        "en": {
            "value": "Refresh checkpoints"
@@ -49,6 +49,44 @@ LOCALES = {
            "value": "刷新断点"
        }
    },
+    "advanced_tab": {
+        "en": {
+            "label": "Advanced configurations"
+        },
+        "zh": {
+            "label": "高级设置"
+        }
+    },
+    "quantization_bit": {
+        "en": {
+            "label": "Quantization bit (optional)",
+            "info": "Enable 4/8-bit model quantization."
+        },
+        "zh": {
+            "label": "量化等级（非必填）",
+            "info": "启用 4/8 比特模型量化。"
+        }
+    },
+    "template": {
+        "en": {
+            "label": "Prompt template",
+            "info": "The template used in constructing prompts."
+        },
+        "zh": {
+            "label": "提示模板",
+            "info": "构建提示词时使用的模板"
+        }
+    },
+    "source_prefix": {
+        "en": {
+            "label": "System prompt (optional)",
+            "info": "A sequence used as the default system prompt."
+        },
+        "zh": {
+            "label": "系统提示词（非必填）",
+            "info": "默认使用的系统提示词"
+        }
+    },
    "dataset_dir": {
        "en": {
            "label": "Data dir",
@@ -99,66 +137,24 @@ LOCALES = {
            "value": "关闭"
        }
    },
-    "max_samples": {
+    "max_source_length": {
        "en": {
-            "label": "Max samples",
-            "info": "Maximum samples per dataset."
+            "label": "Max source length",
+            "info": "Max tokens in source sequence."
        },
        "zh": {
-            "label": "最大样本数",
-            "info": "每个数据集最多使用的样本数。"
+            "label": "输入序列最大长度",
+            "info": "输入序列分词后的最大长度。"
        }
    },
-    "batch_size": {
+    "max_target_length": {
        "en": {
-            "label": "Batch size",
-            "info": "Number of samples to process per GPU."
+            "label": "Max target length",
+            "info": "Max tokens in target sequence."
        },
        "zh": {
-            "label": "批处理大小",
-            "info": "每块 GPU 上处理的样本数量。"
-        }
-    },
-    "quantization_bit": {
-        "en": {
-            "label": "Quantization bit",
-            "info": "Enable 4/8-bit model quantization."
-        },
-        "zh": {
-            "label": "量化",
-            "info": "启用 4/8 比特模型量化。"
-        }
-    },
-    "start_btn": {
-        "en": {
-            "value": "Start"
-        },
-        "zh": {
-            "value": "开始"
-        }
-    },
-    "stop_btn": {
-        "en": {
-            "value": "Abort"
-        },
-        "zh": {
-            "value": "中断"
-        }
-    },
-    "output_box": {
-        "en": {
-            "value": "Ready."
-        },
-        "zh": {
-            "value": "准备就绪。"
-        }
-    },
-    "finetuning_type": {
-        "en": {
-            "label": "Finetuning method"
-        },
-        "zh": {
-            "label": "微调方法"
+            "label": "输出序列最大长度",
+            "info": "输出序列分词后的最大长度。"
        }
    },
    "learning_rate": {
@@ -181,6 +177,26 @@ LOCALES = {
            "info": "需要执行的训练总轮数。"
        }
    },
+    "max_samples": {
+        "en": {
+            "label": "Max samples",
+            "info": "Maximum samples per dataset."
+        },
+        "zh": {
+            "label": "最大样本数",
+            "info": "每个数据集最多使用的样本数。"
+        }
+    },
+    "batch_size": {
+        "en": {
+            "label": "Batch size",
+            "info": "Number of samples to process per GPU."
+        },
+        "zh":{
+            "label": "批处理大小",
+            "info": "每块 GPU 上处理的样本数量。"
+        }
+    },
    "gradient_accumulation_steps": {
        "en": {
            "label": "Gradient accumulation",
@@ -201,20 +217,30 @@ LOCALES = {
            "info": "采用的学习率调节器名称。"
        }
    },
-    "fp16": {
+    "max_grad_norm": {
        "en": {
-            "label": "fp16",
-            "info": "Whether to use fp16 mixed precision training."
+            "label": "Maximum gradient norm",
+            "info": "Norm for gradient clipping.."
        },
        "zh": {
-            "label": "fp16",
-            "info": "是否启用 FP16 混合精度训练。"
+            "label": "最大梯度范数",
+            "info": "用于梯度裁剪的范数。"
+        }
+    },
+    "dev_ratio": {
+        "en": {
+            "label": "Dev ratio",
+            "info": "Proportion of data in the dev set."
+        },
+        "zh": {
+            "label": "验证集比例",
+            "info": "验证集占全部样本的百分比。"
        }
    },
    "logging_steps": {
        "en": {
            "label": "Logging steps",
-            "info": "Number of update steps between two logs."
+            "info": "Number of steps between two logs."
        },
        "zh": {
            "label": "日志间隔",
@@ -224,13 +250,87 @@ LOCALES = {
    "save_steps": {
        "en": {
            "label": "Save steps",
-            "info": "Number of updates steps between two checkpoints."
+            "info": "Number of steps between two checkpoints."
        },
        "zh": {
            "label": "保存间隔",
            "info": "每两次断点保存间的更新步数。"
        }
    },
+    "warmup_steps": {
+        "en": {
+            "label": "Warmup steps",
+            "info": "Number of steps used for warmup."
+        },
+        "zh": {
+            "label": "预热步数",
+            "info": "学习率预热采用的步数。"
+        }
+    },
+    "compute_type": {
+        "en": {
+            "label": "Compute type",
+            "info": "Whether to use fp16 or bf16 mixed precision training."
+        },
+        "zh": {
+            "label": "计算类型",
+            "info": "是否启用 FP16 或 BF16 混合精度训练。"
+        }
+    },
+    "lora_tab": {
+        "en": {
+            "label": "LoRA configurations"
+        },
+        "zh": {
+            "label": "LoRA 参数设置"
+        }
+    },
+    "lora_rank": {
+        "en": {
+            "label": "LoRA rank",
+            "info": "The rank of LoRA matrices."
+        },
+        "zh": {
+            "label": "LoRA 秩",
+            "info": "LoRA 矩阵的秩。"
+        }
+    },
+    "lora_dropout": {
+        "en": {
+            "label": "LoRA Dropout",
+            "info": "Dropout ratio of LoRA weights."
+        },
+        "zh": {
+            "label": "LoRA 随机丢弃",
+            "info": "LoRA 权重随机丢弃的概率。"
+        }
+    },
+    "lora_target": {
+        "en": {
+            "label": "LoRA modules (optional)",
+            "info": "The name(s) of target modules to apply LoRA. Use commas to separate multiple modules."
+        },
+        "zh": {
+            "label": "LoRA 作用层（非必填）",
+            "info": "应用 LoRA 的线性层名称。使用英文逗号分隔多个名称。"
+        }
+    },
+    "start_btn": {
+        "en": {
+            "value": "Start"
+        },
+        "zh": {
+            "value": "开始"
+        }
+    },
+    "stop_btn": {
+        "en": {
+            "value": "Abort"
+        },
+        "zh": {
+            "value": "中断"
+        }
+    },
    "output_dir": {
        "en": {
            "label": "Checkpoint name",
@@ -241,6 +341,14 @@ LOCALES = {
            "info": "保存模型断点的文件夹名称。"
        }
    },
+    "output_box": {
+        "en": {
+            "value": "Ready."
+        },
+        "zh": {
+            "value": "准备就绪。"
+        }
+    },
    "loss_viewer": {
        "en": {
            "label": "Loss"
@@ -257,14 +365,6 @@ LOCALES = {
            "label": "保存预测结果"
        }
    },
-    "info_box": {
-        "en": {
-            "value": "Model unloaded, please load a model first."
-        },
-        "zh": {
-            "value": "模型未加载，请先加载模型。"
-        }
-    },
    "load_btn": {
        "en": {
            "value": "Load model"
@@ -281,6 +381,22 @@ LOCALES = {
            "value": "卸载模型"
        }
    },
+    "info_box": {
+        "en": {
+            "value": "Model unloaded, please load a model first."
+        },
+        "zh": {
+            "value": "模型未加载，请先加载模型。"
+        }
+    },
+    "prefix": {
+        "en": {
+            "placeholder": "System prompt (optional)"
+        },
+        "zh": {
+            "placeholder": "系统提示词（非必填）"
+        }
+    },
    "query": {
        "en": {
            "placeholder": "Input..."
@@ -305,6 +421,14 @@ LOCALES = {
            "value": "清空历史"
        }
    },
+    "max_length": {
+        "en": {
+            "label": "Maximum length"
+        },
+        "zh": {
+            "label": "最大长度"
+        }
+    },
    "max_new_tokens": {
        "en": {
            "label": "Maximum new tokens"
--- a/src/llmtuner/webui/runner.py
+++ b/src/llmtuner/webui/runner.py
@@ -3,10 +3,10 @@ import os
 import threading
 import time
 import transformers
-from typing import Optional, Tuple
+from typing import List, Optional, Tuple

 from llmtuner.extras.callbacks import LogCallback
-from llmtuner.extras.constants import DEFAULT_MODULE # will be deprecated
+from llmtuner.extras.constants import DEFAULT_MODULE
 from llmtuner.extras.logging import LoggerHandler
 from llmtuner.extras.misc import torch_gc
 from llmtuner.tuner import get_train_args, run_sft
@@ -59,10 +59,34 @@ class Runner:
            return finish_info if finish_info is not None else ALERTS["info_finished"][lang]

    def run_train(
-        self, lang, model_name, checkpoints, finetuning_type, template,
-        dataset, dataset_dir, learning_rate, num_train_epochs, max_samples,
-        fp16, quantization_bit, batch_size, gradient_accumulation_steps,
-        lr_scheduler_type, logging_steps, save_steps, output_dir
+        self,
+        lang: str,
+        model_name: str,
+        checkpoints: List[str],
+        finetuning_type: str,
+        quantization_bit: str,
+        template: str,
+        source_prefix: str,
+        dataset_dir: str,
+        dataset: List[str],
+        max_source_length: int,
+        max_target_length: int,
+        learning_rate: str,
+        num_train_epochs: str,
+        max_samples: str,
+        batch_size: int,
+        gradient_accumulation_steps: int,
+        lr_scheduler_type: str,
+        max_grad_norm: str,
+        dev_ratio: float,
+        logging_steps: int,
+        save_steps: int,
+        warmup_steps: int,
+        compute_type: str,
+        lora_rank: int,
+        lora_dropout: float,
+        lora_target: str,
+        output_dir: str
    ):
        model_name_or_path, error, logger_handler, trainer_callback = self.initialize(lang, model_name, dataset)
        if error:
@@ -79,25 +103,40 @@ class Runner:
        args = dict(
            model_name_or_path=model_name_or_path,
            do_train=True,
-            finetuning_type=finetuning_type,
-            lora_target=DEFAULT_MODULE.get(model_name.split("-")[0], None) or "q_proj,v_proj",
-            prompt_template=template,
-            dataset=",".join(dataset),
-            dataset_dir=dataset_dir,
-            max_samples=int(max_samples),
-            output_dir=os.path.join(get_save_dir(model_name), finetuning_type, output_dir),
-            checkpoint_dir=checkpoint_dir,
            overwrite_cache=True,
+            checkpoint_dir=checkpoint_dir,
+            finetuning_type=finetuning_type,
+            quantization_bit=int(quantization_bit) if quantization_bit else None,
+            prompt_template=template,
+            source_prefix=source_prefix,
+            dataset_dir=dataset_dir,
+            dataset=",".join(dataset),
+            max_source_length=max_source_length,
+            max_target_length=max_target_length,
+            learning_rate=float(learning_rate),
+            num_train_epochs=float(num_train_epochs),
+            max_samples=int(max_samples),
            per_device_train_batch_size=batch_size,
            gradient_accumulation_steps=gradient_accumulation_steps,
            lr_scheduler_type=lr_scheduler_type,
+            max_grad_norm=float(max_grad_norm),
            logging_steps=logging_steps,
            save_steps=save_steps,
-            learning_rate=float(learning_rate),
-            num_train_epochs=float(num_train_epochs),
-            fp16=fp16,
-            quantization_bit=int(quantization_bit) if quantization_bit else None
+            warmup_steps=warmup_steps,
+            fp16=(compute_type == "fp16"),
+            bf16=(compute_type == "bf16"),
+            lora_rank=lora_rank,
+            lora_dropout=lora_dropout,
+            lora_target=lora_target or DEFAULT_MODULE.get(model_name.split("-")[0], "q_proj,v_proj"),
+            output_dir=os.path.join(get_save_dir(model_name), finetuning_type, output_dir)
        )
+
+        if dev_ratio > 1e-6:
+            args["dev_ratio"] = dev_ratio
+            args["evaluation_strategy"] = "steps"
+            args["eval_steps"] = save_steps
+            args["load_best_model_at_end"] = True
+
        model_args, data_args, training_args, finetuning_args, _ = get_train_args(args)

        run_args = dict(
@@ -120,8 +159,21 @@ class Runner:
        yield self.finalize(lang)

    def run_eval(
-        self, lang, model_name, checkpoints, finetuning_type, template,
-        dataset, dataset_dir, max_samples, batch_size, quantization_bit, predict
+        self,
+        lang: str,
+        model_name: str,
+        checkpoints: List[str],
+        finetuning_type: str,
+        quantization_bit: str,
+        template: str,
+        source_prefix: str,
+        dataset_dir: str,
+        dataset: List[str],
+        max_source_length: int,
+        max_target_length: int,
+        max_samples: str,
+        batch_size: int,
+        predict: bool
    ):
        model_name_or_path, error, logger_handler, trainer_callback = self.initialize(lang, model_name, dataset)
        if error:
@@ -140,17 +192,20 @@ class Runner:
        args = dict(
            model_name_or_path=model_name_or_path,
            do_eval=True,
-            finetuning_type=finetuning_type,
-            prompt_template=template,
-            dataset=",".join(dataset),
-            dataset_dir=dataset_dir,
-            max_samples=int(max_samples),
-            output_dir=output_dir,
-            checkpoint_dir=checkpoint_dir,
            overwrite_cache=True,
            predict_with_generate=True,
+            checkpoint_dir=checkpoint_dir,
+            finetuning_type=finetuning_type,
+            quantization_bit=int(quantization_bit) if quantization_bit else None,
+            prompt_template=template,
+            source_prefix=source_prefix,
+            dataset_dir=dataset_dir,
+            dataset=",".join(dataset),
+            max_source_length=max_source_length,
+            max_target_length=max_target_length,
+            max_samples=int(max_samples),
            per_device_eval_batch_size=batch_size,
-            quantization_bit=int(quantization_bit) if quantization_bit else None
+            output_dir=output_dir
        )

        if predict:
--- a/src/llmtuner/webui/utils.py
+++ b/src/llmtuner/webui/utils.py
@@ -3,7 +3,7 @@ import json
 import gradio as gr
 import matplotlib.figure
 import matplotlib.pyplot as plt
-from typing import Tuple
+from typing import Any, Dict, Tuple
 from datetime import datetime

 from llmtuner.extras.ploting import smooth
@@ -23,7 +23,7 @@ def get_time() -> str:
    return datetime.now().strftime('%Y-%m-%d-%H-%M-%S')


-def can_preview(dataset_dir: str, dataset: list) -> dict:
+def can_preview(dataset_dir: str, dataset: list) -> Dict[str, Any]:
    with open(os.path.join(dataset_dir, DATA_CONFIG), "r", encoding="utf-8") as f:
        dataset_info = json.load(f)
    if (
@@ -36,7 +36,7 @@ def can_preview(dataset_dir: str, dataset: list) -> dict:
        return gr.update(interactive=False)


-def get_preview(dataset_dir: str, dataset: list) -> Tuple[int, list, dict]:
+def get_preview(dataset_dir: str, dataset: list) -> Tuple[int, list, Dict[str, Any]]:
    with open(os.path.join(dataset_dir, DATA_CONFIG), "r", encoding="utf-8") as f:
        dataset_info = json.load(f)
    data_file = dataset_info[dataset[0]]["file_name"]
@@ -45,6 +45,13 @@ def get_preview(dataset_dir: str, dataset: list) -> Tuple[int, list, dict]:
    return len(data), data[:2], gr.update(visible=True)


+def can_quantize(finetuning_type: str) -> Dict[str, Any]:
+    if finetuning_type != "lora":
+        return gr.update(value="", interactive=False)
+    else:
+        return gr.update(interactive=True)
+
+
 def get_eval_results(path: os.PathLike) -> str:
    with open(path, "r", encoding="utf-8") as f:
        result = json.dumps(json.load(f), indent=4)
@@ -66,6 +73,10 @@ def gen_plot(base_model: str, finetuning_type: str, output_dir: str) -> matplotl
            if log_info.get("loss", None):
                steps.append(log_info["current_steps"])
                losses.append(log_info["loss"])
+
+    if len(losses) == 0:
+        return None
+
    ax.plot(steps, losses, alpha=0.4, label="original")
    ax.plot(steps, smooth(losses), label="smoothed")
    ax.legend()
--- a/src/train_bash.py
+++ b/src/train_bash.py
@@ -1,4 +1,4 @@
-from llmtuner import get_train_args, run_pt, run_sft, run_rm, run_ppo
+from llmtuner.tuner import get_train_args, run_pt, run_sft, run_rm, run_ppo


 def main():
--- a/src/train_web.py
+++ b/src/train_web.py
@@ -1,4 +1,4 @@
-from llmtuner import create_ui
+from llmtuner.webui.interface import create_ui


 def main():
--- a/src/web_demo.py
+++ b/src/web_demo.py
@@ -0,0 +1,36 @@
+# coding=utf-8
+# Implements user interface in browser for fine-tuned models.
+# Usage: python web_demo.py --model_name_or_path path_to_model --checkpoint_dir path_to_checkpoint
+
+import gradio as gr
+from transformers.utils.versions import require_version
+
+from llmtuner.tuner import get_infer_args
+from llmtuner.webui.chat import WebChatModel
+from llmtuner.webui.components.chatbot import create_chat_box
+from llmtuner.webui.manager import Manager
+
+
+require_version("gradio>=3.36.0", "To fix: pip install gradio>=3.36.0")
+
+
+def main():
+    chat_model = WebChatModel(*get_infer_args())
+
+    with gr.Blocks(title="Web Demo") as demo:
+        lang = gr.Dropdown(choices=["en", "zh"], value="en")
+
+        _, _, _, chat_elems = create_chat_box(chat_model, visible=True)
+
+        manager = Manager([{"lang": lang}, chat_elems])
+
+        demo.load(manager.gen_label, [lang], [lang] + list(chat_elems.values()))
+
+        lang.change(manager.gen_label, [lang], [lang] + list(chat_elems.values()))
+
+    demo.queue()
+    demo.launch(server_name="0.0.0.0", share=False, inbrowser=True)
+
+
+if __name__ == "__main__":
+    main()
Author	SHA1	Message	Date
hiyouga	c4e9694c6e	release v0.1.2 Former-commit-id: 04aad91b71cc3a1acaf1bcec4304ce6b2098f7dc	2023-07-20 22:33:59 +08:00
hiyouga	2006a96570	fix api Former-commit-id: 4c3e8be325045e432b31c519132123c7b0689262	2023-07-20 22:14:54 +08:00
hoshi-hiyouga	5dcd95645f	Merge pull request #213 from Ehco1996/patch-1 feat: support pass args before init web app Former-commit-id: b0612c05bc10c281c0a95e08c5517c3fb0a72029	2023-07-20 22:12:07 +08:00
hiyouga	9b3304b054	update UI, fix #212 Former-commit-id: ac92c2bd7c47353759474fad9412f21b38c65501	2023-07-20 22:09:06 +08:00
Ehco	e580d4ef41	feat: support pass args before init web app as title Former-commit-id: 434a5077288927e0be15cd066ca3e562111fad4d	2023-07-20 21:49:26 +08:00
hiyouga	64db4abc68	Update README.md Former-commit-id: 6dc67a495ec7d9fdc2574bae92063ed8a9099725	2023-07-20 17:23:16 +08:00
hiyouga	5ba0b80e5c	simplify code Former-commit-id: d3731754ab7c28ae81f60784e0e4213f279d93fe	2023-07-20 15:08:57 +08:00
hiyouga	7a43ff3d89	tiny fix Former-commit-id: 22b1be7bbb9e7bd863acb88bf7365090b1b8235d	2023-07-19 22:53:46 +08:00
hiyouga	7e1a1d141a	fix #199 Former-commit-id: 7fc778b49bc17688aca39fffe01f9d33e03e0c28	2023-07-19 22:51:29 +08:00
hiyouga	6d881f161b	add datasets Former-commit-id: 02e4b47dea1b25905c61f2ace88bab112610f021	2023-07-19 20:59:15 +08:00
hiyouga	a02b3e6192	fix #196 Former-commit-id: 85fd82926db345a590a7fb32c0e352a1d2f025c3	2023-07-19 17:35:38 +08:00
hiyouga	bcdee9fc19	fix #194 Former-commit-id: 9792921531efefb4bcddbde4380169a78fe064a6	2023-07-19 17:07:33 +08:00
hiyouga	8b688251be	support LLaMA-2 Former-commit-id: 04dfda054855ee9256586aacbd382f8fb0bfed04	2023-07-19 16:42:14 +08:00
hiyouga	718f3382ad	add LLaMA2 template Former-commit-id: 246421bd35cf7bb2203ac4fc924e6cd1c292954d	2023-07-19 00:44:49 +08:00
hiyouga	dc8283d3d7	fix API Former-commit-id: 9b10c9a12e33ab897056ecc61d977d221c19141b	2023-07-19 00:01:14 +08:00
hiyouga	35e76879f5	support dev set in web ui Former-commit-id: fe1370561a9b027d9ebdef52733344f1e3683081	2023-07-18 20:40:49 +08:00
hiyouga	8e4ae0aaac	add web demo Former-commit-id: 25ea647e5ac36b497b8e176b123fdee39be3fd30	2023-07-18 17:21:16 +08:00
hiyouga	5ed2a97056	update baichuan template Former-commit-id: 03520588c39986c98a0515a64993af8c2468b9d0	2023-07-18 16:43:51 +08:00
hiyouga	03eba6f041	fix template Former-commit-id: 729053c9cea6254165ae9c8fd7809479b12f735c	2023-07-18 16:37:23 +08:00
hiyouga	ec166e736a	fix #176 Former-commit-id: 2ae3445b0d28b4ed22ddbb2cfe09089ae0c23fe1	2023-07-18 16:36:24 +08:00
hiyouga	c85a6b83b3	fix webUI, fix #171 #177 Former-commit-id: 3459bb2d35162dbbef79cda05da08a56921aa276	2023-07-18 15:51:48 +08:00
hiyouga	a864a7b395	update webUI, fix #179 Former-commit-id: f9074fed5e22585679661588befcf266a79009f2	2023-07-18 15:35:17 +08:00
hiyouga	fd8c2d4aac	tiny fix Former-commit-id: bcdf5bb55651d639e9f57fd915268137156af9cd	2023-07-18 00:52:31 +08:00
				`@@ -1 +0,0 @@`
				`from llmtuner.webui.interface import create_ui`