week06/asr_ensemble.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "41b0838c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.chdir('../asr/')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "f7982d04",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "NOTE! Installing ujson may make loading annotations faster.\n"
     ]
    }
   ],
   "source": [
    "from typing import Iterable, Tuple, List\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "import sentencepiece\n",
    "from omegaconf import OmegaConf\n",
    "import pytorch_lightning as pl\n",
    "import pandas as pd\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "from src.models import LASModel, CTCModel\n",
    "from src.metrics import WER"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "11eb14ca",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1mdiff --git a/conf/conformer_ctc.yaml b/conf/conformer_ctc_wide.yaml\u001b[m\r\n",
      "\u001b[1mindex 10f8a87..6b58e13 100755\u001b[m\r\n",
      "\u001b[1m--- a/conf/conformer_ctc.yaml\u001b[m\r\n",
      "\u001b[1m+++ b/conf/conformer_ctc_wide.yaml\u001b[m\r\n",
      "\u001b[36m@@ -9,11 +9,11 @@\u001b[m \u001b[mmodel:\u001b[m\r\n",
      "     dropout: 0.0\u001b[m\r\n",
      "     features_num: 64\u001b[m\r\n",
      "     subsampling_stride: 4\u001b[m\r\n",
      "\u001b[31m-    d_model: 256\u001b[m\r\n",
      "\u001b[31m-    n_layers: 10\u001b[m\r\n",
      "\u001b[32m+\u001b[m\u001b[32m    d_model: 320\u001b[m\r\n",
      "\u001b[32m+\u001b[m\u001b[32m    n_layers: 8\u001b[m\r\n",
      "     n_heads: 8\u001b[m\r\n",
      "     ff_exp_factor: 2\u001b[m\r\n",
      "\u001b[31m-    kernel_size: 7\u001b[m\r\n",
      "\u001b[32m+\u001b[m\u001b[32m    kernel_size: 15\u001b[m\r\n",
      "     \u001b[m\r\n",
      " \u001b[m\r\n",
      "   decoder:\u001b[m\r\n"
     ]
    }
   ],
   "source": [
    "!git diff --no-index conf/conformer_ctc.yaml conf/conformer_ctc_wide.yaml\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "37f0fb7c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def init_model(model: pl.LightningModule, ckpt_path: str) -> pl.LightningModule:\n",
    "    ckpt = torch.load(ckpt_path, map_location=\"cpu\")\n",
    "    model.load_state_dict(ckpt)\n",
    "    model.eval()\n",
    "    model.freeze()\n",
    "    return model\n",
    "\n",
    "\n",
    "def compute_wer(refs: Iterable[str], hyps: Iterable[str]) -> float:\n",
    "    wer = WER()\n",
    "    wer.update(refs, hyps)\n",
    "    return wer.compute()[0].item()\n",
    "\n",
    "\n",
    "class GreedyDecoderLAS:\n",
    "    def __init__(self, model: LASModel, max_steps=20):\n",
    "        self.model = model\n",
    "        self.max_steps = max_steps\n",
    "\n",
    "    def __call__(self, encoded: torch.Tensor) -> str:\n",
    "        \n",
    "        tokens = [self.model.decoder.tokenizer.bos_id()]\n",
    "\n",
    "        for _ in range(self.max_steps):\n",
    "            \n",
    "            tokens_batch = torch.tensor(tokens).unsqueeze(0).to(self.model.device)\n",
    "            att_mask = self.model.make_attention_mask(torch.tensor([tokens_batch.size(-1)])).to(self.model.device)\n",
    "            \n",
    "            distribution = self.model.decoder(\n",
    "                encoded=encoded, encoded_pad_mask=None,\n",
    "                target=tokens_batch, target_mask=att_mask, target_pad_mask=None\n",
    "            )\n",
    "        \n",
    "            best_next_token = distribution[0, -1].argmax()\n",
    "            \n",
    "            if best_next_token == self.model.decoder.tokenizer.eos_id():\n",
    "                break\n",
    "\n",
    "            tokens.append(best_next_token.item())\n",
    "        \n",
    "        return self.model.decoder.tokenizer.decode(tokens)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "208a50d9",
   "metadata": {},
   "source": [
    "## Single Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "005fa536",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = 'test_opus/farfield/manifest.jsonl'\n",
    "\n",
    "TOKENIZER_PATH = '../asr/data/checkpoints_and_tokenizer/tokenizer/'\n",
    "LAS_CKPT_PATH = '../asr/data/checkpoints_and_tokenizer/conformer_las_2epochs.ckpt'\n",
    "\n",
    "device = 'cuda:2' if torch.cuda.is_available() else 'cpu'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1be8fcb7",
   "metadata": {},
   "source": [
    "### LAS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "145ae422",
   "metadata": {},
   "outputs": [],
   "source": [
    "conf = OmegaConf.load(\"./conf/conformer_las.yaml\")\n",
    "conf.val_dataloader.dataset.manifest_name = dataset\n",
    "conf.model.decoder.tokenizer_path = os.path.join(TOKENIZER_PATH, \"bpe_1024_bos_eos.model\")\n",
    "\n",
    "conformer_las = init_model(\n",
    "    model=LASModel(conf=conf),\n",
    "    ckpt_path=LAS_CKPT_PATH\n",
    ")\n",
    "conformer_las = conformer_las.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "e0be8c68",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "11121920\n"
     ]
    }
   ],
   "source": [
    "print(sum(p.numel() for p in conformer_las.parameters()))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "e4bdd2f1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b84f2d18795b448cbe545f80bda4e853",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/60 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "las_decoder = GreedyDecoderLAS(conformer_las)\n",
    "\n",
    "refs, hyps_las = [], []\n",
    "\n",
    "for batch in tqdm(conformer_las.val_dataloader()):\n",
    "\n",
    "    features, features_len, targets, target_len = batch\n",
    "\n",
    "    encoded, encoded_len = conformer_las(features.to(device), features_len.to(device))\n",
    "    \n",
    "    for i in range(features.shape[0]):\n",
    "\n",
    "        encoder_states = encoded[\n",
    "            [i],\n",
    "            :encoded_len[i],\n",
    "            :\n",
    "        ]\n",
    "\n",
    "        ref_tokens = targets[i, :target_len[i]].tolist()\n",
    "\n",
    "        refs.append(\n",
    "            conformer_las.decoder.tokenizer.decode(ref_tokens)\n",
    "        )\n",
    "        hyps_las.append(\n",
    "            las_decoder(encoder_states)\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "ed6a77e8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.42290276288986206"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "compute_wer(refs, hyps_las)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8fcd289d",
   "metadata": {},
   "source": [
    "### CTC"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "1e656233",
   "metadata": {},
   "outputs": [],
   "source": [
    "CTC_WIDE_CKPT_PATH = '../asr/data/checkpoints_and_tokenizer/conformer_wide_7epochs_state_dict.ckpt'\n",
    "CTC_DEEP_CKPT_PATH = '../asr/data/checkpoints_and_tokenizer/conformer_7epochs_state_dict.ckpt'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "921279bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: load models, estimate WER"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "9d7cb60a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def decode_ctc_hyps(model: CTCModel) -> Tuple[List[str], List[str]]:\n",
    "    return [], []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "4e39f540",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "9989410\n"
     ]
    }
   ],
   "source": [
    "conf = OmegaConf.load(\"./conf/conformer_ctc.yaml\")\n",
    "conf.val_dataloader.dataset.manifest_name = dataset\n",
    "\n",
    "conformer_ctc = init_model(\n",
    "    model=CTCModel(conf=conf),\n",
    "    ckpt_path=CTC_DEEP_CKPT_PATH\n",
    ")\n",
    "\n",
    "refs, hyps_ctc = decode_ctc_hyps(conformer_ctc)\n",
    "print(sum(p.numel() for p in conformer_ctc.parameters()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "84971367",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "12486114\n"
     ]
    }
   ],
   "source": [
    "conf = OmegaConf.load(\"./conf/conformer_ctc_wide.yaml\")\n",
    "conf.val_dataloader.dataset.manifest_name = dataset\n",
    "\n",
    "conformer_ctc = init_model(\n",
    "    model=CTCModel(conf=conf),\n",
    "    ckpt_path=CTC_WIDE_CKPT_PATH\n",
    ")\n",
    "\n",
    "refs, hyps_ctc = decode_ctc_hyps(conformer_ctc)\n",
    "print(sum(p.numel() for p in conformer_ctc.parameters()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1cee2e7",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "d77ac872",
   "metadata": {},
   "source": [
    "## ROVER: Recognizer Output Voting Error Reduction — 5 points\n",
    "\n",
    "* [A post-processing system to yield reduced word error rates: Recognizer Output Voting Error Reduction (ROVER)](https://ieeexplore.ieee.org/document/659110)\n",
    "* [Improved ROVER using Language Model Information](https://www-tlp.limsi.fr/public/asr00_holger.pdf)\n",
    "\n",
    "Alignment + Voting\n",
    "\n",
    "![](images/rover_table.png)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33759511",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ! pip install crowd-kit\n",
    "from crowdkit.aggregation.texts import ROVER"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "aeabec60",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: aggregate hypotheses, estimate WER\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92346ae0",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34159309",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "59677ddc",
   "metadata": {},
   "source": [
    "## MBR: Minimum Bayes Risk — 5 points\n",
    "\n",
    "* [Minimum Bayes Risk Decoding and System Combination Based on a Recursion for Edit Distance](https://danielpovey.com/files/csl11_consensus.pdf)\n",
    "* [mbr-decoding blog-post](https://suzyahyah.github.io/bayesian%20inference/machine%20translation/2022/02/15/mbr-decoding.html)\n",
    "* [Combination of end-to-end and hybrid models for speech recognition](http://www.interspeech2020.org/uploadfile/pdf/Tue-1-8-4.pdf)\n",
    "\n",
    "![](images/mbr_scheme.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "179a5b5a",
   "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.10.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}