{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Lecture 4: Recurrent Neural Networks\n",
    "\n",
    "Lecture 4 | CMU ANLP Fall 2025 | Instructor: Sean Welleck\n",
    "\n",
    "### Part 1: Recurrent language model\n",
    "\n",
    "This is a notebook for [CMU CS11-711 Advanced NLP](https://cmu-l3.github.io/anlp-fall2025/) that trains a recurrent language model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['emma',\n",
       " 'olivia',\n",
       " 'ava',\n",
       " 'isabella',\n",
       " 'sophia',\n",
       " 'charlotte',\n",
       " 'mia',\n",
       " 'amelia',\n",
       " 'harper',\n",
       " 'evelyn']"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data = open('names.txt').read().splitlines()\n",
    "data[:10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "token_to_index = {tok: i for i, tok in enumerate('abcdefghijklmnopqrstuvwxyz')}\n",
    "# Start/stop token\n",
    "token_to_index['[S]'] = 26\n",
    "# Padding token\n",
    "token_to_index['[PAD]'] = 27\n",
    "\n",
    "index_to_token = {i: tok for tok, i in token_to_index.items()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(25626, 25626)"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "def build_dataset(data):\n",
    "    X, Y = [], []\n",
    "    for item in data:\n",
    "        tokens = ['[S]'] + list(item) + ['[S]']\n",
    "        indices = [token_to_index[token] for token in tokens]\n",
    "        X.append(indices[:-1])\n",
    "        Y.append(indices[1:])\n",
    "    return X, Y\n",
    "\n",
    "# Split into train, dev, test\n",
    "import random\n",
    "random.seed(123)\n",
    "random.shuffle(data)\n",
    "\n",
    "n1 = int(0.8 * len(data))\n",
    "n2 = int(0.9 * len(data))\n",
    "\n",
    "X_train, Y_train = build_dataset(data[:n1])\n",
    "X_dev, Y_dev = build_dataset(data[n1:n2])\n",
    "X_test, Y_test = build_dataset(data[n2:])\n",
    "\n",
    "len(X_train), len(Y_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "([26, 11, 20, 0, 13, 13], 6, [26, 18, 7, 0, 8, 13], 6, 16)"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X_train[0], len(X_train[0]), X_train[1], len(X_train[1]), max(len(x) for x in X_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Write our own RNN cell \n",
    "import torch.nn as nn\n",
    "\n",
    "class RNNCell(torch.nn.Module):\n",
    "    def __init__(self, input_size, hidden_size):\n",
    "        super(RNNCell, self).__init__()\n",
    "        self.input_size = input_size\n",
    "        self.hidden_size = hidden_size\n",
    "        self.Wh = torch.nn.Linear(hidden_size, hidden_size)\n",
    "        self.Wx = torch.nn.Linear(input_size, hidden_size)\n",
    "        self.activation = torch.nn.Tanh()\n",
    "        \n",
    "    def forward(self, x, h):\n",
    "        h = self.activation(self.Wh(h) + self.Wx(x))\n",
    "        return h\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "class RNNLM(nn.Module):\n",
    "    def __init__(self, vocab_size, hidden_size):\n",
    "        super(RNNLM, self).__init__()\n",
    "        self.embedding = nn.Embedding(vocab_size, hidden_size)\n",
    "        self.rnn = RNNCell(hidden_size, hidden_size)\n",
    "        self.output = nn.Linear(hidden_size, vocab_size)\n",
    "        self.hidden_size = hidden_size\n",
    "\n",
    "    def forward(self, x, hidden=None):\n",
    "        if hidden is None:\n",
    "            hidden = self.init_hidden(x.size(0))\n",
    "        \n",
    "        x = self.embedding(x)\n",
    "        \n",
    "        outs = []\n",
    "        for i in range(x.size(1)):\n",
    "            hidden = self.rnn(x[:, i:i+1], hidden)\n",
    "            out = self.output(hidden)\n",
    "            outs.append(out)\n",
    "        \n",
    "        outs = torch.cat(outs, dim=1)\n",
    "        return outs, hidden\n",
    "\n",
    "    def init_hidden(self, batch_size):\n",
    "        return torch.zeros(batch_size, 1, self.hidden_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([1, 6, 28]), torch.Size([1, 1, 32]))"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = RNNLM(len(token_to_index), 32)\n",
    "\n",
    "x = torch.tensor(X_train[:1])\n",
    "\n",
    "output, hidden = model(x, hidden=None)\n",
    "output.size(), hidden.size()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = torch.tensor(X_train[:2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[26, 11, 20,  0, 13, 13, 27, 27, 27, 27],\n",
      "        [26, 18,  7,  0,  8, 13, 27, 27, 27, 27],\n",
      "        [26, 17, 20, 15,  4, 17, 19, 27, 27, 27],\n",
      "        [26, 12, 14, 10, 18,  7,  0,  6, 13,  0]])\n",
      "['[S]', 'l', 'u', 'a', 'n', 'n', '[PAD]', '[PAD]', '[PAD]', '[PAD]']\n",
      "['[S]', 's', 'h', 'a', 'i', 'n', '[PAD]', '[PAD]', '[PAD]', '[PAD]']\n",
      "['[S]', 'r', 'u', 'p', 'e', 'r', 't', '[PAD]', '[PAD]', '[PAD]']\n",
      "['[S]', 'm', 'o', 'k', 's', 'h', 'a', 'g', 'n', 'a']\n"
     ]
    }
   ],
   "source": [
    "def pad_batch(X_batch, Y_batch, pad_index):\n",
    "    max_len = max(len(x) for x in X_batch)\n",
    "    X_padded = torch.zeros(len(X_batch), max_len, dtype=torch.long) + pad_index\n",
    "    Y_padded = torch.zeros(len(Y_batch), max_len, dtype=torch.long) + pad_index\n",
    "    for i, (x, y) in enumerate(zip(X_batch, Y_batch)):\n",
    "        X_padded[i, :len(x)] = torch.tensor(x)\n",
    "        Y_padded[i, :len(y)] = torch.tensor(y)\n",
    "    return X_padded, Y_padded\n",
    "\n",
    "xp, yp = pad_batch(X_train[:4], Y_train[:4], token_to_index['[PAD]'])\n",
    "\n",
    "print(xp)\n",
    "for x in xp:\n",
    "    print([index_to_token[i.item()] for i in x])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([2, 6, 28]), torch.Size([2, 1, 32]))"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X_batch, Y_batch = pad_batch(X_train[:2], Y_train[:2], token_to_index['[PAD]'])\n",
    "\n",
    "output, hidden = model(X_batch, hidden=None)\n",
    "output.size(), hidden.size()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model parameters: 24028\n",
      "Epoch [1/10], Loss: 2.2432\n",
      "Epoch [1/10], Validation Loss: 2.1708\n",
      "Epoch [2/10], Loss: 2.1317\n",
      "Epoch [2/10], Validation Loss: 2.1361\n",
      "Epoch [3/10], Loss: 2.0956\n",
      "Epoch [3/10], Validation Loss: 2.1093\n",
      "Epoch [4/10], Loss: 2.0723\n",
      "Epoch [4/10], Validation Loss: 2.0942\n",
      "Epoch [5/10], Loss: 2.0562\n",
      "Epoch [5/10], Validation Loss: 2.0905\n",
      "Epoch [6/10], Loss: 2.0436\n",
      "Epoch [6/10], Validation Loss: 2.0814\n",
      "Epoch [7/10], Loss: 2.0332\n",
      "Epoch [7/10], Validation Loss: 2.0784\n",
      "Epoch [8/10], Loss: 2.0250\n",
      "Epoch [8/10], Validation Loss: 2.0732\n",
      "Epoch [9/10], Loss: 2.0181\n",
      "Epoch [9/10], Validation Loss: 2.0747\n",
      "Epoch [10/10], Loss: 2.0111\n",
      "Epoch [10/10], Validation Loss: 2.0660\n"
     ]
    }
   ],
   "source": [
    "import torch.optim as optim\n",
    "\n",
    "model = RNNLM(vocab_size=len(token_to_index), hidden_size=96)\n",
    "# Count model parameters\n",
    "print(f\"Model parameters: {sum(p.numel() for p in model.parameters())}\")\n",
    "\n",
    "# Hyperparameters\n",
    "learning_rate = 0.001\n",
    "num_epochs = 10\n",
    "batch_size = 16\n",
    "\n",
    "# Loss function and optimizer\n",
    "# NOTE: We ignore the loss whenever the target token is a padding token\n",
    "criterion = nn.CrossEntropyLoss(ignore_index=token_to_index['[PAD]'])\n",
    "\n",
    "optimizer = optim.Adam(model.parameters(), lr=learning_rate)\n",
    "\n",
    "# Training loop\n",
    "for epoch in range(num_epochs):\n",
    "    # Reshuffle the data\n",
    "    perm = torch.randperm(len(X_train))\n",
    "    X_train = [X_train[i] for i in perm]\n",
    "    Y_train = [Y_train[i] for i in perm]\n",
    "    \n",
    "    model.train()\n",
    "    total_loss = 0\n",
    "    for i in range(0, len(X_train), batch_size):\n",
    "        X_batch = X_train[i:i+batch_size]\n",
    "        Y_batch = Y_train[i:i+batch_size]\n",
    "        X_batch, Y_batch = pad_batch(X_batch, Y_batch, token_to_index['[PAD]'])\n",
    "\n",
    "        # Forward pass\n",
    "        outputs, _ = model(X_batch) # [batch_size, seq_len, vocab_size]\n",
    "        outputs = outputs.view(-1, len(token_to_index)) # [batch_size * seq_len, vocab_size]\n",
    "        Y_batch = Y_batch.view(-1) # [batch_size * seq_len]\n",
    "        loss = criterion(outputs, Y_batch)\n",
    "\n",
    "        # Backward pass and optimization\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        total_loss += loss.item()\n",
    "\n",
    "    avg_loss = total_loss / (len(X_train) // batch_size)\n",
    "    print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {avg_loss:.4f}')\n",
    "\n",
    "    # Evaluate validation loss\n",
    "    eval_loss = 0\n",
    "    model.eval()\n",
    "    with torch.no_grad():\n",
    "        for i in range(0, len(X_dev), batch_size):\n",
    "            X_batch = X_dev[i:i+batch_size]\n",
    "            Y_batch = Y_dev[i:i+batch_size]\n",
    "            X_batch, Y_batch = pad_batch(X_batch, Y_batch, token_to_index['[PAD]'])\n",
    "\n",
    "            outputs, _ = model(X_batch)\n",
    "            outputs = outputs.view(-1, len(token_to_index))\n",
    "            Y_batch = Y_batch.view(-1)\n",
    "            loss = criterion(outputs, Y_batch)\n",
    "\n",
    "            eval_loss += loss.item()\n",
    "    avg_eval_loss = eval_loss / (len(X_dev) // batch_size)\n",
    "    print(f'Epoch [{epoch+1}/{num_epochs}], Validation Loss: {avg_eval_loss:.4f}')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Sample from the model\n",
    "def sample(model, context, max_length=100):\n",
    "    model.eval()\n",
    "    output = []\n",
    "    with torch.no_grad():\n",
    "        x = torch.tensor([[token_to_index['[S]']] + context])\n",
    "        hidden = None\n",
    "        for _ in range(max_length):\n",
    "            y, hidden = model(x, hidden)\n",
    "            y = y[0, -1].softmax(dim=0)\n",
    "            y = torch.multinomial(y, 1)\n",
    "            token = index_to_token[y.item()]\n",
    "            if token == '[S]':\n",
    "                break\n",
    "            output.append(token)\n",
    "            x = y.view(1, 1)\n",
    "    return ''.join(output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "blander\n",
      "javor\n",
      "savina\n",
      "zanlyn\n",
      "edmice\n",
      "korio\n",
      "pailo\n",
      "madalyn\n",
      "hourtal\n",
      "kloie\n"
     ]
    }
   ],
   "source": [
    "for i in range(10):\n",
    "    print(sample(model, []))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "saper\n",
      "surah\n",
      "sayee\n",
      "starun\n",
      "shyala\n",
      "sudelee\n",
      "sillie\n",
      "samsen\n",
      "shiros\n",
      "ses\n"
     ]
    }
   ],
   "source": [
    "prompt = 's'\n",
    "for i in range(10):\n",
    "    out = sample(model, [token_to_index[tok] for tok in prompt])\n",
    "    print(prompt + out)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Suggested Exercises\n",
    "\n",
    "1. Use `nn.RNN` instead of our `RNNCell`. Do you have to change anything in the implementation?\n",
    "2. Change `nn.RNN` to `nn.GRU`. Do you have to change anything else in the implementation? Does the loss improve?\n",
    "3. Change `nn.RNN` to `nn.LSTM`. Do you have to change anything else in the implementation? Does the loss improve?\n",
    "4. Vary the hyperparameters (e.g., hidden size, batch size, learning rate, number of epochs). Can you find any consistent relationships between hyperparameter(s) and the loss?\n",
    "5. When the validation loss begins to increase, and the training loss is decreasing, we have evidence of **overfitting**. Can you induce this overfitting by changing the hyperparameters?\n",
    "6. Train a recurrent model on a more complex dataset. Use a tokenizer learned with BPE (either one that you train your own, or a pre-existing one)."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "prototype",
   "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
