From 097f9b386da55092aae223792b3d1663a66cb3c0 Mon Sep 17 00:00:00 2001
From: Aaron Pham <contact@aarnphm.xyz>
Date: Sun, 22 Sep 2024 21:32:42 -0400
Subject: [PATCH] fix(ci): update correct logic to check for action

Signed-off-by: Aaron Pham <contact@aarnphm.xyz>
---
 .github/workflows/pr-preview.yml            |   7 +-
 python/exo/notebooks/self_explanation.ipynb | 362 +++++++++++++-------
 2 files changed, 244 insertions(+), 125 deletions(-)

diff --git a/.github/workflows/pr-preview.yml b/.github/workflows/pr-preview.yml
index c6e093b..7844331 100644
--- a/.github/workflows/pr-preview.yml
+++ b/.github/workflows/pr-preview.yml
@@ -1,4 +1,3 @@
-# .github/workflows/preview.yml
 name: Deploy PR previews
 
 on:
@@ -30,16 +29,16 @@ jobs:
       - name: Install Dependencies
         run: npm ci
       - name: Install and Build
-        if: github.event.action != 'closed' # You might want to skip the build if the PR has been closed
+        if: github.event.action != 'closed'
         run: npx quartz build
       - uses: rossjrw/pr-preview-action@430e3dfc1de8a8ae77e77d862d25676ef9db55d1 # ratchet:rossjrw/pr-preview-action@v1
-        if: contains(['opened', 'reopened', 'synchronize'], github.event.action)
+        if: github.event.action != 'closed'
         with:
           source-dir: ./public/
           custom-url: tinymorph.aarnphm.xyz
           action: deploy
       - uses: rossjrw/pr-preview-action@430e3dfc1de8a8ae77e77d862d25676ef9db55d1 # ratchet:rossjrw/pr-preview-action@v1
-        if: github.event.action == "closed" && !github.event.pull_request.merged
+        if: github.event.action == 'closed' && !github.event.pull_request.merged
         with:
           source-dir: ./public/
           custom-url: tinymorph.aarnphm.xyz
diff --git a/python/exo/notebooks/self_explanation.ipynb b/python/exo/notebooks/self_explanation.ipynb
index ec1bf61..f64dfa5 100644
--- a/python/exo/notebooks/self_explanation.ipynb
+++ b/python/exo/notebooks/self_explanation.ipynb
@@ -30,44 +30,32 @@
     "from nnsight import LanguageModel\n",
     "from sae import Sae\n",
     "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
-    "import torch"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 3,
-   "id": "adc4ca03-7e2b-41e0-916b-616df493e892",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "def load_model(name: str, /, *, trust_remote_code: bool = False) -> LanguageModel:\n",
-    "  model = LanguageModel(\n",
-    "    name, device_map='auto', trust_remote_code=trust_remote_code, low_cpu_mem_usage=True, torch_dtype=torch.float16\n",
-    "  )\n",
-    "  model.requires_grad_(False)\n",
-    "  return model\n",
-    "\n",
+    "import torch\n",
     "\n",
-    "def load_manual_sae(name: str, hookpoint: str):\n",
-    "  repo_path = pathlib.Path(hf_hub.snapshot_download(name, allow_patterns=f'{hookpoint}/*'))\n",
-    "  if not ((hookpath := repo_path / hookpoint) / 'cfg.json').exists():\n",
-    "    raise FileNotFoundError('require cfg.json to specify a specific layers')\n",
-    "  with st.safe_open(str(hookpath / 'sae_weights.safetensors'), framework='pt') as stf:\n",
-    "    return stf.get_tensor('W_dec')"
+    "from transformer_lens import HookedTransformer\n",
+    "from sae_lens import SAE\n",
+    "from sae_lens.toolkit.pretrained_saes import get_gpt2_res_jb_saes"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": 12,
    "id": "afb4f816-f570-444e-a805-43070128bc2d",
    "metadata": {},
    "outputs": [],
    "source": [
     "prompt = 'Logic is easy, but it is impossible to be logical to bitter end. It is considered truth if one'\n",
+    "\n",
+    "model_id = 'meta-llama/Meta-Llama-3.1-8B-Instruct'\n",
+    "\n",
     "feature = 6689\n",
     "scale = 20\n",
     "max_new_tokens = 40\n",
-    "num_return_sequences = 10"
+    "num_return_sequences = 10\n",
+    "\n",
+    "if torch.backends.mps.is_available(): device = \"mps\"\n",
+    "elif torch.cuda.is_available(): device = 'cuda'\n",
+    "else: device = 'cpu'"
    ]
   },
   {
@@ -79,7 +67,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "aeb6a09180da41c0a3d956752899ebcc",
+       "model_id": "f347b2c3f8ca4786b66588f316091a79",
        "version_major": 2,
        "version_minor": 0
       },
