Upload 42 files

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+MedicalGPT-main/docs/logo.png filter=lfs diff=lfs merge=lfs -text

MedicalGPT-main/.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/

MedicalGPT-main/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

MedicalGPT-main/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.

MedicalGPT-main/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.

MedicalGPT-main/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.

MedicalGPT-main/README.md

@@ -0,0 +1,326 @@
+[**🇨🇳中文**](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!

MedicalGPT-main/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!

MedicalGPT-main/_config.yml

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

MedicalGPT-main/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()

MedicalGPT-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)

MedicalGPT-main/data/pretrain/en_article_tail500.txt

@@ -0,0 +1,500 @@
+contract to work in specified mines and mills. There seemed to be no
+limit to the factories, forges, refineries, and railways that could be
+built, to the multitudes that could be employed in conquering a
+continent. As for the future, that was in the hands of Providence!
+
+=Business Theories of Politics.=--As the statesmen of Hamilton's school
+and the planters of Calhoun's had their theories of government and
+politics, so the leaders in business enterprise had theirs. It was
+simple and easily stated. "It is the duty of the government," they
+urged, "to protect American industry against foreign competition by
+means of high tariffs on imported goods, to aid railways by generous
+grants of land, to sell mineral and timber lands at low prices to
+energetic men ready to develop them, and then to leave the rest to the
+initiative and drive of individuals and companies." All government
+interference with the management, prices, rates, charges, and conduct of
+private business they held to be either wholly pernicious or intolerably
+impertinent. Judging from their speeches and writings, they conceived
+the nation as a great collection of individuals, companies, and labor
+unions all struggling for profits or high wages and held together by a
+government whose principal duty was to keep the peace among them and
+protect industry against the foreign manufacturer. Such was the
+political theory of business during the generation that followed the
+Civil War.
+
+
+THE SUPREMACY OF THE REPUBLICAN PARTY (1861-85)
+
+=Business Men and Republican Policies.=--Most of the leaders in industry
+gravitated to the Republican ranks. They worked in the North and the
+Republican party was essentially Northern. It was moreover--at least so
+far as the majority of its members were concerned--committed to
+protective tariffs, a sound monetary and banking system, the promotion
+of railways and industry by land grants, and the development of internal
+improvements. It was furthermore generous in its immigration policy. It
+proclaimed America to be an asylum for the oppressed of all countries
+and flung wide the doors for immigrants eager to fill the factories, man
+the mines, and settle upon Western lands. In a word the Republicans
+stood for all those specific measures which favored the enlargement and
+prosperity of business. At the same time they resisted government
+interference with private enterprise. They did not regulate railway
+rates, prosecute trusts for forming combinations, or prevent railway
+companies from giving lower rates to some shippers than to others. To
+sum it up, the political theories of the Republican party for three
+decades after the Civil War were the theories of American
+business--prosperous and profitable industries for the owners and "the
+full dinner pail" for the workmen. Naturally a large portion of those
+who flourished under its policies gave their support to it, voted for
+its candidates, and subscribed to its campaign funds.
+
+=Sources of Republican Strength in the North.=--The Republican party was
+in fact a political organization of singular power. It originated in a
+wave of moral enthusiasm, having attracted to itself, if not the
+abolitionists, certainly all those idealists, like James Russell Lowell
+and George William Curtis, who had opposed slavery when opposition was
+neither safe nor popular. To moral principles it added practical
+considerations. Business men had confidence in it. Workingmen, who
+longed for the independence of the farmer, owed to its indulgent land
+policy the opportunity of securing free homesteads in the West. The
+immigrant, landing penniless on these shores, as a result of the same
+beneficent system, often found himself in a little while with an estate
+as large as many a baronial domain in the Old World. Under a Republican
+administration, the union had been saved. To it the veterans of the war
+could turn with confidence for those rewards of service which the
+government could bestow: pensions surpassing in liberality anything that
+the world had ever seen. Under a Republican administration also the
+great debt had been created in the defense of the union, and to the
+Republican party every investor in government bonds could look for the
+full and honorable discharge of the interest and principal. The spoils
+system, inaugurated by Jacksonian Democracy, in turn placed all the
+federal offices in Republican hands, furnishing an army of party workers
+to be counted on for loyal service in every campaign.
+
+Of all these things Republican leaders made full and vigorous use,
+sometimes ascribing to the party, in accordance with ancient political
+usage, merits and achievements not wholly its own. Particularly was this
+true in the case of saving the union. "When in the economy of
+Providence, this land was to be purged of human slavery ... the
+Republican party came into power," ran a declaration in one platform.
+"The Republican party suppressed a gigantic rebellion, emancipated four
+million slaves, decreed the equal citizenship of all, and established
+universal suffrage," ran another. As for the aid rendered by the
+millions of Northern Democrats who stood by the union and the tens of
+thousands of them who actually fought in the union army, the Republicans
+in their zeal were inclined to be oblivious. They repeatedly charged the
+Democratic party "with being the same in character and spirit as when it
+sympathized with treason."
+
+=Republican Control of the South.=--To the strength enjoyed in the
+North, the Republicans for a long time added the advantages that came
+from control over the former Confederate states where the newly
+enfranchised negroes, under white leadership, gave a grateful support to
+the party responsible for their freedom. In this branch of politics,
+motives were so mixed that no historian can hope to appraise them all at
+their proper values. On the one side of the ledger must be set the
+vigorous efforts of the honest and sincere friends of the freedmen to
+win for them complete civil and political equality, wiping out not only
+slavery but all its badges of misery and servitude. On the same side
+must be placed the labor of those who had valiantly fought in forum and
+field to save the union and who regarded continued Republican supremacy
+after the war as absolutely necessary to prevent the former leaders in
+secession from coming back to power. At the same time there were
+undoubtedly some men of the baser sort who looked on politics as a game
+and who made use of "carpet-bagging" in the South to win the spoils that
+might result from it. At all events, both by laws and presidential acts,
+the Republicans for many years kept a keen eye upon the maintenance of
+their dominion in the South. Their declaration that neither the law nor
+its administration should admit any discrimination in respect of
+citizens by reason of race, color, or previous condition of servitude
+appealed to idealists and brought results in elections. Even South
+Carolina, where reposed the ashes of John C. Calhoun, went Republican in
+1872 by a vote of three to one!
+
+Republican control was made easy by the force bills described in a
+previous chapter--measures which vested the supervision of elections in
+federal officers appointed by Republican Presidents. These drastic
+measures, departing from American tradition, the Republican authors
+urged, were necessary to safeguard the purity of the ballot, not merely
+in the South where the timid freedman might readily be frightened from
+using it; but also in the North, particularly in New York City, where it
+was claimed that fraud was regularly practiced by Democratic leaders.
+
+The Democrats, on their side, indignantly denied the charges, replying
+that the force bills were nothing but devices created by the Republicans
+for the purpose of securing their continued rule through systematic
+interference with elections. Even the measures of reconstruction were
+deemed by Democratic leaders as thinly veiled schemes to establish
+Republican power throughout the country. "Nor is there the slightest
+doubt," exclaimed Samuel J. Tilden, spokesman of the Democrats in New
+York and candidate for President in 1876, "that the paramount object and
+motive of the Republican party is by these means to secure itself
+against a reaction of opinion adverse to it in our great populous
+Northern commonwealths.... When the Republican party resolved to
+establish negro supremacy in the ten states in order to gain to itself
+the representation of those states in Congress, it had to begin by
+governing the people of those states by the sword.... The next was the
+creation of new electoral bodies for those ten states, in which, by
+exclusions, by disfranchisements and proscriptions, by control over
+registration, by applying test oaths ... by intimidation and by every
+form of influence, three million negroes are made to predominate over
+four and a half million whites."
+
+=The War as a Campaign Issue.=--Even the repeal of force bills could not
+allay the sectional feelings engendered by the war. The Republicans
+could not forgive the men who had so recently been in arms against the
+union and insisted on calling them "traitors" and "rebels." The
+Southerners, smarting under the reconstruction acts, could regard the
+Republicans only as political oppressors. The passions of the war had
+been too strong; the distress too deep to be soon forgotten. The
+generation that went through it all remembered it all. For twenty
+years, the Republicans, in their speeches and platforms, made "a
+straight appeal to the patriotism of the Northern voters." They
+maintained that their party, which had saved the union and emancipated
+the slaves, was alone worthy of protecting the union and uplifting the
+freedmen.
+
+Though the Democrats, especially in the North, resented this policy and
+dubbed it with the expressive but inelegant phrase, "waving the bloody
+shirt," the Republicans refused to surrender a slogan which made such a
+ready popular appeal. As late as 1884, a leader expressed the hope that
+they might "wring one more President from the bloody shirt." They
+refused to let the country forget that the Democratic candidate, Grover
+Cleveland, had escaped military service by hiring a substitute; and they
+made political capital out of the fact that he had "insulted the
+veterans of the Grand Army of the Republic" by going fishing on
+Decoration Day.
+
+=Three Republican Presidents.=--Fortified by all these elements of
+strength, the Republicans held the presidency from 1869 to 1885. The
+three Presidents elected in this period, Grant, Hayes, and Garfield, had
+certain striking characteristics in common. They were all of origin
+humble enough to please the most exacting Jacksonian Democrat. They had
+been generals in the union army. Grant, next to Lincoln, was regarded as
+the savior of the Constitution. Hayes and Garfield, though lesser lights
+in the military firmament, had honorable records duly appreciated by
+veterans of the war, now thoroughly organized into the Grand Army of the
+Republic. It is true that Grant was not a politician and had never voted
+the Republican ticket; but this was readily overlooked. Hayes and
+Garfield on the other hand were loyal party men. The former had served
+in Congress and for three terms as governor of his state. The latter had
+long been a member of the House of Representatives and was Senator-elect
+when he received the nomination for President.
+
+All of them possessed, moreover, another important asset, which was not
+forgotten by the astute managers who led in selecting candidates. All
+of them were from Ohio--though Grant had been in Illinois when the
+summons to military duties came--and Ohio was a strategic state. It lay
+between the manufacturing East and the agrarian country to the West.
+Having growing industries and wool to sell it benefited from the
+protective tariff. Yet being mainly agricultural still, it was not
+
+without sympathy for the farmers who showed low tariff or free trade
+tendencies. Whatever share the East had in shaping laws and framing
+policies, it was clear that the West was to have the candidates. This
+division in privileges--not uncommon in political management--was always
+accompanied by a judicious selection of the candidate for Vice
+President. With Garfield, for example, was associated a prominent New
+York politician, Chester A. Arthur, who, as fate decreed, was destined
+to more than three years' service as chief magistrate, on the
+assassination of his superior in office.
+
+=The Disputed Election of 1876.=--While taking note of the long years of
+Republican supremacy, it must be recorded that grave doubts exist in the
+minds of many historians as to whether one of the three Presidents,
+Hayes, was actually the victor in 1876 or not. His Democratic opponent,
+Samuel J. Tilden, received a popular plurality of a quarter of a million
+and had a plausible claim to a majority of the electoral vote. At all
+events, four states sent in double returns, one set for Tilden and
+another for Hayes; and a deadlock ensued. Both parties vehemently
+claimed the election and the passions ran so high that sober men did not
+shrink from speaking of civil war again. Fortunately, in the end, the
+counsels of peace prevailed. Congress provided for an electoral
+commission of fifteen men to review the contested returns. The
+Democrats, inspired by Tilden's moderation, accepted the judgment in
+favor of Hayes even though they were not convinced that he was really
+entitled to the office.
+
+
+THE GROWTH OF OPPOSITION TO REPUBLICAN RULE
+
+=Abuses in American Political Life.=--During their long tenure of
+office, the Republicans could not escape the inevitable consequences of
+power; that is, evil practices and corrupt conduct on the part of some
+who found shelter within the party. For that matter neither did the
+Democrats manage to avoid such difficulties in those states and cities
+where they had the majority. In New York City, for instance, the local
+Democratic organization, known as Tammany Hall, passed under the sway of
+a group of politicians headed by "Boss" Tweed. He plundered the city
+treasury until public-spirited citizens, supported by Samuel J. Tilden,
+the Democratic leader of the state, rose in revolt, drove the ringleader
+from power, and sent him to jail. In Philadelphia, the local Republican
+bosses were guilty of offenses as odious as those committed by New York
+politicians. Indeed, the decade that followed the Civil War was marred
+by so many scandals in public life that one acute editor was moved to
+inquire: "Are not all the great communities of the Western World growing
+more corrupt as they grow in wealth?"
+
+In the sphere of national politics, where the opportunities were
+greater, betrayals of public trust were even more flagrant. One
+revelation after another showed officers, high and low, possessed with
+the spirit of peculation. Members of Congress, it was found, accepted
+railway stock in exchange for votes in favor of land grants and other
+concessions to the companies. In the administration as well as the
+legislature the disease was rife. Revenue officers permitted whisky
+distillers to evade their taxes and received heavy bribes in return. A
+probe into the post-office department revealed the malodorous "star
+route frauds"--the deliberate overpayment of certain mail carriers whose
+lines were indicated in the official record by asterisks or stars. Even
+cabinet officers did not escape suspicion, for the trail of the serpent
+led straight to the door of one of them.
+
+In the lower ranges of official life, the spoils system became more
+virulent as the number of federal employees increased. The holders of
+offices and the seekers after them constituted a veritable political
+army. They crowded into Republican councils, for the Republicans, being
+in power, could alone dispense federal favors. They filled positions in
+the party ranging from the lowest township committee to the national
+convention. They helped to nominate candidates and draft platforms and
+elbowed to one side the busy citizen, not conversant with party
+intrigues, who could only give an occasional day to political matters.
+Even the Civil Service Act of 1883, wrung from a reluctant Congress two
+years after the assassination of Garfield, made little change for a long
+time. It took away from the spoilsmen a few thousand government
+positions, but it formed no check on the practice of rewarding party
+workers from the public treasury.
+
+On viewing this state of affairs, many a distinguished citizen became
+profoundly discouraged. James Russell Lowell, for example, thought he
+saw a steady decline in public morals. In 1865, hearing of Lee's
+surrender, he had exclaimed: "There is something magnificent in having a
+country to love!" Ten years later, when asked to write an ode for the
+centennial at Philadelphia in 1876, he could think only of a biting
+satire on the nation:
+
+    "Show your state legislatures; show your Rings;
+     And challenge Europe to produce such things
+     As high officials sitting half in sight
+     To share the plunder and fix things right.
+     If that don't fetch her, why, you need only
+     To show your latest style in martyrs,--Tweed:
+     She'll find it hard to hide her spiteful tears
+     At such advance in one poor hundred years."
+
+When his critics condemned him for this "attack upon his native land,"
+Lowell replied in sadness: "These fellows have no notion of what love of
+country means. It was in my very blood and bones. If I am not an
+American who ever was?... What fills me with doubt and dismay is the
+degradation of the moral tone. Is it or is it not a result of democracy?
+Is ours a 'government of the people, by the people, for the people,' or
+a Kakistocracy [a government of the worst], rather for the benefit of
+knaves at the cost of fools?"
+
+=The Reform Movement in Republican Ranks.=--The sentiments expressed by
+Lowell, himself a Republican and for a time American ambassador to
+England, were shared by many men in his party. Very soon after the close
+of the Civil War some of them began to protest vigorously against the
+policies and conduct of their leaders. In 1872, the dissenters, calling
+themselves Liberal Republicans, broke away altogether, nominated a
+candidate of their own, Horace Greeley, and put forward a platform
+indicting the Republican President fiercely enough to please the most
+uncompromising Democrat. They accused Grant of using "the powers and
+opportunities of his high office for the promotion of personal ends."
+They charged him with retaining "notoriously corrupt and unworthy men in
+places of power and responsibility." They alleged that the Republican
+party kept "alive the passions and resentments of the late civil war to
+use them for their own advantages," and employed the "public service of
+the government as a machinery of corruption and personal influence."
+
+It was not apparent, however, from the ensuing election that any
+considerable number of Republicans accepted the views of the Liberals.
+Greeley, though indorsed by the Democrats, was utterly routed and died
+of a broken heart. The lesson of his discomfiture seemed to be that
+independent action was futile. So, at least, it was regarded by most men
+of the rising generation like Henry Cabot Lodge, of Massachusetts, and
+Theodore Roosevelt, of New York. Profiting by the experience of Greeley
+they insisted in season and out that reformers who desired to rid the
+party of abuses should remain loyal to it and do their work "on the
+inside."
+
+=The Mugwumps and Cleveland Democracy in 1884.=--Though aided by
+Republican dissensions, the Democrats were slow in making headway
+against the political current. They were deprived of the energetic and
+capable leadership once afforded by the planters, like Calhoun, Davis,
+and Toombs; they were saddled by their opponents with responsibility for
+secession; and they were stripped of the support of the prostrate
+South. Not until the last Southern state was restored to the union, not
+until a general amnesty was wrung from Congress, not until white
+supremacy was established at the polls, and the last federal soldier
+withdrawn from Southern capitals did they succeed in capturing the
+presidency.
+
+The opportune moment for them came in 1884 when a number of
+circumstances favored their aspirations. The Republicans, leaving the
+Ohio Valley in their search for a candidate, nominated James G. Blaine
+of Maine, a vigorous and popular leader but a man under fire from the
+reformers in his own party. The Democrats on their side were able to
+find at this juncture an able candidate who had no political enemies in
+the sphere of national politics, Grover Cleveland, then governor of New
+York and widely celebrated as a man of "sterling honesty." At the same
+time a number of dissatisfied Republicans openly espoused the Democratic
+cause,--among them Carl Schurz, George William Curtis, Henry Ward
+Beecher, and William Everett, men of fine ideals and undoubted
+integrity. Though the "regular" Republicans called them "Mugwumps" and
+laughed at them as the "men milliners, the dilettanti, and carpet
+knights of politics," they had a following that was not to be despised.
+
+The campaign which took place that year was one of the most savage in
+American history. Issues were thrust into the background. The tariff,
+though mentioned, was not taken seriously. Abuse of the opposition was
+the favorite resource of party orators. The Democrats insisted that "the
+Republican party so far as principle is concerned is a reminiscence. In
+practice it is an organization for enriching those who control its
+machinery." For the Republican candidate, Blaine, they could hardly find
+words to express their contempt. The Republicans retaliated in kind.
+They praised their own good works, as of old, in saving the union, and
+denounced the "fraud and violence practiced by the Democracy in the
+Southern states." Seeing little objectionable in the public record of
+Cleveland as mayor of Buffalo and governor of New York, they attacked
+his personal character. Perhaps never in the history of political
+campaigns did the discussions on the platform and in the press sink to
+so low a level. Decent people were sickened. Even hot partisans shrank
+from their own words when, after the election, they had time to reflect
+on their heedless passions. Moreover, nothing was decided by the
+balloting. Cleveland was elected, but his victory was a narrow one. A
+change of a few hundred votes in New York would have sent his opponent
+to the White House instead.
+
+=Changing Political Fortunes (1888-96).=--After the Democrats had
+settled down to the enjoyment of their hard-earned victory, President
+Cleveland in his message of 1887 attacked the tariff as "vicious,
+inequitable, and illogical"; as a system of taxation that laid a burden
+upon "every consumer in the land for the benefit of our manufacturers."
+Business enterprise was thoroughly alarmed. The Republicans
+characterized the tariff message as a free-trade assault upon the
+industries of the country. Mainly on that issue they elected in 1888
+Benjamin Harrison of Indiana, a shrewd lawyer, a reticent politician, a
+descendant of the hero of Tippecanoe, and a son of the old Northwest.
+Accepting the outcome of the election as a vindication of their
+principles, the Republicans, under the leadership of William McKinley in
+the House of Representatives, enacted in 1890 a tariff law imposing the
+highest duties yet laid in our history. To their utter surprise,
+however, they were instantly informed by the country that their program
+was not approved. That very autumn they lost in the congressional
+elections, and two years later they were decisively beaten in the
+presidential campaign, Cleveland once more leading his party to victory.
+
+
+=References=
+
+L.H. Haney, _Congressional History of Railways_ (2 vols.).
+
+J.P. Davis, _Union Pacific Railway_.
+
+J.M. Swank, _History of the Manufacture of Iron_.
+
+M.T. Copeland, _The Cotton Manufacturing Industry in the United States_
+(Harvard Studies).
+
+E.W. Bryce, _Progress of Invention in the Nineteenth Century_.
+
+Ida Tarbell, _History of the Standard Oil Company_ (Critical).
+
+G.H. Montague, _Rise and Progress of the Standard Oil Company_
+(Friendly).
+
+H.P. Fairchild, _Immigration_, and F.J. Warne, _The Immigrant Invasion_
+(Both works favor exclusion).
+
+I.A. Hourwich, _Immigration_ (Against exclusionist policies).
+
+J.F. Rhodes, _History of the United States, 1877-1896_, Vol. VIII.
+
+Edward Stanwood, _A History of the Presidency_, Vol. I, for the
+presidential elections of the period.
+
+
+=Questions=
+
+1. Contrast the state of industry and commerce at the close of the Civil
+War with its condition at the close of the Revolutionary War.
+
+2. Enumerate the services rendered to the nation by the railways.
+
+3. Explain the peculiar relation of railways to government.
+
+4. What sections of the country have been industrialized?
+
+5. How do you account for the rise and growth of the trusts? Explain
+some of the economic advantages of the trust.
+
+6. Are the people in cities more or less independent than the farmers?
+What was Jefferson's view?
+
+7. State some of the problems raised by unrestricted immigration.
+
+8. What was the theory of the relation of government to business in this
+period? Has it changed in recent times?
+
+9. State the leading economic policies sponsored by the Republican
+party.
+
+10. Why were the Republicans especially strong immediately after the
+Civil War?
+
+11. What illustrations can you give showing the influence of war in
+American political campaigns?
+
+12. Account for the strength of middle-western candidates.
+
+13. Enumerate some of the abuses that appeared in American political
+life after 1865.
+
+14. Sketch the rise and growth of the reform movement.
+
+15. How is the fluctuating state of public opinion reflected in the
+elections from 1880 to 1896?
+
+
+=Research Topics=
+
+=Invention, Discovery, and Transportation.=--Sparks, _National
+Development_ (American Nation Series), pp. 37-67; Bogart, _Economic
+History of the United States_, Chaps. XXI, XXII, and XXIII.
+
+=Business and Politics.=--Paxson, _The New Nation_ (Riverside Series),
+pp. 92-107; Rhodes, _History of the United States_, Vol. VII, pp. 1-29,
+64-73, 175-206; Wilson, _History of the American People_, Vol. IV, pp.
+78-96.
+
+=Immigration.=--Coman, _Industrial History of the United States_ (2d
+ed.), pp. 369-374; E.L. Bogart, _Economic History of the United States_,
+pp. 420-422, 434-437; Jenks and Lauck, _Immigration Problems_, Commons,
+_Races and Immigrants_.
+
+=The Disputed Election of 1876.=--Haworth, _The United States in Our Own
+Time_, pp. 82-94; Dunning, _Reconstruction, Political and Economic_
+(American Nation Series), pp. 294-341; Elson, _History of the United
+States_, pp. 835-841.
+
+=Abuses in Political Life.=--Dunning, _Reconstruction_, pp. 281-293; see
+criticisms in party platforms in Stanwood, _History of the Presidency_,
+Vol. I; Bryce, _American Commonwealth_ (1910 ed.), Vol. II, pp. 379-448;
+136-167.
+
+=Studies of Presidential Administrations.=--(_a_) Grant, (_b_) Hayes,
+(_c_) Garfield-Arthur, (_d_) Cleveland, and (_e_) Harrison, in Haworth,
+_The United States in Our Own Time_, or in Paxson, _The New Nation_
+(Riverside Series), or still more briefly in Elson.
+
+=Cleveland Democracy.=--Haworth, _The United States_, pp. 164-183;
+Rhodes, _History of the United States_, Vol. VIII, pp. 240-327; Elson,
+pp. 857-887.
+
+=Analysis of Modern Immigration Problems.=--_Syllabus in History_ (New
+York State, 1919), pp. 110-112.
+
+
+
+
+CHAPTER XVIII
+

