App-D - Llama3 Pretrained.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/thedatasense/llm-healthcare/blob/main/App-D%20-%20Llama3%20Pretrained.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "7e64978d33b1636c",
      "metadata": {
        "id": "7e64978d33b1636c"
      },
      "source": [
        "## Llama 3.2 Implementation\n",
        "\n",
        "### Key Highlights\n",
        "\n",
        "1. **Transformer Stack**  \n",
        "   - **Layers**:  \n",
        "     - *LLaMA 3.2 1B:* 16 transformer blocks  \n",
        "     - *LLaMA 3.2 3B:* 28 transformer blocks  \n",
        "   - **Heads**:  \n",
        "     - *1B Model:* 32 attention heads  \n",
        "     - *3B Model:* 24 attention heads  \n",
        "   - **Hidden Dim**: 8,192  \n",
        "   - **Embedding Dim**:  \n",
        "     - *1B Model:* 2,048  \n",
        "     - *3B Model:* 3,072  \n",
        "\n",
        "2. **Attention Mechanism**  \n",
        "   - Uses **masked grouped-query attention** rather than standard multi-head self-attention.  \n",
        "   - Includes **Rotary Positional Embedding (RoPE) rescaling** to support extended context lengths (up to 131K tokens).  \n",
        "\n",
        "3. **Normalization**  \n",
        "   - Employs **RMSNorm** layers (instead of LayerNorm).  \n",
        "   - There is an **RMSNorm** after the attention block (RMSNorm 1), another before or within the feed-forward block (RMSNorm 2), and a **final RMSNorm** before the linear output layer.  \n",
        "\n",
        "4. **Feed-Forward Blocks**  \n",
        "   - Uses **SiLU** activation (sometimes referred to as **Swish**) inside the feed-forward sublayers.  \n",
        "   - Typically consists of linear transformations sandwiching the SiLU nonlinearity.\n",
        "\n",
        "5. **Vocabulary & Embedding**  \n",
        "   - **Vocabulary Size**: 128,256 tokens.  \n",
        "   - **Extended Context**: Up to 131K tokens (significantly larger than GPT-2’s window).  \n",
        "\n",
        "\n",
        "   ### Source\n",
        "\n",
        "- **[Sebastian Raschka's Blog](https://sebastianraschka.com)**  \n",
        "\n",
        "                                                                                                    "
      ]
    },
    {
      "cell_type": "markdown",
      "id": "aafd37d82bbcd144",
      "metadata": {
        "id": "aafd37d82bbcd144"
      },
      "source": [
        "<img src=\"https://sebastianraschka.com/images/LLMs-from-scratch-images/bonus/gpt-to-llama/llama32.webp\" width=\"700px\">\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install blobfile tiktoken"
      ],
      "metadata": {
        "id": "kQEgBfouadDi"
      },
      "id": "kQEgBfouadDi",
      "execution_count": 29,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": 30,
      "id": "initial_id",
      "metadata": {
        "ExecuteTime": {
          "end_time": "2025-01-09T06:42:34.681623Z",
          "start_time": "2025-01-09T06:42:34.675320Z"
        },
        "collapsed": true,
        "jupyter": {
          "outputs_hidden": true
        },
        "id": "initial_id"
      },
      "outputs": [],
      "source": [
        "from importlib.metadata import version\n",
        "\n",
        "pkgs = [\n",
        "    \"blobfile\",         # to download pretrained weights\n",
        "    \"huggingface_hub\",  # to download pretrained weights\n",
        "    \"tiktoken\",         # 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": "code",
      "execution_count": 31,
      "id": "611eefc1",
      "metadata": {
        "ExecuteTime": {
          "end_time": "2025-01-09T06:42:35.511406Z",
          "start_time": "2025-01-09T06:42:34.685915Z"
        },
        "id": "611eefc1"
      },
      "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": 32,
      "id": "3f864415",
      "metadata": {
        "ExecuteTime": {
          "end_time": "2025-01-09T06:42:35.522953Z",
          "start_time": "2025-01-09T06:42:35.518806Z"
        },
        "id": "3f864415"
      },
      "outputs": [],
      "source": [
        "def precompute_rope_params(head_dim, theta_base=10_000, context_length=4096, freq_config=None):\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)[: (head_dim // 2)].float() / head_dim))\n",
        "\n",
        "    # Frequency adjustments\n",
        "    if freq_config is not None:\n",
        "        low_freq_wavelen = freq_config[\"original_context_length\"] / freq_config[\"low_freq_factor\"]\n",
        "        high_freq_wavelen = freq_config[\"original_context_length\"] / freq_config[\"high_freq_factor\"]\n",
        "\n",
        "        wavelen = 2 * torch.pi / inv_freq\n",
        "\n",
        "        inv_freq_llama = torch.where(\n",
        "            wavelen > low_freq_wavelen, inv_freq / freq_config[\"factor\"], inv_freq\n",
        "        )\n",
        "\n",
        "        smooth_factor = (freq_config[\"original_context_length\"] / wavelen - freq_config[\"low_freq_factor\"]) / (\n",
        "            freq_config[\"high_freq_factor\"] - freq_config[\"low_freq_factor\"]\n",
        "        )\n",
        "\n",
        "        smoothed_inv_freq = (\n",
        "            (1 - smooth_factor) * (inv_freq / freq_config[\"factor\"]) + smooth_factor * inv_freq\n",
        "        )\n",
        "\n",
        "        is_medium_freq = (wavelen <= low_freq_wavelen) & (wavelen >= high_freq_wavelen)\n",
        "        inv_freq_llama = torch.where(is_medium_freq, smoothed_inv_freq, inv_freq_llama)\n",
        "        inv_freq = inv_freq_llama\n",
        "\n",
        "    # Generate position indices\n",
        "    positions = torch.arange(context_length)\n",
        "\n",
        "    # Compute the angles\n",
        "    angles = positions[:, None] * inv_freq[None, :]  # 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 compute_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",
        "    return x_rotated.to(dtype=x.dtype)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 33,
      "id": "fa485083",
      "metadata": {
        "ExecuteTime": {
          "end_time": "2025-01-09T06:42:35.535647Z",
          "start_time": "2025-01-09T06:42:35.529985Z"
        },
        "id": "fa485083"
      },
      "outputs": [],
      "source": [
        "class SharedBuffers:\n",
        "    _buffers = {}\n",
        "\n",
        "    @staticmethod\n",
        "    def get_buffers(context_length, head_dim, rope_base, freq_config, dtype=torch.float32):\n",
        "        key = (context_length, head_dim, rope_base, tuple(freq_config.values()) if freq_config else freq_config, dtype)\n",
        "\n",
        "        if key not in SharedBuffers._buffers:\n",
        "            # Create or fetch the buffers\n",
        "            mask = torch.triu(torch.ones(context_length, context_length), diagonal=1)\n",
        "            cos, sin = precompute_rope_params(head_dim, rope_base, context_length, freq_config)\n",
        "            if dtype is not None:\n",
        "                cos = cos.to(dtype)\n",
        "                sin = sin.to(dtype)\n",
        "            SharedBuffers._buffers[key] = (mask, cos, sin)\n",
        "\n",
        "        return SharedBuffers._buffers[key]\n",
        "\n",
        "\n",
        "class GroupedQueryAttention(nn.Module):\n",
        "    def __init__(\n",
        "            self, d_in, d_out, context_length, num_heads,\n",
        "            num_kv_groups,\n",
        "            rope_base=10_000,\n",
        "            rope_config=None,\n",
        "            dtype=None\n",
        "        ):\n",
        "        super().__init__()\n",
        "        assert d_out % num_heads == 0, \"d_out must be divisible by num_heads\"\n",
        "        assert num_heads % num_kv_groups == 0, \"num_heads must be divisible by num_kv_groups\"\n",
        "\n",
        "        self.d_out = d_out\n",
        "        self.num_heads = num_heads\n",
        "        self.head_dim = d_out // num_heads\n",
        "\n",
        "        self.W_key = nn.Linear(d_in, num_kv_groups * self.head_dim, bias=False, dtype=dtype)\n",
        "        self.W_value = nn.Linear(d_in, num_kv_groups * self.head_dim, bias=False, dtype=dtype)\n",
        "        self.num_kv_groups = num_kv_groups\n",
        "        self.group_size = num_heads // num_kv_groups\n",
        "\n",
        "        self.W_query = nn.Linear(d_in, d_out, bias=False, dtype=dtype)\n",
        "        self.out_proj = nn.Linear(d_out, d_out, bias=False, dtype=dtype)\n",
        "\n",
        "        # Fetch buffers using SharedBuffers\n",
        "        mask, cos, sin = SharedBuffers.get_buffers(context_length, self.head_dim, rope_base, rope_config, dtype)\n",
        "        self.register_buffer(\"mask\", mask)\n",
        "\n",
        "        self.register_buffer(\"cos\", cos)\n",
        "        self.register_buffer(\"sin\", sin)\n",
        "\n",
        "    def forward(self, x):\n",
        "        b, num_tokens, d_in = x.shape\n",
        "\n",
        "        queries = self.W_query(x)  # Shape: (b, num_tokens, d_out)\n",
        "        keys = self.W_key(x)  # Shape: (b, num_tokens, num_kv_groups * head_dim)\n",
        "        values = self.W_value(x)  # Shape: (b, num_tokens, num_kv_groups * head_dim)\n",
        "\n",
        "        # Reshape queries, keys, and values\n",
        "        queries = queries.view(b, num_tokens, self.num_heads, self.head_dim)\n",
        "        keys = keys.view(b, num_tokens, self.num_kv_groups, self.head_dim)\n",
        "        values = values.view(b, num_tokens, self.num_kv_groups, self.head_dim)\n",
        "\n",
        "        # Transpose keys, values, and queries\n",
        "        keys = keys.transpose(1, 2)  # Shape: (b, num_heads, num_tokens, head_dim)\n",
        "        values = values.transpose(1, 2)  # Shape: (b, num_heads, num_tokens, head_dim)\n",
        "        queries = queries.transpose(1, 2)  # Shape: (b, num_query_groups, num_tokens, head_dim)\n",
        "\n",
        "        # Apply RoPE\n",
        "        keys = compute_rope(keys, self.cos, self.sin)\n",
        "        queries = compute_rope(queries, self.cos, self.sin)\n",
        "\n",
        "        # Expand keys and values to match the number of heads\n",
        "        # Shape: (b, num_heads, num_tokens, head_dim)\n",
        "        keys = keys.repeat_interleave(self.group_size, dim=1)  # Shape: (b, num_heads, num_tokens, head_dim)\n",
        "        values = values.repeat_interleave(self.group_size, dim=1)  # Shape: (b, num_heads, num_tokens, head_dim)\n",
        "        # For example, before repeat_interleave along dim=1 (query groups):\n",
        "        #   [K1, K2]\n",
        "        # After repeat_interleave (each query group is repeated group_size times):\n",
        "        #   [K1, K1, K2, K2]\n",
        "        # If we used regular repeat instead of repeat_interleave, we'd get:\n",
        "        #   [K1, K2, K1, K2]\n",
        "\n",
        "        # Compute scaled dot-product attention (aka self-attention) with a causal mask\n",
        "        # Shape: (b, num_heads, num_tokens, num_tokens)\n",
        "        attn_scores = queries @ keys.transpose(2, 3)  # Dot product for each head\n",
        "\n",
        "        # Original mask truncated to the number of tokens and converted to boolean\n",
        "        mask_bool = self.mask.bool()[:num_tokens, :num_tokens]\n",
        "\n",
        "        # Use the mask to fill attention scores\n",
        "        attn_scores.masked_fill_(mask_bool, -torch.inf)\n",
        "\n",
        "        attn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)\n",
        "        assert keys.shape[-1] == self.head_dim\n",
        "\n",
        "        # Shape: (b, num_tokens, num_heads, head_dim)\n",
        "        context_vec = (attn_weights @ values).transpose(1, 2)\n",
        "\n",
        "        # Combine heads, where self.d_out = self.num_heads * self.head_dim\n",
        "        context_vec = context_vec.reshape(b, num_tokens, self.d_out)\n",
        "        context_vec = self.out_proj(context_vec)  # optional projection\n",
        "\n",
        "        return context_vec"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 34,
      "id": "9b615bce",
      "metadata": {
        "ExecuteTime": {
          "end_time": "2025-01-09T06:42:35.540698Z",
          "start_time": "2025-01-09T06:42:35.538009Z"
        },
        "id": "9b615bce"
      },
      "outputs": [],
      "source": [
        "class TransformerBlock(nn.Module):\n",
        "    def __init__(self, cfg):\n",
        "        super().__init__()\n",
        "        self.att =  GroupedQueryAttention(\n",
        "            d_in=cfg[\"emb_dim\"],\n",
        "            d_out=cfg[\"emb_dim\"],\n",
        "            context_length=cfg[\"context_length\"],\n",
        "            num_heads=cfg[\"n_heads\"],\n",
        "            num_kv_groups=cfg[\"n_kv_groups\"],\n",
        "            rope_base=cfg[\"rope_base\"],\n",
        "            rope_config=cfg[\"rope_freq\"],\n",
        "            dtype=cfg[\"dtype\"]\n",
        "        )\n",
        "        self.ff = FeedForward(cfg)\n",
        "        self.norm1 = nn.RMSNorm(cfg[\"emb_dim\"], eps=1e-5)\n",
        "        self.norm2 = nn.RMSNorm(cfg[\"emb_dim\"], eps=1e-5)\n",
        "\n",
        "    def forward(self, x):\n",
        "        # Shortcut connection for attention block\n",
        "        shortcut = x\n",
        "        x = self.norm1(x)\n",
        "        x = self.att(x.to(torch.bfloat16))   # 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.to(torch.bfloat16))\n",
        "        x = x + shortcut  # Add the original input back\n",
        "\n",
        "        return x"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 35,
      "id": "e8a2921e",
      "metadata": {
        "ExecuteTime": {
          "end_time": "2025-01-09T06:42:35.547067Z",
          "start_time": "2025-01-09T06:42:35.544616Z"
        },
        "id": "e8a2921e"
      },
      "outputs": [],
      "source": [
        "class Llama3Model(nn.Module):\n",
        "    def __init__(self, cfg):\n",
        "        super().__init__()\n",
        "        self.tok_emb = nn.Embedding(cfg[\"vocab_size\"], cfg[\"emb_dim\"], dtype=cfg[\"dtype\"])\n",
        "\n",
        "        self.trf_blocks = nn.Sequential(\n",
        "            *[TransformerBlock(cfg) for _ in range(cfg[\"n_layers\"])])\n",
        "\n",
        "        self.final_norm = nn.RMSNorm(cfg[\"emb_dim\"], eps=1e-5)\n",
        "        self.out_head = nn.Linear(cfg[\"emb_dim\"], cfg[\"vocab_size\"], bias=False, dtype=cfg[\"dtype\"])\n",
        "\n",
        "    def forward(self, in_idx):\n",
        "        tok_embeds = self.tok_emb(in_idx)\n",
        "        x = tok_embeds\n",
        "        x = self.trf_blocks(x)\n",
        "        x = self.final_norm(x)\n",
        "        logits = self.out_head(x.to(torch.bfloat16))\n",
        "        return logits"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 36,
      "id": "28a7e28e",
      "metadata": {
        "ExecuteTime": {
          "end_time": "2025-01-09T06:42:35.553234Z",
          "start_time": "2025-01-09T06:42:35.551026Z"
        },
        "id": "28a7e28e"
      },
      "outputs": [],
      "source": [
        "# Llama 3.2 1B\n",
        "\n",
        "LLAMA32_CONFIG = {\n",
        "    \"vocab_size\": 128_256,      # Vocabulary size\n",
        "    \"context_length\": 131_072,  # Context length\n",
        "    \"emb_dim\": 2048,            # Embedding dimension\n",
        "    \"n_heads\": 32,              # Number of attention heads\n",
        "    \"n_layers\": 16,             # Number of layers\n",
        "    \"hidden_dim\": 8192,         # Size of the intermediate dimension in FeedForward\n",
        "    \"n_kv_groups\": 8,           # Key-Value groups for grouped-query attention\n",
        "    \"rope_base\": 500_000.0,     # The base in RoPE's \"theta\"\n",
        "    \"dtype\": torch.bfloat16,    # Lower-precision dtype to reduce memory usage\n",
        "    \"rope_freq\": {              # RoPE frequency scaling\n",
        "        \"factor\": 32.0,\n",
        "        \"low_freq_factor\": 1.0,\n",
        "        \"high_freq_factor\": 4.0,\n",
        "        \"original_context_length\": 8192,\n",
        "    }\n",
        "}\n",
        "\n",
        "# Llama 3.2 3B\n",
        "\n",
        "# LLAMA32_CONFIG = {\n",
        "#     \"vocab_size\": 128_256,      # Vocabulary size\n",
        "#     \"context_length\": 131_072,  # Context length\n",
        "#     \"emb_dim\": 3072,            # Embedding dimension\n",
        "#     \"n_heads\": 24,              # Number of attention heads\n",
        "#     \"n_layers\": 28,             # Number of layers\n",
        "#     \"hidden_dim\": 8192,         # Size of the intermediate dimension in FeedForward\n",
        "#     \"n_kv_groups\": 8,           # Key-Value groups for grouped-query attention\n",
        "#     \"rope_base\": 500_000.0,     # The base in RoPE's \"theta\"\n",
        "#     \"dtype\": torch.bfloat16,    # Lower-precision dtype to reduce memory usage\n",
        "#     \"rope_freq\": {              # RoPE frequency scaling\n",
        "#         \"factor\": 32.0,\n",
        "#         \"low_freq_factor\": 1.0,\n",
        "#         \"high_freq_factor\": 4.0,\n",
        "#         \"original_context_length\": 8192,\n",
        "#     }\n",
        "# }\n",
        "\n",
        "LLAMA_SIZE_STR = \"1B\" if LLAMA32_CONFIG[\"emb_dim\"] == 2048 else \"3B\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 37,
      "id": "185457b5",
      "metadata": {
        "ExecuteTime": {
          "end_time": "2025-01-09T06:42:35.559323Z",
          "start_time": "2025-01-09T06:42:35.557114Z"
        },
        "id": "185457b5"
      },
      "outputs": [],
      "source": [
        "old_context_length = LLAMA32_CONFIG[\"context_length\"]\n",
        "LLAMA32_CONFIG[\"context_length\"] = 8192\n",
        "\n",
        "\n",
        "def rescale_theta(theta_old, context_length_old, context_length_new):\n",
        "    scaling_factor = context_length_new / context_length_old\n",
        "    theta_new = theta_old * scaling_factor\n",
        "    return theta_new\n",
        "\n",
        "LLAMA32_CONFIG[\"rope_base\"] = rescale_theta(\n",
        "    LLAMA32_CONFIG[\"rope_base\"],\n",
        "    old_context_length,\n",
        "    LLAMA32_CONFIG[\"context_length\"]\n",
        ")\n",
        "\n",
        "print(\"New RoPE theta:\", LLAMA32_CONFIG[\"rope_base\"])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 38,
      "id": "ea0d1f33",
      "metadata": {
        "ExecuteTime": {
          "end_time": "2025-01-09T06:42:42.174756Z",
          "start_time": "2025-01-09T06:42:35.563403Z"
        },
        "id": "ea0d1f33"
      },
      "outputs": [],
      "source": [
        "model = Llama3Model(LLAMA32_CONFIG)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 39,
      "id": "3db10db2",
      "metadata": {
        "ExecuteTime": {
          "end_time": "2025-01-09T06:42:42.182529Z",
          "start_time": "2025-01-09T06:42:42.179958Z"
        },
        "id": "3db10db2"
      },
      "outputs": [],
      "source": [
        "# Check buffers\n",
        "print(model.trf_blocks[0].att.mask is model.trf_blocks[-1].att.mask)\n",
        "print(model.trf_blocks[0].att.cos is model.trf_blocks[-1].att.cos)\n",
        "print(model.trf_blocks[0].att.sin is model.trf_blocks[-1].att.sin)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 40,
      "id": "044f4f55",
      "metadata": {
        "ExecuteTime": {
          "end_time": "2025-01-09T06:42:42.189832Z",
          "start_time": "2025-01-09T06:42:42.187398Z"
        },
        "id": "044f4f55"
      },
      "outputs": [],
      "source": [
        "total_params = sum(p.numel() for p in model.parameters())\n",
        "print(f\"Total number of parameters: {total_params:,}\")\n",
        "\n",
        "# Account for weight tying\n",
        "total_params_normalized = total_params - model.tok_emb.weight.numel()\n",
        "print(f\"\\nTotal number of unique parameters: {total_params_normalized:,}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 41,
      "id": "8450c848",
      "metadata": {
        "ExecuteTime": {
          "end_time": "2025-01-09T06:42:42.204901Z",
          "start_time": "2025-01-09T06:42:42.201084Z"
        },
        "id": "8450c848"
      },
      "outputs": [],
      "source": [
        "def model_memory_size(model, input_dtype=torch.float32):\n",
        "    total_params = 0\n",
        "    total_grads = 0\n",
        "    for param in model.parameters():\n",
        "        # Calculate total number of elements per parameter\n",
        "        param_size = param.numel()\n",
        "        total_params += param_size\n",
        "        # Check if gradients are stored for this parameter\n",
        "        if param.requires_grad:\n",
        "            total_grads += param_size\n",
        "\n",
        "    # Calculate buffer size (non-parameters that require memory)\n",
        "    total_buffers = sum(buf.numel() for buf in model.buffers())\n",
        "\n",
        "    # Size in bytes = (Number of elements) * (Size of each element in bytes)\n",
        "    # We assume parameters and gradients are stored in the same type as input dtype\n",
        "    element_size = torch.tensor(0, dtype=input_dtype).element_size()\n",
        "    total_memory_bytes = (total_params + total_grads + total_buffers) * element_size\n",
        "\n",
        "    # Convert bytes to gigabytes\n",
        "    total_memory_gb = total_memory_bytes / (1024**3)\n",
        "\n",
        "    return total_memory_gb\n",
        "\n",
        "print(f\"float32 (PyTorch default): {model_memory_size(model, input_dtype=torch.float32):.2f} GB\")\n",
        "print(f\"bfloat16: {model_memory_size(model, input_dtype=torch.bfloat16):.2f} GB\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 42,
      "id": "39c15477",
      "metadata": {
        "ExecuteTime": {
          "end_time": "2025-01-09T06:42:42.826409Z",
          "start_time": "2025-01-09T06:42:42.215526Z"
        },
        "id": "39c15477"
      },
      "outputs": [],
      "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": "code",
      "execution_count": 43,
      "id": "c650d143",
      "metadata": {
        "ExecuteTime": {
          "end_time": "2025-01-09T06:42:47.783548Z",
          "start_time": "2025-01-09T06:42:47.764965Z"
        },
        "id": "c650d143"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "from pathlib import Path\n",
        "\n",
        "import tiktoken\n",
        "from tiktoken.load import load_tiktoken_bpe\n",
        "\n",
        "\n",
        "class Tokenizer:\n",
        "    def __init__(self, model_path):\n",
        "        assert os.path.isfile(model_path), f\"Model file {model_path} not found\"\n",
        "        mergeable_ranks = load_tiktoken_bpe(model_path)\n",
        "\n",
        "        self.special_tokens = {\n",
        "            \"<|begin_of_text|>\": 128000,\n",
        "            \"<|end_of_text|>\": 128001,\n",
        "            \"<|start_header_id|>\": 128006,\n",
        "            \"<|end_header_id|>\": 128007,\n",
        "            \"<|eot_id|>\": 128009,\n",
        "        }\n",
        "        self.special_tokens.update({\n",
        "            f\"<|reserved_{i}|>\": 128002 + i for i in range(256) if (128002 + i) not in self.special_tokens.values()\n",
        "        })\n",
        "\n",
        "        self.model = tiktoken.Encoding(\n",
        "            name=Path(model_path).name,\n",
        "            pat_str=r\"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\\r\\n\\p{L}\\p{N}]?\\p{L}+|\\p{N}{1,3}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+\",\n",
        "            mergeable_ranks=mergeable_ranks,\n",
        "            special_tokens=self.special_tokens\n",
        "        )\n",
        "\n",
        "\n",
        "    def encode(self, text, bos=False, eos=False, allowed_special=set(), disallowed_special=()):\n",
        "        if bos:\n",
        "            tokens = [self.special_tokens[\"<|begin_of_text|>\"]]\n",
        "        else:\n",
        "            tokens = []\n",
        "\n",
        "        tokens += self.model.encode(text, allowed_special=allowed_special, disallowed_special=disallowed_special)\n",
        "\n",
        "        if eos:\n",
        "            tokens.append(self.special_tokens[\"<|end_of_text|>\"])\n",
        "        return tokens\n",
        "\n",
        "    def decode(self, tokens):\n",
        "        return self.model.decode(tokens)\n",
        "\n",
        "\n",
        "class ChatFormat:\n",
        "    def __init__(self, tokenizer):\n",
        "        self.tokenizer = tokenizer\n",
        "\n",
        "    def encode_header(self, message):\n",
        "        tokens = []\n",
        "        tokens.append(self.tokenizer.special_tokens[\"<|start_header_id|>\"])\n",
        "        tokens.extend(self.tokenizer.encode(message[\"role\"], bos=False, eos=False))\n",
        "        tokens.append(self.tokenizer.special_tokens[\"<|end_header_id|>\"])\n",
        "        tokens.extend(self.tokenizer.encode(\"\\n\\n\", bos=False, eos=False))\n",
        "        return tokens\n",
        "\n",
        "    def encode(self, text):\n",
        "        message = {\n",
        "            \"role\": \"user\",\n",
        "            \"content\": text\n",
        "        }\n",
        "\n",
        "        tokens = self.encode_header(message)\n",
        "        tokens.extend(\n",
        "            self.tokenizer.encode(message[\"content\"].strip(), bos=False, eos=False)\n",
        "        )\n",
        "        tokens.append(self.tokenizer.special_tokens[\"<|eot_id|>\"])\n",
        "        return tokens\n",
        "\n",
        "    def decode(self, token_ids):\n",
        "        return self.tokenizer.decode(token_ids)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 44,
      "id": "5055f07a",
      "metadata": {
        "ExecuteTime": {
          "end_time": "2025-01-09T06:42:51.576855Z",
          "start_time": "2025-01-09T06:42:51.569535Z"
        },
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 361,
          "referenced_widgets": [
            "f12454ef44db4269b579fa4f9fc11f13",
            "ce8929113688436583d0eb231f79ff2e",
            "b57f1b3e751c44fa976303988538d72b",
            "c1bbf4103d494ebeb5f199b12a3a11a7",
            "01ead7e64a2549a0bbe473e7c87f5712",
            "ee915247fb9045e7b6169886729709a8",
            "11b9070c2ee84e9499bcda8bdfd584b6",
            "768dcb5cce2b4a01be6697a0d684c163",
            "acf13d32075b4ddb8fc9cc00580d247c",
            "1be24cba5ff948dd94d957cbd895733d",
            "a43a9a987f6344c9a279ed5011f7bb16",
            "56b004a071c7455eb0a1b35c6345e46d",
            "4a5bf6f042234b9294aae03a7ba8dc3b",
            "7b6ec7155905418fbe74d8d50df43fd7",
            "802d1f0923f642b59824e5270bc36878",
            "2fac83563540443283c4af0dcc286de2",
            "03b3b603afbf4e6f906ed246800439bf"
          ]
        },
        "id": "5055f07a",
        "outputId": "46a6fe08-39fd-4c51-c28a-ff884aaad931"
      },
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "VBox(children=(HTML(value='<center> <img\\nsrc=https://huggingface.co/front/assets/huggingface_logo-noborder.sv…"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "f12454ef44db4269b579fa4f9fc11f13"
            }
          },
          "metadata": {}
        }
      ],
      "source": [
        "from huggingface_hub import login\n",
        "\n",
        "login()"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ],
      "metadata": {
        "id": "5Xsd6K39bjs4"
      },
      "id": "5Xsd6K39bjs4",
      "execution_count": 45,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": 46,
      "id": "d8408fe6",
      "metadata": {
        "ExecuteTime": {
          "end_time": "2025-01-09T06:43:15.931768Z",
          "start_time": "2025-01-09T06:43:14.235532Z"
        },
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "d8408fe6",
        "outputId": "0627fea1-6454-4e25-d2f0-22215afde77a"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Couldn't access the Hub to check for update but local file already exists. Defaulting to existing file. (error: 401 Client Error. (Request ID: Root=1-6792a6e1-742c580b35d697382def0dd9;5164c68b-55f1-451f-bb7b-844b1cf08b18)\n",
            "\n",
            "Cannot access gated repo for url https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/resolve/main/original/tokenizer.model.\n",
            "Access to model meta-llama/Llama-3.2-1B-Instruct is restricted. You must have access to it and be authenticated to access it. Please log in.)\n",
            "WARNING:huggingface_hub.file_download:Couldn't access the Hub to check for update but local file already exists. Defaulting to existing file. (error: 401 Client Error. (Request ID: Root=1-6792a6e1-742c580b35d697382def0dd9;5164c68b-55f1-451f-bb7b-844b1cf08b18)\n",
            "\n",
            "Cannot access gated repo for url https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/resolve/main/original/tokenizer.model.\n",
            "Access to model meta-llama/Llama-3.2-1B-Instruct is restricted. You must have access to it and be authenticated to access it. Please log in.)\n"
          ]
        }
      ],
      "source": [
        "from huggingface_hub import hf_hub_download\n",
        "\n",
        "tokenizer_file_path = hf_hub_download(\n",
        "    repo_id=f\"meta-llama/Llama-3.2-{LLAMA_SIZE_STR}-Instruct\",\n",
        "    filename=\"original/tokenizer.model\",\n",
        "    local_dir=f\"/content/drive/MyDrive/PhD- LLM Healthcare/models/Llama-3.2-{LLAMA_SIZE_STR}-Instruct\"\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 47,
      "id": "47283713156dd9c4",
      "metadata": {
        "ExecuteTime": {
          "end_time": "2025-01-09T06:43:19.822713Z",
          "start_time": "2025-01-09T06:43:19.673099Z"
        },
        "id": "47283713156dd9c4"
      },
      "outputs": [],
      "source": [
        "tokenizer = Tokenizer(tokenizer_file_path)\n",
        "chat_tokenizer = ChatFormat(tokenizer)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 48,
      "id": "d5f4c749fbcfef33",
      "metadata": {
        "ExecuteTime": {
          "end_time": "2025-01-09T06:46:04.092656Z",
          "start_time": "2025-01-09T06:46:04.084839Z"
        },
        "id": "d5f4c749fbcfef33"
      },
      "outputs": [],
      "source": [
        "def assign(left, right, tensor_name=\"unknown\"):\n",
        "    if left.shape != right.shape:\n",
        "        raise ValueError(f\"Shape mismatch in tensor '{tensor_name}'. Left: {left.shape}, Right: {right.shape}\")\n",
        "\n",
        "    if isinstance(right, torch.Tensor):\n",
        "        return torch.nn.Parameter(right.clone().detach())\n",
        "    else:\n",
        "        return torch.nn.Parameter(torch.tensor(right))\n",
        "\n",
        "\n",
        "def load_weights_into_llama(model, param_config, params):\n",
        "    model.tok_emb.weight = assign(model.tok_emb.weight, params[\"model.embed_tokens.weight\"], \"model.embed_tokens.weight\")\n",
        "\n",
        "    for l in range(param_config[\"n_layers\"]):\n",
        "\n",
        "        # Load attention weights\n",
        "        model.trf_blocks[l].att.W_query.weight = assign(\n",
        "            model.trf_blocks[l].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",
        "        model.trf_blocks[l].att.W_key.weight = assign(\n",
        "            model.trf_blocks[l].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",
        "        model.trf_blocks[l].att.W_value.weight = assign(\n",
        "            model.trf_blocks[l].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",
        "        model.trf_blocks[l].att.out_proj.weight = assign(\n",
        "            model.trf_blocks[l].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",
        "        model.trf_blocks[l].norm1.weight = assign(\n",
        "            model.trf_blocks[l].norm1.weight,\n",
        "            params[f\"model.layers.{l}.input_layernorm.weight\"],\n",
        "            f\"model.layers.{l}.input_layernorm.weight\"\n",
        "        )\n",
        "\n",
        "        # Load FeedForward weights\n",
        "        model.trf_blocks[l].ff.fc1.weight = assign(\n",
        "            model.trf_blocks[l].ff.fc1.weight,\n",
        "            params[f\"model.layers.{l}.mlp.gate_proj.weight\"],\n",
        "            f\"model.layers.{l}.mlp.gate_proj.weight\"\n",
        "        )\n",
        "        model.trf_blocks[l].ff.fc2.weight = assign(\n",
        "            model.trf_blocks[l].ff.fc2.weight,\n",
        "            params[f\"model.layers.{l}.mlp.up_proj.weight\"],\n",
        "            f\"model.layers.{l}.mlp.up_proj.weight\"\n",
        "        )\n",
        "        model.trf_blocks[l].ff.fc3.weight = assign(\n",
        "            model.trf_blocks[l].ff.fc3.weight,\n",
        "            params[f\"model.layers.{l}.mlp.down_proj.weight\"],\n",
        "            f\"model.layers.{l}.mlp.down_proj.weight\"\n",
        "        )\n",
        "        model.trf_blocks[l].norm2.weight = assign(\n",
        "            model.trf_blocks[l].norm2.weight,\n",
        "            params[f\"model.layers.{l}.post_attention_layernorm.weight\"],\n",
        "            f\"model.layers.{l}.post_attention_layernorm.weight\"\n",
        "        )\n",
        "\n",
        "    # Load output layer weights\n",
        "    model.final_norm.weight = assign(model.final_norm.weight, params[\"model.norm.weight\"], \"model.norm.weight\")\n",
        "\n",
        "    if \"lm_head.weight\" in params.keys():\n",
        "        model.out_head.weight = assign(model.out_head.weight, params[\"lm_head.weight\"], \"lm_head.weight\")\n",
        "    else:\n",
        "        model.out_head.weight = assign(model.out_head.weight, params[\"model.embed_tokens.weight\"], \"model.embed_tokens.weight\")\n",
        "        print(\"Model uses weight tying.\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 49,
      "id": "6167d6d2947a0657",
      "metadata": {
        "ExecuteTime": {
          "end_time": "2025-01-09T06:46:21.371859Z",
          "start_time": "2025-01-09T06:46:18.486188Z"
        },
        "id": "6167d6d2947a0657"
      },
      "outputs": [],
      "source": [
        "!pip install --upgrade safetensors"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 50,
      "id": "6ec0529c4a72469f",
      "metadata": {
        "ExecuteTime": {
          "end_time": "2025-01-09T06:46:50.618055Z",
          "start_time": "2025-01-09T06:46:23.802378Z"
        },
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6ec0529c4a72469f",
        "outputId": "5b48c9e9-846c-4f9b-9c4e-b4945c097162"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Couldn't access the Hub to check for update but local file already exists. Defaulting to existing file. (error: 401 Client Error. (Request ID: Root=1-6792a6e4-294290c34440a5be5ca53be3;0073cfc2-b860-499f-a388-f2e8162b16f0)\n",
            "\n",
            "Cannot access gated repo for url https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/resolve/main/model.safetensors.\n",
            "Access to model meta-llama/Llama-3.2-1B-Instruct is restricted. You must have access to it and be authenticated to access it. Please log in.)\n",
            "WARNING:huggingface_hub.file_download:Couldn't access the Hub to check for update but local file already exists. Defaulting to existing file. (error: 401 Client Error. (Request ID: Root=1-6792a6e4-294290c34440a5be5ca53be3;0073cfc2-b860-499f-a388-f2e8162b16f0)\n",
            "\n",
            "Cannot access gated repo for url https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/resolve/main/model.safetensors.\n",
            "Access to model meta-llama/Llama-3.2-1B-Instruct is restricted. You must have access to it and be authenticated to access it. Please log in.)\n"
          ]
        }
      ],
      "source": [
        "from safetensors.torch import load_file\n",
        "\n",
        "\n",
        "if LLAMA_SIZE_STR == \"1B\":\n",
        "    weights_file = hf_hub_download(\n",
        "        repo_id=f\"meta-llama/Llama-3.2-{LLAMA_SIZE_STR}-Instruct\",\n",
        "        filename=f\"model.safetensors\",\n",
        "        local_dir=f\"/content/drive/MyDrive/PhD- LLM Healthcare/models/Llama-3.2-{LLAMA_SIZE_STR}-Instruct\"\n",
        "    )\n",
        "    combined_weights = load_file(weights_file)\n",
        "\n",
        "\n",
        "else:\n",
        "    combined_weights = {}\n",
        "    for i in range(1, 3):\n",
        "        weights_file = hf_hub_download(\n",
        "            repo_id=f\"meta-llama/Llama-3.2-{LLAMA_SIZE_STR}-Instruct\",\n",
        "            filename=f\"model-0000{i}-of-00002.safetensors\",\n",
        "            local_dir=f\"/content/drive/MyDrive/PhD- LLM Healthcare/models/Llama-3.2-{LLAMA_SIZE_STR}-Instruct\"\n",
        "        )\n",
        "        current_weights = load_file(weights_file)\n",
        "        combined_weights.update(current_weights)\n",
        "\n",
        "\n",
        "load_weights_into_llama(model, LLAMA32_CONFIG, combined_weights)\n",
        "model.to(device)\n",
        "del combined_weights  # free up memory"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 51,
      "id": "32112c51756e6ffe",
      "metadata": {
        "ExecuteTime": {
          "end_time": "2025-01-09T06:39:54.655773Z",
          "start_time": "2025-01-09T06:39:54.649157Z"
        },
        "id": "32112c51756e6ffe"
      },
      "outputs": [],
      "source": [
        "print(\"Weight tying:\", torch.equal(model.tok_emb.weight, model.out_head.weight))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 52,
      "id": "81c6a78f329c173c",
      "metadata": {
        "ExecuteTime": {
          "end_time": "2025-01-09T06:38:34.559344Z",
          "start_time": "2025-01-09T06:38:34.547308Z"
        },
        "id": "81c6a78f329c173c"
      },
      "outputs": [],
      "source": [
        "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",
        "\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",
        "\n",
        "def generate(model, idx, max_new_tokens, context_size, temperature=0.0, top_k=None, eos_id=None):\n",
        "\n",
        "    # For-loop is the same as before: Get logits, and only focus on last time step\n",
        "    for _ in range(max_new_tokens):\n",
        "        idx_cond = idx[:, -context_size:]\n",
        "        with torch.no_grad():\n",
        "            logits = model(idx_cond)\n",
        "        logits = logits[:, -1, :]\n",
        "\n",
        "        # New: Filter logits with top_k sampling\n",
        "        if top_k is not None:\n",
        "            # Keep only top_k values\n",
        "            top_logits, _ = torch.topk(logits, top_k)\n",
        "            min_val = top_logits[:, -1]\n",
        "            logits = torch.where(logits < min_val, torch.tensor(float('-inf')).to(logits.device), logits)\n",
        "\n",
        "        # New: Apply temperature scaling\n",
        "        if temperature > 0.0:\n",
        "            logits = logits / temperature\n",
        "\n",
        "            # Apply softmax to get probabilities\n",
        "            probs = torch.softmax(logits, dim=-1)  # (batch_size, context_len)\n",
        "\n",
        "            # Sample from the distribution\n",
        "            idx_next = torch.multinomial(probs, num_samples=1)  # (batch_size, 1)\n",
        "\n",
        "        # Otherwise same as before: get idx of the vocab entry with the highest logits value\n",
        "        else:\n",
        "            idx_next = torch.argmax(logits, dim=-1, keepdim=True)  # (batch_size, 1)\n",
        "\n",
        "        if idx_next == eos_id:  # Stop generating early if end-of-sequence token is encountered and eos_id is specified\n",
        "            break\n",
        "\n",
        "        # Same as before: append sampled index to the running sequence\n",
        "        idx = torch.cat((idx, idx_next), dim=1)  # (batch_size, num_tokens+1)\n",
        "\n",
        "    return idx"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 28,
      "id": "7ef04e0cfaf55e73",
      "metadata": {
        "ExecuteTime": {
          "end_time": "2025-01-09T06:38:51.551909Z",
          "start_time": "2025-01-09T06:38:38.438873Z"
        },
        "id": "7ef04e0cfaf55e73"
      },
      "outputs": [],
      "source": [
        "PROMPT = \"What do llamas eat?\"\n",
        "\n",
        "torch.manual_seed(123)\n",
        "\n",
        "token_ids = generate(\n",
        "    model=model,\n",
        "    idx=text_to_token_ids(PROMPT, chat_tokenizer).to(device),\n",
        "    max_new_tokens=150,\n",
        "    context_size=LLAMA32_CONFIG[\"context_length\"],\n",
        "    top_k=1,\n",
        "    temperature=0.\n",
        ")\n",
        "\n",
        "output_text = token_ids_to_text(token_ids, tokenizer)\n",
        "\n",
        "\n",
        "def clean_text(text, header_end=\"assistant<|end_header_id|>\\n\\n\"):\n",
        "    # Find the index of the first occurrence of \"<|end_header_id|>\"\n",
        "    index = text.find(header_end)\n",
        "\n",
        "    if index != -1:\n",
        "        # Return the substring starting after \"<|end_header_id|>\"\n",
        "        return text[index + len(header_end):].strip()  # Strip removes leading/trailing whitespace\n",
        "    else:\n",
        "        # If the token is not found, return the original text\n",
        "        return text\n",
        "\n",
        "print(\"Output text:\\n\", clean_text(output_text))"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "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.11.11"
    },
    "colab": {
      "provenance": [],
      "machine_shape": "hm",
      "gpuType": "A100",
      "include_colab_link": true
    },
    "accelerator": "GPU",
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "f12454ef44db4269b579fa4f9fc11f13": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "VBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "VBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "VBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_ce8929113688436583d0eb231f79ff2e",
              "IPY_MODEL_b57f1b3e751c44fa976303988538d72b",
              "IPY_MODEL_c1bbf4103d494ebeb5f199b12a3a11a7",
              "IPY_MODEL_01ead7e64a2549a0bbe473e7c87f5712",
              "IPY_MODEL_ee915247fb9045e7b6169886729709a8"
            ],
            "layout": "IPY_MODEL_11b9070c2ee84e9499bcda8bdfd584b6"
          }
        },
        "ce8929113688436583d0eb231f79ff2e": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_768dcb5cce2b4a01be6697a0d684c163",
            "placeholder": "​",
            "style": "IPY_MODEL_acf13d32075b4ddb8fc9cc00580d247c",
            "value": "<center> <img\nsrc=https://huggingface.co/front/assets/huggingface_logo-noborder.svg\nalt='Hugging Face'> <br> Copy a token from <a\nhref=\"https://huggingface.co/settings/tokens\" target=\"_blank\">your Hugging Face\ntokens page</a> and paste it below. <br> Immediately click login after copying\nyour token or it might be stored in plain text in this notebook file. </center>"
          }
        },
        "b57f1b3e751c44fa976303988538d72b": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "PasswordModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "PasswordModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "PasswordView",
            "continuous_update": true,
            "description": "Token:",
            "description_tooltip": null,
            "disabled": false,
            "layout": "IPY_MODEL_1be24cba5ff948dd94d957cbd895733d",
            "placeholder": "​",
            "style": "IPY_MODEL_a43a9a987f6344c9a279ed5011f7bb16",
            "value": ""
          }
        },
        "c1bbf4103d494ebeb5f199b12a3a11a7": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "CheckboxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "CheckboxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "CheckboxView",
            "description": "Add token as git credential?",
            "description_tooltip": null,
            "disabled": false,
            "indent": true,
            "layout": "IPY_MODEL_56b004a071c7455eb0a1b35c6345e46d",
            "style": "IPY_MODEL_4a5bf6f042234b9294aae03a7ba8dc3b",
            "value": true
          }
        },
        "01ead7e64a2549a0bbe473e7c87f5712": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ButtonModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ButtonModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ButtonView",
            "button_style": "",
            "description": "Login",
            "disabled": false,
            "icon": "",
            "layout": "IPY_MODEL_7b6ec7155905418fbe74d8d50df43fd7",
            "style": "IPY_MODEL_802d1f0923f642b59824e5270bc36878",
            "tooltip": ""
          }
        },
        "ee915247fb9045e7b6169886729709a8": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_2fac83563540443283c4af0dcc286de2",
            "placeholder": "​",
            "style": "IPY_MODEL_03b3b603afbf4e6f906ed246800439bf",
            "value": "\n<b>Pro Tip:</b> If you don't already have one, you can create a dedicated\n'notebooks' token with 'write' access, that you can then easily reuse for all\nnotebooks. </center>"
          }
        },
        "11b9070c2ee84e9499bcda8bdfd584b6": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": "center",
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": "flex",
            "flex": null,
            "flex_flow": "column",
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": "50%"
          }
        },
        "768dcb5cce2b4a01be6697a0d684c163": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "acf13d32075b4ddb8fc9cc00580d247c": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "1be24cba5ff948dd94d957cbd895733d": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "a43a9a987f6344c9a279ed5011f7bb16": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "56b004a071c7455eb0a1b35c6345e46d": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "4a5bf6f042234b9294aae03a7ba8dc3b": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "7b6ec7155905418fbe74d8d50df43fd7": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "802d1f0923f642b59824e5270bc36878": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ButtonStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ButtonStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "button_color": null,
            "font_weight": ""
          }
        },
        "2fac83563540443283c4af0dcc286de2": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "03b3b603afbf4e6f906ed246800439bf": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        }
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}