extracting_activations.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/activation-additions-large-models/myenv/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "from src.utils.extract_utils import gather_activations, gather_activations_from_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch, transformers, accelerate, einops"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from src.utils.model_utils import load_gpt_model_and_tokeniser"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Detected kernel version 5.4.0, which is below the recommended minimum of 5.5.0; this can cause the process to hang. It is recommended to upgrade the kernel to the minimum version or higher.\n"
     ]
    }
   ],
   "source": [
    "model, tokenizer, MODEL_CONFIG = load_gpt_model_and_tokeniser(model_name=\"gpt2-xl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "resid_td = gather_activations(\n",
    "    \"Good day sir, how are you doing today?\",\n",
    "    MODEL_CONFIG['layer_hook_names'],\n",
    "    model,\n",
    "    tokenizer,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(resid_td[MODEL_CONFIG['layer_hook_names'][0]].output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 1, 10, 1600])"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.stack([resid_td['transformer.h.0'].output[0], resid_td['transformer.h.0'].output[0]]).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "attn_td = gather_activations(\n",
    "    \"Good day sir, how are you doing today?\",\n",
    "    MODEL_CONFIG['attn_hook_names'],\n",
    "    model,\n",
    "    tokenizer,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([10, 1600])"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "attn_td['transformer.h.0.attn.c_proj'].output[0].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "ename": "KeyError",
     "evalue": "'transformer.h.0.attn.c_proj'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyError\u001b[0m                                  Traceback (most recent call last)",
      "\u001b[1;32m/root/activation-additions-large-models/extracting_activations.ipynb Cell 8\u001b[0m line \u001b[0;36m2\n\u001b[1;32m      <a href='vscode-notebook-cell://ssh-remote%2Bvastai/root/activation-additions-large-models/extracting_activations.ipynb#X21sdnNjb2RlLXJlbW90ZQ%3D%3D?line=0'>1</a>\u001b[0m \u001b[39m#desired shape: (batch_size, layer, n_tokens, dim)\u001b[39;00m\n\u001b[0;32m----> <a href='vscode-notebook-cell://ssh-remote%2Bvastai/root/activation-additions-large-models/extracting_activations.ipynb#X21sdnNjb2RlLXJlbW90ZQ%3D%3D?line=1'>2</a>\u001b[0m tensor \u001b[39m=\u001b[39m torch\u001b[39m.\u001b[39mstack([resid_td[layer]\u001b[39m.\u001b[39moutput \u001b[39mfor\u001b[39;00m layer \u001b[39min\u001b[39;00m MODEL_CONFIG[\u001b[39m'\u001b[39m\u001b[39mattn_hook_names\u001b[39m\u001b[39m'\u001b[39m]], dim\u001b[39m=\u001b[39m\u001b[39m1\u001b[39m)\n",
      "\u001b[1;32m/root/activation-additions-large-models/extracting_activations.ipynb Cell 8\u001b[0m line \u001b[0;36m2\n\u001b[1;32m      <a href='vscode-notebook-cell://ssh-remote%2Bvastai/root/activation-additions-large-models/extracting_activations.ipynb#X21sdnNjb2RlLXJlbW90ZQ%3D%3D?line=0'>1</a>\u001b[0m \u001b[39m#desired shape: (batch_size, layer, n_tokens, dim)\u001b[39;00m\n\u001b[0;32m----> <a href='vscode-notebook-cell://ssh-remote%2Bvastai/root/activation-additions-large-models/extracting_activations.ipynb#X21sdnNjb2RlLXJlbW90ZQ%3D%3D?line=1'>2</a>\u001b[0m tensor \u001b[39m=\u001b[39m torch\u001b[39m.\u001b[39mstack([resid_td[layer]\u001b[39m.\u001b[39moutput \u001b[39mfor\u001b[39;00m layer \u001b[39min\u001b[39;00m MODEL_CONFIG[\u001b[39m'\u001b[39m\u001b[39mattn_hook_names\u001b[39m\u001b[39m'\u001b[39m]], dim\u001b[39m=\u001b[39m\u001b[39m1\u001b[39m)\n",
      "\u001b[0;31mKeyError\u001b[0m: 'transformer.h.0.attn.c_proj'"
     ]
    }
   ],
   "source": [
    "#desired shape: (batch_size, layer, n_tokens, dim)\n",
    "tensor = torch.stack([resid_td[layer].output for layer in MODEL_CONFIG['attn_hook_names']], dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 48, 10, 1600])"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "einops.rearrange()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([10, 1600])"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "td[MODEL_CONFIG['attn_hook_names'][1]].output[0].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['__abstractmethods__',\n",
       " '__class__',\n",
       " '__class_getitem__',\n",
       " '__delattr__',\n",
       " '__dict__',\n",
       " '__dir__',\n",
       " '__doc__',\n",
       " '__enter__',\n",
       " '__eq__',\n",
       " '__exit__',\n",
       " '__format__',\n",
       " '__ge__',\n",
       " '__getattribute__',\n",
       " '__gt__',\n",
       " '__hash__',\n",
       " '__init__',\n",
       " '__init_subclass__',\n",
       " '__le__',\n",
       " '__lt__',\n",
       " '__module__',\n",
       " '__ne__',\n",
       " '__new__',\n",
       " '__reduce__',\n",
       " '__reduce_ex__',\n",
       " '__repr__',\n",
       " '__setattr__',\n",
       " '__sizeof__',\n",
       " '__slots__',\n",
       " '__str__',\n",
       " '__subclasshook__',\n",
       " '__weakref__',\n",
       " '_abc_impl',\n",
       " 'close',\n",
       " 'input',\n",
       " 'layer',\n",
       " 'registered_hook',\n",
       " 'stop']"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dir(td[MODEL_CONFIG['attn_hook_names'][0]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([48, 10, 1600])\n",
      "torch.Size([48, 10, 1600])\n"
     ]
    }
   ],
   "source": [
    "activations = gather_activations_from_dataset(\n",
    "    [\"Good day sir, how are you doing today?\", \"I am doing well, thank you for asking.\"],\n",
    "    [\"attn_hook_names\"],\n",
    "    model,\n",
    "    tokenizer,\n",
    "    MODEL_CONFIG,\n",
    "    2,\n",
    "    False,\n",
    "    True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([48, 1, 1600])"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "activations['attn_hook_names'][0].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'attn_hook_names': {0: tensor([[[-1.7877e-02,  3.2479e-02, -1.4434e-02,  ..., -6.7941e-02,\n",
       "             7.5504e-02,  1.6733e-02]],\n",
       "  \n",
       "          [[ 1.3076e-01, -6.2840e-02, -4.5099e-02,  ...,  2.1293e-01,\n",
       "            -4.2480e-02,  3.4764e-02]],\n",
       "  \n",
       "          [[ 2.0405e-02,  4.4779e-02, -3.6921e-02,  ..., -3.0827e-02,\n",
       "            -1.4238e-01, -3.7328e-02]],\n",
       "  \n",
       "          ...,\n",
       "  \n",
       "          [[-6.4238e-01, -6.7807e-01, -3.2924e-01,  ..., -1.6132e+01,\n",
       "            -6.1254e-01, -1.1006e+00]],\n",
       "  \n",
       "          [[-1.4672e-01,  9.3154e-02, -1.3612e-01,  ..., -3.6234e+01,\n",
       "            -3.1323e-01, -4.6978e-01]],\n",
       "  \n",
       "          [[ 3.5768e-01,  1.9606e-01,  8.6479e-01,  ..., -1.1085e+01,\n",
       "            -4.4635e-01,  5.1237e-02]]], device='cuda:0',\n",
       "         grad_fn=<ReshapeAliasBackward0>),\n",
       "  1: tensor([[[ 8.4371e-03,  2.3719e-02,  1.0827e-02,  ..., -3.3776e-02,\n",
       "             4.6411e-02, -2.2447e-02]],\n",
       "  \n",
       "          [[-4.3184e-03, -4.8394e-02, -2.3387e-01,  ..., -1.5990e-01,\n",
       "             1.3697e-02,  1.2034e-01]],\n",
       "  \n",
       "          [[ 1.5330e-01, -5.8621e-02,  8.7831e-03,  ..., -5.1612e-02,\n",
       "            -7.5535e-02, -7.2179e-02]],\n",
       "  \n",
       "          ...,\n",
       "  \n",
       "          [[-3.7202e-01, -1.7997e-01, -1.6125e-01,  ..., -1.9993e+01,\n",
       "             5.5155e-01,  1.9151e-01]],\n",
       "  \n",
       "          [[ 5.0026e-01,  5.9802e-01,  5.8433e-01,  ..., -4.3568e+01,\n",
       "            -5.1727e-02,  5.0607e-01]],\n",
       "  \n",
       "          [[ 4.2035e-01,  1.0714e+00,  9.2622e-01,  ..., -9.8997e+00,\n",
       "            -2.6473e-01,  3.8333e-01]]], device='cuda:0',\n",
       "         grad_fn=<ReshapeAliasBackward0>)}}"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "activations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "myenv",
   "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.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}