distill-whisper-large-v2-thai-qlora-train.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip3 install --upgrade transformers peft bitsandbytes datasets accelerate loralib huggingface_hub jiwer evaluate wandb pythainlp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from huggingface_hub import notebook_login\n",
    "\n",
    "notebook_login()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-11-05T07:10:03.850363Z",
     "iopub.status.busy": "2023-11-05T07:10:03.849810Z",
     "iopub.status.idle": "2023-11-05T07:10:17.952413Z",
     "shell.execute_reply": "2023-11-05T07:10:17.951554Z",
     "shell.execute_reply.started": "2023-11-05T07:10:03.850327Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/lib/python3.10/site-packages/scipy/__init__.py:146: UserWarning: A NumPy version >=1.16.5 and <1.23.0 is required for this version of SciPy (detected version 1.23.5\n",
      "  warnings.warn(f\"A NumPy version >={np_minversion} and <{np_maxversion}\"\n"
     ]
    }
   ],
   "source": [
    "from datasets import load_dataset, DatasetDict, load_metric, Audio\n",
    "from transformers import WhisperFeatureExtractor, WhisperTokenizer, WhisperProcessor, WhisperForConditionalGeneration, Seq2SeqTrainingArguments\n",
    "from transformers import Seq2SeqTrainer, TrainerCallback, TrainingArguments, TrainerState, TrainerControl\n",
    "from peft import LoraConfig, PeftModel, LoraModel, LoraConfig, get_peft_model, prepare_model_for_kbit_training\n",
    "\n",
    "from dataclasses import dataclass\n",
    "from typing import Any, Dict, List, Union\n",
    "import os\n",
    "import torch\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import evaluate"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prepare Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "language_abbr = \"th\"\n",
    "dataset_name = \"mozilla-foundation/common_voice_13_0\"\n",
    "\n",
    "common_voice = DatasetDict()\n",
    "\n",
    "common_voice[\"train\"] = load_dataset(dataset_name, language_abbr, split=\"train+validation\", token=True)\n",
    "common_voice[\"test\"] = load_dataset(dataset_name, language_abbr, split=\"test\", token=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "common_voice = common_voice.filter(lambda x: x == 0, input_columns=[\"down_votes\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-11-05T07:15:44.951725Z",
     "iopub.status.busy": "2023-11-05T07:15:44.951439Z",
     "iopub.status.idle": "2023-11-05T07:15:44.963618Z",
     "shell.execute_reply": "2023-11-05T07:15:44.962746Z",
     "shell.execute_reply.started": "2023-11-05T07:15:44.951699Z"
    }
   },
   "outputs": [],
   "source": [
    "common_voice = common_voice.remove_columns(\n",
    "    [\"accent\", \"age\", \"client_id\", \"down_votes\", \"gender\", \"locale\", \"path\", \"segment\", \"up_votes\", \"variant\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-11-05T07:15:44.965526Z",
     "iopub.status.busy": "2023-11-05T07:15:44.965252Z",
     "iopub.status.idle": "2023-11-05T07:15:44.979385Z",
     "shell.execute_reply": "2023-11-05T07:15:44.978625Z",
     "shell.execute_reply.started": "2023-11-05T07:15:44.965504Z"
    }
   },
   "outputs": [],
   "source": [
    "common_voice = common_voice.cast_column(\"audio\", Audio(sampling_rate=16000))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "feature_extractor = WhisperFeatureExtractor.from_pretrained(\"distil-whisper/distil-large-v2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = WhisperTokenizer.from_pretrained(\"distil-whisper/distil-large-v2\", language=\"Thai\", task=\"transcribe\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-11-05T07:15:48.191244Z",
     "iopub.status.busy": "2023-11-05T07:15:48.190970Z",
     "iopub.status.idle": "2023-11-05T07:15:48.197259Z",
     "shell.execute_reply": "2023-11-05T07:15:48.196345Z",
     "shell.execute_reply.started": "2023-11-05T07:15:48.191220Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['audio', 'sentence'],\n",
       "        num_rows: 39140\n",
       "    })\n",
       "    test: Dataset({\n",
       "        features: ['audio', 'sentence'],\n",
       "        num_rows: 9332\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "common_voice"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-11-05T07:15:48.198955Z",
     "iopub.status.busy": "2023-11-05T07:15:48.198583Z",
     "iopub.status.idle": "2023-11-05T07:15:57.405231Z",
     "shell.execute_reply": "2023-11-05T07:15:57.404312Z",
     "shell.execute_reply.started": "2023-11-05T07:15:48.198920Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Input:                 ถ้าทำแบบนี้ ถูกไล่ออก ครอบครัวจะไปอยู่ไหน\n",
      "Decoded w/ special:    <|startoftranscript|><|th|><|transcribe|><|notimestamps|>ถ้าทำแบบนี้ ถูกไล่ออก ครอบครัวจะไปอยู่ไหน<|endoftext|>\n",
      "Decoded w/out special: ถ้าทำแบบนี้ ถูกไล่ออก ครอบครัวจะไปอยู่ไหน\n",
      "Are equal:             True\n"
     ]
    }
   ],
   "source": [
    "input_str = common_voice[\"train\"][0][\"sentence\"]\n",
    "labels = tokenizer(input_str).input_ids\n",
    "decoded_with_special = tokenizer.decode(labels, skip_special_tokens=False)\n",
    "decoded_str = tokenizer.decode(labels, skip_special_tokens=True)\n",
    "\n",
    "print(f\"Input:                 {input_str}\")\n",
    "print(f\"Decoded w/ special:    {decoded_with_special}\")\n",
    "print(f\"Decoded w/out special: {decoded_str}\")\n",
    "print(f\"Are equal:             {input_str == decoded_str}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-11-05T07:15:57.408937Z",
     "iopub.status.busy": "2023-11-05T07:15:57.408208Z",
     "iopub.status.idle": "2023-11-05T07:15:57.767350Z",
     "shell.execute_reply": "2023-11-05T07:15:57.766437Z",
     "shell.execute_reply.started": "2023-11-05T07:15:57.408909Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
     ]
    }
   ],
   "source": [
    "processor = WhisperProcessor.from_pretrained(\"distil-whisper/distil-large-v2\", language=\"Thai\", task=\"transcribe\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-11-05T07:15:57.768891Z",
     "iopub.status.busy": "2023-11-05T07:15:57.768586Z",
     "iopub.status.idle": "2023-11-05T07:15:57.774158Z",
     "shell.execute_reply": "2023-11-05T07:15:57.773196Z",
     "shell.execute_reply.started": "2023-11-05T07:15:57.768864Z"
    }
   },
   "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[\"sentence\"]).input_ids\n",
    "    return batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "common_voice = common_voice.map(prepare_dataset, remove_columns=common_voice.column_names[\"train\"], num_proc=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-11-05T08:16:04.015724Z",
     "iopub.status.busy": "2023-11-05T08:16:04.015078Z",
     "iopub.status.idle": "2023-11-05T08:16:04.027015Z",
     "shell.execute_reply": "2023-11-05T08:16:04.026021Z",
     "shell.execute_reply.started": "2023-11-05T08:16:04.015687Z"
    }
   },
   "outputs": [],
   "source": [
    "@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": 20,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-11-05T08:16:07.027165Z",
     "iopub.status.busy": "2023-11-05T08:16:07.026429Z",
     "iopub.status.idle": "2023-11-05T08:16:07.031176Z",
     "shell.execute_reply": "2023-11-05T08:16:07.030270Z",
     "shell.execute_reply.started": "2023-11-05T08:16:07.027135Z"
    }
   },
   "outputs": [],
   "source": [
    "data_collator = DataCollatorSpeechSeq2SeqWithPadding(processor=processor)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-11-05T08:16:11.403830Z",
     "iopub.status.busy": "2023-11-05T08:16:11.403458Z",
     "iopub.status.idle": "2023-11-05T08:16:12.191923Z",
     "shell.execute_reply": "2023-11-05T08:16:12.191149Z",
     "shell.execute_reply.started": "2023-11-05T08:16:11.403800Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_32/4287149634.py:1: FutureWarning: load_metric is deprecated and will be removed in the next major version of datasets. Use 'evaluate.load' instead, from the new library 🤗 Evaluate: https://huggingface.co/docs/evaluate\n",
      "  cer_metric = load_metric(\"cer\")\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "495d039030464788bd66b64ca104b6bc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading builder script:   0%|          | 0.00/2.16k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "cer_metric = load_metric(\"cer\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-11-05T08:16:17.063367Z",
     "iopub.status.busy": "2023-11-05T08:16:17.062965Z",
     "iopub.status.idle": "2023-11-05T08:16:17.069863Z",
     "shell.execute_reply": "2023-11-05T08:16:17.068487Z",
     "shell.execute_reply.started": "2023-11-05T08:16:17.063332Z"
    }
   },
   "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",
    "    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 = cer_metric.compute(predictions=pred_str, references=label_str)\n",
    "\n",
    "    return {\"cer\": cer}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prepare Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-11-05T08:25:13.510905Z",
     "iopub.status.busy": "2023-11-05T08:25:13.510036Z",
     "iopub.status.idle": "2023-11-05T08:25:16.884300Z",
     "shell.execute_reply": "2023-11-05T08:25:16.883162Z",
     "shell.execute_reply.started": "2023-11-05T08:25:13.510869Z"
    }
   },
   "outputs": [],
   "source": [
    "model = WhisperForConditionalGeneration.from_pretrained(\"distil-whisper/distil-large-v2\", load_in_4bit=True, device_map=\"auto\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-11-05T08:25:16.886375Z",
     "iopub.status.busy": "2023-11-05T08:25:16.886051Z",
     "iopub.status.idle": "2023-11-05T08:25:16.891443Z",
     "shell.execute_reply": "2023-11-05T08:25:16.890501Z",
     "shell.execute_reply.started": "2023-11-05T08:25:16.886340Z"
    }
   },
   "outputs": [],
   "source": [
    "model.config.forced_decoder_ids = None\n",
    "model.config.suppress_tokens = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-11-05T08:25:16.892953Z",
     "iopub.status.busy": "2023-11-05T08:25:16.892622Z",
     "iopub.status.idle": "2023-11-05T08:25:16.917437Z",
     "shell.execute_reply": "2023-11-05T08:25:16.916435Z",
     "shell.execute_reply.started": "2023-11-05T08:25:16.892921Z"
    }
   },
   "outputs": [],
   "source": [
    "model = prepare_model_for_kbit_training(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-11-05T08:25:16.920641Z",
     "iopub.status.busy": "2023-11-05T08:25:16.920132Z",
     "iopub.status.idle": "2023-11-05T08:25:16.960893Z",
     "shell.execute_reply": "2023-11-05T08:25:16.959877Z",
     "shell.execute_reply.started": "2023-11-05T08:25:16.920603Z"
    }
   },
   "outputs": [],
   "source": [
    "config = LoraConfig(\n",
    "    r=32, \n",
    "    lora_alpha=64, \n",
    "    target_modules=\".*decoder.*(self_attn|encoder_attn).*(q_proj|v_proj)$\",#[\"q_proj\", \"v_proj\"],\n",
    "    lora_dropout=0.05,\n",
    "    bias=\"none\"\n",
    ")\n",
    "\n",
    "model = get_peft_model(model, config)\n",
    "model.print_trainable_parameters()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-11-05T11:41:45.468606Z",
     "iopub.status.busy": "2023-11-05T11:41:45.467787Z",
     "iopub.status.idle": "2023-11-05T11:41:45.477702Z",
     "shell.execute_reply": "2023-11-05T11:41:45.476387Z",
     "shell.execute_reply.started": "2023-11-05T11:41:45.468563Z"
    }
   },
   "outputs": [],
   "source": [
    "training_args = Seq2SeqTrainingArguments(\n",
    "    output_dir=\"./distill-whisper-large-v2-thai-qlora\",\n",
    "    per_device_train_batch_size=16,\n",
    "    gradient_accumulation_steps=2,\n",
    "    learning_rate=1e-3,\n",
    "    warmup_steps=100,\n",
    "    max_steps=1000,\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=500,\n",
    "    eval_steps=500,\n",
    "    logging_steps=500,\n",
    "    num_train_epochs=3,\n",
    "    report_to=[\"wandb\"],\n",
    "    load_best_model_at_end=True,\n",
    "    metric_for_best_model=\"cer\",\n",
    "    greater_is_better=False,\n",
    "    push_to_hub=False,\n",
    "    save_total_limit=3,\n",
    "    remove_unused_columns=False,  # required as the PeftModel forward doesn't have the signature of the wrapped model's forward\n",
    "    label_names=[\"labels\"],  # same reason as above\n",
    "    lr_scheduler_type=\"cosine\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-11-05T11:41:50.621005Z",
     "iopub.status.busy": "2023-11-05T11:41:50.620171Z",
     "iopub.status.idle": "2023-11-05T11:41:50.636142Z",
     "shell.execute_reply": "2023-11-05T11:41:50.635101Z",
     "shell.execute_reply.started": "2023-11-05T11:41:50.620971Z"
    }
   },
   "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",
    "    compute_metrics=compute_metrics,\n",
    "    tokenizer=processor.feature_extractor,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "peft_model_id = \"juierror/distill-whisper-large-v2-thai-qlora\"\n",
    "model.push_to_hub(peft_model_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-11-05T14:33:03.808885Z",
     "iopub.status.busy": "2023-11-05T14:33:03.808156Z",
     "iopub.status.idle": "2023-11-05T14:33:05.068556Z",
     "shell.execute_reply": "2023-11-05T14:33:05.067636Z",
     "shell.execute_reply.started": "2023-11-05T14:33:03.808852Z"
    }
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5c8ab008c8d64f0a94ff05b6e3af6b1e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading builder script:   0%|          | 0.00/1.90k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "cer_metric = load_metric(\"cer\")\n",
    "wer_metric = load_metric(\"wer\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "peft_model_id = \"juierror/distill-whisper-large-v2-thai-qlora\"\n",
    "peft_config = PeftConfig.from_pretrained(peft_model_id)\n",
    "model = WhisperForConditionalGeneration.from_pretrained(peft_config.base_model_name_or_path, load_in_4bit=True, device_map=\"auto\")\n",
    "model = PeftModel.from_pretrained(model, peft_model_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-11-05T14:34:47.012006Z",
     "iopub.status.busy": "2023-11-05T14:34:47.011502Z",
     "iopub.status.idle": "2023-11-05T16:34:51.913564Z",
     "shell.execute_reply": "2023-11-05T16:34:51.912287Z",
     "shell.execute_reply.started": "2023-11-05T14:34:47.011962Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 584/584 [2:00:02<00:00, 12.33s/it]  \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "wer: 62.49801320675654\n",
      "cer: 19.122046340140493\n"
     ]
    }
   ],
   "source": [
    "from torch.utils.data import DataLoader\n",
    "from pythainlp.tokenize import word_tokenize\n",
    "from tqdm import tqdm\n",
    "\n",
    "eval_dataloader = DataLoader(common_voice[\"test\"], batch_size=16, collate_fn=data_collator)\n",
    "all_labels = []\n",
    "all_transcription = []\n",
    "all_labels_token = []\n",
    "all_transcription_token = []\n",
    "\n",
    "model.eval()\n",
    "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",
    "                    max_new_tokens=255,\n",
    "                    language=\"Thai\",\n",
    "                    task=\"transcribe\"\n",
    "                )\n",
    "                .cpu()\n",
    "                .numpy()\n",
    "            )\n",
    "            labels = batch[\"labels\"].cpu().numpy()\n",
    "            transcriptions = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)\n",
    "            sentences = tokenizer.batch_decode(labels, skip_special_tokens=True)\n",
    "            all_labels.extend(sentences)\n",
    "            all_transcription.extend(transcriptions)\n",
    "            \n",
    "            cer_metric.add_batch(predictions=[pred_str.replace(\" \", \"\") for pred_str in transcriptions], references=[label_str.replace(\" \", \"\") for label_str in sentences])\n",
    "\n",
    "            pred_str_newmm = [word_tokenize(text=e, engine='newmm', keep_whitespace=False) for e in transcriptions]\n",
    "            label_str_newmm = [word_tokenize(text=e, engine='newmm', keep_whitespace=False) for e in sentences]\n",
    "            all_labels_token.extend(label_str_newmm)\n",
    "            all_transcription_token.extend(pred_str_newmm)\n",
    "            wer_metric.add_batch(predictions=pred_str_newmm, references=label_str_newmm)\n",
    "\n",
    "wer = 100 * wer_metric.compute()\n",
    "cer = 100 * cer_metric.compute()\n",
    "print(f\"wer: {wer}\")\n",
    "print(f\"cer: {cer}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-11-05T16:34:52.223677Z",
     "iopub.status.busy": "2023-11-05T16:34:52.222886Z",
     "iopub.status.idle": "2023-11-05T16:34:52.238471Z",
     "shell.execute_reply": "2023-11-05T16:34:52.237542Z",
     "shell.execute_reply.started": "2023-11-05T16:34:52.223644Z"
    }
   },
   "outputs": [],
   "source": [
    "report_df = pd.DataFrame(data={\n",
    "    \"labels\": all_labels,\n",
    "    \"transcribe\": all_transcription,\n",
    "    \"labels_tokenize\": all_labels_token,\n",
    "    \"transcribe_tokenizer\": all_transcription_token\n",
    "})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-11-05T16:34:52.240673Z",
     "iopub.status.busy": "2023-11-05T16:34:52.240283Z",
     "iopub.status.idle": "2023-11-05T16:34:52.388530Z",
     "shell.execute_reply": "2023-11-05T16:34:52.387276Z",
     "shell.execute_reply.started": "2023-11-05T16:34:52.240635Z"
    }
   },
   "outputs": [],
   "source": [
    "report_df.to_csv(\"report.csv\", index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}