notebooks/seq_to_seq_toy.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Seq2Seq: ?yekruT gnivigsknahT A\n",
    "\n",
    "This notebook illustrates a neural network model whose sole purpose is to learn to reverse a given input word.  Specifically, given a sequence of characters from a word with its characters in reverse order, the model should generate the correct forward sequence of characters representing the original word. For example, if the input is \"yekrut\", the output should be \"turkey\":\n",
    "\n",
    "```\n",
    "input_seq              => output_seq\n",
    "y - e - k - r - u - t  => t - u - r - k - e - y\n",
    "```\n",
    "\n",
    "A bit silly, right?  Of course it is!  But the problem our Seq2Seq model will solve, i.e., reversing words, shares similarities with more complex tasks like machine translation and natural language generation, as tackled by models like ChatGPT.  While the word reversal problem is _much_ simpler, it serves as a stepping stone to understanding the underlying concepts and techniques that apply to those more advanced tasks.\n",
    "\n",
    "\n",
    "\n",
    "## Why build a model to reverse words?\n",
    "\n",
    "Well, this task is simple enough that we can basically nail it with a neural network -- and it's fun to win!  But the deeper reason is because it illustrates a common approach to a very common problem in ML: learning a sequence-to-sequence (seq2seq) mapping.\n",
    "\n",
    "Consider, for example, machine translation, where the goal is to translate a sentence from one language to another while preserving the meaning -- for example, translating the English sentence \"How are you?\" into Spanish as \"¿Como estas?\"  Just like our Seq2Seq model learns to convert the _reversed_ sequence of letters into the _forward_ sequence of letters, a machine translation model learns to convert a sequence of words from the source language into a sequence of words in the target language. The key point is that both tasks involve learning to map an input sequence of \"tokens\" (letters and words, respectively) to an output sequence of tokens based on learning certain rules or patterns.\n",
    "\n",
    "Or take ChatGPT-like models, designed to generate text based on a given prompt. For example, given the prompt \"Once upon a time, there was a\", the model might generate a continuation like \"beautiful castle in a magical land\", or maybe \"guy who liked data science too much, and met a bad end.\" This is structurally similar to the word reversal task: GPT-like models learn to generate an output sequence (response) based on the input sequence (prompt).  The difference is that the tokens in these sequences are words, rather than just the characters a-z.  \n",
    "\n",
    "And that illustrates the key thing that makes our problem simple enough to solve with a fairly basic network: __dimensionality__.  The possible tokens in the input sequences we'll see are limited to a-z characters (including special tokens), which is relatively small compared to the other tasks.In contrast, machine translation and natural language generation tasks have a much, much higher dimensionality. These tasks involve learning complex relationships between words, phrases, and sentences while taking into account the context, semantics, and grammar rules of the languages involved. The vast number of possible words in a language, and the combinations in which they can appear, make these tasks significantly more complex than the word reversal problem.\n",
    "\n",
    "The lower dimensionality in the word reversal task makes it easier for the model to learn the pattern and generalize to new examples. However, this simplicity also serves as a stepping stone for understanding the foundational concepts that apply to more complex tasks. By mastering the basics with a lower-dimensional problem, you can build a solid foundation before tackling more challenging tasks with higher dimensionality, such as machine translation and natural language generation.  After all, there are a lot of specific technical ideas that this Seq2Seq model illustrates:  \n",
    "\n",
    "- **Encoder-Decoder architecture**: The model uses an encoder-decoder architecture, which is a common approach for sequence-to-sequence learning. The encoder processes the input sequence and generates a context vector (hidden state) that represents the input. The decoder then takes this context vector and generates the output sequence.\n",
    "\n",
    "- **Recurrent Neural Networks (RNNs)**: The model employs Gated Recurrent Units (GRUs), which are a type of RNN, to encode and decode the input and output sequences. RNNs are suitable for sequence-to-sequence learning tasks because they can handle sequences of varying lengths and are capable of capturing temporal dependencies in the data.\n",
    "\n",
    "- **One-hot encoding**: The input sequences of characters are represented using one-hot encoding, which is a way to represent categorical variables as binary vectors. This encoding is used to convert the input sequences of integers into a tensor of shape `[batch_size, seq_length, num_chars]`, where each character is represented by a one-hot vector of dimension `num_chars`.\n",
    "\n",
    "- **Custom PyTorch Dataset and Data Loading**: The model demonstrates the creation and usage of a custom PyTorch dataset (`ReverseWordDataset`) to handle the loading, encoding, and decoding of training data. The dataset is used in conjunction with PyTorch's DataLoader to handle batching and shuffling of the data during training.\n",
    "\n",
    "Let's dive in!\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import matplotlib.pyplot as plt\n",
    "import random\n",
    "import string\n",
    "import re\n",
    "from typing import List\n",
    "import nltk\n",
    "from nltk.corpus import wordnet\n",
    "from tqdm import tqdm\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Download and process data"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To get training data, I'm going to download words from the [WordNet](https://wordnet.princeton.edu) database, saving them in `data_dir`.  Change the `data_dir` as appropriate for your local file structure."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[nltk_data] Downloading package wordnet to ../data/...\n",
      "[nltk_data]   Package wordnet is already up-to-date!\n",
      "[nltk_data] Downloading package omw-1.4 to ../data/...\n",
      "[nltk_data]   Package omw-1.4 is already up-to-date!\n"
     ]
    }
   ],
   "source": [
    "# Set download directory for NLTK\n",
    "data_dir = '../data/'\n",
    "if not os.path.exists(data_dir):\n",
    "    os.makedirs(data_dir)\n",
    "\n",
    "# Download NLTK data only if it doesn't already exist\n",
    "if not os.path.exists(os.path.join(data_dir, 'corpora/wordnet')):\n",
    "    nltk.download('wordnet', download_dir=data_dir)\n",
    "if not os.path.exists(os.path.join(data_dir, 'corpora/omw')):\n",
    "    nltk.download('omw-1.4', download_dir=data_dir)\n",
    "\n",
    "# add to the path\n",
    "nltk.data.path.append(os.path.abspath(data_dir))\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The code block below is used to generate a list of words (`word_list`) from the WordNet database, which can then be used as input for the model we'll soon build.   The structure of the WordNet database [is described here](https://www.nltk.org/howto/wordnet.html).  Briefly, WordNet is a large database of English, where words are organized into sets of synonyms called \"synsets\". Each word in a synset is called a \"lemma\".  This is more structure than we need, but to suit our purposes, all we need is a bunch of words.    Basically, if we get every lemma from every synset, we'll get a large training set of English words, which is what I'm after here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create an empty list to hold the words\n",
    "word_list = []\n",
    "letters = set('abcdfeghijklmnopqrstuvwxyz')\n",
    "max_length = 32\n",
    "\n",
    "# Define a regular expression pattern to match standard punctuation\n",
    "pattern = r'[_\\W]+'\n",
    "\n",
    "# Loop through all the synsets in WordNet\n",
    "for synset in wordnet.all_synsets():\n",
    "\n",
    "    # Loop through all the lemma names for each synset\n",
    "    for lemma in synset.lemma_names():\n",
    "\n",
    "        # Split the lemma on '_' and convert to lower case\n",
    "        words = re.split(pattern, lemma.lower())\n",
    "\n",
    "        # Loop through each word in the lemma\n",
    "        for word in words:\n",
    "\n",
    "            # Remove all characters other than the 26 lower-case consonants\n",
    "            word = ''.join(filter(lambda c: c in letters, word))\n",
    "            word = word[:max_length]\n",
    "\n",
    "            # Add the word to the word list\n",
    "            if word:\n",
    "                word_list.append(word)\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The main steps in this code block are as follows:\n",
    "\n",
    "1. **Initialize variables**: An empty list `word_list` is created to store the words, and a set `letters` is defined containing all lower-case English alphabet characters. The `max_length` variable is set to 32, which will be the maximum length for the words in `word_list`. A regular expression pattern `pattern` is defined to match standard punctuation and underscores.\n",
    "\n",
    "2. **Iterate over synsets**: The code block loops through all the synsets (sets of synonyms) in the WordNet database using the `wordnet.all_synsets()` function.\n",
    "\n",
    "3. **Iterate over lemmas**: For each synset, the code block loops through all the lemma names (words or phrases that represent the same concept as the synset) using the `synset.lemma_names()` function.\n",
    "\n",
    "4. **Split and clean lemma**: Each lemma is split on '_' (underscore) and any standard punctuation using the regular expression `pattern` with the `re.split` function. The resulting words are converted to lower case.\n",
    "\n",
    "5. **Filter and truncate words**: For each word in the lemma, the code block filters out any character that is not a lower-case English alphabet character using a lambda function with the `filter` function. The resulting word is then truncated to have a maximum length of `max_length`.\n",
    "\n",
    "6. **Add words to the list**: If the cleaned and truncated word is not empty, it is added to the `word_list`.\n",
    "\n",
    "The resulting `word_list` contains cleaned, lower-case words from the WordNet database, with a maximum length of 32 characters. These words can be used as input for the `ReverseWordDataset` class to train the Seq2Seq model.\n",
    "\n",
    "Let's randomly sample some words from the `word_list`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "296863\n",
      "['limit', 'defense', 'mathematical', 'juncus', 'prelacy', 'diane', 'overnighter', 'lipped']\n",
      "['cap', 'strop', 'newburg', 'cropper', 'city', 'developmental', 'creep', 'of']\n",
      "['tar', 'tail', 'operation', 'sensitively', 'prospectus', 'wee', 'sixth', 'osborne']\n",
      "['dropkick', 'americanus', 'saint', 'angeles', 'highlighting', 'crystalized', 'thickening', 'hibernate']\n",
      "['addlebrained', 'handed', 'thick', 'warfare', 'willis', 'intermission', 'coreid', 'cheerfulness']\n",
      "['leucas', 'snail', 'tier', 'halicarnassus', 'battery', 'sly', 'josef', 'suzerainty']\n",
      "['inconsequential', 'lutecium', 'illegitimacy', 'out', 'helianthemum', 'played', 'fringed', 'aureomycin']\n",
      "['hardpan', 'overwhelm', 'hexachlorophene', 'overmuch', 'racket', 'assyrian', 'product', 'malvaviscus']\n"
     ]
    }
   ],
   "source": [
    "# Print some words\n",
    "print(len(word_list))\n",
    "words_per_line = 8\n",
    "word_sample = random.sample(word_list, 8*8)\n",
    "for i in range(0, len(word_sample), words_per_line):\n",
    "    print(word_sample[i:i+words_per_line])\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you can see, these words cover a lot of ground!"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## A custom Dataset class\n",
    "\n",
    "Next, we need to build a `ReverseWordDataset` class, which will be a custom PyTorch dataset that handles the loading and encoding of the training data for the Seq2Seq model. Its constructed with a list of English words (hey, we've got those!) and provides methods to generate input-output pairs for training, where the input sequence is the reversed word (\"yekrut\"), and the output sequence is the correct word (\"turkey\").  The basic unit of analysis here is a __lower-case English character__: both the input and output strings are conceptualized as sequences of tokens from this limited character set (plus some special tokens, as we see below).  The main work done by this class is representing the sequence of characters as a sequence of integers.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ReverseWordDataset(Dataset):\n",
    "    def __init__(self, word_list, max_length=34):\n",
    "        self.word_list = word_list\n",
    "        self.max_length = max_length\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.word_list)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        word = self.word_list[idx]\n",
    "        reversed_word = word[::-1]\n",
    "        return self.encode_word(reversed_word), self.encode_word(word)\n",
    "\n",
    "    def encode_word(self, word):\n",
    "        encoded = [0]  # Start of sequence character\n",
    "        for char in word[:self.max_length-2]:  # Reserve space for start and end tokens\n",
    "            encoded.append(ord(char) - ord('a') + 1)\n",
    "        encoded.append(27)  # End of sequence character\n",
    "        while len(encoded) < self.max_length:\n",
    "            encoded.append(28)  # Padding character\n",
    "        return torch.tensor(encoded, dtype=torch.long)\n",
    "\n",
    "    def decode_word(self, encoded_word):\n",
    "        decoded = []\n",
    "        for idx in encoded_word:\n",
    "            if idx == 0 or idx == 27:\n",
    "                continue\n",
    "            elif idx == 28:\n",
    "                break\n",
    "            else:\n",
    "                decoded.append(chr(idx + ord('a') - 1))\n",
    "        return ''.join(decoded)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In plan English, the `ReverseWordDataset` class has the following main components.  The first three are typical PyTorch boilerplate for `Dataset` classes; the other two are where the interest is. \n",
    "\n",
    "1. Constructor (`__init__`): The constructor takes a list of English words (`word_list`) and an optional `max_length` parameter, which is set to 34 by default. The `max_length` parameter defines the maximum allowed length for the input and output sequences, including the start of sequence, end of sequence, and padding tokens.\n",
    "\n",
    "2. Length method (`__len__`): This method returns the total number of words in the dataset. It is used by the PyTorch data loader to determine the number of samples in the dataset.\n",
    "\n",
    "3. Get item method (`__getitem__`): This method is called when iterating over the dataset or when accessing individual samples. Given an index `idx`, it returns a tuple containing the input sequence (reversed word) and the output sequence (correct word), both encoded as sequences of integers. To generate the input-output pair, the method first reverses the word at index `idx` in `word_list` and then calls the `encode_word` method for both the reversed and the original word.\n",
    "\n",
    "4. Encode word method (`encode_word`): This method takes a word as input and encodes it as a sequence of integers, where the characters are represented by their positions in the alphabet (1-26), and special tokens are used for start of sequence (0), end of sequence (27), and padding (28). The method first adds the start of sequence token, then encodes each character in the word, and finally adds the end of sequence token. If the length of the encoded word is less than `max_length`, padding tokens are added to reach the desired length. This padding seems superfluous, but is done for the sake of consistency in tensor sizes and to facilitate efficient batch processing during training and inference. By padding all sequences to the same length, we can stack multiple sequences into a single tensor for efficient parallel processing, which can lead to _much_ faster training and inference times.  The encoded word is returned as a tensor of shape `[max_length]`.\n",
    "\n",
    "5. Decode word method (`decode_word`): This method takes an encoded word as input and decodes it back into its original form as a string. It iterates over the integer representations in the input sequence and converts them back into characters, skipping the start of sequence and end of sequence tokens, and stopping at padding tokens. The decoded word is returned as a string.\n",
    "\n",
    "The `ReverseWordDataset` class provides an easy way to handle the loading, encoding, and decoding of the training data for the Seq2Seq model. It ensures that the input and output sequences are properly formatted.\n",
    "\n",
    "One small bookkeeping point: The expression `ord(char) - ord('a')` is used to convert a lower-case English alphabet character to its corresponding position in the alphabet, where 'a' is at position 1, 'b' is at position 2, and so on, up to 'z' at position 26.  `ord(char)` returns the Unicode code point (integer representation) of the given character `char`. For example, `ord('a')` returns 97, and `ord('b')` returns 98.  By subtracting `ord('a')` from `ord(char)`, we effectively normalize the Unicode code points of the lower-case English alphabet characters to the range 1 to 26, corresponding to their positions in the alphabet. For example, ord('a') - ord('a') results in 0, but we add 1 to get the position in the alphabet, so 'a' is represented by 1.  We reserve index 0 for the start-of-sequence token.  "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train/test split\n",
    "\n",
    "Now let's split our original list of training words into training and testing sets and create the appropriate custom data sets using the `ReverseWordDataset` class and PyTorch's default data loaders.  We'll use a smallish testing set here; the main evaluation isn't necessarily quantitative, but to see examples of how well the model reverses particular words."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Randomly sample a set of indices for the train and test sets\n",
    "num_words = len(word_list)\n",
    "test_size = 0.05\n",
    "num_test_words = int(num_words * test_size)\n",
    "np.random.seed(42)\n",
    "test_ind = np.random.choice(num_words, size=num_test_words, replace=False)\n",
    "train_ind = np.array([i for i in range(num_words) if i not in test_ind])\n",
    "\n",
    "# Assign words in word_list to word_list_train and word_list_test based on the sampled indices\n",
    "train_word_list = [word_list[i] for i in train_ind]\n",
    "test_word_list = [word_list[i] for i in test_ind]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "282020\n",
      "14843\n"
     ]
    }
   ],
   "source": [
    "# Create the training and testing datasets\n",
    "train_dataset = ReverseWordDataset(train_word_list)\n",
    "test_dataset = ReverseWordDataset(test_word_list)\n",
    "print(len(train_dataset))\n",
    "print(len(test_dataset))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's take a random sample from the testing set to see the forward (target) word and reversed (input) word."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "hold\n",
      "dloh\n"
     ]
    }
   ],
   "source": [
    "rev, orig = test_dataset.__getitem__(np.random.randint(0, len(test_dataset)))\n",
    "print(test_dataset.decode_word(orig))\n",
    "print(test_dataset.decode_word(rev))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Network structure\n",
    "\n",
    "Let's build a network, starting with an encoder module that will see an input sequence of characters (our reversed word) and \"encode\" that information in a vector of fixed size."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Encoder:**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Encoder(nn.Module):\n",
    "    def __init__(self, num_chars, hidden_size, num_layers=1):\n",
    "        super(Encoder, self).__init__()\n",
    "        self.num_chars = num_chars\n",
    "        self.hidden_size = hidden_size\n",
    "        self.num_layers = num_layers\n",
    "        self.gru = nn.GRU(num_chars, hidden_size, num_layers, batch_first=True)\n",
    "\n",
    "    def forward(self, input_batch):\n",
    "        # One-hot encode the input\n",
    "        one_hot = torch.zeros(input_batch.size(0), input_batch.size(1), self.num_chars).to(input_batch.device)\n",
    "        one_hot.scatter_(2, input_batch.unsqueeze(2), 1)\n",
    "\n",
    "        # Encode the one-hot encoded input using a GRU\n",
    "        _, hidden = self.gru(one_hot)\n",
    "\n",
    "        return hidden\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The Encoder consists of two main components:\n",
    "\n",
    "- One-hot encoding (`self.one_hot_encoding`): Given an input batch of shape `[batch_size, seq_length]` containing sequences of integers, the Encoder applies a one-hot encoding using a linear layer with an identity activation function. This results in a tensor of shape `[batch_size, seq_length, num_chars]`, where num_chars is the number of character classes (29, including special tokens for start/end of sequence and pad).  \n",
    "- Gated Recurrent Unit (GRU) layer (`self.gru`): The Encoder processes the one-hot encoded input sequence using a GRU layer. The GRU is a type of recurrent neural network (RNN) that can capture and model the dependencies between elements in a sequence. The GRU layer takes the one-hot encoded input sequence and outputs a hidden state of shape `[1, batch_size, hidden_size]`. (The `1` as the size of the first dimension corresponds to a single time step, in this case the final one.)  This hidden state is passed to the Decoder module as the initial hidden state.\n",
    "\n",
    "The one-hot encoding in the encoder module is a way to represent the input characters as binary vectors of fixed size. Each character is represented as a vector where only one element is set to 1 (\"hot\"), while the other elements are set to 0 (\"cold\"). The position of the hot element corresponds to the index of the character in our alphabet (which needs also to have positions for the special tokens).\n",
    "\n",
    "For example, let's consider the letter 'b'. Assuming our alphabet consists of lowercase English letters (here omitting the special start/stop/pad tokens for simplicity), 'b' would be the second character in the alphabet. We can represent it using the following one-hot vector:\n",
    "\n",
    "$$\n",
    "\\text{one-hot}(\\text{`b'}) = [0, 1, 0, 0, \\dots, 0]\n",
    "$$\n",
    "\n",
    "Here, the second element of the vector is set to 1, and all other elements are set to 0.  Keep in mind that in our one-hot vectors, we actually set separate positions for \"start of sequence,\" \"end of sequence\", and \"pad\" characters.\n",
    "\n",
    "**One-hot encoding vs embedding.**  One-hot encoding is different from using embedding vectors, where each character would be represented by a dense vector of fixed size, learned by the model during training. Embedding vectors are actually the common practice in natural language applications of seq2seq models.  But while embedding vectors can capture more complex relationships between characters, one-hot encoding is sufficient for our word reversal task, as it is a simpler problem with lower dimensionality. We don't need to capture semantic relationships between characters; we only need a way to represent each character uniquely so that the model can learn the pattern of reversing the input sequence.  With only 29 unique characters including special tokens, our one-hot vectors only need to be dimension 29. This approach would break down, however, if we needed to encode all the words in English, because the dimensionality of the one-hot vectors would increase significantly, making the representation very sparse and inefficient. For each word, we would need a one-hot vector with a dimension equal to the number of unique words in the vocabulary, which could easily reach tens or even hundreds of thousands. This would result in high memory usage and increased computational complexity, making the training process much slower and less efficient. In such cases, using embedding vectors to represent words is more practical, as they provide a dense, lower-dimensional representation that can capture semantic relationships between words while being more computationally efficient.\n",
    "\n",
    "\n",
    "Note that the `forward` method of the GRU module in principle expects that we provide an initial hidden state in order to perform the encoding.  [But the PyTorch API provides a default value of a vector of zeros](https://pytorch.org/docs/stable/generated/torch.nn.GRU.html) for this hidden state, which is perfectly fine, so we don't specify it here.  "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Decoder:**\n",
    "\n",
    "Next we need a Decoder module, which is responsible for generating the target sequence (correct word) one character at a time, conditioned on the initial hidden state from the Encoder and the previously predicted characters. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Decoder(nn.Module):\n",
    "    def __init__(self, num_chars, hidden_size, num_layers=1):\n",
    "        super(Decoder, self).__init__()\n",
    "        self.num_chars = num_chars\n",
    "        self.hidden_size = hidden_size\n",
    "        self.num_layers = num_layers\n",
    "\n",
    "        self.gru = nn.GRU(num_chars, hidden_size, num_layers, batch_first=True)\n",
    "        self.fc = nn.Linear(hidden_size, num_chars)\n",
    "\n",
    "    def forward(self, input_step, hidden):\n",
    "        # One-hot encode the input step\n",
    "        one_hot = torch.zeros(input_step.size(0), input_step.size(1), self.num_chars).to(input_step.device)\n",
    "        one_hot.scatter_(2, input_step.unsqueeze(2), 1)\n",
    "\n",
    "        # Decode the one-hot encoded input step using a GRU\n",
    "        # we must pass a hidden state as well, which we expect to be the \n",
    "        # final hidden state from an encoding of the input sequence\n",
    "        output, hidden = self.gru(one_hot, hidden)\n",
    "\n",
    "        # Output the logits for each character class (0-28)\n",
    "        logits = self.fc(output)\n",
    "\n",
    "        return logits, hidden\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The Decoder consists of two components:\n",
    "\n",
    "1. Gated Recurrent Unit (GRU) layer (`self.gru`): The Decoder processes the input step (one-hot encoded character) and the current hidden state using a GRU layer. The input step is of shape `[batch_size, 1, num_chars]`, and the hidden state is of shape `[1, batch_size, hidden_size]`. (As with the encoder, the `1` as the size of the first dimension corresponds to a single time step.)  The GRU layer updates the hidden state and outputs a new hidden state of the same shape (`[1, batch_size, hidden_size]`).\n",
    "\n",
    "2. Linear layer (`self.fc`): The Decoder applies a linear layer to the output of the GRU to generate logits for each character class (0-28). The linear layer takes the hidden state of shape `[1, batch_size, hidden_size]` and outputs logits of shape `[1, batch_size, num_chars]`.  These logits represent the probabilities of each character being the next character in the target sequence."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Seq2Seq:**\n",
    "\n",
    "The `Seq2Seq` module is a higher-level wrapper that combines the Encoder and Decoder modules to perform sequence-to-sequence learning. The main goal of the Seq2Seq model is to learn the mapping between the input sequence (reversed word) and the target sequence (correct word). The Seq2Seq module consists of an Encoder and a Decoder, which we pass in at initialization.\n",
    "\n",
    "Let's see the code:  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Seq2Seq(nn.Module):\n",
    "    def __init__(self, encoder, decoder):\n",
    "        super(Seq2Seq, self).__init__()\n",
    "        self.encoder = encoder\n",
    "        self.decoder = decoder\n",
    "\n",
    "    def forward(self, input_batch, target_batch):\n",
    "        # Encode the input batch\n",
    "        hidden = self.encoder(input_batch)\n",
    "\n",
    "        # One-hot encode the target sequence (excluding the last character)\n",
    "        target_one_hot = torch.zeros(target_batch.size(0), target_batch.size(1) - 1, self.decoder.num_chars).to(target_batch.device)\n",
    "        target_one_hot.scatter_(2, target_batch[:, :-1].unsqueeze(2), 1)\n",
    "\n",
    "        # Decode the target sequence using a GRU\n",
    "        output, _ = self.decoder.gru(target_one_hot, hidden)\n",
    "\n",
    "        # Output the logits for each character class (0-28)\n",
    "        logits = self.decoder.fc(output)\n",
    "\n",
    "        return logits\n",
    "\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There are three main components here:\n",
    "\n",
    "1. Encoding the input sequence: The Seq2Seq module first encodes the input sequence (reversed word) using the Encoder. The Encoder takes the input batch of shape `[batch_size, seq_length]` and outputs a hidden state of shape `[1, batch_size, hidden_size]`. This hidden state captures the relevant information from the input sequence and is used as the initial hidden state for the Decoder.\n",
    "\n",
    "2. Decoding the target sequence: The Seq2Seq module then generates the target sequence (correct word) one character at a time using the Decoder. The Decoder takes the initial hidden state from the Encoder and the previously predicted characters as inputs, and outputs logits for each character class (0-28). These logits represent the probabilities of each character being the next character in the target sequence.\n",
    "\n",
    "3. Excluding the last character: When generating the target sequence, the Seq2Seq module excludes the last character from the output. The reason for this is that the model is trained to predict the end of sequence token (27) at the end of the target sequence, and this token is not part of the actual word we want to reverse. By excluding the last character, we ensure that the output only contains the characters of the reversed word without the end of sequence token.\n",
    "\n",
    "The combination of the Encoder and Decoder modules in the Seq2Seq model enables it to capture the relevant information from the input sequence and generate the target sequence one character at a time. "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### A helper function to make and view predictions\n",
    "\n",
    "Below we create a `reverse_words` function that generates a predicted output sequence, given an input sequence.  So the way our model is designed, if we input \"tnals\" it should output \"slant\", the reverse of the input.  The key idea in the `reverse_words` function is that we generate the target sequence (reversed input) one character at a time, conditioned on the previously predicted characters and the initial hidden state from the encoder.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "def reverse_words(words, seq2seq, dataset, max_length=34):\n",
    "    # Encode the input words\n",
    "    input_batch = [dataset.encode_word(word) for word in words]\n",
    "    device = next(seq2seq.parameters()).device\n",
    "    # Create a tensor with the encoded words\n",
    "    input_batch = torch.stack(input_batch).to(device)\n",
    "\n",
    "    # Initialize the input step with the start of sequence token (0)\n",
    "    input_step = torch.zeros(input_batch.size(0), 1, dtype=torch.long).to(device)\n",
    "\n",
    "    # Initialize the hidden state with the output from the encoder\n",
    "    hidden = seq2seq.encoder(input_batch)\n",
    "\n",
    "    # Decode the target sequence one character at a time\n",
    "    decoded_words = []\n",
    "    for _ in range(max_length):\n",
    "        logits, hidden = seq2seq.decoder(input_step, hidden)\n",
    "        predictions = torch.argmax(logits, dim=-1)\n",
    "        decoded_words.append(predictions.squeeze(1))\n",
    "        input_step = predictions\n",
    "\n",
    "    # Transpose the list of decoded words\n",
    "    decoded_words = torch.stack(decoded_words).transpose(0, 1)\n",
    "\n",
    "    # Turn the predicted target sequences back into strings\n",
    "    reversed_words = [dataset.decode_word(encoded_word) for encoded_word in decoded_words]\n",
    "\n",
    "    return reversed_words\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There are several important steps here.\n",
    "\n",
    "1. First, we take each word in the input list and encode it using the encode_word method from the dataset object. This method converts each character in the reversed word to its corresponding integer representation and adds special tokens for start of sequence (0), end of sequence (27), and padding (28). Then, we stack the encoded words into a tensor called `input_batch`.  This is moved to whatever device that holds the model we want to use to predict (`seq2seq`).\n",
    "2. Initialize the input step and hidden state.  Here we create a tensor called `input_step` with the same batch size as the input batch and initialize it with the start-of-sequence token (0). This will be used as the initial input to the decoder.  We also encode the input batch using the encoder; you'll recall that the point of this encoder is to process the input data and represent it using a final hidden state vector, thereby obtaining the initial hidden state for the decoder.\n",
    "3. Decode the target sequence one character at a time.  We use a for loop to predict the characters in the target sequence one by one. For each step in the loop:\n",
    "   - Pass the input_step and the current hidden state to the decoder. The decoder processes the input step using its GRU and outputs logits for each character class (0-28). It also updates the hidden state.\n",
    "   - Find the character with the highest logit (probability) using torch.argmax and add it to the list of decoded words.\n",
    "   Update the input_step with the predicted character for the next step.\n",
    "4. Transpose the list of decoded words.  After the loop, we have a list of predicted characters for each word in the batch. We stack these characters into a tensor and transpose it, so that each row corresponds to one decoded (reversed) word.  \n",
    "5. Decode the predicted target sequences back into strings.  Here we use the `decode_word` method from the `dataset` object to convert the integer representations of the characters back into their original form. We then return the list of reversed inputs as output.\n",
    "\n",
    "Notice that in predicting the next token, we follow a naive greedy approach.  To generate the predicted character, we find the character with the highest logit (probability) using `torch.argmax`. We then update the input step with this predicted character, passing it to the next step in the decoding loop. This process is repeated until all characters in the target sequence are predicted.  This greedy, word-at-a-time prediction definitely isn't the smartest way to generate an actual predicted sequence from the logits.  While it works just fine for a simple problem like this one, for more complex tasks like machine translation or natural language generation, this approach could lead to suboptimal outputs, as it doesn't take into account the global context or explore alternative possibilities. In such cases, more sophisticated techniques like beam search or sampling-based methods can be employed to generate higher-quality predictions that better account for the overall sequence structure and token probabilities."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training\n",
    "\n",
    "Now let's set up and train the model.  I have a simple flag to check where a GPU is available on the system and, if so, to train on that."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "# Hyperparameters\n",
    "num_chars = 29  # Including start of sequence (0), end of sequence (27), and padding (28) tokens\n",
    "hidden_size = 128\n",
    "num_layers = 1\n",
    "\n",
    "# Create the model\n",
    "encoder = Encoder(num_chars, hidden_size, num_layers)\n",
    "decoder = Decoder(num_chars, hidden_size, num_layers)\n",
    "seq2seq = Seq2Seq(encoder, decoder).to(device)\n",
    "\n",
    "# Hyperparameters\n",
    "learning_rate = 0.001\n",
    "num_epochs = 10\n",
    "batch_size = 128\n",
    "\n",
    "# Define the loss function and optimizer\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = optim.Adam(seq2seq.parameters(), lr=learning_rate)\n",
    "\n",
    "train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)\n",
    "test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=False)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now for the main training loop.  After each pass through the data, I'll print out the training and test set losses.  I'll also draw a random sample of words from the test set and print out their input/output sequences, so you can a bit more qualitatively how the model is performing.  In the beginning its word reversals are basically garbage, but you'll see they get pretty good pretty quickly.\n",
    "\n",
    "Note that the iterations per second below reflect my training this model on a decent GPU, where I see performance that's 50-100x faster than on the CPU of a Macbook Pro.   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2204it [00:10, 204.33it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/10, Training Loss: 0.2846892178058624, Test Loss: 0.30316397419263574\n",
      "   Input    : ['dioporhtna', 'revoc', 'redro', 'nailartsua', 'lanoitacova', 'egatolip', 'citsilegnave', 'etavitluc']\n",
      "   Target   : ['anthropoid', 'cover', 'order', 'australian', 'avocational', 'pilotage', 'evangelistic', 'cultivate']\n",
      "   Predicted: ['arophitone', 'conte', 'orone', 'antariatia', 'anaconation', 'pintoran', 'ingencorinaly', 'culicalit']\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2204it [00:11, 195.21it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 2/10, Training Loss: 0.05407371371984482, Test Loss: 0.08498447960049942\n",
      "   Input    : ['ecanruf', 'tcerrocni', 'taob', 'livic', 'nemarof', 'gnilles', 'elkcarg', 'yldekan']\n",
      "   Target   : ['furnace', 'incorrect', 'boat', 'civil', 'foramen', 'selling', 'grackle', 'nakedly']\n",
      "   Predicted: ['furnace', 'incorrept', 'boat', 'civil', 'foramio', 'selling', 'surouth', 'nakedly']\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2204it [00:10, 212.05it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 3/10, Training Loss: 0.02194417268037796, Test Loss: 0.03142030179050976\n",
      "   Input    : ['dlohesuoh', 'me', 'yelnats', 'ssessoperp', 'tnarebuxe', 'suneg', 'nus', 'yawgnimeh']\n",
      "   Target   : ['household', 'em', 'stanley', 'prepossess', 'exuberant', 'genus', 'sun', 'hemingway']\n",
      "   Predicted: ['household', 'em', 'stanley', 'prepossess', 'exuncrict', 'genus', 'sun', 'hemingway']\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2204it [00:09, 226.40it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 4/10, Training Loss: 0.007696435321122408, Test Loss: 0.019255242254263883\n",
      "   Input    : ['rolias', 'gib', 'pleh', 'maddog', 'ssenhsarb', 'tnias', 'dnalrag', 'suilofixal']\n",
      "   Target   : ['sailor', 'big', 'help', 'goddam', 'brashness', 'saint', 'garland', 'laxifolius']\n",
      "   Predicted: ['sailor', 'big', 'help', 'goddam', 'brashness', 'saint', 'garland', 'laxifolius']\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2204it [00:09, 227.36it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 5/10, Training Loss: 0.01494721882045269, Test Loss: 0.010845165065457595\n",
      "   Input    : ['noitcudorp', 'sennahoj', 'tej', 'suneg', 'naol', 'teliot', 'nairatinirt', 'esnopser']\n",
      "   Target   : ['production', 'johannes', 'jet', 'genus', 'loan', 'toilet', 'trinitarian', 'response']\n",
      "   Predicted: ['production', 'johannes', 'jet', 'genus', 'loan', 'toilet', 'trinitarian', 'response']\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2204it [00:09, 232.23it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 6/10, Training Loss: 0.006273152306675911, Test Loss: 0.006872830081506279\n",
      "   Input    : ['tinu', 'suneg', 'etirekozo', 'diuqil', 'etagorretni', 'ffo', 'noitareggaxe', 'yawa']\n",
      "   Target   : ['unit', 'genus', 'ozokerite', 'liquid', 'interrogate', 'off', 'exaggeration', 'away']\n",
      "   Predicted: ['unit', 'genus', 'ozokerite', 'liquid', 'interrogate', 'off', 'exaggeration', 'away']\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2204it [00:09, 236.98it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 7/10, Training Loss: 0.0017806236864998937, Test Loss: 0.004037604324021858\n",
      "   Input    : ['eht', 'gabliam', 'segnag', 'ciahcort', 'hgih', 'ronim', 'citsinamuh', 'noitacifirtep']\n",
      "   Target   : ['the', 'mailbag', 'ganges', 'trochaic', 'high', 'minor', 'humanistic', 'petrification']\n",
      "   Predicted: ['the', 'mailbag', 'ganges', 'trochaic', 'high', 'minor', 'humanistic', 'petrification']\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2204it [00:09, 232.56it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 8/10, Training Loss: 0.0038482698146253824, Test Loss: 0.0037217930612627997\n",
      "   Input    : ['qs', 'ksihw', 'noitasitametsys', 'robal', 'pirts', 'ssessoperp', 'hafilahk', 'lairtserretartxe']\n",
      "   Target   : ['sq', 'whisk', 'systematisation', 'labor', 'strip', 'prepossess', 'khalifah', 'extraterrestrial']\n",
      "   Predicted: ['sq', 'whisk', 'systematictian', 'labor', 'strip', 'prepossess', 'khalifah', 'extraterrestrial']\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2204it [00:10, 218.71it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 9/10, Training Loss: 0.0013206122675910592, Test Loss: 0.0038795454225880256\n",
      "   Input    : ['u', 'kcimroccm', 'lenap', 'aiknab', 'annam', 'tsylana', 'noitatic', 'wen']\n",
      "   Target   : ['u', 'mccormick', 'panel', 'bankia', 'manna', 'analyst', 'citation', 'new']\n",
      "   Predicted: ['u', 'mccormick', 'panel', 'bankia', 'manna', 'analyst', 'citation', 'new']\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2204it [00:09, 227.52it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 10/10, Training Loss: 0.004597467835992575, Test Loss: 0.0029354443094203376\n",
      "   Input    : ['evitomotua', 'surusohcirt', 'diuqs', 'lausiv', 'elbaekil', 'epacsretaw', 'egamad', 'dnuf']\n",
      "   Target   : ['automotive', 'trichosurus', 'squid', 'visual', 'likeable', 'waterscape', 'damage', 'fund']\n",
      "   Predicted: ['automotive', 'trichosurus', 'squid', 'visual', 'likeable', 'waterscape', 'damage', 'fund']\n"
     ]
    }
   ],
   "source": [
    "# Main training loop\n",
    "for epoch in range(num_epochs):\n",
    "    # Training phase\n",
    "    seq2seq.train()\n",
    "    for batch_idx, (input_batch, target_batch) in tqdm(enumerate(train_loader)):\n",
    "        # Forward pass\n",
    "        input_batch, target_batch = input_batch.to(device), target_batch.to(device)\n",
    "        logits = seq2seq(input_batch, target_batch)\n",
    "\n",
    "        # Compute the loss\n",
    "        loss = criterion(logits.view(-1, num_chars), target_batch[:, 1:].reshape(-1))\n",
    "\n",
    "        # Backward pass\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "    # Evaluation phase\n",
    "    seq2seq.eval()\n",
    "    test_loss = 0\n",
    "    with torch.no_grad():\n",
    "        for batch_idx, (input_batch, target_batch) in enumerate(test_loader):\n",
    "            # Forward pass\n",
    "            input_batch, target_batch = input_batch.to(device), target_batch.to(device)\n",
    "            logits = seq2seq(input_batch, target_batch)\n",
    "\n",
    "            # Compute the loss\n",
    "            batch_loss = criterion(logits.view(-1, num_chars), target_batch[:, 1:].reshape(-1))\n",
    "            test_loss += batch_loss.item()\n",
    "\n",
    "    # Calculate the average test loss\n",
    "    avg_test_loss = test_loss / len(test_loader)\n",
    "\n",
    "    print(f\"Epoch {epoch+1}/{num_epochs}, Training Loss: {loss.item()}, Test Loss: {avg_test_loss}\")\n",
    "    test_word_sample = random.sample(test_word_list, 8)\n",
    "    input_words_rev = [word[::-1] for word in test_word_sample]\n",
    "    reversed_words = reverse_words(input_words_rev, seq2seq, test_dataset)\n",
    "    print(f\"   Input    : {input_words_rev}\")\n",
    "    print(f\"   Target   : {test_word_sample}\")\n",
    "    print(f\"   Predicted: {reversed_words}\")\n",
    "\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you can see, by the 10th training epoch, the model is nearly perfect!  Let's try a few hand-picked words just for fun, including a whopper:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Input    : ['tsurt', 'citsispilos', 'nacsurte', 'dneherppa', 'atad', 'aimenietorpopilatebrepyh']\n",
      "Target   : ['trust', 'solipsistic', 'etruscan', 'apprehend', 'data', 'hyperbetalipoproteinemia']\n",
      "Predicted: ['trust', 'solipsistic', 'etruscan', 'apprehend', 'data', 'hypripelatpelapeletizean']\n"
     ]
    }
   ],
   "source": [
    "# Example usage\n",
    "input_words = [\"trust\", \"solipsistic\", \"etruscan\", \"apprehend\", \"data\", \"hyperbetalipoproteinemia\"]\n",
    "input_words_rev = [word[::-1] for word in input_words]\n",
    "reversed_words = reverse_words(input_words_rev, seq2seq, test_dataset)\n",
    "\n",
    "print(f\"Input    : {input_words_rev}\")\n",
    "print(f\"Target   : {input_words}\")\n",
    "print(f\"Predicted: {reversed_words}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "OK, so the model can't quite nail [hyperbetalipoproteinemia](https://pubmed.ncbi.nlm.nih.gov/8573733/), but it does really well at nearly everything else!"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Recap\n",
    "\n",
    "In this learning exercise, we have built a simple Seq2Seq model to reverse words using PyTorch. While this model is relatively simple, it serves as a foundation to understand the basic concepts and techniques that can be applied to more complex tasks, such as machine translation or natural language generation.\n",
    "\n",
    "To handle more complex tasks, you can consider the following modifications:\n",
    "\n",
    "1. **Embeddings instead of one-hot encodings**: Using embeddings allows the model to capture semantic relationships between input tokens and results in a more efficient representation. By \"semantic relationships,\" I mean that tokens with similar meaning would end up with similar vector representations, e.g., the words \"dog\" and \"puppy\" would have closer vector representations compared to \"dog\" and \"bacteria,\" or the words \"likable\" and \"pleasant\" would have closer vector representations compared to \"likable\" and \"Elon Musk.\"  One-hot encodings are limited in their ability to capture semantic relationships like this, because they represent each token as a binary vector with a single '1' (hot) in the position corresponding to the token's index and '0' (cold) in all other positions. As a result, any two different tokens by definition must have orthogonal one-hot encodings. This means that the dot product between any two distinct one-hot encoded vectors is always 0, making it impossible to infer any similarity or relationship between them.  You can use PyTorch's [nn.Embedding](https://pytorch.org/docs/stable/generated/torch.nn.Embedding.html) to create and train embedding layers for your Seq2Seq model.\n",
    "\n",
    "2. **Deeper models**: Adding more layers to your encoder and decoder can help capture more complex patterns and dependencies in the input and output sequences. This can be achieved using [nn.GRU](https://pytorch.org/docs/stable/generated/torch.nn.GRU.html) or [nn.LSTM](https://pytorch.org/docs/stable/generated/torch.nn.LSTM.html) with a larger `num_layers` argument.\n",
    "\n",
    "3. **Other encoders**: Instead of using a GRU-based encoder, you could experiment with LSTM or attention-based encoders. LSTMs can be implemented using [nn.LSTM](https://pytorch.org/docs/stable/generated/torch.nn.LSTM.html) in PyTorch, while attention mechanisms can be incorporated using [nn.MultiheadAttention](https://pytorch.org/docs/stable/generated/torch.nn.MultiheadAttention.html). The original paper for LSTMs can be found [here](https://www.mitpressjournals.org/doi/abs/10.1162/neco.1997.9.8.1735), and the paper introducing attention mechanisms is available [here](https://arxiv.org/abs/1409.0473).\n",
    "\n",
    "4. **Better generation strategies.** Greedy, token-at-a-time prediction works fine here, but can generate bland and repetitive text in more typical NLP applications where the tokens are words.  Some smarter alternatives are:\n",
    "    - **Beam search**: Beam search is a popular search algorithm used in many sequence-to-sequence tasks. It maintains a fixed-size set (beam) of the most likely partial sequences (hypotheses) at each step and expands them further. This helps to avoid the limitations of a purely greedy approach by exploring a larger search space, potentially leading to better final predictions. You can find a good [introduction to beam search on Wikipedia](https://en.wikipedia.org/wiki/Beam_search)\n",
    "    - **Sampling-based methods**: Sampling-based methods, such as top-k sampling and nucleus (top-p) sampling, are techniques used in natural language generation tasks to create more diverse and contextually relevant outputs. Instead of always choosing the most likely token, these methods sample tokens from the probability distribution of the logits, introducing some randomness while still favoring higher-probability tokens. Here are some resources for sampling-based methods: Top-k sampling: [Holtzman, A. et al. (2019). The Curious Case of Neural Text Degeneration.](https://arxiv.org/abs/1904.09751); Nucleus (top-p) sampling: [Radford, A. et al. (2019). Language Models are Unsupervised Multitask Learners.](https://d4mucfpksywv.cloudfront.net/better-language-models/language_models_are_unsupervised_multitask_learners.pdf)\n",
    "\n",
    "\n",
    "By exploring these advanced techniques and architectures, you can build more powerful models that can tackle a wide range of sequence-to-sequence learning tasks, including machine translation, natural language generation, and more.\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.10"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}