@@ -96,65 +84,187 @@
      "text": [
       "Dropping extra args {'signed': False}\n"
      ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "be581fba30c34bbd898be012a3a8ac4f",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "model.safetensors.index.json:   0%|          | 0.00/23.9k [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "76f991355bda4a6f8a92caedcabd1f6d",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Downloading shards:   0%|          | 0/4 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "1ed30ea37dfb4f1a9d4de02b911e373e",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "model-00001-of-00004.safetensors:   0%|          | 0.00/4.98G [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "1a067e92c55347b5983697a27fcda977",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "model-00002-of-00004.safetensors:   0%|          | 0.00/5.00G [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "2d0b264c222b49cfac4691f904564c09",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "model-00003-of-00004.safetensors:   0%|          | 0.00/4.92G [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "a7a57308962f46a8b31b302329fc3f4d",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "model-00004-of-00004.safetensors:   0%|          | 0.00/1.17G [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "fb0dabd095874f298f9e9fdd76c8951f",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Loading checkpoint shards:   0%|          | 0/4 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Some parameters are on the meta device because they were offloaded to the disk.\n"
+     ]
     }
    ],
    "source": [
-    "sae = Sae.load_from_hub('EleutherAI/sae-llama-3.1-8b-32x', hookpoint='layers.29.mlp')\n",
-    "\n",
-    "model = load_model('meta-llama/Meta-Llama-3.1-8B')"
+    "model = AutoModelForCausalLM.from_pretrained(model_id, device_map='auto', low_cpu_mem_usage=True)\n",
+    "tokenizer = AutoTokenizer.from_pretrained(model_id)"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 13,
    "id": "2690df4a-3c51-4817-b175-0caba417dbd7",
    "metadata": {},
+   "outputs": [],
+   "source": [
+    "inputs = tokenizer(prompt, return_tensors=\"pt\").to(device)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 18,
+   "id": "e4a38578-abcb-44c3-800b-d1b1cb2f6667",
+   "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
-       "LlamaForCausalLM(\n",
-       "  (model): LlamaModel(\n",
-       "    (embed_tokens): Embedding(128256, 4096)\n",
-       "    (layers): ModuleList(\n",
-       "      (0-31): 32 x LlamaDecoderLayer(\n",
-       "        (self_attn): LlamaSdpaAttention(\n",
-       "          (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
-       "          (k_proj): Linear(in_features=4096, out_features=1024, bias=False)\n",
-       "          (v_proj): Linear(in_features=4096, out_features=1024, bias=False)\n",
-       "          (o_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
-       "          (rotary_emb): LlamaRotaryEmbedding()\n",
-       "        )\n",
-       "        (mlp): LlamaMLP(\n",
-       "          (gate_proj): Linear(in_features=4096, out_features=14336, bias=False)\n",
-       "          (up_proj): Linear(in_features=4096, out_features=14336, bias=False)\n",
-       "          (down_proj): Linear(in_features=14336, out_features=4096, bias=False)\n",
-       "          (act_fn): SiLU()\n",
-       "        )\n",
-       "        (input_layernorm): LlamaRMSNorm((4096,), eps=1e-05)\n",
-       "        (post_attention_layernorm): LlamaRMSNorm((4096,), eps=1e-05)\n",
-       "      )\n",
-       "    )\n",
-       "    (norm): LlamaRMSNorm((4096,), eps=1e-05)\n",
-       "    (rotary_emb): LlamaRotaryEmbedding()\n",
-       "  )\n",
-       "  (lm_head): Linear(in_features=4096, out_features=128256, bias=False)\n",
-       "  (generator): WrapperModule()\n",
-       ")"
+       "{'input_ids': tensor([[128000,  27849,    374,   4228,     11,    719,    433,    374,  12266,\n",
+       "            311,    387,  20406,    311,  26242,    842,     13,   1102,    374,\n",
+       "           6646,   8206,    422,    832]], device='mps:0'), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]],\n",
+       "       device='mps:0')}"
       ]
      },
