{
 "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 Qwen3 on Unlabeled Data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c4fa2aa",
   "metadata": {},
   "source": [
    "- This notebook plugs in the [Qwen3 0.6B from-scratch](../11_qwen3/standalone-qwen3.ipynb) model into (the pretraining portion) of chapter 5\n",
    "- This is to show how to use Qwen3 0.6B as a drop-in replacement for the GTP-2 model used in [chapter 5](../01_main-chapter-code/ch05.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9621e21b-565a-48ae-9432-e8768e991d42",
   "metadata": {},
   "source": [
    "<img src=\"https://sebastianraschka.com/images/LLMs-from-scratch-images/bonus/qwen/qwen-overview.webp\">\n",
    "  "
   ]
  },
  {
   "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",
      "tokenizers version: 0.22.2\n"
     ]
    }
   ],
   "source": [
    "from importlib.metadata import version\n",
    "\n",
    "pkgs = [\n",
    "    \"matplotlib\", \n",
    "    \"numpy\", \n",
    "    \"torch\",\n",
    "    \"tokenizers\",       # to implement the Qwen3 tokenizer\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 Qwen3 to generate text"
   ]
  },
  {
   "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",
    "### Qwen3 Code\n",
    "######################\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",
    "class RMSNorm(nn.Module):\n",
    "    def __init__(self, emb_dim, eps=1e-6, bias=False, qwen3_compatible=True):\n",
    "        super().__init__()\n",
    "        self.eps = eps\n",
    "        self.qwen3_compatible = qwen3_compatible\n",
    "        self.scale = nn.Parameter(torch.ones(emb_dim))\n",
    "        self.shift = nn.Parameter(torch.zeros(emb_dim)) if bias else None\n",
    "\n",
    "    def forward(self, x):\n",
    "        input_dtype = x.dtype\n",
    "\n",
    "        if self.qwen3_compatible:\n",
    "            x = x.to(torch.float32)\n",
    "\n",
    "        variance = x.pow(2).mean(dim=-1, keepdim=True)\n",
    "        norm_x = x * torch.rsqrt(variance + self.eps)\n",
    "        norm_x = norm_x * self.scale\n",
    "\n",
    "        if self.shift is not None:\n",
    "            norm_x = norm_x + self.shift\n",
    "\n",
    "        return norm_x.to(input_dtype)\n",
    "\n",
    "\n",
    "def compute_rope_params(head_dim, theta_base=10_000, context_length=4096, 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",
    "    # 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, num_heads, num_kv_groups, head_dim=None, qk_norm=False, dtype=None\n",
    "    ):\n",
    "        super().__init__()\n",
    "        assert num_heads % num_kv_groups == 0, \"num_heads must be divisible by num_kv_groups\"\n",
    "\n",
    "        self.num_heads = num_heads\n",
    "        self.num_kv_groups = num_kv_groups\n",
    "        self.group_size = num_heads // num_kv_groups\n",
    "\n",
    "        if head_dim is None:\n",
    "            assert d_in % num_heads == 0, \"`d_in` must be divisible by `num_heads` if `head_dim` is not set\"\n",
    "            head_dim = d_in // num_heads\n",
    "\n",
    "        self.head_dim = head_dim\n",
    "        self.d_out = num_heads * head_dim\n",
    "\n",
    "        self.W_query = nn.Linear(d_in, self.d_out, bias=False, dtype=dtype)\n",
    "        self.W_key = nn.Linear(d_in, num_kv_groups * head_dim, bias=False, dtype=dtype)\n",
    "        self.W_value = nn.Linear(d_in, num_kv_groups * head_dim, bias=False, dtype=dtype)\n",
    "\n",
    "        self.out_proj = nn.Linear(self.d_out, d_in, bias=False, dtype=dtype)\n",
    "\n",
    "        if qk_norm:\n",
    "            self.q_norm = RMSNorm(head_dim, eps=1e-6)\n",
    "            self.k_norm = RMSNorm(head_dim, eps=1e-6)\n",
    "        else:\n",
    "            self.q_norm = self.k_norm = None\n",
    "\n",
    "    def forward(self, x, mask, cos, sin):\n",
    "        b, num_tokens, _ = x.shape\n",
    "\n",
    "        # Apply projections\n",
    "        queries = self.W_query(x)  # (b, num_tokens, num_heads * head_dim)\n",
    "        keys = self.W_key(x)       # (b, num_tokens, num_kv_groups * head_dim)\n",
    "        values = self.W_value(x)   # (b, num_tokens, num_kv_groups * head_dim)\n",
    "\n",
    "        # Reshape\n",
    "        queries = queries.view(b, num_tokens, self.num_heads, self.head_dim).transpose(1, 2)\n",
    "        keys = keys.view(b, num_tokens, self.num_kv_groups, self.head_dim).transpose(1, 2)\n",
    "        values = values.view(b, num_tokens, self.num_kv_groups, self.head_dim).transpose(1, 2)\n",
    "\n",
    "        # Optional normalization\n",
    "        if self.q_norm:\n",
    "            queries = self.q_norm(queries)\n",
    "        if self.k_norm:\n",
    "            keys = self.k_norm(keys)\n",
    "\n",
    "        # Apply RoPE\n",
    "        queries = apply_rope(queries, cos, sin)\n",
    "        keys = apply_rope(keys, cos, sin)\n",
    "\n",
    "        # Expand K and V to match number of heads\n",
    "        keys = keys.repeat_interleave(self.group_size, dim=1)\n",
    "        values = values.repeat_interleave(self.group_size, dim=1)\n",
    "\n",
    "        # Attention\n",
    "        attn_scores = queries @ keys.transpose(2, 3)\n",
    "        attn_scores = attn_scores.masked_fill(mask, -torch.inf)\n",
    "        attn_weights = torch.softmax(attn_scores / self.head_dim**0.5, dim=-1)\n",
    "\n",
    "        context = (attn_weights @ values).transpose(1, 2).reshape(b, num_tokens, self.d_out)\n",
    "        return self.out_proj(context)\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",
    "            num_heads=cfg[\"n_heads\"],\n",
    "            head_dim=cfg[\"head_dim\"],\n",
    "            num_kv_groups=cfg[\"n_kv_groups\"],\n",
    "            qk_norm=cfg[\"qk_norm\"],\n",
    "            dtype=cfg[\"dtype\"]\n",
    "        )\n",
    "        self.ff = FeedForward(cfg)\n",
    "        self.norm1 = RMSNorm(cfg[\"emb_dim\"], eps=1e-6)\n",
    "        self.norm2 = RMSNorm(cfg[\"emb_dim\"], eps=1e-6)\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 Qwen3Model(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 = RMSNorm(cfg[\"emb_dim\"])\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",
    "        if cfg[\"head_dim\"] is None:\n",
    "            head_dim = cfg[\"emb_dim\"] // cfg[\"n_heads\"]\n",
    "        else:\n",
    "            head_dim = cfg[\"head_dim\"]\n",
    "        cos, sin = compute_rope_params(\n",
    "            head_dim=head_dim,\n",
    "            theta_base=cfg[\"rope_base\"],\n",
    "            context_length=cfg[\"context_length\"]\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 Qwen3 \n",
    "#######################\n",
    "\n",
    "# 0.6B model\n",
    "QWEN3_CONFIG = {\n",
    "    \"vocab_size\": 151_936,           # Vocabulary size\n",
    "    \"context_length\": 40_960,        # Context length that was used to train the model\n",
    "    \"emb_dim\": 1024,                 # Embedding dimension\n",
    "    \"n_heads\": 16,                   # Number of attention heads\n",
    "    \"n_layers\": 28,                  # Number of layers\n",
    "    \"hidden_dim\": 3072,              # Size of the intermediate dimension in FeedForward\n",
    "    \"head_dim\": 128,                 # Size of the heads in GQA\n",
    "    \"qk_norm\": True,                 # Whether to normalize queries and keys in GQA\n",
    "    \"n_kv_groups\": 8,                # Key-Value groups for grouped-query attention\n",
    "    \"rope_base\": 1_000_000.0,        # The base in RoPE's \"theta\"\n",
    "    \"dtype\": torch.bfloat16,         # Lower-precision dtype to reduce memory usage\n",
    "}\n",
    "\n",
    "QWEN3_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 = Qwen3Model(QWEN3_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 re\n",
    "from tokenizers import Tokenizer\n",
    "\n",
    "class Qwen3Tokenizer:\n",
    "    _SPECIALS = [\n",
    "        \"<|endoftext|>\",\n",
    "        \"<|im_start|>\", \"<|im_end|>\",\n",
    "        \"<|object_ref_start|>\", \"<|object_ref_end|>\",\n",
    "        \"<|box_start|>\", \"<|box_end|>\",\n",
    "        \"<|quad_start|>\", \"<|quad_end|>\",\n",
    "        \"<|vision_start|>\", \"<|vision_end|>\",\n",
    "        \"<|vision_pad|>\", \"<|image_pad|>\", \"<|video_pad|>\",\n",
    "        \"<think>\", \"</think>\"\n",
    "    ]\n",
    "    _SPLIT_RE = re.compile(r\"(<\\|[^>]+?\\|>|<think>|</think>)\")\n",
    "\n",
    "    def __init__(self, tokenizer_file_path=\"tokenizer.json\", repo_id=None,\n",
    "                 apply_chat_template=True, add_generation_prompt=False, add_thinking=False):\n",
    "\n",
    "        self.apply_chat_template = apply_chat_template\n",
    "        self.add_generation_prompt = add_generation_prompt\n",
    "        self.add_thinking = add_thinking\n",
    "\n",
    "        tok_file = Path(tokenizer_file_path)\n",
    "        self._tok = Tokenizer.from_file(str(tok_file))\n",
    "        self._special_to_id = {}\n",
    "        for t in self._SPECIALS:\n",
    "            tid = self._tok.token_to_id(t)\n",
    "            if tid is not None:\n",
    "                self._special_to_id[t] = tid\n",
    "\n",
    "        self.pad_token_id = self._special_to_id[\"<|endoftext|>\"]\n",
    "        self.eos_token_id = self.pad_token_id\n",
    "\n",
    "        if repo_id and \"Base\" not in repo_id:\n",
    "            eos_token = \"<|im_end|>\"\n",
    "        else:\n",
    "            eos_token = \"<|endoftext|>\"\n",
    "        if eos_token in self._special_to_id:\n",
    "            self.eos_token_id = self._special_to_id[eos_token]\n",
    "\n",
    "    def encode(self, text, chat_wrapped=None):\n",
    "        if chat_wrapped is None:\n",
    "            chat_wrapped = self.apply_chat_template\n",
    "\n",
    "        stripped = text.strip()\n",
    "        if stripped in self._special_to_id and \"\\n\" not in stripped:\n",
    "            return [self._special_to_id[stripped]]\n",
    "\n",
    "        if chat_wrapped:\n",
    "            text = self._wrap_chat(text)\n",
    "\n",
    "        ids = []\n",
    "        for part in filter(None, self._SPLIT_RE.split(text)):\n",
    "            if part in self._special_to_id:\n",
    "                ids.append(self._special_to_id[part])\n",
    "            else:\n",
    "                ids.extend(self._tok.encode(part).ids)\n",
    "        return ids\n",
    "\n",
    "    def decode(self, ids):\n",
    "        return self._tok.decode(ids, skip_special_tokens=False)\n",
    "\n",
    "    def _wrap_chat(self, user_msg):\n",
    "        s = f\"<|im_start|>user\\n{user_msg}<|im_end|>\\n\"\n",
    "        if self.add_generation_prompt:\n",
    "            s += \"<|im_start|>assistant\"\n",
    "            if self.add_thinking:\n",
    "                s += \"\\n\"\n",
    "            else:\n",
    "                s += \"\\n<think>\\n\\n</think>\\n\\n\"\n",
    "        return s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "f4b220c0-038b-4c79-9506-a04065331218",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "from huggingface_hub import hf_hub_download\n",
    "\n",
    "repo_id = \"Qwen/Qwen3-0.6B-Base\"\n",
    "tokenizer_file_path = \"tokenizer.json\"\n",
    "local_dir = \".\"\n",
    "\n",
    "hf_hub_download(\n",
    "    repo_id=repo_id,\n",
    "    filename=\"tokenizer.json\",\n",
    "    local_dir=local_dir,\n",
    ")\n",
    "\n",
    "tokenizer = Qwen3Tokenizer(\n",
    "    tokenizer_file_path=tokenizer_file_path,\n",
    "    repo_id=repo_id,\n",
    "    apply_chat_template=False,\n",
    "    add_generation_prompt=False,\n",
    "    add_thinking=False\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "5e062b82-3540-48ce-8eb4-009686d0d16c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output text:\n",
      " Every effort moves you disparities.phasemanuelScheduledpageIndex=zerosuneiapus_VE.Default\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=QWEN3_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": 8,
   "id": "654fde37-b2a9-4a20-a8d3-0206c056e2ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "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": 9,
   "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": 10,
   "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": 11,
   "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: 4943\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": 12,
   "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=QWEN3_CONFIG[\"train_context_length\"],\n",
    "    stride=QWEN3_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=QWEN3_CONFIG[\"train_context_length\"],\n",
    "    stride=QWEN3_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": 13,
   "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": 14,
   "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": 15,
   "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": 16,
   "id": "56f5b0c9-1065-4d67-98b9-010e42fc1e2a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using cuda device.\n",
      "Training loss: 12.09375\n",
      "Validation loss: 12.0625\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": 17,
   "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": 18,
   "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.725, Val loss 11.000\n",
      "Ep 1 (Step 000005): Train loss 8.950, Val loss 9.250\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 2 (Step 000010): Train loss 7.119, Val loss 7.781\n",
      "Ep 2 (Step 000015): Train loss 6.550, Val loss 7.156\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\n",
      "Ep 3 (Step 000020): Train loss 6.419, Val loss 6.938\n",
      "Every effort moves you,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,\n",
      "Ep 4 (Step 000025): Train loss 6.400, Val loss 6.938\n",
      "Ep 4 (Step 000030): Train loss 6.350, Val loss 7.000\n",
      "Every effort moves you..................................................\n",
      "Ep 5 (Step 000035): Train loss 6.181, Val loss 6.969\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\n",
      "Ep 6 (Step 000040): Train loss 6.037, Val loss 7.000\n",
      "Ep 6 (Step 000045): Train loss 6.056, Val loss 6.969\n",
      "Every effort moves you, and 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, and the the the the the the, and the the\n",
      "Ep 7 (Step 000050): Train loss 5.900, Val loss 6.906\n",
      "Ep 7 (Step 000055): Train loss 5.806, Val loss 6.750\n",
      "Every effort moves you, I had I had the the I had I had I had I had I had I had I had the the the the the I had I had I had I had I had I had I had I had I had the the the the the the\n",
      "Ep 8 (Step 000060): Train loss 6.000, Val loss 6.938\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 9 (Step 000065): Train loss 6.062, Val loss 6.875\n",
      "Ep 9 (Step 000070): Train loss 5.981, Val loss 6.875\n",
      "Every effort moves you, and, and, and. Gisburn. Gisburn. Gisburn. Gisburn. Gisburn. Gisburn. Gisburn. Gisburn. Gisburn. Gisburn. Gisburn\n",
      "Ep 10 (Step 000075): Train loss 5.825, Val loss 6.781\n",
      "Every effort moves you, and the fact I had the fact the fact of the fact of the fact of the fact of the fact of the fact of the fact of the fact of the fact of the fact of the fact of the fact of the fact of the fact of\n",
      "Ep 11 (Step 000080): Train loss 5.700, Val loss 6.719\n",
      "Ep 11 (Step 000085): Train loss 5.525, Val loss 6.625\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\n",
      "Ep 12 (Step 000090): Train loss 5.456, Val loss 6.688\n",
      "Ep 12 (Step 000095): Train loss 5.369, Val loss 6.688\n",
      "Every effort moves you--as the fact--as the fact--as I was the fact, and I was the fact, and I was the fact, and I was the fact, and the fact, and the fact, and the fact, and the fact, and\n",
      "Ep 13 (Step 000100): Train loss 5.081, Val loss 6.656\n",
      "Every effort moves you--I had been the last--I had been--I had been--I had been--I had been--I had been--I had been--I had been--I had been--I had been--I had been--I had been\n",
      "Ep 14 (Step 000105): Train loss 4.925, Val loss 6.531\n",
      "Ep 14 (Step 000110): Train loss 4.384, Val loss 6.406\n",
      "Every effort moves you to have been the picture--as and Mrs. And, and established, and untouched, and untouched, and untouched of the picture, and Mrs. And, and untouched, and untouched, and untouched, and untouched, and untouched of the picture,\n",
      "Ep 15 (Step 000115): Train loss 4.131, Val loss 6.312\n",
      "Every effort moves you know to have been his own, and established, and established, and established, and established of the picture--his fair sitters had been denied his glory of the picture.  \"Oh, and established the picture.  \"Oh, and established the picture.  \"Oh\n",
      "Ep 16 (Step 000120): Train loss 3.828, Val loss 6.375\n",
      "Ep 16 (Step 000125): Train loss 3.188, Val loss 6.531\n",
      "Every effort moves you know, and in the fact, and in the fact, and in the fact, and in the fact, and in the fact, and in the fact, and in the fact, and in the fact, and in the fact, and in the\n",
      "Ep 17 (Step 000130): Train loss 2.712, Val loss 6.531\n",
      "Ep 17 (Step 000135): Train loss 2.345, Val loss 6.656\n",
      "Every effort moves you--as a little under a little under a little under a little under a little under a little under a little under a little under a little under a little under a little under a little under a little under a little under a little under a little under\n",
      "Ep 18 (Step 000140): Train loss 2.147, Val loss 6.656\n",
      "Every effort moves you know, and he had been his pictures--and that he had been through, and he had been through, and he had been his pictures--and that he had been through, and he had been through, and he had been no one of the\n",
      "Ep 19 (Step 000145): Train loss 1.778, Val loss 6.969\n",
      "Ep 19 (Step 000150): Train loss 1.627, Val loss 6.938\n",
      "Every effort moves you know. Gisburn had been taken. Gisburn had been taken. Gisburn had been taken. Gisburn had been taken. Gisburn had been taken. Gisburn had been taken. Gisburn had been taken\n",
      "Ep 20 (Step 000155): Train loss 1.670, Val loss 6.906\n",
      "Every effort moves you--as I had to see it, and I had to see a little quickly.  \"Oh, and I had a little quickly.  \"Oh, and I had a little quickly.  \"Oh, and I had a little quickly.  \"Oh, and I had a\n",
      "Ep 21 (Step 000160): Train loss 1.436, Val loss 7.094\n",
      "Ep 21 (Step 000165): Train loss 1.336, Val loss 7.062\n",
      "Every effort moves you know, and said: \"If you can just manage to see it. Gisburn--as such--as such--as such--as such--as such--as Jack's not look upon its like again\"?  Well!--even through it to\n",
      "Ep 22 (Step 000170): Train loss 1.231, Val loss 7.125\n",
      "Ep 22 (Step 000175): Train loss 1.147, Val loss 7.250\n",
      "Every effort moves you ever, and twirling between his part, and straining, and straining, and straining, and said: \"If you stand here you can just manage to see it. Gisburn, and le it over the Riviera, and\n",
      "Ep 23 (Step 000180): Train loss 1.080, Val loss 7.375\n",
      "Every effort moves you know. I had to see the fact that, and he knew just said--I made a little too much longer to see your portrait.  \"Oh, I had been no one of the picture to me. I had the donkey hanging on my eyes\n",
      "Ep 24 (Step 000185): Train loss 0.897, Val loss 7.375\n",
      "Ep 24 (Step 000190): Train loss 0.807, Val loss 7.500\n",
      "Every effort moves you know. Gisburn--as such--because he had married her--because he had been hard to go on painting; but it might be interesting to go on painting because he had given up his painting because he had not led him down, and\n",
      "Ep 25 (Step 000195): Train loss 0.705, Val loss 7.562\n",
      "Every effort moves you know. Gisburn--as such--had not existed till nearly a year after Jack's resolve had been taken. It might be that he had married her--because he liked his ease--because he didn't want to go on painting; but\n",
      "Ep 26 (Step 000200): Train loss 0.615, Val loss 7.531\n",
      "Ep 26 (Step 000205): Train loss 0.574, Val loss 7.625\n",
      "Every effort moves you ever, and said: \"If you stand here you can just manage to see it. Gisburn, and he had always been his fate to see it--his last Chicago sitter--the first portrait of beauty.\"  Poor Jack's I had\n",
      "Ep 27 (Step 000210): Train loss 0.515, Val loss 7.562\n",
      "Ep 27 (Step 000215): Train loss 0.386, Val loss 7.562\n",
      "Every effort moves you ever dabble with a dark plain room at the end of the florid vista. Gisburn--as such--though a good-humoured surprise.  \"Oh, I asked, I asked abruptly.  \"Oh, I asked abruptly.  \"Oh, I\n",
      "Ep 28 (Step 000220): Train loss 0.371, Val loss 7.594\n",
      "Every effort moves you ever dabble with a good fellow--a Stroud!\" I felt as if he was no great surprise, and his close behind his close behind his close grayish beard--as you know, I felt as if he was a little quickly.  \"\n",
      "Ep 29 (Step 000225): Train loss 0.299, Val loss 7.594\n",
      "Ep 29 (Step 000230): Train loss 0.271, Val loss 7.625\n",
      "Every effort moves you ever dabble with a dark plain room at the end of the florid vista. It was square and brown and leathery: no \"effects\"; no bric-a-brac, and brown and leathery: \"effects\"; no br\n",
      "Ep 30 (Step 000235): Train loss 0.288, Val loss 7.656\n",
      "Every effort moves you ever dabble with a shaking hand, and straining, a later day.  \"By Jove!\" I asked, I didn't you know you know you know. Gisburn, who had always been Rome or Florence.)  \"The height of his\n",
      "Ep 31 (Step 000240): Train loss 0.194, Val loss 7.625\n",
      "Ep 31 (Step 000245): Train loss 0.198, Val loss 7.688\n",
      "Every effort moves you know. The women had to see it,\" she began, as if excusing herself. He shrugged his shoulders, still smiling.  \"Oh, he said, he knew me--of forcing it to see a purblind public. And whenever my wonder paid\n",
      "Ep 32 (Step 000250): Train loss 0.177, Val loss 7.750\n",
      "Ep 32 (Step 000255): Train loss 0.129, Val loss 7.750\n",
      "Every effort moves you know, you know. He says they're not fit to have about; he's sent them all away except one--my portrait--and that I have to keep upstairs.\"  His ridiculous modesty--Jack's modesty about his pictures? My curiosity\n",
      "Ep 33 (Step 000260): Train loss 0.111, Val loss 7.781\n",
      "Every effort moves you know. The women had made him, still smiling.  \"Oh, I didn't let a fragment groping in him--and by a dozen lines--so handsome, I didn't--I didn't care a good-humoured shrug.  \"Oh\n",
      "Ep 34 (Step 000265): Train loss 0.092, Val loss 7.844\n",
      "Ep 34 (Step 000270): Train loss 0.066, Val loss 7.875\n",
      "Every effort moves you know.\"  I HAD always thought Jack Gisburn rather a cheap genius--though a good fellow enough--so it was no great surprise to me to hear that, in the height of his glory, he had dropped his painting, married a rich\n",
      "Ep 35 (Step 000275): Train loss 0.052, Val loss 7.875\n",
      "Every effort moves you know. Rickham wanted to see it,\" she began, as if excusing herself. He shrugged his shoulders, still smiling.  \"Oh, Rickham found me out long ago,\" he said lightly; then, passing his arm through mine: \"Come\n",
      "Ep 36 (Step 000280): Train loss 0.052, Val loss 7.875\n",
      "Ep 36 (Step 000285): Train loss 0.040, Val loss 7.938\n",
      "Every effort moves you know, pushed an arm-chair away, and said: \"If you stand here you can just manage to see it. I had it over the mantel-piece, but he wouldn't let it stay.\"  Yes--I could just manage to see it\n",
      "Ep 37 (Step 000290): Train loss 0.027, Val loss 7.938\n",
      "Ep 37 (Step 000295): Train loss 0.029, Val loss 7.938\n",
      "Every effort moves you know, pushed an arm-chair away, and said: \"If you stand here you can just manage to see it. I had it over the mantel-piece, but he wouldn't let it stay.\"  Yes--I could just manage to see it\n",
      "Ep 38 (Step 000300): Train loss 0.022, Val loss 7.938\n",
      "Every effort moves you know.\"  I HAD always thought Jack Gisburn rather a cheap genius--though a good fellow enough--so it was no great surprise to me to hear that, in the height of his glory, he had dropped his painting, married a rich\n",
      "Ep 39 (Step 000305): Train loss 0.016, Val loss 7.969\n",
      "Ep 39 (Step 000310): Train loss 0.010, Val loss 7.969\n",
      "Every effort moves you know.\"  I HAD always thought Jack Gisburn rather a cheap genius--though a good fellow enough--so it was no great surprise to me to hear that, in the height of his glory, he had dropped his painting, married a rich\n",
      "Ep 40 (Step 000315): Train loss 0.011, Val loss 8.000\n",
      "Every effort moves you know, pushed an arm-chair away, and said: \"If you stand here you can just manage to see it. I had it over the mantel-piece, but he wouldn't let it stay.\"  Yes--I could just manage to see it\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 = Qwen3Model(QWEN3_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": 19,
   "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, 33232 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": 20,
   "id": "2734cee0-f6f9-42d5-b71c-fa7e0ef28b6d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output text:\n",
      " Every effort moves you know.\"\n",
      "\n",
      "I HAD always thought Jack Gisburn rather a cheap genius--though a good fellow enough--so it was\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=QWEN3_CONFIG[\"train_context_length\"]\n",
    ")\n",
    "\n",
    "print(\"Output text:\\n\", token_ids_to_text(token_ids, tokenizer))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "bf2e432d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output text:\n",
      " Painting his ridiculous modesty, you know. He says they're not fit to have about; he's sent them all away except\n"
     ]
    }
   ],
   "source": [
    "token_ids = generate_text_simple(\n",
    "    model=model,\n",
    "    idx=text_to_token_ids(\"Painting\", tokenizer).to(inference_device),\n",
    "    max_new_tokens=25,\n",
    "    context_size=QWEN3_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
}
