Generative_Chatbot.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hbajA8_rqFdH"
   },
   "source": [
    "# Develop a Generative Chatbot\n",
    "\n",
    "This notebook outlines the code to train a generative chatbot using the Ubuntu dialogue corpus which can be found here: https://www.kaggle.com/datasets/rtatman/ubuntu-dialogue-corpus.\n",
    "We built 2 sequence to sequence models with one using an attention mechanism that can be either Luong or Bahdanau and another that does not use attention."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "EPagItDmqFdJ"
   },
   "source": [
    "## Step 1: Import Libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "I1wBQ6-aacTF"
   },
   "outputs": [],
   "source": [
    "# Importing the libraries\n",
    "import pandas as pd\n",
    "import re\n",
    "from torch import nn\n",
    "from tqdm import tqdm\n",
    "import torch\n",
    "from torch.nn.utils.rnn import pad_sequence\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import DataLoader\n",
    "from tensorflow.keras.preprocessing.text import Tokenizer\n",
    "import random\n",
    "import sys\n",
    "from sklearn.model_selection import train_test_split\n",
    "import os\n",
    "import numpy as np\n",
    "\n",
    "from nltk.translate.bleu_score import sentence_bleu\n",
    "from sklearn.feature_extraction.text import CountVectorizer\n",
    "from sklearn.metrics.pairwise import cosine_similarity\n",
    "\n",
    "# For METEOR\n",
    "import nltk\n",
    "from nltk.translate.meteor_score import meteor_score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "OJyG8gDuacTF"
   },
   "outputs": [],
   "source": [
    "# Set the seed value all over the place to make this reproducible.\n",
    "SEED = 42\n",
    "random.seed(SEED)\n",
    "torch.manual_seed(SEED)\n",
    "torch.cuda.manual_seed(SEED)\n",
    "torch.cuda.manual_seed_all(SEED)\n",
    "torch.backends.cudnn.deterministic = True\n",
    "torch.backends.cudnn.benchmark = False"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "bAk8G2S8acTG"
   },
   "source": [
    "## Step 2: Data Cleaning and Preparation\n",
    "\n",
    "Before building the training dataset, we must clean the data to handle unexpected characters and extra whitespaces. We also need to add start and end sequence tokens to help the model learn token positions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "F58fXQEvacTG"
   },
   "outputs": [],
   "source": [
    "def clean_text(text: str) -> str:\n",
    "    \"\"\"\n",
    "    Clean the text by removing unnecessary characters and normalizing whitespace.\n",
    "\n",
    "    Args:\n",
    "        input: text (str): the text to be cleaned\n",
    "\n",
    "    Returns:\n",
    "        output: text (str): the cleaned\n",
    "    \"\"\"\n",
    "    if pd.isnull(text):\n",
    "        return \"\"\n",
    "\n",
    "    # lowercase the text\n",
    "    text = text.lower()\n",
    "\n",
    "    # remove URLs if present\n",
    "    text = re.sub(r'http\\S+', '', text)\n",
    "\n",
    "    # Remove special characters and multiple spaces\n",
    "    text = re.sub(r'[^\\w\\s]', '', text)\n",
    "\n",
    "    # Remove extra whitespaces\n",
    "    text = re.sub(r'\\s+', ' ', text).strip()\n",
    "    return text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "F8XkO-XlacTG"
   },
   "outputs": [],
   "source": [
    "# Load the data\n",
    "data = pd.read_csv(\"dialogueText.csv\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "hkEJyW33acTG",
    "outputId": "460c8283-5863-40a9-d8f0-a1c3a9ec7303"
   },
   "outputs": [],
   "source": [
    "# Drop rows with missing values in the 'text' column\n",
    "data = data.dropna(subset=['text'])\n",
    "\n",
    "# Fill missing values in 'from' and 'to' columns with 'unknown'\n",
    "data['from'] = data['from'].fillna('unknown')\n",
    "data['to'] = data['to'].fillna('unknown')\n",
    "\n",
    "# Apply the clean_text function to the 'text' column only\n",
    "data['text'] = data['text'].apply(clean_text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "G6-poy74acTG"
   },
   "outputs": [],
   "source": [
    "# Define the function to add <sos> and <eos> tokens after cleaning\n",
    "def add_sos_eos(text: str) -> str:\n",
    "    \"\"\"\n",
    "    Add <sos> (start of sentence) and <eos> (end of sentence) tokens to the text.\n",
    "\n",
    "    Args:\n",
    "        input: text (str): the text to be modified\n",
    "\n",
    "    Returns:\n",
    "        output: text (str): the text with <sos> and <eos> tokens added\n",
    "    \"\"\"\n",
    "    return f\"<sos> {text} <eos>\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "F1pj5-r8acTG"
   },
   "outputs": [],
   "source": [
    "# Apply the add_sos_eos function to the 'text' column only\n",
    "data['text'] = data['text'].apply(add_sos_eos)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "KJeA5VwyqFdN"
   },
   "source": [
    "## Step 3: Tokenization\n",
    "\n",
    "Sequence-to-sequence models cannot process text directly and require numeric representations. Therefore, we convert all sequences into numbers that map back to their string equivalents. We build a vocabulary from the text corpus, assigning each unique word an ID, and then convert each word to its corresponding ID."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "pxI0lEp0acTG"
   },
   "outputs": [],
   "source": [
    "def tokenize(df: pd.DataFrame, input_column: str, output_column: str) -> tuple[dict, dict, pd.DataFrame]:\n",
    "    \"\"\"\n",
    "    Tokenize the text data in the specified column of a DataFrame.\n",
    "\n",
    "    Args:\n",
    "        df (pd.DataFrame): The input DataFrame containing the text data.\n",
    "        input_column (str): The name of the column in the DataFrame to tokenize.\n",
    "        output_column (str): The name of the column in the DataFrame that contains tokenized text.\n",
    "\n",
    "    Returns:\n",
    "        vocab (dict): A dictionary mapping words to their index.\n",
    "        reverse_vocab (dict): A dictionary mapping indices to their words.\n",
    "        df (pd.DataFrame): The DataFrame with an additional column containing tokenized text.\n",
    "    \"\"\"\n",
    "    # Initialize tokenizer with out-of-vocabulary token set to <unk>\n",
    "    tokenizer = Tokenizer(oov_token=\"<unk>\", filters='!\"#$%&()*+,-./:;=?@[\\\\]^_`{|}~\\t\\n')  # Use <unk> for out-of-vocabulary words\n",
    "\n",
    "    # Fit the tokenizer to the specified column (this will add words to the word_index)\n",
    "    tokenizer.fit_on_texts(df[input_column])\n",
    "    tokenizer.word_index['<pad>'] = 0\n",
    "    tokenizer.index_word[0] = '<pad>'\n",
    "\n",
    "    # Create a vocabulary dictionary\n",
    "    vocab = tokenizer.word_index\n",
    "    reverse_vocab = tokenizer.index_word\n",
    "\n",
    "    # Print the size of the vocabulary\n",
    "    print(f\"Size of the vocabulary in the text: {len(vocab)}\")\n",
    "\n",
    "    # Convert texts to tokenized sequences\n",
    "    tokenized_texts = tokenizer.texts_to_sequences(df[input_column])\n",
    "\n",
    "    # Add the tokenized sequences as a new column\n",
    "    df[f\"{output_column}\"] = tokenized_texts\n",
    "\n",
    "    return vocab, reverse_vocab, df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "OCTC8Pl0acTH",
    "outputId": "5978fe7d-38d6-4309-f7fb-4e299fae6ab8"
   },
   "outputs": [],
   "source": [
    "# Tokenize the data\n",
    "vocab, reverse_vocab, data = tokenize(data, 'text', 'tokenized_text')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "6E8qa9ohacTH",
    "outputId": "a7d5f00a-cfd3-4415-ae7e-969cf95365c0"
   },
   "outputs": [],
   "source": [
    "# Check the vocabulary\n",
    "vocab"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "f6jMPNrMacTH",
    "outputId": "3c00a5c5-6e60-4ec4-9c22-2c415548fa5a"
   },
   "outputs": [],
   "source": [
    "# Check the reverse vocabulary\n",
    "reverse_vocab"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 293
    },
    "id": "eqB5sySuacTH",
    "outputId": "a581fb49-9575-4e96-e944-64b3346fbdf1"
   },
   "outputs": [],
   "source": [
    "# Check the first few rows of the data\n",
    "data[:5]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "xg-xOi6yqFdO"
   },
   "source": [
    "## Step 4: Split the data\n",
    "\n",
    "With the data tokenized, we split it into training, validation, and test sets to train the model on one set and evaluate its performance on unseen data. Our split is 80% for training data, 10% for validation, and 10% for testing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "KXIW_Ee0acTH"
   },
   "outputs": [],
   "source": [
    "# NEW\n",
    "# Perform the dialogueID-aware split\n",
    "unique_ids = data['dialogueID'].unique()\n",
    "\n",
    "# We split the data so that 80% is allocated for the training set and the remaining 20% is then divided evenly between the validation and test set\n",
    "train_ids, temp_ids = train_test_split(unique_ids, test_size=0.2, random_state=42)\n",
    "\n",
    "# Split the temp set into validation and test sets. This means validation set and test set is 10% of all data each\n",
    "val_ids, test_ids = train_test_split(temp_ids, test_size=0.5, random_state=42)\n",
    "\n",
    "# Filter the original DataFrame based on the split IDs\n",
    "train_set = data[data['dialogueID'].isin(train_ids)]\n",
    "val_set = data[data['dialogueID'].isin(val_ids)]\n",
    "test_set = data[data['dialogueID'].isin(test_ids)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "uXLAvF-macTH",
    "outputId": "bddc2425-9c85-4931-e8a0-6ae3469531ea"
   },
   "outputs": [],
   "source": [
    "print(len(train_set))\n",
    "print(len(test_set))\n",
    "print(len(val_set))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "qoZi9bEcqFdO"
   },
   "source": [
    "## Step 5: Create Dialogue Pairs\n",
    "\n",
    "Our goal is to train a chatbot, so it must learn from question-answer pairs to identify questions and respond appropriately. Training on individual sequences without pairing them prevents learning this relationship, making dialogue pairs essential."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Dq9ZCFdvacTH"
   },
   "outputs": [],
   "source": [
    "def create_pairs(df: pd.DataFrame, text_column: str, tokenized_column: str) -> tuple[list[tuple[str, str]], list[tuple[torch.Tensor, torch.Tensor]]]:\n",
    "    \"\"\"\n",
    "    Create dialogue pairs from the DataFrame.\n",
    "\n",
    "    Args:\n",
    "        df (pd.DataFrame): The input DataFrame containing the dialogue data.\n",
    "        text_column (str): The name of the column in the DataFrame containing the text data.\n",
    "        tokenized_column (str): The name of the column in the DataFrame containing the tokenized text data.\n",
    "\n",
    "    Returns:\n",
    "        text_dialogue_pairs (list[tuple[str, str]]): A list of dialogue pairs in text form.\n",
    "        tokenized_dialogue_pairs (list[tuple[torch.Tensor, torch.Tensor]]): A list of dialogue pairs in tensor form.\n",
    "    \"\"\"\n",
    "\n",
    "    # Sort by dialogueID and date\n",
    "    df = df.sort_values(by=['dialogueID', 'date'])\n",
    "\n",
    "    # Create dialogue pairs\n",
    "    text_dialogue_pairs = []\n",
    "    tokenized_dialogue_pairs = []\n",
    "    # We group by the dialogueID so that all sequences related to a specific conversation are grouped together\n",
    "    grouped = df.groupby('dialogueID')\n",
    "\n",
    "    for _, group in tqdm(grouped):\n",
    "        group = group.reset_index(drop=True)\n",
    "        for i in range(len(group) - 1):\n",
    "            # Source is the current row's text\n",
    "            text_source = group.loc[i, text_column]\n",
    "            # Target is the next row's text\n",
    "            text_target = group.loc[i + 1, text_column]\n",
    "\n",
    "            # Tokenized Source is the current row's tokenized text\n",
    "            tokenized_source = group.loc[i, tokenized_column]\n",
    "            # Tokenized Target is the next row's tokenized text\n",
    "            tokenized_target = group.loc[i + 1, tokenized_column]\n",
    "\n",
    "            if text_source and text_target:  # Ensure both source and target have text\n",
    "                text_dialogue_pairs.append((text_source, text_target))\n",
    "\n",
    "            # Ensure both tokenized source and target are present\n",
    "            if tokenized_source and tokenized_target:\n",
    "                # Convert to tensors and add to tokenized dialogue pairs\n",
    "                tokenized_dialogue_pairs.append((torch.tensor(tokenized_source), torch.tensor(tokenized_target)))\n",
    "\n",
    "    return text_dialogue_pairs, tokenized_dialogue_pairs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "KYhTrNwWacTH",
    "outputId": "4ee31d8a-26d6-49d1-bbd6-a7326168a5bc"
   },
   "outputs": [],
   "source": [
    "# Create dialogue pairs\n",
    "train_text_dialogue_pairs, train_tokenized_dialogue_pairs = create_pairs(train_set, 'text', 'tokenized_text')\n",
    "# Test set\n",
    "test_text_dialogue_pairs, test_tokenized_dialogue_pairs = create_pairs(test_set, 'text', 'tokenized_text')\n",
    "# Validation Set\n",
    "val_text_dialogue_pairs, val_tokenized_dialogue_pairs = create_pairs(val_set, 'text', 'tokenized_text')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "UtKxCKfOacTH",
    "outputId": "c22c57ca-dba7-4a7b-a6a4-0aefc8354459"
   },
   "outputs": [],
   "source": [
    "# Display a few text pairs\n",
    "for source, target in train_text_dialogue_pairs[:5]:\n",
    "    print(f\"Source: {source}\")\n",
    "    print(f\"Target: {target}\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "zIm7IcOSacTI",
    "outputId": "f0da784b-a34c-471b-fc0e-7fa7c3586d85"
   },
   "outputs": [],
   "source": [
    "# Display a few tokenized pairs\n",
    "for source, target in train_tokenized_dialogue_pairs[:5]:\n",
    "    print(f\"Source: {source}\")\n",
    "    print(f\"Target: {target}\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "eF0U46JCacTI",
    "outputId": "4c4b739b-319d-46a6-9a8c-17d9d8fb8e62"
   },
   "outputs": [],
   "source": [
    "# Example of a tokenized pair\n",
    "print(f\"Vocabulary Size: {len(vocab)}\")\n",
    "print(f\"Source sentence: {train_text_dialogue_pairs[0][1]}\")\n",
    "print(f\"Source tokens: {train_tokenized_dialogue_pairs[0][1]}\")\n",
    "print(f\"Target sentence: {train_text_dialogue_pairs[1][0]}\")\n",
    "print(f\"Target tokens: {train_tokenized_dialogue_pairs[1][0]}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "-Ba0VtESqFdP"
   },
   "source": [
    "## Step 6: Create our Dataloaders\n",
    "\n",
    "Training models on batches of dialogue pairs, rather than individual ones, normalizes gradient updates and speeds up training. The code below defines an object to load batches of dialogue pairs from the training data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "U3KDL9raacTI"
   },
   "outputs": [],
   "source": [
    "def collate_fn(batch: list[tuple[torch.Tensor, torch.Tensor]]) -> tuple[torch.Tensor, torch.Tensor, list[int]]:\n",
    "    \"\"\"\n",
    "    Custom collate function to pad the sequences in the batch to the same length.\n",
    "\n",
    "    Args:\n",
    "        batch (list[tuple[torch.Tensor, torch.Tensor]]): A list of tuples containing source and target sequences.\n",
    "\n",
    "    Returns:\n",
    "        sources_padded (torch.Tensor): Padded source sequences.\n",
    "        targets_padded (torch.Tensor): Padded target sequences.\n",
    "        input_lengths (list[int]): List of input lengths.\n",
    "    \"\"\"\n",
    "    sources, targets = zip(*batch)\n",
    "\n",
    "    # Extract the input lengths from the raw sources before padding\n",
    "    input_lengths = [len(src) for src in sources]  # Lengths of each sequence\n",
    "\n",
    "    sources_padded = pad_sequence(sources, batch_first=True, padding_value=0)\n",
    "    targets_padded = pad_sequence(targets, batch_first=True, padding_value=0)\n",
    "    return sources_padded, targets_padded, input_lengths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Qpfmmq9uacTI"
   },
   "outputs": [],
   "source": [
    "# Create DataLoader\n",
    "batch_size = 64\n",
    "train_loader = DataLoader(train_tokenized_dialogue_pairs, batch_size=batch_size, collate_fn=collate_fn, shuffle=True) \n",
    "test_loader = DataLoader(test_tokenized_dialogue_pairs, batch_size=batch_size, collate_fn=collate_fn, shuffle=True)\n",
    "val_loader = DataLoader(val_tokenized_dialogue_pairs, batch_size=batch_size, collate_fn=collate_fn, shuffle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "QNN0F4ZAacTI",
    "outputId": "0c0cb5af-80ea-4c51-da53-5979795c8424"
   },
   "outputs": [],
   "source": [
    "# Example of a batch\n",
    "for src, tgt, input_lengths in train_loader:\n",
    "    print(f\"Source batch shape: {src.shape}\")\n",
    "    print(f\"Target batch shape: {tgt.shape}\")\n",
    "    print(f\"Input lengths: {input_lengths}\")\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "TnEdz7lNacTI"
   },
   "outputs": [],
   "source": [
    "def print_progress_bar(epoch: int, batch_idx: int, total_batches: int, batch_loss: float, total_epochs: int):\n",
    "    \"\"\"\n",
    "    Print the progress bar during training.\n",
    "\n",
    "    Args:\n",
    "        epoch (int): The current epoch.\n",
    "        batch_idx (int): The index of the current batch.\n",
    "        total_batches (int): The total number of batches in the dataset.\n",
    "        batch_loss (float): The average loss of the current batch.\n",
    "        total_epochs (int): The total number of epochs.\n",
    "    \"\"\"\n",
    "\n",
    "    progress = (batch_idx + 1) / total_batches\n",
    "    bar_length = 40  # Length of the progress bar\n",
    "    filled_length = int(bar_length * progress)\n",
    "    bar = \"█\" * filled_length + \"-\" * (bar_length - filled_length)\n",
    "    percentage = int(progress * 100)\n",
    "    sys.stdout.write(\n",
    "        f\"\\rEpoch {epoch}/{total_epochs}: |{bar}| {percentage}% \"\n",
    "        f\"({batch_idx + 1}/{total_batches} batches), Loss: {batch_loss:.4f}\"\n",
    "    )\n",
    "    sys.stdout.flush()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "EAAxg1VoqFdQ"
   },
   "source": [
    "## Step 7: Defining the Sequence To Sequence Model Without Attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "oELbi9F4acTN"
   },
   "outputs": [],
   "source": [
    "class EncoderNoAttention(nn.Module):\n",
    "    \"\"\"\n",
    "    Encoder module that encodes the input sequence without attention.\n",
    "    \"\"\"\n",
    "    def __init__(self, input_vocab_len: int, embedding_size: int, hidden_dim: int, n_layers=1, drop_prob=0.1):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            input_vocab_len (int): The input vocabulary size.\n",
    "            embedding_size (int): The size of the word embeddings.\n",
    "            hidden_dim (int): The size of the hidden dimension.\n",
    "            n_layers (int): The number of layers in the LSTM.\n",
    "            drop_prob (float): The dropout probability.\n",
    "        \"\"\"\n",
    "        super(EncoderNoAttention, self).__init__()\n",
    "\n",
    "        self.hidden_dim = hidden_dim\n",
    "        self.n_layers = n_layers\n",
    "\n",
    "        self.embedding = nn.Embedding(input_vocab_len, embedding_size)\n",
    "        self.dropout = nn.Dropout(drop_prob)\n",
    "        self.lstm = nn.LSTM(embedding_size, hidden_dim, n_layers,\n",
    "                            dropout=drop_prob, batch_first=True)\n",
    "\n",
    "    def forward(self, inputs: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:\n",
    "        embedded = self.embedding(inputs)\n",
    "        embedded = self.dropout(embedded)\n",
    "        outputs, (hidden, cell) = self.lstm(embedded)\n",
    "        return hidden, cell"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "yiCCkhXqacTN"
   },
   "outputs": [],
   "source": [
    "class DecoderNoAttention(nn.Module):\n",
    "    \"\"\"\n",
    "    Decoder module that decodes the input sequence without attention.\n",
    "    \"\"\"\n",
    "    def __init__(self, output_vocab_len: int, embedding_size: int, hidden_dim: int, n_layers=1, drop_prob=0.1):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            output_vocab_len (int): The output vocabulary size.\n",
    "            embedding_size (int): The size of the word embeddings.\n",
    "            hidden_dim (int): The size of the hidden dimension.\n",
    "            n_layers (int): The number of layers in the LSTM.\n",
    "            drop_prob (float): The dropout probability.\n",
    "        \"\"\"\n",
    "        super(DecoderNoAttention, self).__init__()\n",
    "\n",
    "        self.hidden_dim = hidden_dim\n",
    "        self.n_layers = n_layers\n",
    "\n",
    "        self.embedding = nn.Embedding(output_vocab_len, embedding_size)\n",
    "        self.dropout = nn.Dropout(drop_prob)\n",
    "        self.lstm = nn.LSTM(embedding_size, hidden_dim, n_layers, dropout=drop_prob, batch_first=True)\n",
    "        self.classifier = nn.Linear(hidden_dim, output_vocab_len)\n",
    "\n",
    "    def forward(self, inputs: torch.tensor, hidden: torch.tensor, cell: torch.tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:\n",
    "        embedded = self.embedding(inputs).unsqueeze(1)\n",
    "        embedded = self.dropout(embedded)\n",
    "\n",
    "        output, (hidden, cell) = self.lstm(embedded, (hidden, cell))\n",
    "\n",
    "        # Pass LSTM outputs through a Linear layer acting as a classifier\n",
    "        output = F.log_softmax(self.classifier(output.squeeze(1)), dim=1)\n",
    "\n",
    "        return output, hidden, cell"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "yl2Dx1dZqFdQ"
   },
   "source": [
    "## Step 8: Defining the function to train the Sequence to Sequence model without attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "yStUNhaeacTN"
   },
   "outputs": [],
   "source": [
    "def train_seq2seq_no_attention(dataloader: DataLoader, vocab: dict, start_of_sequence_token_index: int, padding_token_index: int, embedding_size=100, hidden_dim=256, epochs=10, lr=0.001, teacher_forcing_prob=0.5, device='cpu', folder_path='', model_filename='trained_seq_2_seq_no_attention.pth'):\n",
    "    \"\"\"\n",
    "    Train the sequence-to-sequence model without attention.\n",
    "\n",
    "    Args:\n",
    "        dataloader (DataLoader): The DataLoader object containing the tokenized dialogue pairs.\n",
    "        vocab (dict): The vocabulary dictionary.\n",
    "        start_of_sequence_token_index (int): The index of the <sos> token.\n",
    "        padding_token_index (int): The index of the <pad> token.\n",
    "        embedding_size (int): The size of the word embeddings.\n",
    "        hidden_dim (int): The size of the hidden dimension.\n",
    "        epochs (int): The number of epochs to train the model.\n",
    "        lr (float): The learning rate.\n",
    "        teacher_forcing_prob (float): The probability of using teacher forcing.\n",
    "        device (str): The device to run the model on ('cpu' or 'cuda').\n",
    "        folder_path (str): The folder where the output model parameters will be saved.\n",
    "        model_filename (str): The name of the model parameter file.\n",
    "    \"\"\"\n",
    "\n",
    "    # Initialize the encoder and decoder\n",
    "    encoder = EncoderNoAttention(len(vocab), embedding_size, hidden_dim).to(device)\n",
    "    decoder = DecoderNoAttention(len(vocab), embedding_size, hidden_dim).to(device)\n",
    "\n",
    "    # Initialize the optimizers\n",
    "    encoder_optimizer = optim.Adam(encoder.parameters(), lr=lr)\n",
    "    decoder_optimizer = optim.Adam(decoder.parameters(), lr=lr)\n",
    "\n",
    "    # Define the loss function\n",
    "    loss_fnc = nn.CrossEntropyLoss(ignore_index=padding_token_index)\n",
    "    total_batches = len(dataloader)\n",
    "\n",
    "    # Train the model\n",
    "    for epoch in range(epochs):\n",
    "        epoch_loss = 0.\n",
    "        for batch_idx, batch in enumerate(dataloader):\n",
    "            # Move the batch tensors to the device\n",
    "            q_inputs, a_inputs, _ = batch\n",
    "            q_inputs, a_inputs = q_inputs.to(device), a_inputs.to(device)\n",
    "\n",
    "            # Get batch size and max target length\n",
    "            batch_size = len(q_inputs)\n",
    "            max_target_len = len(a_inputs)\n",
    "\n",
    "            # Zero the gradients to prevent accumulation across batches, which can distort training dynamics.\n",
    "            encoder_optimizer.zero_grad()\n",
    "            decoder_optimizer.zero_grad()\n",
    "\n",
    "            # Hidden and cell states of the encoder are inputs to the decoder\n",
    "            decoder_hidden, decoder_cell = encoder(q_inputs)\n",
    "\n",
    "            # Initialize the decoder input with the start of sequence token\n",
    "            decoder_input = torch.tensor([start_of_sequence_token_index] * batch_size).to(device)\n",
    "            loss = 0.\n",
    "            for t in range(a_inputs.size(1)):\n",
    "                # Forward pass\n",
    "                decoder_output, decoder_hidden, decoder_cell = decoder(decoder_input, decoder_hidden, decoder_cell)\n",
    "                # Get the top prediction\n",
    "                _, top_index = decoder_output.topk(1)\n",
    "                # Teacher forcing: next input is current target\n",
    "                decoder_input = a_inputs[:, t] if random.random() < teacher_forcing_prob else top_index.squeeze().detach()\n",
    "                # Calculate the loss\n",
    "                loss += loss_fnc(decoder_output, a_inputs[:, t])\n",
    "\n",
    "            # Backward pass\n",
    "            loss.backward()\n",
    "\n",
    "\n",
    "            # Clip gradients to prevent exploding gradients\n",
    "            torch.nn.utils.clip_grad_norm_(encoder.parameters(), 1)\n",
    "            torch.nn.utils.clip_grad_norm_(decoder.parameters(), 1)\n",
    "\n",
    "            # Update parameters\n",
    "            encoder_optimizer.step()\n",
    "            decoder_optimizer.step()\n",
    "\n",
    "            # Calculate the average loss\n",
    "            batch_loss = loss.item() / max_target_len\n",
    "            epoch_loss += batch_loss\n",
    "\n",
    "            print_progress_bar(epoch + 1, batch_idx, total_batches, batch_loss, epochs)\n",
    "\n",
    "        print(f\"\\n Epoch {epoch + 1}/{epochs}, Loss: {epoch_loss/total_batches:.4f}\")\n",
    "\n",
    "    # Define the full save path for the model\n",
    "    save_path = os.path.join(folder_path, model_filename)\n",
    "    # Save the model\n",
    "    torch.save({\n",
    "        'encoder': encoder.state_dict(),\n",
    "        'decoder': decoder.state_dict(),\n",
    "        'e_optimizer': encoder_optimizer.state_dict(),\n",
    "        'd_optimizer': decoder_optimizer.state_dict()\n",
    "    }, save_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "aqHnW-XtacTN"
   },
   "outputs": [],
   "source": [
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "folder_path = './'\n",
    "model_filename = \"trained_seq_2_seq_no_attention.pth\"\n",
    "train_seq2seq_no_attention(train_loader, vocab, vocab['<sos>'], vocab['<pad>'], epochs=2, device=device, folder_path=folder_path, model_filename=model_filename)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "y7rV_6-JacTN"
   },
   "source": [
    "## Step 9: Defining the Sequence to Sequence model with Attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "tO4l3MinacTN"
   },
   "outputs": [],
   "source": [
    "class EncoderAttention(nn.Module):\n",
    "    \"\"\"\n",
    "    Encoder module that encodes the input sequence with attention.\n",
    "    \"\"\"\n",
    "    def __init__(self, input_vocab_len: int, embedding_size: int, hidden_size: int, n_layers=1, dropout=0.1):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            input_vocab_len (int): The input vocabulary size.\n",
    "            embedding_size (int): The size of the word embeddings.\n",
    "            hidden_size (int): The size of the hidden dimension.\n",
    "            n_layers (int): The number of layers in the LSTM.\n",
    "            dropout (float): The dropout probability.\n",
    "        \"\"\"\n",
    "        super(EncoderAttention, self).__init__()\n",
    "\n",
    "        self.hidden_size = hidden_size\n",
    "        self.n_layers = n_layers\n",
    "        self.dropout = dropout\n",
    "\n",
    "        self.embedding = nn.Embedding(input_vocab_len, embedding_size)\n",
    "        self.gru = nn.GRU(embedding_size, self.hidden_size, self.n_layers, dropout=(self.dropout if n_layers > 1 else 0), bidirectional=True, batch_first=True)\n",
    "\n",
    "    def forward(self, x: torch.tensor, input_lengths: int, hidden=None, device='cpu') -> tuple[torch.tensor, torch.tensor]:\n",
    "        \"\"\"\n",
    "        Forward pass of the encoder module.\n",
    "\n",
    "        Args:\n",
    "            x (torch.tensor): The input tensor.\n",
    "            input_lengths (int): The lengths of the input sequences.\n",
    "            hidden (torch.tensor): The hidden state tensor.\n",
    "            device (str): The device to run the model on ('cpu' or 'cuda').\n",
    "\n",
    "        Returns:\n",
    "            outputs (torch.tensor): The output tensor.\n",
    "            hidden (torch.tensor): The hidden state tensor.\n",
    "        \"\"\"\n",
    "        # Moves the input to the device prior to processing\n",
    "        x = x.to(device)\n",
    "        # Convert inputs to embeddings\n",
    "        x = self.embedding(x).to(device)\n",
    "        # Pack the sequence for handling variable lengths\n",
    "        packed = nn.utils.rnn.pack_padded_sequence(x, input_lengths, batch_first=True, enforce_sorted=False)\n",
    "        # Pass through GRU\n",
    "        outputs, hidden = self.gru(packed, hidden)\n",
    "        # Unpack the sequence\n",
    "        outputs, _ = nn.utils.rnn.pad_packed_sequence(outputs, batch_first=True)\n",
    "        # Sum bidirectional outputs\n",
    "        outputs = outputs[:, :, :self.hidden_size] + outputs[:, :, self.hidden_size:]\n",
    "        # Move outputs and hidden states to the specified device\n",
    "        return outputs.to(device), hidden.to(device) if hidden is not None else hidden"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "wuO1h9mlacTN"
   },
   "outputs": [],
   "source": [
    "class Attn(nn.Module):\n",
    "    \"\"\"\n",
    "    Attention module that computes the attention weights. This attention class supports Luong's general, Luong's dot, and Bahdanau's concatenation attention methods.\n",
    "    \"\"\"\n",
    "    def __init__(self, method: str, hidden_size: int):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            method (str): The attention method to use.\n",
    "            hidden_size (int): The size of the hidden dimension.\n",
    "        \"\"\"\n",
    "        super(Attn, self).__init__()\n",
    "        self.method = method\n",
    "        self.hidden_size = hidden_size\n",
    "\n",
    "        if self.method not in ['dot', 'general', 'concat']:\n",
    "            raise ValueError(f\"Unsupported attention method: {self.method}\")\n",
    "\n",
    "        # This handles one of Luong's Attention method\n",
    "        if self.method == 'general':\n",
    "            self.attn = nn.Linear(self.hidden_size, self.hidden_size)\n",
    "        # This handles Bahdanau Attention method\n",
    "        elif self.method == 'concat':\n",
    "            self.attn = nn.Linear(self.hidden_size*2, self.hidden_size)\n",
    "            self.v = nn.Parameter(torch.FloatTensor(1, hidden_size))\n",
    "\n",
    "    def forward(self, hidden: torch.tensor, encoder_outputs: torch.tensor) -> torch.tensor:\n",
    "        \"\"\"\n",
    "        Compute the attention weights.\n",
    "        Args:\n",
    "            hidden: Decoder hidden state (1, B, H) or (B, H) depending on implementation\n",
    "            encoder_outputs: Encoder outputs (T, B, H)\n",
    "        Returns:\n",
    "            Attention weights of shape (B, 1, T)\n",
    "        \"\"\"\n",
    "        # Ensure hidden is in (B, H) format\n",
    "        if hidden.dim() == 3:\n",
    "            hidden = hidden.squeeze(0)  # (B, H)\n",
    "        # Compute the attention energies\n",
    "        attn_energies = self.score(hidden, encoder_outputs)  # (B, T)\n",
    "        # Normalize energies to get attention weights\n",
    "        attn_weights = F.softmax(attn_energies, dim=1).unsqueeze(1)  # (B, 1, T)\n",
    "        # Compute the context vector\n",
    "        context = torch.bmm(attn_weights, encoder_outputs)  # (B, 1, H)\n",
    "        return attn_weights, context\n",
    "\n",
    "    def score(self, hidden: torch.tensor, encoder_outputs: torch.tensor) -> torch.tensor:\n",
    "        \"\"\"\n",
    "        Calculate attention scores.\n",
    "        Args:\n",
    "            hidden: Decoder hidden state (1, B, H)\n",
    "            encoder_outputs: Encoder outputs (T, B, H)\n",
    "        Returns:\n",
    "            Attention scores of shape (B, T)\n",
    "        \"\"\"\n",
    "        # Handle \"dot\" attention (Luong)\n",
    "        if self.method == 'dot':\n",
    "            return torch.bmm(hidden.unsqueeze(1), encoder_outputs.transpose(1, 2)).squeeze(1)  # (B, T)\n",
    "\n",
    "        # Handle \"general\" attention (Luong)\n",
    "        elif self.method == 'general':\n",
    "            energy = self.attn(encoder_outputs)  # (B, T, H)\n",
    "            return torch.bmm(hidden.unsqueeze(1), energy.transpose(1, 2)).squeeze(1)  # (B, T)\n",
    "\n",
    "        # Handle \"concat\" attention (Bahdanau)\n",
    "        elif self.method == 'concat':\n",
    "            hidden_expanded = hidden.unsqueeze(1).expand(-1, encoder_outputs.size(1), -1)  # (B, T, H)\n",
    "            energy = self.attn(torch.cat((hidden_expanded, encoder_outputs), dim=2))  # (B, T, H)\n",
    "            energy = torch.tanh(energy)  # (B, T, H)\n",
    "            # Reshape v to (1, H, 1) for matrix multiplication\n",
    "            v_reshaped = self.v.unsqueeze(2)  # (1, H, 1)\n",
    "\n",
    "            # Perform batch matrix multiplication (B, T, H) with (1, H, 1)\n",
    "            attn_scores = torch.bmm(energy, v_reshaped)  # (B, T, 1)\n",
    "            return attn_scores.squeeze(2)  # (B, T)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "REOr2WuxacTO"
   },
   "outputs": [],
   "source": [
    "class DecoderAttention(nn.Module):\n",
    "    \"\"\"\n",
    "    Decoder module that decodes the input sequence with attention. This Decoder module support Luong and Bahdanau's attention which can be changed by passing in the attn_type and attn_approach\n",
    "    \"\"\"\n",
    "    def __init__(self, output_vocab_len:int, embedding_size:int, hidden_size:int, attn_type:str, attn_approach:str, n_layers=1, dropout=0.1):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            output_vocab_len (int): The output vocabulary size.\n",
    "            embedding_size (int): The size of the word embeddings.\n",
    "            hidden_size (int): The size of the hidden dimension.\n",
    "            attn_type (str): The type of attention to use ('luong' or 'bahdanau').\n",
    "            attn_approach (str): The approach to use for attention ('general', 'dot', or 'concat'). Can only use 'general' or 'dot' with 'luong' as the attn_type and can only use 'concat' with 'bahdanau' as the attn_type.\n",
    "            n_layers (int): The number of layers in the LSTM.\n",
    "            dropout (float): The dropout probability.\n",
    "        \"\"\"\n",
    "        super(DecoderAttention, self).__init__()\n",
    "\n",
    "        # Keep for reference\n",
    "        self.attn_type = attn_type\n",
    "        self.attn_approach = attn_approach\n",
    "        self.hidden_size = hidden_size\n",
    "        self.n_layers = n_layers\n",
    "        self.dropout = dropout\n",
    "\n",
    "        # Define layers\n",
    "        self.embedding = nn.Embedding(output_vocab_len, embedding_size)\n",
    "        self.gru = nn.GRU(embedding_size, hidden_size, n_layers, dropout=(self.dropout if n_layers > 1 else 0), batch_first=True)\n",
    "        self.out = nn.Linear(hidden_size, output_vocab_len)\n",
    "        if self.attn_type == 'bahdanau':\n",
    "            self.out = nn.Linear(hidden_size*2, output_vocab_len)\n",
    "\n",
    "\n",
    "        # Attention mechanism\n",
    "        self.attn = Attn(self.attn_approach, hidden_size)\n",
    "        if self.attn_type == 'luong' and self.attn_approach in ['general', 'dot']:\n",
    "            self.concat = nn.Linear(hidden_size * 2, hidden_size)\n",
    "\n",
    "    def forward(self, input_seq: torch.tensor, last_hidden: torch.tensor, encoder_outputs: torch.tensor) -> tuple[torch.tensor, torch.tensor, torch.tensor]:\n",
    "        \"\"\"\n",
    "        Forward pass for the decoder.\n",
    "\n",
    "        Parameters:\n",
    "        - input_seq: Input sequence of shape (B), where B is the batch size.\n",
    "        - last_hidden: Previous hidden state of shape (n_layers, B, H).\n",
    "        - encoder_outputs: Encoder outputs of shape (B, T, H), where T is sequence length.\n",
    "\n",
    "        Returns:\n",
    "        - output: Final output tensor of shape (B, output_size).\n",
    "        - hidden: Updated hidden state of shape (n_layers, B, H).\n",
    "        - attn_weights: Attention weights of shape (B, T).\n",
    "        \"\"\"\n",
    "\n",
    "        # Handle bidirectional hidden states\n",
    "        if last_hidden.size(0) == 2 * self.n_layers:  # For bidirectional hidden state\n",
    "            last_hidden = last_hidden.view(self.n_layers, 2, -1, self.hidden_size).sum(1)  # (n_layers, B, H)\n",
    "            last_hidden = last_hidden.unsqueeze(0)  # Add layer dimension (1, B, H)\n",
    "        # Embedding layer\n",
    "        embedded = self.embedding(input_seq)  # (B, H)\n",
    "        #embedded = self.embedding_dropout(embedded)\n",
    "        embedded = embedded.unsqueeze(1)  # (B, 1, H)\n",
    "        # GRU step\n",
    "\n",
    "        rnn_output, hidden = self.gru(embedded, last_hidden)  # rnn_output: (B, 1, H)\n",
    "\n",
    "        rnn_output = rnn_output.squeeze(1)  # Remove the sequence dimension (B, H)\n",
    "        #hidden = hidden.squeeze(0)  # Remove the layers dimension (B, H). not needed\n",
    "\n",
    "        # Attention mechanism\n",
    "        if self.attn_type == 'luong':\n",
    "            attn_weights, context = self.attn(rnn_output, encoder_outputs)  # (B, T), (B, 1, H)\n",
    "            context = context.squeeze(1)  # (B, H)\n",
    "            concat_input = torch.cat((rnn_output, context), 1)  # (B, 2*H)\n",
    "            concat_output = F.tanh(self.concat(concat_input))  # (B, H)\n",
    "            output = self.out(concat_output)  # (B, output_size)\n",
    "\n",
    "        elif self.attn_type == 'bahdanau':\n",
    "            attn_weights, context = self.attn(hidden, encoder_outputs)  # (B, T), (B, 1, H)\n",
    "            context = context.squeeze(1)  # (B, H)\n",
    "            concatenated_output = torch.cat((rnn_output, context), dim=1)\n",
    "            # After concatenation the shape is [B, hidden_size*2] and self.out has the shape of [hidden_size*2, len(vocab)] so need to transpose self.out to [len(vocab), hidden_size*2] and matmul will result in [B, (len(vocab))]\n",
    "            output = torch.matmul(concatenated_output, self.out.weight.T)\n",
    "\n",
    "        return output, hidden, attn_weights"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "M5REzlvrqFdW"
   },
   "source": [
    "## Step 10: Defining the functions to train the Sequence to Sequence model with attention\n",
    "\n",
    "train_seq2seq_attention is the main function which in turn calls train_step_attention for each batch of training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "NZSUjg85acTO"
   },
   "outputs": [],
   "source": [
    "def train_step_attention(source: torch.tensor, target :torch.tensor, input_lengths: list[int],\n",
    "                         encoder: nn.Module, decoder: nn.Module, encoder_optimizer: torch.optim,\n",
    "                         decoder_optimizer: torch.optim, criterion: nn.Module, device: str, start_of_sequence_token_index: str) -> float:\n",
    "    \"\"\"\n",
    "    Perform a single training step for the sequence-to-sequence model with attention.\n",
    "\n",
    "    Args:\n",
    "        source (torch.Tensor): The source tensor.\n",
    "        target (torch.Tensor): The target tensor.\n",
    "        input_lengths (list[int]): The lengths of the input sequences.\n",
    "        encoder (nn.Module): The encoder model.\n",
    "        decoder (nn.Module): The decoder model.\n",
    "        encoder_optimizer (torch.optim): The encoder optimizer.\n",
    "        decoder_optimizer (torch.optim): The decoder optimizer.\n",
    "        criterion (nn.Module): The loss function.\n",
    "        device (str): The device to run the model on ('cpu' or 'cuda').\n",
    "        start_of_sequence_token_index (int): The index of the <sos> token.\n",
    "\n",
    "    Returns:\n",
    "        float: The loss value.\n",
    "    \"\"\"\n",
    "    batch_size = source.size(0)\n",
    "    max_target_len = target.size(1)\n",
    "\n",
    "    # Move tensors to the specified device\n",
    "    source, target = source.to(device), target.to(device)\n",
    "\n",
    "    # Zero gradients\n",
    "    encoder_optimizer.zero_grad()\n",
    "    decoder_optimizer.zero_grad()\n",
    "\n",
    "    # Encoder forward pass\n",
    "    encoder_outputs, encoder_hidden = encoder(source, input_lengths, device=device)\n",
    "    # Decoder initial state\n",
    "    decoder_input = torch.tensor([start_of_sequence_token_index] * batch_size, device=device)\n",
    "    decoder_hidden = encoder_hidden[:decoder.n_layers]  # Use only the forward hidden states\n",
    "    loss = 0\n",
    "    for t in range(1, max_target_len):\n",
    "        decoder_output, decoder_hidden, _ = decoder(decoder_input, decoder_hidden, encoder_outputs)\n",
    "        loss += criterion(decoder_output, target[:, t])\n",
    "        decoder_input = target[:, t]  # Teacher forcing\n",
    "\n",
    "    # Backpropagation\n",
    "    loss.backward()\n",
    "\n",
    "    # Clip gradients to prevent exploding gradients\n",
    "    torch.nn.utils.clip_grad_norm_(encoder.parameters(), 1)\n",
    "    torch.nn.utils.clip_grad_norm_(decoder.parameters(), 1)\n",
    "\n",
    "    # Update parameters\n",
    "    encoder_optimizer.step()\n",
    "    decoder_optimizer.step()\n",
    "\n",
    "    return loss.item() / max_target_len"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "m0Z2INQMacTO"
   },
   "outputs": [],
   "source": [
    "def train_seq2seq_attention(dataloader: DataLoader, vocab: dict, start_of_sequence_token_index: int, padding_token_index: int, input_dim=100, hidden_dim=256, attn_type='luong', attn_approach='general', epochs=10, lr=0.001, device='cpu', folder_path='', model_filename='trained_seq_2_seq_no_attention.pth'):\n",
    "    \"\"\"\n",
    "    Train the sequence-to-sequence model with attention.\n",
    "\n",
    "    Args:\n",
    "        dataloader (DataLoader): The DataLoader object containing the tokenized dialogue pairs.\n",
    "        vocab (dict): The vocabulary dictionary.\n",
    "        start_of_sequence_token_index (int): The index of the <sos> token.\n",
    "        padding_token_index (int): The index of the <pad> token.\n",
    "        input_dim (int): The size of the word embeddings.\n",
    "        hidden_dim (int): The size of the hidden dimension.\n",
    "        attn_type (str): The attention mechanism which can be either luong or bahdanau.\n",
    "        attn_approach (str): The actual attention calculation which can be either dot or general for luong or concat for bahdanau.\n",
    "        epochs (int): The number of epochs to train the model.\n",
    "        lr (float): The learning rate.\n",
    "        device (str): The device to run the model on ('cpu' or 'cuda').\n",
    "        folder_path (str): The folder where the output model parameters will be saved.\n",
    "        model_filename (str): The name of the model parameter file.\n",
    "    \"\"\"\n",
    "    # Ensure the folder exists, or create it\n",
    "    os.makedirs(folder_path, exist_ok=True)\n",
    "\n",
    "    vocab_size = len(vocab)\n",
    "\n",
    "    encoder = EncoderAttention(vocab_size, input_dim, hidden_dim).to(device)\n",
    "    decoder = DecoderAttention(vocab_size, input_dim, hidden_dim, attn_type, attn_approach).to(device)\n",
    "\n",
    "    encoder_optimizer = optim.Adam(encoder.parameters(), lr=lr)\n",
    "    decoder_optimizer = optim.Adam(decoder.parameters(), lr=lr)\n",
    "\n",
    "    criterion = nn.CrossEntropyLoss(ignore_index=padding_token_index)\n",
    "\n",
    "    # Training loop\n",
    "    for epoch in range(1, epochs + 1):\n",
    "        epoch_loss = 0\n",
    "        total_batches = len(dataloader)\n",
    "        for batch_idx, (source, target, input_lengths) in enumerate(dataloader):\n",
    "            batch_loss = train_step_attention(source, target, input_lengths, encoder, decoder, encoder_optimizer, decoder_optimizer, criterion, device, start_of_sequence_token_index)\n",
    "            epoch_loss += batch_loss\n",
    "            # Update progress bar\n",
    "            print_progress_bar(epoch, batch_idx, total_batches, batch_loss, epochs)\n",
    "            #time.sleep(0.01)  # Optional: Slows the output for demonstration purposes\n",
    "\n",
    "        print(f\"\\n Epoch {epoch}/{epochs}, Loss: {epoch_loss/len(dataloader):.4f}\")\n",
    "\n",
    "    # Define the full save path for the model\n",
    "    save_path = os.path.join(folder_path, model_filename)\n",
    "    # Save the model\n",
    "    torch.save({\n",
    "        'encoder': encoder.state_dict(),\n",
    "        'decoder': decoder.state_dict(),\n",
    "        'e_optimizer': encoder_optimizer.state_dict(),\n",
    "        'd_optimizer': decoder_optimizer.state_dict()\n",
    "    }, save_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "uT0QfzJcacTO"
   },
   "outputs": [],
   "source": [
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "folder_path = './'\n",
    "model_filename = \"trained_seq_2_seq_attention.pth\"\n",
    "train_seq2seq_attention(train_loader, vocab, vocab['<sos>'], vocab['<pad>'], epochs=5, device=device, attn_type='luong', attn_approach='general', folder_path=folder_path, model_filename=model_filename)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6kXjPkm5acTO"
   },
   "source": [
    "## Step 11: Trained no Attention Model Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "VpOdGcaVacTO"
   },
   "outputs": [],
   "source": [
    "def generate_answer_no_attention(question, vocab, reverse_vocab, start_of_sequence_token_index, end_of_sequence_token_index, embedding_size=100, hidden_size=256, max_length=50, device='cpu', saved_model=''):\n",
    "    # Define the encoder and decoder models\n",
    "    encoder = EncoderNoAttention(len(vocab), embedding_size, hidden_size).to(device)\n",
    "    decoder = DecoderNoAttention(len(vocab), embedding_size, hidden_size).to(device)\n",
    "\n",
    "    # If saved_model is not empty, attempt to load it\n",
    "    if saved_model != '':\n",
    "        checkpoint = torch.load(saved_model, map_location=device)\n",
    "        # Load the saved model weights into the models\n",
    "        encoder.load_state_dict(checkpoint['encoder'])\n",
    "        decoder.load_state_dict(checkpoint['decoder'])\n",
    "\n",
    "    # Put them in evaluation mode\n",
    "    encoder.eval()\n",
    "    decoder.eval()\n",
    "\n",
    "    # Prepare the question input\n",
    "    question_tensor = torch.tensor([vocab[word] for word in question.split()]).unsqueeze(0).to(device)\n",
    "\n",
    "    # Encode the question\n",
    "    decoder_hidden, decoder_cell = encoder(question_tensor)\n",
    "\n",
    "    # Initialize the decoder input with the start of sequence token\n",
    "    decoder_input = torch.tensor([start_of_sequence_token_index]).to(device)\n",
    "    generated_answer = []\n",
    "\n",
    "    for _ in range(max_length):\n",
    "        # Pass the latest token (decoder_input), the previous decoder hidden state (initially the encoder's final hidden state), and the decoder cell state to the decoder\n",
    "        decoder_output, decoder_hidden, decoder_cell = decoder(decoder_input, decoder_hidden, decoder_cell)\n",
    "        # Get the token with the highest probability\n",
    "        topv, topi = decoder_output.topk(1)\n",
    "        # If the token is the <eos> token then stop the text generation\n",
    "        if topi.item() == end_of_sequence_token_index:\n",
    "            break\n",
    "        else:\n",
    "            # If not <eos> then add it to the answer array\n",
    "            generated_answer.append(reverse_vocab[topi.item()])\n",
    "        # Get the latest token and store it as the input for the next iteration of the decoder\n",
    "        decoder_input = topi.squeeze().detach().unsqueeze(0)\n",
    "\n",
    "    return ' '.join(generated_answer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "-0D6ivI4qFdX",
    "outputId": "e76e2d70-964f-4304-e016-4ceb134a1f2a"
   },
   "outputs": [],
   "source": [
    "# Example usage\n",
    "model_path = \"trained_seq_2_seq_no_attention.pth\"\n",
    "question = test_text_dialogue_pairs[0][0]\n",
    "print(\"Question:\", question)\n",
    "answer = generate_answer_no_attention(question, vocab, reverse_vocab, vocab[\"<sos>\"], vocab[\"<eos>\"], saved_model=model_path)\n",
    "print(\"Answer:\", answer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Qt76W5gbqFdX"
   },
   "source": [
    "## Step 12: Trained Attention Model Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "OtP6KoZxacTO"
   },
   "outputs": [],
   "source": [
    "def generate_answer_attention(question, vocab, reverse_vocab, start_of_sequence_token_index, end_of_sequence_token_index, embedding_size=100, hidden_size=256, attn_type='luong', attn_approach='general', max_length=50, device='cpu', saved_model=''):\n",
    "    encoder = EncoderAttention(len(vocab), embedding_size, hidden_size).to(device)\n",
    "    decoder = DecoderAttention(len(vocab), embedding_size, hidden_size, attn_type, attn_approach).to(device)\n",
    "\n",
    "    # If saved_model is not empty, attempt to load it\n",
    "    if saved_model != '':\n",
    "        checkpoint = torch.load(saved_model, map_location=device)\n",
    "        # Load the saved model weights into the models\n",
    "        encoder.load_state_dict(checkpoint['encoder'])\n",
    "        decoder.load_state_dict(checkpoint['decoder'])\n",
    "\n",
    "    encoder.eval()\n",
    "    decoder.eval()\n",
    "\n",
    "    # Prepare the question input\n",
    "    question_tensor = torch.tensor([vocab[word] for word in question.split()]).unsqueeze(0).to(device)\n",
    "    input_lengths = torch.tensor([len(question.split())]).to(device)\n",
    "\n",
    "    # Encode the question\n",
    "    encoder_outputs, encoder_hidden = encoder(question_tensor, input_lengths, device=device)\n",
    "\n",
    "    # Initialize the decoder input with the start of sequence token\n",
    "    decoder_input = torch.tensor([start_of_sequence_token_index], device=device)\n",
    "    decoder_hidden = encoder_hidden[:decoder.n_layers]  # Use only the forward hidden states\n",
    "    generated_answer = []\n",
    "\n",
    "    for _ in range(max_length):\n",
    "        # Pass the latest token (decoder_input), the previous decoder hidden state (initially the encoder's final hidden state), and the decoder cell state to the decoder\n",
    "        decoder_output, decoder_hidden, _ = decoder(decoder_input, decoder_hidden, encoder_outputs)\n",
    "        # Get the token with the highest probability\n",
    "        topv, topi = decoder_output.topk(1)\n",
    "        # If the token is the <eos> token then stop the text generation\n",
    "        if topi.item() == end_of_sequence_token_index:\n",
    "            break\n",
    "        else:\n",
    "            # If not <eos> then add it to the answer array\n",
    "            generated_answer.append(reverse_vocab[topi.item()])\n",
    "\n",
    "        # Get the latest token and store it as the input for the next iteration of the decoder\n",
    "        decoder_input = topi.squeeze().detach().unsqueeze(0)\n",
    "\n",
    "    return ' '.join(generated_answer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "b4xkvYUwacTO",
    "outputId": "9c9205ef-f4b9-44fd-8ad7-3c61b1915d4b"
   },
   "outputs": [],
   "source": [
    "# Example usage\n",
    "model_path = \"trained_seq_2_seq_attention.pth\" \n",
    "question = test_text_dialogue_pairs[0][0]\n",
    "print(\"Question:\", question)\n",
    "attention_type = 'luong'\n",
    "attention_approach = 'general'\n",
    "answer = generate_answer_attention(question, vocab, reverse_vocab, vocab[\"<sos>\"], vocab[\"<eos>\"], attn_type=attention_type, attn_approach=attention_approach, saved_model=model_path)\n",
    "print(\"Answer:\", answer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0I0oEaIYqFdX"
   },
   "source": [
    "## Additional model testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "jNEamkqeqFdX",
    "outputId": "7cac813f-c930-4d92-dcc8-a0961d1518c8"
   },
   "outputs": [],
   "source": [
    "# Example usage\n",
    "model_path = \"trained_seq_2_seq_no_attention.pth\"\n",
    "question = test_text_dialogue_pairs[0][0]\n",
    "print(\"Question:\", question)\n",
    "answer = generate_answer_no_attention(question, vocab, reverse_vocab, vocab[\"<sos>\"], vocab[\"<eos>\"], saved_model=model_path)\n",
    "print(\"Answer:\", answer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "zaguWkoCqFdX",
    "outputId": "ea892bec-3355-4d51-dce6-e36ac281b6e6"
   },
   "outputs": [],
   "source": [
    "# Example usage\n",
    "model_path = \"trained_seq_2_seq_attention.pth\"\n",
    "question = test_text_dialogue_pairs[0][0]\n",
    "print(\"Question:\", question)\n",
    "attention_type = 'luong'\n",
    "attention_approach = 'general'\n",
    "answer = generate_answer_attention(question, vocab, reverse_vocab, vocab[\"<sos>\"], vocab[\"<eos>\"], attn_type=attention_type, attn_approach=attention_approach, saved_model=model_path)\n",
    "print(\"Answer:\", answer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "JPTciq0-qFdX"
   },
   "source": [
    "# STEP 14: Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "TP1i1AaNqFdY"
   },
   "outputs": [],
   "source": [
    "def download_nltk_resources():\n",
    "    \"\"\"\n",
    "    Download necessary NLTK resources for evaluation metrics.\n",
    "    \"\"\"\n",
    "    try:\n",
    "        nltk.download('wordnet')\n",
    "        nltk.download('omw-1.4')\n",
    "        nltk.download('punkt')\n",
    "    except Exception as e:\n",
    "        print(f\"Error downloading NLTK resources: {e}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "9YQcdu_NqFdY"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "from nltk.translate.bleu_score import sentence_bleu\n",
    "from nltk.translate.meteor_score import meteor_score\n",
    "from sklearn.feature_extraction.text import CountVectorizer\n",
    "from sklearn.metrics.pairwise import cosine_similarity\n",
    "from tqdm import tqdm\n",
    "\n",
    "class ModelEvaluator:\n",
    "    def __init__(self, model_no_attention, model_attention, test_loader,\n",
    "                 vocab, reverse_vocab, start_token, end_token, device='cuda:0'):\n",
    "        self.model_no_attention = model_no_attention\n",
    "        self.model_attention = model_attention\n",
    "        self.test_loader = test_loader\n",
    "        self.vocab = vocab\n",
    "        self.reverse_vocab = reverse_vocab\n",
    "        self.start_token = start_token\n",
    "        self.end_token = end_token\n",
    "        self.device = device\n",
    "\n",
    "        # Assuming you want to download NLTK resources for METEOR\n",
    "        import nltk\n",
    "        nltk.download('punkt')\n",
    "        nltk.download('wordnet')\n",
    "\n",
    "    def generate_answers_without_attention(self):\n",
    "        references = []\n",
    "        candidates = []\n",
    "\n",
    "        with torch.no_grad():\n",
    "            for batch in tqdm(self.test_loader, desc='Generating answers without attention'):\n",
    "                question, ground_truth = batch\n",
    "                question = question # Move to device if using CUDA\n",
    "\n",
    "                # Generate answer without attention\n",
    "                generated_answer = generate_answer_no_attention(\n",
    "                    question, self.vocab, self.reverse_vocab,\n",
    "                    self.vocab[\"<sos>\"], self.vocab[\"<eos>\"],\n",
    "                    saved_model=self.model_no_attention, device=self.device\n",
    "                )\n",
    "                references.append(ground_truth)\n",
    "                candidates.append(generated_answer)\n",
    "\n",
    "        return references, candidates\n",
    "\n",
    "    def generate_answers_with_attention(self):\n",
    "        references = []\n",
    "        candidates = []\n",
    "\n",
    "        with torch.no_grad():\n",
    "            for batch in tqdm(self.test_loader, desc='Generating answers with attention'):\n",
    "                question, ground_truth = batch\n",
    "                question = question  # Move to device if using CUDA\n",
    "\n",
    "                # Generate answer with attention\n",
    "                attention_type = 'luong'\n",
    "                attention_approach = 'general'\n",
    "                generated_answer = generate_answer_attention(\n",
    "                    question, self.vocab, self.reverse_vocab,\n",
    "                    self.vocab[\"<sos>\"], self.vocab[\"<eos>\"],\n",
    "                    attn_type=attention_type, attn_approach=attention_approach,\n",
    "                    saved_model=self.model_attention, device=self.device\n",
    "                )\n",
    "\n",
    "                references.append(ground_truth)\n",
    "                candidates.append(generated_answer)\n",
    "\n",
    "        return references, candidates\n",
    "\n",
    "    def calculate_bleu_score(self, references, candidates):\n",
    "        bleu_scores = []\n",
    "        # Add tqdm progress bar for BLEU calculation\n",
    "        for ref, cand in tqdm(zip(references, candidates), desc='Calculating BLEU scores', total=len(references)):\n",
    "            ref_tokens = ref.split()\n",
    "            cand_tokens = cand.split()\n",
    "            score = sentence_bleu([ref_tokens], cand_tokens)\n",
    "            bleu_scores.append(score)\n",
    "\n",
    "        return np.mean(bleu_scores)\n",
    "\n",
    "    def calculate_cosine_similarity(self, references, candidates):\n",
    "        vectorizer = CountVectorizer().fit(references + candidates)\n",
    "        ref_vectors = vectorizer.transform(references)\n",
    "        cand_vectors = vectorizer.transform(candidates)\n",
    "        similarities = cosine_similarity(ref_vectors, cand_vectors)\n",
    "        pair_similarities = [similarities[i][i] for i in range(len(references))]\n",
    "        return np.mean(pair_similarities)\n",
    "\n",
    "    def calculate_meteor_score(self, references, candidates):\n",
    "        meteor_scores = []\n",
    "        # Add tqdm progress bar for METEOR calculation\n",
    "        for ref, cand in tqdm(zip(references, candidates), desc='Calculating METEOR scores', total=len(references)):\n",
    "            ref_tokens = ref.split()\n",
    "            cand_tokens = cand.split()\n",
    "            try:\n",
    "                score = meteor_score([ref_tokens], cand_tokens)\n",
    "                meteor_scores.append(score)\n",
    "            except Exception as e:\n",
    "                print(f\"METEOR score calculation error: {e}\")\n",
    "                meteor_scores.append(0.0)\n",
    "\n",
    "        return np.mean(meteor_scores)\n",
    "\n",
    "    def evaluate_models(self):\n",
    "        # Generate answers\n",
    "        attention_refs, attention_cands = self.generate_answers_with_attention()\n",
    "        no_attention_refs, no_attention_cands = self.generate_answers_without_attention()\n",
    "\n",
    "        # Calculate metrics for models without attention\n",
    "        no_attention_metrics = {\n",
    "            'bleu_score': self.calculate_bleu_score(no_attention_refs, no_attention_cands),\n",
    "            'cosine_similarity': self.calculate_cosine_similarity(no_attention_refs, no_attention_cands),\n",
    "            'meteor_score': self.calculate_meteor_score(no_attention_refs, no_attention_cands)\n",
    "        }\n",
    "\n",
    "        # Calculate metrics for models with attention\n",
    "        attention_metrics = {\n",
    "            'bleu_score': self.calculate_bleu_score(attention_refs, attention_cands),\n",
    "            'cosine_similarity': self.calculate_cosine_similarity(attention_refs, attention_cands),\n",
    "            'meteor_score': self.calculate_meteor_score(attention_refs, attention_cands)\n",
    "        }\n",
    "\n",
    "        return {\n",
    "            'no_attention': no_attention_metrics,\n",
    "            'attention': attention_metrics,\n",
    "            'attention_refs': attention_refs,\n",
    "            'attention_cands': attention_cands,\n",
    "            'no_attention_refs': no_attention_refs,\n",
    "            'no_attention_cands':no_attention_cands\n",
    "        }\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "gQFVNdZLqFdY",
    "outputId": "219de9f1-ac89-4e25-c506-63c505c824c5"
   },
   "outputs": [],
   "source": [
    "model_no_attention= \"trained_seq_2_seq_no_attention.pth\"\n",
    "model_attention = \"trained_seq_2_seq_attention.pth\"\n",
    "\n",
    "evaluator = ModelEvaluator(\n",
    "    model_no_attention=model_no_attention,\n",
    "    model_attention=model_attention,\n",
    "    test_loader=test_text_dialogue_pairs[:50],\n",
    "    vocab=vocab,\n",
    "    reverse_vocab=reverse_vocab,\n",
    "    start_token=\"<sos>\",\n",
    "    end_token=\"<eos>\",\n",
    "    device='cpu'\n",
    ")\n",
    "\n",
    "metrics = evaluator.evaluate_models()\n",
    "\n",
    "# Output the metrics\n",
    "print(\"No Attention Model Metrics:\", metrics['no_attention'])\n",
    "print(\"Attention Model Metrics:\", metrics['attention'])"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "author": "Tarek Hathut",
  "colab": {
   "gpuType": "T4",
   "provenance": []
  },
  "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.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}