bengali-sr-wav2vec-v1-bengali-training.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "2dee87c5",
   "metadata": {
    "papermill": {
     "duration": 0.003946,
     "end_time": "2023-10-14T19:42:11.394614",
     "exception": false,
     "start_time": "2023-10-14T19:42:11.390668",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "- This is a training demo, you can run this code locally, using better GPUs.\n",
    "- The inference part is here: [Bengali SR wav2vec_v1_bengali [Inference]](https://www.kaggle.com/takanashihumbert/bengali-sr-wav2vec-v1-bengali-inference), it scores **0.445** on the leaderboard.\n",
    "- Feel free to upvote, thanks!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "27ed5752",
   "metadata": {
    "_cell_guid": "b1076dfc-b9ad-4769-8c92-a6c4dae69d19",
    "_uuid": "8f2839f25d086af736a60e9eeb907d3b93b6e0e5",
    "execution": {
     "iopub.execute_input": "2023-10-14T19:42:11.403431Z",
     "iopub.status.busy": "2023-10-14T19:42:11.403089Z",
     "iopub.status.idle": "2023-10-14T19:43:24.238286Z",
     "shell.execute_reply": "2023-10-14T19:43:24.236948Z"
    },
    "papermill": {
     "duration": 72.842415,
     "end_time": "2023-10-14T19:43:24.240648",
     "exception": false,
     "start_time": "2023-10-14T19:42:11.398233",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "jiwer/\r\n",
      "jiwer/jiwer-2.3.0-py3-none-any.whl\r\n",
      "jiwer/python-Levenshtein-0.12.2.tar.gz\r\n",
      "jiwer/setuptools-65.3.0-py3-none-any.whl\r\n",
      "Looking in links: ./\r\n",
      "Processing ./jiwer/jiwer-2.3.0-py3-none-any.whl\r\n",
      "INFO: pip is looking at multiple versions of jiwer to determine which version is compatible with other requirements. This could take a while.\r\n",
      "\u001b[31mERROR: Could not find a version that satisfies the requirement python-Levenshtein==0.12.2 (from jiwer) (from versions: none)\u001b[0m\u001b[31m\r\n",
      "\u001b[0m\u001b[31mERROR: No matching distribution found for python-Levenshtein==0.12.2\u001b[0m\u001b[31m\r\n",
      "\u001b[0mnormalizer/\r\n",
      "normalizer/bnunicodenormalizer-0.0.24.tar.gz\r\n",
      "Looking in links: ./\r\n",
      "Processing ./normalizer/bnunicodenormalizer-0.0.24.tar.gz\r\n",
      "  Preparing metadata (setup.py) ... \u001b[?25ldone\r\n",
      "\u001b[?25hBuilding wheels for collected packages: bnunicodenormalizer\r\n",
      "  Building wheel for bnunicodenormalizer (setup.py) ... \u001b[?25l-\b \bdone\r\n",
      "\u001b[?25h  Created wheel for bnunicodenormalizer: filename=bnunicodenormalizer-0.0.24-py3-none-any.whl size=17609 sha256=5bf07510bea4df5e2b79cd251c35245fc07aaf21f91e8bd637f70c7311b664db\r\n",
      "  Stored in directory: /root/.cache/pip/wheels/78/d7/75/6986dc3616718f950b80e3bd79a796ef618eaef6cd800e7909\r\n",
      "Successfully built bnunicodenormalizer\r\n",
      "Installing collected packages: bnunicodenormalizer\r\n",
      "Successfully installed bnunicodenormalizer-0.0.24\r\n",
      "pyctcdecode/\r\n",
      "pyctcdecode/hypothesis-6.54.4-py3-none-any.whl\r\n",
      "pyctcdecode/sortedcontainers-2.4.0-py2.py3-none-any.whl\r\n",
      "pyctcdecode/exceptiongroup-1.0.0rc9-py3-none-any.whl\r\n",
      "pyctcdecode/pyctcdecode-0.4.0-py2.py3-none-any.whl\r\n",
      "pyctcdecode/numpy-1.21.6-cp37-cp37m-manylinux_2_12_x86_64.manylinux2010_x86_64.whl\r\n",
      "pyctcdecode/attrs-22.1.0-py2.py3-none-any.whl\r\n",
      "pyctcdecode/pygtrie-2.5.0.tar.gz\r\n",
      "Looking in links: ./\r\n",
      "Processing ./pyctcdecode/attrs-22.1.0-py2.py3-none-any.whl\r\n",
      "Installing collected packages: attrs\r\n",
      "  Attempting uninstall: attrs\r\n",
      "    Found existing installation: attrs 23.1.0\r\n",
      "    Uninstalling attrs-23.1.0:\r\n",
      "      Successfully uninstalled attrs-23.1.0\r\n",
      "Successfully installed attrs-22.1.0\r\n",
      "Looking in links: ./\r\n",
      "Processing ./pyctcdecode/exceptiongroup-1.0.0rc9-py3-none-any.whl\r\n",
      "Installing collected packages: exceptiongroup\r\n",
      "  Attempting uninstall: exceptiongroup\r\n",
      "    Found existing installation: exceptiongroup 1.1.1\r\n",
      "    Uninstalling exceptiongroup-1.1.1:\r\n",
      "      Successfully uninstalled exceptiongroup-1.1.1\r\n",
      "Successfully installed exceptiongroup-1.0.0rc9\r\n",
      "Looking in links: ./\r\n",
      "Processing ./pyctcdecode/hypothesis-6.54.4-py3-none-any.whl\r\n",
      "Installing collected packages: hypothesis\r\n",
      "Successfully installed hypothesis-6.54.4\r\n",
      "Looking in links: ./\r\n",
      "\u001b[31mERROR: numpy-1.21.6-cp37-cp37m-manylinux_2_12_x86_64.manylinux2010_x86_64.whl is not a supported wheel on this platform.\u001b[0m\u001b[31m\r\n",
      "\u001b[0mLooking in links: ./\r\n",
      "Processing ./pyctcdecode/pygtrie-2.5.0.tar.gz\r\n",
      "  Preparing metadata (setup.py) ... \u001b[?25l-\b \bdone\r\n",
      "\u001b[?25hBuilding wheels for collected packages: pygtrie\r\n",
      "  Building wheel for pygtrie (setup.py) ... \u001b[?25l-\b \b\\\b \bdone\r\n",
      "\u001b[?25h  Created wheel for pygtrie: filename=pygtrie-2.5.0-py3-none-any.whl size=20942 sha256=98309434fa463ff1bdd60fafe6845d9a9bdfc264ee5262af5bf8e2f3593810ad\r\n",
      "  Stored in directory: /root/.cache/pip/wheels/78/28/09/b62c97a3e77102645c7ecc78c97580ad57090b1eee5438d6ac\r\n",
      "Successfully built pygtrie\r\n",
      "Installing collected packages: pygtrie\r\n",
      "Successfully installed pygtrie-2.5.0\r\n",
      "Looking in links: ./\r\n",
      "Processing ./pyctcdecode/sortedcontainers-2.4.0-py2.py3-none-any.whl\r\n",
      "sortedcontainers is already installed with the same version as the provided wheel. Use --force-reinstall to force an installation of the wheel.\r\n",
      "Looking in links: ./\r\n",
      "Processing ./pyctcdecode/pyctcdecode-0.4.0-py2.py3-none-any.whl\r\n",
      "Installing collected packages: pyctcdecode\r\n",
      "Successfully installed pyctcdecode-0.4.0\r\n",
      "pypikenlm/\r\n",
      "pypikenlm/pypi-kenlm-0.1.20220713.tar.gz\r\n",
      "Looking in links: ./\r\n",
      "Processing ./pypikenlm/pypi-kenlm-0.1.20220713.tar.gz\r\n",
      "  Preparing metadata (setup.py) ... \u001b[?25ldone\r\n",
      "\u001b[?25hBuilding wheels for collected packages: pypi-kenlm\r\n",
      "  Building wheel for pypi-kenlm (setup.py) ... \u001b[?25l-\b \b\\\b \bdone\r\n",
      "\u001b[?25h  Created wheel for pypi-kenlm: filename=pypi_kenlm-0.1.20220713-cp310-cp310-linux_x86_64.whl size=333233 sha256=6950a7b95fce53d74e3d05eae698a4181f46d7e7207c664f4fc99641618a6d9d\r\n",
      "  Stored in directory: /root/.cache/pip/wheels/1e/7a/db/27645fac296d5d5ba5c461b1af834eebc0ba4643290dbc5476\r\n",
      "Successfully built pypi-kenlm\r\n",
      "Installing collected packages: pypi-kenlm\r\n",
      "Successfully installed pypi-kenlm-0.1.20220713\r\n"
     ]
    }
   ],
   "source": [
    "!cp -r ../input/python-packages2 ./\n",
    "\n",
    "!tar xvfz ./python-packages2/jiwer.tgz\n",
    "!pip install ./jiwer/jiwer-2.3.0-py3-none-any.whl -f ./ --no-index\n",
    "!tar xvfz ./python-packages2/normalizer.tgz\n",
    "!pip install ./normalizer/bnunicodenormalizer-0.0.24.tar.gz -f ./ --no-index\n",
    "!tar xvfz ./python-packages2/pyctcdecode.tgz\n",
    "!pip install ./pyctcdecode/attrs-22.1.0-py2.py3-none-any.whl -f ./ --no-index --no-deps\n",
    "!pip install ./pyctcdecode/exceptiongroup-1.0.0rc9-py3-none-any.whl -f ./ --no-index --no-deps\n",
    "!pip install ./pyctcdecode/hypothesis-6.54.4-py3-none-any.whl -f ./ --no-index --no-deps\n",
    "!pip install ./pyctcdecode/numpy-1.21.6-cp37-cp37m-manylinux_2_12_x86_64.manylinux2010_x86_64.whl -f ./ --no-index --no-deps\n",
    "!pip install ./pyctcdecode/pygtrie-2.5.0.tar.gz -f ./ --no-index --no-deps\n",
    "!pip install ./pyctcdecode/sortedcontainers-2.4.0-py2.py3-none-any.whl -f ./ --no-index --no-deps\n",
    "!pip install ./pyctcdecode/pyctcdecode-0.4.0-py2.py3-none-any.whl -f ./ --no-index --no-deps\n",
    "\n",
    "!tar xvfz ./python-packages2/pypikenlm.tgz\n",
    "!pip install ./pypikenlm/pypi-kenlm-0.1.20220713.tar.gz -f ./ --no-index --no-deps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "5e7b490e",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-10-14T19:43:24.254866Z",
     "iopub.status.busy": "2023-10-14T19:43:24.254509Z",
     "iopub.status.idle": "2023-10-14T19:43:40.171600Z",
     "shell.execute_reply": "2023-10-14T19:43:40.170587Z"
    },
    "papermill": {
     "duration": 15.926711,
     "end_time": "2023-10-14T19:43:40.173855",
     "exception": false,
     "start_time": "2023-10-14T19:43:24.247144",
     "status": "completed"
    },
    "tags": []
   },
   "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": [
    "import torch \n",
    "import torch.nn as nn\n",
    "import torchaudio\n",
    "import torchaudio.transforms as tat\n",
    "from datasets import load_dataset, load_metric, Audio\n",
    "import os\n",
    "\n",
    "import typing as tp\n",
    "from pathlib import Path\n",
    "from functools import partial\n",
    "from dataclasses import dataclass, field\n",
    "from typing import Any, Dict, List, Optional, Union\n",
    "\n",
    "import pandas as pd\n",
    "import pyctcdecode\n",
    "import numpy as np\n",
    "from tqdm.notebook import tqdm\n",
    "\n",
    "import librosa\n",
    "import gc\n",
    "import jiwer\n",
    "import pyctcdecode\n",
    "import kenlm\n",
    "import torch\n",
    "from transformers import Wav2Vec2Processor, Wav2Vec2ProcessorWithLM, Wav2Vec2ForCTC\n",
    "from transformers import TrainingArguments, Trainer, EarlyStoppingCallback\n",
    "from bnunicodenormalizer import Normalizer\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "torchaudio.set_audio_backend(\"soundfile\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "d7b84ff1",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-10-14T19:43:40.187086Z",
     "iopub.status.busy": "2023-10-14T19:43:40.186145Z",
     "iopub.status.idle": "2023-10-14T19:43:40.191409Z",
     "shell.execute_reply": "2023-10-14T19:43:40.190386Z"
    },
    "papermill": {
     "duration": 0.01366,
     "end_time": "2023-10-14T19:43:40.193298",
     "exception": false,
     "start_time": "2023-10-14T19:43:40.179638",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "### hyper-parameters\n",
    "SR = 16000\n",
    "torch.backends.cudnn.benchmark = True\n",
    "output_dir = \"./\"\n",
    "MODEL_PATH = \"/kaggle/input/ai4bharat-indicwav2vec-v1-bengali/indicwav2vec_v1_bengali\"\n",
    "LM_PATH = \"/kaggle/input/arijitx-full-model/wav2vec2-xls-r-300m-bengali/language_model\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "20ea5aa9",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-10-14T19:43:40.205678Z",
     "iopub.status.busy": "2023-10-14T19:43:40.205356Z",
     "iopub.status.idle": "2023-10-14T19:44:15.796728Z",
     "shell.execute_reply": "2023-10-14T19:44:15.795676Z"
    },
    "papermill": {
     "duration": 35.600467,
     "end_time": "2023-10-14T19:44:15.799171",
     "exception": false,
     "start_time": "2023-10-14T19:43:40.198704",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "processor = Wav2Vec2Processor.from_pretrained(MODEL_PATH)\n",
    "vocab_dict = processor.tokenizer.get_vocab()\n",
    "sorted_vocab_dict = {k: v for k, v in sorted(vocab_dict.items(), key=lambda item: item[1])}\n",
    "\n",
    "decoder = pyctcdecode.build_ctcdecoder(\n",
    "    list(sorted_vocab_dict.keys()),\n",
    "    str(LM_PATH+\"/5gram.bin\"),\n",
    "    str(LM_PATH+\"/unigrams.txt\"),\n",
    ")\n",
    "processor_with_lm = Wav2Vec2ProcessorWithLM(\n",
    "    feature_extractor=processor.feature_extractor,\n",
    "    tokenizer=processor.tokenizer,\n",
    "    decoder=decoder\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2d0180d2",
   "metadata": {
    "papermill": {
     "duration": 0.005767,
     "end_time": "2023-10-14T19:44:15.810999",
     "exception": false,
     "start_time": "2023-10-14T19:44:15.805232",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "- From @mbmmurad's [Dataset overlaps with CommonVoice 11 bn](https://www.kaggle.com/code/mbmmurad/dataset-overlaps-with-commonvoice-11-bn), The competition dataset might contain the audios of the mozilla-foundation/common_voice_11_0 dataset. Here I just simply exclude them from the validation set.\n",
    "- Also, I use @UmongSain's normalized data [here](https://www.kaggle.com/code/umongsain/macro-normalization/notebook). Thanks to him!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "468ff1e3",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-10-14T19:44:15.824917Z",
     "iopub.status.busy": "2023-10-14T19:44:15.824241Z",
     "iopub.status.idle": "2023-10-14T19:44:24.996408Z",
     "shell.execute_reply": "2023-10-14T19:44:24.995018Z"
    },
    "papermill": {
     "duration": 9.18191,
     "end_time": "2023-10-14T19:44:24.998717",
     "exception": false,
     "start_time": "2023-10-14T19:44:15.816807",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "963636\n",
      "706689\n"
     ]
    }
   ],
   "source": [
    "sentences = pd.read_csv(\"/kaggle/input/macro-normalization/normalized.csv\")\n",
    "indexes = set(pd.read_csv(\"/kaggle/input/dataset-overlaps-with-commonvoice-11-bn/indexes.csv\")['id'])\n",
    "print(len(sentences))\n",
    "sentences = sentences[~((sentences.index.isin(indexes))&(sentences['split']=='train'))].reset_index(drop=True)\n",
    "print(len(sentences))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c950e768",
   "metadata": {
    "papermill": {
     "duration": 0.005828,
     "end_time": "2023-10-14T19:44:25.011109",
     "exception": false,
     "start_time": "2023-10-14T19:44:25.005281",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "* sample 10% data from \"valid\" part into validation set, 90% into training set.\n",
    "* sample 5% data from \"train\" part, and additionally sample 8% from it into validation set, 92% into training set.\n",
    "* There will be **57776** train data, **5667** valid data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "8586cfdf",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-10-14T19:44:25.027109Z",
     "iopub.status.busy": "2023-10-14T19:44:25.026695Z",
     "iopub.status.idle": "2023-10-14T19:44:25.473758Z",
     "shell.execute_reply": "2023-10-14T19:44:25.472308Z"
    },
    "papermill": {
     "duration": 0.458765,
     "end_time": "2023-10-14T19:44:25.475822",
     "exception": false,
     "start_time": "2023-10-14T19:44:25.017057",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "706689\n",
      "57776\n",
      "5667\n"
     ]
    }
   ],
   "source": [
    "data_0 = sentences.loc[sentences['split']=='valid'].reset_index(drop=True)\n",
    "valid_0 = data_0.sample(frac=0.1, random_state=42)\n",
    "train_0 = data_0[~data_0.index.isin(valid_0.index)]\n",
    "\n",
    "data_1 = sentences.loc[sentences['split']=='train'].reset_index(drop=True).sample(frac=0.05, random_state=42)\n",
    "valid_1 = data_1.sample(frac=0.08, random_state=42)\n",
    "train_1 = data_1[~data_1.index.isin(valid_1.index)]\n",
    "\n",
    "train = pd.concat([train_0, train_1], axis=0).sample(frac=1, random_state=42).reset_index(drop=True)\n",
    "valid = pd.concat([valid_0, valid_1], axis=0).sample(frac=1, random_state=42).reset_index(drop=True)\n",
    "\n",
    "del data_0, data_1, valid_0, valid_1, train_0, train_1\n",
    "all_ids = sentences['id'].to_list()\n",
    "train_ids = train['id'].to_list()\n",
    "valid_ids = valid['id'].to_list()\n",
    "\n",
    "# in kaggle notebook, validating is very time-consuming, so here I use a very small validation set, rather than 5667.\n",
    "valid = valid.sample(n=500, random_state=42)\n",
    "\n",
    "print(len(all_ids))\n",
    "print(len(train_ids))\n",
    "print(len(valid_ids))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "258b6df9",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-10-14T19:44:25.489468Z",
     "iopub.status.busy": "2023-10-14T19:44:25.488831Z",
     "iopub.status.idle": "2023-10-14T19:44:25.599609Z",
     "shell.execute_reply": "2023-10-14T19:44:25.598411Z"
    },
    "papermill": {
     "duration": 0.120232,
     "end_time": "2023-10-14T19:44:25.602039",
     "exception": false,
     "start_time": "2023-10-14T19:44:25.481807",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class W2v2Dataset(torch.utils.data.Dataset):\n",
    "    def __init__(self, df):\n",
    "        self.df = df\n",
    "        self.pathes = df['id'].values\n",
    "        self.sentences = df['normalized'].values\n",
    "        self.resampler = tat.Resample(32000, SR)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        apath = f'/kaggle/input/bengaliai-speech/train_mp3s/{self.pathes[idx]}.mp3'\n",
    "        waveform, sample_rate = torchaudio.load(apath, format=\"mp3\")\n",
    "        waveform = self.resampler(waveform)\n",
    "        batch = dict()\n",
    "        y = processor(waveform.reshape(-1), sampling_rate=SR).input_values[0] \n",
    "        batch[\"input_values\"] = y\n",
    "        with processor.as_target_processor():\n",
    "            batch[\"labels\"] = processor(self.sentences[idx]).input_ids       \n",
    "        \n",
    "        return batch\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.df)\n",
    "\n",
    "train_dataset = W2v2Dataset(train)\n",
    "valid_dataset = W2v2Dataset(valid)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "0fbfbce4",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-10-14T19:44:25.616680Z",
     "iopub.status.busy": "2023-10-14T19:44:25.616346Z",
     "iopub.status.idle": "2023-10-14T19:44:25.626173Z",
     "shell.execute_reply": "2023-10-14T19:44:25.625134Z"
    },
    "papermill": {
     "duration": 0.019192,
     "end_time": "2023-10-14T19:44:25.628105",
     "exception": false,
     "start_time": "2023-10-14T19:44:25.608913",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "@dataclass\n",
    "class DataCollatorCTCWithPadding:\n",
    "    \"\"\"\n",
    "    Data collator that will dynamically pad the inputs received.\n",
    "    Args:\n",
    "        processor (:class:`~transformers.Wav2Vec2Processor`)\n",
    "            The processor used for proccessing the data.\n",
    "        padding (:obj:`bool`, :obj:`str` or :class:`~transformers.tokenization_utils_base.PaddingStrategy`, `optional`, defaults to :obj:`True`):\n",
    "            Select a strategy to pad the returned sequences (according to the model's padding side and padding index)\n",
    "            among:\n",
    "            * :obj:`True` or :obj:`'longest'`: Pad to the longest sequence in the batch (or no padding if only a single\n",
    "              sequence if provided).\n",
    "            * :obj:`'max_length'`: Pad to a maximum length specified with the argument :obj:`max_length` or to the\n",
    "              maximum acceptable input length for the model if that argument is not provided.\n",
    "            * :obj:`False` or :obj:`'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of\n",
    "              different lengths).\n",
    "        max_length (:obj:`int`, `optional`):\n",
    "            Maximum length of the ``input_values`` of the returned list and optionally padding length (see above).\n",
    "        max_length_labels (:obj:`int`, `optional`):\n",
    "            Maximum length of the ``labels`` returned list and optionally padding length (see above).\n",
    "        pad_to_multiple_of (:obj:`int`, `optional`):\n",
    "            If set will pad the sequence to a multiple of the provided value.\n",
    "            This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute capability >=\n",
    "            7.5 (Volta).\n",
    "    \"\"\"\n",
    "\n",
    "    processor: Wav2Vec2Processor\n",
    "    padding: Union[bool, str] = True\n",
    "    max_length: Optional[int] = None\n",
    "    max_length_labels: Optional[int] = None\n",
    "    pad_to_multiple_of: Optional[int] = None\n",
    "    pad_to_multiple_of_labels: Optional[int] = None\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 lenghts and need\n",
    "        # different padding methods\n",
    "        input_features = [{\"input_values\": feature[\"input_values\"]} for feature in features]\n",
    "        label_features = [{\"input_ids\": feature[\"labels\"]} for feature in features]\n",
    "\n",
    "        batch = self.processor.pad(\n",
    "            input_features,\n",
    "            padding=self.padding,\n",
    "            max_length=self.max_length,\n",
    "            pad_to_multiple_of=self.pad_to_multiple_of,\n",
    "            return_tensors=\"pt\",\n",
    "        )\n",
    "        with self.processor.as_target_processor():\n",
    "            labels_batch = self.processor.pad(\n",
    "                label_features,\n",
    "                padding=self.padding,\n",
    "                max_length=self.max_length_labels,\n",
    "                pad_to_multiple_of=self.pad_to_multiple_of_labels,\n",
    "                return_tensors=\"pt\",\n",
    "            )\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",
    "        batch[\"labels\"] = labels\n",
    "\n",
    "        return batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "e383890d",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-10-14T19:44:25.641825Z",
     "iopub.status.busy": "2023-10-14T19:44:25.641342Z",
     "iopub.status.idle": "2023-10-14T19:44:25.645973Z",
     "shell.execute_reply": "2023-10-14T19:44:25.644808Z"
    },
    "papermill": {
     "duration": 0.013793,
     "end_time": "2023-10-14T19:44:25.648020",
     "exception": false,
     "start_time": "2023-10-14T19:44:25.634227",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "data_collator = DataCollatorCTCWithPadding(processor=processor, padding=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d3ad60ff",
   "metadata": {
    "papermill": {
     "duration": 0.005942,
     "end_time": "2023-10-14T19:44:25.660002",
     "exception": false,
     "start_time": "2023-10-14T19:44:25.654060",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "- In kaggle notebook, there is an error: **cannot import name 'compute_measures' from 'jiwer' (unknown location)**. But in my local notebook, there is no such error."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "98e8a7af",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-10-14T19:44:25.674348Z",
     "iopub.status.busy": "2023-10-14T19:44:25.673295Z",
     "iopub.status.idle": "2023-10-14T19:44:25.679837Z",
     "shell.execute_reply": "2023-10-14T19:44:25.678861Z"
    },
    "papermill": {
     "duration": 0.015837,
     "end_time": "2023-10-14T19:44:25.681853",
     "exception": false,
     "start_time": "2023-10-14T19:44:25.666016",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "#wer_metric = load_metric(\"wer\")\n",
    "\n",
    "def compute_metrics(pred):\n",
    "    pred_logits = pred.predictions\n",
    "    pred_ids = np.argmax(pred_logits, axis=-1)\n",
    "\n",
    "    pred.label_ids[pred.label_ids == -100] = processor.tokenizer.pad_token_id\n",
    "\n",
    "    pred_str = processor.batch_decode(pred_ids)\n",
    "    # we do not want to group tokens when computing the metrics\n",
    "    label_str = processor.batch_decode(pred.label_ids, group_tokens=False)\n",
    "\n",
    "    wer = wer_metric.compute(predictions=pred_str, references=label_str)\n",
    "\n",
    "    return {\"wer\": wer}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "3a5c1258",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-10-14T19:44:25.695703Z",
     "iopub.status.busy": "2023-10-14T19:44:25.695375Z",
     "iopub.status.idle": "2023-10-14T19:44:38.571371Z",
     "shell.execute_reply": "2023-10-14T19:44:38.570053Z"
    },
    "papermill": {
     "duration": 12.885525,
     "end_time": "2023-10-14T19:44:38.573743",
     "exception": false,
     "start_time": "2023-10-14T19:44:25.688218",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "model = Wav2Vec2ForCTC.from_pretrained(\n",
    "    MODEL_PATH,\n",
    "    attention_dropout=0.1,\n",
    "    hidden_dropout=0.1,\n",
    "    feat_proj_dropout=0.0,\n",
    "    mask_time_prob=0.05,\n",
    "    layerdrop=0.1,\n",
    "    #gradient_checkpointing=True, \n",
    "    ctc_loss_reduction=\"mean\", \n",
    "    pad_token_id=processor.tokenizer.pad_token_id,\n",
    "    vocab_size=len(processor.tokenizer),\n",
    "    ctc_zero_infinity=True,\n",
    "    diversity_loss_weight=100 \n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "3310c7ac",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-10-14T19:44:38.587573Z",
     "iopub.status.busy": "2023-10-14T19:44:38.587240Z",
     "iopub.status.idle": "2023-10-14T19:44:38.592315Z",
     "shell.execute_reply": "2023-10-14T19:44:38.590932Z"
    },
    "papermill": {
     "duration": 0.014068,
     "end_time": "2023-10-14T19:44:38.594251",
     "exception": false,
     "start_time": "2023-10-14T19:44:38.580183",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# you can freeze some params\n",
    "model.freeze_feature_extractor()\n",
    "#model.freeze_feature_encoder()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7831b5c3",
   "metadata": {
    "papermill": {
     "duration": 0.005792,
     "end_time": "2023-10-14T19:44:38.606170",
     "exception": false,
     "start_time": "2023-10-14T19:44:38.600378",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "- As a demo, \"**num_train_epochs**\", \"**eval_steps**\" and \"**early_stopping_patience**\" are set to very small values, you can make them larger.\n",
    "- If there is no error about jiwer, you can set **metric_for_best_model**=\"wer\", and remember to set **greater_is_better**=False and use **compute_metrics**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "af56ee62",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-10-14T19:44:38.623329Z",
     "iopub.status.busy": "2023-10-14T19:44:38.623001Z",
     "iopub.status.idle": "2023-10-14T19:44:38.710690Z",
     "shell.execute_reply": "2023-10-14T19:44:38.709660Z"
    },
    "papermill": {
     "duration": 0.097575,
     "end_time": "2023-10-14T19:44:38.713112",
     "exception": false,
     "start_time": "2023-10-14T19:44:38.615537",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "training_args = TrainingArguments(\n",
    "    output_dir=output_dir,\n",
    "    overwrite_output_dir=True,\n",
    "    group_by_length=False,\n",
    "    lr_scheduler_type='cosine',\n",
    "    weight_decay=0.01,\n",
    "    per_device_train_batch_size=4,\n",
    "    per_device_eval_batch_size=16,\n",
    "    gradient_accumulation_steps=1,\n",
    "    evaluation_strategy=\"steps\",\n",
    "    save_strategy=\"steps\",\n",
    "    #max_steps=100, # you can change to \"num_train_epochs\"\n",
    "    num_train_epochs=20,\n",
    "    fp16=True,\n",
    "    save_steps=20,\n",
    "    eval_steps=20,\n",
    "    logging_steps=20,\n",
    "    learning_rate=2e-5,\n",
    "    warmup_steps=600,\n",
    "    save_total_limit=1,\n",
    "    load_best_model_at_end=True,\n",
    "    #metric_for_best_model=\"wer\",\n",
    "    #greater_is_better=False,\n",
    "    prediction_loss_only=False,\n",
    "    auto_find_batch_size=True,\n",
    "    report_to=\"none\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "f20e2fc8",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-10-14T19:44:38.726571Z",
     "iopub.status.busy": "2023-10-14T19:44:38.726062Z",
     "iopub.status.idle": "2023-10-14T19:44:44.100833Z",
     "shell.execute_reply": "2023-10-14T19:44:44.099807Z"
    },
    "papermill": {
     "duration": 5.384409,
     "end_time": "2023-10-14T19:44:44.103312",
     "exception": false,
     "start_time": "2023-10-14T19:44:38.718903",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    data_collator=data_collator,\n",
    "    args=training_args,\n",
    "    #compute_metrics=compute_metrics,\n",
    "    train_dataset=train_dataset,\n",
    "    eval_dataset=valid_dataset,\n",
    "    tokenizer=processor.feature_extractor,\n",
    "    callbacks=[EarlyStoppingCallback(early_stopping_patience=1)],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "72f3c67d",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-10-14T19:44:44.120627Z",
     "iopub.status.busy": "2023-10-14T19:44:44.119624Z",
     "iopub.status.idle": "2023-10-14T20:27:02.267388Z",
     "shell.execute_reply": "2023-10-14T20:27:02.266073Z"
    },
    "papermill": {
     "duration": 2538.162117,
     "end_time": "2023-10-14T20:27:02.273327",
     "exception": false,
     "start_time": "2023-10-14T19:44:44.111210",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='680' max='288880' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [   680/288880 42:05 < 298:10:30, 0.27 it/s, Epoch 0/20]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Step</th>\n",
       "      <th>Training Loss</th>\n",
       "      <th>Validation Loss</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>20</td>\n",
       "      <td>2.383300</td>\n",
       "      <td>2.591151</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>40</td>\n",
       "      <td>2.061000</td>\n",
       "      <td>2.586119</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>60</td>\n",
       "      <td>2.400700</td>\n",
       "      <td>2.566355</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>80</td>\n",
       "      <td>2.323500</td>\n",
       "      <td>2.530945</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>100</td>\n",
       "      <td>2.376500</td>\n",
       "      <td>2.470588</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>120</td>\n",
       "      <td>2.296000</td>\n",
       "      <td>2.391886</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>140</td>\n",
       "      <td>2.194800</td>\n",
       "      <td>2.305820</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>160</td>\n",
       "      <td>2.305300</td>\n",
       "      <td>2.217217</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>180</td>\n",
       "      <td>2.104800</td>\n",
       "      <td>2.123956</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>200</td>\n",
       "      <td>1.916400</td>\n",
       "      <td>2.024333</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>220</td>\n",
       "      <td>1.945100</td>\n",
       "      <td>1.919927</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>240</td>\n",
       "      <td>1.821300</td>\n",
       "      <td>1.816791</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>260</td>\n",
       "      <td>2.060600</td>\n",
       "      <td>1.730359</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>280</td>\n",
       "      <td>2.090900</td>\n",
       "      <td>1.634165</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>300</td>\n",
       "      <td>1.894500</td>\n",
       "      <td>1.550624</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>320</td>\n",
       "      <td>1.845700</td>\n",
       "      <td>1.477670</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>340</td>\n",
       "      <td>1.559600</td>\n",
       "      <td>1.413101</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>360</td>\n",
       "      <td>1.688000</td>\n",
       "      <td>1.339739</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>380</td>\n",
       "      <td>1.849000</td>\n",
       "      <td>1.298396</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>400</td>\n",
       "      <td>1.583700</td>\n",
       "      <td>1.266363</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>420</td>\n",
       "      <td>1.573000</td>\n",
       "      <td>1.227807</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>440</td>\n",
       "      <td>1.594600</td>\n",
       "      <td>1.195420</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>460</td>\n",
       "      <td>1.440100</td>\n",
       "      <td>1.165610</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>480</td>\n",
       "      <td>1.568300</td>\n",
       "      <td>1.135737</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>500</td>\n",
       "      <td>1.513700</td>\n",
       "      <td>1.106301</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>520</td>\n",
       "      <td>1.234600</td>\n",
       "      <td>1.078182</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>540</td>\n",
       "      <td>1.603700</td>\n",
       "      <td>1.054003</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>560</td>\n",
       "      <td>1.220500</td>\n",
       "      <td>1.031834</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>580</td>\n",
       "      <td>1.328900</td>\n",
       "      <td>1.022463</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>600</td>\n",
       "      <td>1.307200</td>\n",
       "      <td>1.004617</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>620</td>\n",
       "      <td>1.125200</td>\n",
       "      <td>1.001430</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>640</td>\n",
       "      <td>1.085000</td>\n",
       "      <td>0.983673</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>660</td>\n",
       "      <td>1.299100</td>\n",
       "      <td>0.970705</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>680</td>\n",
       "      <td>1.084300</td>\n",
       "      <td>0.971934</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "trainer.train()\n",
    "trainer.save_model(output_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e92c153",
   "metadata": {
    "papermill": {
     "duration": 0.066229,
     "end_time": "2023-10-14T20:27:02.397011",
     "exception": false,
     "start_time": "2023-10-14T20:27:02.330782",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "- To improve scores you can: \n",
    "    * use different pretrained models\n",
    "    * alter the parameters\n",
    "    * choose more data\n",
    "    * filter data in another way."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.10.12"
  },
  "papermill": {
   "default_parameters": {},
   "duration": 2698.086818,
   "end_time": "2023-10-14T20:27:06.153695",
   "environment_variables": {},
   "exception": null,
   "input_path": "__notebook__.ipynb",
   "output_path": "__notebook__.ipynb",
   "parameters": {},
   "start_time": "2023-10-14T19:42:08.066877",
   "version": "2.4.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}