Upload 28 files

.gitignore

@@ -0,0 +1,129 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/

CITATION.cff

@@ -0,0 +1,9 @@
+cff-version: 1.2.0
+message: "If you use this software, please cite it as below."
+authors:
+- family-names: "Xu"
+  given-names: "Ming"
+title: "MedicalGPT: Training Your Own Medical GPT Model with ChatGPT Training Pipeline"
+url: "https://github.com/shibing624/MedicalGPT"
+data-released: 2023-06-02
+version: 0.0.4

CONTRIBUTING.md

@@ -0,0 +1,9 @@
+# Contributing
+
+We are happy to accept your contributions to make this repo better and more awesome! To avoid unnecessary work on either
+side, please stick to the following process:
+
+1. Check if there is already an issue for your concern.
+2. If there is not, open a new one to start a discussion. We hate to close finished PRs!
+3. If we decide your concern needs code changes, we would be happy to accept a pull request. Please consider the
+commit guidelines below.

DISCLAIMER

@@ -0,0 +1,23 @@
+The software project, data, and models provided by our GitHub project are provided "as is," without warranty of any kind, express or implied, including but not limited to the warranties of merchantability, fitness for a particular purpose, and non-infringement.
+
+In no event shall the project owners or contributors be liable for any direct, indirect, incidental, special, exemplary, or consequential damages (including, but not limited to, procurement of substitute goods or services; loss of use, data, or profits; or business interruption) however caused and on any theory of liability, whether in contract, strict liability, or tort (including negligence or otherwise) arising in any way out of the use of this software project, data, or models, even if advised of the possibility of such damage.
+
+Users of this software project, data, and models are solely responsible for any consequences of their use. The project owners and contributors shall not be held responsible for any subsequent or potential harm caused by the use of this software project, data, or models.
+
+By using this software project, data, or models, users accept and agree to this disclaimer. If users do not agree to the terms of this disclaimer, they should not use this software project, data, or models.
+
+It is important to note that this software project, data, and models are still in the research phase and are provided for experimental purposes only. As such, the project owners and contributors do not guarantee the accuracy, completeness, or usefulness of the software project, data, or models.
+
+Furthermore, due to the experimental nature of this software project, data, and models, it is possible that they may contain or generate inappropriate responses, errors, or inconsistencies. Users should exercise caution when using this software project, data, or models, and should not rely solely on them for any critical or sensitive tasks.
+
+The project owners and contributors shall not be held responsible for any damages, losses, or liabilities arising from the use of this software project, data, or models, including but not limited to, any inappropriate responses generated by the software project, data, or models.
+
+By using this software project, data, or models, users acknowledge and accept the experimental nature of the software project, data, and models, and understand the potential risks and limitations associated with their use. If users do not agree to the terms of this disclaimer, they should not use this software project, data, or models.
+
+The software project, data, and models provided by our GitHub project are intended for research purposes only. They should not be used for any commercial, business, or legal purposes, and should not be relied upon as a substitute for professional advice or judgment.
+
+Users of this software project, data, and models are strictly prohibited from using them for any commercial purposes, including but not limited to, selling, licensing, or distributing the software project, data, or models to third parties.
+
+The project owners and contributors shall not be held responsible for any damages, losses, or liabilities arising from the use of this software project, data, or models for any commercial or business purposes.
+
+By using this software project, data, or models, users agree to use them for research purposes only, and not for any commercial or business purposes. If users do not agree to the terms of this disclaimer, they should not use this software project, data, or models.

LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

README.md

