transformer chatbot.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "corpus_movie_conv = 'cornell movie-dialogs corpus/movie_conversations.txt'\n",
    "corpus_movie_lines = 'cornell movie-dialogs corpus/movie_lines.txt'\n",
    "max_len = 25"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(corpus_movie_conv, 'r') as c:\n",
    "    conv = c.readlines()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(corpus_movie_lines, 'r') as l:\n",
    "    lines = l.readlines()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lines_dic = {}\n",
    "for line in lines:\n",
    "    objects = line.split(\" +++$+++ \")\n",
    "    lines_dic[objects[0]] = objects[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pairs = []\n",
    "for con in conv:\n",
    "    ids = eval(con.split(\" +++$+++ \")[-1])\n",
    "    for i in range(len(ids)):\n",
    "        qa_pairs = []\n",
    "        \n",
    "        if i==len(ids)-1:\n",
    "            break\n",
    "        \n",
    "        first = remove_punc(lines_dic[ids[i]].strip())      \n",
    "        second = remove_punc(lines_dic[ids[i+1]].strip())\n",
    "        qa_pairs.append(first.split()[:max_len])\n",
    "        qa_pairs.append(second.split()[:max_len])\n",
    "        pairs.append(qa_pairs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "word_freq = Counter()\n",
    "for pair in pairs:\n",
    "    word_freq.update(pair[0])\n",
    "    word_freq.update(pair[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "min_word_freq = 5\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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Total words are: {}\".format(len(word_map)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('WORDMAP_corpus.json', 'w') as j:\n",
    "    json.dump(word_map, j)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('pairs_encoded.json', 'w') as p:\n",
    "    json.dump(pairs_encoded, p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_loader = torch.utils.data.DataLoader(Dataset(),\n",
    "                                           batch_size = 100, \n",
    "                                           shuffle=True, \n",
    "                                           pin_memory=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# question, reply = next(iter(train_loader))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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()       "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint = torch.load('checkpoint.pth.tar')\n",
    "transformer = checkpoint['transformer']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}