tiny-stories-gpt.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "#----- imports --------\n",
    "import tqdm\n",
    "import torch\n",
    "# import wandb\n",
    "import os\n",
    "import tokenizers\n",
    "\n",
    "\n",
    "device= 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "torch.set_default_device(device)\n",
    "assert device == 'cuda', \"This notebook is not optimized for CPU\"\n",
    "\n",
    "config = {\n",
    "    \"learning_rate\": 1e-3,\n",
    "    \"eval_interval\": 300,\n",
    "    \"max_iters\": 60000, \n",
    "    \"H\": 32, # hidden dimension size\n",
    "    \"B\": 64,\n",
    "    \"T\": 256,\n",
    "    \"C\": 256,\n",
    "    \"feedforward_factor\": 3,\n",
    "    \"n_heads\": 8,\n",
    "    \"dropout\": 0.0,\n",
    "    \"l2_penalty\": 0.0,\n",
    "    \"n_layers\": 12,\n",
    "    \"tokenizer_vocab_size\": 2**13,\n",
    "    # \"git_hash\": os.popen(\"git rev-parse HEAD\").read().strip()\n",
    "}\n",
    "\n",
    "# initial\n",
    "for k,v in config.items():\n",
    "    locals ()[k] = v\n",
    "\n",
    "\n",
    "#wandb.init(\n",
    "#    project = \"tinystories\",\n",
    "#    config = config,\n",
    "#)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.chdir('multilingual_tinystories')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_sharded_story(shard_no):\n",
    "    return torch.load(f'tokenized/tokenized-{shard_no}.pt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/11 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 11/11 [05:11<00:00, 28.34s/it]\n"
     ]
    }
   ],
   "source": [
    "# load the tokenized stories in parallel using threads\n",
    "# this is faster than loading them sequentially\n",
    "from concurrent.futures import ThreadPoolExecutor\n",
    "with ThreadPoolExecutor() as pool:\n",
    "    stories = list(tqdm.tqdm(pool.map(load_sharded_story, range(11)), total=11))\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_stories = []\n",
    "for story in stories:\n",
    "    all_stories.extend(story)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "length of dataset in stories:  5304143\n",
      "length of stories in tokens 919732177\n"
     ]
    }
   ],
   "source": [
    "print(\"length of dataset in stories: \", len(all_stories))\n",
    "print(\"length of stories in tokens\", sum(len(story) for story in all_stories))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "# stories longer than 256 : 47531 out of 1000000, 4.75%\n"
     ]
    }
   ],
   "source": [
    "# plot a histogram of the lengths of the stories\n",
    "import matplotlib.pyplot as plt\n",
    "# plt.hist([len(story) for story in all_stories[:10000]], bins=50)\n",
    "num_stories_to_check = 1_000_000\n",
    "num_long = sum(len(story) > T for story in all_stories[:num_stories_to_check])\n",
    "print(\n",
    "    f\"# stories longer than {T} : {num_long} out of {num_stories_to_check}, {num_long/num_stories_to_check:.2%}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = tokenizers.ByteLevelBPETokenizer(\n",
    "    \"./tiny-stories-spanish-bpe-vocab.json\", \n",
    "    \"./tiny-stories-spanish-bpe-merges.txt\"\n",
    ")\n",
    "chars_per_token = 3.9 # hack"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[573]\n",
      "Hola\n",
      "vocab size:  8192\n",
      "first story decoded:  Un día, un niño llamado Leo jugó con cartas. Él tomó las cartas y las puso en un montón. Después, él las puso en otro montón. ¡Él inventó un nuevo juego! \n",
      "\n",
      "Leo miró las cartas. Él vio un gato. \"¡Miau!\" dijo el gato. Después, él vio un perro. \"¡Guau!\" dijo el perro. Leo sonrió. Él pensó que su juego era divertido. \n",
      "\n",
      "Después de jugar, Leo se fue a la cama. Él soñaba con el gato y el perro. En su sueño, el gato y el perro corrieron juntos. Ellos eran felices.\n",
      "\n",
      "¡De repente! Leo despertó. Él sintió que algo no estaba bien. Él fue a la ventana. Él vio al gato y al perro jugando afuera. ¡El sueño se había hecho realidad! \n",
      "\n",
      "Leo corrió afuera y jugó con el gato y el perro. Él jugó despacio. Él no quería asustarlos. Él era feliz. Él sabía que su sueño había sido real. \n",
      "\n"
     ]
    }
   ],
   "source": [
    "\n",
    "def encode(text):\n",
    "    return tokenizer.encode(text).ids\n",
    "def decode(encoded_text):\n",
    "    return tokenizer.decode(encoded_text)\n",
    "\n",
    "from tqdm import tqdm\n",
    "\n",
    "def batch_encode(text, batch_size):\n",
    "    tokens = []\n",
    "    for i in tqdm(range(0, len(text), batch_size)):\n",
    "        tokens.extend(encode(text[i:i+batch_size]))\n",
    "    return tokens\n",
    "\n",
    "\n",
    "hello_encoded = encode(\"Hola\")\n",
    "print(hello_encoded)\n",
    "print(decode(hello_encoded))\n",
    "vocab_size = tokenizer.get_vocab_size()\n",
    "print(\"vocab size: \", vocab_size)\n",
    "print('first story decoded: ', decode(all_stories[0].tolist()))\n",
    "PADDING_TOKEN_IDX= encode(\" \")[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "n = int(0.9*len(all_stories))\n",
    "\n",
    "train_data = all_stories[:n]\n",
    "val_data = all_stories[n:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "4773728"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(train_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = train_data[:T]\n",
    "y = train_data[1:T+1]\n",
    "for t in range(T):\n",
    "    context = x[:t+1]\n",
    "    target = y[t]\n",
    "    # print(\"when we see the text\", context, \"we predict the next character is\", target)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([ 220,  463,  281,  420, 1170,   13,  284,  410,  337,  503,  403,  945,\n",
      "        1213,   13,  452, 1082,  330,  436, 3540,   13,  371,  350,  330,  403,\n",
      "         381,  302,   25,  285,  639,   11,  292,  299, 1082,  330,  436, 3540,\n",
      "         315,  403,  352,  394,  749,  553, 3540,   13, 1094,  348, 5732,   13,\n",
      "         301,  198,  639,  311,  889,  323,  325, 1479,  306,  837,  436, 3540,\n",
      "          13, 1094,  348, 4641,   13,  375,  394,  355, 1213,   13,  403, 1369,\n",
      "         291, 1479,  331,  553, 2323,   13,  285,  400,  807,  749, 1243, 3540,\n",
      "         822,  302,  403,   13,  301,  198,  444,  350,  330,  403,  381,  302,\n",
      "          25,  285,  639,   11,  292,  347, 3830,  749, 1543, 3540,  315,  403,\n",
      "         318,  416,   13,  220,  512,  796,  651,  351,  388,  351,  749,  553,\n",
      "        3540,   13,  375, 1369,  436, 3540,  331,  325, 1334,   13,  301,  198,\n",
      "         639, 1802,  261,  749,  553, 3540,   13,  452, 2489,   11,  588,  796,\n",
      "         484,  792,   13, 1847, 2541,   11,  403,  311,  504,  399,   13,  375,\n",
      "        1178,  355, 1213,   13,  403,  880, 1048,  306, 1369,  294, 1697,  331,\n",
      "         553, 2323,   13,  301,  198,  709,  730,  521,  286,  355,  315,  729,\n",
      "         403,   13,  375,  311,  504,  399,   13,  375,  792,  553, 3540,  306,\n",
      "        1042,  523,  355,   13,  301,  220,  220,  220,  220,  220,  220,  220,\n",
      "         220,  220,  220,  220,  220,  220,  220,  220,  220,  220,  220,  220,\n",
      "         220,  220,  220,  220,  220,  220,  220,  220,  220,  220,  220,  220,\n",
      "         220,  220,  220,  220,  220,  220,  220,  220,  220,  220,  220,  220,\n",
      "         220,  220,  220,  220,  220,  220,  220,  220,  220,  220,  220,  220,\n",
      "         220,  220,  220,  220], device='cuda:0')\n",
      "tensor([ 463,  281,  420, 1170,   13,  284,  410,  337,  503,  403,  945, 1213,\n",
      "          13,  452, 1082,  330,  436, 3540,   13,  371,  350,  330,  403,  381,\n",
      "         302,   25,  285,  639,   11,  292,  299, 1082,  330,  436, 3540,  315,\n",
      "         403,  352,  394,  749,  553, 3540,   13, 1094,  348, 5732,   13,  301,\n",
      "         198,  639,  311,  889,  323,  325, 1479,  306,  837,  436, 3540,   13,\n",
      "        1094,  348, 4641,   13,  375,  394,  355, 1213,   13,  403, 1369,  291,\n",
      "        1479,  331,  553, 2323,   13,  285,  400,  807,  749, 1243, 3540,  822,\n",
      "         302,  403,   13,  301,  198,  444,  350,  330,  403,  381,  302,   25,\n",
      "         285,  639,   11,  292,  347, 3830,  749, 1543, 3540,  315,  403,  318,\n",
      "         416,   13,  220,  512,  796,  651,  351,  388,  351,  749,  553, 3540,\n",
      "          13,  375, 1369,  436, 3540,  331,  325, 1334,   13,  301,  198,  639,\n",
      "        1802,  261,  749,  553, 3540,   13,  452, 2489,   11,  588,  796,  484,\n",
      "         792,   13, 1847, 2541,   11,  403,  311,  504,  399,   13,  375, 1178,\n",
      "         355, 1213,   13,  403,  880, 1048,  306, 1369,  294, 1697,  331,  553,\n",
      "        2323,   13,  301,  198,  709,  730,  521,  286,  355,  315,  729,  403,\n",
      "          13,  375,  311,  504,  399,   13,  375,  792,  553, 3540,  306, 1042,\n",
      "         523,  355,   13,  301,  220,  220,  220,  220,  220,  220,  220,  220,\n",
      "         220,  220,  220,  220,  220,  220,  220,  220,  220,  220,  220,  220,\n",
      "         220,  220,  220,  220,  220,  220,  220,  220,  220,  220,  220,  220,\n",
      "         220,  220,  220,  220,  220,  220,  220,  220,  220,  220,  220,  220,\n",
      "         220,  220,  220,  220,  220,  220,  220,  220,  220,  220,  220,  220,\n",
      "         220,  220,  220,  220], device='cuda:0')\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(1339)\n",
    "\n",
    "def get_batch(split):\n",
    "    data = train_data if split == 'train' else val_data\n",
    "    ix = torch.randint(0, len(data), (B,)) \n",
    "    # ix = [i for i in range(B)]\n",
    "    \n",
    "    x = torch.full((B, T), PADDING_TOKEN_IDX, dtype=torch.long)\n",
    "    y = torch.full((B, T), PADDING_TOKEN_IDX, dtype=torch.long)\n",
    "\n",
    "    for sequence_index, random_story_index in enumerate(ix):\n",
    "        story = data[random_story_index].long()[:T - 1]\n",
    "        x[sequence_index][1: story.shape[0]+1] = story\n",
    "        y[sequence_index][: story.shape[0]] = story\n",
    "\n",
    "        \n",
    "    return x, y\n",
    "\n",
    "xb, yb = get_batch('train')\n",
    "\n",
    "print(xb[0])\n",
    "print(yb[0])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.nn import functional as F\n",
    "torch.manual_seed(1337)\n",
    "\n",
    "\n",
    "class Head(nn.Module):\n",
    "    '''One Head of self-attention'''\n",
    "    def __init__(self, H):\n",
    "        super().__init__()\n",
    "        self.query = nn.Linear(C, H, bias=False)\n",
    "        self.key = nn.Linear(C, H, bias=False)\n",
    "        self.value = nn.Linear(C, H, bias=False)\n",
    "        # self.output = nn.Linear(H, C, bias=False) # output matrix\n",
    "        self.register_buffer('tril', torch.tril(torch.ones(T, T)))\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, x):\n",
    "        query_vectors = self.query(x)\n",
    "        key_vectors = self.key(x)\n",
    "\n",
    "\n",
    "        # Attention masking(so we can't look into the past):\n",
    "\n",
    "        tril = self.tril\n",
    "        wei = torch.zeros(T, T) \n",
    "        wei = wei.masked_fill(tril == 0, float('-inf')) # set the upper triangular to -inf\n",
    "        # xbow = wei @ x # apply the mask to the input, bag of words because simple avg.\n",
    "\n",
    "        # multiply the two to get the attention weights\n",
    "        attention_pattern = query_vectors @ key_vectors.transpose(-2, -1) # T, T\n",
    "        attention_pattern = attention_pattern / (H ** 0.5) # scale the attention pattern for numerical stability\n",
    "        attention_weights = F.softmax(attention_pattern + wei, dim=-1) # T, T (the row dimension is the query)\n",
    "        attention_weights = self.dropout(attention_weights)\n",
    "\n",
    "        value_vectors = self.value(x) # the direction we should go in the embedding space for each token (ie more blue) T, H\n",
    "\n",
    "        # apply the attention weights to the value vectors\n",
    "        context = attention_weights @ value_vectors # T, H\n",
    "\n",
    "        # project back into original space from value space\n",
    "        # return self.output(context)\n",
    "        return context\n",
    "\n",
    "x = torch.randn(B,T,C)\n",
    "head = Head(H)\n",
    "# head(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultiHeadAttention(nn.Module):\n",
    "    '''Multiple heads of self-attention'''\n",
    "    def __init__(self, H, C, n_heads): # H is head embedding space size, n_heads is number of heads\n",
    "        super().__init__()\n",
    "        self.heads = nn.ModuleList([Head(H) for _ in range(n_heads)])\n",
    "        self.combine_heads = nn.Linear(H*n_heads, C)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "\n",
    "    def forward(self,x):\n",
    "        x = torch.cat([head(x) for head in self.heads], dim=-1)\n",
    "        x = self.combine_heads(x)  # T, C\n",
    "        return self.dropout(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([64, 256, 32])"
      ]
     },
     "execution_count": 67,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "head = MultiHeadAttention(H, C, n_heads)\n",
    "head.heads[0].forward(x).shape\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [],
   "source": [
    "class FeedForward(nn.Module):\n",
    "    '''Feed-forward neural network'''\n",
    "    def __init__(self, C):\n",
    "        super().__init__()\n",
    "        self.net = nn.Sequential(\n",
    "            nn.Linear(C, C * feedforward_factor),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(C * feedforward_factor, C),\n",
    "            nn.Dropout(dropout)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.net(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LayerNorm(nn.Module):\n",
    "    '''Layer normalization'''\n",
    "    def __init__(self, C, use_affine=True):\n",
    "        super().__init__()\n",
    "        self.gamma = nn.Parameter(torch.ones(C)) if use_affine else None\n",
    "        self.beta = nn.Parameter(torch.zeros(C)) if use_affine else None\n",
    "\n",
    "    def forward(self, x):\n",
    "        mean = x.mean(-1, keepdim=True)\n",
    "        std = x.std(-1, keepdim=True)\n",
    "        if self.gamma is not None and self.beta is not None:\n",
    "            return self.gamma * (x - mean) / (std + 1e-6) + self.beta\n",
    "        else:\n",
    "            return (x - mean) / (std + 1e-6)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Block(nn.Module):\n",
    "    '''Transformer block'''\n",
    "    def __init__(self, H, C, n_heads):\n",
    "        super().__init__()\n",
    "        self.attention = MultiHeadAttention(H, C, n_heads)\n",
    "        self.ff = FeedForward(C)\n",
    "        self.norm1 = LayerNorm(C, use_affine=True)\n",
    "        self.norm2 = LayerNorm(C, use_affine=True)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = x + self.attention(self.norm1(x))\n",
    "        x = x + self.ff(self.norm2(x))\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([64, 256, 8192])\n",
      "tensor(9.4822, device='cuda:0', grad_fn=<NllLossBackward0>)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(tensor([[[ 1.2796,  0.2360,  0.4649,  ..., -0.8759,  0.0897,  0.0632],\n",
       "          [ 2.1053, -0.8288,  0.3834,  ..., -1.1195,  0.1282, -0.5274],\n",
       "          [ 2.2574,  0.5220,  0.4413,  ..., -1.4343, -1.8042, -1.5558],\n",
       "          ...,\n",
       "          [ 2.5515, -0.5823,  0.1534,  ..., -1.5195, -1.1673, -0.1516],\n",
       "          [ 1.2516, -0.4528,  0.2085,  ..., -0.0323, -0.3258, -0.9920],\n",
       "          [ 2.0084, -1.3079,  0.2250,  ...,  0.0589, -1.0247, -1.3288]]],\n",
       "        device='cuda:0', grad_fn=<ViewBackward0>),\n",
       " None)"
      ]
     },
     "execution_count": 71,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class GPT(nn.Module):\n",
    "\n",
    "    def __init__(self, n_layers):\n",
    "        super().__init__()\n",
    "        self.token_embedding_table = nn.Embedding(vocab_size, C) \n",
    "        self.position_embedding_table = nn.Embedding(T, C)\n",
    "        self.lm_head = nn.Linear(C, vocab_size)\n",
    "        self.layers = nn.ModuleList([Block(H, C, n_heads) for _ in range(n_layers)])\n",
    "        self.block = nn.ModuleList([Block(H, C, n_heads)])\n",
    "    \n",
    "    def forward(self, idx, targets=None):\n",
    "        B, T = idx.shape\n",
    "        token_emb = self.token_embedding_table(idx) # batch_dim, sequence_dim, embedding_dim\n",
    "        pos_emb = self.position_embedding_table(torch.arange(T))\n",
    "        x = token_emb + pos_emb # token identities and positions contained\n",
    "\n",
    "        for layer in self.layers:\n",
    "            x = layer(x)\n",
    "\n",
    "        logits = self.lm_head(x) # batch_dim, sequence_dim, vocab_size\n",
    "\n",
    "        batch_dim, sequence_dim, embedding_dim = logits.size()\n",
    "\n",
    "        # loss = F.cross_entropy(logits, targets) this won't work because we need 1d logits and 1d targets\n",
    "        # one-hot-vectors are a line in the x-dimension, so the shape of shape of the logits should be (-1, vocab_size).\n",
    "\n",
    "        if targets is None:\n",
    "            return logits, None\n",
    "        else:\n",
    "            # a list of all the predictions, reguardles of batch.\n",
    "            # xdim: probabilities of each character in the vocab (embedding_dim=vocab_size)\n",
    "            # ydim: all predictions for all batches flattened (batch_dim*sequence_dim)\n",
    "            logits_loss_view = logits.view(-1, vocab_size) \n",
    "            # targets loss view\n",
    "            # xdim: all targets for all batches flattened (batch_dim*sequence_dim)\n",
    "            # so this would be like, [1,4,5,1,2,3, ...]\n",
    "            # where each number is the correct next index of the one hot vector\n",
    "            targets_loss_view = targets.view(-1)\n",
    "            loss = F.cross_entropy(logits_loss_view, targets_loss_view)\n",
    "            return logits, loss\n",
    "\n",
    "    def generate(self, idx, max_new_tokens):\n",
    "        for _ in range(max_new_tokens):\n",
    "            logits, loss = self(idx[:,-T:])\n",
    "            # get the predictions of the last token\n",
    "            last_token_logits = logits[:, -1, :] # all batches, last token, all probabilities\n",
    "            # softmax to get probabilities\n",
    "            probabilities = F.softmax(last_token_logits, dim=-1)\n",
    "            # sample from the probabilities\n",
    "            next_token = torch.multinomial(probabilities, num_samples=1)\n",
    "            # add the new token to the idx tensor\n",
    "            idx = torch.cat((idx, next_token), dim=1)\n",
    "        return idx\n",
    "    def prompt_model(self, prompt, max_new_tokens, temperature=0.5):\n",
    "        autoregressive_seq = encode(prompt)\n",
    "        for _ in range(max_new_tokens):\n",
    "            prediction_index = len(autoregressive_seq)-1\n",
    "\n",
    "            model_input = torch.tensor(autoregressive_seq)\n",
    "            \n",
    "            while model_input.shape[0] < T:\n",
    "                pad_token = torch.tensor(encode(\"\\n\"))\n",
    "                model_input = torch.cat((model_input, pad_token), dim=0)\n",
    "\n",
    "            model_input\n",
    "            model_input = model_input.unsqueeze(0)\n",
    "\n",
    "            logits, loss = model(model_input)\n",
    "            prediction_token = logits[:, prediction_index, :] / temperature\n",
    "            probabilities = F.softmax(prediction_token, dim=-1)\n",
    "            next_token = torch.multinomial(probabilities, num_samples=1)\n",
    "            next_token = next_token.item()\n",
    "\n",
    "            autoregressive_seq.append(next_token)\n",
    "        # get the autoregressive sequence\n",
    "        return decode(autoregressive_seq)\n",
    "\n",
    "\n",
    "\n",
    "    \n",
    "\n",
    "model = GPT(n_layers)\n",
    "logits, loss = model(xb, yb)\n",
    "print(logits.shape)\n",
    "print(loss)\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "test_idx = torch.zeros(1, T).long()\n",
    "model.forward(idx=test_idx)\n",
    "# decode(model.generate(idx=test_idx, max_new_tokens=100)[0].tolist())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "GPT(\n",
       "  (token_embedding_table): Embedding(8192, 256)\n",
       "  (position_embedding_table): Embedding(256, 256)\n",
       "  (lm_head): Linear(in_features=256, out_features=8192, bias=True)\n",
       "  (layers): ModuleList(\n",
       "    (0-11): 12 x Block(\n",
       "      (attention): MultiHeadAttention(\n",
       "        (heads): ModuleList(\n",
       "          (0-7): 8 x Head(\n",
       "            (query): Linear(in_features=256, out_features=32, bias=False)\n",
       "            (key): Linear(in_features=256, out_features=32, bias=False)\n",
       "            (value): Linear(in_features=256, out_features=32, bias=False)\n",
       "            (dropout): Dropout(p=0.0, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (combine_heads): Linear(in_features=256, out_features=256, bias=True)\n",
       "        (dropout): Dropout(p=0.0, inplace=False)\n",
       "      )\n",
       "      (ff): FeedForward(\n",
       "        (net): Sequential(\n",
       "          (0): Linear(in_features=256, out_features=768, bias=True)\n",
       "          (1): ReLU()\n",
       "          (2): Linear(in_features=768, out_features=256, bias=True)\n",
       "          (3): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (norm1): LayerNorm()\n",
       "      (norm2): LayerNorm()\n",
       "    )\n",
       "  )\n",
       "  (block): ModuleList(\n",
       "    (0): Block(\n",
       "      (attention): MultiHeadAttention(\n",
       "        (heads): ModuleList(\n",
       "          (0-7): 8 x Head(\n",
       "            (query): Linear(in_features=256, out_features=32, bias=False)\n",
       "            (key): Linear(in_features=256, out_features=32, bias=False)\n",
       "            (value): Linear(in_features=256, out_features=32, bias=False)\n",
       "            (dropout): Dropout(p=0.0, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (combine_heads): Linear(in_features=256, out_features=256, bias=True)\n",
       "        (dropout): Dropout(p=0.0, inplace=False)\n",
       "      )\n",
       "      (ff): FeedForward(\n",
       "        (net): Sequential(\n",
       "          (0): Linear(in_features=256, out_features=768, bias=True)\n",
       "          (1): ReLU()\n",
       "          (2): Linear(in_features=768, out_features=256, bias=True)\n",
       "          (3): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (norm1): LayerNorm()\n",
       "      (norm2): LayerNorm()\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 72,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "number of parameters in the model:  12817664\n"
     ]
    }
   ],
   "source": [
    "# get the number of parameters in the model\n",
    "def count_parameters(model):\n",
    "    return sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
    "\n",
    "print(\"number of parameters in the model: \", count_parameters(model))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0]], device='cuda:0')"
      ]
     },
     "execution_count": 75,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# logits, loss = self(idx[:,-T:])\n",
    "\n",
    "idx = torch.zeros(1, 1).long()\n",
    "idx[:,-T:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "device(type='cuda', index=0)"
      ]
     },
     "execution_count": 76,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.token_embedding_table.weight.device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_iters = 10\n",
    "eval_interval = 300\n",
    "@torch.no_grad()\n",
    "def estimate_loss(is_last=False):\n",
    "    out = {}\n",
    "    model.eval()\n",
    "    for split in ['train', 'val']:\n",
    "        real_iters = eval_iters\n",
    "        if is_last and split == 'val':  # increase last eval to mitigate noise\n",
    "            real_iters *= 10 \n",
    "        losses = torch.zeros(real_iters)\n",
    "        for k in range(real_iters):\n",
    "            X, Y = get_batch(split)\n",
    "            logits, loss = model(X, Y)\n",
    "            losses[k] = loss.item()\n",
    "        out[split] = losses.mean() / chars_per_token\n",
    "    model.train()\n",
    "    return out\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "parameter_to_data_ratio=2.6850428009304257\n",
      "token_embedding_table.weight: 2097152\n",
      "lm_head.weight: 2097152\n",
      "layers.0.ff.net.0.weight: 196608\n",
      "layers.0.ff.net.2.weight: 196608\n",
      "layers.1.ff.net.0.weight: 196608\n",
      "layers.1.ff.net.2.weight: 196608\n",
      "layers.2.ff.net.0.weight: 196608\n",
      "layers.2.ff.net.2.weight: 196608\n",
      "layers.3.ff.net.0.weight: 196608\n",
      "layers.3.ff.net.2.weight: 196608\n",
      "layers.4.ff.net.0.weight: 196608\n",
      "layers.4.ff.net.2.weight: 196608\n",
      "layers.5.ff.net.0.weight: 196608\n",
      "layers.5.ff.net.2.weight: 196608\n",
      "layers.6.ff.net.0.weight: 196608\n",
      "layers.6.ff.net.2.weight: 196608\n",
      "layers.7.ff.net.0.weight: 196608\n",
      "layers.7.ff.net.2.weight: 196608\n",
      "layers.8.ff.net.0.weight: 196608\n",
      "layers.8.ff.net.2.weight: 196608\n",
      "layers.9.ff.net.0.weight: 196608\n",
      "layers.9.ff.net.2.weight: 196608\n",
      "layers.10.ff.net.0.weight: 196608\n",
      "layers.10.ff.net.2.weight: 196608\n",
      "layers.11.ff.net.0.weight: 196608\n",
      "layers.11.ff.net.2.weight: 196608\n",
      "block.0.ff.net.0.weight: 196608\n",
      "block.0.ff.net.2.weight: 196608\n",
      "position_embedding_table.weight: 65536\n",
      "layers.0.attention.combine_heads.weight: 65536\n",
      "layers.1.attention.combine_heads.weight: 65536\n",
      "layers.2.attention.combine_heads.weight: 65536\n",
      "layers.3.attention.combine_heads.weight: 65536\n",
      "layers.4.attention.combine_heads.weight: 65536\n",
      "layers.5.attention.combine_heads.weight: 65536\n",
      "layers.6.attention.combine_heads.weight: 65536\n",
      "layers.7.attention.combine_heads.weight: 65536\n",
      "layers.8.attention.combine_heads.weight: 65536\n",
      "layers.9.attention.combine_heads.weight: 65536\n",
      "layers.10.attention.combine_heads.weight: 65536\n",
      "layers.11.attention.combine_heads.weight: 65536\n",
      "block.0.attention.combine_heads.weight: 65536\n",
      "lm_head.bias: 8192\n",
      "layers.0.attention.heads.0.query.weight: 8192\n",
      "layers.0.attention.heads.0.key.weight: 8192\n",
      "layers.0.attention.heads.0.value.weight: 8192\n",
      "layers.0.attention.heads.1.query.weight: 8192\n",
      "layers.0.attention.heads.1.key.weight: 8192\n",
      "layers.0.attention.heads.1.value.weight: 8192\n",
      "layers.0.attention.heads.2.query.weight: 8192\n",
      "layers.0.attention.heads.2.key.weight: 8192\n",
      "layers.0.attention.heads.2.value.weight: 8192\n",
      "layers.0.attention.heads.3.query.weight: 8192\n",
      "layers.0.attention.heads.3.key.weight: 8192\n",
      "layers.0.attention.heads.3.value.weight: 8192\n",
      "layers.0.attention.heads.4.query.weight: 8192\n",
      "layers.0.attention.heads.4.key.weight: 8192\n",
      "layers.0.attention.heads.4.value.weight: 8192\n",
      "layers.0.attention.heads.5.query.weight: 8192\n",
      "layers.0.attention.heads.5.key.weight: 8192\n",
      "layers.0.attention.heads.5.value.weight: 8192\n",
      "layers.0.attention.heads.6.query.weight: 8192\n",
      "layers.0.attention.heads.6.key.weight: 8192\n",
      "layers.0.attention.heads.6.value.weight: 8192\n",
      "layers.0.attention.heads.7.query.weight: 8192\n",
      "layers.0.attention.heads.7.key.weight: 8192\n",
      "layers.0.attention.heads.7.value.weight: 8192\n",
      "layers.1.attention.heads.0.query.weight: 8192\n",
      "layers.1.attention.heads.0.key.weight: 8192\n",
      "layers.1.attention.heads.0.value.weight: 8192\n",
      "layers.1.attention.heads.1.query.weight: 8192\n",
      "layers.1.attention.heads.1.key.weight: 8192\n",
      "layers.1.attention.heads.1.value.weight: 8192\n",
      "layers.1.attention.heads.2.query.weight: 8192\n",
      "layers.1.attention.heads.2.key.weight: 8192\n",
      "layers.1.attention.heads.2.value.weight: 8192\n",
      "layers.1.attention.heads.3.query.weight: 8192\n",
      "layers.1.attention.heads.3.key.weight: 8192\n",
      "layers.1.attention.heads.3.value.weight: 8192\n",
      "layers.1.attention.heads.4.query.weight: 8192\n",
      "layers.1.attention.heads.4.key.weight: 8192\n",
      "layers.1.attention.heads.4.value.weight: 8192\n",
      "layers.1.attention.heads.5.query.weight: 8192\n",
      "layers.1.attention.heads.5.key.weight: 8192\n",
      "layers.1.attention.heads.5.value.weight: 8192\n",
      "layers.1.attention.heads.6.query.weight: 8192\n",
      "layers.1.attention.heads.6.key.weight: 8192\n",
      "layers.1.attention.heads.6.value.weight: 8192\n",
      "layers.1.attention.heads.7.query.weight: 8192\n",
      "layers.1.attention.heads.7.key.weight: 8192\n",
      "layers.1.attention.heads.7.value.weight: 8192\n",
      "layers.2.attention.heads.0.query.weight: 8192\n",
      "layers.2.attention.heads.0.key.weight: 8192\n",
      "layers.2.attention.heads.0.value.weight: 8192\n",
      "layers.2.attention.heads.1.query.weight: 8192\n",
      "layers.2.attention.heads.1.key.weight: 8192\n",
      "layers.2.attention.heads.1.value.weight: 8192\n",
      "layers.2.attention.heads.2.query.weight: 8192\n",
      "layers.2.attention.heads.2.key.weight: 8192\n",
      "layers.2.attention.heads.2.value.weight: 8192\n",
      "layers.2.attention.heads.3.query.weight: 8192\n",
      "layers.2.attention.heads.3.key.weight: 8192\n",
      "layers.2.attention.heads.3.value.weight: 8192\n",
      "layers.2.attention.heads.4.query.weight: 8192\n",
      "layers.2.attention.heads.4.key.weight: 8192\n",
      "layers.2.attention.heads.4.value.weight: 8192\n",
      "layers.2.attention.heads.5.query.weight: 8192\n",
      "layers.2.attention.heads.5.key.weight: 8192\n",
      "layers.2.attention.heads.5.value.weight: 8192\n",
      "layers.2.attention.heads.6.query.weight: 8192\n",
      "layers.2.attention.heads.6.key.weight: 8192\n",
      "layers.2.attention.heads.6.value.weight: 8192\n",
      "layers.2.attention.heads.7.query.weight: 8192\n",
      "layers.2.attention.heads.7.key.weight: 8192\n",
      "layers.2.attention.heads.7.value.weight: 8192\n",
      "layers.3.attention.heads.0.query.weight: 8192\n",
      "layers.3.attention.heads.0.key.weight: 8192\n",
      "layers.3.attention.heads.0.value.weight: 8192\n",
      "layers.3.attention.heads.1.query.weight: 8192\n",
      "layers.3.attention.heads.1.key.weight: 8192\n",
      "layers.3.attention.heads.1.value.weight: 8192\n",
      "layers.3.attention.heads.2.query.weight: 8192\n",
      "layers.3.attention.heads.2.key.weight: 8192\n",
      "layers.3.attention.heads.2.value.weight: 8192\n",
      "layers.3.attention.heads.3.query.weight: 8192\n",
      "layers.3.attention.heads.3.key.weight: 8192\n",
      "layers.3.attention.heads.3.value.weight: 8192\n",
      "layers.3.attention.heads.4.query.weight: 8192\n",
      "layers.3.attention.heads.4.key.weight: 8192\n",
      "layers.3.attention.heads.4.value.weight: 8192\n",
      "layers.3.attention.heads.5.query.weight: 8192\n",
      "layers.3.attention.heads.5.key.weight: 8192\n",
      "layers.3.attention.heads.5.value.weight: 8192\n",
      "layers.3.attention.heads.6.query.weight: 8192\n",
      "layers.3.attention.heads.6.key.weight: 8192\n",
      "layers.3.attention.heads.6.value.weight: 8192\n",
      "layers.3.attention.heads.7.query.weight: 8192\n",
      "layers.3.attention.heads.7.key.weight: 8192\n",
      "layers.3.attention.heads.7.value.weight: 8192\n",
      "layers.4.attention.heads.0.query.weight: 8192\n",
      "layers.4.attention.heads.0.key.weight: 8192\n",
      "layers.4.attention.heads.0.value.weight: 8192\n",
      "layers.4.attention.heads.1.query.weight: 8192\n",
      "layers.4.attention.heads.1.key.weight: 8192\n",
      "layers.4.attention.heads.1.value.weight: 8192\n",
      "layers.4.attention.heads.2.query.weight: 8192\n",
      "layers.4.attention.heads.2.key.weight: 8192\n",
      "layers.4.attention.heads.2.value.weight: 8192\n",
      "layers.4.attention.heads.3.query.weight: 8192\n",
      "layers.4.attention.heads.3.key.weight: 8192\n",
      "layers.4.attention.heads.3.value.weight: 8192\n",
      "layers.4.attention.heads.4.query.weight: 8192\n",
      "layers.4.attention.heads.4.key.weight: 8192\n",
      "layers.4.attention.heads.4.value.weight: 8192\n",
      "layers.4.attention.heads.5.query.weight: 8192\n",
      "layers.4.attention.heads.5.key.weight: 8192\n",
      "layers.4.attention.heads.5.value.weight: 8192\n",
      "layers.4.attention.heads.6.query.weight: 8192\n",
      "layers.4.attention.heads.6.key.weight: 8192\n",
      "layers.4.attention.heads.6.value.weight: 8192\n",
      "layers.4.attention.heads.7.query.weight: 8192\n",
      "layers.4.attention.heads.7.key.weight: 8192\n",
      "layers.4.attention.heads.7.value.weight: 8192\n",
      "layers.5.attention.heads.0.query.weight: 8192\n",
      "layers.5.attention.heads.0.key.weight: 8192\n",
      "layers.5.attention.heads.0.value.weight: 8192\n",
      "layers.5.attention.heads.1.query.weight: 8192\n",
      "layers.5.attention.heads.1.key.weight: 8192\n",
      "layers.5.attention.heads.1.value.weight: 8192\n",
      "layers.5.attention.heads.2.query.weight: 8192\n",
      "layers.5.attention.heads.2.key.weight: 8192\n",
      "layers.5.attention.heads.2.value.weight: 8192\n",
      "layers.5.attention.heads.3.query.weight: 8192\n",
      "layers.5.attention.heads.3.key.weight: 8192\n",
      "layers.5.attention.heads.3.value.weight: 8192\n",
      "layers.5.attention.heads.4.query.weight: 8192\n",
      "layers.5.attention.heads.4.key.weight: 8192\n",
      "layers.5.attention.heads.4.value.weight: 8192\n",
      "layers.5.attention.heads.5.query.weight: 8192\n",
      "layers.5.attention.heads.5.key.weight: 8192\n",
      "layers.5.attention.heads.5.value.weight: 8192\n",
      "layers.5.attention.heads.6.query.weight: 8192\n",
      "layers.5.attention.heads.6.key.weight: 8192\n",
      "layers.5.attention.heads.6.value.weight: 8192\n",
      "layers.5.attention.heads.7.query.weight: 8192\n",
      "layers.5.attention.heads.7.key.weight: 8192\n",
      "layers.5.attention.heads.7.value.weight: 8192\n",
      "layers.6.attention.heads.0.query.weight: 8192\n",
      "layers.6.attention.heads.0.key.weight: 8192\n",
      "layers.6.attention.heads.0.value.weight: 8192\n",
      "layers.6.attention.heads.1.query.weight: 8192\n",
      "layers.6.attention.heads.1.key.weight: 8192\n",
      "layers.6.attention.heads.1.value.weight: 8192\n",
      "layers.6.attention.heads.2.query.weight: 8192\n",
      "layers.6.attention.heads.2.key.weight: 8192\n",
      "layers.6.attention.heads.2.value.weight: 8192\n",
      "layers.6.attention.heads.3.query.weight: 8192\n",
      "layers.6.attention.heads.3.key.weight: 8192\n",
      "layers.6.attention.heads.3.value.weight: 8192\n",
      "layers.6.attention.heads.4.query.weight: 8192\n",
      "layers.6.attention.heads.4.key.weight: 8192\n",
      "layers.6.attention.heads.4.value.weight: 8192\n",
      "layers.6.attention.heads.5.query.weight: 8192\n",
      "layers.6.attention.heads.5.key.weight: 8192\n",
      "layers.6.attention.heads.5.value.weight: 8192\n",
      "layers.6.attention.heads.6.query.weight: 8192\n",
      "layers.6.attention.heads.6.key.weight: 8192\n",
      "layers.6.attention.heads.6.value.weight: 8192\n",
      "layers.6.attention.heads.7.query.weight: 8192\n",
      "layers.6.attention.heads.7.key.weight: 8192\n",
      "layers.6.attention.heads.7.value.weight: 8192\n",
      "layers.7.attention.heads.0.query.weight: 8192\n",
      "layers.7.attention.heads.0.key.weight: 8192\n",
      "layers.7.attention.heads.0.value.weight: 8192\n",
      "layers.7.attention.heads.1.query.weight: 8192\n",
      "layers.7.attention.heads.1.key.weight: 8192\n",
      "layers.7.attention.heads.1.value.weight: 8192\n",
      "layers.7.attention.heads.2.query.weight: 8192\n",
      "layers.7.attention.heads.2.key.weight: 8192\n",
      "layers.7.attention.heads.2.value.weight: 8192\n",
      "layers.7.attention.heads.3.query.weight: 8192\n",
      "layers.7.attention.heads.3.key.weight: 8192\n",
      "layers.7.attention.heads.3.value.weight: 8192\n",
      "layers.7.attention.heads.4.query.weight: 8192\n",
      "layers.7.attention.heads.4.key.weight: 8192\n",
      "layers.7.attention.heads.4.value.weight: 8192\n",
      "layers.7.attention.heads.5.query.weight: 8192\n",
      "layers.7.attention.heads.5.key.weight: 8192\n",
      "layers.7.attention.heads.5.value.weight: 8192\n",
      "layers.7.attention.heads.6.query.weight: 8192\n",
      "layers.7.attention.heads.6.key.weight: 8192\n",
      "layers.7.attention.heads.6.value.weight: 8192\n",
      "layers.7.attention.heads.7.query.weight: 8192\n",
      "layers.7.attention.heads.7.key.weight: 8192\n",
      "layers.7.attention.heads.7.value.weight: 8192\n",
      "layers.8.attention.heads.0.query.weight: 8192\n",
      "layers.8.attention.heads.0.key.weight: 8192\n",
      "layers.8.attention.heads.0.value.weight: 8192\n",
      "layers.8.attention.heads.1.query.weight: 8192\n",
      "layers.8.attention.heads.1.key.weight: 8192\n",
      "layers.8.attention.heads.1.value.weight: 8192\n",
      "layers.8.attention.heads.2.query.weight: 8192\n",
      "layers.8.attention.heads.2.key.weight: 8192\n",
      "layers.8.attention.heads.2.value.weight: 8192\n",
      "layers.8.attention.heads.3.query.weight: 8192\n",
      "layers.8.attention.heads.3.key.weight: 8192\n",
      "layers.8.attention.heads.3.value.weight: 8192\n",
      "layers.8.attention.heads.4.query.weight: 8192\n",
      "layers.8.attention.heads.4.key.weight: 8192\n",
      "layers.8.attention.heads.4.value.weight: 8192\n",
      "layers.8.attention.heads.5.query.weight: 8192\n",
      "layers.8.attention.heads.5.key.weight: 8192\n",
      "layers.8.attention.heads.5.value.weight: 8192\n",
      "layers.8.attention.heads.6.query.weight: 8192\n",
      "layers.8.attention.heads.6.key.weight: 8192\n",
      "layers.8.attention.heads.6.value.weight: 8192\n",
      "layers.8.attention.heads.7.query.weight: 8192\n",
      "layers.8.attention.heads.7.key.weight: 8192\n",
      "layers.8.attention.heads.7.value.weight: 8192\n",
      "layers.9.attention.heads.0.query.weight: 8192\n",
      "layers.9.attention.heads.0.key.weight: 8192\n",
      "layers.9.attention.heads.0.value.weight: 8192\n",
      "layers.9.attention.heads.1.query.weight: 8192\n",
      "layers.9.attention.heads.1.key.weight: 8192\n",
      "layers.9.attention.heads.1.value.weight: 8192\n",
      "layers.9.attention.heads.2.query.weight: 8192\n",
      "layers.9.attention.heads.2.key.weight: 8192\n",
      "layers.9.attention.heads.2.value.weight: 8192\n",
      "layers.9.attention.heads.3.query.weight: 8192\n",
      "layers.9.attention.heads.3.key.weight: 8192\n",
      "layers.9.attention.heads.3.value.weight: 8192\n",
      "layers.9.attention.heads.4.query.weight: 8192\n",
      "layers.9.attention.heads.4.key.weight: 8192\n",
      "layers.9.attention.heads.4.value.weight: 8192\n",
      "layers.9.attention.heads.5.query.weight: 8192\n",
      "layers.9.attention.heads.5.key.weight: 8192\n",
      "layers.9.attention.heads.5.value.weight: 8192\n",
      "layers.9.attention.heads.6.query.weight: 8192\n",
      "layers.9.attention.heads.6.key.weight: 8192\n",
      "layers.9.attention.heads.6.value.weight: 8192\n",
      "layers.9.attention.heads.7.query.weight: 8192\n",
      "layers.9.attention.heads.7.key.weight: 8192\n",
      "layers.9.attention.heads.7.value.weight: 8192\n",
      "layers.10.attention.heads.0.query.weight: 8192\n",
      "layers.10.attention.heads.0.key.weight: 8192\n",
      "layers.10.attention.heads.0.value.weight: 8192\n",
      "layers.10.attention.heads.1.query.weight: 8192\n",
      "layers.10.attention.heads.1.key.weight: 8192\n",
      "layers.10.attention.heads.1.value.weight: 8192\n",
      "layers.10.attention.heads.2.query.weight: 8192\n",
      "layers.10.attention.heads.2.key.weight: 8192\n",
      "layers.10.attention.heads.2.value.weight: 8192\n",
      "layers.10.attention.heads.3.query.weight: 8192\n",
      "layers.10.attention.heads.3.key.weight: 8192\n",
      "layers.10.attention.heads.3.value.weight: 8192\n",
      "layers.10.attention.heads.4.query.weight: 8192\n",
      "layers.10.attention.heads.4.key.weight: 8192\n",
      "layers.10.attention.heads.4.value.weight: 8192\n",
      "layers.10.attention.heads.5.query.weight: 8192\n",
      "layers.10.attention.heads.5.key.weight: 8192\n",
      "layers.10.attention.heads.5.value.weight: 8192\n",
      "layers.10.attention.heads.6.query.weight: 8192\n",
      "layers.10.attention.heads.6.key.weight: 8192\n",
      "layers.10.attention.heads.6.value.weight: 8192\n",
      "layers.10.attention.heads.7.query.weight: 8192\n",
      "layers.10.attention.heads.7.key.weight: 8192\n",
      "layers.10.attention.heads.7.value.weight: 8192\n",
      "layers.11.attention.heads.0.query.weight: 8192\n",
      "layers.11.attention.heads.0.key.weight: 8192\n",
      "layers.11.attention.heads.0.value.weight: 8192\n",
      "layers.11.attention.heads.1.query.weight: 8192\n",
      "layers.11.attention.heads.1.key.weight: 8192\n",
      "layers.11.attention.heads.1.value.weight: 8192\n",
      "layers.11.attention.heads.2.query.weight: 8192\n",
      "layers.11.attention.heads.2.key.weight: 8192\n",
      "layers.11.attention.heads.2.value.weight: 8192\n",
      "layers.11.attention.heads.3.query.weight: 8192\n",
      "layers.11.attention.heads.3.key.weight: 8192\n",
      "layers.11.attention.heads.3.value.weight: 8192\n",
      "layers.11.attention.heads.4.query.weight: 8192\n",
      "layers.11.attention.heads.4.key.weight: 8192\n",
      "layers.11.attention.heads.4.value.weight: 8192\n",
      "layers.11.attention.heads.5.query.weight: 8192\n",
      "layers.11.attention.heads.5.key.weight: 8192\n",
      "layers.11.attention.heads.5.value.weight: 8192\n",
      "layers.11.attention.heads.6.query.weight: 8192\n",
      "layers.11.attention.heads.6.key.weight: 8192\n",
      "layers.11.attention.heads.6.value.weight: 8192\n",
      "layers.11.attention.heads.7.query.weight: 8192\n",
      "layers.11.attention.heads.7.key.weight: 8192\n",
      "layers.11.attention.heads.7.value.weight: 8192\n",
      "block.0.attention.heads.0.query.weight: 8192\n",
      "block.0.attention.heads.0.key.weight: 8192\n",
      "block.0.attention.heads.0.value.weight: 8192\n",
      "block.0.attention.heads.1.query.weight: 8192\n",
      "block.0.attention.heads.1.key.weight: 8192\n",
      "block.0.attention.heads.1.value.weight: 8192\n",
      "block.0.attention.heads.2.query.weight: 8192\n",
      "block.0.attention.heads.2.key.weight: 8192\n",
      "block.0.attention.heads.2.value.weight: 8192\n",
      "block.0.attention.heads.3.query.weight: 8192\n",
      "block.0.attention.heads.3.key.weight: 8192\n",
      "block.0.attention.heads.3.value.weight: 8192\n",
      "block.0.attention.heads.4.query.weight: 8192\n",
      "block.0.attention.heads.4.key.weight: 8192\n",
      "block.0.attention.heads.4.value.weight: 8192\n",
      "block.0.attention.heads.5.query.weight: 8192\n",
      "block.0.attention.heads.5.key.weight: 8192\n",
      "block.0.attention.heads.5.value.weight: 8192\n",
      "block.0.attention.heads.6.query.weight: 8192\n",
      "block.0.attention.heads.6.key.weight: 8192\n",
      "block.0.attention.heads.6.value.weight: 8192\n",
      "block.0.attention.heads.7.query.weight: 8192\n",
      "block.0.attention.heads.7.key.weight: 8192\n",
      "block.0.attention.heads.7.value.weight: 8192\n",
      "layers.0.ff.net.0.bias: 768\n",
      "layers.1.ff.net.0.bias: 768\n",
      "layers.2.ff.net.0.bias: 768\n",
      "layers.3.ff.net.0.bias: 768\n",
      "layers.4.ff.net.0.bias: 768\n",
      "layers.5.ff.net.0.bias: 768\n",
      "layers.6.ff.net.0.bias: 768\n",
      "layers.7.ff.net.0.bias: 768\n",
      "layers.8.ff.net.0.bias: 768\n",
      "layers.9.ff.net.0.bias: 768\n",
      "layers.10.ff.net.0.bias: 768\n",
      "layers.11.ff.net.0.bias: 768\n",
      "block.0.ff.net.0.bias: 768\n",
      "layers.0.attention.combine_heads.bias: 256\n",
      "layers.0.ff.net.2.bias: 256\n",
      "layers.0.norm1.gamma: 256\n",
      "layers.0.norm1.beta: 256\n",
      "layers.0.norm2.gamma: 256\n",
      "layers.0.norm2.beta: 256\n",
      "layers.1.attention.combine_heads.bias: 256\n",
      "layers.1.ff.net.2.bias: 256\n",
      "layers.1.norm1.gamma: 256\n",
      "layers.1.norm1.beta: 256\n",
      "layers.1.norm2.gamma: 256\n",
      "layers.1.norm2.beta: 256\n",
      "layers.2.attention.combine_heads.bias: 256\n",
      "layers.2.ff.net.2.bias: 256\n",
      "layers.2.norm1.gamma: 256\n",
      "layers.2.norm1.beta: 256\n",
      "layers.2.norm2.gamma: 256\n",
      "layers.2.norm2.beta: 256\n",
      "layers.3.attention.combine_heads.bias: 256\n",
      "layers.3.ff.net.2.bias: 256\n",
      "layers.3.norm1.gamma: 256\n",
      "layers.3.norm1.beta: 256\n",
      "layers.3.norm2.gamma: 256\n",
      "layers.3.norm2.beta: 256\n",
      "layers.4.attention.combine_heads.bias: 256\n",
      "layers.4.ff.net.2.bias: 256\n",
      "layers.4.norm1.gamma: 256\n",
      "layers.4.norm1.beta: 256\n",
      "layers.4.norm2.gamma: 256\n",
      "layers.4.norm2.beta: 256\n",
      "layers.5.attention.combine_heads.bias: 256\n",
      "layers.5.ff.net.2.bias: 256\n",
      "layers.5.norm1.gamma: 256\n",
      "layers.5.norm1.beta: 256\n",
      "layers.5.norm2.gamma: 256\n",
      "layers.5.norm2.beta: 256\n",
      "layers.6.attention.combine_heads.bias: 256\n",
      "layers.6.ff.net.2.bias: 256\n",
      "layers.6.norm1.gamma: 256\n",
      "layers.6.norm1.beta: 256\n",
      "layers.6.norm2.gamma: 256\n",
      "layers.6.norm2.beta: 256\n",
      "layers.7.attention.combine_heads.bias: 256\n",
      "layers.7.ff.net.2.bias: 256\n",
      "layers.7.norm1.gamma: 256\n",
      "layers.7.norm1.beta: 256\n",
      "layers.7.norm2.gamma: 256\n",
      "layers.7.norm2.beta: 256\n",
      "layers.8.attention.combine_heads.bias: 256\n",
      "layers.8.ff.net.2.bias: 256\n",
      "layers.8.norm1.gamma: 256\n",
      "layers.8.norm1.beta: 256\n",
      "layers.8.norm2.gamma: 256\n",
      "layers.8.norm2.beta: 256\n",
      "layers.9.attention.combine_heads.bias: 256\n",
      "layers.9.ff.net.2.bias: 256\n",
      "layers.9.norm1.gamma: 256\n",
      "layers.9.norm1.beta: 256\n",
      "layers.9.norm2.gamma: 256\n",
      "layers.9.norm2.beta: 256\n",
      "layers.10.attention.combine_heads.bias: 256\n",
      "layers.10.ff.net.2.bias: 256\n",
      "layers.10.norm1.gamma: 256\n",
      "layers.10.norm1.beta: 256\n",
      "layers.10.norm2.gamma: 256\n",
      "layers.10.norm2.beta: 256\n",
      "layers.11.attention.combine_heads.bias: 256\n",
      "layers.11.ff.net.2.bias: 256\n",
      "layers.11.norm1.gamma: 256\n",
      "layers.11.norm1.beta: 256\n",
      "layers.11.norm2.gamma: 256\n",
      "layers.11.norm2.beta: 256\n",
      "block.0.attention.combine_heads.bias: 256\n",
      "block.0.ff.net.2.bias: 256\n",
      "block.0.norm1.gamma: 256\n",
      "block.0.norm1.beta: 256\n",
      "block.0.norm2.gamma: 256\n",
      "block.0.norm2.beta: 256\n"
     ]
    }
   ],
   "source": [
    "# get the number of parameters\n",
    "n_params = sum(p.numel() for p in model.parameters())\n",
    "parameter_to_data_ratio = n_params / len(train_data)\n",
    "print(f\"{parameter_to_data_ratio=}\")\n",
    "\n",
    "parameters = []\n",
    "for name, param in model.named_parameters():\n",
    "    parameters.append({\"name\": name, \"params\": param.numel()})\n",
    "\n",
    "# sort parameters by size\n",
    "sorted_parameters = sorted(parameters, key=lambda x: x[\"params\"], reverse=True)\n",
    "for p in sorted_parameters:\n",
    "    print(f\"{p['name']}: {p['params']}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 1/60000 [00:02<40:54:15,  2.45s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.23402966558933258, 'val': 0.2347315102815628}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  1%|          | 301/60000 [01:58<17:00:04,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.23102733492851257, 'val': 0.23457872867584229}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  1%|          | 601/60000 [03:55<17:05:31,  1.04s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.23514676094055176, 'val': 0.23262785375118256}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  2%|▏         | 901/60000 [05:52<16:59:11,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.22860178351402283, 'val': 0.23220586776733398}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  2%|▏         | 1201/60000 [07:49<16:50:35,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.23102501034736633, 'val': 0.2318619340658188}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  3%|▎         | 1501/60000 [09:45<16:44:23,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.23200327157974243, 'val': 0.22724390029907227}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  3%|▎         | 1801/60000 [11:42<16:39:45,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.22791123390197754, 'val': 0.2276018112897873}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  4%|▎         | 2101/60000 [13:39<16:30:17,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.2298070788383484, 'val': 0.22448444366455078}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  4%|▍         | 2401/60000 [15:35<16:24:02,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.2273319810628891, 'val': 0.23055371642112732}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  5%|▍         | 2701/60000 [17:31<16:18:46,  1.02s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.23374542593955994, 'val': 0.2308010309934616}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  5%|▌         | 3000/60000 [19:25<6:00:45,  2.63it/s] "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.2305615097284317, 'val': 0.23260973393917084}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  6%|▌         | 3301/60000 [21:24<16:08:51,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.23247721791267395, 'val': 0.22675520181655884}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  6%|▌         | 3601/60000 [23:21<16:28:11,  1.05s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.2265218049287796, 'val': 0.2333671897649765}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  7%|▋         | 3901/60000 [25:19<16:10:22,  1.04s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.22779016196727753, 'val': 0.22945493459701538}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  7%|▋         | 4201/60000 [27:17<16:04:49,  1.04s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.22743047773838043, 'val': 0.22484704852104187}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  8%|▊         | 4501/60000 [29:14<15:59:38,  1.04s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.22629421949386597, 'val': 0.21891817450523376}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  8%|▊         | 4801/60000 [31:12<15:52:58,  1.04s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.22761407494544983, 'val': 0.22414278984069824}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  9%|▊         | 5101/60000 [33:09<15:47:15,  1.04s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.22607861459255219, 'val': 0.22540126740932465}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  9%|▉         | 5401/60000 [35:11<15:45:33,  1.04s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.22323814034461975, 'val': 0.22532668709754944}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 10%|▉         | 5701/60000 [37:08<15:35:57,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.22388917207717896, 'val': 0.22756782174110413}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 10%|█         | 6000/60000 [39:03<5:45:51,  2.60it/s] "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.2242693305015564, 'val': 0.23001538217067719}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 11%|█         | 6301/60000 [41:03<15:30:44,  1.04s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.22754281759262085, 'val': 0.22606401145458221}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 11%|█         | 6601/60000 [43:01<15:22:04,  1.04s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.2241056114435196, 'val': 0.2206396609544754}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 12%|█▏        | 6901/60000 [44:58<15:18:24,  1.04s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.22395986318588257, 'val': 0.22235983610153198}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 12%|█▏        | 7201/60000 [46:56<15:09:02,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.22518014907836914, 'val': 0.2252006232738495}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 13%|█▎        | 7501/60000 [48:53<15:06:28,  1.04s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.22738833725452423, 'val': 0.2227017730474472}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 13%|█▎        | 7801/60000 [50:50<14:59:48,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.22759272158145905, 'val': 0.23072466254234314}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 14%|█▎        | 8101/60000 [52:48<14:55:14,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.22645916044712067, 'val': 0.2266225963830948}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 14%|█▍        | 8401/60000 [54:45<14:51:01,  1.04s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.2240157574415207, 'val': 0.22298069298267365}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 15%|█▍        | 8701/60000 [56:42<14:42:20,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.2198890745639801, 'val': 0.22396335005760193}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 15%|█▌        | 9000/60000 [58:37<5:26:30,  2.60it/s] "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.22803305089473724, 'val': 0.2280309647321701}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 16%|█▌        | 9301/60000 [1:00:37<14:35:46,  1.04s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.22462260723114014, 'val': 0.22183598577976227}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 16%|█▌        | 9601/60000 [1:02:34<14:25:31,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.22535118460655212, 'val': 0.2275611013174057}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 17%|█▋        | 9901/60000 [1:04:32<14:25:13,  1.04s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.2175770252943039, 'val': 0.2227979451417923}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 17%|█▋        | 10201/60000 [1:06:29<14:19:23,  1.04s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.22116100788116455, 'val': 0.22420012950897217}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 18%|█▊        | 10501/60000 [1:08:26<14:17:01,  1.04s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.2217835932970047, 'val': 0.22787146270275116}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 18%|█▊        | 10801/60000 [1:10:23<14:05:45,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.22239038348197937, 'val': 0.22207090258598328}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 19%|█▊        | 11101/60000 [1:12:21<14:07:02,  1.04s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.22589954733848572, 'val': 0.22363737225532532}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 19%|█▉        | 11401/60000 [1:14:18<13:59:18,  1.04s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.22317910194396973, 'val': 0.22500072419643402}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 20%|█▉        | 11701/60000 [1:16:15<13:52:15,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.22546085715293884, 'val': 0.22057278454303741}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 20%|██        | 12000/60000 [1:18:09<5:07:45,  2.60it/s] "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.22139465808868408, 'val': 0.22054219245910645}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 21%|██        | 12301/60000 [1:20:10<13:46:42,  1.04s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.2189365029335022, 'val': 0.22736552357673645}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 21%|██        | 12601/60000 [1:22:08<13:36:07,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.2203863263130188, 'val': 0.22525236010551453}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 22%|██▏       | 12901/60000 [1:24:05<13:32:39,  1.04s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.223595529794693, 'val': 0.2181542068719864}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 22%|██▏       | 13201/60000 [1:26:03<13:27:24,  1.04s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.2223498821258545, 'val': 0.22010686993598938}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 23%|██▎       | 13501/60000 [1:28:00<13:19:49,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.21735262870788574, 'val': 0.22116126120090485}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 23%|██▎       | 13801/60000 [1:29:57<13:16:54,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.2216835767030716, 'val': 0.22128503024578094}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 24%|██▎       | 14101/60000 [1:31:54<13:09:01,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.22352482378482819, 'val': 0.22083450853824615}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 24%|██▍       | 14401/60000 [1:33:51<13:05:25,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.21935327351093292, 'val': 0.2197258472442627}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 25%|██▍       | 14701/60000 [1:35:48<13:02:55,  1.04s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.22142884135246277, 'val': 0.22097070515155792}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 25%|██▌       | 15000/60000 [1:37:43<4:48:06,  2.60it/s] "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.21535253524780273, 'val': 0.22567224502563477}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 26%|██▌       | 15301/60000 [1:39:43<12:52:33,  1.04s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.22048071026802063, 'val': 0.22459924221038818}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 26%|██▌       | 15601/60000 [1:41:40<12:47:08,  1.04s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.21865606307983398, 'val': 0.221701979637146}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 27%|██▋       | 15901/60000 [1:43:38<12:40:31,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.21715882420539856, 'val': 0.22263969480991364}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 27%|██▋       | 16201/60000 [1:45:35<12:36:37,  1.04s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.22185018658638, 'val': 0.2247619926929474}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 28%|██▊       | 16501/60000 [1:47:32<12:28:03,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.22173704206943512, 'val': 0.22425131499767303}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 28%|██▊       | 16801/60000 [1:49:29<12:23:56,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.2222270518541336, 'val': 0.2256855070590973}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 29%|██▊       | 17101/60000 [1:51:26<12:18:29,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.22454912960529327, 'val': 0.21832899749279022}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 29%|██▉       | 17401/60000 [1:53:23<12:11:09,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.22053129971027374, 'val': 0.21757467091083527}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 30%|██▉       | 17701/60000 [1:55:20<12:09:46,  1.04s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.218051016330719, 'val': 0.22109462320804596}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 30%|███       | 18001/60000 [1:57:18<12:43:36,  1.09s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.21668575704097748, 'val': 0.2159418910741806}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 31%|███       | 18301/60000 [1:59:15<11:58:00,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.2159796953201294, 'val': 0.21854905784130096}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 31%|███       | 18601/60000 [2:01:12<11:53:54,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.21561938524246216, 'val': 0.22615666687488556}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 32%|███▏      | 18901/60000 [2:03:09<11:46:38,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.2206479012966156, 'val': 0.22147218883037567}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 32%|███▏      | 19201/60000 [2:05:06<11:39:22,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.21714197099208832, 'val': 0.21793168783187866}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 33%|███▎      | 19501/60000 [2:07:02<11:38:42,  1.04s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.22195585072040558, 'val': 0.21947044134140015}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 33%|███▎      | 19801/60000 [2:08:59<11:29:59,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.21641582250595093, 'val': 0.21723683178424835}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 34%|███▎      | 20101/60000 [2:10:56<11:22:43,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.2174321413040161, 'val': 0.21760499477386475}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 34%|███▍      | 20401/60000 [2:12:53<11:19:49,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.2166900634765625, 'val': 0.22116951644420624}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 35%|███▍      | 20701/60000 [2:14:51<11:26:50,  1.05s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.2167970985174179, 'val': 0.22283132374286652}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 35%|███▌      | 21001/60000 [2:16:49<11:48:19,  1.09s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.21980369091033936, 'val': 0.22174105048179626}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 36%|███▌      | 21301/60000 [2:18:46<11:07:53,  1.04s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.22131116688251495, 'val': 0.2174418419599533}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 36%|███▌      | 21601/60000 [2:20:44<10:57:18,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.21663831174373627, 'val': 0.217766672372818}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 37%|███▋      | 21901/60000 [2:22:41<10:57:09,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.21636193990707397, 'val': 0.21460787951946259}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 37%|███▋      | 22201/60000 [2:24:38<10:51:39,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.21644411981105804, 'val': 0.21354049444198608}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 38%|███▊      | 22501/60000 [2:26:35<10:46:49,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.2211376577615738, 'val': 0.22263997793197632}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 38%|███▊      | 22801/60000 [2:28:32<10:39:50,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.21899563074111938, 'val': 0.21904589235782623}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 39%|███▊      | 23101/60000 [2:30:30<10:37:36,  1.04s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.22140035033226013, 'val': 0.22070683538913727}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 39%|███▉      | 23401/60000 [2:32:27<10:24:43,  1.02s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.21216653287410736, 'val': 0.22125546634197235}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 40%|███▉      | 23701/60000 [2:34:23<10:24:58,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.2125297337770462, 'val': 0.21790200471878052}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 40%|████      | 24001/60000 [2:36:20<10:48:39,  1.08s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.21204355359077454, 'val': 0.2202649563550949}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 41%|████      | 24301/60000 [2:38:17<10:13:34,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.21106477081775665, 'val': 0.2110428512096405}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 41%|████      | 24601/60000 [2:40:14<10:07:14,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.22030699253082275, 'val': 0.22033578157424927}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 42%|████▏     | 24901/60000 [2:42:11<9:59:00,  1.02s/it] "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.21886534988880157, 'val': 0.21671023964881897}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 42%|████▏     | 25201/60000 [2:44:08<9:57:48,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.2150259017944336, 'val': 0.21558447182178497}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 43%|████▎     | 25501/60000 [2:46:05<9:52:40,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.2199040651321411, 'val': 0.21858356893062592}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 43%|████▎     | 25801/60000 [2:48:02<9:48:26,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.2181493490934372, 'val': 0.21897730231285095}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 44%|████▎     | 26101/60000 [2:49:59<9:43:25,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.21698586642742157, 'val': 0.21872881054878235}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 44%|████▍     | 26401/60000 [2:51:56<9:37:31,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.21564103662967682, 'val': 0.21340948343276978}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 45%|████▍     | 26701/60000 [2:53:53<9:34:19,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.21633298695087433, 'val': 0.21642987430095673}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 45%|████▌     | 27001/60000 [2:55:50<9:55:36,  1.08s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.21886643767356873, 'val': 0.21931809186935425}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 46%|████▌     | 27301/60000 [2:57:46<9:20:56,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.21687261760234833, 'val': 0.21577580273151398}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 46%|████▌     | 27601/60000 [2:59:43<9:14:02,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.2198447287082672, 'val': 0.22071872651576996}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 47%|████▋     | 27901/60000 [3:01:40<9:09:53,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.2165859341621399, 'val': 0.21137981116771698}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 47%|████▋     | 28201/60000 [3:03:37<9:04:41,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.2148956060409546, 'val': 0.2206512838602066}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 48%|████▊     | 28501/60000 [3:05:34<9:02:27,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.21614733338356018, 'val': 0.2164343297481537}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 48%|████▊     | 28801/60000 [3:07:31<8:55:24,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.21390385925769806, 'val': 0.21845579147338867}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 49%|████▊     | 29101/60000 [3:09:28<8:50:02,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.2187126725912094, 'val': 0.22018644213676453}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 49%|████▉     | 29401/60000 [3:11:25<8:46:49,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.2150123417377472, 'val': 0.21930420398712158}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 50%|████▉     | 29701/60000 [3:13:22<8:41:00,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.2159375697374344, 'val': 0.21935001015663147}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 50%|█████     | 30001/60000 [3:15:19<9:02:36,  1.09s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.2160373479127884, 'val': 0.21181191504001617}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 51%|█████     | 30301/60000 [3:17:15<8:29:34,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.21590515971183777, 'val': 0.21907760202884674}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 51%|█████     | 30601/60000 [3:19:12<8:24:02,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.2152610719203949, 'val': 0.21843081712722778}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 52%|█████▏    | 30901/60000 [3:21:08<8:20:29,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.2185816913843155, 'val': 0.2144077867269516}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 52%|█████▏    | 31201/60000 [3:23:05<8:14:45,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.21750035881996155, 'val': 0.21540626883506775}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 53%|█████▎    | 31501/60000 [3:25:02<8:11:47,  1.04s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.21398994326591492, 'val': 0.21754899621009827}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 53%|█████▎    | 31801/60000 [3:26:59<8:03:43,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.21481259167194366, 'val': 0.21224753558635712}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 54%|█████▎    | 32101/60000 [3:28:56<7:58:48,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.21857668459415436, 'val': 0.21815378963947296}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 54%|█████▍    | 32401/60000 [3:30:54<7:55:03,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.21311573684215546, 'val': 0.2143121212720871}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 55%|█████▍    | 32701/60000 [3:32:50<7:49:01,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.21372349560260773, 'val': 0.21655850112438202}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 55%|█████▌    | 33001/60000 [3:34:47<8:08:15,  1.09s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.21449893712997437, 'val': 0.2158009558916092}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 56%|█████▌    | 33301/60000 [3:36:44<7:38:41,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.20933601260185242, 'val': 0.2223808467388153}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 56%|█████▌    | 33601/60000 [3:38:41<7:32:13,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.21703845262527466, 'val': 0.2172713279724121}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 57%|█████▋    | 33901/60000 [3:40:38<7:26:01,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.21583184599876404, 'val': 0.21799267828464508}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 57%|█████▋    | 34201/60000 [3:42:35<7:23:51,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.21374750137329102, 'val': 0.21263626217842102}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 58%|█████▊    | 34501/60000 [3:44:32<7:17:38,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.2115737795829773, 'val': 0.2191653549671173}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 58%|█████▊    | 34801/60000 [3:46:28<7:12:07,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.21611818671226501, 'val': 0.2129705846309662}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 59%|█████▊    | 35101/60000 [3:48:25<7:06:38,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.21639306843280792, 'val': 0.21427667140960693}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 59%|█████▉    | 35401/60000 [3:50:22<7:03:08,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.213717982172966, 'val': 0.21559777855873108}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 60%|█████▉    | 35701/60000 [3:52:19<6:58:17,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.21386462450027466, 'val': 0.22001028060913086}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 60%|██████    | 36001/60000 [3:54:16<7:11:24,  1.08s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.2101926952600479, 'val': 0.21392837166786194}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 61%|██████    | 36301/60000 [3:56:13<6:49:00,  1.04s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.21436835825443268, 'val': 0.21068131923675537}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 61%|██████    | 36601/60000 [3:58:10<6:42:13,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.21814747154712677, 'val': 0.2107490748167038}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 62%|██████▏   | 36901/60000 [4:00:07<6:38:49,  1.04s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.21134644746780396, 'val': 0.21475271880626678}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 62%|██████▏   | 37201/60000 [4:02:04<6:32:24,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.21544675529003143, 'val': 0.21514268219470978}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 63%|██████▎   | 37501/60000 [4:04:01<6:27:58,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.21239422261714935, 'val': 0.2137480229139328}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 63%|██████▎   | 37801/60000 [4:05:59<6:21:58,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.21066300570964813, 'val': 0.2147558182477951}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 64%|██████▎   | 38101/60000 [4:07:56<6:16:51,  1.03s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.21533919870853424, 'val': 0.21453677117824554}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 64%|██████▍   | 38401/60000 [4:09:53<6:12:55,  1.04s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'tIain': 0.21800558269023895, 'val': 0.21224913001060486}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 64%|██████▍   | 38585/60000 [4:11:04<2:19:21,  2.56it/s]\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[87], line 15\u001b[0m\n\u001b[1;32m     11\u001b[0m \u001b[38;5;66;03m# l2 regularization\u001b[39;00m\n\u001b[1;32m     12\u001b[0m \u001b[38;5;66;03m# l2 = sum(p.pow(2).sum() for p in model.parameters()) / num_params\u001b[39;00m\n\u001b[1;32m     13\u001b[0m loss \u001b[38;5;241m=\u001b[39m loss \u001b[38;5;66;03m# + l2 * l2_penalty\u001b[39;00m\n\u001b[0;32m---> 15\u001b[0m \u001b[43mloss\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbackward\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     16\u001b[0m optimizer\u001b[38;5;241m.\u001b[39mstep()\n\u001b[1;32m     17\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m steps \u001b[38;5;241m%\u001b[39m eval_interval \u001b[38;5;241m==\u001b[39m \u001b[38;5;241m0\u001b[39m:\n",
      "File \u001b[0;32m/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/_tensor.py:513\u001b[0m, in \u001b[0;36mTensor.backward\u001b[0;34m(self, gradient, retain_graph, create_graph, inputs)\u001b[0m\n\u001b[1;32m    466\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124mr\u001b[39m\u001b[38;5;124;03m\"\"\"Computes the gradient of current tensor wrt graph leaves.\u001b[39;00m\n\u001b[1;32m    467\u001b[0m \n\u001b[1;32m    468\u001b[0m \u001b[38;5;124;03mThe graph is differentiated using the chain rule. If the tensor is\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    510\u001b[0m \u001b[38;5;124;03m        used to compute the attr::tensors.\u001b[39;00m\n\u001b[1;32m    511\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    512\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m has_torch_function_unary(\u001b[38;5;28mself\u001b[39m):\n\u001b[0;32m--> 513\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mhandle_torch_function\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    514\u001b[0m \u001b[43m        \u001b[49m\u001b[43mTensor\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbackward\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    515\u001b[0m \u001b[43m        \u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    516\u001b[0m \u001b[43m        \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m    517\u001b[0m \u001b[43m        \u001b[49m\u001b[43mgradient\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mgradient\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    518\u001b[0m \u001b[43m        \u001b[49m\u001b[43mretain_graph\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mretain_graph\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    519\u001b[0m \u001b[43m        \u001b[49m\u001b[43mcreate_graph\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mcreate_graph\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    520\u001b[0m \u001b[43m        \u001b[49m\u001b[43minputs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43minputs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    521\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    522\u001b[0m torch\u001b[38;5;241m.\u001b[39mautograd\u001b[38;5;241m.\u001b[39mbackward(\n\u001b[1;32m    523\u001b[0m     \u001b[38;5;28mself\u001b[39m, gradient, retain_graph, create_graph, inputs\u001b[38;5;241m=\u001b[39minputs\n\u001b[1;32m    524\u001b[0m )\n",
      "File \u001b[0;32m/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/overrides.py:1604\u001b[0m, in \u001b[0;36mhandle_torch_function\u001b[0;34m(public_api, relevant_args, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1600\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m _is_torch_function_mode_enabled():\n\u001b[1;32m   1601\u001b[0m     \u001b[38;5;66;03m# if we're here, the mode must be set to a TorchFunctionStackMode\u001b[39;00m\n\u001b[1;32m   1602\u001b[0m     \u001b[38;5;66;03m# this unsets it and calls directly into TorchFunctionStackMode's torch function\u001b[39;00m\n\u001b[1;32m   1603\u001b[0m     \u001b[38;5;28;01mwith\u001b[39;00m _pop_mode_temporarily() \u001b[38;5;28;01mas\u001b[39;00m mode:\n\u001b[0;32m-> 1604\u001b[0m         result \u001b[38;5;241m=\u001b[39m \u001b[43mmode\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m__torch_function__\u001b[49m\u001b[43m(\u001b[49m\u001b[43mpublic_api\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtypes\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1605\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m result \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28mNotImplemented\u001b[39m:\n\u001b[1;32m   1606\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m result\n",
      "File \u001b[0;32m/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/utils/_device.py:77\u001b[0m, in \u001b[0;36mDeviceContext.__torch_function__\u001b[0;34m(self, func, types, args, kwargs)\u001b[0m\n\u001b[1;32m     75\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m func \u001b[38;5;129;01min\u001b[39;00m _device_constructors() \u001b[38;5;129;01mand\u001b[39;00m kwargs\u001b[38;5;241m.\u001b[39mget(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mdevice\u001b[39m\u001b[38;5;124m'\u001b[39m) \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m     76\u001b[0m     kwargs[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mdevice\u001b[39m\u001b[38;5;124m'\u001b[39m] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdevice\n\u001b[0;32m---> 77\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mfunc\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/_tensor.py:522\u001b[0m, in \u001b[0;36mTensor.backward\u001b[0;34m(self, gradient, retain_graph, create_graph, inputs)\u001b[0m\n\u001b[1;32m    512\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m has_torch_function_unary(\u001b[38;5;28mself\u001b[39m):\n\u001b[1;32m    513\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m handle_torch_function(\n\u001b[1;32m    514\u001b[0m         Tensor\u001b[38;5;241m.\u001b[39mbackward,\n\u001b[1;32m    515\u001b[0m         (\u001b[38;5;28mself\u001b[39m,),\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    520\u001b[0m         inputs\u001b[38;5;241m=\u001b[39minputs,\n\u001b[1;32m    521\u001b[0m     )\n\u001b[0;32m--> 522\u001b[0m \u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mautograd\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbackward\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    523\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mgradient\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mretain_graph\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcreate_graph\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43minputs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43minputs\u001b[49m\n\u001b[1;32m    524\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/autograd/__init__.py:266\u001b[0m, in \u001b[0;36mbackward\u001b[0;34m(tensors, grad_tensors, retain_graph, create_graph, grad_variables, inputs)\u001b[0m\n\u001b[1;32m    261\u001b[0m     retain_graph \u001b[38;5;241m=\u001b[39m create_graph\n\u001b[1;32m    263\u001b[0m \u001b[38;5;66;03m# The reason we repeat the same comment below is that\u001b[39;00m\n\u001b[1;32m    264\u001b[0m \u001b[38;5;66;03m# some Python versions print out the first line of a multi-line function\u001b[39;00m\n\u001b[1;32m    265\u001b[0m \u001b[38;5;66;03m# calls in the traceback and some print out the last line\u001b[39;00m\n\u001b[0;32m--> 266\u001b[0m \u001b[43mVariable\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_execution_engine\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrun_backward\u001b[49m\u001b[43m(\u001b[49m\u001b[43m  \u001b[49m\u001b[38;5;66;43;03m# Calls into the C++ engine to run the backward pass\u001b[39;49;00m\n\u001b[1;32m    267\u001b[0m \u001b[43m    \u001b[49m\u001b[43mtensors\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    268\u001b[0m \u001b[43m    \u001b[49m\u001b[43mgrad_tensors_\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    269\u001b[0m \u001b[43m    \u001b[49m\u001b[43mretain_graph\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    270\u001b[0m \u001b[43m    \u001b[49m\u001b[43mcreate_graph\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    271\u001b[0m \u001b[43m    \u001b[49m\u001b[43minputs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    272\u001b[0m \u001b[43m    \u001b[49m\u001b[43mallow_unreachable\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[1;32m    273\u001b[0m \u001b[43m    \u001b[49m\u001b[43maccumulate_grad\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[1;32m    274\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "import tqdm\n",
    "num_params = sum([p.numel() for p in model.parameters()])\n",
    "\n",
    "dump_model_interval = 1000\n",
    "\n",
    "for steps in tqdm.tqdm(range(max_iters)):\n",
    "    xb, yb = get_batch('train')\n",
    "    # loss\n",
    "    logits, loss = model(xb, yb)\n",
    "    optimizer.zero_grad(set_to_none=True)\n",
    "    # l2 regularization\n",
    "    # l2 = sum(p.pow(2).sum() for p in model.parameters()) / num_params\n",
    "    loss = loss # + l2 * l2_penalty\n",
    "\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    if steps % eval_interval == 0:\n",
    "        losses = estimate_loss()\n",
    "        # wandb.log({\"train\": losses['train'].item(), \"val\": losses['val'].item(), \"l2\":l2})\n",
    "        print({\"tIain\": losses['train'].item(), \"val\": losses['val'].item()})\n",
    "    if steps % dump_model_interval == 0 and steps > 0:\n",
    "        model_no = steps // dump_model_interval\n",
    "        torch.save(model.state_dict(), f'tiny-stories-model-{model_no}.pt')\n",
    "\n",
    "losses = estimate_loss(is_last=True)\n",
    "#wandb.log({\"train\": losses['train'].item(), \"val\": losses['val'].item()})\n",
    "#wandb.finish()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'train': tensor(0.4513, device='cuda:0'),\n",
       " 'val': tensor(0.4492, device='cuda:0')}"
      ]
     },
     "execution_count": 74,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "estimate_loss()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save model\n",
    "torch.save(model.state_dict(), 'tiny-stories-model.pt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 52,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# load the model\n",
    "model.load_state_dict(torch.load('tiny-stories-model.pt'))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[198]"
      ]
     },
     "execution_count": 79,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "encode(\"\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 92,
   "metadata": {},
   "outputs": [
    {
     "ename": "RuntimeError",
     "evalue": "CUDA error: device-side assert triggered\nCUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.\nFor debugging consider passing CUDA_LAUNCH_BLOCKING=1.\nCompile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.\n",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[92], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[43mmodel\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mprompt_model\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43m Un día, un niño\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m200\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtemperature\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m0.5\u001b[39;49m\u001b[43m)\u001b[49m)\n",
      "Cell \u001b[0;32mIn[71], line 59\u001b[0m, in \u001b[0;36mGPT.prompt_model\u001b[0;34m(self, prompt, max_new_tokens, temperature)\u001b[0m\n\u001b[1;32m     56\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m _ \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mrange\u001b[39m(max_new_tokens):\n\u001b[1;32m     57\u001b[0m     prediction_index \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mlen\u001b[39m(autoregressive_seq)\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m\n\u001b[0;32m---> 59\u001b[0m     model_input \u001b[38;5;241m=\u001b[39m \u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtensor\u001b[49m\u001b[43m(\u001b[49m\u001b[43mautoregressive_seq\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     61\u001b[0m     \u001b[38;5;28;01mwhile\u001b[39;00m model_input\u001b[38;5;241m.\u001b[39mshape[\u001b[38;5;241m0\u001b[39m] \u001b[38;5;241m<\u001b[39m T:\n\u001b[1;32m     62\u001b[0m         pad_token \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39mtensor(encode(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[38;5;124m\"\u001b[39m))\n",
      "File \u001b[0;32m/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/utils/_device.py:77\u001b[0m, in \u001b[0;36mDeviceContext.__torch_function__\u001b[0;34m(self, func, types, args, kwargs)\u001b[0m\n\u001b[1;32m     75\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m func \u001b[38;5;129;01min\u001b[39;00m _device_constructors() \u001b[38;5;129;01mand\u001b[39;00m kwargs\u001b[38;5;241m.\u001b[39mget(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mdevice\u001b[39m\u001b[38;5;124m'\u001b[39m) \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m     76\u001b[0m     kwargs[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mdevice\u001b[39m\u001b[38;5;124m'\u001b[39m] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdevice\n\u001b[0;32m---> 77\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mfunc\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "\u001b[0;31mRuntimeError\u001b[0m: CUDA error: device-side assert triggered\nCUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.\nFor debugging consider passing CUDA_LAUNCH_BLOCKING=1.\nCompile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.\n"
     ]
    }
   ],
   "source": [
    "print(model.prompt_model(\" Un día, un niño\", 200, temperature=0.5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\n",
      "The first room was filled with joy. He raised wide and wide-bye and skipped on to the chest with a big smile. Every nowe on the same head. Aneagull was the most beautiful and quiet place! Aeon lilyolog opened the window andœOf course clapped in the lush road. Forward was nothing! Today had been great! She was so happy and so excited that she cried out that they couldn't stop until they gave up. She had done it!\"\n",
      "\"Once upon a time there was an angry volcano. Daisy shouted loud and refused to do anything.\n",
      "\n",
      "Then, bowl was filled with her anger from the volcano. She wished that the volcano would be more careful.\n",
      "\n",
      "But, the volcano didn't ignorant. Gootators were too bossy. It wanted to cause trouble and become angry.\"\n",
      "\"Once upon a time, there was a man who wanted to go for a ride. He was so excited! â€œThomas!â€ he shouted. â€œLet's go!â€\n",
      "\n",
      "Thomas and the man got on the bus. The billboard opened. Out of the window light came on Max but he got so excited. The picture in the picture was there all about him.\n"
     ]
    }
   ],
   "source": [
    "test_idx = torch.zeros(1, T).long() * 198\n",
    "print(decode(\n",
    "    model.generate(idx=test_idx, max_new_tokens=C)[0].tolist()\n",
    ")[T:])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}