Notebook_with_model_code_improved_for_jit_script.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Conala with original data with python embeddings.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "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/Notebook_with_model_code_improved_for_jit_script.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "FwZ6BsshCGZv",
        "outputId": "4b6b888e-6d70-48b0-f52a-e5a09e83431b"
      },
      "source": [
        "!pip install rouge"
      ],
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Requirement already satisfied: rouge in /usr/local/lib/python3.7/dist-packages (1.0.0)\n",
            "Requirement already satisfied: six in /usr/local/lib/python3.7/dist-packages (from rouge) (1.15.0)\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "h3jYaF2NxaFI"
      },
      "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\n",
        "import re\n",
        "\n",
        "from rouge import Rouge\n",
        "from typing import Optional"
      ],
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0OMTNG2PxhOq",
        "outputId": "9294ae75-e6a1-4eda-f42d-f8dae373d2c6"
      },
      "source": [
        "! git clone https://github.com/akashe/Python-Code-Generation"
      ],
      "execution_count": 3,
      "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": "rY-aGuxDxist"
      },
      "source": [
        "import sys\n",
        "sys.path.append(\"/content/Python-Code-Generation/\")"
      ],
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "O9tutyvZxkrg"
      },
      "source": [
        "from data_processing import getTokenizer,getData"
      ],
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "cTvI-Pb_xl1m",
        "outputId": "2dca3df2-96f1-4a1a-b1ad-c58b5cb704fa"
      },
      "source": [
        "!wget -c \"http://www.phontron.com/download/conala-corpus-v1.1.zip\""
      ],
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "--2021-05-28 22:07:50--  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": "pvNkOb54xmQX",
        "outputId": "c357c01f-c59e-4206-a824-2755b5ad0bde"
      },
      "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": "En3A5BQWxn9P",
        "outputId": "731801cf-940c-470d-d3fa-9ffbd0e237d0"
      },
      "source": [
        "!head -20 conala-corpus/conala-test.json"
      ],
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "[\n",
            "  {\n",
            "    \"intent\": \"How can I send a signal from a python program?\",\n",
            "    \"rewritten_intent\": \"send a signal `signal.SIGUSR1` to the current process\",\n",
            "    \"snippet\": \"os.kill(os.getpid(), signal.SIGUSR1)\",\n",
            "    \"question_id\": 15080500\n",
            "  },\n",
            "  {\n",
            "    \"intent\": \"Decode Hex String in Python 3\",\n",
            "    \"rewritten_intent\": \"decode a hex string '4a4b4c' to UTF-8.\",\n",
            "    \"snippet\": \"bytes.fromhex('4a4b4c').decode('utf-8')\",\n",
            "    \"question_id\": 3283984\n",
            "  },\n",
            "  {\n",
            "    \"intent\": \"check if all elements in a list are identical\",\n",
            "    \"rewritten_intent\": \"check if all elements in list `myList` are identical\",\n",
            "    \"snippet\": \"all(x == myList[0] for x in myList)\",\n",
            "    \"question_id\": 3844801\n",
            "  },\n",
            "  {\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zjinuEdDxr4n"
      },
      "source": [
        "questions, answers = [],[]"
      ],
      "execution_count": 9,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-PpIlzVNxtNI"
      },
      "source": [
        "# loading train json\n",
        "f = open(\"/content/conala-corpus/conala-train.json\",\"r\")\n",
        "train_file = json.load(f)\n",
        "f.close()"
      ],
      "execution_count": 10,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QriUuRU4xuiH"
      },
      "source": [
        "for num,i in enumerate(train_file):\n",
        "  if i['intent'] is not None:\n",
        "    questions.append(i['intent'])\n",
        "    answers.append(i['snippet'])"
      ],
      "execution_count": 11,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OK27P2rsxwNE"
      },
      "source": [
        "# loading train json\n",
        "f = open(\"/content/conala-corpus/conala-test.json\",\"r\")\n",
        "test_file = json.load(f)\n",
        "f.close()"
      ],
      "execution_count": 12,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WKWO3SymxxLB"
      },
      "source": [
        "for num,i in enumerate(test_file):\n",
        "  if i['intent'] is not None:\n",
        "    questions.append(i['intent'])\n",
        "    answers.append(i['snippet'])\n"
      ],
      "execution_count": 13,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "oRJSvV_w9qiY"
      },
      "source": [
        "# questions[50:70],answers[50:70]"
      ],
      "execution_count": 14,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fq185vXoxyvq"
      },
      "source": [
        "questions_, answers_ = getData(\"/content/Python-Code-Generation/data/english_python_data_pruned.txt\")"
      ],
      "execution_count": 15,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6sLzaWVxxzMe"
      },
      "source": [
        "# Setting max word len\n",
        "max_word_len = 301"
      ],
      "execution_count": 16,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "j1xBUHYmx0mF",
        "outputId": "bd102e68-b914-4038-ad34-0e1af092b760"
      },
      "source": [
        "# removing examples with len more than max_word_len\n",
        "# Only checking in original data because I checked in CONALA data and answers their dont exceed 301.\n",
        "pruned_questions = []\n",
        "pruned_answers = []\n",
        "for j,i in zip(questions_,answers_):\n",
        "  tokens = getTokenizer(i)\n",
        "  if not len(tokens) > max_word_len:\n",
        "    pruned_answers.append(i)\n",
        "    pruned_questions.append(j) \n",
        "\n",
        "print(len(pruned_answers))\n",
        "answers_ = pruned_answers\n",
        "questions_ = pruned_questions"
      ],
      "execution_count": 17,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Error in tokenization\n",
            "3058\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "t5sUoxQGx1-a"
      },
      "source": [
        "questions = questions + questions_\n",
        "answers = answers + answers_"
      ],
      "execution_count": 18,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "liZsYtBQx3U7",
        "outputId": "b56cdd26-7d25-4b0e-80df-2bfc2b9ac5c1"
      },
      "source": [
        "questions[100],answers[100]"
      ],
      "execution_count": 19,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "('Clicking a link using selenium using python',\n",
              " \"driver.find_element_by_xpath('xpath').click()\")"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 19
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "VB-Z3Yu-x4Yz",
        "outputId": "8d50b74c-0425-4074-87bf-5cf0924ec48e"
      },
      "source": [
        "print(len(questions),len(answers))"
      ],
      "execution_count": 20,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "5937 5937\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QHjZk5zxx5ik"
      },
      "source": [
        "SEED = 1327\n",
        "\n",
        "random.seed(SEED)\n",
        "np.random.seed(SEED)\n",
        "torch.manual_seed(SEED)\n",
        "torch.cuda.manual_seed(SEED)\n",
        "torch.backends.cudnn.deterministic = True"
      ],
      "execution_count": 21,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "iPECgzYox6xt",
        "outputId": "ee77787d-2522-4d67-fb12-25c3a34836b3"
      },
      "source": [
        "!python -m spacy download en"
      ],
      "execution_count": 22,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Requirement already satisfied: en_core_web_sm==2.2.5 from https://github.com/explosion/spacy-models/releases/download/en_core_web_sm-2.2.5/en_core_web_sm-2.2.5.tar.gz#egg=en_core_web_sm==2.2.5 in /usr/local/lib/python3.7/dist-packages (2.2.5)\n",
            "Requirement already satisfied: spacy>=2.2.2 in /usr/local/lib/python3.7/dist-packages (from en_core_web_sm==2.2.5) (2.2.4)\n",
            "Requirement already satisfied: wasabi<1.1.0,>=0.4.0 in /usr/local/lib/python3.7/dist-packages (from spacy>=2.2.2->en_core_web_sm==2.2.5) (0.8.2)\n",
            "Requirement already satisfied: requests<3.0.0,>=2.13.0 in /usr/local/lib/python3.7/dist-packages (from spacy>=2.2.2->en_core_web_sm==2.2.5) (2.23.0)\n",
            "Requirement already satisfied: catalogue<1.1.0,>=0.0.7 in /usr/local/lib/python3.7/dist-packages (from spacy>=2.2.2->en_core_web_sm==2.2.5) (1.0.0)\n",
            "Requirement already satisfied: preshed<3.1.0,>=3.0.2 in /usr/local/lib/python3.7/dist-packages (from spacy>=2.2.2->en_core_web_sm==2.2.5) (3.0.5)\n",
            "Requirement already satisfied: blis<0.5.0,>=0.4.0 in /usr/local/lib/python3.7/dist-packages (from spacy>=2.2.2->en_core_web_sm==2.2.5) (0.4.1)\n",
            "Requirement already satisfied: thinc==7.4.0 in /usr/local/lib/python3.7/dist-packages (from spacy>=2.2.2->en_core_web_sm==2.2.5) (7.4.0)\n",
            "Requirement already satisfied: numpy>=1.15.0 in /usr/local/lib/python3.7/dist-packages (from spacy>=2.2.2->en_core_web_sm==2.2.5) (1.19.5)\n",
            "Requirement already satisfied: plac<1.2.0,>=0.9.6 in /usr/local/lib/python3.7/dist-packages (from spacy>=2.2.2->en_core_web_sm==2.2.5) (1.1.3)\n",
            "Requirement already satisfied: srsly<1.1.0,>=1.0.2 in /usr/local/lib/python3.7/dist-packages (from spacy>=2.2.2->en_core_web_sm==2.2.5) (1.0.5)\n",
            "Requirement already satisfied: setuptools in /usr/local/lib/python3.7/dist-packages (from spacy>=2.2.2->en_core_web_sm==2.2.5) (56.1.0)\n",
            "Requirement already satisfied: tqdm<5.0.0,>=4.38.0 in /usr/local/lib/python3.7/dist-packages (from spacy>=2.2.2->en_core_web_sm==2.2.5) (4.41.1)\n",
            "Requirement already satisfied: cymem<2.1.0,>=2.0.2 in /usr/local/lib/python3.7/dist-packages (from spacy>=2.2.2->en_core_web_sm==2.2.5) (2.0.5)\n",
            "Requirement already satisfied: murmurhash<1.1.0,>=0.28.0 in /usr/local/lib/python3.7/dist-packages (from spacy>=2.2.2->en_core_web_sm==2.2.5) (1.0.5)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.7/dist-packages (from requests<3.0.0,>=2.13.0->spacy>=2.2.2->en_core_web_sm==2.2.5) (2020.12.5)\n",
            "Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.7/dist-packages (from requests<3.0.0,>=2.13.0->spacy>=2.2.2->en_core_web_sm==2.2.5) (1.24.3)\n",
            "Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.7/dist-packages (from requests<3.0.0,>=2.13.0->spacy>=2.2.2->en_core_web_sm==2.2.5) (3.0.4)\n",
            "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.7/dist-packages (from requests<3.0.0,>=2.13.0->spacy>=2.2.2->en_core_web_sm==2.2.5) (2.10)\n",
            "Requirement already satisfied: importlib-metadata>=0.20; python_version < \"3.8\" in /usr/local/lib/python3.7/dist-packages (from catalogue<1.1.0,>=0.0.7->spacy>=2.2.2->en_core_web_sm==2.2.5) (4.0.1)\n",
            "Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.7/dist-packages (from importlib-metadata>=0.20; python_version < \"3.8\"->catalogue<1.1.0,>=0.0.7->spacy>=2.2.2->en_core_web_sm==2.2.5) (3.4.1)\n",
            "Requirement already satisfied: typing-extensions>=3.6.4; python_version < \"3.8\" in /usr/local/lib/python3.7/dist-packages (from importlib-metadata>=0.20; python_version < \"3.8\"->catalogue<1.1.0,>=0.0.7->spacy>=2.2.2->en_core_web_sm==2.2.5) (3.7.4.3)\n",
            "\u001b[38;5;2m✔ Download and installation successful\u001b[0m\n",
            "You can now load the model via spacy.load('en_core_web_sm')\n",
            "\u001b[38;5;2m✔ Linking successful\u001b[0m\n",
            "/usr/local/lib/python3.7/dist-packages/en_core_web_sm -->\n",
            "/usr/local/lib/python3.7/dist-packages/spacy/data/en\n",
            "You can now load the model via spacy.load('en')\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0RKCvzMnx7_v"
      },
      "source": [
        "spacy_en = spacy.load('en')"
      ],
      "execution_count": 23,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5iT71I83x9FP"
      },
      "source": [
        "def tokenize_en(text):\n",
        "    \"\"\"\n",
        "    Tokenizes English text from a string into a list of strings\n",
        "    \"\"\"\n",
        "    return [tok.text for tok in spacy_en.tokenizer(text)]"
      ],
      "execution_count": 24,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BNLt9Bhfx-vp"
      },
      "source": [
        "SRC = Field(tokenize = tokenize_en, \n",
        "            init_token = '<sos>', \n",
        "            eos_token = '<eos>', \n",
        "            lower = True, \n",
        "            batch_first = True)\n",
        "\n",
        "TRG = Field(tokenize = getTokenizer, \n",
        "            init_token = '<sos>', \n",
        "            eos_token = '<eos>', \n",
        "            lower = False, \n",
        "            batch_first = True)"
      ],
      "execution_count": 25,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "q0dwLgsxx_Nl",
        "outputId": "65f8d1e2-539e-4bfb-854d-e43de10208d7"
      },
      "source": [
        "fields = [('src', SRC), ('trg', TRG)]\n",
        "\n",
        "Examples = [data.Example.fromlist([i,j], fields) for i,j in zip(questions,answers)]\n",
        "Dataset = data.Dataset(Examples, fields)"
      ],
      "execution_count": 26,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Error in tokenization\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bAX890DdyA6e"
      },
      "source": [
        "train_data,valid_data = Dataset.split(split_ratio=[0.90,0.10])"
      ],
      "execution_count": 27,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8jz4qsZPyBdw"
      },
      "source": [
        "SRC.build_vocab(train_data, min_freq = 1)\n",
        "TRG.build_vocab(train_data, min_freq = 1)"
      ],
      "execution_count": 28,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "-Fdl1V03yDlw",
        "outputId": "3aa9094b-fc01-48be-fe8b-c387d07cb20f"
      },
      "source": [
        "print(len(SRC.vocab))\n",
        "print(len(TRG.vocab))"
      ],
      "execution_count": 29,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "2890\n",
            "4531\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bKq0T5zAyEsH"
      },
      "source": [
        "# Dumps dicts\n",
        "with open(\"/content/SRC_stio_local\",\"wb\") as f:\n",
        "  pickle.dump(SRC.vocab.stoi,f)\n",
        "with open(\"/content/TRG_itos_local\",\"wb\") as f:\n",
        "  pickle.dump(TRG.vocab.itos,f)"
      ],
      "execution_count": 30,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "AkRq2zvqyJCK",
        "outputId": "c9637ac8-03dd-4503-d667-efee69072372"
      },
      "source": [
        "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
        "device"
      ],
      "execution_count": 31,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "device(type='cuda')"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 31
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JvDEFNU2yKph"
      },
      "source": [
        "BATCH_SIZE = 32\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.trg),\n",
        "     device = device)"
      ],
      "execution_count": 32,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fxfoYfe0yLIy"
      },
      "source": [
        "class PositionalEncodingComponent(nn.Module):\n",
        "  '''\n",
        "  Class to encode positional information to tokens.\n",
        "  \n",
        "\n",
        "  '''\n",
        "  def __init__(self,hid_dim,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 = nn.Parameter(torch.zeros(1,max_len,hid_dim),requires_grad=False)\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",
        "\n",
        "  def forward(self,x):\n",
        "    \n",
        "    x = x + self.positional_encodings[:,:x.size(1)].detach()\n",
        "    return self.dropout(x)"
      ],
      "execution_count": 33,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2ZoZTEBayMS7"
      },
      "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": 34,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cE3EQX0dyNns"
      },
      "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):\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 = nn.Parameter(torch.sqrt(torch.FloatTensor([self.head_dim])),requires_grad=False)\n",
        "\n",
        "  def forward(self,query,key,value,mask:Optional[torch.Tensor]=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": 35,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "x0o8hEs2yOvl"
      },
      "source": [
        "class EncoderLayer(nn.Module):  \n",
        "  '''\n",
        "  Operations of a single layer in an Encoder. An Encoder employs multiple such layers. Each layer contains:\n",
        "  1) multihead attention, folllowed by\n",
        "  2) LayerNorm of addition of multihead attention output and input to the layer, followed by\n",
        "  3) FeedForward connections, followed by\n",
        "  4) LayerNorm of addition of FeedForward outputs and output of previous layerNorm.\n",
        "  '''\n",
        "  def __init__(self, hid_dim,n_heads,pf_dim,dropout):\n",
        "    super().__init__()\n",
        "    \n",
        "    self.self_attn_layer_norm = nn. LayerNorm(hid_dim) #Layer norm after self-attention\n",
        "    self.ff_layer_norm = nn.LayerNorm(hid_dim) # Layer norm after FeedForward component\n",
        "\n",
        "    self.self_attention = MultiHeadedAttentionComponent(hid_dim,n_heads,dropout)\n",
        "    self.feed_forward = FeedForwardComponent(hid_dim,pf_dim,dropout)\n",
        "\n",
        "    self.dropout = nn.Dropout(dropout)\n",
        "    \n",
        "  def forward(self,src,src_mask):\n",
        "    \n",
        "    # src : [batch_size, src_len, hid_dim]\n",
        "    # src_mask : [batch_size, 1, 1, src_len]\n",
        "\n",
        "    # get self-attention\n",
        "    _src, _ = self.self_attention(src,src,src,src_mask)\n",
        "\n",
        "    # LayerNorm after dropout\n",
        "    src = self.self_attn_layer_norm(src + self.dropout(_src))\n",
        "    # src : [batch_size, src_len, hid_dim]\n",
        "\n",
        "    # FeedForward\n",
        "    _src = self.feed_forward(src)\n",
        "\n",
        "    # layerNorm after dropout\n",
        "    src = self.ff_layer_norm(src + self.dropout(_src))\n",
        "    # src: [batch_size, src_len, hid_dim]\n",
        "\n",
        "    return src"
      ],
      "execution_count": 36,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8USmzaqwyP58"
      },
      "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):\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)\n",
        "\n",
        "    # encoder attention\n",
        "    self.encoder_attention = MultiHeadedAttentionComponent(hid_dim,n_heads,dropout)\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, enc_src,trg_mask,src_mask):\n",
        "\n",
        "    #trg : [batch_size, trg_len, hid_dim]\n",
        "    #enc_src : [batch_size, src_len, hid_dim]\n",
        "    #trg_mask : [batch_size, 1, trg_len, trg_len]\n",
        "    #src_mask : [batch_size, 1, 1, src_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",
        "    '''\n",
        "    Encoder attention:\n",
        "    Query: trg\n",
        "    key: enc_src\n",
        "    Value : enc_src\n",
        "    Why? \n",
        "    the idea here is to extract information from encoder outputs. So we use decoder self-attention as a query to find important values from enc_src\n",
        "    and that is why we use src_mask, to avoid getting information from enc_src positions where it is equal to pad-id\n",
        "    After we get necessary infromation from encoder outputs we add them back to decoder self-attention.\n",
        "    '''\n",
        "    _trg, encoder_attn_alpha = self.encoder_attention(trg,enc_src,enc_src,src_mask)\n",
        "\n",
        "        # LayerNorm , residual connection and dropout\n",
        "    trg = self.enc_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, encoder_attn_alpha"
      ],
      "execution_count": 37,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2PoUKoRiyRIV"
      },
      "source": [
        "class Encoder(nn.Module):\n",
        "  '''\n",
        "  An encoder, creates token embeddings and position embeddings and passes them through multiple encoder layers\n",
        "  '''\n",
        "  def __init__(self,input_dim,hid_dim,n_layers,n_heads,pf_dim,dropout,max_length = 5000):\n",
        "    super().__init__()\n",
        "\n",
        "    self.tok_embedding = nn.Embedding(input_dim,hid_dim)\n",
        "    self.pos_embedding = PositionalEncodingComponent(hid_dim,dropout,max_length)\n",
        "\n",
        "    # encoder layers\n",
        "    self.layers = nn.ModuleList([EncoderLayer(hid_dim,n_heads,pf_dim,dropout) for _ in range(n_layers)])\n",
        "\n",
        "    self.dropout = nn.Dropout(dropout)\n",
        "\n",
        "    self.scale = nn.Parameter(torch.sqrt(torch.FloatTensor([hid_dim])),requires_grad=False)\n",
        "\n",
        "  def forward(self,src,src_mask):\n",
        "\n",
        "    # src : [batch_size, src_len]\n",
        "    # src_mask : [batch_size,1,1,src_len]\n",
        "\n",
        "    batch_size = src.shape[0]\n",
        "    src_len = src.shape[1]\n",
        "\n",
        "    tok_embeddings = self.tok_embedding(src)*self.scale\n",
        "\n",
        "    # token plus position embeddings\n",
        "    src  = self.pos_embedding(tok_embeddings)\n",
        "\n",
        "    for layer in self.layers:\n",
        "      src = layer(src,src_mask)\n",
        "    # src : [batch_size, src_len, hid_dim]\n",
        "\n",
        "    return src"
      ],
      "execution_count": 38,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "70Bg3YGPySQ4"
      },
      "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,max_length= 5000):\n",
        "    super().__init__()\n",
        "\n",
        "\n",
        "    self.tok_embedding = nn.Embedding(output_dim,hid_dim)\n",
        "    self.pos_embedding = PositionalEncodingComponent(hid_dim,dropout,max_length)\n",
        "\n",
        "    # decoder layers\n",
        "    self.layers = nn.ModuleList([DecoderLayer(hid_dim,n_heads,pf_dim,dropout) 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 = nn.Parameter(torch.sqrt(torch.FloatTensor([hid_dim])),requires_grad=False)\n",
        "\n",
        "  def forward(self, trg, enc_src,trg_mask,src_mask):\n",
        "    \n",
        "    #trg : [batch_size, trg_len]\n",
        "    #enc_src : [batch_size, src_len, hid_dim]\n",
        "    #trg_mask : [batch_size, 1, trg_len, trg_len]\n",
        "    #src_mask : [batch_size, 1, 1, src_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, encoder_attention = layer(trg,enc_src,trg_mask,src_mask)\n",
        "    \n",
        "    # trg : [batch_size,trg_len,hid_dim]\n",
        "    # encoder_attention :  [batch_size, n_head,trg_len, src_len]\n",
        "\n",
        "    # Convert to outputs\n",
        "    output = self.fc_out(trg)\n",
        "    # output : [batch_size, trg_len, output_dim]\n",
        "    \n",
        "    return output, encoder_attention"
      ],
      "execution_count": 39,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uXzefNw-yTer"
      },
      "source": [
        "\n",
        "class Seq2Seq(nn.Module):\n",
        "  def __init__(self, encoder, decoder, src_pad_idx, trg_pad_idx):\n",
        "    super().__init__()\n",
        "    self.encoder = encoder\n",
        "    self.decoder = decoder\n",
        "    self.src_pad_idx = src_pad_idx\n",
        "    self.trg_pad_idx = trg_pad_idx\n",
        "\n",
        "  def make_src_mask(self,src):\n",
        "    # src : [batch_size, src_len]\n",
        "\n",
        "    # Masking pad values\n",
        "    src_mask = (src != self.src_pad_idx).unsqueeze(1).unsqueeze(2)\n",
        "    # src_mask : [batch_size,1,1,src_len]\n",
        "\n",
        "    return src_mask\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).long()\n",
        "    # trg_pad_mask : [batch_size,1,1, trg_len]\n",
        "\n",
        "    # Masking future values\n",
        "    to_concat = [trg_pad_mask for i in range(trg.shape[-1])]\n",
        "    trg_sub_mask = torch.cat(to_concat,dim=2).tril()\n",
        "    # trg_sub_mask : [batch_size,1,trg_len, trg_len]\n",
        "\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,src,trg):\n",
        "\n",
        "    # src : [batch_size, src_len]\n",
        "    # trg : [batch_size, trg_len]\n",
        "\n",
        "    src_mask = self.make_src_mask(src)\n",
        "    trg_mask = self.make_trg_mask(trg)\n",
        "\n",
        "    # src_mask : [ batch_size, 1,1,src_len]\n",
        "    # trg_mask : [batch_size, 1, trg_len, trg_len]\n",
        "\n",
        "    enc_src = self.encoder(src,src_mask)\n",
        "    #enc_src : [batch_size, src_len, hid_dim]\n",
        "\n",
        "    output, encoder_decoder_attention = self.decoder(trg,enc_src,trg_mask,src_mask)\n",
        "    # output : [batch_size, trg_len, output_dim]\n",
        "    # encoder_decoder_attention : [batch_size, n_heads, trg_len, src_len]\n",
        "\n",
        "    return output, encoder_decoder_attention"
      ],
      "execution_count": 40,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PmQttitayUp2"
      },
      "source": [
        "INPUT_DIM = len(SRC.vocab)\n",
        "OUTPUT_DIM = len(TRG.vocab)\n",
        "HID_DIM = 256\n",
        "ENC_LAYERS = 3\n",
        "DEC_LAYERS = 3\n",
        "ENC_HEADS = 8\n",
        "DEC_HEADS = 8\n",
        "ENC_PF_DIM = 512\n",
        "DEC_PF_DIM = 512\n",
        "ENC_DROPOUT = 0.1\n",
        "DEC_DROPOUT = 0.1\n",
        "\n",
        "enc = Encoder(INPUT_DIM, \n",
        "              HID_DIM, \n",
        "              ENC_LAYERS, \n",
        "              ENC_HEADS, \n",
        "              ENC_PF_DIM, \n",
        "              ENC_DROPOUT, \n",
        "              )\n",
        "\n",
        "dec = Decoder(OUTPUT_DIM, \n",
        "              HID_DIM, \n",
        "              DEC_LAYERS, \n",
        "              DEC_HEADS, \n",
        "              DEC_PF_DIM, \n",
        "              DEC_DROPOUT, \n",
        "              )\n",
        "\n",
        "SRC_PAD_IDX = SRC.vocab.stoi[SRC.pad_token]\n",
        "TRG_PAD_IDX = TRG.vocab.stoi[TRG.pad_token]\n",
        "\n",
        "model = Seq2Seq(enc, dec, SRC_PAD_IDX, TRG_PAD_IDX).to(device)"
      ],
      "execution_count": 41,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fNZkpmLryVx7"
      },
      "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": 42,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "TM4dSNbF3U4j",
        "outputId": "5b433be4-818e-408e-b4bf-34a6a96328bf"
      },
      "source": [
        "model.decoder.tok_embedding"
      ],
      "execution_count": 43,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Embedding(4531, 256)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 43
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "lD4oJ-ch4GQi",
        "outputId": "2a649d43-a8a6-444b-f489-6309bb2279d8"
      },
      "source": [
        "# load python embedding weights\n",
        "pretrained_embeddings = torch.load('/content/Python-Code-Generation/data/python_embedding_weigts.pt')\n",
        "pretrained_embeddings"
      ],
      "execution_count": 44,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Embedding(15018, 256)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 44
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zlA3wWgrz2TY"
      },
      "source": [
        "# Loading pretrained embeddings\n",
        "\n",
        "with open(\"/content/Python-Code-Generation/data/TRG_stio\",\"rb\") as f:\n",
        "  stoi_weights = pickle.load(f)\n",
        "\n",
        "with torch.no_grad():\n",
        "  indexes = []\n",
        "  for i,j in enumerate(TRG.vocab.stoi):\n",
        "    if j in stoi_weights:\n",
        "      model.decoder.tok_embedding.weight[TRG.vocab.stoi[j]] = pretrained_embeddings.weight[stoi_weights[j]]\n",
        "    else:\n",
        "      model.decoder.tok_embedding.weight[TRG.vocab.stoi[j]] = pretrained_embeddings.weight[stoi_weights['<unk>']]\n",
        "\n"
      ],
      "execution_count": 45,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "IuLULqqFyXDk",
        "outputId": "79274848-d92d-4fe6-d353-c6ff485197fb"
      },
      "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": 46,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "The model has 7,017,907 trainable parameters\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tccCxPXjyYKO"
      },
      "source": [
        "LEARNING_RATE = 0.0001\n",
        "\n",
        "optimizer = torch.optim.Adam(model.parameters(), lr = LEARNING_RATE)\n",
        "\n",
        "criterion = nn.CrossEntropyLoss(ignore_index = TRG_PAD_IDX)"
      ],
      "execution_count": 47,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4r3UAlYGEQnc"
      },
      "source": [
        "def get_rouge_score(output,trg,vocab=TRG.vocab.itos):\n",
        "  rouge = Rouge()\n",
        "  rouge_score = 0\n",
        "  argmax_outputs = output.argmax(2)\n",
        "  assert argmax_outputs.shape == trg.shape\n",
        "  with torch.no_grad():\n",
        "    for i,row in enumerate(argmax_outputs):\n",
        "      output_sentence = []\n",
        "      trg_sentence = []\n",
        "      for j,column in enumerate(row):\n",
        "        output_sentence.append(vocab[argmax_outputs[i][j]])\n",
        "        trg_sentence.append(vocab[trg[i][j]])\n",
        "      output_sentence = \"\".join(output_sentence)\n",
        "      trg_sentence = \"\".join(trg_sentence)\n",
        "\n",
        "      rouge_score += rouge.get_scores(output_sentence, trg_sentence)[0][\"rouge-l\"][\"f\"]\n",
        "\n",
        "  return rouge_score/len(output)\n"
      ],
      "execution_count": 48,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7VGzrP2-yeg_"
      },
      "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 = batch.trg\n",
        "        \n",
        "        optimizer.zero_grad()\n",
        "        \n",
        "        output, _ = model(src, trg[:,:-1])\n",
        "                \n",
        "        #output = [batch size, trg len - 1, output dim]\n",
        "        #trg = [batch size, trg len]\n",
        "            \n",
        "        output_dim = output.shape[-1]\n",
        "            \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": 49,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "X-zlXEgDyfk0"
      },
      "source": [
        "def evaluate(model, iterator, criterion):\n",
        "    \n",
        "    model.eval()\n",
        "    \n",
        "    epoch_loss = 0\n",
        "\n",
        "    epoch_rouge = 0\n",
        "    \n",
        "    with torch.no_grad():\n",
        "    \n",
        "        for i, batch in enumerate(iterator):\n",
        "\n",
        "            src = batch.src\n",
        "            trg = batch.trg\n",
        "\n",
        "            output, _ = model(src, trg[:,:-1])\n",
        "            \n",
        "            #output = [batch size, trg len - 1, output dim]\n",
        "            #trg = [batch size, trg len]\n",
        "\n",
        "            rouge_score = get_rouge_score(output, trg[:,1:])\n",
        "\n",
        "            epoch_rouge += rouge_score\n",
        "            \n",
        "            output_dim = output.shape[-1]\n",
        "            \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",
        "        \n",
        "    return epoch_loss / len(iterator) , epoch_rouge / len(iterator)"
      ],
      "execution_count": 50,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "AiwOdJ9dygp0"
      },
      "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": 51,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "BWMHus2iyhwz",
        "outputId": "729d3e25-d518-4054-b4e3-9d9a9d197ea3"
      },
      "source": [
        "N_EPOCHS = 100\n",
        "CLIP = 1\n",
        "\n",
        "best_valid_loss = float('inf')\n",
        "best_train_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, rouge_score = 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(), 'conala_plus_original_data.pt')\n",
        "    \n",
        "    if train_loss < best_train_loss:\n",
        "        best_train_loss = train_loss\n",
        "        torch.save(model.state_dict(), 'best_train_loss.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} | Val. Rouge: {rouge_score:7.3f}')"
      ],
      "execution_count": 52,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Epoch: 01 | Time: 0m 17s\n",
            "\tTrain Loss: 5.639 | Train PPL: 281.265\n",
            "\t Val. Loss: 4.012 |  Val. PPL:  55.281 | Val. Rouge:   0.001\n",
            "Epoch: 02 | Time: 0m 17s\n",
            "\tTrain Loss: 3.629 | Train PPL:  37.678\n",
            "\t Val. Loss: 3.218 |  Val. PPL:  24.976 | Val. Rouge:   0.002\n",
            "Epoch: 03 | Time: 0m 17s\n",
            "\tTrain Loss: 3.132 | Train PPL:  22.917\n",
            "\t Val. Loss: 2.914 |  Val. PPL:  18.424 | Val. Rouge:   0.002\n",
            "Epoch: 04 | Time: 0m 17s\n",
            "\tTrain Loss: 2.891 | Train PPL:  18.007\n",
            "\t Val. Loss: 2.759 |  Val. PPL:  15.787 | Val. Rouge:   0.004\n",
            "Epoch: 05 | Time: 0m 17s\n",
            "\tTrain Loss: 2.729 | Train PPL:  15.323\n",
            "\t Val. Loss: 2.639 |  Val. PPL:  13.993 | Val. Rouge:   0.019\n",
            "Epoch: 06 | Time: 0m 17s\n",
            "\tTrain Loss: 2.604 | Train PPL:  13.521\n",
            "\t Val. Loss: 2.542 |  Val. PPL:  12.703 | Val. Rouge:   0.027\n",
            "Epoch: 07 | Time: 0m 17s\n",
            "\tTrain Loss: 2.501 | Train PPL:  12.197\n",
            "\t Val. Loss: 2.464 |  Val. PPL:  11.755 | Val. Rouge:   0.038\n",
            "Epoch: 08 | Time: 0m 17s\n",
            "\tTrain Loss: 2.410 | Train PPL:  11.139\n",
            "\t Val. Loss: 2.394 |  Val. PPL:  10.958 | Val. Rouge:   0.039\n",
            "Epoch: 09 | Time: 0m 18s\n",
            "\tTrain Loss: 2.333 | Train PPL:  10.310\n",
            "\t Val. Loss: 2.336 |  Val. PPL:  10.336 | Val. Rouge:   0.045\n",
            "Epoch: 10 | Time: 0m 18s\n",
            "\tTrain Loss: 2.261 | Train PPL:   9.595\n",
            "\t Val. Loss: 2.284 |  Val. PPL:   9.816 | Val. Rouge:   0.047\n",
            "Epoch: 11 | Time: 0m 18s\n",
            "\tTrain Loss: 2.198 | Train PPL:   9.004\n",
            "\t Val. Loss: 2.234 |  Val. PPL:   9.338 | Val. Rouge:   0.049\n",
            "Epoch: 12 | Time: 0m 18s\n",
            "\tTrain Loss: 2.136 | Train PPL:   8.468\n",
            "\t Val. Loss: 2.201 |  Val. PPL:   9.038 | Val. Rouge:   0.049\n",
            "Epoch: 13 | Time: 0m 17s\n",
            "\tTrain Loss: 2.083 | Train PPL:   8.029\n",
            "\t Val. Loss: 2.165 |  Val. PPL:   8.711 | Val. Rouge:   0.053\n",
            "Epoch: 14 | Time: 0m 18s\n",
            "\tTrain Loss: 2.034 | Train PPL:   7.643\n",
            "\t Val. Loss: 2.139 |  Val. PPL:   8.490 | Val. Rouge:   0.052\n",
            "Epoch: 15 | Time: 0m 18s\n",
            "\tTrain Loss: 1.984 | Train PPL:   7.268\n",
            "\t Val. Loss: 2.100 |  Val. PPL:   8.167 | Val. Rouge:   0.053\n",
            "Epoch: 16 | Time: 0m 18s\n",
            "\tTrain Loss: 1.937 | Train PPL:   6.939\n",
            "\t Val. Loss: 2.079 |  Val. PPL:   7.995 | Val. Rouge:   0.053\n",
            "Epoch: 17 | Time: 0m 18s\n",
            "\tTrain Loss: 1.892 | Train PPL:   6.631\n",
            "\t Val. Loss: 2.056 |  Val. PPL:   7.814 | Val. Rouge:   0.057\n",
            "Epoch: 18 | Time: 0m 17s\n",
            "\tTrain Loss: 1.850 | Train PPL:   6.362\n",
            "\t Val. Loss: 2.036 |  Val. PPL:   7.659 | Val. Rouge:   0.061\n",
            "Epoch: 19 | Time: 0m 18s\n",
            "\tTrain Loss: 1.813 | Train PPL:   6.126\n",
            "\t Val. Loss: 2.014 |  Val. PPL:   7.491 | Val. Rouge:   0.059\n",
            "Epoch: 20 | Time: 0m 17s\n",
            "\tTrain Loss: 1.774 | Train PPL:   5.893\n",
            "\t Val. Loss: 1.996 |  Val. PPL:   7.362 | Val. Rouge:   0.062\n",
            "Epoch: 21 | Time: 0m 18s\n",
            "\tTrain Loss: 1.741 | Train PPL:   5.704\n",
            "\t Val. Loss: 1.976 |  Val. PPL:   7.216 | Val. Rouge:   0.068\n",
            "Epoch: 22 | Time: 0m 17s\n",
            "\tTrain Loss: 1.703 | Train PPL:   5.490\n",
            "\t Val. Loss: 1.961 |  Val. PPL:   7.110 | Val. Rouge:   0.071\n",
            "Epoch: 23 | Time: 0m 18s\n",
            "\tTrain Loss: 1.671 | Train PPL:   5.320\n",
            "\t Val. Loss: 1.952 |  Val. PPL:   7.040 | Val. Rouge:   0.074\n",
            "Epoch: 24 | Time: 0m 17s\n",
            "\tTrain Loss: 1.633 | Train PPL:   5.118\n",
            "\t Val. Loss: 1.948 |  Val. PPL:   7.012 | Val. Rouge:   0.069\n",
            "Epoch: 25 | Time: 0m 18s\n",
            "\tTrain Loss: 1.608 | Train PPL:   4.992\n",
            "\t Val. Loss: 1.934 |  Val. PPL:   6.919 | Val. Rouge:   0.073\n",
            "Epoch: 26 | Time: 0m 18s\n",
            "\tTrain Loss: 1.574 | Train PPL:   4.828\n",
            "\t Val. Loss: 1.920 |  Val. PPL:   6.823 | Val. Rouge:   0.074\n",
            "Epoch: 27 | Time: 0m 17s\n",
            "\tTrain Loss: 1.545 | Train PPL:   4.689\n",
            "\t Val. Loss: 1.919 |  Val. PPL:   6.815 | Val. Rouge:   0.077\n",
            "Epoch: 28 | Time: 0m 17s\n",
            "\tTrain Loss: 1.516 | Train PPL:   4.554\n",
            "\t Val. Loss: 1.908 |  Val. PPL:   6.740 | Val. Rouge:   0.079\n",
            "Epoch: 29 | Time: 0m 18s\n",
            "\tTrain Loss: 1.486 | Train PPL:   4.418\n",
            "\t Val. Loss: 1.900 |  Val. PPL:   6.684 | Val. Rouge:   0.082\n",
            "Epoch: 30 | Time: 0m 17s\n",
            "\tTrain Loss: 1.461 | Train PPL:   4.309\n",
            "\t Val. Loss: 1.892 |  Val. PPL:   6.631 | Val. Rouge:   0.081\n",
            "Epoch: 31 | Time: 0m 18s\n",
            "\tTrain Loss: 1.435 | Train PPL:   4.198\n",
            "\t Val. Loss: 1.890 |  Val. PPL:   6.619 | Val. Rouge:   0.080\n",
            "Epoch: 32 | Time: 0m 17s\n",
            "\tTrain Loss: 1.409 | Train PPL:   4.094\n",
            "\t Val. Loss: 1.878 |  Val. PPL:   6.542 | Val. Rouge:   0.085\n",
            "Epoch: 33 | Time: 0m 18s\n",
            "\tTrain Loss: 1.385 | Train PPL:   3.997\n",
            "\t Val. Loss: 1.883 |  Val. PPL:   6.572 | Val. Rouge:   0.082\n",
            "Epoch: 34 | Time: 0m 18s\n",
            "\tTrain Loss: 1.356 | Train PPL:   3.879\n",
            "\t Val. Loss: 1.875 |  Val. PPL:   6.521 | Val. Rouge:   0.087\n",
            "Epoch: 35 | Time: 0m 18s\n",
            "\tTrain Loss: 1.335 | Train PPL:   3.799\n",
            "\t Val. Loss: 1.864 |  Val. PPL:   6.447 | Val. Rouge:   0.089\n",
            "Epoch: 36 | Time: 0m 18s\n",
            "\tTrain Loss: 1.313 | Train PPL:   3.717\n",
            "\t Val. Loss: 1.870 |  Val. PPL:   6.491 | Val. Rouge:   0.090\n",
            "Epoch: 37 | Time: 0m 18s\n",
            "\tTrain Loss: 1.290 | Train PPL:   3.633\n",
            "\t Val. Loss: 1.873 |  Val. PPL:   6.509 | Val. Rouge:   0.094\n",
            "Epoch: 38 | Time: 0m 18s\n",
            "\tTrain Loss: 1.265 | Train PPL:   3.543\n",
            "\t Val. Loss: 1.872 |  Val. PPL:   6.503 | Val. Rouge:   0.090\n",
            "Epoch: 39 | Time: 0m 18s\n",
            "\tTrain Loss: 1.243 | Train PPL:   3.464\n",
            "\t Val. Loss: 1.871 |  Val. PPL:   6.494 | Val. Rouge:   0.095\n",
            "Epoch: 40 | Time: 0m 18s\n",
            "\tTrain Loss: 1.226 | Train PPL:   3.407\n",
            "\t Val. Loss: 1.859 |  Val. PPL:   6.415 | Val. Rouge:   0.089\n",
            "Epoch: 41 | Time: 0m 18s\n",
            "\tTrain Loss: 1.202 | Train PPL:   3.325\n",
            "\t Val. Loss: 1.870 |  Val. PPL:   6.486 | Val. Rouge:   0.096\n",
            "Epoch: 42 | Time: 0m 18s\n",
            "\tTrain Loss: 1.180 | Train PPL:   3.256\n",
            "\t Val. Loss: 1.864 |  Val. PPL:   6.449 | Val. Rouge:   0.098\n",
            "Epoch: 43 | Time: 0m 18s\n",
            "\tTrain Loss: 1.164 | Train PPL:   3.202\n",
            "\t Val. Loss: 1.863 |  Val. PPL:   6.441 | Val. Rouge:   0.096\n",
            "Epoch: 44 | Time: 0m 18s\n",
            "\tTrain Loss: 1.147 | Train PPL:   3.147\n",
            "\t Val. Loss: 1.867 |  Val. PPL:   6.471 | Val. Rouge:   0.096\n",
            "Epoch: 45 | Time: 0m 18s\n",
            "\tTrain Loss: 1.126 | Train PPL:   3.083\n",
            "\t Val. Loss: 1.874 |  Val. PPL:   6.515 | Val. Rouge:   0.105\n",
            "Epoch: 46 | Time: 0m 18s\n",
            "\tTrain Loss: 1.108 | Train PPL:   3.027\n",
            "\t Val. Loss: 1.874 |  Val. PPL:   6.516 | Val. Rouge:   0.103\n",
            "Epoch: 47 | Time: 0m 18s\n",
            "\tTrain Loss: 1.089 | Train PPL:   2.973\n",
            "\t Val. Loss: 1.882 |  Val. PPL:   6.566 | Val. Rouge:   0.103\n",
            "Epoch: 48 | Time: 0m 18s\n",
            "\tTrain Loss: 1.074 | Train PPL:   2.927\n",
            "\t Val. Loss: 1.881 |  Val. PPL:   6.562 | Val. Rouge:   0.107\n",
            "Epoch: 49 | Time: 0m 18s\n",
            "\tTrain Loss: 1.061 | Train PPL:   2.889\n",
            "\t Val. Loss: 1.889 |  Val. PPL:   6.614 | Val. Rouge:   0.110\n",
            "Epoch: 50 | Time: 0m 18s\n",
            "\tTrain Loss: 1.042 | Train PPL:   2.836\n",
            "\t Val. Loss: 1.881 |  Val. PPL:   6.563 | Val. Rouge:   0.110\n",
            "Epoch: 51 | Time: 0m 18s\n",
            "\tTrain Loss: 1.027 | Train PPL:   2.793\n",
            "\t Val. Loss: 1.888 |  Val. PPL:   6.607 | Val. Rouge:   0.109\n",
            "Epoch: 52 | Time: 0m 18s\n",
            "\tTrain Loss: 1.009 | Train PPL:   2.742\n",
            "\t Val. Loss: 1.889 |  Val. PPL:   6.611 | Val. Rouge:   0.112\n",
            "Epoch: 53 | Time: 0m 18s\n",
            "\tTrain Loss: 0.996 | Train PPL:   2.708\n",
            "\t Val. Loss: 1.894 |  Val. PPL:   6.645 | Val. Rouge:   0.112\n",
            "Epoch: 54 | Time: 0m 18s\n",
            "\tTrain Loss: 0.983 | Train PPL:   2.673\n",
            "\t Val. Loss: 1.896 |  Val. PPL:   6.657 | Val. Rouge:   0.114\n",
            "Epoch: 55 | Time: 0m 18s\n",
            "\tTrain Loss: 0.965 | Train PPL:   2.625\n",
            "\t Val. Loss: 1.904 |  Val. PPL:   6.710 | Val. Rouge:   0.114\n",
            "Epoch: 56 | Time: 0m 18s\n",
            "\tTrain Loss: 0.954 | Train PPL:   2.597\n",
            "\t Val. Loss: 1.901 |  Val. PPL:   6.695 | Val. Rouge:   0.118\n",
            "Epoch: 57 | Time: 0m 18s\n",
            "\tTrain Loss: 0.938 | Train PPL:   2.555\n",
            "\t Val. Loss: 1.911 |  Val. PPL:   6.757 | Val. Rouge:   0.114\n",
            "Epoch: 58 | Time: 0m 18s\n",
            "\tTrain Loss: 0.924 | Train PPL:   2.519\n",
            "\t Val. Loss: 1.915 |  Val. PPL:   6.788 | Val. Rouge:   0.118\n",
            "Epoch: 59 | Time: 0m 18s\n",
            "\tTrain Loss: 0.912 | Train PPL:   2.490\n",
            "\t Val. Loss: 1.913 |  Val. PPL:   6.776 | Val. Rouge:   0.119\n",
            "Epoch: 60 | Time: 0m 18s\n",
            "\tTrain Loss: 0.898 | Train PPL:   2.454\n",
            "\t Val. Loss: 1.920 |  Val. PPL:   6.822 | Val. Rouge:   0.116\n",
            "Epoch: 61 | Time: 0m 18s\n",
            "\tTrain Loss: 0.888 | Train PPL:   2.430\n",
            "\t Val. Loss: 1.931 |  Val. PPL:   6.898 | Val. Rouge:   0.120\n",
            "Epoch: 62 | Time: 0m 18s\n",
            "\tTrain Loss: 0.874 | Train PPL:   2.397\n",
            "\t Val. Loss: 1.937 |  Val. PPL:   6.937 | Val. Rouge:   0.121\n",
            "Epoch: 63 | Time: 0m 18s\n",
            "\tTrain Loss: 0.863 | Train PPL:   2.369\n",
            "\t Val. Loss: 1.943 |  Val. PPL:   6.979 | Val. Rouge:   0.126\n",
            "Epoch: 64 | Time: 0m 18s\n",
            "\tTrain Loss: 0.850 | Train PPL:   2.340\n",
            "\t Val. Loss: 1.948 |  Val. PPL:   7.014 | Val. Rouge:   0.127\n",
            "Epoch: 65 | Time: 0m 18s\n",
            "\tTrain Loss: 0.840 | Train PPL:   2.317\n",
            "\t Val. Loss: 1.947 |  Val. PPL:   7.011 | Val. Rouge:   0.119\n",
            "Epoch: 66 | Time: 0m 18s\n",
            "\tTrain Loss: 0.831 | Train PPL:   2.296\n",
            "\t Val. Loss: 1.951 |  Val. PPL:   7.035 | Val. Rouge:   0.125\n",
            "Epoch: 67 | Time: 0m 18s\n",
            "\tTrain Loss: 0.817 | Train PPL:   2.265\n",
            "\t Val. Loss: 1.963 |  Val. PPL:   7.122 | Val. Rouge:   0.124\n",
            "Epoch: 68 | Time: 0m 18s\n",
            "\tTrain Loss: 0.810 | Train PPL:   2.248\n",
            "\t Val. Loss: 1.973 |  Val. PPL:   7.194 | Val. Rouge:   0.125\n",
            "Epoch: 69 | Time: 0m 18s\n",
            "\tTrain Loss: 0.798 | Train PPL:   2.220\n",
            "\t Val. Loss: 1.970 |  Val. PPL:   7.168 | Val. Rouge:   0.126\n",
            "Epoch: 70 | Time: 0m 18s\n",
            "\tTrain Loss: 0.786 | Train PPL:   2.196\n",
            "\t Val. Loss: 1.980 |  Val. PPL:   7.241 | Val. Rouge:   0.126\n",
            "Epoch: 71 | Time: 0m 18s\n",
            "\tTrain Loss: 0.780 | Train PPL:   2.181\n",
            "\t Val. Loss: 1.980 |  Val. PPL:   7.239 | Val. Rouge:   0.126\n",
            "Epoch: 72 | Time: 0m 18s\n",
            "\tTrain Loss: 0.764 | Train PPL:   2.148\n",
            "\t Val. Loss: 1.990 |  Val. PPL:   7.312 | Val. Rouge:   0.130\n",
            "Epoch: 73 | Time: 0m 18s\n",
            "\tTrain Loss: 0.759 | Train PPL:   2.137\n",
            "\t Val. Loss: 1.997 |  Val. PPL:   7.368 | Val. Rouge:   0.133\n",
            "Epoch: 74 | Time: 0m 18s\n",
            "\tTrain Loss: 0.750 | Train PPL:   2.117\n",
            "\t Val. Loss: 1.991 |  Val. PPL:   7.322 | Val. Rouge:   0.128\n",
            "Epoch: 75 | Time: 0m 18s\n",
            "\tTrain Loss: 0.739 | Train PPL:   2.093\n",
            "\t Val. Loss: 2.005 |  Val. PPL:   7.423 | Val. Rouge:   0.135\n",
            "Epoch: 76 | Time: 0m 18s\n",
            "\tTrain Loss: 0.732 | Train PPL:   2.078\n",
            "\t Val. Loss: 2.011 |  Val. PPL:   7.474 | Val. Rouge:   0.129\n",
            "Epoch: 77 | Time: 0m 18s\n",
            "\tTrain Loss: 0.724 | Train PPL:   2.062\n",
            "\t Val. Loss: 2.023 |  Val. PPL:   7.561 | Val. Rouge:   0.130\n",
            "Epoch: 78 | Time: 0m 18s\n",
            "\tTrain Loss: 0.715 | Train PPL:   2.045\n",
            "\t Val. Loss: 2.027 |  Val. PPL:   7.595 | Val. Rouge:   0.131\n",
            "Epoch: 79 | Time: 0m 18s\n",
            "\tTrain Loss: 0.705 | Train PPL:   2.023\n",
            "\t Val. Loss: 2.033 |  Val. PPL:   7.638 | Val. Rouge:   0.131\n",
            "Epoch: 80 | Time: 0m 18s\n",
            "\tTrain Loss: 0.701 | Train PPL:   2.015\n",
            "\t Val. Loss: 2.050 |  Val. PPL:   7.766 | Val. Rouge:   0.137\n",
            "Epoch: 81 | Time: 0m 18s\n",
            "\tTrain Loss: 0.690 | Train PPL:   1.993\n",
            "\t Val. Loss: 2.055 |  Val. PPL:   7.808 | Val. Rouge:   0.136\n",
            "Epoch: 82 | Time: 0m 18s\n",
            "\tTrain Loss: 0.680 | Train PPL:   1.974\n",
            "\t Val. Loss: 2.049 |  Val. PPL:   7.760 | Val. Rouge:   0.134\n",
            "Epoch: 83 | Time: 0m 18s\n",
            "\tTrain Loss: 0.676 | Train PPL:   1.967\n",
            "\t Val. Loss: 2.059 |  Val. PPL:   7.835 | Val. Rouge:   0.135\n",
            "Epoch: 84 | Time: 0m 18s\n",
            "\tTrain Loss: 0.664 | Train PPL:   1.942\n",
            "\t Val. Loss: 2.057 |  Val. PPL:   7.823 | Val. Rouge:   0.132\n",
            "Epoch: 85 | Time: 0m 18s\n",
            "\tTrain Loss: 0.658 | Train PPL:   1.932\n",
            "\t Val. Loss: 2.069 |  Val. PPL:   7.919 | Val. Rouge:   0.131\n",
            "Epoch: 86 | Time: 0m 18s\n",
            "\tTrain Loss: 0.652 | Train PPL:   1.919\n",
            "\t Val. Loss: 2.078 |  Val. PPL:   7.987 | Val. Rouge:   0.129\n",
            "Epoch: 87 | Time: 0m 18s\n",
            "\tTrain Loss: 0.645 | Train PPL:   1.906\n",
            "\t Val. Loss: 2.087 |  Val. PPL:   8.064 | Val. Rouge:   0.133\n",
            "Epoch: 88 | Time: 0m 18s\n",
            "\tTrain Loss: 0.636 | Train PPL:   1.889\n",
            "\t Val. Loss: 2.087 |  Val. PPL:   8.061 | Val. Rouge:   0.139\n",
            "Epoch: 89 | Time: 0m 18s\n",
            "\tTrain Loss: 0.631 | Train PPL:   1.880\n",
            "\t Val. Loss: 2.101 |  Val. PPL:   8.175 | Val. Rouge:   0.137\n",
            "Epoch: 90 | Time: 0m 18s\n",
            "\tTrain Loss: 0.620 | Train PPL:   1.858\n",
            "\t Val. Loss: 2.106 |  Val. PPL:   8.212 | Val. Rouge:   0.138\n",
            "Epoch: 91 | Time: 0m 18s\n",
            "\tTrain Loss: 0.615 | Train PPL:   1.850\n",
            "\t Val. Loss: 2.105 |  Val. PPL:   8.210 | Val. Rouge:   0.137\n",
            "Epoch: 92 | Time: 0m 17s\n",
            "\tTrain Loss: 0.611 | Train PPL:   1.843\n",
            "\t Val. Loss: 2.119 |  Val. PPL:   8.322 | Val. Rouge:   0.139\n",
            "Epoch: 93 | Time: 0m 18s\n",
            "\tTrain Loss: 0.604 | Train PPL:   1.829\n",
            "\t Val. Loss: 2.122 |  Val. PPL:   8.345 | Val. Rouge:   0.134\n",
            "Epoch: 94 | Time: 0m 18s\n",
            "\tTrain Loss: 0.599 | Train PPL:   1.820\n",
            "\t Val. Loss: 2.123 |  Val. PPL:   8.356 | Val. Rouge:   0.141\n",
            "Epoch: 95 | Time: 0m 18s\n",
            "\tTrain Loss: 0.592 | Train PPL:   1.807\n",
            "\t Val. Loss: 2.136 |  Val. PPL:   8.466 | Val. Rouge:   0.137\n",
            "Epoch: 96 | Time: 0m 17s\n",
            "\tTrain Loss: 0.583 | Train PPL:   1.791\n",
            "\t Val. Loss: 2.147 |  Val. PPL:   8.560 | Val. Rouge:   0.138\n",
            "Epoch: 97 | Time: 0m 18s\n",
            "\tTrain Loss: 0.579 | Train PPL:   1.785\n",
            "\t Val. Loss: 2.145 |  Val. PPL:   8.540 | Val. Rouge:   0.140\n",
            "Epoch: 98 | Time: 0m 18s\n",
            "\tTrain Loss: 0.574 | Train PPL:   1.775\n",
            "\t Val. Loss: 2.159 |  Val. PPL:   8.665 | Val. Rouge:   0.142\n",
            "Epoch: 99 | Time: 0m 18s\n",
            "\tTrain Loss: 0.567 | Train PPL:   1.764\n",
            "\t Val. Loss: 2.154 |  Val. PPL:   8.617 | Val. Rouge:   0.142\n",
            "Epoch: 100 | Time: 0m 18s\n",
            "\tTrain Loss: 0.561 | Train PPL:   1.753\n",
            "\t Val. Loss: 2.167 |  Val. PPL:   8.731 | Val. Rouge:   0.139\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lvjHMGS2ykod",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "f862c327-fd48-465b-ecbe-98bbd701ccd6"
      },
      "source": [
        "import pickle\n",
        "\n",
        "with open(\"SRC_stio_local\",\"rb\") as f:\n",
        "  stoi = pickle.load(f)\n",
        "with open(\"TRG_itos_local\",\"rb\") as f:\n",
        "  itos = pickle.load(f)\n",
        "\n",
        "# Load model\n",
        "# trained_model = 'conala_plus_original_data.pt'\n",
        "trained_model = '/content/best_train_loss.pt'\n",
        "\n",
        "s = torch.jit.script(model.to('cpu'))\n",
        "torch.jit.save(s,\"model_saved_by_jit.pt\")\n",
        "\n",
        "model.to(device)\n",
        "# model.load_state_dict(torch.load(trained_model));\n",
        "model.eval()"
      ],
      "execution_count": 53,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Seq2Seq(\n",
              "  (encoder): Encoder(\n",
              "    (tok_embedding): Embedding(2890, 256)\n",
              "    (pos_embedding): PositionalEncodingComponent(\n",
              "      (dropout): Dropout(p=0.1, inplace=False)\n",
              "    )\n",
              "    (layers): ModuleList(\n",
              "      (0): EncoderLayer(\n",
              "        (self_attn_layer_norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
              "        (ff_layer_norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
              "        (self_attention): MultiHeadedAttentionComponent(\n",
              "          (fc_q): Linear(in_features=256, out_features=256, bias=True)\n",
              "          (fc_k): Linear(in_features=256, out_features=256, bias=True)\n",
              "          (fc_v): Linear(in_features=256, out_features=256, bias=True)\n",
              "          (fc_o): Linear(in_features=256, out_features=256, bias=True)\n",
              "          (dropout): Dropout(p=0.1, inplace=False)\n",
              "        )\n",
              "        (feed_forward): FeedForwardComponent(\n",
              "          (dropout): Dropout(p=0.1, inplace=False)\n",
              "          (fc1): Linear(in_features=256, out_features=512, bias=True)\n",
              "          (fc2): Linear(in_features=512, out_features=256, bias=True)\n",
              "        )\n",
              "        (dropout): Dropout(p=0.1, inplace=False)\n",
              "      )\n",
              "      (1): EncoderLayer(\n",
              "        (self_attn_layer_norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
              "        (ff_layer_norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
              "        (self_attention): MultiHeadedAttentionComponent(\n",
              "          (fc_q): Linear(in_features=256, out_features=256, bias=True)\n",
              "          (fc_k): Linear(in_features=256, out_features=256, bias=True)\n",
              "          (fc_v): Linear(in_features=256, out_features=256, bias=True)\n",
              "          (fc_o): Linear(in_features=256, out_features=256, bias=True)\n",
              "          (dropout): Dropout(p=0.1, inplace=False)\n",
              "        )\n",
              "        (feed_forward): FeedForwardComponent(\n",
              "          (dropout): Dropout(p=0.1, inplace=False)\n",
              "          (fc1): Linear(in_features=256, out_features=512, bias=True)\n",
              "          (fc2): Linear(in_features=512, out_features=256, bias=True)\n",
              "        )\n",
              "        (dropout): Dropout(p=0.1, inplace=False)\n",
              "      )\n",
              "      (2): EncoderLayer(\n",
              "        (self_attn_layer_norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
              "        (ff_layer_norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
              "        (self_attention): MultiHeadedAttentionComponent(\n",
              "          (fc_q): Linear(in_features=256, out_features=256, bias=True)\n",
              "          (fc_k): Linear(in_features=256, out_features=256, bias=True)\n",
              "          (fc_v): Linear(in_features=256, out_features=256, bias=True)\n",
              "          (fc_o): Linear(in_features=256, out_features=256, bias=True)\n",
              "          (dropout): Dropout(p=0.1, inplace=False)\n",
              "        )\n",
              "        (feed_forward): FeedForwardComponent(\n",
              "          (dropout): Dropout(p=0.1, inplace=False)\n",
              "          (fc1): Linear(in_features=256, out_features=512, bias=True)\n",
              "          (fc2): Linear(in_features=512, out_features=256, bias=True)\n",
              "        )\n",
              "        (dropout): Dropout(p=0.1, inplace=False)\n",
              "      )\n",
              "    )\n",
              "    (dropout): Dropout(p=0.1, inplace=False)\n",
              "  )\n",
              "  (decoder): Decoder(\n",
              "    (tok_embedding): Embedding(4531, 256)\n",
              "    (pos_embedding): PositionalEncodingComponent(\n",
              "      (dropout): Dropout(p=0.1, inplace=False)\n",
              "    )\n",
              "    (layers): ModuleList(\n",
              "      (0): DecoderLayer(\n",
              "        (self_attn_layer_norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
              "        (enc_attn_layer_norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
              "        (ff_layer_norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
              "        (self_attention): MultiHeadedAttentionComponent(\n",
              "          (fc_q): Linear(in_features=256, out_features=256, bias=True)\n",
              "          (fc_k): Linear(in_features=256, out_features=256, bias=True)\n",
              "          (fc_v): Linear(in_features=256, out_features=256, bias=True)\n",
              "          (fc_o): Linear(in_features=256, out_features=256, bias=True)\n",
              "          (dropout): Dropout(p=0.1, inplace=False)\n",
              "        )\n",
              "        (encoder_attention): MultiHeadedAttentionComponent(\n",
              "          (fc_q): Linear(in_features=256, out_features=256, bias=True)\n",
              "          (fc_k): Linear(in_features=256, out_features=256, bias=True)\n",
              "          (fc_v): Linear(in_features=256, out_features=256, bias=True)\n",
              "          (fc_o): Linear(in_features=256, out_features=256, bias=True)\n",
              "          (dropout): Dropout(p=0.1, inplace=False)\n",
              "        )\n",
              "        (feed_forward): FeedForwardComponent(\n",
              "          (dropout): Dropout(p=0.1, inplace=False)\n",
              "          (fc1): Linear(in_features=256, out_features=512, bias=True)\n",
              "          (fc2): Linear(in_features=512, out_features=256, bias=True)\n",
              "        )\n",
              "        (dropout): Dropout(p=0.1, inplace=False)\n",
              "      )\n",
              "      (1): DecoderLayer(\n",
              "        (self_attn_layer_norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
              "        (enc_attn_layer_norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
              "        (ff_layer_norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
              "        (self_attention): MultiHeadedAttentionComponent(\n",
              "          (fc_q): Linear(in_features=256, out_features=256, bias=True)\n",
              "          (fc_k): Linear(in_features=256, out_features=256, bias=True)\n",
              "          (fc_v): Linear(in_features=256, out_features=256, bias=True)\n",
              "          (fc_o): Linear(in_features=256, out_features=256, bias=True)\n",
              "          (dropout): Dropout(p=0.1, inplace=False)\n",
              "        )\n",
              "        (encoder_attention): MultiHeadedAttentionComponent(\n",
              "          (fc_q): Linear(in_features=256, out_features=256, bias=True)\n",
              "          (fc_k): Linear(in_features=256, out_features=256, bias=True)\n",
              "          (fc_v): Linear(in_features=256, out_features=256, bias=True)\n",
              "          (fc_o): Linear(in_features=256, out_features=256, bias=True)\n",
              "          (dropout): Dropout(p=0.1, inplace=False)\n",
              "        )\n",
              "        (feed_forward): FeedForwardComponent(\n",
              "          (dropout): Dropout(p=0.1, inplace=False)\n",
              "          (fc1): Linear(in_features=256, out_features=512, bias=True)\n",
              "          (fc2): Linear(in_features=512, out_features=256, bias=True)\n",
              "        )\n",
              "        (dropout): Dropout(p=0.1, inplace=False)\n",
              "      )\n",
              "      (2): DecoderLayer(\n",
              "        (self_attn_layer_norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
              "        (enc_attn_layer_norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
              "        (ff_layer_norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
              "        (self_attention): MultiHeadedAttentionComponent(\n",
              "          (fc_q): Linear(in_features=256, out_features=256, bias=True)\n",
              "          (fc_k): Linear(in_features=256, out_features=256, bias=True)\n",
              "          (fc_v): Linear(in_features=256, out_features=256, bias=True)\n",
              "          (fc_o): Linear(in_features=256, out_features=256, bias=True)\n",
              "          (dropout): Dropout(p=0.1, inplace=False)\n",
              "        )\n",
              "        (encoder_attention): MultiHeadedAttentionComponent(\n",
              "          (fc_q): Linear(in_features=256, out_features=256, bias=True)\n",
              "          (fc_k): Linear(in_features=256, out_features=256, bias=True)\n",
              "          (fc_v): Linear(in_features=256, out_features=256, bias=True)\n",
              "          (fc_o): Linear(in_features=256, out_features=256, bias=True)\n",
              "          (dropout): Dropout(p=0.1, inplace=False)\n",
              "        )\n",
              "        (feed_forward): FeedForwardComponent(\n",
              "          (dropout): Dropout(p=0.1, inplace=False)\n",
              "          (fc1): Linear(in_features=256, out_features=512, bias=True)\n",
              "          (fc2): Linear(in_features=512, out_features=256, bias=True)\n",
              "        )\n",
              "        (dropout): Dropout(p=0.1, inplace=False)\n",
              "      )\n",
              "    )\n",
              "    (fc_out): Linear(in_features=256, out_features=4531, bias=True)\n",
              "    (dropout): Dropout(p=0.1, inplace=False)\n",
              "  )\n",
              ")"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 53
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Q0q23Cs8yo8f"
      },
      "source": [
        "import spacy\n",
        "spacy_en = spacy.load('en')\n",
        "\n",
        "def encode_inputs(input,vocab):\n",
        "\n",
        "  tokenized_input_ = [tok.text.lower() for tok in spacy_en.tokenizer(input)]\n",
        "  tokenized_input = ['<sos>'] + tokenized_input_ +['<eos>']\n",
        "\n",
        "  numericalized_input = [vocab[i] for i in tokenized_input]\n",
        "\n",
        "  tensor_input = torch.LongTensor([numericalized_input])\n",
        "  \n",
        "  return tensor_input,tokenized_input_\n",
        "\n",
        "def decode_outputs(output,vocab):\n",
        "  # output: [1,1,hid_dim]\n",
        "  predicted_token = output.argmax(-1)\n",
        "  return vocab[predicted_token.item()], predicted_token"
      ],
      "execution_count": 54,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_aM5tB6r2BcI"
      },
      "source": [
        "def display_attention(sentence, translation, attention, n_heads = 8, n_rows = 4, n_cols = 2):\n",
        "    \n",
        "    assert n_rows * n_cols == n_heads\n",
        "    \n",
        "    fig = plt.figure(figsize=(15,25))\n",
        "    \n",
        "    for i in range(n_heads):\n",
        "        \n",
        "        ax = fig.add_subplot(n_rows, n_cols, i+1)\n",
        "        \n",
        "        _attention = attention.squeeze(0)[i].cpu().detach().numpy()\n",
        "\n",
        "        cax = ax.matshow(_attention, cmap='bone')\n",
        "\n",
        "        ax.tick_params(labelsize=12)\n",
        "        ax.set_xticklabels(['']+['<sos>']+[t.lower() for t in sentence]+['<eos>'], \n",
        "                           rotation=45)\n",
        "        ax.set_yticklabels(['']+translation)\n",
        "\n",
        "        ax.xaxis.set_major_locator(ticker.MultipleLocator(1))\n",
        "        ax.yaxis.set_major_locator(ticker.MultipleLocator(1))\n",
        "\n",
        "    plt.show()"
      ],
      "execution_count": 55,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "63t2JTVt2lmA"
      },
      "source": [
        "def variables_names_in_print(matchobj):\n",
        "  statement = matchobj.group(1)\n",
        "  statement = statement.replace(\" \",\"\")\n",
        "  return \"{\"+statement+\"}\"\n",
        "\n",
        "def print_decoder_output(decoder_outputs):\n",
        "  decoder_outputs = [i for i in decoder_outputs if i is not ''] # removing redundant empty token created by tokenizer while identation during tokenization\n",
        "  combined_output = \" \".join(decoder_outputs)\n",
        "  pruned_output = re.sub(r'\\n |\\n  |\\n   ',r'\\n',combined_output) # removing empty lines\n",
        "  pruned_output = re.sub(r'{(.*?)}',variables_names_in_print,pruned_output) # setting printing variable names inside print(f'{}') statements\n",
        "  print(pruned_output)\n"
      ],
      "execution_count": 56,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "3vXUPbhHyqBV",
        "outputId": "3029a842-2873-4b89-c5d6-002df37dfa88"
      },
      "source": [
        " print(\" Enter q or quit to exit.\")\n",
        "\n",
        "answer_max_len = 200\n",
        "\n",
        "while(True):\n",
        "\n",
        "  input_ = input(\"Enter text:\")\n",
        "\n",
        "  if input_=='q' or input_=='quit':\n",
        "    break\n",
        "\n",
        "  src,tokenized_input_ = encode_inputs(input_,stoi)\n",
        "  src = src.to(device)\n",
        "  # src_mask = torch.ones([1,1,1,src.shape[-1]]).to(device)\n",
        "  src_mask = model.make_src_mask(src)\n",
        "\n",
        "  with torch.no_grad():\n",
        "    enc_src = model.encoder(src,src_mask)\n",
        "  \n",
        "  trg = '<sos>'\n",
        "  trg_indexes = [stoi[trg]]\n",
        "  # trg_mask = torch.ones([1,1,1,1]).to(device)\n",
        "\n",
        "  decoder_outputs = []\n",
        "  for i in range(answer_max_len):\n",
        "    trg_tensor = torch.LongTensor(trg_indexes).unsqueeze(0).to(device)\n",
        "    trg_mask = model.make_trg_mask(trg_tensor)\n",
        "    \n",
        "    with torch.no_grad():\n",
        "      decoder_output,encoder_decoder_attention = model.decoder(trg_tensor,enc_src,trg_mask,src_mask)\n",
        "\n",
        "    pred_token = decoder_output.argmax(2)[:,-1].item()\n",
        "\n",
        "    if pred_token == TRG.vocab.stoi[TRG.eos_token]:\n",
        "      break\n",
        "    decoder_outputs.append(itos[pred_token])\n",
        "    trg_indexes.append(pred_token)\n",
        "\n",
        "\n",
        "  print_decoder_output(decoder_outputs)\n",
        "  # display_attention(tokenized_input_,decoder_outputs,encoder_decoder_attention,DEC_HEADS)"
      ],
      "execution_count": 57,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            " Enter q or quit to exit.\n",
            "Enter text:write a Python function to convert binary to Gray codeword\n",
            "def binary_to_gray ( n ) : \n",
            "\t n = int ( n , 2 ) \n",
            "\t n ^= ( n >> 1 ) \n",
            "\t return bin ( n ) [ 2 : ]\n",
            "Enter text:q\n"
          ],
          "name": "stdout"
        }
      ]
    }
  ]
}