scaled_self-attention.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Scaled Self-Attention\n",
    "![Scaled Self-Attention](./images/scaled_self-attention.png)\n",
    "\n",
    "We will build an autoregressive single-headed self-attention 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",
    "In the Bigram model the raw token embedding was used to perform the prediction. In this case we will combine the token embedding with a positional embedding, then process the combined embedding through a single-headed self-attention block, which is then processed through a dense layer to get the logits used for prediction. The scaled self-attention mechanism is performed by the query, key, and value (Q, K, V) following the operations shown in the image above. The self-attention calculations are demonstrated in the 'Self-Attention Mathematical Trick' section, further explanation is provided in 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: \"scaled\" self-attention](https://www.youtube.com/watch?v=kCc8FmEb1nY&t=4616s)\n",
    "- [3] [Attention Is All You Need](https://arxiv.org/abs/1706.03762)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import requests\n",
    "import numpy as np\n",
    "import jax\n",
    "from jax import random\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",
    "\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": [
    "## Self-Attention Mathematical Trick"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "a=\n",
      "[[1.         0.         0.        ]\n",
      " [0.5        0.5        0.        ]\n",
      " [0.33333334 0.33333334 0.33333334]]\n",
      "--\n",
      "b=\n",
      "[[5 7]\n",
      " [3 2]\n",
      " [9 9]]\n",
      "--\n",
      "c=\n",
      "[[5.        7.       ]\n",
      " [4.        4.5      ]\n",
      " [5.666667  6.0000005]]\n"
     ]
    }
   ],
   "source": [
    "# toy example illustrating how matrix multiplication can be used for a \"weighted aggregation\"\n",
    "a = jnp.tril(jnp.ones(shape=(3, 3)))\n",
    "a = a / jnp.sum(a, axis=1, keepdims=1)\n",
    "b = random.randint(key=rng_key, shape=(3, 2), minval=0, maxval=10)\n",
    "c = a @ b\n",
    "print('a=')\n",
    "print(a)\n",
    "print('--')\n",
    "print('b=')\n",
    "print(b)\n",
    "print('--')\n",
    "print('c=')\n",
    "print(c)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(4, 8, 2)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "B,T,C = 4,8,2 # batch, time, channels\n",
    "x = random.normal(key=rng_key, shape=(B, T, C))\n",
    "x.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We want x[b,t] = mean_{i<=t} x[b,i]\n",
    "xbow = np.zeros(shape=(B, T, C))\n",
    "for b in range(B):\n",
    "    for t in range(T):\n",
    "        xprev = x[b,:t+1] # (t,C)\n",
    "        xbow[b,t] = np.mean(xprev, axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(True, dtype=bool)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# version 2: using matrix multiply for a weighted aggregation\n",
    "weights = jnp.tril(jnp.ones(shape=(T, T)))\n",
    "weights = weights / jnp.sum(weights, axis=1, keepdims=1)\n",
    "xbow2 = weights @ x # (B, T, T) @ (B, T, C) ----> (B, T, C)\n",
    "jnp.allclose(xbow, xbow2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(True, dtype=bool)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# version 3: use Softmax\n",
    "tril = jnp.tril(jnp.ones(shape=(T, T), dtype=bool))\n",
    "weights = jnp.zeros(shape=(T, T))\n",
    "weights = masked_fill(tril, weights, -jnp.inf)\n",
    "weights = jax.nn.softmax(weights, axis=-1)\n",
    "xbow3 = weights @ x\n",
    "jnp.allclose(xbow, xbow3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(4, 8, 16)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# version 4: self-attention!\n",
    "B,T,C = 4,8,32 # batch, time, channels\n",
    "x = random.normal(key=rng_key, shape=(B, T, C))\n",
    "\n",
    "# let's see a single Head perform self-attention\n",
    "head_size = 16\n",
    "key = nn.Dense(head_size, use_bias=False)\n",
    "query = nn.Dense(head_size, use_bias=False)\n",
    "value = nn.Dense(head_size, use_bias=False)\n",
    "\n",
    "rng_key, subkey = jax.random.split(rng_key)\n",
    "key_variables = key.init(subkey, x)\n",
    "rng_key, subkey = jax.random.split(rng_key)\n",
    "query_variables = query.init(subkey, x)\n",
    "rng_key, subkey = jax.random.split(rng_key)\n",
    "value_variables = value.init(subkey, x)\n",
    "\n",
    "k = key.apply(key_variables, x)   # (B, T, 16)\n",
    "q = query.apply(query_variables, x) # (B, T, 16)\n",
    "weights =  q @ k.transpose((0, -1, -2)) # (B, T, 16) @ (B, 16, T) ---> (B, T, T)\n",
    "\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",
    "\n",
    "v = value.apply(value_variables, x)\n",
    "out = weights @ v\n",
    "\n",
    "out.shape\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([[1.0000000e+00, 0.0000000e+00, 0.0000000e+00, 0.0000000e+00,\n",
       "        0.0000000e+00, 0.0000000e+00, 0.0000000e+00, 0.0000000e+00],\n",
       "       [8.3365250e-01, 1.6634753e-01, 0.0000000e+00, 0.0000000e+00,\n",
       "        0.0000000e+00, 0.0000000e+00, 0.0000000e+00, 0.0000000e+00],\n",
       "       [2.7693647e-01, 7.2302037e-01, 4.3186730e-05, 0.0000000e+00,\n",
       "        0.0000000e+00, 0.0000000e+00, 0.0000000e+00, 0.0000000e+00],\n",
       "       [1.3745368e-04, 5.7609759e-02, 2.2170721e-03, 9.4003570e-01,\n",
       "        0.0000000e+00, 0.0000000e+00, 0.0000000e+00, 0.0000000e+00],\n",
       "       [9.7506523e-01, 2.3148456e-03, 7.4237287e-03, 9.0312017e-03,\n",
       "        6.1650318e-03, 0.0000000e+00, 0.0000000e+00, 0.0000000e+00],\n",
       "       [1.5689149e-01, 2.0382376e-02, 2.6865533e-01, 2.7895829e-01,\n",
       "        9.8608084e-02, 1.7650440e-01, 0.0000000e+00, 0.0000000e+00],\n",
       "       [9.9750141e-06, 3.4491016e-04, 3.6039501e-02, 1.0690327e-01,\n",
       "        8.4371493e-07, 2.5878797e-05, 8.5667557e-01, 0.0000000e+00],\n",
       "       [1.2620859e-01, 6.4960361e-02, 1.4840560e-03, 1.5524941e-02,\n",
       "        1.7746943e-01, 8.2721516e-02, 5.2967173e-01, 1.9592715e-03]],      dtype=float32)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "weights[0]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notes:\n",
    "- Attention is a **communication mechanism**. Can be seen as nodes in a directed graph looking at each other and aggregating information with a weighted sum from all nodes that point to them, with data-dependent weights.\n",
    "- There is no notion of space. Attention simply acts over a set of vectors. This is why we need to positionally encode tokens.\n",
    "- Each example across batch dimension is of course processed completely independently and never \"talk\" to each other\n",
    "- In an \"encoder\" attention block just delete the single line that does masking with `tril`, allowing all tokens to communicate. This block here is called a \"decoder\" attention block because it has triangular masking, and is usually used in autoregressive settings, like language modeling.\n",
    "- \"self-attention\" just means that the keys and values are produced from the same source as queries. In \"cross-attention\", the queries still get produced from x, but the keys and values come from some other, external source (e.g. an encoder module)\n",
    "- \"Scaled\" attention additionally divides `weights` by 1/sqrt(head_size). This makes it so when input Q,K are unit variance, wei will be unit variance too and Softmax will stay diffuse and not saturate too much. Illustration below"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "rng_key, subkey = jax.random.split(rng_key)\n",
    "k = random.normal(key=subkey, shape=(B, T, C))\n",
    "rng_key, subkey = jax.random.split(rng_key)\n",
    "q = random.normal(key=subkey, shape=(B, T, C))\n",
    "\n",
    "weights =  q @ k.transpose((0, -1, -2)) * C**-0.5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(1.0384574, dtype=float32)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "k.var()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(1.0208337, dtype=float32)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "q.var()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(0.96584, dtype=float32)"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "weights.var()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([0.19249782, 0.1426059 , 0.23511738, 0.1426059 , 0.287173  ],      dtype=float32)"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jax.nn.softmax(jnp.array([0.1, -0.2, 0.3, -0.2, 0.5]), axis=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([0.03260834, 0.00295816, 0.1615102 , 0.00295816, 0.79996514],      dtype=float32)"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jax.nn.softmax(jnp.array([0.1, -0.2, 0.3, -0.2, 0.5])*8, axis=-1) # gets too peaky, converges to one-hot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "# French to English translation example:\n",
    "\n",
    "# <--------- ENCODE ------------------><--------------- DECODE ----------------->\n",
    "# les réseaux de neurones sont géniaux! <START> neural networks are awesome!<END>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Build the Attention Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "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": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "class AttentionLanguageModel(nn.Module):\n",
    "    \"\"\"\n",
    "    Single-headed self-attention 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 a self-attention decoder block, which is then\n",
    "    processed through a dense layer to 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",
    "    head_size: 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",
    "        sa_head = Head(self.head_size)\n",
    "        x = sa_head(x) # apply one head of self-attention (B, T, C)\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": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = AttentionLanguageModel(vocab_size, n_embed, block_size, head_size)\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": 22,
   "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": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generated text: \n",
      "\n",
      "Tvg.'jtMvetQ\n",
      "y;:zYEfUmEhOuyYaXqu,wzhi Sfh,i3qD-'rqjGm&PDy'sja33d&?J3,EEgIdMBOm zu;vZlPkMm.lqbbLmqhFJ\n"
     ]
    }
   ],
   "source": [
    "index_seq = jnp.zeros(shape=(1,1), dtype=jnp.uint16)\n",
    "max_new_tokens = 100\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": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = optax.adamw(learning_rate=1e-2)\n",
    "opt_state = optimizer.init(variables)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 0, Loss: 4.1797\n",
      "Epoch: 1, Loss: 4.1357\n",
      "Epoch: 2, Loss: 4.1036\n",
      "Epoch: 3, Loss: 4.0445\n",
      "Epoch: 4, Loss: 3.9720\n",
      "Epoch: 5, Loss: 3.9320\n",
      "Epoch: 6, Loss: 3.8340\n",
      "Epoch: 7, Loss: 3.7325\n",
      "Epoch: 8, Loss: 3.6364\n",
      "Epoch: 9, Loss: 3.7192\n",
      "Epoch: 10, Loss: 3.6024\n",
      "Epoch: 11, Loss: 3.5051\n",
      "Epoch: 12, Loss: 3.3887\n",
      "Epoch: 13, Loss: 3.3672\n",
      "Epoch: 14, Loss: 3.4806\n",
      "Epoch: 15, Loss: 3.3368\n",
      "Epoch: 16, Loss: 3.3332\n",
      "Epoch: 17, Loss: 3.2800\n",
      "Epoch: 18, Loss: 3.1807\n",
      "Epoch: 19, Loss: 3.2574\n",
      "Epoch: 20, Loss: 3.2909\n",
      "Epoch: 21, Loss: 3.3127\n",
      "Epoch: 22, Loss: 3.2338\n",
      "Epoch: 23, Loss: 3.2703\n",
      "Epoch: 24, Loss: 3.2303\n",
      "Epoch: 25, Loss: 3.2596\n",
      "Epoch: 26, Loss: 3.2458\n",
      "Epoch: 27, Loss: 3.2228\n",
      "Epoch: 28, Loss: 3.2018\n",
      "Epoch: 29, Loss: 3.1379\n",
      "Epoch: 30, Loss: 3.2039\n",
      "Epoch: 31, Loss: 3.0974\n",
      "Epoch: 32, Loss: 3.1259\n",
      "Epoch: 33, Loss: 3.0729\n",
      "Epoch: 34, Loss: 3.1007\n",
      "Epoch: 35, Loss: 3.0092\n",
      "Epoch: 36, Loss: 3.1513\n",
      "Epoch: 37, Loss: 3.0775\n",
      "Epoch: 38, Loss: 3.0723\n",
      "Epoch: 39, Loss: 3.0594\n",
      "Epoch: 40, Loss: 3.3393\n",
      "Epoch: 41, Loss: 3.2729\n",
      "Epoch: 42, Loss: 3.0307\n",
      "Epoch: 43, Loss: 2.9294\n",
      "Epoch: 44, Loss: 3.1625\n",
      "Epoch: 45, Loss: 2.9055\n",
      "Epoch: 46, Loss: 2.8963\n",
      "Epoch: 47, Loss: 2.9345\n",
      "Epoch: 48, Loss: 2.8297\n",
      "Epoch: 49, Loss: 3.0023\n",
      "Epoch: 50, Loss: 2.9626\n",
      "Epoch: 51, Loss: 2.9266\n",
      "Epoch: 52, Loss: 2.9982\n",
      "Epoch: 53, Loss: 2.9798\n",
      "Epoch: 54, Loss: 2.9792\n",
      "Epoch: 55, Loss: 2.8609\n",
      "Epoch: 56, Loss: 2.7133\n",
      "Epoch: 57, Loss: 2.8850\n",
      "Epoch: 58, Loss: 2.9150\n",
      "Epoch: 59, Loss: 2.9828\n",
      "Epoch: 60, Loss: 2.7939\n",
      "Epoch: 61, Loss: 2.9564\n",
      "Epoch: 62, Loss: 2.6885\n",
      "Epoch: 63, Loss: 2.9355\n",
      "Epoch: 64, Loss: 2.7544\n",
      "Epoch: 65, Loss: 2.6792\n",
      "Epoch: 66, Loss: 2.7480\n",
      "Epoch: 67, Loss: 2.7320\n",
      "Epoch: 68, Loss: 2.6350\n",
      "Epoch: 69, Loss: 2.9241\n",
      "Epoch: 70, Loss: 2.8060\n",
      "Epoch: 71, Loss: 2.6282\n",
      "Epoch: 72, Loss: 2.9058\n",
      "Epoch: 73, Loss: 2.6956\n",
      "Epoch: 74, Loss: 2.8139\n",
      "Epoch: 75, Loss: 2.5760\n",
      "Epoch: 76, Loss: 2.6226\n",
      "Epoch: 77, Loss: 2.7261\n",
      "Epoch: 78, Loss: 2.9341\n",
      "Epoch: 79, Loss: 2.8351\n",
      "Epoch: 80, Loss: 2.6299\n",
      "Epoch: 81, Loss: 2.7057\n",
      "Epoch: 82, Loss: 2.7353\n",
      "Epoch: 83, Loss: 2.7719\n",
      "Epoch: 84, Loss: 2.9223\n",
      "Epoch: 85, Loss: 2.6682\n",
      "Epoch: 86, Loss: 2.8670\n",
      "Epoch: 87, Loss: 2.7461\n",
      "Epoch: 88, Loss: 2.6462\n",
      "Epoch: 89, Loss: 2.7124\n",
      "Epoch: 90, Loss: 2.7848\n",
      "Epoch: 91, Loss: 2.7165\n",
      "Epoch: 92, Loss: 2.6420\n",
      "Epoch: 93, Loss: 2.7257\n",
      "Epoch: 94, Loss: 2.7709\n",
      "Epoch: 95, Loss: 2.7682\n",
      "Epoch: 96, Loss: 2.6839\n",
      "Epoch: 97, Loss: 2.6401\n",
      "Epoch: 98, Loss: 2.5804\n",
      "Epoch: 99, Loss: 2.7055\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": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generated text: \n",
      "\n",
      "ag, g hart,\n",
      "T:\n",
      "llo nth anotd, I pe uusurkde,\n",
      "TOurnm, by\n",
      "J\n",
      "O faul?\n",
      "A bthsurthe.\n",
      "Coetilc mun he ur, ht\n"
     ]
    }
   ],
   "source": [
    "index_seq = jnp.zeros(shape=(1,1), dtype=jnp.uint16)\n",
    "max_new_tokens = 100\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": "flax2",
   "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.9.16"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "61a659281f728d55fc033691f61c254a0324a8156d8c726d9ae327f7f6499663"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}