chatbot.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch import optim\n",
    "import torch.nn.functional as F\n",
    "import csv\n",
    "import random\n",
    "import re\n",
    "import os\n",
    "import unicodedata\n",
    "import codecs\n",
    "import itertools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "CUDA = torch.cuda.is_available()\n",
    "device = torch.device(\"cuda\" if CUDA else \"cpu\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data Preprocessing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "corpus_name = \"cornell movie-dialogs corpus\"\n",
    "corpus = os.path.join(\"data\", corpus_name)\n",
    "lines_filepath = os.path.join(corpus,\"movie_lines.txt\")\n",
    "conv_filepath = os.path.join(corpus,\"movie_conversations.txt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "L1045 +++$+++ u0 +++$+++ m0 +++$+++ BIANCA +++$+++ They do not!\n",
      "L1044 +++$+++ u2 +++$+++ m0 +++$+++ CAMERON +++$+++ They do to!\n",
      "L985 +++$+++ u0 +++$+++ m0 +++$+++ BIANCA +++$+++ I hope so.\n",
      "L984 +++$+++ u2 +++$+++ m0 +++$+++ CAMERON +++$+++ She okay?\n",
      "L925 +++$+++ u0 +++$+++ m0 +++$+++ BIANCA +++$+++ Let's go.\n",
      "L924 +++$+++ u2 +++$+++ m0 +++$+++ CAMERON +++$+++ Wow\n",
      "L872 +++$+++ u0 +++$+++ m0 +++$+++ BIANCA +++$+++ Okay -- you're gonna need to learn how to lie.\n",
      "L871 +++$+++ u2 +++$+++ m0 +++$+++ CAMERON +++$+++ No\n",
      "L870 +++$+++ u0 +++$+++ m0 +++$+++ BIANCA +++$+++ I'm kidding.  You know how sometimes you just become this \"persona\"?  And you don't know how to quit?\n",
      "L869 +++$+++ u0 +++$+++ m0 +++$+++ BIANCA +++$+++ Like my fear of wearing pastels?\n"
     ]
    }
   ],
   "source": [
    "# visualizing some lines\n",
    "with open(lines_filepath,'r') as file:\n",
    "    lines = file.readlines()\n",
    "for line in lines[:10]:\n",
    "    print(line.strip())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "# splits each line of the file into a dictionary of fields(lineID,characterID,movieID,character,text)\n",
    "line_fields = [\"lineID\",\"characterID\",\"movieID\",\"character\",\"text\"]\n",
    "lines = {}\n",
    "\n",
    "with open(lines_filepath, 'r', encoding='iso-8859-1') as f:\n",
    "    for line in f:\n",
    "        values = line.split(\" +++$+++ \")\n",
    "        # Extract fields\n",
    "        lineObj = {}\n",
    "        for i, field in enumerate(line_fields):\n",
    "            lineObj[field] = values[i]\n",
    "        lines[lineObj['lineID']] = lineObj"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Grouping fields of lines from the above loaded lines into conversation based on \"movie_conversations.txt\"\n",
    "\n",
    "conv_fields = [\"character1ID\", \"character2ID\", \"movieID\", \"utteranceIDs\"]\n",
    "conversations = []\n",
    "\n",
    "with open(conv_filepath, 'r', encoding='iso-8859-1') as f:\n",
    "    for line in f:\n",
    "        values = line.split(\" +++$+++ \")\n",
    "        # Extract fields\n",
    "        convObj = {}\n",
    "        for i, field in enumerate(conv_fields):\n",
    "            convObj[field] = values[i]\n",
    "        lineIds = eval(convObj[\"utteranceIDs\"])\n",
    "        # Reassemble lines\n",
    "        convObj[\"lines\"] = []\n",
    "        for lineId in lineIds:\n",
    "            convObj[\"lines\"].append(lines[lineId])\n",
    "        conversations.append(convObj)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Extracts pairs of sentences from conversations\n",
    "qa_pairs = []\n",
    "for conversation in conversations:\n",
    "    # Iterate over all the lines of the conversation\n",
    "    for i in range(len(conversation[\"lines\"]) - 1):  # We ignore the last line (no answer for it)\n",
    "        inputLine = conversation[\"lines\"][i][\"text\"].strip()\n",
    "        targetLine = conversation[\"lines\"][i+1][\"text\"].strip()\n",
    "        # Filter wrong samples (if one of the lists is empty)\n",
    "        if inputLine and targetLine:\n",
    "            qa_pairs.append([inputLine, targetLine])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Writing newly formatted file...\n"
     ]
    }
   ],
   "source": [
    "# Define path to new file\n",
    "datafile = os.path.join(corpus, \"formatted_movie_lines.txt\")\n",
    "\n",
    "delimiter = '\\t'\n",
    "# Unescape the delimiter\n",
    "delimiter = str(codecs.decode(delimiter, \"unicode_escape\"))\n",
    "\n",
    "# Writing the conversational pairs into new csv file\n",
    "print(\"\\nWriting newly formatted file...\")\n",
    "with open(datafile, 'w', encoding='utf-8') as outputfile:\n",
    "    writer = csv.writer(outputfile, delimiter=delimiter)\n",
    "    for pair in qa_pairs:\n",
    "        writer.writerow(pair)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "b\"Can we make this quick?  Roxanne Korrine and Andrew Barrett are having an incredibly horrendous public break- up on the quad.  Again.\\tWell, I thought we'd start with pronunciation, if that's okay with you.\"\n",
      "b\"Well, I thought we'd start with pronunciation, if that's okay with you.\\tNot the hacking and gagging and spitting part.  Please.\"\n",
      "b\"Not the hacking and gagging and spitting part.  Please.\\tOkay... then how 'bout we try out some French cuisine.  Saturday?  Night?\"\n",
      "b\"You're asking me out.  That's so cute. What's your name again?\\tForget it.\"\n",
      "b\"No, no, it's my fault -- we didn't have a proper introduction ---\\tCameron.\"\n",
      "b\"Cameron.\\tThe thing is, Cameron -- I'm at the mercy of a particularly hideous breed of loser.  My sister.  I can't date until she does.\"\n",
      "b\"The thing is, Cameron -- I'm at the mercy of a particularly hideous breed of loser.  My sister.  I can't date until she does.\\tSeems like she could get a date easy enough...\"\n",
      "b'Why?\\tUnsolved mystery.  She used to be really popular when she started high school, then it was just like she got sick of it or something.'\n",
      "b\"Unsolved mystery.  She used to be really popular when she started high school, then it was just like she got sick of it or something.\\tThat's a shame.\"\n",
      "b'Gosh, if only we could find Kat a boyfriend...\\tLet me see what I can do.'\n"
     ]
    }
   ],
   "source": [
    "# visualizing some lines\n",
    "with open(datafile,'rb') as file:\n",
    "    lines = file.readlines()\n",
    "for line in lines[:10]:\n",
    "    print(line.strip())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Default word tokens\n",
    "PAD_token = 0  # Used for padding short sentences\n",
    "SOS_token = 1  # Start-of-sentence token\n",
    "EOS_token = 2  # End-of-sentence token\n",
    "\n",
    "class Vocabulary:\n",
    "    def __init__(self, name):\n",
    "        self.name = name\n",
    "        self.trimmed = False\n",
    "        self.word2index = {}\n",
    "        self.word2count = {}\n",
    "        self.index2word = {PAD_token: \"PAD\", SOS_token: \"SOS\", EOS_token: \"EOS\"}\n",
    "        self.num_words = 3  # Counting SOS, EOS, PAD\n",
    "\n",
    "    def addSentence(self, sentence):\n",
    "        for word in sentence.split(' '):\n",
    "            self.addWord(word)\n",
    "\n",
    "    def addWord(self, word):\n",
    "        if word not in self.word2index:\n",
    "            self.word2index[word] = self.num_words\n",
    "            self.word2count[word] = 1\n",
    "            self.index2word[self.num_words] = word\n",
    "            self.num_words += 1\n",
    "        else:\n",
    "            self.word2count[word] += 1\n",
    "\n",
    "    # Removing words below a certain count threshold\n",
    "    def trim(self, min_count):\n",
    "        if self.trimmed:\n",
    "            return\n",
    "        self.trimmed = True\n",
    "\n",
    "        keep_words = []\n",
    "\n",
    "        for k, v in self.word2count.items():\n",
    "            if v >= min_count:\n",
    "                keep_words.append(k)\n",
    "\n",
    "        print('keep_words {} / {} = {:.4f}'.format(\n",
    "            len(keep_words), len(self.word2index), len(keep_words) / len(self.word2index)\n",
    "        ))\n",
    "\n",
    "        # Reinitializing dictionaries\n",
    "        self.word2index = {}\n",
    "        self.word2count = {}\n",
    "        self.index2word = {PAD_token: \"PAD\", SOS_token: \"SOS\", EOS_token: \"EOS\"}\n",
    "        self.num_words = 3 # Counting default tokens\n",
    "\n",
    "        for word in keep_words:\n",
    "            self.addWord(word)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Turn a Unicode string to plain ASCII\n",
    "def unicodeToAscii(s):\n",
    "    return ''.join(c for c in unicodedata.normalize('NFD', s) if unicodedata.category(c) != 'Mn')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'cedillea'"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Testing the function - eliminates special characters\n",
    "unicodeToAscii(\"cédillea\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Lowercase, trim, and remove non-letter characters\n",
    "def normalizeString(s):\n",
    "    \n",
    "    # Unicode string to plain ASCII\n",
    "    s = unicodeToAscii(s.lower().strip())\n",
    "    \n",
    "    # Replacing any .!? by a whitespace plus the character\n",
    "    # ' \\1' means the first bracketed group\n",
    "    # r is not to consider ' \\1' as an individual character\n",
    "    # r in r\" \\1\" is to esccape the backslash\n",
    "    s = re.sub(r\"([.!?])\", r\" \\1\", s)\n",
    "    \n",
    "    # Removing any character that is not a sequence of lower or upper case letters\n",
    "    # + means one or more\n",
    "    s = re.sub(r\"[^a-zA-Z.!?]+\", r\" \", s)\n",
    "    \n",
    "    # Removing a sequence of whitespace characters\n",
    "    s = re.sub(r\"\\s+\", r\" \", s).strip()\n",
    "    \n",
    "    return s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'aaa ! ? . aaa a aaa ?'"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Testing the function\n",
    "normalizeString(\"aaa1103!?.     aaa'a aaa?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Reading and processing file....Please Wait\n",
      "Done Reading!!!\n"
     ]
    }
   ],
   "source": [
    "datafile = os.path.join(corpus, \"formatted_movie_lines.txt\")\n",
    "\n",
    "# Reading the file and splitting into lines\n",
    "print(\"Reading and processing file....Please Wait\")\n",
    "lines = open(datafile,encoding = 'utf-8').read().strip().split('\\n')\n",
    "\n",
    "# Splitting every line into pairs and normalizing them\n",
    "pairs = [[normalizeString(s) for s in pair.split('\\t')] for pair in lines]\n",
    "\n",
    "print(\"Done Reading!!!\")\n",
    "voc = Vocabulary(corpus)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "MAX_LENGTH = 10  # Maximum sentence length to consider\n",
    "\n",
    "# Returns True iff both sentences in a pair 'p' are under the MAX_LENGTH threshold\n",
    "def filterPair(p):\n",
    "    # Input sequences need to preserve the last word for EOS token\n",
    "    if p:\n",
    "        return len(p[0].split(' ')) < MAX_LENGTH and len(p[1].split(' ')) < MAX_LENGTH\n",
    "    else:\n",
    "        return False\n",
    "\n",
    "# Filter pairs using filterPair condition\n",
    "def filterPairs(pairs):\n",
    "    pairs = [pair for pair in pairs if pair != ['']]\n",
    "    return [pair for pair in pairs if filterPair(pair)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "There are 442563 pairs/conversations in the dataset\n",
      "After filtering there are 64271 pairs/conversations\n"
     ]
    }
   ],
   "source": [
    "print(\"There are {} pairs/conversations in the dataset\".format(len(pairs)))\n",
    "pairs = filterPairs(pairs)\n",
    "print(\"After filtering there are {} pairs/conversations\".format(len(pairs)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Counted words: 18008\n",
      "['there .', 'where ?']\n",
      "['you have my word . as a gentleman', 'you re sweet .']\n",
      "['hi .', 'looks like things worked out tonight huh ?']\n",
      "['you know chastity ?', 'i believe we share an art instructor']\n",
      "['have fun tonight ?', 'tons']\n",
      "['well no . . .', 'then that s all you had to say .']\n",
      "['then that s all you had to say .', 'but']\n",
      "['but', 'you always been this selfish ?']\n",
      "['do you listen to this crap ?', 'what crap ?']\n",
      "['what good stuff ?', 'the real you .']\n"
     ]
    }
   ],
   "source": [
    "for pair in pairs:\n",
    "        voc.addSentence(pair[0])\n",
    "        voc.addSentence(pair[1])\n",
    "\n",
    "print(\"Counted words:\", voc.num_words)\n",
    "for pair in pairs[:10]:\n",
    "    print(pair)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "keep_words 7823 / 18005 = 0.4345\n",
      "Trimmed from 64271 pairs to 53165, 0.8272 of total\n"
     ]
    }
   ],
   "source": [
    "MIN_COUNT = 3    # Minimum word count threshold for trimming\n",
    "\n",
    "def trimRareWords(voc, pairs, MIN_COUNT):\n",
    "    # Trim words used under the MIN_COUNT from the voc\n",
    "    voc.trim(MIN_COUNT)\n",
    "    # Filter out pairs with trimmed words\n",
    "    keep_pairs = []\n",
    "    for pair in pairs:\n",
    "        input_sentence = pair[0]\n",
    "        output_sentence = pair[1]\n",
    "        keep_input = True\n",
    "        keep_output = True\n",
    "        # Check input sentence\n",
    "        for word in input_sentence.split(' '):\n",
    "            if word not in voc.word2index:\n",
    "                keep_input = False\n",
    "                break\n",
    "        # Check output sentence\n",
    "        for word in output_sentence.split(' '):\n",
    "            if word not in voc.word2index:\n",
    "                keep_output = False\n",
    "                break\n",
    "\n",
    "        # Only keep pairs that do not contain trimmed word(s) in their input or output sentence\n",
    "        if keep_input and keep_output:\n",
    "            keep_pairs.append(pair)\n",
    "\n",
    "    print(\"Trimmed from {} pairs to {}, {:.4f} of total\".format(len(pairs), len(keep_pairs), len(keep_pairs) / len(pairs)))\n",
    "    return keep_pairs\n",
    "\n",
    "\n",
    "# Trim voc and pairs\n",
    "pairs = trimRareWords(voc, pairs, MIN_COUNT)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "def indexesFromSentence(voc, sentence):\n",
    "    return [voc.word2index[word] for word in sentence.split(' ')] + [EOS_token]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[7, 8, 9, 10, 4, 11, 12, 13, 2]"
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "indexesFromSentence(voc,pairs[1][0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [],
   "source": [
    "def zeroPadding(l, fillvalue=PAD_token):\n",
    "    return list(itertools.zip_longest(*l, fillvalue=fillvalue))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [],
   "source": [
    "def binaryMatrix(l, value=PAD_token):\n",
    "    m = []\n",
    "    for i, seq in enumerate(l):\n",
    "        m.append([])\n",
    "        for token in seq:\n",
    "            if token == PAD_token:\n",
    "                m[i].append(0)\n",
    "            else:\n",
    "                m[i].append(1)\n",
    "    return m"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Returns padded input sequence tensor and lengths\n",
    "def inputVar(l, voc):\n",
    "    indexes_batch = [indexesFromSentence(voc, sentence) for sentence in l]\n",
    "    lengths = torch.tensor([len(indexes) for indexes in indexes_batch])\n",
    "    padList = zeroPadding(indexes_batch)\n",
    "    padVar = torch.LongTensor(padList)\n",
    "    return padVar, lengths\n",
    "\n",
    "# Returns padded target sequence tensor, padding mask, and max target length\n",
    "def outputVar(l, voc):\n",
    "    indexes_batch = [indexesFromSentence(voc, sentence) for sentence in l]\n",
    "    max_target_len = max([len(indexes) for indexes in indexes_batch])\n",
    "    padList = zeroPadding(indexes_batch)\n",
    "    mask = binaryMatrix(padList)\n",
    "    mask = torch.ByteTensor(mask)\n",
    "    padVar = torch.LongTensor(padList)\n",
    "    return padVar, mask, max_target_len"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Returns all items for a given batch of pairs\n",
    "def batch2TrainData(voc, pair_batch):\n",
    "    pair_batch.sort(key=lambda x: len(x[0].split(\" \")), reverse=True)\n",
    "    input_batch, output_batch = [], []\n",
    "    for pair in pair_batch:\n",
    "        input_batch.append(pair[0])\n",
    "        output_batch.append(pair[1])\n",
    "    inp, lengths = inputVar(input_batch, voc)\n",
    "    output, mask, max_target_len = outputVar(output_batch, voc)\n",
    "    return inp, lengths, output, mask, max_target_len"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input_variable: tensor([[  94,   65,   16,  157, 4514],\n",
      "        [ 117,  702, 1119,  266,    4],\n",
      "        [   7,   96,    4,    6,    2],\n",
      "        [ 697,  214,    2,    2,    0],\n",
      "        [  83,    6,    0,    0,    0],\n",
      "        [   6,    2,    0,    0,    0],\n",
      "        [   2,    0,    0,    0,    0]])\n",
      "lengths: tensor([7, 6, 4, 4, 3])\n",
      "target_variable: tensor([[  25,   65,   16,  239,   77],\n",
      "        [ 697,  517, 4607,  488,   37],\n",
      "        [   7,   96,    4,  198, 4514],\n",
      "        [2182,  214,    5,   76,    6],\n",
      "        [   4,    4,    7,  180,    2],\n",
      "        [   2,    2,   44,    4,    0],\n",
      "        [   0,    0,    6,    2,    0],\n",
      "        [   0,    0,    2,    0,    0]])\n",
      "mask: tensor([[1, 1, 1, 1, 1],\n",
      "        [1, 1, 1, 1, 1],\n",
      "        [1, 1, 1, 1, 1],\n",
      "        [1, 1, 1, 1, 1],\n",
      "        [1, 1, 1, 1, 1],\n",
      "        [1, 1, 1, 1, 0],\n",
      "        [0, 0, 1, 1, 0],\n",
      "        [0, 0, 1, 0, 0]], dtype=torch.uint8)\n",
      "max_target_len: 8\n"
     ]
    }
   ],
   "source": [
    "# Example for validation\n",
    "small_batch_size = 5\n",
    "batches = batch2TrainData(voc, [random.choice(pairs) for _ in range(small_batch_size)])\n",
    "input_variable, lengths, target_variable, mask, max_target_len = batches\n",
    "\n",
    "print(\"input_variable:\", input_variable)\n",
    "print(\"lengths:\", lengths)\n",
    "print(\"target_variable:\", target_variable)\n",
    "print(\"mask:\", mask)\n",
    "print(\"max_target_len:\", max_target_len)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [],
   "source": [
    "class EncoderRNN(nn.Module):\n",
    "    def __init__(self, hidden_size, embedding, n_layers=1, dropout=0):\n",
    "        super(EncoderRNN, self).__init__()\n",
    "        self.n_layers = n_layers\n",
    "        self.hidden_size = hidden_size\n",
    "        self.embedding = embedding\n",
    "\n",
    "        # Initialize GRU; the input_size and hidden_size params are both set to 'hidden_size'\n",
    "        #   because our input size is a word embedding with number of features == hidden_size\n",
    "        self.gru = nn.GRU(hidden_size, hidden_size, n_layers,dropout=(0 if n_layers == 1 else dropout), bidirectional=True)\n",
    "\n",
    "    def forward(self, input_seq, input_lengths, hidden=None):\n",
    "        # Convert word indexes to embeddings\n",
    "        embedded = self.embedding(input_seq)\n",
    "        # Pack padded batch of sequences for RNN module\n",
    "        packed = torch.nn.utils.rnn.pack_padded_sequence(embedded, input_lengths)\n",
    "        # Forward pass through GRU\n",
    "        outputs, hidden = self.gru(packed, hidden)\n",
    "        # Unpack padding\n",
    "        outputs, _ = torch.nn.utils.rnn.pad_packed_sequence(outputs)\n",
    "        # Sum bidirectional GRU outputs\n",
    "        outputs = outputs[:, :, :self.hidden_size] + outputs[:, : ,self.hidden_size:]\n",
    "        # Return output and final hidden state\n",
    "        return outputs, hidden"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Luong attention layer\n",
    "class Attn(torch.nn.Module):\n",
    "    def __init__(self, method, hidden_size):\n",
    "        super(Attn, self).__init__()\n",
    "        self.method = method\n",
    "        if self.method not in ['dot', 'general', 'concat']:\n",
    "            raise ValueError(self.method, \"is not an appropriate attention method.\")\n",
    "        self.hidden_size = hidden_size\n",
    "        if self.method == 'general':\n",
    "            self.attn = torch.nn.Linear(self.hidden_size, hidden_size)\n",
    "        elif self.method == 'concat':\n",
    "            self.attn = torch.nn.Linear(self.hidden_size * 2, hidden_size)\n",
    "            self.v = torch.nn.Parameter(torch.FloatTensor(hidden_size))\n",
    "\n",
    "    def dot_score(self, hidden, encoder_output):\n",
    "        return torch.sum(hidden * encoder_output, dim=2)\n",
    "\n",
    "    def general_score(self, hidden, encoder_output):\n",
    "        energy = self.attn(encoder_output)\n",
    "        return torch.sum(hidden * energy, dim=2)\n",
    "\n",
    "    def concat_score(self, hidden, encoder_output):\n",
    "        energy = self.attn(torch.cat((hidden.expand(encoder_output.size(0), -1, -1), encoder_output), 2)).tanh()\n",
    "        return torch.sum(self.v * energy, dim=2)\n",
    "\n",
    "    def forward(self, hidden, encoder_outputs):\n",
    "        # Calculate the attention weights (energies) based on the given method\n",
    "        if self.method == 'general':\n",
    "            attn_energies = self.general_score(hidden, encoder_outputs)\n",
    "        elif self.method == 'concat':\n",
    "            attn_energies = self.concat_score(hidden, encoder_outputs)\n",
    "        elif self.method == 'dot':\n",
    "            attn_energies = self.dot_score(hidden, encoder_outputs)\n",
    "\n",
    "        # Transpose max_length and batch_size dimensions\n",
    "        attn_energies = attn_energies.t()\n",
    "\n",
    "        # Return the softmax normalized probability scores (with added dimension)\n",
    "        return F.softmax(attn_energies, dim=1).unsqueeze(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LuongAttnDecoderRNN(nn.Module):\n",
    "    def __init__(self, attn_model, embedding, hidden_size, output_size, n_layers=1, dropout=0.1):\n",
    "        super(LuongAttnDecoderRNN, self).__init__()\n",
    "\n",
    "        # Keep for reference\n",
    "        self.attn_model = attn_model\n",
    "        self.hidden_size = hidden_size\n",
    "        self.output_size = output_size\n",
    "        self.n_layers = n_layers\n",
    "        self.dropout = dropout\n",
    "\n",
    "        # Define layers\n",
    "        self.embedding = embedding\n",
    "        self.embedding_dropout = nn.Dropout(dropout)\n",
    "        self.gru = nn.GRU(hidden_size, hidden_size, n_layers, dropout=(0 if n_layers == 1 else dropout))\n",
    "        self.concat = nn.Linear(hidden_size * 2, hidden_size)\n",
    "        self.out = nn.Linear(hidden_size, output_size)\n",
    "\n",
    "        self.attn = Attn(attn_model, hidden_size)\n",
    "\n",
    "    def forward(self, input_step, last_hidden, encoder_outputs):\n",
    "        # Note: we run this one step (word) at a time\n",
    "        # Get embedding of current input word\n",
    "        embedded = self.embedding(input_step)\n",
    "        embedded = self.embedding_dropout(embedded)\n",
    "        # Forward through unidirectional GRU\n",
    "        rnn_output, hidden = self.gru(embedded, last_hidden)\n",
    "        # Calculate attention weights from the current GRU output\n",
    "        attn_weights = self.attn(rnn_output, encoder_outputs)\n",
    "        # Multiply attention weights to encoder outputs to get new \"weighted sum\" context vector\n",
    "        context = attn_weights.bmm(encoder_outputs.transpose(0, 1))\n",
    "        # Concatenate weighted context vector and GRU output using Luong eq. 5\n",
    "        rnn_output = rnn_output.squeeze(0)\n",
    "        context = context.squeeze(1)\n",
    "        concat_input = torch.cat((rnn_output, context), 1)\n",
    "        concat_output = torch.tanh(self.concat(concat_input))\n",
    "        # Predict next word using Luong eq. 6\n",
    "        output = self.out(concat_output)\n",
    "        output = F.softmax(output, dim=1)\n",
    "        # Return output and final hidden state\n",
    "        return output, hidden"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [],
   "source": [
    "def maskNLLLoss(inp, target, mask):\n",
    "    nTotal = mask.sum()\n",
    "    crossEntropy = -torch.log(torch.gather(inp, 1, target.view(-1, 1)))\n",
    "    loss = crossEntropy.masked_select(mask).mean()\n",
    "    loss = loss.to(device)\n",
    "    return loss, nTotal.item()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(input_variable, lengths, target_variable, mask, max_target_len, encoder, decoder, embedding,\n",
    "          encoder_optimizer, decoder_optimizer, batch_size, clip, max_length=MAX_LENGTH):\n",
    "\n",
    "    # Zero gradients\n",
    "    encoder_optimizer.zero_grad()\n",
    "    decoder_optimizer.zero_grad()\n",
    "\n",
    "    # Set device options\n",
    "    input_variable = input_variable.to(device)\n",
    "    lengths = lengths.to(device)\n",
    "    target_variable = target_variable.to(device)\n",
    "    mask = mask.to(device)\n",
    "\n",
    "    # Initialize variables\n",
    "    loss = 0\n",
    "    print_losses = []\n",
    "    n_totals = 0\n",
    "\n",
    "    # Forward pass through encoder\n",
    "    encoder_outputs, encoder_hidden = encoder(input_variable, lengths)\n",
    "\n",
    "    # Create initial decoder input (start with SOS tokens for each sentence)\n",
    "    decoder_input = torch.LongTensor([[SOS_token for _ in range(batch_size)]])\n",
    "    decoder_input = decoder_input.to(device)\n",
    "\n",
    "    # Set initial decoder hidden state to the encoder's final hidden state\n",
    "    decoder_hidden = encoder_hidden[:decoder.n_layers]\n",
    "\n",
    "    # Determine if we are using teacher forcing this iteration\n",
    "    use_teacher_forcing = True if random.random() < teacher_forcing_ratio else False\n",
    "\n",
    "    # Forward batch of sequences through decoder one time step at a time\n",
    "    if use_teacher_forcing:\n",
    "        for t in range(max_target_len):\n",
    "            decoder_output, decoder_hidden = decoder(\n",
    "                decoder_input, decoder_hidden, encoder_outputs\n",
    "            )\n",
    "            # Teacher forcing: next input is current target\n",
    "            decoder_input = target_variable[t].view(1, -1)\n",
    "            # Calculate and accumulate loss\n",
    "            mask_loss, nTotal = maskNLLLoss(decoder_output, target_variable[t], mask[t])\n",
    "            loss += mask_loss\n",
    "            print_losses.append(mask_loss.item() * nTotal)\n",
    "            n_totals += nTotal\n",
    "    else:\n",
    "        for t in range(max_target_len):\n",
    "            decoder_output, decoder_hidden = decoder(\n",
    "                decoder_input, decoder_hidden, encoder_outputs\n",
    "            )\n",
    "            # No teacher forcing: next input is decoder's own current output\n",
    "            _, topi = decoder_output.topk(1)\n",
    "            decoder_input = torch.LongTensor([[topi[i][0] for i in range(batch_size)]])\n",
    "            decoder_input = decoder_input.to(device)\n",
    "            # Calculate and accumulate loss\n",
    "            mask_loss, nTotal = maskNLLLoss(decoder_output, target_variable[t], mask[t])\n",
    "            loss += mask_loss\n",
    "            print_losses.append(mask_loss.item() * nTotal)\n",
    "            n_totals += nTotal\n",
    "\n",
    "    # Perform backpropatation\n",
    "    loss.backward()\n",
    "\n",
    "    # Clip gradients: gradients are modified in place\n",
    "    _ = torch.nn.utils.clip_grad_norm_(encoder.parameters(), clip)\n",
    "    _ = torch.nn.utils.clip_grad_norm_(decoder.parameters(), clip)\n",
    "\n",
    "    # Adjust model weights\n",
    "    encoder_optimizer.step()\n",
    "    decoder_optimizer.step()\n",
    "\n",
    "    return sum(print_losses) / n_totals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [],
   "source": [
    "def trainIters(model_name, voc, pairs, encoder, decoder, encoder_optimizer, decoder_optimizer, embedding, encoder_n_layers, decoder_n_layers, save_dir, n_iteration, batch_size, print_every, save_every, clip, corpus_name, loadFilename):\n",
    "\n",
    "    # Load batches for each iteration\n",
    "    training_batches = [batch2TrainData(voc, [random.choice(pairs) for _ in range(batch_size)])\n",
    "                      for _ in range(n_iteration)]\n",
    "\n",
    "    # Initializations\n",
    "    print('Initializing ...')\n",
    "    start_iteration = 1\n",
    "    print_loss = 0\n",
    "    if loadFilename:\n",
    "        start_iteration = checkpoint['iteration'] + 1\n",
    "\n",
    "    # Training loop\n",
    "    print(\"Training...\")\n",
    "    for iteration in range(start_iteration, n_iteration + 1):\n",
    "        training_batch = training_batches[iteration - 1]\n",
    "        # Extract fields from batch\n",
    "        input_variable, lengths, target_variable, mask, max_target_len = training_batch\n",
    "\n",
    "        # Run a training iteration with batch\n",
    "        loss = train(input_variable, lengths, target_variable, mask, max_target_len, encoder,\n",
    "                     decoder, embedding, encoder_optimizer, decoder_optimizer, batch_size, clip)\n",
    "        print_loss += loss\n",
    "\n",
    "        # Print progress\n",
    "        if iteration % print_every == 0:\n",
    "            print_loss_avg = print_loss / print_every\n",
    "            print(\"Iteration: {}; Percent complete: {:.1f}%; Average loss: {:.4f}\".format(iteration, iteration / n_iteration * 100, print_loss_avg))\n",
    "            print_loss = 0\n",
    "\n",
    "        # Save checkpoint\n",
    "        if (iteration % save_every == 0):\n",
    "            directory = os.path.join(save_dir, model_name, corpus_name, '{}-{}_{}'.format(encoder_n_layers, decoder_n_layers, hidden_size))\n",
    "            if not os.path.exists(directory):\n",
    "                os.makedirs(directory)\n",
    "            torch.save({\n",
    "                'iteration': iteration,\n",
    "                'en': encoder.state_dict(),\n",
    "                'de': decoder.state_dict(),\n",
    "                'en_opt': encoder_optimizer.state_dict(),\n",
    "                'de_opt': decoder_optimizer.state_dict(),\n",
    "                'loss': loss,\n",
    "                'voc_dict': voc.__dict__,\n",
    "                'embedding': embedding.state_dict()\n",
    "            }, os.path.join(directory, '{}_{}.tar'.format(iteration, 'checkpoint')))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GreedySearchDecoder(nn.Module):\n",
    "    def __init__(self, encoder, decoder):\n",
    "        super(GreedySearchDecoder, self).__init__()\n",
    "        self.encoder = encoder\n",
    "        self.decoder = decoder\n",
    "\n",
    "    def forward(self, input_seq, input_length, max_length):\n",
    "        # Forward input through encoder model\n",
    "        encoder_outputs, encoder_hidden = self.encoder(input_seq, input_length)\n",
    "        # Prepare encoder's final hidden layer to be first hidden input to the decoder\n",
    "        decoder_hidden = encoder_hidden[:decoder.n_layers]\n",
    "        # Initialize decoder input with SOS_token\n",
    "        decoder_input = torch.ones(1, 1, device=device, dtype=torch.long) * SOS_token\n",
    "        # Initialize tensors to append decoded words\n",
    "        all_tokens = torch.zeros([0], device=device, dtype=torch.long)\n",
    "        all_scores = torch.zeros([0], device=device)\n",
    "        # Iteratively decode one word token at a time\n",
    "        for _ in range(max_length):\n",
    "            # Forward pass through decoder\n",
    "            decoder_output, decoder_hidden = self.decoder(decoder_input, decoder_hidden, encoder_outputs)\n",
    "            # Obtain most likely word token and its softmax score\n",
    "            decoder_scores, decoder_input = torch.max(decoder_output, dim=1)\n",
    "            # Record token and score\n",
    "            all_tokens = torch.cat((all_tokens, decoder_input), dim=0)\n",
    "            all_scores = torch.cat((all_scores, decoder_scores), dim=0)\n",
    "            # Prepare current token to be next decoder input (add a dimension)\n",
    "            decoder_input = torch.unsqueeze(decoder_input, 0)\n",
    "        # Return collections of word tokens and scores\n",
    "        return all_tokens, all_scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate(encoder, decoder, searcher, voc, sentence, max_length=MAX_LENGTH):\n",
    "    ### Format input sentence as a batch\n",
    "    # words -> indexes\n",
    "    indexes_batch = [indexesFromSentence(voc, sentence)]\n",
    "    # Create lengths tensor\n",
    "    lengths = torch.tensor([len(indexes) for indexes in indexes_batch])\n",
    "    # Transpose dimensions of batch to match models' expectations\n",
    "    input_batch = torch.LongTensor(indexes_batch).transpose(0, 1)\n",
    "    # Use appropriate device\n",
    "    input_batch = input_batch.to(device)\n",
    "    lengths = lengths.to(device)\n",
    "    # Decode sentence with searcher\n",
    "    tokens, scores = searcher(input_batch, lengths, max_length)\n",
    "    # indexes -> words\n",
    "    decoded_words = [voc.index2word[token.item()] for token in tokens]\n",
    "    return decoded_words\n",
    "\n",
    "\n",
    "def evaluateInput(encoder, decoder, searcher, voc):\n",
    "    input_sentence = ''\n",
    "    while(1):\n",
    "        try:\n",
    "            # Get input sentence\n",
    "            input_sentence = input('> ')\n",
    "            # Check if it is quit case\n",
    "            if input_sentence == 'q' or input_sentence == 'quit': break\n",
    "            # Normalize sentence\n",
    "            input_sentence = normalizeString(input_sentence)\n",
    "            # Evaluate sentence\n",
    "            output_words = evaluate(encoder, decoder, searcher, voc, input_sentence)\n",
    "            # Format and print response sentence\n",
    "            output_words[:] = [x for x in output_words if not (x == 'EOS' or x == 'PAD')]\n",
    "            print('Bot:', ' '.join(output_words))\n",
    "\n",
    "        except KeyError:\n",
    "            print(\"Error: Encountered unknown word.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Building encoder and decoder ...\n",
      "Models built and ready to go!\n"
     ]
    }
   ],
   "source": [
    "# Configure models\n",
    "model_name = 'cb_model'\n",
    "attn_model = 'dot'\n",
    "#attn_model = 'general'\n",
    "#attn_model = 'concat'\n",
    "hidden_size = 512\n",
    "encoder_n_layers = 2\n",
    "decoder_n_layers = 2\n",
    "dropout = 0.1\n",
    "batch_size = 32\n",
    "save_dir = os.path.join(\"data\", \"save\")\n",
    "\n",
    "# NOTE : Set checkpoint to load from; set to None if starting from scratch\n",
    "# loadFilename = None\n",
    "checkpoint_iter = 8000\n",
    "loadFilename = os.path.join(save_dir, model_name, corpus_name,\n",
    "                           '{}-{}_{}'.format(encoder_n_layers, decoder_n_layers, hidden_size),\n",
    "                           '{}_checkpoint.tar'.format(checkpoint_iter))\n",
    "\n",
    "\n",
    "# Load model if a loadFilename is provided\n",
    "if loadFilename:\n",
    "    # If loading on same machine the model was trained on\n",
    "    checkpoint = torch.load(loadFilename)\n",
    "    # If loading a model trained on GPU to CPU\n",
    "    #checkpoint = torch.load(loadFilename, map_location=torch.device('cpu'))\n",
    "    encoder_sd = checkpoint['en']\n",
    "    decoder_sd = checkpoint['de']\n",
    "    encoder_optimizer_sd = checkpoint['en_opt']\n",
    "    decoder_optimizer_sd = checkpoint['de_opt']\n",
    "    embedding_sd = checkpoint['embedding']\n",
    "    voc.__dict__ = checkpoint['voc_dict']\n",
    "\n",
    "\n",
    "print('Building encoder and decoder ...')\n",
    "# Initialize word embeddings\n",
    "embedding = nn.Embedding(voc.num_words, hidden_size)\n",
    "if loadFilename:\n",
    "    embedding.load_state_dict(embedding_sd)\n",
    "# Initialize encoder & decoder models\n",
    "encoder = EncoderRNN(hidden_size, embedding, encoder_n_layers, dropout)\n",
    "decoder = LuongAttnDecoderRNN(attn_model, embedding, hidden_size, voc.num_words, decoder_n_layers, dropout)\n",
    "if loadFilename:\n",
    "    encoder.load_state_dict(encoder_sd)\n",
    "    decoder.load_state_dict(decoder_sd)\n",
    "# Use appropriate device\n",
    "encoder = encoder.to(device)\n",
    "decoder = decoder.to(device)\n",
    "print('Models built and ready to go!')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "> hello\n",
      "Bot: hello .\n",
      "> how are you\n",
      "Bot: fine .\n",
      "> what is your name\n",
      "Bot: victor frankenstein .\n",
      "> how old are you ?\n",
      "Bot: twelve .\n",
      "> How is the weather ?\n",
      "Bot: gone .\n",
      "> what is the time right now ?\n",
      "Bot: i don t know .\n",
      "> really\n",
      "Bot: yes .\n",
      "> quit\n"
     ]
    }
   ],
   "source": [
    "# Set dropout layers to eval mode\n",
    "encoder.eval()\n",
    "decoder.eval()\n",
    "\n",
    "# Initialize search module\n",
    "searcher = GreedySearchDecoder(encoder, decoder)\n",
    "\n",
    "# Begin chatting \n",
    "evaluateInput(encoder, decoder, searcher, voc)"
   ]
  }
 ],
 "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}