{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VRdIio8gyD5F"
      },
      "source": [
        "# **Assignment 1: Building a Simple Language Model**\n",
        "\n",
        "In this assignment, you will build and extend a tiny neural language model step by step. Starting from a single-vector bag-of-words baseline, you will gradually add positional information, non-linear transformations, MLP layers, residual connections, and gating mechanisms to see how each architectural choice changes the model's ability to predict the next token.\n",
        "\n",
        "From this assignment, you will know how to construct the core modules of a language model, training it on data, and understanding how architectural design choices affect model performance.\n",
        "\n",
        "**Key Components You Will Implement**:\n",
        "\n",
        "**Cross-Entropy Loss**: Measures next-token prediction accuracy.\\\n",
        "**Training Loop**: Forward pass, loss computation, and gradient update.\\\n",
        "**Positional Embeddings**: Encodes token order to enable sequence modeling.\\\n",
        "**Feed-Forward Network (MLP)**: Learning a function of previous token that is more useful than a simple average\\\n",
        "**Residual Connections**: Stabilizes training and enhances gradient flow.\\\n",
        "**Gating Mechanism**: Regulates information flow across layers.\n",
        "\n",
        "You will write code inside the `TODO` blocks and answer short conceptual questions.\n",
        "\n",
        "At then end you will write a paragraph or two discussing your findings\n",
        "\n",
        "We will be using a logging framework called TensorBoard to record metrics during training in order to visualize and compare metrics across different architectures.\n",
        "\n",
        "A GPU is recommended but not required. You can select the  T4 GPU through Colab by going to Runtime > Change Runtime Type > T4 GPU"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MyQeJ7iYyHIp"
      },
      "source": [
        "# **Imports and Dependencies**"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install pygtrie datasets\n",
        "!pip -q install torchviz graphviz"
      ],
      "metadata": {
        "id": "3mzYdObhpEgO"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dK7ek49OuJOv"
      },
      "outputs": [],
      "source": [
        "import itertools\n",
        "import json\n",
        "from collections import Counter\n",
        "from typing import Dict, List\n",
        "import datasets\n",
        "import numpy as np\n",
        "import pygtrie\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torch.optim as optim\n",
        "from torch.utils.data import Dataset, DataLoader, IterableDataset\n",
        "from torch.utils.tensorboard import SummaryWriter\n",
        "from torchviz import make_dot"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c-8yYry-u3kf"
      },
      "source": [
        "# **Tokenizer**\n",
        "\n",
        "Here is the same byte-pairing tokenizer that you implemented in A0, except it has been trained with a vocab size of 65536 on linux text data and a special token has been added to identify the end of a document. Try it on a couple example sentences.\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "class Tokenizer:\n",
        "    def __init__(self, vocab_path=None, vocab=None, max_token_length=30, partial_trie=None, special_tokens = None):\n",
        "        \"\"\"\n",
        "        special_tokens: optional dict like {'eos_token': '<eos>', 'pad_token': '<pad>'}\n",
        "        These will be assigned fresh IDs at the end of the vocab and stored separately.\n",
        "        \"\"\"\n",
        "        print(f\"\\nInitializing Tokenizer:\")\n",
        "        print(f\"  - vocab_path: {vocab_path}\")\n",
        "        print(f\"  - vocab size: {len(vocab) if vocab else 'None'}\")\n",
        "        print(f\"  - max_token_length: {max_token_length}\")\n",
        "        self.vocab_path = vocab_path\n",
        "        self.id_to_tok = []\n",
        "        self.trie = pygtrie.Trie()\n",
        "        self.trie = self.trie if partial_trie is None else partial_trie\n",
        "        self.max_token_length = max_token_length\n",
        "\n",
        "        #adding a special eos token to seperate between rows\n",
        "        self.special_tokens = {}\n",
        "        self.special_id_to_name = {}\n",
        "\n",
        "        if vocab_path or vocab:\n",
        "            self.update_trie(vocab)\n",
        "\n",
        "        if special_tokens:\n",
        "            self.add_special_tokens(special_tokens)\n",
        "    @property\n",
        "    def vocab_size(self):\n",
        "        return len(self.id_to_tok) + len(self.special_tokens)\n",
        "\n",
        "    @property\n",
        "    def eos_token(self):\n",
        "      return self.special_tokens.get('eos_token', (None,None))[0]\n",
        "\n",
        "    @property\n",
        "    def eos_token_id(self):\n",
        "      return self.special_tokens.get('eos_token', (None,None))[1]\n",
        "\n",
        "    def add_special_tokens(self, mapping):\n",
        "        \"\"\"\n",
        "        mapping: {'eos_token': '<eos>', 'pad_token': '<pad>', ...}\n",
        "        Assigns new IDs after the base vocab (not added to trie).\n",
        "        \"\"\"\n",
        "        base = len(self.id_to_tok)\n",
        "        for i, (name, text_token) in enumerate(mapping.items()):\n",
        "            tok_id = base + len(self.special_tokens)\n",
        "            self.special_tokens[name] = (text_token, tok_id)\n",
        "            self.special_id_to_name[tok_id] = name\n",
        "        print(f\"Added special tokens: { {k:v[0] for k,v in self.special_tokens.items()} }\")\n",
        "\n",
        "    def update_trie(self, new_vocab=None):\n",
        "        print(\"\\nUpdating trie...\")\n",
        "        if new_vocab is None and self.vocab_path:\n",
        "            print(f\"Loading from vocab file: {self.vocab_path}\")\n",
        "            with open(self.vocab_path, 'r') as f:\n",
        "                for i, line in enumerate(f):\n",
        "                    token = tuple(json.loads(line)[0])\n",
        "                    self.id_to_tok.append(token)\n",
        "                    self.trie[token] = i\n",
        "        elif new_vocab:\n",
        "            for token in new_vocab:\n",
        "                print(token)\n",
        "                self.id_to_tok.append(token)\n",
        "                self.trie[token] = len(self.trie)\n",
        "\n",
        "    def encode(self, text):\n",
        "        return self._tokenize(text, return_ids=True)\n",
        "\n",
        "    def decode(self, token_ids):\n",
        "      out_bytes = bytearray()\n",
        "      for tid in token_ids:\n",
        "          if 0 <= tid < len(self.id_to_tok):\n",
        "              out_bytes.extend(self.id_to_tok[tid])\n",
        "          # ids outside the range are simply ignored\n",
        "      return out_bytes.decode('utf-8', errors='ignore')\n",
        "\n",
        "\n",
        "    def _tokenize(self, text, return_ids=False):\n",
        "        if isinstance(text, str):\n",
        "            text = text.encode('utf-8', errors='ignore')\n",
        "        tokens = []\n",
        "        while text:\n",
        "            prefix = text[:self.max_token_length]\n",
        "            longest = self.trie.longest_prefix(prefix)\n",
        "            if not longest:\n",
        "                print(f\"Warning: No token found for prefix: {prefix[:10]}...\")\n",
        "                break\n",
        "            tokens.append(longest.value if return_ids else longest.key)\n",
        "            text = text[len(longest.key):]\n",
        "        return tokens"
      ],
      "metadata": {
        "id": "tvUeLwUg5lBh"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import datasets\n",
        "dataset = datasets.load_dataset('coms4705-hewitt/linuxlike-tokenized', 'default', streaming=True)['train']\n",
        "with open('vocab-65k-fw-byte.txt', 'w') as f:\n",
        "    for line in dataset:\n",
        "        f.write(line['text'] + '\\n')\n",
        "tok = Tokenizer(vocab_path='vocab-65k-fw-byte.txt', special_tokens={'eos_token': '<eos>'})\n",
        "print(tok.eos_token, tok.eos_token_id)\n",
        "print(tok.vocab_size)"
      ],
      "metadata": {
        "id": "dwVgkLGcwG3-"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BCqhVlJQzFPZ"
      },
      "source": [
        "# **Single Vector Model**\n",
        "\n",
        "In class you've learned some rudimentary ways in which we can build representations of text. Now, we can use it for next token prediction. Here is a very simplified neural language model, on top of which you will be building your own model. It is equivalent to the model presented in Lecture 1: Text Representation and Language Modeling."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VYnAIRaYzJQy"
      },
      "outputs": [],
      "source": [
        "class SingleVectorWordModel(nn.Module):\n",
        "    def __init__(self, vocab_size, embedding_dim):\n",
        "        \"\"\"\n",
        "        Initialize the model with a single vector per word.\n",
        "\n",
        "        Args:\n",
        "            vocab_size: Size of the vocabulary\n",
        "            embedding_dim: Dimension of word embeddings\n",
        "        \"\"\"\n",
        "        super().__init__()\n",
        "\n",
        "        self.embeddings1 = nn.Embedding(vocab_size, embedding_dim)\n",
        "\n",
        "        # Initialize weights to small values\n",
        "        nn.init.uniform_(self.embeddings1.weight, -0.001, 0.001)\n",
        "\n",
        "    def encode_context(self, context_indices):\n",
        "        emb1 = self.embeddings1(context_indices)  # (bs, seq, d)\n",
        "        return torch.sum(emb1, dim=1) / context_indices.size(1) # (bs, d)\n",
        "\n",
        "    def score(self, h):\n",
        "        return torch.einsum('vd,bd->bv', self.embeddings1.weight, h)\n",
        "\n",
        "    def forward(self, context_indices):\n",
        "        \"\"\"\n",
        "        Forward pass of the model.\n",
        "\n",
        "        Args:\n",
        "            context_indices: Tensor of shape (batch_size, 3) containing indices of context words\n",
        "\n",
        "        Returns:\n",
        "            logits: Tensor of shape (batch_size, vocab_size) containing prediction logits\n",
        "        \"\"\"\n",
        "        h = self.encode_context(context_indices)\n",
        "        logits = self.score(h)\n",
        "        return logits"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UTIJSNWEqBcC"
      },
      "source": [
        "# **Training**\n",
        "\n",
        "Now that we have defined our model, we can start training!\n",
        "**Tasks**\n",
        "1. Question: What does the find_similiar_words method and print_simliar_words method do?\n",
        "2. Implement the function `cross_entropy_loss` using\n",
        "   `F.log_softmax` and tensor indexing (`torch.gather` or advanced indexing).\n",
        "3. Fill in the missing steps in `train_model`:\n",
        "   * move tensors to the correct device\n",
        "   * zero gradients\n",
        "   * forward pass → logits\n",
        "   * compute loss\n",
        "   * backward pass\n",
        "   * optimizer & scheduler steps\n",
        "4. Look at the visualization in Tensorboard and play with the interface.\n",
        "5. Question: What do you notice about the loss of the first epoch across multiple training runs. Does this surprise you? Why or why not?"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "def find_similar_words(model, tokenizer, query_words: List[str], k: int = 5) -> Dict[str, List[str]]:\n",
        "    embeddings_list = {'emb1': model.embeddings1.weight}\n",
        "    if hasattr(model, 'embeddings2'):\n",
        "        embeddings_list['emb2'] = model.embeddings2.weight\n",
        "    for name, embeddings in embeddings_list.items():\n",
        "        embeddings = F.normalize(embeddings, p=2, dim=1)\n",
        "        results = {}\n",
        "        for query in query_words:\n",
        "            query_id = tokenizer.encode(query)[0]\n",
        "            query_embed = embeddings[query_id]\n",
        "            similarities = torch.matmul(embeddings, query_embed)\n",
        "            top_k = torch.topk(similarities, k=k+1)\n",
        "            similar_words = [tokenizer.decode([idx.item()]) for idx in top_k.indices[1:]]\n",
        "            results[query] = similar_words\n",
        "        yield results\n",
        "\n",
        "def print_similar_words(model, tokenizer, query_words: List[str], k: int = 5):\n",
        "    itr = find_similar_words(model, tokenizer, query_words, k)\n",
        "    for results in itr:\n",
        "        print(\"\\nMost similar words:\")\n",
        "        for query, similar in results.items():\n",
        "            print(f\"\\n'{query}':\")\n",
        "            for word in similar:\n",
        "                print(f\"  - {word}\")"
      ],
      "metadata": {
        "id": "KryVEh4MtdMT"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "1. TODO: What does the find_similiar_words method and print_simliar_words method do?\n"
      ],
      "metadata": {
        "id": "iubGSCkwVM0i"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "Answer: response here"
      ],
      "metadata": {
        "id": "tJz-xolhWgBt"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Training Objective\n",
        "\n",
        "For training a language model, the objective is for the model to assign high probability to the correct next word among all possible candidates in the vocabulary. In order to measure how far part the predicted distribution is from the ground-truth, we use cross entropy loss.\n",
        "\n",
        "Cross entropy loss is defined as the negative log-probability of the correct word.\n",
        "\n",
        "Intuitively:\n",
        "- If the model assigns high probability to the correct word, the loss is small.\n",
        "- If the model assigns low probability, the loss is large.\n",
        "\n",
        "Formally, if the model outputs logits of shape $(B, V)$ where\n",
        "- $B = $ batch size,\n",
        "- $V = $ vocabulary size,\n",
        "- $y_i = $ the index of the correct word for example $i$,\n",
        "\n",
        "then the loss is\n",
        "\n",
        "$$\n",
        "\\text{CrossEntropyLoss}(\\text{logits}, \\text{targets})\n",
        "= -\\frac{1}{B} \\sum_{i=1}^{B}\n",
        "\\log \\frac{\\exp\\big(\\text{logits}_{i,\\, y_i}\\big)}\n",
        "{\\sum_{j=1}^{V} \\exp\\big(\\text{logits}_{i,j}\\big)}\n",
        "$$\n",
        "where $j$ runs over all vocabulary items."
      ],
      "metadata": {
        "id": "DYxdASo4zNlj"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#2. TODO: Implement cross entropy loss using F.softmax and torch.gather\n",
        "\n",
        "def cross_entropy_loss(logits, targets):\n",
        "    \"\"\"\n",
        "    Compute cross-entropy loss using log and gather.\n",
        "\n",
        "    Args:\n",
        "        logits: Tensor of shape (batch_size, vocab_size)\n",
        "        targets: Tensor of shape (batch_size,), containing correct word indices\n",
        "\n",
        "    Returns:\n",
        "        loss: Scalar tensor, mean cross-entropy loss over the batch\n",
        "    \"\"\"\n",
        "    # TODO: Implement log softmax over the vocabulary dimension\n",
        "    log_probs = FILL_IN\n",
        "\n",
        "    # TODO: Select the log-probabilities corresponding to the target words\n",
        "    chosen_log_probs = FILL_IN\n",
        "\n",
        "    # TODO: Take the mean negative log likelihood\n",
        "    loss = FILL_IN\n",
        "\n",
        "    return loss"
      ],
      "metadata": {
        "id": "B4M72EQmYNl9"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# 3. TODO: fill in the missing parts of train_model\n",
        "import time\n",
        "import torch\n",
        "import torch.nn.functional as F\n",
        "import torch.optim as optim\n",
        "from torch.utils.tensorboard import SummaryWriter\n",
        "from datetime import datetime\n",
        "\n",
        "def train_model(model, train_loader, learning_rate, tokenizer=None,\n",
        "                max_steps=3000, experiment_name=None):\n",
        "    \"\"\"\n",
        "    Train model with simple max_steps termination.\n",
        "    \"\"\"\n",
        "    #Creates unique experiment name if not provided\n",
        "    if experiment_name is None:\n",
        "        timestamp = datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n",
        "        experiment_name = f\"experiment_{timestamp}\"\n",
        "\n",
        "    #Initializes TensorBoard writer\n",
        "    writer = SummaryWriter(log_dir=f\"runs/{experiment_name}\")\n",
        "\n",
        "    optimizer = optim.Adam(model.parameters(), lr=learning_rate)\n",
        "\n",
        "    scheduler = optim.lr_scheduler.LinearLR(\n",
        "        optimizer, start_factor=1.0, end_factor=0.001, total_iters=max_steps\n",
        "    )\n",
        "\n",
        "    acc_loss = 0.0\n",
        "    acc_count = 0\n",
        "    global_step = 0\n",
        "    best_loss = float('inf')\n",
        "\n",
        "    writer.add_hparams({\n",
        "        'learning_rate': learning_rate,\n",
        "        'max_steps': max_steps,\n",
        "        'batch_size': train_loader.batch_size,\n",
        "        'model_params': sum(p.numel() for p in model.parameters()),\n",
        "    }, {})\n",
        "\n",
        "    t0 = time.time()\n",
        "    model.train()\n",
        "\n",
        "    print(f\"Training for {max_steps:,} steps...\")\n",
        "\n",
        "    train_iter = iter(train_loader)\n",
        "\n",
        "    while global_step < max_steps:\n",
        "        try:\n",
        "        # TODO:\n",
        "            context_indices, target_idx = FILL_IN\n",
        "        except StopIteration:\n",
        "            train_iter = iter(train_loader)\n",
        "            context_indices, target_idx = FILL_IN\n",
        "\n",
        "        context_indices = context_indices.to(model.device)\n",
        "        target_idx = target_idx.to(model.device)\n",
        "\n",
        "        # TODO: forward pass\n",
        "        logits = None\n",
        "\n",
        "        # TODO: compute loss\n",
        "        loss = None\n",
        "\n",
        "        # TODO: backward pass\n",
        "        # (zero gradients, backpropagate, optimizer step, scheduler step)\n",
        "\n",
        "\n",
        "        global_step += 1\n",
        "        acc_loss += loss.item()\n",
        "        acc_count += 1\n",
        "\n",
        "        # Log to TensorBoard every step\n",
        "        writer.add_scalar('Loss/Train_Step', loss.item(), global_step)\n",
        "        writer.add_scalar('Learning_Rate', scheduler.get_last_lr()[0], global_step)\n",
        "\n",
        "        if global_step % 100 == 0:\n",
        "            avg_loss = acc_loss / acc_count\n",
        "            current_lr = scheduler.get_last_lr()[0]\n",
        "            elapsed = time.time() - t0\n",
        "            progress = global_step / max_steps * 100\n",
        "\n",
        "            print(f\"[Step {global_step:,}/{max_steps:,} ({progress:.1f}%)] Loss: {avg_loss:.4f} | LR: {current_lr:.2e} | Time: {elapsed:.1f}s\")\n",
        "\n",
        "            writer.add_scalar('Loss/Train_Avg100', avg_loss, global_step)\n",
        "\n",
        "            if avg_loss < best_loss:\n",
        "                best_loss = avg_loss\n",
        "            writer.add_scalar('Best_Loss', best_loss, global_step)\n",
        "\n",
        "            acc_loss = 0.0\n",
        "            acc_count = 0\n",
        "\n",
        "            if tokenizer is not None and global_step % 500 == 0:\n",
        "                try:\n",
        "                    query_words = [\"linux\", \"python\", \"code\", \"system\", \"file\"]\n",
        "                    print_similar_words(model, tokenizer, query_words, k=3)\n",
        "                except Exception as e:\n",
        "                    print(f\"Similarity check failed: {e}\")\n",
        "\n",
        "        # Save checkpoint occasionally\n",
        "        if global_step % 2000 == 0:\n",
        "            checkpoint_path = f\"model_checkpoint_step_{global_step}.pt\"\n",
        "            torch.save({\n",
        "                'step': global_step,\n",
        "                'model_state_dict': model.state_dict(),\n",
        "                'optimizer_state_dict': optimizer.state_dict(),\n",
        "                'loss': best_loss,\n",
        "            }, checkpoint_path)\n",
        "            print(f\"Saved checkpoint: {checkpoint_path}\")\n",
        "\n",
        "    writer.close()\n",
        "\n",
        "    final_time = time.time() - t0\n",
        "    print(f\"\\nTraining completed!\")\n",
        "    print(f\"Total steps: {global_step:,}\")\n",
        "    print(f\"Total time: {final_time:.1f}s\")\n",
        "    print(f\"Best loss achieved: {best_loss:.4f}\")\n",
        "    print(f\"TensorBoard logs saved to: runs/{experiment_name}\")\n",
        "    print(f\"To view: tensorboard --logdir=runs/{experiment_name}\")\n",
        "\n",
        "    return {\n",
        "        'final_step': global_step,\n",
        "        'final_loss': best_loss,\n",
        "        'total_time': final_time,\n",
        "        'experiment_name': experiment_name\n",
        "    }"
      ],
      "metadata": {
        "id": "AiGLxnACdOgp"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Training Setup"
      ],
      "metadata": {
        "id": "5-j-Rna5VpMj"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "The preprocessing first **tokenizes each document**, optionally truncates it to a maximum length, and appends an **end-of-sequence (EOS)** token to mark document boundaries.  \n",
        "All tokenized documents are **concatenated into a single 1-D array of token IDs**, forming one long text stream.\n",
        "\n",
        "The `PackedLMDataset` then **slides a fixed-length context window** of size `T` across this stream.  \n",
        "For each valid position, it outputs:\n",
        "* **x** – the next `T` tokens (the context)\n",
        "* **y** – the single token immediately following that context (the prediction target)\n",
        "\n",
        "This setup trains the model to predict the next token given a rolling window of prior tokens, even across document boundaries.\n",
        "\n",
        "Preprocessing for 100k documents should be sufficient for training, expect it to take around 6 minutes on colab. No work required here.\n"
      ],
      "metadata": {
        "id": "c7sWOXIDI_uS"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from tqdm import tqdm\n",
        "class PackedLMDataset(Dataset):\n",
        "    \"\"\"\n",
        "    token_ids: 1-D torch.LongTensor of packed tokens (… doc + EOS + doc + EOS …)\n",
        "    context_size: length T of the context window.\n",
        "    Returns (x:[T], y:int)\n",
        "    \"\"\"\n",
        "    def __init__(self, token_ids: torch.Tensor, context_size: int):\n",
        "        assert token_ids.dim() == 1\n",
        "        self.ids = token_ids\n",
        "        self.T = int(context_size)\n",
        "        self.N = max(0, len(self.ids) - self.T)  # valid start positions\n",
        "\n",
        "    def __len__(self):\n",
        "        return self.N\n",
        "\n",
        "    def __getitem__(self, i):\n",
        "        x = self.ids[i : i + self.T]\n",
        "        y = self.ids[i + self.T]\n",
        "        return x.long(), y.long()\n",
        "\n",
        "def pack_rows(rows, tokenizer, max_doc_tokens, eos_id):\n",
        "    \"\"\"\n",
        "    Tokenize each row, truncate to max_doc_tokens, append EOS, and concatenate.\n",
        "    Returns a 1D numpy array of dtype int64.\n",
        "    \"\"\"\n",
        "    pieces = []\n",
        "    append = pieces.append\n",
        "\n",
        "    for r in tqdm(rows, total=len(rows)):\n",
        "        t = r.get(\"text\") if isinstance(r, dict) else r[\"text\"]\n",
        "        if not t:\n",
        "            continue\n",
        "\n",
        "        ids = tokenizer.encode(t)\n",
        "        if not ids:\n",
        "            continue\n",
        "\n",
        "        if max_doc_tokens is not None:\n",
        "            ids = ids[:max_doc_tokens]\n",
        "\n",
        "        append(np.asarray(ids, dtype=np.int64))\n",
        "        append(np.asarray([eos_id], dtype=np.int64))\n",
        "\n",
        "    if not pieces:\n",
        "        return np.zeros((1,), dtype=np.int64)\n",
        "\n",
        "    return np.concatenate(pieces, axis=0)"
      ],
      "metadata": {
        "id": "qg2D-bK62-wZ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Hyperparameters and dataset loading\n",
        "CONTEXT = 10\n",
        "MAX_DOC_TOKENS = 2500\n",
        "VAL_FRACTION = 0.05\n",
        "EOS_ID = len(tok.id_to_tok)\n",
        "VOCAB_SIZE = len(tok.id_to_tok) + 1\n",
        "NUM_DOCS = 100000\n",
        "\n",
        "print(\"Loading dataset...\")\n",
        "ds_full = datasets.load_dataset('coms4705-hewitt/fineweb-linuxlike', 'default', split='train')\n",
        "ds_subset = ds_full.shuffle(seed=123).select(range(NUM_DOCS))\n",
        "splits = ds_subset.train_test_split(test_size=VAL_FRACTION, seed=123)\n",
        "train_rows, val_rows = splits['train'], splits['test']\n",
        "\n",
        "print(\"Packing training data...\")\n",
        "train_ids_np = pack_rows(train_rows, tok, MAX_DOC_TOKENS, EOS_ID)\n",
        "print(\"Packing validation data...\")\n",
        "val_ids_np = pack_rows(val_rows, tok, MAX_DOC_TOKENS, EOS_ID)\n",
        "\n",
        "train_ids = torch.from_numpy(train_ids_np)\n",
        "val_ids = torch.from_numpy(val_ids_np)\n",
        "\n",
        "train_ds = PackedLMDataset(train_ids, context_size=CONTEXT)\n",
        "val_ds = PackedLMDataset(val_ids, context_size=CONTEXT)\n",
        "\n",
        "train_loader = DataLoader(\n",
        "    train_ds, batch_size=512, shuffle=True,\n",
        "    num_workers=0, drop_last=True\n",
        ")\n",
        "val_loader = DataLoader(\n",
        "    val_ds, batch_size=512, shuffle=False,\n",
        "    num_workers=0, drop_last=False\n",
        ")\n",
        "\n",
        "print(f\"Packed tokens: train={len(train_ids):,} val={len(val_ids):,}\")\n",
        "print(f\"Training examples: {len(train_ds):,}\")\n",
        "print(f\"Validation examples: {len(val_ds):,}\")\n",
        "print(f\"Vocab size for the model (include EOS): {VOCAB_SIZE}\")"
      ],
      "metadata": {
        "id": "wAiFcGFBHAIk"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "EMBEDDING_DIM = 100\n",
        "learning_rate = 0.001"
      ],
      "metadata": {
        "id": "Uo551j3TJZ1Q"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# **Time to train your first model!**\n",
        "\n",
        "notes: Please do not change the model names [single_vector_model, positional_model, mlpdeep, residualmlp, gatedresidualmlp]"
      ],
      "metadata": {
        "id": "fFsdJ-XAQkve"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#TRAINING THE MODEL (should approximately take 2 minutes on the T4 GPU for 3000 steps)\n",
        "single_vector_model = SingleVectorWordModel(VOCAB_SIZE, EMBEDDING_DIM)\n",
        "single_vector_model.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "single_vector_model.to(single_vector_model.device)\n",
        "\n",
        "print('Model has {} parameters'.format(sum([np.prod(p.size()) for p in single_vector_model.parameters()])))\n",
        "\n",
        "history = train_model(single_vector_model, train_loader, learning_rate, tokenizer=tok, max_steps=3000, experiment_name = \"SingleVectorWordModel1\")"
      ],
      "metadata": {
        "id": "7Lu_bHuIMilk"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "4. TODO: You can use Tensorboard like this to visualize your training metrics"
      ],
      "metadata": {
        "id": "qz3QFlOYdiZ4"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%load_ext tensorboard\n",
        "%tensorboard --logdir runs"
      ],
      "metadata": {
        "id": "JSj-cAW5Mfpj"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "5. TODO: Try training the model multiple times. What do you notice about the loss of the first step across multiple training runs. Does this surprise you? Why or why not?\n",
        "\n",
        "Answer: response here"
      ],
      "metadata": {
        "id": "4vBwwwBxndto"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# **Improving the model**\n",
        "Let's improve the model by allowing it to capture the positional meaning in a text sequence.\n",
        "$$\n",
        "h_{< i} =\n",
        "\\sigma \\left(\n",
        "    \\frac{1}{\\,i-1\\,}\n",
        "    \\sum_{j=1}^{\\,i-1}\n",
        "        \\sigma\\big(E x_j + p_j\\big)\n",
        "\\right)\n",
        "$$\n",
        "\n",
        "Task: Complete the TODO: sections in the NonLinearPositionalContextModel class"
      ],
      "metadata": {
        "id": "Q7fCeWrj7d-d"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class NonLinearPositionalContextModel(SingleVectorWordModel):\n",
        "    def __init__(self, vocab_size, embedding_dim, max_context):\n",
        "        super().__init__(vocab_size, embedding_dim)\n",
        "        #-------Begin TODO-------\n",
        "\n",
        "        #-------End TODO-----------\n",
        "\n",
        "    def encode_context(self, context_indices):\n",
        "        #--------Begin TODO--------\n",
        "\n",
        "        #--------End TODO----------\n"
      ],
      "metadata": {
        "id": "N1MbwKNC7cxq"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#TRAINING THE MODEL\n",
        "positional_model = NonLinearPositionalContextModel(VOCAB_SIZE, EMBEDDING_DIM, max_context=10)\n",
        "positional_model.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "positional_model.to(positional_model.device)\n",
        "\n",
        "print('Model has {} parameters'.format(sum([np.prod(p.size()) for p in positional_model.parameters()])))\n",
        "\n",
        "history = train_model(positional_model, train_loader, learning_rate, tokenizer=tok, max_steps=3000, experiment_name = \"NonLinearPositionalContextModel1\")"
      ],
      "metadata": {
        "id": "pfnpzVP-pbHx"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### From Positional Model to Context-Level MLP\n",
        "\n",
        "Up to now, the NonLinearPositionalContextModel enriches each token with a\n",
        "positional embedding, applies a non-linearity *independently to every token*,\n",
        "and then **averages** the results before scoring.  \n",
        "This design captures position-specific features inside each token vector, but\n",
        "all cross-token information is lost once the mean is taken.  \n",
        "The final predictor can only work with a *sum of per-token features*.\n",
        "\n",
        "Let's try to remove this averaging bottleneck.  \n",
        "What if we embed each token, concatenate the entire fixed-length context\n",
        "(`L x d` → `L*d`), and feed this single vector through a multi-layer\n",
        "perceptron before scoring.  \n",
        "Because the MLP operates on the *whole* context vector, it can freely mix\n",
        "dimensions from different positions in every layer.\n",
        "\n",
        "#### The motivation for this change\n",
        "* **Broader function class.**  \n",
        "  For a fixed context length, a sufficiently wide MLP can approximate any\n",
        "mapping from the concatenated embeddings to the output, while the positional model was limited to additive “bag-of-tokens” functions.\n",
        "\n",
        "* **Position awareness without extra machinery.**  \n",
        "  The concatenation preserves a consistent slot order\n",
        "(position 1 occupies the first block of `d` features, position 2 the next, etc.).\n",
        "The weight matrices naturally learn different transformations for each slot,\n",
        "so no separate positional embeddings are required.\n",
        "\n",
        "This architecture is still simpler than a Transformer-there is no self-attention\n",
        "or dynamic token-token interaction but it is a meaningful next step in\n",
        "expressivity beyond token-wise pooling.\n",
        "\n",
        "Task: Complete the TODO: sections in the FeedForwardContextModel class\n"
      ],
      "metadata": {
        "id": "rU-G0yF4E3Cf"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "\n",
        "class FeedForwardContextModel(SingleVectorWordModel):\n",
        "    \"\"\"\n",
        "    Context-level MLP (no positional embeddings):\n",
        "      tokens -> embed -> flatten (B, L*d) -> MLP -> score\n",
        "    This model treats the concatenated embeddings of all L tokens\n",
        "    as a single input vector, allowing cross-token interactions.\n",
        "    \"\"\"\n",
        "    def __init__(self,\n",
        "                 vocab_size: int,\n",
        "                 embedding_dim: int,\n",
        "                 ff_hidden_dim: int,\n",
        "                 max_context: int):\n",
        "        super().__init__(vocab_size, embedding_dim)\n",
        "        self.max_context = int(max_context)\n",
        "\n",
        "        # TODO: define MLP: (L*d) -> hidden -> d\n",
        "        self.ff = None\n",
        "\n",
        "        # TODO: initialize MLP weights\n",
        "        # for m in self.ff:\n",
        "        #     if isinstance(m, nn.Linear):\n",
        "        #         ...\n",
        "\n",
        "    def forward(self, context_indices: torch.LongTensor) -> torch.Tensor:\n",
        "        \"\"\"\n",
        "        context_indices: (B, L) with L == max_context\n",
        "        \"\"\"\n",
        "        B, L = context_indices.shape\n",
        "        assert L == self.max_context, f\"Expected context length {self.max_context}, got {L}\"\n",
        "\n",
        "        # TODO: 1) Embed tokens -> (B, L, d)\n",
        "        x = None\n",
        "\n",
        "        # TODO: 2) Flatten entire context -> (B, L*d)\n",
        "        x_flat = None\n",
        "\n",
        "        # TODO: 3) Context-level MLP -> (B, d)\n",
        "        h = None\n",
        "\n",
        "        # TODO: 4) Score with tied vocabulary embeddings -> (B, V)\n",
        "        logits = None\n",
        "\n",
        "        return logits\n"
      ],
      "metadata": {
        "id": "SFZbuLOcE2ay"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "HIDDEN_DIM = 1028\n",
        "mlp = FeedForwardContextModel(VOCAB_SIZE, EMBEDDING_DIM, HIDDEN_DIM, 10)\n",
        "mlp.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "mlp.to(mlp.device)\n",
        "\n",
        "print('Model has {} parameters'.format(sum([np.prod(p.size()) for p in mlp.parameters()])))\n",
        "\n",
        "history = train_model(mlp, train_loader, learning_rate, tokenizer=tok, max_steps=10000, experiment_name = \"FeedForwardContextModel1\")\n"
      ],
      "metadata": {
        "id": "68wPZUBZsRI9"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Going Deeper\n",
        "\n",
        "The current context level mlp applies a **single feed-forward block**\n",
        "to the flattened `L*d` context vector before prediction.  \n",
        "A natural next question is: *how does the model’s **expressive power** change if\n",
        "we stack multiple such blocks?*\n",
        "\n",
        "In function-approximation terms, a one-layer MLP represents functions of the form\n",
        "\n",
        "$$\n",
        "f(x) = \\sigma(Wx + b)\n",
        "$$\n",
        "\n",
        "where \\(\\sigma\\) is a fixed non-linearity.  \n",
        "Adding more layers composes these transformations:\n",
        "\n",
        "$$\n",
        "f(x) = W_k \\,\\sigma\\big(W_{k-1}\\,\\sigma(\\dots \\sigma(W_1 x)\\dots )\\big)\n",
        "$$\n",
        "\n",
        "Classical results in approximation theory show that **compositions of simple\n",
        "nonlinearities can approximate a far richer class of functions** than a single\n",
        "layer of the same width.  \n",
        "Unlike the positional model, the input here is already the full context vector,\n",
        "so deeper stacks directly expand the space of context-level mappings the model\n",
        "can learn.\n",
        "\n",
        "\n",
        "#### Task\n",
        "1. Add a `num_layers` parameter so the feed-forward block can be repeated\n",
        "   multiple times.  \n",
        "   Begin with `num_layers = 1` (the current model) and experiment with `2`, `4`,\n",
        "   and `8`.\n",
        "2. Compute how the parameter count scales with depth.\n",
        "3. Train each version under identical hyper-parameters and compare:\n",
        "   * training loss and validation perplexity\n",
        "   * gradient norms in each layer (use `print_grad_stats`) to monitor stability.\n",
        "\n",
        "\n",
        "\n",
        "#### Things to Observe\n",
        "* **Expressivity vs. optimization** – deeper networks can capture more complex\n",
        "context functions but may be harder to train.\n",
        "* **Gradient flow** – check for vanishing or exploding gradients as depth grows.\n",
        "* **Parameter efficiency** – assess whether additional layers provide measurable\n",
        "gains relative to their extra parameters and compute.\n"
      ],
      "metadata": {
        "id": "-LRW7GNHJAYm"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def print_grad_stats(model):\n",
        "    \"\"\"Print L2 norms of gradients for each layer.\"\"\"\n",
        "    for name, p in model.named_parameters():\n",
        "        if p.grad is not None:\n",
        "            print(f\"{name:<30s} grad_norm={p.grad.norm().item():.6f}\")"
      ],
      "metadata": {
        "id": "QvhmaCFu67fm"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class _ContextMLPBlock(nn.Module):\n",
        "    \"\"\"\n",
        "    Minimal width-preserving block operating on a d-dim vector:\n",
        "      d -> hidden -> d (with an activation in between).\n",
        "    No residuals, no LayerNorm—keeps the step small.\n",
        "    \"\"\"\n",
        "    def __init__(self, d: int, hidden: int, activation=nn.ReLU):\n",
        "        super().__init__()\n",
        "        # TODO: define first linear layer (d -> hidden)\n",
        "        self.fc1 = None\n",
        "        # TODO: activation function\n",
        "        self.act = None\n",
        "        # TODO: define second linear layer (hidden -> d)\n",
        "        self.fc2 = None\n",
        "        # TODO: initialize weights (e.g., xavier) and biases to zero\n",
        "\n",
        "    def forward(self, x: torch.Tensor) -> torch.Tensor:\n",
        "        # TODO: forward pass through fc1 -> activation -> fc2\n",
        "        return None\n",
        "\n",
        "\n",
        "class FeedForwardContextModelDeep(SingleVectorWordModel):\n",
        "    \"\"\"\n",
        "    Context-level MLP (no positional embeddings):\n",
        "      tokens -> embed -> flatten (B, L*d) -> stem -> [block] * num_layers -> score\n",
        "\n",
        "    - With num_layers=1 this matches your original two-layer ff.\n",
        "    - Increasing num_layers deepens only the d→hidden→d mapping.\n",
        "    \"\"\"\n",
        "    def __init__(self,\n",
        "                 vocab_size: int,\n",
        "                 embedding_dim: int,   # d\n",
        "                 ff_hidden_dim: int,   # hidden width inside each block\n",
        "                 max_context: int,     # L\n",
        "                 num_layers: int = 1,\n",
        "                 activation=nn.ReLU):\n",
        "        super().__init__(vocab_size, embedding_dim)\n",
        "        self.max_context = int(max_context)\n",
        "        self.num_layers = int(num_layers)\n",
        "        d = embedding_dim\n",
        "        L = self.max_context\n",
        "\n",
        "        # TODO: Stem layer: (L*d) -> d\n",
        "        self.stem_fc = None\n",
        "        self.stem_act = None\n",
        "        # TODO: initialize stem weights and biases\n",
        "\n",
        "        # TODO: create ModuleList of num_layers ContextMLPBlocks\n",
        "        self.blocks = None\n",
        "\n",
        "    def forward(self, context_indices: torch.LongTensor) -> torch.Tensor:\n",
        "        \"\"\"\n",
        "        context_indices: (B, L) with L == max_context\n",
        "        \"\"\"\n",
        "        B, L = context_indices.shape\n",
        "        assert L == self.max_context, f\"Expected context length {self.max_context}, got {L}\"\n",
        "\n",
        "        # TODO: 1) Embed tokens -> (B, L, d)\n",
        "        x = None\n",
        "\n",
        "        # TODO: 2) Flatten entire context -> (B, L*d)\n",
        "        x_flat = None\n",
        "\n",
        "        # TODO: 3) Stem to d -> (B, d)\n",
        "        h = None\n",
        "\n",
        "        # TODO: 4) Depth: apply each block (d -> hidden -> d)\n",
        "        # for blk in self.blocks:\n",
        "        #     h = blk(h)\n",
        "\n",
        "        # TODO: 5) Tied scoring -> (B, V)\n",
        "        logits = None\n",
        "\n",
        "        return logits\n"
      ],
      "metadata": {
        "id": "lI_I8hHC1wqZ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#approximately 10 mins to train for 4 layers of hidden_dim = 1028, 15k steps.\n",
        "HIDDEN_DIM = 1028\n",
        "mlpdeep = FeedForwardContextModelDeep(VOCAB_SIZE, EMBEDDING_DIM, HIDDEN_DIM, max_context=10, num_layers = 4)\n",
        "mlpdeep.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "mlpdeep.to(mlpdeep.device)\n",
        "\n",
        "print('Model has {} parameters'.format(sum([np.prod(p.size()) for p in mlpdeep.parameters()])))\n",
        "\n",
        "history = train_model(mlpdeep, train_loader, learning_rate, tokenizer=tok, max_steps=15000, experiment_name = \"FeedForwardContextModelDeep1\")\n"
      ],
      "metadata": {
        "id": "GMzSnsO13SBu"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "print_grad_stats(mlpdeep)"
      ],
      "metadata": {
        "id": "6TrZXMrVe-7M"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Residual Connections\n",
        "\n",
        "When we stack many fully-connected layers, training can slow down or become\n",
        "unstable because gradients shrink or grow as they pass through each layer.\n",
        "A **residual connection** helps by adding the input of a block directly to its\n",
        "output:\n",
        "\n",
        "$$\n",
        "\\text{Block}(x) = x + f(x)\n",
        "$$\n",
        "\n",
        "where \\(f(x)\\) is the small network inside the block.\n",
        "\n",
        "**Why this helps**\n",
        "* The gradient can flow along the simple identity path \\(x \\rightarrow x\\),\n",
        "  so learning is less sensitive to depth or initialization.\n",
        "* Each block only needs to learn a *correction* \\(f(x)\\) on top of the input,\n",
        "  instead of the entire mapping from scratch.\n",
        "\n",
        "#### Tasks\n",
        "\n",
        "1. **Implement a `ResidualBlock`**\n",
        "   * Input and output both have size `d`.\n",
        "   * Inside the block: `Linear(d→h)` → activation (e.g., ReLU) → `Linear(h→d)`\n",
        "     to compute \\(f(x)\\).\n",
        "   * Return `x + f(x)`.\n",
        "\n",
        "2. **Build a `ResidualFeedForwardModel`**\n",
        "   * Start from your context-level MLP pipeline\n",
        "     (`embed → flatten → stem` to dimension `d`).\n",
        "   * Stack several `ResidualBlock`s in a `ModuleList`.\n",
        "   * After the stack, use the same `score()` method to produce logits.\n",
        "\n",
        "3. **Experiment**\n",
        "   * Train deep models with and without residuals using the same settings.\n",
        "   * Plot training and validation loss to compare convergence speed.\n",
        "   * Use `print_grad_stats` to inspect gradient norms across layers.\n"
      ],
      "metadata": {
        "id": "n1Z7YAOHffhj"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class ResidualBlock(nn.Module):\n",
        "    \"\"\"\n",
        "    Width-preserving residual block on a d-dim vector:\n",
        "      x -> Linear(d->hidden) -> Act -> Linear(hidden->d) -> + x\n",
        "    \"\"\"\n",
        "    def __init__(self, d: int, hidden: int, activation=nn.ReLU, use_layernorm: bool = False):\n",
        "        super().__init__()\n",
        "        self.use_layernorm = use_layernorm\n",
        "        # TODO: optional LayerNorm or Identity\n",
        "        self.ln = None\n",
        "\n",
        "        # TODO: first linear layer (d -> hidden)\n",
        "        self.fc1 = None\n",
        "        # TODO: activation function\n",
        "        self.act = None\n",
        "        # TODO: second linear layer (hidden -> d)\n",
        "        self.fc2 = None\n",
        "        # TODO: initialize weights (e.g., xavier) and zero biases\n",
        "\n",
        "    def forward(self, x: torch.Tensor) -> torch.Tensor:\n",
        "        # TODO: optional pre-norm, forward through fc1 -> activation -> fc2, add residual\n",
        "        return None\n",
        "\n",
        "\n",
        "class ResidualFeedForwardContextModel(SingleVectorWordModel):\n",
        "    \"\"\"\n",
        "    Context-level residual MLP:\n",
        "      tokens -> embed -> flatten (B, L*d) -> stem (L*d -> d)\n",
        "              -> [ResidualBlock(d, hidden)] * num_layers\n",
        "              -> score (tied embeddings)\n",
        "    \"\"\"\n",
        "    def __init__(self,\n",
        "                 vocab_size: int,\n",
        "                 embedding_dim: int,   # d\n",
        "                 max_context: int,     # L\n",
        "                 hidden_dim: int,      # hidden width inside each residual block\n",
        "                 num_layers: int = 2,  # number of residual blocks\n",
        "                 activation=nn.ReLU,\n",
        "                 use_layernorm: bool = False):\n",
        "        super().__init__(vocab_size, embedding_dim)\n",
        "        self.max_context = int(max_context)\n",
        "        d = embedding_dim\n",
        "        L = self.max_context\n",
        "\n",
        "        # TODO: Stem (L*d -> d) with activation\n",
        "        self.stem = None\n",
        "        # TODO: initialize stem weights/biases\n",
        "\n",
        "        # TODO: stack of residual blocks\n",
        "        self.blocks = None\n",
        "\n",
        "    def forward(self, context_indices: torch.LongTensor) -> torch.Tensor:\n",
        "        \"\"\"\n",
        "        context_indices: (B, L) with L == max_context\n",
        "        \"\"\"\n",
        "        B, L = context_indices.shape\n",
        "        assert L == self.max_context, f\"Expected context length {self.max_context}, got {L}\"\n",
        "\n",
        "        # TODO: 1) Embed tokens -> (B, L, d)\n",
        "        x = None\n",
        "\n",
        "        # TODO: 2) Flatten entire context -> (B, L*d)\n",
        "        x_flat = None\n",
        "\n",
        "        # TODO: 3) Stem to d -> (B, d)\n",
        "        h = None\n",
        "\n",
        "        # TODO: 4) Apply residual blocks\n",
        "        # for blk in self.blocks:\n",
        "        #     h = blk(h)\n",
        "\n",
        "        # TODO: 5) Tied scoring -> (B, V)\n",
        "        logits = None\n",
        "\n",
        "        return logits\n"
      ],
      "metadata": {
        "id": "eysVp08SNmtI"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "residualmlp = ResidualFeedForwardContextModel(VOCAB_SIZE, EMBEDDING_DIM, hidden_dim = HIDDEN_DIM, max_context=10, num_layers=4)\n",
        "residualmlp.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "residualmlp.to(residualmlp.device)\n",
        "\n",
        "print('Model has {} parameters'.format(sum([np.prod(p.size()) for p in residualmlp.parameters()])))\n",
        "\n",
        "history = train_model(residualmlp, train_loader, learning_rate, tokenizer=tok, max_steps=15000, experiment_name = \"ResidualFeedForwardContextModel1\")\n"
      ],
      "metadata": {
        "id": "Ykh5g9uZxJfr"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Gating\n",
        "\n",
        "Gating is a learned, elementwise mask that modulates a signal.  \n",
        "Given a representation \\(x\\), a gate \\(g(x)\\) produces per-dimension weights in \\([0,1]\\), and the modulated output is\n",
        "$$\n",
        "y \\;=\\; \\sigma(g(x)) \\odot h(x)\n",
        "$$\n",
        "where \\(h(x)\\) is any transformation of \\(x\\), \\(\\sigma\\) is a sigmoid, and \\(\\odot\\) is elementwise multiply.\n",
        "\n",
        "**Why use it**\n",
        "- **Selective flow:** the model can amplify useful features and suppress noisy ones.\n",
        "- **Dynamic sparsity:** encourages the network to use only the channels it needs for a given input."
      ],
      "metadata": {
        "id": "-7bo3eCZSRH_"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class GatedResidualBlock(nn.Module):\n",
        "    \"\"\"\n",
        "    Width-preserving gated residual block on a d-dim vector:\n",
        "      y = x + sigmoid(g(x)) * f(x)\n",
        "      where f(x): d -> h -> d, g(x): d -> d\n",
        "    \"\"\"\n",
        "    def __init__(self, d: int, h: int, activation=nn.ReLU, use_layernorm: bool=False):\n",
        "        super().__init__()\n",
        "        # TODO: optional LayerNorm or Identity\n",
        "        self.ln = None\n",
        "\n",
        "        # TODO: define f1 (d -> h), activation, f2 (h -> d)\n",
        "        self.f1 = None\n",
        "        self.act = None\n",
        "        self.f2 = None\n",
        "\n",
        "        # TODO: define gating layer g (d -> d)\n",
        "        self.g = None\n",
        "\n",
        "        # TODO: initialize all weights/biases (e.g., xavier) and set gate bias to start closed\n",
        "\n",
        "    def forward(self, x: torch.Tensor) -> torch.Tensor:\n",
        "        # TODO: apply layernorm (if used), compute f(x), compute gate(x),\n",
        "        #       and return x + sigmoid(gate) * f(x)\n",
        "        return None\n",
        "\n",
        "\n",
        "class GatedResidualFeedForwardContextModel(SingleVectorWordModel):\n",
        "    \"\"\"\n",
        "    Context-level gated residual MLP (no positional embeddings):\n",
        "      tokens -> embed -> flatten (B, L*d) -> stem (L*d -> d)\n",
        "             -> [GatedResidualBlock(d,h)] * num_layers\n",
        "             -> score()  (tied embeddings from SingleVectorWordModel)\n",
        "    \"\"\"\n",
        "    def __init__(self,\n",
        "                 vocab_size: int,\n",
        "                 embedding_dim: int,   # d\n",
        "                 max_context: int,     # L\n",
        "                 hidden_dim: int,      # h inside each block\n",
        "                 num_layers: int = 2,\n",
        "                 activation=nn.ReLU,\n",
        "                 use_layernorm: bool=False):\n",
        "        super().__init__(vocab_size, embedding_dim)\n",
        "        self.max_context = int(max_context)\n",
        "        d, L = embedding_dim, self.max_context\n",
        "\n",
        "        # TODO: Stem layer (L*d -> d) with activation\n",
        "        self.stem = None\n",
        "        # TODO: initialize stem weights/biases\n",
        "\n",
        "        # TODO: stack of gated residual blocks\n",
        "        self.blocks = None\n",
        "\n",
        "    def forward(self, context_indices: torch.LongTensor) -> torch.Tensor:\n",
        "        \"\"\"\n",
        "        context_indices: (B, L) with L == max_context\n",
        "        \"\"\"\n",
        "        B, L = context_indices.shape\n",
        "        assert L == self.max_context, f\"Expected context length {self.max_context}, got {L}\"\n",
        "\n",
        "        # TODO: 1) Embed tokens -> (B, L, d)\n",
        "        x = None\n",
        "\n",
        "        # TODO: 2) Flatten -> (B, L*d)\n",
        "        x_flat = None\n",
        "\n",
        "        # TODO: 3) Stem -> (B, d)\n",
        "        h = None\n",
        "\n",
        "        # TODO: 4) Apply each gated residual block\n",
        "        # for blk in self.blocks:\n",
        "        #     h = blk(h)\n",
        "\n",
        "        # TODO: 5) Tied scoring -> (B, V)\n",
        "        logits = None\n",
        "\n",
        "        return logits\n"
      ],
      "metadata": {
        "id": "9usvzp_rRNwz"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "gatedresidualmlp = GatedResidualFeedForwardContextModel(VOCAB_SIZE, EMBEDDING_DIM, hidden_dim = HIDDEN_DIM, num_layers=4, max_context=10)\n",
        "gatedresidualmlp.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "gatedresidualmlp.to(gatedresidualmlp.device)\n",
        "\n",
        "print('Model has {} parameters'.format(sum([np.prod(p.size()) for p in gatedresidualmlp.parameters()])))\n",
        "\n",
        "history = train_model(gatedresidualmlp, train_loader, learning_rate, tokenizer=tok, max_steps=15000, experiment_name = \"GatedResidualFeedForwardContextModel1\")"
      ],
      "metadata": {
        "id": "0W79DHdeaKws"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# **Take a look at how the models compare!**"
      ],
      "metadata": {
        "id": "iv1cV1qDz5h3"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%load_ext tensorboard\n",
        "%tensorboard --logdir runs"
      ],
      "metadata": {
        "id": "PoccIOW7qcHL"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Discuss the results that you gathered throughout this assignment in two short paragraphs. You might consider topics such as:\n",
        "\n",
        "1. How the different model architectures performed relative to each other.\n",
        "\n",
        "2. Trends you observed in training (loss curves, convergence speed, overfitting, etc.).\n",
        "\n",
        "3. The impact of hyper-parameters.\n",
        "\n",
        "4. Any unexpected behaviors or challenges you encountered.\n",
        "\n",
        "Support your discussion with specific observations or metrics where possible."
      ],
      "metadata": {
        "id": "mPjb57AR0KFy"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Response here:\n"
      ],
      "metadata": {
        "id": "ppMpkRf50wqy"
      }
    }
  ],
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4"
    },
    "kernelspec": {
      "display_name": "Python 3",
      "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.9.7"
    },
    "accelerator": "GPU"
  },
  "nbformat": 4,
  "nbformat_minor": 0
}