{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "# Lecture 3: Language Modeling Fundamentals\n\nLecture 3 | CMU ANLP Spring 2026 | Instructor: Sean Welleck\n\n#### Part 2: Feedforward neural language model\n\nThis is a notebook for [CMU CS11-711 Advanced NLP](https://cmu-l3.github.io/anlp-spring2026/) that trains a feedforward language model, i.e. one based on [Bengio et al 2003, A Neural Probabilistic Language Model](https://www.jmlr.org/papers/volume3/bengio03a/bengio03a.pdf)."
  },
  {
   "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",
    "token_to_index['[S]'] = 26\n",
    "index_to_token = {i: tok for tok, i in token_to_index.items()}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Build the dataset\n",
    "\n",
    "Our dataset consists of $x,y$ pairs, where $x$ is a $(n-1)$-token context, and $y$ is a token."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([182427, 5]), torch.Size([182427]))"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "context_size = 5\n",
    "\n",
    "def build_dataset(data):\n",
    "    X, Y = [], []\n",
    "    for item in data:\n",
    "        context = [token_to_index['[S]']] * context_size\n",
    "        tokens = list(item) + ['[S]']\n",
    "        for token in tokens:\n",
    "            X.append(context)\n",
    "            Y.append(token_to_index[token])\n",
    "            context = context[1:] + [token_to_index[token]]\n",
    "\n",
    "    X = torch.tensor(X)\n",
    "    Y = torch.tensor(Y)\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",
    "X_train.shape, Y_train.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn as nn\n",
    "\n",
    "class MLPLM(nn.Module):\n",
    "    def __init__(self, vocab_size, context_size, embedding_size, hidden_size):\n",
    "        super(MLPLM, self).__init__()\n",
    "        self.embedding = nn.Embedding(vocab_size, embedding_size)\n",
    "        self.fc1 = nn.Linear(context_size * embedding_size, hidden_size)\n",
    "        self.fc2 = nn.Linear(hidden_size, vocab_size)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.embedding(x)       # (batch_size, context_size, hidden_size)\n",
    "        x = x.view(x.shape[0], -1)  # (batch_size, context_size * hidden_size)\n",
    "        x = torch.relu(self.fc1(x)) # (batch_size, hidden_size)\n",
    "        x = self.fc2(x)             # (batch_size, vocab_size)\n",
    "        return x\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[26, 26, 26, 26, 26],\n",
       "        [26, 26, 26, 26, 11]])"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = MLPLM(len(token_to_index), context_size, 64, 64)\n",
    "\n",
    "x = X_train[:2]\n",
    "x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 27])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.forward(x).shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model parameters: 24027\n",
      "Epoch [1/10], Loss: 2.2374\n",
      "Epoch [2/10], Loss: 2.1336\n",
      "Epoch [3/10], Loss: 2.1027\n",
      "Epoch [4/10], Loss: 2.0855\n",
      "Epoch [5/10], Loss: 2.0739\n",
      "Epoch [6/10], Loss: 2.0660\n",
      "Epoch [7/10], Loss: 2.0598\n",
      "Epoch [8/10], Loss: 2.0554\n",
      "Epoch [9/10], Loss: 2.0519\n",
      "Epoch [10/10], Loss: 2.0481\n"
     ]
    }
   ],
   "source": [
    "import torch.optim as optim\n",
    "\n",
    "model = MLPLM(len(token_to_index), context_size, 64, 64)\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 = 32\n",
    "\n",
    "# Loss function and optimizer\n",
    "optimizer = optim.Adam(model.parameters(), lr=learning_rate)\n",
    "criterion = nn.CrossEntropyLoss()\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[perm]\n",
    "    Y_train = Y_train[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",
    "\n",
    "        # Forward pass\n",
    "        outputs = model(X_batch)\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}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "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",
    "        context = torch.tensor(context).unsqueeze(0)\n",
    "        for i in range(max_length):\n",
    "            logits = model(context)\n",
    "            probs = torch.softmax(logits, dim=-1)\n",
    "            token = torch.multinomial(probs, num_samples=1)\n",
    "            context = torch.cat([context[:, 1:], token], dim=1)\n",
    "\n",
    "            output.append(index_to_token[token.item()])\n",
    "            if index_to_token[token.item()] == '[S]':\n",
    "                return ''.join(output)\n",
    "    return ''.join(output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tahi[S]\n",
      "araniko[S]\n",
      "calen[S]\n",
      "milezian[S]\n",
      "malayie[S]\n",
      "terrohette[S]\n",
      "keven[S]\n",
      "medanie[S]\n",
      "famoum[S]\n",
      "leion[S]\n"
     ]
    }
   ],
   "source": [
    "for i in range(10):\n",
    "    print(sample(model, [token_to_index['[S]']] * context_size))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Conditional generation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "shyanna[S]\n",
      "serai[S]\n",
      "soaddon[S]\n",
      "soadi[S]\n",
      "sureel[S]\n",
      "suer[S]\n",
      "shell[S]\n",
      "stere[S]\n",
      "shaeree[S]\n",
      "saralynneeson[S]\n"
     ]
    }
   ],
   "source": [
    "prompt = 's'\n",
    "for i in range(10):\n",
    "    out = sample(model, ([token_to_index['[S]']] * (context_size-len(prompt))) + [token_to_index[c] for c in prompt])\n",
    "    print(prompt + out)"
   ]
  }
 ],
 "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
}