lucia_and_lovis.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded 2 autoencoders\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using pad_token, but it is not set yet.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded pretrained model EleutherAI/pythia-160m into HookedTransformer\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from transformer_lens import HookedTransformer\n",
    "import numpy as np \n",
    "\n",
    "autoencoder_filename = \"pythia160m/output_sweep_tied_mlpout_l8_r4/_19/learned_dicts.pt\"\n",
    "# autoencoder_filename = \"pythia160m/output_hoagy_dense_sweep_tied_residual_l8_r4/_29/learned_dicts.pt\"\n",
    "auto_num = 0 # Selects which specific autoencoder to use\n",
    "all_autoencoders = torch.load(autoencoder_filename)\n",
    "num_dictionaries = len(all_autoencoders)\n",
    "print(f\"Loaded {num_dictionaries} autoencoders\")\n",
    "autoencoder, hyperparams = all_autoencoders[auto_num]\n",
    "l1_alpha = hyperparams['l1_alpha']\n",
    "autoencoder2, hyperparams2 = all_autoencoders[auto_num+1]\n",
    "smaller_dict = autoencoder.get_learned_dict()\n",
    "larger_dict = autoencoder2.get_learned_dict()\n",
    "\n",
    "#TODO: Lucis & Lovis\n",
    "#Change these settings to load the correct autoencoder #PLEEEEEAAASE Do this. Both layer & setting\n",
    "layer = 8\n",
    "# setting = \"residual\"\n",
    "# setting = \"attention\"\n",
    "# setting = \"mlp\"\n",
    "setting = \"mlp_out\"\n",
    "# model_name = \"EleutherAI/pythia-70m-deduped\"\n",
    "model_name = \"EleutherAI/pythia-160m\"\n",
    "\n",
    "device = torch.device(\"cuda:3\" if torch.cuda.is_available() else \"cpu\")\n",
    "model = HookedTransformer.from_pretrained(model_name, device=device)\n",
    "\n",
    "if setting == \"residual\":\n",
    "    cache_name = f\"blocks.{layer}.hook_resid_post\"\n",
    "    neurons = model.cfg.d_model\n",
    "elif setting == \"mlp\":\n",
    "    cache_name = f\"blocks.{layer}.mlp.hook_post\"\n",
    "    neurons = model.cfg.d_mlp\n",
    "elif setting == \"attention\":\n",
    "    cache_name = f\"blocks.{layer}.hook_attn_out\"\n",
    "    neurons = model.cfg.d_model\n",
    "elif setting == \"mlp_out\":\n",
    "    cache_name = f\"blocks.{layer}.hook_mlp_out\"\n",
    "    neurons = model.cfg.d_model\n",
    "else:\n",
    "    raise NotImplementedError"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "({'dict_size': 3072, 'l1_alpha': 0.0010000000474974513},\n",
       " {'dict_size': 3072, 'l1_alpha': 0.0007999999797903001})"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Two dicts: trained w/ different l1_alpha values\n",
    "hyperparams, hyperparams2"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MCS\n",
    "Max cosine similarity between one dictionary & another one. If they learned the same feature, then they'll have high cosine similarity. \n",
    "\n",
    "If two dictionaries learned it, it's probably a real feature."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "('# of features above 0.9:', 514)\n"
     ]
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAigAAAGdCAYAAAA44ojeAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/bCgiHAAAACXBIWXMAAA9hAAAPYQGoP6dpAAApl0lEQVR4nO3dfXRUdX7H8U8eyBAeZmKQzJAaHnRViMK6C2sY0W5XUyJEq4e4ym7KxpWFXQy0kBUhlQcBJWxq1WKBVOsSeoRS6VnsGh40hApVhgcj9CAg6oKbWJiJls0MYJk83f7Rw90dwZUJmeQ34f06557j3N/v3vu9PzPOx9/ceyfBsixLAAAABkns6gIAAAC+jIACAACMQ0ABAADGIaAAAADjEFAAAIBxCCgAAMA4BBQAAGAcAgoAADBOclcX0B5tbW06ceKE+vbtq4SEhK4uBwAAXALLsnT69GllZmYqMfGPz5HEZUA5ceKEsrKyuroMAADQDvX19brmmmv+aJ+4DCh9+/aV9P8n6HQ6u7gaAABwKUKhkLKysuzP8T8mLgPK+a91nE4nAQUAgDhzKZdncJEsAAAwDgEFAAAYJ6qA0traqvnz52vIkCFKTU3VddddpyVLlsiyLLuPZVlasGCBBgwYoNTUVOXm5uqjjz6K2M+pU6dUWFgop9OptLQ0TZ48WWfOnOmYMwIAAHEvqoDyi1/8QqtWrdI//MM/6MiRI/rFL36h8vJyvfDCC3af8vJyLV++XBUVFdqzZ4969+6tvLw8nTt3zu5TWFioQ4cOqbq6WlVVVdq5c6emTp3acWcFAADiWoL1h9MfX+Oee+6R2+3Wyy+/bK8rKChQamqqXnnlFVmWpczMTP385z/XY489JkkKBoNyu92qrKzUxIkTdeTIEWVnZ2vfvn0aNWqUJGnr1q0aP368Pv30U2VmZn5tHaFQSC6XS8FgkItkAQCIE9F8fkc1g3LbbbeppqZGH374oSTpv/7rv/T2229r3LhxkqTjx4/L7/crNzfX3sblciknJ0c+n0+S5PP5lJaWZocTScrNzVViYqL27NkTTTkAAKCbiuo247lz5yoUCmno0KFKSkpSa2urnn76aRUWFkqS/H6/JMntdkds53a77Ta/36+MjIzIIpKTlZ6ebvf5snA4rHA4bL8OhULRlA0AAOJMVDMor776qtauXat169bpvffe05o1a/TMM89ozZo1sapPklRWViaXy2UvPEUWAIDuLaqAMnv2bM2dO1cTJ07U8OHDNWnSJM2aNUtlZWWSJI/HI0kKBAIR2wUCAbvN4/GooaEhor2lpUWnTp2y+3xZaWmpgsGgvdTX10dTNgAAiDNRBZQvvvjigh/3SUpKUltbmyRpyJAh8ng8qqmpsdtDoZD27Nkjr9crSfJ6vWpsbFRtba3dZ/v27Wpra1NOTs5Fj+twOOynxvL0WAAAur+orkG599579fTTT2vgwIG66aabtH//fj377LN65JFHJP3/o2tnzpypp556Stdff72GDBmi+fPnKzMzU/fff78kadiwYbr77rs1ZcoUVVRUqLm5WdOnT9fEiRMv6Q4eAADQ/UUVUF544QXNnz9fjz76qBoaGpSZmamf/vSnWrBggd3n8ccf19mzZzV16lQ1Njbq9ttv19atW9WzZ0+7z9q1azV9+nTdddddSkxMVEFBgZYvX95xZwUAAOJaVM9BMQXPQQEAIP7E7DkoAAAAnSGqr3gAxJ/BczfFbN+fLMuP2b4BXNmYQQEAAMYhoAAAAOMQUAAAgHEIKAAAwDgEFAAAYBwCCgAAMA4BBQAAGIeAAgAAjENAAQAAxiGgAAAA4xBQAACAcQgoAADAOAQUAABgHAIKAAAwDgEFAAAYh4ACAACMQ0ABAADGIaAAAADjEFAAAIBxCCgAAMA4BBQAAGAcAgoAADAOAQUAABiHgAIAAIxDQAEAAMYhoAAAAOMQUAAAgHEIKAAAwDgEFAAAYBwCCgAAMA4BBQAAGCe5qwsAEL8Gz90Uk/1+siw/JvsFED+imkEZPHiwEhISLliKi4slSefOnVNxcbH69eunPn36qKCgQIFAIGIfdXV1ys/PV69evZSRkaHZs2erpaWl484IAADEvagCyr59+3Ty5El7qa6uliR9//vflyTNmjVLr7/+ujZs2KAdO3boxIkTmjBhgr19a2ur8vPz1dTUpF27dmnNmjWqrKzUggULOvCUAABAvEuwLMtq78YzZ85UVVWVPvroI4VCIfXv31/r1q3TAw88IEn64IMPNGzYMPl8Po0ePVpbtmzRPffcoxMnTsjtdkuSKioqNGfOHH322WdKSUm5pOOGQiG5XC4Fg0E5nc72lg9cEWL1NUws8RUP0D1F8/nd7otkm5qa9Morr+iRRx5RQkKCamtr1dzcrNzcXLvP0KFDNXDgQPl8PkmSz+fT8OHD7XAiSXl5eQqFQjp06FB7SwEAAN1Muy+Sfe2119TY2KiHH35YkuT3+5WSkqK0tLSIfm63W36/3+7zh+HkfPv5tq8SDocVDoft16FQqL1lAwCAONDuGZSXX35Z48aNU2ZmZkfWc1FlZWVyuVz2kpWVFfNjAgCArtOugPLb3/5W27Zt009+8hN7ncfjUVNTkxobGyP6BgIBeTweu8+X7+o5//p8n4spLS1VMBi0l/r6+vaUDQAA4kS7Asrq1auVkZGh/PzfX8g2cuRI9ejRQzU1Nfa6o0ePqq6uTl6vV5Lk9Xp18OBBNTQ02H2qq6vldDqVnZ39lcdzOBxyOp0RCwAA6L6ivgalra1Nq1evVlFRkZKTf7+5y+XS5MmTVVJSovT0dDmdTs2YMUNer1ejR4+WJI0dO1bZ2dmaNGmSysvL5ff7NW/ePBUXF8vhcHTcWQEAgLgWdUDZtm2b6urq9Mgjj1zQ9txzzykxMVEFBQUKh8PKy8vTypUr7fakpCRVVVVp2rRp8nq96t27t4qKirR48eLLOwsAANCtXNZzULoKz0EBLh3PQQFgik55DgoAAECsEFAAAIBxCCgAAMA4BBQAAGAcAgoAADAOAQUAABiHgAIAAIxDQAEAAMYhoAAAAOMQUAAAgHEIKAAAwDgEFAAAYBwCCgAAMA4BBQAAGIeAAgAAjENAAQAAxiGgAAAA4xBQAACAcQgoAADAOAQUAABgHAIKAAAwDgEFAAAYh4ACAACMQ0ABAADGIaAAAADjEFAAAIBxCCgAAMA4BBQAAGAcAgoAADAOAQUAABiHgAIAAIxDQAEAAMYhoAAAAOMQUAAAgHGiDij//d//rb/8y79Uv379lJqaquHDh+vdd9+12y3L0oIFCzRgwAClpqYqNzdXH330UcQ+Tp06pcLCQjmdTqWlpWny5Mk6c+bM5Z8NAADoFqIKKL/73e80ZswY9ejRQ1u2bNHhw4f1d3/3d7rqqqvsPuXl5Vq+fLkqKiq0Z88e9e7dW3l5eTp37pzdp7CwUIcOHVJ1dbWqqqq0c+dOTZ06tePOCgAAxLUEy7KsS+08d+5cvfPOO/rP//zPi7ZblqXMzEz9/Oc/12OPPSZJCgaDcrvdqqys1MSJE3XkyBFlZ2dr3759GjVqlCRp69atGj9+vD799FNlZmZ+bR2hUEgul0vBYFBOp/NSyweuSIPnburqEqL2ybL8ri4BQAxE8/kd1QzKr3/9a40aNUrf//73lZGRoW9961t66aWX7Pbjx4/L7/crNzfXXudyuZSTkyOfzydJ8vl8SktLs8OJJOXm5ioxMVF79uyJphwAANBNRRVQjh07plWrVun666/XG2+8oWnTpumv/uqvtGbNGkmS3++XJLnd7ojt3G633eb3+5WRkRHRnpycrPT0dLvPl4XDYYVCoYgFAAB0X8nRdG5ra9OoUaO0dOlSSdK3vvUtvf/++6qoqFBRUVFMCpSksrIyLVq0KGb7BwAAZolqBmXAgAHKzs6OWDds2DDV1dVJkjwejyQpEAhE9AkEAnabx+NRQ0NDRHtLS4tOnTpl9/my0tJSBYNBe6mvr4+mbAAAEGeiCihjxozR0aNHI9Z9+OGHGjRokCRpyJAh8ng8qqmpsdtDoZD27Nkjr9crSfJ6vWpsbFRtba3dZ/v27Wpra1NOTs5Fj+twOOR0OiMWAADQfUX1Fc+sWbN02223aenSpXrwwQe1d+9evfjii3rxxRclSQkJCZo5c6aeeuopXX/99RoyZIjmz5+vzMxM3X///ZL+f8bl7rvv1pQpU1RRUaHm5mZNnz5dEydOvKQ7eAAAQPcXVUD5zne+o40bN6q0tFSLFy/WkCFD9Pzzz6uwsNDu8/jjj+vs2bOaOnWqGhsbdfvtt2vr1q3q2bOn3Wft2rWaPn267rrrLiUmJqqgoEDLly/vuLMCAABxLarnoJiC56AAl47noAAwRcyegwIAANAZCCgAAMA4BBQAAGAcAgoAADAOAQUAABiHgAIAAIxDQAEAAMYhoAAAAOMQUAAAgHEIKAAAwDgEFAAAYBwCCgAAMA4BBQAAGIeAAgAAjENAAQAAxiGgAAAA4xBQAACAcQgoAADAOAQUAABgHAIKAAAwDgEFAAAYh4ACAACMQ0ABAADGIaAAAADjEFAAAIBxCCgAAMA4BBQAAGAcAgoAADAOAQUAABiHgAIAAIxDQAEAAMYhoAAAAOMQUAAAgHEIKAAAwDhRBZQnn3xSCQkJEcvQoUPt9nPnzqm4uFj9+vVTnz59VFBQoEAgELGPuro65efnq1evXsrIyNDs2bPV0tLSMWcDAAC6heRoN7jpppu0bdu23+8g+fe7mDVrljZt2qQNGzbI5XJp+vTpmjBhgt555x1JUmtrq/Lz8+XxeLRr1y6dPHlSP/rRj9SjRw8tXbq0A04HAAB0B1EHlOTkZHk8ngvWB4NBvfzyy1q3bp3uvPNOSdLq1as1bNgw7d69W6NHj9abb76pw4cPa9u2bXK73brlllu0ZMkSzZkzR08++aRSUlIu/4wAAEDci/oalI8++kiZmZm69tprVVhYqLq6OklSbW2tmpublZuba/cdOnSoBg4cKJ/PJ0ny+XwaPny43G633ScvL0+hUEiHDh263HMBAADdRFQzKDk5OaqsrNSNN96okydPatGiRbrjjjv0/vvvy+/3KyUlRWlpaRHbuN1u+f1+SZLf748IJ+fbz7d9lXA4rHA4bL8OhULRlA0AAOJMVAFl3Lhx9j+PGDFCOTk5GjRokF599VWlpqZ2eHHnlZWVadGiRTHbPwAAMMtl3WaclpamG264QR9//LE8Ho+amprU2NgY0ScQCNjXrHg8ngvu6jn/+mLXtZxXWlqqYDBoL/X19ZdTNgAAMNxlBZQzZ87oN7/5jQYMGKCRI0eqR48eqqmpsduPHj2quro6eb1eSZLX69XBgwfV0NBg96murpbT6VR2dvZXHsfhcMjpdEYsAACg+4rqK57HHntM9957rwYNGqQTJ05o4cKFSkpK0g9+8AO5XC5NnjxZJSUlSk9Pl9Pp1IwZM+T1ejV69GhJ0tixY5Wdna1JkyapvLxcfr9f8+bNU3FxsRwOR0xOEAAAxJ+oAsqnn36qH/zgB/qf//kf9e/fX7fffrt2796t/v37S5Kee+45JSYmqqCgQOFwWHl5eVq5cqW9fVJSkqqqqjRt2jR5vV717t1bRUVFWrx4cceeFQAAiGsJlmVZXV1EtEKhkFwul4LBIF/3AF9j8NxNXV1C1D5Zlt/VJQCIgWg+v/ktHgAAYBwCCgAAMA4BBQAAGIeAAgAAjENAAQAAxiGgAAAA4xBQAACAcQgoAADAOAQUAABgHAIKAAAwDgEFAAAYh4ACAACMQ0ABAADGIaAAAADjEFAAAIBxCCgAAMA4BBQAAGAcAgoAADAOAQUAABiHgAIAAIxDQAEAAMYhoAAAAOMQUAAAgHEIKAAAwDgEFAAAYBwCCgAAMA4BBQAAGIeAAgAAjENAAQAAxiGgAAAA4xBQAACAcQgoAADAOAQUAABgHAIKAAAwzmUFlGXLlikhIUEzZ8601507d07FxcXq16+f+vTpo4KCAgUCgYjt6urqlJ+fr169eikjI0OzZ89WS0vL5ZQCAAC6kXYHlH379ukf//EfNWLEiIj1s2bN0uuvv64NGzZox44dOnHihCZMmGC3t7a2Kj8/X01NTdq1a5fWrFmjyspKLViwoP1nAQAAupV2BZQzZ86osLBQL730kq666ip7fTAY1Msvv6xnn31Wd955p0aOHKnVq1dr165d2r17tyTpzTff1OHDh/XKK6/olltu0bhx47RkyRKtWLFCTU1NHXNWAAAgrrUroBQXFys/P1+5ubkR62tra9Xc3ByxfujQoRo4cKB8Pp8kyefzafjw4XK73XafvLw8hUIhHTp06KLHC4fDCoVCEQsAAOi+kqPdYP369Xrvvfe0b9++C9r8fr9SUlKUlpYWsd7tdsvv99t9/jCcnG8/33YxZWVlWrRoUbSlAgCAOBXVDEp9fb3++q//WmvXrlXPnj1jVdMFSktLFQwG7aW+vr7Tjg0AADpfVAGltrZWDQ0N+va3v63k5GQlJydrx44dWr58uZKTk+V2u9XU1KTGxsaI7QKBgDwejyTJ4/FccFfP+dfn+3yZw+GQ0+mMWAAAQPcVVUC56667dPDgQR04cMBeRo0apcLCQvufe/TooZqaGnubo0ePqq6uTl6vV5Lk9Xp18OBBNTQ02H2qq6vldDqVnZ3dQacFAADiWVTXoPTt21c333xzxLrevXurX79+9vrJkyerpKRE6enpcjqdmjFjhrxer0aPHi1JGjt2rLKzszVp0iSVl5fL7/dr3rx5Ki4ulsPh6KDTAgAA8Szqi2S/znPPPafExEQVFBQoHA4rLy9PK1eutNuTkpJUVVWladOmyev1qnfv3ioqKtLixYs7uhQAABCnEizLsrq6iGiFQiG5XC4Fg0GuRwG+xuC5m7q6hKh9siy/q0sAEAPRfH7zWzwAAMA4BBQAAGAcAgoAADAOAQUAABiHgAIAAIxDQAEAAMYhoAAAAOMQUAAAgHEIKAAAwDgEFAAAYJwO/y0eAADQOWL5UxZd/ZMTzKAAAADjEFAAAIBxCCgAAMA4BBQAAGAcAgoAADAOAQUAABiHgAIAAIxDQAEAAMYhoAAAAOMQUAAAgHEIKAAAwDgEFAAAYBwCCgAAMA4BBQAAGIeAAgAAjENAAQAAxiGgAAAA4xBQAACAcQgoAADAOAQUAABgHAIKAAAwDgEFAAAYh4ACAACME1VAWbVqlUaMGCGn0ymn0ymv16stW7bY7efOnVNxcbH69eunPn36qKCgQIFAIGIfdXV1ys/PV69evZSRkaHZs2erpaWlY84GAAB0C1EFlGuuuUbLli1TbW2t3n33Xd1555267777dOjQIUnSrFmz9Prrr2vDhg3asWOHTpw4oQkTJtjbt7a2Kj8/X01NTdq1a5fWrFmjyspKLViwoGPPCgAAxLUEy7Ksy9lBenq6/vZv/1YPPPCA+vfvr3Xr1umBBx6QJH3wwQcaNmyYfD6fRo8erS1btuiee+7RiRMn5Ha7JUkVFRWaM2eOPvvsM6WkpFzSMUOhkFwul4LBoJxO5+WUD3R7g+du6uoSovbJsvyuLgGIC7F8f8fifRjN53e7r0FpbW3V+vXrdfbsWXm9XtXW1qq5uVm5ubl2n6FDh2rgwIHy+XySJJ/Pp+HDh9vhRJLy8vIUCoXsWZiLCYfDCoVCEQsAAOi+og4oBw8eVJ8+feRwOPSzn/1MGzduVHZ2tvx+v1JSUpSWlhbR3+12y+/3S5L8fn9EODnffr7tq5SVlcnlctlLVlZWtGUDAIA4EnVAufHGG3XgwAHt2bNH06ZNU1FRkQ4fPhyL2mylpaUKBoP2Ul9fH9PjAQCArpUc7QYpKSn6xje+IUkaOXKk9u3bp7//+7/XQw89pKamJjU2NkbMogQCAXk8HkmSx+PR3r17I/Z3/i6f830uxuFwyOFwRFsqAACIU5f9HJS2tjaFw2GNHDlSPXr0UE1Njd129OhR1dXVyev1SpK8Xq8OHjyohoYGu091dbWcTqeys7MvtxQAANBNRDWDUlpaqnHjxmngwIE6ffq01q1bp7feektvvPGGXC6XJk+erJKSEqWnp8vpdGrGjBnyer0aPXq0JGns2LHKzs7WpEmTVF5eLr/fr3nz5qm4uJgZEgAAYIsqoDQ0NOhHP/qRTp48KZfLpREjRuiNN97Qn//5n0uSnnvuOSUmJqqgoEDhcFh5eXlauXKlvX1SUpKqqqo0bdo0eb1e9e7dW0VFRVq8eHHHnhUAAIhrl/0clK7Ac1CAS8dzUIDui+egAAAAdCICCgAAMA4BBQAAGIeAAgAAjENAAQAAxiGgAAAA4xBQAACAcQgoAADAOAQUAABgHAIKAAAwDgEFAAAYh4ACAACMQ0ABAADGIaAAAADjEFAAAIBxCCgAAMA4BBQAAGAcAgoAADAOAQUAABiHgAIAAIxDQAEAAMYhoAAAAOMQUAAAgHEIKAAAwDgEFAAAYBwCCgAAMA4BBQAAGIeAAgAAjENAAQAAxiGgAAAA4xBQAACAcQgoAADAOAQUAABgHAIKAAAwTlQBpaysTN/5znfUt29fZWRk6P7779fRo0cj+pw7d07FxcXq16+f+vTpo4KCAgUCgYg+dXV1ys/PV69evZSRkaHZs2erpaXl8s8GAAB0C1EFlB07dqi4uFi7d+9WdXW1mpubNXbsWJ09e9buM2vWLL3++uvasGGDduzYoRMnTmjChAl2e2trq/Lz89XU1KRdu3ZpzZo1qqys1IIFCzrurAAAQFxLsCzLau/Gn332mTIyMrRjxw796Z/+qYLBoPr3769169bpgQcekCR98MEHGjZsmHw+n0aPHq0tW7bonnvu0YkTJ+R2uyVJFRUVmjNnjj777DOlpKR87XFDoZBcLpeCwaCcTmd7yweuCIPnburqEqL2ybL8ri4BiAuxfH/H4n0Yzef3ZV2DEgwGJUnp6emSpNraWjU3Nys3N9fuM3ToUA0cOFA+n0+S5PP5NHz4cDucSFJeXp5CoZAOHTp00eOEw2GFQqGIBQAAdF/tDihtbW2aOXOmxowZo5tvvlmS5Pf7lZKSorS0tIi+brdbfr/f7vOH4eR8+/m2iykrK5PL5bKXrKys9pYNAADiQLsDSnFxsd5//32tX7++I+u5qNLSUgWDQXupr6+P+TEBAEDXSW7PRtOnT1dVVZV27typa665xl7v8XjU1NSkxsbGiFmUQCAgj8dj99m7d2/E/s7f5XO+z5c5HA45HI72lAoAAOJQVDMolmVp+vTp2rhxo7Zv364hQ4ZEtI8cOVI9evRQTU2Nve7o0aOqq6uT1+uVJHm9Xh08eFANDQ12n+rqajmdTmVnZ1/OuQAAgG4iqhmU4uJirVu3Tv/+7/+uvn372teMuFwupaamyuVyafLkySopKVF6erqcTqdmzJghr9er0aNHS5LGjh2r7OxsTZo0SeXl5fL7/Zo3b56Ki4uZJQEAAJKiDCirVq2SJP3Zn/1ZxPrVq1fr4YcfliQ999xzSkxMVEFBgcLhsPLy8rRy5Uq7b1JSkqqqqjRt2jR5vV717t1bRUVFWrx48eWdCQAA6DaiCiiX8siUnj17asWKFVqxYsVX9hk0aJA2b94czaEBAMAVhN/iAQAAxiGgAAAA4xBQAACAcQgoAADAOAQUAABgHAIKAAAwDgEFAAAYh4ACAACMQ0ABAADGIaAAAADjEFAAAIBxCCgAAMA4BBQAAGAcAgoAADAOAQUAABiHgAIAAIyT3NUFAADQnQ2eu6mrS4hLzKAAAADjEFAAAIBxCCgAAMA4BBQAAGAcAgoAADAOAQUAABiHgAIAAIxDQAEAAMYhoAAAAOMQUAAAgHEIKAAAwDj8Fk8niuXvMXyyLD9m+wYAoLMRUAAD8GNiABCJr3gAAIBxCCgAAMA4BBQAAGAcAgoAADBO1AFl586duvfee5WZmamEhAS99tprEe2WZWnBggUaMGCAUlNTlZubq48++iiiz6lTp1RYWCin06m0tDRNnjxZZ86cuawTAQAA3UfUAeXs2bP65je/qRUrVly0vby8XMuXL1dFRYX27Nmj3r17Ky8vT+fOnbP7FBYW6tChQ6qurlZVVZV27typqVOntv8sAABAtxL1bcbjxo3TuHHjLtpmWZaef/55zZs3T/fdd58k6Z//+Z/ldrv12muvaeLEiTpy5Ii2bt2qffv2adSoUZKkF154QePHj9czzzyjzMzMyzgdAADQHXToNSjHjx+X3+9Xbm6uvc7lciknJ0c+n0+S5PP5lJaWZocTScrNzVViYqL27Nlz0f2Gw2GFQqGIBQAAdF8dGlD8fr8kye12R6x3u912m9/vV0ZGRkR7cnKy0tPT7T5fVlZWJpfLZS9ZWVkdWTYAADBMXNzFU1paqmAwaC/19fVdXRIAAIihDg0oHo9HkhQIBCLWBwIBu83j8aihoSGivaWlRadOnbL7fJnD4ZDT6YxYAABA99WhAWXIkCHyeDyqqamx14VCIe3Zs0der1eS5PV61djYqNraWrvP9u3b1dbWppycnI4sBwAAxKmo7+I5c+aMPv74Y/v18ePHdeDAAaWnp2vgwIGaOXOmnnrqKV1//fUaMmSI5s+fr8zMTN1///2SpGHDhunuu+/WlClTVFFRoebmZk2fPl0TJ07kDh4AACCpHQHl3Xff1fe+9z37dUlJiSSpqKhIlZWVevzxx3X27FlNnTpVjY2Nuv3227V161b17NnT3mbt2rWaPn267rrrLiUmJqqgoEDLly/vgNMBAADdQYJlWVZXFxGtUCgkl8ulYDAYV9ejDJ67KWb7/mRZfsz2jdiL5d8GIvFeQWeL1/d3LN4r0Xx+x8VdPAAA4MpCQAEAAMaJ+hoUmClWU4hMhwMAugIzKAAAwDjMoAAA4gY3G1w5mEEBAADGYQYFAADF7+3A3RUzKAAAwDgEFAAAYBwCCgAAMA7XoADAFYo7YmAyZlAAAIBxCCgAAMA4fMUDRIHbEAGgcxBQAAAdjjCPy0VAuQjeWL/HRXRA1+O/SbgSEVAAoAMQIoCOxUWyAADAOMygoMvE6v84+eoIfwwzHUB8YAYFAAAYh4ACAACMQ0ABAADG4RoUdDtcYwAA8Y8ZFAAAYBwCCgAAMA4BBQAAGIeAAgAAjENAAQAAxiGgAAAA4xBQAACAcQgoAADAOAQUAABgHAIKAAAwDgEFAAAYp0sDyooVKzR48GD17NlTOTk52rt3b1eWAwAADNFlAeVf//VfVVJSooULF+q9997TN7/5TeXl5amhoaGrSgIAAIbosoDy7LPPasqUKfrxj3+s7OxsVVRUqFevXvrlL3/ZVSUBAABDJHfFQZuamlRbW6vS0lJ7XWJionJzc+Xz+S7oHw6HFQ6H7dfBYFCSFAqFYlJfW/iLmOwXAIB4EYvP2PP7tCzra/t2SUD5/PPP1draKrfbHbHe7Xbrgw8+uKB/WVmZFi1adMH6rKysmNUIAMCVzPV87PZ9+vRpuVyuP9qnSwJKtEpLS1VSUmK/bmtr06lTp9SvXz8lJCR0yDFCoZCysrJUX18vp9PZIfvEhRjnzsNYdx7GunMwzp0nVmNtWZZOnz6tzMzMr+3bJQHl6quvVlJSkgKBQMT6QCAgj8dzQX+HwyGHwxGxLi0tLSa1OZ1O/vA7AePceRjrzsNYdw7GufPEYqy/bubkvC65SDYlJUUjR45UTU2Nva6trU01NTXyer1dURIAADBIl33FU1JSoqKiIo0aNUq33nqrnn/+eZ09e1Y//vGPu6okAABgiC4LKA899JA+++wzLViwQH6/X7fccou2bt16wYWzncXhcGjhwoUXfJWEjsU4dx7GuvMw1p2Dce48Jox1gnUp9/oAAAB0In6LBwAAGIeAAgAAjENAAQAAxiGgAAAA41xRAWXFihUaPHiwevbsqZycHO3du/eP9t+wYYOGDh2qnj17avjw4dq8eXMnVRrfohnnl156SXfccYeuuuoqXXXVVcrNzf3afy/4vWj/ps9bv369EhISdP/998e2wG4i2nFubGxUcXGxBgwYIIfDoRtuuIH/flyiaMf6+eef14033qjU1FRlZWVp1qxZOnfuXCdVG5927type++9V5mZmUpISNBrr732tdu89dZb+va3vy2Hw6FvfOMbqqysjHmdsq4Q69evt1JSUqxf/vKX1qFDh6wpU6ZYaWlpViAQuGj/d955x0pKSrLKy8utw4cPW/PmzbN69OhhHTx4sJMrjy/RjvMPf/hDa8WKFdb+/futI0eOWA8//LDlcrmsTz/9tJMrjz/RjvV5x48ft/7kT/7EuuOOO6z77ruvc4qNY9GOczgctkaNGmWNHz/eevvtt63jx49bb731lnXgwIFOrjz+RDvWa9eutRwOh7V27Vrr+PHj1htvvGENGDDAmjVrVidXHl82b95sPfHEE9avfvUrS5K1cePGP9r/2LFjVq9evaySkhLr8OHD1gsvvGAlJSVZW7dujWmdV0xAufXWW63i4mL7dWtrq5WZmWmVlZVdtP+DDz5o5efnR6zLycmxfvrTn8a0zngX7Th/WUtLi9W3b19rzZo1sSqx22jPWLe0tFi33Xab9U//9E9WUVERAeUSRDvOq1atsq699lqrqamps0rsNqId6+LiYuvOO++MWFdSUmKNGTMmpnV2J5cSUB5//HHrpptuilj30EMPWXl5eTGszLKuiK94mpqaVFtbq9zcXHtdYmKicnNz5fP5LrqNz+eL6C9JeXl5X9kf7RvnL/viiy/U3Nys9PT0WJXZLbR3rBcvXqyMjAxNnjy5M8qMe+0Z51//+tfyer0qLi6W2+3WzTffrKVLl6q1tbWzyo5L7Rnr2267TbW1tfbXQMeOHdPmzZs1fvz4Tqn5StFVn4dx8WvGl+vzzz9Xa2vrBU+pdbvd+uCDDy66jd/vv2h/v98fszrjXXvG+cvmzJmjzMzMC94MiNSesX777bf18ssv68CBA51QYffQnnE+duyYtm/frsLCQm3evFkff/yxHn30UTU3N2vhwoWdUXZcas9Y//CHP9Tnn3+u22+/XZZlqaWlRT/72c/0N3/zN51R8hXjqz4PQ6GQ/vd//1epqakxOe4VMYOC+LBs2TKtX79eGzduVM+ePbu6nG7l9OnTmjRpkl566SVdffXVXV1Ot9bW1qaMjAy9+OKLGjlypB566CE98cQTqqio6OrSup233npLS5cu1cqVK/Xee+/pV7/6lTZt2qQlS5Z0dWnoAFfEDMrVV1+tpKQkBQKBiPWBQEAej+ei23g8nqj6o33jfN4zzzyjZcuWadu2bRoxYkQsy+wWoh3r3/zmN/rkk09077332uva2tokScnJyTp69Kiuu+662BYdh9rzNz1gwAD16NFDSUlJ9rphw4bJ7/erqalJKSkpMa05XrVnrOfPn69JkybpJz/5iSRp+PDhOnv2rKZOnaonnnhCiYn8P3hH+KrPQ6fTGbPZE+kKmUFJSUnRyJEjVVNTY69ra2tTTU2NvF7vRbfxer0R/SWpurr6K/ujfeMsSeXl5VqyZIm2bt2qUaNGdUapcS/asR46dKgOHjyoAwcO2Mtf/MVf6Hvf+54OHDigrKysziw/brTnb3rMmDH6+OOP7QAoSR9++KEGDBhAOPkj2jPWX3zxxQUh5HwwtPiZuQ7TZZ+HMb0E1yDr16+3HA6HVVlZaR0+fNiaOnWqlZaWZvn9fsuyLGvSpEnW3Llz7f7vvPOOlZycbD3zzDPWkSNHrIULF3Kb8SWIdpyXLVtmpaSkWP/2b/9mnTx50l5Onz7dVacQN6Id6y/jLp5LE+0419XVWX379rWmT59uHT161KqqqrIyMjKsp556qqtOIW5EO9YLFy60+vbta/3Lv/yLdezYMevNN9+0rrvuOuvBBx/sqlOIC6dPn7b2799v7d+/35JkPfvss9b+/fut3/72t5ZlWdbcuXOtSZMm2f3P32Y8e/Zs68iRI9aKFSu4zbijvfDCC9bAgQOtlJQU69Zbb7V2795tt333u9+1ioqKIvq/+uqr1g033GClpKRYN910k7Vp06ZOrjg+RTPOgwYNsiRdsCxcuLDzC49D0f5N/yECyqWLdpx37dpl5eTkWA6Hw7r22mutp59+2mppaenkquNTNGPd3NxsPfnkk9Z1111n9ezZ08rKyrIeffRR63e/+13nFx5H/uM//uOi/909P7ZFRUXWd7/73Qu2ueWWW6yUlBTr2muvtVavXh3zOhMsi3kwAABgliviGhQAABBfCCgAAMA4BBQAAGAcAgoAADAOAQUAABiHgAIAAIxDQAEAAMYhoAAAAOMQUAAAgHEIKAAAwDgEFAAAYBwCCgAAMM7/Af3NGUggHVFaAAAAAElFTkSuQmCC",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from scipy.optimize import linear_sum_assignment\n",
    "import matplotlib.pyplot as plt\n",
    "#Dictionary Comparison\n",
    "smaller_dict_features, _ = smaller_dict.shape\n",
    "larger_dict_features, _ = larger_dict.shape\n",
    "larger_dict = larger_dict.to(device)\n",
    "# Hungary algorithm\n",
    "# Calculate all cosine similarities and store in a 2D array\n",
    "cos_sims = np.zeros((smaller_dict_features, larger_dict_features))\n",
    "for idx, vector in enumerate(smaller_dict):\n",
    "    cos_sims[idx] = torch.nn.functional.cosine_similarity(vector.to(device), larger_dict, dim=1).cpu().numpy()\n",
    "# Convert to a minimization problem\n",
    "cos_sims = 1 - cos_sims\n",
    "# Use the Hungarian algorithm to solve the assignment problem\n",
    "row_ind, col_ind = linear_sum_assignment(cos_sims)\n",
    "# Retrieve the max cosine similarities and corresponding indices\n",
    "max_cosine_similarities = 1 - cos_sims[row_ind, col_ind]\n",
    "\n",
    "# Get the indices of the max cosine similarities in descending order\n",
    "max_indices = np.argsort(max_cosine_similarities)[::-1]\n",
    "max_cosine_similarities[max_indices][:20]\n",
    "print((\"# of features above 0.9:\", (max_cosine_similarities > .9).sum()))\n",
    "# Plot histogram of max_cosine_similarities\n",
    "plt.hist(max_cosine_similarities, bins=20)\n",
    "plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Model activations & Dictionary Activations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Found cached dataset parquet (/home/mchorse/.cache/huggingface/datasets/NeelNanda___parquet/NeelNanda--pile-10k-72f566e9f7c464ab/0.0.0/2a3b91fbd88a2c90d1dbbb32b460cf621d31bd5b05b934492fdef7d8d6f236ec)\n",
      "Loading cached processed dataset at /home/mchorse/.cache/huggingface/datasets/NeelNanda___parquet/NeelNanda--pile-10k-72f566e9f7c464ab/0.0.0/2a3b91fbd88a2c90d1dbbb32b460cf621d31bd5b05b934492fdef7d8d6f236ec/cache-f5a27b62a97e2a58.arrow\n",
      "Loading cached processed dataset at /home/mchorse/.cache/huggingface/datasets/NeelNanda___parquet/NeelNanda--pile-10k-72f566e9f7c464ab/0.0.0/2a3b91fbd88a2c90d1dbbb32b460cf621d31bd5b05b934492fdef7d8d6f236ec/cache-631c582ca5f19f42.arrow\n",
      "Loading cached processed dataset at /home/mchorse/.cache/huggingface/datasets/NeelNanda___parquet/NeelNanda--pile-10k-72f566e9f7c464ab/0.0.0/2a3b91fbd88a2c90d1dbbb32b460cf621d31bd5b05b934492fdef7d8d6f236ec/cache-7b1c265e5dfd20cb.arrow\n"
     ]
    }
   ],
   "source": [
    "# Downnload dataset\n",
    "from datasets import Dataset, load_dataset\n",
    "dataset_name = \"NeelNanda/pile-10k\"\n",
    "token_amount= 40\n",
    "dataset = load_dataset(dataset_name, split=\"train\").map(\n",
    "    lambda x: model.tokenizer(x['text']),\n",
    "    batched=True,\n",
    ").filter(\n",
    "    lambda x: len(x['input_ids']) > token_amount\n",
    ").map(\n",
    "    lambda x: {'input_ids': x['input_ids'][:token_amount]}\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 309/309 [00:13<00:00, 23.64it/s]\n"
     ]
    }
   ],
   "source": [
    "# Now we can use the model to get the activations\n",
    "from torch.utils.data import DataLoader\n",
    "from tqdm.auto import tqdm\n",
    "from einops import rearrange\n",
    "# neurons = model.W_in.shape[-1]\n",
    "neurons = model.cfg.d_model\n",
    "datapoints = dataset.num_rows\n",
    "batch_size = 32\n",
    "neuron_activations = torch.zeros((datapoints*token_amount, neurons))\n",
    "dictionary_activations = torch.zeros((datapoints*token_amount, smaller_dict_features))\n",
    "smaller_auto_encoder = autoencoder\n",
    "smaller_auto_encoder.to_device(device)\n",
    "\n",
    "with torch.no_grad(), dataset.formatted_as(\"pt\"):\n",
    "    dl = DataLoader(dataset[\"input_ids\"], batch_size=batch_size)\n",
    "    for i, batch in enumerate(tqdm(dl)):\n",
    "        _, cache = model.run_with_cache(batch.to(device))\n",
    "        batched_neuron_activations = rearrange(cache[cache_name], \"b s n -> (b s) n\" )\n",
    "        neuron_activations[i*batch_size*token_amount:(i+1)*batch_size*token_amount,:] = batched_neuron_activations.cpu()\n",
    "        batched_dictionary_activations = smaller_auto_encoder.encode(batched_neuron_activations)\n",
    "        dictionary_activations[i*batch_size*token_amount:(i+1)*batch_size*token_amount,:] = batched_dictionary_activations.cpu()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(37.0351)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Sparsity: or features/token\n",
    "dictionary_activations[:100000].count_nonzero(dim=1).float().mean()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Feature Activation Analysis"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Useful Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "from circuitsvis.activations import text_neuron_activations\n",
    "# Get the activations for the best dict features\n",
    "def get_feature_datapoints(feature_index, dictionary_activations, dataset, k=10, setting=\"max\"):\n",
    "    best_feature_activations = dictionary_activations[:, feature_index]\n",
    "    # Sort the features by activation, get the indices\n",
    "    if setting==\"max\":\n",
    "        found_indices = torch.argsort(best_feature_activations, descending=True)[:k]\n",
    "    elif setting==\"uniform\":\n",
    "        min_value = torch.min(best_feature_activations)\n",
    "        max_value = torch.max(best_feature_activations)\n",
    "\n",
    "        # Define the number of bins\n",
    "        num_bins = k\n",
    "\n",
    "        # Calculate the bin boundaries as linear interpolation between min and max\n",
    "        bin_boundaries = torch.linspace(min_value, max_value, num_bins + 1)\n",
    "\n",
    "        # Assign each activation to its respective bin\n",
    "        bins = torch.bucketize(best_feature_activations, bin_boundaries)\n",
    "\n",
    "        # Initialize a list to store the sampled indices\n",
    "        sampled_indices = []\n",
    "\n",
    "        # Sample from each bin\n",
    "        for bin_idx in torch.unique(bins):\n",
    "            # Get the indices corresponding to the current bin\n",
    "            bin_indices = torch.nonzero(bins == bin_idx, as_tuple=False).squeeze(dim=1)\n",
    "            \n",
    "            # Randomly sample from the current bin\n",
    "            sampled_indices.extend(np.random.choice(bin_indices, size=1, replace=False))\n",
    "\n",
    "        # Convert the sampled indices to a PyTorch tensor & reverse order\n",
    "        found_indices = torch.tensor(sampled_indices).long().flip(dims=[0])\n",
    "    else: # random\n",
    "        # get nonzero indices\n",
    "        nonzero_indices = torch.nonzero(best_feature_activations)[:, 0]\n",
    "        # shuffle\n",
    "        shuffled_indices = nonzero_indices[torch.randperm(nonzero_indices.shape[0])]\n",
    "        found_indices = shuffled_indices[:k]\n",
    "    datapoint_indices =[np.unravel_index(i, (datapoints, token_amount)) for i in found_indices]\n",
    "    text_list = []\n",
    "    full_text = []\n",
    "    token_list = []\n",
    "    full_token_list = []\n",
    "    for md, s_ind in datapoint_indices:\n",
    "        md = int(md)\n",
    "        s_ind = int(s_ind)\n",
    "        full_tok = torch.tensor(dataset[md][\"input_ids\"])\n",
    "        full_text.append(model.tokenizer.decode(full_tok))\n",
    "        tok = dataset[md][\"input_ids\"][:s_ind+1]\n",
    "        text = model.tokenizer.decode(tok)\n",
    "        text_list.append(text)\n",
    "        token_list.append(tok)\n",
    "        full_token_list.append(full_tok)\n",
    "    return text_list, full_text, token_list, full_token_list\n",
    "\n",
    "def get_neuron_activation(token, feature, model, setting=\"dictionary_basis\"):\n",
    "    with torch.no_grad():\n",
    "        _, cache = model.run_with_cache(token.to(model.cfg.device))\n",
    "        neuron_act_batch = cache[cache_name]\n",
    "        if setting==\"dictionary_basis\":\n",
    "            neuron_act_batch = rearrange(neuron_act_batch, \"b s n -> (b s) n\" )\n",
    "            act = smaller_auto_encoder.encode(neuron_act_batch)\n",
    "            return act[:, feature].tolist()\n",
    "        else: # neuron/residual basis\n",
    "            return neuron_act_batch[0, :, feature].tolist()\n",
    "\n",
    "def ablate_text(text, feature, model, setting=\"dictionary_basis\"):\n",
    "    if isinstance(text, str):\n",
    "        text = [text]\n",
    "    display_text_list = []\n",
    "    activation_list = []\n",
    "    for t in text:\n",
    "        # Convert text into tokens\n",
    "        if isinstance(t, str): # If the text is a list of tokens\n",
    "            split_text = model.to_str_tokens(t, prepend_bos=False)\n",
    "            tokens = model.to_tokens(t, prepend_bos=False)\n",
    "        else: # t equals tokens\n",
    "            tokens = t\n",
    "            split_text = model.to_str_tokens(t, prepend_bos=False)\n",
    "        seq_size = tokens.shape[1]\n",
    "        if(seq_size == 1): # If the text is a single token, we can't ablate it\n",
    "            continue\n",
    "        original = get_neuron_activation(tokens, feature, model)[-1]\n",
    "        changed_activations = torch.zeros(seq_size, device=device).cpu()\n",
    "        for i in range(seq_size):\n",
    "            # Remove the i'th token from the input\n",
    "            ablated_tokens = torch.cat((tokens[:,:i], tokens[:,i+1:]), dim=1)\n",
    "            changed_activations[i] += get_neuron_activation(ablated_tokens, feature, model, setting)[-1]\n",
    "        changed_activations -= original\n",
    "        display_text_list += [x.replace('\\n', '\\\\newline') for x in split_text] + [\"\\n\"]\n",
    "        activation_list += changed_activations.tolist() + [0.0]\n",
    "    activation_list = torch.tensor(activation_list).reshape(-1,1,1)\n",
    "    return text_neuron_activations(tokens=display_text_list, activations=activation_list)\n",
    "\n",
    "def visualize_text(text, feature, model, setting=\"dictionary_basis\", max_activation = None):\n",
    "    if isinstance(text, str):\n",
    "        text = [text]\n",
    "    if isinstance(feature, int):\n",
    "        feature = [feature]\n",
    "    display_text_list = []\n",
    "    act_list = []\n",
    "    for t in text:\n",
    "        if isinstance(t, str): # If the text is a list of tokens\n",
    "            split_text = model.to_str_tokens(t, prepend_bos=False)\n",
    "            token = model.to_tokens(t, prepend_bos=False)\n",
    "        else: # t are tokens\n",
    "            token = t\n",
    "            split_text = model.to_str_tokens(t, prepend_bos=False)\n",
    "        for f in feature:\n",
    "            display_text_list += [x.replace('\\n', '\\\\newline') for x in split_text] + [\"\\n\"]\n",
    "            act_list += get_neuron_activation(token, f, model, setting) + [0.0]\n",
    "    act_list = torch.tensor(act_list).reshape(-1,1,1)\n",
    "    if(max_activation is not None):\n",
    "        act_list = torch.clamp(act_list, max=max_activation)\n",
    "    return text_neuron_activations(tokens=display_text_list, activations=act_list)\n",
    "# Ablate the feature direction of the tokens\n",
    "# token_list is a list of tokens, convert to tensor of shape (batch_size, seq_len)\n",
    "from einops import rearrange\n",
    "def ablate_feature_direction(tokens, feature, model, autoencoder):\n",
    "    def mlp_ablation_hook(value, hook):\n",
    "        # Rearrange to fit autoencoder\n",
    "        int_val = rearrange(value, 'b s h -> (b s) h')\n",
    "\n",
    "        # Run through the autoencoder\n",
    "        act = autoencoder.encode(int_val)\n",
    "        feature_to_ablate = feature # TODO: bring this out of the function\n",
    "\n",
    "        # Subtract value with feature direction*act_of_feature\n",
    "        dictionary_for_this_autoencoder = autoencoder.get_learned_dict()\n",
    "        feature_direction = torch.outer(act[:, feature_to_ablate].squeeze(), dictionary_for_this_autoencoder[feature_to_ablate].squeeze())\n",
    "        batch, seq_len, hidden_size = value.shape\n",
    "        feature_direction = rearrange(feature_direction, '(b s) h -> b s h', b=batch, s=seq_len)\n",
    "        value -= feature_direction\n",
    "        return value\n",
    "\n",
    "    return model.run_with_hooks(tokens, \n",
    "        fwd_hooks=[(\n",
    "            cache_name, \n",
    "            mlp_ablation_hook\n",
    "            )]\n",
    "        )\n",
    "def add_feature_direction(tokens, feature, model, autoencoder, scalar=1.0):\n",
    "    def residual_add_hook(value, hook):\n",
    "        feature_direction = autoencoder.decoder.weight[:, feature].squeeze()\n",
    "        value += scalar*feature_direction\n",
    "        return value\n",
    "\n",
    "    return model.run_with_hooks(tokens, \n",
    "        fwd_hooks=[(\n",
    "            cache_name,\n",
    "            residual_add_hook\n",
    "            )]\n",
    "        )\n",
    "def ablate_feature_direction_display(text, features=None, setting=\"true_tokens\", verbose=False):\n",
    "\n",
    "    if features==None:\n",
    "        features = torch.tensor([best_feature])\n",
    "    if isinstance(features, int):\n",
    "        features = torch.tensor([features])\n",
    "    if isinstance(features, list):\n",
    "        features = torch.tensor(features)\n",
    "    if isinstance(text, str):\n",
    "        text = [text]\n",
    "    text_list = []\n",
    "    logit_list = []\n",
    "    for t in text:\n",
    "        tokens = model.to_tokens(t, prepend_bos=False)\n",
    "        with torch.no_grad():\n",
    "            original_logits = model(tokens).log_softmax(-1).cpu()\n",
    "            ablated_logits = ablate_feature_direction(tokens, features, model, smaller_auto_encoder).log_softmax(-1).cpu()\n",
    "        diff_logits = ablated_logits  - original_logits# ablated > original -> negative diff\n",
    "        tokens = tokens.cpu()\n",
    "        if setting == \"true_tokens\":\n",
    "            split_text = model.to_str_tokens(t, prepend_bos=False)\n",
    "            gather_tokens = rearrange(tokens[:,1:], \"b s -> b s 1\") # TODO: verify this is correct\n",
    "            # Gather the logits for the true tokens\n",
    "            diff = rearrange(diff_logits[:, :-1].gather(-1,gather_tokens), \"b s n -> (b s n)\")\n",
    "        elif setting == \"max\":\n",
    "            # Negate the diff_logits to see which tokens have the largest effect on the neuron\n",
    "            val, ind = (-1*diff_logits).max(-1)\n",
    "            diff = rearrange(val[:, :-1], \"b s -> (b s)\")\n",
    "            diff*= -1 # Negate the values gathered\n",
    "            split_text = model.to_str_tokens(ind, prepend_bos=False)\n",
    "            gather_tokens = rearrange(ind[:,1:], \"1 s -> 1 s 1\")\n",
    "        split_text = split_text[1:] # Remove the first token since we're not predicting it\n",
    "        if(verbose):\n",
    "            text_list += [x.replace('\\n', '\\\\newline') for x in split_text] + [\"\\n\"]\n",
    "            text_list += [x.replace('\\n', '\\\\newline') for x in split_text] + [\"\\n\"]\n",
    "            orig = rearrange(original_logits[:, :-1].gather(-1, gather_tokens), \"b s n -> (b s n)\")\n",
    "            ablated = rearrange(ablated_logits[:, :-1].gather(-1, gather_tokens), \"b s n -> (b s n)\")\n",
    "            logit_list += orig.tolist() + [0.0]\n",
    "            logit_list += ablated.tolist() + [0.0]\n",
    "        text_list += [x.replace('\\n', '\\\\newline') for x in split_text] + [\"\\n\"]\n",
    "        logit_list += diff.tolist() + [0.0]\n",
    "    logit_list = torch.tensor(logit_list).reshape(-1,1,1)\n",
    "    if verbose:\n",
    "        print(f\"Max & Min logit-diff: {logit_list.max().item():.2f} & {logit_list.min().item():.2f}\")\n",
    "    return text_neuron_activations(tokens=text_list, activations=logit_list)\n",
    "def generate_text(input_text, num_tokens, model, autoencoder, feature, temperature=0.7, setting=\"add\", scalar=1.0):\n",
    "    # Convert input text to tokens\n",
    "    input_ids = model.tokenizer.encode(input_text, return_tensors='pt').to(device)\n",
    "\n",
    "    for _ in range(num_tokens):\n",
    "        # Generate logits\n",
    "        with torch.no_grad():\n",
    "            if(setting==\"add\"):\n",
    "                logits = add_feature_direction(input_ids, feature, model, autoencoder, scalar=scalar)\n",
    "            else:\n",
    "                logits = model(input_ids)\n",
    "\n",
    "        # Apply temperature\n",
    "        logits = logits / temperature\n",
    "\n",
    "        # Sample from the distribution\n",
    "        probs = torch.nn.functional.softmax(logits[:, -1, :], dim=-1)\n",
    "        predicted_token = torch.multinomial(probs, num_samples=1)\n",
    "\n",
    "        # Append predicted token to input_ids\n",
    "        input_ids = torch.cat((input_ids, predicted_token), dim=-1)\n",
    "\n",
    "    # Decode the tokens to text\n",
    "    output_text = model.tokenizer.decode(input_ids[0])\n",
    "\n",
    "    return output_text\n",
    "\n",
    "# Logit Lens\n",
    "def logit_lens(model, best_feature, smaller_dict, layer):\n",
    "    with torch.no_grad():\n",
    "        # There are never-used tokens, which have high norm. We want to ignore these.\n",
    "        bad_ind = (model.W_U.norm(dim=0) > 20)\n",
    "        feature_direction = smaller_dict[best_feature].to(device)\n",
    "        # feature_direction = torch.matmul(feature_direction, model.W_out[layer]) # if MLP\n",
    "        logits = torch.matmul(feature_direction, model.W_U).cpu()\n",
    "    # Don't include bad indices\n",
    "    logits[bad_ind] = -1000\n",
    "    topk_values, topk_indices = torch.topk(logits, 20)\n",
    "    top_text = model.to_str_tokens(topk_indices)\n",
    "    print(f\"{top_text}\")\n",
    "    print(topk_values)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Feature Search\n",
    "Find Features that activate on the last token activation! Plug those into \"Best Feature\" below!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Activations: [7.89, 4.1, 2.38, 2.22, 2.16, 1.35, 0.75, 0.73, 0.52, 0.49]\n",
      "Feature_ids [920, 2709, 2987, 1384, 2468, 726, 235, 1550, 3042, 2935]\n"
     ]
    }
   ],
   "source": [
    "# t = \" I do like a\"\n",
    "# t = \" He had a first one (1), and then a second (2\"\n",
    "t = \"Das Berghain ist ein bekannter Nachtclub in Berlin, der für\"\n",
    "split_text = model.to_str_tokens(t, prepend_bos=False)\n",
    "token = model.to_tokens(t, prepend_bos=False)\n",
    "_, cache = model.run_with_cache(token.to(model.cfg.device))\n",
    "neuron_act_batch = cache[cache_name]\n",
    "act = smaller_auto_encoder.encode(neuron_act_batch.squeeze())\n",
    "# neg = \" He had a first one (1), and then a second 2\"\n",
    "# split_text = model.to_str_tokens(neg, prepend_bos=False)\n",
    "# token = model.to_tokens(neg, prepend_bos=False)\n",
    "# _, cache = model.run_with_cache(token.to(model.cfg.device))\n",
    "# neuron_act_batch = cache[cache_name]\n",
    "# act[-1, :] -= smaller_auto_encoder.encode(neuron_act_batch.squeeze())[-1,:]\n",
    "v, i = act[-1, :].topk(10)\n",
    "\n",
    "print(\"Activations:\",[round(val,2) for val in v.tolist()])\n",
    "print(\"Feature_ids\", i.tolist())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Feature Interp\n",
    "Investigate the example sentences the activate this feature.\n",
    "\n",
    "Max: show max activating (tokens,contexts)\n",
    "\n",
    "Uniform: Show range of activations from each bin (e.g. sample an example from 1-2, 2-3, etc). \n",
    "[Note: if a feature is monosemantic, then the full range of activations should be that feature, not just max-activating ones]\n",
    "\n",
    "Full_text: shows the full text example\n",
    "\n",
    "Text_list: shows up to the most activating example (try w/ max activating on a couple of examples to see)\n",
    "\n",
    "ablate_text: remove the context one token at a time, and show the decrease/increase in activation of that feature\n",
    "\n",
    "ablate_feature_direction: removes feature direction from model's activation mid-inference, showing the logit diff in the output for every token.\n",
    "\n",
    "logit_lens: show the logit lens for that feature. If matches ablate_feature_direction, then the computation path is through the residual stream, else, it's through future layers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Feature index: 920\n",
      "MCS: 0.9888160228729248\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div id=\"circuits-vis-d8f3b345-239b\" style=\"margin: 15px 0;\"/>\n",
       "    <script crossorigin type=\"module\">\n",
       "    import { render, TextNeuronActivations } from \"https://unpkg.com/circuitsvis@1.40.1/dist/cdn/esm.js\";\n",
       "    render(\n",
       "      \"circuits-vis-d8f3b345-239b\",\n",
       "      TextNeuronActivations,\n",
       "      {\"tokens\": [\"In\", \" O\", \"er\", \"-\", \"Er\", \"k\", \"ensch\", \"wick\", \" hat\", \" es\", \" an\", \" einer\", \"\\n\", \"Die\", \" Bundes\", \"reg\", \"ierung\", \" will\", \" T\", \"aus\", \"ende\", \" As\", \"yl\", \"b\", \"ew\", \"er\", \"ber\", \" nach\", \" Afghanistan\", \" abs\", \"chie\", \"ben\", \".\", \" Ein\", \"ige\", \"\\n\", \"Il\", \"leg\", \"ale\", \" Ein\", \"w\", \"ande\", \"rer\", \" soll\", \"ten\", \" fest\", \"ge\", \"hal\", \"ten\", \" und\", \" ab\", \"ges\", \"ch\", \"ob\", \"en\", \" werden\", \"\\n\", \"Q\", \":\", \"\\\\newline\", \"\\\\newline\", \"W\", \"ie\", \" sp\", \"richt\", \" man\", \" von\", \" der\", \" deut\", \"schen\", \"\\n\", \"W\", \"ie\", \" bet\", \"re\", \"ib\", \"t\", \" man\", \" die\", \" D\", \"ek\", \"ol\", \"on\", \"is\", \"ierung\", \" einer\", \" F\", \"ach\", \"ze\", \"its\", \"chr\", \"ift\", \"?\", \" Der\", \" Che\", \"fred\", \"ak\", \"te\", \"ur\", \" der\", \" \\u201e\", \"American\", \" Historical\", \" Review\", \"\\u201c\", \" er\", \"\\n\", \"----------\", \" Forward\", \"ed\", \" message\", \" ----------\", \"\\\\newline\", \"From\", \":\", \" Google\", \" Group\", \" Inc\", \".\", \" <\", \"sch\", \"ac\", \"hen\", \"\\n\", \"Wa\", \"\\u00df\", \"\\n\", \"An\", \"ze\", \"ige\", \"\\\\newline\", \"\\\\newline\", \"Das\", \" B\", \"Af\", \"\\u00f6\", \"\\n\", \"Fried\", \"rich\", \" Sch\", \"\\u00fc\", \"tz\", \"\\\\newline\", \"\\\\newline\", \"Fried\", \"\\n\", \"D\", \"arius\", \"z\", \" Kos\", \"zyk\", \"owski\", \"\\\\newline\", \"\\\\newline\", \"D\", \"arius\", \"z\", \" Kos\", \"zyk\", \"\\n\", \"[\", \"Sign\", \"ificance\", \" of\", \" serum\", \" CD\", \"62\", \"p\", \" and\", \" CD\", \"\\n\"], \"activations\": [[[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[2.761497974395752]], [[2.556185245513916]], [[6.779508113861084]], [[6.402298450469971]], [[6.393479824066162]], [[9.202953338623047]], [[0.0]], [[0.0]], [[2.8261609077453613]], [[2.8808608055114746]], [[4.256248474121094]], [[4.907643795013428]], [[4.657352924346924]], [[4.431193828582764]], [[7.116590976715088]], [[3.074345111846924]], [[4.69190788269043]], [[3.887798309326172]], [[3.1577510833740234]], [[3.822205066680908]], [[6.355565547943115]], [[7.044996738433838]], [[4.094240665435791]], [[5.3558244705200195]], [[3.9581379890441895]], [[6.04323148727417]], [[0.0]], [[6.969540119171143]], [[7.96804666519165]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[2.4952797889709473]], [[2.3113036155700684]], [[2.9278173446655273]], [[4.072571277618408]], [[6.238897323608398]], [[8.556170463562012]], [[5.919338703155518]], [[5.079263210296631]], [[5.494926452636719]], [[6.613473415374756]], [[7.357731342315674]], [[5.631309986114502]], [[4.43295431137085]], [[2.4008240699768066]], [[4.086171627044678]], [[5.165167331695557]], [[6.983290195465088]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[5.8058366775512695]], [[3.8297910690307617]], [[5.3074212074279785]], [[5.48158073425293]], [[4.259832382202148]], [[4.862842082977295]], [[2.689751148223877]], [[5.608889102935791]], [[0.0]], [[0.0]], [[1.7032147645950317]], [[3.8614373207092285]], [[5.3016252517700195]], [[4.0391926765441895]], [[5.155192852020264]], [[5.059718608856201]], [[6.404574871063232]], [[3.6709961891174316]], [[2.8125243186950684]], [[1.0952092409133911]], [[2.739896774291992]], [[2.440690517425537]], [[6.269233703613281]], [[6.910330295562744]], [[3.2476134300231934]], [[4.677999973297119]], [[3.7331504821777344]], [[4.047740936279297]], [[2.406341552734375]], [[6.365929126739502]], [[4.197783946990967]], [[6.945300102233887]], [[2.6104397773742676]], [[2.691013813018799]], [[1.7540613412857056]], [[2.5533275604248047]], [[5.964476585388184]], [[7.070001125335693]], [[3.487975597381592]], [[0.0]], [[0.0]], [[0.0]], [[5.235105991363525]], [[5.276565074920654]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.04280674457550049]], [[0.0]], [[3.8751773834228516]], [[0.0]], [[0.0]], [[3.292965888977051]], [[0.0]], [[0.0]], [[0.0]], [[4.694018363952637]], [[0.0]], [[2.4566402435302734]], [[6.074641227722168]], [[3.7707462310791016]], [[1.6781708002090454]], [[2.516864776611328]], [[0.0]], [[0.0]], [[0.6254352331161499]], [[1.971948266029358]], [[1.5496550798416138]], [[1.5544191598892212]], [[0.0]], [[1.5334964990615845]], [[1.0842658281326294]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.040796875953674316]], [[0.0]], [[0.0]], [[0.10650455951690674]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]]], \"firstDimensionName\": \"Layer\", \"secondDimensionName\": \"Neuron\"}\n",
       "    )\n",
       "    </script>"
      ],
      "text/plain": [
       "<circuitsvis.utils.render.RenderedHTML at 0x7f0e62e6b370>"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# N is sorted by MCS, so maybe they're cool and good, but\n",
    "# N = 106 # 106\n",
    "# best_feature = int(max_indices[N])\n",
    "\n",
    "# German ones?: 920, 2709, 2987 (maybe just 920)\n",
    "best_feature = 920 # Change this one for global index (N is sorted by MCS)\n",
    "\n",
    "print(f\"Feature index: {best_feature}\")\n",
    "print(f\"MCS: {max_cosine_similarities[best_feature]}\")\n",
    "text_list, full_text, token_list, full_token_list = get_feature_datapoints(best_feature, dictionary_activations, dataset, setting=\"uniform\")\n",
    "# text_list, full_text, token_list, full_token_list = get_feature_datapoints(best_feature, dictionary_activations, dataset, setting=\"max\")\n",
    "visualize_text(text_list, best_feature, model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div id=\"circuits-vis-65b156b8-e8c6\" style=\"margin: 15px 0;\"/>\n",
       "    <script crossorigin type=\"module\">\n",
       "    import { render, TextNeuronActivations } from \"https://unpkg.com/circuitsvis@1.40.1/dist/cdn/esm.js\";\n",
       "    render(\n",
       "      \"circuits-vis-65b156b8-e8c6\",\n",
       "      TextNeuronActivations,\n",
       "      {\"tokens\": [\"Die\", \" Bundes\", \"reg\", \"ierung\", \" will\", \" T\", \"aus\", \"ende\", \" As\", \"yl\", \"b\", \"ew\", \"er\", \"ber\", \" nach\", \" Afghanistan\", \" abs\", \"chie\", \"ben\", \".\", \" Ein\", \"ige\", \" von\", \" ihn\", \"en\", \" bez\", \"icht\", \"igen\", \" sich\", \" des\", \"hal\", \"b\", \" nun\", \" sel\", \"bst\", \" als\", \" Terror\", \"isten\", \" \\u2013\", \" und\", \"\\n\", \"W\", \"ie\", \" bet\", \"re\", \"ib\", \"t\", \" man\", \" die\", \" D\", \"ek\", \"ol\", \"on\", \"is\", \"ierung\", \" einer\", \" F\", \"ach\", \"ze\", \"its\", \"chr\", \"ift\", \"?\", \" Der\", \" Che\", \"fred\", \"ak\", \"te\", \"ur\", \" der\", \" \\u201e\", \"American\", \" Historical\", \" Review\", \"\\u201c\", \" er\", \"kl\", \"\\u00e4\", \"rt\", \",\", \" wie\", \"\\n\", \"Die\", \" \\u201e\", \"Sch\", \"reck\", \"ens\", \"n\", \"acht\", \"\\u201c\", \" von\", \" Im\", \"st\", \"\\\\newline\", \"\\\\newline\", \"En\", \"de\", \" April\", \" 1938\", \" haben\", \"\\n\", \"\\\\newline\", \"\\\\newline\", \"In\", \" T\", \"ans\", \"ania\", \" ab\", \" so\", \"fort\", \" ver\", \"bot\", \"en\", \":\", \" Gle\", \"it\", \"gel\", \"\\\\newline\", \"\\\\newline\", \"23\", \".\", \" Jul\", \"i\", \" 2016\", \",\", \" 16\", \":\", \"33\", \"h\", \",\", \"\\\\newline\", \"\\\\newline\", \"F\", \"\\u00fcr\", \" die\", \" Reg\", \"ierung\", \"\\n\", \"Q\", \":\", \"\\\\newline\", \"\\\\newline\", \"W\", \"ie\", \" sp\", \"richt\", \" man\", \" von\", \" der\", \" deut\", \"schen\", \" Spr\", \"ache\", \"?\", \"\\\\newline\", \"\\\\newline\", \"Es\", \" k\", \"ling\", \"t\", \" fast\", \"\\n\", \"Die\", \" Bundes\", \"reg\", \"ierung\", \" will\", \" T\", \"aus\", \"ende\", \" As\", \"yl\", \"b\", \"ew\", \"er\", \"ber\", \" nach\", \" Afghanistan\", \" abs\", \"chie\", \"ben\", \".\", \" Ein\", \"ige\", \" von\", \" ihn\", \"en\", \" bez\", \"icht\", \"igen\", \" sich\", \" des\", \"hal\", \"b\", \" nun\", \" sel\", \"bst\", \" als\", \" Terror\", \"\\n\", \"He\", \"i\", \",\", \" den\", \"ne\", \"\\n\", \"Da\", \" den\", \" f\", \"\\n\", \"K\", \"alt\", \"br\", \"\\n\", \"P\", \"rene\", \"z\", \" note\", \" que\", \" cet\", \" article\", \" pub\", \"li\", \"\\u00e9\", \" en\", \" 2016\", \" pour\", \"rait\", \"\\n\", \"Ox\", \"id\", \"ative\", \" stress\", \" in\", \" stable\", \" cystic\", \" fibrosis\", \" patients\", \":\", \" do\", \" we\", \" need\", \" higher\", \" antioxidant\", \" plasma\", \" levels\", \"?\", \"\\n\"], \"activations\": [[[-0.023637771606445312]], [[-0.051506996154785156]], [[0.04055213928222656]], [[0.0015716552734375]], [[-0.07532405853271484]], [[-0.02707958221435547]], [[-0.03195476531982422]], [[0.014710426330566406]], [[-0.03145408630371094]], [[-0.022602081298828125]], [[-0.048148155212402344]], [[-0.007037162780761719]], [[0.03299140930175781]], [[-0.014740943908691406]], [[-0.18398284912109375]], [[0.030517578125]], [[-0.009102821350097656]], [[0.032296180725097656]], [[0.0013341903686523438]], [[-0.43611812591552734]], [[-0.03895092010498047]], [[0.17428970336914062]], [[0.1356954574584961]], [[-0.13155746459960938]], [[0.10892772674560547]], [[-0.2001638412475586]], [[-0.051502227783203125]], [[-0.09539413452148438]], [[-0.05119037628173828]], [[-0.023145675659179688]], [[0.000335693359375]], [[0.09727859497070312]], [[-0.13712406158447266]], [[-0.28638172149658203]], [[-0.16753292083740234]], [[0.06489753723144531]], [[-0.3724374771118164]], [[0.12169933319091797]], [[-1.7097086906433105]], [[-1.577406406402588]], [[0.0]], [[-0.09592056274414062]], [[-0.061183929443359375]], [[0.07899665832519531]], [[0.0314178466796875]], [[-0.032283782958984375]], [[0.05222320556640625]], [[0.19234466552734375]], [[-0.2124309539794922]], [[-0.09919357299804688]], [[-0.1547985076904297]], [[-0.01975536346435547]], [[-0.04128074645996094]], [[0.0006685256958007812]], [[-0.33336639404296875]], [[-0.5867748260498047]], [[-0.060516357421875]], [[-0.12253665924072266]], [[0.003589630126953125]], [[-0.030580520629882812]], [[-0.24455642700195312]], [[-0.31314563751220703]], [[0.0756235122680664]], [[-0.14393043518066406]], [[-0.04239177703857422]], [[-0.012620925903320312]], [[-0.03628826141357422]], [[-0.030408859252929688]], [[-0.05785369873046875]], [[-0.035119056701660156]], [[-0.35732460021972656]], [[0.031007766723632812]], [[0.060821533203125]], [[0.05432319641113281]], [[0.12416839599609375]], [[0.5324392318725586]], [[0.6094355583190918]], [[-0.022995948791503906]], [[0.033184051513671875]], [[0.025758743286132812]], [[-0.32021045684814453]], [[0.0]], [[0.3095884323120117]], [[-0.016898632049560547]], [[-0.1406879425048828]], [[0.016869068145751953]], [[0.010786056518554688]], [[-0.15294790267944336]], [[-0.017946720123291016]], [[-0.2882261276245117]], [[0.2579636573791504]], [[-0.0800313949584961]], [[0.03359508514404297]], [[0.42479991912841797]], [[0.42479991912841797]], [[1.0462865829467773]], [[0.6097078323364258]], [[-0.19437599182128906]], [[-0.5789093971252441]], [[-4.481317043304443]], [[0.0]], [[0.23646163940429688]], [[0.23646163940429688]], [[-0.016844749450683594]], [[-0.024759769439697266]], [[-0.024598121643066406]], [[0.08219051361083984]], [[-0.016562461853027344]], [[-0.025472640991210938]], [[0.0011744499206542969]], [[-0.11406421661376953]], [[-0.1020956039428711]], [[0.015277862548828125]], [[-0.14716672897338867]], [[-0.02135181427001953]], [[-0.01723480224609375]], [[0.013921737670898438]], [[-0.012501716613769531]], [[-0.012501716613769531]], [[0.03267955780029297]], [[-0.037877559661865234]], [[-0.08435249328613281]], [[0.017656326293945312]], [[-0.037520408630371094]], [[0.07842636108398438]], [[-0.025156497955322266]], [[0.0642848014831543]], [[-0.023447036743164062]], [[-0.03663301467895508]], [[0.04711627960205078]], [[-0.11603260040283203]], [[-0.11603260040283203]], [[-0.8972272872924805]], [[-0.6904926300048828]], [[-0.5606675148010254]], [[1.393712043762207]], [[-2.6099729537963867]], [[0.0]], [[0.8881540298461914]], [[0.31980323791503906]], [[-0.3688349723815918]], [[-0.3688349723815918]], [[-0.5791721343994141]], [[-0.5661487579345703]], [[-0.38187313079833984]], [[-0.7041912078857422]], [[0.2030339241027832]], [[-0.22793340682983398]], [[-0.5695877075195312]], [[-0.2479419708251953]], [[-0.4424715042114258]], [[-0.09918880462646484]], [[-0.11435842514038086]], [[-0.04607677459716797]], [[-0.22595787048339844]], [[-0.22595787048339844]], [[-1.1960477828979492]], [[-0.24495267868041992]], [[-0.575014591217041]], [[0.8008108139038086]], [[1.9008417129516602]], [[0.0]], [[0.06685256958007812]], [[0.057578086853027344]], [[0.05002927780151367]], [[0.005375862121582031]], [[-0.013236045837402344]], [[-0.04496192932128906]], [[0.0046100616455078125]], [[0.02403736114501953]], [[0.03704071044921875]], [[-0.011187076568603516]], [[-0.031964778900146484]], [[-0.06307029724121094]], [[-0.00243377685546875]], [[-0.019817829132080078]], [[-0.09124088287353516]], [[0.11089086532592773]], [[0.027091026306152344]], [[0.0068759918212890625]], [[-0.007462501525878906]], [[0.2672882080078125]], [[-0.1105194091796875]], [[-0.04537391662597656]], [[0.04446125030517578]], [[-0.04476308822631836]], [[0.027647972106933594]], [[0.0016603469848632812]], [[0.10020160675048828]], [[0.0163421630859375]], [[-0.019906997680664062]], [[0.01024484634399414]], [[0.08594512939453125]], [[0.006433963775634766]], [[-0.012852668762207031]], [[-0.03166770935058594]], [[-0.4375171661376953]], [[-0.02010059356689453]], [[3.7599925994873047]], [[0.0]], [[0.8641700744628906]], [[0.6467180252075195]], [[-1.92446768283844]], [[-3.300853729248047]], [[-1.4697235822677612]], [[0.0]], [[0.09003877639770508]], [[-2.2223849296569824]], [[0.10798454284667969]], [[0.0]], [[-1.7188316583633423]], [[-1.7188316583633423]], [[-1.4913793802261353]], [[0.0]], [[0.08318686485290527]], [[0.06662595272064209]], [[-0.015944480895996094]], [[-0.015944480895996094]], [[0.1279517412185669]], [[-0.011188507080078125]], [[9.334087371826172e-05]], [[-0.015944480895996094]], [[0.007233619689941406]], [[-0.015944480895996094]], [[-0.015944480895996094]], [[-0.015944480895996094]], [[-0.015944480895996094]], [[-0.015944480895996094]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]]], \"firstDimensionName\": \"Layer\", \"secondDimensionName\": \"Neuron\"}\n",
       "    )\n",
       "    </script>"
      ],
      "text/plain": [
       "<circuitsvis.utils.render.RenderedHTML at 0x7f0e94fc2410>"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ablate_text(text_list, best_feature, model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div id=\"circuits-vis-932b68f9-0eb3\" style=\"margin: 15px 0;\"/>\n",
       "    <script crossorigin type=\"module\">\n",
       "    import { render, TextNeuronActivations } from \"https://unpkg.com/circuitsvis@1.40.1/dist/cdn/esm.js\";\n",
       "    render(\n",
       "      \"circuits-vis-932b68f9-0eb3\",\n",
       "      TextNeuronActivations,\n",
       "      {\"tokens\": [\" Bundes\", \"reg\", \"ierung\", \" will\", \" T\", \"aus\", \"ende\", \" As\", \"yl\", \"b\", \"ew\", \"er\", \"ber\", \" nach\", \" Afghanistan\", \" abs\", \"chie\", \"ben\", \".\", \" Ein\", \"ige\", \" von\", \" ihn\", \"en\", \" bez\", \"icht\", \"igen\", \" sich\", \" des\", \"hal\", \"b\", \" nun\", \" sel\", \"bst\", \" als\", \" Terror\", \"isten\", \" \\u2013\", \" und\", \"\\n\", \"ie\", \" bet\", \"re\", \"ib\", \"t\", \" man\", \" die\", \" D\", \"ek\", \"ol\", \"on\", \"is\", \"ierung\", \" einer\", \" F\", \"ach\", \"ze\", \"its\", \"chr\", \"ift\", \"?\", \" Der\", \" Che\", \"fred\", \"ak\", \"te\", \"ur\", \" der\", \" \\u201e\", \"American\", \" Historical\", \" Review\", \"\\u201c\", \" er\", \"kl\", \"\\u00e4\", \"rt\", \",\", \" wie\", \"\\n\", \" \\u201e\", \"Sch\", \"reck\", \"ens\", \"n\", \"acht\", \"\\u201c\", \" von\", \" Im\", \"st\", \"\\\\newline\", \"\\\\newline\", \"En\", \"de\", \" April\", \" 1938\", \" haben\", \" Nazis\", \" in\", \" der\", \" Ti\", \"rol\", \"er\", \" Kle\", \"inst\", \"ad\", \"t\", \" Im\", \"st\", \" einen\", \" G\", \"ew\", \"al\", \"tex\", \"z\", \"ess\", \" gegen\", \" Aust\", \"ro\", \"\\n\", \"\\\\newline\", \"In\", \" T\", \"ans\", \"ania\", \" ab\", \" so\", \"fort\", \" ver\", \"bot\", \"en\", \":\", \" Gle\", \"it\", \"gel\", \"\\\\newline\", \"\\\\newline\", \"23\", \".\", \" Jul\", \"i\", \" 2016\", \",\", \" 16\", \":\", \"33\", \"h\", \",\", \"\\\\newline\", \"\\\\newline\", \"F\", \"\\u00fcr\", \" die\", \" Reg\", \"ierung\", \" ist\", \" das\", \" Ver\", \"k\", \"\\n\", \":\", \"\\\\newline\", \"\\\\newline\", \"W\", \"ie\", \" sp\", \"richt\", \" man\", \" von\", \" der\", \" deut\", \"schen\", \" Spr\", \"ache\", \"?\", \"\\\\newline\", \"\\\\newline\", \"Es\", \" k\", \"ling\", \"t\", \" fast\", \" wie\", \" eine\", \" Meta\", \"fr\", \"age\", \",\", \" aber\", \" wie\", \" sp\", \"richt\", \" man\", \" e\", \"igent\", \"lich\", \" von\", \" der\", \" deut\", \"\\n\", \" Bundes\", \"reg\", \"ierung\", \" will\", \" T\", \"aus\", \"ende\", \" As\", \"yl\", \"b\", \"ew\", \"er\", \"ber\", \" nach\", \" Afghanistan\", \" abs\", \"chie\", \"ben\", \".\", \" Ein\", \"ige\", \" von\", \" ihn\", \"en\", \" bez\", \"icht\", \"igen\", \" sich\", \" des\", \"hal\", \"b\", \" nun\", \" sel\", \"bst\", \" als\", \" Terror\", \"isten\", \" \\u2013\", \" und\", \"\\n\", \"i\", \",\", \" den\", \"ne\", \" art\", \"ik\", \"kel\", \"en\", \" er\", \" over\", \" ett\", \" \\u00e5r\", \" g\", \"amm\", \"el\", \" og\", \" kan\", \" in\", \"ne\", \"hol\", \"de\", \" ut\", \"d\", \"ater\", \"t\", \" inform\", \"as\", \"jon\", \"\\\\newline\", \"\\\\newline\", \"R\", \"une\", \" Gra\", \"hn\", \" (\", \"66\", \"),\", \" re\", \"kt\", \"\\n\", \" den\", \" f\", \"yn\", \"ske\", \" opt\", \"iker\", \" Finn\", \" Jun\", \"cker\", \" to\", \"g\", \" til\", \" Sen\", \"egal\", \" for\", \" at\", \" v\\u00e6re\", \" med\", \" til\", \" at\", \" ud\", \"de\", \"le\", \" b\", \"rug\", \"te\", \" br\", \"iller\", \" til\", \" be\", \"folk\", \"ningen\", \",\", \" var\", \" han\", \" op\", \"fy\", \"ld\", \"t\", \"\\n\", \"alt\", \"br\", \"unn\", \" railway\", \" station\", \"\\\\newline\", \"\\\\newline\", \"K\", \"alt\", \"br\", \"unn\", \" railway\", \" station\", \" is\", \" a\", \" railway\", \" station\", \" situated\", \" in\", \" the\", \" municipality\", \" of\", \" K\", \"alt\", \"br\", \"unn\", \" in\", \" the\", \" Swiss\", \" cant\", \"on\", \" of\", \" St\", \".\", \" Gall\", \"en\", \".\", \" It\", \" is\", \"\\n\", \"rene\", \"z\", \" note\", \" que\", \" cet\", \" article\", \" pub\", \"li\", \"\\u00e9\", \" en\", \" 2016\", \" pour\", \"rait\", \" cont\", \"en\", \"ir\", \" des\", \" inform\", \"ations\", \" qui\", \" ne\", \" sont\", \" plus\", \" \\u00e0\", \" jour\", \".\", \"\\\\newline\", \"\\\\newline\", \"L\", \"'\", \"h\", \"umor\", \"iste\", \" rim\", \"ous\", \"ko\", \"is\", \" Fred\", \" Dub\", \"\\n\", \"id\", \"ative\", \" stress\", \" in\", \" stable\", \" cystic\", \" fibrosis\", \" patients\", \":\", \" do\", \" we\", \" need\", \" higher\", \" antioxidant\", \" plasma\", \" levels\", \"?\", \"\\\\newline\", \"Ox\", \"id\", \"ative\", \" stress\", \" plays\", \" an\", \" important\", \" role\", \" in\", \" cystic\", \" fibrosis\", \" (\", \"CF\", \").\", \" However\", \",\", \" there\", \" is\", \" a\", \" lack\", \" of\", \"\\n\"], \"activations\": [[[0.0]], [[-0.3218259811401367]], [[-0.05874896049499512]], [[0.31967639923095703]], [[-0.8733654022216797]], [[-1.0546834468841553]], [[-0.03999197483062744]], [[0.1558375358581543]], [[0.06452858448028564]], [[-0.3899409770965576]], [[0.14168167114257812]], [[-0.07440585643053055]], [[-0.06836798042058945]], [[-1.0900664329528809]], [[-0.37645721435546875]], [[-0.4290647506713867]], [[-1.0491008758544922]], [[-0.016864173114299774]], [[0.06154584884643555]], [[-1.1704912185668945]], [[0.2922399044036865]], [[-0.6397190093994141]], [[-0.15118515491485596]], [[0.0016124295070767403]], [[0.291780948638916]], [[-0.9136462211608887]], [[0.38416504859924316]], [[-0.7679028511047363]], [[-0.710059642791748]], [[-0.5874987244606018]], [[0.0005374866304919124]], [[-1.1508822441101074]], [[0.06333112716674805]], [[-0.12583771347999573]], [[-0.22880077362060547]], [[-0.866055965423584]], [[-0.29585063457489014]], [[0.18523263931274414]], [[0.08344364166259766]], [[0.0]], [[0.0]], [[-0.41448020935058594]], [[-0.10072016716003418]], [[-0.2648963928222656]], [[0.02682371437549591]], [[-0.1695566177368164]], [[-0.6481101512908936]], [[-1.0115604400634766]], [[0.31163883209228516]], [[0.11135578155517578]], [[0.17723298072814941]], [[0.36434268951416016]], [[-0.7242796421051025]], [[-1.0796265602111816]], [[-0.23223638534545898]], [[-0.06881237030029297]], [[0.040450096130371094]], [[-0.26040518283843994]], [[0.43692779541015625]], [[0.009894296526908875]], [[0.04834747314453125]], [[-0.6105313301086426]], [[-0.09868621826171875]], [[0.1952601671218872]], [[0.25025129318237305]], [[-0.1227681040763855]], [[0.012392893433570862]], [[-1.1551249027252197]], [[-0.6751809120178223]], [[0.27086734771728516]], [[0.3318214416503906]], [[0.05504465103149414]], [[0.02987605333328247]], [[-0.2867889404296875]], [[-0.4499160051345825]], [[0.04514756053686142]], [[-0.001943914219737053]], [[-0.08317756652832031]], [[0.11937379837036133]], [[0.0]], [[0.0]], [[-0.17901325225830078]], [[-0.5197529792785645]], [[0.04530525207519531]], [[-0.19847869873046875]], [[0.1607210636138916]], [[-0.08080863952636719]], [[-0.3497505187988281]], [[-0.4134054183959961]], [[-0.3150639533996582]], [[0.10925769805908203]], [[-0.10661578178405762]], [[0.3711094856262207]], [[-0.09821867942810059]], [[-0.2879619598388672]], [[0.19746732711791992]], [[-0.6024656295776367]], [[1.1798110008239746]], [[-0.163069486618042]], [[-0.5850539207458496]], [[0.1254281997680664]], [[-0.047180354595184326]], [[0.056365013122558594]], [[0.21611881256103516]], [[-0.5246355533599854]], [[0.35551851987838745]], [[0.017551345750689507]], [[-0.6034889221191406]], [[0.3360788822174072]], [[-0.47838401794433594]], [[-0.449007511138916]], [[-0.7239668369293213]], [[0.5881268978118896]], [[0.17605924606323242]], [[0.18179607391357422]], [[0.6484737396240234]], [[-0.08983945846557617]], [[0.9738092422485352]], [[0.18219566345214844]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[-0.23151874542236328]], [[-0.3183095455169678]], [[-0.0693395733833313]], [[-0.056485652923583984]], [[-0.03947162628173828]], [[-0.260988712310791]], [[-0.19884109497070312]], [[0.79736328125]], [[0.0025788545608520508]], [[-0.015984535217285156]], [[-0.0012897253036499023]], [[0.0038700103759765625]], [[-0.01642453670501709]], [[0.005591869354248047]], [[-0.015954017639160156]], [[0.00802922248840332]], [[-0.0004787333309650421]], [[-0.000591278076171875]], [[0.03276538848876953]], [[0.007204771041870117]], [[0.017396211624145508]], [[0.0021700412034988403]], [[0.008937835693359375]], [[-0.42462992668151855]], [[-1.0638998746871948]], [[-0.45436859130859375]], [[-0.1977512240409851]], [[0.16856622695922852]], [[-0.8910989761352539]], [[-0.7217168807983398]], [[0.41931962966918945]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[-0.1952981948852539]], [[-0.5261383056640625]], [[-0.19075965881347656]], [[-0.7603964805603027]], [[-0.5770342350006104]], [[-0.5774350166320801]], [[-0.1009913980960846]], [[-1.3511791229248047]], [[0.015484459698200226]], [[0.10318517684936523]], [[0.01935294270515442]], [[0.015796512365341187]], [[-0.2008523941040039]], [[0.10736560821533203]], [[-0.26402968168258667]], [[-0.2325124591588974]], [[-1.1293392181396484]], [[-0.32776594161987305]], [[-0.6501502990722656]], [[-0.19797801971435547]], [[-0.13251447677612305]], [[-0.04382915049791336]], [[0.042588233947753906]], [[-0.12176847457885742]], [[-0.604215145111084]], [[0.5770561695098877]], [[-0.009601429104804993]], [[0.011026926338672638]], [[-0.17394161224365234]], [[-1.8102266788482666]], [[-0.004016394726932049]], [[0.4441053867340088]], [[-0.017999231815338135]], [[-0.07211154699325562]], [[0.0]], [[0.0]], [[-0.3218259811401367]], [[-0.05874896049499512]], [[0.31967639923095703]], [[-0.8733654022216797]], [[-1.0546834468841553]], [[-0.03999197483062744]], [[0.1558375358581543]], [[0.06452858448028564]], [[-0.3899409770965576]], [[0.14168167114257812]], [[-0.07440585643053055]], [[-0.06836798042058945]], [[-1.0900664329528809]], [[-0.37645721435546875]], [[-0.4290647506713867]], [[-1.0491008758544922]], [[-0.016864173114299774]], [[0.06154584884643555]], [[-1.1704912185668945]], [[0.2922399044036865]], [[-0.6397190093994141]], [[-0.15118515491485596]], [[0.0016124295070767403]], [[0.291780948638916]], [[-0.9136462211608887]], [[0.38416504859924316]], [[-0.7679028511047363]], [[-0.710059642791748]], [[-0.5874987244606018]], [[0.0005374866304919124]], [[-1.1508822441101074]], [[0.06333112716674805]], [[-0.12583771347999573]], [[-0.22880077362060547]], [[-0.866055965423584]], [[-0.29585063457489014]], [[0.18523263931274414]], [[0.08344364166259766]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[-0.04432964324951172]], [[1.2239303588867188]], [[0.47414302825927734]], [[-0.22148799896240234]], [[0.2704048156738281]], [[-0.049974024295806885]], [[-0.23661518096923828]], [[0.004257917404174805]], [[-0.00022459030151367188]], [[5.237641744315624e-06]], [[0.03453585505485535]], [[5.959300324320793e-07]], [[-0.00023164227604866028]], [[-0.007128894329071045]], [[0.0020613670349121094]], [[-0.00017996132373809814]], [[0.0010520368814468384]], [[0.00032951822504401207]], [[-0.0009657144546508789]], [[-0.0030884742736816406]], [[0.005077242851257324]], [[0.01538577675819397]], [[0.0003580451011657715]], [[-1.2023025192320347e-05]], [[9.957700967788696e-05]], [[-0.00022596120834350586]], [[-0.000981450080871582]], [[0.0011143684387207031]], [[0.0006127357482910156]], [[0.0002117156982421875]], [[0.0004200935363769531]], [[-0.0035669803619384766]], [[0.0015687942504882812]], [[-0.0011185407638549805]], [[-0.00035762786865234375]], [[0.00014865398406982422]], [[0.0]], [[0.0]], [[0.2694540023803711]], [[0.2994985580444336]], [[-0.00664520263671875]], [[0.1991558074951172]], [[-0.25322580337524414]], [[0.12588977813720703]], [[-0.10116291046142578]], [[0.04668474197387695]], [[0.39623212814331055]], [[0.019811630249023438]], [[0.19306468963623047]], [[-0.015414237976074219]], [[-0.005452156066894531]], [[0.05266904830932617]], [[-0.02033233642578125]], [[0.03507232666015625]], [[-0.04171633720397949]], [[0.002676248550415039]], [[0.012113094329833984]], [[0.03576850891113281]], [[0.004972934722900391]], [[-0.018936634063720703]], [[0.020246028900146484]], [[0.0031375885009765625]], [[-0.0017704963684082031]], [[-0.008822441101074219]], [[-0.0033087730407714844]], [[-0.024227142333984375]], [[-0.01262521743774414]], [[-0.0020264387130737305]], [[0.0003393888473510742]], [[-0.0029549598693847656]], [[0.006681919097900391]], [[0.04571247100830078]], [[0.05498790740966797]], [[-0.0006237030029296875]], [[-0.00010104849934577942]], [[0.0003539975732564926]], [[0.0]], [[0.0]], [[-0.04947948455810547]], [[-0.2176990509033203]], [[0.004530906677246094]], [[2.9325485229492188e-05]], [[-7.104873657226562e-05]], [[-0.001272439956665039]], [[-1.1920928955078125e-05]], [[0.0057405829429626465]], [[-0.0011014044284820557]], [[0.0018322467803955078]], [[7.796287536621094e-05]], [[1.0523945093154907e-06]], [[-1.239776611328125e-05]], [[-9.036064147949219e-05]], [[-3.540515899658203e-05]], [[1.4789402484893799e-06]], [[0.00010728836059570312]], [[1.7881393432617188e-06]], [[0.0004565715789794922]], [[0.0037713050842285156]], [[4.188716411590576e-05]], [[-0.0012238025665283203]], [[-0.0002270340919494629]], [[-0.0019274502992630005]], [[0.0015996471047401428]], [[-0.0008099079132080078]], [[-3.6597251892089844e-05]], [[-0.0006375312805175781]], [[-5.778670310974121e-05]], [[-9.208917617797852e-06]], [[-2.1532177925109863e-05]], [[-0.1324324607849121]], [[-0.006254225969314575]], [[-0.004974745213985443]], [[0.00010584807023406029]], [[7.593631744384766e-05]], [[-1.8715858459472656e-05]], [[-9.417533874511719e-06]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[5.872175097465515e-05]], [[-2.956390380859375e-05]], [[-5.0067901611328125e-05]], [[1.3828277587890625e-05]], [[3.981590270996094e-05]], [[0.0018463134765625]], [[0.016333788633346558]], [[-3.57162207365036e-07]], [[-5.245208740234375e-06]], [[-9.775161743164062e-06]], [[4.7031790018081665e-07]], [[9.775161743164062e-06]], [[1.621246337890625e-05]], [[-1.609325408935547e-05]], [[0.07371640205383301]], [[-0.0029659271240234375]], [[0.0022079944610595703]], [[-0.0033664703369140625]], [[-0.00081634521484375]], [[-3.094226121902466e-05]], [[4.4345855712890625e-05]], [[-3.3974647521972656e-06]], [[5.7220458984375e-06]], [[0.0004830360412597656]], [[5.9485435485839844e-05]], [[-5.340576171875e-05]], [[-0.014880180358886719]], [[-0.0020351409912109375]], [[0.010648488998413086]], [[0.005695343017578125]], [[-0.000102996826171875]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]]], \"firstDimensionName\": \"Layer\", \"secondDimensionName\": \"Neuron\"}\n",
       "    )\n",
       "    </script>"
      ],
      "text/plain": [
       "<circuitsvis.utils.render.RenderedHTML at 0x7f0e95954070>"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ablate_feature_direction_display(full_text, best_feature)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[' nun', 'unter', ' sein', ' durch', ' einen', 'ungen', ' belie', ' keine', ' und', ' auf', ' dieser', ' bei', 'recht', ' sehr', ' mö', ' unter', ' zum', 'zw', 'igent', 'zeit']\n",
      "tensor([1.3300, 1.2846, 1.2522, 1.2375, 1.2310, 1.1838, 1.1765, 1.1692, 1.1676,\n",
      "        1.1579, 1.1498, 1.1394, 1.1206, 1.1199, 1.1132, 1.1127, 1.1083, 1.1004,\n",
      "        1.0921, 1.0913])\n"
     ]
    }
   ],
   "source": [
    "logit_lens(model,best_feature, smaller_dict, layer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div id=\"circuits-vis-7dddbb60-0c2e\" style=\"margin: 15px 0;\"/>\n",
       "    <script crossorigin type=\"module\">\n",
       "    import { render, TextNeuronActivations } from \"https://unpkg.com/circuitsvis@1.40.1/dist/cdn/esm.js\";\n",
       "    render(\n",
       "      \"circuits-vis-7dddbb60-0c2e\",\n",
       "      TextNeuronActivations,\n",
       "      {\"tokens\": [\"Include\", \" your\", \" own\", \" German\", \" text\", \".\", \" Ja\", \",\", \" das\", \" ist\", \" gut\", \".\", \"\\n\"], \"activations\": [[[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[2.570918560028076]], [[3.604458808898926]], [[7.990434169769287]], [[8.273178100585938]], [[7.29265832901001]], [[4.7365875244140625]], [[0.0]]], \"firstDimensionName\": \"Layer\", \"secondDimensionName\": \"Neuron\"}\n",
       "    )\n",
       "    </script>"
      ],
      "text/plain": [
       "<circuitsvis.utils.render.RenderedHTML at 0x7f0e951f35e0>"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "custom_text = [\n",
    "    \"Include your own German text. Ja, das ist gut!\",\n",
    "]\n",
    "visualize_text(custom_text, best_feature, model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# NOTE TO LUCIA & LOVIS\n",
    "I haven't checked the next part of code at all, so beware!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data Centric Viewpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Go through datapoints & see if the features that activate on them make sense.\n",
    "d_point = 0\n",
    "# text = tokens_dataset[d_point]\n",
    "data_ind, sequence_pos = np.unravel_index(d_point, (datapoints, token_amount))\n",
    "feature_val, feature_ind = dictionary_activations[d_point].topk(10)\n",
    "data_ind = int(data_ind)\n",
    "sequence_pos = int(sequence_pos)\n",
    "full_tok = torch.tensor(dataset[data_ind][\"input_ids\"])\n",
    "full_text = []\n",
    "full_text.append(model.tokenizer.decode(full_tok))\n",
    "visualize_text(full_text, feature_ind, model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Check the neuron/residual basis\n",
    "When we look at the weights of a feature, we are seeing the literal dimensions from the residual stream/neurons being read from the feature. \n",
    "\n",
    "Here I'm visualizing the weight values for the residual stream. If there are outliers, then it's mainly reading from that dimension."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(weights*max_activation).topk(20), (weights*max_activation).topk(20, largest=False).values, (weights*max_activation > 0.2).sum()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Prepend/Append tokens\n",
    "We can iterate over all tokens to check which ones activate a feature a lot to more rigorously test a hypothesis on what a feature means.\n",
    "\n",
    "Note: I'm literately running the model through all 50k tokens prepended to the text here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def prepend_all_tokens_and_get_feature_activation(model, minimal_activating_example, feature, setting=\"prepend\"):\n",
    "    tokens = model.to_tokens(minimal_activating_example, prepend_bos=False)\n",
    "\n",
    "    # Run through every number up to vocab size\n",
    "    vocab_size = model.cfg.d_vocab\n",
    "    batch_size = 256*2 # Define your desired batch size\n",
    "\n",
    "    dollar_feature_activations = torch.zeros(vocab_size)\n",
    "    for start in range(0, vocab_size, batch_size):\n",
    "        end = min(start + batch_size, vocab_size)\n",
    "\n",
    "        token_prep = torch.arange(start, end).to(device)\n",
    "        token_prep = token_prep.unsqueeze(1)  # Add a dimension for concatenation\n",
    "\n",
    "        # 1. Prepend to the tokens\n",
    "        if setting == \"prepend\":\n",
    "            tokens_catted = torch.cat((token_prep, tokens.repeat(end - start, 1)), dim=1).long()\n",
    "        elif setting == \"append\":\n",
    "            tokens_catted = torch.cat((tokens.repeat(end - start, 1), token_prep), dim=1).long()\n",
    "        else:\n",
    "            raise ValueError(f\"Unknown setting: {setting}\")\n",
    "\n",
    "        # 2. Run through the model\n",
    "        with torch.no_grad():\n",
    "            _, cache = model.run_with_cache(tokens_catted.to(device))\n",
    "            neuron_act_batch = cache[cache_name]\n",
    "            _, act = smaller_auto_encoder(neuron_act_batch)\n",
    "\n",
    "        # 3. Get the feature\n",
    "        dollar_feature_activations[start:end] = act[:, -1, feature].cpu().squeeze()\n",
    "\n",
    "    k = 20\n",
    "    k_increasing_val, k_increasing_ind = dollar_feature_activations.topk(k)\n",
    "    k_decreasing_val, k_decreasing_ind = dollar_feature_activations.topk(k, largest=False)\n",
    "    if(setting == \"prepend\"):\n",
    "        print(f\"[token]{minimal_activating_example}\")\n",
    "    elif(setting == \"append\"):\n",
    "        print(f\"{minimal_activating_example}[token]\")\n",
    "    else:\n",
    "        raise ValueError(f\"Unknown setting: {setting}\")\n",
    "    # Print indices converted to tokens\n",
    "    print(f\"Top-{k} increasing: {model.to_str_tokens(k_increasing_ind)}\")\n",
    "    # Print values\n",
    "    print(f\"Top-{k} increasing: {[f'{val:.2f}' for val in k_increasing_val]}\")\n",
    "    print(f\"Top-{k} decreasing: {model.to_str_tokens(k_decreasing_ind)}\")\n",
    "    print(f\"Top-{k} decreasing: {[f'{val:.2f}' for val in k_decreasing_val]}\")\n",
    "    print(f\"Number of 0 activations: {torch.sum(dollar_feature_activations == 0)}\")\n",
    "    if(setting == \"prepend\"):\n",
    "        best_text = \"\".join(model.to_str_tokens(dollar_feature_activations.argmax()) + [minimal_activating_example])\n",
    "    else:\n",
    "        best_text = \"\".join([minimal_activating_example] + model.to_str_tokens(dollar_feature_activations.argmax()))\n",
    "    return best_text\n",
    "\n",
    "best_text = \"\"\n",
    "for x in range(3):\n",
    "    # best_text = prepend_all_tokens_and_get_feature_activation(model, best_text, best_feature, setting=\"prepend\")\n",
    "    best_text = prepend_all_tokens_and_get_feature_activation(model, best_text, best_feature, setting=\"append\")\n",
    "# prepend_all_tokens_and_get_feature_activation(model, \" for all $\", best_feature, setting=\"prepend\")\n",
    "# prepend_all_tokens_and_get_feature_activation(model, \" tree\", best_feature, setting=\"prepend\")\n",
    "# prepend_all_tokens_and_get_feature_activation(model, \" tree\", best_feature, setting=\"append\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "best_text = \"\"\n",
    "for x in range(3):\n",
    "    best_text = prepend_all_tokens_and_get_feature_activation(model, best_text, best_feature, setting=\"prepend\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prepend_all_tokens_and_get_feature_activation(model, \" for all $\", best_feature, setting=\"prepend\")\n",
    "# prepend_all_tokens_and_get_feature_activation(model, \"The\", best_feature, setting=\"append\")\n",
    "# prepend_all_tokens_and_get_feature_activation(model, \" tree\", best_feature, setting=\"append\")"
   ]
  }
 ],
 "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.10.6"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}