8_huggingface_whisper_finetuning.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fine-tuning Whisper For Korean ASR with 🤗 Transformers\n",
    "\n",
    "> 본 튜토리얼은 [Fine-Tune Whisper For Multilingual ASR with 🤗 Transformers](https://colab.research.google.com/github/sanchit-gandhi/notebooks/blob/main/fine_tune_whisper.ipynb)을 참고하여 제작되었습니다."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Whisper\n",
    "\n",
    "- https://openai.com/blog/whisper/  \n",
    "- trained on 680,000 hours of multilingual and multitask supervised data collected from the web.\n",
    "- multiple languages (https://github.com/openai/whisper/blob/main/whisper/tokenizer.py)\n",
    "\n",
    "<img src=\"https://cdn.openai.com/whisper/asr-summary-of-model-architecture-desktop.svg\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load Dataset\n",
    "- zeroth_korean 의 한국어 데이터셋 사용"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Import wandb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install wandb\n",
    "\n",
    "import wandb\n",
    "wandb.login()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Install library"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install accelerate\n",
    "!pip install datasets\n",
    "!pip install transformers==4.26.1\n",
    "!pip install librosa\n",
    "!pip install evaluate\n",
    "!pip install jiwer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Huggingface login"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from huggingface_hub import notebook_login\n",
    "\n",
    "# notebook_login()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !apt install git-lfs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### load dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "train_dataset = load_dataset(\"kresnik/zeroth_korean\", split='train[:1800]')\n",
    "test_dataset = load_dataset(\"kresnik/zeroth_korean\", split='test[:200]')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(train_dataset)\n",
    "print(test_dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Remove unnecessary columns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_ds = train_dataset.remove_columns([\"speaker_id\", \"chapter_id\", \"id\"])\n",
    "test_ds = test_dataset.remove_columns([\"speaker_id\", \"chapter_id\", \"id\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(train_ds)\n",
    "print(test_ds)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prepare Feature Extractor, Tokenizer and Data\n",
    "1. A feature extractor which pre-processes the raw audio-inputs \n",
    "2. The model which performs the sequence-to-sequence mapping\n",
    "3. A tokenizer which post-processes the model outputs to text format"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load WhisperFeatureExtractor\n",
    "\n",
    "1. 오디오 입력을 30초로 pads/truncate\n",
    "    - 30초보다 짧은 오디오 입력은 무음(0)으로 30초로 pad하고, 30초보다 긴 오디오 입력은 30초로 truncate\n",
    "2. 오디오 입력 log-Mel spectorgram으로 변환"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import WhisperFeatureExtractor\n",
    "\n",
    "feature_extractor = WhisperFeatureExtractor.from_pretrained(\"openai/whisper-tiny\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load WhisperTokenizer\n",
    "\n",
    "위스퍼 모델은 token id sequence를 결과로 출력하고 Tokenizer는 이러한 각 token id를 를 해당 텍스트 문자열에 매핑합니다. \n",
    "\n",
    "tokenizer에 language와 task를 전달 할 수 있다. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import WhisperTokenizer\n",
    "\n",
    "tokenizer = WhisperTokenizer.from_pretrained(\"openai/whisper-tiny\", language=\"korean\", task=\"transcribe\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Combine To Create A WhisperProcessor\n",
    "\n",
    "feauter extractor 와 tokenizer는 WhisperProcessor 클래스로 wrap되어 train에서는 `processor` 와 `model`만 사용"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import WhisperProcessor\n",
    "\n",
    "processor = WhisperProcessor.from_pretrained(\"openai/whisper-tiny\", language=\"korean\", task=\"transcribe\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Prepare Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_ds[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Resample to 16kHz\n",
    "\n",
    "`cast_column` : 오디오를 제자리에서 변경하는 것이 아니라, 오디오 샘플을 처음 load할 때 즉시 resample되도록 Dataset에 신호를 보낸다. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import Audio\n",
    "\n",
    "train_ds = train_ds.cast_column(\"audio\", Audio(sampling_rate=16000))\n",
    "test_ds = test_ds.cast_column(\"audio\", Audio(sampling_rate=16000))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "0번 index의 audio sample을 reload하면 16kHz로 resample"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_ds[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Audio Sample 확인"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import IPython.display as ipd\n",
    "import numpy as np\n",
    "import random\n",
    "\n",
    "rand_int = random.randint(0, len(train_ds)-1)\n",
    "\n",
    "ipd.display(ipd.Audio(data=train_ds[rand_int][\"audio\"][\"array\"], autoplay=True, rate=16000))\n",
    "\n",
    "print(\"Target text:\", train_ds[rand_int][\"text\"])\n",
    "print(\"Input array shape:\", train_ds[rand_int][\"audio\"][\"array\"].shape)\n",
    "print(\"Sampling rate:\", train_ds[rand_int][\"audio\"][\"sampling_rate\"])\n",
    "print(\"Input Duration:\", train_ds[rand_int][\"audio\"][\"array\"].shape[0]/train_ds[rand_int][\"audio\"][\"sampling_rate\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "1. `batch[\"audio\"]`를 호출하여 오디오 데이터를 로드하고 16kHz로 resampling  \n",
    "\n",
    "2. 로드된 오디오 파일에서 `input_values`값을 추출, 이 단계에서 Log-Mel Spectrogram\n",
    "\n",
    "3. transcription을 인코딩하여 label ID를 지정"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def prepare_dataset(batch):\n",
    "    # load and resample audio data from 48 to 16kHz\n",
    "    audio = batch[\"audio\"]\n",
    "\n",
    "    # compute log-Mel input features from input audio array \n",
    "    batch[\"input_features\"] = feature_extractor(audio[\"array\"], sampling_rate=audio[\"sampling_rate\"]).input_features[0]\n",
    "\n",
    "    # encode target text to label ids \n",
    "    batch[\"labels\"] = tokenizer(batch[\"text\"]).input_ids\n",
    "    return batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_ds = train_ds.map(prepare_dataset, remove_columns=train_ds.column_names, num_proc=4)\n",
    "test_ds = test_ds.map(prepare_dataset, remove_columns=test_ds.column_names, num_proc=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_ds[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_ds[0].keys()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training and Evaluation\n",
    "\n",
    "🤗 Trainer\n",
    "\n",
    "1. define Data Collector : Data Collector는 pre-processed data를 가져와서 model에 사용할 수 있는 PyTorch tensor를 준비\n",
    "\n",
    "2. Evaluation metrics: Evaluation 중에 CER 메트릭을 사용하여 model을 평가. 이 계산을 처리하는 compute_metrics 함수를 정의\n",
    "\n",
    "3. Load a pre-trained checkpoint: pre-trained checkpoint를 load하고 학습을 위해 올바르게 구성\n",
    "\n",
    "4. Define the training configuration: 🤗 Trainer의 Training schedule을 정의"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define a Data Collector\n",
    "\n",
    "1. input_features\n",
    "    - 이미 `input_features`는 feature extractor로 30s로 pad되고 log-Mel spectrogram으로 변환되었음\n",
    "    - `input_features`를 PyTorch tensor로 변환을 진행 (`return_tensors=pt`)\n",
    "2. labels\n",
    "    - `labels` 는 아직 un-padded. \n",
    "    - `labels`를 sequence의 최대 길이까지 패딩\n",
    "    - `labels`를 PyTorch tensor로 변환을 진행 (`return_tensors=pt`)\n",
    "    - 패딩된 token은 -100으로 대체되어 손실을 계산할 때 해당 token을 고려하지 않도록 함.\n",
    "    - Train 과정에서 BOS token (begin of sequence)을 추가하기 때문에 sequence의 시작 부분에서 BOS token을 잘라낸다. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "from dataclasses import dataclass\n",
    "from typing import Any, Dict, List, Union\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",
    "        # split inputs and labels since they have to be of different lengths and need different padding methods\n",
    "        # first treat the audio inputs by simply returning torch tensors\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",
    "        # get the tokenized label sequences\n",
    "        label_features = [{\"input_ids\": feature[\"labels\"]} for feature in features]\n",
    "        # pad the labels to max length\n",
    "        labels_batch = self.processor.tokenizer.pad(label_features, return_tensors=\"pt\")\n",
    "\n",
    "        # replace padding with -100 to ignore loss correctly\n",
    "        labels = labels_batch[\"input_ids\"].masked_fill(labels_batch.attention_mask.ne(1), -100)\n",
    "\n",
    "        # if bos token is appended in previous tokenization step,\n",
    "        # cut bos token here as it's append later anyways\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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_collator = DataCollatorSpeechSeq2SeqWithPadding(processor=processor)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Evaluation Matrics\n",
    "\n",
    "- Character Error Rate(CER)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import evaluate\n",
    "\n",
    "metric = evaluate.load(\"cer\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_metrics(pred):\n",
    "    pred_ids = pred.predictions\n",
    "    label_ids = pred.label_ids\n",
    "\n",
    "    # replace -100 with the pad_token_id\n",
    "    label_ids[label_ids == -100] = tokenizer.pad_token_id\n",
    "\n",
    "    # we do not want to group tokens when computing the metrics\n",
    "    pred_str = tokenizer.batch_decode(pred_ids, skip_special_tokens=True)\n",
    "    label_str = tokenizer.batch_decode(label_ids, skip_special_tokens=True)\n",
    "\n",
    "    cer = metric.compute(predictions=pred_str, references=label_str)\n",
    "\n",
    "    return {\"cer\": cer}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load a Pre-Trained Checkpoint\n",
    "\n",
    "- pre-trained `whisper-tiny` checkpoint 로드 \n",
    "- 자유롭게 output sequence를 생성하도록 model.config 수정"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import WhisperForConditionalGeneration\n",
    "\n",
    "model = WhisperForConditionalGeneration.from_pretrained(\n",
    "    \"openai/whisper-tiny\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.config.forced_decoder_ids = None      # no tokens are forced as decoder outputs \n",
    "model.config.suppress_tokens = []           # no tokens are suppressed during generation "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "model.config.use_cache = False, when using gradient checkpointing during training.\n",
    "but during inference make sure to set it back to True."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define the Training Configuration\n",
    "\n",
    "Seq2SeqTrainingArguments [docs](https://huggingface.co/docs/transformers/main_classes/trainer#transformers.Seq2SeqTrainingArguments)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import Seq2SeqTrainingArguments\n",
    "\n",
    "training_args = Seq2SeqTrainingArguments(\n",
    "    output_dir=\"./whisper-tiny-korean\",  # change to a repo name of your choice\n",
    "    per_device_train_batch_size=8,\n",
    "    gradient_accumulation_steps=2,  # increase by 2x for every 2x decrease in batch size\n",
    "    learning_rate=1e-5,\n",
    "    warmup_steps=500,\n",
    "    max_steps=4000,\n",
    "    gradient_checkpointing=True,\n",
    "    fp16=True,\n",
    "    evaluation_strategy=\"steps\",\n",
    "    per_device_eval_batch_size=8,\n",
    "    predict_with_generate=True,\n",
    "    generation_max_length=225,\n",
    "    save_steps=100,\n",
    "    eval_steps=100,\n",
    "    logging_steps=25,\n",
    "    report_to=\"wandb\",\n",
    "    run_name=\"whisper-tiny-korean\",\n",
    "    load_best_model_at_end=True,\n",
    "    metric_for_best_model=\"cer\",\n",
    "    greater_is_better=False,\n",
    "    # push_to_hub=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import Seq2SeqTrainer\n",
    "\n",
    "trainer = Seq2SeqTrainer(\n",
    "    args=training_args,\n",
    "    model=model,\n",
    "    train_dataset=train_ds,\n",
    "    eval_dataset=test_ds,\n",
    "    data_collator=data_collator,\n",
    "    compute_metrics=compute_metrics,\n",
    "    tokenizer=processor.feature_extractor,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# tokenizer.push_to_hub(training_args.output_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "processor.save_pretrained(training_args.output_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# trainer.push_to_hub()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wandb.finish()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Inference\n",
    "- spow12/whisper-medium-zeroth_korean [link](https://huggingface.co/spow12/whisper-medium-zeroth_korean)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import WhisperProcessor, WhisperForConditionalGeneration\n",
    "from datasets import load_dataset\n",
    "import soundfile as sf\n",
    "import torch\n",
    "from jiwer import cer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "processor = WhisperProcessor.from_pretrained(\"openai/whisper-medium\", language=\"ko\", task=\"transcribe\")\n",
    "model = WhisperForConditionalGeneration.from_pretrained('spow12/whisper-medium-zeroth_korean').to('cuda')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = load_dataset(\"kresnik/zeroth_korean\", \"clean\")\n",
    "\n",
    "test_ds = ds['test']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def map_to_array(batch):\n",
    "    speech, _ = sf.read(batch[\"file\"])\n",
    "    batch[\"speech\"] = speech\n",
    "    return batch\n",
    "\n",
    "test_ds = test_ds.map(map_to_array)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def map_to_pred(batch):\n",
    "    input_features = processor(batch[\"speech\"], sampling_rate=16000, return_tensors=\"pt\").input_features.cuda()\n",
    "    \n",
    "    predicted_ids = model.generate(input_features)\n",
    "\n",
    "    transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True)\n",
    "    batch[\"transcription\"] = transcription\n",
    "    \n",
    "    return batch\n",
    "\n",
    "result = test_ds.map(map_to_pred,batched=True, batch_size=16, remove_columns=[\"speech\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ipd.display(ipd.Audio(result[0]['file']))\n",
    "print(\"ref  : \", result[0][\"text\"])\n",
    "print(\"trans: \", result[0]['transcription'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"CER:\", cer(result[\"text\"], result[\"transcription\"]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "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.12"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}