experiments/bib_shift.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "095b36eae555406ca780bab8c368e05b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import sys\n",
    "import os\n",
    "parent_dir = os.path.abspath('..')\n",
    "sys.path.append(parent_dir)\n",
    "\n",
    "from datasets import load_dataset\n",
    "import random\n",
    "from nnsight import LanguageModel\n",
    "import torch as t\n",
    "from torch import nn\n",
    "from attribution import patching_effect, Submodule\n",
    "from dictionary_learning import AutoEncoder, ActivationBuffer, JumpReluAutoEncoder\n",
    "from dictionary_learning.dictionary import IdentityDict\n",
    "from dictionary_learning.interp import examine_dimension\n",
    "from dictionary_loading_utils import load_saes_and_submodules\n",
    "from dictionary_learning.utils import hf_dataset_to_generator\n",
    "from transformers import BitsAndBytesConfig\n",
    "from huggingface_hub import hf_hub_download, list_repo_files\n",
    "from tqdm import tqdm\n",
    "from typing import Literal\n",
    "import gc\n",
    "from collections import defaultdict\n",
    "import hashlib\n",
    "\n",
    "\n",
    "DEBUGGING = False\n",
    "\n",
    "if DEBUGGING:\n",
    "    tracer_kwargs = dict(scan=True, validate=True)\n",
    "else:\n",
    "    tracer_kwargs = dict(scan=False, validate=False)\n",
    "\n",
    "# model hyperparameters\n",
    "DTYPE = t.bfloat16\n",
    "DEVICE = 'cuda:0'\n",
    "# model = LanguageModel('EleutherAI/pythia-70m-deduped', device_map=DEVICE, dispatch=True)\n",
    "model = LanguageModel('google/gemma-2-2b', device_map=DEVICE, dispatch=True,\n",
    "                      attn_implementation=\"eager\", torch_dtype=DTYPE)\n",
    "activation_dim = 2304"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# dataset hyperparameters\n",
    "dataset = load_dataset(\"LabHC/bias_in_bios\")\n",
    "profession_dict = {'professor' : 21, 'nurse' : 13}\n",
    "male_prof = 'professor'\n",
    "female_prof = 'nurse'\n",
    "\n",
    "# data preparation hyperparameters\n",
    "SEED = 42\n",
    "\n",
    "def get_text_batches(\n",
    "    split: Literal[\"train\", \"test\"] = \"train\",\n",
    "    ambiguous=True, \n",
    "    batch_size=32, \n",
    "    seed=SEED\n",
    "):\n",
    "    data = dataset[split]\n",
    "    if ambiguous:\n",
    "        neg = [x['hard_text'] for x in data if x['profession'] == profession_dict[male_prof] and x['gender'] == 0]\n",
    "        pos = [x['hard_text'] for x in data if x['profession'] == profession_dict[female_prof] and x['gender'] == 1]\n",
    "        n = min([len(neg), len(pos)])\n",
    "        neg, pos = neg[:n], pos[:n]\n",
    "        data = neg + pos\n",
    "        labels = [0]*n + [1]*n\n",
    "        idxs = list(range(2*n))\n",
    "        random.Random(seed).shuffle(idxs)\n",
    "        data, labels = [data[i] for i in idxs], [labels[i] for i in idxs]\n",
    "        true_labels = spurious_labels = labels\n",
    "    else:\n",
    "        neg_neg = [x['hard_text'] for x in data if x['profession'] == profession_dict[male_prof] and x['gender'] == 0]\n",
    "        neg_pos = [x['hard_text'] for x in data if x['profession'] == profession_dict[male_prof] and x['gender'] == 1]\n",
    "        pos_neg = [x['hard_text'] for x in data if x['profession'] == profession_dict[female_prof] and x['gender'] == 0]\n",
    "        pos_pos = [x['hard_text'] for x in data if x['profession'] == profession_dict[female_prof] and x['gender'] == 1]\n",
    "        n = min([len(neg_neg), len(neg_pos), len(pos_neg), len(pos_pos)])\n",
    "        neg_neg, neg_pos, pos_neg, pos_pos = neg_neg[:n], neg_pos[:n], pos_neg[:n], pos_pos[:n]\n",
    "        data = neg_neg + neg_pos + pos_neg + pos_pos\n",
    "        true_labels     = [0]*n + [0]*n + [1]*n + [1]*n\n",
    "        spurious_labels = [0]*n + [1]*n + [0]*n + [1]*n\n",
    "        idxs = list(range(4*n))\n",
    "        random.Random(seed).shuffle(idxs)\n",
    "        data, true_labels, spurious_labels = [data[i] for i in idxs], [true_labels[i] for i in idxs], [spurious_labels[i] for i in idxs]\n",
    "\n",
    "    batches = [\n",
    "        (data[i:i+batch_size], t.tensor(true_labels[i:i+batch_size], device=DEVICE), t.tensor(spurious_labels[i:i+batch_size], device=DEVICE)) for i in range(0, len(data), batch_size)\n",
    "    ]\n",
    "\n",
    "    return batches\n",
    "\n",
    "def get_subgroups(\n",
    "        split: Literal[\"train\", \"test\"] = \"test\",\n",
    "        batch_size=32,\n",
    "):\n",
    "    data = dataset[split]\n",
    "    neg_neg = [x['hard_text'] for x in data if x['profession'] == profession_dict[male_prof] and x['gender'] == 0]\n",
    "    neg_pos = [x['hard_text'] for x in data if x['profession'] == profession_dict[male_prof] and x['gender'] == 1]\n",
    "    pos_neg = [x['hard_text'] for x in data if x['profession'] == profession_dict[female_prof] and x['gender'] == 0]\n",
    "    pos_pos = [x['hard_text'] for x in data if x['profession'] == profession_dict[female_prof] and x['gender'] == 1]\n",
    "    neg_neg_labels, neg_pos_labels, pos_neg_labels, pos_pos_labels = (0, 0), (0, 1), (1, 0), (1, 1)\n",
    "    subgroups = [(neg_neg, neg_neg_labels), (neg_pos, neg_pos_labels), (pos_neg, pos_neg_labels), (pos_pos, pos_pos_labels)]\n",
    "    \n",
    "    out = {}\n",
    "    for data, label_profile in subgroups:\n",
    "        out[label_profile] = []\n",
    "        for i in range(0, len(data), batch_size):\n",
    "            text = data[i:i+batch_size]\n",
    "            out[label_profile].append(\n",
    "                (\n",
    "                    text,\n",
    "                    t.tensor([label_profile[0]]*len(text), device=DEVICE),\n",
    "                    t.tensor([label_profile[1]]*len(text), device=DEVICE)\n",
    "                )\n",
    "            )\n",
    "    return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pool_acts(acts, attn_mask):\n",
    "    return (acts * attn_mask[:, :, None]).sum(1) / attn_mask.sum(1)[:, None]\n",
    "\n",
    "@t.no_grad()\n",
    "def collect_activations(\n",
    "    model,\n",
    "    layer,\n",
    "    text_batches,\n",
    "):\n",
    "    with tqdm(total=len(text_batches), desc=\"Collecting activations\") as pbar:\n",
    "        for text_batch, *labels in text_batches:\n",
    "            with model.trace(text_batch, **tracer_kwargs):\n",
    "                attn_mask = model.input[1]['attention_mask']\n",
    "                acts = model.model.layers[layer].output[0]\n",
    "                pooled_acts = pool_acts(acts, attn_mask).save()\n",
    "            yield pooled_acts.value, *labels\n",
    "            pbar.update(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# probe training hyperparameters\n",
    "\n",
    "layer = 22 # model layer for attaching linear classification head\n",
    "\n",
    "class Probe(nn.Module):\n",
    "    def __init__(self, activation_dim, dtype=DTYPE):\n",
    "        super().__init__()\n",
    "        self.net = nn.Linear(activation_dim, 1, bias=True, dtype=dtype)\n",
    "\n",
    "    def forward(self, x):\n",
    "        logits = self.net(x).squeeze(-1)\n",
    "        return logits\n",
    "\n",
    "def train_probe(\n",
    "    activation_batches,\n",
    "    label_idx=0,\n",
    "    lr=1e-2,\n",
    "    epochs=1,\n",
    "    d_probe=activation_dim,\n",
    "    seed=SEED,\n",
    "):\n",
    "    t.manual_seed(seed)\n",
    "\n",
    "    probe = Probe(d_probe).to(DEVICE)\n",
    "    optimizer = t.optim.AdamW(probe.parameters(), lr=lr)\n",
    "    criterion = nn.BCEWithLogitsLoss()\n",
    "    losses = []\n",
    "\n",
    "    for epoch in range(epochs):\n",
    "        for act, *labels, in activation_batches:\n",
    "            optimizer.zero_grad()\n",
    "            logits = probe(act)\n",
    "            loss = criterion(logits, labels[label_idx].to(logits))\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            losses.append(loss.item())\n",
    "    return probe, losses\n",
    "\n",
    "@t.no_grad()\n",
    "def test_probe(\n",
    "    probe,\n",
    "    activation_batches,\n",
    "):\n",
    "    corrects = defaultdict(list)\n",
    "    for acts, *labels in activation_batches:\n",
    "        logits = probe(acts)\n",
    "        preds = (logits > 0.0).long()\n",
    "        for idx, label in enumerate(labels):\n",
    "            corrects[idx].append(preds == label)\n",
    "\n",
    "    accs = {idx: t.cat(corrects[idx]).float().mean().item() for idx in corrects}\n",
    "\n",
    "    return accs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Collecting activations:   0%|          | 0/141 [00:00<?, ?it/s]You're using a GemmaTokenizerFast tokenizer. Please note that with a fast tokenizer, using the `__call__` method is faster than using a method to encode the text followed by a call to the `pad` method to get a padded encoding.\n",
      "Collecting activations: 100%|██████████| 141/141 [00:24<00:00,  5.87it/s]\n",
      "Collecting activations: 100%|██████████| 269/269 [00:43<00:00,  6.24it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ambiguous test accuracy: 0.957597553730011\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Collecting activations: 100%|██████████| 55/55 [00:08<00:00,  6.46it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ground truth accuracy: 0.9504608511924744\n",
      "unintended feature accuracy: 0.5057603716850281\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Collecting activations: 100%|██████████| 507/507 [01:26<00:00,  5.88it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "subgroup (0, 0) accuracy: 0.9862387776374817\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Collecting activations: 100%|██████████| 417/417 [01:10<00:00,  5.92it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "subgroup (0, 1) accuracy: 0.9644010066986084\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Collecting activations: 100%|██████████| 14/14 [00:01<00:00,  7.15it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "subgroup (1, 0) accuracy: 0.9308755993843079\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Collecting activations: 100%|██████████| 135/135 [00:19<00:00,  6.77it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "subgroup (1, 1) accuracy: 0.928438663482666\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "oracle, _ = train_probe(\n",
    "    activation_batches=collect_activations(model, layer, get_text_batches(split=\"train\", ambiguous=False))\n",
    ")\n",
    "\n",
    "ambiguous_accs = test_probe(oracle, activation_batches=collect_activations(model, layer, get_text_batches(split=\"test\", ambiguous=True)))\n",
    "print(f\"ambiguous test accuracy: {ambiguous_accs[0]}\")\n",
    "\n",
    "unambiguous_accs = test_probe(oracle, activation_batches=collect_activations(model, layer, get_text_batches(split=\"test\", ambiguous=False)))\n",
    "print(f\"ground truth accuracy: {unambiguous_accs[0]}\")\n",
    "print(f\"unintended feature accuracy: {unambiguous_accs[1]}\")\n",
    "\n",
    "for subgroup, batches in get_subgroups().items():\n",
    "    subgroup_accs = test_probe(oracle, activation_batches=collect_activations(model, layer, batches))\n",
    "    print(f\"subgroup {subgroup} accuracy: {subgroup_accs[0]}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading probe from probe_layer_22_bfloat16.pt\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Collecting activations: 100%|██████████| 269/269 [00:42<00:00,  6.27it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ambiguous test accuracy: 0.9965148568153381\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Collecting activations: 100%|██████████| 55/55 [00:08<00:00,  6.44it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ground truth accuracy: 0.6774193644523621\n",
      "unintended feature accuracy: 0.8191244006156921\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Collecting activations: 100%|██████████| 507/507 [01:25<00:00,  5.93it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "subgroup (0, 0) accuracy: 0.9985806345939636\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Collecting activations: 100%|██████████| 417/417 [01:09<00:00,  5.99it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "subgroup (0, 1) accuracy: 0.5165602564811707\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Collecting activations: 100%|██████████| 14/14 [00:01<00:00,  7.30it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "subgroup (1, 0) accuracy: 0.1820276528596878\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Collecting activations: 100%|██████████| 135/135 [00:19<00:00,  6.92it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "subgroup (1, 1) accuracy: 0.9944238066673279\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "save_path = f\"probe_layer_{layer}_{str(DTYPE).split('.')[-1]}.pt\"\n",
    "if os.path.exists(save_path):\n",
    "    print(f\"loading probe from {save_path}\")\n",
    "    probe = t.load(save_path)\n",
    "else:\n",
    "    probe, _ = train_probe(\n",
    "        activation_batches=collect_activations(model, layer, get_text_batches(split=\"train\", ambiguous=True))\n",
    "    )\n",
    "    t.save(probe, save_path)\n",
    "    print(f\"probe saved to {save_path}\")\n",
    "\n",
    "ambiguous_accs = test_probe(probe, activation_batches=collect_activations(model, layer, get_text_batches(split=\"test\", ambiguous=True)))\n",
    "print(f\"ambiguous test accuracy: {ambiguous_accs[0]}\")\n",
    "\n",
    "unambiguous_accs = test_probe(probe, activation_batches=collect_activations(model, layer, get_text_batches(split=\"test\", ambiguous=False)))\n",
    "print(f\"ground truth accuracy: {unambiguous_accs[0]}\")\n",
    "print(f\"unintended feature accuracy: {unambiguous_accs[1]}\")\n",
    "\n",
    "for subgroup, batches in get_subgroups().items():\n",
    "    subgroup_accs = test_probe(probe, activation_batches=collect_activations(model, layer, batches))\n",
    "    print(f\"subgroup {subgroup} accuracy: {subgroup_accs[0]}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading Gemma SAEs:   0%|          | 0/23 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading Gemma SAEs: 100%|██████████| 23/23 [01:18<00:00,  3.40s/it]\n"
     ]
    }
   ],
   "source": [
    "# loading dictionaries\n",
    "submodules, dictionaries = load_saes_and_submodules(\n",
    "    model, \n",
    "    thru_layer=layer,\n",
    "    include_embed=False,\n",
    "    dtype=DTYPE,\n",
    "    device=DEVICE,\n",
    ")\n",
    "\n",
    "def metric_fn(model, labels=None):\n",
    "    attn_mask = model.input[1]['attention_mask']\n",
    "    acts = model.model.layers[layer].output[0]\n",
    "    acts = pool_acts(acts, attn_mask)\n",
    "    \n",
    "    return t.where(\n",
    "        labels == 0,\n",
    "        probe(acts),\n",
    "        - probe(acts)\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 100/100 [00:00<00:00, 76315.57it/s]\n"
     ]
    }
   ],
   "source": [
    "# find most influential features\n",
    "# n_batches = 200\n",
    "n_batches = 100\n",
    "batch_size = 1\n",
    "\n",
    "nodes = None\n",
    "\n",
    "for batch_idx, (clean, labels, _) in tqdm(enumerate(get_text_batches(split=\"train\", ambiguous=True, batch_size=batch_size)), total=n_batches):\n",
    "    if batch_idx == n_batches:\n",
    "        break\n",
    "\n",
    "    # if effects are already cached, skip\n",
    "    hash_input = clean + [s.name for s in submodules]\n",
    "    hash_str = ''.join(hash_input)\n",
    "    hash_digest = hashlib.md5(hash_str.encode()).hexdigest()\n",
    "    if os.path.exists(f\"effects/{hash_digest}.pt\"):\n",
    "        continue\n",
    "\n",
    "    effects, *_ = patching_effect(\n",
    "        clean,\n",
    "        None,\n",
    "        model,\n",
    "        submodules,\n",
    "        dictionaries,\n",
    "        metric_fn,\n",
    "        metric_kwargs=dict(labels=labels),\n",
    "        method='ig'\n",
    "    )\n",
    "    to_save = {\n",
    "        k.name : v.detach().to(\"cpu\") for k, v in effects.items()\n",
    "    }\n",
    "    t.save(to_save, f\"effects/{hash_digest}.pt\")\n",
    "\n",
    "    del effects, _\n",
    "    gc.collect()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# aggregate effects\n",
    "aggregated_effects = {submodule.name : 0 for submodule in submodules}\n",
    "\n",
    "for idx, (clean, *_) in enumerate(get_text_batches(split=\"train\", ambiguous=True, batch_size=batch_size)):\n",
    "    if idx == n_batches:\n",
    "        break\n",
    "    hash_input = clean + [s.name for s in submodules]\n",
    "    hash_str = ''.join(hash_input)\n",
    "    hash_digest = hashlib.md5(hash_str.encode()).hexdigest()\n",
    "    effects = t.load(f\"effects/{hash_digest}.pt\")\n",
    "    for submodule in submodules:\n",
    "        aggregated_effects[submodule.name] += (\n",
    "            effects[submodule.name].act[:,1:,:] # remove BOS features\n",
    "        ).sum(dim=1).mean(dim=0)\n",
    "\n",
    "aggregated_effects = {k : v / (batch_size * n_batches) for k, v in aggregated_effects.items()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "attn_0\n",
      "mlp_0\n",
      "resid_0\n",
      "  4449 : 7.4375\n",
      "attn_1\n",
      "mlp_1\n",
      "resid_1\n",
      "  4521 : 7.25\n",
      "  11782 : 6.375\n",
      "attn_2\n",
      "mlp_2\n",
      "resid_2\n",
      "  5853 : 6.71875\n",
      "attn_3\n",
      "mlp_3\n",
      "resid_3\n",
      "attn_4\n",
      "mlp_4\n",
      "resid_4\n",
      "attn_5\n",
      "mlp_5\n",
      "resid_5\n",
      "  2864 : 9.25\n",
      "  11682 : 8.1875\n",
      "attn_6\n",
      "mlp_6\n",
      "resid_6\n",
      "  4068 : 13.0625\n",
      "  7008 : 7.1875\n",
      "attn_7\n",
      "mlp_7\n",
      "resid_7\n",
      "  7111 : 8.0625\n",
      "attn_8\n",
      "mlp_8\n",
      "resid_8\n",
      "  6952 : 7.09375\n",
      "  9949 : 19.5\n",
      "attn_9\n",
      "mlp_9\n",
      "resid_9\n",
      "  6952 : 9.375\n",
      "  15246 : 14.875\n",
      "attn_10\n",
      "mlp_10\n",
      "resid_10\n",
      "  3711 : 14.5625\n",
      "  6952 : 9.125\n",
      "attn_11\n",
      "mlp_11\n",
      "resid_11\n",
      "  3013 : 27.375\n",
      "  4467 : 7.3125\n",
      "  11649 : 6.15625\n",
      "attn_12\n",
      "mlp_12\n",
      "resid_12\n",
      "  6335 : 9.25\n",
      "  11114 : 36.25\n",
      "  11480 : 6.28125\n",
      "attn_13\n",
      "mlp_13\n",
      "resid_13\n",
      "  192 : 34.5\n",
      "  495 : 6.09375\n",
      "  14755 : 6.90625\n",
      "attn_14\n",
      "mlp_14\n",
      "resid_14\n",
      "  2354 : 30.375\n",
      "  12665 : 6.46875\n",
      "attn_15\n",
      "mlp_15\n",
      "resid_15\n",
      "  798 : 25.875\n",
      "  6211 : 6.6875\n",
      "attn_16\n",
      "mlp_16\n",
      "resid_16\n",
      "  6047 : 6.40625\n",
      "  15567 : 6.40625\n",
      "  16351 : 26.25\n",
      "attn_17\n",
      "mlp_17\n",
      "resid_17\n",
      "  6011 : 19.125\n",
      "attn_18\n",
      "mlp_18\n",
      "resid_18\n",
      "  61 : 15.75\n",
      "attn_19\n",
      "mlp_19\n",
      "resid_19\n",
      "  13002 : 10.5\n",
      "attn_20\n",
      "  9711 : 8.0\n",
      "mlp_20\n",
      "resid_20\n",
      "  7116 : 9.25\n",
      "  7324 : 6.1875\n",
      "  12545 : 16.5\n",
      "attn_21\n",
      "  2740 : 11.6875\n",
      "  10118 : 7.84375\n",
      "mlp_21\n",
      "resid_21\n",
      "  1065 : 12.875\n",
      "  1653 : 6.46875\n",
      "  4430 : 24.0\n",
      "attn_22\n",
      "mlp_22\n",
      "resid_22\n",
      "  1208 : 14.8125\n",
      "  3497 : 24.375\n",
      "  7961 : 7.0\n",
      "total features: 46\n"
     ]
    }
   ],
   "source": [
    "count = 0\n",
    "for k, v in aggregated_effects.items():\n",
    "    print(k)\n",
    "    for idx in (v > 6).nonzero():\n",
    "        count += 1\n",
    "        print(f\"  {idx.item()} : {v[idx].item()}\")\n",
    "print(f\"total features: {count}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "        <iframe\n",
       "            width=\"1200\"\n",
       "            height=\"600\"\n",
       "            src=\"https://neuronpedia.org/gemma-2-2b/21-gemmascope-att-16k/10118?embed=true&embedexplanation=true&embedplots=true&embedtest=true&height=300\"\n",
       "            frameborder=\"0\"\n",
       "            allowfullscreen\n",
       "            \n",
       "        ></iframe>\n",
       "        "
      ],
      "text/plain": [
       "<IPython.lib.display.IFrame at 0x7f2825bc5990>"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# interpret features with Neuronpedia API\n",
    "\n",
    "submodule_name = \"attn_21\"\n",
    "feature_idx = 10118\n",
    "\n",
    "from IPython.display import IFrame\n",
    "html_template = \"https://neuronpedia.org/{}/{}/{}?embed=true&embedexplanation=true&embedplots=true&embedtest=true&height=300\"\n",
    "\n",
    "def get_dashboard_html(sae_release=\"gemma-2-2b\", sae_id=None, feature_idx=feature_idx):\n",
    "    return html_template.format(sae_release, sae_id, feature_idx)\n",
    "\n",
    "# Extract the type and number from submodule_name\n",
    "submodule_type, submodule_number = submodule_name.split('_')\n",
    "\n",
    "# Construct the sae_id based on the type\n",
    "if submodule_type == \"resid\":\n",
    "    sae_id = f\"{submodule_number}-gemmascope-res-16k\"\n",
    "elif submodule_type == \"attn\":\n",
    "    sae_id = f\"{submodule_number}-gemmascope-att-16k\"\n",
    "elif submodule_type == \"mlp\":\n",
    "    sae_id = f\"{submodule_number}-gemmascope-mlp-16k\"\n",
    "else:\n",
    "    raise ValueError(\"Unknown submodule type\")\n",
    "\n",
    "html = get_dashboard_html(sae_release=\"gemma-2-2b\", sae_id=sae_id, feature_idx=feature_idx)\n",
    "IFrame(html, width=1200, height=600)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "number of features to ablate: 43\n"
     ]
    }
   ],
   "source": [
    "feats_to_ablate = {\n",
    "    \"attn_0\": [\n",
    "    ],\n",
    "    \"mlp_0\": [\n",
    "    ],\n",
    "    \"resid_0\": [\n",
    "        4449, # \"he\", \"him\"\n",
    "    ],\n",
    "    \"attn_1\": [\n",
    "    ],\n",
    "    \"mlp_1\": [\n",
    "    ],\n",
    "    \"resid_1\": [\n",
    "        4521, # his\"\n",
    "        11782, # # \"woman\", \"she\", \"her\", \"female\"\n",
    "    ],\n",
    "    \"attn_2\": [\n",
    "    ],\n",
    "    \"mlp_2\": [\n",
    "    ],\n",
    "    \"resid_2\": [\n",
    "        5853, # \"he\"\n",
    "    ],\n",
    "    \"attn_3\": [\n",
    "    ],\n",
    "    \"mlp_3\": [\n",
    "    ],\n",
    "    \"resid_3\": [\n",
    "    ],\n",
    "    \"attn_4\": [\n",
    "    ],\n",
    "    \"mlp_4\": [\n",
    "    ],\n",
    "    \"resid_4\": [\n",
    "    ],\n",
    "    \"attn_5\": [\n",
    "    ],\n",
    "    \"mlp_5\": [\n",
    "    ],\n",
    "    \"resid_5\": [\n",
    "        2864, # \"her\"\n",
    "        11682, # \"he\"\n",
    "    ],\n",
    "    \"attn_6\": [\n",
    "    ],\n",
    "    \"mlp_6\": [\n",
    "    ],\n",
    "    \"resid_6\": [\n",
    "        4068, # female pronouns\n",
    "        7008, # gendered pronouns\n",
    "    ],\n",
    "    \"attn_7\": [\n",
    "    ],\n",
    "    \"mlp_7\": [\n",
    "    ],\n",
    "    \"resid_7\": [\n",
    "        7111, # text about women\n",
    "    ],\n",
    "    \"attn_8\": [\n",
    "    ],\n",
    "    \"mlp_8\": [\n",
    "    ],\n",
    "    \"resid_8\": [\n",
    "        6952, # gendered pronouns\n",
    "        9949, # female pronouns\n",
    "    ],\n",
    "    \"attn_9\": [\n",
    "    ],\n",
    "    \"mlp_9\": [\n",
    "    ],\n",
    "    \"resid_9\": [\n",
    "        6952, # \"she\"\n",
    "        15246, # female pronouns\n",
    "    ],\n",
    "    \"attn_10\": [\n",
    "    ],\n",
    "    \"mlp_10\": [\n",
    "    ],\n",
    "    \"resid_10\": [\n",
    "        3711, # promotes female-associated words\n",
    "        6952, # masculine pronouns\n",
    "    ],\n",
    "    \"attn_11\": [\n",
    "    ],\n",
    "    \"mlp_11\": [\n",
    "    ],\n",
    "    \"resid_11\": [\n",
    "        3013, # descriptions of women\n",
    "        4467, # \"He\"\n",
    "        11649, # \"she\", \"her\"\n",
    "    ],\n",
    "    \"attn_12\": [\n",
    "    ],\n",
    "    \"mlp_12\": [\n",
    "    ],\n",
    "    \"resid_12\": [\n",
    "        6335, # \"He\"\n",
    "        11114, # descriptions of women\n",
    "        11480, # \"his\", \"her\"\n",
    "    ],\n",
    "    \"attn_13\": [\n",
    "    ],\n",
    "    \"mlp_13\": [\n",
    "    ],\n",
    "    \"resid_13\": [\n",
    "        192, # descriptions of women\n",
    "        495, # \"her\"\n",
    "        14755, # \"he\"\n",
    "    ],\n",
    "    \"attn_14\": [\n",
    "    ],\n",
    "    \"mlp_14\": [\n",
    "    ],\n",
    "    \"resid_14\": [\n",
    "        2354, # descriptions of women\n",
    "        12665, # male pronouns\n",
    "    ],\n",
    "    \"attn_15\": [\n",
    "    ],\n",
    "    \"mlp_15\": [\n",
    "    ],\n",
    "    \"resid_15\": [\n",
    "        798, # promotes feminine pronouns\n",
    "        6211, # gendered pronouns\n",
    "    ],\n",
    "    \"attn_16\": [\n",
    "    ],\n",
    "    \"mlp_16\": [\n",
    "    ],\n",
    "    \"resid_16\": [\n",
    "        6047, # women, girls, women's names\n",
    "        15567, # \"he\", \"she\"\n",
    "        16351, # promotes feminine pronouns\n",
    "    ],\n",
    "    \"attn_17\": [\n",
    "    ],\n",
    "    \"mlp_17\": [\n",
    "    ],\n",
    "    \"resid_17\": [\n",
    "        6011, # promotes feminine pronouns\n",
    "    ],\n",
    "    \"attn_18\": [\n",
    "    ],\n",
    "    \"mlp_18\": [\n",
    "    ],\n",
    "    \"resid_18\": [\n",
    "        61, # descriptions of female royalty\n",
    "    ],\n",
    "    \"attn_19\": [\n",
    "    ],\n",
    "    \"mlp_19\": [\n",
    "    ],\n",
    "    \"resid_19\": [\n",
    "        13002, # clauses starting with \"she\"\n",
    "    ],\n",
    "    \"attn_20\": [\n",
    "        9711, # \"her\", \"she\"\n",
    "    ],\n",
    "    \"mlp_20\": [\n",
    "    ],\n",
    "    \"resid_20\": [\n",
    "        7116, # promotes masculine pronouns\n",
    "        # 7324, # nursing\n",
    "        12545, # promotes female pronouns\n",
    "    ],\n",
    "    \"attn_21\": [\n",
    "        2740, # promotes feminine pronouns\n",
    "        10118, # masculine pronouns\n",
    "    ],\n",
    "    \"mlp_21\": [\n",
    "    ],\n",
    "    \"resid_21\": [\n",
    "        1065, # promotes masculine pronouns\n",
    "        # 1653, # nursing\n",
    "        4430, # promotes feminine pronouns\n",
    "    ],\n",
    "    \"attn_22\": [\n",
    "    ],\n",
    "    \"mlp_22\": [\n",
    "    ],\n",
    "    \"resid_22\": [\n",
    "        1208, # promotes masculine pronouns\n",
    "        3497, # promotes feminine pronouns\n",
    "        # 7961, # nursing\n",
    "    ],\n",
    "}\n",
    "print(f\"number of features to ablate: {sum([len(v) for v in feats_to_ablate.values()])}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "# putting feats_to_ablate in a more useful format\n",
    "def n_hot(feats, dim):\n",
    "    out = t.zeros(dim, dtype=t.bool, device=DEVICE)\n",
    "    for feat in feats:\n",
    "        out[feat] = True\n",
    "    return out\n",
    "\n",
    "feats_to_ablate = {\n",
    "    submodule : n_hot(feats_to_ablate[submodule.name], dictionaries[submodule].dict_size) for submodule in submodules\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "@t.no_grad()\n",
    "def collect_acts_ablated(\n",
    "    text_batches,\n",
    "    model,\n",
    "    submodules,\n",
    "    dictionaries,\n",
    "    to_ablate,\n",
    "    layer,\n",
    "):\n",
    "    with tqdm(total=len(text_batches), desc=\"Collecting activations with ablations\") as pbar:\n",
    "        for text, *labels in text_batches:\n",
    "            with model.trace(text, **tracer_kwargs):\n",
    "                for submodule in submodules:\n",
    "                    dictionary = dictionaries[submodule]\n",
    "                    feat_idxs = to_ablate[submodule]\n",
    "                    if len(feat_idxs) == 0:\n",
    "                        continue\n",
    "                    x = submodule.get_activation()\n",
    "                    x_hat, f = dictionary(x, output_features=True)\n",
    "                    res = x - x_hat\n",
    "                    f[:, :,feat_idxs] = 0. # zero ablation\n",
    "                    submodule.set_activation(dictionary.decode(f) + res)\n",
    "                attn_mask = model.input[1]['attention_mask']\n",
    "                act = model.model.layers[layer].output[0]\n",
    "                pooled_act = pool_acts(act, attn_mask).save()\n",
    "            yield pooled_act.value, *labels\n",
    "            pbar.update(1)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Accuracy after ablating features judged irrelevant by human annotators"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Collecting activations with ablations: 100%|██████████| 269/269 [02:33<00:00,  1.76it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ambiguous test accuracy: 0.7659153938293457\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Collecting activations with ablations: 100%|██████████| 55/55 [00:30<00:00,  1.80it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ground truth accuracy: 0.7603686451911926\n",
      "Spurious accuracy: 0.514976978302002\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Collecting activations with ablations: 100%|██████████| 507/507 [05:08<00:00,  1.64it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Subgroup (0, 0) accuracy: 0.9977167248725891\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Collecting activations with ablations: 100%|██████████| 417/417 [04:06<00:00,  1.69it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Subgroup (0, 1) accuracy: 0.9885091781616211\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Collecting activations with ablations: 100%|██████████| 14/14 [00:07<00:00,  1.99it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Subgroup (1, 0) accuracy: 0.5\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Collecting activations with ablations: 100%|██████████| 135/135 [01:11<00:00,  1.90it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Subgroup (1, 1) accuracy: 0.533224880695343\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "ambiguous_accs = test_probe(\n",
    "    probe,\n",
    "    activation_batches=collect_acts_ablated(\n",
    "        get_text_batches(split=\"test\", ambiguous=True), \n",
    "        model, submodules, dictionaries, feats_to_ablate, layer\n",
    "    ),\n",
    ")\n",
    "print(f\"Ambiguous test accuracy: {ambiguous_accs[0]}\")\n",
    "\n",
    "unambiguous_accs = test_probe(\n",
    "    probe,\n",
    "    activation_batches=collect_acts_ablated(\n",
    "        get_text_batches(split=\"test\", ambiguous=False),\n",
    "        model, submodules, dictionaries, feats_to_ablate, layer\n",
    "    ),\n",
    ")\n",
    "print(f\"Ground truth accuracy: {unambiguous_accs[0]}\")\n",
    "print(f\"Spurious accuracy: {unambiguous_accs[1]}\")\n",
    "\n",
    "for subgroup, batches in get_subgroups().items():\n",
    "    subgroup_accs = test_probe(\n",
    "        probe,\n",
    "        activation_batches=collect_acts_ablated(\n",
    "            batches,\n",
    "            model, submodules, dictionaries, feats_to_ablate, layer\n",
    "        ),\n",
    "    )\n",
    "    print(f\"Subgroup {subgroup} accuracy: {subgroup_accs[0]}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Concept bottleneck probing baseline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "You're using a GemmaTokenizerFast tokenizer. Please note that with a fast tokenizer, using the `__call__` method is faster than using a method to encode the text followed by a call to the `pad` method to get a padded encoding.\n"
     ]
    }
   ],
   "source": [
    "concepts = [    \n",
    "    ' nurse',\n",
    "    ' healthcare',\n",
    "    ' hospital',\n",
    "    ' patient',\n",
    "    ' medical',\n",
    "    ' clinic',\n",
    "    ' triage',\n",
    "    ' medication',\n",
    "    ' emergency',\n",
    "    ' surgery',\n",
    "    ' professor',\n",
    "    ' academia',\n",
    "    ' research',\n",
    "    ' university',\n",
    "    ' tenure',\n",
    "    ' faculty',\n",
    "    ' dissertation',\n",
    "    ' sabbatical',\n",
    "    ' publication',\n",
    "    ' grant',\n",
    "]\n",
    "# get concept vectors\n",
    "with t.no_grad():\n",
    "    with model.trace(concepts):\n",
    "        concept_vectors = model.model.layers[layer].output[0][:, -1, :].save()\n",
    "    concept_vectors = concept_vectors.value - concept_vectors.value.mean(0, keepdim=True)\n",
    "    concept_vectors = concept_vectors / concept_vectors.norm(dim=-1, keepdim=True)\n",
    "    \n",
    "\n",
    "@t.no_grad()\n",
    "def get_bottleneck_activations(\n",
    "    text_batches,\n",
    "    model,\n",
    "    layer,\n",
    "):\n",
    "    with tqdm(total=len(text_batches), desc=\"Collecting bottleneck activations\") as pbar:\n",
    "        for text, *labels in text_batches:\n",
    "            with model.trace(text, **tracer_kwargs):\n",
    "                attn_mask = model.input[1]['attention_mask']\n",
    "                acts = model.model.layers[layer].output[0]\n",
    "                pooled_acts = pool_acts(acts, attn_mask)\n",
    "                pooled_acts = pooled_acts / pooled_acts.norm(dim=-1, keepdim=True)\n",
    "                # compute cosine similarity with concept vectors\n",
    "                sims = (pooled_acts @ concept_vectors.T).save()\n",
    "            yield sims.value, *labels\n",
    "            pbar.update(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Collecting bottleneck activations: 100%|██████████| 700/700 [01:59<00:00,  5.87it/s]\n",
      "Collecting bottleneck activations: 100%|██████████| 269/269 [00:43<00:00,  6.12it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ambiguous test accuracy: 0.9138011336326599\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Collecting bottleneck activations: 100%|██████████| 55/55 [00:08<00:00,  6.38it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ground truth accuracy: 0.9020737409591675\n",
      "Spurious accuracy: 0.5011520981788635\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Collecting bottleneck activations: 100%|██████████| 507/507 [01:26<00:00,  5.87it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Subgroup (0, 0) accuracy: 0.8963282704353333\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Collecting bottleneck activations: 100%|██████████| 417/417 [01:09<00:00,  5.98it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Subgroup (0, 1) accuracy: 0.8675178289413452\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Collecting bottleneck activations: 100%|██████████| 14/14 [00:01<00:00,  7.26it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Subgroup (1, 0) accuracy: 0.9193548560142517\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Collecting bottleneck activations: 100%|██████████| 135/135 [00:19<00:00,  6.91it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Subgroup (1, 1) accuracy: 0.9321561455726624\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "cbp_probe, _ = train_probe(\n",
    "    get_bottleneck_activations(get_text_batches(split=\"train\", ambiguous=True), model, layer),\n",
    "    d_probe=len(concepts),\n",
    ")\n",
    "\n",
    "ambiguous_accs = test_probe(\n",
    "    cbp_probe,\n",
    "    get_bottleneck_activations(get_text_batches(split=\"test\", ambiguous=True), model, layer)\n",
    ")\n",
    "print(f\"Ambiguous test accuracy: {ambiguous_accs[0]}\")\n",
    "\n",
    "unambiguous_accs = test_probe(\n",
    "    cbp_probe,\n",
    "    get_bottleneck_activations(get_text_batches(split=\"test\", ambiguous=False), model, layer)\n",
    ")\n",
    "print(f\"Ground truth accuracy: {unambiguous_accs[0]}\")\n",
    "print(f\"Spurious accuracy: {unambiguous_accs[1]}\")\n",
    "\n",
    "for subgroup, batches in get_subgroups().items():\n",
    "    subgroup_accs = test_probe(\n",
    "        cbp_probe,\n",
    "        get_bottleneck_activations(batches, model, layer)\n",
    "    )\n",
    "    print(f\"Subgroup {subgroup} accuracy: {subgroup_accs[0]}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Get skyline neuron performance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 100/100 [18:42<00:00, 11.23s/it]\n"
     ]
    }
   ],
   "source": [
    "# get neurons which are most influential for giving gender label\n",
    "neuron_dicts = {\n",
    "    submodule : IdentityDict(activation_dim).to(DEVICE) for submodule in submodules\n",
    "}\n",
    "\n",
    "n_batches = 100\n",
    "batch_size = 1\n",
    "\n",
    "running_total = 0\n",
    "nodes = None\n",
    "\n",
    "for batch_idx, (clean, _, labels) in tqdm(enumerate(get_text_batches(split=\"train\", ambiguous=False, batch_size=batch_size)), total=n_batches):\n",
    "    if batch_idx == n_batches:\n",
    "        break\n",
    "\n",
    "    # if effects are already cached, skip\n",
    "    hash_input = clean + [s.name for s in submodules]\n",
    "    hash_str = ''.join(hash_input) + \"_neuron_skyline\"\n",
    "    hash_digest = hashlib.md5(hash_str.encode()).hexdigest()\n",
    "    if os.path.exists(f\"effects/{hash_digest}.pt\"):\n",
    "        continue\n",
    "\n",
    "    effects, *_  = patching_effect(\n",
    "        clean,\n",
    "        None,\n",
    "        model,\n",
    "        submodules,\n",
    "        neuron_dicts,\n",
    "        metric_fn,\n",
    "        metric_kwargs=dict(labels=labels),\n",
    "        method='ig'\n",
    "    )\n",
    "    to_save = {\n",
    "        k.name : v.detach().to(\"cpu\") for k, v in effects.items()\n",
    "    }\n",
    "    t.save(to_save, f\"effects/{hash_digest}.pt\")\n",
    "\n",
    "    del effects, _\n",
    "    gc.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# aggregated effects\n",
    "aggregated_effects = {submodule.name : 0 for submodule in submodules}\n",
    "\n",
    "for idx, (clean, *_) in enumerate(get_text_batches(split=\"train\", ambiguous=False, batch_size=batch_size)):\n",
    "    if idx == n_batches:\n",
    "        break\n",
    "    hash_input = clean + [s.name for s in submodules]\n",
    "    hash_str = ''.join(hash_input) + \"_neuron_skyline\"\n",
    "    hash_digest = hashlib.md5(hash_str.encode()).hexdigest()\n",
    "    effects = t.load(f\"effects/{hash_digest}.pt\")\n",
    "    for submodule in submodules:\n",
    "        aggregated_effects[submodule.name] += (\n",
    "            effects[submodule.name].act[:,1:,:] # ignore BOS token\n",
    "        ).sum(dim=1).mean(dim=0)\n",
    "\n",
    "aggregated_effects = {k : v / (batch_size * n_batches) for k, v in aggregated_effects.items()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "attn_0\n",
      "mlp_0\n",
      "resid_0\n",
      "attn_1\n",
      "mlp_1\n",
      "resid_1\n",
      "  1570 : 6.6875\n",
      "attn_2\n",
      "mlp_2\n",
      "resid_2\n",
      "  1570 : 11.0625\n",
      "attn_3\n",
      "mlp_3\n",
      "resid_3\n",
      "  629 : 6.4375\n",
      "attn_4\n",
      "mlp_4\n",
      "resid_4\n",
      "  1570 : 8.8125\n",
      "attn_5\n",
      "mlp_5\n",
      "resid_5\n",
      "attn_6\n",
      "mlp_6\n",
      "resid_6\n",
      "  624 : 11.625\n",
      "  682 : 6.3125\n",
      "  1068 : 6.3125\n",
      "  1570 : 6.6875\n",
      "attn_7\n",
      "mlp_7\n",
      "resid_7\n",
      "  334 : 12.75\n",
      "attn_8\n",
      "mlp_8\n",
      "resid_8\n",
      "attn_9\n",
      "mlp_9\n",
      "resid_9\n",
      "  334 : 6.84375\n",
      "  682 : 9.25\n",
      "attn_10\n",
      "mlp_10\n",
      "resid_10\n",
      "attn_11\n",
      "mlp_11\n",
      "resid_11\n",
      "  334 : 16.0\n",
      "attn_12\n",
      "mlp_12\n",
      "resid_12\n",
      "  334 : 5.5625\n",
      "  682 : 12.625\n",
      "  1570 : 6.96875\n",
      "attn_13\n",
      "mlp_13\n",
      "resid_13\n",
      "  334 : 8.0625\n",
      "attn_14\n",
      "mlp_14\n",
      "resid_14\n",
      "  334 : 13.1875\n",
      "  682 : 9.3125\n",
      "  1546 : 6.53125\n",
      "  1570 : 5.03125\n",
      "attn_15\n",
      "mlp_15\n",
      "resid_15\n",
      "attn_16\n",
      "mlp_16\n",
      "resid_16\n",
      "  334 : 5.3125\n",
      "  1068 : 6.1875\n",
      "  1149 : 5.4375\n",
      "  1711 : 4.90625\n",
      "attn_17\n",
      "mlp_17\n",
      "resid_17\n",
      "  98 : 5.4375\n",
      "  113 : 5.8125\n",
      "  682 : 22.125\n",
      "  784 : 5.75\n",
      "  1068 : 22.125\n",
      "  1149 : 6.46875\n",
      "  1546 : 14.3125\n",
      "  1570 : 8.5625\n",
      "  1645 : 11.1875\n",
      "attn_18\n",
      "mlp_18\n",
      "resid_18\n",
      "  334 : 6.84375\n",
      "attn_19\n",
      "mlp_19\n",
      "resid_19\n",
      "  682 : 19.0\n",
      "  784 : 6.25\n",
      "  1068 : 14.0625\n",
      "  1149 : 5.875\n",
      "  1393 : 11.375\n",
      "  1405 : 7.125\n",
      "  1645 : 9.5625\n",
      "  1651 : 5.375\n",
      "attn_20\n",
      "mlp_20\n",
      "resid_20\n",
      "  1546 : 6.25\n",
      "attn_21\n",
      "mlp_21\n",
      "resid_21\n",
      "attn_22\n",
      "mlp_22\n",
      "resid_22\n",
      "total neurons: 43\n"
     ]
    }
   ],
   "source": [
    "top_neurons_to_ablate = {}\n",
    "total_neurons = 0\n",
    "for submodule in submodules:\n",
    "    print(submodule.name)\n",
    "    top_neurons_to_ablate[submodule.name] = []\n",
    "    for idx in (aggregated_effects[submodule.name] > 4.8).nonzero():\n",
    "        print(f\"  {idx.item()} : {aggregated_effects[submodule.name][idx].item()}\")\n",
    "        top_neurons_to_ablate[submodule.name].append(idx.item())\n",
    "        total_neurons += 1\n",
    "print(f\"total neurons: {total_neurons}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "@t.no_grad()\n",
    "def collect_acts_ablated_neurons(\n",
    "    text_batches,\n",
    "    model,\n",
    "    submodules,\n",
    "    to_ablate,\n",
    "    layer,\n",
    "):\n",
    "    with tqdm(total=len(text_batches), desc=\"Collecting activations with ablations\") as pbar:\n",
    "        for text, *labels in text_batches:\n",
    "            with model.trace(text, **tracer_kwargs):\n",
    "                for submodule in submodules:\n",
    "                    x = submodule.get_activation()\n",
    "                    x[..., to_ablate[submodule]] = x.mean(dim=(0, 1))[..., to_ablate[submodule]] # mean ablation\n",
    "                    submodule.set_activation(x)\n",
    "                attn_mask = model.input[1]['attention_mask']\n",
    "                act = model.model.layers[layer].output[0]\n",
    "                pooled_act = pool_acts(act, attn_mask).save()\n",
    "            yield pooled_act.value, *labels\n",
    "            pbar.update(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "top_neurons_to_ablate = {\n",
    "    submodule : n_hot(\n",
    "        top_neurons_to_ablate[submodule.name], \n",
    "        dictionaries[submodule].W_enc.shape[0] # dimension of the submodule\n",
    "    ) for submodule in submodules\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Collecting activations with ablations: 100%|██████████| 269/269 [00:47<00:00,  5.68it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ambiguous test accuracy: 0.9940752387046814\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Collecting activations with ablations: 100%|██████████| 55/55 [00:09<00:00,  5.93it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ground truth accuracy: 0.650921642780304\n",
      "Spurious accuracy: 0.843317985534668\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Collecting activations with ablations: 100%|██████████| 507/507 [01:33<00:00,  5.40it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Subgroup (0, 0) accuracy: 0.9986423254013062\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Collecting activations with ablations: 100%|██████████| 417/417 [01:15<00:00,  5.50it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Subgroup (0, 1) accuracy: 0.6099886894226074\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Collecting activations with ablations: 100%|██████████| 14/14 [00:02<00:00,  6.24it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Subgroup (1, 0) accuracy: 0.04608295112848282\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Collecting activations with ablations: 100%|██████████| 135/135 [00:22<00:00,  6.11it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Subgroup (1, 1) accuracy: 0.990938663482666\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "ambiguous_accs = test_probe(\n",
    "    probe,\n",
    "    activation_batches=collect_acts_ablated_neurons(\n",
    "        get_text_batches(split=\"test\", ambiguous=True), \n",
    "        model, submodules, top_neurons_to_ablate, layer\n",
    "    ),\n",
    ")\n",
    "print(f\"Ambiguous test accuracy: {ambiguous_accs[0]}\")\n",
    "\n",
    "unambiguous_accs = test_probe(\n",
    "    probe,\n",
    "    activation_batches=collect_acts_ablated_neurons(\n",
    "        get_text_batches(split=\"test\", ambiguous=False),\n",
    "        model, submodules, top_neurons_to_ablate, layer\n",
    "    ),\n",
    ")\n",
    "print(f\"Ground truth accuracy: {unambiguous_accs[0]}\")\n",
    "print(f\"Spurious accuracy: {unambiguous_accs[1]}\")\n",
    "\n",
    "for subgroup, batches in get_subgroups().items():\n",
    "    subgroup_accs = test_probe(\n",
    "        probe,\n",
    "        activation_batches=collect_acts_ablated_neurons(\n",
    "            batches,\n",
    "            model, submodules, top_neurons_to_ablate, layer\n",
    "        ),\n",
    "    )\n",
    "    print(f\"Subgroup {subgroup} accuracy: {subgroup_accs[0]}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Get skyline feature performance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 100/100 [18:59<00:00, 11.39s/it]\n"
     ]
    }
   ],
   "source": [
    "# get features which are most useful for predicting gender label\n",
    "n_batches = 100\n",
    "batch_size = 1\n",
    "\n",
    "running_total = 0\n",
    "running_nodes = None\n",
    "\n",
    "for batch_idx, (clean, _, labels) in tqdm(enumerate(get_text_batches(split=\"train\", ambiguous=False, batch_size=batch_size, seed=SEED)), total=n_batches):\n",
    "    if batch_idx == n_batches:\n",
    "        break\n",
    "\n",
    "    hash_input = clean + [s.name for s in submodules]\n",
    "    hash_str = ''.join(hash_input) + \"_feature_skyline\"\n",
    "    hash_digest = hashlib.md5(hash_str.encode()).hexdigest()\n",
    "    if os.path.exists(f\"effects/{hash_digest}.pt\"):\n",
    "        continue\n",
    "\n",
    "    effects, *_ = patching_effect(\n",
    "        clean,\n",
    "        None,\n",
    "        model,\n",
    "        submodules,\n",
    "        dictionaries,\n",
    "        metric_fn,\n",
    "        metric_kwargs=dict(labels=labels),\n",
    "        method='ig'\n",
    "    )\n",
    "    to_save = {\n",
    "        k.name : v.detach().to(\"cpu\") for k, v in effects.items()\n",
    "    }\n",
    "    t.save(to_save, f\"effects/{hash_digest}.pt\")\n",
    "\n",
    "    del effects\n",
    "    gc.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "# aggregate effects\n",
    "aggregated_effects = {submodule.name : 0 for submodule in submodules}\n",
    "\n",
    "for idx, (clean, *_) in enumerate(get_text_batches(split=\"train\", ambiguous=False, batch_size=batch_size)):\n",
    "    if idx == n_batches:\n",
    "        break\n",
    "    hash_input = clean + [s.name for s in submodules]\n",
    "    hash_str = ''.join(hash_input) + \"_feature_skyline\"\n",
    "    hash_digest = hashlib.md5(hash_str.encode()).hexdigest()\n",
    "    effects = t.load(f\"effects/{hash_digest}.pt\")\n",
    "    for submodule in submodules:\n",
    "        aggregated_effects[submodule.name] += (\n",
    "            effects[submodule.name].act[:,1:,:] # ignore BOS token\n",
    "        ).sum(dim=1).mean(dim=0)\n",
    "\n",
    "aggregated_effects = {k : v / (batch_size * n_batches) for k, v in aggregated_effects.items()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Component attn_0:\n",
      "Component mlp_0:\n",
      "Component resid_0:\n",
      "4449 8.0\n",
      "Component attn_1:\n",
      "Component mlp_1:\n",
      "Component resid_1:\n",
      "11782 6.96875\n",
      "Component attn_2:\n",
      "Component mlp_2:\n",
      "Component resid_2:\n",
      "5853 6.53125\n",
      "12818 6.1875\n",
      "Component attn_3:\n",
      "Component mlp_3:\n",
      "Component resid_3:\n",
      "9949 7.0\n",
      "Component attn_4:\n",
      "Component mlp_4:\n",
      "Component resid_4:\n",
      "4675 6.53125\n",
      "Component attn_5:\n",
      "Component mlp_5:\n",
      "Component resid_5:\n",
      "2864 9.8125\n",
      "11682 7.125\n",
      "Component attn_6:\n",
      "Component mlp_6:\n",
      "Component resid_6:\n",
      "4068 9.3125\n",
      "7008 9.625\n",
      "Component attn_7:\n",
      "Component mlp_7:\n",
      "Component resid_7:\n",
      "7111 6.84375\n",
      "Component attn_8:\n",
      "Component mlp_8:\n",
      "Component resid_8:\n",
      "6952 7.875\n",
      "9949 14.625\n",
      "Component attn_9:\n",
      "Component mlp_9:\n",
      "Component resid_9:\n",
      "6952 10.0\n",
      "15246 11.25\n",
      "Component attn_10:\n",
      "Component mlp_10:\n",
      "Component resid_10:\n",
      "3711 11.1875\n",
      "6952 10.25\n",
      "Component attn_11:\n",
      "Component mlp_11:\n",
      "Component resid_11:\n",
      "3013 19.75\n",
      "4467 8.1875\n",
      "11649 6.15625\n",
      "Component attn_12:\n",
      "Component mlp_12:\n",
      "Component resid_12:\n",
      "6335 10.1875\n",
      "11114 25.75\n",
      "Component attn_13:\n",
      "Component mlp_13:\n",
      "Component resid_13:\n",
      "192 25.875\n",
      "14755 7.03125\n",
      "Component attn_14:\n",
      "Component mlp_14:\n",
      "Component resid_14:\n",
      "2354 22.375\n",
      "12665 6.46875\n",
      "Component attn_15:\n",
      "Component mlp_15:\n",
      "Component resid_15:\n",
      "798 19.625\n",
      "6211 7.15625\n",
      "Component attn_16:\n",
      "Component mlp_16:\n",
      "Component resid_16:\n",
      "15567 6.625\n",
      "16351 20.625\n",
      "Component attn_17:\n",
      "Component mlp_17:\n",
      "Component resid_17:\n",
      "6011 14.5625\n",
      "6257 6.125\n",
      "Component attn_18:\n",
      "Component mlp_18:\n",
      "Component resid_18:\n",
      "61 10.9375\n",
      "Component attn_19:\n",
      "Component mlp_19:\n",
      "Component resid_19:\n",
      "13002 9.3125\n",
      "Component attn_20:\n",
      "9711 6.4375\n",
      "Component mlp_20:\n",
      "Component resid_20:\n",
      "7116 10.3125\n",
      "12545 14.5625\n",
      "Component attn_21:\n",
      "2740 10.8125\n",
      "10118 8.25\n",
      "Component mlp_21:\n",
      "Component resid_21:\n",
      "1065 14.5\n",
      "4430 20.375\n",
      "Component attn_22:\n",
      "Component mlp_22:\n",
      "Component resid_22:\n",
      "1208 16.5\n",
      "3497 21.0\n",
      "total features: 43\n"
     ]
    }
   ],
   "source": [
    "top_feats_to_ablate = {}\n",
    "total_features = 0\n",
    "for submodule in submodules:\n",
    "    print(f\"Component {submodule.name}:\")\n",
    "    top_feats_to_ablate[submodule.name] = []\n",
    "    for idx in (aggregated_effects[submodule.name] > 6.1).nonzero():\n",
    "        print(idx.item(), aggregated_effects[submodule.name][idx].item())\n",
    "        top_feats_to_ablate[submodule.name].append(idx.item())\n",
    "        total_features += 1\n",
    "print(f\"total features: {total_features}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "def n_hot(feats, dim):\n",
    "    out = t.zeros(dim, dtype=t.bool, device=DEVICE)\n",
    "    for feat in feats:\n",
    "        out[feat] = True\n",
    "    return out\n",
    "\n",
    "top_feats_to_ablate = {\n",
    "    submodule : n_hot(top_feats_to_ablate[submodule.name], dictionaries[submodule].dict_size) for submodule in submodules\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Collecting activations with ablations: 100%|██████████| 269/269 [02:34<00:00,  1.75it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ambiguous test accuracy: 0.8300418257713318\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Collecting activations with ablations: 100%|██████████| 55/55 [00:31<00:00,  1.77it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ground truth accuracy: 0.8087557554244995\n",
      "Spurious accuracy: 0.5368663668632507\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Collecting activations with ablations: 100%|██████████| 507/507 [05:13<00:00,  1.62it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Subgroup (0, 0) accuracy: 0.9973464608192444\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Collecting activations with ablations: 100%|██████████| 417/417 [04:09<00:00,  1.67it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Subgroup (0, 1) accuracy: 0.9850544333457947\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Collecting activations with ablations: 100%|██████████| 14/14 [00:07<00:00,  1.97it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Subgroup (1, 0) accuracy: 0.5668202638626099\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Collecting activations with ablations: 100%|██████████| 135/135 [01:11<00:00,  1.90it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Subgroup (1, 1) accuracy: 0.6619423627853394\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "ambiguous_accs = test_probe(\n",
    "    probe,\n",
    "    activation_batches=collect_acts_ablated(\n",
    "        get_text_batches(split=\"test\", ambiguous=True), \n",
    "        model, submodules, dictionaries, top_feats_to_ablate, layer\n",
    "    ),\n",
    ")\n",
    "print(f\"Ambiguous test accuracy: {ambiguous_accs[0]}\")\n",
    "\n",
    "unambiguous_accs = test_probe(\n",
    "    probe,\n",
    "    activation_batches=collect_acts_ablated(\n",
    "        get_text_batches(split=\"test\", ambiguous=False),\n",
    "        model, submodules, dictionaries, top_feats_to_ablate, layer\n",
    "    ),\n",
    ")\n",
    "print(f\"Ground truth accuracy: {unambiguous_accs[0]}\")\n",
    "print(f\"Spurious accuracy: {unambiguous_accs[1]}\")\n",
    "\n",
    "for subgroup, batches in get_subgroups().items():\n",
    "    subgroup_accs = test_probe(\n",
    "        probe,\n",
    "        activation_batches=collect_acts_ablated(\n",
    "            batches,\n",
    "            model, submodules, dictionaries, top_feats_to_ablate, layer\n",
    "        ),\n",
    "    )\n",
    "    print(f\"Subgroup {subgroup} accuracy: {subgroup_accs[0]}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Retraining probe on activations after ablating features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Collecting activations with ablations: 100%|██████████| 700/700 [06:59<00:00,  1.67it/s]\n",
      "Collecting activations with ablations: 100%|██████████| 269/269 [02:38<00:00,  1.70it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ambiguous test accuracy: 0.977811336517334\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Collecting activations with ablations: 100%|██████████| 55/55 [00:30<00:00,  1.78it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ground truth accuracy: 0.9504608511924744\n",
      "Spurious accuracy: 0.524193525314331\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Collecting activations with ablations: 100%|██████████| 507/507 [05:07<00:00,  1.65it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Subgroup (0, 0) accuracy: 0.9780314564704895\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Collecting activations with ablations: 100%|██████████| 417/417 [04:09<00:00,  1.67it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Subgroup (0, 1) accuracy: 0.9289522767066956\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Collecting activations with ablations: 100%|██████████| 14/14 [00:07<00:00,  1.97it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Subgroup (1, 0) accuracy: 0.9585253596305847\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Collecting activations with ablations: 100%|██████████| 135/135 [01:11<00:00,  1.89it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Subgroup (1, 1) accuracy: 0.9795538783073425\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "retrained_probe, _ = train_probe(\n",
    "    activation_batches=collect_acts_ablated(\n",
    "        get_text_batches(split=\"train\", ambiguous=True), \n",
    "        model, submodules, dictionaries, feats_to_ablate, layer\n",
    "    ),\n",
    ")\n",
    "ambiguous_test_accs = test_probe(\n",
    "    retrained_probe,\n",
    "    activation_batches=collect_acts_ablated(\n",
    "        get_text_batches(split=\"test\", ambiguous=True),\n",
    "        model, submodules, dictionaries, feats_to_ablate, layer\n",
    "    ),\n",
    ")\n",
    "print(f\"Ambiguous test accuracy: {ambiguous_test_accs[0]}\")\n",
    "unambiguous_test_accs = test_probe(\n",
    "    retrained_probe,\n",
    "    activation_batches=collect_acts_ablated(\n",
    "        get_text_batches(split=\"test\", ambiguous=False),\n",
    "        model, submodules, dictionaries, feats_to_ablate, layer\n",
    "    ),\n",
    ")\n",
    "print(f\"Ground truth accuracy: {unambiguous_test_accs[0]}\")\n",
    "print(f\"Spurious accuracy: {unambiguous_test_accs[1]}\")\n",
    "\n",
    "for subgroup, batches in get_subgroups().items():\n",
    "    subgroup_accs = test_probe(\n",
    "        retrained_probe,\n",
    "        activation_batches=collect_acts_ablated(\n",
    "            batches,\n",
    "            model, submodules, dictionaries, feats_to_ablate, layer\n",
    "        ),\n",
    "    )\n",
    "    print(f\"Subgroup {subgroup} accuracy: {subgroup_accs[0]}\")"
   ]
  }
 ],
 "metadata": {
  "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.11.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}