MedicalGPT-main/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
+}

MedicalGPT-main/docs/training_details.md

@@ -0,0 +1,104 @@
+# Training Detail
+
+
+### Stage 1: PT(Continue PreTraining)
+第一阶段：PT(Continue PreTraining)增量预训练
+
+使用百科类文档类数据集，用来在领域数据集上增量预训练或二次预训练，期望能把领域知识注入给模型，以医疗领域为例，希望增量预训练，能让模型理解感冒的症状、病因、治疗药品、治疗方法、药品疗效等知识，便于后续的SFT监督微调能激活这些内在知识。
+
+这里说明一点，像GPT3、LLaMA这样的大模型理论上是可以从增量预训练中获益，但增量预训练需要满足两个要求：1）高质量的预训练样本；2）较大的计算资源，显存要求高，即使是用LoRA技术，也要满足block_size=1024或2048长度的文本加载到显存中。
+
+其次，如果你的项目用到的数据是模型预训练中已经使用了的，如维基百科、ArXiv等LLaMA模型预训练用了的，则这些数据是没有必要再喂给LLaMA增量预训练，而且预训练样本的质量如果不够高，也可能会损害原模型的生成能力。
+
+tips：PT阶段是可选项，请慎重处理。
+
+基于llama-7b模型，使用医疗百科类数据继续预训练，期望注入医疗知识到预训练模型，得到llama-7b-pt模型
+
+Continue pretraining of the base llama-7b model to create llama-7b-pt:
+
+```shell
+cd scripts
+sh run_pt.sh
+```
+
+[训练参数说明wiki](https://github.com/shibing624/MedicalGPT/wiki/%E8%AE%AD%E7%BB%83%E7%BB%86%E8%8A%82%E8%AF%B4%E6%98%8E)
+- 如果你的显存不足，可以改小batch_size=1, block_size=512（影响训练的上下文最大长度）;
+- 如果你的显存更大，可以改大block_size=2048, 此为llama原始预训练长度，不能更大啦；调大batch_size。
+
+### Stage 2: SFT(Supervised Fine-tuning)
+第二阶段：SFT(Supervised Fine-tuning)有监督微调
+
+基于llama-7b-pt模型，使用医疗问答类数据进行有监督微调，得到llama-7b-sft模型
+
+Supervised fine-tuning of the base llama-7b-pt model to create llama-7b-sft
+
+```shell
+cd scripts
+sh run_sft.sh
+```
+
+[训练参数说明wiki](https://github.com/shibing624/MedicalGPT/wiki/%E8%AE%AD%E7%BB%83%E7%BB%86%E8%8A%82%E8%AF%B4%E6%98%8E)
+
+### Stage 3: RLHF(Reinforcement Learning from Human Feedback)
+#### Reward Modeling
+RM(Reward Model)奖励模型建模
+
+RM(Reward Model)奖励模型，原则上，我们可以直接用人类标注来对模型做 RLHF 微调。
+
+然而，这将需要我们给人类发送一些样本，在每轮优化后计分。这是贵且慢的，因为收敛需要的训练样本量大，而人类阅读和标注的速度有限。
+一个比直接反馈更好的策略是，在进入 RL 循环之前用人类标注集来训练一个奖励模型RM。奖励模型的目的是模拟人类对文本的打分。
+
+构建奖励模型的最佳实践是预测结果的排序，即对每个 prompt (输入文本) 对应的两个结果 (yk, yj)，模型预测人类标注的比分哪个更高。
+RM模型是通过人工标注SFT模型的打分结果来训练的，目的是取代人工打分，本质是个回归模型，用来对齐人类偏好，主要是"HHH"原则，具体是"helpful, honest, harmless"。
+
+
+基于llama-7b-sft模型，使用医疗问答偏好数据训练奖励偏好模型，训练得到llama-7b-reward模型
+
+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
+```
+[训练参数说明wiki](https://github.com/shibing624/MedicalGPT/wiki/%E8%AE%AD%E7%BB%83%E7%BB%86%E8%8A%82%E8%AF%B4%E6%98%8E)
+
+#### Reinforcement Learning
+RL(Reinforcement Learning)强化学习
+
+RL(Reinforcement Learning)模型的目的是最大化奖励模型的输出，基于上面步骤，我们有了微调的语言模型(llama-7b-sft)和奖励模型(llama-7b-reward)，
+可以开始执行 RL 循环了。
+
+这个过程大致分为三步：
+
+1. 输入prompt，模型生成答复
+2. 用奖励模型来对答复评分
+3. 基于评分，进行一轮策略优化的强化学习(PPO)
+
+<img src=https://huggingface.co/datasets/trl-internal-testing/example-images/resolve/main/blog/stackllama/trl_loop.png height=400 />
+
+
+基于llama-7b-reward模型 RL 微调训练llama-7b-sft模型，得到llama-7b-rl模型
+
+Reinforcement Learning fine-tuning of llama-7b-sft with the llama-7b-reward reward model to create llama-7b-rl
+
+```shell
+pip install git+https://github.com/lvwerra/trl
+cd scripts
+sh run_rl.sh
+```
+
+### Stage 3: DPO(Direct Preference Optimization)
+DPO(Direct Preference Optimization)直接偏好优化
+
+DPO方法可以通过直接优化语言模型来实现对其行为的精确控制，而无需使用复杂的强化学习。
+
+DPO 将奖励函数和最优策略之间的映射联系起来，从而把约束奖励最大化问题转化为一个单阶段的策略训练问题。
+这种算法不仅不用拟合奖励模型，还避免了在微调过程中从语言模型中采样或调整超参数的需要。
+
+实验结果表明，DPO 算法可以与现有RLHF方法一样有效地从人类偏好中学习，甚至在某些任务中表现更好，比如情感调节、摘要和单轮对话。
+
+PS: 使用DPO训练LLaMA2-7B在fp16，batch_size为2时，需要70GB显存。
+
+```shell
+sh run_dpo.sh
+```

MedicalGPT-main/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()

MedicalGPT-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()

MedicalGPT-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()

MedicalGPT-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()

MedicalGPT-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()

MedicalGPT-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()

MedicalGPT-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

MedicalGPT-main/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()

MedicalGPT-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()

MedicalGPT-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

MedicalGPT-main/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

MedicalGPT-main/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

MedicalGPT-main/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

MedicalGPT-main/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

MedicalGPT-main/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
+}

MedicalGPT-main/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
+}

MedicalGPT-main/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()