{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Lecture 5: Transformers\n",
    "\n",
    "Lecture 5 | CMU ANLP Fall 2025 | Instructor: Sean Welleck\n",
    "\n",
    "\n",
    "This is a notebook for [CMU CS11-711 Advanced NLP](https://cmu-l3.github.io/anlp-fall2025/) that trains a simple transformer language model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "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": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data = open('names.txt').read().splitlines()\n",
    "data[:10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "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": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(25626, 25626)"
      ]
     },
     "execution_count": 3,
     "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": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "([26, 11, 20, 0, 13, 13], 6, [26, 18, 7, 0, 8, 13], 6, 16)"
      ]
     },
     "execution_count": 4,
     "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": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn as nn\n",
    "\n",
    "class Block(nn.Module):\n",
    "    def __init__(self, d_model, nhead, dim_ff=64, max_len=128):\n",
    "        super(Block, self).__init__()\n",
    "        self.attn = nn.MultiheadAttention(d_model, nhead, dropout=0.0, batch_first=True)\n",
    "        self.ff1 = nn.Linear(d_model, dim_ff)\n",
    "        self.ff2 = nn.Linear(dim_ff, d_model)\n",
    "        self.ln1 = nn.LayerNorm(d_model)\n",
    "        self.ln2 = nn.LayerNorm(d_model)\n",
    "        self.act = nn.ReLU()\n",
    "        self.register_buffer('mask', torch.triu(torch.ones(max_len, max_len), diagonal=1).bool())\n",
    "\n",
    "    def forward(self, x):\n",
    "        B, T, D = x.size()\n",
    "        # Pre-normalization\n",
    "        x = self.ln1(x)\n",
    "        # Self-attention\n",
    "        x2 = self.attn(x, x, x, is_causal=True, attn_mask=self.mask[:T,:T])[0]\n",
    "        # Residual connection\n",
    "        x = x + x2\n",
    "        # Pre-normalization\n",
    "        x = self.ln2(x)\n",
    "        # Feed-forward\n",
    "        x2 = self.ff2(self.act(self.ff1(x)))\n",
    "        # Residual connection\n",
    "        x = x + x2\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([10, 32, 10])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# test out the block\n",
    "block = Block(10, 2)\n",
    "x = torch.randn(10, 32, 10)\n",
    "block(x).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TransformerLM(nn.Module):\n",
    "    def __init__(self, vocab_size, d_model, nhead, num_layers, dim_ff, max_len=128):\n",
    "        super(TransformerLM, self).__init__()\n",
    "        self.embedding = nn.Embedding(vocab_size, d_model)\n",
    "        self.pos_encoder = nn.Embedding(max_len, d_model)\n",
    "        self.blocks = nn.ModuleList([\n",
    "            Block(d_model, nhead, dim_ff) for _ in range(num_layers)\n",
    "        ])\n",
    "        self.fc = nn.Linear(d_model, vocab_size)\n",
    "        self.d_model = d_model\n",
    "\n",
    "    def forward(self, x):\n",
    "        pos = torch.arange(x.size(1), device=x.device).unsqueeze(0)\n",
    "        x = self.embedding(x) + self.pos_encoder(pos)\n",
    "        for block in self.blocks:\n",
    "            x = block(x)\n",
    "        logits = self.fc(x)\n",
    "        return logits\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 6, 28])"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = TransformerLM(len(token_to_index), 64, 2, 2, 64)\n",
    "\n",
    "x = torch.tensor(X_train[:1])\n",
    "\n",
    "logits = model(x)\n",
    "logits.size()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "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": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 6, 28])"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X_batch, Y_batch = pad_batch(X_train[:2], Y_train[:2], token_to_index['[PAD]'])\n",
    "\n",
    "logits = model(X_batch)\n",
    "logits.size()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model parameters: 62236\n",
      "Epoch [1/10], Loss: 2.2521\n",
      "Epoch [1/10], Validation Loss: 2.1835\n",
      "Epoch [2/10], Loss: 2.1471\n",
      "Epoch [2/10], Validation Loss: 2.1452\n",
      "Epoch [3/10], Loss: 2.1106\n",
      "Epoch [3/10], Validation Loss: 2.1222\n",
      "Epoch [4/10], Loss: 2.0859\n",
      "Epoch [4/10], Validation Loss: 2.1071\n",
      "Epoch [5/10], Loss: 2.0657\n",
      "Epoch [5/10], Validation Loss: 2.0989\n",
      "Epoch [6/10], Loss: 2.0511\n",
      "Epoch [6/10], Validation Loss: 2.0830\n",
      "Epoch [7/10], Loss: 2.0378\n",
      "Epoch [7/10], Validation Loss: 2.0900\n",
      "Epoch [8/10], Loss: 2.0274\n",
      "Epoch [8/10], Validation Loss: 2.0700\n",
      "Epoch [9/10], Loss: 2.0168\n",
      "Epoch [9/10], Validation Loss: 2.0717\n",
      "Epoch [10/10], Loss: 2.0109\n",
      "Epoch [10/10], Validation Loss: 2.0653\n"
     ]
    }
   ],
   "source": [
    "import torch.optim as optim\n",
    "\n",
    "model = TransformerLM(len(token_to_index), 64, 2, 2, 64)\n",
    "\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": 16,
   "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",
    "        for _ in range(max_length):\n",
    "            logits = model(x)\n",
    "            y = torch.softmax(logits[0, -1], 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 = torch.cat([x, y.unsqueeze(0)], dim=1)\n",
    "    return ''.join(output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "akayah\n",
      "fallianus\n",
      "bela\n",
      "deklyn\n",
      "amelyn\n",
      "bilsea\n",
      "ramib\n",
      "paylin\n",
      "ellison\n",
      "wenn\n"
     ]
    }
   ],
   "source": [
    "for i in range(10):\n",
    "    print(sample(model, []))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "sariel\n",
      "samiyah\n",
      "sendry\n",
      "spore\n",
      "syergison\n",
      "souand\n",
      "shrah\n",
      "sannayah\n",
      "shanet\n",
      "sthyva\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": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "anlp",
   "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.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}