Forward_pass_from_scratch.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2b9bcf4359344e76b618e6eeedc2f4f7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Dropdown(description='Select Model:', options=('gpt2-small', 'gpt2-medium', 'Qwen 2-0.5B', 'Qwen 2-0.5B-Instru…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import torch\n",
    "import ipywidgets as widgets\n",
    "import transformer_lens\n",
    "import activation_analysis\n",
    "import model_loader\n",
    "import forward_pass\n",
    "from forward_pass import list_activation_tensors, list_weight_matrices\n",
    "from activation_analysis import activation_agg_sim, compare_activation_similarity, display_activation_similarity_tables, display_activation_similarity_plots\n",
    "from model_loader import load_model, print_available_models, gpu_mem_check,list_available_models\n",
    "import math\n",
    "import re\n",
    "import torch.nn.functional as F\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "# SELECT MODEL #\n",
    "model_dropdown = widgets.Dropdown(\n",
    "    options=list_available_models(),\n",
    "    value=list_available_models()[0],  # Set default to first model\n",
    "    description='Select Model:',\n",
    "    disabled=False,\n",
    ")\n",
    "display(model_dropdown)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading model gpt2-medium from gpt2-medium...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "c:\\Users\\akozlo\\AppData\\Local\\miniconda3\\envs\\akpy24\\Lib\\site-packages\\transformers\\tokenization_utils_base.py:1601: FutureWarning: `clean_up_tokenization_spaces` was not set. It will be set to `True` by default. This behavior will be depracted in transformers v4.45, and will be then set to `False` by default. For more details check this issue: https://github.com/huggingface/transformers/issues/31884\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded pretrained model gpt2-medium into HookedTransformer\n",
      "Model loaded successfully on cuda.\n",
      "Total GPU Memory: 11.99 GB\n",
      "Allocated GPU Memory: 1.54 GB\n",
      "Cached GPU Memory: 1.55 GB\n"
     ]
    }
   ],
   "source": [
    "# Load Model #\n",
    "model_name = model_dropdown.value\n",
    "model = load_model(model_name)\n",
    "gpu_mem_check()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Available activation tensors:\n",
      "hook_embed                       torch.Size([3, 1024])\n",
      "hook_pos_embed                   torch.Size([3, 1024])\n",
      "blocks.n.hook_resid_pre          torch.Size([3, 1024])\n",
      "blocks.n.ln1.hook_scale          torch.Size([3, 1])\n",
      "blocks.n.ln1.hook_normalized     torch.Size([3, 1024])\n",
      "blocks.n.attn.hook_q             torch.Size([3, 16, 64])\n",
      "blocks.n.attn.hook_k             torch.Size([3, 16, 64])\n",
      "blocks.n.attn.hook_v             torch.Size([3, 16, 64])\n",
      "blocks.n.attn.hook_attn_scores   torch.Size([16, 3, 3])\n",
      "blocks.n.attn.hook_pattern       torch.Size([16, 3, 3])\n",
      "blocks.n.attn.hook_z             torch.Size([3, 16, 64])\n",
      "blocks.n.hook_attn_out           torch.Size([3, 1024])\n",
      "blocks.n.hook_resid_mid          torch.Size([3, 1024])\n",
      "blocks.n.ln2.hook_scale          torch.Size([3, 1])\n",
      "blocks.n.ln2.hook_normalized     torch.Size([3, 1024])\n",
      "blocks.n.mlp.hook_pre            torch.Size([3, 4096])\n",
      "blocks.n.mlp.hook_post           torch.Size([3, 4096])\n",
      "blocks.n.hook_mlp_out            torch.Size([3, 1024])\n",
      "blocks.n.hook_resid_post         torch.Size([3, 1024])\n",
      "ln_final.hook_scale              torch.Size([3, 1])\n",
      "ln_final.hook_normalized         torch.Size([3, 1024])\n",
      "\n",
      "Call with cache['{tensor_name}'][0]\n",
      "\n",
      "Example tensor call:\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([  4.0842,   4.6890,   0.9476,  ...,  -1.1796, -10.8282,   3.6351],\n",
       "       device='cuda:0')"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "prompt = \"Hello world\"\n",
    "list_activation_tensors(model, prompt)\n",
    "\n",
    "#example tensor call\n",
    "logits, cache = model.run_with_cache(prompt) \n",
    "print(\"\\nExample tensor call:\")\n",
    "cache['blocks.23.hook_resid_post'][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model state_dict keys:\n",
      "embed.W_E               torch.Size([50257, 1024])\n",
      "pos_embed.W_pos         torch.Size([1024, 1024])\n",
      "blocks[n].ln1.w         torch.Size([1024])\n",
      "blocks[n].ln1.b         torch.Size([1024])\n",
      "blocks[n].ln2.w         torch.Size([1024])\n",
      "blocks[n].ln2.b         torch.Size([1024])\n",
      "blocks[n].attn.W_Q      torch.Size([16, 1024, 64])\n",
      "blocks[n].attn.W_O      torch.Size([16, 64, 1024])\n",
      "blocks[n].attn.b_Q      torch.Size([16, 64])\n",
      "blocks[n].attn.b_O      torch.Size([1024])\n",
      "blocks[n].attn.W_K      torch.Size([16, 1024, 64])\n",
      "blocks[n].attn.W_V      torch.Size([16, 1024, 64])\n",
      "blocks[n].attn.b_K      torch.Size([16, 64])\n",
      "blocks[n].attn.b_V      torch.Size([16, 64])\n",
      "blocks[n].attn.mask     torch.Size([1024, 1024])\n",
      "blocks[n].attn.IGNORE   torch.Size([])\n",
      "blocks[n].mlp.W_in      torch.Size([1024, 4096])\n",
      "blocks[n].mlp.b_in      torch.Size([4096])\n",
      "blocks[n].mlp.W_out     torch.Size([4096, 1024])\n",
      "blocks[n].mlp.b_out     torch.Size([1024])\n",
      "ln_final.w              torch.Size([1024])\n",
      "ln_final.b              torch.Size([1024])\n",
      "unembed.W_U             torch.Size([1024, 50257])\n",
      "unembed.b_U             torch.Size([50257])\n",
      "\n",
      "Call with model.{weight_name}\n",
      "\n",
      "Example tensor call:\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Parameter containing:\n",
       "tensor([[ 0.3126, -0.1634,  0.1243,  ...,  0.0084,  0.0499,  0.2266],\n",
       "        [ 0.2040, -0.0206,  0.0207,  ...,  0.1752,  0.2064, -0.0457],\n",
       "        [-0.0872,  0.0126,  0.2485,  ...,  0.1453, -0.1659, -0.1803],\n",
       "        ...,\n",
       "        [ 0.0513, -0.0420, -0.2821,  ...,  0.3503,  0.2665, -0.0642],\n",
       "        [ 0.1500,  0.0376,  0.0267,  ..., -0.0566, -0.1505,  0.1074],\n",
       "        [-0.0872, -0.0176,  0.0117,  ...,  0.4246,  0.0556,  0.0542]],\n",
       "       device='cuda:0', requires_grad=True)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "list_weight_matrices(model)\n",
    "\n",
    "#example tensor call\n",
    "print(\"\\nExample tensor call:\")\n",
    "model.blocks[23].mlp.W_out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0.2016], device='cuda:0')"
      ]
     },
     "execution_count": 78,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Example of testing cosine similarity between two activation vectors\n",
    "act1 = cache['blocks.5.hook_resid_post'][0]\n",
    "act2 = cache['blocks.23.hook_resid_post'][0]\n",
    "final_row_act1 = act1[-1, :]\n",
    "final_row_act2 = act2[-1, :]\n",
    "\n",
    "# Calculate cosine similarity between final_row_act1 and final_row_act2\n",
    "F.cosine_similarity(final_row_act1.unsqueeze(0), final_row_act2.unsqueeze(0))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 1.) We begin by specifying our prompt and retrieving the relevant tokens, subwords, and embeddings\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TokenID\tSubword\n",
      "50256\t<|endoftext|>\n",
      "818\tIn\n",
      "262\t the\n",
      "3726\t beginning\n",
      "11\t,\n",
      "1793\t God\n",
      "2727\t created\n",
      "262\t the\n"
     ]
    }
   ],
   "source": [
    "prompt = \"In the beginning, God created the\"\n",
    "\n",
    "# Run the prompt through the hooked transformer model #\n",
    "logits, cache = model.run_with_cache(prompt)          # This stores the activations in a dictionary called cache\n",
    "#######################################################\n",
    "\n",
    "tokens = model.to_tokens(prompt)            #Extract token ids from input\n",
    "subwords = model.to_str_tokens(tokens)      #Convert token ids to subwords\n",
    "\n",
    "print(\"TokenID\\tSubword\")\n",
    "for token, subword in zip(tokens.squeeze().tolist(), subwords):\n",
    "    print(f\"{token}\\t{subword}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Shape of prompt word vectors: torch.Size([8, 1024])\n",
      "\n",
      "Is input_vectors the same as hook_embed? True\n"
     ]
    }
   ],
   "source": [
    "# We pull the word vectors from the embedding matrix by token id\n",
    "# We then confirm that this is the same as the hook_embed tensor\n",
    "\n",
    "input_vectors = model.W_E[tokens.squeeze()]\n",
    "print(f\"Shape of prompt word vectors: {input_vectors.shape}\")\n",
    "\n",
    "# Confirm that word_vectors is identical to cache['hook_embed']\n",
    "print(\"\\nIs input_vectors the same as hook_embed?\", torch.allclose(input_vectors, cache['hook_embed'][0], atol=1e-6))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 2.) Extract the first *n* position embeddings and **add** them to the respective token embeddings, producing *resid_pre*\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Is input_vectors_with_pos the same as blocks.0.hook_resid_pre? True\n"
     ]
    }
   ],
   "source": [
    "# Positional embedding is model.W_pos -- model.W_E_pos is the W_E matrix concatenated with the W_pos matrix. idk its purpose.\n",
    "# hook_pos_embed identical to W_pos. Again, idk why.\n",
    "\n",
    "position_embeddings = model.W_pos[:len(tokens[0])]                   # Extract n vectors from the position embedding matrix when n = prompt length\n",
    "input_vectors_with_pos = input_vectors + position_embeddings           # Add position embeddings to token embeddings elementwise\n",
    "\n",
    "# Confirm that word_vectors_with_pos is identical to cache['blocks.0.hook_resid_pre']\n",
    "print(\"\\nIs input_vectors_with_pos the same as blocks.0.hook_resid_pre?\", torch.allclose(input_vectors_with_pos, cache['blocks.0.hook_resid_pre'][0], atol=1e-5))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 3.) Layernorm 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Is std the same as blocks.0.ln1.hook_scale? True\n",
      "\n",
      "Is ln1_output the same as blocks.0.ln1.hook_normalized? True\n"
     ]
    }
   ],
   "source": [
    "# 1. Normalize\n",
    "mean = input_vectors_with_pos.mean(dim=-1, keepdim=True)\n",
    "std = input_vectors_with_pos.std(dim=-1, keepdim=True)\n",
    "\n",
    "# Confirm that std is identical to cache['blocks.0.ln1.hook_scale'] (I can't find ln1.hook_bias anywhere)\n",
    "print(\"\\nIs std the same as blocks.0.ln1.hook_scale?\", torch.allclose(std, cache['blocks.0.ln1.hook_scale'][0], atol=1e-4))\n",
    "\n",
    "normalized_resid_pre = (input_vectors_with_pos - mean) / cache['blocks.0.ln1.hook_scale'][0]\n",
    "\n",
    "# 2. Apply weight (gamma) and bias (beta) (Even after normalizing, we still need to apply the learned layernorm weights)\n",
    "ln1_output = normalized_resid_pre * model.blocks[0].ln1.w + model.blocks[0].ln1.b\n",
    "\n",
    "# Compare with the output of the first layer norm in the model\n",
    "print(\"\\nIs ln1_output the same as blocks.0.ln1.hook_normalized?\", torch.allclose(ln1_output, cache['blocks.0.ln1.hook_normalized'][0], atol=1e-4))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Is Q_heads_aligned the same as blocks.0.attn.hook_q? True\n",
      "\n",
      "Is K_heads_aligned the same as blocks.0.attn.hook_k? True\n",
      "\n",
      "Is V_heads_aligned the same as blocks.0.attn.hook_v? True\n"
     ]
    }
   ],
   "source": [
    "# Initialize an empty list to store the results\n",
    "Query_heads = []\n",
    "Key_heads = []\n",
    "Value_heads = []\n",
    "\n",
    "# Multiply ln1_output by each of the heads in W_Q, W_K, and W_V\n",
    "for i in range(model.cfg.n_heads):\n",
    "    Query_head = ln1_output @ model.blocks[0].attn.W_Q[i] + model.blocks[0].attn.b_Q[i]\n",
    "    Query_heads.append(Query_head)\n",
    "    Key_head = ln1_output @ model.blocks[0].attn.W_K[i] + model.blocks[0].attn.b_K[i]\n",
    "    Key_heads.append(Key_head)\n",
    "    Value_head = ln1_output @ model.blocks[0].attn.W_V[i] + model.blocks[0].attn.b_V[i]\n",
    "    Value_heads.append(Value_head)\n",
    "\n",
    "# Stack results into single tensors\n",
    "Q_stacked_heads = torch.stack(Query_heads)\n",
    "Q_heads_aligned = Q_stacked_heads.permute(1, 0, 2)          # Reshape to (seq_len, num_heads, head_dim)\n",
    "K_stacked_heads = torch.stack(Key_heads)\n",
    "K_heads_aligned = K_stacked_heads.permute(1, 0, 2)          # Reshape to (seq_len, num_heads, head_dim)\n",
    "V_stacked_heads = torch.stack(Value_heads)\n",
    "V_heads_aligned = V_stacked_heads.permute(1, 0, 2)          # Reshape to (seq_len, num_heads, head_dim)\n",
    "\n",
    "print(\"\\nIs Q_heads_aligned the same as blocks.0.attn.hook_q?\", torch.allclose(Q_heads_aligned, cache['blocks.0.attn.hook_q'][0], atol=1e-4))\n",
    "print(\"\\nIs K_heads_aligned the same as blocks.0.attn.hook_k?\", torch.allclose(K_heads_aligned, cache['blocks.0.attn.hook_k'][0], atol=1e-4))\n",
    "print(\"\\nIs V_heads_aligned the same as blocks.0.attn.hook_v?\", torch.allclose(V_heads_aligned, cache['blocks.0.attn.hook_v'][0], atol=1e-4))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Shape of attention_scores_masked: torch.Size([16, 8, 8])\n",
      "\n",
      "Is attention_scores_masked the same as blocks.0.attn.hook_attn_scores? True\n"
     ]
    }
   ],
   "source": [
    "attention_heads = []\n",
    "# Compute attention scores\n",
    "for i in range(model.cfg.n_heads):\n",
    "    attention_head = Query_heads[i] @ Key_heads[i].transpose(0, 1)\n",
    "    attention_heads.append(attention_head)\n",
    "\n",
    "attention_scores = torch.stack(attention_heads)\n",
    "d_k = model.cfg.d_head\n",
    "attention_scores_scaled = attention_scores / math.sqrt(d_k)\n",
    "\n",
    "# Create causal mask\n",
    "seq_len = attention_scores_scaled.size(-1)\n",
    "causal_mask = torch.triu(torch.ones(seq_len, seq_len, device=attention_scores_scaled.device), diagonal=1).bool()\n",
    "\n",
    "# Apply causal mask\n",
    "attention_scores_masked = attention_scores_scaled.masked_fill(causal_mask, float('-inf'))\n",
    "\n",
    "print(\"\\nShape of attention_scores_masked:\", attention_scores_masked.shape)\n",
    "print(\"\\nIs attention_scores_masked the same as blocks.0.attn.hook_attn_scores?\", torch.allclose(attention_scores_masked, cache['blocks.0.attn.hook_attn_scores'][0], atol=1e-4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Shape of attention_weights: torch.Size([16, 8, 8])\n",
      "\n",
      "Is attention_weights the same as blocks.0.attn.hook_pattern? True\n"
     ]
    }
   ],
   "source": [
    "# Apply softmax to get attention weights\n",
    "attention_weights = F.softmax(attention_scores_masked, dim=-1)\n",
    "\n",
    "print(\"\\nShape of attention_weights:\", attention_weights.shape)\n",
    "print(\"\\nIs attention_weights the same as blocks.0.attn.hook_pattern?\", torch.allclose(attention_weights, cache['blocks.0.attn.hook_pattern'][0], atol=1e-4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Is weighted_values_final the same as blocks.0.attn.hook_z? True\n"
     ]
    }
   ],
   "source": [
    "weighted_values = []\n",
    "\n",
    "for i in range(model.cfg.n_heads):\n",
    "    weighted_values.append(attention_weights[i] @ Value_heads[i])\n",
    "\n",
    "weighted_values_stacked = torch.stack(weighted_values)\n",
    "weighted_values_final = weighted_values_stacked.permute(1, 0, 2)\n",
    "\n",
    "print(\"\\nIs weighted_values_final the same as blocks.0.attn.hook_z?\", torch.allclose(weighted_values_final, cache['blocks.0.attn.hook_z'][0], atol=1e-4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Is attention_outputs_final the same as blocks.0.hook_attn_out? True\n"
     ]
    }
   ],
   "source": [
    "attention_output_heads = []\n",
    "for i in range(model.cfg.n_heads):\n",
    "    attention_output_heads.append(weighted_values_stacked[i] @ model.W_O[0][i])\n",
    "\n",
    "attention_outputs= torch.stack(attention_output_heads)              # Attention outputs are the weighted Values multiplied by the W_O weights\n",
    "\n",
    "attention_out_summed = torch.sum(attention_outputs, dim=0)          # Sum the attention outputs across heads\n",
    "attention_outputs_final  = attention_out_summed + model.b_O[0]      # Add the bias to the summed attention outputs\n",
    "\n",
    "print(\"\\nIs attention_outputs_final the same as blocks.0.hook_attn_out?\", torch.allclose(attention_outputs_final, cache['blocks.0.hook_attn_out'][0], atol=1e-4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Is resid_mid the same as blocks.0.hook_resid_mid? True\n"
     ]
    }
   ],
   "source": [
    "resid_mid = cache['blocks.0.hook_resid_pre'][0] + attention_outputs_final\n",
    "print(\"\\nIs resid_mid the same as blocks.0.hook_resid_mid?\", torch.allclose(resid_mid, cache['blocks.0.hook_resid_mid'][0], atol=1e-4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Is std the same as blocks.0.ln2.hook_scale? True\n",
      "\n",
      "Is ln2_output the same as blocks.0.ln2.hook_normalized? True\n"
     ]
    }
   ],
   "source": [
    "# 1. Normalize\n",
    "mean = resid_mid.mean(dim=-1, keepdim=True)\n",
    "std = resid_mid.std(dim=-1, keepdim=True)\n",
    "\n",
    "# Confirm that std is identical to cache['blocks.0.ln2.hook_scale'] (I can't find ln1.hook_bias anywhere)\n",
    "print(\"\\nIs std the same as blocks.0.ln2.hook_scale?\", torch.allclose(std, cache['blocks.0.ln2.hook_scale'][0], atol=1e-4))\n",
    "\n",
    "normalized_resid_mid = (resid_mid - mean) / cache['blocks.0.ln2.hook_scale'][0]\n",
    "\n",
    "# 2. Apply weight (gamma) and bias (beta) (Even after normalizing, we still need to apply the learned layernorm weights)\n",
    "ln2_output = normalized_resid_mid * model.blocks[0].ln2.w + model.blocks[0].ln2.b\n",
    "\n",
    "# Compare with the output of the first layer norm in the model\n",
    "print(\"\\nIs ln2_output the same as blocks.0.ln2.hook_normalized?\", torch.allclose(ln2_output, cache['blocks.0.ln2.hook_normalized'][0], atol=1e-4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Is mlp_in the same as blocks.0.mlp.hook_pre? True\n"
     ]
    }
   ],
   "source": [
    "mlp_in = ln2_output @ model.W_in[0] + model.b_in[0]\n",
    "\n",
    "print(\"\\nIs mlp_in the same as blocks.0.mlp.hook_pre?\", torch.allclose(mlp_in, cache['blocks.0.mlp.hook_pre'][0], atol=1e-4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Is mlp_gelu the same as blocks.0.mlp.hook_post? True\n"
     ]
    }
   ],
   "source": [
    "# Apply GELU activation function\n",
    "mlp_gelu = torch.nn.functional.gelu(mlp_in)\n",
    "\n",
    "print(\"Is mlp_gelu the same as blocks.0.mlp.hook_post?\", torch.allclose(mlp_gelu, cache['blocks.0.mlp.hook_post'][0], atol=1e-3))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Is mlp_out the same as blocks.0.hook_mlp_out? True\n"
     ]
    }
   ],
   "source": [
    "mlp_out = cache['blocks.0.mlp.hook_post'][0] @ model.W_out[0] + model.b_out[0]\n",
    "\n",
    "print(\"\\nIs mlp_out the same as blocks.0.hook_mlp_out?\", torch.allclose(mlp_out, cache['blocks.0.hook_mlp_out'][0], atol=1e-4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Is resid_post the same as blocks.0.hook_resid_post? True\n",
      "\n",
      "Is resid_post_0 the same as resid_pre_1? True\n"
     ]
    }
   ],
   "source": [
    "resid_post = resid_mid + mlp_out\n",
    "print(\"\\nIs resid_post the same as blocks.0.hook_resid_post?\", torch.allclose(resid_post, cache['blocks.0.hook_resid_post'][0], atol=1e-4))\n",
    "print(\"\\nIs resid_post_0 the same as resid_pre_1?\", torch.allclose(cache['blocks.0.hook_resid_post'][0],cache['blocks.1.hook_resid_pre'][0], atol=1e-4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Is final_ln_output equivalent to cache['ln_final.hook_normalized']? True\n",
      "\n",
      "Is manual_logits equivalent to logits? True\n"
     ]
    }
   ],
   "source": [
    "# Get the final residual stream output\n",
    "final_residual = cache['blocks.23.hook_resid_post'][0]  # Assuming 24 layers (0-23)\n",
    "\n",
    "# Apply final layer normalization\n",
    "final_ln_mean = final_residual.mean(dim=-1, keepdim=True)\n",
    "final_ln_std = cache['ln_final.hook_scale'][0]  ## ln2.hook_scale is sd of resid_post\n",
    "final_ln_normalized = (final_residual - final_ln_mean) / final_ln_std\n",
    "\n",
    "# Apply learned scale and bias of the final layer norm\n",
    "final_ln_output = final_ln_normalized * model.ln_final.w + model.ln_final.b\n",
    "\n",
    "# Check if final_ln_output is equivalent to cache['ln_final.hook_normalized']\n",
    "is_equivalent = torch.allclose(final_ln_output, cache['ln_final.hook_normalized'][0], atol=1e-4)\n",
    "print(\"\\nIs final_ln_output equivalent to cache['ln_final.hook_normalized']?\", is_equivalent)\n",
    "\n",
    "# Apply the unembedding (usually just a linear transformation)\n",
    "manual_logits = final_ln_output @ model.W_U + model.b_U   # for gpt2-med b_U is a vector of 0s\n",
    "print(\"\\nIs manual_logits equivalent to logits?\", torch.allclose(manual_logits, logits, atol=1e-4))\n",
    "\n",
    "probs = F.softmax(manual_logits, dim=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Top 5 most likely next words:\n",
      "-----------------------------\n",
      "Word\t\tProbability\n",
      "-----------------------------\n",
      " heavens       28.98%\n",
      " world         13.31%\n",
      " earth         10.56%\n",
      " universe      5.70%\n",
      " heaven        4.76%\n"
     ]
    }
   ],
   "source": [
    "# If you want to get the prediction for the next token after the sequence:\n",
    "next_token_probs = probs[-1, :]  # Shape: (batch_size, vocab_size)\n",
    "\n",
    "# Get the top 5 probabilities and their indices\n",
    "top_5_probs, top_5_indices = torch.topk(next_token_probs, k=5)\n",
    "\n",
    "# Convert probabilities to percentages\n",
    "top_5_percentages = top_5_probs * 100\n",
    "\n",
    "# Create a list of tuples (word, probability)\n",
    "top_5_words_and_probs = [(model.to_string(idx.item()), prob.item()) for idx, prob in zip(top_5_indices, top_5_percentages)]\n",
    "\n",
    "# Print the table\n",
    "print(\"Top 5 most likely next words:\")\n",
    "print(\"-----------------------------\")\n",
    "print(\"Word\\t\\tProbability\")\n",
    "print(\"-----------------------------\")\n",
    "for word, prob in top_5_words_and_probs:\n",
    "    print(f\"{word:<15}{prob:.2f}%\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "akpy24",
   "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}