{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c219841f-493c-40f9-a6c9-3700f0c525d0",
   "metadata": {},
   "source": [
    "# PEFT 库 LoRA 实战 - OpenAI Whisper-large-v2\n",
    "\n",
    "本教程使用 LoRA 在`OpenAI Whisper-large-v2`模型上实现`语音识别(ASR)`任务的微调训练。\n",
    "\n",
    "我们还结合了`int8` 量化进一步降低训练过程资源开销，同时保证了精度几乎不受影响。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d0a1e23-ea71-45d6-82d6-453077cf2d29",
   "metadata": {},
   "source": [
    "## 全局参数设置"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "19d11aa3-9a73-4ce9-b6c5-a65a2fcb07c3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_DEVICE_ORDER\"]=\"PCI_BUS_ID\"   # see issue #152\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"2\"\n",
    "\n",
    "import torch\n",
    "torch.cuda.device_count()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "ccd00402-d821-485e-8703-fb16bcb56a9e",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name_or_path = \"openai/whisper-large-v2\"\n",
    "language = \"Chinese (China)\"\n",
    "language_abbr = \"zh-CN\"\n",
    "task = \"transcribe\"\n",
    "dataset_name = \"mozilla-foundation/common_voice_11_0\"\n",
    "\n",
    "batch_size=64"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cfffa1df-e51e-4026-9817-1cebddf0061a",
   "metadata": {},
   "source": [
    "## 下载数据集 Common Voice\n",
    "\n",
    "Common Voice 11.0 数据集包含许多不同语言的录音，总时长达数小时。\n",
    "\n",
    "本教程以中文数据为例，展示如何使用 LoRA 在 Whisper-large-v2 上进行微调训练。\n",
    "\n",
    "首先，初始化一个DatasetDict结构，并将训练集（将训练+验证拆分为训练集）和测试集拆分好，按照中文数据集构建配置加载到内存中："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "21ff42f4-f3ec-46d3-b0c0-dd9ffbf7b50b",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'client_id': '95368aab163e0387e4fd4991b4f2d8ccfbd4364bf656c860230501fd27dcedf087773e4695a6cf5de9c4f1d406d582283190d065cdfa36b0e2b060cffaca977e',\n",
       " 'path': '/store/jxzhang/.cache/huggingface/datasets/downloads/extracted/edf8cf7fef3457433a3a59929c4c4809972172377467a8f189ac185f3d5e4b53/zh-CN_train_0/common_voice_zh-CN_33211332.mp3',\n",
       " 'audio': {'path': '/store/jxzhang/.cache/huggingface/datasets/downloads/extracted/edf8cf7fef3457433a3a59929c4c4809972172377467a8f189ac185f3d5e4b53/zh-CN_train_0/common_voice_zh-CN_33211332.mp3',\n",
       "  'array': array([-6.82121026e-13, -2.27373675e-12, -2.27373675e-12, ...,\n",
       "          1.21667399e-05,  3.23003678e-06, -2.43066324e-07]),\n",
       "  'sampling_rate': 48000},\n",
       " 'sentence': '性喜温暖润湿气候且耐寒。',\n",
       " 'up_votes': 2,\n",
       " 'down_votes': 0,\n",
       " 'age': '',\n",
       " 'gender': '',\n",
       " 'accent': '',\n",
       " 'locale': 'zh-CN',\n",
       " 'segment': ''}"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from datasets import load_dataset\n",
    "from datasets import load_dataset, DatasetDict\n",
    "\n",
    "common_voice = DatasetDict()\n",
    "\n",
    "common_voice[\"train\"] = load_dataset(dataset_name, language_abbr, split=\"train+validation\", trust_remote_code=True)\n",
    "common_voice[\"test\"] = load_dataset(dataset_name, language_abbr, split=\"test\", trust_remote_code=True)\n",
    "common_voice[\"train\"][0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c81faa4-d8fe-4cc7-afe6-4c2615b9050f",
   "metadata": {},
   "source": [
    "## 预处理训练数据集\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "5822025f-7f8e-4141-8bfe-d8822d0da20f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoFeatureExtractor, AutoTokenizer, AutoProcessor\n",
    "\n",
    "feature_extractor = AutoFeatureExtractor.from_pretrained(model_name_or_path)\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(\n",
    "    model_name_or_path, language=language, task=task)\n",
    "\n",
    "processor = AutoProcessor.from_pretrained(\n",
    "    model_name_or_path, language=language, task=task)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f394e5cd-23b8-413e-8bde-88c3542b84fa",
   "metadata": {},
   "source": [
    "#### 移除数据集中不必要的字段"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "1690dc5a-c1f7-4556-9be3-d31ad888e52e",
   "metadata": {},
   "outputs": [],
   "source": [
    "common_voice = common_voice.remove_columns(\n",
    "    [\"accent\", \"age\", \"client_id\", \"down_votes\", \"gender\", \"locale\", \"path\", \"segment\", \"up_votes\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "309aff16-ea26-4474-af54-7ef244783999",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'audio': {'path': '/store/jxzhang/.cache/huggingface/datasets/downloads/extracted/edf8cf7fef3457433a3a59929c4c4809972172377467a8f189ac185f3d5e4b53/zh-CN_train_0/common_voice_zh-CN_33211332.mp3',\n",
       "  'array': array([-6.82121026e-13, -2.27373675e-12, -2.27373675e-12, ...,\n",
       "          1.21667399e-05,  3.23003678e-06, -2.43066324e-07]),\n",
       "  'sampling_rate': 48000},\n",
       " 'sentence': '性喜温暖润湿气候且耐寒。'}"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "common_voice[\"train\"][0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "881546ab-72e4-4bcf-852f-a8be736164b7",
   "metadata": {},
   "source": [
    "#### 降采样音频数据\n",
    "\n",
    "查看`common_voice` 数据集介绍，你会发现其音频是以48kHz的采样率进行采样的.\n",
    "\n",
    "而`Whisper`模型是在16kHZ的音频输入上预训练的，因此我们需要将音频输入降采样以匹配模型预训练时使用的采样率。\n",
    "\n",
    "通过在音频列上使用`cast_column`方法，并将`sampling_rate`设置为16kHz来对音频进行降采样。\n",
    "\n",
    "下次调用时，音频输入将实时重新取样："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "5fc451cc-e21e-473c-a702-d7d6ed098f91",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import Audio\n",
    "\n",
    "common_voice = common_voice.cast_column(\"audio\", Audio(sampling_rate=16000))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "cc3d7fcc-7c34-41c8-9857-5a6e883f6115",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'audio': {'path': '/store/jxzhang/.cache/huggingface/datasets/downloads/extracted/edf8cf7fef3457433a3a59929c4c4809972172377467a8f189ac185f3d5e4b53/zh-CN_train_0/common_voice_zh-CN_33211332.mp3',\n",
       "  'array': array([ 5.09317033e-11, -7.27595761e-12, -6.54836185e-11, ...,\n",
       "         -5.96661994e-06,  2.71382887e-05,  1.29687978e-05]),\n",
       "  'sampling_rate': 16000},\n",
       " 'sentence': '性喜温暖润湿气候且耐寒。'}"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "common_voice[\"train\"][0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee55908f-3ea3-4aee-8062-6f8d3a6573b9",
   "metadata": {},
   "source": [
    "### 整合以上数据处理为一个函数\n",
    "\n",
    "该数据预处理函数应该包括：\n",
    "- 通过加载音频列将音频输入重新采样为16kHZ。\n",
    "- 使用特征提取器从音频数组计算输入特征。\n",
    "- 将句子列标记化为输入标签。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "58f42c35-35ba-4d6b-9d15-095963cec67c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def prepare_dataset(batch):\n",
    "    audio = batch[\"audio\"]\n",
    "    batch[\"input_features\"] = feature_extractor(audio[\"array\"], sampling_rate=audio[\"sampling_rate\"]).input_features[0]\n",
    "    batch[\"labels\"] = tokenizer(batch[\"sentence\"]).input_ids\n",
    "    return batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "392f7856-a720-40a7-af7e-40e185fc315b",
   "metadata": {},
   "outputs": [],
   "source": [
    "common_voice = common_voice.map(\n",
    "    prepare_dataset, remove_columns=common_voice.column_names[\"train\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84ec184e-d840-40b6-99af-d11392273442",
   "metadata": {},
   "source": [
    "创建一个`DataCollator`类来将每个批次中的`attention_mask`填充到最大长度，并用`-100`替换填充值，以便在损失函数中被忽略。\n",
    "\n",
    "然后初始化数据收集器的实例："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "4c89ffcf-c805-48c2-b7d3-ae01b687178c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "from dataclasses import dataclass\n",
    "from typing import Any, Dict, List, Union\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class DataCollatorSpeechSeq2SeqWithPadding:\n",
    "    processor: Any\n",
    "\n",
    "    def __call__(self, features: List[Dict[str, Union[List[int], torch.Tensor]]]) -> Dict[str, torch.Tensor]:\n",
    "        input_features = [{\"input_features\": feature[\"input_features\"]} for feature in features]\n",
    "        batch = self.processor.feature_extractor.pad(input_features, return_tensors=\"pt\")\n",
    "\n",
    "        label_features = [{\"input_ids\": feature[\"labels\"]} for feature in features]\n",
    "        labels_batch = self.processor.tokenizer.pad(label_features, return_tensors=\"pt\")\n",
    "\n",
    "        labels = labels_batch[\"input_ids\"].masked_fill(labels_batch.attention_mask.ne(1), -100)\n",
    "\n",
    "        if (labels[:, 0] == self.processor.tokenizer.bos_token_id).all().cpu().item():\n",
    "            labels = labels[:, 1:]\n",
    "\n",
    "        batch[\"labels\"] = labels\n",
    "\n",
    "        return batch\n",
    "\n",
    "\n",
    "data_collator = DataCollatorSpeechSeq2SeqWithPadding(processor=processor)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "80ecd4bc-01fd-4286-afe5-fe2639ae15a1",
   "metadata": {},
   "source": [
    "## 训练模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "f9fcb121-fa5c-4c30-8bdc-9ab08ab75427",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoModelForSpeechSeq2Seq\n",
    "\n",
    "model = AutoModelForSpeechSeq2Seq.from_pretrained(model_name_or_path, load_in_8bit=True, device_map=\"auto\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "2cb016f1-e6e9-4fd8-9c8b-72fd23be92d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.config.forced_decoder_ids = None\n",
    "model.config.suppress_tokens = []"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25ba1fa0-ea15-48d9-8c16-70df9f0b60b1",
   "metadata": {},
   "source": [
    "为了准备模型进行int8量化，使用 `prepare_model_for_int8_training` 函数来处理模型：\n",
    "- 将所有非int8模块转换为完全精度（fp32）以保持稳定性\n",
    "- 在输入嵌入层上添加前向钩子，计算输入隐藏状态的梯度\n",
    "- 启用渐变检查点以进行更高效的内存训练"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "1ee34359-fe1b-48f1-827c-6a8ec4a53af7",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/data/miniconda3/envs/jxzhang/lib/python3.11/site-packages/peft/utils/other.py:141: FutureWarning: prepare_model_for_int8_training is deprecated and will be removed in a future version. Use prepare_model_for_kbit_training instead.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "from peft import prepare_model_for_int8_training\n",
    "\n",
    "model = prepare_model_for_int8_training(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24b6f8a2-867f-4ed5-bad5-15ca9fd9547c",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "cdf6bc9c-6d2c-4dbf-b09e-a89cb1041c46",
   "metadata": {},
   "outputs": [],
   "source": [
    "from peft import LoraConfig, PeftModel, LoraModel, LoraConfig, get_peft_model\n",
    "\n",
    "config = LoraConfig(\n",
    "    r=8,\n",
    "    lora_alpha=64,\n",
    "    target_modules=[\"q_proj\", \"v_proj\"],\n",
    "    lora_dropout=0.05,\n",
    "    bias=\"none\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "b74c7508-e6f4-42d8-8aaf-fe83c5977c35",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "trainable params: 3,932,160 || all params: 1,547,237,120 || trainable%: 0.25414074863974306\n"
     ]
    }
   ],
   "source": [
    "model = get_peft_model(model, config)\n",
    "model.print_trainable_parameters()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1cc6b26a-3e54-4a46-9b36-a048b40a37d7",
   "metadata": {},
   "source": [
    "### 演示需要，只训练了100 steps。建议同学改为默认的 3个 epochs 完整训练一个中文语音识别模型。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "11f259c8-dbcf-4a7f-bbb5-821ab104efee",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import Seq2SeqTrainingArguments\n",
    "\n",
    "# 设置序列到序列模型训练的参数\n",
    "training_args = Seq2SeqTrainingArguments(\n",
    "    output_dir=\"models/whisper-large-v2-asr-int8\",  # 指定模型输出和保存的目录\n",
    "    per_device_train_batch_size=batch_size,  # 每个设备上的训练批量大小\n",
    "    gradient_accumulation_steps=1,  # 梯度累积步数，在每次优化器步骤之前累积的更新步数\n",
    "    learning_rate=1e-3,  # 学习率\n",
    "    warmup_steps=50,  # 在训练初期增加学习率的步数，有助于稳定训练\n",
    "    # max_steps=100, # 训练总步数\n",
    "    num_train_epochs=3,  # 训练的总轮数\n",
    "    evaluation_strategy=\"epoch\",  # 设置评估策略，这里是在每个epoch结束时进行评估\n",
    "    fp16=True,  # 启用混合精度训练，可以提高训练速度，同时减少内存使用\n",
    "    per_device_eval_batch_size=batch_size,  # 每个设备上的评估批量大小\n",
    "    generation_max_length=128,  # 生成任务的最大长度\n",
    "    logging_steps=25,  # 指定日志记录的步骤，用于跟踪训练进度\n",
    "    remove_unused_columns=False,  # 是否删除不使用的列，以减少数据处理开销\n",
    "    label_names=[\"labels\"],  # 指定标签列的名称，用于训练过程中\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c57ee183-b16f-4313-97f6-0df6c0f5f467",
   "metadata": {},
   "source": [
    "#### 训练过程保存状态的回调，长时期训练建议使用"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "2ce443d9-f309-4c03-bd74-c6842292b713",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from transformers.trainer_utils import PREFIX_CHECKPOINT_DIR\n",
    "from transformers import Seq2SeqTrainer, TrainerCallback, Seq2SeqTrainingArguments, TrainerState, TrainerControl\n",
    "\n",
    "class SavePeftModelCallback(TrainerCallback):\n",
    "    def on_save(\n",
    "        self,\n",
    "        args: Seq2SeqTrainingArguments,\n",
    "        state: TrainerState,\n",
    "        control: TrainerControl,\n",
    "        **kwargs,\n",
    "    ):\n",
    "        checkpoint_folder = os.path.join(args.output_dir, f\"{PREFIX_CHECKPOINT_DIR}-{state.global_step}\")\n",
    "\n",
    "        peft_model_path = os.path.join(checkpoint_folder, \"adapter_model\")\n",
    "        kwargs[\"model\"].save_pretrained(peft_model_path)\n",
    "\n",
    "        pytorch_model_path = os.path.join(checkpoint_folder, \"pytorch_model.bin\")\n",
    "        if os.path.exists(pytorch_model_path):\n",
    "            os.remove(pytorch_model_path)\n",
    "        return control"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "f8a52ed7-cae0-4aba-818e-87717430d908",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = Seq2SeqTrainer(\n",
    "    args=training_args,\n",
    "    model=model,\n",
    "    train_dataset=common_voice[\"train\"],\n",
    "    eval_dataset=common_voice[\"test\"],\n",
    "    data_collator=data_collator,\n",
    "    tokenizer=processor.feature_extractor,\n",
    "    callbacks=[SavePeftModelCallback],\n",
    ")\n",
    "model.config.use_cache = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "6973bed7-8f53-4d55-966c-f037941e5ef3",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/data/miniconda3/envs/jxzhang/lib/python3.11/site-packages/torch/utils/checkpoint.py:429: UserWarning: torch.utils.checkpoint: please pass in use_reentrant=True or use_reentrant=False explicitly. The default value of use_reentrant will be updated to be False in the future. To maintain current behavior, pass use_reentrant=True. It is recommended that you use use_reentrant=False. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n",
      "/data/miniconda3/envs/jxzhang/lib/python3.11/site-packages/torch/utils/checkpoint.py:61: UserWarning: None of the inputs have requires_grad=True. Gradients will be None\n",
      "  warnings.warn(\n",
      "/data/miniconda3/envs/jxzhang/lib/python3.11/site-packages/bitsandbytes/autograd/_functions.py:322: UserWarning: MatMul8bitLt: inputs will be cast from torch.float32 to float16 during quantization\n",
      "  warnings.warn(f\"MatMul8bitLt: inputs will be cast from {A.dtype} to float16 during quantization\")\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='1860' max='1860' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [1860/1860 7:48:09, Epoch 3/3]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Epoch</th>\n",
       "      <th>Training Loss</th>\n",
       "      <th>Validation Loss</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>1</td>\n",
       "      <td>0.341900</td>\n",
       "      <td>0.264468</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2</td>\n",
       "      <td>0.259600</td>\n",
       "      <td>0.248249</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3</td>\n",
       "      <td>0.214400</td>\n",
       "      <td>0.248773</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/data/miniconda3/envs/jxzhang/lib/python3.11/site-packages/torch/utils/checkpoint.py:429: UserWarning: torch.utils.checkpoint: please pass in use_reentrant=True or use_reentrant=False explicitly. The default value of use_reentrant will be updated to be False in the future. To maintain current behavior, pass use_reentrant=True. It is recommended that you use use_reentrant=False. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n",
      "/data/miniconda3/envs/jxzhang/lib/python3.11/site-packages/torch/utils/checkpoint.py:61: UserWarning: None of the inputs have requires_grad=True. Gradients will be None\n",
      "  warnings.warn(\n",
      "/data/miniconda3/envs/jxzhang/lib/python3.11/site-packages/bitsandbytes/autograd/_functions.py:322: UserWarning: MatMul8bitLt: inputs will be cast from torch.float32 to float16 during quantization\n",
      "  warnings.warn(f\"MatMul8bitLt: inputs will be cast from {A.dtype} to float16 during quantization\")\n",
      "/data/miniconda3/envs/jxzhang/lib/python3.11/site-packages/torch/utils/checkpoint.py:429: UserWarning: torch.utils.checkpoint: please pass in use_reentrant=True or use_reentrant=False explicitly. The default value of use_reentrant will be updated to be False in the future. To maintain current behavior, pass use_reentrant=True. It is recommended that you use use_reentrant=False. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n",
      "/data/miniconda3/envs/jxzhang/lib/python3.11/site-packages/torch/utils/checkpoint.py:61: UserWarning: None of the inputs have requires_grad=True. Gradients will be None\n",
      "  warnings.warn(\n",
      "/data/miniconda3/envs/jxzhang/lib/python3.11/site-packages/bitsandbytes/autograd/_functions.py:322: UserWarning: MatMul8bitLt: inputs will be cast from torch.float32 to float16 during quantization\n",
      "  warnings.warn(f\"MatMul8bitLt: inputs will be cast from {A.dtype} to float16 during quantization\")\n",
      "/data/miniconda3/envs/jxzhang/lib/python3.11/site-packages/torch/utils/checkpoint.py:429: UserWarning: torch.utils.checkpoint: please pass in use_reentrant=True or use_reentrant=False explicitly. The default value of use_reentrant will be updated to be False in the future. To maintain current behavior, pass use_reentrant=True. It is recommended that you use use_reentrant=False. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n",
      "/data/miniconda3/envs/jxzhang/lib/python3.11/site-packages/torch/utils/checkpoint.py:61: UserWarning: None of the inputs have requires_grad=True. Gradients will be None\n",
      "  warnings.warn(\n",
      "/data/miniconda3/envs/jxzhang/lib/python3.11/site-packages/bitsandbytes/autograd/_functions.py:322: UserWarning: MatMul8bitLt: inputs will be cast from torch.float32 to float16 during quantization\n",
      "  warnings.warn(f\"MatMul8bitLt: inputs will be cast from {A.dtype} to float16 during quantization\")\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "TrainOutput(global_step=1860, training_loss=0.32258521408163093, metrics={'train_runtime': 28110.1202, 'train_samples_per_second': 4.23, 'train_steps_per_second': 0.066, 'total_flos': 2.531417002463232e+20, 'train_loss': 0.32258521408163093, 'epoch': 3.0})"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "620992c3-64f5-48f9-8e66-fdc5f6a27427",
   "metadata": {},
   "source": [
    "### 保存 LoRA 模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "53310565-7313-46a7-acf1-215970fd4f8e",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.save_pretrained(\"models/whisper-large-v2-asr-int8\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dcfe9611-eee5-462f-8cb8-fed86eec76e0",
   "metadata": {},
   "source": [
    "### 使用 Pipiline 加载 LoRA 模型，实现自动语音识别任务"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "426d7520-62cb-42bb-a4dd-000aa607b105",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "5.763907432556152"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from transformers import AutoModelForSpeechSeq2Seq\n",
    "\n",
    "my_model_name_or_path = \"yqzhangjx/whisper-large-v2-asr-int8\"\n",
    "\n",
    "model = AutoModelForSpeechSeq2Seq.from_pretrained(my_model_name_or_path, device_map=\"auto\")\n",
    "\n",
    "model.get_memory_footprint()/1024**3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "7536c488-0526-4c12-baaf-b4f7c7075be6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoFeatureExtractor, AutoTokenizer, AutoProcessor\n",
    "\n",
    "feature_extractor = AutoFeatureExtractor.from_pretrained(model_name_or_path)\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, language=language, task=task)\n",
    "processor = AutoProcessor.from_pretrained(model_name_or_path, language=language, task=task)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "18181692-a143-44ee-b56c-e754d308e0ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_audio = \"data/audio/test_zh.flac\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "9d494647-082c-4e48-9486-7945618ae679",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutomaticSpeechRecognitionPipeline\n",
    "\n",
    "pipeline = AutomaticSpeechRecognitionPipeline(model=model, tokenizer=tokenizer, feature_extractor=feature_extractor)\n",
    "\n",
    "forced_decoder_ids = processor.get_decoder_prompt_ids(language=\"chinese\", task=task)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "c3eac486-169f-41ad-b9c3-69f2c27c3e1f",
   "metadata": {},
   "outputs": [],
   "source": [
    "with torch.cuda.amp.autocast():\n",
    "    text = pipeline(test_audio, generate_kwargs={\"forced_decoder_ids\": forced_decoder_ids}, max_new_tokens=255)[\"text\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "cc4b24b2-c65b-4e63-8f16-26c72039a38d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'这是一段测试用于Whisper Large V2模型的自动语音识别测试。'"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89f49787-6ab4-4bc1-91b8-a1c104c9feaf",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "0285dd19-229e-4241-b680-71e25ab51dde",
   "metadata": {},
   "source": [
    "#### Homework 1: 为中文语料的训练过程增加过程评估，观察 Train Loss 和 Validation Loss 变化；\n",
    "#### Homework 2: LoRA 模型训练完成后，使用测试集进行完整的模型评估"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c90ad4c-70eb-43d1-96ec-cc74c4bae345",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "b24fccce-fec3-48a3-b43c-b9077788521d",
   "metadata": {},
   "source": [
    "## 评估模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "b021c6a8-645d-44f7-970b-f410180787a6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5874cd7f7cba4c1d9d89092923e0b8a5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading builder script:   0%|          | 0.00/4.49k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import evaluate\n",
    "\n",
    "# 词错误率（WER）是评估ASR模型常用的指标。从 Evaluate加载 WER 指标\n",
    "metric = evaluate.load(\"wer\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "5156ce16-0e04-4e41-b308-52c6c9b2a20d",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "PeftModel(\n",
       "  (base_model): LoraModel(\n",
       "    (model): WhisperForConditionalGeneration(\n",
       "      (model): WhisperModel(\n",
       "        (encoder): WhisperEncoder(\n",
       "          (conv1): Conv1d(80, 1280, kernel_size=(3,), stride=(1,), padding=(1,))\n",
       "          (conv2): Conv1d(1280, 1280, kernel_size=(3,), stride=(2,), padding=(1,))\n",
       "          (embed_positions): Embedding(1500, 1280)\n",
       "          (layers): ModuleList(\n",
       "            (0-31): 32 x WhisperEncoderLayer(\n",
       "              (self_attn): WhisperSdpaAttention(\n",
       "                (k_proj): Linear8bitLt(in_features=1280, out_features=1280, bias=False)\n",
       "                (v_proj): lora.Linear8bitLt(\n",
       "                  (base_layer): Linear8bitLt(in_features=1280, out_features=1280, bias=True)\n",
       "                  (lora_dropout): ModuleDict(\n",
       "                    (default): Dropout(p=0.05, inplace=False)\n",
       "                  )\n",
       "                  (lora_A): ModuleDict(\n",
       "                    (default): Linear(in_features=1280, out_features=8, bias=False)\n",
       "                  )\n",
       "                  (lora_B): ModuleDict(\n",
       "                    (default): Linear(in_features=8, out_features=1280, bias=False)\n",
       "                  )\n",
       "                  (lora_embedding_A): ParameterDict()\n",
       "                  (lora_embedding_B): ParameterDict()\n",
       "                )\n",
       "                (q_proj): lora.Linear8bitLt(\n",
       "                  (base_layer): Linear8bitLt(in_features=1280, out_features=1280, bias=True)\n",
       "                  (lora_dropout): ModuleDict(\n",
       "                    (default): Dropout(p=0.05, inplace=False)\n",
       "                  )\n",
       "                  (lora_A): ModuleDict(\n",
       "                    (default): Linear(in_features=1280, out_features=8, bias=False)\n",
       "                  )\n",
       "                  (lora_B): ModuleDict(\n",
       "                    (default): Linear(in_features=8, out_features=1280, bias=False)\n",
       "                  )\n",
       "                  (lora_embedding_A): ParameterDict()\n",
       "                  (lora_embedding_B): ParameterDict()\n",
       "                )\n",
       "                (out_proj): Linear8bitLt(in_features=1280, out_features=1280, bias=True)\n",
       "              )\n",
       "              (self_attn_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "              (activation_fn): GELUActivation()\n",
       "              (fc1): Linear8bitLt(in_features=1280, out_features=5120, bias=True)\n",
       "              (fc2): Linear8bitLt(in_features=5120, out_features=1280, bias=True)\n",
       "              (final_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "            )\n",
       "          )\n",
       "          (layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "        )\n",
       "        (decoder): WhisperDecoder(\n",
       "          (embed_tokens): Embedding(51865, 1280, padding_idx=50257)\n",
       "          (embed_positions): WhisperPositionalEmbedding(448, 1280)\n",
       "          (layers): ModuleList(\n",
       "            (0-31): 32 x WhisperDecoderLayer(\n",
       "              (self_attn): WhisperSdpaAttention(\n",
       "                (k_proj): Linear8bitLt(in_features=1280, out_features=1280, bias=False)\n",
       "                (v_proj): lora.Linear8bitLt(\n",
       "                  (base_layer): Linear8bitLt(in_features=1280, out_features=1280, bias=True)\n",
       "                  (lora_dropout): ModuleDict(\n",
       "                    (default): Dropout(p=0.05, inplace=False)\n",
       "                  )\n",
       "                  (lora_A): ModuleDict(\n",
       "                    (default): Linear(in_features=1280, out_features=8, bias=False)\n",
       "                  )\n",
       "                  (lora_B): ModuleDict(\n",
       "                    (default): Linear(in_features=8, out_features=1280, bias=False)\n",
       "                  )\n",
       "                  (lora_embedding_A): ParameterDict()\n",
       "                  (lora_embedding_B): ParameterDict()\n",
       "                )\n",
       "                (q_proj): lora.Linear8bitLt(\n",
       "                  (base_layer): Linear8bitLt(in_features=1280, out_features=1280, bias=True)\n",
       "                  (lora_dropout): ModuleDict(\n",
       "                    (default): Dropout(p=0.05, inplace=False)\n",
       "                  )\n",
       "                  (lora_A): ModuleDict(\n",
       "                    (default): Linear(in_features=1280, out_features=8, bias=False)\n",
       "                  )\n",
       "                  (lora_B): ModuleDict(\n",
       "                    (default): Linear(in_features=8, out_features=1280, bias=False)\n",
       "                  )\n",
       "                  (lora_embedding_A): ParameterDict()\n",
       "                  (lora_embedding_B): ParameterDict()\n",
       "                )\n",
       "                (out_proj): Linear8bitLt(in_features=1280, out_features=1280, bias=True)\n",
       "              )\n",
       "              (activation_fn): GELUActivation()\n",
       "              (self_attn_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "              (encoder_attn): WhisperSdpaAttention(\n",
       "                (k_proj): Linear8bitLt(in_features=1280, out_features=1280, bias=False)\n",
       "                (v_proj): lora.Linear8bitLt(\n",
       "                  (base_layer): Linear8bitLt(in_features=1280, out_features=1280, bias=True)\n",
       "                  (lora_dropout): ModuleDict(\n",
       "                    (default): Dropout(p=0.05, inplace=False)\n",
       "                  )\n",
       "                  (lora_A): ModuleDict(\n",
       "                    (default): Linear(in_features=1280, out_features=8, bias=False)\n",
       "                  )\n",
       "                  (lora_B): ModuleDict(\n",
       "                    (default): Linear(in_features=8, out_features=1280, bias=False)\n",
       "                  )\n",
       "                  (lora_embedding_A): ParameterDict()\n",
       "                  (lora_embedding_B): ParameterDict()\n",
       "                )\n",
       "                (q_proj): lora.Linear8bitLt(\n",
       "                  (base_layer): Linear8bitLt(in_features=1280, out_features=1280, bias=True)\n",
       "                  (lora_dropout): ModuleDict(\n",
       "                    (default): Dropout(p=0.05, inplace=False)\n",
       "                  )\n",
       "                  (lora_A): ModuleDict(\n",
       "                    (default): Linear(in_features=1280, out_features=8, bias=False)\n",
       "                  )\n",
       "                  (lora_B): ModuleDict(\n",
       "                    (default): Linear(in_features=8, out_features=1280, bias=False)\n",
       "                  )\n",
       "                  (lora_embedding_A): ParameterDict()\n",
       "                  (lora_embedding_B): ParameterDict()\n",
       "                )\n",
       "                (out_proj): Linear8bitLt(in_features=1280, out_features=1280, bias=True)\n",
       "              )\n",
       "              (encoder_attn_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "              (fc1): Linear8bitLt(in_features=1280, out_features=5120, bias=True)\n",
       "              (fc2): Linear8bitLt(in_features=5120, out_features=1280, bias=True)\n",
       "              (final_layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "            )\n",
       "          )\n",
       "          (layer_norm): LayerNorm((1280,), eps=1e-05, elementwise_affine=True)\n",
       "        )\n",
       "      )\n",
       "      (proj_out): Linear(in_features=1280, out_features=51865, bias=False)\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from torch.utils.data import DataLoader\n",
    "from tqdm import tqdm\n",
    "import numpy as np\n",
    "import gc\n",
    "\n",
    "eval_dataloader = DataLoader(common_voice[\"test\"], batch_size=8, collate_fn=data_collator)\n",
    "\n",
    "model.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "9120279e-b10a-44ea-9275-2152ec204fae",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/1323 [00:00<?, ?it/s]/data/miniconda3/envs/jxzhang/lib/python3.11/site-packages/bitsandbytes/autograd/_functions.py:322: UserWarning: MatMul8bitLt: inputs will be cast from torch.float32 to float16 during quantization\n",
      "  warnings.warn(f\"MatMul8bitLt: inputs will be cast from {A.dtype} to float16 during quantization\")\n",
      "100%|██████████| 1323/1323 [3:13:07<00:00,  8.76s/it] \n"
     ]
    }
   ],
   "source": [
    "for step, batch in enumerate(tqdm(eval_dataloader)):\n",
    "    with torch.cuda.amp.autocast():\n",
    "        with torch.no_grad():\n",
    "            generated_tokens = (\n",
    "                model.generate(\n",
    "                    input_features=batch[\"input_features\"].to(\"cuda\"),\n",
    "                    decoder_input_ids=batch[\"labels\"][:, :4].to(\"cuda\"),\n",
    "                    max_new_tokens=255,\n",
    "                )\n",
    "                .cpu()\n",
    "                .numpy()\n",
    "            )\n",
    "            labels = batch[\"labels\"].cpu().numpy()\n",
    "            labels = np.where(labels != -100, labels, tokenizer.pad_token_id)\n",
    "            decoded_preds = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)\n",
    "            decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)\n",
    "            metric.add_batch(\n",
    "                predictions=decoded_preds,\n",
    "                references=decoded_labels,\n",
    "            )\n",
    "    del generated_tokens, labels, batch\n",
    "    gc.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "aad4d7e8-ee0d-484a-9ae7-490c8b9898a0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "wer=54.73445473445473\n"
     ]
    }
   ],
   "source": [
    "wer = 100 * metric.compute()\n",
    "print(f\"{wer=}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "4120d924-c684-44af-8a1d-76aaf506bd08",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "65c2e25ce2434db2a59617adf7439064",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "adapter_model.safetensors:   0%|          | 0.00/15.8M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "CommitInfo(commit_url='https://huggingface.co/yqzhangjx/whisper-large-v2-finetune-for-common_voice_11_0/commit/e9a78de42addc1cfb814188de95de3e97f9281c6', commit_message='Upload model', commit_description='', oid='e9a78de42addc1cfb814188de95de3e97f9281c6', pr_url=None, pr_revision=None, pr_num=None)"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.push_to_hub('whisper-large-v2-asr-int8')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b740c690-35ff-4375-b59f-3c6de5155ec0",
   "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.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}