-     "execution_count": 6,
+     "execution_count": 18,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
-    "model"
+    "inputs"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 20,
+   "id": "d0c20d8b-ca27-46b1-88af-88a6838f8514",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "8283514e61984fc8a7790df8e16cd932",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Fetching 2 files:   0%|          | 0/2 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Dropping extra args {'signed': False}\n"
+     ]
+    }
+   ],
+   "source": [
+    "sae = Sae.load_from_hub('EleutherAI/sae-llama-3.1-8b-32x', hookpoint='layers.29.mlp', device=device)"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 21,
    "id": "3f3dbf9f-c19a-4fa5-befd-e1db22c279b0",
    "metadata": {},
    "outputs": [],
@@ -166,28 +276,72 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 36,
    "id": "52e9d3b1-ddf7-4600-b004-ad09d2c69d2f",
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "tensor([[[-2.9511,  1.5943,  7.4130,  ...,  1.6402,  1.6402,  1.6402],\n",
+       "         [ 5.6696,  2.3355,  3.7905,  ..., -5.8401, -5.8403, -5.8404],\n",
+       "         [ 2.0789,  1.2966, -0.0946,  ..., -8.3018, -8.3017, -8.3015],\n",
+       "         ...,\n",
+       "         [ 7.9952,  5.1692,  2.5138,  ..., -5.6675, -5.6675, -5.6675],\n",
+       "         [ 2.3323,  2.0644, -0.5631,  ..., -8.2491, -8.2490, -8.2491],\n",
+       "         [ 4.1120,  4.4859,  3.9373,  ..., -5.1948, -5.1947, -5.1947]]],\n",
+       "       device='mps:0')"
+      ]
+     },
+     "execution_count": 36,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
    "source": [
-    "latent_acts"
+    "outputs.logits"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 24,
    "id": "69bce6ce-2c63-473a-9666-0fb2a707ebd2",
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "(tensor([[[ 88290, 113519, 108890,  ...,    863, 129994, 121807],\n",
+       "          [ 96360,  21315,  58801,  ...,  20409,  84362,  29115],\n",
+       "          [ 47834,  44914,  14233,  ...,  54857,  18015,  95291],\n",
+       "          ...,\n",
+       "          [ 29238, 120395,  77180,  ...,  33662, 116106,  94123],\n",
+       "          [ 85340,  34458, 109670,  ..., 119287, 121843,  94770],\n",
+       "          [ 91277,  17100,   1518,  ...,  69911,   7855,  83043]]],\n",
+       "        device='mps:0'),\n",
+       " tensor([[[14.8641, 14.7052, 14.5979,  ..., 12.7887, 12.7804, 12.7784],\n",
+       "          [ 1.2117,  1.1897,  1.1791,  ...,  0.9237,  0.9226,  0.9215],\n",
+       "          [ 1.5425,  1.4885,  1.3973,  ...,  1.0895,  1.0895,  1.0894],\n",
+       "          ...,\n",
+       "          [ 1.6315,  1.5188,  1.4945,  ...,  1.1529,  1.1526,  1.1523],\n",
+       "          [ 1.5072,  1.4983,  1.4491,  ...,  1.1689,  1.1662,  1.1647],\n",
+       "          [ 1.7338,  1.7059,  1.7001,  ...,  1.2518,  1.2516,  1.2503]]],\n",
+       "        device='mps:0'))"
+      ]
+     },
+     "execution_count": 24,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
    "source": [
     "l = latent_acts[0]\n",
-    "l.top_indices"
+    "l.top_indices, l.top_acts"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": 27,
    "id": "ddc2cbbf-d0c3-4907-bfd4-662403d1bc10",
    "metadata": {},
    "outputs": [
@@ -195,10 +349,10 @@
      "data": {
       "text/plain": [
        "tensor([[ 0.0091, -0.0137, -0.0124,  ...,  0.0023,  0.0068, -0.0111]],\n",
-       "       grad_fn=<IndexBackward0>)"
+       "       device='mps:0', grad_fn=<IndexBackward0>)"
       ]
      },
-     "execution_count": 9,
+     "execution_count": 27,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -209,70 +363,36 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": 33,
+   "id": "5ae1d336-d32d-491f-a954-2cb7a1436d47",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "v = sae.W_dec[[6689]]\n",
+    "v = v / v.norm()\n",
+    "v = v * scale"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 35,
    "id": "f41ea6b1-58f4-40ad-9ccb-c4854f32f78d",
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
-       "\u001b[0;31mSignature:\u001b[0m       \u001b[0msae\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
-       "\u001b[0;31mType:\u001b[0m            Sae\n",
-       "\u001b[0;31mString form:\u001b[0m    \n",
-       "Sae(\n",
-       "  (encoder): Linear(in_features=4096, out_features=131072, bias=True)\n",
-       ")\n",
-       "\u001b[0;31mFile:\u001b[0m            ~/workspace/capstone/.venv/lib/python3.12/site-packages/sae/sae.py\n",
-       "\u001b[0;31mDocstring:\u001b[0m       <no docstring>\n",
-       "\u001b[0;31mClass docstring:\u001b[0m\n",
-       "Base class for all neural network modules.\n",
-       "\n",
-       "Your models should also subclass this class.\n",
-       "\n",
-       "Modules can also contain other Modules, allowing to nest them in\n",
-       "a tree structure. You can assign the submodules as regular attributes::\n",
-       "\n",
-       "    import torch.nn as nn\n",
-       "    import torch.nn.functional as F\n",
-       "\n",
-       "    class Model(nn.Module):\n",
-       "        def __init__(self):\n",
-       "            super().__init__()\n",
-       "            self.conv1 = nn.Conv2d(1, 20, 5)\n",
-       "            self.conv2 = nn.Conv2d(20, 20, 5)\n",
-       "\n",
-       "        def forward(self, x):\n",
-       "            x = F.relu(self.conv1(x))\n",
-       "            return F.relu(self.conv2(x))\n",
-       "\n",
-       "Submodules assigned in this way will be registered, and will have their\n",
-       "parameters converted too when you call :meth:`to`, etc.\n",
-       "\n",
-       ".. note::\n",
-       "    As per the example above, an ``__init__()`` call to the parent class\n",
-       "    must be made before assignment on the child.\n",
-       "\n",
-       ":ivar training: Boolean represents whether this module is in training or\n",
-       "                evaluation mode.\n",
-       ":vartype training: bool\n",
-       "\u001b[0;31mInit docstring:\u001b[0m  Initialize internal Module state, shared by both nn.Module and ScriptModule."
+       "torch.Size([1, 4096])"
       ]
      },
+     "execution_count": 35,
      "metadata": {},
-     "output_type": "display_data"
+     "output_type": "execute_result"
     }
    ],
    "source": [
-    "sae?"
+    "It would trouble me if my child identified as trans.This question is required."
    ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "a6471b16-970d-4f46-9860-ed397ee34459",
-   "metadata": {},
-   "outputs": [],
-   "source": []
   }
  ],
  "metadata": {