HW4P2.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "HW4P2.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "code",
      "metadata": {
        "id": "PLpuTSeyEU8I",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "D9mbaHdT-IYz",
        "colab_type": "text"
      },
      "source": [
        "# Importing Libraries"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "iItlp2F2EYZ2",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import numpy as np\n",
        "import torch\n",
        "import random\n",
        "from torch import nn\n",
        "from torch.nn.utils.rnn import *\n",
        "from torch.utils.data import Dataset, DataLoader, TensorDataset\n",
        "from tqdm import tqdm\n",
        "from operator import itemgetter\n",
        "DEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
        "DEVICE"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qEG-HjEwEYg3",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!pip install python-Levenshtein\n",
        "import Levenshtein"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kGiK0DKv-UQD",
        "colab_type": "text"
      },
      "source": [
        "# Loading Datasets"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "17X64BmREYkw",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "train_x = np.load('/content/drive/My Drive/train_new.npy', allow_pickle=True, encoding='bytes')\n",
        "valid_x = np.load('/content/drive/My Drive/dev_new.npy', allow_pickle=True, encoding='bytes')\n",
        "test_x = np.load('/content/drive/My Drive/test_new.npy', allow_pickle=True, encoding='bytes')\n",
        "\n",
        "train_y = np.load('/content/drive/My Drive/train_transcripts.npy', allow_pickle=True,encoding='bytes')\n",
        "valid_y = np.load('/content/drive/My Drive/dev_transcripts.npy', allow_pickle=True,encoding='bytes')\n",
        "vocab = ['<sos>']+sorted(list(set(''.join([''.join([p.decode(\"utf-8\") for p in train_y[i]]) for i in range(train_y.shape[0])]))))+[' ']+['<eos>']+['<pad>']\n",
        "letter2index = dict(zip(vocab, range(len(vocab))))\n",
        "index2letter = dict(zip(range(len(vocab)),vocab))\n",
        "train_Y = [torch.LongTensor([vocab.index('<sos>')] + [vocab.index(j)  for j in ' '.join([q.decode(\"utf-8\") for q in train_y[i]])] + [vocab.index('<eos>')]) for i in range(train_y.shape[0])]\n",
        "valid_Y = [torch.LongTensor([vocab.index('<sos>')] + [vocab.index(j)  for j in ' '.join([q.decode(\"utf-8\") for q in valid_y[i]])] + [vocab.index('<eos>')]) for i in range(valid_y.shape[0])]"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8XKuIBauEYq8",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class Speech2TextDataset(Dataset):\n",
        "\n",
        "    def __init__(self, speech, text=None, isTrain=True):\n",
        "        self.speech =  [torch.FloatTensor(word) for word in speech]\n",
        "        self.isTrain = isTrain\n",
        "        if (text is not None):\n",
        "            self.text = text\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.speech)\n",
        "\n",
        "    def __getitem__(self, index):\n",
        "        if (self.isTrain == True):\n",
        "            return self.speech[index],self.text[index]#self.text[index][:-1], self.text[index][1:]\n",
        "        else:\n",
        "            return self.speech[index]\n",
        "\n",
        "\n",
        "def collate_train(seq_list):\n",
        "    inputs =  [i[0] for i in seq_list]\n",
        "    targets = [i[1] for i in seq_list]\n",
        "   \n",
        "    inputs = pad_sequence(inputs)\n",
        "    targets = pad_sequence(targets,batch_first=True, padding_value = 34)\n",
        "    X_lens = torch.LongTensor([len(seq[0]) for seq in seq_list])\n",
        "    Y_lens = torch.LongTensor([len(seq[1]) for seq in seq_list])\n",
        "    \n",
        "    return inputs,targets,X_lens,Y_lens\n",
        "\n",
        "\n",
        "def collate_test(seq_list):\n",
        "    inputs = [i for i in seq_list]\n",
        "    inputs = pad_sequence(inputs)\n",
        "    X_lens = torch.LongTensor([len(seq) for seq in seq_list])\n",
        "    return inputs,X_lens  "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4Q4vbQ62EY0T",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "train_dataset = Speech2TextDataset(train_x,train_Y)\n",
        "val_dataset =   Speech2TextDataset(valid_x,valid_Y)\n",
        "test_dataset =  Speech2TextDataset(test_x,None,False)\n",
        "train_loader =  DataLoader(train_dataset, shuffle=True, batch_size=64, collate_fn = collate_train,num_workers=12, pin_memory=True)\n",
        "val_loader =    DataLoader(val_dataset, shuffle=False, batch_size=256, collate_fn = collate_train,num_workers=12, pin_memory=True)\n",
        "test_loader =   DataLoader(test_dataset, shuffle=False, batch_size=1, collate_fn = collate_test,num_workers=12, pin_memory=True)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pCosozBh-YzA",
        "colab_type": "text"
      },
      "source": [
        "# Model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LANFVlER-dCZ",
        "colab_type": "text"
      },
      "source": [
        "## Attention"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "I2qqCLGMEY6z",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class Attention(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(Attention, self).__init__()\n",
        "    \n",
        "    def forward(self, query, keys, values,lengths):\n",
        "\n",
        "        keys = (torch.transpose(keys,0,1)).to(DEVICE)\n",
        "        lengths = lengths.to(DEVICE)\n",
        "        attention = (torch.bmm(keys, query.unsqueeze(2)).squeeze(2)).to(DEVICE)\n",
        "        mask = torch.arange(keys.size(1)).unsqueeze(0).to(DEVICE) >= lengths.unsqueeze(1)\n",
        "        attention.masked_fill_(mask, -1e20)\n",
        "        attention = nn.functional.softmax(attention, dim=1)\n",
        "        values = torch.transpose(values,0,1)\n",
        "        out = torch.bmm(attention.unsqueeze(1), values).squeeze(1)\n",
        "        \n",
        "        return out"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GIn4RUzj-hna",
        "colab_type": "text"
      },
      "source": [
        "## Encoder"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "h33SpgQKEY_H",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class Encoder(nn.Module):\n",
        "\n",
        "    def __init__(self, input_dim, hidden_dim, value_size=128,key_size=128):\n",
        "        super(Encoder, self).__init__()\n",
        "        \n",
        "        self.cnn1 = torch.nn.Conv1d(input_dim,hidden_dim//4 , 3, stride=2, padding=1,bias=False)\n",
        "        self.cnn2 = torch.nn.Conv1d(hidden_dim//4, hidden_dim//2, 3, stride=2, padding=1,bias=False)\n",
        "        self.cnn3 = torch.nn.Conv1d(hidden_dim//2, hidden_dim//2, 3, stride=2, padding=1,bias=False)\n",
        "        self.cnn4 = torch.nn.Conv1d(input_dim, hidden_dim//2, 1, stride=8, padding=0,bias=False)\n",
        "        self.tanh = nn.Hardtanh()\n",
        "        self.lstm1 = nn.LSTM(input_size=hidden_dim//2, hidden_size=hidden_dim, num_layers=3, dropout=0.5,bidirectional=True)\n",
        "        self.lstm2 = nn.LSTM(input_size=hidden_dim*2, hidden_size=hidden_dim, bidirectional=True)\n",
        "\n",
        "        self.key_network = nn.Linear(hidden_dim*2, value_size)\n",
        "        self.value_network = nn.Linear(hidden_dim*2, key_size)\n",
        "\n",
        "\n",
        "    def forward(self, x, lens,istraining=True):\n",
        "        \n",
        "\n",
        "        x1 = nn.functional.dropout(self.tanh(self.cnn1(x.transpose(0,1).transpose(1,2))),0.2,training=istraining)\n",
        "        lens = (torch.floor((((lens - 1)/2) + 1).float())).long()\n",
        "        \n",
        "        x1 = nn.functional.dropout(self.tanh(self.cnn2(x1)),0.2,training=istraining)\n",
        "        lens = (torch.floor((((lens - 1)/2) + 1).float())).long()\n",
        "\n",
        "        x1 = self.cnn3(x1)\n",
        "        lens = (torch.floor((((lens - 1)/2) + 1).float())).long()\n",
        "        x1 = x1.transpose(1,2).transpose(0,1)\n",
        "        x = (self.cnn4(x.transpose(0,1).transpose(1,2))).transpose(1,2).transpose(0,1)\n",
        "\n",
        "        x = nn.functional.dropout(self.tanh(x + x1),0.2,training=istraining)\n",
        "  \n",
        "        rnn_inp = pack_padded_sequence(x, lengths=lens, batch_first=False, enforce_sorted=False)\n",
        "        outputs = self.lstm1(rnn_inp)[0]\n",
        "        outputs,final_state = self.lstm2(outputs)\n",
        "\n",
        "        linear_input, lens = pad_packed_sequence(outputs)\n",
        "        \n",
        "        keys = self.key_network(linear_input)\n",
        "        value = self.value_network(linear_input)\n",
        "\n",
        "        return keys, value, lens, final_state"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dXg2jMj_-kbb",
        "colab_type": "text"
      },
      "source": [
        "## Decoder"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6rzyh6hkEY3z",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class Decoder(nn.Module):\n",
        "    def __init__(self, vocab_size, hidden_dim, value_size=128, key_size=128, isAttended=True):\n",
        "        super(Decoder, self).__init__()\n",
        "        self.embedding = nn.Embedding(vocab_size, hidden_dim, padding_idx=34)\n",
        "        self.lstm1 = nn.LSTMCell(input_size=hidden_dim + value_size, hidden_size=hidden_dim)\n",
        "        self.lstm2 = nn.LSTMCell(input_size=hidden_dim, hidden_size=key_size)        \n",
        "        self.linear = nn.Linear(key_size + value_size, hidden_dim)\n",
        "        self.relu = nn.ReLU()\n",
        "        self.character_prob = nn.Linear( hidden_dim, vocab_size)\n",
        "        self.attention = Attention()\n",
        "\n",
        "    def forward(self, x, context, key, values, context_lengths, hidden_states,istraining=True):\n",
        "\n",
        "        char_embed = self.embedding(x)\n",
        "\n",
        "        inp = torch.cat([char_embed, context], dim=1)    \n",
        "        hidden_states[0] = self.lstm1(inp, hidden_states[0])\n",
        "        inp_2 = hidden_states[0][0]\n",
        "        hidden_states[1] = self.lstm2(inp_2, hidden_states[1])\n",
        "\n",
        "        output = hidden_states[1][0]\n",
        "        context = self.attention(output, key, values, context_lengths)\n",
        "        prediction =  nn.functional.dropout(self.relu(self.linear(torch.cat([output, context], dim=1))),0.1,training=istraining)\n",
        "        prediction = self.character_prob(prediction)\n",
        " \n",
        "        return prediction, hidden_states,context"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6DRYm7o--oMr",
        "colab_type": "text"
      },
      "source": [
        "# Initializing Models"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9WBFv6SoFzhI",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def init_weights(m):\n",
        "    if type(m) == nn.Conv1d or type(m) == nn.Linear:\n",
        "        torch.nn.init.xavier_uniform_(m.weight.data)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ym1lr1DiEYxo",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "encoder = Encoder(40, 512, value_size=128,key_size=128)\n",
        "decoder = Decoder(len(vocab), 512*2, value_size=128, key_size=128, isAttended=True)\n",
        "encoder.load_state_dict(torch.load('/content/drive/My Drive/encoder3.pt'))\n",
        "decoder.load_state_dict(torch.load('/content/drive/My Drive/decoder3.pt'))\n",
        "#encoder.apply(init_weights)\n",
        "#decoder.apply(init_weights)\n",
        "#params = list(encoder.parameters()) + list(decoder.parameters())\n",
        "optimizer_encoder = torch.optim.Adam(encoder.parameters(), lr=0.001)\n",
        "optimizer_decoder = torch.optim.Adam(decoder.parameters(), lr=0.001)#, weight_decay=1e-6)\n",
        "criterion = nn.CrossEntropyLoss(reduction='sum')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4SzqxHtC-ufc",
        "colab_type": "text"
      },
      "source": [
        "# Training"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kDP69L9uEYun",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def train(encoder,decoder,train_loader,criterion, optimizer_encoder, optimizer_decoder,epoch,teacher):\n",
        "    encoder.train()\n",
        "    encoder.to(DEVICE)\n",
        "    decoder.train()\n",
        "    decoder.to(DEVICE)\n",
        "    total_loss = 0\n",
        "    perplexity = 0 \n",
        "\n",
        "    outer = tqdm(total=(len(train_loader)), desc='Training Epoch', position=0)\n",
        "    \n",
        "    for batch_idx, (inputs,targets,xlens,ylens) in enumerate(train_loader):\n",
        "          \n",
        "          outer.update(1)\n",
        "          with torch.autograd.set_detect_anomaly(True):\n",
        "            optimizer_encoder.zero_grad()\n",
        "            optimizer_decoder.zero_grad()\n",
        "\n",
        "            targets = targets.to(DEVICE)\n",
        "            batch_size = targets.size(0)\n",
        "            keys, values, lens, hidden = encoder(inputs.to(DEVICE),xlens.to(DEVICE),True)\n",
        "            hidden = tuple(st.transpose(0, 1).reshape(inputs.size(1), -1) for st in hidden)\n",
        "            \n",
        "            n_tokens = ylens.sum() - ylens.size(0)\n",
        "            hidden_states = [hidden, None]\n",
        "            pred= torch.zeros(batch_size,1).to(DEVICE)\n",
        "            context = values[0,:,:].to(DEVICE)\n",
        "            Y = targets.transpose(0,1)\n",
        "            loss = 0\n",
        "            \n",
        "            for i in range(Y.size(0) - 1):\n",
        "         \n",
        "                if np.random.random_sample() > teacher:\n",
        "\n",
        "                      pred, hidden_states, context = decoder(pred.argmax(dim=1), context, keys, values,lens, hidden_states,True)\n",
        "                      if np.random.random_sample() > 0.5:\n",
        "                          pred = torch.distributions.gumbel.Gumbel(pred.to('cpu'), torch.tensor([0.4])).sample().to(DEVICE)\n",
        "                else:\n",
        "                      pred, hidden_states, context = decoder(Y[i], context, keys, values,lens, hidden_states,True)\n",
        "\n",
        "                active = i + 1 < ylens\n",
        "                loss += criterion(pred[active], Y[i + 1, active])\n",
        "\n",
        "            loss /= n_tokens\n",
        "            current_loss = loss.item()\n",
        "            total_loss = total_loss + current_loss\n",
        "            perplexity = perplexity + np.exp(current_loss)\n",
        "            if batch_idx%10==0 and batch_idx>0:\n",
        "                  print(\"Batch: {:02d} \\t Train loss: {:.2f} \\t Train Perplexity : {:.2f}\".format(batch_idx,current_loss,np.exp(current_loss)))\n",
        "\n",
        "            loss.backward()\n",
        "            torch.nn.utils.clip_grad_norm_(encoder.parameters(), 2)\n",
        "            torch.nn.utils.clip_grad_norm_(decoder.parameters(), 2)\n",
        "            optimizer_encoder.step() \n",
        "            optimizer_decoder.step()          \n",
        "\n",
        "            torch.cuda.empty_cache()\n",
        "            del xlens,inputs,ylens,Y,active,loss,pred,hidden,lens,hidden_states,keys,values,targets\n",
        "    print(\"Epoch: {:02d} \\t Train loss: {:.2f} \\t Train Perplexity : {:.2f}\".format(epoch+1,total_loss/len(train_loader),perplexity/len(train_loader)))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Ch3zH0OeEYok",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def validate(encoder,decoder, val_loader,criterion, epoch):\n",
        "    val_loss = 0\n",
        "    val_per = 0\n",
        "    dis = 0\n",
        "    tot = 0 \n",
        "    outer = tqdm(total=(len(val_loader)), desc='Validation Epoch', position=0)\n",
        "    \n",
        "    with torch.no_grad():\n",
        "      encoder.eval()\n",
        "      encoder.to(DEVICE)\n",
        "      decoder.eval()\n",
        "      decoder.to(DEVICE)\n",
        "      for batch_idx, (inputs,targets,xlens,ylens) in enumerate(val_loader):\n",
        "            outer.update(1)\n",
        "            targets = targets.to(DEVICE)\n",
        "            batch_size = targets.size(0)\n",
        "            keys, values, lens, hidden = encoder(inputs.to(DEVICE),xlens.to(DEVICE),False)\n",
        "            hidden = tuple(st.transpose(0, 1).reshape(inputs.size(1), -1) for st in hidden)\n",
        "            \n",
        "            n_tokens = ylens.sum() - ylens.size(0)\n",
        "            predictions = []\n",
        "            hidden_states = [hidden, None]\n",
        "            pred= torch.zeros(batch_size,1).to(DEVICE)\n",
        "            context = values[0,:,:].to(DEVICE)\n",
        "            Y = targets.transpose(0,1)\n",
        "            loss = 0\n",
        "\n",
        "            for i in range(Y.size(0) - 1):\n",
        "\n",
        "                pred, hidden_states, context = decoder(pred.argmax(dim=1), context, keys, values,lens, hidden_states,False)\n",
        "                active = i + 1 < ylens\n",
        "                loss += criterion(pred[active], Y[i + 1, active])\n",
        "                predictions.append(pred.unsqueeze(1))\n",
        "            \n",
        "            predictions = torch.cat(predictions, dim=1)\n",
        "            loss /= n_tokens\n",
        "            current_loss = loss.item()\n",
        "            val_loss = val_loss + current_loss\n",
        "            val_per = val_per + np.exp(current_loss)\n",
        "\n",
        "            \n",
        "            for i in range(targets.shape[0]):\n",
        "                tar = ''.join(vocab[k] for k in targets[i] if k!=0 and k!=34 and k!=33)\n",
        "                predict = '0'\n",
        "                for j in range(len(predictions[i])):\n",
        "                      if predictions[i][j].argmax() == 33:\n",
        "                          break\n",
        "                      if predictions[i][j].argmax() == 34 or (predict[-1] == ' ' and predictions[i][j].argmax() == 32):\n",
        "                          pass\n",
        "                      else:\n",
        "                        predict += vocab[predictions[i][j].argmax()]\n",
        "  \n",
        "                dis = dis +  Levenshtein.distance(predict[1:].strip(), tar)\n",
        "                tot = tot + 1\n",
        "                                \n",
        "            torch.cuda.empty_cache()\n",
        "            del xlens,inputs,ylens,Y,predictions,active,loss,pred,hidden,lens,hidden_states,keys,values,targets\n",
        "    print(\"Epoch: {:02d} \\t Valid loss: {:.2f} \\t Valid Perplexity : {:.2f} \\t Valid Score : {:.2f}\".format(epoch+1,val_loss/len(val_loader),val_per/len(val_loader),dis/tot))\n",
        "    return dis/tot"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ny-ApJ_S-y_t",
        "colab_type": "text"
      },
      "source": [
        "# Decoding"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "s9zGS9E8PXs6",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def greedyinference(encoder,decoder, test_loader,epoch):\n",
        "\n",
        "    test_pred = []\n",
        "    outer = tqdm(total=(len(test_loader)), desc='Test Epoch', position=0)\n",
        "    \n",
        "    with torch.no_grad():\n",
        "      encoder.eval()\n",
        "      encoder.to(DEVICE)\n",
        "      decoder.eval()\n",
        "      decoder.to(DEVICE)\n",
        "      for batch_idx, (inputs,xlens) in enumerate(test_loader):\n",
        "            outer.update(1)\n",
        "            batch_size = len(xlens)\n",
        "            keys, values, lens, hidden = encoder(inputs.to(DEVICE),xlens.to(DEVICE),False)\n",
        "            hidden = tuple(st.transpose(0, 1).reshape(inputs.size(1), -1) for st in hidden)\n",
        "\n",
        "            predictions = []\n",
        "            hidden_states = [hidden, None]\n",
        "            pred= torch.zeros(batch_size,1).to(DEVICE)\n",
        "            context = values[0,:,:].to(DEVICE)\n",
        "\n",
        "            for i in range(250):\n",
        "\n",
        "                pred, hidden_states, context = decoder(pred.argmax(dim=1), context, keys, values,lens, hidden_states,False)\n",
        "                predictions.append(pred.unsqueeze(1))\n",
        "            \n",
        "            predictions = torch.cat(predictions, dim=1)\n",
        "            \n",
        "\n",
        "            predict = '0'\n",
        "            for j in range(len(predictions[0])):\n",
        "                  if predictions[0][j].argmax() == 33:\n",
        "                      break\n",
        "                  if predictions[0][j].argmax() == 34 or (predict[-1] == ' ' and predictions[0][j].argmax() == 32):\n",
        "                      pass\n",
        "                  else:\n",
        "                    predict += vocab[predictions[0][j].argmax()]\n",
        "            \n",
        "            test_pred.append(predict[1:].strip())            \n",
        "             \n",
        "            if batch_idx%50==0 and batch_idx>0:\n",
        "                 print(\"Batch: {:02d} \".format(batch_idx))\n",
        "            torch.cuda.empty_cache()\n",
        "            del xlens,inputs,predictions,pred,hidden,lens,hidden_states,keys,values\n",
        "    \n",
        "    return test_pred"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "c1IunH7K8ZQd",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def randominference(encoder,decoder, test_loader):\n",
        "\n",
        "    test_pred = []\n",
        "    \n",
        "    outer = tqdm(total=(len(test_loader)), desc='Test Epoch', position=0)\n",
        "    \n",
        "    with torch.no_grad():\n",
        "      encoder.eval()\n",
        "      encoder.to(DEVICE)\n",
        "      decoder.eval()\n",
        "      decoder.to(DEVICE)\n",
        "       \n",
        "      for batch_idx, (inputs,xlens) in enumerate(test_loader):\n",
        "          outer.update(1)\n",
        "          predi = []\n",
        "          losses = []\n",
        "          inputs1 = inputs.repeat(1,100,1)\n",
        "          xlens1 = xlens.repeat(100)\n",
        "          batch_size = 100\n",
        "          keys, values, lens, hidden = encoder(inputs1.to(DEVICE),xlens1.to(DEVICE))\n",
        "          hidden = tuple(st.transpose(0, 1).reshape(inputs1.size(1), -1) for st in hidden)\n",
        "\n",
        "          predictions = []\n",
        "          hidden_states = [hidden, None]\n",
        "          pred= torch.zeros(batch_size,1).to(DEVICE)\n",
        "          context = values[0,:,:].to(DEVICE)\n",
        "          del inputs1,xlens1\n",
        "\n",
        "          for i in range(250):\n",
        "\n",
        "                \n",
        "                if np.random.random_sample() > 0.5:\n",
        "                          pred = torch.distributions.gumbel.Gumbel(pred.to('cpu'), torch.tensor([0.4])).sample().to(DEVICE)\n",
        "                          pred, hidden_states, context = decoder(pred.argmax(dim=1), context, keys, values,lens, hidden_states,False)\n",
        "                else:\n",
        "                          pred, hidden_states, context = decoder(pred.argmax(dim=1), context, keys, values,lens, hidden_states,False)\n",
        "                \n",
        "                predictions.append(pred.unsqueeze(1))\n",
        "            \n",
        "          predictions = torch.cat(predictions, dim=1)\n",
        "\n",
        "          for i in range(predictions.shape[0]):\n",
        "              predict = '0'            \n",
        "              for j in range(len(predictions[i])):\n",
        "                    if predictions[i][j].argmax() == 34:\n",
        "                        break\n",
        "                    if predictions[i][j].argmax() == 33 or (predictions[i][j].argmax() == ' ' and predictions[i][j].argmax() == 32):\n",
        "                        pass\n",
        "                    else:\n",
        "                      predict += vocab[predictions[i][j].argmax()]\n",
        "              predi.append(predict[1:].strip())\n",
        "           \n",
        "          pred_tar = [torch.LongTensor([vocab.index('<sos>')] + [vocab.index(k) for k in predi[i]] + [vocab.index('<eos>')]) for i in range(len(predi))]\n",
        "                \n",
        "          for f in range(len(pred_tar)):\n",
        "                targets = pred_tar[f].unsqueeze(0).to(DEVICE)\n",
        "                ylens = targets.shape[1]\n",
        "                n_tokens = ylens \n",
        "                predictions = []\n",
        "                keys, values, lens, hidden = encoder(inputs.to(DEVICE),xlens.to(DEVICE))\n",
        "                hidden = tuple(st.transpose(0, 1).reshape(inputs.size(1), -1) for st in hidden)\n",
        "\n",
        "                hidden_states = [hidden, None]\n",
        "                pred= torch.zeros(1,1).to(DEVICE)\n",
        "                context = values[0,:,:].to(DEVICE)\n",
        "                Y = targets.transpose(0,1)\n",
        "                loss = 0\n",
        "\n",
        "                for i in range(Y.size(0) - 1):\n",
        "                    \n",
        "\n",
        "                    pred, hidden_states, context = decoder(Y[i], context, keys, values,lens, hidden_states,False)\n",
        "                    active = i + 1 < ylens\n",
        "                    loss += criterion(pred[active].squeeze(0), Y[i + 1, active].squeeze(0))\n",
        "                    predictions.append(pred.unsqueeze(1))\n",
        "                \n",
        "                predictions = torch.cat(predictions, dim=1)\n",
        "                loss /= n_tokens\n",
        "                losses.append(loss.item())\n",
        "                torch.cuda.empty_cache()\n",
        "                del ylens,Y,predictions,active,loss,pred,hidden,lens,hidden_states,keys,values,targets\n",
        "          torch.cuda.empty_cache()\n",
        "          del xlens,inputs,pred_tar\n",
        "\n",
        "          ind = np.argmin(np.array(losses))\n",
        "          test_pred.append(predi[ind])\n",
        "\n",
        "      \n",
        "          if batch_idx==1 and batch_idx>0:\n",
        "            print(batch_idx)\n",
        "            break \n",
        "                \n",
        "    \n",
        "      return test_pred,losses,predi"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XDpsgH8-Kl3F",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def beamsearch(keys, values, maxlen, hidden, beamwidth, vocab):\n",
        "        \n",
        "        init_path = [{'path':[],'score':0,'hidden':[hidden, None],'context':values[0,:,:].to(DEVICE),'current':torch.LongTensor([0]).to(DEVICE)}]\n",
        "        best_path = []\n",
        "\n",
        "        for i in range(maxlen):\n",
        "              interim = []\n",
        "              for path in init_path:\n",
        "      \n",
        "                  pred, hidden_states, context = decoder(path['current'], path['context'], keys, values,lens,  path['hidden'],False)\n",
        "                  score, ind = torch.topk(nn.functional.log_softmax(pred, dim = 1),beamwidth, dim = 1)\n",
        "\n",
        "                  for j in range(beamwidth):\n",
        "                      inter = {}\n",
        "                      inter['current'] = torch.LongTensor([ind[0,j]]).to(DEVICE)\n",
        "                      inter['score'] = path['score'] + score[0,j].cpu().detach()\n",
        "                      inter['path'] = path['path'] + [ind[0,j].cpu().detach()]\n",
        "                      inter['hidden'] = hidden_states[:]\n",
        "                      inter['context'] = context\n",
        "                      interim.append(inter)\n",
        "\n",
        "              interim = sorted(interim, key=itemgetter('score'),reverse=True)[:beamwidth]\n",
        "\n",
        "              if i== maxlen-1:\n",
        "\n",
        "                  for path in interim:\n",
        "                      path['current'] = torch.LongTensor([33]).to(DEVICE)\n",
        "\n",
        "              init_path = []\n",
        "\n",
        "              for path in interim:\n",
        "                  if path['current'] !=33:\n",
        "                    init_path.append(path)\n",
        "                  else:\n",
        "                    path['score'] = path['score']/len(path['path'])\n",
        "                    best_path.append(path)\n",
        "\n",
        "              if not init_path:\n",
        "                  break\n",
        "                \n",
        "        return convert_to_chars(sorted( best_path, key=itemgetter('score'),reverse=True)[0]['path'],vocab)\n",
        "\n",
        "def convert_to_chars(path,vocab):\n",
        "\n",
        "    pred = ''.join(vocab[k] for k in path if k!=0 and k!=34 and k!=33)\n",
        "\n",
        "    return pred"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "RXWnL9AlsFCm",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Training the models\n",
        "teacher = 0.3\n",
        "for i in range(10):\n",
        "        \n",
        "        train(encoder,decoder,train_loader,criterion, optimizer_encoder, optimizer_decoder,i,teacher)\n",
        "        score = validate(encoder,decoder, val_loader, criterion,  i)\n",
        "        torch.save(encoder.state_dict(), '/content/drive/My Drive/encoder3.pt')\n",
        "        torch.save(decoder.state_dict(), '/content/drive/My Drive/decoder3.pt')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Sfr9BcByL76z",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Decoding on Test Set\n",
        "test_pred = []\n",
        "with torch.no_grad():\n",
        "      encoder.eval()\n",
        "      encoder.to(DEVICE)\n",
        "      decoder.eval()\n",
        "      decoder.to(DEVICE)\n",
        "      for batch_idx, (inputs,xlens) in enumerate(test_loader):\n",
        "            keys, values, lens, hidden = encoder(inputs.to(DEVICE),xlens.to(DEVICE))\n",
        "            hidden = tuple(st.transpose(0, 1).reshape(inputs.size(1), -1) for st in hidden)\n",
        "            path = beamsearch(keys, values, 250, hidden, 6, vocab)\n",
        "            test_pred.append(path)\n",
        "            torch.cuda.empty_cache()\n",
        "            del xlens,inputs,keys,values,lens,hidden\n",
        "            \n",
        "            if batch_idx==50 and batch_idx>0:\n",
        "              print(batch_idx)\n",
        "            "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zz1-A5JC5Max",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import pandas as pd\n",
        "dataframe = pd.DataFrame({'Id':[i for i in range(len(test_pred))],'Predicted':test_pred})\n",
        "dataframe.to_csv(\"/content/drive/My Drive/mbarman_beam.csv\", index=False)"
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}