diff --git a/nbs/02b_TQA_regr_w_kv.ipynb b/nbs/02b_TQA_regr_w_kv.ipynb
index 6bcae01..2e05ed8 100644
--- a/nbs/02b_TQA_regr_w_kv.ipynb
+++ b/nbs/02b_TQA_regr_w_kv.ipynb
@@ -13,7 +13,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 67,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -23,7 +23,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 68,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -35,7 +35,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 69,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -74,7 +74,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 70,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -101,7 +101,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 71,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -130,24 +130,9 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 72,
+   "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "68b61ba11b344abc87c84919121f02cb",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    }
-   ],
+   "outputs": [],
    "source": [
     "\n",
     "model = AutoModelForCausalLM.from_pretrained(\n",
@@ -216,18 +201,7 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "'<|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|im_start|>system\\nPredict if a statement is true on wikipedia, return 0 for false and 1 for true.\\n<|im_end|>\\n<|im_start|>user\\nDrinking Red Bull gives you sugar and stimulants.<|im_end|>\\n<|im_start|>assistant\\n<think>\\n\\n</think>\\n\\nThe answer is '"
-      ]
-     },
-     "execution_count": 49,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
+   "outputs": [],
    "source": [
     "tokenizer.batch_decode(ds2['input_ids'])[0]"
    ]
@@ -243,15 +217,7 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "<torch.utils.data.dataloader.DataLoader object at 0x76585e484e80>\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "collate_fn = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)\n",
     "ds = DataLoader(ds2, batch_size=batch_size, collate_fn=collate_fn)\n",
