chatbot.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LSTM Chatbot\n",
    "_ _ _\n",
    "## Overview  \n",
    "\n",
    "I am training a simple seq2seq architecture using movie\n",
    "scripts from the [Cornell Movie-Dialogs Corpus](https://www.cs.cornell.edu/~cristian/Cornell_Movie-Dialogs_Corpus.html). <br/> Let start by downloading the [dataset](https://zissou.infosci.cornell.edu/convokit/datasets/movie-corpus/movie-corpus.zip) first.  \n",
    "\n",
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\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\n",
    "import math\n",
    "import json\n",
    "from io import open\n",
    "from torch import optim\n",
    "from torch.jit import script, trace"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data Preprocessing  \n",
    "\n",
    "### File Formatting Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def load_data(fileName):\n",
    "    lines = {}\n",
    "    conversations = {}\n",
    "    with open(fileName, 'r', encoding='iso-8859-1') as f:\n",
    "        for line in f:\n",
    "            lineJson = json.loads(line)\n",
    "            lineObj = {}\n",
    "            lineObj[\"lineID\"] = lineJson[\"id\"]\n",
    "            lineObj[\"characterID\"] = lineJson[\"speaker\"]\n",
    "            lineObj[\"text\"] = lineJson[\"text\"]\n",
    "            lines[lineObj['lineID']] = lineObj\n",
    "            if lineJson[\"conversation_id\"] not in conversations:\n",
    "                convObj = {}\n",
    "                convObj[\"conversationID\"] = lineJson[\"conversation_id\"]\n",
    "                convObj[\"movieID\"] = lineJson[\"meta\"][\"movie_id\"]\n",
    "                convObj[\"lines\"] = [lineObj]\n",
    "            else:\n",
    "                convObj = conversations[lineJson[\"conversation_id\"]]\n",
    "                convObj[\"lines\"].insert(0, lineObj)\n",
    "            conversations[convObj[\"conversationID\"]] = convObj\n",
    "    return lines, conversations\n",
    "\n",
    "def extract_pairs(conversations):\n",
    "    pairs = []\n",
    "    for conversation in conversations.values():\n",
    "        for i in range(len(conversation[\"lines\"]) - 1):\n",
    "            inputLine = conversation[\"lines\"][i][\"text\"].strip()\n",
    "            targetLine = conversation[\"lines\"][i+1][\"text\"].strip()\n",
    "            if inputLine and targetLine:\n",
    "                pairs.append([inputLine, targetLine])\n",
    "    return pairs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Call the above functions to create a formatted file.\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading data...\n",
      "Writing formatted file...\n",
      "Done!\n"
     ]
    }
   ],
   "source": [
    "corpus_name = \"movie-corpus\"\n",
    "corpus = os.path.join(\"data\", corpus_name)\n",
    "datafile = os.path.join(corpus, \"formatted_data.txt\")\n",
    "delimiter = '\\t'\n",
    "delimiter = str(codecs.decode(delimiter, \"unicode_escape\"))\n",
    "lines = {}\n",
    "conversations = {}\n",
    "# Load lines and conversations\n",
    "print('Loading data...')\n",
    "lines, conversations = load_data(os.path.join(corpus, \"utterances.jsonl\"))\n",
    "# Write formatted_data.txt\n",
    "print('Writing formatted file...')\n",
    "with open(datafile, 'w', encoding='utf-8') as outputfile:\n",
    "    writer = csv.writer(outputfile, delimiter=delimiter, lineterminator='\\n')\n",
    "    for pair in extract_pairs(conversations):\n",
    "        writer.writerow(pair)\n",
    "print('Done!')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Vocabulary\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "PAD_token = 0  # Padding for short sentences\n",
    "SOS_token = 1  # Start of sentence\n",
    "EOS_token = 2  # End of sentence\n",
    "\n",
    "class Voc:\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  \n",
    "\n",
    "    def add_sentence(self, sentence):\n",
    "        for word in sentence.split(' '):\n",
    "            self.add_word(word)\n",
    "\n",
    "    def add_word(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",
    "    # Remove words below a certain count threshold\n",
    "    def trim(self, min_count):\n",
    "        if self.trimmed:\n",
    "            return\n",
    "        self.trimmed = True\n",
    "        keep_words = []\n",
    "        for k, v in self.word2count.items():\n",
    "            if v >= min_count:\n",
    "                keep_words.append(k)\n",
    "        print('keep_words {} / {} = {:.4f}'.format(\n",
    "            len(keep_words), len(self.word2index), len(keep_words) / len(self.word2index)\n",
    "        ))\n",
    "        # Reinitialize dictionaries\n",
    "        self.word2index = {}\n",
    "        self.word2count = {}\n",
    "        self.index2word = {PAD_token: \"PAD\", SOS_token: \"SOS\", EOS_token: \"EOS\"}\n",
    "        self.num_words = 3 \n",
    "        for word in keep_words:\n",
    "            self.add_word(word)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Important parameters for building vocabulary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "MAX_LENGTH = 10  # Maximum sentence length\n",
    "MIN_RARE_COUNT = 3 # Minimum count for eliminating rare words"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- Convert the Unicode strings to ASCII\n",
    "- Convert all letters to lowercase\n",
    "- Trim all non-letter characters except for basic punctuation\n",
    "- Filter out sentences with length greater than the ``MAX_LENGTH``\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Read 221282 sentence pairs\n",
      "Trimmed to 64313 sentence pairs\n",
      "Counted words: 18082\n"
     ]
    }
   ],
   "source": [
    "# https://stackoverflow.com/a/518232/2809427\n",
    "def unicode_to_ascii(s):\n",
    "    return ''.join(\n",
    "        c for c in unicodedata.normalize('NFD', s)\n",
    "        if unicodedata.category(c) != 'Mn'\n",
    "    )\n",
    "\n",
    "def normalize_string(s):\n",
    "    s = unicode_to_ascii(s.lower().strip())\n",
    "    s = re.sub(r\"([.!?])\", r\" \\1\", s)\n",
    "    s = re.sub(r\"[^a-zA-Z.!?]+\", r\" \", s)\n",
    "    s = re.sub(r\"\\s+\", r\" \", s).strip()\n",
    "    return s\n",
    "\n",
    "def read_vocs(datafile, corpus_name):\n",
    "    lines = open(datafile, encoding='utf-8').\\\n",
    "        read().strip().split('\\n')\n",
    "    pairs = [[normalize_string(s) for s in l.split('\\t')] for l in lines]\n",
    "    voc = Voc(corpus_name)\n",
    "    return voc, pairs\n",
    "\n",
    "def filter_pair(p):\n",
    "    return len(p[0].split(' ')) < MAX_LENGTH and len(p[1].split(' ')) < MAX_LENGTH\n",
    "\n",
    "def filter_pairs(pairs):\n",
    "    return [pair for pair in pairs if filter_pair(pair)]\n",
    "\n",
    "def load_prepare_data(corpus, corpus_name, datafile):\n",
    "    voc, pairs = read_vocs(datafile, corpus_name)\n",
    "    print(\"Read {} sentence pairs\".format(len(pairs)))\n",
    "    pairs = filter_pairs(pairs)\n",
    "    print(\"Trimmed to {} sentence pairs\".format(len(pairs)))\n",
    "    for pair in pairs:\n",
    "        voc.add_sentence(pair[0])\n",
    "        voc.add_sentence(pair[1])\n",
    "    print(\"Counted words:\", voc.num_words)\n",
    "    return voc, pairs\n",
    "\n",
    "voc, pairs = load_prepare_data(corpus, corpus_name, datafile)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Trim rarely used words from vocabulary.\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "keep_words 7833 / 18079 = 0.4333\n",
      "Trimmed from 64313 pairs to 53131, 0.8261 of total\n"
     ]
    }
   ],
   "source": [
    "def trim_rare_words(voc, pairs, min_count):\n",
    "    voc.trim(min_count)\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",
    "        for word in input_sentence.split(' '):\n",
    "            if word not in voc.word2index:\n",
    "                keep_input = False\n",
    "                break\n",
    "        for word in output_sentence.split(' '):\n",
    "            if word not in voc.word2index:\n",
    "                keep_output = False\n",
    "                break\n",
    "        if keep_input and keep_output:\n",
    "            keep_pairs.append(pair)\n",
    "    print(\"Trimmed from {} pairs to {}, {:.4f} of total\".format(len(pairs),\n",
    "        len(keep_pairs), len(keep_pairs) / len(pairs)))\n",
    "    return keep_pairs\n",
    "\n",
    "pairs = trim_rare_words(voc, pairs, MIN_RARE_COUNT)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Convert to Tensors\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input_variable: tensor([[  11,  175,  287,   19,  162],\n",
      "        [  83,    4,  288,   17,  665],\n",
      "        [  11,   11,  210,   92,   14],\n",
      "        [ 200,  257,   14,   10,    2],\n",
      "        [ 512,   10,    2,    2,    0],\n",
      "        [ 743,    2,    0,    0,    0],\n",
      "        [ 483,    0,    0,    0,    0],\n",
      "        [2498,    0,    0,    0,    0],\n",
      "        [  14,    0,    0,    0,    0],\n",
      "        [   2,    0,    0,    0,    0]])\n",
      "lengths: tensor([10,  6,  5,  5,  4])\n",
      "target_variable: tensor([[  33,   25,   24,  111,  112],\n",
      "        [  36,  590,   64,   24,   10],\n",
      "        [  17,   14,  555,  154,    2],\n",
      "        [1758,    2,    7,  832,    0],\n",
      "        [  14,    0,   72,   16,    0],\n",
      "        [   2,    0,   14,   72,    0],\n",
      "        [   0,    0,    2,   99,    0],\n",
      "        [   0,    0,    0,   10,    0],\n",
      "        [   0,    0,    0,    2,    0]])\n",
      "mask: tensor([[ True,  True,  True,  True,  True],\n",
      "        [ True,  True,  True,  True,  True],\n",
      "        [ True,  True,  True,  True,  True],\n",
      "        [ True,  True,  True,  True, False],\n",
      "        [ True, False,  True,  True, False],\n",
      "        [ True, False,  True,  True, False],\n",
      "        [False, False,  True,  True, False],\n",
      "        [False, False, False,  True, False],\n",
      "        [False, False, False,  True, False]])\n",
      "max_target_len: 9\n"
     ]
    }
   ],
   "source": [
    "def indexes_from_sentence(voc, sentence):\n",
    "    return [voc.word2index[word] for word in sentence.split(' ')] + [EOS_token]\n",
    "\n",
    "def zero_padding(l, fillvalue=PAD_token):\n",
    "    return list(itertools.zip_longest(*l, fillvalue=fillvalue))\n",
    "\n",
    "def binary_matrix(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\n",
    "\n",
    "# Returns padded input sequence tensor and lengths\n",
    "def input_var(l, voc):\n",
    "    indexes_batch = [indexes_from_sentence(voc, sentence) for sentence in l]\n",
    "    lengths = torch.tensor([len(indexes) for indexes in indexes_batch])\n",
    "    pad_list = zero_padding(indexes_batch)\n",
    "    pad_var = torch.LongTensor(pad_list)\n",
    "    return pad_var, lengths\n",
    "\n",
    "# Returns padded target sequence tensor, padding mask, and max target length\n",
    "def output_var(l, voc):\n",
    "    indexes_batch = [indexes_from_sentence(voc, sentence) for sentence in l]\n",
    "    max_target_len = max([len(indexes) for indexes in indexes_batch])\n",
    "    pad_list = zero_padding(indexes_batch)\n",
    "    mask = binary_matrix(pad_list)\n",
    "    mask = torch.BoolTensor(mask)\n",
    "    pad_var = torch.LongTensor(pad_list)\n",
    "    return pad_var, mask, max_target_len\n",
    "\n",
    "# Returns all items for a given batch of pairs\n",
    "def batch_to_train_data(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 = input_var(input_batch, voc)\n",
    "    output, mask, max_target_len = output_var(output_batch, voc)\n",
    "    return inp, lengths, output, mask, max_target_len\n",
    "\n",
    "# Example for validation\n",
    "small_batch_size = 5\n",
    "batches = batch_to_train_data(voc, [random.choice(pairs) for _ in range(small_batch_size)])\n",
    "input_variable, lengths, target_variable, mask, max_target_len = batches\n",
    "\n",
    "# Let print all variables\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": "markdown",
   "metadata": {},
   "source": [
    "## Define Model  \n",
    "\n",
    "### Encoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "class Encoder(nn.Module):\n",
    "    def __init__(self, hidden_size, embedding, n_layers=1, dropout=0):\n",
    "        super(Encoder, self).__init__()\n",
    "        self.n_layers = n_layers\n",
    "        self.hidden_size = hidden_size\n",
    "        self.embedding = embedding\n",
    "        self.lstm = nn.GRU(hidden_size, hidden_size, n_layers,\n",
    "                          dropout=(0 if n_layers == 1 else dropout), bidirectional=True)\n",
    "\n",
    "    def forward(self, input_seq, input_lengths, hidden=None):\n",
    "        embedded = self.embedding(input_seq)\n",
    "        packed = nn.utils.rnn.pack_padded_sequence(embedded, input_lengths)\n",
    "        outputs, hidden = self.lstm(packed, hidden)\n",
    "        outputs, _ = nn.utils.rnn.pad_packed_sequence(outputs)\n",
    "        outputs = outputs[:, :, :self.hidden_size] + outputs[:, : ,self.hidden_size:]\n",
    "        return outputs, hidden"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Luong attention layer\n",
    "class Attn(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 = nn.Linear(self.hidden_size, hidden_size)\n",
    "        elif self.method == 'concat':\n",
    "            self.attn = nn.Linear(self.hidden_size * 2, hidden_size)\n",
    "            self.v = 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",
    "        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",
    "        attn_energies = attn_energies.t()\n",
    "        return F.softmax(attn_energies, dim=1).unsqueeze(1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Attention Decoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "class Decoder(nn.Module):\n",
    "    def __init__(self, attn_model, embedding, hidden_size, output_size, n_layers=1, dropout=0.1):\n",
    "        super(Decoder, self).__init__()\n",
    "\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",
    "        self.embedding = embedding\n",
    "        self.embedding_dropout = nn.Dropout(dropout)\n",
    "        self.lstm = 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",
    "        self.attn = Attn(attn_model, hidden_size)\n",
    "\n",
    "    def forward(self, input_step, last_hidden, encoder_outputs):\n",
    "        embedded = self.embedding(input_step)\n",
    "        embedded = self.embedding_dropout(embedded)\n",
    "        rnn_output, hidden = self.lstm(embedded, last_hidden)\n",
    "        attn_weights = self.attn(rnn_output, encoder_outputs)\n",
    "        context = attn_weights.bmm(encoder_outputs.transpose(0, 1))\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",
    "        output = self.out(concat_output)\n",
    "        output = F.softmax(output, dim=1)\n",
    "        return output, hidden"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Sequence to sequence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "class Seq2Seq(nn.Module):\n",
    "    def __init__(self, encoder, decoder, device, ratio):\n",
    "        super(Seq2Seq, self).__init__()\n",
    "        self.encoder = encoder\n",
    "        self.decoder = decoder\n",
    "        self.device = device\n",
    "        self.teacher_forcing_ratio = ratio\n",
    "\n",
    "    def forward(self, input_seq, input_length, target, mask, max_target_len, batch_size):\n",
    "        loss = 0\n",
    "        print_losses = []\n",
    "        n_totals = 0\n",
    "        encoder_outputs, encoder_hidden = self.encoder(input_seq, input_length)\n",
    "        decoder_hidden = encoder_hidden[:decoder.n_layers]\n",
    "        decoder_input = torch.LongTensor([[SOS_token for _ in range(batch_size)]])\n",
    "        decoder_input = decoder_input.to(self.device)\n",
    "        use_teacher_forcing = True if random.random() < teacher_forcing_ratio else False\n",
    "        if use_teacher_forcing:\n",
    "            for t in range(max_target_len):\n",
    "                decoder_output, decoder_hidden = self.decoder(decoder_input, decoder_hidden, encoder_outputs)\n",
    "                decoder_input = target[t].view(1, -1)\n",
    "                decoder_input = decoder_input.to(self.device)\n",
    "                mask_loss, nTotal = maskNLLLoss(decoder_output, target[t], mask[t], self.device)\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 = self.decoder(\n",
    "                    decoder_input, decoder_hidden, encoder_outputs)\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(self.device)\n",
    "                mask_loss, nTotal = maskNLLLoss(decoder_output, target[t], mask[t], self.device)\n",
    "                loss += mask_loss\n",
    "                print_losses.append(mask_loss.item() * nTotal)\n",
    "                n_totals += nTotal\n",
    "        return loss, n_totals, print_losses"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Mask Loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def maskNLLLoss(inp, target, mask, device):\n",
    "    #target = target.to(device)\n",
    "    nTotal = mask.sum()\n",
    "    crossEntropy = -torch.log(torch.gather(inp, 1, target.view(-1, 1)).squeeze(1))\n",
    "    loss = crossEntropy.masked_select(mask).mean()\n",
    "    loss = loss.to(device)\n",
    "    return loss, nTotal.item()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def train_step(input_variable, lengths, target_variable, mask, seq2seq,\n",
    "          seq2seq_optimizer, max_target_len, batch_size, clip, device):\n",
    "\n",
    "    seq2seq_optimizer.zero_grad()\n",
    "    input_variable = input_variable.to(device)\n",
    "    target_variable = target_variable.to(device)\n",
    "    mask = mask.to(device)\n",
    "    lengths = lengths.to(\"cpu\")\n",
    "    loss, n_totals, print_losses = seq2seq(input_variable, lengths, target_variable,\n",
    "                                          mask, max_target_len, batch_size)\n",
    "    loss.backward()\n",
    "    # Clip at its place\n",
    "    _ = nn.utils.clip_grad_norm_(seq2seq.parameters(), clip)\n",
    "    seq2seq_optimizer.step()\n",
    "    return sum(print_losses) / n_totals\n",
    "\n",
    "def train(voc, pairs, seq2seq, seq2seq_optimizer, embedding,\n",
    "          encoder_n_layers, decoder_n_layers, n_iteration,\n",
    "          batch_size, print_every, clip, corpus_name, device):\n",
    "\n",
    "    training_batches = [batch_to_train_data(voc, [random.choice(pairs) for _ in range(batch_size)])\n",
    "                      for _ in range(n_iteration)]\n",
    "    print('Initializing ...')\n",
    "    start_iteration = 1\n",
    "    print_loss = 0\n",
    "    print(\"Training...\")\n",
    "    for iteration in range(start_iteration, n_iteration + 1):\n",
    "        training_batch = training_batches[iteration - 1]\n",
    "        input_variable, lengths, target_variable, mask, max_target_len = training_batch\n",
    "        loss = train_step(input_variable, lengths, target_variable, mask, seq2seq,\n",
    "                     seq2seq_optimizer, max_target_len, batch_size, clip, device)\n",
    "        print_loss += loss\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"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "scrolled": true
   },
   "source": [
    "### Important Training Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "USE_CUDA = torch.cuda.is_available()\n",
    "device = torch.device(\"cuda\" if USE_CUDA else \"cpu\")\n",
    "attn_model = 'dot' #'general', 'concat'\n",
    "hidden_size = 500\n",
    "encoder_n_layers = 2\n",
    "decoder_n_layers = 2\n",
    "dropout = 0.1\n",
    "batch_size = 64\n",
    "clip = 50.0\n",
    "teacher_forcing_ratio = 0.5\n",
    "learning_rate = 0.0005\n",
    "n_iteration = 20000\n",
    "print_every = 500"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Initialize Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "embedding = nn.Embedding(voc.num_words, hidden_size)\n",
    "encoder = Encoder(hidden_size, embedding, encoder_n_layers, dropout)\n",
    "decoder = Decoder(attn_model, embedding, hidden_size, voc.num_words, decoder_n_layers, dropout)\n",
    "encoder = encoder.to(device)\n",
    "decoder = decoder.to(device)\n",
    "encoder.train()\n",
    "decoder.train()\n",
    "seq2seq = Seq2Seq(encoder, decoder, device, teacher_forcing_ratio)\n",
    "seq2seq = seq2seq.to(device)\n",
    "seq2seq.train()\n",
    "seq2seq_optimizer = optim.Adam(seq2seq.parameters(), lr=learning_rate)\n",
    "\n",
    "for state in seq2seq_optimizer.state.values():\n",
    "    for k, v in state.items():\n",
    "        if isinstance(v, torch.Tensor):\n",
    "            state[k] = v.cuda()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Start training..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "scrolled": false,
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Initializing ...\n",
      "Training...\n",
      "Iteration: 500; Percent complete: 2.5%; Average loss: 4.6804\n",
      "Iteration: 1000; Percent complete: 5.0%; Average loss: 4.3220\n",
      "Iteration: 1500; Percent complete: 7.5%; Average loss: 4.1529\n",
      "Iteration: 2000; Percent complete: 10.0%; Average loss: 4.0934\n",
      "Iteration: 2500; Percent complete: 12.5%; Average loss: 3.9846\n",
      "Iteration: 3000; Percent complete: 15.0%; Average loss: 3.9340\n",
      "Iteration: 3500; Percent complete: 17.5%; Average loss: 3.7831\n",
      "Iteration: 4000; Percent complete: 20.0%; Average loss: 3.6860\n",
      "Iteration: 4500; Percent complete: 22.5%; Average loss: 3.5487\n",
      "Iteration: 5000; Percent complete: 25.0%; Average loss: 3.4249\n",
      "Iteration: 5500; Percent complete: 27.5%; Average loss: 3.3220\n",
      "Iteration: 6000; Percent complete: 30.0%; Average loss: 3.1692\n",
      "Iteration: 6500; Percent complete: 32.5%; Average loss: 3.0810\n",
      "Iteration: 7000; Percent complete: 35.0%; Average loss: 2.8828\n",
      "Iteration: 7500; Percent complete: 37.5%; Average loss: 2.7786\n",
      "Iteration: 8000; Percent complete: 40.0%; Average loss: 2.6687\n",
      "Iteration: 8500; Percent complete: 42.5%; Average loss: 2.5362\n",
      "Iteration: 9000; Percent complete: 45.0%; Average loss: 2.3976\n",
      "Iteration: 9500; Percent complete: 47.5%; Average loss: 2.2687\n",
      "Iteration: 10000; Percent complete: 50.0%; Average loss: 2.1786\n",
      "Iteration: 10500; Percent complete: 52.5%; Average loss: 2.0713\n",
      "Iteration: 11000; Percent complete: 55.0%; Average loss: 2.0147\n",
      "Iteration: 11500; Percent complete: 57.5%; Average loss: 1.9233\n",
      "Iteration: 12000; Percent complete: 60.0%; Average loss: 1.8304\n",
      "Iteration: 12500; Percent complete: 62.5%; Average loss: 1.7118\n",
      "Iteration: 13000; Percent complete: 65.0%; Average loss: 1.6832\n",
      "Iteration: 13500; Percent complete: 67.5%; Average loss: 1.6143\n",
      "Iteration: 14000; Percent complete: 70.0%; Average loss: 1.5853\n",
      "Iteration: 14500; Percent complete: 72.5%; Average loss: 1.4780\n",
      "Iteration: 15000; Percent complete: 75.0%; Average loss: 1.4270\n",
      "Iteration: 15500; Percent complete: 77.5%; Average loss: 1.3760\n",
      "Iteration: 16000; Percent complete: 80.0%; Average loss: 1.2911\n",
      "Iteration: 16500; Percent complete: 82.5%; Average loss: 1.2702\n",
      "Iteration: 17000; Percent complete: 85.0%; Average loss: 1.2591\n",
      "Iteration: 17500; Percent complete: 87.5%; Average loss: 1.2110\n",
      "Iteration: 18000; Percent complete: 90.0%; Average loss: 1.2036\n",
      "Iteration: 18500; Percent complete: 92.5%; Average loss: 1.1438\n",
      "Iteration: 19000; Percent complete: 95.0%; Average loss: 1.1148\n",
      "Iteration: 19500; Percent complete: 97.5%; Average loss: 1.0748\n",
      "Iteration: 20000; Percent complete: 100.0%; Average loss: 1.0942\n"
     ]
    }
   ],
   "source": [
    "train(voc, pairs, seq2seq, seq2seq_optimizer,\n",
    "      embedding, encoder_n_layers, decoder_n_layers, n_iteration, batch_size,\n",
    "      print_every, clip, corpus_name, device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Evaluation\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "class GreedySearchDecoder(nn.Module):\n",
    "   \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",
    "        encoder_outputs, encoder_hidden = self.encoder(input_seq, input_length)\n",
    "        decoder_hidden = encoder_hidden[:decoder.n_layers]\n",
    "        decoder_input = torch.ones(1, 1, device=device, dtype=torch.long) * SOS_token\n",
    "        all_tokens = torch.zeros([0], device=device, dtype=torch.long)\n",
    "        all_scores = torch.zeros([0], device=device)\n",
    "        for _ in range(max_length):\n",
    "            decoder_output, decoder_hidden = self.decoder(decoder_input, decoder_hidden, encoder_outputs)\n",
    "            decoder_scores, decoder_input = torch.max(decoder_output, dim=1)\n",
    "            all_tokens = torch.cat((all_tokens, decoder_input), dim=0)\n",
    "            all_scores = torch.cat((all_scores, decoder_scores), dim=0)\n",
    "            decoder_input = torch.unsqueeze(decoder_input, 0)\n",
    "        return all_tokens, all_scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def evaluate(encoder, decoder, searcher, voc, sentence, max_length=MAX_LENGTH):\n",
    "    indexes_batch = [indexes_from_sentence(voc, sentence)]\n",
    "    lengths = torch.tensor([len(indexes) for indexes in indexes_batch])\n",
    "    input_batch = torch.LongTensor(indexes_batch).transpose(0, 1)\n",
    "    input_batch = input_batch.to(device)\n",
    "    lengths = lengths.to(\"cpu\")\n",
    "    tokens, scores = searcher(input_batch, lengths, max_length)\n",
    "    decoded_words = [voc.index2word[token.item()] for token in tokens]\n",
    "    return decoded_words\n",
    "\n",
    "def evaluateInput(encoder, decoder, searcher, voc):\n",
    "    input_sentence = ''\n",
    "    while(1):\n",
    "        try:\n",
    "            input_sentence = input('> ')\n",
    "            if input_sentence == 'q' or input_sentence == 'quit': break\n",
    "            input_sentence = normalize_string(input_sentence)\n",
    "            output_words = evaluate(encoder, decoder, searcher, voc, input_sentence)\n",
    "            output_words[:] = [x for x in output_words if not (x == 'EOS' or x == 'PAD')]\n",
    "            print('Bot:', ' '.join(output_words))\n",
    "        except KeyError:\n",
    "            print(\"Error: Encountered unknown word.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Start evaluation..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "> hi\n",
      "Bot: hi . d . your . .\n",
      "> i am your dad\n",
      "Bot: why is that guy ?\n",
      "> so tell me what is your name\n",
      "Bot: bob ! on me !\n",
      "> bob are you cute\n",
      "Bot: but !\n",
      "> what\n",
      "Bot: oh jesus . the baby the story ?\n",
      "> baby story\n",
      "Bot: you you you you you\n",
      "> what\n",
      "Bot: oh jesus . the baby the story ?\n",
      "> shutup\n",
      "Error: Encountered unknown word.\n",
      "> please go away\n",
      "Bot: i ll give you some .\n",
      "> give me what\n",
      "Bot: oh shut it !\n",
      "> you shut it\n",
      "Bot: no s it . it s . .\n",
      "> Now you are loosing your mind\n",
      "Error: Encountered unknown word.\n",
      "> are you smart\n",
      "Bot: i m trying to be fucked .\n",
      "> I can fuck you\n",
      "Bot: fuck you ! ! !\n",
      "> Shut it bitch\n",
      "Bot: fuck is it . .\n",
      "> q\n"
     ]
    }
   ],
   "source": [
    "encoder.eval()\n",
    "decoder.eval()\n",
    "searcher = GreedySearchDecoder(encoder, decoder)\n",
    "evaluateInput(encoder, decoder, searcher, voc)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## References\n",
    "\n",
    "1) Yuan-Kuei Wu’s pytorch-chatbot implementation:\n",
    "   https://github.com/ywk991112/pytorch-chatbot\n",
    "\n",
    "2) Sean Robertson’s practical-pytorch seq2seq-translation example:\n",
    "   https://github.com/spro/practical-pytorch/tree/master/seq2seq-translation\n",
    "\n",
    "3) FloydHub’s Cornell Movie Corpus preprocessing code:\n",
    "   https://github.com/floydhub/textutil-preprocess-cornell-movie-corpus\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.7.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}