deepspeech_example_basic.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tutorial for deepSpeech.pytorch\n",
    "\n",
    "this jupyter notebooks is explaining how deepspeech.pytorch work.\n",
    "I put every componento on this notebooks, so you can check each step from data preprocessing to model to training.\n",
    "\n",
    "## What is deepspeech.pytorch ?\n",
    "deepspeech.pytorch is an awesome implentation of Baidu's DeepSpeech 2 paper.\n",
    "\n",
    "https://github.com/SeanNaren/deepspeech.pytorch\n",
    "\n",
    "## Requirements\n",
    "### Warp-CTC\n",
    "you need Warp-CTC for loss function.\n",
    "please follow the instruction at https://github.com/SeanNaren/deepspeech.pytorch to install Warp-CTC.  \n",
    "```\n",
    "git clone https://github.com/SeanNaren/warp-ctc.git\n",
    "cd warp-ctc; mkdir build; cd build; cmake ..; make\n",
    "export CUDA_HOME=\"/usr/local/cuda\"\n",
    "cd ../pytorch_binding && python setup.py install\n",
    "```\n",
    "\n",
    "### other dependancy:  \n",
    "scipy, numpy, soundfile,  python-levenshtein, torch, torchelastic, wget, librosa, tqdm, matplotlib, sox, sklearn\n",
    "\n",
    "I remove apex part so that you can test this code even with CPU."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Google colab environment\n",
    "If you are using Google Colab environment, you need to restart runtime after installing war-ctc.  \n",
    "\n",
    "Install warp-ctc, then **restart tuntime**.\n",
    "```\n",
    "!git clone https://github.com/SeanNaren/warp-ctc.git\n",
    "!cd warp-ctc; mkdir build; cd build; cmake ..;make\n",
    "!cd warp-ctc/pytorch_binding;python setup.py install\n",
    "```\n",
    "\n",
    "Also you need to install other dependancies.  \n",
    "```\n",
    "! pip install soundfile\n",
    "! sudo apt install sox\n",
    "! pip install sox\n",
    "! pip install python-levenshtein\n",
    "! pip install wget\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# this is just a jupyter things\n",
    "%reload_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import basic library\n",
    "import json\n",
    "import os\n",
    "import random\n",
    "import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import deep learning related\n",
    "import torch\n",
    "from warpctc_pytorch import CTCLoss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Let's define modules "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define SpectrogramDataset, DSRandomSampler, AudioDataLoader\n",
    "This is equivalent of ```from data.data_loader import SpectrogramDataset, DSRandomSampler, AudioDataLoader```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# pip install numba==0.48"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import os\n",
    "from tempfile import NamedTemporaryFile\n",
    "import librosa\n",
    "import numpy as np\n",
    "import scipy.signal\n",
    "import soundfile as sf\n",
    "import sox\n",
    "import torch\n",
    "from torch.utils.data import Dataset, Sampler, DataLoader\n",
    "\n",
    "windows = {\n",
    "    'hamming': scipy.signal.hamming,\n",
    "    'hann': scipy.signal.hann,\n",
    "    'blackman': scipy.signal.blackman,\n",
    "    'bartlett': scipy.signal.bartlett\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def augment_audio_with_sox(path, sample_rate, tempo, gain):\n",
    "    \"\"\"\n",
    "    Changes tempo and gain of the recording with sox and loads it.\n",
    "    \"\"\"\n",
    "    with NamedTemporaryFile(suffix=\".wav\") as augmented_file:\n",
    "        augmented_filename = augmented_file.name\n",
    "        sox_augment_params = [\"tempo\", \"{:.3f}\".format(tempo), \"gain\", \"{:.3f}\".format(gain)]\n",
    "        sox_params = \"sox \\\"{}\\\" -r {} -c 1 -b 16 -e si {} {} >/dev/null 2>&1\".format(path, sample_rate,\n",
    "                                                                                      augmented_filename,\n",
    "                                                                                      \" \".join(sox_augment_params))\n",
    "        os.system(sox_params)\n",
    "        y = load_audio(augmented_filename)\n",
    "        return y\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_randomly_augmented_audio(path, sample_rate=16000, tempo_range=(0.85, 1.15),\n",
    "                                  gain_range=(-6, 8)):\n",
    "    \"\"\"\n",
    "    Picks tempo and gain uniformly, applies it to the utterance by using sox utility.\n",
    "    Returns the augmented utterance.\n",
    "    \"\"\"\n",
    "    low_tempo, high_tempo = tempo_range\n",
    "    tempo_value = np.random.uniform(low=low_tempo, high=high_tempo)\n",
    "    low_gain, high_gain = gain_range\n",
    "    gain_value = np.random.uniform(low=low_gain, high=high_gain)\n",
    "    audio = augment_audio_with_sox(path=path, sample_rate=sample_rate,\n",
    "                                   tempo=tempo_value, gain=gain_value)\n",
    "    return audio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "class AudioParser(object):\n",
    "    def parse_transcript(self, transcript_path):\n",
    "        \"\"\"\n",
    "        :param transcript_path: Path where transcript is stored from the manifest file\n",
    "        :return: Transcript in training/testing format\n",
    "        \"\"\"\n",
    "        raise NotImplementedError\n",
    "\n",
    "    def parse_audio(self, audio_path):\n",
    "        \"\"\"\n",
    "        :param audio_path: Path where audio is stored from the manifest file\n",
    "        :return: Audio in training/testing format\n",
    "        \"\"\"\n",
    "        raise NotImplementedError"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "class SpectrogramParser(AudioParser):\n",
    "    def __init__(self, audio_conf, normalize=False, speed_volume_perturb=False, spec_augment=False):\n",
    "        \"\"\"\n",
    "        Parses audio file into spectrogram with optional normalization and various augmentations\n",
    "        :param audio_conf: Dictionary containing the sample rate, window and the window length/stride in seconds\n",
    "        :param normalize(default False):  Apply standard mean and deviation normalization to audio tensor\n",
    "        :param speed_volume_perturb(default False): Apply random tempo and gain perturbations\n",
    "        :param spec_augment(default False): Apply simple spectral augmentation to mel spectograms\n",
    "        \"\"\"\n",
    "        super(SpectrogramParser, self).__init__()\n",
    "        self.window_stride = audio_conf['window_stride']\n",
    "        self.window_size = audio_conf['window_size']\n",
    "        self.sample_rate = audio_conf['sample_rate']\n",
    "        # self.window = windows.get(audio_conf['window'], windows['hamming'])\n",
    "        self.window = windows.get(audio_conf['window'], windows['hamming'])\n",
    "        self.normalize = normalize\n",
    "        self.speed_volume_perturb = speed_volume_perturb\n",
    "        self.spec_augment = spec_augment\n",
    "        self.noiseInjector = NoiseInjection(audio_conf['noise_dir'], self.sample_rate,\n",
    "                                            audio_conf['noise_levels']) if audio_conf.get(\n",
    "            'noise_dir') is not None else None\n",
    "        self.noise_prob = audio_conf.get('noise_prob')\n",
    "\n",
    "    def parse_audio(self, audio_path):\n",
    "        if self.speed_volume_perturb:\n",
    "            y = load_randomly_augmented_audio(audio_path, self.sample_rate)\n",
    "        else:\n",
    "            y = load_audio(audio_path)\n",
    "        if self.noiseInjector:\n",
    "            add_noise = np.random.binomial(1, self.noise_prob)\n",
    "            if add_noise:\n",
    "                y = self.noiseInjector.inject_noise(y)\n",
    "        n_fft = int(self.sample_rate * self.window_size)\n",
    "        win_length = n_fft\n",
    "        hop_length = int(self.sample_rate * self.window_stride)\n",
    "        # STFT\n",
    "        D = librosa.stft(y, n_fft=n_fft, hop_length=hop_length,\n",
    "                         win_length=win_length, window=self.window)\n",
    "        spect, phase = librosa.magphase(D)\n",
    "        # S = log(S+1)\n",
    "        spect = np.log1p(spect)\n",
    "        spect = torch.FloatTensor(spect)\n",
    "        if self.normalize:\n",
    "            mean = spect.mean()\n",
    "            std = spect.std()\n",
    "            spect.add_(-mean)\n",
    "            spect.div_(std)\n",
    "\n",
    "        if self.spec_augment:\n",
    "            spect = spec_augment(spect)\n",
    "\n",
    "        return spect\n",
    "\n",
    "    def parse_transcript(self, transcript_path):\n",
    "        raise NotImplementedError"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SpectrogramDataset(Dataset, SpectrogramParser):\n",
    "    def __init__(self, audio_conf, manifest_filepath, labels, normalize=False, speed_volume_perturb=False,\n",
    "                 spec_augment=False):\n",
    "        \"\"\"\n",
    "        Dataset that loads tensors via a csv containing file paths to audio files and transcripts separated by\n",
    "        a comma. Each new line is a different sample. Example below:\n",
    "\n",
    "        /path/to/audio.wav,/path/to/audio.txt\n",
    "        ...\n",
    "\n",
    "        :param audio_conf: Dictionary containing the sample rate, window and the window length/stride in seconds\n",
    "        :param manifest_filepath: Path to manifest csv as describe above\n",
    "        :param labels: String containing all the possible characters to map to\n",
    "        :param normalize: Apply standard mean and deviation normalization to audio tensor\n",
    "        :param speed_volume_perturb(default False): Apply random tempo and gain perturbations\n",
    "        :param spec_augment(default False): Apply simple spectral augmentation to mel spectograms\n",
    "        \"\"\"\n",
    "        with open(manifest_filepath) as f:\n",
    "            ids = f.readlines()\n",
    "        ids = [x.strip().split(',') for x in ids]\n",
    "        self.ids = ids\n",
    "        self.size = len(ids)\n",
    "        self.labels_map = dict([(labels[i], i) for i in range(len(labels))])\n",
    "        super(SpectrogramDataset, self).__init__(audio_conf, normalize, speed_volume_perturb, spec_augment)\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        sample = self.ids[index]\n",
    "        audio_path, transcript_path = sample[0], sample[1]\n",
    "        spect = self.parse_audio(audio_path)\n",
    "        transcript = self.parse_transcript(transcript_path)\n",
    "        return spect, transcript\n",
    "\n",
    "    def parse_transcript(self, transcript_path):\n",
    "        with open(transcript_path, 'r', encoding='utf8') as transcript_file:\n",
    "            transcript = transcript_file.read().replace('\\n', '')\n",
    "        transcript = list(filter(None, [self.labels_map.get(x) for x in list(transcript)]))\n",
    "        return transcript\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.size\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DSRandomSampler(Sampler):\n",
    "    \"\"\"\n",
    "    Implementation of a Random Sampler for sampling the dataset.\n",
    "    Added to ensure we reset the start index when an epoch is finished.\n",
    "    This is essential since we support saving/loading state during an epoch.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, dataset, batch_size=1, start_index=0):\n",
    "        super().__init__(data_source=dataset)\n",
    "\n",
    "        self.dataset = dataset\n",
    "        self.start_index = start_index\n",
    "        self.batch_size = batch_size\n",
    "        ids = list(range(len(self.dataset)))\n",
    "        self.bins = [ids[i:i + self.batch_size] for i in range(0, len(ids), self.batch_size)]\n",
    "\n",
    "    def __iter__(self):\n",
    "        # deterministically shuffle based on epoch\n",
    "        g = torch.Generator()\n",
    "        g.manual_seed(self.epoch)\n",
    "        indices = (\n",
    "            torch.randperm(len(self.bins) - self.start_index, generator=g)\n",
    "                .add(self.start_index)\n",
    "                .tolist()\n",
    "        )\n",
    "        for x in indices:\n",
    "            batch_ids = self.bins[x]\n",
    "            np.random.shuffle(batch_ids)\n",
    "            yield batch_ids\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.bins) - self.start_index\n",
    "\n",
    "    def set_epoch(self, epoch):\n",
    "        self.epoch = epoch\n",
    "\n",
    "    def reset_training_step(self, training_step):\n",
    "        self.start_index = training_step"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _collate_fn(batch):\n",
    "\n",
    "    def func(p):\n",
    "        return p[0].size(1)\n",
    "\n",
    "    batch = sorted(batch, key=lambda sample: sample[0].size(1), reverse=True)\n",
    "    longest_sample = max(batch, key=func)[0]\n",
    "    freq_size = longest_sample.size(0)\n",
    "    minibatch_size = len(batch)\n",
    "    max_seqlength = longest_sample.size(1)\n",
    "    inputs = torch.zeros(minibatch_size, 1, freq_size, max_seqlength)\n",
    "    input_percentages = torch.FloatTensor(minibatch_size)\n",
    "    target_sizes = torch.IntTensor(minibatch_size)\n",
    "    targets = []\n",
    "    # audio_paths = []\n",
    "    for x in range(minibatch_size):\n",
    "        sample = batch[x]\n",
    "        tensor = sample[0]\n",
    "        target = sample[1]\n",
    "        seq_length = tensor.size(1)\n",
    "        inputs[x][0].narrow(1, 0, seq_length).copy_(tensor)\n",
    "        input_percentages[x] = seq_length / float(max_seqlength)\n",
    "        target_sizes[x] = len(target)\n",
    "        targets.extend(target)\n",
    "        # audio_path = sample[2]\n",
    "        # audio_paths.append(audio_path)\n",
    "    targets = torch.IntTensor(targets)\n",
    "    return inputs, targets, input_percentages, target_sizes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "class AudioDataLoader(DataLoader):\n",
    "    def __init__(self, *args, **kwargs):\n",
    "        \"\"\"\n",
    "        Creates a data loader for AudioDatasets.\n",
    "        \"\"\"\n",
    "        super(AudioDataLoader, self).__init__(*args, **kwargs)\n",
    "        self.collate_fn = _collate_fn"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### define GreedyDecoder, decoder\n",
    "this is equivalent of ```from decoder import GreedyDecoder```\n",
    "\n",
    "You can use *BeamDecoder*. But you need to install ctcdecode at that case."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import Levenshtein as Lev\n",
    "import torch\n",
    "from six.moves import xrange"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Decoder(object):\n",
    "    \"\"\"\n",
    "    Basic decoder class from which all other decoders inherit. Implements several\n",
    "    helper functions. Subclasses should implement the decode() method.\n",
    "\n",
    "    Arguments:\n",
    "        labels (list): mapping from integers to characters.\n",
    "        blank_index (int, optional): index for the blank '_' character. Defaults to 0.\n",
    "        space_index (int, optional): index for the space ' ' character. Defaults to 28.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, labels, blank_index=0):\n",
    "        self.labels = labels\n",
    "        self.int_to_char = dict([(i, c) for (i, c) in enumerate(labels)])\n",
    "        self.blank_index = blank_index\n",
    "        space_index = len(labels)  # To prevent errors in decode, we add an out of bounds index for the space\n",
    "        if ' ' in labels:\n",
    "            space_index = labels.index(' ')\n",
    "        self.space_index = space_index\n",
    "\n",
    "    def wer(self, s1, s2):\n",
    "        \"\"\"\n",
    "        Computes the Word Error Rate, defined as the edit distance between the\n",
    "        two provided sentences after tokenizing to words.\n",
    "        Arguments:\n",
    "            s1 (string): space-separated sentence\n",
    "            s2 (string): space-separated sentence\n",
    "        \"\"\"\n",
    "\n",
    "        # build mapping of words to integers\n",
    "        b = set(s1.split() + s2.split())\n",
    "        word2char = dict(zip(b, range(len(b))))\n",
    "\n",
    "        # map the words to a char array (Levenshtein packages only accepts\n",
    "        # strings)\n",
    "        w1 = [chr(word2char[w]) for w in s1.split()]\n",
    "        w2 = [chr(word2char[w]) for w in s2.split()]\n",
    "\n",
    "        return Lev.distance(''.join(w1), ''.join(w2))\n",
    "\n",
    "    def cer(self, s1, s2):\n",
    "        \"\"\"\n",
    "        Computes the Character Error Rate, defined as the edit distance.\n",
    "\n",
    "        Arguments:\n",
    "            s1 (string): space-separated sentence\n",
    "            s2 (string): space-separated sentence\n",
    "        \"\"\"\n",
    "        s1, s2, = s1.replace(' ', ''), s2.replace(' ', '')\n",
    "        return Lev.distance(s1, s2)\n",
    "\n",
    "    def decode(self, probs, sizes=None):\n",
    "        \"\"\"\n",
    "        Given a matrix of character probabilities, returns the decoder's\n",
    "        best guess of the transcription\n",
    "\n",
    "        Arguments:\n",
    "            probs: Tensor of character probabilities, where probs[c,t]\n",
    "                            is the probability of character c at time t\n",
    "            sizes(optional): Size of each sequence in the mini-batch\n",
    "        Returns:\n",
    "            string: sequence of the model's best guess for the transcription\n",
    "        \"\"\"\n",
    "        raise NotImplementedError\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GreedyDecoder(Decoder):\n",
    "    def __init__(self, labels, blank_index=0):\n",
    "        super(GreedyDecoder, self).__init__(labels, blank_index)\n",
    "\n",
    "    def convert_to_strings(self, sequences, sizes=None, remove_repetitions=False, return_offsets=False):\n",
    "        \"\"\"Given a list of numeric sequences, returns the corresponding strings\"\"\"\n",
    "        strings = []\n",
    "        offsets = [] if return_offsets else None\n",
    "        for x in xrange(len(sequences)):\n",
    "            seq_len = sizes[x] if sizes is not None else len(sequences[x])\n",
    "            string, string_offsets = self.process_string(sequences[x], seq_len, remove_repetitions)\n",
    "            strings.append([string])  # We only return one path\n",
    "            if return_offsets:\n",
    "                offsets.append([string_offsets])\n",
    "        if return_offsets:\n",
    "            return strings, offsets\n",
    "        else:\n",
    "            return strings\n",
    "\n",
    "    def process_string(self, sequence, size, remove_repetitions=False):\n",
    "        string = ''\n",
    "        offsets = []\n",
    "        for i in range(size):\n",
    "            char = self.int_to_char[sequence[i].item()]\n",
    "            if char != self.int_to_char[self.blank_index]:\n",
    "                # if this char is a repetition and remove_repetitions=true, then skip\n",
    "                if remove_repetitions and i != 0 and char == self.int_to_char[sequence[i - 1].item()]:\n",
    "                    pass\n",
    "                elif char == self.labels[self.space_index]:\n",
    "                    string += ' '\n",
    "                    offsets.append(i)\n",
    "                else:\n",
    "                    string = string + char\n",
    "                    offsets.append(i)\n",
    "        return string, torch.tensor(offsets, dtype=torch.int)\n",
    "\n",
    "    def decode(self, probs, sizes=None):\n",
    "        \"\"\"\n",
    "        Returns the argmax decoding given the probability matrix. Removes\n",
    "        repeated elements in the sequence, as well as blanks.\n",
    "\n",
    "        Arguments:\n",
    "            probs: Tensor of character probabilities from the network. Expected shape of batch x seq_length x output_dim\n",
    "            sizes(optional): Size of each sequence in the mini-batch\n",
    "        Returns:\n",
    "            strings: sequences of the model's best guess for the transcription on inputs\n",
    "            offsets: time step per character predicted\n",
    "        \"\"\"\n",
    "        _, max_probs = torch.max(probs, 2)\n",
    "        strings, offsets = self.convert_to_strings(max_probs.view(max_probs.size(0), max_probs.size(1)), sizes,\n",
    "                                                   remove_repetitions=True, return_offsets=True)\n",
    "        return strings, offsets\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define DeepSpeech, and supported_rnns\n",
    "This is equivalent of ```from model import DeepSpeech, supported_rnns```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "from collections import OrderedDict\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class InferenceBatchSoftmax(nn.Module):\n",
    "    def forward(self, input_):\n",
    "        if not self.training:\n",
    "            return F.softmax(input_, dim=-1)\n",
    "        else:\n",
    "            return input_\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class BatchRNN(nn.Module):\n",
    "    def __init__(self, input_size, hidden_size, rnn_type=nn.LSTM, bidirectional=False, batch_norm=True):\n",
    "        super(BatchRNN, self).__init__()\n",
    "        self.input_size = input_size\n",
    "        self.hidden_size = hidden_size\n",
    "        self.bidirectional = bidirectional\n",
    "        self.batch_norm = SequenceWise(nn.BatchNorm1d(input_size)) if batch_norm else None\n",
    "        self.rnn = rnn_type(input_size=input_size, hidden_size=hidden_size,\n",
    "                            bidirectional=bidirectional, bias=True)\n",
    "        self.num_directions = 2 if bidirectional else 1\n",
    "\n",
    "    def flatten_parameters(self):\n",
    "        self.rnn.flatten_parameters()\n",
    "\n",
    "    def forward(self, x, output_lengths):\n",
    "        if self.batch_norm is not None:\n",
    "            x = self.batch_norm(x)\n",
    "        x = nn.utils.rnn.pack_padded_sequence(x, output_lengths)\n",
    "        x, h = self.rnn(x)\n",
    "        x, _ = nn.utils.rnn.pad_packed_sequence(x)\n",
    "        if self.bidirectional:\n",
    "            x = x.view(x.size(0), x.size(1), 2, -1).sum(2).view(x.size(0), x.size(1), -1)  # (TxNxH*2) -> (TxNxH) by sum\n",
    "        return x\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "class SequenceWise(nn.Module):\n",
    "    def __init__(self, module):\n",
    "        \"\"\"\n",
    "        Collapses input of dim T*N*H to (T*N)*H, and applies to a module.\n",
    "        Allows handling of variable sequence lengths and minibatch sizes.\n",
    "        :param module: Module to apply input to.\n",
    "        \"\"\"\n",
    "        super(SequenceWise, self).__init__()\n",
    "        self.module = module\n",
    "\n",
    "    def forward(self, x):\n",
    "        t, n = x.size(0), x.size(1)\n",
    "        x = x.view(t * n, -1)\n",
    "        x = self.module(x)\n",
    "        x = x.view(t, n, -1)\n",
    "        return x\n",
    "\n",
    "    def __repr__(self):\n",
    "        tmpstr = self.__class__.__name__ + ' (\\n'\n",
    "        tmpstr += self.module.__repr__()\n",
    "        tmpstr += ')'\n",
    "        return tmpstr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MaskConv(nn.Module):\n",
    "    def __init__(self, seq_module):\n",
    "        \"\"\"\n",
    "        Adds padding to the output of the module based on the given lengths. This is to ensure that the\n",
    "        results of the model do not change when batch sizes change during inference.\n",
    "        Input needs to be in the shape of (BxCxDxT)\n",
    "        :param seq_module: The sequential module containing the conv stack.\n",
    "        \"\"\"\n",
    "        super(MaskConv, self).__init__()\n",
    "        self.seq_module = seq_module\n",
    "\n",
    "    def forward(self, x, lengths):\n",
    "        \"\"\"\n",
    "        :param x: The input of size BxCxDxT\n",
    "        :param lengths: The actual length of each sequence in the batch\n",
    "        :return: Masked output from the module\n",
    "        \"\"\"\n",
    "        for module in self.seq_module:\n",
    "            x = module(x)\n",
    "            mask = torch.BoolTensor(x.size()).fill_(0)\n",
    "            if x.is_cuda:\n",
    "                mask = mask.cuda()\n",
    "            for i, length in enumerate(lengths):\n",
    "                length = length.item()\n",
    "                if (mask[i].size(2) - length) > 0:\n",
    "                    mask[i].narrow(2, length, mask[i].size(2) - length).fill_(1)\n",
    "            x = x.masked_fill(mask, 0)\n",
    "        return x, lengths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DeepSpeech(nn.Module):\n",
    "    def __init__(self, rnn_type, labels, rnn_hidden_size, nb_layers, audio_conf,\n",
    "                 bidirectional, context=20):\n",
    "        super(DeepSpeech, self).__init__()\n",
    "\n",
    "        self.hidden_size = rnn_hidden_size\n",
    "        self.hidden_layers = nb_layers\n",
    "        self.rnn_type = rnn_type\n",
    "        self.audio_conf = audio_conf\n",
    "        self.labels = labels\n",
    "        self.bidirectional = bidirectional\n",
    "\n",
    "        sample_rate = self.audio_conf[\"sample_rate\"]\n",
    "        window_size = self.audio_conf[\"window_size\"]\n",
    "        num_classes = len(self.labels)\n",
    "\n",
    "        self.conv = MaskConv(nn.Sequential(\n",
    "            nn.Conv2d(1, 32, kernel_size=(41, 11), stride=(2, 2), padding=(20, 5)),\n",
    "            nn.BatchNorm2d(32),\n",
    "            nn.Hardtanh(0, 20, inplace=True),\n",
    "            nn.Conv2d(32, 32, kernel_size=(21, 11), stride=(2, 1), padding=(10, 5)),\n",
    "            nn.BatchNorm2d(32),\n",
    "            nn.Hardtanh(0, 20, inplace=True)\n",
    "        ))\n",
    "        # Based on above convolutions and spectrogram size using conv formula (W - F + 2P)/ S+1\n",
    "        rnn_input_size = int(math.floor((sample_rate * window_size) / 2) + 1)\n",
    "        rnn_input_size = int(math.floor(rnn_input_size + 2 * 20 - 41) / 2 + 1)\n",
    "        rnn_input_size = int(math.floor(rnn_input_size + 2 * 10 - 21) / 2 + 1)\n",
    "        rnn_input_size *= 32\n",
    "\n",
    "        rnns = []\n",
    "        rnn = BatchRNN(input_size=rnn_input_size, hidden_size=rnn_hidden_size, rnn_type=rnn_type,\n",
    "                       bidirectional=bidirectional, batch_norm=False)\n",
    "        rnns.append(('0', rnn))\n",
    "        for x in range(nb_layers - 1):\n",
    "            rnn = BatchRNN(input_size=rnn_hidden_size, hidden_size=rnn_hidden_size, rnn_type=rnn_type,\n",
    "                           bidirectional=bidirectional)\n",
    "            rnns.append(('%d' % (x + 1), rnn))\n",
    "        self.rnns = nn.Sequential(OrderedDict(rnns))\n",
    "        self.lookahead = nn.Sequential(\n",
    "            # consider adding batch norm?\n",
    "            Lookahead(rnn_hidden_size, context=context),\n",
    "            nn.Hardtanh(0, 20, inplace=True)\n",
    "        ) if not bidirectional else None\n",
    "\n",
    "        fully_connected = nn.Sequential(\n",
    "            nn.BatchNorm1d(rnn_hidden_size),\n",
    "            nn.Linear(rnn_hidden_size, num_classes, bias=False)\n",
    "        )\n",
    "        self.fc = nn.Sequential(\n",
    "            SequenceWise(fully_connected),\n",
    "        )\n",
    "        self.inference_softmax = InferenceBatchSoftmax()\n",
    "\n",
    "    def forward(self, x, lengths):\n",
    "        lengths = lengths.cpu().int()\n",
    "        output_lengths = self.get_seq_lens(lengths)\n",
    "        x, _ = self.conv(x, output_lengths)\n",
    "\n",
    "        sizes = x.size()\n",
    "        x = x.view(sizes[0], sizes[1] * sizes[2], sizes[3])  # Collapse feature dimension\n",
    "        x = x.transpose(1, 2).transpose(0, 1).contiguous()  # TxNxH\n",
    "\n",
    "        for rnn in self.rnns:\n",
    "            x = rnn(x, output_lengths)\n",
    "\n",
    "        if not self.bidirectional:  # no need for lookahead layer in bidirectional\n",
    "            x = self.lookahead(x)\n",
    "\n",
    "        x = self.fc(x)\n",
    "        x = x.transpose(0, 1)\n",
    "        # identity in training mode, softmax in eval mode\n",
    "        x = self.inference_softmax(x)\n",
    "        return x, output_lengths\n",
    "\n",
    "    def get_seq_lens(self, input_length):\n",
    "        \"\"\"\n",
    "        Given a 1D Tensor or Variable containing integer sequence lengths, return a 1D tensor or variable\n",
    "        containing the size sequences that will be output by the network.\n",
    "        :param input_length: 1D Tensor\n",
    "        :return: 1D Tensor scaled by model\n",
    "        \"\"\"\n",
    "        seq_len = input_length\n",
    "        for m in self.conv.modules():\n",
    "            if type(m) == nn.modules.conv.Conv2d:\n",
    "                seq_len = ((seq_len + 2 * m.padding[1] - m.dilation[1] * (m.kernel_size[1] - 1) - 1) / m.stride[1] + 1)\n",
    "        return seq_len.int()\n",
    "\n",
    "    @classmethod\n",
    "    def load_model(cls, path):\n",
    "        package = torch.load(path, map_location=lambda storage, loc: storage)\n",
    "        model = DeepSpeech.load_model_package(package)\n",
    "        return model\n",
    "\n",
    "    @classmethod\n",
    "    def load_model_package(cls, package):\n",
    "        model = cls(rnn_hidden_size=package['hidden_size'],\n",
    "                    nb_layers=package['hidden_layers'],\n",
    "                    labels=package['labels'],\n",
    "                    audio_conf=package['audio_conf'],\n",
    "                    rnn_type=supported_rnns[package['rnn_type']],\n",
    "                    bidirectional=package.get('bidirectional', True))\n",
    "        model.load_state_dict(package['state_dict'])\n",
    "        return model\n",
    "\n",
    "    def serialize_state(self):\n",
    "        return {\n",
    "            'hidden_size': self.hidden_size,\n",
    "            'hidden_layers': self.hidden_layers,\n",
    "            'rnn_type': supported_rnns_inv.get(self.rnn_type, self.rnn_type.__name__.lower()),\n",
    "            'audio_conf': self.audio_conf,\n",
    "            'labels': self.labels,\n",
    "            'state_dict': self.state_dict(),\n",
    "            'bidirectional': self.bidirectional,\n",
    "        }\n",
    "\n",
    "    @staticmethod\n",
    "    def get_param_size(model):\n",
    "        params = 0\n",
    "        for p in model.parameters():\n",
    "            tmp = 1\n",
    "            for x in p.size():\n",
    "                tmp *= x\n",
    "            params += tmp\n",
    "        return params\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "supported_rnns = {\n",
    "    'lstm': nn.LSTM,\n",
    "    'rnn': nn.RNN,\n",
    "    'gru': nn.GRU\n",
    "}\n",
    "supported_rnns_inv = dict((v, k) for k, v in supported_rnns.items())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define TrainingState\n",
    "this is equivalent of ```from state import TrainingState```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def remove_parallel_wrapper(model):\n",
    "    \"\"\"\n",
    "    Return the model or extract the model out of the parallel wrapper\n",
    "    :param model: The training model\n",
    "    :return: The model without parallel wrapper\n",
    "    \"\"\"\n",
    "    # Take care of distributed/data-parallel wrapper\n",
    "    model_no_wrapper = model.module if hasattr(model, \"module\") else model\n",
    "    return model_no_wrapper"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ResultState:\n",
    "    def __init__(self,\n",
    "                 loss_results,\n",
    "                 wer_results,\n",
    "                 cer_results):\n",
    "        self.loss_results = loss_results\n",
    "        self.wer_results = wer_results\n",
    "        self.cer_results = cer_results\n",
    "\n",
    "    def add_results(self,\n",
    "                    epoch,\n",
    "                    loss_result,\n",
    "                    wer_result,\n",
    "                    cer_result):\n",
    "        self.loss_results[epoch] = loss_result\n",
    "        self.wer_results[epoch] = wer_result\n",
    "        self.cer_results[epoch] = cer_result\n",
    "\n",
    "    def serialize_state(self):\n",
    "        return {\n",
    "            'loss_results': self.loss_results,\n",
    "            'wer_results': self.wer_results,\n",
    "            'cer_results': self.cer_results\n",
    "        }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TrainingState:\n",
    "    def __init__(self,\n",
    "                 model,\n",
    "                 result_state=None,\n",
    "                 optim_state=None,\n",
    "                 amp_state=None,\n",
    "                 best_wer=None,\n",
    "                 avg_loss=0,\n",
    "                 epoch=0,\n",
    "                 training_step=0):\n",
    "        \"\"\"\n",
    "        Wraps around training model and states for saving/loading convenience.\n",
    "        For backwards compatibility there are more states being saved than necessary.\n",
    "        \"\"\"\n",
    "        self.model = model\n",
    "        self.result_state = result_state\n",
    "        self.optim_state = optim_state\n",
    "        self.amp_state = amp_state\n",
    "        self.best_wer = best_wer\n",
    "        self.avg_loss = avg_loss\n",
    "        self.epoch = epoch\n",
    "        self.training_step = training_step\n",
    "\n",
    "    def track_optim_state(self, optimizer):\n",
    "        self.optim_state = optimizer.state_dict()\n",
    "\n",
    "    def track_amp_state(self, amp):\n",
    "        self.amp_state = amp.state_dict()\n",
    "\n",
    "    def init_results_tracking(self, epochs):\n",
    "        self.result_state = ResultState(loss_results=torch.IntTensor(epochs),\n",
    "                                        wer_results=torch.IntTensor(epochs),\n",
    "                                        cer_results=torch.IntTensor(epochs))\n",
    "\n",
    "    def add_results(self,\n",
    "                    epoch,\n",
    "                    loss_result,\n",
    "                    wer_result,\n",
    "                    cer_result):\n",
    "        self.result_state.add_results(epoch=epoch,\n",
    "                                      loss_result=loss_result,\n",
    "                                      wer_result=wer_result,\n",
    "                                      cer_result=cer_result)\n",
    "\n",
    "    def init_finetune_states(self, epochs):\n",
    "        \"\"\"\n",
    "        Resets the training environment, but keeps model specific states in tact.\n",
    "        This is when fine-tuning a model on another dataset where training is to be reset but the model\n",
    "        weights are to be loaded\n",
    "        :param epochs: Number of epochs fine-tuning.\n",
    "        \"\"\"\n",
    "        self.init_results_tracking(epochs)\n",
    "        self._reset_amp_state()\n",
    "        self._reset_optim_state()\n",
    "        self._reset_epoch()\n",
    "        self.reset_training_step()\n",
    "        self._reset_best_wer()\n",
    "\n",
    "    def serialize_state(self, epoch, iteration):\n",
    "        model = remove_parallel_wrapper(self.model)\n",
    "        model_dict = model.serialize_state()\n",
    "        training_dict = self._serialize_training_state(epoch=epoch,\n",
    "                                                       iteration=iteration)\n",
    "        results_dict = self.result_state.serialize_state()\n",
    "        # Ensure flat structure for backwards compatibility\n",
    "        state_dict = {**model_dict, **training_dict, **results_dict}  # Combine dicts\n",
    "        return state_dict\n",
    "\n",
    "    def _serialize_training_state(self, epoch, iteration):\n",
    "        return {\n",
    "            'optim_dict': self.optim_state,\n",
    "            'amp': self.amp_state,\n",
    "            'avg_loss': self.avg_loss,\n",
    "            'best_wer': self.best_wer,\n",
    "            'epoch': epoch + 1,  # increment for readability\n",
    "            'iteration': iteration,\n",
    "        }\n",
    "\n",
    "    @classmethod\n",
    "    def load_state(cls, state_path):\n",
    "        print(\"Loading state from model %s\" % state_path)\n",
    "        state = torch.load(state_path, map_location=lambda storage, loc: storage)\n",
    "        model = DeepSpeech.load_model_package(state)\n",
    "        optim_state = state['optim_dict']\n",
    "        amp_state = state['amp']\n",
    "        epoch = int(state.get('epoch', 1)) - 1  # Index start at 0 for training\n",
    "        training_step = state.get('iteration', None)\n",
    "        if training_step is None:\n",
    "            epoch += 1  # We saved model after epoch finished, start at the next epoch.\n",
    "            training_step = 0\n",
    "        else:\n",
    "            training_step += 1\n",
    "        avg_loss = int(state.get('avg_loss', 0))\n",
    "        loss_results = state['loss_results']\n",
    "        cer_results = state['cer_results']\n",
    "        wer_results = state['wer_results']\n",
    "        best_wer = state.get('best_wer')\n",
    "\n",
    "        result_state = ResultState(loss_results=loss_results,\n",
    "                                   cer_results=cer_results,\n",
    "                                   wer_results=wer_results)\n",
    "        return cls(optim_state=optim_state,\n",
    "                   amp_state=amp_state,\n",
    "                   model=model,\n",
    "                   result_state=result_state,\n",
    "                   best_wer=best_wer,\n",
    "                   avg_loss=avg_loss,\n",
    "                   epoch=epoch,\n",
    "                   training_step=training_step)\n",
    "\n",
    "    def set_epoch(self, epoch):\n",
    "        self.epoch = epoch\n",
    "\n",
    "    def set_best_wer(self, wer):\n",
    "        self.best_wer = wer\n",
    "\n",
    "    def set_training_step(self, training_step):\n",
    "        self.training_step = training_step\n",
    "\n",
    "    def reset_training_step(self):\n",
    "        self.training_step = 0\n",
    "\n",
    "    def reset_avg_loss(self):\n",
    "        self.avg_loss = 0\n",
    "\n",
    "    def _reset_amp_state(self):\n",
    "        self.amp_state = None\n",
    "\n",
    "    def _reset_optim_state(self):\n",
    "        self.optim_state = None\n",
    "\n",
    "    def _reset_epoch(self):\n",
    "        self.epoch = 0\n",
    "\n",
    "    def _reset_best_wer(self):\n",
    "        self.best_wer = None\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define evaluate\n",
    "this is equivalent of ```from test import evaluate```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate(test_loader, device, model, decoder, target_decoder, save_output=None, verbose=False, half=False):\n",
    "    model.eval()\n",
    "    total_cer, total_wer, num_tokens, num_chars = 0, 0, 0, 0\n",
    "    output_data = []\n",
    "    for i, (data) in tqdm(enumerate(test_loader), total=len(test_loader)):\n",
    "        \n",
    "        inputs, targets, input_percentages, target_sizes = data\n",
    "\n",
    "        input_sizes = input_percentages.mul_(int(inputs.size(3))).int()\n",
    "        inputs = inputs.to(device)\n",
    "        if half:\n",
    "            inputs = inputs.half()\n",
    "        # unflatten targets\n",
    "        split_targets = []\n",
    "        offset = 0\n",
    "        for size in target_sizes:\n",
    "            split_targets.append(targets[offset:offset + size])\n",
    "            offset += size\n",
    "\n",
    "        out, output_sizes = model(inputs, input_sizes)\n",
    "\n",
    "        decoded_output, _ = decoder.decode(out, output_sizes)\n",
    "        target_strings = target_decoder.convert_to_strings(split_targets)\n",
    "\n",
    "        if save_output is not None:\n",
    "            # add output to data array, and continue\n",
    "            output_data.append((out.cpu(), output_sizes, target_strings))\n",
    "        for x in range(len(target_strings)):\n",
    "            transcript, reference = decoded_output[x][0], target_strings[x][0]\n",
    "            wer_inst = decoder.wer(transcript, reference)\n",
    "            cer_inst = decoder.cer(transcript, reference)\n",
    "            total_wer += wer_inst\n",
    "            total_cer += cer_inst\n",
    "            num_tokens += len(reference.split())\n",
    "            num_chars += len(reference.replace(' ', ''))\n",
    "            if verbose:\n",
    "                print(\"Ref:\", reference.lower())\n",
    "                print(\"Hyp:\", transcript.lower())\n",
    "                print(\"WER:\", float(wer_inst) / len(reference.split()),\n",
    "                      \"CER:\", float(cer_inst) / len(reference.replace(' ', '')), \"\\n\")\n",
    "    wer = float(total_wer) / num_tokens\n",
    "    cer = float(total_cer) / num_chars\n",
    "    return wer * 100, cer * 100, output_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define check_loss, CheckpointHandler\n",
    "this is equivalent of ```from utils import check_loss, CheckpointHandler```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from pathlib import Path\n",
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def check_loss(loss, loss_value):\n",
    "    \"\"\"\n",
    "    Check that warp-ctc loss is valid and will not break training\n",
    "    :return: Return if loss is valid, and the error in case it is not\n",
    "    \"\"\"\n",
    "    loss_valid = True\n",
    "    error = ''\n",
    "    if loss_value == float(\"inf\") or loss_value == float(\"-inf\"):\n",
    "        loss_valid = False\n",
    "        error = \"WARNING: received an inf loss\"\n",
    "    elif torch.isnan(loss).sum() > 0:\n",
    "        loss_valid = False\n",
    "        error = 'WARNING: received a nan loss, setting loss value to 0'\n",
    "    elif loss_value < 0:\n",
    "        loss_valid = False\n",
    "        error = \"WARNING: received a negative loss\"\n",
    "    return loss_valid, error"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CheckpointHandler:\n",
    "    def __init__(self,\n",
    "                 save_folder: str,\n",
    "                 best_val_model_name: str,\n",
    "                 checkpoint_per_iteration: int,\n",
    "                 save_n_recent_models: int):\n",
    "        self.save_folder = Path(save_folder)\n",
    "        self.save_folder.mkdir(parents=True, exist_ok=True)  # Ensure save folder exists\n",
    "        self.checkpoint_prefix = 'deepspeech_checkpoint_'  # TODO do we want to expose this?\n",
    "        self.checkpoint_prefix_path = self.save_folder / self.checkpoint_prefix\n",
    "        self.best_val_path = self.save_folder / best_val_model_name\n",
    "        self.checkpoint_per_iteration = checkpoint_per_iteration\n",
    "        self.save_n_recent_models = save_n_recent_models\n",
    "\n",
    "    def find_latest_checkpoint(self):\n",
    "        \"\"\"\n",
    "        Finds the latest checkpoint in a folder based on the timestamp of the file.\n",
    "        If there are no checkpoints, returns None.\n",
    "        :return: The latest checkpoint path, or None if no checkpoints are found.\n",
    "        \"\"\"\n",
    "        paths = list(self.save_folder.rglob(self.checkpoint_prefix + '*'))\n",
    "        if paths:\n",
    "            paths.sort(key=os.path.getctime)\n",
    "            latest_checkpoint_path = paths[-1]\n",
    "            return latest_checkpoint_path\n",
    "        else:\n",
    "            return None\n",
    "\n",
    "    def check_and_delete_oldest_checkpoint(self):\n",
    "        paths = list(self.save_folder.rglob(self.checkpoint_prefix + '*'))\n",
    "        if paths and len(paths) >= self.save_n_recent_models:\n",
    "            paths.sort(key=os.path.getctime)\n",
    "            print(\"Deleting old checkpoint %s\" % str(paths[0]))\n",
    "            os.remove(paths[0])\n",
    "\n",
    "    def save_checkpoint_model(self, epoch, state, i=None):\n",
    "        if self.save_n_recent_models > 0:\n",
    "            self.check_and_delete_oldest_checkpoint()\n",
    "        model_path = self._create_checkpoint_path(epoch=epoch,\n",
    "                                                  i=i)\n",
    "        print(\"Saving checkpoint model to %s\" % model_path)\n",
    "        torch.save(obj=state.serialize_state(epoch=epoch,\n",
    "                                             iteration=i),\n",
    "                   f=model_path)\n",
    "\n",
    "    def save_iter_checkpoint_model(self, epoch, state, i):\n",
    "        if self.checkpoint_per_iteration > 0 and i > 0 and (i + 1) % self.checkpoint_per_iteration == 0:\n",
    "            self.save_checkpoint_model(epoch=epoch,\n",
    "                                       state=state,\n",
    "                                       i=i)\n",
    "\n",
    "    def save_best_model(self, epoch, state):\n",
    "        print(\"Found better validated model, saving to %s\" % self.best_val_path)\n",
    "        torch.save(obj=state.serialize_state(epoch=epoch,\n",
    "                                             iteration=None),\n",
    "                   f=self.best_val_path)\n",
    "\n",
    "    def _create_checkpoint_path(self, epoch, i=None):\n",
    "        \"\"\"\n",
    "        Creates path to save checkpoint.\n",
    "        We automatically iterate the epoch and iteration for readibility.\n",
    "        :param epoch: The epoch (index starts at 0).\n",
    "        :param i: The iteration (index starts at 0).\n",
    "        :return: The path to save the model\n",
    "        \"\"\"\n",
    "        if i:\n",
    "            checkpoint_path = str(self.checkpoint_prefix_path) + 'epoch_%d_iter_%d.pth' % (epoch + 1, i + 1)\n",
    "        else:\n",
    "            checkpoint_path = str(self.checkpoint_prefix_path) + 'epoch_%d.pth' % (epoch + 1)\n",
    "        return checkpoint_path"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Finally We're ready to start"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Prepare datasets\n",
    "we use an4 audio dataset, which is relatively small and easy to experiment  \n",
    "http://www.speech.cs.cmu.edu/databases/an4/\n",
    "\n",
    "Input is wav file  \n",
    "Outout is text file which contains text of course."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import io\n",
    "import shutil\n",
    "import tarfile\n",
    "\n",
    "from sklearn.model_selection import train_test_split\n",
    "import wget\n",
    "\n",
    "# from data.utils import create_manifest\n",
    "import fnmatch\n",
    "import io\n",
    "import os\n",
    "import subprocess\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_manifest(data_path, output_path, min_duration=None, max_duration=None):\n",
    "    file_paths = [os.path.join(dirpath, f)\n",
    "                  for dirpath, dirnames, files in os.walk(data_path)\n",
    "                  for f in fnmatch.filter(files, '*.wav')]\n",
    "\n",
    "\n",
    "    # print(\"file_paths: \", file_paths)\n",
    "\n",
    "    file_paths = order_and_prune_files(file_paths, min_duration, max_duration)\n",
    "    with io.FileIO(output_path, \"w\") as file:\n",
    "        for wav_path in tqdm(file_paths, total=len(file_paths)):\n",
    "            transcript_path = wav_path.replace('/wav/', '/txt/').replace('.wav', '.txt')\n",
    "            sample = os.path.abspath(wav_path) + ',' + os.path.abspath(transcript_path) + '\\n'\n",
    "            file.write(sample.encode('utf-8'))\n",
    "    print('\\n')\n",
    "    \n",
    "def order_and_prune_files(file_paths, min_duration, max_duration):\n",
    "    print(\"Sorting manifests...\")\n",
    "    duration_file_paths = [(path, float(subprocess.check_output(\n",
    "        ['soxi -D \\\"%s\\\"' % path.strip()], shell=True))) for path in file_paths]\n",
    "    if min_duration and max_duration:\n",
    "        print(\"Pruning manifests between %d and %d seconds\" % (min_duration, max_duration))\n",
    "        duration_file_paths = [(path, duration) for path, duration in duration_file_paths if\n",
    "                               min_duration <= duration <= max_duration]\n",
    "\n",
    "    def func(element):\n",
    "        return element[1]\n",
    "\n",
    "    duration_file_paths.sort(key=func)\n",
    "    return [x[0] for x in duration_file_paths]  # Remove durations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _format_test_data(root_path):\n",
    "    wav_path = root_path + 'wav/'\n",
    "    file_ids_path = root_path + 'etc/an4_test.fileids'\n",
    "    transcripts_path = root_path + 'etc/an4_test.transcription'\n",
    "    root_wav_path = wav_path + 'an4test_clstk'\n",
    "\n",
    "    _convert_audio_to_wav(root_wav_path)\n",
    "    file_ids, transcripts = _retrieve_file_ids_and_transcripts(file_ids_path, transcripts_path)\n",
    "\n",
    "    _save_wav_transcripts(data_type='test',\n",
    "                          file_ids=file_ids,\n",
    "                          transcripts=transcripts,\n",
    "                          wav_dir=wav_path)\n",
    "\n",
    "\n",
    "def _process_transcript(transcript):\n",
    "    \"\"\"\n",
    "    Removes tags found in AN4.\n",
    "    \"\"\"\n",
    "    extracted_transcript = transcript.split('(')[0].strip(\"<s>\").split('<')[0].strip().upper()\n",
    "    return extracted_transcript\n",
    "\n",
    "def _save_files(file_ids, transcripts, wav_dir, new_wav_dir, new_transcript_dir):\n",
    "    for file_id, transcript in zip(file_ids, transcripts):\n",
    "        path = wav_dir + file_id.strip() + '.wav'\n",
    "        filename = path.split('/')[-1]\n",
    "        extracted_transcript = _process_transcript(transcript)\n",
    "        new_path = new_wav_dir + filename\n",
    "        text_path = new_transcript_dir + filename.replace('.wav', '.txt')\n",
    "        with io.FileIO(text_path, \"w\") as file:\n",
    "            file.write(extracted_transcript.encode('utf-8'))\n",
    "        current_path = os.path.abspath(path)\n",
    "        os.rename(current_path, new_path)\n",
    "\n",
    "\n",
    "def _save_wav_transcripts(data_type, file_ids, transcripts, wav_dir):\n",
    "    \n",
    "    target_dir = './data/an4_dataset/'\n",
    "    \n",
    "    data_path = target_dir + data_type + '/an4/'\n",
    "    new_transcript_dir = data_path + '/txt/'\n",
    "    new_wav_dir = data_path + '/wav/'\n",
    "\n",
    "    os.makedirs(new_transcript_dir)\n",
    "    os.makedirs(new_wav_dir)\n",
    "\n",
    "    _save_files(file_ids=file_ids,\n",
    "                transcripts=transcripts,\n",
    "                wav_dir=wav_dir,\n",
    "                new_wav_dir=new_wav_dir,\n",
    "                new_transcript_dir=new_transcript_dir)\n",
    "\n",
    "\n",
    "def _retrieve_file_ids_and_transcripts(file_id_path, transcripts_path):\n",
    "    with open(file_id_path, 'r') as f:\n",
    "        file_ids = f.readlines()\n",
    "    with open(transcripts_path, 'r') as t:\n",
    "        transcripts = t.readlines()\n",
    "    return file_ids, transcripts\n",
    "\n",
    "def _convert_audio_to_wav(an4_audio_path):\n",
    "    with os.popen('find %s -type f -name \"*.raw\"' % an4_audio_path) as pipe:\n",
    "        for line in pipe:\n",
    "            raw_path = line.strip()\n",
    "            new_path = line.replace('.raw', '.wav').strip()\n",
    "            cmd = 'sox -t raw -r %d -b 16 -e signed-integer -B -c 1 \\\"%s\\\" \\\"%s\\\"' % (\n",
    "                16000, raw_path, new_path)\n",
    "            os.system(cmd)\n",
    "\n",
    "def _format_training_data(root_path):\n",
    "    \n",
    "    val_fraction = 0.1 \n",
    "    \n",
    "    wav_path = root_path + 'wav/'\n",
    "    file_ids_path = root_path + 'etc/an4_train.fileids'\n",
    "    transcripts_path = root_path + 'etc/an4_train.transcription'\n",
    "    root_wav_path = wav_path + 'an4_clstk'\n",
    "\n",
    "    _convert_audio_to_wav(root_wav_path)\n",
    "    file_ids, transcripts = _retrieve_file_ids_and_transcripts(file_ids_path, transcripts_path)\n",
    "\n",
    "    split_files = train_test_split(file_ids, transcripts, test_size=val_fraction)\n",
    "    train_file_ids, val_file_ids, train_transcripts, val_transcripts = split_files\n",
    "\n",
    "    _save_wav_transcripts(data_type='train',\n",
    "                          file_ids=train_file_ids,\n",
    "                          transcripts=train_transcripts,\n",
    "                          wav_dir=wav_path)\n",
    "    _save_wav_transcripts(data_type='val',\n",
    "                          file_ids=val_file_ids,\n",
    "                          transcripts=val_transcripts,\n",
    "                          wav_dir=wav_path)\n",
    "\n",
    "\n",
    "def parepare_dataset():\n",
    "    \n",
    "    target_dir = './data/an4_dataset/'\n",
    "    min_duration = 2\n",
    "    max_duration = 15\n",
    "#     val_fraction = 0.1 \n",
    "    \n",
    "    root_path = 'data/an4/'\n",
    "    \n",
    "    raw_tar_path = './data/an4_raw.bigendian.tar.gz'\n",
    "\n",
    "    if not os.path.exists(raw_tar_path):\n",
    "        wget.download('http://www.speech.cs.cmu.edu/databases/an4/an4_raw.bigendian.tar.gz')\n",
    "        os.rename(\"an4_raw.bigendian.tar.gz\", \"data/an4_raw.bigendian.tar.gz\")\n",
    "\n",
    "    tar = tarfile.open(raw_tar_path)\n",
    "    tar.extractall()\n",
    "\n",
    "    # move\n",
    "    os.rename(\"an4\", \"data/an4\")\n",
    "\n",
    "    os.makedirs(target_dir)\n",
    "    _format_training_data(root_path=root_path)\n",
    "    _format_test_data(root_path=root_path)\n",
    "    shutil.rmtree(root_path)\n",
    "    # os.remove(raw_tar_path)\n",
    "    train_path = target_dir + '/train/'\n",
    "    val_path = target_dir + '/val/'\n",
    "    test_path = target_dir + '/test/'\n",
    "    print('Creating manifests...')\n",
    "    create_manifest(train_path, './data/an4_train_manifest.csv', min_duration, max_duration)\n",
    "    create_manifest(val_path, './data/an4_val_manifest.csv', min_duration, max_duration)\n",
    "    create_manifest(test_path, './data/an4_test_manifest.csv')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# do it once\n",
    "!mkdir data\n",
    "parepare_dataset()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# see audio example \n",
    "import IPython.display"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "audio_example_path = \"data/an4_dataset/test/an4/wav/an391-mjwl-b.wav\"\n",
    "IPython.display.Audio(audio_example_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# label for this\n",
    "! cat \"data/an4_dataset/test/an4/txt/an391-mjwl-b.txt\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# let's visualize it\n",
    "\n",
    "def visualization_spectrogram(mel_spectrogram, title):\n",
    "    \"\"\"visualizing result of SpecAugment\n",
    "    # Arguments:\n",
    "      mel_spectrogram(ndarray): mel_spectrogram to visualize.\n",
    "      title(String): plot figure's title\n",
    "    \"\"\"\n",
    "    \n",
    "    # Show mel-spectrogram using librosa's specshow.\n",
    "    plt.figure(figsize=(10, 4))\n",
    "    # librosa.display.specshow(librosa.power_to_db(mel_spectrogram[0, :, :], ref=np.max), y_axis='mel', fmax=8000, x_axis='time')\n",
    "    librosa.display.specshow(librosa.power_to_db(mel_spectrogram, ref=np.max), y_axis='mel', fmax=8000, x_axis='time')\n",
    "\n",
    "    # plt.colorbar(format='%+2.0f dB')\n",
    "    plt.title(title)\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load audio\n",
    "from scipy.io.wavfile import read\n",
    "import librosa\n",
    "import librosa.display # need this for ver 0.5.0+\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "def load_audio(path):\n",
    "\n",
    "    try:\n",
    "        sample_rate, sound = read(path)\n",
    "    except:\n",
    "        print(\"path: \", path)\n",
    "    \n",
    "    sound = sound.astype('float32') / 32767  # normalize audio\n",
    "    if len(sound.shape) > 1:\n",
    "        if sound.shape[1] == 1:\n",
    "            sound = sound.squeeze()\n",
    "        else:\n",
    "            sound = sound.mean(axis=1)  # multiple channels, average\n",
    "    return sound"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# pase audio \n",
    "\n",
    "sample_rate = 16000\n",
    "window_size = .02\n",
    "window_stride = .01\n",
    "window = 'hamming'\n",
    "\n",
    "def parse_audio(audio_path, spec_augment_flg=True):\n",
    "\n",
    "    y = load_audio(audio_path)\n",
    "\n",
    "    n_fft = int(sample_rate * window_size)\n",
    "    win_length = n_fft\n",
    "    hop_length = int(sample_rate * window_stride)\n",
    "    # STFT\n",
    "    D = librosa.stft(y, n_fft=n_fft, hop_length=hop_length,\n",
    "                     win_length=win_length, window=window)\n",
    "    spect, phase = librosa.magphase(D)\n",
    "    # S = log(S+1)\n",
    "    spect = np.log1p(spect)\n",
    "    spect = torch.FloatTensor(spect)\n",
    "    \n",
    "    # if self.normalize:\n",
    "    # True\n",
    "    mean = spect.mean()\n",
    "    std = spect.std()\n",
    "    spect.add_(-mean)\n",
    "    spect.div_(std)\n",
    "\n",
    "#     if spec_augment_flg is True:\n",
    "#         spect= spec_augment(spect)\n",
    "\n",
    "    return spect"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "spect_example = parse_audio(audio_example_path)\n",
    "spect_example.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%reload_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# lets visualized it\n",
    "visualization_spectrogram(spect_example, 'mel_spectrogram example')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# let's see wave form too\n",
    "\n",
    "import librosa\n",
    "x, sr = librosa.load(audio_example_path)\n",
    "print(x, sr)\n",
    "\n",
    "plt.figure(figsize=(14, 5))\n",
    "librosa.display.waveplot(x, sr=sr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Traininfg process"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# params for training...\n",
    "\n",
    "class Parser:\n",
    "    \"\"\"Parameters for training\"\"\"\n",
    "    def __init__(self):\n",
    "        self.lr = 3e-4\n",
    "        # self.lr = 4e-4\n",
    "        self.epochs = 25\n",
    "        self.batch_size = 16\n",
    "        self.num_workers = 1\n",
    "        self.hidden_size = 1024\n",
    "        self.hidden_layers = 5\n",
    "        self.bidirectional = True\n",
    "        self.seed = 42\n",
    "        self.momentum = 0.9\n",
    "        self.wd = 1e-5\n",
    "        self.opt_level = 'O0'\n",
    "        self.keep_batchnorm_fp32 = None\n",
    "        self.loss_scale = 1\n",
    "        # self.max_norm = 400\n",
    "        self.max_norm = 300\n",
    "        self.checkpoint_per_iteration = 0\n",
    "        self.checkpoint = False\n",
    "        self.save_folder = 'snapshots'\n",
    "        self.best_val_model_name = \"deepspeech_final.pth\"\n",
    "        self.save_n_recent_models = 0\n",
    "        self.learning_anneal = 1.01\n",
    "        self.eps = 1e-8 # for adamw\n",
    "        self.betas = (0.9, 0.999) # for adamw\n",
    "    \n",
    "args = Parser()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### labels\n",
    "tokens are alphabets and \"_\" and \"'\" and \" \""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # labels\n",
    "# labels_path = \"labels.json\"\n",
    "# with open(labels_path) as label_file:\n",
    "#     labels = json.load(label_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# labels to map token to index\n",
    "labels = [\n",
    "  \"_\",\n",
    "  \"'\",\n",
    "  \"A\",\n",
    "  \"B\",\n",
    "  \"C\",\n",
    "  \"D\",\n",
    "  \"E\",\n",
    "  \"F\",\n",
    "  \"G\",\n",
    "  \"H\",\n",
    "  \"I\",\n",
    "  \"J\",\n",
    "  \"K\",\n",
    "  \"L\",\n",
    "  \"M\",\n",
    "  \"N\",\n",
    "  \"O\",\n",
    "  \"P\",\n",
    "  \"Q\",\n",
    "  \"R\",\n",
    "  \"S\",\n",
    "  \"T\",\n",
    "  \"U\",\n",
    "  \"V\",\n",
    "  \"W\",\n",
    "  \"X\",\n",
    "  \"Y\",\n",
    "  \"Z\",\n",
    "  \" \"\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# labels[:10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# audio conf\n",
    "audio_conf = dict(sample_rate=16000,\n",
    "                  window_size=.02,\n",
    "                  window_stride=.01,\n",
    "                  window='hamming',\n",
    "                  noise_dir=None,\n",
    "                  noise_prob=0.4,\n",
    "                  noise_levels=(0.0, 0.5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# csv file for input wav and output text file \n",
    "train_manifest = 'data/an4_train_manifest.csv'\n",
    "val_manifest = 'data/an4_val_manifest.csv'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Dataset\n",
    "train_dataset = SpectrogramDataset(audio_conf=audio_conf,\n",
    "                                   manifest_filepath=train_manifest,\n",
    "                                   labels=labels,\n",
    "                                   normalize=True,\n",
    "                                   speed_volume_perturb=True,\n",
    "                                   spec_augment=False)\n",
    "test_dataset = SpectrogramDataset(audio_conf=audio_conf,\n",
    "                                  manifest_filepath=val_manifest,\n",
    "                                  labels=labels,\n",
    "                                  normalize=True,\n",
    "                                  speed_volume_perturb=False,\n",
    "                                  spec_augment=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# sampler to use bins\n",
    "train_sampler = DSRandomSampler(dataset=train_dataset,\n",
    "                                batch_size=args.batch_size,\n",
    "                                start_index=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# data loader\n",
    "\n",
    "train_loader = AudioDataLoader(dataset=train_dataset,\n",
    "                               num_workers=args.num_workers,\n",
    "                               batch_sampler=train_sampler)\n",
    "test_loader = AudioDataLoader(dataset=test_dataset,\n",
    "                              num_workers=args.num_workers,\n",
    "                              batch_size=args.batch_size)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### model "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set rnn type\n",
    "rnn_type = 'lstm'\n",
    "\n",
    "# prepare model\n",
    "model = DeepSpeech(rnn_hidden_size=args.hidden_size,\n",
    "                   nb_layers=args.hidden_layers,\n",
    "                   labels=labels,\n",
    "                   rnn_type=supported_rnns[rnn_type],\n",
    "                   audio_conf=audio_conf,\n",
    "                   bidirectional=args.bidirectional)\n",
    "\n",
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# use TrainingState to track training state\n",
    "state = TrainingState(model=model)\n",
    "state.init_results_tracking(epochs=args.epochs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# use greedy decoder. you can use BeamDecoder also\n",
    "evaluation_decoder = GreedyDecoder(model.labels)  # Decoder used for validation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# device = torch.device(\"cpu\" if args.no_cuda else \"cuda\")\n",
    "# ctc-loos works only with gpu?\n",
    "device = torch.device(\"cuda\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# optim\n",
    "model = model.to(device)\n",
    "parameters = model.parameters()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# loss go inf if i use SGD, use AdamW\n",
    "\n",
    "# optimizer = torch.optim.SGD(parameters,\n",
    "#                             lr=args.lr,\n",
    "#                             momentum=args.momentum,\n",
    "#                             nesterov=True,\n",
    "#                             weight_decay=args.wd)\n",
    "\n",
    "optimizer = torch.optim.AdamW(parameters,\n",
    "                              lr=args.lr,\n",
    "                              betas=args.betas,\n",
    "                              eps=args.eps,\n",
    "                              weight_decay=args.wd)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Track states for optimizer\n",
    "state.track_optim_state(optimizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class AverageMeter(object):\n",
    "    \"\"\"Computes and stores the average and current value\"\"\"\n",
    "\n",
    "    def __init__(self):\n",
    "        self.reset()\n",
    "\n",
    "    def reset(self):\n",
    "        self.val = 0\n",
    "        self.avg = 0\n",
    "        self.sum = 0\n",
    "        self.count = 0\n",
    "\n",
    "    def update(self, val, n=1):\n",
    "        self.val = val\n",
    "        self.sum += val * n\n",
    "        self.count += n\n",
    "        self.avg = self.sum / self.count\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# use CTC for our loss\n",
    "criterion = CTCLoss()\n",
    "batch_time = AverageMeter()\n",
    "data_time = AverageMeter()\n",
    "losses = AverageMeter()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# check step count\n",
    "len(train_loader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# prepare check point handler\n",
    "checkpoint_handler = CheckpointHandler(save_folder=args.save_folder,\n",
    "                                       best_val_model_name=args.best_val_model_name,\n",
    "                                       checkpoint_per_iteration=args.checkpoint_per_iteration,\n",
    "                                       save_n_recent_models=args.save_n_recent_models)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### finally training loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# training loop\n",
    "\n",
    "main_proc = True\n",
    "\n",
    "for epoch in range(state.epoch, args.epochs):\n",
    "    model.train()\n",
    "    end = time.time()\n",
    "    start_epoch_time = time.time()\n",
    "    state.set_epoch(epoch=epoch)\n",
    "    train_sampler.set_epoch(epoch=epoch)\n",
    "    train_sampler.reset_training_step(training_step=state.training_step)\n",
    "    for idx, (data) in enumerate(train_loader, start=state.training_step):\n",
    "        state.set_training_step(training_step=idx)\n",
    "        inputs, targets, input_percentages, target_sizes = data\n",
    "        input_sizes = input_percentages.mul_(int(inputs.size(3))).int()\n",
    "        # measure data loading time\n",
    "        data_time.update(time.time() - end)\n",
    "        inputs = inputs.to(device)\n",
    "\n",
    "        out, output_sizes = model(inputs, input_sizes)\n",
    "        out = out.transpose(0, 1)  # TxNxH\n",
    "\n",
    "        float_out = out.float()  # ensure float32 for loss\n",
    "        loss = criterion(float_out, targets, output_sizes, target_sizes).to(device)\n",
    "        loss = loss / inputs.size(0)  # average the loss by minibatch\n",
    "        loss_value = loss.item()\n",
    "\n",
    "        # Check to ensure valid loss was calculated\n",
    "        valid_loss, error = check_loss(loss, loss_value)\n",
    "        if valid_loss:\n",
    "            optimizer.zero_grad()\n",
    "            \n",
    "            # compute gradient\n",
    "            loss.backward()\n",
    "    \n",
    "            torch.nn.utils.clip_grad_norm_(parameters, args.max_norm)\n",
    "            optimizer.step()\n",
    "        else:\n",
    "            # print(error)\n",
    "            # print('Skipping grad update')\n",
    "            loss_value = 0\n",
    "\n",
    "        state.avg_loss += loss_value\n",
    "        losses.update(loss_value, inputs.size(0))\n",
    "\n",
    "        # measure elapsed time\n",
    "        batch_time.update(time.time() - end)\n",
    "        end = time.time()\n",
    "        if idx % 5 == 0:\n",
    "            print('Epoch: [{0}][{1}/{2}]\\t'\n",
    "                  'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\\t'\n",
    "                  'Data {data_time.val:.3f} ({data_time.avg:.3f})\\t'\n",
    "                  'Loss {loss.val:.4f} ({loss.avg:.4f})\\t'.format(\n",
    "                (epoch + 1), (idx + 1), len(train_loader), batch_time=batch_time, data_time=data_time, loss=losses))\n",
    "\n",
    "        if main_proc and args.checkpoint_per_iteration:\n",
    "            checkpoint_handler.save_iter_checkpoint_model(epoch=epoch, i=idx, state=state)\n",
    "            \n",
    "        del loss, out, float_out\n",
    "\n",
    "    state.avg_loss /= len(train_dataset)\n",
    "\n",
    "    epoch_time = time.time() - start_epoch_time\n",
    "    print('Training Summary Epoch: [{0}]\\t'\n",
    "          'Time taken (s): {epoch_time:.0f}\\t'\n",
    "          'Average Loss {loss:.3f}\\t'.format(epoch + 1, epoch_time=epoch_time, loss=state.avg_loss))\n",
    "\n",
    "    with torch.no_grad():\n",
    "        wer, cer, output_data = evaluate(test_loader=test_loader,\n",
    "                                         device=device,\n",
    "                                         model=model,\n",
    "                                         decoder=evaluation_decoder,\n",
    "                                         target_decoder=evaluation_decoder)\n",
    "\n",
    "    state.add_results(epoch=epoch,\n",
    "                      loss_result=state.avg_loss,\n",
    "                      wer_result=wer,\n",
    "                      cer_result=cer)\n",
    "\n",
    "    print('Validation Summary Epoch: [{0}]\\t'\n",
    "          'Average WER {wer:.3f}\\t'\n",
    "          'Average CER {cer:.3f}\\t'.format(epoch + 1, wer=wer, cer=cer))\n",
    "\n",
    "    if main_proc and args.checkpoint:  # Save epoch checkpoint\n",
    "        checkpoint_handler.save_checkpoint_model(epoch=epoch, state=state)\n",
    "    # anneal lr\n",
    "    for g in optimizer.param_groups:\n",
    "        g['lr'] = g['lr'] / args.learning_anneal\n",
    "    print('Learning rate annealed to: {lr:.6f}'.format(lr=g['lr']))\n",
    "\n",
    "    if main_proc and (state.best_wer is None or state.best_wer > wer):\n",
    "        checkpoint_handler.save_best_model(epoch=epoch, state=state)\n",
    "        state.set_best_wer(wer)\n",
    "        state.reset_avg_loss()\n",
    "    state.reset_training_step()  # Reset training step for next epoch\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## let's evaluate the model we trained just now"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# function to get inference\n",
    "\n",
    "def transcribe(audio_path, spect_parser, model, decoder, device, use_half):\n",
    "    spect = spect_parser.parse_audio(audio_path).contiguous()\n",
    "    spect = spect.view(1, 1, spect.size(0), spect.size(1))\n",
    "    spect = spect.to(device)\n",
    "    if use_half:\n",
    "        spect = spect.half()\n",
    "    input_sizes = torch.IntTensor([spect.size(3)]).int()\n",
    "    out, output_sizes = model(spect, input_sizes)\n",
    "    decoded_output, decoded_offsets = decoder.decode(out, output_sizes)\n",
    "    return decoded_output, decoded_offsets\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# decode inference\n",
    "\n",
    "def decode_results(decoded_output, decoded_offsets):\n",
    "    results = {\n",
    "        \"output\": [],\n",
    "        \"_meta\": {\n",
    "            \"acoustic_model\": {\n",
    "                \"name\": os.path.basename(\"snapshots/deepspeech_final.pth\")\n",
    "            },\n",
    "            \"language_model\": {\n",
    "                \"name\": None,\n",
    "            },\n",
    "            \"decoder\": {\n",
    "                \"lm\": None,\n",
    "                \"alpha\": None,\n",
    "                \"beta\": None,\n",
    "                \"type\": \"GreedyDecoder\",\n",
    "            }\n",
    "        }\n",
    "    }\n",
    "\n",
    "    for b in range(len(decoded_output)):\n",
    "        for pi in range(min(1, len(decoded_output[b]))):\n",
    "            result = {'transcription': decoded_output[b][pi]}\n",
    "            results['output'].append(result)\n",
    "    return results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# this is sample audio \n",
    "audio_test_path = \"data/an4_dataset/test/an4/wav/cen2-mjwl-b.wav\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from utils import load_model\n",
    "def load_model(device, model_path, use_half):\n",
    "    model = DeepSpeech.load_model(model_path)\n",
    "    model.eval()\n",
    "    model = model.to(device)\n",
    "    if use_half:\n",
    "        model = model.half()\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load model\n",
    "model = load_model(\"cpu\", \"./snapshots/deepspeech_final.pth\", False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "decoder = GreedyDecoder(model.labels, blank_index=model.labels.index('_'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "spect_parser = SpectrogramParser(model.audio_conf, normalize=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "decoded_output, decoded_offsets = transcribe(audio_path=audio_test_path,\n",
    "                                             spect_parser=spect_parser,\n",
    "                                             model=model,\n",
    "                                             decoder=decoder,\n",
    "                                             device=\"cpu\",\n",
    "                                             use_half=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "final_result = decode_results(decoded_output, decoded_offsets)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "final_result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "IPython.display.Audio(audio_test_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! cat data/an4_dataset/test/an4/txt/cen2-mjwl-b.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}