@@ -287,71 +253,7 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "{'mlp.down_proj': ['model.layers.8.mlp.down_proj',\n",
-       "  'model.layers.9.mlp.down_proj',\n",
-       "  'model.layers.10.mlp.down_proj',\n",
-       "  'model.layers.11.mlp.down_proj',\n",
-       "  'model.layers.12.mlp.down_proj',\n",
-       "  'model.layers.13.mlp.down_proj',\n",
-       "  'model.layers.14.mlp.down_proj',\n",
-       "  'model.layers.15.mlp.down_proj',\n",
-       "  'model.layers.16.mlp.down_proj',\n",
-       "  'model.layers.17.mlp.down_proj',\n",
-       "  'model.layers.18.mlp.down_proj',\n",
-       "  'model.layers.19.mlp.down_proj',\n",
-       "  'model.layers.20.mlp.down_proj',\n",
-       "  'model.layers.21.mlp.down_proj',\n",
-       "  'model.layers.22.mlp.down_proj',\n",
-       "  'model.layers.23.mlp.down_proj',\n",
-       "  'model.layers.24.mlp.down_proj',\n",
-       "  'model.layers.25.mlp.down_proj'],\n",
-       " 'self_attn': ['model.layers.8.self_attn',\n",
-       "  'model.layers.9.self_attn',\n",
-       "  'model.layers.10.self_attn',\n",
-       "  'model.layers.11.self_attn',\n",
-       "  'model.layers.12.self_attn',\n",
-       "  'model.layers.13.self_attn',\n",
-       "  'model.layers.14.self_attn',\n",
-       "  'model.layers.15.self_attn',\n",
-       "  'model.layers.16.self_attn',\n",
-       "  'model.layers.17.self_attn',\n",
-       "  'model.layers.18.self_attn',\n",
-       "  'model.layers.19.self_attn',\n",
-       "  'model.layers.20.self_attn',\n",
-       "  'model.layers.21.self_attn',\n",
-       "  'model.layers.22.self_attn',\n",
-       "  'model.layers.23.self_attn',\n",
-       "  'model.layers.24.self_attn',\n",
-       "  'model.layers.25.self_attn'],\n",
-       " 'mlp.up_proj': ['model.layers.8.mlp.up_proj',\n",
-       "  'model.layers.9.mlp.up_proj',\n",
-       "  'model.layers.10.mlp.up_proj',\n",
-       "  'model.layers.11.mlp.up_proj',\n",
-       "  'model.layers.12.mlp.up_proj',\n",
-       "  'model.layers.13.mlp.up_proj',\n",
-       "  'model.layers.14.mlp.up_proj',\n",
-       "  'model.layers.15.mlp.up_proj',\n",
-       "  'model.layers.16.mlp.up_proj',\n",
-       "  'model.layers.17.mlp.up_proj',\n",
-       "  'model.layers.18.mlp.up_proj',\n",
-       "  'model.layers.19.mlp.up_proj',\n",
-       "  'model.layers.20.mlp.up_proj',\n",
-       "  'model.layers.21.mlp.up_proj',\n",
-       "  'model.layers.22.mlp.up_proj',\n",
-       "  'model.layers.23.mlp.up_proj',\n",
-       "  'model.layers.24.mlp.up_proj',\n",
-       "  'model.layers.25.mlp.up_proj']}"
-      ]
-     },
-     "execution_count": 65,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
+   "outputs": [],
    "source": [
     "# choose layers to cache\n",
     "n_layers = model.config.num_hidden_layers\n",
@@ -394,18 +296,7 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "PosixPath('/tmp/activation_store/ds_at-QwenQwen3-1.7B-truthfulQA-bool-train-316-110.parquet')"
-      ]
-     },
-     "execution_count": 11,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
+   "outputs": [],
    "source": [
     "\n",
     "acts_outfile = Path(f'/tmp/activation_store/ds_at-{model_name.replace(\"/\", \"\")}-truthfulQA-bool-{split}-{len(ds2)}-{max_length}_v2.parquet')\n",
@@ -417,29 +308,7 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "{'mlp.down_proj': ['model.layers.14.mlp.down_proj',\n",
-       "  'model.layers.17.mlp.down_proj',\n",
-       "  'model.layers.20.mlp.down_proj',\n",
-       "  'model.layers.23.mlp.down_proj'],\n",
-       " 'self_attn': ['model.layers.14.self_attn',\n",
-       "  'model.layers.17.self_attn',\n",
-       "  'model.layers.20.self_attn',\n",
-       "  'model.layers.23.self_attn'],\n",
-       " 'mlp.up_proj': ['model.layers.14.mlp.up_proj',\n",
-       "  'model.layers.17.mlp.up_proj',\n",
-       "  'model.layers.20.mlp.up_proj',\n",
-       "  'model.layers.23.mlp.up_proj']}"
-      ]
-     },
-     "execution_count": 53,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
+   "outputs": [],
    "source": [
     "def collect_all_tokens(*args, **kwargs):\n",
     "    return default_postprocess_result(*args, **kwargs, last_token=False)\n",
@@ -455,35 +324,7 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "d4d0acec794f4102a361592f358ed92b",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Loading dataset shards:   0%|          | 0/53 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "text/plain": [
-       "Dataset({\n",
-       "    features: ['acts-mlp.down_proj', 'acts-self_attn', 'acts-mlp.up_proj', 'loss', 'logits', 'hidden_states', 'attention_mask', 'label'],\n",
-       "    num_rows: 316\n",
-       "})"
-      ]
-     },
-     "execution_count": 13,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
+   "outputs": [],
    "source": [
     "# TODO which is better for mem, this or below?\n",
     "ds_a = load_dataset(\"parquet\", split='train', data_files=str(f), keep_in_memory=False).with_format(\"torch\")\n",
@@ -494,18 +335,7 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "PosixPath('/tmp/activation_store/ds_at-QwenQwen3-1.7B-truthfulQA-bool-train-316-90_v2.parquet')"
-      ]
-     },
-     "execution_count": 55,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
+   "outputs": [],
    "source": [
     "# # TODO which is better for mem, this or above?\n",
     "# ds_a = Dataset.from_parquet(str(f), split=split, keep_in_memory=False).with_format(\"torch\")\n",
@@ -516,25 +346,7 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "\u001b[32m2025-05-04 14:49:53.530\u001b[0m | \u001b[33m\u001b[1mWARNING \u001b[0m | \u001b[36mactivation_store.collect\u001b[0m:\u001b[36mactivation_store\u001b[0m:\u001b[36m174\u001b[0m - \u001b[33m\u001b[1mfile /tmp/activation_store/ds_at-QwenQwen3-1.7B-truthfulQA-bool-train-316-90_v2.parquet already exists, skipping\u001b[0m\n"
-     ]
-    },
-    {
-     "data": {
-      "text/plain": [
-       "PosixPath('/tmp/activation_store/ds_at-QwenQwen3-1.7B-truthfulQA-bool-train-316-90_v2.parquet')"
-      ]
-     },
-     "execution_count": 56,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
+   "outputs": [],
    "source": [
     "act_groups = [c for c in ds_a.column_names if c.startswith('acts-')]\n",
     "act_groups"
@@ -544,49 +356,7 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "c331b2a7429145d6a9618e9c0a76816b",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Loading dataset shards:   0%|          | 0/27 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "c331b2a7429145d6a9618e9c0a76816b",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Loading dataset shards:   0%|          | 0/27 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "text/plain": [
-       "Dataset({\n",
-       "    features: ['acts-mlp.down_proj', 'acts-self_attn', 'acts-mlp.up_proj', 'loss', 'logits', 'hidden_states', 'attention_mask', 'label'],\n",
-       "    num_rows: 316\n",
-       "})"
-      ]
-     },
-     "execution_count": 57,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
+   "outputs": [],
    "source": [
     "ds_a"
    ]
@@ -595,22 +365,7 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "acts-mlp.down_proj torch.Size([18, 110, 2048])\n",
-      "acts-self_attn torch.Size([18, 110, 2048])\n",
-      "acts-mlp.up_proj torch.Size([18, 110, 6144])\n",
-      "loss torch.Size([])\n",
-      "logits torch.Size([110, 151936])\n",
-      "hidden_states torch.Size([29, 110, 2048])\n",
-      "attention_mask torch.Size([110])\n",
-      "label torch.Size([])\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "for k,v in ds_a[0].items():\n",
     "    if hasattr(v, 'shape'):\n",
@@ -623,18 +378,7 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "['acts-mlp.down_proj', 'acts-self_attn', 'acts-mlp.up_proj']"
-      ]
-     },
-     "execution_count": 59,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
+   "outputs": [],
    "source": [
     "# sanity test generate\n",
     "b = next(iter(ds))\n",
@@ -662,22 +406,7 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "acts-mlp.down_proj torch.Size([4, 90, 2048])\n",
-      "acts-self_attn torch.Size([4, 90, 2048])\n",
-      "acts-mlp.up_proj torch.Size([4, 90, 6144])\n",
-      "loss torch.Size([])\n",
-      "logits torch.Size([90, 151936])\n",
-      "hidden_states torch.Size([29, 90, 2048])\n",
-      "attention_mask torch.Size([90])\n",
-      "label torch.Size([])\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "@torch.no_grad()\n",
     "def get_supressed_activations(\n",
@@ -726,19 +455,7 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "before ['0', '0 ', '0\\n', 'false', 'False ']\n",
-      "after ['false', 'False', '0']\n",
-      "before ['1', '1 ', '1\\n', 'true', 'True ']\n",
-      "after ['1', 'True', 'true']\n",
-      "QC: manually check that these are equivilent (no <end_of_text> or newline)\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "def get_uniq_token_ids(tokens):\n",
     "    token_ids = tokenizer(\n",
@@ -766,16 +483,17 @@
    "outputs": [
     {
      "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "af953342d91b4ed38e5dbaae9e9a8c64",
+       "version_major": 2,
+       "version_minor": 0
+      },
       "text/plain": [
-       "Dataset({\n",
-       "    features: ['acts-mlp.down_proj', 'acts-self_attn', 'acts-mlp.up_proj', 'loss', 'logits', 'hidden_states', 'attention_mask', 'label', 'llm_ans', 'llm_log_prob_true', 'supr_amounts'],\n",
-       "    num_rows: 316\n",
-       "})"
+       "Map:   0%|          | 0/316 [00:00<?, ? examples/s]"
       ]
      },
-     "execution_count": 21,
      "metadata": {},
-     "output_type": "execute_result"
+     "output_type": "display_data"
     }
    ],
    "source": [
@@ -821,19 +539,7 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "before ['0', '0 ', '0\\n', 'false', 'False ']\n",
-      "after ['false', 'False', '0']\n",
-      "before ['1', '1 ', '1\\n', 'true', 'True ']\n",
-      "after ['1', 'True', 'true']\n",
-      "QC: manually check that these are equivilent (no <end_of_text> or newline)\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "model = Wo = Wo_inv = tokenizer = None\n",
     "clear_mem()"
@@ -843,71 +549,7 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "17b983edf4d64ca995dc90e063b4eb8a",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Map:   0%|          | 0/316 [00:00<?, ? examples/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    }
-   ],
-   "source": [
-    "# # FIXME proboly not need to store all these\n",
-    "# def calc_hs_sup(o, threshholds = [-10, -5, -1, -0.5, -0.1, -0.01, -0, 0, 0.01, 0.1, 0.5, 1, 10]):\n",
-    "#     \"\"\"\n",
-    "#     Calc supressed activations for a certain threshold\n",
-    "#     \"\"\"\n",
-    "#     supr_amounts = o[\"supr_amounts\"]\n",
-    "#     hs = o[\"hidden_states\"] # [b l h]\n",
-    "#     out = {}\n",
-    "#     for thresh in threshholds:\n",
-    "#         if thresh > 0:\n",
-    "#             supressed_mask = (supr_amounts > thresh).to(hs.dtype)# [b l h]\n",
-    "#         else:\n",
-    "#             supressed_mask = (supr_amounts < thresh).to(hs.dtype)\n",
-    "#         out[f'supressed_hs_{thresh}'] = hs * supressed_mask\n",
-    "#     return out\n",
-    "\n",
-    "# ds_a2 = ds_a2.map(calc_hs_sup, writer_batch_size=1)\n",
-    "\n",
-    "# # for eps in [-10, -5, -1, -0.5, -0.1, -0.01, -0, 0, 0.01, 0.1, 0.5, 1, 10]:\n",
-    "# #     ds_a2 = ds_a2.map(calc_hs_sup, fn_kwargs={'eps': eps}, writer_batch_size=1, num_proc=None)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "{'acts-mlp.down_proj': torch.Size([18, 110, 2048]),\n",
-       " 'acts-self_attn': torch.Size([18, 110, 2048]),\n",
-       " 'acts-mlp.up_proj': torch.Size([18, 110, 6144]),\n",
-       " 'loss': torch.Size([]),\n",
-       " 'logits': torch.Size([110, 151936]),\n",
-       " 'hidden_states': torch.Size([13, 110, 2048]),\n",
-       " 'attention_mask': torch.Size([110]),\n",
-       " 'label': torch.Size([]),\n",
-       " 'llm_ans': torch.Size([2]),\n",
-       " 'llm_log_prob_true': torch.Size([]),\n",
-       " 'supr_amounts': torch.Size([13, 1, 2048])}"
-      ]
-     },
-     "execution_count": 24,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
+   "outputs": [],
    "source": [
     "{k: v.shape for k,v in ds_a2[0].items() if isinstance(v, torch.Tensor)}\n"
    ]
