Demos/06. 尝试微调 LLM：让它会写唐诗.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c9d7a78c-068d-4daa-80f0-b46436ef3268",
   "metadata": {},
   "source": [
    "# 尝试微调 LLM：让它会写唐诗\n",
    "> [GenAI HW5: LLM Fine-tuning](https://colab.research.google.com/drive/1nB3jwRJVKXSDDNO-pbURrao0N2MpqHl8?usp=sharing#scrollTo=uh5rwbr4q5Nw) 中文镜像版\n",
    ">\n",
    "> 指导文章：[08. 尝试微调 LLM：让它会写唐诗](https://github.com/Hoper-J/LLM-Guide-and-Demos-zh_CN/blob/master/Guide/08.%20尝试微调%20LLM：让它会写唐诗.md)\n",
    "\n",
    "在此作业中，你将微调自己的 LLM，使其能够写唐诗。有关详细信息，请参阅 [作业PDF](https://github.com/Hoper-J/LLM-Guide-and-Demos-zh_CN/blob/master/GenAI_PDF/HW5.pdf)\n",
    "\n",
    "**TODOs**\n",
    "1. 阅读幻灯片并确保你理解了作业的目标。\n",
    "2. 按照此笔记本中的步骤微调你的 LLM。\n",
    "\n",
    "在这里你不需要去从 0 开始写你的微调代码，所有的函数主体都已经给出。\n",
    "\n",
    "在线链接：[Kaggle](https://www.kaggle.com/code/aidemos/06-llm) | [Colab](https://colab.research.google.com/drive/1u2xgN5gWnZCwPlH2R-0gz_CpI9gzLdiy?usp=sharing)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23a96b52-6163-40db-a4ae-cc0edadf40c8",
   "metadata": {},
   "source": [
    "## 安装必要的库\n",
    "我们安装并导入一些由其他人编写的优秀库，以便简化微调过程。\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66e2bfe8-7311-4f17-844a-b2ec1e8782a6",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!pip install bitsandbytes\n",
    "!pip install datasets\n",
    "!pip install transformers\n",
    "!pip install peft\n",
    "!pip install sentencepiece\n",
    "!pip install colorama\n",
    "!pip install fsspec\n",
    "!pip install accelerate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00ad15d0-3d1d-4200-9ec5-a73a9dc1a1c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "import argparse\n",
    "import json\n",
    "import warnings\n",
    "import logging\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import bitsandbytes as bnb\n",
    "from datasets import load_dataset, load_from_disk\n",
    "import transformers\n",
    "from peft import PeftModel\n",
    "from colorama import Fore, Style\n",
    "\n",
    "from tqdm import tqdm\n",
    "from transformers import (\n",
    "    AutoTokenizer,\n",
    "    AutoConfig,\n",
    "    AutoModelForCausalLM,\n",
    "    BitsAndBytesConfig,\n",
    "    GenerationConfig\n",
    ")\n",
    "from peft import (\n",
    "    LoraConfig,\n",
    "    get_peft_model,\n",
    "    get_peft_model_state_dict,\n",
    "    prepare_model_for_kbit_training\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6858a26-4a85-4816-98f9-6d2bd0d5736f",
   "metadata": {},
   "source": [
    "## 下载用于微调的数据集"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f534a8fd-13af-401d-ae63-f6d5320c5cba",
   "metadata": {},
   "outputs": [],
   "source": [
    "!git clone https://github.com/CheeEn-Yu/GenAI-Hw5.git"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f551dec-6c8d-4b62-acbf-6733fc30fbec",
   "metadata": {},
   "source": [
    "## 固定随机种子\n",
    "为了使微调结果可重复，我们固定随机种子。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c15ad9c4-8ac7-4551-96ee-f3b36d0e054a",
   "metadata": {},
   "outputs": [],
   "source": [
    "seed = 42\n",
    "torch.backends.cudnn.deterministic = True\n",
    "torch.backends.cudnn.benchmark = False\n",
    "torch.manual_seed(seed)\n",
    "if torch.cuda.is_available():\n",
    "    torch.cuda.manual_seed_all(seed)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11476e6d-ffd7-4042-9265-eb86faa1afd4",
   "metadata": {},
   "source": [
    "## 定义一些有用的函数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0881a399-70fa-4d01-95e9-3195bee74855",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 生成训练数据\n",
    "def generate_training_data(data_point):\n",
    "    \"\"\"\n",
    "    将输入和输出文本转换为模型可读取的 tokens。\n",
    "\n",
    "    参数：\n",
    "    - data_point: 包含 \"instruction\"、\"input\" 和 \"output\" 字段的字典。\n",
    "\n",
    "    返回：\n",
    "    - 包含模型输入 IDs、标签和注意力掩码的字典。\n",
    "\n",
    "    示例:\n",
    "    - 如果你构建了一个字典 data_point_1，并包含字段 \"instruction\"、\"input\" 和 \"output\"，你可以像这样使用函数：\n",
    "        generate_training_data(data_point_1)\n",
    "    \"\"\"\n",
    "    # 构建完整的输入提示词\n",
    "    prompt = f\"\"\"\\\n",
    "[INST] <<SYS>>\n",
    "You are a helpful assistant and good at writing Tang poem. 你是一個樂於助人的助手且擅長寫唐詩。\n",
    "<</SYS>>\n",
    "\n",
    "{data_point[\"instruction\"]}\n",
    "{data_point[\"input\"]}\n",
    "[/INST]\"\"\"\n",
    "\n",
    "    # 计算用户提示词的 token 数量\n",
    "    len_user_prompt_tokens = (\n",
    "        len(\n",
    "            tokenizer(\n",
    "                prompt,\n",
    "                truncation=True,\n",
    "                max_length=CUTOFF_LEN + 1,\n",
    "                padding=\"max_length\",\n",
    "            )[\"input_ids\"]\n",
    "        ) - 1\n",
    "    )\n",
    "\n",
    "    # 将完整的输入和输出转换为 tokens\n",
    "    full_tokens = tokenizer(\n",
    "        prompt + \" \" + data_point[\"output\"] + \"</s>\",\n",
    "        truncation=True,\n",
    "        max_length=CUTOFF_LEN + 1,\n",
    "        padding=\"max_length\",\n",
    "    )[\"input_ids\"][:-1]\n",
    "\n",
    "    return {\n",
    "        \"input_ids\": full_tokens,\n",
    "        \"labels\": [-100] * len_user_prompt_tokens + full_tokens[len_user_prompt_tokens:],\n",
    "        \"attention_mask\": [1] * len(full_tokens),\n",
    "    }\n",
    "\n",
    "# 评估模型生成的回复\n",
    "def evaluate(instruction, generation_config, max_len, input_text=\"\", verbose=True):\n",
    "    \"\"\"\n",
    "    获取模型在给定输入下的生成结果。\n",
    "\n",
    "    参数：\n",
    "    - instruction: 描述任务的字符串。\n",
    "    - generation_config: 模型生成配置。\n",
    "    - max_len: 最大生成长度。\n",
    "    - input_text: 输入文本，默认为空字符串。\n",
    "    - verbose: 是否打印生成结果。\n",
    "\n",
    "    返回：\n",
    "    - output: 模型生成的文本。\n",
    "    \"\"\"\n",
    "    # 构建完整的输入提示词\n",
    "    prompt = f\"\"\"\\\n",
    "[INST] <<SYS>>\n",
    "You are a helpful assistant and good at writing Tang poem. 你是一個樂於助人的助手且擅長寫唐詩。\n",
    "<</SYS>>\n",
    "\n",
    "{instruction}\n",
    "{input_text}\n",
    "[/INST]\"\"\"\n",
    "\n",
    "    # 将提示词转换为模型所需的 token 格式\n",
    "    inputs = tokenizer(prompt, return_tensors=\"pt\")\n",
    "    input_ids = inputs[\"input_ids\"].cuda()\n",
    "    \n",
    "    # 使用模型生成回复\n",
    "    generation_output = model.generate(\n",
    "        input_ids=input_ids,\n",
    "        generation_config=generation_config,\n",
    "        return_dict_in_generate=True,\n",
    "        output_scores=True,\n",
    "        max_new_tokens=max_len,\n",
    "    )\n",
    "    \n",
    "    # 解码并打印生成的回复\n",
    "    for s in generation_output.sequences:\n",
    "        output = tokenizer.decode(s)\n",
    "        output = output.split(\"[/INST]\")[1].replace(\"</s>\", \"\").replace(\"<s>\", \"\").replace(\"Assistant:\", \"\").replace(\"Assistant\", \"\").strip()\n",
    "        if verbose:\n",
    "            print(output)\n",
    "    \n",
    "    return output"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3059e858-1356-4cc7-b662-5f998989d953",
   "metadata": {},
   "source": [
    "## 下载模型并在微调前进行推理"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f30ac92-21ae-4f97-95a0-5566c8bfa296",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\" 你可以（但不一定需要）更改 LLM 模型 \"\"\"\n",
    "\n",
    "model_name = \"MediaTek-Research/Breeze-7B-Instruct-v0_1\"  # 默认选择 MediaTek Breeze 7B 模型\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5561de8b-2f15-4ca4-a7a7-a7b34879f543",
   "metadata": {},
   "source": [
    "### 微调前的推理\n",
    "让我们先看看在未进行微调的情况下，模型的表现如何。\n",
    "\n",
    "#### 加载模型\n",
    "注意，下面这段代码会占用大概 5.*G 的显存"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61a17e79-2c37-4cdb-ac64-37ba6201efa5",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\" 建议不要更改此单元格中的代码 \"\"\"\n",
    "\n",
    "cache_dir = \"./cache\"\n",
    "\n",
    "nf4_config = BitsAndBytesConfig(\n",
    "   load_in_4bit=True,\n",
    "   bnb_4bit_quant_type=\"nf4\",\n",
    "   bnb_4bit_use_double_quant=True,\n",
    "   bnb_4bit_compute_dtype=torch.bfloat16\n",
    ")\n",
    "\n",
    "# 从指定模型名称或路径加载预训练语言模型\n",
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    model_name,\n",
    "    cache_dir=cache_dir,\n",
    "    quantization_config=nf4_config,\n",
    "    low_cpu_mem_usage=True\n",
    ")\n",
    "\n",
    "# 创建 tokenizer 并设置结束符号 (eos_token)\n",
    "logging.getLogger('transformers').setLevel(logging.ERROR)\n",
    "tokenizer = AutoTokenizer.from_pretrained(\n",
    "    model_name,\n",
    "    add_eos_token=True,\n",
    "    cache_dir=cache_dir,\n",
    "    quantization_config=nf4_config\n",
    ")\n",
    "tokenizer.pad_token = tokenizer.eos_token\n",
    "\n",
    "# 设置模型推理时的解码参数\n",
    "max_len = 128\n",
    "generation_config = GenerationConfig(\n",
    "    do_sample=True,\n",
    "    temperature=0.1,\n",
    "    num_beams=1,\n",
    "    top_p=0.3,\n",
    "    no_repeat_ngram_size=3,\n",
    "    pad_token_id=2,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "087f9dbf-78d0-401c-af29-063d55b91707",
   "metadata": {},
   "source": [
    "#### 初始表现"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2523080b-8126-4869-8836-7d70002f7395",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\" 建议不要更改此单元格中的代码，样例和Prompt都保持繁体 \"\"\"\n",
    "\n",
    "# 测试样例\n",
    "test_tang_list = [\n",
    "    '相見時難別亦難，東風無力百花殘。',\n",
    "    '重帷深下莫愁堂，臥後清宵細細長。',\n",
    "    '芳辰追逸趣，禁苑信多奇。'\n",
    "]\n",
    "\n",
    "# 获取每个样例的模型输出\n",
    "demo_before_finetune = []\n",
    "for tang in test_tang_list:\n",
    "    demo_before_finetune.append(\n",
    "        f'模型輸入:\\n以下是一首唐詩的第一句話，請用你的知識判斷並完成整首詩。{tang}\\n\\n模型輸出:\\n' +\n",
    "        evaluate('以下是一首唐詩的第一句話，請用你的知識判斷並完成整首詩。', generation_config, max_len, tang, verbose=False)\n",
    "    )\n",
    "\n",
    "# 打印并将输出存储到文本文件\n",
    "for idx in range(len(demo_before_finetune)):\n",
    "    print(f\"Example {idx + 1}:\")\n",
    "    print(demo_before_finetune[idx])\n",
    "    print(\"-\" * 80)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "47321fad-4af5-4f3b-87bf-a5e17b292dbe",
   "metadata": {},
   "source": [
    "## 设置用于微调的超参数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11382564-149a-4628-bb2d-2b9f8ad0dfca",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\" 强烈建议你尝试调整这个参数 \"\"\"\n",
    "\n",
    "num_train_data = 1040  # 设置用于训练的数据量，最大值为5000。通常，训练数据越多越好，模型会见到更多样化的诗句，从而提高生成质量，但也会增加训练时间。\n",
    "                      # 使用默认参数(1040)：微调大约需要25分钟，完整运行所有单元大约需要50分钟。\n",
    "                      # 使用最大值(5000)：微调大约需要100分钟，完整运行所有单元大约需要120分钟。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "baa7ff2e-390f-438c-b839-2de169e88f9f",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\" 你可以（但不一定需要）更改这些超参数 \"\"\"\n",
    "\n",
    "output_dir = \"./output\"  # 设置作业结果输出目录。\n",
    "ckpt_dir = \"./exp1\"  # 设置 model checkpoint 保存目录（如果想将 model checkpoints 保存到其他目录下，可以修改这里）。\n",
    "num_epoch = 1  # 设置训练的总 Epoch 数（数值越高，训练时间越长，若使用免费版的 Colab 需要注意时间太长可能会断线，本地运行不需要担心）。\n",
    "LEARNING_RATE = 3e-4  # 设置学习率"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dfbdc6bf-bc98-4626-a111-e05a5a26a61b",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\" 建议不要更改此单元格中的代码 \"\"\"\n",
    "\n",
    "cache_dir = \"./cache\"  # 设置缓存目录路径\n",
    "from_ckpt = False  # 是否从 checkpoint 加载模型权重，默认值为否\n",
    "ckpt_name = None  # 加载特定 checkpoint 时使用的文件名，默认值为无\n",
    "dataset_dir = \"./GenAI-Hw5/Tang_training_data.json\"  # 设置数据集目录或文件路径\n",
    "logging_steps = 20  # 定义训练过程中每隔多少步骤输出一次日志\n",
    "save_steps = 65  # 定义训练过程中每隔多少步骤保存一次模型\n",
    "save_total_limit = 3  # 控制最多保留多少个模型 checkpoint\n",
    "report_to = \"none\"  # 设置不上报实验指标，也可以设置为 \"wandb\"，此时需要获取对应的 API，见：https://github.com/Hoper-J/AI-Guide-and-Demos-zh_CN/pull/5\n",
    "MICRO_BATCH_SIZE = 4  # 定义微批次大小\n",
    "BATCH_SIZE = 16  # 定义一个批次的大小\n",
    "GRADIENT_ACCUMULATION_STEPS = BATCH_SIZE // MICRO_BATCH_SIZE  # 计算每个微批次累积的梯度步骤\n",
    "CUTOFF_LEN = 256  # 设置文本截断的最大长度\n",
    "LORA_R = 8  # 设置 LORA（Layer-wise Random Attention）的 R 值\n",
    "LORA_ALPHA = 16  # 设置 LORA 的 Alpha 值\n",
    "LORA_DROPOUT = 0.05  # 设置 LORA 的 Dropout 率\n",
    "VAL_SET_SIZE = 0  # 设置验证集的大小，默认值为无\n",
    "TARGET_MODULES = [\"q_proj\", \"up_proj\", \"o_proj\", \"k_proj\", \"down_proj\", \"gate_proj\", \"v_proj\"]  # 设置目标模块，这些模块的权重将被保存为 checkpoint。\n",
    "device_map = \"auto\"  # 设置设备映射，默认值为 \"auto\"\n",
    "world_size = int(os.environ.get(\"WORLD_SIZE\", 1))  # 获取环境变量 \"WORLD_SIZE\" 的值，若未设置则默认为 1\n",
    "ddp = world_size != 1  # 根据 world_size 判断是否使用分布式数据处理(DDP)，若 world_size 为 1 则不使用 DDP\n",
    "if ddp:\n",
    "    device_map = {\"\": int(os.environ.get(\"LOCAL_RANK\") or 0)}\n",
    "    GRADIENT_ACCUMULATION_STEPS = GRADIENT_ACCUMULATION_STEPS // world_size"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f8203b78-ffb4-4504-9d52-12db44e53562",
   "metadata": {},
   "source": [
    "## 开始微调\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67626ef4-2636-487c-a490-b08b869040ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\" 建议不要更改此单元格中的代码 \"\"\"\n",
    "\n",
    "# 设置TOKENIZERS_PARALLELISM为false，这里简单禁用并行性以避免报错\n",
    "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n",
    "\n",
    "# 创建指定的输出目录\n",
    "os.makedirs(output_dir, exist_ok=True)\n",
    "os.makedirs(ckpt_dir, exist_ok=True)\n",
    "\n",
    "# 根据 from_ckpt 标志，从 checkpoint 加载模型权重\n",
    "if from_ckpt:\n",
    "    model = PeftModel.from_pretrained(model, ckpt_name)\n",
    "\n",
    "# 准备模型以使用 INT8 进行训练\n",
    "model = prepare_model_for_kbit_training(model)\n",
    "\n",
    "# 使用 LoraConfig 配置 LORA 模型\n",
    "config = LoraConfig(\n",
    "    r=LORA_R,\n",
    "    lora_alpha=LORA_ALPHA,\n",
    "    target_modules=TARGET_MODULES,\n",
    "    lora_dropout=LORA_DROPOUT,\n",
    "    bias=\"none\",\n",
    "    task_type=\"CAUSAL_LM\",\n",
    ")\n",
    "model = get_peft_model(model, config)\n",
    "\n",
    "# 将 tokenizer 的填充 token 设置为 0\n",
    "tokenizer.pad_token_id = 0\n",
    "\n",
    "# 加载并处理训练数据\n",
    "with open(dataset_dir, \"r\", encoding=\"utf-8\") as f:\n",
    "    data_json = json.load(f)\n",
    "with open(\"tmp_dataset.json\", \"w\", encoding=\"utf-8\") as f:\n",
    "    json.dump(data_json[:num_train_data], f, indent=2, ensure_ascii=False)\n",
    "\n",
    "data = load_dataset('json', data_files=\"tmp_dataset.json\", download_mode=\"force_redownload\")\n",
    "\n",
    "# 将训练数据分为训练集和验证集（若 VAL_SET_SIZE 大于 0）\n",
    "if VAL_SET_SIZE > 0:\n",
    "    train_val = data[\"train\"].train_test_split(\n",
    "        test_size=VAL_SET_SIZE, shuffle=True, seed=42\n",
    "    )\n",
    "    train_data = train_val[\"train\"].shuffle().map(generate_training_data)\n",
    "    val_data = train_val[\"test\"].shuffle().map(generate_training_data)\n",
    "else:\n",
    "    train_data = data['train'].shuffle().map(generate_training_data)\n",
    "    val_data = None\n",
    "\n",
    "# 使用 Transformers Trainer 进行模型训练\n",
    "trainer = transformers.Trainer(\n",
    "    model=model,\n",
    "    train_dataset=train_data,\n",
    "    eval_dataset=val_data,\n",
    "    args=transformers.TrainingArguments(\n",
    "        per_device_train_batch_size=MICRO_BATCH_SIZE,\n",
    "        gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS,\n",
    "        warmup_steps=50,\n",
    "        num_train_epochs=num_epoch,\n",
    "        learning_rate=LEARNING_RATE,\n",
    "        fp16=True,  # 使用混合精度训练\n",
    "        logging_steps=logging_steps,\n",
    "        save_strategy=\"steps\",\n",
    "        save_steps=save_steps,\n",
    "        output_dir=ckpt_dir,\n",
    "        save_total_limit=save_total_limit,\n",
    "        ddp_find_unused_parameters=False if ddp else None,  # 是否使用 DDP，控制梯度更新策略\n",
    "        report_to=report_to,\n",
    "    ),\n",
    "    data_collator=transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False),\n",
    ")\n",
    "\n",
    "# 禁用模型的缓存功能\n",
    "model.config.use_cache = False\n",
    "\n",
    "# 若使用 PyTorch 2.0 以上版本且非 Windows 系统，编译模型\n",
    "if torch.__version__ >= \"2\" and sys.platform != 'win32':\n",
    "    model = torch.compile(model)\n",
    "\n",
    "# 开始模型训练\n",
    "trainer.train()\n",
    "\n",
    "# 将训练好的模型保存到指定目录\n",
    "model.save_pretrained(ckpt_dir)\n",
    "\n",
    "# 打印训练过程中可能出现的缺失权重警告信息\n",
    "print(\"\\n 如果上方有关于缺少键的警告，请忽略 :)\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a046650-5f92-4109-94b3-0c0e029eb16f",
   "metadata": {},
   "source": [
    "## 测试\n",
    "微调过程完成后，我们希望测试模型能否完成之前失败的任务。\n",
    "\n",
    "我们首先需要加载保存的微调模型的 checkpoint。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78799ffd-4720-49d8-9f1a-b680cb6da9d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\" 建议不要更改此单元格中的代码 \"\"\"\n",
    "\n",
    "# 查找所有可用的 checkpoints\n",
    "ckpts = []\n",
    "for ckpt in os.listdir(ckpt_dir):\n",
    "    if ckpt.startswith(\"checkpoint-\"):\n",
    "        ckpts.append(ckpt)\n",
    "\n",
    "# 列出所有的 checkpoints\n",
    "ckpts = sorted(ckpts, key=lambda ckpt: int(ckpt.split(\"-\")[-1]))\n",
    "print(\"所有可用的 checkpoints：\")\n",
    "print(\" id: checkpoint 名称\")\n",
    "for (i, ckpt) in enumerate(ckpts):\n",
    "    print(f\"{i:>3}: {ckpt}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77e2c0cb-70a1-4fe0-8524-4e9a9fcc3730",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\" 你可以（但不一定需要）更改 checkpoint \"\"\"\n",
    "\n",
    "id_of_ckpt_to_use = -1  # 要用于推理的 checkpoint 的 id（对应上一单元格的输出结果）。\n",
    "                        # 默认值 -1 表示使用列出的最后一个 checkpoint。\n",
    "                        # 如果你想选择其他 checkpoint，可以将 -1 更改为列出的 checkpoint id 中的某一个。\n",
    "\n",
    "ckpt_name = os.path.join(ckpt_dir, ckpts[id_of_ckpt_to_use])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c41c045f-304d-4b52-b836-a3a6f0cb2f7a",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\" 你可以（但不一定需要）更改解码参数 \"\"\"\n",
    "# 你可以在此处调整解码参数，解码参数的详细解释请见作业幻灯片。\n",
    "max_len = 128  # 生成回复的最大长度\n",
    "temperature = 0.1  # 设置生成回复的随机度，值越小生成的回复越稳定。\n",
    "top_p = 0.3  # Top-p (nucleus) 采样的概率阈值，用于控制生成回复的多样性。\n",
    "# top_k = 5  # 调整 Top-k 值，以增加生成回复的多样性并避免生成重复的词汇。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00d68b83-99fe-4b84-87e0-e2e8ce8be15b",
   "metadata": {},
   "source": [
    "### 释放显存\n",
    "\n",
    "增加了下面的代码释放之前占用的显存，防止你遇到 OutOfMemoryError。（如果显存足够大可以注释）"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7515e2dd-72ed-4a7f-9b15-4a2e2395f00e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import gc\n",
    "\n",
    "# 删除模型和 tokenizer 对象\n",
    "del model\n",
    "del tokenizer\n",
    "\n",
    "# 调用垃圾回收机制，强制释放未使用的内存\n",
    "gc.collect()\n",
    "\n",
    "# 清理 GPU 缓存\n",
    "torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f3e3ef0-b905-4e32-8cb9-f8684d2b559c",
   "metadata": {},
   "source": [
    "### 加载模型和分词器"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7236095d-d605-478a-b50b-7e5a761728a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_data_path = \"GenAI-Hw5/Tang_testing_data.json\"  # 测试数据集的路径\n",
    "output_path = os.path.join(output_dir, \"results.txt\")  # 生成结果的输出路径\n",
    "\n",
    "cache_dir = \"./cache\"  # 设置缓存目录\n",
    "seed = 42  # 设置随机种子以重现结果\n",
    "no_repeat_ngram_size = 3  # 设置禁止重复 Ngram 的大小，避免生成重复片段\n",
    "\n",
    "# 配置模型的量化设置，使用 4 位精度\n",
    "nf4_config = BitsAndBytesConfig(\n",
    "   load_in_4bit=True,\n",
    "   bnb_4bit_quant_type=\"nf4\",\n",
    "   bnb_4bit_use_double_quant=True,\n",
    "   bnb_4bit_compute_dtype=torch.bfloat16\n",
    ")\n",
    "\n",
    "# 从预训练模型加载 tokenizer\n",
    "tokenizer = AutoTokenizer.from_pretrained(\n",
    "    model_name,\n",
    "    cache_dir=cache_dir,\n",
    "    quantization_config=nf4_config\n",
    ")\n",
    "\n",
    "# 从预训练模型加载语言模型，使用量化配置并指定设备\n",
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    model_name,\n",
    "    quantization_config=nf4_config,\n",
    "    device_map={'': 0},  # 指定使用的设备，GPU 0\n",
    "    cache_dir=cache_dir\n",
    ")\n",
    "\n",
    "# 从 checkpoint 加载已保存的模型权重\n",
    "model = PeftModel.from_pretrained(model, ckpt_name, device_map={'': 0})\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab3fcaf1-63aa-408f-a54c-7f4d9ce996fc",
   "metadata": {},
   "source": [
    "### 生成测试结果"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b89508bf-2b44-4040-9792-2aa16f2e5314",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\" 建议不要更改此单元格中的代码 \"\"\"\n",
    "\n",
    "results = []\n",
    "\n",
    "# 设置生成配置，包括随机度、束搜索等参数\n",
    "generation_config = GenerationConfig(\n",
    "    do_sample=True,\n",
    "    temperature=temperature,\n",
    "    num_beams=1,\n",
    "    top_p=top_p,\n",
    "    # top_k=top_k,  # 如果需要使用 top-k，可以在此设置\n",
    "    no_repeat_ngram_size=no_repeat_ngram_size,\n",
    "    pad_token_id=2\n",
    ")\n",
    "\n",
    "# 读取测试数据集\n",
    "with open(test_data_path, \"r\", encoding=\"utf-8\") as f:\n",
    "    test_datas = json.load(f)\n",
    "\n",
    "# 对每个测试样例生成预测，并保存结果\n",
    "with open(output_path, \"w\", encoding=\"utf-8\") as f:\n",
    "    for (i, test_data) in enumerate(test_datas):\n",
    "        predict = evaluate(test_data[\"instruction\"], generation_config, max_len, test_data[\"input\"], verbose=False)\n",
    "        f.write(f\"{i+1}. \" + test_data[\"input\"] + predict + \"\\n\")\n",
    "        print(f\"{i+1}. \" + test_data[\"input\"] + predict)\n",
    "      "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a0a0aae-405d-43ef-a7fc-936a0e213df0",
   "metadata": {},
   "source": [
    "## 查看微调后的模型与微调前的对比\n",
    "\n",
    "我们现在检查微调后的模型在之前看到的例子上的表现。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2198a01-0856-412c-8d3f-4cf7bfd38f6b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 使用之前的测试例子\n",
    "test_tang_list = [\n",
    "    '相見時難別亦難，東風無力百花殘。',\n",
    "    '重帷深下莫愁堂，臥後清宵細細長。',\n",
    "    '芳辰追逸趣，禁苑信多奇。'\n",
    "]\n",
    "\n",
    "# 使用微调后的模型进行推理\n",
    "demo_after_finetune = []\n",
    "for tang in test_tang_list:\n",
    "    demo_after_finetune.append(\n",
    "        f'模型輸入:\\n以下是一首唐詩的第一句話，請用你的知識判斷並完成整首詩。{tang}\\n\\n模型輸出:\\n' +\n",
    "        evaluate('以下是一首唐詩的第一句話，請用你的知識判斷並完成整首詩。', generation_config, max_len, tang, verbose=False)\n",
    "    )\n",
    "\n",
    "# 打印输出结果\n",
    "for idx in range(len(demo_after_finetune)):\n",
    "    print(f\"Example {idx + 1}:\")\n",
    "    print(demo_after_finetune[idx])\n",
    "    print(\"-\" * 80)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64120aa4-1a72-42b3-bedd-a223c18506cc",
   "metadata": {},
   "source": [
    "## 参考链接\n",
    "\n",
    "[唐诗数据集](https://github.com/chinese-poetry/chinese-poetry/tree/master/%E5%85%A8%E5%94%90%E8%AF%97?fbclid=IwAR2bM14S42T-VtrvMi3wywCqKfYJraBtMl7QVTo0qyPMjX9jj9Vj3JepFBA)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "113376f7-8e5e-4efc-8ff7-97848e5d65c2",
   "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
}