Transformer_chatbot.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.6.5"
    },
    "colab": {
      "name": "chatbot.ipynb",
      "provenance": []
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "code",
      "metadata": {
        "id": "_y-dMnFzDnCE",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 122
        },
        "outputId": "8170cf4e-18f2-4e27-8321-3e1318ec1512"
      },
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ],
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Go to this URL in a browser: https://accounts.google.com/o/oauth2/auth?client_id=947318989803-6bn6qk8qdgf4n4g3pfee6491hc0brc4i.apps.googleusercontent.com&redirect_uri=urn%3aietf%3awg%3aoauth%3a2.0%3aoob&response_type=code&scope=email%20https%3a%2f%2fwww.googleapis.com%2fauth%2fdocs.test%20https%3a%2f%2fwww.googleapis.com%2fauth%2fdrive%20https%3a%2f%2fwww.googleapis.com%2fauth%2fdrive.photos.readonly%20https%3a%2f%2fwww.googleapis.com%2fauth%2fpeopleapi.readonly\n",
            "\n",
            "Enter your authorization code:\n",
            "··········\n",
            "Mounted at /content/drive\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "AK3mrVEYDVHm",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "cf97c304-b4c1-44c3-c7cd-28440fd77eb9"
      },
      "source": [
        "from collections import Counter\n",
        "import json\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "from torch.utils.data import Dataset\n",
        "import torch.utils.data\n",
        "import math\n",
        "import torch.nn.functional as F\n",
        "\n",
        "torch.cuda.is_available()"
      ],
      "execution_count": 7,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "True"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 7
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mQN8iNdWDVHr",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "corpus_conv = '/content/drive/My Drive/ijcnlp_dailydialog/train/dialogues_train.txt'\n",
        "max_len = 30"
      ],
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pF_mXSA8DVH0",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "with open(corpus_conv, 'r') as c:\n",
        "    conv = c.readlines()"
      ],
      "execution_count": 9,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_GRZ2pFbDVIE",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def remove_punc(string):\n",
        "    punctuations = '''!()-[]{};:'\"\\,<>./?@#$%^&*_~'''\n",
        "    no_punct = \"\"\n",
        "    for char in string:\n",
        "        if char not in punctuations:\n",
        "            no_punct = no_punct + char  # space is also a character\n",
        "    return no_punct.lower()"
      ],
      "execution_count": 10,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xj9PPwuyDVII",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "pairs = []\n",
        "for con in conv:\n",
        "    con=con.strip('\\n')\n",
        "    values = con.split(\"__eou__\")\n",
        "    for i in range(len(values)):\n",
        "        qa_pairs = []\n",
        "        \n",
        "        if i==len(values)-1:\n",
        "            break\n",
        "        \n",
        "        first = remove_punc(values[i].strip())      \n",
        "        second = remove_punc(values[i+1].strip())\n",
        "        qa_pairs.append(first.split()[:max_len])\n",
        "        qa_pairs.append(second.split()[:max_len])\n",
        "        pairs.append(qa_pairs)       "
      ],
      "execution_count": 11,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "AsA8wtbODVIM",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "word_freq = Counter()\n",
        "for pair in pairs:\n",
        "    word_freq.update(pair[0])\n",
        "    word_freq.update(pair[1])"
      ],
      "execution_count": 12,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7UyIo9rCDVIR",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "min_word_freq = 3\n",
        "words = [w for w in word_freq.keys() if word_freq[w] > min_word_freq]\n",
        "word_map = {k: v + 1 for v, k in enumerate(words)}\n",
        "word_map['<unk>'] = len(word_map) + 1\n",
        "word_map['<start>'] = len(word_map) + 1\n",
        "word_map['<end>'] = len(word_map) + 1\n",
        "word_map['<pad>'] = 0"
      ],
      "execution_count": 16,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "b9fMYUrADVIW",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "219d2ff4-0cbb-480e-dd95-0e33a779bc87"
      },
      "source": [
        "print(\"Total words are: {}\".format(len(word_map)))"
      ],
      "execution_count": 17,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Total words are: 12207\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "35Sa0BtQDVIZ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "with open('WORDMAP_corpus.json', 'w') as j:\n",
        "    json.dump(word_map, j)"
      ],
      "execution_count": 18,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jm0Wz3iGDVIe",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def encode_question(words, word_map):\n",
        "    enc_c = [word_map.get(word, word_map['<unk>']) for word in words] + [word_map['<pad>']] * (max_len - len(words))\n",
        "    return enc_c"
      ],
      "execution_count": 19,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "b5rAx290DVIi",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def encode_reply(words, word_map):\n",
        "    enc_c = [word_map['<start>']] + [word_map.get(word, word_map['<unk>']) for word in words] + \\\n",
        "    [word_map['<end>']] + [word_map['<pad>']] * (max_len - len(words))\n",
        "    return enc_c"
      ],
      "execution_count": 20,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0cOKXxNdDVIn",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "pairs_encoded = []\n",
        "for pair in pairs:\n",
        "    qus = encode_question(pair[0], word_map)\n",
        "    ans = encode_reply(pair[1], word_map)\n",
        "    pairs_encoded.append([qus, ans])"
      ],
      "execution_count": 21,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "T_kKlbZWDVIp",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "with open('pairs_encoded.json', 'w') as p:\n",
        "    json.dump(pairs_encoded, p)"
      ],
      "execution_count": 22,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nfyD5EiMDVIu",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# rev_word_map = {v: k for k, v in word_map.items()}\n",
        "# ' '.join([rev_word_map[v] for v in pairs_encoded[1][0]])"
      ],
      "execution_count": 33,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1LAEWLVxDVIz",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class Dataset(Dataset):\n",
        "\n",
        "    def __init__(self):\n",
        "\n",
        "        self.pairs = json.load(open('pairs_encoded.json'))\n",
        "        self.dataset_size = len(self.pairs)\n",
        "\n",
        "    def __getitem__(self, i):\n",
        "        \n",
        "        question = torch.LongTensor(self.pairs[i][0])\n",
        "        reply = torch.LongTensor(self.pairs[i][1])\n",
        "            \n",
        "        return question, reply\n",
        "\n",
        "    def __len__(self):\n",
        "        return self.dataset_size"
      ],
      "execution_count": 23,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "u8X9-UnvDVI2",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "train_loader = torch.utils.data.DataLoader(Dataset(),\n",
        "                                           batch_size = 100, \n",
        "                                           shuffle=True, \n",
        "                                           pin_memory=True)"
      ],
      "execution_count": 24,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5UtZgDlUDVI8",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# question, reply = next(iter(train_loader))"
      ],
      "execution_count": 25,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ugcQb0kIDVI-",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def create_masks(question, reply_input, reply_target):\n",
        "    \n",
        "    def subsequent_mask(size):\n",
        "        mask = torch.triu(torch.ones(size, size)).transpose(0, 1).type(dtype=torch.uint8)\n",
        "        return mask.unsqueeze(0)\n",
        "    \n",
        "    question_mask = question!=0\n",
        "    question_mask = question_mask.to(device)\n",
        "    question_mask = question_mask.unsqueeze(1).unsqueeze(1)         # (batch_size, 1, 1, max_words)\n",
        "     \n",
        "    reply_input_mask = reply_input!=0\n",
        "    reply_input_mask = reply_input_mask.unsqueeze(1)  # (batch_size, 1, max_words)\n",
        "    reply_input_mask = reply_input_mask & subsequent_mask(reply_input.size(-1)).type_as(reply_input_mask.data) \n",
        "    reply_input_mask = reply_input_mask.unsqueeze(1) # (batch_size, 1, max_words, max_words)\n",
        "    reply_target_mask = reply_target!=0              # (batch_size, max_words)\n",
        "    \n",
        "    return question_mask, reply_input_mask, reply_target_mask"
      ],
      "execution_count": 26,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EW6f4bN2DVJA",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class Embeddings(nn.Module):\n",
        "    \"\"\"\n",
        "    Implements embeddings of the words and adds their positional encodings. \n",
        "    \"\"\"\n",
        "    def __init__(self, vocab_size, d_model, max_len = 50):\n",
        "        super(Embeddings, self).__init__()\n",
        "        self.d_model = d_model\n",
        "        self.dropout = nn.Dropout(0.1)\n",
        "        self.embed = nn.Embedding(vocab_size, d_model)\n",
        "        self.pe = self.create_positinal_encoding(max_len, self.d_model)\n",
        "        self.dropout = nn.Dropout(0.1)\n",
        "        \n",
        "    def create_positinal_encoding(self, max_len, d_model):\n",
        "        pe = torch.zeros(max_len, d_model).to(device)\n",
        "        for pos in range(max_len):   # for each position of the word\n",
        "            for i in range(0, d_model, 2):   # for each dimension of the each position\n",
        "                pe[pos, i] = math.sin(pos / (10000 ** ((2 * i)/d_model)))\n",
        "                pe[pos, i + 1] = math.cos(pos / (10000 ** ((2 * (i + 1))/d_model)))\n",
        "        pe = pe.unsqueeze(0)   # include the batch size\n",
        "        return pe\n",
        "        \n",
        "    def forward(self, encoded_words):\n",
        "        embedding = self.embed(encoded_words) * math.sqrt(self.d_model)\n",
        "        embedding += self.pe[:, :embedding.size(1)]   # pe will automatically be expanded with the same batch size as encoded_words\n",
        "        embedding = self.dropout(embedding)\n",
        "        return embedding"
      ],
      "execution_count": 27,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "r1xAGCHhDVJF",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class MultiHeadAttention(nn.Module):\n",
        "    \n",
        "    def __init__(self, heads, d_model):\n",
        "        \n",
        "        super(MultiHeadAttention, self).__init__()\n",
        "        assert d_model % heads == 0\n",
        "        self.d_k = d_model // heads\n",
        "        self.heads = heads\n",
        "        self.dropout = nn.Dropout(0.1)\n",
        "        self.query = nn.Linear(d_model, d_model)\n",
        "        self.key = nn.Linear(d_model, d_model)\n",
        "        self.value = nn.Linear(d_model, d_model)\n",
        "        self.concat = nn.Linear(d_model, d_model)\n",
        "        \n",
        "    def forward(self, query, key, value, mask):\n",
        "        \"\"\"\n",
        "        query, key, value of shape: (batch_size, max_len, 512)\n",
        "        mask of shape: (batch_size, 1, 1, max_words)\n",
        "        \"\"\"\n",
        "        # (batch_size, max_len, 512)\n",
        "        query = self.query(query)\n",
        "        key = self.key(key)        \n",
        "        value = self.value(value)   \n",
        "        \n",
        "        # (batch_size, max_len, 512) --> (batch_size, max_len, h, d_k) --> (batch_size, h, max_len, d_k)\n",
        "        query = query.view(query.shape[0], -1, self.heads, self.d_k).permute(0, 2, 1, 3)   \n",
        "        key = key.view(key.shape[0], -1, self.heads, self.d_k).permute(0, 2, 1, 3)  \n",
        "        value = value.view(value.shape[0], -1, self.heads, self.d_k).permute(0, 2, 1, 3)  \n",
        "        \n",
        "        # (batch_size, h, max_len, d_k) matmul (batch_size, h, d_k, max_len) --> (batch_size, h, max_len, max_len)\n",
        "        scores = torch.matmul(query, key.permute(0,1,3,2)) / math.sqrt(query.size(-1))\n",
        "        scores = scores.masked_fill(mask == 0, -1e9)    # (batch_size, h, max_len, max_len)\n",
        "        weights = F.softmax(scores, dim = -1)           # (batch_size, h, max_len, max_len)\n",
        "        weights = self.dropout(weights)\n",
        "        # (batch_size, h, max_len, max_len) matmul (batch_size, h, max_len, d_k) --> (batch_size, h, max_len, d_k)\n",
        "        context = torch.matmul(weights, value)\n",
        "        # (batch_size, h, max_len, d_k) --> (batch_size, max_len, h, d_k) --> (batch_size, max_len, h * d_k)\n",
        "        context = context.permute(0,2,1,3).contiguous().view(context.shape[0], -1, self.heads * self.d_k)\n",
        "        # (batch_size, max_len, h * d_k)\n",
        "        interacted = self.concat(context)\n",
        "        return interacted "
      ],
      "execution_count": 28,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rWlhHpNoDVJH",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class FeedForward(nn.Module):\n",
        "\n",
        "    def __init__(self, d_model, middle_dim = 2048):\n",
        "        super(FeedForward, self).__init__()\n",
        "        \n",
        "        self.fc1 = nn.Linear(d_model, middle_dim)\n",
        "        self.fc2 = nn.Linear(middle_dim, d_model)\n",
        "        self.dropout = nn.Dropout(0.1)\n",
        "\n",
        "    def forward(self, x):\n",
        "        out = F.relu(self.fc1(x))\n",
        "        out = self.fc2(self.dropout(out))\n",
        "        return out"
      ],
      "execution_count": 29,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OspsSuhZDVJK",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class EncoderLayer(nn.Module):\n",
        "\n",
        "    def __init__(self, d_model, heads):\n",
        "        super(EncoderLayer, self).__init__()\n",
        "        self.layernorm = nn.LayerNorm(d_model)\n",
        "        self.self_multihead = MultiHeadAttention(heads, d_model)\n",
        "        self.feed_forward = FeedForward(d_model)\n",
        "        self.dropout = nn.Dropout(0.1)\n",
        "\n",
        "    def forward(self, embeddings, mask):\n",
        "        interacted = self.dropout(self.self_multihead(embeddings, embeddings, embeddings, mask))\n",
        "        interacted = self.layernorm(interacted + embeddings)\n",
        "        feed_forward_out = self.dropout(self.feed_forward(interacted))\n",
        "        encoded = self.layernorm(feed_forward_out + interacted)\n",
        "        return encoded"
      ],
      "execution_count": 30,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Uc-9Wc8CDVJN",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class DecoderLayer(nn.Module):\n",
        "    \n",
        "    def __init__(self, d_model, heads):\n",
        "        super(DecoderLayer, self).__init__()\n",
        "        self.layernorm = nn.LayerNorm(d_model)\n",
        "        self.self_multihead = MultiHeadAttention(heads, d_model)\n",
        "        self.src_multihead = MultiHeadAttention(heads, d_model)\n",
        "        self.feed_forward = FeedForward(d_model)\n",
        "        self.dropout = nn.Dropout(0.1)\n",
        "        \n",
        "    def forward(self, embeddings, encoded, src_mask, target_mask):\n",
        "        query = self.dropout(self.self_multihead(embeddings, embeddings, embeddings, target_mask))\n",
        "        query = self.layernorm(query + embeddings)\n",
        "        interacted = self.dropout(self.src_multihead(query, encoded, encoded, src_mask))\n",
        "        interacted = self.layernorm(interacted + query)\n",
        "        feed_forward_out = self.dropout(self.feed_forward(interacted))\n",
        "        decoded = self.layernorm(feed_forward_out + interacted)\n",
        "        return decoded"
      ],
      "execution_count": 31,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nbiixOk0DVJR",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class Transformer(nn.Module):\n",
        "    \n",
        "    def __init__(self, d_model, heads, num_layers, word_map):\n",
        "        super(Transformer, self).__init__()\n",
        "        \n",
        "        self.d_model = d_model\n",
        "        self.vocab_size = len(word_map)\n",
        "        self.embed = Embeddings(self.vocab_size, d_model)\n",
        "        self.encoder = nn.ModuleList([EncoderLayer(d_model, heads) for _ in range(num_layers)])\n",
        "        self.decoder = nn.ModuleList([DecoderLayer(d_model, heads) for _ in range(num_layers)])\n",
        "        self.logit = nn.Linear(d_model, self.vocab_size)\n",
        "        \n",
        "    def encode(self, src_words, src_mask):\n",
        "        src_embeddings = self.embed(src_words)\n",
        "        for layer in self.encoder:\n",
        "            src_embeddings = layer(src_embeddings, src_mask)\n",
        "        return src_embeddings\n",
        "    \n",
        "    def decode(self, target_words, target_mask, src_embeddings, src_mask):\n",
        "        tgt_embeddings = self.embed(target_words)\n",
        "        for layer in self.decoder:\n",
        "            tgt_embeddings = layer(tgt_embeddings, src_embeddings, src_mask, target_mask)\n",
        "        return tgt_embeddings\n",
        "        \n",
        "    def forward(self, src_words, src_mask, target_words, target_mask):\n",
        "        encoded = self.encode(src_words, src_mask)\n",
        "        decoded = self.decode(target_words, target_mask, encoded, src_mask)\n",
        "        out = F.log_softmax(self.logit(decoded), dim = 2)\n",
        "        return out"
      ],
      "execution_count": 32,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "eMrM7zX1DVJU",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class AdamWarmup:\n",
        "    \n",
        "    def __init__(self, model_size, warmup_steps, optimizer):\n",
        "        \n",
        "        self.model_size = model_size\n",
        "        self.warmup_steps = warmup_steps\n",
        "        self.optimizer = optimizer\n",
        "        self.current_step = 0\n",
        "        self.lr = 0\n",
        "        \n",
        "    def get_lr(self):\n",
        "        return self.model_size ** (-0.5) * min(self.current_step ** (-0.5), self.current_step * self.warmup_steps ** (-1.5))\n",
        "        \n",
        "    def step(self):\n",
        "        # Increment the number of steps each time we call the step function\n",
        "        self.current_step += 1\n",
        "        lr = self.get_lr()\n",
        "        for param_group in self.optimizer.param_groups:\n",
        "            param_group['lr'] = lr\n",
        "        # update the learning rate\n",
        "        self.lr = lr\n",
        "        self.optimizer.step()       "
      ],
      "execution_count": 33,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3pEADemZDVJa",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class LossWithLS(nn.Module):\n",
        "\n",
        "    def __init__(self, size, smooth):\n",
        "        super(LossWithLS, self).__init__()\n",
        "        self.criterion = nn.KLDivLoss(size_average=False, reduce=False)\n",
        "        self.confidence = 1.0 - smooth\n",
        "        self.smooth = smooth\n",
        "        self.size = size\n",
        "        \n",
        "    def forward(self, prediction, target, mask):\n",
        "        \"\"\"\n",
        "        prediction of shape: (batch_size, max_words, vocab_size)\n",
        "        target and mask of shape: (batch_size, max_words)\n",
        "        \"\"\"\n",
        "        prediction = prediction.view(-1, prediction.size(-1))   # (batch_size * max_words, vocab_size)\n",
        "        target = target.contiguous().view(-1)   # (batch_size * max_words)\n",
        "        mask = mask.float()\n",
        "        mask = mask.view(-1)       # (batch_size * max_words)\n",
        "        labels = prediction.data.clone()\n",
        "        labels.fill_(self.smooth / (self.size - 1))\n",
        "        labels.scatter_(1, target.data.unsqueeze(1), self.confidence)\n",
        "        loss = self.criterion(prediction, labels)    # (batch_size * max_words, vocab_size)\n",
        "        loss = (loss.sum(1) * mask).sum() / mask.sum()\n",
        "        return loss"
      ],
      "execution_count": 34,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mE9LPZfQDVJc",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 71
        },
        "outputId": "6c592107-87d1-44ce-c80e-38d48e3bee0b"
      },
      "source": [
        "d_model = 512\n",
        "heads = 8\n",
        "num_layers = 3\n",
        "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "epochs = 10\n",
        "\n",
        "with open('WORDMAP_corpus.json', 'r') as j:\n",
        "    word_map = json.load(j)\n",
        "    \n",
        "transformer = Transformer(d_model = d_model, heads = heads, num_layers = num_layers, word_map = word_map)\n",
        "transformer = transformer.to(device)\n",
        "adam_optimizer = torch.optim.Adam(transformer.parameters(), lr=0, betas=(0.9, 0.98), eps=1e-9)\n",
        "transformer_optimizer = AdamWarmup(model_size = d_model, warmup_steps = 4000, optimizer = adam_optimizer)\n",
        "criterion = LossWithLS(len(word_map), 0.1)"
      ],
      "execution_count": 35,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "/usr/local/lib/python3.6/dist-packages/torch/nn/_reduction.py:43: UserWarning: size_average and reduce args will be deprecated, please use reduction='none' instead.\n",
            "  warnings.warn(warning.format(ret))\n"
          ],
          "name": "stderr"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9NH-WHyODVJe",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def train(train_loader, transformer, criterion, epoch):\n",
        "    \n",
        "    transformer.train()\n",
        "    sum_loss = 0\n",
        "    count = 0\n",
        "\n",
        "    for i, (question, reply) in enumerate(train_loader):\n",
        "        \n",
        "        samples = question.shape[0]\n",
        "\n",
        "        # Move to device\n",
        "        question = question.to(device)\n",
        "        reply = reply.to(device)\n",
        "\n",
        "        # Prepare Target Data\n",
        "        reply_input = reply[:, :-1]\n",
        "        reply_target = reply[:, 1:]\n",
        "\n",
        "        # Create mask and add dimensions\n",
        "        question_mask, reply_input_mask, reply_target_mask = create_masks(question, reply_input, reply_target)\n",
        "\n",
        "        # Get the transformer outputs\n",
        "        out = transformer(question, question_mask, reply_input, reply_input_mask)\n",
        "\n",
        "        # Compute the loss\n",
        "        loss = criterion(out, reply_target, reply_target_mask)\n",
        "        \n",
        "        # Backprop\n",
        "        transformer_optimizer.optimizer.zero_grad()\n",
        "        loss.backward()\n",
        "        transformer_optimizer.step()\n",
        "        \n",
        "        sum_loss += loss.item() * samples\n",
        "        count += samples\n",
        "        \n",
        "        if i % 100 == 0:\n",
        "            print(\"Epoch [{}][{}/{}]\\tLoss: {:.3f}\".format(epoch, i, len(train_loader), sum_loss/count))"
      ],
      "execution_count": 36,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EXIKAP5qDVJh",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def evaluate(transformer, question, question_mask, max_len, word_map):\n",
        "    \"\"\"\n",
        "    Performs Greedy Decoding with a batch size of 1\n",
        "    \"\"\"\n",
        "    rev_word_map = {v: k for k, v in word_map.items()}\n",
        "    transformer.eval()\n",
        "    start_token = word_map['<start>']\n",
        "    encoded = transformer.encode(question, question_mask)\n",
        "    words = torch.LongTensor([[start_token]]).to(device)\n",
        "    \n",
        "    for step in range(max_len - 1):\n",
        "        size = words.shape[1]\n",
        "        target_mask = torch.triu(torch.ones(size, size)).transpose(0, 1).type(dtype=torch.uint8)\n",
        "        target_mask = target_mask.to(device).unsqueeze(0).unsqueeze(0)\n",
        "        decoded = transformer.decode(words, target_mask, encoded, question_mask)\n",
        "        predictions = transformer.logit(decoded[:, -1])\n",
        "        _, next_word = torch.max(predictions, dim = 1)\n",
        "        next_word = next_word.item()\n",
        "        if next_word == word_map['<end>']:\n",
        "            break\n",
        "        words = torch.cat([words, torch.LongTensor([[next_word]]).to(device)], dim = 1)   # (1,step+2)\n",
        "        \n",
        "    # Construct Sentence\n",
        "    if words.dim() == 2:\n",
        "        words = words.squeeze(0)\n",
        "        words = words.tolist()\n",
        "        \n",
        "    sen_idx = [w for w in words if w not in {word_map['<start>']}]\n",
        "    sentence = ' '.join([rev_word_map[sen_idx[k]] for k in range(len(sen_idx))])\n",
        "    \n",
        "    return sentence"
      ],
      "execution_count": 37,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "C13mS_E4DVJk",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "outputId": "603f7129-6633-4e9b-fd2e-1f52a6113e0e"
      },
      "source": [
        "for epoch in range(epochs):\n",
        "    \n",
        "    train(train_loader, transformer, criterion, epoch)\n",
        "    \n",
        "    state = {'epoch': epoch, 'transformer': transformer, 'transformer_optimizer': transformer_optimizer}\n",
        "    torch.save(state, 'checkpoint_' + str(epoch) + '.pth.tar')"
      ],
      "execution_count": 38,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Epoch [0][0/872]\tLoss: 8.336\n",
            "Epoch [0][100/872]\tLoss: 7.556\n",
            "Epoch [0][200/872]\tLoss: 6.890\n",
            "Epoch [0][300/872]\tLoss: 6.392\n",
            "Epoch [0][400/872]\tLoss: 6.040\n",
            "Epoch [0][500/872]\tLoss: 5.788\n",
            "Epoch [0][600/872]\tLoss: 5.593\n",
            "Epoch [0][700/872]\tLoss: 5.443\n",
            "Epoch [0][800/872]\tLoss: 5.316\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "/usr/local/lib/python3.6/dist-packages/torch/serialization.py:402: UserWarning: Couldn't retrieve source code for container of type Transformer. It won't be checked for correctness upon loading.\n",
            "  \"type \" + obj.__name__ + \". It won't be checked \"\n",
            "/usr/local/lib/python3.6/dist-packages/torch/serialization.py:402: UserWarning: Couldn't retrieve source code for container of type Embeddings. It won't be checked for correctness upon loading.\n",
            "  \"type \" + obj.__name__ + \". It won't be checked \"\n",
            "/usr/local/lib/python3.6/dist-packages/torch/serialization.py:402: UserWarning: Couldn't retrieve source code for container of type EncoderLayer. It won't be checked for correctness upon loading.\n",
            "  \"type \" + obj.__name__ + \". It won't be checked \"\n",
            "/usr/local/lib/python3.6/dist-packages/torch/serialization.py:402: UserWarning: Couldn't retrieve source code for container of type MultiHeadAttention. It won't be checked for correctness upon loading.\n",
            "  \"type \" + obj.__name__ + \". It won't be checked \"\n",
            "/usr/local/lib/python3.6/dist-packages/torch/serialization.py:402: UserWarning: Couldn't retrieve source code for container of type FeedForward. It won't be checked for correctness upon loading.\n",
            "  \"type \" + obj.__name__ + \". It won't be checked \"\n",
            "/usr/local/lib/python3.6/dist-packages/torch/serialization.py:402: UserWarning: Couldn't retrieve source code for container of type DecoderLayer. It won't be checked for correctness upon loading.\n",
            "  \"type \" + obj.__name__ + \". It won't be checked \"\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "Epoch [1][0/872]\tLoss: 4.381\n",
            "Epoch [1][100/872]\tLoss: 4.270\n",
            "Epoch [1][200/872]\tLoss: 4.255\n",
            "Epoch [1][300/872]\tLoss: 4.230\n",
            "Epoch [1][400/872]\tLoss: 4.213\n",
            "Epoch [1][500/872]\tLoss: 4.192\n",
            "Epoch [1][600/872]\tLoss: 4.176\n",
            "Epoch [1][700/872]\tLoss: 4.158\n",
            "Epoch [1][800/872]\tLoss: 4.142\n",
            "Epoch [2][0/872]\tLoss: 3.932\n",
            "Epoch [2][100/872]\tLoss: 3.888\n",
            "Epoch [2][200/872]\tLoss: 3.881\n",
            "Epoch [2][300/872]\tLoss: 3.880\n",
            "Epoch [2][400/872]\tLoss: 3.878\n",
            "Epoch [2][500/872]\tLoss: 3.874\n",
            "Epoch [2][600/872]\tLoss: 3.871\n",
            "Epoch [2][700/872]\tLoss: 3.867\n",
            "Epoch [2][800/872]\tLoss: 3.864\n",
            "Epoch [3][0/872]\tLoss: 3.654\n",
            "Epoch [3][100/872]\tLoss: 3.654\n",
            "Epoch [3][200/872]\tLoss: 3.659\n",
            "Epoch [3][300/872]\tLoss: 3.668\n",
            "Epoch [3][400/872]\tLoss: 3.671\n",
            "Epoch [3][500/872]\tLoss: 3.675\n",
            "Epoch [3][600/872]\tLoss: 3.677\n",
            "Epoch [3][700/872]\tLoss: 3.680\n",
            "Epoch [3][800/872]\tLoss: 3.684\n",
            "Epoch [4][0/872]\tLoss: 3.454\n",
            "Epoch [4][100/872]\tLoss: 3.487\n",
            "Epoch [4][200/872]\tLoss: 3.515\n",
            "Epoch [4][300/872]\tLoss: 3.532\n",
            "Epoch [4][400/872]\tLoss: 3.547\n",
            "Epoch [4][500/872]\tLoss: 3.553\n",
            "Epoch [4][600/872]\tLoss: 3.560\n",
            "Epoch [4][700/872]\tLoss: 3.568\n",
            "Epoch [4][800/872]\tLoss: 3.575\n",
            "Epoch [5][0/872]\tLoss: 3.396\n",
            "Epoch [5][100/872]\tLoss: 3.376\n",
            "Epoch [5][200/872]\tLoss: 3.392\n",
            "Epoch [5][300/872]\tLoss: 3.400\n",
            "Epoch [5][400/872]\tLoss: 3.403\n",
            "Epoch [5][500/872]\tLoss: 3.411\n",
            "Epoch [5][600/872]\tLoss: 3.418\n",
            "Epoch [5][700/872]\tLoss: 3.423\n",
            "Epoch [5][800/872]\tLoss: 3.428\n",
            "Epoch [6][0/872]\tLoss: 3.081\n",
            "Epoch [6][100/872]\tLoss: 3.189\n",
            "Epoch [6][200/872]\tLoss: 3.207\n",
            "Epoch [6][300/872]\tLoss: 3.218\n",
            "Epoch [6][400/872]\tLoss: 3.230\n",
            "Epoch [6][500/872]\tLoss: 3.238\n",
            "Epoch [6][600/872]\tLoss: 3.247\n",
            "Epoch [6][700/872]\tLoss: 3.257\n",
            "Epoch [6][800/872]\tLoss: 3.263\n",
            "Epoch [7][0/872]\tLoss: 2.935\n",
            "Epoch [7][100/872]\tLoss: 3.009\n",
            "Epoch [7][200/872]\tLoss: 3.036\n",
            "Epoch [7][300/872]\tLoss: 3.053\n",
            "Epoch [7][400/872]\tLoss: 3.069\n",
            "Epoch [7][500/872]\tLoss: 3.080\n",
            "Epoch [7][600/872]\tLoss: 3.089\n",
            "Epoch [7][700/872]\tLoss: 3.100\n",
            "Epoch [7][800/872]\tLoss: 3.107\n",
            "Epoch [8][0/872]\tLoss: 2.941\n",
            "Epoch [8][100/872]\tLoss: 2.874\n",
            "Epoch [8][200/872]\tLoss: 2.889\n",
            "Epoch [8][300/872]\tLoss: 2.909\n",
            "Epoch [8][400/872]\tLoss: 2.923\n",
            "Epoch [8][500/872]\tLoss: 2.937\n",
            "Epoch [8][600/872]\tLoss: 2.950\n",
            "Epoch [8][700/872]\tLoss: 2.960\n",
            "Epoch [8][800/872]\tLoss: 2.969\n",
            "Epoch [9][0/872]\tLoss: 2.734\n",
            "Epoch [9][100/872]\tLoss: 2.716\n",
            "Epoch [9][200/872]\tLoss: 2.742\n",
            "Epoch [9][300/872]\tLoss: 2.764\n",
            "Epoch [9][400/872]\tLoss: 2.782\n",
            "Epoch [9][500/872]\tLoss: 2.799\n",
            "Epoch [9][600/872]\tLoss: 2.813\n",
            "Epoch [9][700/872]\tLoss: 2.823\n",
            "Epoch [9][800/872]\tLoss: 2.833\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GARxdSyTDVJn",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "checkpoint = torch.load('checkpoint_9.pth.tar')\n",
        "transformer = checkpoint['transformer']"
      ],
      "execution_count": 39,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IPaI39LZDVJq",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "while(1):\n",
        "    question = input(\"Question: \") \n",
        "    if question == 'quit':\n",
        "        break\n",
        "    max_len = input(\"Maximum Reply Length: \")\n",
        "    enc_qus = [word_map.get(word, word_map['<unk>']) for word in question.split()]\n",
        "    question = torch.LongTensor(enc_qus).to(device).unsqueeze(0)\n",
        "    question_mask = (question!=0).to(device).unsqueeze(1).unsqueeze(1)  \n",
        "    sentence = evaluate(transformer, question, question_mask, int(max_len), word_map)\n",
        "    print(sentence)"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}