Python_Embeddings_on_CoNaLa_mined_data.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Python Embeddings on CoNaLa mined data.ipynb",
      "provenance": [],
      "authorship_tag": "ABX9TyMou7fnLd45bZrQY4XLmHm2",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/akashe/Python-Code-Generation/blob/main/Python_Embeddings_on_CoNaLa_mined_data.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FgIWQ4-wITMz"
      },
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim\n",
        "\n",
        "import torchtext\n",
        "from torchtext.legacy import data\n",
        "from torchtext.legacy.data import Field, BucketIterator\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "import matplotlib.ticker as ticker\n",
        "\n",
        "import spacy\n",
        "import numpy as np\n",
        "\n",
        "import random\n",
        "import math\n",
        "import time\n",
        "import pickle\n",
        "import json"
      ],
      "execution_count": 1,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "cAHy9n_wIhpc",
        "outputId": "1ecff6ee-f60b-4ed7-fc91-701e72ffca50"
      },
      "source": [
        "! git clone https://github.com/akashe/Python-Code-Generation"
      ],
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "fatal: destination path 'Python-Code-Generation' already exists and is not an empty directory.\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1DBw1iz_IiOX"
      },
      "source": [
        "import sys\n",
        "sys.path.append(\"/content/Python-Code-Generation/\")"
      ],
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xNoN8DivIlqU"
      },
      "source": [
        "from data_processing import getTokenizer"
      ],
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8DTA_S7VIoVM"
      },
      "source": [
        "# checking CoNALA data"
      ],
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "je1nsFwXIqsS",
        "outputId": "64a10dbb-4b98-4a73-dabd-e5981b9e30ba"
      },
      "source": [
        "!wget -c \"http://www.phontron.com/download/conala-corpus-v1.1.zip\""
      ],
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "--2021-03-12 11:25:40--  http://www.phontron.com/download/conala-corpus-v1.1.zip\n",
            "Resolving www.phontron.com (www.phontron.com)... 208.113.196.149\n",
            "Connecting to www.phontron.com (www.phontron.com)|208.113.196.149|:80... connected.\n",
            "HTTP request sent, awaiting response... 416 Requested Range Not Satisfiable\n",
            "\n",
            "    The file is already fully retrieved; nothing to do.\n",
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "lrQIODl9I5r7",
        "outputId": "3a767771-4d1a-4fdc-e637-2be08b26d354"
      },
      "source": [
        "!unzip -n conala-corpus-v1.1.zip"
      ],
      "execution_count": 7,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Archive:  conala-corpus-v1.1.zip\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "cy9yNVjhI_xd",
        "outputId": "28adeb42-6650-488d-f213-3ed4f81dd3bc"
      },
      "source": [
        "!head -2 conala-corpus/conala-mined.jsonl"
      ],
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "{\"parent_answer_post_id\": 34705233, \"prob\": 0.8690001442846342, \"snippet\": \"sorted(l, key=lambda x: (-int(x[1]), x[0]))\", \"intent\": \"Sort a nested list by two elements\", \"id\": \"34705205_34705233_0\", \"question_id\": 34705205}\n",
            "{\"parent_answer_post_id\": 13905946, \"prob\": 0.8526701436370034, \"snippet\": \"[int(x) for x in str(num)]\", \"intent\": \"converting integer to list in python\", \"id\": \"13905936_13905946_0\", \"question_id\": 13905936}\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "JzDAx4piJQyx",
        "outputId": "b4b55b4d-03b6-45f9-ddb2-6fe59f36dac7"
      },
      "source": [
        "!head -20 conala-corpus/conala-train.json"
      ],
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "[\n",
            "  {\n",
            "    \"intent\": \"How to convert a list of multiple integers into a single integer?\",\n",
            "    \"rewritten_intent\": \"Concatenate elements of a list 'x' of multiple integers to a single integer\",\n",
            "    \"snippet\": \"sum(d * 10 ** i for i, d in enumerate(x[::-1]))\",\n",
            "    \"question_id\": 41067960\n",
            "  },\n",
            "  {\n",
            "    \"intent\": \"How to convert a list of multiple integers into a single integer?\",\n",
            "    \"rewritten_intent\": \"convert a list of integers into a single integer\",\n",
            "    \"snippet\": \"r = int(''.join(map(str, x)))\",\n",
            "    \"question_id\": 41067960\n",
            "  },\n",
            "  {\n",
            "    \"intent\": \"how to convert a datetime string back to datetime object?\",\n",
            "    \"rewritten_intent\": \"convert a DateTime string back to a DateTime object of format '%Y-%m-%d %H:%M:%S.%f'\",\n",
            "    \"snippet\": \"datetime.strptime('2010-11-13 10:33:54.227806', '%Y-%m-%d %H:%M:%S.%f')\",\n",
            "    \"question_id\": 4170655\n",
            "  },\n",
            "  {\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "L_HrLxZvKwiL"
      },
      "source": [
        "questions, answers = [],[]"
      ],
      "execution_count": 10,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qv5pZG-iJzVQ"
      },
      "source": [
        "# using examples from conala-mined with probability greater than 80%\n",
        "prob = 0.0\n",
        "f = open(\"/content/conala-corpus/conala-mined.jsonl\")\n",
        "for i,line in enumerate(f):\n",
        "  dict_ = json.loads(line)\n",
        "  if dict_['prob'] > prob:\n",
        "    questions.append(dict_['intent'])\n",
        "    answers.append(dict_['snippet'])\n"
      ],
      "execution_count": 11,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "V2ZtBx1TLY6p",
        "outputId": "9510459c-cfcb-481d-9d8e-8729ed430e86"
      },
      "source": [
        "print(f'Total examples from mined examples with probability grater than {prob*100}% is {len(questions)}')"
      ],
      "execution_count": 12,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Total examples from mined examples with probability grater than 0.0% is 593891\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5WJzDN1_R8ZK"
      },
      "source": [
        "assert len(questions)==len(answers)"
      ],
      "execution_count": 13,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8xBx_Xn3SBep"
      },
      "source": [
        "# Setting max word len\n",
        "max_word_len = 301"
      ],
      "execution_count": 14,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "f1M5E97Ch7sk",
        "outputId": "4450ffd0-1653-4e87-fa1d-b797b4f5b05a"
      },
      "source": [
        "# removing examples with len more than max_word_len\n",
        "pruned_answers = []\n",
        "for i in answers:\n",
        "  tokens = getTokenizer(i)\n",
        "  if not len(tokens) > max_word_len:\n",
        "    pruned_answers.append(tokens)\n",
        "\n",
        "print(len(pruned_answers))\n",
        "answers = pruned_answers"
      ],
      "execution_count": 15,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "590763\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1YW-NSQ_SKiW"
      },
      "source": [
        "SRC = Field(tokenize = None, \n",
        "            init_token = '<sos>', \n",
        "            eos_token = '<eos>', \n",
        "            lower = False, \n",
        "            batch_first = True)"
      ],
      "execution_count": 16,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "W_jPPOZKSlYq"
      },
      "source": [
        "fields = [('src', SRC)]\n",
        "\n",
        "Examples = [data.Example.fromlist([i], fields) for i in answers]"
      ],
      "execution_count": 17,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "sqGR23KQTO8u",
        "outputId": "816b623e-4464-4d7c-d1af-9440130b4388"
      },
      "source": [
        "print(len(Examples))\n",
        "vars(Examples[0])"
      ],
      "execution_count": 18,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "590763\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'src': ['sorted',\n",
              "  '(',\n",
              "  'l',\n",
              "  ',',\n",
              "  'key',\n",
              "  '=',\n",
              "  'lambda',\n",
              "  'x',\n",
              "  ':',\n",
              "  '(',\n",
              "  '-',\n",
              "  'int',\n",
              "  '(',\n",
              "  'x',\n",
              "  '[',\n",
              "  '1',\n",
              "  ']',\n",
              "  ')',\n",
              "  ',',\n",
              "  'x',\n",
              "  '[',\n",
              "  '0',\n",
              "  ']',\n",
              "  ')',\n",
              "  ')',\n",
              "  '']}"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 18
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vFNtIDtoTxbK"
      },
      "source": [
        "Dataset = data.Dataset(Examples, fields)"
      ],
      "execution_count": 19,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "x9x2angfU0S2"
      },
      "source": [
        "train_data,valid_data = Dataset.split(split_ratio=[0.90,0.10])"
      ],
      "execution_count": 20,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "iZrGCWLyU3wl"
      },
      "source": [
        "SRC.build_vocab(train_data, min_freq = 15)"
      ],
      "execution_count": 21,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "T4sEEXl-VUMl",
        "outputId": "c8c439df-3b62-4a55-a61b-f46beb6f4162"
      },
      "source": [
        "len(SRC.vocab)"
      ],
      "execution_count": 22,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "15018"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 22
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cKg4NTTFVsW_"
      },
      "source": [
        "# Dumps dicts\n",
        "with open(\"/content/SRC_stio\",\"wb\") as f:\n",
        "  pickle.dump(SRC.vocab.stoi,f)\n",
        "with open(\"/content/SRC_itos\",\"wb\") as f:\n",
        "  pickle.dump(SRC.vocab.itos,f)"
      ],
      "execution_count": 23,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "WpKzCCI9VxRG",
        "outputId": "670218d1-6a19-4f4e-c2b1-2f6c1b34c3a9"
      },
      "source": [
        "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
        "device"
      ],
      "execution_count": 24,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "device(type='cuda')"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 24
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xWYuJxZnV0jj"
      },
      "source": [
        "BATCH_SIZE = 64\n",
        "\n",
        "train_iterator, valid_iterator = BucketIterator.splits(\n",
        "    (train_data, valid_data), \n",
        "     batch_size = BATCH_SIZE,\n",
        "     sort_key = lambda x: len(x.src),\n",
        "     device = device)"
      ],
      "execution_count": 25,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yMKj10-6V2gv"
      },
      "source": [
        "class PositionalEncodingComponent(nn.Module):\n",
        "  '''\n",
        "  Class to encode positional information to tokens.\n",
        "  \n",
        "\n",
        "  '''\n",
        "  def __init__(self,hid_dim,device,dropout=0.2,max_len=5000):\n",
        "    super().__init__()\n",
        "\n",
        "    assert hid_dim%2==0 # If not, it will result error in allocation to positional_encodings[:,1::2] later\n",
        "\n",
        "    self.dropout = nn.Dropout(dropout)\n",
        "\n",
        "    self.positional_encodings = torch.zeros(max_len,hid_dim)\n",
        "\n",
        "    pos = torch.arange(0,max_len).unsqueeze(1) # pos : [max_len,1]\n",
        "    div_term  = torch.exp(-torch.arange(0,hid_dim,2)*math.log(10000.0)/hid_dim) # Calculating value of 1/(10000^(2i/hid_dim)) in log space and then exponentiating it\n",
        "    # div_term: [hid_dim//2]\n",
        "\n",
        "    self.positional_encodings[:,0::2] = torch.sin(pos*div_term) # pos*div_term [max_len,hid_dim//2]\n",
        "    self.positional_encodings[:,1::2] = torch.cos(pos*div_term) \n",
        "\n",
        "    self.positional_encodings = self.positional_encodings.unsqueeze(0) # To account for batch_size in inputs\n",
        "\n",
        "    self.device = device\n",
        "\n",
        "  def forward(self,x):\n",
        "    x = x + self.positional_encodings[:,:x.size(1)].detach().to(self.device)\n",
        "    return self.dropout(x)"
      ],
      "execution_count": 26,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Zodr8VcRWHVK"
      },
      "source": [
        "class FeedForwardComponent(nn.Module):\n",
        "  '''\n",
        "  Class for pointwise feed forward connections\n",
        "  '''\n",
        "  def __init__(self,hid_dim,pf_dim,dropout):\n",
        "    super().__init__()\n",
        "\n",
        "    self.dropout = nn.Dropout(dropout)\n",
        "\n",
        "    self.fc1 = nn.Linear(hid_dim,pf_dim)\n",
        "    self.fc2 = nn.Linear(pf_dim,hid_dim)\n",
        "\n",
        "  def forward(self,x):\n",
        "\n",
        "    # x : [batch_size,seq_len,hid_dim]\n",
        "    x = self.dropout(torch.relu(self.fc1(x)))\n",
        "\n",
        "    # x : [batch_size,seq_len,pf_dim]\n",
        "    x = self.fc2(x)\n",
        "\n",
        "    # x : [batch_size,seq_len,hid_dim]\n",
        "    return x"
      ],
      "execution_count": 27,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3Puau8kvWIsA"
      },
      "source": [
        "class MultiHeadedAttentionComponent(nn.Module):\n",
        "  '''\n",
        "  Multiheaded attention Component. This implementation also supports mask. \n",
        "  The reason for mask that in Decoder, we don't want attention mechanism to get\n",
        "  important information from future tokens.\n",
        "  '''\n",
        "  def __init__(self,hid_dim, n_heads, dropout, device):\n",
        "    super().__init__()\n",
        "\n",
        "    assert hid_dim % n_heads == 0 # Since we split hid_dims into n_heads\n",
        "\n",
        "    self.hid_dim = hid_dim\n",
        "    self.n_heads = n_heads # no of heads in 'multiheaded' attention\n",
        "    self.head_dim = hid_dim//n_heads # dims of each head\n",
        "\n",
        "    # Transformation from source vector to query vector\n",
        "    self.fc_q = nn.Linear(hid_dim,hid_dim)\n",
        "\n",
        "    # Transformation from source vector to key vector\n",
        "    self.fc_k = nn.Linear(hid_dim,hid_dim)\n",
        "\n",
        "    # Transformation from source vector to value vector\n",
        "    self.fc_v = nn.Linear(hid_dim,hid_dim)\n",
        "\n",
        "    self.fc_o = nn.Linear(hid_dim,hid_dim)\n",
        "\n",
        "    self.dropout = nn.Dropout(dropout)\n",
        "\n",
        "    # Used in self attention for smoother gradients\n",
        "    self.scale = torch.sqrt(torch.FloatTensor([self.head_dim])).to(device)\n",
        "\n",
        "  def forward(self,query,key,value,mask=None):\n",
        "\n",
        "    #query : [batch_size, query_len, hid_dim]\n",
        "    #key : [batch_size, key_len, hid_dim]\n",
        "    #value : [batch_size, value_len, hid_dim]\n",
        "\n",
        "    batch_size = query.shape[0]\n",
        "\n",
        "    # Transforming quey,key,values\n",
        "    Q = self.fc_q(query)\n",
        "    K = self.fc_k(key)\n",
        "    V = self.fc_v(value)\n",
        "\n",
        "    #Q : [batch_size, query_len, hid_dim]\n",
        "    #K : [batch_size, key_len, hid_dim]\n",
        "    #V : [batch_size, value_len,hid_dim]\n",
        "\n",
        "    # Changing shapes to acocmadate n_heads information\n",
        "    Q = Q.view(batch_size, -1, self.n_heads, self.head_dim).permute(0, 2, 1, 3)\n",
        "    K = K.view(batch_size, -1, self.n_heads, self.head_dim).permute(0, 2, 1, 3)\n",
        "    V = V.view(batch_size, -1, self.n_heads, self.head_dim).permute(0, 2, 1, 3)\n",
        "\n",
        "    #Q : [batch_size, n_heads, query_len, head_dim]\n",
        "    #K : [batch_size, n_heads, key_len, head_dim]\n",
        "    #V : [batch_size, n_heads, value_len, head_dim]\n",
        "\n",
        "    # Calculating alpha\n",
        "    score = torch.matmul(Q,K.permute(0,1,3,2))/self.scale\n",
        "    # score : [batch_size, n_heads, query_len, key_len]\n",
        "\n",
        "    if mask is not None:\n",
        "      score = score.masked_fill(mask==0,-1e10)\n",
        "\n",
        "    alpha = torch.softmax(score,dim=-1)\n",
        "    # alpha : [batch_size, n_heads, query_len, key_len]\n",
        "\n",
        "    # Get the final self-attention  vector\n",
        "    x = torch.matmul(self.dropout(alpha),V)\n",
        "    # x : [batch_size, n_heads, query_len, head_dim]\n",
        "\n",
        "    # Reshaping self attention vector to concatenate\n",
        "    x = x.permute(0,2,1,3).contiguous()\n",
        "    # x : [batch_size, query_len, n_heads, head_dim]\n",
        "\n",
        "    x = x.view(batch_size,-1,self.hid_dim)\n",
        "    # x: [batch_size, query_len, hid_dim]\n",
        "\n",
        "    # Transforming concatenated outputs \n",
        "    x = self.fc_o(x)\n",
        "    #x : [batch_size, query_len, hid_dim] \n",
        "\n",
        "    return x, alpha"
      ],
      "execution_count": 28,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XrAlZqsGWM-e"
      },
      "source": [
        "class DecoderLayer(nn.Module):\n",
        "  '''\n",
        "  Operations of a single layer in an Decoder. An Decoder employs multiple such layers. Each layer contains:\n",
        "  1) masked decoder self attention, followed by\n",
        "  2) LayerNorm of addition of previous attention output and input to the layer,, followed by\n",
        "  3) encoder self attention, followed by\n",
        "  4) LayerNorm of addition of result of encoder self attention and its input, followed by\n",
        "  5) FeedForward connections, followed by\n",
        "  6) LayerNorm of addition of Feedforward results and its input.\n",
        "  '''\n",
        "  def __init__(self,hid_dim,n_heads,pf_dim,dropout,device):\n",
        "    super().__init__()\n",
        "\n",
        "    self.self_attn_layer_norm = nn.LayerNorm(hid_dim)\n",
        "    self.enc_attn_layer_norm = nn.LayerNorm(hid_dim)\n",
        "    self.ff_layer_norm = nn.LayerNorm(hid_dim)\n",
        "\n",
        "    # decoder self attention\n",
        "    self.self_attention = MultiHeadedAttentionComponent(hid_dim,n_heads,dropout,device)\n",
        "\n",
        "    # FeedForward\n",
        "    self.feed_forward = FeedForwardComponent(hid_dim,pf_dim,dropout)\n",
        "\n",
        "    self.dropout = nn.Dropout(dropout)\n",
        "\n",
        "  def forward(self,trg,trg_mask):\n",
        "\n",
        "    #trg : [batch_size, trg_len, hid_dim]\n",
        "    #trg_mask : [batch_size, 1, trg_len, trg_len]\n",
        "\n",
        "    '''\n",
        "    Decoder self-attention\n",
        "    trg_mask is to force decoder to look only into past tokens and not get information from future tokens.\n",
        "    Since we apply mask before doing softmax, the final self attention vector gets no information from future tokens.\n",
        "    '''\n",
        "    _trg, _ = self.self_attention(trg,trg,trg,trg_mask)\n",
        "\n",
        "    # LayerNorm and dropout with resdiual connection\n",
        "    trg = self.self_attn_layer_norm(trg + self.dropout(_trg))\n",
        "    # trg : [batch_size, trg_len, hid_dim]\n",
        "\n",
        "    # Feed Forward\n",
        "    _trg = self.feed_forward(trg)\n",
        "\n",
        "    # LayerNorm, residual connection and dropout\n",
        "    trg = self.ff_layer_norm(trg + self.dropout(_trg))\n",
        "\n",
        "    return trg"
      ],
      "execution_count": 29,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WGEweAFDWwuz"
      },
      "source": [
        "class Decoder(nn.Module):\n",
        "  '''\n",
        "  An decoder, creates token embeddings and position embeddings and passes them through multiple decoder layers\n",
        "  '''\n",
        "  def __init__(self,output_dim,hid_dim,n_layers,n_heads,pf_dim,dropout,device,max_length= 350):\n",
        "    super().__init__()\n",
        "\n",
        "    self.device = device\n",
        "\n",
        "    self.tok_embedding = nn.Embedding(output_dim,hid_dim)\n",
        "    self.pos_embedding = PositionalEncodingComponent(hid_dim,device,dropout,max_length)\n",
        "\n",
        "    # decoder layers\n",
        "    self.layers = nn.ModuleList([DecoderLayer(hid_dim,n_heads,pf_dim,dropout,device) for _ in range(n_layers)])\n",
        "\n",
        "    # convert decoder outputs to real outputs\n",
        "    self.fc_out = nn.Linear(hid_dim,output_dim)\n",
        "\n",
        "    self.dropout = nn.Dropout(dropout)\n",
        "\n",
        "    self.scale = torch.sqrt(torch.FloatTensor([hid_dim])).to(device)\n",
        "\n",
        "  def forward(self, trg,trg_mask):\n",
        "    \n",
        "    #trg : [batch_size, trg_len]\n",
        "    #trg_mask : [batch_size, 1, trg_len, trg_len]\n",
        "\n",
        "    batch_size = trg.shape[0]\n",
        "    trg_len = trg.shape[1]\n",
        "\n",
        "    tok_embeddings = self.tok_embedding(trg)*self.scale\n",
        "\n",
        "    # token plus pos embeddings\n",
        "    trg = self.pos_embedding(tok_embeddings)\n",
        "    # trg : [batch_size, trg_len, hid_dim]\n",
        "\n",
        "    # Pass trg thorugh decoder layers\n",
        "    for layer in self.layers:\n",
        "      trg= layer(trg,trg_mask)\n",
        "    \n",
        "    # trg : [batch_size,trg_len,hid_dim]\n",
        "\n",
        "    # Convert to outputs\n",
        "    output = self.fc_out(trg)\n",
        "    # output : [batch_size, trg_len, output_dim]\n",
        "    \n",
        "    return output"
      ],
      "execution_count": 30,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Vu30TD_UXJe9"
      },
      "source": [
        "class TrainEmbeddings(nn.Module):\n",
        "  def __init__(self, decoder, trg_pad_idx, device):\n",
        "    super().__init__()\n",
        "    self.decoder = decoder\n",
        "    self.trg_pad_idx = trg_pad_idx\n",
        "    self.device = device\n",
        "\n",
        "  def make_trg_mask(self,trg):\n",
        "    # trg : [batch_size, trg_len]\n",
        "\n",
        "    # Masking pad values\n",
        "    trg_pad_mask = (trg != self.trg_pad_idx).unsqueeze(1).unsqueeze(2)\n",
        "    # trg_pad_mask : [batch_size,1,1, trg_len]\n",
        "\n",
        "    # Masking future values\n",
        "    trg_len = trg.shape[1]\n",
        "    trg_sub_mask = torch.tril(torch.ones((trg_len,trg_len),device= self.device)).bool()\n",
        "    # trg_sub_mask : [trg_len, trg_len]\n",
        "\n",
        "    # combine both masks\n",
        "    trg_mask = trg_pad_mask & trg_sub_mask\n",
        "    # trg_mask = [batch_size,1,trg_len,trg_len]\n",
        "\n",
        "    return trg_mask\n",
        "\n",
        "  def forward(self,trg):\n",
        "\n",
        "    # trg : [batch_size, trg_len]\n",
        "\n",
        "    trg_mask = self.make_trg_mask(trg)\n",
        "\n",
        "    # trg_mask : [batch_size, 1, trg_len, trg_len]\n",
        "\n",
        "    output = self.decoder(trg,trg_mask)\n",
        "    # output : [batch_size, trg_len, output_dim]\n",
        "\n",
        "    return output"
      ],
      "execution_count": 31,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "k__zm3vtXi9R"
      },
      "source": [
        "INPUT_DIM = len(SRC.vocab)\n",
        "OUTPUT_DIM = len(SRC.vocab)\n",
        "HID_DIM = 256\n",
        "DEC_LAYERS = 2\n",
        "DEC_HEADS = 8\n",
        "DEC_PF_DIM = 256\n",
        "DEC_DROPOUT = 0.1\n",
        "\n",
        "dec = Decoder(OUTPUT_DIM, \n",
        "              HID_DIM, \n",
        "              DEC_LAYERS, \n",
        "              DEC_HEADS, \n",
        "              DEC_PF_DIM, \n",
        "              DEC_DROPOUT, \n",
        "              device)\n",
        "\n",
        "TRG_PAD_IDX = SRC.vocab.stoi[SRC.pad_token]\n",
        "\n",
        "model = TrainEmbeddings( dec, TRG_PAD_IDX, device).to(device)"
      ],
      "execution_count": 32,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JOQOUbj8YYbF"
      },
      "source": [
        "def initialize_weights(m):\n",
        "    if hasattr(m, 'weight') and m.weight.dim() > 1:\n",
        "        nn.init.xavier_uniform_(m.weight.data)\n",
        "\n",
        "model.apply(initialize_weights);"
      ],
      "execution_count": 33,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-CCVYvjrJhJM"
      },
      "source": [
        "# # Loading trained model\n",
        "# trained_model = 'python_embedding.pt'\n",
        "# model.load_state_dict(torch.load(trained_model))"
      ],
      "execution_count": 34,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "YelNW8qSYaIm",
        "outputId": "9b1c8341-55ea-4924-bc43-f09a4da9b7b0"
      },
      "source": [
        "def count_parameters(model):\n",
        "    return sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
        "\n",
        "print(f'The model has {count_parameters(model):,} trainable parameters')"
      ],
      "execution_count": 35,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "The model has 8,496,810 trainable parameters\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NtXhm3icYbXs"
      },
      "source": [
        "LEARNING_RATE = 0.0005\n",
        "\n",
        "optimizer = torch.optim.Adam(model.parameters(), lr = LEARNING_RATE)\n",
        "\n",
        "criterion = nn.CrossEntropyLoss(ignore_index = TRG_PAD_IDX)"
      ],
      "execution_count": 36,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PbRTCE5PYd2q"
      },
      "source": [
        "def train(model, iterator, optimizer, criterion, clip):\n",
        "    \n",
        "    model.train()\n",
        "    \n",
        "    epoch_loss = 0\n",
        "    \n",
        "    for i, batch in enumerate(iterator):\n",
        "        \n",
        "        src = batch.src\n",
        "        trg = src\n",
        "        \n",
        "        optimizer.zero_grad()\n",
        "        \n",
        "        output = model(src)\n",
        "                \n",
        "        #output = [batch size, trg len, output dim]\n",
        "        #trg = [batch size, trg len]\n",
        "   \n",
        "        output_dim = output.shape[-1]\n",
        "\n",
        "        output = output[:,:-1,:]             \n",
        "        output = output.contiguous().view(-1, output_dim)\n",
        "        trg = trg[:,1:].contiguous().view(-1)\n",
        "                \n",
        "        #output = [batch size * trg len - 1, output dim]\n",
        "        #trg = [batch size * trg len - 1]\n",
        "            \n",
        "        loss = criterion(output, trg)\n",
        "        \n",
        "        loss.backward()\n",
        "\n",
        "        torch.nn.utils.clip_grad_norm_(model.parameters(), clip)\n",
        "        \n",
        "        optimizer.step()\n",
        "        \n",
        "        epoch_loss += loss.item()\n",
        "        \n",
        "    return epoch_loss / len(iterator)"
      ],
      "execution_count": 37,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "z9F_cdugaaau"
      },
      "source": [
        "def evaluate(model, iterator, criterion):\n",
        "    \n",
        "    model.eval()\n",
        "    \n",
        "    epoch_loss = 0\n",
        "    \n",
        "    with torch.no_grad():\n",
        "    \n",
        "        for i, batch in enumerate(iterator):\n",
        "\n",
        "            src = batch.src\n",
        "            trg = src\n",
        "\n",
        "            output = model(src)\n",
        "            \n",
        "            #output = [batch size, trg len, output dim]\n",
        "            #trg = [batch size, trg len]\n",
        "            \n",
        "            output_dim = output.shape[-1]\n",
        "            \n",
        "            output = output[:,:-1,:] \n",
        "            output = output.contiguous().view(-1, output_dim)\n",
        "            trg = trg[:,1:].contiguous().view(-1)\n",
        "            \n",
        "            #output = [batch size * trg len - 1, output dim]\n",
        "            #trg = [batch size * trg len - 1]\n",
        "            \n",
        "            loss = criterion(output, trg)\n",
        "\n",
        "            epoch_loss += loss.item()\n",
        "        \n",
        "    return epoch_loss / len(iterator)"
      ],
      "execution_count": 38,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uBqbA4ZOYfEA"
      },
      "source": [
        "def epoch_time(start_time, end_time):\n",
        "    elapsed_time = end_time - start_time\n",
        "    elapsed_mins = int(elapsed_time / 60)\n",
        "    elapsed_secs = int(elapsed_time - (elapsed_mins * 60))\n",
        "    return elapsed_mins, elapsed_secs"
      ],
      "execution_count": 39,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-YbVpDHvYgOS",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "c32f8bd1-a2be-481e-94dd-95bf3f796e58"
      },
      "source": [
        "N_EPOCHS = 10\n",
        "CLIP = 1\n",
        "\n",
        "best_valid_loss = float('inf')\n",
        "\n",
        "for epoch in range(N_EPOCHS):\n",
        "    \n",
        "    start_time = time.time()\n",
        "    \n",
        "    train_loss = train(model, train_iterator, optimizer, criterion, CLIP)\n",
        "    valid_loss = evaluate(model, valid_iterator, criterion)\n",
        "    \n",
        "    end_time = time.time()\n",
        "    \n",
        "    epoch_mins, epoch_secs = epoch_time(start_time, end_time)\n",
        "    \n",
        "    if valid_loss < best_valid_loss:\n",
        "        best_valid_loss = valid_loss\n",
        "        torch.save(model.state_dict(), 'python_embedding.pt')\n",
        "    \n",
        "    print(f'Epoch: {epoch+1:02} | Time: {epoch_mins}m {epoch_secs}s')\n",
        "    print(f'\\tTrain Loss: {train_loss:.3f} | Train PPL: {math.exp(train_loss):7.3f}')\n",
        "    print(f'\\t Val. Loss: {valid_loss:.3f} |  Val. PPL: {math.exp(valid_loss):7.3f}')"
      ],
      "execution_count": 40,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Epoch: 01 | Time: 28m 58s\n",
            "\tTrain Loss: 2.216 | Train PPL:   9.173\n",
            "\t Val. Loss: 1.825 |  Val. PPL:   6.203\n",
            "Epoch: 02 | Time: 28m 59s\n",
            "\tTrain Loss: 1.776 | Train PPL:   5.907\n",
            "\t Val. Loss: 1.722 |  Val. PPL:   5.595\n",
            "Epoch: 03 | Time: 29m 0s\n",
            "\tTrain Loss: 1.668 | Train PPL:   5.299\n",
            "\t Val. Loss: 1.676 |  Val. PPL:   5.342\n",
            "Epoch: 04 | Time: 29m 2s\n",
            "\tTrain Loss: 1.611 | Train PPL:   5.008\n",
            "\t Val. Loss: 1.646 |  Val. PPL:   5.186\n",
            "Epoch: 05 | Time: 29m 1s\n",
            "\tTrain Loss: 1.576 | Train PPL:   4.834\n",
            "\t Val. Loss: 1.631 |  Val. PPL:   5.110\n",
            "Epoch: 06 | Time: 28m 58s\n",
            "\tTrain Loss: 1.550 | Train PPL:   4.710\n",
            "\t Val. Loss: 1.617 |  Val. PPL:   5.035\n",
            "Epoch: 07 | Time: 29m 5s\n",
            "\tTrain Loss: 1.530 | Train PPL:   4.619\n",
            "\t Val. Loss: 1.604 |  Val. PPL:   4.972\n",
            "Epoch: 08 | Time: 28m 59s\n",
            "\tTrain Loss: 1.515 | Train PPL:   4.548\n",
            "\t Val. Loss: 1.598 |  Val. PPL:   4.944\n",
            "Epoch: 09 | Time: 28m 58s\n",
            "\tTrain Loss: 1.503 | Train PPL:   4.493\n",
            "\t Val. Loss: 1.588 |  Val. PPL:   4.895\n",
            "Epoch: 10 | Time: 29m 3s\n",
            "\tTrain Loss: 1.492 | Train PPL:   4.446\n",
            "\t Val. Loss: 1.583 |  Val. PPL:   4.869\n"
          ],
          "name": "stdout"
        }
      ]
    }
  ]
}