main.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generative Chatbot\n",
    "\n",
    "- Prepared by Zhu Lin Ch'ng, Avinash Fernando, Matthew Stevenson and Thomas Whiteley   for the University of Liverpool CSCK507 Natural Language Processing and Understanding June 2023 group project.\n",
    "- This notebook implements a seq-2-seq recurrent neural network architecture comparing the performance with and without the use of Luong Attention.\n",
    "- The detailed training logs can be accessed at https://tensorboard.dev/experiment/5hXpB3jpRWqqWu3HnrfXGw/ \n",
    "- The model weights and demonstration can be accessed at https://huggingface.co/spaces/csck507/Seq2Seq \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Table of Contents\n",
    "\n",
    "- [Importing Libraries](#Importing-Libraries)\n",
    "- [Data Preprocessing](#Data-Preprocessing)\n",
    "    - [Importing the dataset](#Importing-the-dataset)\n",
    "    - [Preparing the dataset](#Preparing-the-dataset)\n",
    "        - [Read Data](#Read-Data)\n",
    "        - [Preprocess Data and structure it](#Preprocess-Data-and-structure-it)\n",
    "        - [Load the structured data into dataframe and index it](#Load-the-structured-data-into-dataframe-and-index-it)\n",
    "        - [Create tensors](#Create-tensors)\n",
    "        - [Split and batch the data](#Split-and-batch-the-data)\n",
    "- [Model Development](#Model-Development)\n",
    "    - [Building the seq2seq model](#Building-the-seq2seq-model)\n",
    "    - [Training the seq2seq model](#Training-the-seq2seq-model)\n",
    "- [Model Evaluation](#Model-Evaluation)\n",
    "    - [Using the seq2seq models](#Using-the-seq2seq-models)\n",
    "    - [Evaluating the seq2seq models](#Evaluating-the-seq2seq-models)\n",
    "        - [BLEU Score](#BLEU-Score)\n",
    "        - [Cosine Similarity](#Cosine-Similarity)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Importing Libraries\n",
    "\n",
    "This notebook uses the following 3rd party libraries:\n",
    "- [pytorch](https://pytorch.org/): machine learning library\n",
    "- [pandas](https://pandas.pydata.org): data analysis library\n",
    "- [numpy](https://numpy.org): mathematical function library \n",
    "- [spacy](https://spacy.io/): natural language processing library\n",
    "- [nltk](https://www.nltk.org/): natural language processing library"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import io\n",
    "import json\n",
    "import time\n",
    "import os\n",
    "import random\n",
    "import re\n",
    "import tarfile\n",
    "import unicodedata\n",
    "import zipfile\n",
    "from io import open\n",
    "from typing import Tuple\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import requests\n",
    "import spacy\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import DataLoader, TensorDataset, random_split\n",
    "from torch.utils.tensorboard import SummaryWriter\n",
    "from torch import optim\n",
    "from nltk.translate.bleu_score import corpus_bleu\n",
    "from sentence_transformers import SentenceTransformer\n",
    "from scipy.spatial.distance import cosine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialise the torch device, favouring GPU if available\n",
    "spacy.prefer_gpu()\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "\n",
    "# load english dictionary for spacy\n",
    "try:\n",
    "    spacy.load('en_core_web_sm')\n",
    "except LookupError:\n",
    "    print('Run: python -m spacy download en_core_web_sm')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data Preprocessing\n",
    "\n",
    "The [Ubuntu dialogue corpus](https://www.kaggle.com/datasets/rtatman/ubuntu-dialogue-corpus) is selected to provide conversational data for training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Importing the dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def download_file(url, dir):\n",
    "    \"\"\"\n",
    "    Download file from url\n",
    "    :param url: url of file\n",
    "    :param filename: name of file\n",
    "    :return: None\n",
    "    \"\"\"\n",
    "    r = requests.get(url)\n",
    "    if url.endswith('.tar.gz'):\n",
    "        z = tarfile.open(fileobj=io.BytesIO(r.content), mode=\"r:gz\")\n",
    "        z.extractall(dir)\n",
    "        z.close()\n",
    "    elif url.endswith('.zip'):\n",
    "        z = zipfile.ZipFile(io.BytesIO(r.content))\n",
    "        z.extractall(dir)\n",
    "    else:\n",
    "        print('Unknown file type')\n",
    "    return None\n",
    "\n",
    "def extract_zip(filename, dir):\n",
    "    \"\"\"\n",
    "    Extract zip file\n",
    "    :param filename: name of file\n",
    "    :return: None\n",
    "    \"\"\"\n",
    "    z = zipfile.ZipFile(filename)\n",
    "    z.extractall(dir)\n",
    "    return None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "datasets = {'ubuntu-dialogue': 'data/ubuntu dialogue.zip'}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ubuntu-dialogue already exists\n"
     ]
    }
   ],
   "source": [
    "# Create directory\n",
    "if not os.path.exists('data'):\n",
    "    os.makedirs('data')\n",
    "\n",
    "# Check if data is already downloaded\n",
    "for dataset, source in datasets.items():\n",
    "    if os.path.exists('data/' + dataset):\n",
    "        print(dataset + ' already exists')\n",
    "    elif dataset == 'ubuntu-dialogue':\n",
    "        ubuntu = 'data/ubuntu dialogue'\n",
    "        if os.path.exists(source):\n",
    "            os.makedirs(ubuntu)\n",
    "            extract_zip(source, ubuntu)\n",
    "            os.remove(source)\n",
    "            print(dataset + ' extracted')\n",
    "        elif os.path.exists(ubuntu):\n",
    "            print(dataset + ' already exists')\n",
    "        else:\n",
    "            kag = 'https://www.kaggle.com/datasets/rtatman/ubuntu-dialogue-corpus/download?datasetVersionNumber=2'\n",
    "            print(f'Manually download ubuntu dialogue dataset from {kag} and place in data folder')\n",
    "    else:\n",
    "        download_file(source, 'data')\n",
    "        print(dataset + ' downloaded')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Preparing the dataset\n",
    "\n",
    "\n",
    "The Ubuntu dataset is presented as dialogues with 3 entries with varying formats of questions (Q) and answers (A) including:\n",
    "- QQQ\n",
    "- AAA\n",
    "- QQA\n",
    "- QAA\n",
    "- QAQ\n",
    "\n",
    "These formats are parsed to produce quesion answer pairs as follows:\n",
    "- QQQ -> ignored\n",
    "- AAA -> ignored\n",
    "- QQA -> (Q+Q) - A\n",
    "- QAA -> Q - (A+A)\n",
    "- QAQ -> Q - A, A - Q"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Read Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(1038324, 6)\n"
     ]
    }
   ],
   "source": [
    "variants = {'small':'',\n",
    "            'medium':'_196',\n",
    "            'large':'_301'}\n",
    "ubuntufile = f'data/ubuntu dialogue/ubuntu-dialogue-corpus/dialogueText{variants[\"small\"]}.csv'\n",
    "text_df = pd.read_csv(ubuntufile)\n",
    "text_df['dialogueID'] = text_df['dialogueID'].apply(lambda x: int(x.split('.')[0]))\n",
    "print(text_df.shape)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>folder</th>\n",
       "      <th>dialogueID</th>\n",
       "      <th>date</th>\n",
       "      <th>from</th>\n",
       "      <th>to</th>\n",
       "      <th>text</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>3</td>\n",
       "      <td>126125</td>\n",
       "      <td>2008-04-23T14:55:00.000Z</td>\n",
       "      <td>bad_image</td>\n",
       "      <td>NaN</td>\n",
       "      <td>Hello folks, please help me a bit with the fol...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>3</td>\n",
       "      <td>126125</td>\n",
       "      <td>2008-04-23T14:56:00.000Z</td>\n",
       "      <td>bad_image</td>\n",
       "      <td>NaN</td>\n",
       "      <td>Did I choose a bad channel? I ask because you ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>3</td>\n",
       "      <td>126125</td>\n",
       "      <td>2008-04-23T14:57:00.000Z</td>\n",
       "      <td>lordleemo</td>\n",
       "      <td>bad_image</td>\n",
       "      <td>the second sentence is better english   and we...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>3</td>\n",
       "      <td>64545</td>\n",
       "      <td>2009-08-01T06:22:00.000Z</td>\n",
       "      <td>mechtech</td>\n",
       "      <td>NaN</td>\n",
       "      <td>Sock Puppe?t</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>3</td>\n",
       "      <td>64545</td>\n",
       "      <td>2009-08-01T06:22:00.000Z</td>\n",
       "      <td>mechtech</td>\n",
       "      <td>NaN</td>\n",
       "      <td>WTF?</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   folder  dialogueID                      date       from         to  \\\n",
       "0       3      126125  2008-04-23T14:55:00.000Z  bad_image        NaN   \n",
       "1       3      126125  2008-04-23T14:56:00.000Z  bad_image        NaN   \n",
       "2       3      126125  2008-04-23T14:57:00.000Z  lordleemo  bad_image   \n",
       "3       3       64545  2009-08-01T06:22:00.000Z   mechtech        NaN   \n",
       "4       3       64545  2009-08-01T06:22:00.000Z   mechtech        NaN   \n",
       "\n",
       "                                                text  \n",
       "0  Hello folks, please help me a bit with the fol...  \n",
       "1  Did I choose a bad channel? I ask because you ...  \n",
       "2  the second sentence is better english   and we...  \n",
       "3                                       Sock Puppe?t  \n",
       "4                                               WTF?  "
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# preview text from ubuntu dialogue dataset\n",
    "text_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# reduce df size\n",
    "# text_df = text_df[:50000]\n",
    "# text_df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Preprocess Data and structure it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "nlp = spacy.load('en_core_web_sm') # load spacy model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def unicodetoascii(text):\n",
    "    \"\"\"\n",
    "    Turn a Unicode string to plain ASCII\n",
    "\n",
    "    :param text: text to be converted\n",
    "    :return: text in ascii format\n",
    "    \"\"\"\n",
    "    normalized_text = unicodedata.normalize('NFKD', str(text))\n",
    "    ascii_text = ''.join(char for char in normalized_text if unicodedata.category(char) != 'Mn')\n",
    "    return ascii_text\n",
    "\n",
    "def preprocess_text(text, fn=unicodetoascii):\n",
    "\n",
    "    text = fn(text)\n",
    "    text = text.lower()\n",
    "    text = re.sub(r'http\\S+', '', text)\n",
    "    text = re.sub(r'[^\\x00-\\x7F]+', \"\", text) # Remove non-ASCII characters\n",
    "    text = re.sub(r\"(\\w)[!?]+(\\w)\", r'\\1\\2', text) # Remove !? between words\n",
    "    text = re.sub(r\"\\s\\s+\", r\" \", text).strip() # Remove extra spaces\n",
    "    return text\n",
    "\n",
    "def parse_dialogue(data):\n",
    "    dialogues = {}\n",
    "    df = data.copy()\n",
    "    df.reset_index(inplace=True)\n",
    "    # Group by dialogueID\n",
    "    for dialogue_id, group in df.groupby('dialogueID'):\n",
    "        sentence_pairs = {}\n",
    "        context = ''\n",
    "        previous_direction = (None, None)\n",
    "        for i, row in group.iterrows():\n",
    "            idx = row['index']\n",
    "            sender = row['from']\n",
    "            recipient = row['to']\n",
    "            response = str(row['text'])\n",
    "            direction = (sender, recipient)\n",
    "\n",
    "            if direction == previous_direction:\n",
    "                # add to the response to the previous message if the current message is consecutive\n",
    "                prev_idx = idx - 1\n",
    "                while prev_idx not in sentence_pairs:\n",
    "                    prev_idx -= 1\n",
    "                response = context + ' ' + response\n",
    "                sentence_pairs[prev_idx] = (sentence_pairs[prev_idx][0], response)\n",
    "                # sentence_pairs[-1] = (sentence_pairs[-1][0], response)\n",
    "            elif (direction == previous_direction[::-1]) or (previous_direction[1] == None) and (direction[1] == previous_direction[0]):\n",
    "                # if the current message is from the previous recipient to the previous sender\n",
    "                # if the previous message did not have a recipient, but the current message is to the previous sender\n",
    "                sentence_pairs[idx]=(context, response)\n",
    "            else:\n",
    "                sentence_pairs[idx]=(context, response)\n",
    "            \n",
    "            previous_direction = tuple(direction)\n",
    "            context = str(response) # response is the context for the next message\n",
    "        # remove the sentence pairs that does not have context but only responses\n",
    "        sentence_pairs = {k: v for k, v in sentence_pairs.items() if v[0] != ''}\n",
    "        dialogues[dialogue_id] = sentence_pairs\n",
    "\n",
    "    return dialogues"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>folder</th>\n",
       "      <th>dialogueID</th>\n",
       "      <th>date</th>\n",
       "      <th>from</th>\n",
       "      <th>to</th>\n",
       "      <th>text</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>3</td>\n",
       "      <td>126125</td>\n",
       "      <td>2008-04-23T14:55:00.000Z</td>\n",
       "      <td>bad_image</td>\n",
       "      <td>NaN</td>\n",
       "      <td>hello folks, please help me a bit with the fol...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>3</td>\n",
       "      <td>126125</td>\n",
       "      <td>2008-04-23T14:56:00.000Z</td>\n",
       "      <td>bad_image</td>\n",
       "      <td>NaN</td>\n",
       "      <td>did i choose a bad channel? i ask because you ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>3</td>\n",
       "      <td>126125</td>\n",
       "      <td>2008-04-23T14:57:00.000Z</td>\n",
       "      <td>lordleemo</td>\n",
       "      <td>bad_image</td>\n",
       "      <td>the second sentence is better english and we a...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>3</td>\n",
       "      <td>64545</td>\n",
       "      <td>2009-08-01T06:22:00.000Z</td>\n",
       "      <td>mechtech</td>\n",
       "      <td>NaN</td>\n",
       "      <td>sock puppet</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>3</td>\n",
       "      <td>64545</td>\n",
       "      <td>2009-08-01T06:22:00.000Z</td>\n",
       "      <td>mechtech</td>\n",
       "      <td>NaN</td>\n",
       "      <td>wtf?</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   folder  dialogueID                      date       from         to  \\\n",
       "0       3      126125  2008-04-23T14:55:00.000Z  bad_image        NaN   \n",
       "1       3      126125  2008-04-23T14:56:00.000Z  bad_image        NaN   \n",
       "2       3      126125  2008-04-23T14:57:00.000Z  lordleemo  bad_image   \n",
       "3       3       64545  2009-08-01T06:22:00.000Z   mechtech        NaN   \n",
       "4       3       64545  2009-08-01T06:22:00.000Z   mechtech        NaN   \n",
       "\n",
       "                                                text  \n",
       "0  hello folks, please help me a bit with the fol...  \n",
       "1  did i choose a bad channel? i ask because you ...  \n",
       "2  the second sentence is better english and we a...  \n",
       "3                                        sock puppet  \n",
       "4                                               wtf?  "
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "text_df['text'] = text_df['text'].apply(preprocess_text)\n",
    "text_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "dialogues = parse_dialogue(text_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>dialogueID</th>\n",
       "      <th>index</th>\n",
       "      <th>context</th>\n",
       "      <th>response</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>456667</td>\n",
       "      <td>also guys, i'm trying to get into my firefox p...</td>\n",
       "      <td>are you logged in as 'root' ?</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1</td>\n",
       "      <td>456668</td>\n",
       "      <td>are you logged in as 'root' ?</td>\n",
       "      <td>no.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>2</td>\n",
       "      <td>936173</td>\n",
       "      <td>lifeless : no, but i have had trouble printing...</td>\n",
       "      <td>arhh, i should point my windows machine at the...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>3</td>\n",
       "      <td>788303</td>\n",
       "      <td>gosh it's late my brains not working what's th...</td>\n",
       "      <td>tar xf blah.tar</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>4</td>\n",
       "      <td>670258</td>\n",
       "      <td>i have to install hoary in server mode because...</td>\n",
       "      <td>6^</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   dialogueID   index                                            context  \\\n",
       "0           1  456667  also guys, i'm trying to get into my firefox p...   \n",
       "1           1  456668                      are you logged in as 'root' ?   \n",
       "2           2  936173  lifeless : no, but i have had trouble printing...   \n",
       "3           3  788303  gosh it's late my brains not working what's th...   \n",
       "4           4  670258  i have to install hoary in server mode because...   \n",
       "\n",
       "                                            response  \n",
       "0                      are you logged in as 'root' ?  \n",
       "1                                                no.  \n",
       "2  arhh, i should point my windows machine at the...  \n",
       "3                                    tar xf blah.tar  \n",
       "4                                                 6^  "
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# convert nested dictionary to dataframe\n",
    "def dict_to_df(data):\n",
    "    rows = []\n",
    "    for dialogue_id, sentence_pairs in data.items():\n",
    "        for idx, pair in sentence_pairs.items():\n",
    "            rows.append([dialogue_id, idx, pair[0], pair[1]])\n",
    "    df = pd.DataFrame(rows, columns=['dialogueID', 'index', 'context', 'response'])\n",
    "    return df\n",
    "\n",
    "dialogue_df = dict_to_df(dialogues)\n",
    "dialogue_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'pandas.core.frame.DataFrame'>\n",
      "RangeIndex: 432502 entries, 0 to 432501\n",
      "Data columns (total 4 columns):\n",
      " #   Column      Non-Null Count   Dtype \n",
      "---  ------      --------------   ----- \n",
      " 0   dialogueID  432502 non-null  int64 \n",
      " 1   index       432502 non-null  int64 \n",
      " 2   context     432502 non-null  object\n",
      " 3   response    432502 non-null  object\n",
      "dtypes: int64(2), object(2)\n",
      "memory usage: 13.2+ MB\n"
     ]
    }
   ],
   "source": [
    "dialogue_df.info()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load the structured data into dataframe and index it"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- As the model contains a layer of embedding, lemmatization is not required. (words that are similar will have vectors that are close to each other)\n",
    "  - https://aclanthology.org/2021.nodalida-main.25/\n",
    "  - https://aclanthology.org/2021.nodalida-main.25.pdf\n",
    "- As one of the models will be using attention, removing stop words is not required. (attention will learn to ignore them)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "# use torch text to create vocabulary\n",
    "def tokenize(text, nlp=nlp):\n",
    "    \"\"\"\n",
    "    Tokenize text\n",
    "    :param text: text to be tokenized\n",
    "    :return: list of tokens\n",
    "    \"\"\"\n",
    "    return [tok.text for tok in nlp.tokenizer(text)]\n",
    "\n",
    "def create_mapping(df, tokenize=tokenize):\n",
    "    \"\"\"\n",
    "    Create vocabulary mapping from context and response dataframes\n",
    "    :param df_context: context dataframe\n",
    "    :param df_response: response dataframe\n",
    "    :param tokenize: tokenization function\n",
    "    :return: vocabulary mapping\n",
    "    \"\"\"\n",
    "    # Create vocabulary mapping\n",
    "    vocab = set()\n",
    "    default_tokens = ['<pad>', '<bos>', '<eos>', '<unk>']\n",
    "    start_index = len(default_tokens)\n",
    "    for context, response in zip(df['context'], df['response']):\n",
    "        vocab.update(tokenize(context))\n",
    "        vocab.update(tokenize(response))\n",
    "    word2idx = {word: start_index+idx for idx, word in enumerate(vocab)}\n",
    "    idx2word = {start_index+idx: word for idx, word in enumerate(vocab)}\n",
    "    for idx, token in enumerate(default_tokens):\n",
    "        word2idx[token] = idx\n",
    "        idx2word[idx] = token\n",
    "    return word2idx, idx2word\n",
    "\n",
    "def lookup_words(idx2word, indices):\n",
    "    \"\"\"\n",
    "    Lookup words from indices\n",
    "    :param idx2word: index to word mapping\n",
    "    :param indices: indices to be converted\n",
    "    :return: list of words\n",
    "    \"\"\"\n",
    "    return [idx2word[idx] for idx in indices]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "word2idx, idx2word = create_mapping(dialogue_df)\n",
    "word2idx['<pad>']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Map words to indices\n",
    "dialogue_df['context_idx'] = dialogue_df['context'].apply(lambda x: [word2idx[word] for word in tokenize(x)])\n",
    "dialogue_df['response_idx'] = dialogue_df['response'].apply(lambda x: [word2idx[word] for word in tokenize(x)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "# add bos and eos tokens to context and response\n",
    "bos = word2idx['<bos>']\n",
    "eos = word2idx['<eos>']\n",
    "dialogue_df['context_idx'] = dialogue_df['context_idx'].apply(lambda x: [bos] + x + [eos])\n",
    "dialogue_df['response_idx'] = dialogue_df['response_idx'].apply(lambda x: [bos] + x + [eos])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>dialogueID</th>\n",
       "      <th>index</th>\n",
       "      <th>context</th>\n",
       "      <th>response</th>\n",
       "      <th>context_idx</th>\n",
       "      <th>response_idx</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>456667</td>\n",
       "      <td>also guys, i'm trying to get into my firefox p...</td>\n",
       "      <td>are you logged in as 'root' ?</td>\n",
       "      <td>[1, 52110, 106586, 141735, 19079, 32908, 94235...</td>\n",
       "      <td>[1, 24873, 112112, 101731, 154669, 159790, 105...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1</td>\n",
       "      <td>456668</td>\n",
       "      <td>are you logged in as 'root' ?</td>\n",
       "      <td>no.</td>\n",
       "      <td>[1, 24873, 112112, 101731, 154669, 159790, 105...</td>\n",
       "      <td>[1, 116460, 50180, 2]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>2</td>\n",
       "      <td>936173</td>\n",
       "      <td>lifeless : no, but i have had trouble printing...</td>\n",
       "      <td>arhh, i should point my windows machine at the...</td>\n",
       "      <td>[1, 144960, 30968, 116460, 141735, 96323, 1907...</td>\n",
       "      <td>[1, 137852, 141735, 19079, 70554, 141721, 1405...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>3</td>\n",
       "      <td>788303</td>\n",
       "      <td>gosh it's late my brains not working what's th...</td>\n",
       "      <td>tar xf blah.tar</td>\n",
       "      <td>[1, 43384, 22524, 29425, 36278, 140538, 75262,...</td>\n",
       "      <td>[1, 4174, 141889, 116729, 2]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>4</td>\n",
       "      <td>670258</td>\n",
       "      <td>i have to install hoary in server mode because...</td>\n",
       "      <td>6^</td>\n",
       "      <td>[1, 19079, 154214, 124236, 84473, 61065, 15466...</td>\n",
       "      <td>[1, 137030, 2]</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   dialogueID   index                                            context  \\\n",
       "0           1  456667  also guys, i'm trying to get into my firefox p...   \n",
       "1           1  456668                      are you logged in as 'root' ?   \n",
       "2           2  936173  lifeless : no, but i have had trouble printing...   \n",
       "3           3  788303  gosh it's late my brains not working what's th...   \n",
       "4           4  670258  i have to install hoary in server mode because...   \n",
       "\n",
       "                                            response  \\\n",
       "0                      are you logged in as 'root' ?   \n",
       "1                                                no.   \n",
       "2  arhh, i should point my windows machine at the...   \n",
       "3                                    tar xf blah.tar   \n",
       "4                                                 6^   \n",
       "\n",
       "                                         context_idx  \\\n",
       "0  [1, 52110, 106586, 141735, 19079, 32908, 94235...   \n",
       "1  [1, 24873, 112112, 101731, 154669, 159790, 105...   \n",
       "2  [1, 144960, 30968, 116460, 141735, 96323, 1907...   \n",
       "3  [1, 43384, 22524, 29425, 36278, 140538, 75262,...   \n",
       "4  [1, 19079, 154214, 124236, 84473, 61065, 15466...   \n",
       "\n",
       "                                        response_idx  \n",
       "0  [1, 24873, 112112, 101731, 154669, 159790, 105...  \n",
       "1                              [1, 116460, 50180, 2]  \n",
       "2  [1, 137852, 141735, 19079, 70554, 141721, 1405...  \n",
       "3                       [1, 4174, 141889, 116729, 2]  \n",
       "4                                     [1, 137030, 2]  "
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dialogue_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save word2idx and idx2word\n",
    "files_to_save = ['vocab/word2idx.json', 'vocab/idx2word.json']\n",
    "if not os.path.exists('vocab'):\n",
    "    os.makedirs('vocab')\n",
    "\n",
    "for file in files_to_save:\n",
    "    if os.path.exists(file):\n",
    "        os.remove(file)\n",
    "    with open(file, 'w') as f:\n",
    "        if file == 'vocab/word2idx.json':\n",
    "            json.dump(word2idx, f)\n",
    "        elif file == 'vocab/idx2word.json':\n",
    "            json.dump(idx2word, f)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create tensors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create tensors with sos, pad, eos tokens\n",
    "def create_tensors(df, max_len=20, min_len=1):\n",
    "    \"\"\"\n",
    "    Create tensors with sos, pad, eos tokens\n",
    "    :param df: dataframe with context and response\n",
    "    :param max_len: maximum length of sequence\n",
    "    :return: tensors with sos, pad, eos tokens\n",
    "    \"\"\"\n",
    "    # Create tensors\n",
    "    context_tensor = torch.zeros((len(df), max_len), dtype=torch.long)\n",
    "    response_tensor = torch.zeros((len(df), max_len), dtype=torch.long)\n",
    "    for i, (context, response) in enumerate(zip(df['context_idx'], df['response_idx'])):\n",
    "        # Trim context and response\n",
    "        if len(context) < max_len and len(context) >= min_len and len(response) < max_len and len(response) >= min_len:\n",
    "            context_tensor[i, :len(context)] = torch.tensor(context, dtype=torch.long)\n",
    "            response_tensor[i, :len(response)] = torch.tensor(response, dtype=torch.long)\n",
    "    # remove rows with all zeros\n",
    "    context_tensor = context_tensor[~(context_tensor == 0).all(1)]\n",
    "    response_tensor = response_tensor[~(response_tensor == 0).all(1)]\n",
    "    \n",
    "    return context_tensor, response_tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Maximum length of sequence: 12 Minimum length of sequence: 3\n"
     ]
    }
   ],
   "source": [
    "# get max length of context and response\n",
    "max_len_context = max(dialogue_df['context_idx'].apply(len))\n",
    "max_len_response = max(dialogue_df['response_idx'].apply(len))\n",
    "max_len = max(max_len_context, max_len_response)\n",
    "max_len = 12 # override max_len to reduce training time\n",
    "min_len = 3\n",
    "print(f'Maximum length of sequence: {max_len}', f'Minimum length of sequence: {min_len}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- A filter is applied to remove sentences that are too long or too short.\n",
    "- Remove sentences that are too long is required as it will take a long time to train the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([71866, 12]) torch.Size([71866, 12])\n",
      "Total data size: 71866\n"
     ]
    }
   ],
   "source": [
    "context_tensor, response_tensor = create_tensors(dialogue_df, max_len=max_len, min_len=min_len)\n",
    "print(context_tensor.shape, response_tensor.shape)\n",
    "print(f'Total data size: {context_tensor.shape[0]}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Split and batch the data\n",
    "\n",
    "Data is batched to reduce memory overhead and improve the speed of training. Data presented in the tensor is padded to be of equal lenght as demonstrated below:\n",
    "\n",
    "\n",
    "```\n",
    "[<bos>,can,you,help,me,with,a,support,issue,<eos>,<pad>,<pad>,<pad>,<pad>,<pad>]\n",
    "```\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ContextReponseBatch:\n",
    "    def __init__(self, data):\n",
    "        transposed_data = list(zip(*data))\n",
    "        self.input = torch.stack(transposed_data[0], 0)\n",
    "        # self.input_mask = (self.input != 0)\n",
    "        self.target = torch.stack(transposed_data[1], 0)\n",
    "        # self.target_mask = (self.target != 0)\n",
    "\n",
    "    def pin_memory(self):\n",
    "        \"\"\"\n",
    "        Pin memory for faster data transfer to GPU\n",
    "        :return: self\n",
    "        \"\"\"\n",
    "        self.input = self.input.pin_memory()\n",
    "        # self.input_mask = self.input_mask.pin_memory()\n",
    "        self.target = self.target.pin_memory()\n",
    "        # self.target_mask = self.target_mask.pin_memory()\n",
    "        return self\n",
    "\n",
    "def collate_wrapper(batch):\n",
    "    \"\"\"\n",
    "    Wrapper for collate function\n",
    "    :param batch: batch of data\n",
    "    :return: ContextReponseBatch object\n",
    "    \"\"\"\n",
    "    return ContextReponseBatch(batch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Split data into train, validation, and test sets\n",
    "def split_data(context_tensor, response_tensor, train_ratio=0.8, val_ratio=0.1, test_ratio=0.1):\n",
    "    \"\"\"\n",
    "    Split data into train, validation, and test sets\n",
    "    :param context_tensor: context tensor\n",
    "    :param response_tensor: response tensor\n",
    "    :param train_ratio: ratio of train set\n",
    "    :param val_ratio: ratio of validation set\n",
    "    :param test_ratio: ratio of test set\n",
    "    :return: train, validation, and test sets\n",
    "    \"\"\"\n",
    "    # Split data into train, validation, and test sets\n",
    "    dataset = TensorDataset(context_tensor, response_tensor)\n",
    "    train_size = int(train_ratio * len(dataset))\n",
    "    val_size = int(val_ratio * len(dataset))\n",
    "    test_size = len(dataset) - train_size - val_size\n",
    "    train_set, val_set, test_set = random_split(dataset, [train_size, val_size, test_size], generator=torch.Generator().manual_seed(42))\n",
    "    return train_set, val_set, test_set\n",
    "\n",
    "# Batch data\n",
    "def batch_data(train_set, val_set, test_set, batch_size=8, fn=collate_wrapper):\n",
    "    \"\"\"\n",
    "    Batch data\n",
    "    :param train_set: train set\n",
    "    :param val_set: validation set\n",
    "    :param test_set: test set\n",
    "    :param batch_size: batch size\n",
    "    :return: train, validation, and test loaders\n",
    "    \"\"\"\n",
    "    # Batch data\n",
    "    train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True, collate_fn=collate_wrapper)\n",
    "    val_loader = DataLoader(val_set, batch_size=batch_size, shuffle=True, collate_fn=collate_wrapper)\n",
    "    test_loader = DataLoader(test_set, batch_size=batch_size, shuffle=True, collate_fn=collate_wrapper)\n",
    "    return train_loader, val_loader, test_loader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create train, validation, and test sets\n",
    "train_set, val_set, test_set = split_data(context_tensor, response_tensor)\n",
    "\n",
    "# Batch data\n",
    "train_loader, val_loader, test_loader = batch_data(train_set, val_set, test_set, batch_size=16)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([16, 12])"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# preview shape of batch\n",
    "next(iter(train_loader)).input.shape # (batch_size, max_len)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Model Development\n",
    "\n",
    "The seq-2-seq model is based on a reference implementation provided by pytorch for language transaltion https://pytorch.org/tutorials/beginner/torchtext_translation_tutorial.html. It include the following key components:\n",
    "- Gated recurrent unit (GRU) RNN encoder\n",
    "- GRU RNN decoder\n",
    "- Luong attention\n",
    "\n",
    "Gradient clipping is used to avoid exploding gradients, and dropout layers are used to avoid overfitting."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Building the seq2seq model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "def init_weights(m):\n",
    "    \"\"\"\n",
    "    Initialize weights\n",
    "    :param m: model\n",
    "    :return: None\n",
    "    \"\"\"\n",
    "    for name, param in m.named_parameters():\n",
    "        if 'weight' in name:\n",
    "            nn.init.xavier_uniform_(param.data)\n",
    "        elif 'bias' in name:\n",
    "            nn.init.constant_(param.data, 0)\n",
    "    return None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Encoder(nn.Module):\n",
    "    \"\"\"\n",
    "    GRU RNN Encoder\n",
    "    \"\"\"\n",
    "    def __init__(self,\n",
    "                 input_dim: int,\n",
    "                 emb_dim: int,\n",
    "                 enc_hid_dim: int,\n",
    "                 dec_hid_dim: int,\n",
    "                 dropout: float = 0):\n",
    "        super(Encoder, self).__init__()\n",
    "\n",
    "        # dimension of imput\n",
    "        self.input_dim = input_dim\n",
    "        # dimension of embedding layer\n",
    "        self.emb_dim = emb_dim\n",
    "        # dimension of encoding hidden layer\n",
    "        self.enc_hid_dim = enc_hid_dim\n",
    "        # dimension of decoding hidden layer\n",
    "        self.dec_hid_dim = dec_hid_dim\n",
    "\n",
    "        # create embedding layer use to train embedding representations of the corpus\n",
    "        self.embedding = nn.Embedding(input_dim, emb_dim)\n",
    "\n",
    "        # use GRU for RNN\n",
    "        self.rnn = nn.GRU(emb_dim, enc_hid_dim, bidirectional=True, batch_first=False, num_layers=1)\n",
    "        self.fc = nn.Linear(enc_hid_dim * 2, dec_hid_dim)\n",
    "        # create dropout layer which will help produce a more generalisable model\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, src: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:\n",
    "        # apply dropout to the embedding layer\n",
    "        embedded = self.dropout(self.embedding(src))\n",
    "        # generate an output and hidden layer from the rnn\n",
    "        outputs, hidden = self.rnn(embedded)\n",
    "        hidden = torch.tanh(self.fc(torch.cat((hidden[-2, :, :], hidden[-1, :, :]), dim=1)))\n",
    "        return outputs, hidden\n",
    "\n",
    "\n",
    "class Attention(nn.Module):\n",
    "    \"\"\"\n",
    "    Luong attention\n",
    "    \"\"\"\n",
    "    def __init__(self,\n",
    "                 enc_hid_dim: int,\n",
    "                 dec_hid_dim: int,\n",
    "                 attn_dim: int):\n",
    "        super(Attention, self).__init__()\n",
    "\n",
    "        # dimension of encoding hidden layer\n",
    "        self.enc_hid_dim = enc_hid_dim\n",
    "        # dimension of decoding hidden layer\n",
    "        self.dec_hid_dim = dec_hid_dim\n",
    "        self.attn_in = (enc_hid_dim * 2) + dec_hid_dim\n",
    "\n",
    "        self.attn = nn.Linear(self.attn_in, attn_dim)\n",
    "\n",
    "    def forward(self,\n",
    "                decoder_hidden: torch.Tensor,\n",
    "                encoder_outputs: torch.Tensor) -> torch.Tensor:\n",
    "\n",
    "        src_len = encoder_outputs.shape[0]\n",
    "        repeated_decoder_hidden = decoder_hidden.unsqueeze(1).repeat(1, src_len, 1)\n",
    "        encoder_outputs = encoder_outputs.permute(1, 0, 2)\n",
    "        # Luong attention\n",
    "        energy = torch.tanh(self.attn(torch.cat((repeated_decoder_hidden, encoder_outputs), dim=2)))\n",
    "        attention = torch.sum(energy, dim=2)\n",
    "\n",
    "        return F.softmax(attention, dim=1)\n",
    "\n",
    "\n",
    "class AttnDecoder(nn.Module):\n",
    "    \"\"\"\n",
    "    GRU RNN Decoder with attention\n",
    "    \"\"\"\n",
    "    def __init__(self,\n",
    "                 output_dim: int,\n",
    "                 emb_dim: int,\n",
    "                 enc_hid_dim: int,\n",
    "                 dec_hid_dim: int,\n",
    "                 attention: nn.Module,\n",
    "                 dropout: float = 0):\n",
    "        super(AttnDecoder, self).__init__()\n",
    "\n",
    "        # dimention of output layer\n",
    "        self.output_dim = output_dim\n",
    "        # dimention of embedding layer\n",
    "        self.emb_dim = emb_dim\n",
    "        # dimention of encoding hidden layer\n",
    "        self.enc_hid_dim = enc_hid_dim\n",
    "        # dimention of decoding hidden layer\n",
    "        self.dec_hid_dim = dec_hid_dim\n",
    "        # drouput rate\n",
    "        self.dropout = dropout\n",
    "        # attention layer\n",
    "        self.attention = attention\n",
    "\n",
    "        # create embedding layer use to train embedding representations of the corpus\n",
    "        self.embedding = nn.Embedding(output_dim, emb_dim)\n",
    "        # use GRU for RNN\n",
    "        self.rnn = nn.GRU((enc_hid_dim * 2) + emb_dim, dec_hid_dim, batch_first=False, num_layers=1)\n",
    "        self.out = nn.Linear(self.attention.attn_in + emb_dim, output_dim)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def encode_attention(self,\n",
    "                              decoder_hidden: torch.Tensor,\n",
    "                              encoder_outputs: torch.Tensor) -> torch.Tensor:\n",
    "\n",
    "        a = self.attention(decoder_hidden, encoder_outputs)\n",
    "        a = a.unsqueeze(1)\n",
    "        encoder_outputs = encoder_outputs.permute(1, 0, 2)\n",
    "        weighted_encoder_rep = torch.bmm(a, encoder_outputs)\n",
    "        weighted_encoder_rep = weighted_encoder_rep.permute(1, 0, 2)\n",
    "        return weighted_encoder_rep\n",
    "\n",
    "    def forward(self,\n",
    "                input: torch.Tensor,\n",
    "                decoder_hidden: torch.Tensor,\n",
    "                encoder_outputs: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:\n",
    "\n",
    "        input = input.unsqueeze(0)\n",
    "        # apply dropout to embedding layer\n",
    "        embedded = self.dropout(self.embedding(input))\n",
    "        weighted_encoder = self.encode_attention(decoder_hidden, encoder_outputs)\n",
    "        \n",
    "        # generate an output and hidden layer from the rnn\n",
    "        rnn_input = torch.cat((embedded, weighted_encoder), dim=2)\n",
    "        output, decoder_hidden = self.rnn(rnn_input, decoder_hidden.unsqueeze(0))\n",
    "\n",
    "        embedded = embedded.squeeze(0)\n",
    "        output = output.squeeze(0)\n",
    "        weighted_encoder = weighted_encoder.squeeze(0)\n",
    "        output = self.out(torch.cat((output, weighted_encoder, embedded), dim=1))\n",
    "        return output, decoder_hidden.squeeze(0)\n",
    "\n",
    "class Decoder(nn.Module):\n",
    "    \"\"\"\n",
    "    GRU RNN Decoder without attention\n",
    "    \"\"\"\n",
    "    def __init__(self,\n",
    "                 output_dim: int,\n",
    "                 emb_dim: int,\n",
    "                 enc_hid_dim: int,\n",
    "                 dec_hid_dim: int,\n",
    "                 dropout: float = 0):\n",
    "        super(Decoder, self).__init__()\n",
    "\n",
    "        # dimention of output layer\n",
    "        self.output_dim = output_dim\n",
    "        # dimention of embedding layer\n",
    "        self.emb_dim = emb_dim\n",
    "        # dimention of encoding hidden layer\n",
    "        self.enc_hid_dim = enc_hid_dim\n",
    "        # dimention of decoding hidden layer\n",
    "        self.dec_hid_dim = dec_hid_dim\n",
    "        # drouput rate\n",
    "        self.dropout = dropout\n",
    "\n",
    "        # create embedding layer use to train embedding representations of the corpus\n",
    "        self.embedding = nn.Embedding(output_dim, emb_dim)\n",
    "        # GRU RNN\n",
    "        self.rnn = nn.GRU((enc_hid_dim * 2) + emb_dim, dec_hid_dim, batch_first=False, num_layers=1)\n",
    "        self.out = nn.Linear((enc_hid_dim * 2) + dec_hid_dim + emb_dim, output_dim)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self,\n",
    "                input: torch.Tensor,\n",
    "                decoder_hidden: torch.Tensor,\n",
    "                encoder_outputs: torch.Tensor) -> Tuple[torch.Tensor\n",
    "                                                        , torch.Tensor]:\n",
    "        \n",
    "        input = input.unsqueeze(0)\n",
    "        # apply dropout to embedding layer\n",
    "        embedded = self.dropout(self.embedding(input))\n",
    "        context = encoder_outputs[-1,:,:]\n",
    "        context = context.repeat(embedded.shape[0], 1, 1)\n",
    "        embs_and_context = torch.cat((embedded, context), -1)\n",
    "        # generate an output and hidden layer from the rnn\n",
    "        output, decoder_hidden = self.rnn(embs_and_context, decoder_hidden.unsqueeze(0))\n",
    "        embedded = embedded.squeeze(0)\n",
    "        output = output.squeeze(0)\n",
    "        context = context.squeeze(0)\n",
    "        output = self.out(torch.cat((output, embedded, context), -1))\n",
    "        return output, decoder_hidden.squeeze(0)\n",
    "\n",
    "class Seq2Seq(nn.Module):\n",
    "    \"\"\"\n",
    "    Seq-2-Seq model combining RNN encoder and RNN decoder\n",
    "    \"\"\"\n",
    "    def __init__(self,\n",
    "                 encoder: nn.Module,\n",
    "                 decoder: nn.Module,\n",
    "                 device: torch.device):\n",
    "        super(Seq2Seq, self).__init__()\n",
    "\n",
    "        self.encoder = encoder\n",
    "        self.decoder = decoder\n",
    "        self.device = device\n",
    "\n",
    "    def forward(self,\n",
    "                src: torch.Tensor,\n",
    "                trg: torch.Tensor,\n",
    "                teacher_forcing_ratio: float = 0.5) -> torch.Tensor:\n",
    "        src = src.transpose(0, 1) # (max_len, batch_size)\n",
    "        trg = trg.transpose(0, 1) # (max_len, batch_size)\n",
    "        batch_size = src.shape[1]\n",
    "        max_len = trg.shape[0]\n",
    "        trg_vocab_size = self.decoder.output_dim\n",
    "\n",
    "        outputs = torch.zeros(max_len, batch_size, trg_vocab_size).to(self.device)\n",
    "        encoder_outputs, hidden = self.encoder(src)\n",
    "\n",
    "        # first input to the decoder is the <sos> token\n",
    "        output = trg[0,:]\n",
    "\n",
    "        for t in range(1, max_len):\n",
    "            output, hidden = self.decoder(output, hidden, encoder_outputs)\n",
    "            outputs[t] = output\n",
    "            teacher_force = random.random() < teacher_forcing_ratio\n",
    "            top1 = output.max(1)[1]\n",
    "            output = trg[t] if teacher_force else top1\n",
    "\n",
    "        return outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([16, 12]), torch.Size([16, 12]))"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test_batch = next(iter(train_loader))\n",
    "test_batch.input.shape, test_batch.target.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "# enc = Encoder(input_dim=len(word2idx), emb_dim=256, enc_hid_dim=512, dec_hid_dim=512)\n",
    "# attn = Attention(enc_hid_dim=512, dec_hid_dim=512, attn_dim=64)\n",
    "# dec = AttnDecoder(output_dim=len(word2idx), emb_dim=256, enc_hid_dim=512, dec_hid_dim=512, attention=attn)\n",
    "# model = Seq2Seq(encoder=enc, decoder=dec, device=device)\n",
    "# model.apply(init_weights)\n",
    "# model.to(device)\n",
    "# model.train()\n",
    "# optimizer = optim.Adam(model.parameters(), lr=0.001)\n",
    "# optimizer.zero_grad()\n",
    "# model(test_batch.input.to(device), test_batch.target.to(device), teacher_forcing_ratio=0.5).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "162004\n"
     ]
    }
   ],
   "source": [
    "print(len(word2idx))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "# enc = Encoder(input_dim=len(word2idx), emb_dim=256, enc_hid_dim=512, dec_hid_dim=512)\n",
    "# dec = Decoder(output_dim=len(word2idx), emb_dim=256, enc_hid_dim=512, dec_hid_dim=512)\n",
    "# model = Seq2Seq(encoder=enc, decoder=dec, device=device)\n",
    "# model.apply(init_weights)\n",
    "# model.to(device)\n",
    "# model.train()\n",
    "# optimizer = optim.Adam(model.parameters(), lr=0.001)\n",
    "# optimizer.zero_grad()\n",
    "# model(test_batch.input.to(device), test_batch.target.to(device), teacher_forcing_ratio=0.5).shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training the seq2seq model\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "PAD_INDEX = word2idx['<pad>']\n",
    "criterion = nn.CrossEntropyLoss(ignore_index=PAD_INDEX)\n",
    "\n",
    "def train_model(model, train_loader, val_loader, optimizer,\n",
    "                run_name = 'seq2seq', init_weights=init_weights, device=device,\n",
    "                n_epochs=10, clip=1, criterion=criterion,\n",
    "                teacher_forcing_ratio=0.5, params=None):\n",
    "    \"\"\"\n",
    "    Train model\n",
    "    :param model: model\n",
    "    :param train_loader: train loader\n",
    "    :param val_loader: validation loader\n",
    "    :param optimizer: optimizer\n",
    "    :param n_epochs: number of epochs\n",
    "    :param clip: clip\n",
    "    :param criterion: loss function\n",
    "    :param PAD_INDEX: index for pad token\n",
    "    :param teacher_forcing_ratio: teacher forcing ratio\n",
    "    :return: model, train loss, validation loss\n",
    "    \"\"\"\n",
    "    if not os.path.exists('models'):\n",
    "        os.makedirs('models')\n",
    "    model = model.to(device)\n",
    "    model.apply(init_weights)\n",
    "    hyperparams = {'n_epochs': n_epochs,\n",
    "                    'clip': clip,\n",
    "                    'teacher_forcing_ratio': teacher_forcing_ratio}\n",
    "    if params:\n",
    "        hyperparams.update(params)\n",
    "    writer = SummaryWriter(f'runs/{run_name}')\n",
    "    train_loss = []\n",
    "    val_loss = []\n",
    "    for epoch in range(n_epochs):\n",
    "        model.train()\n",
    "        start_time = time.time()\n",
    "        epoch_loss = 1\n",
    "        val_epoch_loss = 1\n",
    "        for i, batch in enumerate(train_loader):\n",
    "            src = batch.input.to(device)\n",
    "            trg = batch.target.to(device)\n",
    "\n",
    "            optimizer.zero_grad()\n",
    "            output = model(src, trg, teacher_forcing_ratio)\n",
    "            \n",
    "            output_dim = output.shape[-1]\n",
    "            output = output[1:].view(-1, output_dim)\n",
    "            trg = trg.transpose(0, 1)\n",
    "            trg = trg[1:].reshape(-1)\n",
    "\n",
    "            loss = criterion(output, trg)\n",
    "            loss.to(device).backward()\n",
    "            torch.nn.utils.clip_grad_norm_(model.parameters(), clip) # clip gradients\n",
    "            optimizer.step() # update parameters\n",
    "            epoch_loss += loss.item() # update epoch loss\n",
    "            writer.add_scalar('Train Loss', loss.item(), epoch * len(train_loader) + i)\n",
    "        train_loss.append(epoch_loss / len(train_loader))\n",
    "        # save model with datetime and epoch\n",
    "        torch.save(model.state_dict(), f'models/{run_name}_epoch{epoch+1}.pt')\n",
    "        # remove previous model\n",
    "        if epoch > 0:\n",
    "            os.remove(f'models/{run_name}_epoch{epoch}.pt')\n",
    "        model.eval()\n",
    "        with torch.no_grad():\n",
    "            for i, batch in enumerate(val_loader):\n",
    "                src = batch.input.to(device)\n",
    "                trg = batch.target.to(device)\n",
    "                \n",
    "                output = model(src, trg, teacher_forcing_ratio)\n",
    "                output_dim = output.shape[-1]\n",
    "                output = output[1:].view(-1, output_dim)\n",
    "                trg = trg.transpose(0, 1)\n",
    "                trg = trg[1:].reshape(-1)\n",
    "                \n",
    "                loss = criterion(output, trg)\n",
    "                val_epoch_loss += loss.item()\n",
    "        val_loss.append(val_epoch_loss / len(val_loader))\n",
    "        writer.add_scalar('Validation Loss', val_epoch_loss / len(val_loader), epoch)\n",
    "        duration = time.time()-start_time\n",
    "        remaining = duration * (n_epochs - epoch - 1)\n",
    "        print(f'Epoch: {epoch+1:02} | Time: {duration:.3f}s | Train Loss: {epoch_loss/len(train_loader):.3f} | Val Loss: {val_epoch_loss/len(val_loader):.3f} | Remaining: {remaining//60:.0f}m {remaining%60:.0f}s')\n",
    "    writer.add_hparams(hyperparams, {'hparam/train_loss': train_loss[-1], 'hparam/val_loss': val_loss[-1]})\n",
    "    writer.close()\n",
    "    return model, train_loss, val_loss\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "params = {'input_dim': len(word2idx),\n",
    "            'emb_dim': 192,\n",
    "            'enc_hid_dim': 256,\n",
    "            'dec_hid_dim': 256,\n",
    "            'dropout': 0.5,\n",
    "            'attn_dim': 64,\n",
    "            'teacher_forcing_ratio': 0.5,\n",
    "            'epochs': 35}\n",
    "\n",
    "batch_size = 18"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training set size: 57492 Validation set size: 7186 Test set size: 7188\n"
     ]
    }
   ],
   "source": [
    "# Create train, validation, and test sets\n",
    "train_set, val_set, test_set = split_data(context_tensor, response_tensor)\n",
    "\n",
    "# Batch data\n",
    "train_loader, val_loader, test_loader = batch_data(train_set, val_set, test_set, batch_size=batch_size)\n",
    "\n",
    "print(f'Training set size: {len(train_set)}', f'Validation set size: {len(val_set)}', f'Test set size: {len(test_set)}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model has 219,456,724 trainable parameters\n",
      "Epoch: 01 | Time: 409.775s | Train Loss: 5.456 | Val Loss: 5.150 | Remaining: 232m 12s\n",
      "Epoch: 02 | Time: 409.705s | Train Loss: 4.895 | Val Loss: 5.067 | Remaining: 225m 20s\n",
      "Epoch: 03 | Time: 410.125s | Train Loss: 4.655 | Val Loss: 5.019 | Remaining: 218m 44s\n",
      "Epoch: 04 | Time: 410.241s | Train Loss: 4.430 | Val Loss: 5.063 | Remaining: 211m 57s\n",
      "Epoch: 05 | Time: 410.065s | Train Loss: 4.207 | Val Loss: 5.141 | Remaining: 205m 2s\n",
      "Epoch: 06 | Time: 409.643s | Train Loss: 3.973 | Val Loss: 5.240 | Remaining: 197m 60s\n",
      "Epoch: 07 | Time: 409.830s | Train Loss: 3.774 | Val Loss: 5.392 | Remaining: 191m 15s\n",
      "Epoch: 08 | Time: 409.824s | Train Loss: 3.586 | Val Loss: 5.486 | Remaining: 184m 25s\n",
      "Epoch: 09 | Time: 409.562s | Train Loss: 3.419 | Val Loss: 5.561 | Remaining: 177m 29s\n",
      "Epoch: 10 | Time: 410.796s | Train Loss: 3.287 | Val Loss: 5.660 | Remaining: 171m 10s\n",
      "Epoch: 11 | Time: 410.472s | Train Loss: 3.189 | Val Loss: 5.777 | Remaining: 164m 11s\n",
      "Epoch: 12 | Time: 410.119s | Train Loss: 3.082 | Val Loss: 5.810 | Remaining: 157m 13s\n",
      "Epoch: 13 | Time: 410.192s | Train Loss: 3.009 | Val Loss: 5.880 | Remaining: 150m 24s\n",
      "Epoch: 14 | Time: 410.212s | Train Loss: 2.933 | Val Loss: 5.933 | Remaining: 143m 34s\n",
      "Epoch: 15 | Time: 409.900s | Train Loss: 2.877 | Val Loss: 5.978 | Remaining: 136m 38s\n",
      "Epoch: 16 | Time: 410.961s | Train Loss: 2.818 | Val Loss: 5.996 | Remaining: 130m 8s\n",
      "Epoch: 17 | Time: 410.931s | Train Loss: 2.784 | Val Loss: 6.052 | Remaining: 123m 17s\n",
      "Epoch: 18 | Time: 411.111s | Train Loss: 2.726 | Val Loss: 6.120 | Remaining: 116m 29s\n",
      "Epoch: 19 | Time: 411.122s | Train Loss: 2.690 | Val Loss: 6.130 | Remaining: 109m 38s\n",
      "Epoch: 20 | Time: 411.253s | Train Loss: 2.657 | Val Loss: 6.201 | Remaining: 102m 49s\n",
      "Epoch: 21 | Time: 410.999s | Train Loss: 2.624 | Val Loss: 6.244 | Remaining: 95m 54s\n",
      "Epoch: 22 | Time: 411.017s | Train Loss: 2.598 | Val Loss: 6.247 | Remaining: 89m 3s\n",
      "Epoch: 23 | Time: 410.866s | Train Loss: 2.550 | Val Loss: 6.295 | Remaining: 82m 10s\n",
      "Epoch: 24 | Time: 410.783s | Train Loss: 2.537 | Val Loss: 6.284 | Remaining: 75m 19s\n",
      "Epoch: 25 | Time: 410.660s | Train Loss: 2.518 | Val Loss: 6.342 | Remaining: 68m 27s\n",
      "Epoch: 26 | Time: 410.904s | Train Loss: 2.481 | Val Loss: 6.387 | Remaining: 61m 38s\n",
      "Epoch: 27 | Time: 411.135s | Train Loss: 2.468 | Val Loss: 6.429 | Remaining: 54m 49s\n",
      "Epoch: 28 | Time: 411.056s | Train Loss: 2.453 | Val Loss: 6.412 | Remaining: 47m 57s\n",
      "Epoch: 29 | Time: 411.269s | Train Loss: 2.430 | Val Loss: 6.471 | Remaining: 41m 8s\n",
      "Epoch: 30 | Time: 410.948s | Train Loss: 2.432 | Val Loss: 6.443 | Remaining: 34m 15s\n",
      "Epoch: 31 | Time: 411.371s | Train Loss: 2.408 | Val Loss: 6.512 | Remaining: 27m 25s\n",
      "Epoch: 32 | Time: 410.756s | Train Loss: 2.382 | Val Loss: 6.515 | Remaining: 20m 32s\n",
      "Epoch: 33 | Time: 410.915s | Train Loss: 2.365 | Val Loss: 6.514 | Remaining: 13m 42s\n",
      "Epoch: 34 | Time: 411.123s | Train Loss: 2.349 | Val Loss: 6.544 | Remaining: 6m 51s\n",
      "Epoch: 35 | Time: 411.183s | Train Loss: 2.341 | Val Loss: 6.583 | Remaining: 0m 0s\n"
     ]
    }
   ],
   "source": [
    "enc = Encoder(input_dim=params['input_dim'], emb_dim=params['emb_dim'], enc_hid_dim=params['enc_hid_dim'], dec_hid_dim=params['dec_hid_dim'], dropout=params['dropout'])\n",
    "dec = Decoder(output_dim=params['input_dim'], emb_dim=params['emb_dim'], enc_hid_dim=params['enc_hid_dim'], dec_hid_dim=params['dec_hid_dim'], dropout=params['dropout'])\n",
    "model = Seq2Seq(encoder=enc, decoder=dec, device=device)\n",
    "num_params = sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
    "print(f'Model has {num_params:,} trainable parameters')\n",
    "optimizer = optim.Adam(model.parameters())\n",
    "\n",
    "norm_model, train_loss, val_loss = train_model(model, train_loader, val_loader, optimizer,\n",
    "                                               run_name=f'NormSeq2Seq-{int(num_params/1e6)}M', n_epochs=params['epochs'],\n",
    "                                               teacher_forcing_ratio=params['teacher_forcing_ratio'],\n",
    "                                               params=params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model has 219,505,940 trainable parameters\n",
      "Epoch: 01 | Time: 416.459s | Train Loss: 5.428 | Val Loss: 5.075 | Remaining: 235m 60s\n",
      "Epoch: 02 | Time: 416.565s | Train Loss: 4.868 | Val Loss: 5.010 | Remaining: 229m 7s\n",
      "Epoch: 03 | Time: 416.400s | Train Loss: 4.580 | Val Loss: 5.030 | Remaining: 222m 5s\n",
      "Epoch: 04 | Time: 416.621s | Train Loss: 4.318 | Val Loss: 5.074 | Remaining: 215m 15s\n",
      "Epoch: 05 | Time: 416.376s | Train Loss: 4.043 | Val Loss: 5.193 | Remaining: 208m 11s\n",
      "Epoch: 06 | Time: 415.696s | Train Loss: 3.771 | Val Loss: 5.344 | Remaining: 200m 55s\n",
      "Epoch: 07 | Time: 415.791s | Train Loss: 3.520 | Val Loss: 5.454 | Remaining: 194m 2s\n",
      "Epoch: 08 | Time: 415.745s | Train Loss: 3.321 | Val Loss: 5.593 | Remaining: 187m 5s\n",
      "Epoch: 09 | Time: 416.094s | Train Loss: 3.150 | Val Loss: 5.690 | Remaining: 180m 18s\n",
      "Epoch: 10 | Time: 416.417s | Train Loss: 3.024 | Val Loss: 5.756 | Remaining: 173m 30s\n",
      "Epoch: 11 | Time: 416.098s | Train Loss: 2.905 | Val Loss: 5.847 | Remaining: 166m 26s\n",
      "Epoch: 12 | Time: 416.880s | Train Loss: 2.821 | Val Loss: 5.909 | Remaining: 159m 48s\n",
      "Epoch: 13 | Time: 416.091s | Train Loss: 2.736 | Val Loss: 5.990 | Remaining: 152m 34s\n",
      "Epoch: 14 | Time: 416.460s | Train Loss: 2.673 | Val Loss: 6.080 | Remaining: 145m 46s\n",
      "Epoch: 15 | Time: 416.391s | Train Loss: 2.612 | Val Loss: 6.117 | Remaining: 138m 48s\n",
      "Epoch: 16 | Time: 415.961s | Train Loss: 2.565 | Val Loss: 6.174 | Remaining: 131m 43s\n",
      "Epoch: 17 | Time: 415.869s | Train Loss: 2.513 | Val Loss: 6.236 | Remaining: 124m 46s\n",
      "Epoch: 18 | Time: 415.797s | Train Loss: 2.476 | Val Loss: 6.268 | Remaining: 117m 49s\n",
      "Epoch: 19 | Time: 416.003s | Train Loss: 2.437 | Val Loss: 6.315 | Remaining: 110m 56s\n",
      "Epoch: 20 | Time: 416.273s | Train Loss: 2.406 | Val Loss: 6.390 | Remaining: 104m 4s\n",
      "Epoch: 21 | Time: 414.519s | Train Loss: 2.373 | Val Loss: 6.410 | Remaining: 96m 43s\n",
      "Epoch: 22 | Time: 413.022s | Train Loss: 2.348 | Val Loss: 6.478 | Remaining: 89m 29s\n",
      "Epoch: 23 | Time: 413.122s | Train Loss: 2.330 | Val Loss: 6.454 | Remaining: 82m 37s\n",
      "Epoch: 24 | Time: 413.658s | Train Loss: 2.304 | Val Loss: 6.520 | Remaining: 75m 50s\n",
      "Epoch: 25 | Time: 413.456s | Train Loss: 2.286 | Val Loss: 6.562 | Remaining: 68m 55s\n",
      "Epoch: 26 | Time: 413.281s | Train Loss: 2.260 | Val Loss: 6.541 | Remaining: 61m 60s\n",
      "Epoch: 27 | Time: 413.164s | Train Loss: 2.245 | Val Loss: 6.612 | Remaining: 55m 5s\n",
      "Epoch: 28 | Time: 413.426s | Train Loss: 2.228 | Val Loss: 6.630 | Remaining: 48m 14s\n",
      "Epoch: 29 | Time: 414.064s | Train Loss: 2.208 | Val Loss: 6.661 | Remaining: 41m 24s\n",
      "Epoch: 30 | Time: 416.876s | Train Loss: 2.206 | Val Loss: 6.720 | Remaining: 34m 44s\n",
      "Epoch: 31 | Time: 416.488s | Train Loss: 2.170 | Val Loss: 6.720 | Remaining: 27m 46s\n",
      "Epoch: 32 | Time: 416.957s | Train Loss: 2.173 | Val Loss: 6.707 | Remaining: 20m 51s\n",
      "Epoch: 33 | Time: 417.002s | Train Loss: 2.160 | Val Loss: 6.755 | Remaining: 13m 54s\n",
      "Epoch: 34 | Time: 417.000s | Train Loss: 2.147 | Val Loss: 6.824 | Remaining: 6m 57s\n",
      "Epoch: 35 | Time: 417.396s | Train Loss: 2.138 | Val Loss: 6.802 | Remaining: 0m 0s\n"
     ]
    }
   ],
   "source": [
    "enc = Encoder(input_dim=params['input_dim'], emb_dim=params['emb_dim'], enc_hid_dim=params['enc_hid_dim'], dec_hid_dim=params['dec_hid_dim'], dropout=params['dropout'])\n",
    "attn = Attention(enc_hid_dim=params['enc_hid_dim'], dec_hid_dim=params['dec_hid_dim'], attn_dim=params['attn_dim'])\n",
    "dec = AttnDecoder(output_dim=params['input_dim'], emb_dim=params['emb_dim'], enc_hid_dim=params['enc_hid_dim'], dec_hid_dim=params['dec_hid_dim'], attention=attn, dropout=params['dropout'])\n",
    "model = Seq2Seq(encoder=enc, decoder=dec, device=device)\n",
    "num_params = sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
    "print(f'Model has {num_params:,} trainable parameters')\n",
    "optimizer = optim.Adam(model.parameters())\n",
    "\n",
    "attn_model, train_loss, val_loss = train_model(model, train_loader, val_loader, optimizer,\n",
    "                                               run_name=f'AttnSeq2Seq-{int(num_params/1e6)}M', n_epochs=params['epochs'],\n",
    "                                               teacher_forcing_ratio=params['teacher_forcing_ratio'],\n",
    "                                               params=params)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Model Evaluation\n",
    "\n",
    "The model evaluated against the following metrics:\n",
    "- Bleu score: measures the similarity of the input text to the validation set\n",
    "- Cosine similarity: measures the cosine similarity of the embedding representation of text to the validation set"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using the seq2seq models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate(model, sentence, max_len=30, word2idx=word2idx, idx2word=idx2word, device=device, tokenize=tokenize, preprocess_text=preprocess_text, lookup_words=lookup_words):\n",
    "    \"\"\"\n",
    "    Generate response\n",
    "    :param model: model\n",
    "    :param sentence: sentence\n",
    "    :param max_len: maximum length of sequence\n",
    "    :param word2idx: word to index mapping\n",
    "    :param idx2word: index to word mapping\n",
    "    :return: response\n",
    "    \"\"\"\n",
    "    model.eval()\n",
    "    sentence = preprocess_text(sentence)\n",
    "    tokens = tokenize(sentence)\n",
    "    tokens = [word2idx[token] if token in word2idx else word2idx['<unk>'] for token in tokens]\n",
    "    tokens = [word2idx['<bos>']] + tokens + [word2idx['<eos>']]\n",
    "    tokens = torch.tensor(tokens, dtype=torch.long).unsqueeze(1).to(device)\n",
    "    outputs = [word2idx['<bos>']]\n",
    "    with torch.no_grad():\n",
    "        encoder_outputs, hidden = model.encoder(tokens)\n",
    "    for t in range(max_len):\n",
    "        output, hidden = model.decoder(torch.tensor([outputs[-1]], dtype=torch.long).to(device), hidden, encoder_outputs)\n",
    "        top1 = output.max(1)[1]\n",
    "        outputs.append(top1.item())\n",
    "        if top1.item() == word2idx['<eos>']:\n",
    "            break\n",
    "    response = lookup_words(idx2word, outputs)\n",
    "    return response"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "# enc = Encoder(input_dim=params['input_dim'], emb_dim=params['emb_dim'], enc_hid_dim=params['enc_hid_dim'], dec_hid_dim=params['dec_hid_dim'], dropout=params['dropout'])\n",
    "# dec = Decoder(output_dim=params['input_dim'], emb_dim=params['emb_dim'], enc_hid_dim=params['enc_hid_dim'], dec_hid_dim=params['dec_hid_dim'], dropout=params['dropout'])\n",
    "# norm_model = Seq2Seq(encoder=enc, decoder=dec, device=device)\n",
    "# norm_model.load_state_dict(torch.load('models/NormSeq2Seq-188M_epoch35.pt'))\n",
    "# norm_model.to(device)\n",
    "# print(f'Model has {sum(p.numel() for p in model.parameters() if p.requires_grad):,} trainable parameters')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'i i do i i massive swap ?'"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test = 'is ubuntu good?'\n",
    "norm_test = generate(norm_model, test, max_len=12, word2idx=word2idx, idx2word=idx2word, device=device)\n",
    "' '.join(norm_test).replace('<bos>', '').replace('<eos>', '').strip()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'maybe sent you'"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test = 'is ubuntu good?'\n",
    "attn_test = generate(attn_model, test, max_len=12, word2idx=word2idx, idx2word=idx2word, device=device)\n",
    "' '.join(attn_test).replace('<bos>', '').replace('<eos>', '').strip()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluating the seq2seq models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "def tensor_to_text(tensor, idx2word=idx2word):\n",
    "    \"\"\"\n",
    "    Convert tensor to text\n",
    "    :param tensor: tensor\n",
    "    :param idx2word: index to word mapping\n",
    "    :return: text\n",
    "    \"\"\"\n",
    "    textlist = tensor.tolist()\n",
    "    text = lookup_words(idx2word, textlist)\n",
    "    # remove default tokens\n",
    "    text = [word for word in text if word not in ['<bos>', '<eos>', '<pad>', '<unk>']]\n",
    "    return text\n",
    "\n",
    "def predict_from_tensor(model, input_tensor, max_len=12, word2idx=word2idx, idx2word=idx2word, device=device, lookup_words=lookup_words):\n",
    "    \"\"\"\n",
    "    Predict response\n",
    "    :param model: model\n",
    "    :param input_tensor: input tensor\n",
    "    :param max_len: maximum length of sequence\n",
    "    :param word2idx: word to index mapping\n",
    "    :param idx2word: index to word mapping\n",
    "    :return: response\n",
    "    \"\"\"\n",
    "    model.eval()\n",
    "    input_tensor = input_tensor.unsqueeze(1).to(device)\n",
    "    outputs = [word2idx['<bos>']]\n",
    "    with torch.no_grad():\n",
    "        encoder_outputs, hidden = model.encoder(input_tensor)\n",
    "    for t in range(max_len):\n",
    "        output, hidden = model.decoder(torch.tensor([outputs[-1]], dtype=torch.long).to(device), hidden, encoder_outputs)\n",
    "        top1 = output.max(1)[1]\n",
    "        outputs.append(top1.item())\n",
    "        if top1.item() == word2idx['<eos>']:\n",
    "            break\n",
    "    response = lookup_words(idx2word, outputs)\n",
    "    return response"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### BLEU Score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "def eval_pair(model, data, max_len=12,\n",
    "             word2idx=word2idx,\n",
    "             idx2word=idx2word, device=device,\n",
    "             tensor_to_text=tensor_to_text,\n",
    "             predict_from_tensor=predict_from_tensor):\n",
    "    predicted = []\n",
    "    actual = []\n",
    "    for c, r in data.dataset:\n",
    "        predict = predict_from_tensor(model, c, max_len=max_len, word2idx=word2idx, idx2word=idx2word, device=device)\n",
    "        predict = [word for word in predict if word not in ['<bos>', '<eos>', '<pad>', '<unk>']]\n",
    "        predicted.append([predict])\n",
    "        actual.append(tensor_to_text(r))\n",
    "    return predicted, actual"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [],
   "source": [
    "attn_predicted, attn_actual = eval_pair(attn_model, test_loader, max_len=12)\n",
    "norm_predicted, norm_actual = eval_pair(norm_model, test_loader, max_len=12)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "BLEU score for Seq2Seq with Attention: 1.46\n"
     ]
    }
   ],
   "source": [
    "attn_bleu = corpus_bleu(attn_predicted, attn_actual)\n",
    "print(f'BLEU score for Seq2Seq with Attention: {attn_bleu*100:.2f}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "BLEU score for Seq2Seq without Attention: 1.29\n"
     ]
    }
   ],
   "source": [
    "norm_bleu = corpus_bleu(norm_predicted, norm_actual)\n",
    "print(f'BLEU score for Seq2Seq without Attention: {norm_bleu*100:.2f}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Cosine Similarity\n",
    "- Cosine similarity is used to measure the similarity between two sentences.\n",
    "- First, sentence embeddings are created by encoding using a pre-trained model (BERT).\n",
    "- Next, the cosine similarity is calculated between the sentence embeddings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [],
   "source": [
    "def list_to_sent(list_of_words):\n",
    "    \"\"\"\n",
    "    Convert list of words to sentence\n",
    "    :param list_of_words: list of words\n",
    "    :return: sentence\n",
    "    \"\"\"\n",
    "    return ' '.join(list_of_words)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [],
   "source": [
    "sbert = SentenceTransformer('bert-base-nli-mean-tokens')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_cosines(list_pred, act, sbert=sbert, list_to_sent=list_to_sent):\n",
    "    \"\"\"\n",
    "    Get cosine similarity scores\n",
    "    :param list_pred: list of list of predicted sentences [[a],[b]]\n",
    "    :param act: list of actual sentences [a,b]\n",
    "    :param sbert: sentence transformer model\n",
    "    :param list_to_sent: function to convert list of words to sentence\n",
    "    :return: cosine similarity scores\n",
    "    \"\"\"\n",
    "    sent_pred = [val for sublist in list_pred for val in sublist]\n",
    "    sent_pred = [list_to_sent(sent) for sent in sent_pred]\n",
    "    sent_act = [list_to_sent(sent) for sent in act]\n",
    "    cosine_scores = []\n",
    "    for pred, act in zip(sent_pred, sent_act):\n",
    "        pred = sbert.encode(pred)\n",
    "        act = sbert.encode(act)\n",
    "        cosine_scores.append(cosine(pred, act))\n",
    "    cosine_scores = np.array(cosine_scores)\n",
    "    return cosine_scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Average cosine similarity for Seq2Seq with Attention: 0.40\n",
      "Average cosine similarity for Seq2Seq without Attention: 0.40\n"
     ]
    }
   ],
   "source": [
    "attn_cosine = get_cosines(attn_predicted, attn_actual)\n",
    "norm_cosine = get_cosines(norm_predicted, norm_actual)\n",
    "\n",
    "print(f'Average cosine similarity for Seq2Seq with Attention: {np.mean(attn_cosine):.2f}')\n",
    "print(f'Average cosine similarity for Seq2Seq without Attention: {np.mean(norm_cosine):.2f}')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "liv",
   "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.11.4"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}