Transformer_LM.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "4PuLlMEIzqWE",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "4PuLlMEIzqWE",
    "outputId": "47e637c1-3f33-47f3-b0d5-9e39d07dc7bb"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: transformers in /csehome/heatz123/.conda/envs/vision/lib/python3.10/site-packages (4.24.0)\n",
      "Requirement already satisfied: datasets in /csehome/heatz123/.conda/envs/vision/lib/python3.10/site-packages (2.7.0)\n",
      "Requirement already satisfied: jiwer in /csehome/heatz123/.conda/envs/vision/lib/python3.10/site-packages (2.5.1)\n",
      "Collecting soundfile\n",
      "  Downloading soundfile-0.11.0-py2.py3-none-any.whl (23 kB)\n",
      "Requirement already satisfied: requests in /csehome/heatz123/.conda/envs/vision/lib/python3.10/site-packages (from transformers) (2.28.1)\n",
      "Requirement already satisfied: packaging>=20.0 in /csehome/heatz123/.conda/envs/vision/lib/python3.10/site-packages (from transformers) (21.3)\n",
      "Requirement already satisfied: regex!=2019.12.17 in /csehome/heatz123/.conda/envs/vision/lib/python3.10/site-packages (from transformers) (2022.10.31)\n",
      "Requirement already satisfied: filelock in /csehome/heatz123/.conda/envs/vision/lib/python3.10/site-packages (from transformers) (3.8.0)\n",
      "Requirement already satisfied: tokenizers!=0.11.3,<0.14,>=0.11.1 in /csehome/heatz123/.conda/envs/vision/lib/python3.10/site-packages (from transformers) (0.13.2)\n",
      "Requirement already satisfied: huggingface-hub<1.0,>=0.10.0 in /csehome/heatz123/.conda/envs/vision/lib/python3.10/site-packages (from transformers) (0.11.0)\n",
      "Requirement already satisfied: tqdm>=4.27 in /csehome/heatz123/.conda/envs/vision/lib/python3.10/site-packages (from transformers) (4.64.1)\n",
      "Requirement already satisfied: pyyaml>=5.1 in /csehome/heatz123/.conda/envs/vision/lib/python3.10/site-packages (from transformers) (6.0)\n",
      "Requirement already satisfied: numpy>=1.17 in /csehome/heatz123/.conda/envs/vision/lib/python3.10/site-packages (from transformers) (1.23.4)\n",
      "Requirement already satisfied: xxhash in /csehome/heatz123/.conda/envs/vision/lib/python3.10/site-packages (from datasets) (3.1.0)\n",
      "Requirement already satisfied: fsspec[http]>=2021.11.1 in /csehome/heatz123/.conda/envs/vision/lib/python3.10/site-packages (from datasets) (2022.11.0)\n",
      "Requirement already satisfied: responses<0.19 in /csehome/heatz123/.conda/envs/vision/lib/python3.10/site-packages (from datasets) (0.18.0)\n",
      "Requirement already satisfied: multiprocess in /csehome/heatz123/.conda/envs/vision/lib/python3.10/site-packages (from datasets) (0.70.14)\n",
      "Requirement already satisfied: aiohttp in /csehome/heatz123/.conda/envs/vision/lib/python3.10/site-packages (from datasets) (3.8.3)\n",
      "Requirement already satisfied: pyarrow>=6.0.0 in /csehome/heatz123/.conda/envs/vision/lib/python3.10/site-packages (from datasets) (10.0.0)\n",
      "Requirement already satisfied: pandas in /csehome/heatz123/.conda/envs/vision/lib/python3.10/site-packages (from datasets) (1.5.1)\n",
      "Requirement already satisfied: dill<0.3.7 in /csehome/heatz123/.conda/envs/vision/lib/python3.10/site-packages (from datasets) (0.3.6)\n",
      "Requirement already satisfied: levenshtein==0.20.2 in /csehome/heatz123/.conda/envs/vision/lib/python3.10/site-packages (from jiwer) (0.20.2)\n",
      "Requirement already satisfied: rapidfuzz<3.0.0,>=2.3.0 in /csehome/heatz123/.conda/envs/vision/lib/python3.10/site-packages (from levenshtein==0.20.2->jiwer) (2.13.2)\n",
      "Requirement already satisfied: cffi>=1.0 in /csehome/heatz123/.conda/envs/vision/lib/python3.10/site-packages (from soundfile) (1.15.1)\n",
      "Requirement already satisfied: pycparser in /csehome/heatz123/.conda/envs/vision/lib/python3.10/site-packages (from cffi>=1.0->soundfile) (2.21)\n",
      "Requirement already satisfied: attrs>=17.3.0 in /csehome/heatz123/.local/lib/python3.10/site-packages (from aiohttp->datasets) (22.1.0)\n",
      "Requirement already satisfied: yarl<2.0,>=1.0 in /csehome/heatz123/.conda/envs/vision/lib/python3.10/site-packages (from aiohttp->datasets) (1.8.1)\n",
      "Requirement already satisfied: charset-normalizer<3.0,>=2.0 in /csehome/heatz123/.conda/envs/vision/lib/python3.10/site-packages (from aiohttp->datasets) (2.0.4)\n",
      "Requirement already satisfied: async-timeout<5.0,>=4.0.0a3 in /csehome/heatz123/.conda/envs/vision/lib/python3.10/site-packages (from aiohttp->datasets) (4.0.2)\n",
      "Requirement already satisfied: multidict<7.0,>=4.5 in /csehome/heatz123/.conda/envs/vision/lib/python3.10/site-packages (from aiohttp->datasets) (6.0.2)\n",
      "Requirement already satisfied: aiosignal>=1.1.2 in /csehome/heatz123/.conda/envs/vision/lib/python3.10/site-packages (from aiohttp->datasets) (1.3.1)\n",
      "Requirement already satisfied: frozenlist>=1.1.1 in /csehome/heatz123/.conda/envs/vision/lib/python3.10/site-packages (from aiohttp->datasets) (1.3.3)\n",
      "Requirement already satisfied: typing-extensions>=3.7.4.3 in /csehome/heatz123/.conda/envs/vision/lib/python3.10/site-packages (from huggingface-hub<1.0,>=0.10.0->transformers) (4.3.0)\n",
      "Requirement already satisfied: pyparsing!=3.0.5,>=2.0.2 in /csehome/heatz123/.conda/envs/vision/lib/python3.10/site-packages (from packaging>=20.0->transformers) (3.0.9)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /csehome/heatz123/.conda/envs/vision/lib/python3.10/site-packages (from requests->transformers) (3.4)\n",
      "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /csehome/heatz123/.conda/envs/vision/lib/python3.10/site-packages (from requests->transformers) (1.26.12)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /csehome/heatz123/.conda/envs/vision/lib/python3.10/site-packages (from requests->transformers) (2022.9.24)\n",
      "Requirement already satisfied: python-dateutil>=2.8.1 in /csehome/heatz123/.conda/envs/vision/lib/python3.10/site-packages (from pandas->datasets) (2.8.2)\n",
      "Requirement already satisfied: pytz>=2020.1 in /csehome/heatz123/.conda/envs/vision/lib/python3.10/site-packages (from pandas->datasets) (2022.6)\n",
      "Requirement already satisfied: six>=1.5 in /csehome/heatz123/.conda/envs/vision/lib/python3.10/site-packages (from python-dateutil>=2.8.1->pandas->datasets) (1.16.0)\n",
      "Installing collected packages: soundfile\n",
      "Successfully installed soundfile-0.11.0\n"
     ]
    }
   ],
   "source": [
    "!pip install transformers datasets jiwer soundfile"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "084d826c-7b1c-40b7-bb46-110f085488b4",
   "metadata": {
    "id": "084d826c-7b1c-40b7-bb46-110f085488b4"
   },
   "outputs": [],
   "source": [
    "from transformers import Wav2Vec2Processor, Wav2Vec2ForCTC\n",
    "from datasets import load_dataset\n",
    "import torch\n",
    "import numpy as np\n",
    "import soundfile as sf"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "nZv1Wqw6t-Dd",
   "metadata": {
    "id": "nZv1Wqw6t-Dd"
   },
   "source": [
    "# Load Acoustic Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "K6jw6thWpgdj",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "K6jw6thWpgdj",
    "outputId": "75e442ee-f876-4720-c537-e4e3dba04aa7"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using custom data configuration Siyong--speech_timit-e00a2cf2b9b45bd5\n",
      "Found cached dataset parquet (/csehome/heatz123/.cache/huggingface/datasets/Siyong___parquet/Siyong--speech_timit-e00a2cf2b9b45bd5/0.0.0/2a3b91fbd88a2c90d1dbbb32b460cf621d31bd5b05b934492fdef7d8d6f236ec)\n",
      "Loading cached processed dataset at /csehome/heatz123/.cache/huggingface/datasets/Siyong___parquet/Siyong--speech_timit-e00a2cf2b9b45bd5/0.0.0/2a3b91fbd88a2c90d1dbbb32b460cf621d31bd5b05b934492fdef7d8d6f236ec/cache-6fb23cde0307bff6.arrow\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "reference: she had your dark suit in greasy wash water all year\n",
      "predicted: she had your dark suit in greasy wash water all year\n",
      "--\n",
      "reference: she had your dark suit in greasy wash water all year\n",
      "predicted: pshe had your dark suit in greasy wash water all year\n",
      "--\n",
      "reference: there are more obvious nymphomaniacs on any privateeye series\n",
      "predicted: ther are more obvious nimfom many acts on any privite eye series\n",
      "--\n",
      "reference: only the best players enjoy popularity\n",
      "predicted: only the best players enjoy popularity\n",
      "--\n",
      "reference: december and january are nice months to spend in miami\n",
      "predicted: tdisember and january ar nice months tospend in my amy\n",
      "--\n",
      "reference: keep the thermometer under your tongue\n",
      "predicted: keep the thermometer under your tongue\n",
      "--\n",
      "reference: you're a taxpayer householder landlord\n",
      "predicted: you'r a taxpayer householder landlord\n",
      "--\n",
      "reference: does creole cooking use curry\n",
      "predicted: does creole coking use curry\n",
      "--\n",
      "reference: scholastic aptitude is judged by standardized tests\n",
      "predicted: scholastic aptitude is judged by standardized tests\n",
      "--\n",
      "reference: are you looking for employment\n",
      "predicted: are oking for imploymnt\n",
      "--\n"
     ]
    }
   ],
   "source": [
    "model_name = \"elgeish/wav2vec2-base-timit-asr\"\n",
    "processor = Wav2Vec2Processor.from_pretrained(model_name)\n",
    "model = Wav2Vec2ForCTC.from_pretrained(model_name)\n",
    "model.eval()\n",
    "\n",
    "dataset = load_dataset(\"Siyong/speech_timit\", split=\"test\").select(range(10))\n",
    "char_translations = str.maketrans({\"-\": \" \", \",\": \"\", \".\": \"\", \"?\": \"\"})\n",
    "\n",
    "def prepare_example(example):\n",
    "    example[\"speech\"] = example[\"audio\"][\"array\"]\n",
    "    example[\"text\"] = example[\"sentence\"].translate(char_translations)\n",
    "    example[\"text\"] = \" \".join(example[\"text\"].split())  # clean up whitespaces\n",
    "    example[\"text\"] = example[\"text\"].lower()\n",
    "    return example\n",
    "\n",
    "dataset = dataset.map(prepare_example)\n",
    "inputs = processor(dataset[\"speech\"], sampling_rate=16000, return_tensors=\"pt\", padding=\"longest\")\n",
    "\n",
    "with torch.no_grad():\n",
    "    predicted_ids = torch.argmax(model(inputs.input_values).logits, dim=-1)\n",
    "predicted_ids[predicted_ids == -100] = processor.tokenizer.pad_token_id\n",
    "predicted_transcripts = processor.tokenizer.batch_decode(predicted_ids)\n",
    "\n",
    "for reference, predicted in zip(dataset[\"text\"], predicted_transcripts):\n",
    "    print(\"reference:\", reference)\n",
    "    print(\"predicted:\", predicted)\n",
    "    print(\"--\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "cbcf36de-0b53-4046-af8d-5f86573c6e42",
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = model.to(\"cuda\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "QABbAEutVcip",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 600,
     "referenced_widgets": [
      "5616be9997024ad59789bd4517799e47",
      "1a76658e85d54d5cab8d52cce3b1ce66",
      "9d4552f6f2b94ef49217a4bceb66a3b2",
      "cb3945d494574bd6b58b171933e00714",
      "23b464b2e2da4cddbe85b926e5116057",
      "ee6a5765a74b4a088749d6b3485f18a2",
      "8c3406cc0914438a9c94b9adef8d83b9",
      "b0b68a7a21354f3096ecde41eedfc22c",
      "88edb334108a42509e704bfc129158ef",
      "b952f16dc875486d876d60c1e6dbe97a",
      "b7c692a87957437dbd3f09c77a0f1adc"
     ]
    },
    "id": "QABbAEutVcip",
    "outputId": "2ef249fd-dc5d-44e5-9222-19db00e8060b"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████████████| 10/10 [00:00<00:00, 20.26ex/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "label: she had your dark suit in greasy wash water all year \n",
      "pred: she had your dark suit in greasy wash water all year\n",
      "-----------\n",
      "label: she had your dark suit in greasy wash water all year \n",
      "pred: she had your dark suit in greasy wash water all year\n",
      "-----------\n",
      "label: there are more obvious nymphomaniacs on any privateeye series \n",
      "pred: ther are more obvious nimpfom many acts on any privite eye series\n",
      "-----------\n",
      "label: only the best players enjoy popularity \n",
      "pred: only the best players enjoy popularity\n",
      "-----------\n",
      "label: december and january are nice months to spend in miami \n",
      "pred: tdisember and january ar nice months toespend in my amy\n",
      "-----------\n",
      "label: keep the thermometer under your tongue \n",
      "pred: keep the thermometer under your tongue\n",
      "-----------\n",
      "label: you're a taxpayer householder landlord \n",
      "pred: you'r a taxpayer householder landlord\n",
      "-----------\n",
      "label: does creole cooking use curry \n",
      "pred: does creole coking use curry\n",
      "-----------\n",
      "label: scholastic aptitude is judged by standardized tests \n",
      "pred: scholastic aptitude is judged by standardized tests\n",
      "-----------\n",
      "label: are you looking for employment \n",
      "pred: are oloking for iemployment\n",
      "-----------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "def map_to_pred(batch):\n",
    "    input_values = processor(batch[\"audio\"][\"array\"], return_tensors=\"pt\", padding=\"longest\", sampling_rate=16_000).input_values\n",
    "    with torch.no_grad():\n",
    "        logits = model(input_values.to(\"cuda\")).logits\n",
    "\n",
    "    transcription = processor.decode(logits[0].argmax(dim=-1).cpu().numpy())\n",
    "    batch[\"transcription\"] = transcription\n",
    "    return batch\n",
    "\n",
    "result = dataset.map(map_to_pred, batched=False, batch_size=1, remove_columns=[\"audio\"])\n",
    "for label, pred in zip(result[\"sentence\"], result[\"transcription\"]):\n",
    "    print('label:', label)\n",
    "    print('pred:', pred)\n",
    "    print('-----------')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "QdgeBB9718X1",
   "metadata": {
    "id": "QdgeBB9718X1"
   },
   "source": [
    "# Load GPT as Language Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "3vccy7N_GODD",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "3vccy7N_GODD",
    "outputId": "2c3bc4ba-af3b-4fc5-96f1-b7ea59f97649"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "processor.tokenizer.word_delimiter_token_id\n",
    "processor.tokenizer.pad_token_id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "KL35I2Bm18Ap",
   "metadata": {
    "id": "KL35I2Bm18Ap"
   },
   "outputs": [],
   "source": [
    "from transformers import GPT2Tokenizer, GPT2LMHeadModel\n",
    "\n",
    "gpt_tokenizer = GPT2Tokenizer.from_pretrained('gpt2')\n",
    "gpt = GPT2LMHeadModel.from_pretrained('gpt2').to(\"cuda\").eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "8e0739e5-d16c-4438-ba53-faf7189761bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import GPTNeoForCausalLM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "uCjt1ZLt2T3L",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "uCjt1ZLt2T3L",
    "outputId": "66e11653-05b2-4287-b44d-c759629f050a"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['. the with a means, want like.\\n']"
      ]
     },
     "execution_count": 52,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "text = \"Replace me by any text you'd like.\"\n",
    "encoded_input = gpt_tokenizer(text, return_tensors='pt').to(\"cuda\")\n",
    "output = gpt(**encoded_input)\n",
    "gpt_tokenizer.batch_decode(output.logits.argmax(dim=-1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "XsWdnF1QSSRD",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "XsWdnF1QSSRD",
    "outputId": "954f8fd8-bac2-405e-df53-55a108acbe0b"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using pad_token, but it is not set yet.\n"
     ]
    }
   ],
   "source": [
    "if not gpt_tokenizer.pad_token:\n",
    "  gpt_tokenizer.pad_token = gpt_tokenizer.eos_token"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "PC2u_czP-06d",
   "metadata": {
    "id": "PC2u_czP-06d"
   },
   "source": [
    "# Transformer Decoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 105,
   "id": "2sw8inFFZ2lY",
   "metadata": {
    "id": "2sw8inFFZ2lY"
   },
   "outputs": [],
   "source": [
    "class TransformerDecoder:\n",
    "    def __init__(self, processor, gpt, gpt_tokenizer, beam_width=50, num_topk=10, alpha=0.1, beta=0., score_interval=50):\n",
    "        self.processor = processor\n",
    "        self.gpt = gpt\n",
    "        self.gpt_tokenizer = gpt_tokenizer\n",
    "        self.B = beam_width\n",
    "        self.VK = num_topk\n",
    "        self.alpha = alpha\n",
    "        self.beta = beta\n",
    "        self.score_interval = score_interval\n",
    "\n",
    "        assert gpt_tokenizer.pad_token is not None\n",
    "\n",
    "        self.processor_vocab = processor.tokenizer.get_vocab()\n",
    "        self.i2c = {v: k for k, v in self.processor_vocab.items()} # index to char\n",
    "\n",
    "        self.pad_token = processor.tokenizer.pad_token_id\n",
    "        self.word_delimiter_token = processor.tokenizer.word_delimiter_token_id\n",
    "        self.i2c[self.word_delimiter_token] = ' '\n",
    "\n",
    "    def decode(self, logits, **kwargs):\n",
    "        return self._beam_search(torch.from_numpy(logits))\n",
    "\n",
    "    def _beam_search(self, audio_logit):\n",
    "        B, VK, alpha, beta = self.B, self.VK, self.alpha, self.beta\n",
    "\n",
    "        beam_scores = np.ones(B) * (-np.log(B))\n",
    "        \n",
    "        beam_infos = [\n",
    "            {\n",
    "                'text': '', # char level vocab doesn't need starting char (convert this to <s> when using kogpt2)\n",
    "                'last_token': self.pad_token,\n",
    "                'last_lm_score': 0,\n",
    "            } for i in range(B)\n",
    "        ]\n",
    "        \n",
    "        logits = audio_logit - audio_logit.logsumexp(dim=-1, keepdim=True)\n",
    "\n",
    "        logits_topk, tokens_topk = torch.topk(logits, k=VK, dim=-1)\n",
    "        tokens_topk = tokens_topk.numpy()\n",
    "        logits_topk = logits_topk.numpy()\n",
    "        chars_topk = [[self.i2c[i] for i in tokens] for tokens in tokens_topk]\n",
    "        N, V = audio_logit.shape\n",
    "        for n in range(N):\n",
    "            \n",
    "            logits, tokens = logits_topk[n], tokens_topk[n]\n",
    "            chars = chars_topk[n]\n",
    "            \n",
    "            # get next scores first\n",
    "            beam_hyp_scores = beam_scores[:, None].repeat(VK, axis=1)\n",
    "            beam_hyp_scores += logits[None, :] # (B, VK) + (1, VK)\n",
    "            beam_hyp_scores = beam_hyp_scores.reshape(-1)\n",
    "            \n",
    "            # get next beams by score\n",
    "            indices_topk = np.argsort(beam_hyp_scores, axis=-1)[::-1][:B+5]\n",
    "            beam_scores = beam_hyp_scores[indices_topk] # B*2, 1\n",
    "            # print(beam_scores)\n",
    "            # get next beam_info by scores\n",
    "            new_beam_infos = [] # B, VK\n",
    "            cache = {}\n",
    "            for i, idx in enumerate(indices_topk):\n",
    "                b, vk = idx // VK, idx % VK\n",
    "                beam_info = beam_infos[b]\n",
    "                \n",
    "                new_last_token = tokens[vk]\n",
    "                new_text = beam_info['text'] + chars[vk] if new_last_token != self.pad_token and beam_info['last_token'] != new_last_token \\\n",
    "                    else beam_info['text']\n",
    "                last_lm_score = beam_info['last_lm_score']\n",
    "                \n",
    "                if (new_text, new_last_token) in cache:\n",
    "                    target_i = cache[(new_text, new_last_token)]\n",
    "                    beam_scores[target_i] = np.log(np.exp(beam_scores[target_i]) + np.exp(beam_scores[i]))\n",
    "                    \n",
    "                    beam_scores[i] = -1e9\n",
    "                else:\n",
    "                    cache[(new_text, new_last_token)] = i # beam index\n",
    "                \n",
    "                beam_hyp = {\n",
    "                    'text': new_text,\n",
    "                    'last_token': new_last_token,\n",
    "                    'last_lm_score': last_lm_score,\n",
    "                }\n",
    "                \n",
    "                new_beam_infos.append(beam_hyp)\n",
    "\n",
    "            # now that we got top B*x scores and indices,\n",
    "            # let's get final beams after merging scores\n",
    "            \n",
    "            indices_topk = np.argsort(beam_scores, axis=0)[::-1][:B]\n",
    "            beam_scores = beam_scores[indices_topk]\n",
    "            beam_infos = [new_beam_infos[i] for i in indices_topk]\n",
    "            \n",
    "            if n % self.score_interval == self.score_interval - 1:\n",
    "                beam_scores = self._add_lm_scores(beam_infos, beam_scores, alpha, beta)\n",
    "\n",
    "        beam_scores = self._add_lm_scores(beam_infos, beam_scores, alpha, beta)\n",
    "        selected_beam_idx = beam_scores.argmax(axis=-1)\n",
    "\n",
    "\n",
    "        return beam_infos[selected_beam_idx]['text'].strip()\n",
    "\n",
    "    def _query_lm(self, texts, beta=0.):\n",
    "        inputs = self.gpt_tokenizer(texts, return_tensors='pt', padding=True)\n",
    "       \n",
    "        B, N = inputs.input_ids.shape\n",
    "        if N < 2:\n",
    "            return torch.zeros(B)\n",
    "        inputs = inputs.to(self.gpt.device)\n",
    "        text_input_ids = inputs.input_ids\n",
    "        B, N = text_input_ids.shape\n",
    "        with torch.no_grad():\n",
    "            labels = text_input_ids[:, :]\n",
    "            text_logits = self.gpt(**inputs).logits\n",
    "            \n",
    "            shift_logits = text_logits[..., :-1, :].contiguous()\n",
    "            shift_labels = labels[..., 1:].contiguous()\n",
    "            \n",
    "            loss_fct = torch.nn.CrossEntropyLoss(reduction='none')\n",
    "            lm_loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))\n",
    "            \n",
    "            res = -lm_loss.reshape(B, -1)\n",
    "            res = (res * inputs.attention_mask[:, 1:])\n",
    "    \n",
    "            scores = res.sum(dim=1) # N, 1\n",
    "        \n",
    "        scores = scores.cpu()\n",
    "        \n",
    "        input_lengths = np.array(list(map(len, texts)), dtype=np.float32)\n",
    "        scores += input_lengths * beta\n",
    "\n",
    "        return scores\n",
    "\n",
    "    def _add_lm_scores(self, beam_infos, beam_scores, alpha, beta):\n",
    "        beam_last_lm_scores = np.array(list(map(lambda x: x['last_lm_score'], beam_infos)), dtype=np.float32)\n",
    "        beam_scores -= beam_last_lm_scores        \n",
    "        \n",
    "        texts = list(map(lambda x: x['text'], beam_infos))\n",
    "        lm_scores = self._query_lm(texts, beta).numpy() * alpha\n",
    "        \n",
    "        for i, x in enumerate(beam_infos):\n",
    "             x['last_lm_score'] = lm_scores[i]\n",
    "        \n",
    "        final_beam_scores = beam_scores + lm_scores\n",
    "        return final_beam_scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 109,
   "id": "Hwwv0-1lv35U",
   "metadata": {
    "id": "Hwwv0-1lv35U"
   },
   "outputs": [],
   "source": [
    "transformer_decoder = TransformerDecoder(\n",
    "    processor=processor,\n",
    "    gpt=gpt,\n",
    "    gpt_tokenizer=gpt_tokenizer,\n",
    "    beam_width=200,\n",
    "    num_topk=10,\n",
    "    alpha=1.0,\n",
    "    beta=2.0,\n",
    "    score_interval=50\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 110,
   "id": "xFE9XHDL1pu4",
   "metadata": {
    "id": "xFE9XHDL1pu4"
   },
   "outputs": [],
   "source": [
    "import gc\n",
    "gc.collect()\n",
    "torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 111,
   "id": "d9nC4vHOuKnK",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 967,
     "referenced_widgets": [
      "9229066bf480480d9792645e4b1d41db",
      "cd65bce74af043bea90d819511b3f007",
      "dae3d68d0f2440829a55f5bc79d4d34d",
      "821bbd190c7f4949aea7793dac0522be",
      "032a7a83c211427e80e786d06a9e1757",
      "6e5e3856d8c1442896fed34397a8b6a7",
      "338a089213be4bc6aba2224e00a1f893",
      "01f1d72232944b199bf4f262953af1a1",
      "8547714ec1e14b7fb2539dc276d4b2ed",
      "067eb0f85bac4c88b9655d01bf841f28",
      "efc5289070b9415390219f0dc1ae63f9"
     ]
    },
    "id": "d9nC4vHOuKnK",
    "outputId": "c2809fa4-002e-41a9-e710-20145223b88a"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████████████| 10/10 [00:07<00:00,  1.25ex/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "label   :\t she had your dark suit in greasy wash water all year \n",
      "pred_raw:\t she had your dark suit in greasy wash water all year\n",
      "pred_gpt:\t she had your dark suit in greasy wash water all year\n",
      "-----------\n",
      "label   :\t she had your dark suit in greasy wash water all year \n",
      "pred_raw:\t she had your dark suit in greasy wash water all year\n",
      "pred_gpt:\t she had your dark suit in greasy wash water all year\n",
      "-----------\n",
      "label   :\t there are more obvious nymphomaniacs on any privateeye series \n",
      "pred_raw:\t ther are more obvious nimpfom many acts on any privite eye series\n",
      "pred_gpt:\t there are more obvious nimfom many acts on any private eye series\n",
      "-----------\n",
      "label   :\t only the best players enjoy popularity \n",
      "pred_raw:\t only the best players enjoy popularity\n",
      "pred_gpt:\t only the best players enjoy popularity\n",
      "-----------\n",
      "label   :\t december and january are nice months to spend in miami \n",
      "pred_raw:\t tdisember and january ar nice months toespend in my amy\n",
      "pred_gpt:\t disember and january are nice months to spend in my amy\n",
      "-----------\n",
      "label   :\t keep the thermometer under your tongue \n",
      "pred_raw:\t keep the thermometer under your tongue\n",
      "pred_gpt:\t keep the thermometer under your tongue\n",
      "-----------\n",
      "label   :\t you're a taxpayer householder landlord \n",
      "pred_raw:\t you'r a taxpayer householder landlord\n",
      "pred_gpt:\t you're a taxpayer householder landlord\n",
      "-----------\n",
      "label   :\t does creole cooking use curry \n",
      "pred_raw:\t does creole coking use curry\n",
      "pred_gpt:\t does creole cooking use curry\n",
      "-----------\n",
      "label   :\t scholastic aptitude is judged by standardized tests \n",
      "pred_raw:\t scholastic aptitude is judged by standardized tests\n",
      "pred_gpt:\t scholastic aptitude is judged by standardized tests\n",
      "-----------\n",
      "label   :\t are you looking for employment \n",
      "pred_raw:\t are oloking for iemployment\n",
      "pred_gpt:\t are looking for employment\n",
      "-----------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "def map_to_pred(batch):\n",
    "    input_values = processor(batch[\"audio\"][\"array\"], return_tensors=\"pt\", padding=\"longest\", sampling_rate=16_000).input_values\n",
    "    with torch.no_grad():\n",
    "        logits = model(input_values.to(\"cuda\")).logits\n",
    "\n",
    "    batch[\"pred_raw\"] = processor.decode(logits[0].argmax(dim=-1).cpu().numpy()).lower()\n",
    "    batch[\"pred_transformer_decoder\"] = transformer_decoder.decode(logits[0].cpu().numpy()).lower()\n",
    "    return batch\n",
    "\n",
    "result = dataset.map(map_to_pred, batched=False, batch_size=1, remove_columns=[\"audio\"])\n",
    "for label, p, p1 in zip(result[\"sentence\"], result[\"pred_raw\"], result[\"pred_transformer_decoder\"]):\n",
    "    print('label   :\\t', label)\n",
    "    print('pred_raw:\\t', p)\n",
    "    print('pred_gpt:\\t', p1)\n",
    "    print('-----------')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1edcb16c-8e7e-4493-b0d9-ca027e20491a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "machine_shape": "hm",
   "provenance": []
  },
  "environment": {
   "kernel": "python3",
   "name": "pytorch-gpu.1-12.m95",
   "type": "gcloud",
   "uri": "gcr.io/deeplearning-platform-release/pytorch-gpu.1-12:m95"
  },
  "gpuClass": "standard",
  "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.8"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "01f1d72232944b199bf4f262953af1a1": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "032a7a83c211427e80e786d06a9e1757": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "067eb0f85bac4c88b9655d01bf841f28": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "1a246b5e180648008a442bff746d1c4e": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "1a76658e85d54d5cab8d52cce3b1ce66": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_ee6a5765a74b4a088749d6b3485f18a2",
      "placeholder": "​",
      "style": "IPY_MODEL_8c3406cc0914438a9c94b9adef8d83b9",
      "value": "100%"
     }
    },
    "23b464b2e2da4cddbe85b926e5116057": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "2ce3984e38ef46afb8de2023571e884b": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "ProgressStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "ProgressStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "bar_color": null,
      "description_width": ""
     }
    },
    "338a089213be4bc6aba2224e00a1f893": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "3add3c0c7900416d99846b7b534af892": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "4b2c4ee01c1a4623b75906caa83fa693": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "FloatProgressModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "FloatProgressModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "ProgressView",
      "bar_style": "success",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_dd8dbf4c874641f3b150380937693db1",
      "max": 2620,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_2ce3984e38ef46afb8de2023571e884b",
      "value": 2620
     }
    },
    "549a86ec991c46cca6129faa192a37e5": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HBoxModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HBoxModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HBoxView",
      "box_style": "",
      "children": [
       "IPY_MODEL_7894a42341864caf85172073258dc745",
       "IPY_MODEL_4b2c4ee01c1a4623b75906caa83fa693",
       "IPY_MODEL_c627779a6d284b1ebcb8480c51c197ff"
      ],
      "layout": "IPY_MODEL_1a246b5e180648008a442bff746d1c4e"
     }
    },
    "5616be9997024ad59789bd4517799e47": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HBoxModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HBoxModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HBoxView",
      "box_style": "",
      "children": [
       "IPY_MODEL_1a76658e85d54d5cab8d52cce3b1ce66",
       "IPY_MODEL_9d4552f6f2b94ef49217a4bceb66a3b2",
       "IPY_MODEL_cb3945d494574bd6b58b171933e00714"
      ],
      "layout": "IPY_MODEL_23b464b2e2da4cddbe85b926e5116057"
     }
    },
    "6e5e3856d8c1442896fed34397a8b6a7": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "7894a42341864caf85172073258dc745": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_3add3c0c7900416d99846b7b534af892",
      "placeholder": "​",
      "style": "IPY_MODEL_e1976bd9b57a489781d2660469f06593",
      "value": "100%"
     }
    },
    "821bbd190c7f4949aea7793dac0522be": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_067eb0f85bac4c88b9655d01bf841f28",
      "placeholder": "​",
      "style": "IPY_MODEL_efc5289070b9415390219f0dc1ae63f9",
      "value": " 10/10 [00:04&lt;00:00,  2.80ex/s]"
     }
    },
    "8547714ec1e14b7fb2539dc276d4b2ed": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "ProgressStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "ProgressStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "bar_color": null,
      "description_width": ""
     }
    },
    "88edb334108a42509e704bfc129158ef": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "ProgressStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "ProgressStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "bar_color": null,
      "description_width": ""
     }
    },
    "899164ea0836481c8ff71cf788ff0daa": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "8c3406cc0914438a9c94b9adef8d83b9": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "9229066bf480480d9792645e4b1d41db": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HBoxModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HBoxModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HBoxView",
      "box_style": "",
      "children": [
       "IPY_MODEL_cd65bce74af043bea90d819511b3f007",
       "IPY_MODEL_dae3d68d0f2440829a55f5bc79d4d34d",
       "IPY_MODEL_821bbd190c7f4949aea7793dac0522be"
      ],
      "layout": "IPY_MODEL_032a7a83c211427e80e786d06a9e1757"
     }
    },
    "9d4552f6f2b94ef49217a4bceb66a3b2": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "FloatProgressModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "FloatProgressModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "ProgressView",
      "bar_style": "success",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_b0b68a7a21354f3096ecde41eedfc22c",
      "max": 10,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_88edb334108a42509e704bfc129158ef",
      "value": 10
     }
    },
    "b0b68a7a21354f3096ecde41eedfc22c": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "b7c692a87957437dbd3f09c77a0f1adc": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "b952f16dc875486d876d60c1e6dbe97a": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "c627779a6d284b1ebcb8480c51c197ff": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_ffb93bf16fdf4999b50f345a3d5864b5",
      "placeholder": "​",
      "style": "IPY_MODEL_899164ea0836481c8ff71cf788ff0daa",
      "value": " 2620/2620 [14:16&lt;00:00,  1.50ex/s]"
     }
    },
    "cb3945d494574bd6b58b171933e00714": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_b952f16dc875486d876d60c1e6dbe97a",
      "placeholder": "​",
      "style": "IPY_MODEL_b7c692a87957437dbd3f09c77a0f1adc",
      "value": " 10/10 [00:01&lt;00:00,  7.71ex/s]"
     }
    },
    "cd65bce74af043bea90d819511b3f007": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_6e5e3856d8c1442896fed34397a8b6a7",
      "placeholder": "​",
      "style": "IPY_MODEL_338a089213be4bc6aba2224e00a1f893",
      "value": "100%"
     }
    },
    "dae3d68d0f2440829a55f5bc79d4d34d": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "FloatProgressModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "FloatProgressModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "ProgressView",
      "bar_style": "success",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_01f1d72232944b199bf4f262953af1a1",
      "max": 10,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_8547714ec1e14b7fb2539dc276d4b2ed",
      "value": 10
     }
    },
    "dd8dbf4c874641f3b150380937693db1": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "e1976bd9b57a489781d2660469f06593": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "ee6a5765a74b4a088749d6b3485f18a2": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "efc5289070b9415390219f0dc1ae63f9": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "ffb93bf16fdf4999b50f345a3d5864b5": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    }
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}