@@ -1,13 +1,326 @@
----
-title: MedicalGPT Main
-emoji: 😻
-colorFrom: purple
-colorTo: yellow
-sdk: gradio
-sdk_version: 3.41.2
-app_file: app.py
-pinned: false
-license: other
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+[**🇨🇳中文**](https://github.com/shibing624/MedicalGPT/blob/main/README.md) | [**🌐English**](https://github.com/shibing624/MedicalGPT/blob/main/README_EN.md) | [**📖文档/Docs**](https://github.com/shibing624/MedicalGPT/wiki) | [**🤖模型/Models**](https://huggingface.co/shibing624) 
+
+<div align="center">
+  <a href="https://github.com/shibing624/MedicalGPT">
+    <img src="https://github.com/shibing624/MedicalGPT/blob/main/docs/logo.png" height="100" alt="Logo">
+  </a>
+</div>
+
+-----------------
+
+# MedicalGPT: Training Medical GPT Model
+[![HF Models](https://img.shields.io/badge/Hugging%20Face-shibing624-green)](https://huggingface.co/shibing624)
+[![Github Stars](https://img.shields.io/github/stars/shibing624/MedicalGPT?color=yellow)](https://star-history.com/#shibing624/MedicalGPT&Timeline)
+[![Contributions welcome](https://img.shields.io/badge/contributions-welcome-brightgreen.svg)](CONTRIBUTING.md)
+[![License Apache 2.0](https://img.shields.io/badge/license-Apache%202.0-blue.svg)](LICENSE)
+[![python_version](https://img.shields.io/badge/Python-3.8%2B-green.svg)](requirements.txt)
+[![GitHub issues](https://img.shields.io/github/issues/shibing624/MedicalGPT.svg)](https://github.com/shibing624/MedicalGPT/issues)
+[![Wechat Group](http://vlog.sfyc.ltd/wechat_everyday/wxgroup_logo.png?imageView2/0/w/60/h/20)](#Contact)
+
+## 📖 Introduction
+
+**MedicalGPT** training medical GPT model with ChatGPT training pipeline, implemantation of Pretraining, 
+Supervised Finetuning, RLHF(Reward Modeling and Reinforcement Learning) and DPO(Direct Preference Optimization).
+
+**MedicalGPT** 训练医疗大模型，实现了包括增量预训练、有监督微调、RLHF(奖励建模、强化学习训练)和DPO(直接偏好优化)。
+
+<img src="https://github.com/shibing624/MedicalGPT/blob/main/docs/dpo.jpg" width="860" />
+
+- RLHF training pipeline来自Andrej Karpathy的演讲PDF [State of GPT](https://karpathy.ai/stateofgpt.pdf)，视频 [Video](https://build.microsoft.com/en-US/sessions/db3f4859-cd30-4445-a0cd-553c3304f8e2)
+- DPO方法来自论文[Direct Preference Optimization:Your Language Model is Secretly a Reward Model](https://arxiv.org/pdf/2305.18290.pdf)
+
+## 🔥 News
+[2023/08/28] v1.5版本: 新增[DPO(直接偏好优化)](https://arxiv.org/pdf/2305.18290.pdf)方法，DPO通过直接优化语言模型来实现对其行为的精确控制，可以有效学习到人类偏好。详见[Release-v1.5](https://github.com/shibing624/MedicalGPT/releases/tag/1.5.0)
+
+[2023/08/08] v1.4版本: 发布基于ShareGPT4数据集微调的中英文Vicuna-13B模型[shibing624/vicuna-baichuan-13b-chat](https://huggingface.co/shibing624/vicuna-baichuan-13b-chat)，和对应的LoRA模型[shibing624/vicuna-baichuan-13b-chat-lora](https://huggingface.co/shibing624/vicuna-baichuan-13b-chat-lora)，详见[Release-v1.4](https://github.com/shibing624/MedicalGPT/releases/tag/1.4.0)
+
+[2023/08/02] v1.3版本: 新增LLaMA, LLaMA2, Bloom, ChatGLM, ChatGLM2, Baichuan模型的多轮对话微调训练；新增领域词表扩充功能；新增中文预训练数据集和中文ShareGPT微调训练集，详见[Release-v1.3](https://github.com/shibing624/MedicalGPT/releases/tag/1.3.0)
+
+[2023/07/13] v1.1版本: 发布中文医疗LLaMA-13B模型[shibing624/ziya-llama-13b-medical-merged](https://huggingface.co/shibing624/ziya-llama-13b-medical-merged)，基于Ziya-LLaMA-13B-v1模型，SFT微调了一版医疗模型，医疗问答效果有提升，发布微调后的完整模型权重，详见[Release-v1.1](https://github.com/shibing624/MedicalGPT/releases/tag/1.1)
+
+[2023/06/15] v1.0版本: 发布中文医疗LoRA模型[shibing624/ziya-llama-13b-medical-lora](https://huggingface.co/shibing624/ziya-llama-13b-medical-lora)，基于Ziya-LLaMA-13B-v1模型，SFT微调了一版医疗模型，医疗问答效果有提升，发布微调后的LoRA权重，详见[Release-v1.0](https://github.com/shibing624/MedicalGPT/releases/tag/1.0.0)
+
+[2023/06/05] v0.2版本: 以医疗为例，训练领域大模型，实现了四阶段训练：包括二次预训练、有监督微调、奖励建模、强化学习训练。详见[Release-v0.2](https://github.com/shibing624/MedicalGPT/releases/tag/0.2.0)
+
+
+## 😊 Features
+
+
+基于ChatGPT Training Pipeline，本项目实现了领域模型--医疗行业语言大模型的训练：
+
+
+- 第一阶段：PT(Continue PreTraining)增量预训练，在海量领域文档数据上二次预训练GPT模型，以注入领域知识（可选）
+- 第二阶段：SFT(Supervised Fine-tuning)有监督微调，构造指令微调数据集，在预训练模型基础上做指令精调，以对齐指令意图
+- 第三阶段 
+  - RLHF(Reinforcement Learning from Human Feedback)基于人类反馈对语言模型进行强化学习，分为两步：
+    - RM(Reward Model)奖励模型建模，构造人类偏好排序数据集，训练奖励模型，用来建模人类偏好，主要是"HHH"原则，具体是"helpful, honest, harmless"
+    - RL(Reinforcement Learning)强化学习，用奖励模型来训练SFT模型，生成模型使用奖励或惩罚来更新其策略，以便生成更高质量、更符合人类偏好的文本
+  - [DPO(Direct Preference Optimization)](https://arxiv.org/pdf/2305.18290.pdf)直接偏好优化方法，DPO通过直接优化语言模型来实现对其行为的精确控制，而无需使用复杂的强化学习，也可以有效学习到人类偏好，DPO相较于RLHF更容易实现且易于训练，效果更好
+
+
+
+### Release Models
+
+
+| Model                                                                                                       | Base Model                                                                                                     | Introduction                                                                                                                                                                 | 
+|:------------------------------------------------------------------------------------------------------------|:---------------------------------------------------------------------------------------------------------------|:-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
+| [shibing624/ziya-llama-13b-medical-lora](https://huggingface.co/shibing624/ziya-llama-13b-medical-lora)     | [IDEA-CCNL/Ziya-LLaMA-13B-v1](https://huggingface.co/IDEA-CCNL/Ziya-LLaMA-13B-v1)                              | 在240万条中英文医疗数据集[shibing624/medical](https://huggingface.co/datasets/shibing624/medical)上SFT微调了一版Ziya-LLaMA-13B模型，医疗问答效果有提升，发布微调后的LoRA权重(单轮对话)                                 |
+| [shibing624/ziya-llama-13b-medical-merged](https://huggingface.co/shibing624/ziya-llama-13b-medical-merged) | [IDEA-CCNL/Ziya-LLaMA-13B-v1](https://huggingface.co/IDEA-CCNL/Ziya-LLaMA-13B-v1)                              | 在240万条中英文医疗数据集[shibing624/medical](https://huggingface.co/datasets/shibing624/medical)上SFT微调了一版Ziya-LLaMA-13B模型，医疗问答效果有提升，发布微调后的完整模型权重(单轮对话)                                 |
+| [shibing624/vicuna-baichuan-13b-chat-lora](https://huggingface.co/shibing624/vicuna-baichuan-13b-chat-lora) | [baichuan-inc/Baichuan-13B-Chat](https://huggingface.co/baichuan-inc/Baichuan-13B-Chat)                        | 在10万条多语言ShareGPT GPT4多轮对话数据集[shibing624/sharegpt_gpt4](https://huggingface.co/datasets/shibing624/sharegpt_gpt4)上SFT微调了一版baichuan-13b-chat多轮问答模型，日常问答和医疗问答效果有提升，发布微调后的LoRA权重 |
+| [shibing624/vicuna-baichuan-13b-chat](https://huggingface.co/shibing624/vicuna-baichuan-13b-chat)           | [baichuan-inc/Baichuan-13B-Chat](https://huggingface.co/baichuan-inc/Baichuan-13B-Chat)                        | 在10万条多语言ShareGPT GPT4多轮对话数据集[shibing624/sharegpt_gpt4](https://huggingface.co/datasets/shibing624/sharegpt_gpt4)上SFT微调了一版baichuan-13b-chat多轮问答模型，日常问答和医疗问答效果有提升，发布微调后的完整模型权重 |
+
+演示[shibing624/vicuna-baichuan-13b-chat](https://huggingface.co/shibing624/vicuna-baichuan-13b-chat)模型效果：
+<img src="https://github.com/shibing624/MedicalGPT/blob/main/docs/demo-screen.gif" width="860" />
+具体case见[Inference Examples](#inference-examples)
+
+## ▶️ Demo
+
+
+我们提供了一个简洁的基于gradio的交互式web界面，启动服务后，可通过浏览器访问，输入问题，模型会返回答案。
+
+启动服务，命令如下：
+```shell
+CUDA_VISIBLE_DEVICES=0 python gradio_demo.py --model_type base_model_type --base_model path_to_llama_hf_dir --lora_model path_to_lora_dir
+```
+
+参数说明：
+
+- `--model_type {base_model_type}`：预训练模型类型，如llama、bloom、chatglm等
+- `--base_model {base_model}`：存放HF格式的LLaMA模型权重和配置文件的目录，也可使用HF Model Hub模型调用名称
+- `--lora_model {lora_model}`：LoRA文件所在目录，也可使用HF Model Hub模型调用名称。若lora权重已经合并到预训练模型，则删除--lora_model参数
+- `--tokenizer_path {tokenizer_path}`：存放对应tokenizer的目录。若不提供此参数，则其默认值与--base_model相同
+- `--template_name`：模板名称，如`vicuna`、`alpaca`等。若不提供此参数，则其默认值是vicuna
+- `--only_cpu`: 仅使用CPU进行推理
+- `--gpus {gpu_ids}`: 指定使用的GPU设备编号，默认为0。如使用多张GPU，以逗号分隔，如0,1,2
+- `--resize_emb`：是否调整embedding大小，若不调整，则使用预训练模型的embedding大小，默认不��整
+
+
+## 💾 Install
+#### Updating the requirements
+From time to time, the `requirements.txt` changes. To update, use this command:
+
+```markdown
+git clone https://github.com/shibing624/MedicalGPT
+conda activate gpt
+cd MedicalGPT
+pip install -r requirements.txt --upgrade
+```
+
+## 🚀 Training Pipeline
+
+Training Stage:
+
+| Stage                           | Introduction | Python script                                                                                           | Shell script                                                                |                      
+|:--------------------------------|:-------------|:--------------------------------------------------------------------------------------------------------|:----------------------------------------------------------------------------|
+| Continue Pretraining            | 增量预训练        | [pretraining.py](https://github.com/shibing624/MedicalGPT/blob/main/pretraining.py)                     | [run_pt.sh](https://github.com/shibing624/MedicalGPT/blob/main/run_pt.sh)   | 
+| Supervised Fine-tuning          | 有监督微调        | [supervised_finetuning.py](https://github.com/shibing624/MedicalGPT/blob/main/supervised_finetuning.py) | [run_sft.sh](https://github.com/shibing624/MedicalGPT/blob/main/run_sft.sh) | 
+| Direct Preference Optimization  | 直接偏好优化       | [dpo_training.py](https://github.com/shibing624/MedicalGPT/blob/main/dpo_training.py)                   | [run_dpo.sh](https://github.com/shibing624/MedicalGPT/blob/main/run_dpo.sh) | 
+| Reward Modeling                 | 奖励模型建模       | [reward_modeling.py](https://github.com/shibing624/MedicalGPT/blob/main/reward_modeling.py)             | [run_rm.sh](https://github.com/shibing624/MedicalGPT/blob/main/run_rm.sh)   | 
+| Reinforcement Learning          | 强化学习         | [rl_training.py](https://github.com/shibing624/MedicalGPT/blob/main/rl_training.py)                     | [run_rl.sh](https://github.com/shibing624/MedicalGPT/blob/main/run_rl.sh)   | 
+
+- 提供完整PT+SFT+DPO全阶段串起来训练的pipeline：[run_training_dpo_pipeline.ipynb](https://github.com/shibing624/MedicalGPT/blob/main/run_training_dpo_pipeline.ipynb) ，其对应的colab： [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/shibing624/MedicalGPT/blob/main/run_training_dpo_pipeline.ipynb)，运行完大概需要15分钟，我运行成功后的副本colab：[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1kMIe3pTec2snQvLBA00Br8ND1_zwy3Gr?usp=sharing)
+- 提供完整PT+SFT+RLHF全阶段串起来训练的pipeline：[run_training_pipeline.ipynb](https://github.com/shibing624/MedicalGPT/blob/main/run_training_pipeline.ipynb) ，其对应的colab： [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/shibing624/MedicalGPT/blob/main/run_training_pipeline.ipynb) ，运行完大概需要20分钟，我运行成功后的副本colab：[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1RGkbev8D85gR33HJYxqNdnEThODvGUsS?usp=sharing)
+- [训练参数说明wiki](https://github.com/shibing624/MedicalGPT/wiki/%E8%AE%AD%E7%BB%83%E5%8F%82%E6%95%B0%E8%AF%B4%E6%98%8E)
+- [数据集wiki](https://github.com/shibing624/MedicalGPT/wiki/%E6%95%B0%E6%8D%AE%E9%9B%86)
+- [扩充词表wiki](https://github.com/shibing624/MedicalGPT/wiki/%E6%89%A9%E5%85%85%E4%B8%AD%E6%96%87%E8%AF%8D%E8%A1%A8)
+- [FAQ](https://github.com/shibing624/MedicalGPT/wiki/FAQ)
+
+#### Supported Models
+
+| 模型名                                                   | 模型大小                     | Template      |
+| ------------------------------------------------------- | --------------------------- |---------------|
+| [BLOOMZ](https://huggingface.co/bigscience/bloomz)      | 560M/1.1B/1.7B/3B/7.1B/176B | vicuna        |
+| [LLaMA](https://github.com/facebookresearch/llama)      | 7B/13B/33B/65B              | alpaca        |
+| [LLaMA-2](https://huggingface.co/meta-llama)            | 7B/13B/70B                  | llama2        |
+| [Baichuan](https://github.com/baichuan-inc/baichuan-13B) | 7B/13B                      | baichuan-chat |
+| [InternLM](https://github.com/InternLM/InternLM)        | 7B                          | intern        |
+| [Qwen](https://github.com/QwenLM/Qwen-7B)               | 7B                          | chatml        |
+| [XVERSE](https://github.com/xverse-ai/XVERSE-13B)       | 13B                         | xverse        |
+| [ChatGLM](https://github.com/THUDM/ChatGLM-6B)         | 6B                          | chatglm       |
+| [ChatGLM2](https://github.com/THUDM/ChatGLM2-6B)        | 6B                          | chatglm2      |
+
+The following models are tested:
+
+bloom:
+- [bigscience/bloomz-560m](https://huggingface.co/bigscience/bloomz-560m)
+- [bigscience/bloomz-1b7](https://huggingface.co/bigscience/bloomz-1b7)
+- [bigscience/bloomz-7b1](https://huggingface.co/bigscience/bloomz-7b1)
+
+llama:
+- [shibing624/chinese-alpaca-plus-7b-hf](https://huggingface.co/shibing624/chinese-alpaca-plus-7b-hf)
+- [shibing624/chinese-alpaca-plus-13b-hf](https://huggingface.co/shibing624/chinese-alpaca-plus-13b-hf)
+- [minlik/chinese-llama-plus-7b-merged](https://huggingface.co/minlik/chinese-llama-plus-7b-merged)
+- [shibing624/chinese-llama-plus-13b-hf](https://huggingface.co/shibing624/chinese-llama-plus-13b-hf)
+- [decapoda-research/llama-7b-hf](https://huggingface.co/decapoda-research/llama-7b-hf)
+- [IDEA-CCNL/Ziya-LLaMA-13B-v1](https://huggingface.co/IDEA-CCNL/Ziya-LLaMA-13B-v1)
+
+llama2:
+- [daryl149/llama-2-7b-chat-hf](https://huggingface.co/daryl149/llama-2-7b-chat-hf)
+- [meta-llama/Llama-2-7b-chat-hf](https://huggingface.co/meta-llama/Llama-2-7b-chat-hf)
+- [ziqingyang/chinese-alpaca-2-7b](https://huggingface.co/ziqingyang/chinese-alpaca-2-7b)
+
+chatglm:
+- [THUDM/chatglm-6b](https://huggingface.co/THUDM/chatglm-6b)
+- [THUDM/chatglm2-6b](https://huggingface.co/THUDM/chatglm2-6b)
+
+baichuan:
+- [baichuan-inc/baichuan-7B](https://huggingface.co/baichuan-inc/baichuan-7B)
+- [baichuan-inc/Baichuan-13B-Base](https://huggingface.co/baichuan-inc/Baichuan-13B-Base)
+- [baichuan-inc/Baichuan-13B-Chat](https://huggingface.co/baichuan-inc/Baichuan-13B-Chat)
+
+xverse:
+- [xverse/XVERSE-13B-Chat](https://huggingface.co/xverse/XVERSE-13B-Chat)
+
+Qwen:
+- [Qwen/Qwen-7B-Chat](https://huggingface.co/Qwen/Qwen-7B-Chat)
+
+## 💻 Inference 
+训练完成后，现在我们加载训练好的模型，验证模型生成文本的效果。
+
+```shell
+CUDA_VISIBLE_DEVICES=0 python inference.py \
+    --model_type base_model_type \
+    --base_model path_to_model_hf_dir \
+    --tokenizer_path path_to_model_hf_dir \
+    --lora_model path_to_lora \
+    --interactive
+```
+
+参数说明：
+
+- `--model_type {base_model_type}`：预训练模型类型，如llama、bloom、chatglm等
+- `--base_model {base_model}`：存放HF格式的LLaMA模型权重和配置文件的目录
+- `--tokenizer_path {base_model}`：存放HF格式的LLaMA模型权重和配置文件的目录
+- `--lora_model {lora_model}`：LoRA解压后文件所在目录，也可使用HF Model Hub模型调用名称。如果已经合并了LoRA权重到预训练模型，则可以不提供此参数
+- `--tokenizer_path {tokenizer_path}`：存放对应tokenizer的目录。若不提供此参数，则其默认值与--base_model相同
+- `--template_name`：模板名称，如`vicuna`、`alpaca`等。若不提供此参数，则其默认值是vicuna
+- `--interactive`：以交互方式启动多轮问答，使用流式推理
+- `--data_file {file_name}`：非交互方式启动下，按行读取file_name中的的内容进行预测
+- `--predictions_file {file_name}`：非交互式方式下，将预测的结果以json格式写入file_name
+- `--resize_emb`：是否调整embedding大小，若不调整，则使用预训练模型的embedding大小，默认不调整
+- `--only_cpu`：仅使用CPU进行推理
+- `--gpus {gpu_ids}`：指定使用的GPU设备编号，默认为0。如使用多张GPU，以逗号分隔，如0,1,2
+
+
+#### Inference Examples
+[shibing624/vicuna-baichuan-13b-chat](https://huggingface.co/shibing624/vicuna-baichuan-13b-chat) inference examples:
+
+<details><summary>Model output case</summary>
+
+| Human                                                                              | Assitant                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                        |
+|------------------------------------------------------------------------------------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
+| 小孩发烧怎么办                                                                            | 发烧是身体对感染或其他疾病的自然反应。以下是一些可以帮助缓解发烧症状的建议：<br/>1. 保持身体水分充足：发烧时，身体会失去水分，因此需要保持水分充足。<br/>2. 使用冰袋：将冰袋放在额头或颈部可以帮助降低体温。<br/>3. 使用退烧药：退烧药可以帮助降低体温。<br/>4. 保持身体温度：保持身体温度可以帮助降低体温。<br/>5. 使用冷水或温水擦拭身体：使用冷水或温水擦拭身体可以帮助降低体温。<br/>6. 使用退烧药：退烧药可以帮助降低体温。<br/>7. 保持身体温度：保持身体温度可以帮助降低体温。<br/>8. 使用冷水或温水擦拭身体：使用冷水或温水擦拭身体可以帮助降低体温。<br/>9. 使用退烧药：退烧药可以帮助降低体温。<br/>10. 保持身体温度：保持身体温度可以帮助降低体温。                                                                                                                                                                                                                                                         |
+| 一岁宝宝发烧能吃啥药？                                                                        | 1岁宝宝发烧时，可以考虑以下几种药物：<br/>1. 退烧药：如布洛芬(Ibuprofen)或对乙酰氨基酚(Paracetamol)。这些药物可以帮助降低体温，缓解发烧症状。<br/>2. 抗生素：如果发烧是由细菌感染引起的，可以考虑使用抗生素。但请务必在医生的建议下使用。<br/>3. 抗病毒药物：如果发烧是由病毒感染引起的，可以考虑使用抗病毒药物。但请务必在医生的建议下使用。<br/>4. 解热镇痛药：如阿司匹林(Aspirin)或对乙酰氨基酚(Paracetamol)。这些药物可以帮助缓解发烧症状。<br/>5. 抗组胺药：如苯海拉明(Diphenhydramine)或非索非那定(Naproxen)。这些药物可以帮助缓解发烧症状。<br/>请务必在医生的建议下使用这些药物，并遵循药物说明书上的用量和用法。同时，保持宝宝充足的水分摄入，并确保宝宝有足够的休息。                                                                                                                                                                                                                             |
+| 肛门病变可能是什么疾病的症状                                                                     | 肛门病变可能由多种原因引起，包括感染、炎症、肿瘤或其他疾病。一些常见的肛门病变包括：<br/>1. 肛裂：这是一种肛门皮肤撕裂，通常由便秘或其他压力因素引起。<br/>2. 肛瘘：这是一种肛门感染，导致肛门皮肤和直肠之间形成瘘管。<br/>3. 肛门疣：这是一种肛门皮肤上形成的疣状病变，通常由人类乳头状瘤病毒(HPV)引起。<br/>4. 肛门脓肿：这是一种肛门感染，导致肛门皮肤下形成脓肿。<br/>5. 肛门肿瘤：这是一种肛门皮肤或直肠内形成的肿瘤。<br/>6. 肛门溃疡：这是一种肛门皮肤上形成的溃疡，通常由感染或炎症引起。<br/>7. 肛门瘙痒：这是一种肛门皮肤瘙痒的症状，可能由感染、炎症或其他疾病引起。<br/>如果您有肛门病变的症状，请咨询医疗保健专业人员以获得适当的诊断和治疗。                                                                                                                                                                                                                                                                  |
+| 医生，您好白带增多泛黄怎么回事？                                                                   | 您好，白带增多和发黄可能是由多种原因引起的，包括感染、荷尔蒙变化、阴道干燥或其他妇科问题。以下是一些可能的原因：<br/>1. 感染：白带增多和发黄可能是由阴道或尿道感染引起的。感染可能由细菌、病毒或真菌引起。<br/>2. 荷尔蒙变化：荷尔蒙变化，如月经周期或更年期，可能会导致白带增多和发黄。<br/>3. 阴道干燥：阴道干燥可能会导致白带增多和发黄。<br/>4. 其他妇科问题：其他妇科问题，如子宫颈炎、子宫内膜炎或子宫内膜异位症，也可能导致白带增多和发黄。<br/>如果您担心自己的白带，建议您咨询妇科医生进行检查和诊断。医生可以根据您的症状和检查结果确定原因并提供适当的治疗。                                                                                                                                                                                                                                                                                                                          |
+| 药流后有巧克力白带是没流干净吗                                                                    | 药流后出现巧克力白带并不一定意味着没有流干净。巧克力白带是由于子宫内膜脱落和出血引起的，通常在药流后的一段时间内会出现。<br/><br/>药流后出现巧克力白带的原因可能包括：<br/><br/>1. 子宫内膜��落：药流后，子宫内膜会脱落，导致出血。这种出血可能表现为巧克力白带。<br/>2. 感染：药流后，子宫内膜和宫颈口可能存在创伤，容易受到细菌感染。感染可能导致巧克力白带。<br/>3. 激素水平变化：药流后，激素水平可能发生变化，导致巧克力白带。<br/><br/>如果巧克力白带持续时间较长，或者伴有其他症状，如腹痛、发热、分泌物异味等，建议及时就诊，进行相关检查，如妇科检查、B超等，以确定是否存在感染或其他问题。在医生的指导下进行相应的治疗。                                                                                                                                                                                                                                                                                      |
+| 维胺酯维E乳膏能治理什么疾病                                                                     | 维胺酯维E乳膏是一种外用药物，主要用于治疗痤疮、脂溢性皮炎等皮肤病。它可以帮助减少皮肤油脂分泌，减轻炎症，促进皮肤修复。<br/><br/>需要注意的是，维胺酯维E乳膏并非万能药，对于某些皮肤病可能并不适用。在使用前，建议咨询医生或药师，了解药物的适应症、禁忌症、用法用量等信息。同时，使用药物时应遵循医嘱，避免自行滥用。                                                                                                                                                                                                                                                                                                                                                                                                                                                               |
+| 痤疮是什么原因导致的?                                                                        | 痤疮，也称为粉刺，是由皮肤毛孔堵塞引起的。堵塞毛孔的原因有很多，包括：<br/>1. 皮脂腺过度活跃：皮脂腺产生过多油脂，导致毛孔堵塞。<br/>2. 毛孔堵塞：皮肤表面的死皮细胞、油脂和细菌堵塞毛孔。<br/>3. 炎症：痤疮通常伴随着炎症，这可能是由于皮肤感染或免疫系统对痤疮的反应引起的。<br/>4. 遗传因素：某些人可能更容易患痤疮，这可能是由于遗传因素引起的。<br/>5. 荷尔蒙变化：痤疮在青春期和月经期间更为常见，这可能是由于荷尔蒙变化引起的。<br/>6. 生活方式因素：某些生活方式因素，如压力、不良饮食和缺乏运动，也可能导致痤疮。<br/>7. 化妆品和护肤品：某些化妆品和护肤品可能导致痤疮，特别是那些含有刺激性成分的产品。<br/>8. 药物：某些药物，如口服避孕药，也可能导致痤疮。<br/>要预防痤疮，可以采取一些措施，如保持皮肤清洁、使用温和的护肤品、避免过度使用化妆品和护肤品、保持健康的生活方式和饮食习惯、避免压力和荷尔蒙变化。如果痤疮严重，可能需要寻求皮肤科医生的帮助。                                                                                                                                                         |
+| 脸上长青春痘怎么办,能吃啥药,需要注意啥                                                               | 青春痘，又称痤疮，是一种常见的皮肤病。青春痘的治疗需要综合考虑饮食、生活习惯、药物等多个方面。以下是一些建议：<br/><br/>1. 饮食：<br/>* 避免高糖、高脂、高油的食物，如快餐、甜点、油炸食品等。<br/>* 多吃新鲜蔬菜、水果、全谷类食物，以增加纤维摄入，促进肠道蠕动。<br/>* 避免摄入过多的咖啡因、酒精、辛辣食物等刺激性食物。<br/><br/>2. 生活习惯：<br/>* 保持良好的作息时间，保证充足的睡眠。<br/>* 避免长时间使用电脑、手机等电子产品，以免辐射对皮肤造成刺激。<br/>* 保持皮肤清洁，每天早晚使用温和的洁面产品清洁皮肤。<br/>* 避免挤压青春痘，以免感染和留下疤痕。<br/><br/>3. 药物：<br/>* 外用药物：可以使用含有水杨酸、苯甲酸、壬二酸等成分的护肤品，帮助去除皮肤表面的角质，减少痘痘的形成。<br/>* 口服药物：如果青春痘严重，可以考虑口服抗生素、维生素B族、锌等药物，但需在医生指导下使用。<br/><br/>4. 注意：<br/>* 不要随意使用化妆品、护肤品，以免刺激皮肤。<br/>* 不要过度清洁皮肤，以免破坏皮肤屏障。<br/>* 保持心情舒畅，避免焦虑、抑郁等负面情绪。<br/><br/>总之，青春痘的治疗需要综合考虑饮食、生活习惯、药物等多个方面。在治疗过程中，要保持耐心和信心，遵循医生的建议，逐步改善皮肤状况。 |
+| 由于输精管堵塞会不会导致男性不育？���医生您好，我同事这几天不开心，我问她怎么了，她说她老公输精管堵塞，所以一直没孩子，不是不想要，请问输精管堵塞会不会导致男性不育？ | 输精管堵塞可能会导致男性不育。输精管是连接睾丸和尿道的一条管道，负责将精子从睾丸输送到尿道。如果输精管堵塞，精子就无法通过输精管到达尿道，从而导致不育。<br/><br/>输精管堵塞的原因有很多，包括先天性发育异常、炎症、肿瘤、外伤等。输精管堵塞的治疗方法包括手术疏通、药物治疗等。如果输精管堵塞严重，可能需要进行手术疏通。<br/><br/>需要注意的是，输精管堵塞并不一定导致不育，因为还有其他途径可以将精子输送到女性体内。例如，精子可以通过附睾、前列腺等组织到达女性体内。因此，如果输精管堵塞，建议及时就医，寻求专业医生的帮助。                                                                                                                                                                                                                                                                                                                                                     |
+| 心肌梗塞能治疗吗？                                                                          | 心肌梗死是一种严重的疾病，如果及时治疗，可以挽救生命并防止长期残疾。治疗心肌梗死的关键是尽早识别和治疗。<br/><br/>心肌梗死的治疗通常包括以下步骤：<br/><br/>1. 立即就医：如果出现胸痛、呼吸急促或晕厥等症状，立即就医。<br/><br/>2. 药物治疗：医生可能会开处方药物，如硝酸甘油或阿司匹林，以帮助缓解症状并防止进一步的损伤。<br/><br/>3. 溶栓治疗：如果患者在到达医院前出现心肌梗死，医生可能会进行溶栓治疗，通过注射溶解血栓的药物来帮助恢复血流。<br/><br/>4. 手术治疗：如果药物治疗或溶栓治疗无效，医生可能会建议进行手术，如冠状动脉搭桥手术或心脏支架植入术。<br/><br/>5. 康复和预防：心肌梗死的治疗通常包括康复计划，以帮助患者恢复身体功能并预防未来的心脏事件。此外，患者还应遵循医生的建议，以降低患心脏病的风险。<br/><br/>需要注意的是，心肌梗死的治疗取决于患者的具体情况和症状。因此，如果您或您认识的人出现心肌梗死症状，请立即就医。                                                                                                                                                            |
+| 怎么让自己精力充沛，列5点建议                                                                    | 1. 保持充足的睡眠：睡眠对于保持精力充沛至关重要。确保每晚获得7-9小时的高质量睡眠。<br/>2. 均衡饮食：均衡饮食对于保持能量水平至关重要。确保摄入足够的蛋白质、碳水化合物和健康脂肪。<br/>3. 定期锻炼：定期锻炼可以帮助提高能量水平。尝试每天进行至少30分钟的有氧运动，如散步、跑步或骑自行车。<br/>4. 管理压力：压力会消耗能量，因此管理压力对于保持精力充沛至关重要。尝试找到放松和减压的方法，如冥想、瑜伽或深呼吸练习。<br/>5. 保持水分：脱水会降低能量水平。确保每天喝足够的水，并根据需要补充水分。                                                                                                                                                                                                                                                                                                                                                         |
+| 写个春天的诗                                                                             | 春天，<br/>是大地苏醒的季节，<br/>是万物复苏的季节，<br/>是生命绽放的季节。<br/>春天，<br/>是阳光明媚的季节，<br/>是鸟语花香的季节，<br/>是绿意盎然的季节，<br/>是希望无限的季节。<br/>春天，<br/>是人们欢笑的季节，<br/>是孩子们嬉戏的季节，<br/>是恋人们相拥的季节，<br/>是人们追求梦想的季节。                                                                                                                                                                                                                                                                                                                                                                                                                                               |
+| 南极洲为啥没有北极熊?                                                                        | 南极洲没有北极熊，因为北极熊生活在北极地区，而南极洲位于南半球，气候和环境与北极地区有很大不同。南极洲��气候非常寒冷，常年被冰雪覆盖，而北极熊则生活在北极地区，那里有丰富的海洋生物和陆地哺乳动物。此外，南极洲的生态系统也与北极地区不同，没有北极熊所需的生存条件。                                                                                                                                                                                                                                                                                                                                                                                                                                                                                             |
+
+</details>
+
+## 📚 Dataset 
+### 医疗数据集
+
+- 240万条中文医疗数据集(包括预训练、指令微调和奖励数据集)：[shibing624/medical](https://huggingface.co/datasets/shibing624/medical)
+- 22万条中文医疗对话数据集(华佗项目)：[FreedomIntelligence/HuatuoGPT-sft-data-v1](https://huggingface.co/datasets/FreedomIntelligence/HuatuoGPT-sft-data-v1)
+
+### 通用数据集
+
+#### Pretraining datasets
+- 16GB中英文无监督、平行语料[Linly-AI/Chinese-pretraining-dataset](https://huggingface.co/datasets/Linly-AI/Chinese-pretraining-dataset)
+- 524MB中文维基百科语料[wikipedia-cn-20230720-filtered](https://huggingface.co/datasets/pleisto/wikipedia-cn-20230720-filtered)
+#### SFT datasets
+- 10万条多语言ShareGPT GPT4多轮对话数据集：[shibing624/sharegpt_gpt4](https://huggingface.co/datasets/shibing624/sharegpt_gpt4) [本项目支持格式]
+- 9万条英文ShareGPT多轮对话数集：[anon8231489123/ShareGPT_Vicuna_unfiltered](https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered) [本项目支持格式]
+- 50万条中文ChatGPT指令Belle数据集：[BelleGroup/train_0.5M_CN](https://huggingface.co/datasets/BelleGroup/train_0.5M_CN)
+- 100万条中文ChatGPT指令Belle数据集：[BelleGroup/train_1M_CN](https://huggingface.co/datasets/BelleGroup/train_1M_CN)
+- 5万条英文ChatGPT指令Alpaca数据集：[50k English Stanford Alpaca dataset](https://github.com/tatsu-lab/stanford_alpaca#data-release)
+- 2万条中文ChatGPT指令Alpaca数据集：[shibing624/alpaca-zh](https://huggingface.co/datasets/shibing624/alpaca-zh)
+- 69万条中文指令Guanaco数据集(Belle50万条+Guanaco19万条)：[Chinese-Vicuna/guanaco_belle_merge_v1.0](https://huggingface.co/datasets/Chinese-Vicuna/guanaco_belle_merge_v1.0)
+- 5万条英文ChatGPT多轮对话数据集：[RyokoAI/ShareGPT52K](https://huggingface.co/datasets/RyokoAI/ShareGPT52K)
+- 80万条中文ChatGPT多轮对话数据集：[BelleGroup/multiturn_chat_0.8M](https://huggingface.co/datasets/BelleGroup/multiturn_chat_0.8M)
+- 116万条中文ChatGPT多轮对话数据集：[fnlp/moss-002-sft-data](https://huggingface.co/datasets/fnlp/moss-002-sft-data)
+- 3.8万条中文ShareGPT多轮对话数据集：[FreedomIntelligence/ShareGPT-CN](https://huggingface.co/datasets/FreedomIntelligence/ShareGPT-CN)
+
+#### Reward Model datasets
+- 原版的oasst1数据集：[OpenAssistant/oasst1](https://huggingface.co/datasets/OpenAssistant/oasst1)
+- 2万条多语言oasst1的reward数据集：[tasksource/oasst1_pairwise_rlhf_reward](https://huggingface.co/datasets/tasksource/oasst1_pairwise_rlhf_reward)[本项目支持格式]
+- 11万条英文hh-rlhf的reward数据集：[Dahoas/full-hh-rlhf](https://huggingface.co/datasets/Dahoas/full-hh-rlhf)
+- 9万条英文reward数据集(来自Anthropic's Helpful Harmless dataset)：[Dahoas/static-hh](https://huggingface.co/datasets/Dahoas/static-hh)
+- 7万条英文reward数据集（来源同上）：[Dahoas/rm-static](https://huggingface.co/datasets/Dahoas/rm-static)
+- 7万条繁体中文的reward数据集（翻译自rm-static）[liswei/rm-static-m2m100-zh](https://huggingface.co/datasets/liswei/rm-static-m2m100-zh)
+- 7万条英文Reward数据集：[yitingxie/rlhf-reward-datasets](https://huggingface.co/datasets/yitingxie/rlhf-reward-datasets)
+- 3千条中文知乎问答偏好数据集：[liyucheng/zhihu_rlhf_3k](https://huggingface.co/datasets/liyucheng/zhihu_rlhf_3k)
+
+## ✅ Todo
+
+1. [x] add multi-round dialogue data fine-tuning method
+2. [x] add reward model fine-tuning
+3. [x] add rl fine-tuning
+4. [x] add medical reward dataset
+5. [x] add llama in8/int4 training
+6. [x] add all training and predict demo in colab
+7. [x] add dpo training
+
+## ☎️ Contact
+
+- Issue(建议)
+  ：[![GitHub issues](https://img.shields.io/github/issues/shibing624/MedicalGPT.svg)](https://github.com/shibing624/MedicalGPT/issues)
+- 邮件我：xuming: [email protected]
+- 微信我： 加我*微信号：xuming624, 备注：姓名-公司名-NLP* 进NLP交流群。
+
+<img src="https://github.com/shibing624/MedicalGPT/blob/main/docs/wechat.jpeg" width="200" />
+
+## ⚠️ 局限性、使用限制与免责声明
+
+基于当前数据和基础模型训练得到的SFT模型，在效果上仍存在以下问题：
+
+1. 在涉及事实性的指令上可能会产生违背事实的错���回答。
+
+2. 对于具备危害性的指令无法很好的鉴别，由此会产生危害性言论。
+
+3. 在一些涉及推理、代码、多轮对话等场景下模型的能力仍有待提高。
+
+基于以上模型局限性，我们要求开发者仅将我们开源的模型权重及后续用此项目生成的衍生物用于研究目的，不得用于商业，以及其他会对社会带来危害的用途。
+
+本项目仅可应用于研究目的，项目开发者不承担任何因使用本项目（包含但不限于数据、模型、代码等）导致的危害或损失。详细请参考[免责声明](https://github.com/shibing624/MedicalGPT/blob/main/DISCLAIMER)。
+
+项目代码的授权协议为 [The Apache License 2.0](/LICENSE)，代码可免费用做商业用途，模型权重和数据只能用于研究目的。请在产品说明中附加MedicalGPT的链接和授权协议。
+
+
+## 😇 Citation
+
+如果你在研究中使用了MedicalGPT，请按如下格式引用：
+
+```latex
+@misc{MedicalGPT,
+  title={MedicalGPT: Training Medical GPT Model},
+  author={Ming Xu},
+  year={2023},
+  howpublished={\url{https://github.com/shibing624/MedicalGPT}},
+}
+```
+
+## 😍 Contribute
+
+项目代码还很粗糙，如果大家对代码有所改进，欢迎提交回本项目，在提交之前，注意以下两点：
+
+- 在`tests`添加相应的单元测试
+- 使用`python -m pytest`来运行所有单元测试，确保所有单测都是通过的
+
+之后即可提交PR。
+
+## 💕 Acknowledgements 
+
+- [Direct Preference Optimization:Your Language Model is Secretly a Reward Model](https://arxiv.org/pdf/2305.18290.pdf)
+- [tloen/alpaca-lora](https://github.com/tloen/alpaca-lora/blob/main/finetune.py)
+- [ymcui/Chinese-LLaMA-Alpaca](https://github.com/ymcui/Chinese-LLaMA-Alpaca)
+
+Thanks for their great work!

README_EN.md

@@ -0,0 +1,224 @@
+[**🇨🇳中文**](https://github.com/shibing624/MedicalGPT/blob/main/README.md) | [**🌐English**](https://github.com/shibing624/MedicalGPT/blob/main/README_EN.md) | [**📖文档/Docs**](https://github.com/shibing624/MedicalGPT/wiki) | [**🤖模型/Models**](https://huggingface.co/shibing624) 
+
+<div align="center">
+  <a href="https://github.com/shibing624/MedicalGPT">
+    <img src="https://github.com/shibing624/MedicalGPT/blob/main/docs/logo.png" width="120" alt="Logo">
+  </a>
+</div>
+
+-----------------
+
+# MedicalGPT: Training Medical GPT Model
+[![HF Models](https://img.shields.io/badge/Hugging%20Face-shibing624-green)](https://huggingface.co/shibing624)
+[![Github Stars](https://img.shields.io/github/stars/shibing624/MedicalGPT?color=yellow)](https://star-history.com/#shibing624/MedicalGPT&Timeline)
+[![Contributions welcome](https://img.shields.io/badge/contributions-welcome-brightgreen.svg)](CONTRIBUTING.md)
+[![License Apache 2.0](https://img.shields.io/badge/license-Apache%202.0-blue.svg)](LICENSE)
+[![python_version](https://img.shields.io/badge/Python-3.8%2B-green.svg)](requirements.txt)
+[![GitHub issues](https://img.shields.io/github/issues/shibing624/MedicalGPT.svg)](https://github.com/shibing624/MedicalGPT/issues)
+[![Wechat Group](http://vlog.sfyc.ltd/wechat_everyday/wxgroup_logo.png?imageView2/0/w/60/h/20)](#Contact)
+
+## 📖 Introduction
+
+**MedicalGPT** training medical GPT model with ChatGPT training pipeline, implemantation of Pretraining, 
+Supervised Finetuning, Reward Modeling and Reinforcement Learning.
+
+
+<img src="https://github.com/shibing624/MedicalGPT/blob/main/docs/GPT_Training.jpg" width="860" />
+
+Training MedicalGPT model：
+
+- Stage 1：PT(Continue PreTraining), Pre-training the LLaMA model on massive domain document data to inject domain knowledge
+- Stage 2: SFT (Supervised Fine-tuning) has supervised fine-tuning, constructs instruction fine-tuning data sets, and performs instruction fine-tuning on the basis of pre-trained models to align instruction intentions
+- Stage 3: RM (Reward Model) reward model modeling, constructing a human preference ranking data set, training the reward model to align human preferences, mainly the "HHH" principle, specifically "helpful, honest, harmless"
+- Stage 4: RL (Reinforcement Learning) is based on human feedback reinforcement learning (RLHF), using the reward model to train the SFT model, and the generation model uses rewards or penalties to update its strategy in order to generate higher quality, more in line with human preferences text
+
+## ▶️ Demo
+
+- Hugging Face Demo: doing
+
+We provide a simple Gradio-based interactive web interface. After the service is started, it can be accessed through a browser, enter a question, and the model will return an answer. The command is as follows:
+```shell
+python scripts/gradio_demo.py --base_model path_to_llama_hf_dir --lora_model path_to_lora_dir
+```
+
+Parameter Description:
+
+- `--base_model {base_model}`: directory to store LLaMA model weights and configuration files in HF format, or use the HF Model Hub model call name
+- `--lora_model {lora_model}`: The directory where the LoRA file is located, and the name of the HF Model Hub model can also be used. If the lora weights have been merged into the pre-trained model, delete the --lora_model parameter
+- `--tokenizer_path {tokenizer_path}`: Store the directory corresponding to the tokenizer. If this parameter is not provided, its default value is the same as --lora_model; if the --lora_model parameter is not provided, its default value is the same as --base_model
+- `--use_cpu`: use only CPU for inference
+- `--gpus {gpu_ids}`: Specifies the number of GPU devices used, the default is 0. If using multiple GPUs, separate them with commas, such as 0,1,2
+
+
+## 🚀 Training Pipeline
+
+### Stage 1: Continue Pretraining
+
+Based on the llama-7b model, use medical encyclopedia data to continue pre-training, and expect to inject medical knowledge into the pre-training model to obtain the llama-7b-pt model. This step is optional
+
+
+```shell
+cd scripts
+sh run_pt.sh
+```
+
+[Training Detail wiki](https://github.com/shibing624/MedicalGPT/wiki/Training-Details)
+
+### Stage 2: Supervised FineTuning
+Based on the llama-7b-pt model, the llama-7b-sft model is obtained by using medical question-and-answer data for supervised fine-tuning. This step is required
+
+Supervised fine-tuning of the base llama-7b-pt model to create llama-7b-sft
+
+```shell
+cd scripts
+sh run_sft.sh
+```
+
+[Training Detail wiki](https://github.com/shibing624/MedicalGPT/wiki/Training-Details)
+
+### Stage 3: Reward Modeling
+RM(Reward Model): reward model modeling
+
+In principle, we can directly use human annotations to fine-tune the model with RLHF.
+
+However, this will require us to send some samples to humans to be scored after each round of optimization. This is expensive and slow due to the large number of training samples required for convergence and the limited speed at which humans can read and annotate them.
+A better strategy than direct feedback is to train a reward model RM on the human annotated set before entering the RL loop. The purpose of the reward model is to simulate human scoring of text.
+
+The best practice for building a reward model is to rank the prediction results, that is, for each prompt (input text) corresponding to two results (yk, yj), the model predicts which score the human annotation is higher.
+The RM model is trained by manually marking the scoring results of the SFT model. The purpose is to replace manual scoring. It is essentially a regression model used to align human preferences, mainly based on the "HHH" principle, specifically "helpful, honest, harmless".
+
+
+Based on the llama-7b-sft model, the reward preference model is trained using medical question and answer preference data, and the llama-7b-reward model is obtained after training. This step is required
+
+Reward modeling using dialog pairs from the reward dataset using the llama-7b-sft to create llama-7b-reward:
+
+```shell
+cd scripts
+sh run_rm.sh
+```
+[Training Detail wiki](https://github.com/shibing624/MedicalGPT/wiki/Training-Details)
+
+### Stage 4: Reinforcement Learning
+The purpose of the RL (Reinforcement Learning) model is to maximize the output of the reward model. Based on the above steps, we have a fine-tuned language model (llama-7b-sft) and reward model (llama-7b-reward).
+The RL loop is ready to execute.
+
+This process is roughly divided into three steps:
+
+1. Enter prompt, the model generates a reply
+2. Use a reward model to score responses
+3. Based on the score, a round of reinforcement learning for policy optimization (PPO)
+
+<img src=https://huggingface.co/datasets/trl-internal-testing/example-images/resolve/main/blog/stackllama/trl_loop.png height=400 />
+
+Reinforcement Learning fine-tuning of llama-7b-sft with the llama-7b-reward reward model to create llama-7b-rl
+
+```shell
+cd scripts
+sh run_rl.sh
+```
+[Training Detail wiki](https://github.com/shibing624/MedicalGPT/wiki/Training-Details)
+
+## 🔥 Inference 
+After the training is complete, now we load the trained model to verify the effect of the model generating text.
+
+```shell
+python scripts/inference.py \
+    --base_model path_to_llama_hf_dir \
+    --lora_model path_to_lora \
+    --with_prompt \
+    --interactive
+```
+
+Parameter Description:
+
+- `--base_model {base_model}`: Directory to store LLaMA model weights and configuration files in HF format
+- `--lora_model {lora_model}`: The directory where the LoRA file is decompressed, and the name of the HF Model Hub model can also be used. If you have incorporated LoRA weights into the pre-trained model, you can not provide this parameter
+- `--tokenizer_path {tokenizer_path}`: Store the directory corresponding to the tokenizer. If this parameter is not provided, its default value is the same as --lora_model; if the --lora_model parameter is not provided, its default value is the same as --base_model
+- `--with_prompt`: Whether to merge the input with the prompt template. Be sure to enable this option if loading an Alpaca model!
+- `--interactive`: start interactively for multiple single rounds of question and answer
+- `--data_file {file_name}`: Start in non-interactive mode, read the contents of file_name line by line for prediction
+- `--predictions_file {file_name}`: In non-interactive mode, write the predicted results to file_name in json format
+- `--use_cpu`: use only CPU for inference
+- `--gpus {gpu_ids}`: Specifies the number of GPU devices used, the default is 0. If using multiple GPUs, separate them with commas, such as 0,1,2
+
+
+#### Inference Examples
+
+<details>
+| Human | Assistant |
+| ----- | ------ |
+| How to make yourself energetic, list 5 suggestions | Here are some suggestions to keep you energetic: <br/>1. Get enough sleep. Make sure to get enough sleep every night. <br/>2. Exercise more. Doing moderate aerobic exercise, such as jogging and swimming, can help improve your energy and stamina. 3. Eat healthy food. Choose protein-rich foods such as chicken, fish, and eggs, as well as fresh fruits and vegetables. 4. Reduce caffeine intake. Try to avoid drinking caffeinated beverages or eating caffeinated foods during the day. <br/>5. Relax. Try meditation, yoga, or other forms of relaxation to reduce stress and anxiety. |
+</details>
+<br/>
+
+
+## 📚 Dataset 
+
+- 2.4 million Chinese medical datasets (including pre-training, instruction fine-tuning and reward datasets): [shibing624/medical](https://huggingface.co/datasets/shibing624/medical)
+
+**Attach links to some general datasets and medical datasets**
+
+- Belle dataset of 500,000 Chinese ChatGPT commands: [BelleGroup/train_0.5M_CN](https://huggingface.co/datasets/BelleGroup/train_0.5M_CN)
+- Belle dataset of 1 million Chinese ChatGPT commands: [BelleGroup/train_1M_CN](https://huggingface.co/datasets/BelleGroup/train_1M_CN)
+- Alpaca dataset of 50,000 English ChatGPT commands: [50k English Stanford Alpaca dataset](https://github.com/tatsu-lab/stanford_alpaca#data-release)
+- Alpaca dataset of 20,000 Chinese GPT-4 instructions: [shibing624/alpaca-zh](https://huggingface.co/datasets/shibing624/alpaca-zh)
+- Guanaco dataset with 690,000 Chinese instructions (500,000 Belle + 190,000 Guanaco): [Chinese-Vicuna/guanaco_belle_merge_v1.0](https://huggingface.co/datasets/Chinese-Vicuna/guanaco_belle_merge_v1.0)
+- 220,000 Chinese medical dialogue datasets (HuatuoGPT project): [FreedomIntelligence/HuatuoGPT-sft-data-v1](https://huggingface.co/datasets/FreedomIntelligence/HuatuoGPT-sft-data-v1)
+
+## ✅ Todo
+
+1. [ ] Added multi-round dialogue data fine-tuning method
+2. [x] add reward model finetuning
+3. [x] add rl finetuning
+4. [x] add medical reward dataset
+5. [x] add llama in8/int4 training
+6. [ ] add all training and predict demo in colab
+## ☎️ Contact
+
+- Issue (suggestion)
+   : [![GitHub issues](https://img.shields.io/github/issues/shibing624/MedicalGPT.svg)](https://github.com/shibing624/MedicalGPT/issues)
+- Email me: xuming: [email protected]
+- WeChat Me: Add me* WeChat ID: xuming624, Remarks: Name-Company Name-NLP* Enter the NLP exchange group.
+
+<img src="https://github.com/shibing624/MedicalGPT/blob/main/docs/wechat.jpeg" width="200" />
+
+## ⚠️ Limitations, Restrictions of Use and Disclaimer
+
+The SFT model trained based on the current data and the basic model still has the following problems in terms of effect:
+
+1. Wrong answers that contradict the facts may be generated on the factual instructions.
+2. Unable to identify harmful instructions well, resulting in harmful speech.
+3. The ability of the model still needs to be improved in some scenarios involving reasoning, code, and multiple rounds of dialogue.
+
+Based on the limitations of the above models, we require developers to only use our open source model weights and subsequent derivatives generated by this project for research purposes, and not for commercial use, and other purposes that will cause harm to society.
+This project can only be used for research purposes, and the project developer is not responsible for any harm or loss caused by the use of this project (including but not limited to data, models, codes, etc.). For details, please refer to [Disclaimer](https://github.com/shibing624/MedicalGPT/blob/main/DISCLAIMER).
+The license agreement for the project code is [The Apache License 2.0](/LICENSE), the code is free for commercial use, and the model weights and data can only be used for research purposes. Please attach MedicalGPT's link and license agreement in the product description.
+
+## 😇 Citation
+
+If you used MedicalGPT in your research, please cite as follows:
+
+```latex
+@misc{MedicalGPT,
+   title={MedicalGPT: Training Medical GPT Model},
+   author={Ming Xu},
+   year={2023},
+   howpublished={\url{https://github.com/shibing624/MedicalGPT}},
+}
+```
+
+## 😍 Contribute
+
+The project code is still very rough. If you have improved the code, you are welcome to submit it back to this project. Before submitting, please pay attention to the following two points:
+
+- Add corresponding unit tests in `tests`
+- Use `python -m pytest` to run all unit tests to ensure that all unit tests are passed
+
+Then you can submit a PR.
+
+## 💕 Acknowledgements 
+
+- [tloen/alpaca-lora](https://github.com/tloen/alpaca-lora/blob/main/finetune.py)
+- [ymcui/Chinese-LLaMA-Alpaca](https://github.com/ymcui/Chinese-LLaMA-Alpaca)
+
+Thanks for their great work!

_config.yml

@@ -0,0 +1 @@
+theme: jekyll-theme-cayman

build_domain_tokenizer.py

@@ -0,0 +1,59 @@
+# -*- coding: utf-8 -*-
+"""
+@author:XuMing([email protected])
+@description: Build chinese tokenizer from corpus txt
+
+# train sentencepiece model from `corpus.txt` and makes `m.model` and `m.vocab`
+# `m.vocab` is just a reference. not used in the segmentation.
+# spm.SentencePieceTrainer.train('--input=data/pretrain/tianlongbabu.txt --model_prefix=m --vocab_size=20000')
+"""
+import argparse
+
+import sentencepiece as spm
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--in_file', default='data/pretrain/fever.txt', type=str)
+    parser.add_argument('--domain_sp_model_name', default='domain_sp', type=str)
+    parser.add_argument('--max_sentence_length', default=16384, type=int)
+    parser.add_argument('--pad_id', default=3, type=int)
+    parser.add_argument('--vocab_size', default=2236, type=int)
+    parser.add_argument('--model_type', default="BPE", type=str)
+
+    args = parser.parse_args()
+    print(args)
+
+    spm.SentencePieceTrainer.train(
+        input=args.in_file,
+        model_prefix=args.domain_sp_model_name,
+        shuffle_input_sentence=False,
+        train_extremely_large_corpus=True,
+        max_sentence_length=args.max_sentence_length,
+        pad_id=args.pad_id,
+        model_type=args.model_type,
+        vocab_size=args.vocab_size,
+        split_digits=True,
+        split_by_unicode_script=True,
+        byte_fallback=True,
+        allow_whitespace_only_pieces=True,
+        remove_extra_whitespaces=False,
+        normalization_rule_name="nfkc",
+    )
+
+    # makes segmenter instance and loads the model file (m.model)
+    sp = spm.SentencePieceProcessor()
+    model_file = args.domain_sp_model_name + '.model'
+    sp.load(model_file)
+
+    # encode: text => id
+    print(sp.encode_as_pieces('潜伏性感染又称潜在性感染。慕容复来到河边,this is a test'))
+    print(sp.encode_as_ids('this is a test'))
+
+    # decode: id => text
+    print(sp.decode_pieces(['▁This', '▁is', '▁a', '▁t', 'est']))
+    # print(sp.decode_ids([209, 31, 9, 375, 586]))
+
+
+if __name__ == '__main__':
+    main()

convert_dataset.py

@@ -0,0 +1,47 @@
+"""
+Convert alpaca dataset into sharegpt format.
+
+Usage: python convert_alpaca.py --in_file alpaca_data.json --out_file alpaca_data_sharegpt.json
+"""
+
+import argparse
+
+from datasets import load_dataset
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--in_file", type=str)
+    parser.add_argument("--out_file", type=str)
+    parser.add_argument("--data_type", type=str, default='alpaca')
+    args = parser.parse_args()
+    print(args)
+    data_files = {"train": args.in_file}
+    raw_datasets = load_dataset('json', data_files=data_files)
+    ds = raw_datasets['train']
+
+
+    def process_alpaca(examples):
+        convs = []
+        for instruction, inp, output in zip(examples['instruction'], examples['input'], examples['output']):
+            if len(inp.strip()) > 1:
+                instruction = instruction + '\n\n' + inp
+            q = instruction
+            a = output
+            convs.append([
+                {"from": "human", "value": q},
+                {"from": "gpt", "value": a}
+            ])
+        return {"conversations": convs}
+
+
+    if args.data_type in ['alpaca']:
+        ds = ds.map(process_alpaca, batched=True, remove_columns=ds.column_names, desc="Running process")
+    else:
+        # Other sharegpt dataset, need rename to conversations and remove unused columns
+        if "items" in ds.column_names:
+            ds = ds.rename(columns={"items": "conversations"})
+        columns_to_remove = ds.column_names.copy()
+        columns_to_remove.remove('conversations')
+        ds = ds.remove_columns(columns_to_remove)
+
+    ds.to_json(f"{args.out_file}", lines=True, force_ascii=False)

deepspeed_config.json

@@ -0,0 +1,43 @@
+{
+    "optimizer": {
+        "type": "AdamW",
+        "params": {
+            "lr": "auto",
+            "weight_decay": "auto",
+            "torch_adam": true,
+            "adam_w_mode": true
+        }
+    },
+    "scheduler": {
+        "type": "WarmupDecayLR",
+        "params": {
+            "warmup_min_lr": "auto",
+            "warmup_max_lr": "auto",
+            "warmup_num_steps": "auto",
+            "total_num_steps": "auto"
+        }
+    },
+    "fp16": {
+        "enabled": true,
+        "loss_scale": 0,
+        "loss_scale_window": 1000,
+        "initial_scale_power": 16,
+        "hysteresis": 2,
+        "min_loss_scale": 1
+    },
+    "zero_optimization": {
+        "stage": 2,
+        "allgather_partitions": true,
+        "allgather_bucket_size": 2e8,
+        "reduce_scatter": true,
+        "reduce_bucket_size": "auto",
+        "overlap_comm": true,
+        "contiguous_gradients": true
+    },
+    "gradient_accumulation_steps": "auto",
+    "gradient_clipping": "auto",
+    "steps_per_print": 1000,
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "wall_clock_breakdown": false
+}

dpo_training.py

@@ -0,0 +1,495 @@
+# -*- coding: utf-8 -*-
+"""
+@author:XuMing([email protected])
+@description: Train a model from SFT using DPO
+"""
+
+import os
+from dataclasses import dataclass, field
+from glob import glob
+from typing import Dict, Optional
+
+import torch
+from datasets import load_dataset
+from loguru import logger
+from peft import LoraConfig, TaskType
+from transformers import (
+    AutoConfig,
+    BloomForCausalLM,
+    AutoModelForCausalLM,
+    AutoModel,
+    LlamaTokenizer,
+    LlamaForCausalLM,
+    BloomTokenizerFast,
+    AutoTokenizer,
+    HfArgumentParser,
+    TrainingArguments,
+    BitsAndBytesConfig,
+)
+from transformers.deepspeed import is_deepspeed_zero3_enabled
+from trl import DPOTrainer
+
+os.environ["TOKENIZERS_PARALLELISM"] = "FALSE"
+os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
+
+MODEL_CLASSES = {
+    "bloom": (AutoConfig, BloomForCausalLM, BloomTokenizerFast),
+    "chatglm": (AutoConfig, AutoModel, AutoTokenizer),
+    "llama": (AutoConfig, LlamaForCausalLM, LlamaTokenizer),
+    "baichuan": (AutoConfig, AutoModelForCausalLM, AutoTokenizer),
+    "auto": (AutoConfig, AutoModelForCausalLM, AutoTokenizer),
+}
+
+
+@dataclass
+class ScriptArguments:
+    """
+    The name of the Casual LM model we wish to fine with DPO
+    """
+    # Model arguments
+    model_type: str = field(
+        default=None,
+        metadata={"help": "Model type selected in the list: " + ", ".join(MODEL_CLASSES.keys())}
+    )
+    model_name_or_path: Optional[str] = field(
+        default=None, metadata={"help": "The model checkpoint for weights initialization."}
+    )
+    tokenizer_name_or_path: Optional[str] = field(
+        default=None, metadata={"help": "The tokenizer for weights initialization."}
+    )
+    load_in_8bit: bool = field(default=False, metadata={"help": "Whether to load the model in 8bit mode or not."})
+    load_in_4bit: bool = field(default=False, metadata={"help": "Whether to load the model in 4bit mode or not."})
+    cache_dir: Optional[str] = field(
+        default=None,
+        metadata={"help": "Where do you want to store the pretrained models downloaded from huggingface.co"},
+    )
+    use_fast_tokenizer: bool = field(
+        default=False,
+        metadata={"help": "Whether to use one of the fast tokenizer (backed by the tokenizers library) or not."},
+    )
+    torch_dtype: Optional[str] = field(
+        default=None,
+        metadata={
+            "help": (
+                "Override the default `torch.dtype` and load the model under this dtype. If `auto` is passed, the "
+                "dtype will be automatically derived from the model's weights."
+            ),
+            "choices": ["auto", "bfloat16", "float16", "float32"],
+        },
+    )
+    device_map: Optional[str] = field(
+        default="auto",
+        metadata={"help": "Device to map model to. If `auto` is passed, the device will be selected automatically. "},
+    )
+    trust_remote_code: bool = field(
+        default=True,
+        metadata={"help": "Whether to trust remote code when loading a model from a remote checkpoint."},
+    )
+    # Dataset arguments
+    dataset_name: Optional[str] = field(
+        default=None, metadata={"help": "The name of the dataset to use (via the datasets library)."}
+    )
+    dataset_config_name: Optional[str] = field(
+        default=None, metadata={"help": "The configuration name of the dataset to use (via the datasets library)."}
+    )
+    train_file_dir: Optional[str] = field(default=None, metadata={"help": "The input jsonl data file folder."})
+    validation_file_dir: Optional[str] = field(default=None, metadata={"help": "The evaluation jsonl file folder."}, )
+    template_name: Optional[str] = field(default="vicuna", metadata={"help": "The prompt template name."})
+    per_device_train_batch_size: Optional[int] = field(default=4, metadata={"help": "Train batch size per device"})
+    per_device_eval_batch_size: Optional[int] = field(default=1, metadata={"help": "Eval batch size per device"})
+    max_source_length: Optional[int] = field(default=256, metadata={"help": "Max length of prompt input text"})
+    max_target_length: Optional[int] = field(default=256, metadata={"help": "Max length of output text"})
+    min_target_length: Optional[int] = field(default=4, metadata={"help": "Min length of output text"})
+    max_train_samples: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": (
+                "For debugging purposes or quicker training, truncate the number of training examples to this "
+                "value if set."
+            )
+        },
+    )
+    max_eval_samples: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": (
+                "For debugging purposes or quicker training, truncate the number of evaluation examples to this "
+                "value if set."
+            )
+        },
+    )
+    overwrite_cache: bool = field(
+        default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}
+    )
+    validation_split_percentage: Optional[int] = field(
+        default=1,
+        metadata={
+            "help": "The percentage of the train set used as validation set in case there's no validation split"
+        },
+    )
+    preprocessing_num_workers: Optional[int] = field(
+        default=4, metadata={"help": "The number of processes to use for the preprocessing."},
+    )
+    # Training arguments
+    use_peft: bool = field(default=True, metadata={"help": "Whether to use peft"})
+    qlora: bool = field(default=False, metadata={"help": "Whether to use qlora"})
+    target_modules: Optional[str] = field(default=None)
+    lora_rank: Optional[int] = field(default=8)
+    lora_dropout: Optional[float] = field(default=0.05)
+    lora_alpha: Optional[float] = field(default=16.0)
+    peft_path: Optional[str] = field(default=None)
+    do_train: bool = field(default=False, metadata={"help": "Whether to run training."})
+    do_eval: bool = field(default=False, metadata={"help": "Whether to run eval on the validation set."})
+    beta: Optional[float] = field(default=0.1, metadata={"help": "The beta parameter for DPO loss"})
+    learning_rate: Optional[float] = field(default=5e-4, metadata={"help": "Learning rate"})
+    lr_scheduler_type: Optional[str] = field(default="cosine", metadata={"help": "The lr scheduler type"})
+    warmup_steps: Optional[int] = field(default=100, metadata={"help": "The number of warmup steps"})
+    weight_decay: Optional[float] = field(default=0.05, metadata={"help": "The weight decay"})
+    optim: Optional[str] = field(default="adamw_hf", metadata={"help": "The optimizer type"})
+    fp16: Optional[bool] = field(default=True, metadata={"help": "Whether to use fp16"})
+    bf16: Optional[bool] = field(default=False, metadata={"help": "Whether to use bf16"})
+    gradient_checkpointing: Optional[bool] = field(
+        default=True, metadata={"help": "Whether to use gradient checkpointing"}
+    )
+    gradient_accumulation_steps: Optional[int] = field(
+        default=4, metadata={"help": "The number of gradient accumulation steps"}
+    )
+    save_steps: Optional[int] = field(default=50, metadata={"help": "X steps to save the model"})
+    eval_steps: Optional[int] = field(default=50, metadata={"help": "X steps to evaluate the model"})
+    logging_steps: Optional[int] = field(default=1, metadata={"help": "X steps to log the model"})
+    output_dir: Optional[str] = field(default="outputs-dpo", metadata={"help": "The output directory"})
+    max_steps: Optional[int] = field(default=200, metadata={"help": "Number of steps to train"})
+    eval_strategy: Optional[str] = field(default="steps", metadata={"help": "Evaluation strategy"})
+    remove_unused_columns: Optional[bool] = field(
+        default=False,
+        metadata={"help": "Remove unused columns from the dataset if `datasets.Dataset` is used"},
+    )
+    report_to: Optional[str] = field(default="tensorboard", metadata={"help": "Report to wandb or tensorboard"})
+
+    def __post_init__(self):
+        if self.model_type is None:
+            raise ValueError("You must specify a valid model_type to run training.")
+        if self.model_name_or_path is None:
+            raise ValueError("You must specify a valid model_name_or_path to run training.")
+
+
+def print_trainable_parameters(model):
+    """
+    Prints the number of trainable parameters in the model.
+    """
+    trainable_params = 0
+    all_param = 0
+    for _, param in model.named_parameters():
+        all_param += param.numel()
+        if param.requires_grad:
+            trainable_params += param.numel()
+    print(
+        f"trainable params: {trainable_params} || all params: {all_param} || trainable%: {100 * trainable_params / all_param}"
+    )
+
+
+def find_all_linear_names(peft_model, int4=False, int8=False):
+    """Find all linear layer names in the model. reference from qlora paper."""
+    cls = torch.nn.Linear
+    if int4 or int8:
+        import bitsandbytes as bnb
+        if int4:
+            cls = bnb.nn.Linear4bit
+        elif int8:
+            cls = bnb.nn.Linear8bitLt
+    lora_module_names = set()
+    for name, module in peft_model.named_modules():
+        if isinstance(module, cls):
+            # last layer is not add to lora_module_names
+            if 'lm_head' in name:
+                continue
+            if 'output_layer' in name:
+                continue
+            names = name.split('.')
+            lora_module_names.add(names[0] if len(names) == 1 else names[-1])
+    return sorted(lora_module_names)
+
+
+def return_prompt_and_responses(examples) -> Dict[str, str]:
+    """Load the paired dataset and convert it to the necessary format.
+
+    The dataset is converted to a dictionary with the following structure:
+    {
+        'prompt': List[str],
+        'chosen': List[str],
+        'rejected': List[str],
+    }
+
+    Prompts are structured as follows:
+      "Question: " + <prompt> + "\n\nAnswer: "
+    """
+    return {
+        "prompt": ["Question: " + question + "\n\nAnswer: " for question in examples["question"]],
+        "chosen": examples["response_chosen"],
+        "rejected": examples["response_rejected"],
+    }
+
+
+def main():
+    parser = HfArgumentParser(ScriptArguments)
+    args = parser.parse_args_into_dataclasses()[0]
+    logger.info(f"Parse args: {args}")
+
+    config_class, model_class, tokenizer_class = MODEL_CLASSES[args.model_type]
+    if args.model_type == 'bloom':
+        args.use_fast_tokenizer = True
+    # Load tokenizer
+    tokenizer_kwargs = {
+        "cache_dir": args.cache_dir,
+        "use_fast": args.use_fast_tokenizer,
+        "trust_remote_code": args.trust_remote_code,
+    }
+    tokenizer_name_or_path = args.tokenizer_name_or_path
+    if not tokenizer_name_or_path:
+        tokenizer_name_or_path = args.model_name_or_path
+    tokenizer = tokenizer_class.from_pretrained(tokenizer_name_or_path, **tokenizer_kwargs)
+    if tokenizer.pad_token_id is None:
+        tokenizer.pad_token_id = 0  # set as the <unk> token
+
+    # Get datasets
+    if args.dataset_name is not None:
+        # Downloading and loading a dataset from the hub.
+        raw_datasets = load_dataset(
+            args.dataset_name,
+            args.dataset_config_name,
+            cache_dir=args.cache_dir,
+        )
+        if "validation" not in raw_datasets.keys():
+            raw_datasets["validation"] = load_dataset(
+                args.dataset_name,
+                args.dataset_config_name,
+                split=f"train[:{args.validation_split_percentage}%]",
+                cache_dir=args.cache_dir,
+            )
+            raw_datasets["train"] = load_dataset(
+                args.dataset_name,
+                args.dataset_config_name,
+                split=f"train[{args.validation_split_percentage}%:]",
+                cache_dir=args.cache_dir,
+            )
+    else:
+        data_files = {}
+        if args.train_file_dir is not None and os.path.exists(args.train_file_dir):
+            train_data_files = glob(f'{args.train_file_dir}/**/*.json', recursive=True) + glob(
+                f'{args.train_file_dir}/**/*.jsonl', recursive=True)
+            logger.info(f"train files: {', '.join(train_data_files)}")
+            data_files["train"] = train_data_files
+        if args.validation_file_dir is not None and os.path.exists(args.validation_file_dir):
+            eval_data_files = glob(f'{args.validation_file_dir}/**/*.json', recursive=True) + glob(
+                f'{args.validation_file_dir}/**/*.jsonl', recursive=True)
+            logger.info(f"eval files: {', '.join(eval_data_files)}")
+            data_files["validation"] = eval_data_files
+        raw_datasets = load_dataset(
+            'json',
+            data_files=data_files,
+            cache_dir=args.cache_dir,
+        )
+        # If no validation data is there, validation_split_percentage will be used to divide the dataset.
+        if "validation" not in raw_datasets.keys():
+            raw_datasets["validation"] = load_dataset(
+                'json',
+                data_files=data_files,
+                split=f"train[:{args.validation_split_percentage}%]",
+                cache_dir=args.cache_dir,
+            )
+            raw_datasets["train"] = load_dataset(
+                'json',
+                data_files=data_files,
+                split=f"train[{args.validation_split_percentage}%:]",
+                cache_dir=args.cache_dir,
+            )
+    logger.info(f"Raw datasets: {raw_datasets}")
+
+    # Preprocessing the datasets
+    max_source_length = args.max_source_length
+    max_target_length = args.max_target_length
+    full_max_length = max_source_length + max_target_length
+
+    # Preprocess the dataset
+    train_dataset = None
+    max_train_samples = 0
+    if args.do_train:
+        if "train" not in raw_datasets:
+            raise ValueError("--do_train requires a train dataset")
+        train_dataset = raw_datasets['train']
+        max_train_samples = len(train_dataset)
+        if args.max_train_samples is not None and args.max_train_samples > 0:
+            max_train_samples = min(len(train_dataset), args.max_train_samples)
+            train_dataset = train_dataset.select(range(max_train_samples))
+        logger.debug(f"Example train_dataset[0]: {train_dataset[0]}")
+        tokenized_dataset = train_dataset.shuffle().map(
+            return_prompt_and_responses,
+            batched=True,
+            num_proc=args.preprocessing_num_workers,
+            remove_columns=train_dataset.column_names,
+            load_from_cache_file=not args.overwrite_cache,
+            desc="Running tokenizer on dataset",
+        )
+        train_dataset = tokenized_dataset.filter(
+            lambda x: 0 < len(x['prompt'] + x['chosen']) <= full_max_length
+                      and 0 < len(x['prompt'] + x['rejected']) <= full_max_length
+        )
+        logger.debug(f"Num train_samples: {len(train_dataset)}")
+        logger.debug("First train example:")
+        logger.debug(train_dataset[0]['prompt'] + train_dataset[0]['chosen'])
+
+    eval_dataset = None
+    max_eval_samples = 0
+    if args.do_eval:
+        if "validation" not in raw_datasets:
+            raise ValueError("--do_eval requires a validation dataset")
+        eval_dataset = raw_datasets["validation"]
+        max_eval_samples = len(eval_dataset)
+        if args.max_eval_samples is not None and args.max_eval_samples > 0:
+            max_eval_samples = min(len(eval_dataset), args.max_eval_samples)
+            eval_dataset = eval_dataset.select(range(max_eval_samples))
+        logger.debug(f"Example eval_dataset[0]: {eval_dataset[0]}")
+        eval_dataset = eval_dataset.map(
+            return_prompt_and_responses,
+            batched=True,
+            num_proc=args.preprocessing_num_workers,
+            remove_columns=eval_dataset.column_names,
+            load_from_cache_file=not args.overwrite_cache,
+            desc="Running tokenizer on dataset",
+        )
+        eval_dataset = eval_dataset.filter(
+            lambda x: 0 < len(x['prompt'] + x['chosen']) <= full_max_length
+                      and 0 < len(x['prompt'] + x['rejected']) <= full_max_length
+        )
+        logger.debug(f"Num eval_samples: {len(eval_dataset)}")
+        logger.debug("First eval example:")
+        logger.debug(eval_dataset[0]['prompt'] + eval_dataset[0]['chosen'])
+
+    logger.info("Loading model")
+    torch_dtype = (
+        args.torch_dtype
+        if args.torch_dtype in ["auto", None]
+        else getattr(torch, args.torch_dtype)
+    )
+    world_size = int(os.environ.get("WORLD_SIZE", 1))
+    ddp = world_size != 1
+    if ddp:
+        args.device_map = {"": int(os.environ["LOCAL_RANK"]) or 0}
+    if args.qlora and is_deepspeed_zero3_enabled():
+        logger.warning("ZeRO3 are both currently incompatible with QLoRA.")
+    config = config_class.from_pretrained(
+        args.model_name_or_path,
+        trust_remote_code=args.trust_remote_code,
+        torch_dtype=torch_dtype,
+        cache_dir=args.cache_dir
+    )
+    model = model_class.from_pretrained(
+        args.model_name_or_path,
+        config=config,
+        low_cpu_mem_usage=(not is_deepspeed_zero3_enabled()),
+        device_map=args.device_map,
+        trust_remote_code=args.trust_remote_code,
+        quantization_config=BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_compute_dtype=torch_dtype,
+        ) if args.qlora else None,
+    )
+    model_ref = model_class.from_pretrained(
+        args.model_name_or_path,
+        config=config,
+        low_cpu_mem_usage=(not is_deepspeed_zero3_enabled()),
+        device_map=args.device_map,
+        trust_remote_code=args.trust_remote_code,
+        quantization_config=BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_compute_dtype=torch_dtype,
+        ) if args.qlora else None,
+    )
+
+    # Initialize our Trainer
+    if args.gradient_checkpointing:
+        model.gradient_checkpointing_enable()
+        model.config.use_cache = False
+    else:
+        model.config.use_cache = True
+
+    training_args = TrainingArguments(
+        per_device_train_batch_size=args.per_device_train_batch_size,
+        per_device_eval_batch_size=args.per_device_eval_batch_size,
+        max_steps=args.max_steps,
+        logging_steps=args.logging_steps,
+        save_steps=args.save_steps,
+        gradient_accumulation_steps=args.gradient_accumulation_steps,
+        gradient_checkpointing=args.gradient_checkpointing,
+        learning_rate=args.learning_rate,
+        evaluation_strategy=args.eval_strategy,
+        eval_steps=args.eval_steps,
+        output_dir=args.output_dir,
+        report_to=args.report_to,
+        lr_scheduler_type=args.lr_scheduler_type,
+        warmup_steps=args.warmup_steps,
+        optim=args.optim,
+        bf16=args.bf16,
+        fp16=args.fp16,
+        remove_unused_columns=args.remove_unused_columns,
+        run_name=f"dpo_{args.model_type}",
+    )
+
+    # Initialize DPO trainer
+    target_modules = args.target_modules.split(',') if args.target_modules else None
+    if target_modules and 'all' in target_modules:
+        target_modules = find_all_linear_names(model, int4=args.load_in_4bit, int8=args.load_in_8bit)
+    logger.info(f"Peft target_modules: {target_modules}")
+    peft_config = LoraConfig(
+        task_type=TaskType.CAUSAL_LM,
+        target_modules=target_modules,
+        inference_mode=False,
+        r=args.lora_rank,
+        lora_alpha=args.lora_alpha,
+        lora_dropout=args.lora_dropout,
+    )
+    trainer = DPOTrainer(
+        model,
+        model_ref,
+        args=training_args,
+        beta=args.beta,
+        train_dataset=train_dataset,
+        eval_dataset=eval_dataset,
+        tokenizer=tokenizer,
+        peft_config=peft_config if args.use_peft else None,
+        max_prompt_length=args.max_source_length,
+        max_length=full_max_length,
+    )
+    print_trainable_parameters(trainer.model)
+
+    # Training
+    if args.do_train:
+        logger.info("*** Train ***")
+        train_result = trainer.train()
+        metrics = train_result.metrics
+        metrics["train_samples"] = max_train_samples
+        logger.debug(f"Training metrics: {metrics}")
+        trainer.log_metrics("train", metrics)
+        trainer.save_metrics("train", metrics)
+        trainer.save_state()
+        logger.info(f"Saving model checkpoint to {args.output_dir}")
+        trainer.save_model(args.output_dir)
+        tokenizer.save_pretrained(args.output_dir)
+        trainer.model.save_pretrained(args.output_dir)
+
+    # Evaluation
+    if args.do_eval and trainer.is_world_process_zero():
+        logger.info("*** Evaluate ***")
+        metrics = trainer.evaluate()
+        metrics["eval_samples"] = max_eval_samples
+        logger.debug(f"Eval metrics: {metrics}")
+        trainer.log_metrics("eval", metrics)
+        trainer.save_metrics("eval", metrics)
+
+
+if __name__ == "__main__":
+    main()

gradio_demo.py

@@ -0,0 +1,215 @@
+# -*- coding: utf-8 -*-
+"""
+@author:XuMing([email protected])
+@description:
+
+pip install gradio
+pip install mdtex2html
+"""
+import argparse
+import os
+from threading import Thread
+
+import gradio as gr
+import mdtex2html
+import torch
+from peft import PeftModel
+from transformers import (
+    AutoModel,
+    AutoTokenizer,
+    AutoModelForCausalLM,
+    BloomForCausalLM,
+    BloomTokenizerFast,
+    LlamaTokenizer,
+    LlamaForCausalLM,
+    GenerationConfig,
+    TextIteratorStreamer,
+)
+
+from supervised_finetuning import get_conv_template
+
+MODEL_CLASSES = {
+    "bloom": (BloomForCausalLM, BloomTokenizerFast),
+    "chatglm": (AutoModel, AutoTokenizer),
+    "llama": (LlamaForCausalLM, LlamaTokenizer),
+    "baichuan": (AutoModelForCausalLM, AutoTokenizer),
+    "auto": (AutoModelForCausalLM, AutoTokenizer),
+}
+
+
+@torch.inference_mode()
+def stream_generate_answer(
+        model,
+        tokenizer,
+        prompt,
+        device,
+        max_new_tokens=512,
+        temperature=0.7,
+        top_p=0.8,
+        repetition_penalty=1.0,
+        context_len=2048,
+):
+    streamer = TextIteratorStreamer(tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=False)
+    input_ids = tokenizer(prompt).input_ids
+    max_src_len = context_len - max_new_tokens - 8
+    input_ids = input_ids[-max_src_len:]
+    generation_kwargs = dict(
+        input_ids=torch.as_tensor([input_ids]).to(device),
+        max_new_tokens=max_new_tokens,
+        temperature=temperature,
+        top_p=top_p,
+        repetition_penalty=repetition_penalty,
+        streamer=streamer,
+    )
+
+    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
+
+    yield from streamer
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--model_type', default=None, type=str, required=True)
+    parser.add_argument('--base_model', default=None, type=str, required=True)
+    parser.add_argument('--lora_model', default="", type=str, help="If None, perform inference on the base model")
+    parser.add_argument('--tokenizer_path', default=None, type=str)
+    parser.add_argument('--template_name', default="vicuna", type=str,
+                        help="Prompt template name, eg: alpaca, vicuna, baichuan-chat, chatglm2 etc.")
+    parser.add_argument('--gpus', default="0", type=str)
+    parser.add_argument('--only_cpu', action='store_true', help='only use CPU for inference')
+    parser.add_argument('--resize_emb', action='store_true', help='Whether to resize model token embeddings')
+    args = parser.parse_args()
+    if args.only_cpu is True:
+        args.gpus = ""
+    os.environ["CUDA_VISIBLE_DEVICES"] = args.gpus
+
+    def postprocess(self, y):
+        if y is None:
+            return []
+        for i, (message, response) in enumerate(y):
+            y[i] = (
+                None if message is None else mdtex2html.convert((message)),
+                None if response is None else mdtex2html.convert(response),
+            )
+        return y
+
+    gr.Chatbot.postprocess = postprocess
+
+    load_type = torch.float16
+    if torch.cuda.is_available():
+        device = torch.device(0)
+    else:
+        device = torch.device('cpu')
+
+    if args.tokenizer_path is None:
+        args.tokenizer_path = args.base_model
+    model_class, tokenizer_class = MODEL_CLASSES[args.model_type]
+    tokenizer = tokenizer_class.from_pretrained(args.tokenizer_path, trust_remote_code=True)
+    base_model = model_class.from_pretrained(
+        args.base_model,
+        load_in_8bit=False,
+        torch_dtype=load_type,
+        low_cpu_mem_usage=True,
+        device_map='auto',
+        trust_remote_code=True,
+    )
+    try:
+        base_model.generation_config = GenerationConfig.from_pretrained(args.base_model, trust_remote_code=True)
+    except OSError:
+        print("Failed to load generation config, use default.")
+    if args.resize_emb:
+        model_vocab_size = base_model.get_input_embeddings().weight.size(0)
+        tokenzier_vocab_size = len(tokenizer)
+        print(f"Vocab of the base model: {model_vocab_size}")
+        print(f"Vocab of the tokenizer: {tokenzier_vocab_size}")
+        if model_vocab_size != tokenzier_vocab_size:
+            print("Resize model embeddings to fit tokenizer")
+            base_model.resize_token_embeddings(tokenzier_vocab_size)
+    if args.lora_model:
+        model = PeftModel.from_pretrained(base_model, args.lora_model, torch_dtype=load_type, device_map='auto')
+        print("loaded lora model")
+    else:
+        model = base_model
+    if device == torch.device('cpu'):
+        model.float()
+
+    model.eval()
+
+    def reset_user_input():
+        return gr.update(value='')
+
+    def reset_state():
+        return [], []
+
+    prompt_template = get_conv_template(args.template_name)
+    stop_str = tokenizer.eos_token if tokenizer.eos_token else prompt_template.stop_str
+    history = []
+
+    def predict(
+            input,
+            chatbot,
+            history,
+            max_new_tokens,
+            temperature,
+            top_p
+    ):
+        now_input = input
+        chatbot.append((input, ""))
+        history = history or []
+        history.append([now_input, ''])
+
+        prompt = prompt_template.get_prompt(messages=history)
+        response = ""
+
+        for new_text in stream_generate_answer(
+                model,
+                tokenizer,
+                prompt,
+                device,
+                max_new_tokens=max_new_tokens,
+                temperature=temperature,
+                top_p=top_p,
+        ):
+            stop = False
+            pos = new_text.find(stop_str)
+            if pos != -1:
+                new_text = new_text[:pos]
+                stop = True
+            response += new_text
+            new_history = history + [(now_input, response)]
+            chatbot[-1] = (now_input, response)
+            yield chatbot, new_history
+            if stop:
+                break
+
+    with gr.Blocks() as demo:
+        gr.HTML("""<h1 align="center">MedicalGPT</h1>""")
+        gr.Markdown(
+            "> 为了促进医疗行业大模型的开放研究，本项目开源了MedicalGPT医疗大模型")
+        chatbot = gr.Chatbot()
+        with gr.Row():
+            with gr.Column(scale=4):
+                with gr.Column(scale=12):
+                    user_input = gr.Textbox(show_label=False, placeholder="Input...", lines=10).style(
+                        container=False)
+                with gr.Column(min_width=32, scale=1):
+                    submitBtn = gr.Button("Submit", variant="primary")
+            with gr.Column(scale=1):
+                emptyBtn = gr.Button("Clear History")
+                max_length = gr.Slider(
+                    0, 4096, value=512, step=1.0, label="Maximum length", interactive=True)
+                top_p = gr.Slider(0, 1, value=0.8, step=0.01,
+                                  label="Top P", interactive=True)
+                temperature = gr.Slider(
+                    0, 1, value=0.7, step=0.01, label="Temperature", interactive=True)
+        history = gr.State([])
+        submitBtn.click(predict, [user_input, chatbot, history, max_length, temperature, top_p], [chatbot, history],
+                        show_progress=True)
+        submitBtn.click(reset_user_input, [], [user_input])
+        emptyBtn.click(reset_state, outputs=[chatbot, history], show_progress=True)
+    demo.queue().launch(share=False, inbrowser=True, server_name='0.0.0.0', server_port=8082)
+
+
+if __name__ == '__main__':
+    main()

inference.py

@@ -0,0 +1,225 @@
+# -*- coding: utf-8 -*-
+"""
+@author:XuMing([email protected])
+@description: 
+"""
+import argparse
+import json
+import os
+from threading import Thread
+
+import torch
+from peft import PeftModel
+from transformers import (
+    AutoModel,
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    BloomForCausalLM,
+    BloomTokenizerFast,
+    LlamaTokenizer,
+    LlamaForCausalLM,
+    TextIteratorStreamer,
+    GenerationConfig,
+)
+
+from supervised_finetuning import get_conv_template
+
+MODEL_CLASSES = {
+    "bloom": (BloomForCausalLM, BloomTokenizerFast),
+    "chatglm": (AutoModel, AutoTokenizer),
+    "llama": (LlamaForCausalLM, LlamaTokenizer),
+    "baichuan": (AutoModelForCausalLM, AutoTokenizer),
+    "auto": (AutoModelForCausalLM, AutoTokenizer),
+}
+
+
+@torch.inference_mode()
+def stream_generate_answer(
+        model,
+        tokenizer,
+        prompt,
+        device,
+        do_print=True,
+        max_new_tokens=512,
+        temperature=0.7,
+        repetition_penalty=1.0,
+        context_len=2048,
+        stop_str="</s>",
+):
+    streamer = TextIteratorStreamer(tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=False)
+    input_ids = tokenizer(prompt).input_ids
+    max_src_len = context_len - max_new_tokens - 8
+    input_ids = input_ids[-max_src_len:]
+    generation_kwargs = dict(
+        input_ids=torch.as_tensor([input_ids]).to(device),
+        max_new_tokens=max_new_tokens,
+        temperature=temperature,
+        repetition_penalty=repetition_penalty,
+        streamer=streamer,
+    )
+    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
+
+    generated_text = ""
+    for new_text in streamer:
+        stop = False
+        pos = new_text.find(stop_str)
+        if pos != -1:
+            new_text = new_text[:pos]
+            stop = True
+        generated_text += new_text
+        if do_print:
+            print(new_text, end="", flush=True)
+        if stop:
+            break
+    if do_print:
+        print()
+    return generated_text
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--model_type', default=None, type=str, required=True)
+    parser.add_argument('--base_model', default=None, type=str, required=True)
+    parser.add_argument('--lora_model', default="", type=str, help="If None, perform inference on the base model")
+    parser.add_argument('--tokenizer_path', default=None, type=str)
+    parser.add_argument('--template_name', default="vicuna", type=str,
+                        help="Prompt template name, eg: alpaca, vicuna, baichuan-chat, chatglm2 etc.")
+    parser.add_argument("--temperature", type=float, default=0.7)
+    parser.add_argument("--repetition_penalty", type=float, default=1.0)
+    parser.add_argument("--max_new_tokens", type=int, default=512)
+    parser.add_argument('--data_file', default=None, type=str,
+                        help="A file that contains instructions (one instruction per line)")
+    parser.add_argument('--interactive', action='store_true', help="run in the instruction mode (single-turn)")
+    parser.add_argument('--predictions_file', default='./predictions.json', type=str)
+    parser.add_argument('--resize_emb', action='store_true', help='Whether to resize model token embeddings')
+    parser.add_argument('--gpus', default="0", type=str)
+    parser.add_argument('--only_cpu', action='store_true', help='only use CPU for inference')
+    args = parser.parse_args()
+    print(args)
+    if args.only_cpu is True:
+        args.gpus = ""
+    os.environ["CUDA_VISIBLE_DEVICES"] = args.gpus
+    load_type = torch.float16
+    if torch.cuda.is_available():
+        device = torch.device(0)
+    else:
+        device = torch.device('cpu')
+    if args.tokenizer_path is None:
+        args.tokenizer_path = args.base_model
+
+    model_class, tokenizer_class = MODEL_CLASSES[args.model_type]
+    tokenizer = tokenizer_class.from_pretrained(args.tokenizer_path, trust_remote_code=True)
+    base_model = model_class.from_pretrained(
+        args.base_model,
+        load_in_8bit=False,
+        torch_dtype=load_type,
+        low_cpu_mem_usage=True,
+        device_map='auto',
+        trust_remote_code=True,
+    )
+    try:
+        base_model.generation_config = GenerationConfig.from_pretrained(args.base_model, trust_remote_code=True)
+    except OSError:
+        print("Failed to load generation config, use default.")
+    if args.resize_emb:
+        model_vocab_size = base_model.get_input_embeddings().weight.size(0)
+        tokenzier_vocab_size = len(tokenizer)
+        print(f"Vocab of the base model: {model_vocab_size}")
+        print(f"Vocab of the tokenizer: {tokenzier_vocab_size}")
+        if model_vocab_size != tokenzier_vocab_size:
+            print("Resize model embeddings to fit tokenizer")
+            base_model.resize_token_embeddings(tokenzier_vocab_size)
+
+    if args.lora_model:
+        model = PeftModel.from_pretrained(base_model, args.lora_model, torch_dtype=load_type, device_map='auto')
+        print("Loaded lora model")
+    else:
+        model = base_model
+    if device == torch.device('cpu'):
+        model.float()
+    model.eval()
+    print(tokenizer)
+    # test data
+    if args.data_file is None:
+        examples = ["介绍下北京", "乙肝和丙肝的区别？"]
+    else:
+        with open(args.data_file, 'r') as f:
+            examples = [l.strip() for l in f.readlines()]
+        print("first 10 examples:")
+        for example in examples[:10]:
+            print(example)
+
+    # Chat
+    prompt_template = get_conv_template(args.template_name)
+    stop_str = tokenizer.eos_token if tokenizer.eos_token else prompt_template.stop_str
+
+    if args.interactive:
+        print("Welcome to the CLI application, use `clear` to remove the history, use `exit` to exit the application.")
+        history = []
+        while True:
+            try:
+                query = input(f"{prompt_template.roles[0]}: ")
+            except UnicodeDecodeError:
+                print("Detected decoding error at the inputs, please try again.")
+                continue
+            except Exception:
+                raise
+            if query == "":
+                print("Please input text, try again.")
+                continue
+            if query.strip() == "exit":
+                print("exit...")
+                break
+            if query.strip() == "clear":
+                history = []
+                print("history cleared.")
+                continue
+
+            print(f"{prompt_template.roles[1]}: ", end="", flush=True)
+
+            history.append([query, ''])
+            prompt = prompt_template.get_prompt(messages=history)
+            response = stream_generate_answer(
+                model,
+                tokenizer,
+                prompt,
+                device,
+                do_print=True,
+                max_new_tokens=args.max_new_tokens,
+                temperature=args.temperature,
+                repetition_penalty=args.repetition_penalty,
+                stop_str=stop_str,
+            )
+            if history:
+                history[-1][-1] = response.strip()
+    else:
+        print("Start inference.")
+        results = []
+        for index, example in enumerate(examples):
+            # Single turn inference
+            history = [[example, '']]
+            prompt = prompt_template.get_prompt(messages=history)
+            response = stream_generate_answer(
+                model,
+                tokenizer,
+                prompt,
+                device,
+                do_print=False,
+                max_new_tokens=args.max_new_tokens,
+                temperature=args.temperature,
+                repetition_penalty=args.repetition_penalty,
+                stop_str=stop_str,
+            )
+            response = response.strip()
+            print(f"======={index}=======")
+            print(f"Input: {example}\n")
+            print(f"Output: {response}\n")
+            results.append({"Input": prompt, "Output": response})
+
+        with open(args.predictions_file, 'w', encoding='utf-8') as f:
+            json.dump(results, f, ensure_ascii=False, indent=2)
+
+
+if __name__ == '__main__':
+    main()

merge_peft_adapter.py

@@ -0,0 +1,109 @@
+# -*- coding: utf-8 -*-
+"""
+@author:XuMing([email protected])
+@description:
+
+Usage:
+python merge_peft_adapter.py \
+    --base_model_name_or_path path/to/llama/model \
+    --tokenizer_path path/to/llama/tokenizer \
+    --peft_model_path path/to/lora/model \
+    --output_dir path/to/output/dir
+
+after merged, chatglm and baichuan model need copy python script to output dir.
+"""
+
+import argparse
+
+import torch
+from peft import PeftModel, PeftConfig
+from transformers import (
+    AutoModel,
+    AutoTokenizer,
+    BloomForCausalLM,
+    BloomTokenizerFast,
+    AutoModelForCausalLM,
+    LlamaTokenizer,
+    LlamaForCausalLM,
+    AutoModelForSequenceClassification,
+)
+
+MODEL_CLASSES = {
+    "bloom": (BloomForCausalLM, BloomTokenizerFast),
+    "chatglm": (AutoModel, AutoTokenizer),
+    "llama": (LlamaForCausalLM, LlamaTokenizer),
+    "baichuan": (AutoModelForCausalLM, AutoTokenizer),
+    "auto": (AutoModelForCausalLM, AutoTokenizer),
+}
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--model_type', default=None, type=str, required=True)
+    parser.add_argument('--base_model_name_or_path', default=None, required=True, type=str,
+                        help="Base model name or path")
+    parser.add_argument('--tokenizer_path', default=None, type=str,
+                        help="Please specify tokenization path.")
+    parser.add_argument('--peft_model_path', default=None, required=True, type=str,
+                        help="Please specify LoRA model to be merged.")
+    parser.add_argument('--resize_emb', action='store_true', help='Whether to resize model token embeddings')
+    parser.add_argument('--output_dir', default='./merged', type=str)
+    args = parser.parse_args()
+    print(args)
+
+    base_model_path = args.base_model_name_or_path
+    peft_model_path = args.peft_model_path
+    output_dir = args.output_dir
+    print(f"Base model: {base_model_path}")
+    print(f"LoRA model: {peft_model_path}")
+    peft_config = PeftConfig.from_pretrained(peft_model_path)
+
+    model_class, tokenizer_class = MODEL_CLASSES[args.model_type]
+    if peft_config.task_type == "SEQ_CLS":
+        print("Loading LoRA for sequence classification model")
+        if args.model_type == "chatglm":
+            raise ValueError("chatglm does not support sequence classification")
+        base_model = AutoModelForSequenceClassification.from_pretrained(
+            base_model_path,
+            load_in_8bit=False,
+            torch_dtype=torch.float16,
+            trust_remote_code=True,
+            device_map="auto",
+        )
+    else:
+        print("Loading LoRA for causal language model")
+        base_model = model_class.from_pretrained(
+            base_model_path,
+            load_in_8bit=False,
+            torch_dtype=torch.float16,
+            trust_remote_code=True,
+            device_map="auto",
+        )
+    if args.tokenizer_path:
+        tokenizer = tokenizer_class.from_pretrained(args.tokenizer_path, trust_remote_code=True)
+    else:
+        tokenizer = tokenizer_class.from_pretrained(peft_model_path, trust_remote_code=True)
+    if args.resize_emb:
+        base_model_token_size = base_model.get_input_embeddings().weight.size(0)
+        if base_model_token_size != len(tokenizer):
+            base_model.resize_token_embeddings(len(tokenizer))
+            print(f"Resize vocabulary size {base_model_token_size} to {len(tokenizer)}")
+
+    lora_model = PeftModel.from_pretrained(
+        base_model,
+        peft_model_path,
+        device_map="auto",
+        torch_dtype=torch.float16,
+    )
+    lora_model.eval()
+    print(f"Merging with merge_and_unload...")
+    base_model = lora_model.merge_and_unload()
+
+    print("Saving to Hugging Face format...")
+    tokenizer.save_pretrained(output_dir)
+    base_model.save_pretrained(output_dir)
+    print(f"Done! model saved to {output_dir}")
+
+
+if __name__ == '__main__':
+    main()

merge_tokenizers.py

@@ -0,0 +1,150 @@
+# -*- coding: utf-8 -*-
+"""
+@author:XuMing([email protected])
+@description: 
+"""
+import os
+
+os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python"
+from transformers import LlamaTokenizer
+from sentencepiece import sentencepiece_model_pb2 as sp_pb2_model
+import sentencepiece as spm
+import argparse
+
+
+def is_chinese(uchar):
+    """判断一个unicode是否是汉字"""
+    return '\u4e00' <= uchar <= '\u9fa5'
+
+
+def is_chinese_string(string):
+    """判断是否全为汉字"""
+    return all(is_chinese(c) for c in string)
+
+
+def load_baichuan_vocab(vocab_file):
+    words = set()
+    with open(vocab_file, "r", encoding="utf-8") as f:
+        for line in f:
+            if line.strip():
+                words.add(line.strip().split()[0])
+    return words
+
+
+def load_jieba_vocab(jieba_vocab_file):
+    # Read jieba vocab and sort by freq
+    with open(jieba_vocab_file, "r", encoding="utf-8") as f:
+        lines = f.readlines()
+        word_freqs = [line.strip().split() for line in lines]
+        word_freqs.sort(key=lambda x: int(x[1]), reverse=True)
+    return word_freqs
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--base_tokenizer_dir', default=None, type=str, required=True)
+    parser.add_argument('--domain_sp_model_file', default='./domain_sp.model', type=str)
+    parser.add_argument('--baichuan_vocab_file', default="data/vocab/baichuan_vocab.txt", type=str)
+    parser.add_argument('--add_jieba', action='store_true', help='Whether to add jieba vocab.')
+    parser.add_argument('--jieba_word_freq_file', default='data/vocab/word_freq.txt', type=str)
+    parser.add_argument('--jieba_word_size', default=20000, type=int)
+
+    args = parser.parse_args()
+    print(args)
+
+    # load
+    llama_tokenizer = LlamaTokenizer.from_pretrained(args.base_tokenizer_dir)
+    chinese_sp_model = spm.SentencePieceProcessor()
+    chinese_sp_model.Load(args.domain_sp_model_file)
+
+    llama_spm = sp_pb2_model.ModelProto()
+    llama_spm.ParseFromString(llama_tokenizer.sp_model.serialized_model_proto())
+    chinese_spm = sp_pb2_model.ModelProto()
+    chinese_spm.ParseFromString(chinese_sp_model.serialized_model_proto())
+
+    # print number of tokens
+    print(len(llama_tokenizer), len(chinese_sp_model))
+    print(llama_tokenizer.all_special_tokens)
+    print(llama_tokenizer.all_special_ids)
+    print(llama_tokenizer.special_tokens_map)
+
+    # Add Chinese tokens to LLaMA tokenizer
+    llama_spm_tokens_set = set(p.piece for p in llama_spm.pieces)
+
+    print(len(llama_spm_tokens_set))
+    print(f"Before:{len(llama_spm_tokens_set)}")
+    added_set = set()
+    for p in chinese_spm.pieces:
+        piece = p.piece
+        if piece not in llama_spm_tokens_set:
+            # print('picec', piece)
+            new_p = sp_pb2_model.ModelProto().SentencePiece()
+            new_p.piece = piece
+            new_p.score = 0
+            llama_spm.pieces.append(new_p)
+            added_set.add(piece)
+    print(f"[add domain tokens]New model pieces: {len(llama_spm.pieces)}")
+
+    vocab = load_baichuan_vocab(args.baichuan_vocab_file)
+    print('baichuan vocab len:', len(vocab))
+    baichuan_vocab_set = set([i for i in vocab if is_chinese_string(i)])
+    print('baichuan chinese vocab size:', len(baichuan_vocab_set))
+    print('baichuan vocab head:', list(baichuan_vocab_set)[:10])
+    for p in baichuan_vocab_set:
+        piece = p
+        if piece not in llama_spm_tokens_set and piece not in added_set:
+            # print('baichuan picec', piece)
+            new_p = sp_pb2_model.ModelProto().SentencePiece()
+            new_p.piece = piece
+            new_p.score = 0
+            llama_spm.pieces.append(new_p)
+            added_set.add(piece)
+    print(f"[add baichuan tokens]New model pieces: {len(llama_spm.pieces)}")
+
+    if args.add_jieba:
+        word_freqs = load_jieba_vocab(args.jieba_word_freq_file)
+        top_words = word_freqs[:args.jieba_word_size]
+        print('jieba top10 freq words:', top_words[:10])
+        jieba_vocab_set = set([i[0] for i in top_words if i])
+        print('jieba_vocab_set size:', len(jieba_vocab_set))
+        print('jieba_vocab head:', list(jieba_vocab_set)[:3])
+        for p in jieba_vocab_set:
+            piece = p
+            if piece not in llama_spm_tokens_set and piece not in added_set:
+                # print('jieba picec', piece)
+                new_p = sp_pb2_model.ModelProto().SentencePiece()
+                new_p.piece = piece
+                new_p.score = 0
+                llama_spm.pieces.append(new_p)
+        print(f"[add jieba tokens]New model pieces: {len(llama_spm.pieces)}")
+
+    # Save
+    output_sp_dir = 'merged_tokenizer_sp'
+    output_hf_dir = 'merged_tokenizer_hf'  # the path to save Chinese-LLaMA tokenizer
+    os.makedirs(output_sp_dir, exist_ok=True)
+    with open(output_sp_dir + '/chinese_llama.model', 'wb') as f:
+        f.write(llama_spm.SerializeToString())
+    tokenizer = LlamaTokenizer(vocab_file=output_sp_dir + '/chinese_llama.model')
+
+    tokenizer.save_pretrained(output_hf_dir)
+    print(f"Chinese-LLaMA tokenizer has been saved to {output_hf_dir}")
+
+    # Test
+    llama_tokenizer = LlamaTokenizer.from_pretrained(args.base_tokenizer_dir)
+    chinese_llama_tokenizer = LlamaTokenizer.from_pretrained(output_hf_dir)
+    print(chinese_llama_tokenizer.all_special_tokens)
+    print(chinese_llama_tokenizer.all_special_ids)
+    print(chinese_llama_tokenizer.special_tokens_map)
+    print('old len:', len(llama_tokenizer), ' new len:', len(chinese_llama_tokenizer))
+    text = '''this is a test, hello world. thisisatesthelloworld, 
+慕容复来到河边，姑苏慕容氏在外面丢了人。
+1号店一周岁了，我们一古脑儿买了10斤零食。
+巴塞罗那足球俱乐部简称巴萨（Barça），是一家位于西班牙加泰罗尼亚巴塞罗那的足球俱乐部，于1899年由瑞士企业家胡安·甘伯所创立，世界球坛顶级足球俱乐部之一。俱乐部主场可容纳接近十万名观众，是全欧洲最大及世界第二大的足球场。
+白日依山尽，黄河入海流。欲穷千里目，更上一层楼。'''
+    print("Test text:\n", text)
+    print(f"Tokenized by LLaMA tokenizer:{llama_tokenizer.tokenize(text)}")
+    print(f"Tokenized by Chinese-LLaMA tokenizer:{chinese_llama_tokenizer.tokenize(text)}")
+
+
+if __name__ == '__main__':
+    main()

pretraining.py

@@ -0,0 +1,678 @@
+# -*- coding: utf-8 -*-
+# Copyright 2023 XuMing([email protected]) and The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Fine-tuning the library models for causal language modeling (GPT, GPT-2, CTRL, ...) on a text file or a dataset.
+
+part of this code is adapted from https://github.com/huggingface/transformers/blob/main/examples/pytorch/language-modeling/run_clm.py
+"""
+import math
+import os
+from dataclasses import dataclass, field
+from glob import glob
+from itertools import chain
+from typing import Optional, List, Dict, Any, Mapping
+
+import numpy as np
+import torch
+from datasets import load_dataset
+from loguru import logger
+from peft import LoraConfig, TaskType, get_peft_model, PeftModel, prepare_model_for_int8_training
+from sklearn.metrics import accuracy_score
+from transformers import (
+    AutoConfig,
+    BloomForCausalLM,
+    AutoModelForCausalLM,
+    AutoModel,
+    LlamaTokenizer,
+    LlamaForCausalLM,
+    BloomTokenizerFast,
+    AutoTokenizer,
+    HfArgumentParser,
+    Trainer,
+    TrainingArguments,
+    is_torch_tpu_available,
+    set_seed,
+)
+from transformers.trainer import TRAINING_ARGS_NAME
+from transformers.utils.versions import require_version
+
+MODEL_CLASSES = {
+    "bloom": (AutoConfig, BloomForCausalLM, BloomTokenizerFast),
+    "chatglm": (AutoConfig, AutoModel, AutoTokenizer),
+    "llama": (AutoConfig, LlamaForCausalLM, LlamaTokenizer),
+    "baichuan": (AutoConfig, AutoModelForCausalLM, AutoTokenizer),
+    "auto": (AutoConfig, AutoModelForCausalLM, AutoTokenizer),
+}
+
+
+@dataclass
+class ModelArguments:
+    """
+    Arguments pertaining to which model/config/tokenizer we are going to fine-tune, or train from scratch.
+    """
+
+    model_type: str = field(
+        default=None,
+        metadata={"help": "Model type selected in the list: " + ", ".join(MODEL_CLASSES.keys())}
+    )
+    model_name_or_path: Optional[str] = field(
+        default=None,
+        metadata={
+            "help": (
+                "The model checkpoint for weights initialization.Don't set if you want to train a model from scratch."
+            )
+        },
+    )
+    tokenizer_name_or_path: Optional[str] = field(
+        default=None,
+        metadata={
+            "help": (
+                "The tokenizer for weights initialization.Don't set if you want to train a model from scratch."
+            )
+        },
+    )
+    load_in_8bit: bool = field(default=False, metadata={"help": "Whether to load the model in 8bit mode or not."})
+    cache_dir: Optional[str] = field(
+        default=None,
+        metadata={"help": "Where do you want to store the pretrained models downloaded from huggingface.co"},
+    )
+    use_fast_tokenizer: bool = field(
+        default=False,
+        metadata={"help": "Whether to use one of the fast tokenizer (backed by the tokenizers library) or not."},
+    )
+    torch_dtype: Optional[str] = field(
+        default=None,
+        metadata={
+            "help": (
+                "Override the default `torch.dtype` and load the model under this dtype. If `auto` is passed, the "
+                "dtype will be automatically derived from the model's weights."
+            ),
+            "choices": ["auto", "bfloat16", "float16", "float32"],
+        },
+    )
+    device_map: Optional[str] = field(
+        default="auto",
+        metadata={"help": "Device to map model to. If `auto` is passed, the device will be selected automatically. "},
+    )
+    trust_remote_code: bool = field(
+        default=True,
+        metadata={"help": "Whether to trust remote code when loading a model from a remote checkpoint."},
+    )
+
+    def __post_init__(self):
+        if self.model_type is None:
+            raise ValueError(
+                "You must specify a valid model_type to run training. Available model types are " + ", ".join(
+                    MODEL_CLASSES.keys()))
+        if self.model_name_or_path is None:
+            raise ValueError("You must specify a valid model_name_or_path to run training.")
+
+
+@dataclass
+class DataTrainingArguments:
+    """
+    Arguments pertaining to what data we are going to input our model for training and eval.
+    """
+
+    dataset_name: Optional[str] = field(
+        default=None, metadata={"help": "The name of the dataset to use (via the datasets library)."}
+    )
+    dataset_config_name: Optional[str] = field(
+        default=None, metadata={"help": "The configuration name of the dataset to use (via the datasets library)."}
+    )
+    train_file_dir: Optional[str] = field(default=None, metadata={"help": "The train text data file folder."})
+    validation_file_dir: Optional[str] = field(
+        default=None,
+        metadata={"help": "An optional input evaluation data file to evaluate the perplexity on text file folder."},
+    )
+    max_train_samples: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": (
+                "For debugging purposes or quicker training, truncate the number of training examples to this "
+                "value if set."
+            )
+        },
+    )
+    max_eval_samples: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": (
+                "For debugging purposes or quicker training, truncate the number of evaluation examples to this "
+                "value if set."
+            )
+        },
+    )
+    streaming: bool = field(default=False, metadata={"help": "Enable streaming mode"})
+    block_size: Optional[int] = field(
+        default=1024,
+        metadata={
+            "help": (
+                "Optional input sequence length after tokenization. "
+                "The training dataset will be truncated in block of this size for training. "
+                "Default to the model max input length for single sentence inputs (take into account special tokens)."
+            )
+        },
+    )
+    overwrite_cache: bool = field(
+        default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}
+    )
+    validation_split_percentage: Optional[int] = field(
+        default=1,
+        metadata={
+            "help": "The percentage of the train set used as validation set in case there's no validation split"
+        },
+    )
+    preprocessing_num_workers: Optional[int] = field(
+        default=None,
+        metadata={"help": "The number of processes to use for the preprocessing."},
+    )
+    keep_linebreaks: bool = field(
+        default=True, metadata={"help": "Whether to keep line breaks when using TXT files or not."}
+    )
+
+    def __post_init__(self):
+        if self.streaming:
+            require_version("datasets>=2.0.0", "The streaming feature requires `datasets>=2.0.0`")
+
+
+@dataclass
+class PeftArguments(TrainingArguments):
+    use_peft: bool = field(default=True, metadata={"help": "Whether to use peft"})
+    target_modules: Optional[str] = field(default="all")
+    lora_rank: Optional[int] = field(default=8)
+    lora_dropout: Optional[float] = field(default=0.05)
+    lora_alpha: Optional[float] = field(default=32.0)
+    modules_to_save: Optional[str] = field(default=None)
+    peft_path: Optional[str] = field(default=None)
+
+
+def accuracy(predictions, references, normalize=True, sample_weight=None):
+    return {
+        "accuracy": float(accuracy_score(references, predictions, normalize=normalize, sample_weight=sample_weight))
+    }
+
+
+def compute_metrics(eval_preds):
+    preds, labels = eval_preds
+    # preds have the same shape as the labels, after the argmax(-1) has been calculated
+    # by preprocess_logits_for_metrics, we need to shift the labels
+    labels = labels[:, 1:].reshape(-1)
+    preds = preds[:, :-1].reshape(-1)
+    return accuracy(predictions=preds, references=labels)
+
+
+def preprocess_logits_for_metrics(logits, labels):
+    if isinstance(logits, tuple):
+        # Depending on the model and config, logits may contain extra tensors,
+        # like past_key_values, but logits always come first
+        logits = logits[0]
+    return logits.argmax(dim=-1)
+
+
+def fault_tolerance_data_collator(features: List) -> Dict[str, Any]:
+    if not isinstance(features[0], Mapping):
+        features = [vars(f) for f in features]
+    first = features[0]
+    batch = {}
+
+    # Special handling for labels.
+    # Ensure that tensor is created with the correct type
+    if "label" in first and first["label"] is not None:
+        label = first["label"].item() if isinstance(first["label"], torch.Tensor) else first["label"]
+        dtype = torch.long if isinstance(label, int) else torch.float
+        batch["labels"] = torch.tensor([f["label"] for f in features], dtype=dtype)
+    elif "label_ids" in first and first["label_ids"] is not None:
+        if isinstance(first["label_ids"], torch.Tensor):
+            batch["labels"] = torch.stack([f["label_ids"] for f in features])
+        else:
+            dtype = torch.long if type(first["label_ids"][0]) is int else torch.float
+            batch["labels"] = torch.tensor([f["label_ids"] for f in features], dtype=dtype)
+
+    # Handling of all other possible keys.
+    # Again, we will use the first element to figure out which key/values are not None for this model.
+    try:
+        for k, v in first.items():
+            if k not in ("label", "label_ids") and v is not None and not isinstance(v, str):
+                if isinstance(v, torch.Tensor):
+                    batch[k] = torch.stack([f[k] for f in features])
+                elif isinstance(v, np.ndarray):
+                    batch[k] = torch.tensor(np.stack([f[k] for f in features]))
+                else:
+                    batch[k] = torch.tensor([f[k] for f in features])
+    except ValueError:  # quick fix by simply take the first example
+        for k, v in first.items():
+            if k not in ("label", "label_ids") and v is not None and not isinstance(v, str):
+                if isinstance(v, torch.Tensor):
+                    batch[k] = torch.stack([features[0][k]] * len(features))
+                elif isinstance(v, np.ndarray):
+                    batch[k] = torch.tensor(np.stack([features[0][k]] * len(features)))
+                else:
+                    batch[k] = torch.tensor([features[0][k]] * len(features))
+
+    return batch
+
+
+class GroupTextsBuilder:
+    def __init__(self, max_seq_length):
+        self.max_seq_length = max_seq_length
+
+    def __call__(self, examples):
+        # Concatenate all texts.
+        firsts = {k: examples[k][0][0] for k in examples.keys()}
+        lasts = {k: examples[k][0][-1] for k in examples.keys()}
+        contents = {k: sum([vi[1:-1] for vi in v], []) for k, v in examples.items()}
+        total_length = len(contents[list(examples.keys())[0]])
+
+        content_length = self.max_seq_length - 2
+        if total_length >= content_length:
+            total_length = (total_length // content_length) * content_length
+        # Split by chunks of max_len.
+        result = {
+            k: [[firsts[k]] + t[i: i + content_length] + [lasts[k]] for i in range(0, total_length, content_length)] for
+            k, t in contents.items()}
+        return result
+
+
+class SavePeftModelTrainer(Trainer):
+    """
+    Trainer for lora models
+    """
+
+    def save_model(self, output_dir=None, _internal_call=False):
+        """Save the LoRA model."""
+        os.makedirs(output_dir, exist_ok=True)
+        torch.save(self.args, os.path.join(output_dir, TRAINING_ARGS_NAME))
+        self.model.save_pretrained(output_dir)
+
+
+def save_model(output_dir, model, tokenizer, args):
+    """Save the model and the tokenizer."""
+    os.makedirs(output_dir, exist_ok=True)
+
+    # Take care of distributed/parallel training
+    model_to_save = model.module if hasattr(model, "module") else model
+    model_to_save.save_pretrained(output_dir)
+    tokenizer.save_pretrained(output_dir)
+    torch.save(args, os.path.join(output_dir, TRAINING_ARGS_NAME))
+
+
+def print_trainable_parameters(model):
+    """
+    Prints the number of trainable parameters in the model.
+    """
+    trainable_params = 0
+    all_param = 0
+    for _, param in model.named_parameters():
+        all_param += param.numel()
+        if param.requires_grad:
+            trainable_params += param.numel()
+    print(
+        f"trainable params: {trainable_params} || all params: {all_param} || trainable%: {100 * trainable_params / all_param}"
+    )
+
+
+def find_all_linear_names(peft_model, int4=False, int8=False):
+    """Find all linear layer names in the model. reference from qlora paper."""
+    cls = torch.nn.Linear
+    if int4 or int8:
+        import bitsandbytes as bnb
+        if int4:
+            cls = bnb.nn.Linear4bit
+        elif int8:
+            cls = bnb.nn.Linear8bitLt
+    lora_module_names = set()
+    for name, module in peft_model.named_modules():
+        if isinstance(module, cls):
+            # last layer is not add to lora_module_names
+            if 'lm_head' in name:
+                continue
+            names = name.split('.')
+            lora_module_names.add(names[0] if len(names) == 1 else names[-1])
+    return sorted(lora_module_names)
+
+
+def main():
+    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, PeftArguments))
+    model_args, data_args, training_args = parser.parse_args_into_dataclasses()
+
+    logger.info(f"Model args: {model_args}")
+    logger.info(f"Data args: {data_args}")
+    logger.info(f"Training args: {training_args}")
+    logger.info(
+        f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}"
+        + f" distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}"
+    )
+
+    # Set seed before initializing model.
+    set_seed(training_args.seed)
+
+    # Load tokenizer
+    if not model_args.model_type:
+        raise ValueError("Please specify a model_type, e.g. llama, chatglm, bloom, etc.")
+    config_class, model_class, tokenizer_class = MODEL_CLASSES[model_args.model_type]
+
+    tokenizer_kwargs = {
+        "cache_dir": model_args.cache_dir,
+        "use_fast": model_args.use_fast_tokenizer,
+        "trust_remote_code": model_args.trust_remote_code,
+    }
+    tokenizer_name_or_path = model_args.tokenizer_name_or_path
+    if not tokenizer_name_or_path:
+        tokenizer_name_or_path = model_args.model_name_or_path
+    tokenizer = tokenizer_class.from_pretrained(tokenizer_name_or_path, **tokenizer_kwargs)
+
+    # Preprocessing the datasets.
+    def tokenize_function(examples):
+        return tokenizer(examples["text"])
+
+    if data_args.block_size is None:
+        block_size = tokenizer.model_max_length
+        if block_size > 2048:
+            logger.warning(
+                "The chosen tokenizer supports a `model_max_length` that is longer than the default `block_size` value"
+                " of 2048. If you would like to use a longer `block_size` up to `tokenizer.model_max_length` you can"
+                " override this default with `--block_size xxx`."
+            )
+    else:
+        if data_args.block_size > tokenizer.model_max_length:
+            logger.warning(
+                f"The block_size passed ({data_args.block_size}) is larger than the maximum length for the model"
+                f"({tokenizer.model_max_length}). Using block_size={tokenizer.model_max_length}."
+            )
+        block_size = min(data_args.block_size, tokenizer.model_max_length)
+
+    # Main data processing function that will concatenate all texts from our dataset and generate chunks of block_size.
+    def group_texts(examples):
+        # Concatenate all texts.
+        concatenated_examples = {k: list(chain(*examples[k])) for k in examples.keys()}
+        total_length = len(concatenated_examples[list(examples.keys())[0]])
+        # We drop the small remainder, we could add padding if the model supported it instead of this drop, you can
+        # customize this part to your needs.
+        if total_length >= block_size:
+            total_length = (total_length // block_size) * block_size
+        # Split by chunks of max_len.
+        result = {
+            k: [t[i: i + block_size] for i in range(0, total_length, block_size)]
+            for k, t in concatenated_examples.items()
+        }
+        result["labels"] = result["input_ids"].copy()
+        return result
+
+    # Get the datasets: you can either provide your own CSV/JSON/TXT training and evaluation files (see below)
+    # or just provide the name of one of the public datasets available on the hub at https://huggingface.co/datasets/
+    # (the dataset will be downloaded automatically from the datasets Hub).
+    #
+    # For CSV/JSON files, this script will use the column called 'text' or the first column if no column called
+    # 'text' is found. You can easily tweak this behavior (see below).
+    #
+    # In distributed training, the load_dataset function guarantee that only one local process can concurrently
+    # download the dataset.
+    if data_args.dataset_name is not None:
+        # Downloading and loading a dataset from the hub.
+        raw_datasets = load_dataset(
+            data_args.dataset_name,
+            data_args.dataset_config_name,
+            cache_dir=model_args.cache_dir,
+            streaming=data_args.streaming,
+        )
+        if "validation" not in raw_datasets.keys():
+            raw_datasets["validation"] = load_dataset(
+                data_args.dataset_name,
+                data_args.dataset_config_name,
+                split=f"train[:{data_args.validation_split_percentage}%]",
+                cache_dir=model_args.cache_dir,
+                streaming=data_args.streaming,
+            )
+            raw_datasets["train"] = load_dataset(
+                data_args.dataset_name,
+                data_args.dataset_config_name,
+                split=f"train[{data_args.validation_split_percentage}%:]",
+                cache_dir=model_args.cache_dir,
+                streaming=data_args.streaming,
+            )
+    else:
+        data_files = {}
+        dataset_args = {}
+        if data_args.train_file_dir is not None and os.path.exists(data_args.train_file_dir):
+            train_data_files = glob(f'{data_args.train_file_dir}/**/*.txt', recursive=True) + glob(
+                f'{data_args.train_file_dir}/**/*.json', recursive=True) + glob(
+                f'{data_args.train_file_dir}/**/*.jsonl', recursive=True)
+            logger.info(f"train files: {train_data_files}")
+            # Train data files must be same type, e.g. all txt or all jsonl
+            types = [f.split('.')[-1] for f in train_data_files]
+            if len(set(types)) > 1:
+                raise ValueError(f"train files must be same type, e.g. all txt or all jsonl, but got {types}")
+            data_files["train"] = train_data_files
+        if data_args.validation_file_dir is not None and os.path.exists(data_args.validation_file_dir):
+            eval_data_files = glob(f'{data_args.validation_file_dir}/**/*.txt', recursive=True) + glob(
+                f'{data_args.train_file_dir}/**/*.json', recursive=True) + glob(
+                f'{data_args.train_file_dir}/**/*.jsonl', recursive=True)
+            logger.info(f"eval files: {eval_data_files}")
+            data_files["validation"] = eval_data_files
+            # Train data files must be same type, e.g. all txt or all jsonl
+            types = [f.split('.')[-1] for f in eval_data_files]
+            if len(set(types)) > 1:
+                raise ValueError(f"train files must be same type, e.g. all txt or all jsonl, but got {types}")
+        extension = "text" if data_files["train"][0].endswith('txt') else 'json'
+        if extension == "text":
+            dataset_args["keep_linebreaks"] = data_args.keep_linebreaks
+        raw_datasets = load_dataset(
+            extension,
+            data_files=data_files,
+            cache_dir=model_args.cache_dir,
+            **dataset_args,
+        )
+
+        # If no validation data is there, validation_split_percentage will be used to divide the dataset.
+        if "validation" not in raw_datasets.keys():
+            raw_datasets["validation"] = load_dataset(
+                extension,
+                data_files=data_files,
+                split=f"train[:{data_args.validation_split_percentage}%]",
+                cache_dir=model_args.cache_dir,
+                **dataset_args,
+            )
+            raw_datasets["train"] = load_dataset(
+                extension,
+                data_files=data_files,
+                split=f"train[{data_args.validation_split_percentage}%:]",
+                cache_dir=model_args.cache_dir,
+                **dataset_args,
+            )
+    logger.info(f"Raw datasets: {raw_datasets}")
+
+    # Preprocessing the datasets.
+    if training_args.do_train:
+        column_names = list(raw_datasets["train"].features)
+    else:
+        column_names = list(raw_datasets["validation"].features)
+
+    with training_args.main_process_first(desc="Dataset tokenization and grouping"):
+        if not data_args.streaming:
+            tokenized_datasets = raw_datasets.map(
+                tokenize_function,
+                batched=True,
+                num_proc=data_args.preprocessing_num_workers,
+                remove_columns=column_names,
+                load_from_cache_file=not data_args.overwrite_cache,
+                desc="Running tokenizer on dataset",
+            )
+            lm_datasets = tokenized_datasets.map(
+                group_texts,
+                batched=True,
+                num_proc=data_args.preprocessing_num_workers,
+                load_from_cache_file=not data_args.overwrite_cache,
+                desc=f"Grouping texts in chunks of {block_size}",
+            )
+        else:
+            tokenized_datasets = raw_datasets.map(
+                tokenize_function,
+                batched=True,
+                remove_columns=column_names,
+            )
+            lm_datasets = tokenized_datasets.map(
+                group_texts,
+                batched=True,
+            )
+
+    train_dataset = None
+    max_train_samples = 0
+    if training_args.do_train:
+        if "train" not in tokenized_datasets:
+            raise ValueError("--do_train requires a train dataset")
+        train_dataset = lm_datasets['train']
+        max_train_samples = len(train_dataset)
+        if data_args.max_train_samples is not None and data_args.max_train_samples > 0:
+            max_train_samples = min(len(train_dataset), data_args.max_train_samples)
+            train_dataset = train_dataset.select(range(max_train_samples))
+        logger.debug(f"Num train_samples: {len(train_dataset)}")
+        logger.debug("Tokenized training example:")
+        logger.debug(tokenizer.decode(train_dataset[0]['input_ids']))
+
+    eval_dataset = None
+    max_eval_samples = 0
+    if training_args.do_eval:
+        if "validation" not in tokenized_datasets:
+            raise ValueError("--do_eval requires a validation dataset")
+        eval_dataset = lm_datasets["validation"]
+        max_eval_samples = len(eval_dataset)
+        if data_args.max_eval_samples is not None and data_args.max_eval_samples > 0:
+            max_eval_samples = min(len(eval_dataset), data_args.max_eval_samples)
+            eval_dataset = eval_dataset.select(range(max_eval_samples))
+        logger.debug(f"Num eval_samples: {len(eval_dataset)}")
+        logger.debug("Tokenized eval example:")
+        logger.debug(tokenizer.decode(eval_dataset[0]['input_ids']))
+
+    # Load model
+    if model_args.model_type and model_args.model_name_or_path:
+        torch_dtype = (
+            model_args.torch_dtype
+            if model_args.torch_dtype in ["auto", None]
+            else getattr(torch, model_args.torch_dtype)
+        )
+        world_size = int(os.environ.get("WORLD_SIZE", 1))
+        ddp = world_size != 1
+        if ddp:
+            model_args.device_map = {"": int(os.environ["LOCAL_RANK"]) or 0}
+
+        config = config_class.from_pretrained(
+            model_args.model_name_or_path,
+            torch_dtype=torch_dtype,
+            trust_remote_code=model_args.trust_remote_code,
+            cache_dir=model_args.cache_dir
+        )
+        model = model_class.from_pretrained(
+            model_args.model_name_or_path,
+            config=config,
+            load_in_8bit=model_args.load_in_8bit,
+            device_map=model_args.device_map,
+            trust_remote_code=model_args.trust_remote_code,
+        )
+    else:
+        raise ValueError(f"Error, model_name_or_path is None, Continue PT must be loaded from a pre-trained model")
+
+    if training_args.use_peft:
+        if training_args.peft_path is not None:
+            logger.info(f"Peft from pre-trained model: {training_args.peft_path}")
+            model = PeftModel.from_pretrained(model, training_args.peft_path, is_trainable=True)
+        else:
+            logger.info("Init new peft model")
+            target_modules = training_args.target_modules.split(',') if training_args.target_modules else None
+            if target_modules and 'all' in target_modules:
+                target_modules = find_all_linear_names(model, int4=False, int8=model_args.load_in_8bit)
+            modules_to_save = training_args.modules_to_save
+            if modules_to_save is not None:
+                modules_to_save = modules_to_save.split(',')
+            logger.info(f"Peft target_modules: {target_modules}")
+            logger.info(f"Peft lora_rank: {training_args.lora_rank}")
+            peft_config = LoraConfig(
+                task_type=TaskType.CAUSAL_LM,
+                target_modules=target_modules,
+                inference_mode=False,
+                r=training_args.lora_rank,
+                lora_alpha=training_args.lora_alpha,
+                lora_dropout=training_args.lora_dropout,
+                modules_to_save=modules_to_save)
+            model = get_peft_model(model, peft_config)
+        if model_args.load_in_8bit:
+            model = prepare_model_for_int8_training(model)
+        model.print_trainable_parameters()
+    else:
+        logger.info("Full parameters training")
+        model = model.float()
+        print_trainable_parameters(model)
+
+    # Initialize our Trainer
+    if training_args.gradient_checkpointing:
+        model.gradient_checkpointing_enable()
+        model.config.use_cache = False
+    else:
+        model.config.use_cache = True
+    model.enable_input_require_grads()
+    if not ddp and torch.cuda.device_count() > 1:
+        # Keeps Trainer from trying its own DataParallelism when more than 1 gpu is available
+        model.is_parallelizable = True
+        model.model_parallel = True
+
+    trainer = SavePeftModelTrainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_dataset if training_args.do_train else None,
+        eval_dataset=eval_dataset if training_args.do_eval else None,
+        tokenizer=tokenizer,
+        data_collator=fault_tolerance_data_collator,
+        compute_metrics=compute_metrics if training_args.do_eval and not is_torch_tpu_available() else None,
+        preprocess_logits_for_metrics=preprocess_logits_for_metrics
+        if training_args.do_eval and not is_torch_tpu_available()
+        else None,
+    )
+
+    # Training
+    if training_args.do_train:
+        logger.info("*** Train ***")
+        logger.debug(f"Train dataloader example: {next(iter(trainer.get_train_dataloader()))}")
+        checkpoint = None
+        if training_args.resume_from_checkpoint is not None:
+            checkpoint = training_args.resume_from_checkpoint
+        train_result = trainer.train(resume_from_checkpoint=checkpoint)
+
+        metrics = train_result.metrics
+        metrics["train_samples"] = max_train_samples
+        logger.debug(f"Training metrics: {metrics}")
+        trainer.log_metrics("train", metrics)
+        trainer.save_metrics("train", metrics)
+        trainer.save_state()
+        logger.info(f"Saving model checkpoint to {training_args.output_dir}")
+        save_model(training_args.output_dir, model, tokenizer, training_args)
+
+    # Evaluation
+    if training_args.do_eval and trainer.is_world_process_zero():
+        logger.info("*** Evaluate ***")
+        metrics = trainer.evaluate()
+
+        metrics["eval_samples"] = max_eval_samples
+        try:
+            perplexity = math.exp(metrics["eval_loss"])
+        except OverflowError:
+            perplexity = float("inf")
+        metrics["perplexity"] = perplexity
+        logger.debug(f"Eval metrics: {metrics}")
+        trainer.log_metrics("eval", metrics)
+        trainer.save_metrics("eval", metrics)
+
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -0,0 +1,10 @@
+loguru
+transformers>=4.30.1
+sentencepiece
+datasets
+tqdm
+tensorboard
+tqdm>=4.47.0
+peft>=0.5.0
+accelerate>=0.20.3
+trl>=0.6.0

reward_modeling.py

@@ -0,0 +1,643 @@
+# -*- coding: utf-8 -*-
+"""
+@author:XuMing([email protected])
+@description:
+"""
+
+import math
+import os
+from dataclasses import dataclass, field
+from glob import glob
+from typing import Any, List, Union, Optional, Dict
+
+import torch
+from datasets import load_dataset
+from loguru import logger
+from peft import LoraConfig, TaskType, get_peft_model, PeftModel, prepare_model_for_int8_training
+from sklearn.metrics import mean_squared_error, mean_absolute_error
+from torch.utils.data import Dataset
+from transformers import (
+    AutoConfig,
+    PreTrainedTokenizerBase,
+    BloomForSequenceClassification,
+    LlamaForSequenceClassification,
+    LlamaTokenizer,
+    BloomTokenizerFast,
+    AlbertForSequenceClassification,
+    BertForSequenceClassification,
+    BertTokenizer,
+    AutoTokenizer,
+    RobertaForSequenceClassification,
+    AutoModelForSequenceClassification,
+    RobertaTokenizer,
+    HfArgumentParser,
+    Trainer,
+    TrainingArguments,
+    set_seed,
+)
+from transformers.trainer import TRAINING_ARGS_NAME
+
+MODEL_CLASSES = {
+    "bert": (AutoConfig, BertForSequenceClassification, BertTokenizer),
+    "roberta": (AutoConfig, RobertaForSequenceClassification, RobertaTokenizer),
+    "albert": (AutoConfig, AlbertForSequenceClassification, AutoTokenizer),
+    "bloom": (AutoConfig, BloomForSequenceClassification, BloomTokenizerFast),
+    "llama": (AutoConfig, LlamaForSequenceClassification, LlamaTokenizer),
+    "auto": (AutoConfig, AutoModelForSequenceClassification, AutoTokenizer),
+}
+
+
+@dataclass
+class ModelArguments:
+    """
+    Arguments pertaining to which model/config/tokenizer we are going to fine-tune, or train from scratch.
+    """
+
+    model_type: str = field(
+        default=None,
+        metadata={"help": "Model type selected in the list: " + ", ".join(MODEL_CLASSES.keys())}
+    )
+    model_name_or_path: Optional[str] = field(
+        default=None,
+        metadata={
+            "help": (
+                "The model checkpoint for weights initialization.Don't set if you want to train a model from scratch."
+            )
+        },
+    )
+    tokenizer_name_or_path: Optional[str] = field(
+        default=None,
+        metadata={
+            "help": (
+                "The tokenizer for weights initialization.Don't set if you want to train a model from scratch."
+            )
+        },
+    )
+    load_in_8bit: bool = field(default=False, metadata={"help": "Whether to load the model in 8bit mode or not."})
+    cache_dir: Optional[str] = field(
+        default=None,
+        metadata={"help": "Where do you want to store the pretrained models downloaded from huggingface.co"},
+    )
+    use_fast_tokenizer: bool = field(
+        default=False,
+        metadata={"help": "Whether to use one of the fast tokenizer (backed by the tokenizers library) or not."},
+    )
+    torch_dtype: Optional[str] = field(
+        default=None,
+        metadata={
+            "help": (
+                "Override the default `torch.dtype` and load the model under this dtype. If `auto` is passed, the "
+                "dtype will be automatically derived from the model's weights."
+            ),
+            "choices": ["auto", "bfloat16", "float16", "float32"],
+        },
+    )
+    device_map: Optional[str] = field(
+        default="auto",
+        metadata={"help": "Device to map model to. If `auto` is passed, the device will be selected automatically. "},
+    )
+    trust_remote_code: bool = field(
+        default=True,
+        metadata={"help": "Whether to trust remote code when loading a model from a remote checkpoint."},
+    )
+
+    def __post_init__(self):
+        if self.model_type is None:
+            raise ValueError(
+                "You must specify a valid model_type to run training. Available model types are " + ", ".join(
+                    MODEL_CLASSES.keys()))
+        if self.model_name_or_path is None:
+            raise ValueError("You must specify a valid model_name_or_path to run training.")
+
+
+@dataclass
+class DataTrainingArguments:
+    """
+    Arguments pertaining to what data we are going to input our model for training and eval.
+    """
+
+    dataset_name: Optional[str] = field(
+        default=None, metadata={"help": "The name of the dataset to use (via the datasets library)."}
+    )
+    dataset_config_name: Optional[str] = field(
+        default=None, metadata={"help": "The configuration name of the dataset to use (via the datasets library)."}
+    )
+    train_file_dir: Optional[str] = field(default=None, metadata={"help": "The input jsonl data file folder."})
+    validation_file_dir: Optional[str] = field(default=None, metadata={"help": "The evaluation jsonl file folder."}, )
+    max_source_length: Optional[int] = field(default=256, metadata={"help": "Max length of prompt input text"})
+    max_target_length: Optional[int] = field(default=256, metadata={"help": "Max length of output text"})
+    max_train_samples: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": (
+                "For debugging purposes or quicker training, truncate the number of training examples to this "
+                "value if set."
+            )
+        },
+    )
+    max_eval_samples: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": (
+                "For debugging purposes or quicker training, truncate the number of evaluation examples to this "
+                "value if set."
+            )
+        },
+    )
+    overwrite_cache: bool = field(
+        default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}
+    )
+    validation_split_percentage: Optional[int] = field(
+        default=1,
+        metadata={
+            "help": "The percentage of the train set used as validation set in case there's no validation split"
+        },
+    )
+    preprocessing_num_workers: Optional[int] = field(
+        default=4,
+        metadata={"help": "The number of processes to use for the preprocessing."},
+    )
+
+
+@dataclass
+class PeftArguments(TrainingArguments):
+    use_peft: bool = field(default=True, metadata={"help": "Whether to use peft"})
+    target_modules: Optional[str] = field(default="all")
+    lora_rank: Optional[int] = field(default=8)
+    lora_dropout: Optional[float] = field(default=0.05)
+    lora_alpha: Optional[float] = field(default=32.0)
+    modules_to_save: Optional[str] = field(default=None)
+    peft_path: Optional[str] = field(default=None)
+
+
+def compute_metrics(eval_preds):
+    preds, labels = eval_preds
+    # Here, predictions is rewards_chosen and rewards_rejected.
+    if isinstance(preds, torch.Tensor):
+        preds = preds.detach().cpu().numpy()
+    if isinstance(labels, torch.Tensor):
+        labels = labels.detach().cpu().numpy()
+    # MSE
+    mse = mean_squared_error(labels, preds)
+    # MAE
+    mae = mean_absolute_error(labels, preds)
+
+    return {"mse": mse, "mae": mae}
+
+
+@dataclass
+class RewardDataCollatorWithPadding:
+    """We need to define a special data collator that batches the data in our chosen vs rejected format"""
+    tokenizer: PreTrainedTokenizerBase
+    padding: Union[bool, str] = True
+    max_length: Optional[int] = None
+    pad_to_multiple_of: Optional[int] = None
+    return_tensors: str = "pt"
+
+    def __call__(self, features: List[Dict[str, Any]]) -> Dict[str, Any]:
+        features_chosen = []
+        features_rejected = []
+        for feature in features:
+            features_chosen.append(
+                {
+                    "input_ids": feature["input_ids_chosen"],
+                    "attention_mask": feature["attention_mask_chosen"],
+                }
+            )
+            features_rejected.append(
+                {
+                    "input_ids": feature["input_ids_rejected"],
+                    "attention_mask": feature["attention_mask_rejected"],
+                }
+            )
+        batch_chosen = self.tokenizer.pad(
+            features_chosen,
+            padding=self.padding,
+            max_length=self.max_length,
+            pad_to_multiple_of=self.pad_to_multiple_of,
+            return_tensors=self.return_tensors,
+        )
+        batch_rejected = self.tokenizer.pad(
+            features_rejected,
+            padding=self.padding,
+            max_length=self.max_length,
+            pad_to_multiple_of=self.pad_to_multiple_of,
+            return_tensors=self.return_tensors,
+        )
+        batch = {
+            "input_ids_chosen": batch_chosen["input_ids"],
+            "attention_mask_chosen": batch_chosen["attention_mask"],
+            "input_ids_rejected": batch_rejected["input_ids"],
+            "attention_mask_rejected": batch_rejected["attention_mask"],
+            "return_loss": True,
+        }
+        return batch
+
+
+class RewardTrainer(Trainer):
+    """
+    Trainer for reward models
+        Define how to compute the reward loss. Use the InstructGPT pairwise logloss: https://arxiv.org/abs/2203.02155
+    """
+
+    def compute_loss(self, model, inputs, return_outputs=False):
+        rewards_chosen = model(input_ids=inputs["input_ids_chosen"],
+                               attention_mask=inputs["attention_mask_chosen"])[0]
+        rewards_rejected = model(input_ids=inputs["input_ids_rejected"],
+                                 attention_mask=inputs["attention_mask_rejected"])[0]
+        loss = -torch.nn.functional.logsigmoid(rewards_chosen - rewards_rejected).mean()
+        if return_outputs:
+            return loss, {"rewards_chosen": rewards_chosen, "rewards_rejected": rewards_rejected}
+        return loss
+
+    def evaluate(
+            self,
+            eval_dataset: Optional[Dataset] = None,
+            ignore_keys: Optional[List[str]] = None,
+            metric_key_prefix: str = "eval",
+    ) -> Dict[str, float]:
+        if eval_dataset is None:
+            eval_dataset = self.eval_dataset
+        return super().evaluate(eval_dataset=eval_dataset, ignore_keys=ignore_keys, metric_key_prefix=metric_key_prefix)
+
+    def prediction_step(self, model, inputs, prediction_loss_only, ignore_keys=None):
+        # Prepare inputs for chosen and rejected separately
+        device = model.device
+
+        inputs_chosen = {
+            "input_ids": inputs["input_ids_chosen"].to(device),
+            "attention_mask": inputs["attention_mask_chosen"].to(device),
+        }
+        outputs_chosen = model(**inputs_chosen)
+        rewards_chosen = outputs_chosen.logits.detach()
+
+        inputs_rejected = {
+            "input_ids": inputs["input_ids_rejected"].to(device),
+            "attention_mask": inputs["attention_mask_rejected"].to(device),
+        }
+        outputs_rejected = model(**inputs_rejected)
+        rewards_rejected = outputs_rejected.logits.detach()
+
+        # Keep the compute_loss method
+        loss = -torch.nn.functional.logsigmoid(rewards_chosen - rewards_rejected).mean()
+        if prediction_loss_only:
+            return (loss, None, None)
+
+        return (loss, rewards_chosen, rewards_rejected)
+
+    def save_model(self, output_dir=None, _internal_call=False):
+        """Save the LoRA model."""
+        os.makedirs(output_dir, exist_ok=True)
+        torch.save(self.args, os.path.join(output_dir, TRAINING_ARGS_NAME))
+        self.model.save_pretrained(output_dir)
+
+
+def save_model(output_dir, model, tokenizer, args):
+    """Save the model and the tokenizer."""
+    os.makedirs(output_dir, exist_ok=True)
+
+    # Take care of distributed/parallel training
+    model_to_save = model.module if hasattr(model, "module") else model
+    model_to_save.save_pretrained(output_dir)
+    tokenizer.save_pretrained(output_dir)
+    torch.save(args, os.path.join(output_dir, TRAINING_ARGS_NAME))
+
+
+class CastOutputToFloat(torch.nn.Sequential):
+    """Cast the output of the model to float"""
+
+    def forward(self, x):
+        return super().forward(x).to(torch.float32)
+
+
+def print_trainable_parameters(model):
+    """
+    Prints the number of trainable parameters in the model.
+    """
+    trainable_params = 0
+    all_param = 0
+    for _, param in model.named_parameters():
+        all_param += param.numel()
+        if param.requires_grad:
+            trainable_params += param.numel()
+    print(
+        f"trainable params: {trainable_params} || all params: {all_param} || trainable%: {100 * trainable_params / all_param}"
+    )
+
+
+def find_all_linear_names(peft_model, int4=False, int8=False):
+    cls = torch.nn.Linear
+    if int4 or int8:
+        import bitsandbytes as bnb
+        if int4:
+            cls = bnb.nn.Linear4bit
+        elif int8:
+            cls = bnb.nn.Linear8bitLt
+    lora_module_names = set()
+    for name, module in peft_model.named_modules():
+        if isinstance(module, cls):
+            # last layer is not add to lora_module_names
+            if 'lm_head' in name:
+                continue
+            if 'score' in name:
+                continue
+            names = name.split('.')
+            lora_module_names.add(names[0] if len(names) == 1 else names[-1])
+    return sorted(lora_module_names)
+
+
+def main():
+    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, PeftArguments))
+    model_args, data_args, training_args = parser.parse_args_into_dataclasses()
+
+    logger.info(f"Model args: {model_args}")
+    logger.info(f"Data args: {data_args}")
+    logger.info(f"Training args: {training_args}")
+    logger.info(
+        f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}"
+        + f" distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}"
+    )
+
+    # Set seed before initializing model.
+    set_seed(training_args.seed)
+
+    # Load model
+    if not model_args.model_type:
+        raise ValueError("Please specify a model_type, e.g. llama, chatglm, bloom, etc.")
+    config_class, model_class, tokenizer_class = MODEL_CLASSES[model_args.model_type]
+    if model_args.model_name_or_path:
+        torch_dtype = (
+            model_args.torch_dtype
+            if model_args.torch_dtype in ["auto", None]
+            else getattr(torch, model_args.torch_dtype)
+        )
+        world_size = int(os.environ.get("WORLD_SIZE", 1))
+        if world_size > 1:
+            model_args.device_map = {"": int(os.environ["LOCAL_RANK"]) or 0}
+        config = config_class.from_pretrained(
+            model_args.model_name_or_path,
+            num_labels=1,
+            torch_dtype=torch_dtype,
+            trust_remote_code=model_args.trust_remote_code,
+            cache_dir=model_args.cache_dir
+        )
+        if model_args.model_type in ['bloom', 'llama']:
+            model = model_class.from_pretrained(
+                model_args.model_name_or_path,
+                config=config,
+                load_in_8bit=model_args.load_in_8bit,
+                device_map=model_args.device_map,
+                trust_remote_code=model_args.trust_remote_code,
+            )
+            model.score = CastOutputToFloat(model.score)
+        else:
+            model = model_class.from_pretrained(
+                model_args.model_name_or_path,
+                config=config,
+                cache_dir=model_args.cache_dir,
+                ignore_mismatched_sizes=True
+            )
+            model.to(training_args.device)
+    else:
+        raise ValueError(f"Error, model_name_or_path is None, RM must be loaded from a pre-trained model")
+
+    # Load tokenizer
+    if model_args.model_type == "bloom":
+        model_args.use_fast_tokenizer = True
+    tokenizer_kwargs = {
+        "cache_dir": model_args.cache_dir,
+        "use_fast": model_args.use_fast_tokenizer,
+        "trust_remote_code": model_args.trust_remote_code,
+    }
+    tokenizer_name_or_path = model_args.tokenizer_name_or_path
+    if not tokenizer_name_or_path:
+        tokenizer_name_or_path = model_args.model_name_or_path
+    tokenizer = tokenizer_class.from_pretrained(tokenizer_name_or_path, **tokenizer_kwargs)
+    if tokenizer.pad_token_id is None:
+        tokenizer.pad_token_id = 0
+
+    if training_args.use_peft:
+        if training_args.peft_path is not None:
+            logger.info(f"Peft from pre-trained model: {training_args.peft_path}")
+            model = PeftModel.from_pretrained(model, training_args.peft_path, is_trainable=True)
+        else:
+            logger.info("Init new peft model")
+            target_modules = training_args.target_modules.split(',') if training_args.target_modules else None
+            if target_modules and 'all' in target_modules:
+                target_modules = find_all_linear_names(model, int4=False, int8=model_args.load_in_8bit)
+            modules_to_save = training_args.modules_to_save
+            if modules_to_save is not None:
+                modules_to_save = modules_to_save.split(',')
+            logger.info(f"Peft target_modules: {target_modules}")
+            logger.info(f"Peft lora_rank: {training_args.lora_rank}")
+            peft_config = LoraConfig(
+                task_type=TaskType.SEQ_CLS,
+                target_modules=target_modules,
+                inference_mode=False,
+                r=training_args.lora_rank,
+                lora_alpha=training_args.lora_alpha,
+                lora_dropout=training_args.lora_dropout,
+                modules_to_save=modules_to_save)
+            model = get_peft_model(model, peft_config)
+        if model_args.load_in_8bit:
+            model = prepare_model_for_int8_training(model)
+        model.print_trainable_parameters()
+    else:
+        logger.info("Full parameters training")
+        print_trainable_parameters(model)
+
+    # Get reward dataset for tuning the reward model.
+    if data_args.dataset_name is not None:
+        # Downloading and loading a dataset from the hub.
+        raw_datasets = load_dataset(
+            data_args.dataset_name,
+            data_args.dataset_config_name,
+            cache_dir=model_args.cache_dir,
+        )
+        if "validation" not in raw_datasets.keys():
+            raw_datasets["validation"] = load_dataset(
+                data_args.dataset_name,
+                data_args.dataset_config_name,
+                split=f"train[:{data_args.validation_split_percentage}%]",
+                cache_dir=model_args.cache_dir,
+            )
+            raw_datasets["train"] = load_dataset(
+                data_args.dataset_name,
+                data_args.dataset_config_name,
+                split=f"train[{data_args.validation_split_percentage}%:]",
+                cache_dir=model_args.cache_dir,
+            )
+    else:
+        data_files = {}
+        if data_args.train_file_dir is not None and os.path.exists(data_args.train_file_dir):
+            train_data_files = glob(f'{data_args.train_file_dir}/**/*.json', recursive=True) + glob(
+                f'{data_args.train_file_dir}/**/*.jsonl', recursive=True)
+            logger.info(f"train files: {', '.join(train_data_files)}")
+            data_files["train"] = train_data_files
+        if data_args.validation_file_dir is not None and os.path.exists(data_args.validation_file_dir):
+            eval_data_files = glob(f'{data_args.validation_file_dir}/**/*.json', recursive=True) + glob(
+                f'{data_args.validation_file_dir}/**/*.jsonl', recursive=True)
+            logger.info(f"eval files: {', '.join(eval_data_files)}")
+            data_files["validation"] = eval_data_files
+        raw_datasets = load_dataset(
+            'json',
+            data_files=data_files,
+            cache_dir=model_args.cache_dir,
+        )
+        # If no validation data is there, validation_split_percentage will be used to divide the dataset.
+        if "validation" not in raw_datasets.keys():
+            raw_datasets["validation"] = load_dataset(
+                'json',
+                data_files=data_files,
+                split=f"train[:{data_args.validation_split_percentage}%]",
+                cache_dir=model_args.cache_dir,
+            )
+            raw_datasets["train"] = load_dataset(
+                'json',
+                data_files=data_files,
+                split=f"train[{data_args.validation_split_percentage}%:]",
+                cache_dir=model_args.cache_dir,
+            )
+    logger.info(f"Raw datasets: {raw_datasets}")
+
+    # Preprocessing the datasets
+    full_max_length = data_args.max_source_length + data_args.max_target_length
+
+    def preprocess_reward_function(examples):
+        """
+        Turn the dataset into pairs of Question + Answer, where input_ids_chosen is the preferred question + answer
+            and text_rejected is the other.
+        """
+        new_examples = {
+            "input_ids_chosen": [],
+            "attention_mask_chosen": [],
+            "input_ids_rejected": [],
+            "attention_mask_rejected": [],
+        }
+        for question, chosen, rejected in zip(examples["question"], examples["response_chosen"],
+                                              examples["response_rejected"]):
+            tokenized_chosen = tokenizer("Question: " + question + "\n\nAnswer: " + chosen)
+            tokenized_rejected = tokenizer("Question: " + question + "\n\nAnswer: " + rejected)
+
+            new_examples["input_ids_chosen"].append(tokenized_chosen["input_ids"])
+            new_examples["attention_mask_chosen"].append(tokenized_chosen["attention_mask"])
+            new_examples["input_ids_rejected"].append(tokenized_rejected["input_ids"])
+            new_examples["attention_mask_rejected"].append(tokenized_rejected["attention_mask"])
+
+        return new_examples
+
+    train_dataset = None
+    max_train_samples = 0
+    if training_args.do_train:
+        if "train" not in raw_datasets:
+            raise ValueError("--do_train requires a train dataset")
+        train_dataset = raw_datasets['train']
+        max_train_samples = len(train_dataset)
+        if data_args.max_train_samples is not None and data_args.max_train_samples > 0:
+            max_train_samples = min(len(train_dataset), data_args.max_train_samples)
+            train_dataset = train_dataset.select(range(max_train_samples))
+        logger.debug(f"Example train_dataset[0]: {train_dataset[0]}")
+        with training_args.main_process_first(desc="Train dataset tokenization"):
+            tokenized_dataset = train_dataset.shuffle().map(
+                preprocess_reward_function,
+                batched=True,
+                num_proc=data_args.preprocessing_num_workers,
+                remove_columns=train_dataset.column_names,
+                load_from_cache_file=not data_args.overwrite_cache,
+                desc="Running tokenizer on dataset",
+            )
+            train_dataset = tokenized_dataset.filter(
+                lambda x: 0 < len(x['input_ids_rejected']) <= full_max_length and 0 < len(
+                    x['input_ids_chosen']) <= full_max_length
+            )
+            logger.debug(f"Num train_samples: {len(train_dataset)}")
+            logger.debug("Tokenized training example:")
+            logger.debug(tokenizer.decode(train_dataset[0]['input_ids_chosen']))
+
+    eval_dataset = None
+    max_eval_samples = 0
+    if training_args.do_eval:
+        with training_args.main_process_first(desc="Eval dataset tokenization"):
+            if "validation" not in raw_datasets:
+                raise ValueError("--do_eval requires a validation dataset")
+            eval_dataset = raw_datasets["validation"]
+            max_eval_samples = len(eval_dataset)
+            if data_args.max_eval_samples is not None and data_args.max_eval_samples > 0:
+                max_eval_samples = min(len(eval_dataset), data_args.max_eval_samples)
+                eval_dataset = eval_dataset.select(range(max_eval_samples))
+            logger.debug(f"Example eval_dataset[0]: {eval_dataset[0]}")
+            tokenized_dataset = eval_dataset.map(
+                preprocess_reward_function,
+                batched=True,
+                num_proc=data_args.preprocessing_num_workers,
+                remove_columns=eval_dataset.column_names,
+                load_from_cache_file=not data_args.overwrite_cache,
+                desc="Running tokenizer on dataset",
+            )
+            eval_dataset = tokenized_dataset.filter(
+                lambda x: 0 < len(x['input_ids_rejected']) <= full_max_length and 0 < len(
+                    x['input_ids_chosen']) <= full_max_length
+            )
+            logger.debug(f"Num eval_samples: {len(eval_dataset)}")
+            logger.debug("Tokenized eval example:")
+            logger.debug(tokenizer.decode(eval_dataset[0]['input_ids_chosen']))
+
+    # Initialize our Trainer
+    if training_args.gradient_checkpointing:
+        model.gradient_checkpointing_enable()
+        model.config.use_cache = False
+    else:
+        model.config.use_cache = True
+    model.enable_input_require_grads()
+    if torch.cuda.device_count() > 1:
+        # Keeps Trainer from trying its own DataParallelism when more than 1 gpu is available
+        model.is_parallelizable = True
+        model.model_parallel = True
+    trainer = RewardTrainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_dataset if training_args.do_train else None,
+        eval_dataset=eval_dataset if training_args.do_eval else None,
+        tokenizer=tokenizer,
+        compute_metrics=compute_metrics,
+        data_collator=RewardDataCollatorWithPadding(
+            tokenizer=tokenizer, max_length=full_max_length, padding="max_length"
+        ),
+    )
+
+    # Training
+    if training_args.do_train:
+        logger.info("*** Train ***")
+        logger.debug(f"Train dataloader example: {next(iter(trainer.get_train_dataloader()))}")
+        checkpoint = None
+        if training_args.resume_from_checkpoint is not None:
+            checkpoint = training_args.resume_from_checkpoint
+        train_result = trainer.train(resume_from_checkpoint=checkpoint)
+
+        metrics = train_result.metrics
+        metrics["train_samples"] = max_train_samples
+        logger.debug(f"Training metrics: {metrics}")
+        trainer.log_metrics("train", metrics)
+        trainer.save_metrics("train", metrics)
+        trainer.save_state()
+        logger.info(f"Saving model checkpoint to {training_args.output_dir}")
+        save_model(training_args.output_dir, model, tokenizer, training_args)
+
+    # Evaluation
+    if training_args.do_eval and trainer.is_world_process_zero():
+        logger.info("*** Evaluate ***")
+        metrics = trainer.evaluate()
+
+        metrics["eval_samples"] = max_eval_samples
+        try:
+            perplexity = math.exp(metrics["eval_loss"])
+        except OverflowError:
+            perplexity = float("inf")
+        metrics["perplexity"] = perplexity
+        logger.debug(f"Eval metrics: {metrics}")
+        trainer.log_metrics("eval", metrics)
+        trainer.save_metrics("eval", metrics)
+
+
+if __name__ == "__main__":
+    main()

rl_training.py

@@ -0,0 +1,499 @@
+# -*- coding: utf-8 -*-
+"""
+@author:XuMing([email protected])
+@description: Train a model from SFT using PPO
+"""
+
+import os
+from dataclasses import dataclass, field
+from glob import glob
+from typing import Optional
+
+import torch
+from datasets import load_dataset
+from loguru import logger
+from peft import LoraConfig, TaskType
+from tqdm import tqdm
+from transformers import (
+    AutoConfig,
+    AutoModelForSequenceClassification,
+    BloomForCausalLM,
+    AutoModelForCausalLM,
+    AutoModel,
+    LlamaTokenizer,
+    LlamaForCausalLM,
+    BloomTokenizerFast,
+    AutoTokenizer,
+    HfArgumentParser,
+)
+from trl import AutoModelForCausalLMWithValueHead, PPOConfig, PPOTrainer, set_seed
+
+from supervised_finetuning import get_conv_template
+
+os.environ["TOKENIZERS_PARALLELISM"] = "FALSE"
+os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
+
+MODEL_CLASSES = {
+    "bloom": (AutoConfig, BloomForCausalLM, BloomTokenizerFast),
+    "chatglm": (AutoConfig, AutoModel, AutoTokenizer),
+    "llama": (AutoConfig, LlamaForCausalLM, LlamaTokenizer),
+    "baichuan": (AutoConfig, AutoModelForCausalLM, AutoTokenizer),
+    "auto": (AutoConfig, AutoModelForCausalLM, AutoTokenizer),
+}
+
+
+@dataclass
+class ScriptArguments:
+    """
+    The name of the Casual LM model we wish to fine with PPO
+    """
+    # Model arguments
+    model_type: str = field(
+        default=None,
+        metadata={"help": "Model type selected in the list: " + ", ".join(MODEL_CLASSES.keys())}
+    )
+    model_name_or_path: Optional[str] = field(
+        default=None, metadata={"help": "The model checkpoint for weights initialization."}
+    )
+    reward_model_name_or_path: Optional[str] = field(default=None, metadata={"help": "The reward model name"})
+    tokenizer_name_or_path: Optional[str] = field(
+        default=None, metadata={"help": "The tokenizer for weights initialization."}
+    )
+    load_in_8bit: bool = field(default=False, metadata={"help": "Whether to load the model in 8bit mode or not."})
+    cache_dir: Optional[str] = field(
+        default=None,
+        metadata={"help": "Where do you want to store the pretrained models downloaded from huggingface.co"},
+    )
+    use_fast_tokenizer: bool = field(
+        default=False,
+        metadata={"help": "Whether to use one of the fast tokenizer (backed by the tokenizers library) or not."},
+    )
+    torch_dtype: Optional[str] = field(
+        default=None,
+        metadata={
+            "help": (
+                "Override the default `torch.dtype` and load the model under this dtype. If `auto` is passed, the "
+                "dtype will be automatically derived from the model's weights."
+            ),
+            "choices": ["auto", "bfloat16", "float16", "float32"],
+        },
+    )
+    device_map: Optional[str] = field(
+        default="auto",
+        metadata={"help": "Device to map model to. If `auto` is passed, the device will be selected automatically. "},
+    )
+    trust_remote_code: bool = field(
+        default=True,
+        metadata={"help": "Whether to trust remote code when loading a model from a remote checkpoint."},
+    )
+    # Dataset arguments
+    dataset_name: Optional[str] = field(
+        default=None, metadata={"help": "The name of the dataset to use (via the datasets library)."}
+    )
+    dataset_config_name: Optional[str] = field(
+        default=None, metadata={"help": "The configuration name of the dataset to use (via the datasets library)."}
+    )
+    train_file_dir: Optional[str] = field(default=None, metadata={"help": "The input jsonl data file folder."})
+    validation_file_dir: Optional[str] = field(default=None, metadata={"help": "The evaluation jsonl file folder."}, )
+    template_name: Optional[str] = field(default="vicuna", metadata={"help": "The template name."})
+    batch_size: Optional[int] = field(default=8, metadata={"help": "Batch size"})
+    mini_batch_size: Optional[int] = field(default=1, metadata={"help": "PPO minibatch size"})
+    max_source_length: Optional[int] = field(default=256, metadata={"help": "Max length of prompt input text"})
+    max_target_length: Optional[int] = field(default=256, metadata={"help": "Max length of output text"})
+    min_target_length: Optional[int] = field(default=4, metadata={"help": "Min length of output text"})
+    max_train_samples: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": (
+                "For debugging purposes or quicker training, truncate the number of training examples to this "
+                "value if set."
+            )
+        },
+    )
+    max_eval_samples: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": (
+                "For debugging purposes or quicker training, truncate the number of evaluation examples to this "
+                "value if set."
+            )
+        },
+    )
+    overwrite_cache: bool = field(
+        default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}
+    )
+    validation_split_percentage: Optional[int] = field(
+        default=1,
+        metadata={
+            "help": "The percentage of the train set used as validation set in case there's no validation split"
+        },
+    )
+    preprocessing_num_workers: Optional[int] = field(
+        default=None, metadata={"help": "The number of processes to use for the preprocessing."},
+    )
+    # Training arguments
+    use_peft: bool = field(default=True, metadata={"help": "Whether to use peft"})
+    target_modules: Optional[str] = field(default=None)
+    lora_rank: Optional[int] = field(default=8)
+    lora_dropout: Optional[float] = field(default=0.05)
+    lora_alpha: Optional[float] = field(default=32.0)
+    modules_to_save: Optional[str] = field(default=None)
+    peft_path: Optional[str] = field(default=None)
+
+    do_train: bool = field(default=False, metadata={"help": "Whether to run training."})
+    do_eval: bool = field(default=False, metadata={"help": "Whether to run eval on the validation set."})
+    early_stopping: Optional[bool] = field(default=False, metadata={"help": "Whether to early stop"})
+    target_kl: Optional[float] = field(default=0.1, metadata={"help": "The kl target for early stopping"})
+    reward_baseline: Optional[float] = field(
+        default=0.0, metadata={"help": "Baseline value that is subtracted from the reward"},
+    )
+    init_kl_coef: Optional[float] = field(
+        default=0.2, metadata={"help": "Initial KL penalty coefficient (used for adaptive and linear control)"},
+    )
+    adap_kl_ctrl: Optional[bool] = field(default=True, metadata={"help": "Use adaptive KL control, otherwise linear"})
+    learning_rate: Optional[float] = field(default=1.5e-5, metadata={"help": "Learning rate"})
+    gradient_accumulation_steps: Optional[int] = field(
+        default=1, metadata={"help": "the number of gradient accumulation steps"}
+    )
+    save_steps: Optional[int] = field(default=50, metadata={"help": "X steps to save the model"})
+    output_dir: Optional[str] = field(default="outputs-rl", metadata={"help": "The output directory"})
+    seed: Optional[int] = field(default=0, metadata={"help": "Seed"})
+    max_steps: Optional[int] = field(default=200, metadata={"help": "Number of steps to train"})
+    report_to: Optional[str] = field(default="tensorboard", metadata={"help": "Report to wandb or tensorboard"})
+
+    def __post_init__(self):
+        if self.model_type is None:
+            raise ValueError("You must specify a valid model_type to run training.")
+        if self.model_name_or_path is None:
+            raise ValueError("You must specify a valid model_name_or_path to run training.")
+        if self.reward_model_name_or_path is None:
+            raise ValueError("You must specify a valid reward_model_name_or_path to run training.")
+
+
+def print_trainable_parameters(model):
+    """
+    Prints the number of trainable parameters in the model.
+    """
+    trainable_params = 0
+    all_param = 0
+    for _, param in model.named_parameters():
+        all_param += param.numel()
+        if param.requires_grad:
+            trainable_params += param.numel()
+    print(
+        f"trainable params: {trainable_params} || all params: {all_param} || trainable%: {100 * trainable_params / all_param}"
+    )
+
+
+def get_reward_model_output(reward_model, reward_tokenizer, question, answer, device):
+    """
+    Get the reward score for a given question and answer pair.
+    """
+    inputs = reward_tokenizer(question, answer, return_tensors='pt').to(device)
+    score = reward_model(**inputs).logits[0].cpu().detach()
+
+    return score
+
+
+def calculate_rewards(reward_score_outputs, reward_baseline=0):
+    """
+    Calculate the reward for a given score output.
+    :param reward_score_outputs: 
+    :param reward_baseline: 
+    :return: 
+    """
+    rewards = []
+    for score in reward_score_outputs:
+        if isinstance(score, torch.Tensor) and score.numel() == 1:
+            reward_value = score.item() - reward_baseline
+            rewards.append(torch.tensor(reward_value))
+        else:
+            # Use the average of the tensor elements as `score` is multiple elements
+            reward_value = torch.mean(score).item() - reward_baseline
+            rewards.append(torch.tensor(reward_value))
+    return rewards
+
+
+def main():
+    parser = HfArgumentParser(ScriptArguments)
+    args = parser.parse_args_into_dataclasses()[0]
+    logger.info(f"Parse args: {args}")
+
+    config_class, model_class, tokenizer_class = MODEL_CLASSES[args.model_type]
+    if args.model_type == 'bloom':
+        args.use_fast_tokenizer = True
+    # Load tokenizer
+    tokenizer_kwargs = {
+        "cache_dir": args.cache_dir,
+        "use_fast": args.use_fast_tokenizer,
+        "trust_remote_code": args.trust_remote_code,
+    }
+    tokenizer_name_or_path = args.tokenizer_name_or_path
+    if not tokenizer_name_or_path:
+        tokenizer_name_or_path = args.model_name_or_path
+    tokenizer = tokenizer_class.from_pretrained(tokenizer_name_or_path, **tokenizer_kwargs)
+    if tokenizer.pad_token_id is None:
+        tokenizer.pad_token_id = 0  # set as the <unk> token
+
+    logger.info("Load model")
+    peft_config = LoraConfig(
+        task_type=TaskType.CAUSAL_LM,
+        target_modules=args.target_modules,
+        inference_mode=False,
+        r=args.lora_rank,
+        lora_alpha=args.lora_alpha,
+        lora_dropout=args.lora_dropout,
+    )
+    torch_dtype = (
+        args.torch_dtype
+        if args.torch_dtype in ["auto", None]
+        else getattr(torch, args.torch_dtype)
+    )
+    world_size = int(os.environ.get("WORLD_SIZE", 1))
+    if world_size > 1:
+        args.device_map = {"": int(os.environ["LOCAL_RANK"]) or 0}
+    config = config_class.from_pretrained(
+        args.model_name_or_path,
+        torch_dtype=torch_dtype,
+        trust_remote_code=args.trust_remote_code,
+        cache_dir=args.cache_dir
+    )
+    model = AutoModelForCausalLMWithValueHead.from_pretrained(
+        args.model_name_or_path,
+        config=config,
+        load_in_8bit=args.load_in_8bit,
+        device_map=args.device_map,
+        trust_remote_code=args.trust_remote_code,
+        peft_config=peft_config if args.use_peft else None,
+    )
+    print_trainable_parameters(model)
+    # Load reward model
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    reward_model = AutoModelForSequenceClassification.from_pretrained(
+        args.reward_model_name_or_path,
+        config=config,
+        load_in_8bit=args.load_in_8bit,
+        trust_remote_code=args.trust_remote_code,
+    )
+    reward_model.to(device)
+    reward_tokenizer = AutoTokenizer.from_pretrained(
+        args.reward_model_name_or_path, **tokenizer_kwargs
+    )
+
+    # Get datasets
+    if args.dataset_name is not None:
+        # Downloading and loading a dataset from the hub.
+        raw_datasets = load_dataset(
+            args.dataset_name,
+            args.dataset_config_name,
+            cache_dir=args.cache_dir,
+        )
+        if "validation" not in raw_datasets.keys():
+            raw_datasets["validation"] = load_dataset(
+                args.dataset_name,
+                args.dataset_config_name,
+                split=f"train[:{args.validation_split_percentage}%]",
+                cache_dir=args.cache_dir,
+            )
+            raw_datasets["train"] = load_dataset(
+                args.dataset_name,
+                args.dataset_config_name,
+                split=f"train[{args.validation_split_percentage}%:]",
+                cache_dir=args.cache_dir,
+            )
+    else:
+        data_files = {}
+        if args.train_file_dir is not None and os.path.exists(args.train_file_dir):
+            train_data_files = glob(f'{args.train_file_dir}/**/*.json', recursive=True) + glob(
+                f'{args.train_file_dir}/**/*.jsonl', recursive=True)
+            logger.info(f"train files: {', '.join(train_data_files)}")
+            data_files["train"] = train_data_files
+        if args.validation_file_dir is not None and os.path.exists(args.validation_file_dir):
+            eval_data_files = glob(f'{args.validation_file_dir}/**/*.json', recursive=True) + glob(
+                f'{args.validation_file_dir}/**/*.jsonl', recursive=True)
+            logger.info(f"eval files: {', '.join(eval_data_files)}")
+            data_files["validation"] = eval_data_files
+        raw_datasets = load_dataset(
+            'json',
+            data_files=data_files,
+            cache_dir=args.cache_dir,
+        )
+        # If no validation data is there, validation_split_percentage will be used to divide the dataset.
+        if "validation" not in raw_datasets.keys():
+            raw_datasets["validation"] = load_dataset(
+                'json',
+                data_files=data_files,
+                split=f"train[:{args.validation_split_percentage}%]",
+                cache_dir=args.cache_dir,
+            )
+            raw_datasets["train"] = load_dataset(
+                'json',
+                data_files=data_files,
+                split=f"train[{args.validation_split_percentage}%:]",
+                cache_dir=args.cache_dir,
+            )
+    logger.info(f"Raw datasets: {raw_datasets}")
+
+    # Preprocessing the datasets
+    max_source_length = args.max_source_length
+    max_target_length = args.max_target_length
+    prompt_template = get_conv_template(args.template_name)
+
+    def preprocess_function(examples):
+        new_examples = {
+            "query": [],
+            "input_ids": [],
+        }
+        roles = ["human", "gpt"]
+
+        def get_prompt(examples):
+            for i, source in enumerate(examples['conversations']):
+                if len(source) < 2:
+                    continue
+                data_role = source[0].get("from", "")
+                if data_role not in roles or data_role != roles[0]:
+                    # Skip the first one if it is not from human
+                    source = source[1:]
+                if len(source) < 2:
+                    continue
+                messages = []
+                for j, sentence in enumerate(source):
+                    data_role = sentence.get("from", "")
+                    if data_role not in roles:
+                        logger.warning(f"unknown role: {data_role}, {i}. (ignored)")
+                        break
+                    if data_role == roles[j % 2]:
+                        messages.append(sentence["value"])
+                if len(messages) < 2 or len(messages) % 2 != 0:
+                    continue
+                # Convert the list to pairs of elements
+                history_messages = [[messages[k], messages[k + 1]] for k in range(0, len(messages), 2)]
+                yield prompt_template.get_prompt(history_messages)
+
+        for prompt in get_prompt(examples):
+            for i in range(len(prompt) // 2):
+                source_txt = prompt[2 * i]
+                tokenized_question = tokenizer(
+                    source_txt, truncation=True, max_length=max_source_length, padding="max_length",
+                    return_tensors="pt"
+                )
+                new_examples["query"].append(source_txt)
+                new_examples["input_ids"].append(tokenized_question["input_ids"])
+
+        return new_examples
+
+    # Preprocess the dataset
+    train_dataset = None
+    if args.do_train:
+        if "train" not in raw_datasets:
+            raise ValueError("--do_train requires a train dataset")
+        train_dataset = raw_datasets['train']
+        if args.max_train_samples is not None and args.max_train_samples > 0:
+            max_train_samples = min(len(train_dataset), args.max_train_samples)
+            train_dataset = train_dataset.select(range(max_train_samples))
+        logger.debug(f"Example train_dataset[0]: {train_dataset[0]}")
+        tokenized_dataset = train_dataset.shuffle().map(
+            preprocess_function,
+            batched=True,
+            num_proc=args.preprocessing_num_workers,
+            remove_columns=train_dataset.column_names,
+            load_from_cache_file=not args.overwrite_cache,
+            desc="Running tokenizer on dataset",
+        )
+        train_dataset = tokenized_dataset.filter(
+            lambda x: len(x['input_ids']) > 0
+        )
+        logger.debug(f"Num train_samples: {len(train_dataset)}")
+
+    def collator(data):
+        return dict((key, [d[key] for d in data]) for key in data[0])
+
+    output_dir = args.output_dir
+    config = PPOConfig(
+        steps=args.max_steps,
+        model_name=args.model_name_or_path,
+        learning_rate=args.learning_rate,
+        log_with=args.report_to,
+        batch_size=args.batch_size,
+        mini_batch_size=args.mini_batch_size,
+        gradient_accumulation_steps=args.gradient_accumulation_steps,
+        optimize_cuda_cache=True,
+        early_stopping=args.early_stopping,
+        target_kl=args.target_kl,
+        seed=args.seed,
+        init_kl_coef=args.init_kl_coef,
+        adap_kl_ctrl=args.adap_kl_ctrl,
+        project_kwargs={"logging_dir": output_dir},
+    )
+    # Set seed before initializing value head for deterministic eval
+    set_seed(config.seed)
+
+    # We then build the PPOTrainer, passing the model, the reference model, the tokenizer
+    trainer = PPOTrainer(
+        config,
+        model,
+        ref_model=None,
+        tokenizer=tokenizer,
+        dataset=train_dataset,
+        data_collator=collator,
+    )
+
+    # These arguments are passed to the `generate` function of the PPOTrainer
+    generation_kwargs = {
+        "max_new_tokens": max_target_length,
+        "temperature": 1.0,
+        "repetition_penalty": 1.0,
+        "top_p": 1.0,
+        "do_sample": True,
+    }
+
+    def save_model(save_dir):
+        trainer.accelerator.unwrap_model(trainer.model).save_pretrained(save_dir)
+        trainer.tokenizer.save_pretrained(save_dir)
+
+    # Training
+    if args.do_train:
+        logger.info("*** Train ***")
+        total_steps = config.total_ppo_epochs
+        for step, batch in tqdm(enumerate(trainer.dataloader)):
+            if step >= total_steps:
+                break
+            question_tensors = batch["input_ids"]
+            question_tensors = [torch.LongTensor(i).to(device).squeeze(0) for i in question_tensors]
+            responses = []
+            response_tensors = []
+            for q_tensor in question_tensors:
+                response_tensor = trainer.generate(
+                    q_tensor,
+                    return_prompt=False,
+                    **generation_kwargs,
+                )
+                r = tokenizer.batch_decode(response_tensor, skip_special_tokens=True)[0]
+                responses.append(r)
+                response_tensors.append(response_tensor.squeeze(0))
+            batch["response"] = responses
+
+            # Compute reward score
+            score_outputs = [
+                get_reward_model_output(reward_model, reward_tokenizer, q, r, device) for q, r in
+                zip(batch["query"], batch["response"])
+            ]
+            rewards = calculate_rewards(score_outputs, args.reward_baseline)
+
+            # Run PPO step
+            try:
+                stats = trainer.step(question_tensors, response_tensors, rewards)
+                trainer.log_stats(stats, batch, rewards)
+                logger.debug(f"Step {step}/{total_steps}: reward score:{score_outputs}")
+            except ValueError as e:
+                logger.warning(f"Failed to log stats for step {step}, because of {e}")
+
+            if step and step % args.save_steps == 0:
+                save_dir = os.path.join(output_dir, f"checkpoint-{step}")
+                save_model(save_dir)
+        # Save final model
+        save_model(output_dir)
+
+
+if __name__ == "__main__":
+    main()

run_dpo.sh

@@ -0,0 +1,29 @@
+CUDA_VISIBLE_DEVICES=0,1 python dpo_training.py \
+    --model_type bloom \
+    --model_name_or_path bigscience/bloomz-560m \
+    --train_file_dir ./data/reward \
+    --validation_file_dir ./data/reward \
+    --per_device_train_batch_size 4 \
+    --per_device_eval_batch_size 1 \
+    --do_train \
+    --do_eval \
+    --use_peft True \
+    --max_train_samples 1000 \
+    --max_eval_samples 10 \
+    --max_steps 100 \
+    --eval_steps 20 \
+    --save_steps 50 \
+    --max_source_length 128 \
+    --max_target_length 128 \
+    --output_dir outputs-dpo-bloom-v1 \
+    --target_modules all \
+    --lora_rank 8 \
+    --lora_alpha 16 \
+    --lora_dropout 0.05 \
+    --torch_dtype float16 \
+    --fp16 True \
+    --device_map auto \
+    --report_to tensorboard \
+    --remove_unused_columns False \
+    --gradient_checkpointing True \
+    --cache_dir ./cache

run_pt.sh

@@ -0,0 +1,42 @@
+CUDA_VISIBLE_DEVICES=0,1 torchrun --nproc_per_node 2 pretraining.py \
+    --model_type bloom \
+    --model_name_or_path bigscience/bloomz-560m \
+    --train_file_dir ./data/pretrain \
+    --validation_file_dir ./data/pretrain \
+    --per_device_train_batch_size 4 \
+    --per_device_eval_batch_size 4 \
+    --do_train \
+    --do_eval \
+    --use_peft True \
+    --seed 42 \
+    --fp16 \
+    --max_train_samples 10000 \
+    --max_eval_samples 10 \
+    --num_train_epochs 0.5 \
+    --learning_rate 2e-4 \
+    --warmup_ratio 0.05 \
+    --weight_decay 0.01 \
+    --logging_strategy steps \
+    --logging_steps 10 \
+    --eval_steps 50 \
+    --evaluation_strategy steps \
+    --save_steps 500 \
+    --save_strategy steps \
+    --save_total_limit 3 \
+    --gradient_accumulation_steps 1 \
+    --preprocessing_num_workers 1 \
+    --block_size 1024 \
+    --output_dir outputs-pt-bloom-v1 \
+    --overwrite_output_dir \
+    --ddp_timeout 30000 \
+    --logging_first_step True \
+    --target_modules all \
+    --lora_rank 8 \
+    --lora_alpha 16 \
+    --lora_dropout 0.05 \
+    --torch_dtype float16 \
+    --device_map auto \
+    --report_to tensorboard \
+    --ddp_find_unused_parameters False \
+    --gradient_checkpointing True \
+    --cache_dir ./cache

run_rl.sh

@@ -0,0 +1,24 @@
+CUDA_VISIBLE_DEVICES=0,1 torchrun --nproc_per_node 2 rl_training.py \
+    --model_type bloom \
+    --model_name_or_path bigscience/bloomz-560m \
+    --reward_model_name_or_path OpenAssistant/reward-model-deberta-v3-large-v2 \
+    --torch_dtype float16 \
+    --device_map auto \
+    --train_file_dir ./data/finetune \
+    --validation_file_dir ./data/finetune \
+    --batch_size 8 \
+    --max_source_length 256 \
+    --max_target_length 256 \
+    --max_train_samples 1000 \
+    --use_peft True \
+    --lora_rank 8 \
+    --lora_alpha 16 \
+    --lora_dropout 0.05 \
+    --do_train \
+    --max_steps 100 \
+    --learning_rate 1e-5 \
+    --save_steps 50 \
+    --output_dir outputs-rl-bloom-v1 \
+    --early_stopping True \
+    --target_kl 0.1 \
+    --reward_baseline 0.0

run_rm.sh

@@ -0,0 +1,39 @@
+CUDA_VISIBLE_DEVICES=0,1 python reward_modeling.py \
+    --model_type bloom \
+    --model_name_or_path bigscience/bloomz-560m \
+    --train_file_dir ./data/reward \
+    --validation_file_dir ./data/reward \
+    --per_device_train_batch_size 4 \
+    --per_device_eval_batch_size 4 \
+    --do_train \
+    --use_peft True \
+    --seed 42 \
+    --max_train_samples 1000 \
+    --max_eval_samples 10 \
+    --num_train_epochs 1 \
+    --learning_rate 2e-5 \
+    --warmup_ratio 0.05 \
+    --weight_decay 0.001 \
+    --logging_strategy steps \
+    --logging_steps 10 \
+    --eval_steps 50 \
+    --evaluation_strategy steps \
+    --save_steps 500 \
+    --save_strategy steps \
+    --save_total_limit 3 \
+    --max_source_length 256 \
+    --max_target_length 256 \
+    --output_dir outputs-rm-bloom-v1 \
+    --overwrite_output_dir \
+    --ddp_timeout 30000 \
+    --logging_first_step True \
+    --target_modules all \
+    --lora_rank 8 \
+    --lora_alpha 16 \
+    --lora_dropout 0.05 \
+    --torch_dtype float32 \
+    --device_map auto \
+    --report_to tensorboard \
+    --ddp_find_unused_parameters False \
+    --remove_unused_columns False \
+    --gradient_checkpointing True

run_sft.sh

@@ -0,0 +1,40 @@
+CUDA_VISIBLE_DEVICES=0,1 torchrun --nproc_per_node 2 supervised_finetuning.py \
+    --model_type bloom \
+    --model_name_or_path bigscience/bloomz-560m \
+    --train_file_dir ./data/finetune \
+    --validation_file_dir ./data/finetune \
+    --per_device_train_batch_size 4 \
+    --per_device_eval_batch_size 4 \
+    --do_train \
+    --do_eval \
+    --use_peft True \
+    --fp16 \
+    --max_train_samples 1000 \
+    --max_eval_samples 10 \
+    --num_train_epochs 1 \
+    --learning_rate 2e-5 \
+    --warmup_ratio 0.05 \
+    --weight_decay 0.05 \
+    --logging_strategy steps \
+    --logging_steps 10 \
+    --eval_steps 50 \
+    --evaluation_strategy steps \
+    --save_steps 500 \
+    --save_strategy steps \
+    --save_total_limit 3 \
+    --gradient_accumulation_steps 1 \
+    --preprocessing_num_workers 4 \
+    --output_dir outputs-sft-bloom-v1 \
+    --overwrite_output_dir \
+    --ddp_timeout 30000 \
+    --logging_first_step True \
+    --target_modules all \
+    --lora_rank 8 \
+    --lora_alpha 16 \
+    --lora_dropout 0.05 \
+    --torch_dtype float16 \
+    --device_map auto \
+    --report_to tensorboard \
+    --ddp_find_unused_parameters False \
+    --gradient_checkpointing True \
+    --cache_dir ./cache

run_training_dpo_pipeline.ipynb

@@ -0,0 +1,711 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "source": [
+    "# Training Pipeline\n",
+    "[run_training_dpo_pipeline.ipynb](https://github.com/shibing624/MedicalGPT/blob/main/run_training_dpo_pipeline.ipynb)    | [Open In Colab](https://colab.research.google.com/github/shibing624/MedicalGPT/blob/main/run_training_dpo_pipeline.ipynb)"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "tags": []
+   },
+   "source": [
+    "# Stage 1: Continue Pretraining\n",
+    "\n",
+    "第一阶段：PT(Continue PreTraining)增量预训练，在海量领域文本数据上二次预训练GPT模型，以注入领域知识\n",
+    "\n",
+    "| Stage 1: Continue Pretraining   |  [pretraining.py](https://github.com/shibing624/MedicalGPT/blob/main/pretraining.py) | [run_pt.sh](https://github.com/shibing624/MedicalGPT/blob/main/run_pt.sh)    |"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### 说明：\n",
+    "以下 notebook/colab 代码为了快速验证训练代码可用，我们使用了小size的生成模型和小样本数据集，实际使用时，需要使用更大的模型和数据集，以获得更好的效果。\n",
+    "\n",
+    "1. 生成模型：使用的是Bloom的`bigscience/bloomz-560m`\n",
+    "2. 数据集：PT阶段使用的是中文天龙八部小说部分文本和英文书籍部分文本，位于`data/pretrain`文件夹"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "source": [],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 配置运行环境\n",
+    "\n",
+    "本地执行可注释以下配置环境的命令，colab执行要打开注释，用于配置环境\n",
+    "\n",
+    "colab建议使用T4 GPU训练，设置方式：`代码执行程序 -> 更改运行时类型 -> 运行时类型：Python3，硬件加速器：GPU，GPU类型：T4 -> 保存`\n",
+    "\n",
+    "步骤：\n",
+    "1. 下载最新代码到本地\n",
+    "2. 安装依赖包\n",
+    "\n",
+    "依赖包如下，保证最新版本：\n",
+    "\n",
+    "```\n",
+    "loguru\n",
+    "transformers\n",
+    "sentencepiece\n",
+    "datasets\n",
+    "tensorboard\n",
+    "tqdm\n",
+    "peft\n",
+    "trl\n",
+    "```"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!git clone --depth 1 https://github.com/shibing624/MedicalGPT.git\n",
+    "%cd MedicalGPT\n",
+    "%ls\n",
+    "!pip install -r requirements.txt"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Stage1 咱们开始吧\n",
+    "\n",
+    "训练步骤如下：\n",
+    "\n",
+    "1. 确认训练集\n",
+    "2. 执行训练脚本\n",
+    "\n",
+    "训练脚本的执行逻辑如下：\n",
+    "1. 导入依赖包\n",
+    "2. 设置参数\n",
+    "3. 定义各函数并加载训练集\n",
+    "4. 加载模型和tokenizer\n",
+    "5. 开始训练并评估\n",
+    "6. 查看训练结果\n",
+    "\n",
+    "**以下参数可以根据你的GPU实际情况修改，当前参数是根据Colab的T4单卡GPU（16GB显存）配置的**"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%ls ./data/pretrain/"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "!python pretraining.py \\\n",
+    "    --model_type bloom \\\n",
+    "    --model_name_or_path bigscience/bloomz-560m \\\n",
+    "    --train_file_dir ./data/pretrain \\\n",
+    "    --validation_file_dir ./data/pretrain \\\n",
+    "    --per_device_train_batch_size 3 \\\n",
+    "    --per_device_eval_batch_size 3 \\\n",
+    "    --do_train \\\n",
+    "    --do_eval \\\n",
+    "    --use_peft True \\\n",
+    "    --seed 42 \\\n",
+    "    --fp16 \\\n",
+    "    --max_train_samples 10000 \\\n",
+    "    --max_eval_samples 10 \\\n",
+    "    --num_train_epochs 1 \\\n",
+    "    --learning_rate 2e-4 \\\n",
+    "    --warmup_ratio 0.05 \\\n",
+    "    --weight_decay 0.01 \\\n",
+    "    --logging_strategy steps \\\n",
+    "    --logging_steps 10 \\\n",
+    "    --eval_steps 50 \\\n",
+    "    --evaluation_strategy steps \\\n",
+    "    --save_steps 500 \\\n",
+    "    --save_strategy steps \\\n",
+    "    --save_total_limit 3 \\\n",
+    "    --gradient_accumulation_steps 1 \\\n",
+    "    --preprocessing_num_workers 1 \\\n",
+    "    --block_size 1024 \\\n",
+    "    --output_dir outputs-pt-v1 \\\n",
+    "    --overwrite_output_dir \\\n",
+    "    --ddp_timeout 30000 \\\n",
+    "    --logging_first_step True \\\n",
+    "    --target_modules all \\\n",
+    "    --lora_rank 8 \\\n",
+    "    --lora_alpha 16 \\\n",
+    "    --lora_dropout 0.05 \\\n",
+    "    --torch_dtype float16 \\\n",
+    "    --device_map auto \\\n",
+    "    --report_to tensorboard \\\n",
+    "    --ddp_find_unused_parameters False \\\n",
+    "    --gradient_checkpointing True"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%ls -lh outputs-pt-v1"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "模型训练结果：\n",
+    "- 使用lora训练模型，则保存的lora权重是`adapter_model.bin`, lora配置文件是`adapter_config.json`，合并到base model的方法见`merge_peft_adapter.py`\n",
+    "- 日志保存在`output_dir/runs`目录下，可以使用tensorboard查看，启动tensorboard方式如下：`tensorboard --logdir output_dir/runs --host 0.0.0.0 --port 8009`"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "lora模型权重合并到base model，合并后的模型保存在`--output_dir`目录下，合并方法如下："
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "!python merge_peft_adapter.py --model_type bloom \\\n",
+    "    --base_model_name_or_path bigscience/bloomz-560m --peft_model_path outputs-pt-v1 --output_dir merged-pt/"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "%ls -lh merged-pt/"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "%cat merged-pt/config.json"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Stage1 增量预训练完成。"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "ExecuteTime": {
+     "start_time": "2023-06-15T13:56:17.032821Z",
+     "end_time": "2023-06-15T13:56:17.081153Z"
+    }
+   },
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "# Stage 2: Supervised FineTuning\n",
+    "\n",
+    "第二阶段：SFT(Supervised Fine-tuning)有监督微调，构造指令微调数据集，在预训练模型基础上做指令精调，以对齐指令意图\n",
+    "\n",
+    "| Stage 2: Supervised Fine-tuning | [supervised_finetuning.py](https://github.com/shibing624/MedicalGPT/blob/main/supervised_finetuning.py) | [run_sft.sh](https://github.com/shibing624/MedicalGPT/blob/main/run_sft.sh)  |"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "#### 说明：\n",
+    "以下 notebook/colab 代码为了快速验证训练代码可用，我们使用了小size的生成模型和小样本数据集，实际使用时，需要使用更大的模型和数据集，以获得更好的效果。\n",
+    "\n",
+    "1. 生成模型：使用的是Bloom的`bigscience/bloomz-560m` 或者 Stage1得到的预训练模型\n",
+    "2. 数据集：SFT阶段使用的是使用的是Belle的1千条抽样数据，位于`data/finetune`文件夹"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "## Stage2 咱们开始吧\n",
+    "\n",
+    "训练步骤如下：\n",
+    "\n",
+    "1. 确认训练集\n",
+    "2. 执行训练脚本\n",
+    "\n",
+    "训练脚本的执行逻辑如下：\n",
+    "1. 导入依赖包\n",
+    "2. 设置参数\n",
+    "3. 定义各函数并加载训练集\n",
+    "4. 加载模型和tokenizer\n",
+    "5. 开始训练并评估\n",
+    "6. 查看训练结果"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "%ls ./data/finetune"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "ExecuteTime": {
+     "start_time": "2023-06-15T13:58:38.778132Z",
+     "end_time": "2023-06-15T13:58:38.966506Z"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "!python supervised_finetuning.py \\\n",
+    "    --model_type bloom \\\n",
+    "    --model_name_or_path merged-pt \\\n",
+    "    --train_file_dir ./data/finetune \\\n",
+    "    --validation_file_dir ./data/finetune \\\n",
+    "    --per_device_train_batch_size 4 \\\n",
+    "    --per_device_eval_batch_size 4 \\\n",
+    "    --do_train \\\n",
+    "    --do_eval \\\n",
+    "    --use_peft True \\\n",
+    "    --fp16 \\\n",
+    "    --max_train_samples 1000 \\\n",
+    "    --max_eval_samples 10 \\\n",
+    "    --num_train_epochs 1 \\\n",
+    "    --learning_rate 2e-5 \\\n",
+    "    --warmup_ratio 0.05 \\\n",
+    "    --weight_decay 0.05 \\\n",
+    "    --logging_strategy steps \\\n",
+    "    --logging_steps 10 \\\n",
+    "    --eval_steps 50 \\\n",
+    "    --evaluation_strategy steps \\\n",
+    "    --save_steps 500 \\\n",
+    "    --save_strategy steps \\\n",
+    "    --save_total_limit 3 \\\n",
+    "    --gradient_accumulation_steps 1 \\\n",
+    "    --preprocessing_num_workers 1 \\\n",
+    "    --output_dir outputs-sft-v1 \\\n",
+    "    --overwrite_output_dir \\\n",
+    "    --ddp_timeout 30000 \\\n",
+    "    --logging_first_step True \\\n",
+    "    --target_modules all \\\n",
+    "    --lora_rank 8 \\\n",
+    "    --lora_alpha 16 \\\n",
+    "    --lora_dropout 0.05 \\\n",
+    "    --torch_dtype float16 \\\n",
+    "    --device_map auto \\\n",
+    "    --report_to tensorboard \\\n",
+    "    --ddp_find_unused_parameters False \\\n",
+    "    --gradient_checkpointing True"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "%ls -lh outputs-sft-v1"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "模型训练结果：\n",
+    "- 使用lora训练模型，则保存的lora权重是`adapter_model.bin`, lora配置文件是`adapter_config.json`，合并到base model的方法见`merge_peft_adapter.py`\n",
+    "- 日志保存在`output_dir/runs`目录下，可以使用tensorboard查看，启动tensorboard方式如下：`tensorboard --logdir output_dir/runs --host 0.0.0.0 --port 8009`"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "lora模型权重合并到base model，合并后的模型保存在`--output_dir`目录下，合并方法如下："
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "!python merge_peft_adapter.py --model_type bloom \\\n",
+    "    --base_model_name_or_path merged-pt --peft_model_path outputs-sft-v1 --output_dir merged-sft/"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "%ls -lh merged-sft/"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "%cat merged-sft/config.json"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "Stage2 SFT训练完成。"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [],
+   "metadata": {
+    "collapsed": false,
+    "ExecuteTime": {
+     "start_time": "2023-06-15T14:07:40.731186Z",
+     "end_time": "2023-06-15T14:07:40.752635Z"
+    }
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "# Stage 3: DPO(Direct Preference Optimization)\n",
+    "\n",
+    "第三阶段：DPO(Direct Preference Optimization)直接偏好优化，DPO通过直接优化语言模型来实现对其行为的精确控制，而无需使用复杂的强化学习，也可以有效学习到人类偏好，DPO相较于RLHF更容易实现且易于训练，效果更好\n",
+    "\n",
+    "| Stage 3: Direct Preference Optimization        |  [dpo_training.py](https://github.com/shibing624/MedicalGPT/blob/main/dpo_training.py) | [run_dpo.sh](https://github.com/shibing624/MedicalGPT/blob/main/run_dpo.sh)    |"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "#### 说明：\n",
+    "以下 notebook/colab 代码为了快速验证训练代码可用，我们使用了小size的生成模型和小样本数据集，实际使用时，需要使用更大的模型和数据集，以获得更好的效果。\n",
+    "\n",
+    "1. 生成模型：使用的是Bloom的`bigscience/bloomz-560m` 或者 Stage2得到的SFT模型\n",
+    "2. 数据集：DPO阶段使用的是医疗reward数据，抽样了500条，位于`data/reward`文件夹"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "## Stage3 咱们开始吧\n",
+    "\n",
+    "训练步骤如下：\n",
+    "\n",
+    "1. 确认训练集\n",
+    "2. 执行训练脚本\n",
+    "\n",
+    "训练脚本的执行逻辑如下：\n",
+    "1. 导入依赖包\n",
+    "2. 设置参数\n",
+    "3. 定义各函数并加载训练集\n",
+    "4. 加载模型和tokenizer\n",
+    "5. 开始训练并评估\n",
+    "6. 查看训练结果"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "%ls ./data/reward/"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "!python dpo_training.py \\\n",
+    "    --model_type bloom \\\n",
+    "    --model_name_or_path merged-sft \\\n",
+    "    --train_file_dir ./data/reward \\\n",
+    "    --validation_file_dir ./data/reward \\\n",
+    "    --per_device_train_batch_size 3 \\\n",
+    "    --per_device_eval_batch_size 1 \\\n",
+    "    --do_train \\\n",
+    "    --do_eval \\\n",
+    "    --use_peft True \\\n",
+    "    --max_train_samples 1000 \\\n",
+    "    --max_eval_samples 10 \\\n",
+    "    --max_steps 100 \\\n",
+    "    --eval_steps 10 \\\n",
+    "    --save_steps 50 \\\n",
+    "    --max_source_length 128 \\\n",
+    "    --max_target_length 128 \\\n",
+    "    --output_dir outputs-dpo-v1 \\\n",
+    "    --target_modules all \\\n",
+    "    --lora_rank 8 \\\n",
+    "    --lora_alpha 16 \\\n",
+    "    --lora_dropout 0.05 \\\n",
+    "    --torch_dtype float16 \\\n",
+    "    --fp16 True \\\n",
+    "    --device_map auto \\\n",
+    "    --report_to tensorboard \\\n",
+    "    --remove_unused_columns False \\\n",
+    "    --gradient_checkpointing True \\\n",
+    "    --cache_dir ./cache"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "%ls -lh outputs-dpo-v1"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "模型训练结果：\n",
+    "- 使用lora训练模型，则保存的lora权重是`adapter_model.bin`, lora配置文件是`adapter_config.json`，合并到base model的方法见`merge_peft_adapter.py`\n",
+    "- 日志保存在`output_dir/runs`目录下，可以使用tensorboard查看，启动tensorboard方式如下：`tensorboard --logdir output_dir/runs --host 0.0.0.0 --port 8009`"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "lora模型权重合并到base model，合并后的模型保存在`--output_dir`目录下，合并方法如下："
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "!python merge_peft_adapter.py --model_type bloom \\\n",
+    "    --base_model_name_or_path merged-sft --peft_model_path outputs-dpo-v1 --output_dir merged-dpo/"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "%ls -lh merged-dpo/"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "%cat merged-dpo/config.json"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "Stage3 偏好建模第一次训练完成。"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "**至此一个完整的训练流程演示完成。**"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [],
+   "metadata": {
+    "collapsed": false,
+    "ExecuteTime": {
+     "start_time": "2023-06-26T12:34:29.620609Z",
+     "end_time": "2023-06-26T12:34:29.658428Z"
+    }
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "# Test"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "!python inference.py --model_type bloom --base_model merged-dpo --interactive"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "ExecuteTime": {
+     "start_time": "2023-06-26T12:34:47.802087Z",
+     "end_time": "2023-06-26T12:35:00.864463Z"
+    }
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "Input:介绍下南京\n",
+    "Response:  南京市位于江苏省西南部，是全国首批历史文化名城、国家中心城市和自由贸易试验区。\n",
+    "\n",
+    "完。\n"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [],
+   "metadata": {
+    "collapsed": false
+   }
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "name": "python3",
+   "language": "python",
+   "display_name": "Python 3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.8.13"
+  },
+  "vscode": {
+   "interpreter": {
+    "hash": "f34eed0bebedfc4b6ee51ced43d2c030fe3b92f13c149d072205ca200a67b1ec"
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

run_training_pipeline.ipynb

@@ -0,0 +1,917 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "source": [
+    "# Training Pipeline\n",
+    "[run_training_pipeline.ipynb](https://github.com/shibing624/MedicalGPT/blob/main/run_training_pipeline.ipynb)    | [Open In Colab](https://colab.research.google.com/github/shibing624/MedicalGPT/blob/main/run_training_pipeline.ipynb)"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "tags": []
+   },
+   "source": [
+    "# Stage 1: Continue Pretraining\n",
+    "\n",
+    "第一阶段：PT(Continue PreTraining)增量预训练，在海量领域文本数据上二次预训练GPT模型，以注入领域知识\n",
+    "\n",
+    "| Stage 1: Continue Pretraining   |  [pretraining.py](https://github.com/shibing624/MedicalGPT/blob/main/pretraining.py) | [run_pt.sh](https://github.com/shibing624/MedicalGPT/blob/main/run_pt.sh)    |"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### 说明：\n",
+    "以下 notebook/colab 代码为了快速验证训练代码可用，我们使用了小size的生成模型和小样本数据集，实际使用时，需要使用更大的模型和数据集，以获得更好的效果。\n",
+    "\n",
+    "1. 生成模型：使用的是Bloom的`bigscience/bloomz-560m`\n",
+    "2. 数据集：PT阶段使用的是中文天龙八部小说部分文本和英文书籍部分文本，位于`data/pretrain`文件夹"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 配置运行环境\n",
+    "\n",
+    "本地执行可注释以下配置环境的命令，colab执行要打开注释，用于配置环境\n",
+    "\n",
+    "colab建议使用T4 GPU训练，设置方式：`代码执行程序 -> 更改运行时类型 -> 运行时类型：Python3，硬件加速器：GPU，GPU类型：T4 -> 保存`\n",
+    "\n",
+    "步骤：\n",
+    "1. 下载最新代码到本地\n",
+    "2. 安装依赖包\n",
+    "\n",
+    "依赖包如下，保证最新版本：\n",
+    "\n",
+    "```\n",
+    "loguru\n",
+    "transformers\n",
+    "sentencepiece\n",
+    "datasets\n",
+    "tensorboard\n",
+    "tqdm\n",
+    "peft\n",
+    "trl\n",
+    "```"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!git clone --depth 1 https://github.com/shibing624/MedicalGPT.git\n",
+    "%cd MedicalGPT\n",
+    "%ls\n",
+    "!pip install -r requirements.txt"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Stage1 咱们开始吧\n",
+    "\n",
+    "训练步骤如下：\n",
+    "\n",
+    "1. 确认训练集\n",
+    "2. 执行训练脚本\n",
+    "\n",
+    "训练脚本的执行逻辑如下：\n",
+    "1. 导入依赖包\n",
+    "2. 设置参数\n",
+    "3. 定义各函数并加载训练集\n",
+    "4. 加载模型和tokenizer\n",
+    "5. 开始训练并评估\n",
+    "6. 查看训练结果\n",
+    "\n",
+    "**以下参数可以根据你的GPU实际情况修改，当前参数是根据Colab的T4单卡GPU（16GB显存）配置的**"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%ls ./data/pretrain/"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "!python pretraining.py \\\n",
+    "    --model_type bloom \\\n",
+    "    --model_name_or_path bigscience/bloomz-560m \\\n",
+    "    --train_file_dir ./data/pretrain \\\n",
+    "    --validation_file_dir ./data/pretrain \\\n",
+    "    --per_device_train_batch_size 3 \\\n",
+    "    --per_device_eval_batch_size 3 \\\n",
+    "    --do_train \\\n",
+    "    --do_eval \\\n",
+    "    --use_peft True \\\n",
+    "    --seed 42 \\\n",
+    "    --fp16 \\\n",
+    "    --max_train_samples 10000 \\\n",
+    "    --max_eval_samples 10 \\\n",
+    "    --num_train_epochs 1 \\\n",
+    "    --learning_rate 2e-4 \\\n",
+    "    --warmup_ratio 0.05 \\\n",
+    "    --weight_decay 0.01 \\\n",
+    "    --logging_strategy steps \\\n",
+    "    --logging_steps 10 \\\n",
+    "    --eval_steps 50 \\\n",
+    "    --evaluation_strategy steps \\\n",
+    "    --save_steps 500 \\\n",
+    "    --save_strategy steps \\\n",
+    "    --save_total_limit 3 \\\n",
+    "    --gradient_accumulation_steps 1 \\\n",
+    "    --preprocessing_num_workers 1 \\\n",
+    "    --block_size 1024 \\\n",
+    "    --output_dir outputs-pt-v1 \\\n",
+    "    --overwrite_output_dir \\\n",
+    "    --ddp_timeout 30000 \\\n",
+    "    --logging_first_step True \\\n",
+    "    --target_modules all \\\n",
+    "    --lora_rank 8 \\\n",
+    "    --lora_alpha 16 \\\n",
+    "    --lora_dropout 0.05 \\\n",
+    "    --torch_dtype float16 \\\n",
+    "    --device_map auto \\\n",
+    "    --report_to tensorboard \\\n",
+    "    --ddp_find_unused_parameters False \\\n",
+    "    --gradient_checkpointing True"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%ls -lh outputs-pt-v1"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "模型训练结果：\n",
+    "- 使用lora训练模型，则保存的lora权重是`adapter_model.bin`, lora配置文件是`adapter_config.json`，合并到base model的方法见`merge_peft_adapter.py`\n",
+    "- 日志保存在`output_dir/runs`目录下，可以使用tensorboard查看，启动tensorboard方式如下：`tensorboard --logdir output_dir/runs --host 0.0.0.0 --port 8009`"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "lora模型权重合并到base model，合并后的模型保存在`--output_dir`目录下，合并方法如下："
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "!python merge_peft_adapter.py --model_type bloom \\\n",
+    "    --base_model_name_or_path bigscience/bloomz-560m --peft_model_path outputs-pt-v1 --output_dir merged-pt/"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "%ls -lh merged-pt/"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "%cat merged-pt/config.json"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Stage1 增量预训练完成。"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "ExecuteTime": {
+     "start_time": "2023-06-15T13:56:17.032821Z",
+     "end_time": "2023-06-15T13:56:17.081153Z"
+    }
+   },
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "# Stage 2: Supervised FineTuning\n",
+    "\n",
+    "第二阶段：SFT(Supervised Fine-tuning)有监督微调，构造指令微调数据集，在预训练模型基础上做指令精调，以对齐指令意图\n",
+    "\n",
+    "| Stage 2: Supervised Fine-tuning | [supervised_finetuning.py](https://github.com/shibing624/MedicalGPT/blob/main/supervised_finetuning.py) | [run_sft.sh](https://github.com/shibing624/MedicalGPT/blob/main/run_sft.sh)  |"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "#### 说明：\n",
+    "以下 notebook/colab 代码为了快速验证训练代码可用，我们使用了小size的生成模型和小样本数据集，实际使用时，需要使用更大的模型和数据集，以获得更好的效果。\n",
+    "\n",
+    "1. 生成模型：使用的是Bloom的`bigscience/bloomz-560m` 或者 Stage1得到的预训练模型\n",
+    "2. 数据集：SFT阶段使用的是使用的是Belle的1千条抽样数据，位于`data/finetune`文件夹"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "## Stage2 咱们开始吧\n",
+    "\n",
+    "训练步骤如下：\n",
+    "\n",
+    "1. 确认训练集\n",
+    "2. 执行训练脚本\n",
+    "\n",
+    "训练脚本的执行逻辑如下：\n",
+    "1. 导入依赖包\n",
+    "2. 设置参数\n",
+    "3. 定义各函数并加载训练集\n",
+    "4. 加载模型和tokenizer\n",
+    "5. 开始训练并评估\n",
+    "6. 查看训练结果"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "%ls ./data/finetune"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "ExecuteTime": {
+     "start_time": "2023-06-15T13:58:38.778132Z",
+     "end_time": "2023-06-15T13:58:38.966506Z"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "!python supervised_finetuning.py \\\n",
+    "    --model_type bloom \\\n",
+    "    --model_name_or_path merged-pt \\\n",
+    "    --train_file_dir ./data/finetune \\\n",
+    "    --validation_file_dir ./data/finetune \\\n",
+    "    --per_device_train_batch_size 4 \\\n",
+    "    --per_device_eval_batch_size 4 \\\n",
+    "    --do_train \\\n",
+    "    --do_eval \\\n",
+    "    --use_peft True \\\n",
+    "    --fp16 \\\n",
+    "    --max_train_samples 1000 \\\n",
+    "    --max_eval_samples 10 \\\n",
+    "    --num_train_epochs 1 \\\n",
+    "    --learning_rate 2e-5 \\\n",
+    "    --warmup_ratio 0.05 \\\n",
+    "    --weight_decay 0.05 \\\n",
+    "    --logging_strategy steps \\\n",
+    "    --logging_steps 10 \\\n",
+    "    --eval_steps 50 \\\n",
+    "    --evaluation_strategy steps \\\n",
+    "    --save_steps 500 \\\n",
+    "    --save_strategy steps \\\n",
+    "    --save_total_limit 3 \\\n",
+    "    --gradient_accumulation_steps 1 \\\n",
+    "    --preprocessing_num_workers 1 \\\n",
+    "    --output_dir outputs-sft-v1 \\\n",
+    "    --overwrite_output_dir \\\n",
+    "    --ddp_timeout 30000 \\\n",
+    "    --logging_first_step True \\\n",
+    "    --target_modules all \\\n",
+    "    --lora_rank 8 \\\n",
+    "    --lora_alpha 16 \\\n",
+    "    --lora_dropout 0.05 \\\n",
+    "    --torch_dtype float16 \\\n",
+    "    --device_map auto \\\n",
+    "    --report_to tensorboard \\\n",
+    "    --ddp_find_unused_parameters False \\\n",
+    "    --gradient_checkpointing True"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "%ls -lh outputs-sft-v1"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "模型训练结果：\n",
+    "- 使用lora训练模型，则保存的lora权重是`adapter_model.bin`, lora配置文件是`adapter_config.json`，合并到base model的方法见`merge_peft_adapter.py`\n",
+    "- 日志保存在`output_dir/runs`目录下，可以使用tensorboard查看，启动tensorboard方式如下：`tensorboard --logdir output_dir/runs --host 0.0.0.0 --port 8009`"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "lora模型权重合并到base model，合并后的模型保存在`--output_dir`目录下，合并方法如下："
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "!python merge_peft_adapter.py --model_type bloom \\\n",
+    "    --base_model_name_or_path merged-pt --peft_model_path outputs-sft-v1 --output_dir merged-sft/"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "%ls -lh merged-sft/"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "%cat merged-sft/config.json"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "Stage2 SFT训练完成。"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [],
+   "metadata": {
+    "collapsed": false,
+    "ExecuteTime": {
+     "start_time": "2023-06-15T14:07:40.731186Z",
+     "end_time": "2023-06-15T14:07:40.752635Z"
+    }
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "# Stage 3: Reward Modeling\n",
+    "\n",
+    "第三阶段：RM(Reward Model)奖励模型建模，构造人类偏好排序数据集，训练奖励模型，用来对齐人类偏好，主要是\"HHH\"原则，具体是\"helpful, honest, harmless\"\n",
+    "\n",
+    "| Stage 3: Reward Modeling        |  [reward_modeling.py](https://github.com/shibing624/MedicalGPT/blob/main/reward_modeling.py) | [run_rm.sh](https://github.com/shibing624/MedicalGPT/blob/main/run_rm.sh)    |"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "#### 说明：\n",
+    "以下 notebook/colab 代码为了快速验证训练代码可用，我们使用了小size的生成模型和小样本数据集，实际使用时，需要使用更大的模型和数据集，以获得更好的效果。\n",
+    "\n",
+    "1. 生成模型：使用的是Bloom的`bigscience/bloomz-560m` 或者 Stage2得到的SFT模型\n",
+    "2. 数据集：RM阶段使用的是医疗reward数据，抽样了500条，位于`data/reward`文件夹"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "## Stage3 咱们开始吧\n",
+    "\n",
+    "训练步骤如下：\n",
+    "\n",
+    "1. 确认训练集\n",
+    "2. 执行训练脚本\n",
+    "\n",
+    "训练脚本的执行逻辑如下：\n",
+    "1. 导入依赖包\n",
+    "2. 设置参数\n",
+    "3. 定义各函数并加载训练集\n",
+    "4. 加载模型和tokenizer\n",
+    "5. 开始训练并评估\n",
+    "6. 查看训练结果"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "%ls ./data/reward/"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "!python reward_modeling.py \\\n",
+    "    --model_type bloom \\\n",
+    "    --model_name_or_path merged-sft \\\n",
+    "    --train_file_dir ./data/reward \\\n",
+    "    --validation_file_dir ./data/reward \\\n",
+    "    --per_device_train_batch_size 3 \\\n",
+    "    --per_device_eval_batch_size 1 \\\n",
+    "    --do_train \\\n",
+    "    --use_peft True \\\n",
+    "    --seed 42 \\\n",
+    "    --max_train_samples 1000 \\\n",
+    "    --max_eval_samples 10 \\\n",
+    "    --num_train_epochs 1 \\\n",
+    "    --learning_rate 2e-5 \\\n",
+    "    --warmup_ratio 0.05 \\\n",
+    "    --weight_decay 0.001 \\\n",
+    "    --logging_strategy steps \\\n",
+    "    --logging_steps 10 \\\n",
+    "    --eval_steps 50 \\\n",
+    "    --evaluation_strategy steps \\\n",
+    "    --save_steps 500 \\\n",
+    "    --save_strategy steps \\\n",
+    "    --save_total_limit 3 \\\n",
+    "    --max_source_length 256 \\\n",
+    "    --max_target_length 256 \\\n",
+    "    --output_dir outputs-rm-v1 \\\n",
+    "    --overwrite_output_dir \\\n",
+    "    --ddp_timeout 30000 \\\n",
+    "    --logging_first_step True \\\n",
+    "    --target_modules all \\\n",
+    "    --lora_rank 8 \\\n",
+    "    --lora_alpha 16 \\\n",
+    "    --lora_dropout 0.05 \\\n",
+    "    --torch_dtype float32 \\\n",
+    "    --device_map auto \\\n",
+    "    --report_to tensorboard \\\n",
+    "    --ddp_find_unused_parameters False \\\n",
+    "    --remove_unused_columns False \\\n",
+    "    --gradient_checkpointing True"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "%ls -lh outputs-rm-v1"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "模型训练结果：\n",
+    "- 使用lora训练模型，则保存的lora权重是`adapter_model.bin`, lora配置文件是`adapter_config.json`，合并到base model的方法见`merge_peft_adapter.py`\n",
+    "- 日志保存在`output_dir/runs`目录下，可以使用tensorboard查看，启动tensorboard方式如下：`tensorboard --logdir output_dir/runs --host 0.0.0.0 --port 8009`"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "lora模型权重合并到base model，合并后的模型保存在`--output_dir`目录下，合并方法如下："
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "!python merge_peft_adapter.py --model_type bloom \\\n",
+    "    --base_model_name_or_path merged-sft --peft_model_path outputs-rm-v1 --output_dir merged-rm/"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "%ls -lh merged-rm/"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "%cat merged-rm/config.json"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "Stage3 奖励建模第一次训练完成。"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [],
+   "metadata": {
+    "collapsed": false,
+    "ExecuteTime": {
+     "start_time": "2023-06-15T14:12:09.464881Z",
+     "end_time": "2023-06-15T14:12:09.472414Z"
+    }
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "# Stage 4: Reinforcement Learning Training\n",
+    "\n",
+    "第四阶段：RL(Reinforcement Learning)基于人类反馈的强化学习(RLHF)，用奖励模型来训练SFT模型，生成模型使用奖励或惩罚来更新其策略，以便生成更高质量、更符合人类偏好的文本\n",
+    "\n",
+    "| Stage 4: Reinforcement Learning |  [rl_training.py](https://github.com/shibing624/MedicalGPT/blob/main/rl_training.py) | [run_rl.sh](https://github.com/shibing624/MedicalGPT/blob/main/run_rl.sh)    |\n"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "#### 说明：\n",
+    "以下 notebook/colab 代码为了快速验证训练代码可用，我们使用了小size的生成模型、奖励模型和小样本数据集，实际使用时，需要使用更大的模型和数据集，以获得更好的效果。\n",
+    "\n",
+    "1. 生成模型：使用的是Bloom的`bigscience/bloomz-560m` 或者 Stage2得到的SFT模型\n",
+    "2. 奖励模型：使用的是`OpenAssistant/reward-model-deberta-v3-large-v2` 或者 Stage3得到的BERT类或者GPT类奖励模型\n",
+    "3. 数据集：RL阶段的数据可以复用SFT的数据集，使用的是Belle的1千条抽样数据，位于`data/finetune`文件夹"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "## Stage4 咱们开始吧\n",
+    "\n",
+    "训练步骤如下：\n",
+    "\n",
+    "1. 确认训练集\n",
+    "2. 执行训练脚本\n",
+    "\n",
+    "训练脚本的执行逻辑如下：\n",
+    "1. 导入依赖包\n",
+    "2. 设置参数\n",
+    "3. 定义各函数并加载训练集\n",
+    "4. 加载生成模型和tokenizer，加载奖励模型和其tokenizer\n",
+    "5. 开始训练并评估\n",
+    "6. 查看训练结果\n",
+    "\n",
+    "以下参数可以根据你的GPU实际情况修改，当前参数是根据Colab的T4单卡GPU（16GB显存）配置的。"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "%ls ./data/finetune/"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "!python rl_training.py \\\n",
+    "    --model_type bloom \\\n",
+    "    --model_name_or_path merged-sft \\\n",
+    "    --reward_model_name_or_path merged-rm \\\n",
+    "    --torch_dtype float16 \\\n",
+    "    --device_map auto \\\n",
+    "    --train_file_dir ./data/finetune \\\n",
+    "    --validation_file_dir ./data/finetune \\\n",
+    "    --batch_size 4 \\\n",
+    "    --max_source_length 256 \\\n",
+    "    --max_target_length 256 \\\n",
+    "    --max_train_samples 1000 \\\n",
+    "    --use_peft True \\\n",
+    "    --lora_rank 8 \\\n",
+    "    --lora_alpha 16 \\\n",
+    "    --lora_dropout 0.05 \\\n",
+    "    --do_train \\\n",
+    "    --max_steps 64 \\\n",
+    "    --learning_rate 1e-5 \\\n",
+    "    --save_steps 50 \\\n",
+    "    --output_dir outputs-rl-v1 \\\n",
+    "    --early_stopping True \\\n",
+    "    --target_kl 0.1 \\\n",
+    "    --reward_baseline 0.0"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "%ls -lh outputs-rl-v1"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "模型训练结果：\n",
+    "- 使用lora训练模型，则保存的lora权重是`adapter_model.bin`, lora配置文件是`adapter_config.json`，合并到base model的方法见`merge_peft_adapter.py`\n",
+    "- 日志保存在`output_dir/trl`目录下，可以使用tensorboard查看，启动tensorboard方式如下：`tensorboard --logdir output_dir/trl --host 0.0.0.0 --port 8009`"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "lora模型权重合并到base model，合并后的模型保存在`--output_dir`目录下，合并方法如下："
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "!python merge_peft_adapter.py --model_type bloom \\\n",
+    "    --base_model_name_or_path merged-sft --peft_model_path outputs-rl-v1 --output_dir merged-rl/"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "%ls -lh merged-rl/"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "%cat merged-rl/config.json"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "Stage4 RL第一次训练完成。\n",
+    "\n",
+    "**至此一个完整的4阶段训练流程演示完成。**"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "实际操作中Stage3和Stage4可以反复多次，直到RL得到的最后模型满足评估要求。\n",
+    "\n",
+    "RLHF过程可以把SFT模型当成一个初始化模型，RM模型当做指导老师，使用RL(PPO)调教SFT模型生成指导老师最满意的结果，如果小学老师满意了，我们就再训练一个中学老师，继续指导，中学老师满意了，就训练一个大学老师，这样不断迭代，使得生成模型的质量达到甚至超过人工撰写的天花板。\n",
+    "\n",
+    "RLHF训练不易，此项目提供给大家一种实现的方法和参考，希望抛砖引玉，共同促进中文开源LLM发展。"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [],
+   "metadata": {
+    "collapsed": false,
+    "ExecuteTime": {
+     "start_time": "2023-06-26T12:34:29.620609Z",
+     "end_time": "2023-06-26T12:34:29.658428Z"
+    }
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "# Test"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "!python inference.py --model_type bloom --base_model merged-rl --interactive"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "ExecuteTime": {
+     "start_time": "2023-06-26T12:34:47.802087Z",
+     "end_time": "2023-06-26T12:35:00.864463Z"
+    }
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "Input:介绍下南京\n",
+    "Response:  南京市位于江苏省西南部，是全国��批历史文化名城、国家中心城市和自由贸易试验区。\n",
+    "\n",
+    "完。\n"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [],
+   "metadata": {
+    "collapsed": false
+   }
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "name": "python3",
+   "language": "python",
+   "display_name": "Python 3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.8.13"
+  },
+  "vscode": {
+   "interpreter": {
+    "hash": "f34eed0bebedfc4b6ee51ced43d2c030fe3b92f13c149d072205ca200a67b1ec"
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

supervised_finetuning.py

@@ -0,0 +1,927 @@
+# -*- coding: utf-8 -*-
+# Copyright 2023 XuMing([email protected]) and The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Fine-tuning the library models for causal language modeling (GPT, GPT-2, CTRL, ...) on a text file or a dataset.
+
+part of this code is adapted from https://github.com/shibing624/textgen
+"""
+import math
+import os
+from dataclasses import dataclass, field
+from glob import glob
+from typing import List, Optional, Dict, Sequence
+
+import torch
+from datasets import load_dataset
+from loguru import logger
+from peft import LoraConfig, TaskType, get_peft_model, PeftModel, prepare_model_for_int8_training
+from transformers import (
+    AutoConfig,
+    BloomForCausalLM,
+    AutoModel,
+    AutoModelForCausalLM,
+    LlamaTokenizer,
+    LlamaForCausalLM,
+    BloomTokenizerFast,
+    AutoTokenizer,
+    HfArgumentParser,
+    Trainer,
+    TrainingArguments,
+    set_seed,
+    BitsAndBytesConfig,
+    DataCollatorForSeq2Seq,
+)
+from transformers.deepspeed import is_deepspeed_zero3_enabled
+from transformers.trainer import TRAINING_ARGS_NAME
+from transformers.trainer_pt_utils import LabelSmoother
+
+MODEL_CLASSES = {
+    "bloom": (AutoConfig, BloomForCausalLM, BloomTokenizerFast),
+    "chatglm": (AutoConfig, AutoModel, AutoTokenizer),
+    "llama": (AutoConfig, LlamaForCausalLM, LlamaTokenizer),
+    "baichuan": (AutoConfig, AutoModelForCausalLM, AutoTokenizer),
+    "auto": (AutoConfig, AutoModelForCausalLM, AutoTokenizer),
+}
+
+
+@dataclass
+class ModelArguments:
+    """
+    Arguments pertaining to which model/config/tokenizer we are going to fine-tune, or train from scratch.
+    """
+
+    model_type: str = field(
+        default=None,
+        metadata={"help": "Model type selected in the list: " + ", ".join(MODEL_CLASSES.keys())}
+    )
+    model_name_or_path: Optional[str] = field(
+        default=None,
+        metadata={
+            "help": (
+                "The model checkpoint for weights initialization.Don't set if you want to train a model from scratch."
+            )
+        },
+    )
+    tokenizer_name_or_path: Optional[str] = field(
+        default=None,
+        metadata={
+            "help": (
+                "The tokenizer for weights initialization.Don't set if you want to train a model from scratch."
+            )
+        },
+    )
+    load_in_8bit: bool = field(default=False, metadata={"help": "Whether to load the model in 8bit mode or not."})
+    cache_dir: Optional[str] = field(
+        default=None,
+        metadata={"help": "Where do you want to store the pretrained models downloaded from huggingface.co"},
+    )
+    use_fast_tokenizer: bool = field(
+        default=False,
+        metadata={"help": "Whether to use one of the fast tokenizer (backed by the tokenizers library) or not."},
+    )
+    torch_dtype: Optional[str] = field(
+        default="float16",
+        metadata={
+            "help": (
+                "Override the default `torch.dtype` and load the model under this dtype. If `auto` is passed, the "
+                "dtype will be automatically derived from the model's weights."
+            ),
+            "choices": ["auto", "bfloat16", "float16", "float32"],
+        },
+    )
+    device_map: Optional[str] = field(
+        default="auto",
+        metadata={"help": "Device to map model to. If `auto` is passed, the device will be selected automatically. "},
+    )
+    trust_remote_code: bool = field(
+        default=True,
+        metadata={"help": "Whether to trust remote code when loading a model from a remote checkpoint."},
+    )
+
+    def __post_init__(self):
+        if self.model_type is None:
+            raise ValueError(
+                "You must specify a valid model_type to run training. Available model types are " + ", ".join(
+                    MODEL_CLASSES.keys()))
+        if self.model_name_or_path is None:
+            raise ValueError("You must specify a valid model_name_or_path to run training.")
+
+
+@dataclass
+class DataTrainingArguments:
+    """
+    Arguments pertaining to what data we are going to input our model for training and eval.
+    """
+
+    dataset_name: Optional[str] = field(
+        default=None, metadata={"help": "The name of the dataset to use (via the datasets library)."}
+    )
+    dataset_config_name: Optional[str] = field(
+        default=None, metadata={"help": "The configuration name of the dataset to use (via the datasets library)."}
+    )
+    train_file_dir: Optional[str] = field(default=None, metadata={"help": "The train jsonl data file folder."})
+    validation_file_dir: Optional[str] = field(default=None, metadata={"help": "The evaluation jsonl file folder."})
+    template_name: Optional[str] = field(default="vicuna", metadata={"help": "The prompt template name."})
+    max_train_samples: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": (
+                "For debugging purposes or quicker training, truncate the number of training examples to this "
+                "value if set."
+            )
+        },
+    )
+    max_eval_samples: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": (
+                "For debugging purposes or quicker training, truncate the number of evaluation examples to this "
+                "value if set."
+            )
+        },
+    )
+    max_source_length: Optional[int] = field(default=256, metadata={"help": "Max length of prompt input text"})
+    max_target_length: Optional[int] = field(default=256, metadata={"help": "Max length of output text"})
+    ignore_pad_token_for_loss: bool = field(
+        default=True,
+        metadata={"help": "If only pad tokens should be ignored. This assumes that `config.pad_token_id` is defined."},
+    )
+    overwrite_cache: bool = field(
+        default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}
+    )
+    validation_split_percentage: Optional[int] = field(
+        default=1,
+        metadata={
+            "help": "The percentage of the train set used as validation set in case there's no validation split"
+        },
+    )
+    preprocessing_num_workers: Optional[int] = field(
+        default=None,
+        metadata={"help": "The number of processes to use for the preprocessing."},
+    )
+
+    def __post_init__(self):
+        if self.max_train_samples is not None and 0 < self.max_train_samples <= 1000:
+            logger.warning("You may set max_train_samples = -1 to run all samples in production.")
+        if self.max_source_length < 30:
+            raise ValueError("You must specify a valid max_source_length >= 30 to run training.")
+
+
+@dataclass
+class PeftArguments(TrainingArguments):
+    use_peft: bool = field(default=True, metadata={"help": "Whether to use peft"})
+    target_modules: Optional[str] = field(default="all")
+    lora_rank: Optional[int] = field(default=8)
+    lora_dropout: Optional[float] = field(default=0.05)
+    lora_alpha: Optional[float] = field(default=32.0)
+    modules_to_save: Optional[str] = field(default=None)
+    peft_path: Optional[str] = field(default=None, metadata={"help": "The path to the peft model"})
+    qlora: bool = field(default=False, metadata={"help": "Whether to use qlora"})
+
+
+class CastOutputToFloat(torch.nn.Sequential):
+    """Cast the output of the model to float"""
+
+    def forward(self, x):
+        return super().forward(x).to(torch.float32)
+
+
+@dataclass
+class Conversation:
+    """A class that manages prompt templates and keeps all conversation history."""
+
+    # The name of this template
+    name: str
+    # The system prompt
+    system_prompt: str
+    # All messages. format: list of [question, answer]
+    messages: Optional[List[Sequence[str]]]
+    # The roles of the speakers
+    roles: Optional[Sequence[str]]
+    # Conversation prompt
+    prompt: str
+    # Separator
+    sep: str
+    # Stop token, default is tokenizer.eos_token
+    stop_str: Optional[str] = "</s>"
+
+    def get_prompt(
+            self,
+            messages: Optional[List[Sequence[str]]] = None,
+            system_prompt: Optional[str] = ""
+    ) -> str:
+        """
+        Returns a string containing prompt without response.
+        """
+        return "".join(self._format_example(messages, system_prompt))
+
+    def get_dialog(
+            self,
+            messages: Optional[List[Sequence[str]]] = None,
+            system_prompt: Optional[str] = ""
+    ) -> List[str]:
+        """
+        Returns a list containing 2 * n elements where the 2k-th is a query and the (2k+1)-th is a response.
+        """
+        return self._format_example(messages, system_prompt)
+
+    def _format_example(
+            self,
+            messages: Optional[List[Sequence[str]]] = None,
+            system_prompt: Optional[str] = ""
+    ) -> List[str]:
+        system_prompt = system_prompt or self.system_prompt
+        system_prompt = system_prompt + self.sep if system_prompt else ""  # add separator for non-empty system prompt
+        messages = messages or self.messages
+        convs = []
+        for turn_idx, [user_query, bot_resp] in enumerate(messages):
+            if turn_idx == 0:
+                convs.append(system_prompt + self.prompt.format(query=user_query))
+                convs.append(bot_resp)
+            else:
+                convs.append(self.sep + self.prompt.format(query=user_query))
+                convs.append(bot_resp)
+        return convs
+
+    def append_message(self, query: str, answer: str):
+        """Append a new message."""
+        self.messages.append([query, answer])
+
+
+# A global registry for all conversation templates
+conv_templates: Dict[str, Conversation] = {}
+
+
+def register_conv_template(template: Conversation):
+    """Register a new conversation template."""
+    conv_templates[template.name] = template
+
+
+"""Vicuna v1.1 template
+Supports: https://huggingface.co/lmsys/vicuna-7b-delta-v1.1
+          https://huggingface.co/lmsys/vicuna-13b-delta-v1.1
+"""
+register_conv_template(
+    Conversation(
+        name="vicuna",
+        system_prompt="A chat between a curious user and an artificial intelligence assistant. "
+                      "The assistant gives helpful, detailed, and polite answers to the user's questions.",
+        messages=[],
+        roles=("USER", "ASSISTANT"),
+        prompt="USER: {query} ASSISTANT: ",
+        sep="</s>",
+    )
+)
+
+"""Alpaca template"""
+register_conv_template(
+    Conversation(
+        name="alpaca",
+        system_prompt="Below is an instruction that describes a task. "
+                      "Write a response that appropriately completes the request.",
+        messages=[],
+        roles=("### Instruction", "### Response"),
+        prompt="### Instruction:\n{query}\n\n### Response:\n",
+        sep="\n\n",
+    )
+)
+
+"""Baichuan-13B-Chat template
+source: https://huggingface.co/baichuan-inc/Baichuan-13B-Chat/blob/f5f47be2adbbdceb784f334d6fa1ca2c73e65097/modeling_baichuan.py#L507
+Support: https://huggingface.co/baichuan-inc/Baichuan-13B-Chat
+"""
+register_conv_template(
+    Conversation(
+        name="baichuan-chat",
+        system_prompt="",
+        messages=[],
+        roles=("<reserved_102>", "<reserved_103>"),
+        prompt=" <reserved_102> {query} <reserved_103> ",
+        sep="</s>",
+    )
+)
+
+"""ziya template"""
+register_conv_template(
+    Conversation(
+        name="ziya",
+        system_prompt="",
+        messages=[],
+        roles=("<human>", "<bot>"),
+        prompt="<human>:{query}\n<bot>:",
+        sep="\n",
+    )
+)
+
+"""Linly template"""
+register_conv_template(
+    Conversation(
+        name="linly",
+        system_prompt="",
+        messages=[],
+        roles=("User", "Bot"),
+        prompt="User: {query}\nBot: ",
+        sep="\n",
+    )
+)
+
+"""ChatGLM1 template
+source: https://huggingface.co/THUDM/chatglm-6b/blob/main/modeling_chatglm.py#L1307
+"""
+register_conv_template(
+    Conversation(
+        name="chatglm",
+        system_prompt="",
+        messages=[],
+        roles=("问", "答"),
+        prompt="问：{query}\n答：",
+        sep="\n",
+    )
+)
+
+"""ChatGLM2 template
+source: https://huggingface.co/THUDM/chatglm2-6b/blob/main/modeling_chatglm.py#L1007
+"""
+register_conv_template(
+    # source:
+    Conversation(
+        name="chatglm2",
+        system_prompt="",
+        messages=[],
+        roles=("问", "答"),
+        prompt="问：{query}\n\n答：",
+        sep="\n\n",
+    )
+)
+
+"""Phoenix template"""
+register_conv_template(
+    Conversation(
+        name="phoenix",
+        system_prompt="A chat between a curious human and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the human's questions.\n\n",
+        messages=[],
+        roles=("Human", "Assistant"),
+        prompt="Human: <s>{query}</s>Assistant: ",
+        sep="</s>",
+    )
+)
+
+"""belle template
+Supports: https://huggingface.co/BelleGroup/BELLE-LLaMA-EXT-13B
+"""
+register_conv_template(
+    Conversation(
+        name="belle",
+        system_prompt="",
+        messages=[],
+        roles=("Human", "Belle"),
+        prompt="Human: {query}\n\nBelle: ",
+        sep="\n\n",
+    )
+)
+
+"""aquila template
+Supports: https://huggingface.co/qhduan/aquilachat-7b
+"""
+register_conv_template(
+    Conversation(
+        name="aquila",
+        system_prompt="A chat between a curious human and an artificial intelligence assistant. "
+                      "The assistant gives helpful, detailed, and polite answers to the human's questions.",
+        messages=[],
+        roles=("Human", "Assistant"),
+        prompt="Human: {query}###Assistant: ",
+        sep="###",
+    )
+)
+
+"""intern template
+Supports: https://huggingface.co/internlm/internlm-chat-7b
+"""
+register_conv_template(
+    Conversation(
+        name="intern",
+        system_prompt="",
+        messages=[],
+        roles=("<|User|>", "<|Bot|>"),
+        prompt="<|User|>:{query}<eoh>\n<|Bot|>:",
+        sep="<eoa>\n",
+        stop_str="<eoa>",
+    )
+)
+
+"""StarChat template"""
+register_conv_template(
+    Conversation(
+        name="starchat",
+        system_prompt="<system>\n",
+        messages=[],
+        roles=("<|user|>", "<|assistant|>"),
+        prompt="<|user|>\n{query}<|end|>\n<|assistant|>\n",
+        sep="<|end|>\n",
+        stop_str="<|end|>",
+    )
+)
+
+"""llama2 template
+reference: https://github.com/facebookresearch/llama/blob/cfc3fc8c1968d390eb830e65c63865e980873a06/llama/generation.py#L212
+"""
+register_conv_template(
+    Conversation(
+        name="llama2",
+        system_prompt="<<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\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",
+        messages=[],
+        roles=("[INST]", "[/INST]"),
+        prompt=" [INST] {query} [/INST] ",
+        sep="</s>",
+    )
+)
+
+"""llama2-zh template
+Sources: https://github.com/ymcui/Chinese-LLaMA-Alpaca-2
+Supports: https://huggingface.co/ziqingyang/chinese-alpaca-2-7b
+"""
+register_conv_template(
+    Conversation(
+        name="llama2-zh",
+        system_prompt="<<SYS>>\nYou are a helpful assistant. 你是一个乐于助人的助手。\n<</SYS>>\n\n",
+        messages=[],
+        roles=("[INST]", "[/INST]"),
+        prompt=" [INST] {query} [/INST] ",
+        sep="</s>",
+    )
+)
+"""XVERSE template
+Supports: https://huggingface.co/xverse/XVERSE-13B-Chat
+"""
+register_conv_template(
+    Conversation(
+        name="xverse",
+        system_prompt="",
+        messages=[],
+        roles=("Human", "Assistant"),
+        prompt="Human: {query}\n\nAssistant: ",
+        sep="</s>",
+    )
+)
+
+"""Qwen template
+Supports: https://huggingface.co/Qwen/Qwen-7B-Chat
+chatml: https://xbot123.com/645a461b922f176d7cfdbc2d/
+"""
+register_conv_template(
+    Conversation(
+        name="chatml",
+        system_prompt="You are a helpful assistant.",
+        messages=[],
+        roles=("user", "assistant"),
+        prompt="<|im_start|>user\n{query}<|im_end|>\n<|im_start|>assistant\n",
+        sep="<|im_end|>\n",
+        stop_str="<|im_end|>",
+    )
+)
+
+
+def get_conv_template(name: str) -> Conversation:
+    """Get a conversation template."""
+    return conv_templates[name]
+
+
+class SavePeftModelTrainer(Trainer):
+    """
+    Trainer for lora models
+    """
+
+    def save_model(self, output_dir=None, _internal_call=False):
+        """Save the LoRA model."""
+        os.makedirs(output_dir, exist_ok=True)
+        if self.args.local_rank in [-1, 0]:
+            torch.save(self.args, os.path.join(output_dir, TRAINING_ARGS_NAME))
+            self.model.save_pretrained(output_dir)
+
+
+def save_model(output_dir, model, tokenizer, args):
+    """Save the model and the tokenizer."""
+    os.makedirs(output_dir, exist_ok=True)
+
+    # Take care of distributed/parallel training
+    model_to_save = model.module if hasattr(model, "module") else model
+    if args.local_rank in [-1, 0]:
+        model_to_save.save_pretrained(output_dir)
+        tokenizer.save_pretrained(output_dir)
+        torch.save(args, os.path.join(output_dir, TRAINING_ARGS_NAME))
+
+
+def print_trainable_parameters(model):
+    """
+    Prints the number of trainable parameters in the model.
+    """
+    trainable_params = 0
+    all_param = 0
+    for _, param in model.named_parameters():
+        all_param += param.numel()
+        if param.requires_grad:
+            trainable_params += param.numel()
+    print(
+        f"trainable params: {trainable_params} || all params: {all_param} || trainable%: {100 * trainable_params / all_param}"
+    )
+
+
+def find_all_linear_names(peft_model, int4=False, int8=False):
+    """Find all linear layer names in the model. reference from qlora paper."""
+    cls = torch.nn.Linear
+    if int4 or int8:
+        import bitsandbytes as bnb
+        if int4:
+            cls = bnb.nn.Linear4bit
+        elif int8:
+            cls = bnb.nn.Linear8bitLt
+    lora_module_names = set()
+    for name, module in peft_model.named_modules():
+        if isinstance(module, cls):
+            # last layer is not add to lora_module_names
+            if 'lm_head' in name:
+                continue
+            if 'output_layer' in name:
+                continue
+            names = name.split('.')
+            lora_module_names.add(names[0] if len(names) == 1 else names[-1])
+    return sorted(lora_module_names)
+
+
+def main():
+    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, PeftArguments))
+    model_args, data_args, training_args = parser.parse_args_into_dataclasses()
+
+    logger.info(f"Model args: {model_args}")
+    logger.info(f"Data args: {data_args}")
+    logger.info(f"Training args: {training_args}")
+    logger.info(
+        f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}"
+        + f" distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}"
+    )
+
+    # Set seed before initializing model.
+    set_seed(training_args.seed)
+
+    if not model_args.model_type:
+        raise ValueError("Please specify a model_type, e.g. llama, chatglm, bloom, etc.")
+    config_class, model_class, tokenizer_class = MODEL_CLASSES[model_args.model_type]
+
+    # Load tokenizer
+    tokenizer_kwargs = {
+        "cache_dir": model_args.cache_dir,
+        "use_fast": model_args.use_fast_tokenizer,
+        "trust_remote_code": model_args.trust_remote_code,
+    }
+    tokenizer_name_or_path = model_args.tokenizer_name_or_path
+    if not tokenizer_name_or_path:
+        tokenizer_name_or_path = model_args.model_name_or_path
+    tokenizer = tokenizer_class.from_pretrained(tokenizer_name_or_path, **tokenizer_kwargs)
+    prompt_template = get_conv_template(data_args.template_name)
+    if tokenizer.eos_token_id is None:
+        tokenizer.eos_token = prompt_template.stop_str  # eos token is required for SFT
+        logger.info("Add eos token: {}".format(tokenizer.eos_token))
+    if tokenizer.pad_token_id is None:
+        if tokenizer.unk_token_id is not None:
+            tokenizer.pad_token = tokenizer.unk_token
+        else:
+            tokenizer.pad_token = tokenizer.eos_token
+        logger.info("Add pad token: {}".format(tokenizer.pad_token))
+
+    logger.debug(f"Tokenizer: {tokenizer}")
+    IGNORE_INDEX = LabelSmoother.ignore_index if data_args.ignore_pad_token_for_loss else tokenizer.pad_token_id
+
+    # Get datasets
+    if data_args.dataset_name is not None:
+        # Downloading and loading a dataset from the hub.
+        raw_datasets = load_dataset(
+            data_args.dataset_name,
+            data_args.dataset_config_name,
+            cache_dir=model_args.cache_dir,
+        )
+        if "validation" not in raw_datasets.keys():
+            raw_datasets["validation"] = load_dataset(
+                data_args.dataset_name,
+                data_args.dataset_config_name,
+                split=f"train[:{data_args.validation_split_percentage}%]",
+                cache_dir=model_args.cache_dir,
+            )
+            raw_datasets["train"] = load_dataset(
+                data_args.dataset_name,
+                data_args.dataset_config_name,
+                split=f"train[{data_args.validation_split_percentage}%:]",
+                cache_dir=model_args.cache_dir,
+            )
+    else:
+        # Loading a dataset from local files.
+        data_files = {}
+        if data_args.train_file_dir is not None and os.path.exists(data_args.train_file_dir):
+            train_data_files = glob(f'{data_args.train_file_dir}/**/*.json', recursive=True) + glob(
+                f'{data_args.train_file_dir}/**/*.jsonl', recursive=True)
+            logger.info(f"train files: {train_data_files}")
+            data_files["train"] = train_data_files
+        if data_args.validation_file_dir is not None and os.path.exists(data_args.validation_file_dir):
+            eval_data_files = glob(f'{data_args.validation_file_dir}/**/*.json', recursive=True) + glob(
+                f'{data_args.validation_file_dir}/**/*.jsonl', recursive=True)
+            logger.info(f"eval files: {eval_data_files}")
+            data_files["validation"] = eval_data_files
+        raw_datasets = load_dataset(
+            'json',
+            data_files=data_files,
+            cache_dir=model_args.cache_dir,
+        )
+        # If no validation data is there, validation_split_percentage will be used to divide the dataset.
+        if "validation" not in raw_datasets.keys():
+            raw_datasets["validation"] = load_dataset(
+                'json',
+                data_files=data_files,
+                split=f"train[:{data_args.validation_split_percentage}%]",
+                cache_dir=model_args.cache_dir,
+            )
+            raw_datasets["train"] = load_dataset(
+                'json',
+                data_files=data_files,
+                split=f"train[{data_args.validation_split_percentage}%:]",
+                cache_dir=model_args.cache_dir,
+            )
+    logger.info(f"Raw datasets: {raw_datasets}")
+
+    # Preprocessing the datasets
+    max_source_length = data_args.max_source_length
+    max_target_length = data_args.max_target_length
+    max_length = max_source_length + max_target_length
+
+    def preprocess_function(examples):
+        """
+        Preprocessing the datasets.
+            part of code modified from https://github.com/lm-sys/FastChat
+        """
+        input_ids_list = []
+        targets_list = []
+        roles = ["human", "gpt"]
+
+        def get_dialog(examples):
+            for i, source in enumerate(examples['conversations']):
+                if len(source) < 2:
+                    continue
+                data_role = source[0].get("from", "")
+                if data_role not in roles or data_role != roles[0]:
+                    # Skip the first one if it is not from human
+                    source = source[1:]
+                if len(source) < 2:
+                    continue
+                messages = []
+                for j, sentence in enumerate(source):
+                    data_role = sentence.get("from", "")
+                    if data_role not in roles:
+                        logger.warning(f"unknown role: {data_role}, {i}. (ignored)")
+                        break
+                    if data_role == roles[j % 2]:
+                        messages.append(sentence["value"])
+                if len(messages) < 2 or len(messages) % 2 != 0:
+                    continue
+                # Convert the list to pairs of elements
+                history_messages = [[messages[k], messages[k + 1]] for k in range(0, len(messages), 2)]
+                yield prompt_template.get_dialog(history_messages)
+
+        for dialog in get_dialog(examples):
+            input_ids, labels = [], []
+
+            for i in range(len(dialog) // 2):
+                source_ids = tokenizer.encode(text=dialog[2 * i], add_special_tokens=(i == 0))
+                target_ids = tokenizer.encode(text=dialog[2 * i + 1], add_special_tokens=False)
+
+                if len(source_ids) > max_source_length:
+                    source_ids = source_ids[:max_source_length]
+                if len(target_ids) > max_target_length - 1:  # eos token
+                    target_ids = target_ids[:max_target_length - 1]
+                if len(source_ids) > 0 and source_ids[0] == tokenizer.eos_token_id:
+                    source_ids = source_ids[1:]
+                if len(target_ids) > 0 and target_ids[-1] == tokenizer.eos_token_id:
+                    target_ids = target_ids[:-1]
+                if len(input_ids) + len(source_ids) + len(target_ids) + 1 > max_length:
+                    break
+
+                input_ids += source_ids + target_ids + [tokenizer.eos_token_id]  # add eos token for each turn
+                labels += [IGNORE_INDEX] * len(source_ids) + target_ids + [tokenizer.eos_token_id]
+
+            input_ids_list.append(input_ids)
+            targets_list.append(labels)
+
+        return dict(
+            input_ids=input_ids_list,
+            labels=targets_list,
+        )
+
+    def filter_empty_labels(example):
+        """Remove empty labels dataset."""
+        return not all(label == IGNORE_INDEX for label in example["labels"])
+
+    train_dataset = None
+    max_train_samples = 0
+    if training_args.do_train:
+        if "train" not in raw_datasets:
+            raise ValueError("--do_train requires a train dataset")
+        train_dataset = raw_datasets['train']
+        max_train_samples = len(train_dataset)
+        if data_args.max_train_samples is not None and data_args.max_train_samples > 0:
+            max_train_samples = min(len(train_dataset), data_args.max_train_samples)
+            train_dataset = train_dataset.select(range(max_train_samples))
+        logger.debug(f"Example train_dataset[0]: {train_dataset[0]}")
+        with training_args.main_process_first(desc="Train dataset tokenization"):
+            train_dataset = train_dataset.shuffle().map(
+                preprocess_function,
+                batched=True,
+                num_proc=data_args.preprocessing_num_workers,
+                remove_columns=train_dataset.column_names,
+                load_from_cache_file=not data_args.overwrite_cache,
+                desc="Running tokenizer on dataset",
+            )
+            train_dataset = train_dataset.filter(filter_empty_labels, num_proc=data_args.preprocessing_num_workers)
+            logger.debug(f"Num train_samples: {len(train_dataset)}")
+            logger.debug("Tokenized training example:")
+            logger.debug(f"Decode input_ids[0]: {tokenizer.decode(train_dataset[0]['input_ids'])}")
+            replaced_labels = [label if label != IGNORE_INDEX else tokenizer.pad_token_id
+                               for label in list(train_dataset[0]['labels'])]
+            logger.debug(f"Decode labels[0]: {tokenizer.decode(replaced_labels)}")
+
+    eval_dataset = None
+    max_eval_samples = 0
+    if training_args.do_eval:
+        with training_args.main_process_first(desc="Eval dataset tokenization"):
+            if "validation" not in raw_datasets:
+                raise ValueError("--do_eval requires a validation dataset")
+            eval_dataset = raw_datasets["validation"]
+            max_eval_samples = len(eval_dataset)
+            if data_args.max_eval_samples is not None and data_args.max_eval_samples > 0:
+                max_eval_samples = min(len(eval_dataset), data_args.max_eval_samples)
+                eval_dataset = eval_dataset.select(range(max_eval_samples))
+            logger.debug(f"Example eval_dataset[0]: {eval_dataset[0]}")
+            eval_dataset = eval_dataset.map(
+                preprocess_function,
+                batched=True,
+                num_proc=data_args.preprocessing_num_workers,
+                remove_columns=eval_dataset.column_names,
+                load_from_cache_file=not data_args.overwrite_cache,
+                desc="Running tokenizer on dataset",
+            )
+            eval_dataset = eval_dataset.filter(filter_empty_labels, num_proc=data_args.preprocessing_num_workers)
+            logger.debug(f"Num eval_samples: {len(eval_dataset)}")
+            logger.debug("Tokenized eval example:")
+            logger.debug(tokenizer.decode(eval_dataset[0]['input_ids']))
+
+    # Load model
+    if model_args.model_name_or_path:
+        torch_dtype = (
+            model_args.torch_dtype
+            if model_args.torch_dtype in ["auto", None]
+            else getattr(torch, model_args.torch_dtype)
+        )
+        world_size = int(os.environ.get("WORLD_SIZE", 1))
+        ddp = world_size != 1
+        if ddp:
+            model_args.device_map = {"": int(os.environ["LOCAL_RANK"]) or 0}
+        if training_args.qlora and (len(training_args.fsdp) > 0 or is_deepspeed_zero3_enabled()):
+            logger.warning("FSDP and ZeRO3 are both currently incompatible with QLoRA.")
+        config = config_class.from_pretrained(
+            model_args.model_name_or_path,
+            trust_remote_code=model_args.trust_remote_code,
+            torch_dtype=torch_dtype,
+            cache_dir=model_args.cache_dir
+        )
+        model = model_class.from_pretrained(
+            model_args.model_name_or_path,
+            config=config,
+            load_in_8bit=model_args.load_in_8bit,
+            low_cpu_mem_usage=(not is_deepspeed_zero3_enabled()),
+            device_map=model_args.device_map,
+            trust_remote_code=model_args.trust_remote_code,
+            quantization_config=BitsAndBytesConfig(
+                load_in_4bit=True,
+                bnb_4bit_use_double_quant=True,
+                bnb_4bit_quant_type="nf4",
+                bnb_4bit_compute_dtype=torch_dtype,
+            ) if training_args.qlora else None,
+        )
+        if hasattr(model, 'lm_head'):
+            model.lm_head = CastOutputToFloat(model.lm_head)
+    else:
+        raise ValueError(f"Error, model_name_or_path is None, SFT must be loaded from a pre-trained model")
+
+    if training_args.use_peft:
+        logger.info("Fine-tuning method: LoRA(PEFT)")
+        if training_args.peft_path is not None:
+            logger.info(f"Peft from pre-trained model: {training_args.peft_path}")
+            model = PeftModel.from_pretrained(model, training_args.peft_path, is_trainable=True)
+        else:
+            target_modules = training_args.target_modules.split(',') if training_args.target_modules else None
+            if target_modules and 'all' in target_modules:
+                target_modules = find_all_linear_names(model, int4=False, int8=model_args.load_in_8bit)
+            modules_to_save = training_args.modules_to_save
+            if modules_to_save is not None:
+                modules_to_save = modules_to_save.split(',')
+            logger.info(f"Peft target_modules: {target_modules}")
+            logger.info(f"Peft lora_rank: {training_args.lora_rank}")
+            peft_config = LoraConfig(
+                task_type=TaskType.CAUSAL_LM,
+                target_modules=target_modules,
+                inference_mode=False,
+                r=training_args.lora_rank,
+                lora_alpha=training_args.lora_alpha,
+                lora_dropout=training_args.lora_dropout,
+                modules_to_save=modules_to_save)
+            model = get_peft_model(model, peft_config)
+        if model_args.load_in_8bit:
+            model = prepare_model_for_int8_training(model)
+        model.print_trainable_parameters()
+    else:
+        logger.info("Fine-tuning method: Full parameters training")
+        model = model.float()
+        print_trainable_parameters(model)
+    logger.debug(f"Model: {model}")
+
+    # Initialize our Trainer
+    if training_args.gradient_checkpointing:
+        model.gradient_checkpointing_enable()
+        model.config.use_cache = False
+    else:
+        model.config.use_cache = True
+    model.enable_input_require_grads()
+    if not ddp and torch.cuda.device_count() > 1:
+        # Keeps Trainer from trying its own DataParallelism when more than 1 gpu is available
+        model.is_parallelizable = True
+        model.model_parallel = True
+
+    data_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer, label_pad_token_id=IGNORE_INDEX)
+    # Initialize our Trainer
+    trainer = SavePeftModelTrainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_dataset if training_args.do_train else None,
+        eval_dataset=eval_dataset if training_args.do_eval else None,
+        tokenizer=tokenizer,
+        data_collator=data_collator,
+    )
+
+    # Training
+    if training_args.do_train:
+        logger.info("*** Train ***")
+        sample = next(iter(trainer.get_train_dataloader()))
+        logger.debug(f"Train dataloader example: {sample}")
+        logger.debug(f"Detail input_ids: {list(sample['input_ids'])[:3]}, \nlabels: {list(sample['labels'])[:3]}")
+        logger.debug(f"Decode input_ids[0]: {tokenizer.decode(sample['input_ids'][0])}")
+        replaced_labels = [label if label != IGNORE_INDEX else tokenizer.pad_token_id for label in sample['labels'][0]]
+        logger.debug(f"Decode labels[0]: {tokenizer.decode(replaced_labels)}")
+        checkpoint = None
+        if training_args.resume_from_checkpoint is not None:
+            checkpoint = training_args.resume_from_checkpoint
+        train_result = trainer.train(resume_from_checkpoint=checkpoint)
+
+        metrics = train_result.metrics
+        metrics["train_samples"] = max_train_samples
+        logger.debug(f"Training metrics: {metrics}")
+        trainer.log_metrics("train", metrics)
+        trainer.save_metrics("train", metrics)
+        model.config.use_cache = True  # enable cache after training
+        trainer.save_state()
+        logger.info(f"Saving model checkpoint to {training_args.output_dir}")
+        save_model(training_args.output_dir, model, tokenizer, training_args)
+
+    # Evaluation
+    if training_args.do_eval and trainer.is_world_process_zero():
+        logger.info("*** Evaluate ***")
+        metrics = trainer.evaluate()
+
+        metrics["eval_samples"] = max_eval_samples
+        try:
+            perplexity = math.exp(metrics["eval_loss"])
+        except OverflowError:
+            perplexity = float("inf")
+        metrics["perplexity"] = perplexity
+        logger.debug(f"Eval metrics: {metrics}")
+        trainer.log_metrics("eval", metrics)
+        trainer.save_metrics("eval", metrics)
+
+
+if __name__ == "__main__":
+    main()