{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e1b280ab-b61f-4d1a-bf7e-44e5f9ed3a5c",
   "metadata": {
    "id": "e1b280ab-b61f-4d1a-bf7e-44e5f9ed3a5c"
   },
   "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": "efde77f2-6af3-4781-8597-89ecd3f41a52",
   "metadata": {
    "id": "efde77f2-6af3-4781-8597-89ecd3f41a52"
   },
   "source": [
    "# Olmo 3 From Scratch (A Standalone Notebook)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55cdef4d-de59-4a65-89f9-fa2a8ef3471d",
   "metadata": {
    "id": "55cdef4d-de59-4a65-89f9-fa2a8ef3471d"
   },
   "source": [
    "- This notebook is purposefully minimal and focuses on the code to re-implement Olmo 3 7B and 32 models from Allen AI in pure PyTorch without relying on other external LLM libraries; Olmo 3 is interesting because it is currently the leading fully open-source model\n",
    "- For more information, see the official [Olmo 3 announcement](https://allenai.org/blog/olmo3) and model cards:\n",
    "  - [Olmo-3-1025-7B](https://huggingface.co/allenai/Olmo-3-1025-7B) (base model)\n",
    "  - [Olmo-3-7B-Instruct](https://huggingface.co/allenai/Olmo-3-7B-Instruct)\n",
    "  - [Olmo-3-7B-Think](https://huggingface.co/allenai/Olmo-3-7B-Think)\n",
    "- Note that there are also 32B versions, which are not listed above for brevity; you can find a complete list [here](https://huggingface.co/collections/allenai/olmo-3-post-training)\n",
    "- Below is a side-by-side comparison with Qwen3 8B as a reference model; if you are interested in the Qwen3 0.6B standalone notebook, you can find it [here](../11_qwen3)\n",
    "<br>\n",
    "\n",
    "<img src=\"https://sebastianraschka.com/images/LLMs-from-scratch-images/bonus/olmo3/olmo3.webp\">\n",
    "  \n",
    "<img src=\"https://sebastianraschka.com/images/LLMs-from-scratch-images/bonus/olmo3/olmo3-pipeline.webp\">\n",
    "  \n",
    "  \n",
    "- About the code:\n",
    "  - all code is my own code, mapping the Olmo 3 architecture onto the model code implemented in my [Build A Large Language Model (From Scratch)](http://mng.bz/orYv) book; the code is released under a permissive open-source Apache 2.0 license (see [LICENSE.txt](https://github.com/rasbt/LLMs-from-scratch/blob/main/LICENSE.txt))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "7c201adb-747e-437b-9a62-442802941e01",
   "metadata": {},
   "outputs": [],
   "source": [
    "# pip install -r https://raw.githubusercontent.com/rasbt/LLMs-from-scratch/refs/heads/main/ch05/07_gpt_to_llama/requirements-extra.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "dd1b65a8-4301-444a-bd7c-a6f2bd1df9df",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "dd1b65a8-4301-444a-bd7c-a6f2bd1df9df",
    "outputId": "4f762354-e0a3-4cc2-e5d4-e61a227a202c"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "huggingface_hub version: 0.35.0\n",
      "tokenizers version: 0.22.1\n",
      "torch version: 2.9.1+cu130\n"
     ]
    }
   ],
   "source": [
    "from importlib.metadata import version\n",
    "\n",
    "pkgs = [\n",
    "    \"huggingface_hub\",  # to download pretrained weights\n",
    "    \"tokenizers\",       # to implement the tokenizer\n",
    "    \"torch\",            # to implement the model\n",
    "]\n",
    "for p in pkgs:\n",
    "    print(f\"{p} version: {version(p)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07e96fbb-8e16-4f6d-835f-c6159321280b",
   "metadata": {},
   "source": [
    "- Note that there are three model types, and each of the four model types comes in a 7B and 32B size:\n",
    "1. Base (`Olmo-3-1025-7B` and `Olmo-3-1125-32B`)\n",
    "2. Instruct (`Olmo-3-7B/32B-Think`)\n",
    "3. Reasoning (`Olmo-3-32B/7B-Think`)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "70a90338-624a-4706-aa55-6b4358070194",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Select which model to use\n",
    "\n",
    "# USE_MODEL = \"Olmo-3-1025-7B\"\n",
    "# USE_MODEL = \"Olmo-3-1125-32B\"\n",
    "USE_MODEL = \"Olmo-3-7B-Instruct\"\n",
    "# USE_MODEL = \"Olmo-3-32B-Instruct\"\n",
    "# USE_MODEL = \"Olmo-3-7B-Think\"\n",
    "# USE_MODEL = \"Olmo-3-32B-Think\"\n",
    "# USE_MODEL = \"Olmo-3-7B-RLZero-IF\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f258cb74-1c4e-4880-8772-3c85fb920811",
   "metadata": {},
   "source": [
    "- In addition to the checkpoints listed above, you can also use the intermediate checkpoints listed [here](https://huggingface.co/collections/allenai/olmo-3-post-training); since they all have the same architecture, they are all compatible with this notebook"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "653410a6-dd2b-4eb2-a722-23d9782e726d",
   "metadata": {
    "id": "653410a6-dd2b-4eb2-a722-23d9782e726d"
   },
   "source": [
    "&nbsp;\n",
    "# 1. Architecture code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "82076c21-9331-4dcd-b017-42b046cf1a60",
   "metadata": {
    "id": "82076c21-9331-4dcd-b017-42b046cf1a60"
   },
   "outputs": [],
   "source": [
    "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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "56715760-37e1-433e-89da-04864c139a9e",
   "metadata": {},
   "outputs": [],
   "source": [
    "class RMSNorm(nn.Module):\n",
    "    def __init__(self, emb_dim, eps=1e-6):\n",
    "        super().__init__()\n",
    "        self.eps = eps\n",
    "        self.weight = nn.Parameter(torch.ones(emb_dim))\n",
    "\n",
    "    def forward(self, x):\n",
    "        input_dtype = x.dtype\n",
    "        x_f = x.float()\n",
    "        var = x_f.pow(2).mean(dim=-1, keepdim=True)\n",
    "        x_norm = x_f * torch.rsqrt(var + self.eps)\n",
    "        return (self.weight * x_norm).to(input_dtype)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "4b9a346f-5826-4083-9162-abd56afc03f0",
   "metadata": {
    "id": "4b9a346f-5826-4083-9162-abd56afc03f0"
   },
   "outputs": [],
   "source": [
    "import math\n",
    "\n",
    "\n",
    "def compute_rope_params(head_dim, theta_base=10_000, context_length=4096, attention_factor=1.0, rope_type=\"default\", rope_factor=1.0, rope_orig_max=8192, beta_fast=32.0, beta_slow=1.0, dtype=torch.float32):\n",
    "    assert head_dim % 2 == 0, \"Embedding dimension must be even\"\n",
    "\n",
    "    if rope_type == \"yarn\":\n",
    "        # Compute YaRN-style frequency scaling (as per https://huggingface.co/papers/2309.00071)\n",
    "\n",
    "        def find_correction_dim(num_rotations, dim, base, max_position_embeddings):\n",
    "            \"\"\"Inverse dimension formula to find the dimension based on the number of rotations\"\"\"\n",
    "            return (dim * math.log(max_position_embeddings / (num_rotations * 2 * math.pi))) / (2 * math.log(base))\n",
    "\n",
    "        def find_correction_range(low_rot, high_rot, dim, base, max_position_embeddings):\n",
    "            \"\"\"Find dimension range bounds based on rotations\"\"\"\n",
    "            low = find_correction_dim(low_rot, dim, base, max_position_embeddings)\n",
    "            high = find_correction_dim(high_rot, dim, base, max_position_embeddings)\n",
    "            low = math.floor(low)\n",
    "            high = math.ceil(high)\n",
    "            return max(low, 0), min(high, dim - 1)\n",
    "\n",
    "        def linear_ramp_factor(min_val, max_val, dim):\n",
    "            if min_val == max_val:\n",
    "                max_val += 0.001  # Prevent singularity\n",
    "            linear_func = (torch.arange(dim, dtype=torch.float32) - min_val) / (max_val - min_val)\n",
    "            ramp_func = torch.clamp(linear_func, 0, 1)\n",
    "            return ramp_func\n",
    "\n",
    "        # Base frequencies\n",
    "        pos_freqs = theta_base ** (torch.arange(0, head_dim, 2, dtype=dtype) / head_dim)\n",
    "        inv_freq_extrapolation = 1.0 / pos_freqs  # No scaling (extrapolation)\n",
    "        inv_freq_interpolation = 1.0 / (rope_factor * pos_freqs)  # With scaling (interpolation)\n",
    "\n",
    "        # Find the range where we blend between interpolation and extrapolation\n",
    "        low, high = find_correction_range(beta_fast, beta_slow, head_dim, theta_base, rope_orig_max)\n",
    "\n",
    "        # Get n-dimensional rotational scaling corrected for extrapolation\n",
    "        inv_freq_extrapolation_factor = 1 - linear_ramp_factor(low, high, head_dim // 2).to(dtype=dtype)\n",
    "        inv_freq = (\n",
    "            inv_freq_interpolation * (1 - inv_freq_extrapolation_factor)\n",
    "            + inv_freq_extrapolation * inv_freq_extrapolation_factor\n",
    "        )\n",
    "    else:\n",
    "        # Default RoPE\n",
    "        inv_freq = 1.0 / (\n",
    "            theta_base ** (\n",
    "                torch.arange(0, head_dim, 2, dtype=dtype)[: head_dim // 2].float()\n",
    "                / head_dim\n",
    "            )\n",
    "        )\n",
    "\n",
    "    # Generate position indices\n",
    "    positions = torch.arange(context_length, dtype=dtype)\n",
    "\n",
    "    # Compute the base angles (shape: [context_length, head_dim // 2])\n",
    "    angles = positions.unsqueeze(1) * inv_freq.unsqueeze(0)\n",
    "\n",
    "    # Expand to full head_dim (shape: [context_length, head_dim])\n",
    "    angles = torch.cat([angles, angles], dim=1)\n",
    "\n",
    "    # Precompute sine and cosine\n",
    "    cos = torch.cos(angles) * attention_factor\n",
    "    sin = torch.sin(angles) * attention_factor\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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "e8169ab5-f976-4222-a2e1-eb1cabf267cb",
   "metadata": {
    "id": "e8169ab5-f976-4222-a2e1-eb1cabf267cb"
   },
   "outputs": [],
   "source": [
    "class GroupedQueryAttention(nn.Module):\n",
    "    def __init__(self, d_in, num_heads, num_kv_groups, head_dim, attention_bias=False, dtype=None, sliding_window=None, attn_type=\"full_attention\"):\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",
    "        self.head_dim = head_dim\n",
    "        self.d_out = num_heads * head_dim\n",
    "        self.attn_type = attn_type\n",
    "        self.sliding_window = sliding_window if attn_type == \"sliding_attention\" else None\n",
    "\n",
    "        # Projections\n",
    "        self.W_query = nn.Linear(d_in, self.d_out, bias=attention_bias, dtype=dtype)\n",
    "        self.W_key = nn.Linear(d_in, num_kv_groups * head_dim, bias=attention_bias, dtype=dtype)\n",
    "        self.W_value = nn.Linear(d_in, num_kv_groups * head_dim, bias=attention_bias, dtype=dtype)\n",
    "        self.out_proj = nn.Linear(self.d_out, d_in, bias=attention_bias, dtype=dtype)\n",
    "\n",
    "        # Olmo3-style RMSNorm over the flattened projections\n",
    "        self.q_norm = RMSNorm(self.d_out)\n",
    "        self.k_norm = RMSNorm(num_kv_groups * head_dim)\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",
    "        # Normalize q and k\n",
    "        queries = self.q_norm(queries)\n",
    "        keys = self.k_norm(keys)\n",
    "\n",
    "        # Reshape to (b, heads, seq, head_dim)\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",
    "        # Apply RoPE\n",
    "        queries = apply_rope(queries, cos, sin)\n",
    "        keys = apply_rope(keys, cos, sin)\n",
    "\n",
    "        # Expand KV groups to full head count\n",
    "        if self.group_size > 1:\n",
    "            keys = keys.repeat_interleave(self.group_size, dim=1)\n",
    "            values = values.repeat_interleave(self.group_size, dim=1)\n",
    "\n",
    "        # Scaling before the matmul seems to be a bit more stable for Olmo\n",
    "        scale = self.head_dim ** -0.5  # Python float\n",
    "        queries = queries * scale\n",
    "        \n",
    "        # Attention\n",
    "        attn_scores = queries @ keys.transpose(2, 3)\n",
    "        if mask is not None:\n",
    "            attn_scores = attn_scores.masked_fill(mask, -torch.inf)\n",
    "\n",
    "        attn_weights = torch.softmax(attn_scores, dim=-1)\n",
    "        context = (attn_weights @ values).transpose(1, 2).reshape(b, num_tokens, self.d_out)\n",
    "        return self.out_proj(context)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "457cb2f8-50c1-4045-8a74-f181bfb5fea9",
   "metadata": {
    "id": "457cb2f8-50c1-4045-8a74-f181bfb5fea9"
   },
   "outputs": [],
   "source": [
    "class TransformerBlock(nn.Module):\n",
    "    def __init__(self, cfg, attn_type):\n",
    "        super().__init__()\n",
    "        self.attn_type = attn_type\n",
    "        self.att = GroupedQueryAttention(\n",
    "            d_in=cfg[\"emb_dim\"],\n",
    "            num_heads=cfg[\"n_heads\"],\n",
    "            num_kv_groups=cfg[\"n_kv_heads\"],\n",
    "            head_dim=cfg[\"head_dim\"],\n",
    "            attention_bias=cfg[\"attention_bias\"],\n",
    "            dtype=cfg[\"dtype\"],\n",
    "            sliding_window=cfg[\"sliding_window\"],\n",
    "            attn_type=attn_type,\n",
    "        )\n",
    "        self.ff = FeedForward(cfg)\n",
    "        self.post_attention_layernorm = RMSNorm(cfg[\"emb_dim\"], eps=cfg[\"rms_norm_eps\"])\n",
    "        self.post_feedforward_layernorm = RMSNorm(cfg[\"emb_dim\"], eps=cfg[\"rms_norm_eps\"])\n",
    "\n",
    "    def forward(self, x, mask_global, mask_local, cos, sin):\n",
    "        attn_mask = mask_local if self.attn_type == \"sliding_attention\" else mask_global\n",
    "\n",
    "        shortcut = x\n",
    "        x_attn = self.att(x, attn_mask, cos, sin)\n",
    "        x_attn = self.post_attention_layernorm(x_attn)\n",
    "        x = shortcut + x_attn\n",
    "\n",
    "        shortcut = x\n",
    "        x_ffn = self.ff(x)\n",
    "        x_ffn = self.post_feedforward_layernorm(x_ffn)\n",
    "        x = shortcut + x_ffn\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "e88de3e3-9f07-42cc-816b-28dbd46e96c4",
   "metadata": {
    "id": "e88de3e3-9f07-42cc-816b-28dbd46e96c4"
   },
   "outputs": [],
   "source": [
    "class Olmo3Model(nn.Module):\n",
    "    def __init__(self, cfg):\n",
    "        super().__init__()\n",
    "        assert cfg[\"layer_types\"] is not None and len(cfg[\"layer_types\"]) == cfg[\"n_layers\"]\n",
    "\n",
    "        self.tok_emb = nn.Embedding(cfg[\"vocab_size\"], cfg[\"emb_dim\"], dtype=cfg[\"dtype\"])\n",
    "        self.blocks = nn.ModuleList([TransformerBlock(cfg, attn_type) for attn_type in cfg[\"layer_types\"]])\n",
    "        self.final_norm = RMSNorm(cfg[\"emb_dim\"], eps=cfg[\"rms_norm_eps\"])\n",
    "        self.out_head = nn.Linear(cfg[\"emb_dim\"], cfg[\"vocab_size\"], bias=False, dtype=cfg[\"dtype\"])\n",
    "        self.cfg = cfg\n",
    "\n",
    "        cos, sin = compute_rope_params(\n",
    "            head_dim=cfg[\"head_dim\"],\n",
    "            context_length=cfg[\"context_length\"],\n",
    "            theta_base=cfg[\"rope_base\"],\n",
    "            attention_factor=cfg[\"rope_attention_factor\"],\n",
    "            rope_type=cfg[\"rope_type\"],\n",
    "            rope_factor=cfg[\"rope_factor\"],\n",
    "            rope_orig_max=cfg[\"rope_orig_max\"],\n",
    "            dtype=torch.float32,\n",
    "        )\n",
    "        self.register_buffer(\"cos\", cos, persistent=False)\n",
    "        self.register_buffer(\"sin\", sin, persistent=False)\n",
    "\n",
    "    def create_masks(self, seq_len, device):\n",
    "        ones = torch.ones((seq_len, seq_len), dtype=torch.bool, device=device)\n",
    "\n",
    "        # mask_global (future is masked: j > i)\n",
    "        #     j:  0 1 2 3 4 5 6 7\n",
    "        #  i\n",
    "        #     0:  0 1 1 1 1 1 1 1\n",
    "        #     1:  0 0 1 1 1 1 1 1\n",
    "        #     2:  0 0 0 1 1 1 1 1\n",
    "        #     3:  0 0 0 0 1 1 1 1\n",
    "        #     4:  0 0 0 0 0 1 1 1\n",
    "        #     5:  0 0 0 0 0 0 1 1\n",
    "        #     6:  0 0 0 0 0 0 0 1\n",
    "        #     7:  0 0 0 0 0 0 0 0\n",
    "        mask_global = torch.triu(ones, diagonal=1)\n",
    "\n",
    "        # far_past (too far back is masked: i - j >= sliding_window)\n",
    "        # where sliding_window = 4\n",
    "        #     j:  0 1 2 3 4 5 6 7\n",
    "        #  i\n",
    "        #     0:  0 0 0 0 0 0 0 0\n",
    "        #     1:  0 0 0 0 0 0 0 0\n",
    "        #     2:  0 0 0 0 0 0 0 0\n",
    "        #     3:  0 0 0 0 0 0 0 0\n",
    "        #     4:  1 0 0 0 0 0 0 0\n",
    "        #     5:  1 1 0 0 0 0 0 0\n",
    "        #     6:  1 1 1 0 0 0 0 0\n",
    "        #     7:  1 1 1 1 0 0 0 0\n",
    "        far_past = torch.triu(ones, diagonal=self.cfg[\"sliding_window\"]).T\n",
    "\n",
    "        # Local (sliding_window) = future OR far-past\n",
    "        # mask_local\n",
    "        #     j:  0 1 2 3 4 5 6 7\n",
    "        # i\n",
    "        # 0:      0 1 1 1 1 1 1 1\n",
    "        # 1:      0 0 1 1 1 1 1 1\n",
    "        # 2:      0 0 0 1 1 1 1 1\n",
    "        # 3:      0 0 0 0 1 1 1 1\n",
    "        # 4:      1 0 0 0 0 1 1 1\n",
    "        # 5:      1 1 0 0 0 0 1 1\n",
    "        # 6:      1 1 1 0 0 0 0 1\n",
    "        # 7:      1 1 1 1 0 0 0 0\n",
    "        mask_local = mask_global | far_past\n",
    "        return mask_global, mask_local\n",
    "\n",
    "    def forward(self, input_ids):\n",
    "        b, seq_len = input_ids.shape\n",
    "        x = self.tok_emb(input_ids)\n",
    "        mask_global, mask_local = self.create_masks(seq_len, x.device)\n",
    "\n",
    "        cos = self.cos[:seq_len, :].to(x.device)\n",
    "        sin = self.sin[:seq_len, :].to(x.device)\n",
    "\n",
    "        for block in self.blocks:\n",
    "            x = block(x, mask_global, mask_local, cos, sin)\n",
    "\n",
    "        x = self.final_norm(x)\n",
    "        logits = self.out_head(x.to(self.cfg[\"dtype\"]))\n",
    "        return logits"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be2d201f-74ad-4d63-ab9c-601b00674a48",
   "metadata": {
    "id": "be2d201f-74ad-4d63-ab9c-601b00674a48"
   },
   "source": [
    "&nbsp;\n",
    "# 2. Initialize model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "caa142fa-b375-4e78-b392-2072ced666f3",
   "metadata": {
    "id": "caa142fa-b375-4e78-b392-2072ced666f3"
   },
   "outputs": [],
   "source": [
    "OLMO3_CONFIG_7B = {\n",
    "    \"vocab_size\": 100_278,\n",
    "    \"context_length\": 65_536,\n",
    "    \"emb_dim\": 4_096,\n",
    "    \"n_heads\": 32,\n",
    "    \"n_layers\": 32,\n",
    "    \"hidden_dim\": 11_008,\n",
    "    \"head_dim\": 128,\n",
    "    \"n_kv_heads\": 32,\n",
    "    \"attention_bias\": False,\n",
    "    \"attention_dropout\": 0.0,\n",
    "    \"sliding_window\": 4_096,\n",
    "    \"layer_types\": [\n",
    "        \"sliding_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"full_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"full_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"full_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"full_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"full_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"full_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"full_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"full_attention\",\n",
    "    ],\n",
    "    \"rope_base\": 500_000.0,\n",
    "    \"rope_attention_factor\": 1.2079441541679836,\n",
    "    \"rope_type\": \"yarn\",\n",
    "    \"rope_factor\": 8.0,\n",
    "    \"rope_orig_max\": 8_192,\n",
    "    \"beta_fast\": 32.0,\n",
    "    \"beta_slow\": 1.0,\n",
    "    \"rms_norm_eps\": 1e-6,\n",
    "    \"dtype\": torch.bfloat16,\n",
    "    \"eos_token_id\": 100_257,\n",
    "    \"pad_token_id\": 100_277,\n",
    "}\n",
    "\n",
    "OLMO3_CONFIG_32B = {\n",
    "    \"vocab_size\": 100_278,\n",
    "    \"context_length\": 65_536,\n",
    "    \"emb_dim\": 5_120,\n",
    "    \"n_heads\": 40,\n",
    "    \"n_layers\": 64,\n",
    "    \"hidden_dim\": 27_648,\n",
    "    \"head_dim\": 128,\n",
    "    \"n_kv_heads\": 8,\n",
    "    \"attention_bias\": False,\n",
    "    \"attention_dropout\": 0.0,\n",
    "    \"sliding_window\": 4_096,\n",
    "    \"layer_types\": [\n",
    "        \"sliding_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"full_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"full_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"full_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"full_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"full_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"full_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"full_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"full_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"full_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"full_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"full_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"full_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"full_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"full_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"sliding_attention\",\n",
    "        \"full_attention\",\n",
    "    ],\n",
    "    \"rope_base\": 500_000.0,\n",
    "    \"rope_attention_factor\": 1.2079441541679836,\n",
    "    \"rope_type\": \"yarn\",\n",
    "    \"rope_factor\": 8.0,\n",
    "    \"rope_orig_max\": 8_192,\n",
    "    \"beta_fast\": 32.0,\n",
    "    \"beta_slow\": 1.0,\n",
    "    \"rms_norm_eps\": 1e-6,\n",
    "    \"dtype\": torch.bfloat16,\n",
    "    \"eos_token_id\": 100_257,\n",
    "    \"pad_token_id\": 100_277,\n",
    "}\n",
    "\n",
    "OLMO3_CONFIG = OLMO3_CONFIG_32B if \"32B\" in USE_MODEL else OLMO3_CONFIG_7B"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "156253fe-aacd-4da2-8f13-705f05c4b11e",
   "metadata": {
    "id": "156253fe-aacd-4da2-8f13-705f05c4b11e"
   },
   "outputs": [],
   "source": [
    "torch.manual_seed(123)\n",
    "model = Olmo3Model(OLMO3_CONFIG)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "eaf86265-4e9d-4024-9ed0-99076944e304",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Olmo3Model(\n",
       "  (tok_emb): Embedding(100278, 4096)\n",
       "  (blocks): ModuleList(\n",
       "    (0-31): 32 x TransformerBlock(\n",
       "      (att): GroupedQueryAttention(\n",
       "        (W_query): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        (W_key): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        (W_value): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "        (q_norm): RMSNorm()\n",
       "        (k_norm): RMSNorm()\n",
       "      )\n",
       "      (ff): FeedForward(\n",
       "        (fc1): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "        (fc2): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "        (fc3): Linear(in_features=11008, out_features=4096, bias=False)\n",
       "      )\n",
       "      (post_attention_layernorm): RMSNorm()\n",
       "      (post_feedforward_layernorm): RMSNorm()\n",
       "    )\n",
       "  )\n",
       "  (final_norm): RMSNorm()\n",
       "  (out_head): Linear(in_features=4096, out_features=100278, bias=False)\n",
       ")"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "90aca91d-4bee-45ce-993a-4ec5393abe2b",
   "metadata": {},
   "source": [
    "- A quick check that the forward pass works before continuing:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "adf0a6b7-b688-42c9-966e-c223d34db99f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[ 0.3867, -0.6328, -0.2734,  ...,  1.1484,  0.4258,  0.0400],\n",
       "         [ 1.2734,  0.0040,  0.5000,  ...,  0.5625, -0.2383,  0.1855],\n",
       "         [ 0.5859, -0.0540,  0.7930,  ...,  0.3262, -0.5430, -0.1494]]],\n",
       "       dtype=torch.bfloat16, grad_fn=<UnsafeViewBackward0>)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model(torch.tensor([1, 2, 3]).unsqueeze(0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "31f12baf-f79b-499f-85c0-51328a6a20f5",
   "metadata": {
    "id": "31f12baf-f79b-499f-85c0-51328a6a20f5"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/rasbt/jupyterlab/reasoning/.venv/lib/python3.12/site-packages/torch/cuda/__init__.py:283: UserWarning: \n",
      "    Found GPU0 NVIDIA GB10 which is of cuda capability 12.1.\n",
      "    Minimum and Maximum cuda capability supported by this version of PyTorch is\n",
      "    (8.0) - (12.0)\n",
      "    \n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "if torch.cuda.is_available():\n",
    "    device = torch.device(\"cuda\")\n",
    "elif torch.backends.mps.is_available():\n",
    "    device = torch.device(\"mps\")\n",
    "else:\n",
    "    device = torch.device(\"cpu\")\n",
    "\n",
    "model.to(device);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c172f89f-d301-439f-b809-46169e5f5945",
   "metadata": {
    "id": "c172f89f-d301-439f-b809-46169e5f5945"
   },
   "source": [
    "&nbsp;\n",
    "# 3. Load pretrained weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "75166128-5899-4995-9b88-9672e135650e",
   "metadata": {
    "id": "75166128-5899-4995-9b88-9672e135650e"
   },
   "outputs": [],
   "source": [
    "def load_weights_into_olmo(model, param_config, params):\n",
    "    def assign(left, right, tensor_name=\"unknown\"):\n",
    "        if left.shape != right.shape:\n",
    "            raise ValueError(\n",
    "                f\"Shape mismatch in tensor '{tensor_name}'. \"\n",
    "                f\"Left: {left.shape}, Right: {right.shape}\"\n",
    "            )\n",
    "        \n",
    "        with torch.no_grad():\n",
    "            if isinstance(right, torch.Tensor):\n",
    "                left.copy_(right)\n",
    "            else:\n",
    "                left.copy_(torch.as_tensor(right, dtype=left.dtype, device=left.device))\n",
    "        \n",
    "        return left\n",
    "\n",
    "    # Token embedding\n",
    "    if \"model.embed_tokens.weight\" in params:\n",
    "        model.tok_emb.weight = assign(\n",
    "            model.tok_emb.weight,\n",
    "            params[\"model.embed_tokens.weight\"],\n",
    "            \"model.embed_tokens.weight\",\n",
    "        )\n",
    "\n",
    "    for l in range(param_config[\"n_layers\"]):\n",
    "        block = model.blocks[l]\n",
    "        att = block.att\n",
    "\n",
    "        # Q, K, V projections\n",
    "        att.W_query.weight = assign(\n",
    "            att.W_query.weight,\n",
    "            params[f\"model.layers.{l}.self_attn.q_proj.weight\"],\n",
    "            f\"model.layers.{l}.self_attn.q_proj.weight\",\n",
    "        )\n",
    "        att.W_key.weight = assign(\n",
    "            att.W_key.weight,\n",
    "            params[f\"model.layers.{l}.self_attn.k_proj.weight\"],\n",
    "            f\"model.layers.{l}.self_attn.k_proj.weight\",\n",
    "        )\n",
    "        att.W_value.weight = assign(\n",
    "            att.W_value.weight,\n",
    "            params[f\"model.layers.{l}.self_attn.v_proj.weight\"],\n",
    "            f\"model.layers.{l}.self_attn.v_proj.weight\",\n",
    "        )\n",
    "\n",
    "        # Output projection\n",
    "        att.out_proj.weight = assign(\n",
    "            att.out_proj.weight,\n",
    "            params[f\"model.layers.{l}.self_attn.o_proj.weight\"],\n",
    "            f\"model.layers.{l}.self_attn.o_proj.weight\",\n",
    "        )\n",
    "\n",
    "        # QK norms\n",
    "        att.q_norm.weight = assign(\n",
    "            att.q_norm.weight,\n",
    "            params[f\"model.layers.{l}.self_attn.q_norm.weight\"],\n",
    "            f\"model.layers.{l}.self_attn.q_norm.weight\",\n",
    "        )\n",
    "        att.k_norm.weight = assign(\n",
    "            att.k_norm.weight,\n",
    "            params[f\"model.layers.{l}.self_attn.k_norm.weight\"],\n",
    "            f\"model.layers.{l}.self_attn.k_norm.weight\",\n",
    "        )\n",
    "\n",
    "        # Feedforward weights\n",
    "        block.ff.fc1.weight = assign(\n",
    "            block.ff.fc1.weight,\n",
    "            params[f\"model.layers.{l}.mlp.gate_proj.weight\"],\n",
    "            f\"model.layers.{l}.mlp.gate_proj.weight\",\n",
    "        )\n",
    "        block.ff.fc2.weight = assign(\n",
    "            block.ff.fc2.weight,\n",
    "            params[f\"model.layers.{l}.mlp.up_proj.weight\"],\n",
    "            f\"model.layers.{l}.mlp.up_proj.weight\",\n",
    "        )\n",
    "        block.ff.fc3.weight = assign(\n",
    "            block.ff.fc3.weight,\n",
    "            params[f\"model.layers.{l}.mlp.down_proj.weight\"],\n",
    "            f\"model.layers.{l}.mlp.down_proj.weight\",\n",
    "        )\n",
    "\n",
    "        # Post-attention and post norms\n",
    "        block.post_attention_layernorm.weight = assign(\n",
    "            block.post_attention_layernorm.weight,\n",
    "            params[f\"model.layers.{l}.post_attention_layernorm.weight\"],\n",
    "            f\"model.layers.{l}.post_attention_layernorm.weight\",\n",
    "        )\n",
    "        block.post_feedforward_layernorm.weight = assign(\n",
    "            block.post_feedforward_layernorm.weight,\n",
    "            params[f\"model.layers.{l}.post_feedforward_layernorm.weight\"],\n",
    "            f\"model.layers.{l}.post_feedforward_layernorm.weight\",\n",
    "        )\n",
    "\n",
    "    # Final normalization and output head\n",
    "    if \"model.norm.weight\" in params:\n",
    "        model.final_norm.weight = assign(\n",
    "            model.final_norm.weight,\n",
    "            params[\"model.norm.weight\"],\n",
    "            \"model.norm.weight\",\n",
    "        )\n",
    "\n",
    "    if \"lm_head.weight\" in params:\n",
    "        model.out_head.weight = assign(\n",
    "            model.out_head.weight,\n",
    "            params[\"lm_head.weight\"],\n",
    "            \"lm_head.weight\",\n",
    "        )\n",
    "    else:\n",
    "        model.out_head.weight = model.tok_emb.weight\n",
    "        print(\"Model uses weight tying.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "699cb1b8-a67d-49fb-80a6-0dad9d81f392",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 17,
     "referenced_widgets": [
      "9881b6995c3f49dc89e6992fd9ab660b",
      "17a3174e65c54476b2e0d1faf8f011ca",
      "1bbf2e62c0754d1593beb4105a7f1ac1",
      "b82112e1dec645d98aa1c1ba64abcb61",
      "271e2bd6a35e4a8b92de8697f7c0be5f",
      "90a79523187446dfa692723b2e5833a7",
      "431ffb83b8c14bf182f0430e07ea6154",
      "a8f1b72a33dd4b548de23fbd95e0da18",
      "25cc36132d384189acfbecc59483134b",
      "bfd06423ad544218968648016e731a46",
      "d029630b63ff44cf807ade428d2eb421"
     ]
    },
    "id": "699cb1b8-a67d-49fb-80a6-0dad9d81f392",
    "outputId": "55b2f28c-142f-4698-9d23-d27456d3ed6d"
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ac81f5bc2063498b98e2c8956f0598be",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Fetching 14 files:   0%|          | 0/14 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import json\n",
    "import os\n",
    "from pathlib import Path\n",
    "from safetensors.torch import load_file\n",
    "from huggingface_hub import snapshot_download\n",
    "\n",
    "repo_id = f\"allenai/{USE_MODEL}\"\n",
    "local_dir = Path(repo_id).parts[-1]\n",
    "\n",
    "repo_dir = snapshot_download(repo_id=repo_id, local_dir=local_dir)\n",
    "index_path = os.path.join(repo_dir, \"model.safetensors.index.json\")\n",
    "with open(index_path, \"r\") as f:\n",
    "    index = json.load(f)\n",
    "\n",
    "weights_dict = {}\n",
    "for filename in sorted(set(index[\"weight_map\"].values())):\n",
    "    shard_path = os.path.join(repo_dir, filename)\n",
    "    shard = load_file(shard_path)\n",
    "    weights_dict.update(shard)\n",
    "\n",
    "load_weights_into_olmo(model, OLMO3_CONFIG, weights_dict)\n",
    "model.to(device)\n",
    "del weights_dict"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b345491-3510-4397-92d3-cd0a3fa3deee",
   "metadata": {},
   "source": [
    "&nbsp;\n",
    "# 4. Load tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "b68ab489-48e5-471e-a814-56cda2d60f81",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tokenizers import Tokenizer\n",
    "from huggingface_hub import hf_hub_download\n",
    "\n",
    "\n",
    "class OlmoTokenizer:\n",
    "    def __init__(self, tokenizer_file_path, eos_token_id, pad_token_id):\n",
    "        tok_file = Path(tokenizer_file_path)\n",
    "        self._tok = Tokenizer.from_file(str(tok_file))\n",
    "        eos_from_tok = (\n",
    "            self._tok.token_to_id(\"<|endoftext|>\")\n",
    "            or self._tok.token_to_id(\"<end_of_turn>\")\n",
    "        )\n",
    "        self.eos_token_id = eos_from_tok if eos_from_tok is not None else eos_token_id\n",
    "        pad_from_tok = (\n",
    "            self._tok.token_to_id(\"<|pad|>\")\n",
    "            or self._tok.token_to_id(\"<pad>\")\n",
    "        )\n",
    "        self.pad_token_id = pad_from_tok if pad_from_tok is not None else pad_token_id\n",
    "\n",
    "    def encode(self, text):\n",
    "        return self._tok.encode(text).ids\n",
    "\n",
    "    def decode(self, ids):\n",
    "        return self._tok.decode(ids, skip_special_tokens=False)\n",
    "\n",
    "\n",
    "def apply_chat_template(user_text):\n",
    "    return (\n",
    "        \"<|im_start|>user\\n\"\n",
    "        f\"{user_text}\\n\"\n",
    "        \"<|im_end|>\\n\"\n",
    "        \"<|im_start|>assistant\\n\"\n",
    "    )\n",
    "\n",
    "\n",
    "tokenizer_file_path = os.path.join(local_dir, \"tokenizer.json\")\n",
    "if not os.path.exists(tokenizer_file_path):\n",
    "    try:\n",
    "        tokenizer_file_path = hf_hub_download(repo_id=repo_id, filename=\"tokenizer.json\", local_dir=local_dir)\n",
    "    except Exception as e:\n",
    "        print(f\"Warning: failed to download tokenizer.json: {e}\")\n",
    "        tokenizer_file_path = \"tokenizer.json\"\n",
    "\n",
    "tokenizer = OlmoTokenizer(\n",
    "    tokenizer_file_path=tokenizer_file_path,\n",
    "    eos_token_id=OLMO3_CONFIG[\"eos_token_id\"],\n",
    "    pad_token_id=OLMO3_CONFIG[\"pad_token_id\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "7b6df8bc-7308-468e-93ce-2d5529ea7866",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'<|im_start|>user\\nGive me a short intro to large language models in 3 sentences.\\n<|im_end|>\\n<|im_start|>assistant\\n'"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "prompt = apply_chat_template(\"Give me a short intro to large language models in 3 sentences.\")\n",
    "\n",
    "input_token_ids = tokenizer.encode(prompt)\n",
    "text = tokenizer.decode(input_token_ids)\n",
    "text"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57d07df1-4401-4792-b549-7c4cc5632323",
   "metadata": {
    "id": "57d07df1-4401-4792-b549-7c4cc5632323"
   },
   "source": [
    "&nbsp;\n",
    "# 5. Generate text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "7b8401c6-e244-4cb7-9849-2ba71ce758d5",
   "metadata": {
    "id": "7b8401c6-e244-4cb7-9849-2ba71ce758d5"
   },
   "outputs": [],
   "source": [
    "def generate_text_basic_stream(model, token_ids, max_new_tokens, eos_token_id=None):\n",
    "\n",
    "    model.eval()\n",
    "    with torch.no_grad():\n",
    "        for _ in range(max_new_tokens):\n",
    "            out = model(token_ids)[:, -1]\n",
    "            next_token = torch.argmax(out, dim=-1, keepdim=True)\n",
    "\n",
    "            if (eos_token_id is not None\n",
    "                   and torch.all(next_token == eos_token_id)):\n",
    "               break\n",
    "\n",
    "            yield next_token\n",
    "            \n",
    "            token_ids = torch.cat([token_ids, next_token], dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "1c7a04fa-6aac-416b-8f63-f1e19227633d",
   "metadata": {
    "id": "1c7a04fa-6aac-416b-8f63-f1e19227633d"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Large language models are advanced AI systems trained on vast amounts of text to understand and generate human-like language. They can perform a wide range of tasks, from answering questions to writing essays or code. These models have transformed natural language processing and are now integral to many modern technologies.\n",
      "\n",
      "GPU memory used: 13.70 GB\n"
     ]
    }
   ],
   "source": [
    "input_token_ids_tensor = torch.tensor(input_token_ids, device=device).unsqueeze(0)\n",
    "\n",
    "\n",
    "if torch.cuda.is_available():\n",
    "    torch.cuda.reset_peak_memory_stats()\n",
    "\n",
    "\n",
    "for token in generate_text_basic_stream(\n",
    "    model=model,\n",
    "    token_ids=input_token_ids_tensor,\n",
    "    max_new_tokens=500,\n",
    "    eos_token_id=tokenizer.eos_token_id\n",
    "):\n",
    "    token_id = token.squeeze(0).tolist()\n",
    "    print(\n",
    "        tokenizer.decode(token_id),\n",
    "        end=\"\",\n",
    "        flush=True\n",
    "    )\n",
    "\n",
    "if torch.cuda.is_available():\n",
    "    def calc_gpu_gb(x):\n",
    "        return f\"{x / 1024 / 1024 / 1024:.2f} GB\"\n",
    "    \n",
    "    print(f\"\\n\\nGPU memory used: {calc_gpu_gb(torch.cuda.max_memory_allocated())}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "549324d6-5c71-4147-ae21-2e67675faa3d",
   "metadata": {
    "id": "549324d6-5c71-4147-ae21-2e67675faa3d"
   },
   "source": [
    "&nbsp;\n",
    "# What's next?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6edaaae-2de1-406c-8ffa-897cdfa3808c",
   "metadata": {
    "id": "e6edaaae-2de1-406c-8ffa-897cdfa3808c"
   },
   "source": [
    "- For those interested in a comprehensive guide on building a large language model from scratch and gaining a deeper understanding of its mechanics, you might like my [Build a Large Language Model (From Scratch)](http://mng.bz/orYv)\n",
    "\n",
    "<a href=\"http://mng.bz/orYv\"><img src=\"https://sebastianraschka.com/images/LLMs-from-scratch-images/cover-small.webp\" width=\"100px\"></a>"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "A100",
   "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
}
