transformer_decoder.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Transformer Decoder\n",
    "![Attention Transformer Block](./images/attention_transformer_block.png)\n",
    "\n",
    "We will build an autoregressive transformer decoder language model for predicting the next token in a sequence of text. The shakespeare_char dataset will be used for this demonstration, which can be found in the data folder.  \n",
    "\n",
    "We will be implementing the transformer decoder block, which is the right block in the diagram above. The implementation won't be exactly the same as in the diagram, these are the changes we will be making:\n",
    "1. The norm layer will now come before the transformation layers. This improves the models performance.\n",
    "2. We won't be implementing the encoder block, we don't need it for this type of problem. The encoder block is used for tasks like text translation.\n",
    "3. We wont be using the cross attention module (second multi-head attention module in decoder block). We can't do cross attention because we dont have an encoder block as stated in 2.\n",
    "\n",
    "The transformer decoder combines the communication mechanism of the multi-head attention in [3] with the computation mechanism of the feed forward multi-layer perceptron. It also includes other innovations in deep learning, like skip connections [4], and layer normalization [5]. The skip/residual connections allow us to create very deep neural networks using many transformer decoder blocks. This enabled the creation of the large language models used today. Layer norm normalizes the data distribution, this enables faster training since the trained layers now essentially don't have a moving target to deal with. For further explanation on the transformer decoder see Andrej Karpathy's video [2].\n",
    "\n",
    "\n",
    "### References:\n",
    "- [1] [GPT colab notebook](https://colab.research.google.com/drive/1JMLa53HDuA-i7ZBmqV7ZnA3c_fvtXnx-?usp=sharing)\n",
    "- [2] [Video: feedforward layers of transformer block](https://www.youtube.com/watch?v=kCc8FmEb1nY&t=5065s)\n",
    "- [3] [Attention Is All You Need](https://arxiv.org/abs/1706.03762)\n",
    "- [4] [Deep Residual Learning for Image Recognition](https://arxiv.org/abs/1512.03385)\n",
    "- [5] [Layer Normalization](https://arxiv.org/abs/1607.06450)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import requests\n",
    "import numpy as np\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import flax.linen as nn\n",
    "import optax\n",
    "from jax import value_and_grad\n",
    "\n",
    "from helper_funcs import get_batch, generate, masked_fill, loss_fn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_embed = 32 # Number of embedding dimensions\n",
    "batch_size = 32 # How many independent sequences will we process in parallel?\n",
    "block_size = 8 # What is the maximum context length for predictions?\n",
    "num_heads = 4 # Number of heads in the multi-headed block\n",
    "\n",
    "rng_key = jax.random.PRNGKey(128)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data Preparation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "length of dataset in characters: 1,115,394\n",
      "all the unique characters: \n",
      " !$&',-.3:;?ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz\n",
      "vocab size: 65\n",
      "train has 1,003,854 tokens\n",
      "val has 111,540 tokens\n"
     ]
    }
   ],
   "source": [
    "# download the tiny shakespeare dataset\n",
    "input_file_path = os.path.join('./data/shakespeare_char/input.txt')\n",
    "if not os.path.exists(input_file_path):\n",
    "    data_url = 'https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt'\n",
    "    with open(input_file_path, 'w') as f:\n",
    "        f.write(requests.get(data_url).text)\n",
    "\n",
    "with open(input_file_path, 'r') as f:\n",
    "    data = f.read()\n",
    "print(f\"length of dataset in characters: {len(data):,}\")\n",
    "\n",
    "# get all the unique characters that occur in this text\n",
    "chars = sorted(list(set(data)))\n",
    "vocab_size = len(chars)\n",
    "print(\"all the unique characters:\", ''.join(chars))\n",
    "print(f\"vocab size: {vocab_size:,}\")\n",
    "\n",
    "# create a mapping from characters to integers\n",
    "stoi = { ch:i for i,ch in enumerate(chars) }\n",
    "itos = { i:ch for i,ch in enumerate(chars) }\n",
    "def encode(s):\n",
    "    return [stoi[c] for c in s] # encoder: take a string, output a list of integers\n",
    "def decode(l):\n",
    "    l = np.array(l)\n",
    "    return ''.join([itos[i] for i in l]) # decoder: take a list of integers, output a string\n",
    "\n",
    "# create the train and test splits\n",
    "n = len(data)\n",
    "train_data = data[:int(n*0.9)]\n",
    "val_data = data[int(n*0.9):]\n",
    "\n",
    "# encode both to integers\n",
    "train_ids = encode(train_data)\n",
    "val_ids = encode(val_data)\n",
    "print(f\"train has {len(train_ids):,} tokens\")\n",
    "print(f\"val has {len(val_ids):,} tokens\")\n",
    "\n",
    "train_ids = jnp.array(train_ids, dtype=jnp.uint16)\n",
    "val_ids = jnp.array(val_ids, dtype=jnp.uint16)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "First Citizen:\n",
      "Before we proceed any further, hear me speak.\n",
      "\n",
      "All:\n",
      "Speak, speak.\n",
      "\n",
      "First Citizen:\n",
      "You\n"
     ]
    }
   ],
   "source": [
    "print(decode(train_ids[:100]))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Build the Attention Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class FeedForward(nn.Module):\n",
    "    \"\"\"\n",
    "    A feed forward multi-layer perceptron network.\n",
    "    \"\"\"\n",
    "    n_embed: int\n",
    "\n",
    "    @nn.compact\n",
    "    def __call__(self, x):\n",
    "        net = nn.Sequential([\n",
    "            nn.Dense(4 * self.n_embed),\n",
    "            jax.nn.relu,\n",
    "            nn.Dense(self.n_embed)\n",
    "        ])\n",
    "        x = net(x)\n",
    "\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Head(nn.Module):\n",
    "    \"\"\"\n",
    "    A single-headed self-attention decoder block.\n",
    "    Takes the combined token and position embedding as input,\n",
    "    then calculates the key and query values.\n",
    "    The key and query are multiplied to calculate the \n",
    "    attention scores/affinities. The future weights are\n",
    "    then altered to have zero affinity, this ensures the \n",
    "    model can't \"cheat\". The input is then used to calculate\n",
    "    the values, which are then aggregated by multiplying \n",
    "    them with the weights.\n",
    "    \"\"\"\n",
    "    head_size: int\n",
    "\n",
    "    @nn.compact\n",
    "    def __call__(self, x):\n",
    "        B,T,C = x.shape\n",
    "        key = nn.Dense(self.head_size, use_bias=False)\n",
    "        k = key(x) # (B,T,C)\n",
    "        query = nn.Dense(self.head_size, use_bias=False)\n",
    "        q = query(x) # (B,T,C)\n",
    "        # compute attention scores (\"affinities\")\n",
    "        weights =  q @ k.transpose((0, -1, -2)) * self.head_size**-0.5 # (B, T, C) @ (B, C, T) ---> (B, T, T)\n",
    "        tril = jnp.tril(jnp.ones(shape=(T, T), dtype=bool))\n",
    "        tril = jnp.repeat(tril[None, ...], repeats=B, axis=0)\n",
    "        weights = masked_fill(tril, weights, -jnp.inf)\n",
    "        weights = jax.nn.softmax(weights, axis=-1)\n",
    "        # perform the weighted aggregation of the values\n",
    "        value = nn.Dense(self.head_size, use_bias=False)\n",
    "        v = value(x)\n",
    "        out = weights @ v\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultiHeadedAttention(nn.Module):\n",
    "    \"\"\"\n",
    "    Combines multiple heads of scaled self-attention \n",
    "    in parallel, then concatenates the heads outputs.\n",
    "    \"\"\"\n",
    "    num_heads: int\n",
    "    head_size: int\n",
    "    n_embed: int\n",
    "\n",
    "    @nn.compact\n",
    "    def __call__(self, x):\n",
    "        # Create a list of num_heads heads\n",
    "        heads = [Head(self.head_size) for _ in range(self.num_heads)]\n",
    "        # Provide the same input for each head\n",
    "        heads_out = [h(x) for h in heads]\n",
    "        combined_logits = jnp.concatenate(heads_out, axis=-1)\n",
    "        # Perform a linear projection of the self-attention\n",
    "        proj = nn.Dense(self.n_embed)\n",
    "        logits = proj(combined_logits)\n",
    "        return logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Block(nn.Module):\n",
    "    \"\"\"\n",
    "    Transformer decoder block.\n",
    "    It combines communication and computation.\n",
    "    The communication is performed by the \n",
    "    multi-headed attention layer.\n",
    "    Then the computation is performed by \n",
    "    the feed forward block.\n",
    "    Skip connections are used to make the block scalable \n",
    "    and layer norm is used to speed up training.\n",
    "    \"\"\"\n",
    "    n_embed: int\n",
    "    num_heads: int\n",
    "\n",
    "    @nn.compact\n",
    "    def __call__(self, x):\n",
    "        head_size = self.n_embed // self.num_heads\n",
    "        sa_heads = MultiHeadedAttention(self.num_heads, head_size, self.n_embed)\n",
    "        # Using skip connections with x + heads\n",
    "        x = x + sa_heads(nn.LayerNorm()(x)) # apply one head of self-attention (B, T, C)\n",
    "        ffwd = FeedForward(self.n_embed)\n",
    "        x = x + ffwd(nn.LayerNorm()(x))\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "class AttentionLanguageModel(nn.Module):\n",
    "    \"\"\"\n",
    "    Attention decoder language model.\n",
    "    Uses the previous tokens in the sequence to \n",
    "    determine the probabilities of the next token.\n",
    "    Processes the combined position and token embedding\n",
    "    through multiple transformer decoder blocks, \n",
    "    which is then processed through a dense layer to \n",
    "    aquire the token logits.\n",
    "    The logits can then be processed through a softmax\n",
    "    function to calculate the token probabilities.\n",
    "    \"\"\"\n",
    "    vocab_size: int\n",
    "    n_embed: int\n",
    "    block_size: int\n",
    "    num_heads: int\n",
    "    \n",
    "    @nn.compact\n",
    "    def __call__(self, index_seq):\n",
    "        B, T = index_seq.shape\n",
    "\n",
    "        # Each token directly reads off the logits for the next token from a lookup table\n",
    "        token_embedding_table = nn.Embed(num_embeddings=self.vocab_size, features=self.n_embed) \n",
    "        token_emb = token_embedding_table(index_seq) # (B, T, C)\n",
    "\n",
    "        position_embedding_table = nn.Embed(num_embeddings=self.block_size, features=self.n_embed) \n",
    "        pos_emb = position_embedding_table(jnp.arange(T)) # (T, C)\n",
    "\n",
    "        x = token_emb + pos_emb # (B, T, C)\n",
    "\n",
    "        blocks = nn.Sequential([\n",
    "            Block(self.n_embed, num_heads=self.num_heads),\n",
    "            Block(self.n_embed, num_heads=self.num_heads),\n",
    "            Block(self.n_embed, num_heads=self.num_heads),\n",
    "            nn.LayerNorm()\n",
    "        ])\n",
    "        x = blocks(x)\n",
    "\n",
    "        lm_head = nn.Dense(self.vocab_size)\n",
    "        logits = lm_head(x) # (B, T, vocab_size)\n",
    "\n",
    "        return logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = AttentionLanguageModel(vocab_size, n_embed, block_size, num_heads)\n",
    "dummy_x = jnp.zeros(shape=(batch_size, block_size), dtype=jnp.uint16)\n",
    "variables = model.init(rng_key, dummy_x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(32, 8, 65)\n"
     ]
    }
   ],
   "source": [
    "out = model.apply(variables, dummy_x)\n",
    "print(out.shape)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Text Generation Pre-Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generated text: \n",
      "\n",
      "GdPkiWte;-gpPzhvR zA;zUEhVkElPvzSbWks3vzeg$Sdg3k:s\n"
     ]
    }
   ],
   "source": [
    "index_seq = jnp.zeros(shape=(1,1), dtype=jnp.uint16)\n",
    "max_new_tokens = 50\n",
    "\n",
    "generated_indices = generate(variables, model.apply, index_seq, rng_key, vocab_size, 1, block_size, max_new_tokens)\n",
    "generated_indices = list(np.array(generated_indices[0]))\n",
    "print(\"Generated text: \")\n",
    "print(decode(generated_indices))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train the Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = optax.adamw(learning_rate=1e-2)\n",
    "opt_state = optimizer.init(variables)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 0, Loss: 4.6198\n",
      "Epoch: 1, Loss: 4.2264\n",
      "Epoch: 2, Loss: 3.6980\n",
      "Epoch: 3, Loss: 3.5706\n",
      "Epoch: 4, Loss: 3.6013\n",
      "Epoch: 5, Loss: 3.3840\n",
      "Epoch: 6, Loss: 3.3821\n",
      "Epoch: 7, Loss: 3.2056\n",
      "Epoch: 8, Loss: 3.2752\n",
      "Epoch: 9, Loss: 3.3728\n",
      "Epoch: 10, Loss: 3.3380\n",
      "Epoch: 11, Loss: 3.4187\n",
      "Epoch: 12, Loss: 3.1114\n",
      "Epoch: 13, Loss: 3.1586\n",
      "Epoch: 14, Loss: 3.6710\n",
      "Epoch: 15, Loss: 3.5170\n",
      "Epoch: 16, Loss: 3.3354\n",
      "Epoch: 17, Loss: 3.2102\n",
      "Epoch: 18, Loss: 3.1993\n",
      "Epoch: 19, Loss: 3.3897\n",
      "Epoch: 20, Loss: 3.2154\n",
      "Epoch: 21, Loss: 3.1864\n",
      "Epoch: 22, Loss: 3.1117\n",
      "Epoch: 23, Loss: 3.1442\n",
      "Epoch: 24, Loss: 3.0538\n",
      "Epoch: 25, Loss: 3.1941\n",
      "Epoch: 26, Loss: 3.2109\n",
      "Epoch: 27, Loss: 3.0869\n",
      "Epoch: 28, Loss: 3.1904\n",
      "Epoch: 29, Loss: 3.1063\n",
      "Epoch: 30, Loss: 3.1947\n",
      "Epoch: 31, Loss: 3.2487\n",
      "Epoch: 32, Loss: 3.1228\n",
      "Epoch: 33, Loss: 3.0457\n",
      "Epoch: 34, Loss: 3.0480\n",
      "Epoch: 35, Loss: 3.0954\n",
      "Epoch: 36, Loss: 2.9428\n",
      "Epoch: 37, Loss: 2.9732\n",
      "Epoch: 38, Loss: 2.9597\n",
      "Epoch: 39, Loss: 3.0434\n",
      "Epoch: 40, Loss: 2.8481\n",
      "Epoch: 41, Loss: 2.9712\n",
      "Epoch: 42, Loss: 3.0033\n",
      "Epoch: 43, Loss: 2.9335\n",
      "Epoch: 44, Loss: 2.9867\n",
      "Epoch: 45, Loss: 3.3182\n",
      "Epoch: 46, Loss: 3.1432\n",
      "Epoch: 47, Loss: 2.8075\n",
      "Epoch: 48, Loss: 2.8615\n",
      "Epoch: 49, Loss: 3.1502\n",
      "Epoch: 50, Loss: 2.8608\n",
      "Epoch: 51, Loss: 2.9046\n",
      "Epoch: 52, Loss: 2.8904\n",
      "Epoch: 53, Loss: 2.7666\n",
      "Epoch: 54, Loss: 2.8829\n",
      "Epoch: 55, Loss: 2.8673\n",
      "Epoch: 56, Loss: 2.8579\n",
      "Epoch: 57, Loss: 2.8610\n",
      "Epoch: 58, Loss: 2.8795\n",
      "Epoch: 59, Loss: 2.9398\n",
      "Epoch: 60, Loss: 2.8480\n",
      "Epoch: 61, Loss: 2.7014\n",
      "Epoch: 62, Loss: 2.8627\n",
      "Epoch: 63, Loss: 2.7842\n",
      "Epoch: 64, Loss: 2.9687\n",
      "Epoch: 65, Loss: 2.7167\n",
      "Epoch: 66, Loss: 2.9327\n",
      "Epoch: 67, Loss: 2.6728\n",
      "Epoch: 68, Loss: 2.7890\n",
      "Epoch: 69, Loss: 2.6587\n",
      "Epoch: 70, Loss: 2.6515\n",
      "Epoch: 71, Loss: 2.6870\n",
      "Epoch: 72, Loss: 2.6874\n",
      "Epoch: 73, Loss: 2.5787\n",
      "Epoch: 74, Loss: 2.8392\n",
      "Epoch: 75, Loss: 2.7139\n",
      "Epoch: 76, Loss: 2.6013\n",
      "Epoch: 77, Loss: 2.8684\n",
      "Epoch: 78, Loss: 2.6029\n",
      "Epoch: 79, Loss: 2.7861\n",
      "Epoch: 80, Loss: 2.5331\n",
      "Epoch: 81, Loss: 2.5486\n",
      "Epoch: 82, Loss: 2.6647\n",
      "Epoch: 83, Loss: 2.8498\n",
      "Epoch: 84, Loss: 2.7732\n",
      "Epoch: 85, Loss: 2.5246\n",
      "Epoch: 86, Loss: 2.6781\n",
      "Epoch: 87, Loss: 2.6529\n",
      "Epoch: 88, Loss: 2.7233\n",
      "Epoch: 89, Loss: 2.9014\n",
      "Epoch: 90, Loss: 2.5716\n",
      "Epoch: 91, Loss: 2.7425\n",
      "Epoch: 92, Loss: 2.6585\n",
      "Epoch: 93, Loss: 2.6008\n",
      "Epoch: 94, Loss: 2.6931\n",
      "Epoch: 95, Loss: 2.7012\n",
      "Epoch: 96, Loss: 2.7106\n",
      "Epoch: 97, Loss: 2.6079\n",
      "Epoch: 98, Loss: 2.6693\n",
      "Epoch: 99, Loss: 2.7242\n"
     ]
    }
   ],
   "source": [
    "steps = 100\n",
    "\n",
    "for step in range(steps):\n",
    "    rng_key, subkey = jax.random.split(rng_key)\n",
    "    xb, yb = get_batch(train_ids, subkey, batch_size, block_size)\n",
    "\n",
    "    loss, grads = value_and_grad(loss_fn, argnums=(0))(\n",
    "        variables, \n",
    "        model.apply,\n",
    "        xb, \n",
    "        yb\n",
    "    )\n",
    "    updates, opt_state = optimizer.update(grads, opt_state, variables)\n",
    "    variables = optax.apply_updates(variables, updates) \n",
    "\n",
    "    print(f\"Epoch: {step}, Loss: {loss :.4f}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Text Generation Post-Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generated text: \n",
      "\n",
      "Sou,\n",
      "ic;\n",
      "EEI mod satthiceth anot.\n",
      " HAfaatthth,\n",
      "NCB\n"
     ]
    }
   ],
   "source": [
    "index_seq = jnp.zeros(shape=(1,1), dtype=jnp.uint16)\n",
    "max_new_tokens = 50\n",
    "\n",
    "generated_indices = generate(variables, model.apply, index_seq, rng_key, vocab_size, 1, block_size, max_new_tokens)\n",
    "generated_indices = list(np.array(generated_indices[0]))\n",
    "print(\"Generated text: \")\n",
    "print(decode(generated_indices))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "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.12.3"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}