model_code.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e5044d81",
   "metadata": {},
   "source": [
    "### Import Libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "0b114dd0",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-07-20T03:03:47.391283Z",
     "iopub.status.busy": "2025-07-20T03:03:47.390701Z",
     "iopub.status.idle": "2025-07-20T03:03:53.759920Z",
     "shell.execute_reply": "2025-07-20T03:03:53.759287Z",
     "shell.execute_reply.started": "2025-07-20T03:03:47.391257Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F \n",
    "import os\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import Dataset ,  DataLoader\n",
    "from torch.amp import autocast\n",
    "from torch.cuda.amp import GradScaler\n",
    "import torch.amp\n",
    "\n",
    "\n",
    "import math\n",
    "import pickle\n",
    "import regex\n",
    "from tqdm import tqdm\n",
    "from huggingface_hub import hf_hub_download\n",
    "from graphviz import Digraph\n",
    "import wandb"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c89e7be",
   "metadata": {},
   "source": [
    "### Simple BPE-Algo Class"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2947d1dc",
   "metadata": {},
   "source": [
    "- A BPE class implementation which follows a simple approach"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "da34fb2e",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-07-20T03:03:58.923271Z",
     "iopub.status.busy": "2025-07-20T03:03:58.922432Z",
     "iopub.status.idle": "2025-07-20T03:03:58.930795Z",
     "shell.execute_reply": "2025-07-20T03:03:58.930043Z",
     "shell.execute_reply.started": "2025-07-20T03:03:58.923237Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "class SimpleBytePairEncoding:\n",
    "    def __init__(self, *, pat_str:str, mergeable_ranks: dict[bytes, int])->None:\n",
    "        self.pat_str = pat_str\n",
    "        self.mergeable_ranks = mergeable_ranks\n",
    "        self.decoder = {token:token_bytes for token_bytes, token in mergeable_ranks.items()}\n",
    "        self._path = regex.compile(pat_str)\n",
    "    \n",
    "    def decode(self, tokens: list[int])->str:\n",
    "        return b\"\".join(self.decoder[token] for token in tokens).decode(\"utf-8\", errors=\"replace\")\n",
    "    \n",
    "    @staticmethod\n",
    "    def from_hub(repo_id : str, filename: str):\n",
    "        local_path = hf_hub_download(repo_id=repo_id, filename=filename, token = False)\n",
    "        with open(local_path, 'rb') as file:\n",
    "            data = pickle.load(file)\n",
    "        return SimpleBytePairEncoding(pat_str=data['pat_str'],mergeable_ranks=data['mergeable_ranks'])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33c4a9af",
   "metadata": {},
   "source": [
    "### Dataset Class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1348e4b",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-07-20T03:04:07.051298Z",
     "iopub.status.busy": "2025-07-20T03:04:07.051036Z",
     "iopub.status.idle": "2025-07-20T03:04:07.057694Z",
     "shell.execute_reply": "2025-07-20T03:04:07.056960Z",
     "shell.execute_reply.started": "2025-07-20T03:04:07.051277Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "class LoadDataset(Dataset):\n",
    "    def __init__(self, tokens, seq_len = 512, vocab_size = None):\n",
    "        super().__init__()\n",
    "        self.tokens = torch.tensor(tokens, dtype = torch.long)\n",
    "        self.seq_len = seq_len\n",
    "        self.vocab_size = vocab_size\n",
    "        \n",
    "        #Discard Invalid Tokens\n",
    "        if vocab_size is not None:\n",
    "            self.tokens = [tok for tok in tokens if 0 <= tok < vocab_size]\n",
    "            print(f\"Bounded toknes to valid range [0, {vocab_size - 1}]\")\n",
    "        \n",
    "    def __len__(self):\n",
    "        return (len(self.tokens) - self.seq_len)\n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        \n",
    "        #Return Input Sequences and Target {shifted by 1}\n",
    "        x = self.tokens[idx : idx + self.seq_len]\n",
    "        y = self.tokens[idx + 1 : idx + self.seq_len + 1]\n",
    "        \n",
    "        #Check all the tokens are inside the Vocabulary-Range\n",
    "        if self.vocab_size is not None:\n",
    "            x = torch.clamp(input= x, min=0, max=self.vocab_size - 1)\n",
    "            y = torch.clamp(input= y, min=0, max=self.vocab_size - 1)\n",
    "            \n",
    "        return x, y  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d97d46c",
   "metadata": {},
   "source": [
    "## Model Architecture"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b37e563b",
   "metadata": {},
   "source": [
    "### Rotary Positional Embedding"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1a9bc9d",
   "metadata": {},
   "source": [
    "This class implements **Rotary Positional Embeddings**:-\n",
    "- Which encodes positional information via rotating vectors in the complex plane\n",
    "- Which are applied before the Attention computation; Not added to input embeddings like in traditional Transformers\n",
    "- Which helps in Longer Context Handling and better Generalization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "95a3fcc9",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-07-20T03:04:11.951218Z",
     "iopub.status.busy": "2025-07-20T03:04:11.950962Z",
     "iopub.status.idle": "2025-07-20T03:04:11.957782Z",
     "shell.execute_reply": "2025-07-20T03:04:11.957140Z",
     "shell.execute_reply.started": "2025-07-20T03:04:11.951198Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "class RoPE(nn.Module):\n",
    "    def __init__(self, dim: int, max_seq_len: int):\n",
    "        super().__init__()\n",
    "        \n",
    "        #Generate a vector of having exponentially decreasing freqs in range (dim//4)\n",
    "        angular_freq = (1 / 1024) ** torch.linspace(start=0, end=1, steps= dim//4, dtype=torch.float32)\n",
    "        \n",
    "        #Add zeros so that the angular_freq length becomes {dim//2}\n",
    "        angular_freq = torch.cat([angular_freq, angular_freq.new_zeros(dim // 4)])\n",
    "        \n",
    "        #These are the positions \n",
    "        t = torch.arange(max_seq_len, dtype=torch.float32)\n",
    "        \n",
    "        #theta[i][j] = t[i] * angular_freq[j]\n",
    "        theta = torch.einsum(\"i, j -> ij\", t, angular_freq)\n",
    "        \n",
    "        #Store sin & cos as non-learnable buffers\n",
    "        self.register_buffer('cos', theta.cos(), persistent=False)\n",
    "        self.register_buffer('sin', theta.sin(), persistent=False)\n",
    "        \n",
    "    def forward(self, x:torch.Tensor):\n",
    "        assert self.cos.size(0) >= x.size(-3), \"The seq_len (i.e dim(-3) of the input) must be smaller than max_seq_len (i.e dim(0) of self.cos)\"\n",
    "        cos = self.cos[None, :x.size(-3), None, :] #New_Shape:- [1, seq_len, 1, dim//2]\n",
    "        sin = self.sin[None, :x.size(-3), None, :] #New_Shape:- [1, seq_len, 1, dim//2]\n",
    "        \n",
    "        #Split the dim into 2 parts along the last dimension\n",
    "        x1, x2 = x.to(dtype=torch.float32).chunk(2, dim= - 1)\n",
    "        y1 = x1 * cos + x2 * sin\n",
    "        y2 = x1 * (-sin) + x2 * cos\n",
    "        \n",
    "        #Return concatenated Positional_Embeddings\n",
    "        return torch.cat((y1, y2), 3).type_as(x)\n",
    "        \n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63771417",
   "metadata": {},
   "source": [
    "### Mult-Head Attention "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "75fec220",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-07-20T03:04:17.687737Z",
     "iopub.status.busy": "2025-07-20T03:04:17.686989Z",
     "iopub.status.idle": "2025-07-20T03:04:17.694279Z",
     "shell.execute_reply": "2025-07-20T03:04:17.693631Z",
     "shell.execute_reply.started": "2025-07-20T03:04:17.687712Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "class MultiHeadAttention(nn.Module):\n",
    "    ''' \n",
    "        This class is for computing MultiHead-Attention\n",
    "        Input:-\n",
    "            - dim_model: Total Embedding Dimension\n",
    "            - num_heads: Number of attention heads\n",
    "            - max_seq_len: For computing RoPE embeddings\n",
    "\n",
    "        Important Methodologies:-\n",
    "            - Uses fused QKV projection instead of three_separate projections(W_q, W_k, W_v)\n",
    "            - Uses Flash Attention which is faster and memory-efficient\n",
    "    \n",
    "    '''\n",
    "    def __init__(self, dim_model, num_heads, max_seq_len):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.dim_model = dim_model\n",
    "        self.num_heads = num_heads\n",
    "        self.dim_k = dim_model // num_heads\n",
    "        \n",
    "        #Fused Query,Key,Value projection for better memory bandwidth\n",
    "        self.qkv = nn.Linear(dim_model , dim_model * 3, bias= False)\n",
    "        self.w_o = nn.Linear(dim_model, dim_model, bias= False)\n",
    "        \n",
    "        #Add Rotary Positional Embeddings\n",
    "        self.rope = RoPE(self.dim_k, max_seq_len)\n",
    "        \n",
    "    def forward(self, x, mask = None):\n",
    "        batch_size, seq_len = x.size(0), x.size(1)\n",
    "        \n",
    "        #Single Query,Key,Value Projection\n",
    "        qkv = self.qkv(x) #Shape:- [B, Seq_len, Dim_model]\n",
    "        \n",
    "        qkv = qkv.reshape(batch_size, seq_len, 3, self.num_heads, self.dim_k) #Shape: [B, S, D]-------->[B, S, 3, N_heads, D_k]\n",
    "        qkv = qkv.permute(2, 0, 3, 1, 4) #Shape:- [B, S, 3, N_HEADS, DIM_k]----->[3, B, N_HEADS, S, DIM_K]\n",
    "        Q, K, V = qkv[0], qkv[1], qkv[2]\n",
    "        \n",
    "        #Apply Rotary Positional Embeddings to Query & Key\n",
    "        Q = self.rope(Q)\n",
    "        K = self.rope(K)\n",
    "        \n",
    "        #--------Flash Attention---------\n",
    "        #is_causal = True for Auto-Regressive Training\n",
    "        attn_output = F.scaled_dot_product_attention(\n",
    "            query= Q,\n",
    "            key= K,\n",
    "            value= V,\n",
    "            attn_mask= mask,\n",
    "            is_causal= True,\n",
    "            dropout_p= 0.0\n",
    "        )\n",
    "        \n",
    "        attn_output = attn_output.transpose(1, 2).reshape(batch_size, seq_len, self.dim_model)\n",
    "        return self.w_o(attn_output)\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3feb7911",
   "metadata": {},
   "source": [
    "#### Visualize the Multi-Head Attention Flow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "469fe3a8",
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-07-20T03:03:07.260363Z",
     "iopub.status.idle": "2025-07-20T03:03:07.260577Z",
     "shell.execute_reply": "2025-07-20T03:03:07.260486Z",
     "shell.execute_reply.started": "2025-07-20T03:03:07.260476Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "from graphviz import Digraph\n",
    "\n",
    "def visualize_mha():\n",
    "    dot = Digraph(format=\"png\")\n",
    "    \n",
    "    #Nodes \n",
    "    dot.node('x', 'Input x\\n[B,Seq_len,Dim_Model]')\n",
    "    dot.node('qkv', label='QKV Projection\\n[B,Seq_len,3*Dim_Model]')\n",
    "    dot.node('reshape', label='Reshape\\n[3,B,H,Seq_len,Dim_k]')\n",
    "    dot.node('Q', 'Query\\n[B,H,Seq_len,Dim_Key]')\n",
    "    dot.node('K', 'Key\\n[B,H,Seq_len,Dim_Key]')\n",
    "    dot.node('V', 'Value\\n[B,H,Seq_len,Dim_Key]')\n",
    "    dot.node('rope', 'Rotary Positional Embedding')\n",
    "    dot.node('attn', 'Flash Attention\\n[B, H, Seq_len, Dim_k]')\n",
    "    dot.node('merge', 'Merge Heads\\n[B, Seq_len, Dim_model]')\n",
    "    dot.node('out_proj', 'Output Linear\\n[B, Seq_len, Dim_model]')\n",
    "    \n",
    "    \n",
    "    #Define Edges of the Graph\n",
    "    dot.edge('x', 'qkv')\n",
    "    dot.edge('qkv', 'reshape')\n",
    "    dot.edge('reshape', 'Q')\n",
    "    dot.edge('reshape', 'K')\n",
    "    dot.edge('reshape', 'V')\n",
    "    dot.edge('Q', 'rope')\n",
    "    dot.edge('K', 'rope')\n",
    "    dot.edge('rope', 'attn', label='Q, K')\n",
    "    dot.edge('V', 'attn', label = 'V')\n",
    "    dot.edge('attn', 'merge')\n",
    "    dot.edge('merge', 'out_proj')\n",
    "    \n",
    "    os.makedirs(\"./saves/mha\", exist_ok=True)\n",
    "    dot.render('multihead_attention_flow', view=True, directory=\"./saves/mha\", )\n",
    "\n",
    "visualize_mha()    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "808ae614",
   "metadata": {},
   "source": [
    "![Multi-Head Attention Flow](./saves/mha/multihead_attention_flow.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba83b958",
   "metadata": {},
   "source": [
    "### Feed-Forward Network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "e4f9351b",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-07-20T03:04:22.167191Z",
     "iopub.status.busy": "2025-07-20T03:04:22.166963Z",
     "iopub.status.idle": "2025-07-20T03:04:22.172271Z",
     "shell.execute_reply": "2025-07-20T03:04:22.171601Z",
     "shell.execute_reply.started": "2025-07-20T03:04:22.167174Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "class FeedForward(nn.Module):\n",
    "    ''' \n",
    "        Input(Args):\n",
    "            - dim_model: Dimension of the Input&Output Embeddings\n",
    "            - dim_ff: Dimension of the FeedForward layers{Generally larger than dim_model}\n",
    "        \n",
    "        Methodologies:\n",
    "            - Take an Input project it onto Higher_dim space to extract features from it\n",
    "            - Apply gated Activation Function like SwiGLU for better expressiveness and efficient performance\n",
    "            - Project back to the original dimension(dim_model)\n",
    "    '''\n",
    "    def __init__(self, dim_model , dim_ff):\n",
    "        super().__init__()\n",
    "        self.fc1 = nn.Linear(in_features= dim_model, out_features=dim_ff, bias=False)\n",
    "        self.fc2 = nn.Linear(in_features=dim_ff // 2, out_features=dim_model, bias=False)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.fc1(x)\n",
    "        x1, x2 = x.chunk(2, dim = -1)\n",
    "        #Add a Activation function (SwiGLU)\n",
    "        x = F.silu(x1) * x2\n",
    "        x = self.fc2(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0bf2d393",
   "metadata": {},
   "source": [
    "#### Visualize Feed-Forward Network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "486330dc",
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-07-20T03:03:07.262261Z",
     "iopub.status.idle": "2025-07-20T03:03:07.262588Z",
     "shell.execute_reply": "2025-07-20T03:03:07.262445Z",
     "shell.execute_reply.started": "2025-07-20T03:03:07.262431Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "def visualize_feed_forward():\n",
    "    dot = Digraph(format=\"png\")\n",
    "    \n",
    "    \n",
    "    dot.node('Input', 'Input\\n[B,Seq_Len,Dim_Model]', shape = 'box')\n",
    "    dot.node('FC1', 'Linear\\n(dim_model---->dim_ff)', shape = 'box')\n",
    "    dot.node('Split', 'Split(x--->x1, x2)', shape = 'box')\n",
    "    dot.node('SiLU', 'SiLU(x1)', shape = 'box')\n",
    "    dot.node('Gate', 'SiLU(x1) * x2\\n(SwiGLU)', shape = 'box')\n",
    "    dot.node('FC2', 'Linear\\n(dim_ff---->dim_model)', shape = 'box')\n",
    "    dot.node('Output', 'Output\\n[B,Seq_len,Dim_model]', shape = 'box')\n",
    "    \n",
    "    dot.edge('Input', 'FC1')\n",
    "    dot.edge('FC1', 'Split')\n",
    "    dot.edge('Split', 'SiLU')\n",
    "    dot.edge('Split', 'Gate')\n",
    "    dot.edge('SiLU', 'Gate')\n",
    "    dot.edge('Gate', 'FC2')\n",
    "    dot.edge('FC2', 'Output')\n",
    "    \n",
    "    os.makedirs(\"./saves/ffnn\", exist_ok=True)\n",
    "    dot.render('FeedForward_Network', view=True, directory=\"./saves/ffnn\")\n",
    "    \n",
    "    \n",
    "    \n",
    "visualize_feed_forward() \n",
    "    \n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a4f2e55",
   "metadata": {},
   "source": [
    "![Feed-Forward Network Flow](./saves/ffnn/FeedForward_Network.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d7ddab2e",
   "metadata": {},
   "source": [
    "### Transformer Block"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "52f71711",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-07-20T03:04:28.036031Z",
     "iopub.status.busy": "2025-07-20T03:04:28.035511Z",
     "iopub.status.idle": "2025-07-20T03:04:28.041184Z",
     "shell.execute_reply": "2025-07-20T03:04:28.040535Z",
     "shell.execute_reply.started": "2025-07-20T03:04:28.036006Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "class TransformerBlock(nn.Module):\n",
    "    ''' \n",
    "        Input(Args):\n",
    "            - dim_Model: The dimension of the Input & Output Embeddings\n",
    "            - dim_ff: The dimension of the Feed_Forward Layers\n",
    "            - max_seq_len: The maximum sequence length\n",
    "            \n",
    "        Methodology:\n",
    "            - Add a Normalization to the Input\n",
    "            - Project the Input through the Attention Mechanism First\n",
    "            - Add Residual Connection for better generalized results\n",
    "            - Normalize the Input before Projecting through the Feed_Forward Layer \n",
    "            - Project through the Feed_Forward layers\n",
    "            - Add Residual_Connections\n",
    "    '''\n",
    "    def __init__(self, dim_model , num_heads, dim_ff, max_seq_len):\n",
    "        super().__init__()\n",
    "        self.attention = MultiHeadAttention(dim_model, num_heads, max_seq_len)\n",
    "        self.feed_forward = FeedForward(dim_model, dim_ff)\n",
    "        self.norm1 = nn.RMSNorm(dim_model)\n",
    "        self.norm2 = nn.RMSNorm(dim_model)\n",
    "        \n",
    "    def forward(self, x, mask = None):\n",
    "        #Copy Input to a Residual Variable for Residual Connection\n",
    "        residual = x\n",
    "        \n",
    "        #Project Input through the Self-Attention \n",
    "        attn_out = self.attention(self.norm1(x), mask)\n",
    "        \n",
    "        #Add Residual Connection\n",
    "        x = residual + attn_out\n",
    "        \n",
    "        #Again copy Input to be used for Resiudal Connection after Feed_Forward pass\n",
    "        residual = x\n",
    "        \n",
    "        ff_out = self.feed_forward(self.norm2(x))\n",
    "        \n",
    "        x = residual + ff_out\n",
    "        \n",
    "        return x\n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7c7aa22",
   "metadata": {},
   "source": [
    "#### Visualize Transformer Block"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba4369f9",
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-07-20T03:03:07.264660Z",
     "iopub.status.idle": "2025-07-20T03:03:07.264861Z",
     "shell.execute_reply": "2025-07-20T03:03:07.264762Z",
     "shell.execute_reply.started": "2025-07-20T03:03:07.264754Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "def visualize_transformer():\n",
    "    dot = Digraph(format = \"png\")\n",
    "    \n",
    "    dot.node('Input', 'Input\\n[B,Seq_len,D_Model]', shape = 'box')\n",
    "    dot.node('Norm1', 'RMSNorm', shape = 'ellipse')\n",
    "    dot.node('MHA', 'Multi-Head Attention',shape = 'box')\n",
    "    dot.node('Res1', 'Residual Connection',shape = 'diamond')\n",
    "    dot.node('Norm2', 'RMSNorm',shape = 'ellipse')\n",
    "    dot.node('FFNN', 'Feed-Forward\\n(SwiGLU)',shape = 'box')\n",
    "    dot.node('Res2', 'Residual Connection',shape = 'diamond')\n",
    "    dot.node('Output', 'Output\\n[B,Seq_len,Dim_model]', shape = 'box')\n",
    "    \n",
    "    dot.edge('Input','Norm1')\n",
    "    dot.edge('Input', 'Res1')\n",
    "    dot.edge('Norm1', 'MHA')\n",
    "    dot.edge('MHA', 'Res1')\n",
    "    dot.edge('Res1', 'Norm2')\n",
    "    dot.edge('Res1', 'Res2')\n",
    "    dot.edge('Norm2', 'FFNN')\n",
    "    dot.edge('FFNN','Res2')\n",
    "    dot.edge('Res2', 'Output')\n",
    "    \n",
    "    os.makedirs(\"./saves/transformer\", exist_ok=True)\n",
    "    dot.render('Transformer_Block', view = True, directory=\"./saves/transformer\")\n",
    "    \n",
    "visualize_transformer()\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b19ffe47",
   "metadata": {},
   "source": [
    "![Transformer Flow](./saves/transformer/Transformer_Block.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d94a223",
   "metadata": {},
   "source": [
    "### Model Class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "1b3e73a0",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-07-20T03:04:31.561928Z",
     "iopub.status.busy": "2025-07-20T03:04:31.561095Z",
     "iopub.status.idle": "2025-07-20T03:04:31.572752Z",
     "shell.execute_reply": "2025-07-20T03:04:31.572026Z",
     "shell.execute_reply.started": "2025-07-20T03:04:31.561898Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "class SmoLLM(nn.Module):\n",
    "    ''' \n",
    "        Input(Args):\n",
    "            - vocab_size: Total number of tokens in the vocabulary\n",
    "            - dim_model: Dimension of the token embeddings and model internal layers\n",
    "            - num_heads: Number of attention heads in each transformer block\n",
    "            - num_layers: Number of Transformer Blocks\n",
    "            - dim_ff: Dimension of the FeedForward layers\n",
    "            - max_len: Max sequence length supported by the model\n",
    "            \n",
    "        Methodologies:\n",
    "            - Embed Input_Tokens into high-dimensional space using learned embeddings\n",
    "            - Multiply embeddings \n",
    "            - Generate Causal_Mask to allow model to see the next_tokens for better and quick training\n",
    "            - Pass the embeddings through transformer blocks\n",
    "            - Normalize the Final Output embeddings\n",
    "            - Custom weight initialization\n",
    "            - Project final embeddings onto vocabulary_space using weights from the token_embedding layer\n",
    "            - Generate text by doing sampling using softmax\n",
    "            - Control the Uniqueness of the Generated_Text by scaling with Temperature_scale\n",
    "    '''\n",
    "    \n",
    "    def __init__(self, vocab_size, dim_model, num_heads, num_layers, dim_ff, max_len):\n",
    "        super().__init__()\n",
    "        self.dim_model = dim_model\n",
    "        self.max_len = max_len\n",
    "        self.vocab_size = vocab_size\n",
    "        \n",
    "        self.token_embedding = nn.Embedding(vocab_size, dim_model)\n",
    "        \n",
    "        self.trf_blocks = nn.ModuleList(\n",
    "            [TransformerBlock(dim_model, num_heads, dim_ff, max_len) for _ in range(num_layers)]\n",
    "        )\n",
    "        \n",
    "        self.final_norm = nn.RMSNorm(dim_model)\n",
    "        \n",
    "        #To initialize every module(like Linear, Embedding) with custom weight\n",
    "        self.apply(self._init_weights)\n",
    "    \n",
    "    #Custom\n",
    "    def _init_weights(self, module):\n",
    "        if isinstance(module, nn.Linear):\n",
    "            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)\n",
    "            if module.bias is not None:\n",
    "                torch.nn.init.zeros_(module.bias)\n",
    "        elif isinstance(module, nn.Embedding):\n",
    "            torch.nn.init.normal_(module.weight, mean=0.0, std = 0.02)\n",
    "        \n",
    "    def forward(self, x, use_causal_mask = True):\n",
    "        \n",
    "        seq_len = x.size(1)\n",
    "        \n",
    "        x = self.token_embedding(x) * math.sqrt(self.dim_model)\n",
    "        \n",
    "        #Create Causal_Mask \n",
    "        mask = None\n",
    "        if use_causal_mask:\n",
    "            mask = torch.tril(torch.ones(seq_len, seq_len)).unsqueeze(0).unsqueeze(0).to(x.device) #Shape: [seq_len, seq_len]----->[1, seq_len ,seq_len]----->[1, 1, seq_len, seq_len]\n",
    "            \n",
    "        #TransformerBlocks\n",
    "        for block in self.trf_blocks:\n",
    "            x = block(x, mask)\n",
    "        \n",
    "        #Output\n",
    "        x = self.final_norm(x)\n",
    "        #Use the Embedding layer weights for output projection\n",
    "        logits = F.linear(x, self.token_embedding.weight)\n",
    "        \n",
    "        return logits\n",
    "\n",
    "    def generate(self, tokenizer, prompt_tokens, max_new_tokens = 50, temperature = 0.7):\n",
    "        \n",
    "        self.eval()\n",
    "        \n",
    "        with torch.no_grad():\n",
    "            tokens = prompt_tokens.copy()\n",
    "            \n",
    "            for _ in range(max_new_tokens):\n",
    "                input_tokens = tokens[-self.max_len:]\n",
    "                x = torch.tensor(input_tokens).unsqueeze(0)\n",
    "                \n",
    "                if torch.cuda.is_available():\n",
    "                    x = x.cuda()\n",
    "                    \n",
    "                #Clamp tokens in the range[0, Vocab_size-1]\n",
    "                x = torch.clamp(x, 0, self.vocab_size - 1)\n",
    "                \n",
    "                #Forward Pass \n",
    "                with torch.amp.autocast('cuda'):\n",
    "                    logits = self(x)\n",
    "                \n",
    "                logits = logits[0, -1, :] / temperature\n",
    "                \n",
    "                #Sample Next Token\n",
    "                probs = F.softmax(logits, dim = -1)\n",
    "                next_token = torch.multinomial(probs, 1).item()\n",
    "                \n",
    "                next_token = min(max(next_token , 0), self.vocab_size - 1)\n",
    "                tokens.append(next_token)\n",
    "                \n",
    "        return tokenizer.decode(tokens)\n",
    "        \n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16d05143",
   "metadata": {},
   "source": [
    "#### Visualize the Flow of our Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "473d7a91",
   "metadata": {
    "execution": {
     "iopub.status.busy": "2025-07-20T03:03:07.266822Z",
     "iopub.status.idle": "2025-07-20T03:03:07.267024Z",
     "shell.execute_reply": "2025-07-20T03:03:07.266938Z",
     "shell.execute_reply.started": "2025-07-20T03:03:07.266929Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "def visualize_smollm():\n",
    "    dot = Digraph(format = \"png\")\n",
    "    dot.attr(size = '15')\n",
    "    \n",
    "    dot.node('Tokens', 'Input Tokens\\n[B, Seq_len]', shape = 'box')\n",
    "    dot.node('Embed', 'Token Embedding\\n[B, Seq_len, Dim_model]', shape = 'ellipse')\n",
    "    dot.node('Scale', 'Scale by sqrt(dim_model)', shape = 'ellipse')\n",
    "    dot.node('Mask', 'Causal Mask\\n[seq_len * seq_len]', shape = 'diamond')\n",
    "    dot.node('Transformer_Stack', f'{len(\"X\"*8)} Transformer Blocks', shape = 'box3d')\n",
    "    dot.node('Norm', 'RMSNorm', shape='ellipse')\n",
    "    dot.node('Output', 'Linear Projection\\n[B, Seq_len, Vocab_size]', shape = 'box')\n",
    "    dot.node('Logits', 'Logits', shape = 'box')\n",
    "    \n",
    "    dot.edge(\"Tokens\", \"Embed\")\n",
    "    dot.edge('Embed','Scale')\n",
    "    dot.edge('Scale','Transformer_Stack')\n",
    "    dot.edge('Mask','Transformer_Stack', style = 'dashed', label = '(Optional)')\n",
    "    dot.edge('Transformer_Stack', 'Norm')\n",
    "    dot.edge('Norm', 'Output')\n",
    "    dot.edge('Output', 'Logits')\n",
    "    \n",
    "    #Save & Render\n",
    "    os.makedirs(\"./saves/smollm\", exist_ok=True)\n",
    "    dot.render('smollm_visualization', view=True, directory = \"./saves/smollm/\")\n",
    "\n",
    "visualize_smollm()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c72f6d9",
   "metadata": {},
   "source": [
    "![SmoLLM Flow](./saves/smollm/smollm_visualization.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "86921ba1",
   "metadata": {},
   "source": [
    "## Model Functionality"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9cc751c",
   "metadata": {},
   "source": [
    "### Config Class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "931f6c16",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-07-20T04:10:54.948725Z",
     "iopub.status.busy": "2025-07-20T04:10:54.948408Z",
     "iopub.status.idle": "2025-07-20T04:10:54.954661Z",
     "shell.execute_reply": "2025-07-20T04:10:54.953848Z",
     "shell.execute_reply.started": "2025-07-20T04:10:54.948701Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "class Config:\n",
    "    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "    seq_len = 512 #Sequnce Length of the Input\n",
    "    dim_model = 384 #Dimension of the Input and Output embeddings and layers of the model\n",
    "    num_heads = 6   #Number of Heads in a Transformer Blocks\n",
    "    num_layers = 8  #Number of Transformer Blocks\n",
    "    dim_ff = 384 * 4  #Feed-Forward Dimension\n",
    "    batch_size = 4    \n",
    "    learning_rate = 3e-4\n",
    "    epochs = 3\n",
    "    warmup_steps = 1000  \n",
    "    cooldown_frac = 0.4\n",
    "    save_every = 10000    #Save checkpoint step\n",
    "    use_amp = True\n",
    "    save_checkpoint_path = \"./saves/checkpoint\"\n",
    "    tokenizer_save_path = \"./saves/tokenizer\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc6e8447",
   "metadata": {},
   "source": [
    "### Load Tokenizer & Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "d7d3a05b",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-07-20T03:04:44.381260Z",
     "iopub.status.busy": "2025-07-20T03:04:44.380985Z",
     "iopub.status.idle": "2025-07-20T03:04:44.385513Z",
     "shell.execute_reply": "2025-07-20T03:04:44.384917Z",
     "shell.execute_reply.started": "2025-07-20T03:04:44.381238Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "os.makedirs(Config.tokenizer_save_path, exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ccd03fd6",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-07-20T03:04:46.437018Z",
     "iopub.status.busy": "2025-07-20T03:04:46.436313Z",
     "iopub.status.idle": "2025-07-20T03:04:46.812028Z",
     "shell.execute_reply": "2025-07-20T03:04:46.811328Z",
     "shell.execute_reply.started": "2025-07-20T03:04:46.436994Z"
    },
    "trusted": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "======== Loading Tokenizer and Data ========\n"
     ]
    }
   ],
   "source": [
    "print(\"=\"*8, \"Loading Tokenizer and Data\",\"=\"*8)\n",
    "repo_id = \"vukrosic/essential-web-16k-tokenizer\"\n",
    "\n",
    "#Download Tokneizer File\n",
    "tokenizer_path = hf_hub_download(\n",
    "    repo_id= repo_id,\n",
    "    filename=\"bpe_tokenizer_16k_n1000000.pkl\",\n",
    "    token = False,``\n",
    "    local_dir=Config.tokenizer_save_path,\n",
    ")\n",
    "\n",
    "#Download Tokens File\n",
    "tokens_path = hf_hub_download(\n",
    "    repo_id=repo_id,\n",
    "    filename = \"essential_web_500k_tokens.pkl\",\n",
    "    token = False,\n",
    "    local_dir=Config.tokenizer_save_path\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "c3712229",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-07-20T03:04:49.761145Z",
     "iopub.status.busy": "2025-07-20T03:04:49.760882Z",
     "iopub.status.idle": "2025-07-20T03:04:49.789802Z",
     "shell.execute_reply": "2025-07-20T03:04:49.789039Z",
     "shell.execute_reply.started": "2025-07-20T03:04:49.761122Z"
    },
    "trusted": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded: 123,327 tokens\n",
      "Vocabulary Size:- 16,384\n",
      "========== Checking Token Range...... ==========\n",
      "Token Range: [10, 16383]\n",
      "No Tokens to Discard as all are Valid Tokens: 123327\n"
     ]
    }
   ],
   "source": [
    "#-----------Load Tokenizer----------- \n",
    "with open(tokenizer_path, \"rb\") as f:\n",
    "    data = pickle.load(f)\n",
    "tokenizer = SimpleBytePairEncoding(\n",
    "    pat_str=data[\"pat_str\"],\n",
    "    mergeable_ranks=data[\"mergeable_ranks\"]\n",
    ")\n",
    "#-----------Load Tokens-----------\n",
    "with open(tokens_path, \"rb\") as f:\n",
    "    tokens = pickle.load(f)\n",
    "\n",
    "print(f\"Loaded: {len(tokens):,} tokens\")\n",
    "\n",
    "#-----------Get Vocab Size-----------\n",
    "vocab_size = len(tokenizer.mergeable_ranks)\n",
    "print(f\"Vocabulary Size:- {vocab_size:,}\")\n",
    "\n",
    "#-----------Check the token range-----------\n",
    "print(\"=\"*10,f\"Checking Token Range......\",\"=\"*10)\n",
    "min_token = min(tokens)\n",
    "max_token = max(tokens)\n",
    "print(f\"Token Range: [{min_token}, {max_token}]\")\n",
    "\n",
    "#-----------Discard Invalid Tokens-----------\n",
    "valid_tokens = [t for t in tokens if 0 <= t < vocab_size]\n",
    "if(len(tokens) != len(valid_tokens)):\n",
    "    print(f\"Discarded from:{len(tokens):,} to {len(valid_tokens):,}\")\n",
    "else:\n",
    "    print(f\"No Tokens to Discard as all are Valid Tokens: {len(valid_tokens)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "9acdf0c1",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-07-20T03:05:04.706361Z",
     "iopub.status.busy": "2025-07-20T03:05:04.705889Z",
     "iopub.status.idle": "2025-07-20T03:05:04.717363Z",
     "shell.execute_reply": "2025-07-20T03:05:04.716664Z",
     "shell.execute_reply.started": "2025-07-20T03:05:04.706335Z"
    },
    "trusted": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Bounded toknes to valid range [0, 16383]\n",
      "Dataset Size: 122,815 sequences\n",
      "Total Training Steps Per Epoch:92,112\n"
     ]
    }
   ],
   "source": [
    "train_dataset = LoadDataset(valid_tokens, seq_len=Config.seq_len, vocab_size=vocab_size)\n",
    "train_loader = DataLoader(train_dataset, batch_size=Config.batch_size, shuffle = True)\n",
    "total_steps = len(train_loader) * Config.epochs\n",
    "\n",
    "print(f\"Dataset Size: {len(train_dataset):,} sequences\")\n",
    "print(f\"Total Training Steps Per Epoch:{total_steps:,}\") "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "367edb2a",
   "metadata": {},
   "source": [
    "### Initialize Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "d27e2672",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-07-20T04:10:58.632450Z",
     "iopub.status.busy": "2025-07-20T04:10:58.632190Z",
     "iopub.status.idle": "2025-07-20T04:10:58.910948Z",
     "shell.execute_reply": "2025-07-20T04:10:58.910266Z",
     "shell.execute_reply.started": "2025-07-20T04:10:58.632428Z"
    },
    "trusted": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using Device: cpu\n"
     ]
    }
   ],
   "source": [
    "model = SmoLLM(\n",
    "    vocab_size= vocab_size,\n",
    "    dim_model= Config.dim_model,\n",
    "    num_heads= Config.num_heads,\n",
    "    num_layers= Config.num_layers,\n",
    "    dim_ff= Config.dim_ff,\n",
    "    max_len= Config.seq_len\n",
    ")\n",
    "\n",
    "\n",
    "#Move the Model to Device\n",
    "model = model.to(Config.device)\n",
    "print(f\"Using Device: {Config.device}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f070941",
   "metadata": {},
   "source": [
    "- Compile The Model Before Training for better Efficiency and Faster Computations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "15fbad6d",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-07-20T04:11:01.187335Z",
     "iopub.status.busy": "2025-07-20T04:11:01.187071Z",
     "iopub.status.idle": "2025-07-20T04:11:01.195600Z",
     "shell.execute_reply": "2025-07-20T04:11:01.194947Z",
     "shell.execute_reply.started": "2025-07-20T04:11:01.187315Z"
    },
    "trusted": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Compiling model.......\n",
      "Model Compiled Successfully\n"
     ]
    }
   ],
   "source": [
    "try:\n",
    "    print(\"Compiling model.......\")\n",
    "    model = torch.compile(model, mode = 'default')\n",
    "    print(\"Model Compiled Successfully\")\n",
    "except Exception as e:\n",
    "    print(f\"Failed to Compile the model: {e}\")\n",
    "    print(f\"Continuing without Compiling\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e4a6f2d5",
   "metadata": {},
   "source": [
    "- Number of Parameters in our Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "7f8304ad-484e-4f47-86ca-4bc6d4c98271",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-07-20T04:10:47.747874Z",
     "iopub.status.busy": "2025-07-20T04:10:47.747607Z",
     "iopub.status.idle": "2025-07-20T04:10:47.757040Z",
     "shell.execute_reply": "2025-07-20T04:10:47.756286Z",
     "shell.execute_reply.started": "2025-07-20T04:10:47.747852Z"
    },
    "trusted": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "======== Number of Params ========\n",
      "Total Parameters:18,094,464\n",
      "Trainable Params:18,094,464\n"
     ]
    }
   ],
   "source": [
    "total_params = sum(p.numel() for p in model.parameters())\n",
    "trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
    "print(\"=\"*8, \"Number of Params\", \"=\"*8)\n",
    "print(f\"Total Parameters:{total_params:,}\")\n",
    "print(f\"Trainable Params:{trainable_params:,}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cd1e3988",
   "metadata": {},
   "source": [
    "- Get the Optimizer & Grad_Scaler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "3d69db69",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-07-20T04:11:05.431231Z",
     "iopub.status.busy": "2025-07-20T04:11:05.430649Z",
     "iopub.status.idle": "2025-07-20T04:11:05.436636Z",
     "shell.execute_reply": "2025-07-20T04:11:05.435983Z",
     "shell.execute_reply.started": "2025-07-20T04:11:05.431208Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "optimizer = optim.AdamW(\n",
    "    model.parameters(),\n",
    "    lr= Config.learning_rate,\n",
    "    weight_decay=0.1,\n",
    "    fused = True,\n",
    ")\n",
    "\n",
    "scaler = torch.amp.GradScaler() if Config.use_amp else None"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec4e3e95",
   "metadata": {},
   "source": [
    "- Learning Rate Scheduler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "ebba257f",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-07-20T04:11:08.254401Z",
     "iopub.status.busy": "2025-07-20T04:11:08.253840Z",
     "iopub.status.idle": "2025-07-20T04:11:08.260038Z",
     "shell.execute_reply": "2025-07-20T04:11:08.259507Z",
     "shell.execute_reply.started": "2025-07-20T04:11:08.254379Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "def get_lr(step: int):\n",
    "    x = step / total_steps \n",
    "    assert 0 <= x <= 1\n",
    "    if x < 1 - Config.cooldown_frac:\n",
    "        #Stable Phase Learning Rate stays at full-scale(1.0)\n",
    "        return 1.0\n",
    "    else:\n",
    "        #Cool_Down Phase; Learning_Rate gradually decreases\n",
    "        w = (1 - x) / Config.cooldown_frac\n",
    "        return w * 1.0 + (1 - w) * 0.1\n",
    "    \n",
    "def lr_lambda(step):\n",
    "    \n",
    "    #Warmup-Phase for early_steps; Learning Rate gradually increases\n",
    "    if step < Config.warmup_steps:\n",
    "        return step / Config.warmup_steps\n",
    "    else:\n",
    "        #Apply the Advanced Schedule after warmup\n",
    "        return get_lr(step - Config.warmup_steps)\n",
    "\n",
    "scheduler = optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "173ae8e7",
   "metadata": {},
   "source": [
    "### Train Function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "96dd983c-b3bc-4bbf-9acc-eb27d28faa72",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-07-20T04:11:11.441336Z",
     "iopub.status.busy": "2025-07-20T04:11:11.440729Z",
     "iopub.status.idle": "2025-07-20T04:11:11.445620Z",
     "shell.execute_reply": "2025-07-20T04:11:11.444843Z",
     "shell.execute_reply.started": "2025-07-20T04:11:11.441310Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "import torch._dynamo\n",
    "torch._dynamo.config.suppress_errors = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "385464c0-a4cf-4fd9-b976-6c7a551a4c20",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-07-20T04:11:13.691883Z",
     "iopub.status.busy": "2025-07-20T04:11:13.691617Z",
     "iopub.status.idle": "2025-07-20T04:11:13.698556Z",
     "shell.execute_reply": "2025-07-20T04:11:13.697645Z",
     "shell.execute_reply.started": "2025-07-20T04:11:13.691864Z"
    },
    "trusted": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Trainable Params: True\n"
     ]
    }
   ],
   "source": [
    "trainable_p = any(p.requires_grad for p in model.parameters())\n",
    "print(f\"Trainable Params: {trainable_p}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "ff546318-6016-4c6a-ab2e-0502f987307d",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-07-20T03:29:33.441556Z",
     "iopub.status.busy": "2025-07-20T03:29:33.440970Z",
     "iopub.status.idle": "2025-07-20T03:29:39.789929Z",
     "shell.execute_reply": "2025-07-20T03:29:39.789305Z",
     "shell.execute_reply.started": "2025-07-20T03:29:33.441521Z"
    },
    "trusted": true
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "Finishing previous runs because reinit is set to 'default'."
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<br>    <style><br>        .wandb-row {<br>            display: flex;<br>            flex-direction: row;<br>            flex-wrap: wrap;<br>            justify-content: flex-start;<br>            width: 100%;<br>        }<br>        .wandb-col {<br>            display: flex;<br>            flex-direction: column;<br>            flex-basis: 100%;<br>            flex: 1;<br>            padding: 10px;<br>        }<br>    </style><br><div class=\"wandb-row\"><div class=\"wandb-col\"><h3>Run history:</h3><br/><table class=\"wandb\"><tr><td>learning_rate</td><td>▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁</td></tr><tr><td>step_loss</td><td>▂▄▄▂▁▂▃▂▅▄▂▄▁█▅▁▄▁▄▂▄▄▃▂▃▄▁▄▆▄▂▃▄▂▁▄▅▄▂▃</td></tr></table><br/></div><div class=\"wandb-col\"><h3>Run summary:</h3><br/><table class=\"wandb\"><tr><td>learning_rate</td><td>0.0003</td></tr><tr><td>step_loss</td><td>0.04355</td></tr></table><br/></div></div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View run <strong style=\"color:#cdcd00\">Train_Model_3</strong> at: <a href='https://wandb.ai/atharv3105-dr-a-p-j-abdul-kalam-technical-university/Small_LM/runs/iqu2fy5b' target=\"_blank\">https://wandb.ai/atharv3105-dr-a-p-j-abdul-kalam-technical-university/Small_LM/runs/iqu2fy5b</a><br> View project at: <a href='https://wandb.ai/atharv3105-dr-a-p-j-abdul-kalam-technical-university/Small_LM' target=\"_blank\">https://wandb.ai/atharv3105-dr-a-p-j-abdul-kalam-technical-university/Small_LM</a><br>Synced 5 W&B file(s), 1 media file(s), 2 artifact file(s) and 0 other file(s)"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Find logs at: <code>./wandb/run-20250720_032654-iqu2fy5b/logs</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Tracking run with wandb version 0.20.1"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Run data is saved locally in <code>/kaggle/working/wandb/run-20250720_032933-74fczcc2</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Syncing run <strong><a href='https://wandb.ai/atharv3105-dr-a-p-j-abdul-kalam-technical-university/Small_LM/runs/74fczcc2' target=\"_blank\">Train_Model_4</a></strong> to <a href='https://wandb.ai/atharv3105-dr-a-p-j-abdul-kalam-technical-university/Small_LM' target=\"_blank\">Weights & Biases</a> (<a href='https://wandb.me/developer-guide' target=\"_blank\">docs</a>)<br>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View project at <a href='https://wandb.ai/atharv3105-dr-a-p-j-abdul-kalam-technical-university/Small_LM' target=\"_blank\">https://wandb.ai/atharv3105-dr-a-p-j-abdul-kalam-technical-university/Small_LM</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View run at <a href='https://wandb.ai/atharv3105-dr-a-p-j-abdul-kalam-technical-university/Small_LM/runs/74fczcc2' target=\"_blank\">https://wandb.ai/atharv3105-dr-a-p-j-abdul-kalam-technical-university/Small_LM/runs/74fczcc2</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "wandb.init(project = \"Small_LM\",name= \"Train_Model_4\")\n",
    "table = wandb.Table(columns = [\"Step\", \"Prompt\", \"Generated\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3428058",
   "metadata": {},
   "source": [
    "### For further Training "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3160f150",
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint_path = \"./saves/checkpoint/num_steps30000.pt\"\n",
    "print(\"=\"*10,f\"Loading checkpoint from {checkpoint_path} to resume training...\", \"=\"*10)\n",
    "checkpoint = torch.load(checkpoint_path, map_location=Config.device)\n",
    "\n",
    "#-------Restore the State-Dictionaries of all networks---------\n",
    "model.load_state_dict(checkpoint['model_state_dict'])\n",
    "optimizer.load_state_dict(checkpoint['optimizer_state_dict'])\n",
    "scheduler.load_state_dict(checkpoint['scheduler_state_dict'])\n",
    "scaler.load_state_dict(scaler['scaler_state_dict'])\n",
    "step = checkpoint['step']\n",
    "\n",
    "#-----Set the Model to Train------\n",
    "model.train()\n",
    "print(f\"State_Dictionaries successfully restored; Continue training from step:{step}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0b4d81c",
   "metadata": {},
   "source": [
    "### Training Loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "685cbb3d",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-07-20T04:11:22.187401Z",
     "iopub.status.busy": "2025-07-20T04:11:22.187136Z",
     "iopub.status.idle": "2025-07-20T05:00:38.722589Z",
     "shell.execute_reply": "2025-07-20T05:00:38.721816Z",
     "shell.execute_reply.started": "2025-07-20T04:11:22.187380Z"
    },
    "trusted": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "============================== Starting Training ==============================\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch:[1/3]:  33%|███▎      | 9996/30704 [05:19<10:59, 31.39it/s, Loss=0.1195, Lr=3.00e-04, Step=1e+4]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Checkpoint saved at step: 10000\n",
      "========== Sample Text Generation ==========\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch:[1/3]:  33%|███▎      | 10004/30704 [05:20<26:34, 12.99it/s, Loss=0.1300, Lr=3.00e-04, Step=1e+4]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generated: Wednesday, May 25, 2011\n",
      "\n",
      "The foliina Niemenkari, 2011 Muslims are two of the whole wise to small cadets, a job. The\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch:[1/3]:  65%|██████▌   | 19996/30704 [10:37<05:36, 31.81it/s, Loss=0.0938, Lr=3.00e-04, Step=2e+4] "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Checkpoint saved at step: 20000\n",
      "========== Sample Text Generation ==========\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m You are mutating a Table with log_mode='IMMUTABLE' that has been logged already. Subsequent log() calls will have no effect. Set log_mode='MUTABLE' to enable re-logging after mutations\n",
      "Epoch:[1/3]:  65%|██████▌   | 20004/30704 [10:38<13:17, 13.42it/s, Loss=0.0446, Lr=3.00e-04, Step=2e+4]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generated: Wednesday, May 25, 2011\n",
      "\n",
      "SHINY SUITING THEORY\n",
      "\n",
      "\n",
      "\n",
      "Vintage Wayf Rick Owens, Laitinen\n",
      "\n",
      "Silk   Coded by\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch:[1/3]:  98%|█████████▊| 29996/30704 [15:54<00:22, 31.67it/s, Loss=0.0502, Lr=3.00e-04, Step=3e+4] "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Checkpoint saved at step: 30000\n",
      "========== Sample Text Generation ==========\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch:[1/3]:  98%|█████████▊| 30004/30704 [15:55<00:43, 15.94it/s, Loss=0.0359, Lr=3.00e-04, Step=3e+4]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generated: Wednesday, May 25, 2011\n",
      "\n",
      "Getting a job is increasingly about who you know — and that’s making the playing field less even for African-Americans. In a new study published by the\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch:[1/3]: 100%|██████████| 30704/30704 [16:50<00:00, 30.39it/s, Loss=0.0399, Lr=3.00e-04, Step=30704]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch:[1/3] Completed; Avg_Loss: 0.4236\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch:[2/3]:  30%|███       | 9292/30704 [04:55<11:16, 31.67it/s, Loss=0.0330, Lr=3.00e-04, Step=4e+4] "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Checkpoint saved at step: 40000\n",
      "========== Sample Text Generation ==========\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch:[2/3]:  30%|███       | 9300/30704 [04:55<22:04, 16.16it/s, Loss=0.0269, Lr=3.00e-04, Step=4e+4]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generated: Wednesday, May 25, 2011\n",
      "\n",
      "SHOOK ONES\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "First photo Karoliina Niemenkari, second Sinikka Konttinen.\n",
      "\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch:[2/3]:  63%|██████▎   | 19292/30704 [10:12<06:04, 31.29it/s, Loss=0.0342, Lr=3.00e-04, Step=5e+4] "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Checkpoint saved at step: 50000\n",
      "========== Sample Text Generation ==========\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch:[2/3]:  63%|██████▎   | 19300/30704 [10:13<12:12, 15.56it/s, Loss=0.0206, Lr=3.00e-04, Step=5e+4]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generated: Wednesday, May 25, 2011\n",
      "\n",
      "SHINY SUIT THEORY\n",
      "\n",
      "\n",
      "\n",
      "Vintage Wayfarers, Laitinen jacket, print shirt 2nd hand\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch:[2/3]:  95%|█████████▌| 29292/30704 [15:29<00:44, 31.93it/s, Loss=0.0414, Lr=2.73e-04, Step=6e+4] "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Checkpoint saved at step: 60000\n",
      "========== Sample Text Generation ==========\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch:[2/3]:  95%|█████████▌| 29300/30704 [15:30<01:24, 16.56it/s, Loss=0.0289, Lr=2.73e-04, Step=6e+4]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generated: Wednesday, May 25, 2011\n",
      "\n",
      "SHINY SUIT THEORSYNTA[The sunglasses already by Martin Margiela, vice dean of then in case of\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch:[2/3]: 100%|██████████| 30704/30704 [16:14<00:00, 31.51it/s, Loss=0.0428, Lr=2.62e-04, Step=61408]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch:[2/3] Completed; Avg_Loss: 0.0408\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch:[3/3]:  28%|██▊       | 8588/30704 [04:31<11:38, 31.65it/s, Loss=0.0376, Lr=1.99e-04, Step=7e+4] "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Checkpoint saved at step: 70000\n",
      "========== Sample Text Generation ==========\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch:[3/3]:  28%|██▊       | 8596/30704 [04:32<22:26, 16.42it/s, Loss=0.0306, Lr=1.99e-04, Step=7e+4]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generated: Wednesday, May 25, 2011\n",
      "\n",
      "SHINY SUIT THEORY\n",
      "\n",
      "\n",
      "\n",
      "Vintage Wayfarers, Laitinen jacket, print shirt 2nd hand\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch:[3/3]:  61%|██████    | 18588/30704 [09:48<06:23, 31.56it/s, Loss=0.0261, Lr=1.26e-04, Step=8e+4] "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Checkpoint saved at step: 80000\n",
      "========== Sample Text Generation ==========\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch:[3/3]:  61%|██████    | 18596/30704 [09:48<12:21, 16.34it/s, Loss=0.0172, Lr=1.26e-04, Step=8e+4]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generated: Wednesday, May 25, 2011\n",
      "\n",
      "NEVER GROW UP\n",
      "\n",
      "\n",
      "\n",
      "These are two persons I am always happy to meet, not only because they have amazing style.\n",
      "\n",
      "The\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch:[3/3]:  93%|█████████▎| 28588/30704 [15:04<01:06, 31.80it/s, Loss=0.0187, Lr=5.28e-05, Step=9e+4] "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Checkpoint saved at step: 90000\n",
      "========== Sample Text Generation ==========\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch:[3/3]:  93%|█████████▎| 28596/30704 [15:05<02:07, 16.51it/s, Loss=0.0269, Lr=5.28e-05, Step=9e+4]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generated: Wednesday, May 25, 2011\n",
      "\n",
      "NEVER GROW UP\n",
      "\n",
      "\n",
      "\n",
      "These are two persons I am always happy to meet, not only because they have amazing style.\n",
      "\n",
      "The\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch:[3/3]: 100%|██████████| 30704/30704 [16:11<00:00, 31.60it/s, Loss=0.0209, Lr=3.73e-05, Step=92112]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch:[3/3] Completed; Avg_Loss: 0.0230\n",
      "============================== Training Completed ==============================\n"
     ]
    }
   ],
   "source": [
    "model.train()\n",
    "step = 0\n",
    "print(\"=\"*30, \"Starting Training\", \"=\"*30)\n",
    "for epoch in range(Config.epochs):\n",
    "    \n",
    "    epoch_loss = 0 \n",
    "    loop = tqdm(train_loader, desc=f\"Epoch:[{epoch+1}/{Config.epochs}]\")\n",
    "    \n",
    "    for batch_idx, (x, y) in enumerate(loop):\n",
    "        x, y = x.to(Config.device), y.to(Config.device)\n",
    "        \n",
    "        x = torch.clamp(x, 0, vocab_size - 1)\n",
    "        y = torch.clamp(y, 0, vocab_size - 1)\n",
    "        \n",
    "        optimizer.zero_grad()\n",
    "        \n",
    "        if Config.use_amp:\n",
    "            #Mixed Precision Forward Pass\n",
    "            with torch.amp.autocast('cuda'):\n",
    "                logits = model(x)\n",
    "                \n",
    "                loss = F.cross_entropy(logits.view(-1, vocab_size), y.view(-1))\n",
    "            \n",
    "            #Mixed Precision Backward Pass\n",
    "            scaler.scale(loss).backward()\n",
    "            # print(f\"Gradient of first parameter is: {next(model.parameters()).grad}\")\n",
    "            scaler.unscale_(optimizer)\n",
    "            nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n",
    "            scaler.step(optimizer)\n",
    "            scaler.update()\n",
    "            scheduler.step()\n",
    "        \n",
    "        else:\n",
    "            #Standar Precision Forward Pass\n",
    "            logits = model(x)\n",
    "            loss = F.cross_entropy(logits.view(-1, vocab_size), y.view(-1))\n",
    "            \n",
    "            #Standard Precision Backward Pass\n",
    "            loss.backward()\n",
    "            nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n",
    "            optimizer.step()\n",
    "            scheduler.step()\n",
    "            \n",
    "        #Update Metrics\n",
    "        epoch_loss += loss.item()\n",
    "        step += 1\n",
    "\n",
    "        #WandB Loggin\n",
    "        wandb.log(\n",
    "            {\n",
    "                \"step_loss\":loss.item(),\n",
    "                \"learning_rate\":scheduler.get_last_lr()[0],\n",
    "            }\n",
    "            , step = step)\n",
    "        \n",
    "        #Add tqdm details\n",
    "        loop.set_postfix(\n",
    "            {\n",
    "                'Loss':f'{loss.item():.4f}',\n",
    "                'Lr': f'{scheduler.get_last_lr()[0]:.2e}',\n",
    "                'Step': step\n",
    "            }\n",
    "        )\n",
    "        \n",
    "        #Save checkpoint & Generate Sample\n",
    "        if step % Config.save_every == 0:\n",
    "            checkpoint = {\n",
    "                'model_state_dict':model.state_dict(),\n",
    "                'optimizer_state_dict': optimizer.state_dict(),\n",
    "                'scheduler_state_dict': scheduler.state_dict(),\n",
    "                'scaler_state_dict': scaler.state_dict() if scaler else None,\n",
    "                'step': step,\n",
    "            }\n",
    "            os.makedirs(\"./saves/checkpoint\", exist_ok=True)\n",
    "            torch.save(checkpoint, f\"./saves/checkpoint/num_steps_{step}.pt\")\n",
    "            print(f\"Checkpoint saved at step: {step}\")\n",
    "            \n",
    "            #Generate Sample Text\n",
    "            print(\"=\"*10, \"Sample Text Generation\", \"=\"*10)\n",
    "            sample_tokens = valid_tokens[:10] \n",
    "            try:\n",
    "                generated_text = model.generate(tokenizer, sample_tokens, max_new_tokens = 30)\n",
    "                print(f\"Generated: {generated_text}\")\n",
    "                prompt_text = tokenizer.decode(sample_tokens)\n",
    "                table.add_data(step, prompt_text, generated_text)\n",
    "                wandb.log({\"text_samples\":table}, step = step)\n",
    "            except Exception as e:\n",
    "                print(f\"Generation Failed: {e}\")\n",
    "    avg_loss = epoch_loss / len(train_loader)\n",
    "    print(f\"\\nEpoch:[{epoch+1}/{Config.epochs}] Completed; Avg_Loss: {avg_loss:.4f}\")\n",
    "    wandb.log({\"Avg_Epoch_Loss\":avg_loss}, step = step)\n",
    "    \n",
    "    \n",
    "    \n",
    "#Save Final Path\n",
    "torch.save({\n",
    "    'model_state_dict':model.state_dict(),\n",
    "    'vocab_size':vocab_size,\n",
    "    'config':Config\n",
    "}, \"./saves/checkpoint/final_model.pt\")\n",
    "\n",
    "print(\"=\"*30, \"Training Completed\", \"=\"*30)\n",
    "\n",
    "            \n",
    "                \n",
    "        \n",
    "            \n",
    "            "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "56a868ca",
   "metadata": {},
   "source": [
    "## Model Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "7b17287e",
   "metadata": {
    "trusted": true
   },
   "outputs": [],
   "source": [
    "import regex\n",
    "import glob\n",
    "import random\n",
    "import colorsys\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from IPython.display import HTML, display"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96c9fa78",
   "metadata": {},
   "source": [
    "- Since our model was wrapped with nn.parallel.DistributedDataParallel hence in our state_dictionary every key starts with \"_orig_mod.\" so to load the state_dictionary for Inference we need to create a new dictionary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "24fd90be",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading model from ./saves/checkpoint/final_model.pt for Inference\n",
      "========== Model loaded successfully for Inference ==========\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\athar\\AppData\\Local\\Temp\\ipykernel_17856\\342875502.py:2: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
      "  checkpoint = torch.load(checkpoint_path, map_location=\"cpu\")\n"
     ]
    }
   ],
   "source": [
    "checkpoint_path = \"./saves/checkpoint/final_model.pt\"\n",
    "checkpoint = torch.load(checkpoint_path, map_location=\"cpu\")\n",
    "original_state_dict = checkpoint['model_state_dict']\n",
    "new_state_dict = {}\n",
    "\n",
    "for key,val in original_state_dict.items():\n",
    "    if key.startswith(\"_orig_mod.\"):\n",
    "        new_key = key.replace(\"_orig_mod.\", \"\")\n",
    "        new_state_dict[new_key] = val\n",
    "    else:\n",
    "        new_state_dict[key] = val \n",
    "    \n",
    "print(f\"Loading model from {checkpoint_path} for Inference\")\n",
    "model.load_state_dict(new_state_dict)\n",
    "\n",
    "model.eval()\n",
    "print(\"=\"*10,\"Model loaded successfully for Inference\", \"=\"*10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c62716af",
   "metadata": {
    "trusted": true
   },
   "outputs": [],
   "source": [
    "# checkpoint_path = \"./saves/checkpoint/num_steps_90000.pt\"\n",
    "# print(f\"Loading model from {checkpoint_path} for Inference\")\n",
    "# checkpoint = torch.load(checkpoint_path, map_location=\"cpu\")\n",
    "# model.load_state_dict(checkpoint['model_state_dict'])\n",
    "\n",
    "# #Set the model to eval\n",
    "# model.eval()\n",
    "# print(\"=\"*10,\"Model loaded successfully for Inference\", \"=\"*10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66d43771",
   "metadata": {
    "trusted": true
   },
   "outputs": [],
   "source": [
    "def find_text_in_tokens(target_txt, sample_tokens, tokenizer, max_attempts = 500):\n",
    "    ''' \n",
    "        Function which will try to find a sequence of tokens that decode somewhat similar to target\n",
    "    '''\n",
    "    best_match = []\n",
    "    best_similarity = 0\n",
    "    \n",
    "    for _ in range(max_attempts):\n",
    "        seq_len = random.randint(10,50)\n",
    "        start_idx = random.randint(0, len(sample_tokens) - seq_len)\n",
    "        candidate = sample_tokens[start_idx : start_idx + seq_len]\n",
    "        \n",
    "        try:\n",
    "            decoded = tokenizer.decode(candidate).lower()\n",
    "            #Do similarity_check \n",
    "            keywords = target_txt.lower().split()\n",
    "            matches = sum(1 for keyword in keywords if keyword in decoded)\n",
    "            \n",
    "            similarity = matches / len(keywords)\n",
    "            \n",
    "            if similarity > best_similarity:\n",
    "                best_similarity = similarity\n",
    "                best_match = candidate\n",
    "            \n",
    "            if similarity > 0.5:\n",
    "                return candidate\n",
    "        except:\n",
    "            continue\n",
    "        \n",
    "    return best_match if best_match else sample_tokens[:20]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "0cb9bd3f",
   "metadata": {
    "trusted": true
   },
   "outputs": [],
   "source": [
    "def create_meaningful_prompts(tokenizer, sample_tokens):\n",
    "    prompts = []\n",
    "    \n",
    "    prompt_targets = [\n",
    "        \"Future of Artificial Intelligence\",\n",
    "        \"Import numpy as np\",\n",
    "        \"def calculate\",\n",
    "        \"In the beginning\",\n",
    "        \"Machine Learning\",\n",
    "        \"Deep Learning\",\n",
    "        \"Breaking news\",\n",
    "        \"The weather today\",\n",
    "        \"Scientists have discovered\"\n",
    "    ]\n",
    "    \n",
    "    print(\"=\"*10, \"Creating meaningful prompts by searching token sequences....\", \"=\"*10)\n",
    "    \n",
    "    for target in prompt_targets:\n",
    "        #Find tokens that decode to something similar to target\n",
    "        tokens = find_text_in_tokens(target, sample_tokens, tokenizer)\n",
    "        \n",
    "        tokens = tokens[:30]\n",
    "        \n",
    "        try:\n",
    "            decoded = tokenizer.decode(tokens)\n",
    "            prompts.append({\n",
    "                'target':target,\n",
    "                'tokens':tokens,\n",
    "                'decoded':decoded\n",
    "            })\n",
    "            print(\"-\"*10,f\"Found Prompt Aprroximating '{target}':{decoded[:50]}....\", \"-\"*10)\n",
    "        except:\n",
    "            print(\"-\"*10,f\"Failed to created prompt for '{target}'\",\"-\"*10)\n",
    "            \n",
    "    return prompts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "6bbe9b5f",
   "metadata": {
    "trusted": true
   },
   "outputs": [],
   "source": [
    "def token2color(token_id, vocab_size):\n",
    "    ''' \n",
    "        Create Unique color for each token\n",
    "    '''\n",
    "    hue = (token_id / vocab_size) * 0.8\n",
    "    saturation = 0.3 + (token_id % 100) / 200\n",
    "    value = 0.9\n",
    "    \n",
    "    rgb = colorsys.hsv_to_rgb(hue, saturation, value)\n",
    "    return f\"rgb({int(rgb[0]*255)}, {int(rgb[1]*255)}, {int(rgb[2]*255)})\"\n",
    "\n",
    "\n",
    "def display_colored_tokens(tokens, tokenizer, title = \"\"):\n",
    "    ''' \n",
    "        Display tokens with colored background\n",
    "    '''\n",
    "    vocab_size = len(tokenizer.mergeable_ranks)\n",
    "    html = f\"<h3>{title}</h3><div style = 'font-family:monospace; font-size:14px;line-height: 1.8;'>\"\n",
    "    \n",
    "    for token_id in tokens:\n",
    "        token_id = min(max(0, token_id), vocab_size - 1)\n",
    "        \n",
    "        #Decode a single token\n",
    "        try:\n",
    "            token_text = tokenizer.decode([token_id])\n",
    "        except:\n",
    "            token_text = f\"[{token_id}]\"\n",
    "            \n",
    "        token_text = token_text.replace('&', '&amp;').replace('<', '&lt;').replace('>', '&gt;')\n",
    "        token_text = token_text.replace(' ', '&nbsp;').replace('\\n', '<br>')\n",
    "        \n",
    "        #Manage  empty or whitespace tokens\n",
    "        if not token_text or token_text.isspace():\n",
    "            token_text = '.'\n",
    "        \n",
    "        #Generate color\n",
    "        bg_color = token2color(token_id, vocab_size)\n",
    "        \n",
    "        html += f\"<span style='background-color: {bg_color}; padding: 2px 4px; margin: 1px; border-radius: 3px; cursor: help;' title='Token ID: {token_id}'>{token_text}</span>\"\n",
    "    html += \"<div>\"\n",
    "    display(HTML(html))\n",
    "    \n",
    "    \n",
    "\n",
    "def visualize_attention_pattern(seq_len, current_pos, title= \"Attention Pattern\"):\n",
    "    #Create Causal Mask\n",
    "    mask = np.zeros((seq_len, seq_len))\n",
    "    for i in range(seq_len):\n",
    "        for j in range(i + 1):\n",
    "            mask[i,j] = 1\n",
    "    \n",
    "    #Highlight current pos \n",
    "    if current_pos < seq_len:\n",
    "        mask[current_pos, :] = mask[current_pos, :]*2\n",
    "        \n",
    "    plt.figure(figsize=(8,6))\n",
    "    plt.imshow(mask, cmap='Blues', aspect='auto')\n",
    "    plt.colorbar(label = \"Attention Weight\")\n",
    "    plt.xlabel('Key Position')\n",
    "    plt.ylabel('Query Position')\n",
    "    plt.title(title)\n",
    "    \n",
    "    plt.grid(True, alpha = 0.3, linestyle = '--')\n",
    "    \n",
    "    if current_pos < seq_len:\n",
    "        plt.axhline(y = current_pos, color = 'red', linestyle = '--', alpha = 0.5)\n",
    "        plt.text(seq_len * 0.7, current_pos + 0.5 , f'Current: {current_pos}', color='red')\n",
    "        \n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "def generate_with_attention_visualization(model, tokenizer, prompt_tokens, max_new_tokens = 20, temperature = 0.8):\n",
    "    model.eval()\n",
    "    device = Config.device\n",
    "    vocab_size = len(tokenizer.mergeable_ranks)\n",
    "    \n",
    "    tokens = prompt_tokens.copy()\n",
    "    \n",
    "    print(\"=\"*10,f\"Auto-Regressive Generation with ATTENTION VISUALIZATION\", \"=\"*10)\n",
    "    print(\"=\"*80)\n",
    "    print(\"Each position can onnly attend to previous positions. \\n\")\n",
    "    \n",
    "    display_colored_tokens(tokens, tokenizer, \"Initial_Prompt\")\n",
    "    visualize_attention_pattern(len(tokens)+5, len(tokens) - 1, \"Initial Attention Pattern\")\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        for step in range(min(max_new_tokens, 10)):\n",
    "            print(f\"\\n{'='*60}\")\n",
    "            print(f\"Generation Step: {step + 1}\")\n",
    "            print(f\"{'='*60}\")\n",
    "            \n",
    "            #Get last max_len tokens\n",
    "            input_tokens = tokens[-model.max_len:]\n",
    "            x = torch.tensor(input_tokens).unsqueeze(0).to(device)\n",
    "            x = torch.clamp(x, 0, vocab_size-1)\n",
    "            \n",
    "            #Show what the Model Sees\n",
    "            print(f\"Model Input:Last {len(input_tokens)} tokens\")\n",
    "            print(f\"Predicting token at position: {len(tokens)}\")\n",
    "            \n",
    "            #Visualize attention_pattern for current step\n",
    "            visualize_attention_pattern(\n",
    "                seq_len= min(len(tokens)+1, 20),\n",
    "                current_pos=len(tokens) if len(tokens) < 20 else 19,\n",
    "                title = f\"Step:{step + 1}-->Causal Attention Pattern\"\n",
    "            )\n",
    "            \n",
    "            #Do a Forward Pass\n",
    "            with torch.amp.autocast('cuda' if device == 'cuda' else 'cpu'):\n",
    "                logits = model(x)\n",
    "            logits = logits[0, -1, :] / temperature\n",
    "\n",
    "            # Get top predictions\n",
    "            probs = F.softmax(logits, dim=-1)\n",
    "            top_probs, top_indices = torch.topk(probs, 5)\n",
    "\n",
    "            print(\"\\n🎲 Top 5 predictions:\")\n",
    "            for i, (prob, idx) in enumerate(zip(top_probs, top_indices)):\n",
    "                try:\n",
    "                    token_text = tokenizer.decode([idx.item()])\n",
    "                    print(f\"  {i+1}. Token {idx.item()}: '{token_text}' (prob: {prob.item():.3f})\")\n",
    "                except:\n",
    "                    print(f\"  {i+1}. Token {idx.item()} (prob: {prob.item():.3f})\")\n",
    "\n",
    "            # Sample next token\n",
    "            next_token = torch.multinomial(probs, 1).item()\n",
    "            next_token = min(max(next_token, 0), vocab_size - 1)\n",
    "            tokens.append(next_token)\n",
    "\n",
    "            print(f\"\\n✅ Selected: Token {next_token}\")\n",
    "\n",
    "            # Show updated sequence\n",
    "            display_colored_tokens(tokens, tokenizer, f\"After Step {step + 1}\")\n",
    "\n",
    "            # Show the autoregressive nature\n",
    "            print(f\"\\n📝 Explanation: Token at position {len(tokens)-1} was predicted\")\n",
    "            print(f\"   using all tokens from positions 0 to {len(tokens)-2}\")\n",
    "\n",
    "    return tokens\n",
    "    \n",
    "\n",
    "def generate_multiple_samples(model, tokenizer, prompts):\n",
    "    \"\"\"Generate samples from meaningful prompts\"\"\"\n",
    "    print(\"🎲 Generating samples from meaningful prompts...\\n\")\n",
    "\n",
    "    for i, prompt_data in enumerate(prompts[:5]):  # Limit to 5 samples\n",
    "        print(f\"\\n{'='*80}\")\n",
    "        print(f\"🎯 Sample {i+1}/5\")\n",
    "        print(f\"{'='*80}\")\n",
    "\n",
    "        print(f\"📝 Target prompt: '{prompt_data['target']}'\")\n",
    "        print(f\"🔤 Actual prompt: '{prompt_data['decoded'][:100]}...'\")\n",
    "        print(f\"📊 Prompt length: {len(prompt_data['tokens'])} tokens\")\n",
    "\n",
    "        # Display prompt\n",
    "        display_colored_tokens(prompt_data['tokens'], tokenizer, \"Prompt Tokens\")\n",
    "\n",
    "        # Generate with temperature variation\n",
    "        temperature = random.choice([0.7, 0.8, 0.9])\n",
    "        print(f\"🌡️ Temperature: {temperature}\")\n",
    "\n",
    "        try:\n",
    "            # Generate\n",
    "            model.eval()\n",
    "            vocab_size = len(tokenizer.mergeable_ranks)\n",
    "            device = next(model.parameters()).device\n",
    "            generated_tokens = prompt_data['tokens'].copy()\n",
    "\n",
    "            with torch.no_grad():\n",
    "                for _ in range(50):  # Generate 50 new tokens\n",
    "                    input_tokens = generated_tokens[-model.max_len:]\n",
    "                    x = torch.tensor(input_tokens).unsqueeze(0).to(device)\n",
    "                    x = torch.clamp(x, 0, vocab_size - 1)\n",
    "\n",
    "                    with torch.amp.autocast('cuda' if device.type == 'cuda' else 'cpu'):\n",
    "                        logits = model(x)\n",
    "                    logits = logits[0, -1, :] / temperature\n",
    "\n",
    "                    probs = F.softmax(logits, dim=-1)\n",
    "                    next_token = torch.multinomial(probs, 1).item()\n",
    "                    next_token = min(max(next_token, 0), vocab_size - 1)\n",
    "                    generated_tokens.append(next_token)\n",
    "\n",
    "            # Display generated tokens\n",
    "            new_tokens = generated_tokens[len(prompt_data['tokens']):]\n",
    "            display_colored_tokens(new_tokens, tokenizer, \"Generated Tokens\")\n",
    "\n",
    "            # Show full text\n",
    "            print(f\"\\n📄 Full generated text:\")\n",
    "            try:\n",
    "                full_text = tokenizer.decode(generated_tokens)\n",
    "                print(full_text)\n",
    "            except:\n",
    "                print(\"[Unable to decode]\")\n",
    "\n",
    "        except Exception as e:\n",
    "            print(f\"❌ Generation failed: {e}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c213475",
   "metadata": {
    "trusted": true
   },
   "outputs": [],
   "source": [
    "token_path = \"./saves/tokenizer/essential_web_500k_tokens.pkl\"\n",
    "with open(token_path, 'rb') as f:\n",
    "    sample_tokens = pickle.load(f)\n",
    "sample_tokens = [t for t in sample_tokens[:50000] if 0 <= t < vocab_size]\n",
    "\n",
    "prompts = create_meaningful_prompts(tokenizer, sample_tokens)\n",
    "#Generate Multiple Samples\n",
    "print(\"=\"*30,f\"Multiple Sample Generation\",\"=\"*30)\n",
    "generate_multiple_samples(model, tokenizer, prompts)\n",
    "\n",
    "#Show Generation with Attention Visualization\n",
    "print(\"=\"*30,f\"Sample Generation with Attention Pattern Visualization\",\"=\"*30)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14d856f7",
   "metadata": {},
   "source": [
    "### HuggingfaceHub Deployment Code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "3eff6856",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b156e3f4be80474e902d2896d78f333f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "final_model.pt:   0%|          | 0.00/72.4M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8ae1110f7fcd45be98cf585e2883c7b6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "bpe_tokenizer_16k_n1000000.pkl:   0%|          | 0.00/194k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "CommitInfo(commit_url='https://huggingface.co/Atharva31/smallm-from-scratch/commit/be46848a635ebf233d7d79fd3903c6bac33e99d9', commit_message='Upload model_code.py with huggingface_hub', commit_description='', oid='be46848a635ebf233d7d79fd3903c6bac33e99d9', pr_url=None, repo_url=RepoUrl('https://huggingface.co/Atharva31/smallm-from-scratch', endpoint='https://huggingface.co', repo_type='model', repo_id='Atharva31/smallm-from-scratch'), pr_revision=None, pr_num=None)"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from huggingface_hub import HfApi\n",
    "import torch \n",
    "\n",
    "api = HfApi()\n",
    "repo_id = \"Atharva31/smallm-from-scratch\"\n",
    "api.create_repo(repo_id = repo_id, exist_ok = True)\n",
    "\n",
    "api.upload_file(\n",
    "    path_or_fileobj=\"./saves/checkpoint/final_model.pt\",\n",
    "    path_in_repo = \"pytorch_model.bin\",\n",
    "    repo_id = repo_id,\n",
    ")\n",
    "\n",
    "api.upload_file(\n",
    "    path_or_fileobj = \"./saves/tokenizer/bpe_tokenizer_16k_n1000000.pkl\",\n",
    "    path_in_repo = \"tokenizer.pkl\",\n",
    "    repo_id = repo_id\n",
    ")\n",
    "\n",
    "api.upload_file(\n",
    "    path_or_fileobj=\"small-lm_trained.ipynb\",\n",
    "    path_in_repo=\"model_code.py\",\n",
    "    repo_id = repo_id\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e681aa3",
   "metadata": {},
   "outputs": [],
   "source": [
    "api.upload_file(\n",
    "    path_or_fileobj=\"small-lm_trained.ipynb\",\n",
    "    path_in_repo=\"model_code.ipynb\",\n",
    "    repo_id = repo_id\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f828d7c",
   "metadata": {
    "trusted": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0c3f7c7",
   "metadata": {
    "trusted": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kaggle": {
   "accelerator": "nvidiaTeslaT4",
   "dataSources": [],
   "dockerImageVersionId": 31090,
   "isGpuEnabled": true,
   "isInternetEnabled": true,
   "language": "python",
   "sourceType": "notebook"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}