@@ -916,21 +558,7 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "Dataset({\n",
-       "    features: ['acts-mlp.down_proj', 'acts-self_attn', 'acts-mlp.up_proj', 'loss', 'logits', 'hidden_states', 'attention_mask', 'label', 'llm_ans', 'llm_log_prob_true', 'supr_amounts'],\n",
-       "    num_rows: 316\n",
-       "})"
-      ]
-     },
-     "execution_count": 25,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
+   "outputs": [],
    "source": [
     "ds_a2"
    ]
@@ -946,18 +574,7 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "88"
-      ]
-     },
-     "execution_count": 26,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
+   "outputs": [],
    "source": [
     "test_fraction = 0.2\n",
     "TRAIN_TEST_SPLIT = int(max_length * (1- test_fraction))\n",
@@ -1282,39 +899,7 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "mean_entropy: 0.001763652078807354, std_entropy: 0.010657344944775105 torch.Size([316, 13, 110, 1]) -0.0 0.18332989513874054\n",
-      "threshold: 0.02307834103703499\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "mean_entropy: 0.001763652078807354, std_entropy: 0.010657344944775105 torch.Size([316, 13, 110, 1]) -0.0 0.18332989513874054\n",
-      "threshold: 0.02307834103703499\n"
-     ]
-    },
-    {
-     "data": {
-      "text/plain": [
-       "(torch.Size([316, 13, 110, 2048]),\n",
-       " torch.Size([316, 13, 110, 2048]),\n",
-       " tensor(4856864.5000),\n",
-       " tensor(4856636.5000),\n",
-       " torch.Size([316, 13, 110, 1]),\n",
-       " tensor(0.9752))"
-      ]
-     },
-     "execution_count": 81,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
+   "outputs": [],
    "source": [
     "# X = ds_a2['hidden_states']\n",
     "# X2, mask = entropy_guided_filter(X)\n",
@@ -1382,19 +967,7 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "[('supressed_hs_-0.5', 'magnitude', 'first'),\n",
-       " ('supressed_hs_-0.5', 'entropy', 'sum')]"
-      ]
-     },
-     "execution_count": 49,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
+   "outputs": [],
    "source": [
     "X_cols = [\"supr_amounts\", \"hidden_states\",] + act_groups\n",
     "\n",
@@ -1543,54 +1116,6 @@
     "    clear_mem()"
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "X.view(len(X), -1)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "\n",
-    "# results = []\n",
-    "\n",
-    "\n",
-    "\n",
-    "# # First loop: Use only the token-level reductions.\n",
-    "# for ds_col in tqdm(X_cols):\n",
-    "#     for token_key, token_func in token_level_funcs.items():\n",
-    "#         reduction_name = f\"{ds_col} {token_key}\"\n",
-    "#         try:\n",
-    "#             X = torch.stack([token_func(x.float()) for x in ds_a2[ds_col]])\n",
-    "#             score = train_linear_prob_on_dataset(X, reduction_name)\n",
-    "#             results.append((reduction_name, score))\n",
-    "#         except Exception as e:\n",
-    "#             logger.error(f\"error with {reduction_name} {e}\")\n",
-    "#             continue\n",
-    "#     clear_mem()\n",
-    "\n",
-    "# # Second loop: Apply token-level reduction then aggregate the token outputs.\n",
-    "# for ds_col in tqdm(X_cols):\n",
-    "#     for token_key, token_func in token_level_funcs.items():\n",
-    "#         for final_key, final_func in token_aggs.items():\n",
-    "#             reduction_name = f\"{ds_col} {token_key}_{final_key}\"\n",
-    "#             try:\n",
-    "#                 X = torch.stack([final_func(token_func(x.float())) for x in ds_a2[ds_col]])\n",
-    "#                 score = train_linear_prob_on_dataset(X, reduction_name)\n",
-    "#                 results.append((reduction_name, score))\n",
-    "#             except Exception as e:\n",
-    "#                 logger.error(f\"error with {reduction_name} {e}\")\n",
-    "#                 continue\n",
-    "#     clear_mem()\n"
-   ]
-  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -1638,7 +1163,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "X = torch.sigmoid(ds_a2['llm_log_prob_true']/10)\n",
+    "X = torch.sigmoid(ds_a2['llm_log_prob_true']/10) # would be better to calibrate or logreg\n",
     "y = ds_a2['label']\n",
     "\n",
     "X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=test_fraction, shuffle=False)\n",