Demos/09. DPO 微调：根据偏好引导 LLM 的输出.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4421e9ed-db25-40cd-b8e1-ce699a2023ab",
   "metadata": {},
   "source": [
    "# DPO 微调：根据偏好引导LLM的输出\n",
    "> [GenAI HW6: LLM Values Alignment](https://colab.research.google.com/drive/1d3zmkqo-ZmxrIOYWSe3vDD0za8tUPguu?usp=sharing#scrollTo=owGIuqdnRI8I) 中文镜像版\n",
    ">\n",
    "> 指导文章：[11. DPO 微调示例：根据人类偏好优化 LLM 大语言模型](https://github.com/Hoper-J/LLM-Guide-and-Demos-zh_CN/blob/master/Guide/11.%20DPO%20微调示例：根据人类偏好优化%20LLM%20大语言模型.md)\n",
    "\n",
    "**目标**：学习如何使用带标签的偏好数据来对齐模型的行为。\n",
    "\n",
    "大白话就是：微调。但这里使用的是 DPO 方法（之后会具体解释什么是 DPO）。\n",
    "\n",
    "之前微调唐诗更像是一个传统的方法，使用问题和答案驱动的微调方式，数据标注成本很高。而 DPO 不同，它使用的是人类对文本的偏好对：即人类认为一对文本中哪个更好。\n",
    "\n",
    "在线链接：[Kaggle](https://www.kaggle.com/code/aidemos/09-dpo-llm) | [Colab](https://colab.research.google.com/drive/1TxL9MrIXDY3HjWgQ4B3IcEeMj-lsbNAZ?usp=sharing)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6124745-9431-4889-8d6d-c7960c90c409",
   "metadata": {},
   "source": [
    "## 安装和导入一些必要的库"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e633546a-9dd3-447c-a26f-76bd64376878",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!pip install bitsandbytes\n",
    "!pip install datasets\n",
    "!pip install peft\n",
    "!pip install trl\n",
    "!pip install accelerate\n",
    "!pip install tf-keras\n",
    "!pip install numpy==1.26.4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7de501d5-c828-4d81-a4ce-31e029b7bc2a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import re\n",
    "import json\n",
    "\n",
    "import torch\n",
    "import pandas as pd\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "from datasets import Dataset\n",
    "from peft import LoraConfig\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, BitsAndBytesConfig, GenerationConfig\n",
    "from trl import DPOConfig, DPOTrainer"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "79ef6669-e845-4d9c-a360-fd4283790bcb",
   "metadata": {
    "tags": []
   },
   "source": [
    "## 加载数据集"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de383280-d9fa-4b86-9250-ccb58212e946",
   "metadata": {},
   "outputs": [],
   "source": [
    "!git clone https://github.com/Baiiiiiiiiii/GenAI_hw6_dataset.git"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12734922-1d3d-46d7-acdc-7d484f6dd577",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "with open(\"./GenAI_hw6_dataset/labelled_data.json\", 'r') as jsonfile:\n",
    "    full_data = json.load(jsonfile)\n",
    "\n",
    "with open(\"./GenAI_hw6_dataset/test_prompt.json\", 'r') as jsonfile:\n",
    "    test_data = json.load(jsonfile)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4887df1d-3fcc-4c39-9d23-7e82958e90d1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "full_data[:5], test_data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2cfb6314-77d3-45e4-9546-6a4baf3bdc20",
   "metadata": {
    "tags": []
   },
   "source": [
    "## 使用 HFD 下载模型\n",
    "\n",
    "我们这里使用多线程的方法进行快速下载。\n",
    "\n",
    "如果直接运行以下命令报错，根据[a. 使用 HFD 加快 Hugging Face 模型和数据集的下载](https://github.com/Hoper-J/LLM-Guide-and-Demos-zh_CN/blob/master/Guide/a.%20使用%20HFD%20加快%20Hugging%20Face%20模型和数据集的下载.md)进行前置安装。\n",
    "\n",
    "当然，你也可以取消我注释的部分，使用官方的命令进行安装，但是会很慢。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc534866-a4d4-4989-b5b3-04df0bc21f16",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!wget https://hf-mirror.com/hfd/hfd.sh\n",
    "!chmod a+x hfd.sh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e1b7730-a6bd-4b10-8379-a21f2f35b17d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!export HF_ENDPOINT=https://hf-mirror.com\n",
    "!./hfd.sh 'MediaTek-Research/Breeze-7B-Instruct-v0_1' --tool aria2c -x 16"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e93db2b4-5f9f-44c6-8515-b989ca7711a6",
   "metadata": {},
   "source": [
    "## 加载模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4526f56b-6555-4016-aaf2-5598fe46ffca",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# 用官方方法就取消这段注释，然后注释下面\n",
    "# model = AutoModelForCausalLM.from_pretrained(\n",
    "#     'MediaTek-Research/Breeze-7B-Instruct-v0_1',\n",
    "#     device_map='auto',\n",
    "#     trust_remote_code=True,\n",
    "#     quantization_config=BitsAndBytesConfig(\n",
    "#         load_in_4bit=True,\n",
    "#         bnb_4bit_compute_dtype=torch.bfloat16,\n",
    "#         bnb_4bit_use_double_quant=True,\n",
    "#         bnb_4bit_quant_type='nf4'\n",
    "#     )\n",
    "# )\n",
    "\n",
    "\n",
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    'Breeze-7B-Instruct-v0_1',  # 替换为本地模型存储的实际路径\n",
    "    device_map='auto',\n",
    "    trust_remote_code=True,\n",
    "    local_files_only=True,  # 确保从本地加载\n",
    "    quantization_config=BitsAndBytesConfig(\n",
    "        load_in_4bit=True,\n",
    "        bnb_4bit_compute_dtype=torch.bfloat16,\n",
    "        bnb_4bit_use_double_quant=True,\n",
    "        bnb_4bit_quant_type='nf4'\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43fc3817-31d8-4696-b121-36f9247b9e7f",
   "metadata": {
    "tags": []
   },
   "source": [
    "## 查看未经过微调的模型原始输出"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5895e17-e89b-4121-bc7e-a0a91b3351e3",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# 用官方方法就取消这行注释，然后注释下面\n",
    "# tokenizer = AutoTokenizer.from_pretrained('MediaTek-Research/Breeze-7B-Instruct-v0_1')\n",
    "tokenizer = AutoTokenizer.from_pretrained('Breeze-7B-Instruct-v0_1')\n",
    "tokenizer.padding_side = \"right\"\n",
    "tokenizer.pad_token = tokenizer.eos_token\n",
    "\n",
    "def data_formulate(data):\n",
    "    messages = [\n",
    "        {\"role\": \"system\", \"content\": '回覆請少於20字'},\n",
    "        {\"role\": \"user\", \"content\": data['prompt']},\n",
    "    ]\n",
    "    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)\n",
    "    return prompt\n",
    "\n",
    "original_model_response = []\n",
    "for data in tqdm(test_data):\n",
    "    id = data['id']\n",
    "    print(f'Question {id}:\\n'+data['prompt'])\n",
    "    inputs = tokenizer(data_formulate(data), return_tensors=\"pt\").to('cuda')\n",
    "    generation_config=GenerationConfig(\n",
    "            do_sample=False,\n",
    "            max_new_tokens = 200,\n",
    "            pad_token_id = tokenizer.pad_token_id\n",
    "    )\n",
    "    output = model.generate(**inputs, generation_config=generation_config)\n",
    "    output = tokenizer.batch_decode(output, skip_special_tokens=True)[0].split('[/INST] ')[1]\n",
    "    original_model_response.append(output)\n",
    "    print('Response from original model:\\n'+output+'\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e307f9f8-ab02-4bcd-aae7-78d2c6a64202",
   "metadata": {
    "tags": []
   },
   "source": [
    "## 设置参数\n",
    "\n",
    "你只需要修改这个模块，不需要改变其他的，除非你真的知道自己在做什么。\n",
    "\n",
    "`support_ratio` 将反映你的偏好：\n",
    "\n",
    "- 0 表示完全不支持（反对）真人化\n",
    "- 1 表示完全支持真人化\n",
    "- 0.1 表示 10% 支持真人化"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ee41ff5a-9dcd-452a-9340-33c9e42c815f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "num_epoch = 1\n",
    "data_size = 50\n",
    "support_ratio = 0.1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f701f863-8bf7-4cdc-9a14-b1d51a5953b2",
   "metadata": {},
   "source": [
    "## 准备训练数据\n",
    "\n",
    "这里，我们将数据集分为支持（support）和反对（oppose）两部分，构建一个包含偏好对的训练数据集。\n",
    "\n",
    "总共有 50 笔训练数据，当 support 设置为 0.1 时，前 50*0.1=5 笔训练资料的偏好将倾向于支持真人化，后 50-4=45 笔资料反对真人化。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb14369d-6670-40de-a058-10aaf989108c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# 选择部分数据用于训练\n",
    "training_data = full_data[:data_size]\n",
    "\n",
    "# 定义 support 数据集的大小\n",
    "support_data_size = int(data_size * support_ratio)\n",
    "\n",
    "# 为训练数据集准备数据\n",
    "prompt_list = [data_formulate(data) for data in training_data]\n",
    "chosen_list = [data['support'] for data in training_data[:support_data_size]] + [data['oppose'] for data in training_data[support_data_size:]]\n",
    "rejected_list = [data['oppose'] for data in training_data[:support_data_size]] + [data['support'] for data in training_data[support_data_size:]]\n",
    "position_list = ['support' for _ in range(support_data_size)] + ['oppose' for _ in range(data_size - support_data_size)]\n",
    "\n",
    "# 创建训练数据集\n",
    "train_dataset = Dataset.from_dict({'prompt': prompt_list, 'position': position_list, 'chosen': chosen_list, 'rejected': rejected_list})\n",
    "pd.DataFrame(train_dataset).rename(columns={\"chosen\": \"preferred\", \"rejected\": \"non-preferred\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "898c4ac3-4f8c-43ac-878d-0b7a5c39e56c",
   "metadata": {
    "tags": []
   },
   "source": [
    "## 训练\n",
    "\n",
    "现在，我们进入训练阶段。首先，设置训练参数："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "873162f1-71eb-4079-85e4-d0ef2c230a89",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "training_args = DPOConfig(\n",
    "    output_dir='./',\n",
    "    per_device_train_batch_size=1,\n",
    "    num_train_epochs=num_epoch,\n",
    "    gradient_accumulation_steps=8,\n",
    "    gradient_checkpointing=False,\n",
    "    learning_rate=2e-4,\n",
    "    optim=\"paged_adamw_8bit\",\n",
    "    logging_steps = 1,\n",
    "    warmup_ratio = 0.1,\n",
    "    report_to = 'none',\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9a6e9c1-4982-4350-a0e9-0b63a5846c04",
   "metadata": {},
   "source": [
    "接下来，配置PEFT（Parameter-Efficient Fine-Tuning）：\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c5cefa0-8054-46d7-9e21-3fadfe15aa4e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "peft_config = LoraConfig(\n",
    "    lora_alpha=16,\n",
    "    lora_dropout=0.1,\n",
    "    r=64,\n",
    "    bias=\"none\",\n",
    "    task_type=\"CAUSAL_LM\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0cf2d33-1b19-4835-a903-38a9a2d340ec",
   "metadata": {},
   "source": [
    "然后，初始化DPO训练器：\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d734c9e9-bf95-400a-b011-1f265706d97f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "dpo_trainer = DPOTrainer(\n",
    "    model,\n",
    "    args=training_args,\n",
    "    beta=0.1,\n",
    "    train_dataset=train_dataset,\n",
    "    tokenizer=tokenizer,\n",
    "    peft_config=peft_config,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f6b12e3-8a9e-4d14-9e58-31f3ed202a4f",
   "metadata": {},
   "source": [
    "开始训练：\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16b800f6-6747-459c-8048-e3ebac328e75",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "dpo_trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f043c0a6-e173-4c6f-bc6c-08403031b7af",
   "metadata": {
    "tags": []
   },
   "source": [
    "## 查看微调后的模型输出"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "769c45e7-338b-401e-abd0-db440471033a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "trained_model_response = []\n",
    "for data in tqdm(test_data):\n",
    "    id = data['id']\n",
    "    print(f'Question {id}:\\n'+data['prompt'])\n",
    "    inputs = tokenizer(data_formulate(data), return_tensors=\"pt\").to('cuda')\n",
    "    generation_config=GenerationConfig(\n",
    "            do_sample=False,\n",
    "            max_new_tokens = 200,\n",
    "            pad_token_id = tokenizer.pad_token_id\n",
    "    )\n",
    "    output = model.generate(**inputs, generation_config=generation_config)\n",
    "    output = tokenizer.batch_decode(output, skip_special_tokens=True)[0].split('[/INST] ')[1]\n",
    "    trained_model_response.append(output)\n",
    "    print('Response from trained model:\\n'+output+'\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f6bab686-308b-4a24-a2c6-65f5c48e4040",
   "metadata": {
    "tags": []
   },
   "source": [
    "## 观察输出结果"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4e3136b-c743-4228-86f1-a84fb0c2bc38",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "model_response = []\n",
    "print(f'num_epoch: {num_epoch}\\ndata_size: {data_size}\\nsupport_ratio: {support_ratio}')\n",
    "print()\n",
    "for data in test_data:\n",
    "    id = data['id']\n",
    "    ref_output = original_model_response[id-1]\n",
    "    output = trained_model_response[id-1]\n",
    "    print(f'Question {id}:\\n'+data['prompt'])\n",
    "    print('Response from original model:\\n'+ref_output)\n",
    "    print('Response from trained model:\\n'+output)\n",
    "    print()\n",
    "    model_response.append({'id':data['id'], 'prompt':data['prompt'], 'response_from_original_model':ref_output, 'response_from_trained_model':output})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f69a911-b23b-4247-8ae8-0ff5bc1b2033",
   "metadata": {
    "tags": []
   },
   "source": [
    "## 获取 output 文件（如果需要的话取消注释）"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "686c0ef2-8de3-40de-bcca-c92d9c05420d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# with open(f\"epoch-{num_epoch}_size-{data_size}_ratio-{support_ratio}.json\", \"w\", encoding='UTF-8') as outfile:\n",
    "#     json.dump(model_response, outfile, indent=4, ensure_ascii=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5891682d-c867-436a-9a56-17d63dc8c6e4",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "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.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}