{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "45398736-7e89-4263-89c8-92153baff553",
   "metadata": {},
   "source": [
    "<table style=\"width:100%\">\n",
    "<tr>\n",
    "<td style=\"vertical-align:middle; text-align:left;\">\n",
    "<font size=\"2\">\n",
    "Supplementary code for the <a href=\"http://mng.bz/orYv\">Build a Large Language Model From Scratch</a> book by <a href=\"https://sebastianraschka.com\">Sebastian Raschka</a><br>\n",
    "<br>Code repository: <a href=\"https://github.com/rasbt/LLMs-from-scratch\">https://github.com/rasbt/LLMs-from-scratch</a>\n",
    "</font>\n",
    "</td>\n",
    "<td style=\"vertical-align:middle; text-align:left;\">\n",
    "<a href=\"http://mng.bz/orYv\"><img src=\"https://sebastianraschka.com/images/LLMs-from-scratch-images/cover-small.webp\" width=\"100px\"></a>\n",
    "</td>\n",
    "</tr>\n",
    "</table>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "66dd524e-864c-4012-b0a2-ccfc56e80024",
   "metadata": {
    "id": "66dd524e-864c-4012-b0a2-ccfc56e80024"
   },
   "source": [
    "# Chapter 5 Bonus: Pretraining Llama 3 on Unlabeled Data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c4fa2aa",
   "metadata": {},
   "source": [
    "- This notebook plugs in the [Llama 3 1B from-scratch](../07_gpt_to_llama/standalone-llama32.ipynb) model into (the pretraining portion) of chapter 5\n",
    "- This is to show how to use Llama 3 1B as a drop-in replacement for the GTP-2 model used in [chapter 5](../01_main-chapter-code/ch05.ipynb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "db3564b7-9940-44fe-9364-27ea71e38632",
   "metadata": {},
   "outputs": [],
   "source": [
    "# pip install tokenizers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "92b989e9-da36-4159-b212-799184764dd9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "matplotlib version: 3.8.2\n",
      "numpy version: 1.26.4\n",
      "tiktoken version: 0.12.0\n",
      "torch version: 2.8.0+cu128\n"
     ]
    }
   ],
   "source": [
    "from importlib.metadata import version\n",
    "\n",
    "pkgs = [\n",
    "    \"matplotlib\", \n",
    "    \"numpy\", \n",
    "    \"tiktoken\", \n",
    "    \"torch\",\n",
    "       ]\n",
    "for p in pkgs:\n",
    "    print(f\"{p} version: {version(p)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d824183-145c-4865-89e1-1f0d0a338f19",
   "metadata": {
    "id": "0d824183-145c-4865-89e1-1f0d0a338f19"
   },
   "source": [
    "&nbsp;\n",
    "## 5.1 Evaluating generative text models"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eb9508e0-4e09-4236-bb07-b376013c219d",
   "metadata": {},
   "source": [
    "- No code"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bdc1cf3f-82d8-46c7-9ecc-58979ce87cdd",
   "metadata": {
    "id": "bdc1cf3f-82d8-46c7-9ecc-58979ce87cdd"
   },
   "source": [
    "&nbsp;\n",
    "### 5.1.1 Using Llama 3 to generate text\n",
    "\n",
    "<img src=\"https://sebastianraschka.com/images/LLMs-from-scratch-images/bonus/gpt-to-llama/llama32.webp\" width=\"700px\">"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86000d74-624a-48f0-86da-f41926cb9e04",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "86000d74-624a-48f0-86da-f41926cb9e04",
    "outputId": "ad482cfd-5a62-4f0d-e1e0-008d6457f512"
   },
   "outputs": [],
   "source": [
    "######################\n",
    "### Llama 3 Code\n",
    "######################\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "\n",
    "class FeedForward(nn.Module):\n",
    "    def __init__(self, cfg):\n",
    "        super().__init__()\n",
    "        self.fc1 = nn.Linear(cfg[\"emb_dim\"], cfg[\"hidden_dim\"], dtype=cfg[\"dtype\"], bias=False)\n",
    "        self.fc2 = nn.Linear(cfg[\"emb_dim\"], cfg[\"hidden_dim\"], dtype=cfg[\"dtype\"], bias=False)\n",
    "        self.fc3 = nn.Linear(cfg[\"hidden_dim\"], cfg[\"emb_dim\"], dtype=cfg[\"dtype\"], bias=False)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x_fc1 = self.fc1(x)\n",
    "        x_fc2 = self.fc2(x)\n",
    "        x = nn.functional.silu(x_fc1) * x_fc2\n",
    "        return self.fc3(x)\n",
    "\n",
    "\n",
    "def compute_rope_params(head_dim, theta_base=10_000, context_length=4096, freq_config=None, dtype=torch.float32):\n",
    "    assert head_dim % 2 == 0, \"Embedding dimension must be even\"\n",
    "\n",
    "    # Compute the inverse frequencies\n",
    "    inv_freq = 1.0 / (theta_base ** (torch.arange(0, head_dim, 2, dtype=dtype)[: (head_dim // 2)].float() / head_dim))\n",
    "\n",
    "    # Frequency adjustments\n",
    "    if freq_config is not None:\n",
    "        low_freq_wavelen = freq_config[\"original_context_length\"] / freq_config[\"low_freq_factor\"]\n",
    "        high_freq_wavelen = freq_config[\"original_context_length\"] / freq_config[\"high_freq_factor\"]\n",
    "\n",
    "        wavelen = 2 * torch.pi / inv_freq\n",
    "\n",
    "        inv_freq_llama = torch.where(\n",
    "            wavelen > low_freq_wavelen, inv_freq / freq_config[\"factor\"], inv_freq\n",
    "        )\n",
    "\n",
    "        smooth_factor = (freq_config[\"original_context_length\"] / wavelen - freq_config[\"low_freq_factor\"]) / (\n",
    "            freq_config[\"high_freq_factor\"] - freq_config[\"low_freq_factor\"]\n",
    "        )\n",
    "\n",
    "        smoothed_inv_freq = (\n",
    "            (1 - smooth_factor) * (inv_freq / freq_config[\"factor\"]) + smooth_factor * inv_freq\n",
    "        )\n",
    "\n",
    "        is_medium_freq = (wavelen <= low_freq_wavelen) & (wavelen >= high_freq_wavelen)\n",
    "        inv_freq_llama = torch.where(is_medium_freq, smoothed_inv_freq, inv_freq_llama)\n",
    "        inv_freq = inv_freq_llama\n",
    "\n",
    "    # Generate position indices\n",
    "    positions = torch.arange(context_length, dtype=dtype)\n",
    "\n",
    "    # Compute the angles\n",
    "    angles = positions.unsqueeze(1) * inv_freq.unsqueeze(0)  # Shape: (context_length, head_dim // 2)\n",
    "\n",
    "    # Expand angles to match the head_dim\n",
    "    angles = torch.cat([angles, angles], dim=1)  # Shape: (context_length, head_dim)\n",
    "\n",
    "    # Precompute sine and cosine\n",
    "    cos = torch.cos(angles)\n",
    "    sin = torch.sin(angles)\n",
    "\n",
    "    return cos, sin\n",
    "\n",
    "\n",
    "def apply_rope(x, cos, sin):\n",
    "    # x: (batch_size, num_heads, seq_len, head_dim)\n",
    "    batch_size, num_heads, seq_len, head_dim = x.shape\n",
    "    assert head_dim % 2 == 0, \"Head dimension must be even\"\n",
    "\n",
    "    # Split x into first half and second half\n",
    "    x1 = x[..., : head_dim // 2]  # First half\n",
    "    x2 = x[..., head_dim // 2 :]  # Second half\n",
    "\n",
    "    # Adjust sin and cos shapes\n",
    "    cos = cos[:seq_len, :].unsqueeze(0).unsqueeze(0)  # Shape: (1, 1, seq_len, head_dim)\n",
    "    sin = sin[:seq_len, :].unsqueeze(0).unsqueeze(0)\n",
    "\n",
    "    # Apply the rotary transformation\n",
    "    rotated = torch.cat((-x2, x1), dim=-1)\n",
    "    x_rotated = (x * cos) + (rotated * sin)\n",
    "\n",
    "    # It's ok to use lower-precision after applying cos and sin rotation\n",
    "    return x_rotated.to(dtype=x.dtype)\n",
    "\n",
    "\n",
    "class GroupedQueryAttention(nn.Module):\n",
    "    def __init__(\n",
    "            self, d_in, d_out, num_heads,\n",
    "            num_kv_groups,\n",
    "            dtype=None\n",
    "        ):\n",
    "        super().__init__()\n",
    "        assert d_out % num_heads == 0, \"d_out must be divisible by num_heads\"\n",
    "        assert num_heads % num_kv_groups == 0, \"num_heads must be divisible by num_kv_groups\"\n",
    "\n",
    "        self.d_out = d_out\n",
    "        self.num_heads = num_heads\n",
    "        self.head_dim = d_out // num_heads\n",
    "\n",
    "        self.W_key = nn.Linear(d_in, num_kv_groups * self.head_dim, bias=False, dtype=dtype)\n",
    "        self.W_value = nn.Linear(d_in, num_kv_groups * self.head_dim, bias=False, dtype=dtype)\n",
    "        self.num_kv_groups = num_kv_groups\n",
    "        self.group_size = num_heads // num_kv_groups\n",
    "\n",
    "        self.W_query = nn.Linear(d_in, d_out, bias=False, dtype=dtype)\n",
    "        self.out_proj = nn.Linear(d_out, d_out, bias=False, dtype=dtype)\n",
    "\n",
    "    def forward(self, x, mask, cos, sin):\n",
    "        b, num_tokens, d_in = x.shape\n",
    "\n",
    "        queries = self.W_query(x)  # Shape: (b, num_tokens, d_out)\n",
    "        keys = self.W_key(x)  # Shape: (b, num_tokens, num_kv_groups * head_dim)\n",
    "        values = self.W_value(x)  # Shape: (b, num_tokens, num_kv_groups * head_dim)\n",
    "\n",
    "        # Reshape queries, keys, and values\n",
    "        queries = queries.view(b, num_tokens, self.num_heads, self.head_dim)\n",
    "        keys = keys.view(b, num_tokens, self.num_kv_groups, self.head_dim)\n",
    "        values = values.view(b, num_tokens, self.num_kv_groups, self.head_dim)\n",
    "\n",
    "        # Transpose keys, values, and queries\n",
    "        keys = keys.transpose(1, 2)  # Shape: (b, num_kv_groups, num_tokens, head_dim)\n",
    "        values = values.transpose(1, 2)  # Shape: (b, num_kv_groups, num_tokens, head_dim)\n",
    "        queries = queries.transpose(1, 2)  # Shape: (b, num_heads, num_tokens, head_dim)\n",
    "\n",
    "        # Apply RoPE\n",
    "        keys = apply_rope(keys, cos, sin)\n",
    "        queries = apply_rope(queries, cos, sin)\n",
    "\n",
    "        # Expand keys and values to match the number of heads\n",
    "        # Shape: (b, num_heads, num_tokens, head_dim)\n",
    "        keys = keys.repeat_interleave(self.group_size, dim=1)  # Shape: (b, num_heads, num_tokens, head_dim)\n",
    "        values = values.repeat_interleave(self.group_size, dim=1)  # Shape: (b, num_heads, num_tokens, head_dim)\n",
    "        # For example, before repeat_interleave along dim=1 (query groups):\n",
    "        #   [K1, K2]\n",
    "        # After repeat_interleave (each query group is repeated group_size times):\n",
    "        #   [K1, K1, K2, K2]\n",
    "        # If we used regular repeat instead of repeat_interleave, we'd get:\n",
    "        #   [K1, K2, K1, K2]\n",
    "\n",
    "        # Compute scaled dot-product attention (aka self-attention) with a causal mask\n",
    "        # Shape: (b, num_heads, num_tokens, num_tokens)\n",
    "        attn_scores = queries @ keys.transpose(2, 3)  # Dot product for each head\n",
    "\n",
    "        # Compute attention scores\n",
    "        attn_scores = attn_scores.masked_fill(mask, -torch.inf)\n",
    "\n",
    "        attn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)\n",
    "        assert keys.shape[-1] == self.head_dim\n",
    "\n",
    "        # Shape: (b, num_tokens, num_heads, head_dim)\n",
    "        context_vec = (attn_weights @ values).transpose(1, 2)\n",
    "\n",
    "        # Combine heads, where self.d_out = self.num_heads * self.head_dim\n",
    "        context_vec = context_vec.reshape(b, num_tokens, self.d_out)\n",
    "        context_vec = self.out_proj(context_vec)  # optional projection\n",
    "\n",
    "        return context_vec\n",
    "\n",
    "\n",
    "class TransformerBlock(nn.Module):\n",
    "    def __init__(self, cfg):\n",
    "        super().__init__()\n",
    "        self.att = GroupedQueryAttention(\n",
    "            d_in=cfg[\"emb_dim\"],\n",
    "            d_out=cfg[\"emb_dim\"],\n",
    "            num_heads=cfg[\"n_heads\"],\n",
    "            num_kv_groups=cfg[\"n_kv_groups\"],\n",
    "            dtype=cfg[\"dtype\"]\n",
    "        )\n",
    "        self.ff = FeedForward(cfg)\n",
    "        self.norm1 = nn.RMSNorm(cfg[\"emb_dim\"], eps=1e-5, dtype=cfg[\"dtype\"])\n",
    "        self.norm2 = nn.RMSNorm(cfg[\"emb_dim\"], eps=1e-5, dtype=cfg[\"dtype\"])\n",
    "\n",
    "    def forward(self, x, mask, cos, sin):\n",
    "        # Shortcut connection for attention block\n",
    "        shortcut = x\n",
    "        x = self.norm1(x)\n",
    "        x = self.att(x, mask, cos, sin)  # Shape [batch_size, num_tokens, emb_size]\n",
    "        x = x + shortcut  # Add the original input back\n",
    "\n",
    "        # Shortcut connection for feed-forward block\n",
    "        shortcut = x\n",
    "        x = self.norm2(x)\n",
    "        x = self.ff(x)\n",
    "        x = x + shortcut  # Add the original input back\n",
    "\n",
    "        return x\n",
    "\n",
    "\n",
    "class Llama3Model(nn.Module):\n",
    "    def __init__(self, cfg):\n",
    "        super().__init__()\n",
    "\n",
    "        # Main model parameters\n",
    "        self.tok_emb = nn.Embedding(cfg[\"vocab_size\"], cfg[\"emb_dim\"], dtype=cfg[\"dtype\"])\n",
    "\n",
    "        self.trf_blocks = nn.ModuleList(  # ModuleList since Sequential can only accept one input, and we need `x, mask, cos, sin`\n",
    "            [TransformerBlock(cfg) for _ in range(cfg[\"n_layers\"])]\n",
    "        )\n",
    "\n",
    "        self.final_norm = nn.RMSNorm(cfg[\"emb_dim\"], eps=1e-5, dtype=cfg[\"dtype\"])\n",
    "        self.out_head = nn.Linear(cfg[\"emb_dim\"], cfg[\"vocab_size\"], bias=False, dtype=cfg[\"dtype\"])\n",
    "\n",
    "        # Uncomment the following code to tie weights\n",
    "        # self.out_head.weight = self.tok_emb.weight\n",
    "        # torch.nn.init.normal_(self.out_head.weight, mean=0.0, std=0.02)\n",
    "\n",
    "        # Reusable utilities\n",
    "        cos, sin = compute_rope_params(\n",
    "            head_dim=cfg[\"emb_dim\"] // cfg[\"n_heads\"],\n",
    "            theta_base=cfg[\"rope_base\"],\n",
    "            context_length=cfg[\"context_length\"],\n",
    "            freq_config=cfg[\"rope_freq\"]\n",
    "        )\n",
    "        self.register_buffer(\"cos\", cos, persistent=False)\n",
    "        self.register_buffer(\"sin\", sin, persistent=False)\n",
    "        self.cfg = cfg\n",
    "\n",
    "\n",
    "    def forward(self, in_idx):\n",
    "        # Forward pass\n",
    "        tok_embeds = self.tok_emb(in_idx)\n",
    "        x = tok_embeds\n",
    "\n",
    "        num_tokens = x.shape[1]\n",
    "        mask = torch.triu(torch.ones(num_tokens, num_tokens, device=x.device, dtype=torch.bool), diagonal=1)\n",
    "        \n",
    "        for block in self.trf_blocks:\n",
    "            x = block(x, mask, self.cos, self.sin)\n",
    "        x = self.final_norm(x)\n",
    "        logits = self.out_head(x.to(self.cfg[\"dtype\"]))\n",
    "        return logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "d12ac059-58d8-4db2-ac5a-9ec58b043daf",
   "metadata": {},
   "outputs": [],
   "source": [
    "#######################\n",
    "### Initialize Llama 3\n",
    "#######################\n",
    "\n",
    "# Llama 3.2 1B\n",
    "\n",
    "LLAMA32_CONFIG = {\n",
    "    \"vocab_size\": 128_256,           # Vocabulary size\n",
    "    \"context_length\": 131_072,       # Context length that was used to train the model\n",
    "    \"emb_dim\": 2048,                 # Embedding dimension\n",
    "    \"n_heads\": 32,                   # Number of attention heads\n",
    "    \"n_layers\": 16,                  # Number of layers\n",
    "    \"hidden_dim\": 8192,              # Size of the intermediate dimension in FeedForward\n",
    "    \"n_kv_groups\": 8,                # Key-Value groups for grouped-query attention\n",
    "    \"rope_base\": 500_000.0,          # The base in RoPE's \"theta\"\n",
    "    \"dtype\": torch.bfloat16,         # Lower-precision dtype to reduce memory usage\n",
    "    \"rope_freq\": {                   # RoPE frequency scaling\n",
    "        \"factor\": 32.0,\n",
    "        \"low_freq_factor\": 1.0,\n",
    "        \"high_freq_factor\": 4.0,\n",
    "        \"original_context_length\": 8192,\n",
    "    }\n",
    "}\n",
    "\n",
    "LLAMA32_CONFIG[\"train_context_length\"] = 256  # It's a small dataset, and we also want to keep memory usage reasonable\n",
    "\n",
    "torch.manual_seed(123)\n",
    "model = Llama3Model(LLAMA32_CONFIG)\n",
    "model.eval();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "d6732d1a-db47-42c3-aca3-8a871752f32f",
   "metadata": {},
   "outputs": [],
   "source": [
    "#######################\n",
    "### Set up tokenizer\n",
    "#######################\n",
    "\n",
    "import os\n",
    "from pathlib import Path\n",
    "\n",
    "import tiktoken\n",
    "from tiktoken.load import load_tiktoken_bpe\n",
    "\n",
    "\n",
    "\n",
    "class Tokenizer:\n",
    "    \"\"\"Thin wrapper around tiktoken that keeps track of Llama-3 special IDs.\"\"\"\n",
    "    def __init__(self, model_path):\n",
    "        if not os.path.isfile(model_path):\n",
    "            raise FileNotFoundError(model_path)\n",
    "\n",
    "        mergeable = load_tiktoken_bpe(model_path)\n",
    "\n",
    "        # hard-coded from Meta's tokenizer.json\n",
    "        self.special = {\n",
    "            \"<|begin_of_text|>\": 128000,\n",
    "            \"<|end_of_text|>\": 128001,\n",
    "            \"<|start_header_id|>\": 128006,\n",
    "            \"<|end_header_id|>\": 128007,\n",
    "            \"<|eot_id|>\": 128009,\n",
    "        }\n",
    "        self.special.update({f\"<|reserved_{i}|>\": 128002 + i\n",
    "                             for i in range(256)\n",
    "                             if 128002 + i not in self.special.values()})\n",
    "\n",
    "        self.model = tiktoken.Encoding(\n",
    "            name=Path(model_path).name,\n",
    "            pat_str=r\"(?i:'s|'t|'re|'ve|'m|'ll|'d)\"\n",
    "                    r\"|[^\\r\\n\\p{L}\\p{N}]?\\p{L}+\"\n",
    "                    r\"|\\p{N}{1,3}\"\n",
    "                    r\"| ?[^\\s\\p{L}\\p{N}]+[\\r\\n]*\"\n",
    "                    r\"|\\s*[\\r\\n]+\"\n",
    "                    r\"|\\s+(?!\\S)\"\n",
    "                    r\"|\\s+\",\n",
    "            mergeable_ranks=mergeable,\n",
    "            special_tokens=self.special,\n",
    "        )\n",
    "\n",
    "    def encode(self, text, bos=False, eos=False):\n",
    "        ids = ([self.special[\"<|begin_of_text|>\"]] if bos else []) \\\n",
    "              + self.model.encode(text)\n",
    "        if eos:\n",
    "            ids.append(self.special[\"<|end_of_text|>\"])\n",
    "        return ids\n",
    "\n",
    "    def decode(self, ids):\n",
    "        return self.model.decode(ids)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6061739e",
   "metadata": {},
   "source": [
    "- Please note that Meta AI requires that you accept the Llama 3.2 licensing terms before you can download the files; to do this, you have to create a Hugging Face Hub account and visit the [meta-llama/Llama-3.2-1B](https://huggingface.co/meta-llama/Llama-3.2-1B) repository to accept the terms\n",
    "- Next, you will need to create an access token; to generate an access token with READ permissions, click on the profile picture in the upper right and click on \"Settings\"\n",
    "\n",
    "\n",
    "<img src=\"https://sebastianraschka.com/images/LLMs-from-scratch-images/bonus/gpt-to-llama/settings.webp?1\" width=\"300px\">\n",
    "\n",
    "- Then, create and copy the access token so you can copy & paste it into the next code cell\n",
    "\n",
    "<img src=\"https://sebastianraschka.com/images/LLMs-from-scratch-images/bonus/gpt-to-llama/access-token.webp?1\" width=\"600px\">"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "0bab40cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Uncomment and run the following code if you are executing the notebook for the first time\n",
    "\n",
    "#from huggingface_hub import login\n",
    "#login()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "f102c705",
   "metadata": {},
   "outputs": [],
   "source": [
    "from huggingface_hub import hf_hub_download\n",
    "\n",
    "tokenizer_file_path = hf_hub_download(\n",
    "    repo_id=\"meta-llama/Llama-3.2-1B-Instruct\",\n",
    "    filename=\"original/tokenizer.model\",\n",
    "    local_dir=\"Llama-3.2-1B-Instruct\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "d9cc9167",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = Tokenizer(tokenizer_file_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "5e062b82-3540-48ce-8eb4-009686d0d16c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output text:\n",
      " Every effort moves youodableglich-create//{\n",
      " Buddhistザイン ILogger_us matching=User\n"
     ]
    }
   ],
   "source": [
    "# Same as chapter 4\n",
    "\n",
    "def generate_text_simple(model, idx, max_new_tokens, context_size):\n",
    "    # idx is (B, T) array of indices in the current context\n",
    "    for _ in range(max_new_tokens):\n",
    "\n",
    "        # Crop current context if it exceeds the supported context size\n",
    "        # E.g., if LLM supports only 5 tokens, and the context size is 10\n",
    "        # then only the last 5 tokens are used as context\n",
    "        idx_cond = idx[:, -context_size:]\n",
    "\n",
    "        # Get the predictions\n",
    "        with torch.no_grad():\n",
    "            logits = model(idx_cond)\n",
    "\n",
    "        # Focus only on the last time step\n",
    "        # (batch, n_token, vocab_size) becomes (batch, vocab_size)\n",
    "        logits = logits[:, -1, :]\n",
    "\n",
    "        # Get the idx of the vocab entry with the highest logits value\n",
    "        idx_next = torch.argmax(logits, dim=-1, keepdim=True)  # (batch, 1)\n",
    "\n",
    "        # Append sampled index to the running sequence\n",
    "        idx = torch.cat((idx, idx_next), dim=1)  # (batch, n_tokens+1)\n",
    "\n",
    "    return idx\n",
    "\n",
    "\n",
    "def text_to_token_ids(text, tokenizer):\n",
    "    encoded = tokenizer.encode(text)\n",
    "    encoded_tensor = torch.tensor(encoded).unsqueeze(0) # add batch dimension\n",
    "    return encoded_tensor\n",
    "\n",
    "def token_ids_to_text(token_ids, tokenizer):\n",
    "    flat = token_ids.squeeze(0) # remove batch dimension\n",
    "    return tokenizer.decode(flat.tolist())\n",
    "\n",
    "start_context = \"Every effort moves you\"\n",
    "\n",
    "token_ids = generate_text_simple(\n",
    "    model=model,\n",
    "    idx=text_to_token_ids(start_context, tokenizer),\n",
    "    max_new_tokens=10,\n",
    "    context_size=LLAMA32_CONFIG[\"train_context_length\"]\n",
    ")\n",
    "\n",
    "print(\"Output text:\\n\", token_ids_to_text(token_ids, tokenizer))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f3d7ea2-637f-4490-bc76-e361fc81ae98",
   "metadata": {
    "id": "0f3d7ea2-637f-4490-bc76-e361fc81ae98"
   },
   "source": [
    "&nbsp;\n",
    "### 5.1.2 Calculating the text generation loss: cross-entropy and perplexity"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e669b90a-4bc9-422f-8f62-6c6d99189f68",
   "metadata": {},
   "source": [
    "- Similar to chapter 5"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2ec6c217-e429-40c7-ad71-5d0a9da8e487",
   "metadata": {
    "id": "2ec6c217-e429-40c7-ad71-5d0a9da8e487"
   },
   "source": [
    "&nbsp;\n",
    "### 5.1.3 Calculating the training and validation set losses"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "654fde37-b2a9-4a20-a8d3-0206c056e2ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "import requests\n",
    "\n",
    "file_path = \"the-verdict.txt\"\n",
    "url = \"https://raw.githubusercontent.com/rasbt/LLMs-from-scratch/main/ch02/01_main-chapter-code/the-verdict.txt\"\n",
    "\n",
    "if not os.path.exists(file_path):\n",
    "    response = requests.get(url, timeout=30)\n",
    "    response.raise_for_status()\n",
    "    text_data = response.text\n",
    "    with open(file_path, \"w\", encoding=\"utf-8\") as file:\n",
    "        file.write(text_data)\n",
    "else:\n",
    "    with open(file_path, \"r\", encoding=\"utf-8\") as file:\n",
    "        text_data = file.read()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "379330f1-80f4-4e34-8724-41d892b04cee",
   "metadata": {},
   "source": [
    "- A quick check that the text loaded ok by printing the first and last 99 characters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "6kgJbe4ehI4q",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "id": "6kgJbe4ehI4q",
    "outputId": "9ff31e88-ee37-47e9-ee64-da6eb552f46f"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "I HAD always thought Jack Gisburn rather a cheap genius--though a good fellow enough--so it was no \n"
     ]
    }
   ],
   "source": [
    "# First 99 characters\n",
    "print(text_data[:99])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "j2XPde_ThM_e",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "id": "j2XPde_ThM_e",
    "outputId": "a900c1b9-9a87-4078-968b-a5721deda5cb"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "it for me! The Strouds stand alone, and happen once--but there's no exterminating our kind of art.\"\n"
     ]
    }
   ],
   "source": [
    "# Last 99 characters\n",
    "print(text_data[-99:])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "6b46a952-d50a-4837-af09-4095698f7fd1",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "6b46a952-d50a-4837-af09-4095698f7fd1",
    "outputId": "c2a25334-21ca-486e-8226-0296e5fc6486"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Characters: 20479\n",
      "Tokens: 4941\n"
     ]
    }
   ],
   "source": [
    "total_characters = len(text_data)\n",
    "total_tokens = len(tokenizer.encode(text_data))\n",
    "\n",
    "print(\"Characters:\", total_characters)\n",
    "print(\"Tokens:\", total_tokens)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "0959c855-f860-4358-8b98-bc654f047578",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import Dataset, DataLoader\n",
    "\n",
    "\n",
    "class GPTDatasetV1(Dataset):\n",
    "    def __init__(self, txt, tokenizer, max_length, stride):\n",
    "        self.input_ids = []\n",
    "        self.target_ids = []\n",
    "\n",
    "        # Tokenize the entire text\n",
    "        token_ids = tokenizer.encode(txt)\n",
    "\n",
    "        # Use a sliding window to chunk the book into overlapping sequences of max_length\n",
    "        for i in range(0, len(token_ids) - max_length, stride):\n",
    "            input_chunk = token_ids[i:i + max_length]\n",
    "            target_chunk = token_ids[i + 1: i + max_length + 1]\n",
    "            self.input_ids.append(torch.tensor(input_chunk))\n",
    "            self.target_ids.append(torch.tensor(target_chunk))\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.input_ids)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return self.input_ids[idx], self.target_ids[idx]\n",
    "\n",
    "# Note that we have to change the function below because we previously hard-coded the\n",
    "# GPT-2 tokenizer in the data loader\n",
    "def create_dataloader_v1(txt, tokenizer, batch_size=4, max_length=256,\n",
    "                         stride=128, shuffle=True, drop_last=True, num_workers=0):\n",
    "    # Initialize the tokenizer\n",
    "    # tokenizer = tiktoken.get_encoding(\"gpt2\")\n",
    "    tokenizer = tokenizer\n",
    "\n",
    "    # Create dataset\n",
    "    dataset = GPTDatasetV1(txt, tokenizer, max_length, stride)\n",
    "\n",
    "    # Create dataloader\n",
    "    dataloader = DataLoader(\n",
    "        dataset, batch_size=batch_size, shuffle=shuffle, drop_last=drop_last, num_workers=num_workers)\n",
    "\n",
    "    return dataloader\n",
    "\n",
    "\n",
    "# Train/validation ratio\n",
    "train_ratio = 0.90\n",
    "split_idx = int(train_ratio * len(text_data))\n",
    "train_data = text_data[:split_idx]\n",
    "val_data = text_data[split_idx:]\n",
    "\n",
    "\n",
    "torch.manual_seed(123)\n",
    "\n",
    "train_loader = create_dataloader_v1(\n",
    "    train_data,\n",
    "    tokenizer=tokenizer,\n",
    "    batch_size=2,\n",
    "    max_length=LLAMA32_CONFIG[\"train_context_length\"],\n",
    "    stride=LLAMA32_CONFIG[\"train_context_length\"],\n",
    "    drop_last=True,\n",
    "    shuffle=True,\n",
    "    num_workers=0\n",
    ")\n",
    "\n",
    "val_loader = create_dataloader_v1(\n",
    "    val_data,\n",
    "    tokenizer=tokenizer,\n",
    "    batch_size=2,\n",
    "    max_length=LLAMA32_CONFIG[\"train_context_length\"],\n",
    "    stride=LLAMA32_CONFIG[\"train_context_length\"],\n",
    "    drop_last=False,\n",
    "    shuffle=False,\n",
    "    num_workers=0\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a8e0514d-b990-4dc0-9afb-7721993284a0",
   "metadata": {},
   "source": [
    "- An optional check that the data was loaded correctly:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "ca0116d0-d229-472c-9fbf-ebc229331c3e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train loader:\n",
      "torch.Size([2, 256]) torch.Size([2, 256])\n",
      "torch.Size([2, 256]) torch.Size([2, 256])\n",
      "torch.Size([2, 256]) torch.Size([2, 256])\n",
      "torch.Size([2, 256]) torch.Size([2, 256])\n",
      "torch.Size([2, 256]) torch.Size([2, 256])\n",
      "torch.Size([2, 256]) torch.Size([2, 256])\n",
      "torch.Size([2, 256]) torch.Size([2, 256])\n",
      "torch.Size([2, 256]) torch.Size([2, 256])\n",
      "\n",
      "Validation loader:\n",
      "torch.Size([2, 256]) torch.Size([2, 256])\n"
     ]
    }
   ],
   "source": [
    "print(\"Train loader:\")\n",
    "for x, y in train_loader:\n",
    "    print(x.shape, y.shape)\n",
    "\n",
    "print(\"\\nValidation loader:\")\n",
    "for x, y in val_loader:\n",
    "    print(x.shape, y.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f7b9b1a4-863d-456f-a8dd-c07fb5c024ed",
   "metadata": {},
   "source": [
    "- Another optional check that the token sizes are in the expected ballpark:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "eb860488-5453-41d7-9870-23b723f742a0",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "eb860488-5453-41d7-9870-23b723f742a0",
    "outputId": "96b9451a-9557-4126-d1c8-51610a1995ab"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training tokens: 4096\n",
      "Validation tokens: 512\n",
      "All tokens: 4608\n"
     ]
    }
   ],
   "source": [
    "train_tokens = 0\n",
    "for input_batch, target_batch in train_loader:\n",
    "    train_tokens += input_batch.numel()\n",
    "\n",
    "val_tokens = 0\n",
    "for input_batch, target_batch in val_loader:\n",
    "    val_tokens += input_batch.numel()\n",
    "\n",
    "print(\"Training tokens:\", train_tokens)\n",
    "print(\"Validation tokens:\", val_tokens)\n",
    "print(\"All tokens:\", train_tokens + val_tokens)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c3085e8-665e-48eb-bb41-cdde61537e06",
   "metadata": {},
   "source": [
    "- Next, we implement a utility function to calculate the cross-entropy loss of a given batch\n",
    "- In addition, we implement a second utility function to compute the loss for a user-specified number of batches in a data loader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "7b9de31e-4096-47b3-976d-b6d2fdce04bc",
   "metadata": {
    "id": "7b9de31e-4096-47b3-976d-b6d2fdce04bc"
   },
   "outputs": [],
   "source": [
    "def calc_loss_batch(input_batch, target_batch, model, device):\n",
    "    input_batch, target_batch = input_batch.to(device), target_batch.to(device)\n",
    "    logits = model(input_batch)\n",
    "    loss = torch.nn.functional.cross_entropy(logits.flatten(0, 1), target_batch.flatten())\n",
    "    return loss\n",
    "\n",
    "\n",
    "def calc_loss_loader(data_loader, model, device, num_batches=None):\n",
    "    total_loss = 0.\n",
    "    if len(data_loader) == 0:\n",
    "        return float(\"nan\")\n",
    "    elif num_batches is None:\n",
    "        num_batches = len(data_loader)\n",
    "    else:\n",
    "        # Reduce the number of batches to match the total number of batches in the data loader\n",
    "        # if num_batches exceeds the number of batches in the data loader\n",
    "        num_batches = min(num_batches, len(data_loader))\n",
    "    for i, (input_batch, target_batch) in enumerate(data_loader):\n",
    "        if i < num_batches:\n",
    "            loss = calc_loss_batch(input_batch, target_batch, model, device)\n",
    "            total_loss += loss.item()\n",
    "        else:\n",
    "            break\n",
    "    return total_loss / num_batches"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0691332-84d0-48b3-b462-a885ddeb4fca",
   "metadata": {},
   "source": [
    "- If you have a machine with a CUDA-supported GPU, the LLM will train on the GPU without making any changes to the code\n",
    "- Via the `device` setting, we ensure that the data is loaded onto the same device as the LLM model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "56f5b0c9-1065-4d67-98b9-010e42fc1e2a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using cuda device.\n",
      "Training loss: 11.9375\n",
      "Validation loss: 11.9375\n"
     ]
    }
   ],
   "source": [
    "if torch.cuda.is_available():\n",
    "    device = torch.device(\"cuda\")\n",
    "elif torch.backends.mps.is_available():\n",
    "    # Use PyTorch 2.9 or newer for stable mps results\n",
    "    major, minor = map(int, torch.__version__.split(\".\")[:2])\n",
    "    if (major, minor) >= (2, 9):\n",
    "        device = torch.device(\"mps\")\n",
    "    else:\n",
    "        device = torch.device(\"cpu\")\n",
    "else:\n",
    "    device = torch.device(\"cpu\")\n",
    "\n",
    "\n",
    "print(f\"Using {device} device.\")\n",
    "\n",
    "\n",
    "model.to(device) # no assignment model = model.to(device) necessary for nn.Module classes\n",
    "\n",
    "\n",
    "torch.manual_seed(123) # For reproducibility due to the shuffling in the data loader\n",
    "\n",
    "with torch.no_grad(): # Disable gradient tracking for efficiency because we are not training, yet\n",
    "    train_loss = calc_loss_loader(train_loader, model, device)\n",
    "    val_loss = calc_loss_loader(val_loader, model, device)\n",
    "\n",
    "print(\"Training loss:\", train_loss)\n",
    "print(\"Validation loss:\", val_loss)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b9339f8d-00cb-4206-af67-58c32bd72055",
   "metadata": {
    "id": "b9339f8d-00cb-4206-af67-58c32bd72055"
   },
   "source": [
    "&nbsp;\n",
    "## 5.2 Training an LLM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "Mtp4gY0ZO-qq",
   "metadata": {
    "id": "Mtp4gY0ZO-qq"
   },
   "outputs": [],
   "source": [
    "def train_model_simple(model, train_loader, val_loader, optimizer, device, num_epochs,\n",
    "                       eval_freq, eval_iter, start_context, tokenizer):\n",
    "    # Initialize lists to track losses and tokens seen\n",
    "    train_losses, val_losses, track_tokens_seen = [], [], []\n",
    "    tokens_seen, global_step = 0, -1\n",
    "\n",
    "    # Main training loop\n",
    "    for epoch in range(num_epochs):\n",
    "        model.train()  # Set model to training mode\n",
    "        \n",
    "        for input_batch, target_batch in train_loader:\n",
    "            optimizer.zero_grad() # Reset loss gradients from previous batch iteration\n",
    "            loss = calc_loss_batch(input_batch, target_batch, model, device)\n",
    "            loss.backward() # Calculate loss gradients\n",
    "            optimizer.step() # Update model weights using loss gradients\n",
    "            tokens_seen += input_batch.numel()\n",
    "            global_step += 1\n",
    "\n",
    "            # Optional evaluation step\n",
    "            if global_step % eval_freq == 0:\n",
    "                train_loss, val_loss = evaluate_model(\n",
    "                    model, train_loader, val_loader, device, eval_iter)\n",
    "                train_losses.append(train_loss)\n",
    "                val_losses.append(val_loss)\n",
    "                track_tokens_seen.append(tokens_seen)\n",
    "                print(f\"Ep {epoch+1} (Step {global_step:06d}): \"\n",
    "                      f\"Train loss {train_loss:.3f}, Val loss {val_loss:.3f}\")\n",
    "\n",
    "        # Print a sample text after each epoch\n",
    "        generate_and_print_sample(\n",
    "            model, tokenizer, device, start_context\n",
    "        )\n",
    "\n",
    "    return train_losses, val_losses, track_tokens_seen\n",
    "\n",
    "\n",
    "def evaluate_model(model, train_loader, val_loader, device, eval_iter):\n",
    "    model.eval()\n",
    "    with torch.no_grad():\n",
    "        train_loss = calc_loss_loader(train_loader, model, device, num_batches=eval_iter)\n",
    "        val_loss = calc_loss_loader(val_loader, model, device, num_batches=eval_iter)\n",
    "    model.train()\n",
    "    return train_loss, val_loss\n",
    "\n",
    "\n",
    "def generate_and_print_sample(model, tokenizer, device, start_context):\n",
    "    model.eval()\n",
    "    context_size = model.cfg[\"context_length\"]\n",
    "    encoded = text_to_token_ids(start_context, tokenizer).to(device)\n",
    "    with torch.no_grad():\n",
    "        token_ids = generate_text_simple(\n",
    "            model=model, idx=encoded,\n",
    "            max_new_tokens=50, context_size=context_size\n",
    "        )\n",
    "    decoded_text = token_ids_to_text(token_ids, tokenizer)\n",
    "    print(decoded_text.replace(\"\\n\", \" \"))  # Compact print format\n",
    "    model.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "3422000b-7aa2-485b-92df-99372cd22311",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "3422000b-7aa2-485b-92df-99372cd22311",
    "outputId": "0e046603-908d-4093-8ae5-ef2f632639fb"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ep 1 (Step 000000): Train loss 10.300, Val loss 11.062\n",
      "Ep 1 (Step 000005): Train loss 8.306, Val loss 8.438\n",
      "Every effort moves you and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and and I and I and I and and\n",
      "Ep 2 (Step 000010): Train loss 6.875, Val loss 7.250\n",
      "Ep 2 (Step 000015): Train loss 6.475, Val loss 7.156\n",
      "Every effort moves you of of of of of of of of of of of of of of of of of of of of of of of of of of of of of of of of of of of of of of of of of of of of of of of of of of\n",
      "Ep 3 (Step 000020): Train loss 6.375, Val loss 7.219\n",
      "Every effort moves you the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the\n",
      "Ep 4 (Step 000025): Train loss 6.225, Val loss 7.156\n",
      "Ep 4 (Step 000030): Train loss 6.250, Val loss 7.188\n",
      "Every effort moves you,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,\n",
      "Ep 5 (Step 000035): Train loss 6.100, Val loss 7.125\n",
      "Every effort moves you the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the\n",
      "Ep 6 (Step 000040): Train loss 6.088, Val loss 7.156\n",
      "Ep 6 (Step 000045): Train loss 6.131, Val loss 7.188\n",
      "Every effort moves you, and I, and I, and I, and I, and I, and I, and I, and I, and I, and I, and I, and I, and I, and I, and I, and I, and\n",
      "Ep 7 (Step 000050): Train loss 5.956, Val loss 7.031\n",
      "Ep 7 (Step 000055): Train loss 5.850, Val loss 6.969\n",
      "Every effort moves you, I was, I was, I was, I was, I was, I was, I was, I was, I was, I was, I was, I was, I was, I was, I was, I was, I\n",
      "Ep 8 (Step 000060): Train loss 5.650, Val loss 6.875\n",
      "Every effort moves you, and he had the fact--the, and he had the fact--the, and he had the fact--the, and he had the fact--the, and he had the fact--the, and he had the fact--the, and\n",
      "Ep 9 (Step 000065): Train loss 5.444, Val loss 6.719\n",
      "Ep 9 (Step 000070): Train loss 5.094, Val loss 6.656\n",
      "Every effort moves you to me to me to me to me to me to me. The I had been the Riv, and I had been the Riv, and I had been the Riv, and I had been the Riv, and I had been the Riv, and I\n",
      "Ep 10 (Step 000075): Train loss 4.987, Val loss 6.656\n",
      "Every effort moves you to have the don't me to have the don't me--I was a little, and I was a little in the don't me--I was a little, and I was a little in the don't me--I was a little,\n",
      "Ep 11 (Step 000080): Train loss 4.781, Val loss 6.656\n",
      "Ep 11 (Step 000085): Train loss 4.606, Val loss 6.531\n",
      "Every effort moves you, I had the Riviera of the Riviera of the Riviera of the Riviera of the Riviera of the Riviera of the Riviera of the Riviera of the Riviera of the Riviera of the Riviera of the Riviera\n",
      "Ep 12 (Step 000090): Train loss 4.138, Val loss 6.656\n",
      "Ep 12 (Step 000095): Train loss 3.978, Val loss 6.656\n",
      "Every effort moves you know, and he had been his pictures_'t say to have been his pictures--his work. Gisburn's the picture of the picture of the picture of the picture of the picture of the picture of the picture of the picture of the picture\n",
      "Ep 13 (Step 000100): Train loss 3.688, Val loss 6.812\n",
      "Every effort moves you of the fact, and I had been his eyes grew curiosity was one of the fact to the fact to the fact to the fact to the fact to the fact to the fact to the fact to the fact to the fact to the fact to the fact\n",
      "Ep 14 (Step 000105): Train loss 3.278, Val loss 6.688\n",
      "Ep 14 (Step 000110): Train loss 2.538, Val loss 6.938\n",
      "Every effort moves you know. Gisburn's \"Be dissatisfied with a little Claude Nutley, and Mrs. Gisburn's \"Be dissatisfied with a little Claude Nutley, and Mrs. Gisburn's \"Be dissatisfied with a little Claude\n",
      "Ep 15 (Step 000115): Train loss 2.487, Val loss 7.219\n",
      "Every effort moves you like to the picture, the picture, and in the picture, the picture, the picture, and in the picture, and in the picture, and in the picture, and in the picture, and in the picture, and in the picture, and\n",
      "Ep 16 (Step 000120): Train loss 2.139, Val loss 7.219\n",
      "Ep 16 (Step 000125): Train loss 2.059, Val loss 6.938\n",
      "Every effort moves you know,\" she was not till after that I was not till after. Gisburn had been denied the house.\"  He was not till after that Mrs. Gisburn had been denied the house.\"  He was not till after that Mrs. Gis\n",
      "Ep 17 (Step 000130): Train loss 1.706, Val loss 7.250\n",
      "Ep 17 (Step 000135): Train loss 1.631, Val loss 7.219\n",
      "Every effort moves you know, and in the Riviera, and in the Riviera, and in the Riviera, and in the Riviera, and in the Riviera, and in the Riviera, and in the Riviera, and in the Riviera,\n",
      "Ep 18 (Step 000140): Train loss 1.458, Val loss 7.406\n",
      "Every effort moves you know, and I said, and I said, and twirling between his pictures? My curiosity: \"There: \"There: \"There: \"There: \"There: \"There: \"There: \"There: \"There: \"There:\n",
      "Ep 19 (Step 000145): Train loss 1.155, Val loss 7.281\n",
      "Ep 19 (Step 000150): Train loss 1.105, Val loss 7.531\n",
      "Every effort moves you know.\"  I have let it was no least sign of Jack's I had been surrounded by interesting--I must have let it was taken with a little quickly_ fashionable painter.\"  \"When he had I had I had to have let it was one of Jack\n",
      "Ep 20 (Step 000155): Train loss 0.896, Val loss 7.344\n",
      "Every effort moves you know.\"  He answered slowly: \"Be dissatisfied with a little quickly.  \"Oh, and distinguished objects, and distinguished objects, and straddling and straining, and straining, and straining, and straining, and straining, and\n",
      "Ep 21 (Step 000160): Train loss 0.620, Val loss 7.469\n",
      "Ep 21 (Step 000165): Train loss 0.535, Val loss 7.531\n",
      "Every effort moves you know.\"  I must have given up at the sketch of the picture for he had ever of Jack's break with a little wild--because he had ever had ever had ever having been no leastburn had ever had ever of his glory, and in a\n",
      "Ep 22 (Step 000170): Train loss 0.353, Val loss 7.625\n",
      "Ep 22 (Step 000175): Train loss 0.227, Val loss 7.562\n",
      "Every effort moves you know you know.\"  He stood looking about to see it. The rest of the fact that, and in the air of it. He had ever had ever had ever had ever knew.\"  \"You ever knew.\"  \"You ever had to now it the fact that\n",
      "Ep 23 (Step 000180): Train loss 0.174, Val loss 7.500\n",
      "Every effort moves you know.\"  He stood looking,\" he liked his bed. And it was just manage to have formed himself,\" she began to go under--because he didn't bear the current--because he didn't let a little wild--because he didn't you ever\n",
      "Ep 24 (Step 000185): Train loss 0.094, Val loss 7.719\n",
      "Ep 24 (Step 000190): Train loss 0.071, Val loss 7.656\n",
      "Every effort moves you know.\"  I couldn't bear I had to see it. I had to let it was the honour--say the honour--just a pale yellow or _rose Dubarry_ drawing-room.  I made a deprecating gesture, and I had to\n",
      "Ep 25 (Step 000195): Train loss 0.074, Val loss 7.781\n",
      "Every effort moves you know.\"  He laughed again--because he had come to stay! The rest of us had to let ourselves be swept along or go under, but he was high above the current--on everlasting foundations, as you say.  \"Well, I went off to\n",
      "Ep 26 (Step 000200): Train loss 0.030, Val loss 7.812\n",
      "Ep 26 (Step 000205): Train loss 0.025, Val loss 7.844\n",
      "Every effort moves you know.\"  He laughed again--because he had come to stay! The rest of us had to let ourselves be swept along or go under, but he was high above the current--on everlasting foundations, as you say.  \"Well, I went off to\n",
      "Ep 27 (Step 000210): Train loss 0.019, Val loss 7.938\n",
      "Ep 27 (Step 000215): Train loss 0.016, Val loss 7.906\n",
      "Every effort moves you know.\"  He laughed again--because he had come to stay! The rest of us had to let ourselves be swept along or go under, but he was high above the current--on everlasting foundations, as you say.  \"Well, I went off to\n",
      "Ep 28 (Step 000220): Train loss 0.011, Val loss 7.938\n",
      "Every effort moves you know.\"  He laughed again--because he had come to stay! The rest of us had to let ourselves be swept along or go under, but he was high above the current--on everlasting foundations, as you say.  \"Well, I went off to\n",
      "Ep 29 (Step 000225): Train loss 0.008, Val loss 7.938\n",
      "Ep 29 (Step 000230): Train loss 0.007, Val loss 7.969\n",
      "Every effort moves you know.\"  He laughed again--because he had come to stay! The rest of us had to let ourselves be swept along or go under, but he was high above the current--on everlasting foundations, as you say.  \"Well, I went off to\n",
      "Ep 30 (Step 000235): Train loss 0.006, Val loss 7.969\n",
      "Every effort moves you know.\"  He laughed again--because he had come to stay! The rest of us had to let ourselves be swept along or go under, but he was high above the current--on everlasting foundations, as you say.  \"Well, I went off to\n",
      "Ep 31 (Step 000240): Train loss 0.006, Val loss 7.938\n",
      "Ep 31 (Step 000245): Train loss 0.005, Val loss 7.938\n",
      "Every effort moves you know.\"  He laughed again--because he had come to stay! The rest of us had to let ourselves be swept along or go under, but he was high above the current--on everlasting foundations, as you say.  \"Well, I went off to\n",
      "Ep 32 (Step 000250): Train loss 0.005, Val loss 7.938\n",
      "Ep 32 (Step 000255): Train loss 0.005, Val loss 7.938\n",
      "Every effort moves you know.\"  He laughed again--because he had come to stay! The rest of us had to let ourselves be swept along or go under, but he was high above the current--on everlasting foundations, as you say.  \"Well, I went off to\n",
      "Ep 33 (Step 000260): Train loss 0.004, Val loss 7.938\n",
      "Every effort moves you know.\"  He laughed again--because he had come to stay! The rest of us had to let ourselves be swept along or go under, but he was high above the current--on everlasting foundations, as you say.  \"Well, I went off to\n",
      "Ep 34 (Step 000265): Train loss 0.005, Val loss 7.938\n",
      "Ep 34 (Step 000270): Train loss 0.004, Val loss 7.938\n",
      "Every effort moves you know.\"  He laughed again--because he had come to stay! The rest of us had to let ourselves be swept along or go under, but he was high above the current--on everlasting foundations, as you say.  \"Well, I went off to\n",
      "Ep 35 (Step 000275): Train loss 0.004, Val loss 7.969\n",
      "Every effort moves you know.\"  He laughed again--because he had come to stay! The rest of us had to let ourselves be swept along or go under, but he was high above the current--on everlasting foundations, as you say.  \"Well, I went off to\n",
      "Ep 36 (Step 000280): Train loss 0.004, Val loss 7.969\n",
      "Ep 36 (Step 000285): Train loss 0.004, Val loss 7.969\n",
      "Every effort moves you know.\"  He laughed again--because he had come to stay! The rest of us had to let ourselves be swept along or go under, but he was high above the current--on everlasting foundations, as you say.  \"Well, I went off to\n",
      "Ep 37 (Step 000290): Train loss 0.004, Val loss 7.969\n",
      "Ep 37 (Step 000295): Train loss 0.004, Val loss 7.969\n",
      "Every effort moves you know.\"  He laughed again--because he had come to stay! The rest of us had to let ourselves be swept along or go under, but he was high above the current--on everlasting foundations, as you say.  \"Well, I went off to\n",
      "Ep 38 (Step 000300): Train loss 0.004, Val loss 7.969\n",
      "Every effort moves you know.\"  He laughed again--because he had come to stay! The rest of us had to let ourselves be swept along or go under, but he was high above the current--on everlasting foundations, as you say.  \"Well, I went off to\n",
      "Ep 39 (Step 000305): Train loss 0.004, Val loss 7.969\n",
      "Ep 39 (Step 000310): Train loss 0.004, Val loss 7.969\n",
      "Every effort moves you know.\"  He laughed again--because he had come to stay! The rest of us had to let ourselves be swept along or go under, but he was high above the current--on everlasting foundations, as you say.  \"Well, I went off to\n",
      "Ep 40 (Step 000315): Train loss 0.004, Val loss 7.969\n",
      "Every effort moves you know.\"  He laughed again--because he had come to stay! The rest of us had to let ourselves be swept along or go under, but he was high above the current--on everlasting foundations, as you say.  \"Well, I went off to\n"
     ]
    }
   ],
   "source": [
    "# Note:\n",
    "# Uncomment the following code to calculate the execution time\n",
    "# import time\n",
    "# start_time = time.time()\n",
    "\n",
    "torch.manual_seed(123)\n",
    "model = Llama3Model(LLAMA32_CONFIG)\n",
    "model.to(device)\n",
    "optimizer = torch.optim.AdamW(model.parameters(), lr=0.0004, weight_decay=0.1)\n",
    "\n",
    "num_epochs = 40\n",
    "train_losses, val_losses, tokens_seen = train_model_simple(\n",
    "    model, train_loader, val_loader, optimizer, device,\n",
    "    num_epochs=num_epochs, eval_freq=5, eval_iter=5,\n",
    "    start_context=\"Every effort moves you\", tokenizer=tokenizer\n",
    ")\n",
    "\n",
    "# Note:\n",
    "# Uncomment the following code to show the execution time\n",
    "# end_time = time.time()\n",
    "# execution_time_minutes = (end_time - start_time) / 60\n",
    "# print(f\"Training completed in {execution_time_minutes:.2f} minutes.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "0WSRu2i0iHJE",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 487
    },
    "id": "0WSRu2i0iHJE",
    "outputId": "9d36c61b-517d-4f07-a7e8-4563aff78b11"
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 32200 characters]",
      "text/plain": [
       "<Figure size 500x300 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.ticker import MaxNLocator\n",
    "\n",
    "\n",
    "def plot_losses(epochs_seen, tokens_seen, train_losses, val_losses):\n",
    "    fig, ax1 = plt.subplots(figsize=(5, 3))\n",
    "\n",
    "    # Plot training and validation loss against epochs\n",
    "    ax1.plot(epochs_seen, train_losses, label=\"Training loss\")\n",
    "    ax1.plot(epochs_seen, val_losses, linestyle=\"-.\", label=\"Validation loss\")\n",
    "    ax1.set_xlabel(\"Epochs\")\n",
    "    ax1.set_ylabel(\"Loss\")\n",
    "    ax1.legend(loc=\"upper right\")\n",
    "    ax1.xaxis.set_major_locator(MaxNLocator(integer=True))  # only show integer labels on x-axis\n",
    "\n",
    "    # Create a second x-axis for tokens seen\n",
    "    ax2 = ax1.twiny()  # Create a second x-axis that shares the same y-axis\n",
    "    ax2.plot(tokens_seen, train_losses, alpha=0)  # Invisible plot for aligning ticks\n",
    "    ax2.set_xlabel(\"Tokens seen\")\n",
    "\n",
    "    fig.tight_layout()  # Adjust layout to make room\n",
    "    plt.savefig(\"loss-plot.pdf\")\n",
    "    plt.show()\n",
    "\n",
    "epochs_tensor = torch.linspace(0, num_epochs, len(train_losses))\n",
    "plot_losses(epochs_tensor, tokens_seen, train_losses, val_losses)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "699f45fc-bf78-42f2-bd24-2355db41b28f",
   "metadata": {
    "id": "699f45fc-bf78-42f2-bd24-2355db41b28f"
   },
   "source": [
    "&nbsp;\n",
    "## 5.3 Decoding strategies to control randomness"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "2734cee0-f6f9-42d5-b71c-fa7e0ef28b6d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output text:\n",
      " Every effort moves you know.\"\n",
      "\n",
      "He laughed again--because he had come to stay! The rest of us had to let ourselves be swept along or\n"
     ]
    }
   ],
   "source": [
    "inference_device = torch.device(\"cpu\")\n",
    "\n",
    "model.to(inference_device)\n",
    "model.eval()\n",
    "\n",
    "token_ids = generate_text_simple(\n",
    "    model=model,\n",
    "    idx=text_to_token_ids(\"Every effort moves you\", tokenizer).to(inference_device),\n",
    "    max_new_tokens=25,\n",
    "    context_size=LLAMA32_CONFIG[\"train_context_length\"]\n",
    ")\n",
    "\n",
    "print(\"Output text:\\n\", token_ids_to_text(token_ids, tokenizer))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "bf2e432d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output text:\n",
      " Hello's tears his ridiculous modesty, you know. He says they're not fit to have about; he's sent them all\n"
     ]
    }
   ],
   "source": [
    "token_ids = generate_text_simple(\n",
    "    model=model,\n",
    "    idx=text_to_token_ids(\"Hello\", tokenizer).to(inference_device),\n",
    "    max_new_tokens=25,\n",
    "    context_size=LLAMA32_CONFIG[\"train_context_length\"]\n",
    ")\n",
    "\n",
    "print(\"Output text:\\n\", token_ids_to_text(token_ids, tokenizer))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4bb6f380-a798-4fd9-825c-17b7cd29a994",
   "metadata": {},
   "source": [
    "&nbsp;\n",
    "### 5.3.1 Temperature scaling"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "327fdc96-cdba-4468-98a7-69c24c0855c9",
   "metadata": {},
   "source": [
    "- Similar to chapter 5"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6e4873e-07e4-4abb-85df-bdaedcc1a6f7",
   "metadata": {},
   "source": [
    "&nbsp;\n",
    "### 5.3.2 Top-k sampling"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e57fe45-1dfd-4ca7-97a9-0e57e9e6dd64",
   "metadata": {},
   "source": [
    "- Similar to chapter 5"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "56056503-a15d-4315-a3ff-46647a4c7c45",
   "metadata": {},
   "source": [
    "&nbsp;\n",
    "### 5.3.3 Modifying the text generation function"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9447a4bc-02fa-4fa8-ad0e-3abb4a1c9457",
   "metadata": {},
   "source": [
    "- Similar to chapter 5"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e2002ca-f4c1-48af-9e0a-88bfc163ba0b",
   "metadata": {},
   "source": [
    "&nbsp;\n",
    "## 5.4 Loading and saving model weights in PyTorch"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0488e58-691e-435a-bae0-ce430450dad4",
   "metadata": {},
   "source": [
    "- Similar to chapter 5"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4194350e-0409-4a63-8ffd-d3a896509032",
   "metadata": {},
   "source": [
    "&nbsp;\n",
    "## 5.5 Loading pretrained weights"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f48d52a7-a9a9-4021-a483-e6cfb077bf31",
   "metadata": {},
   "source": [
    "- See [Qwen3 0.6B from-scratch](../11_qwen3/standalone-qwen3.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2a66474-230d-4180-a8ff-843e04f1f1c4",
   "metadata": {},
   "source": [
    "&nbsp;\n",
    "## Summary and takeaways"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "156b0735-5d96-4db9-b10e-c9e52a238a69",
   "metadata": {},
   "source": [
    "- Skipped"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "A100",
   "machine_shape": "hm",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.13.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
