Interactive_Neuroscope_for GPT-2 finetuning.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "k4dKWN1xPGLw"
      },
      "source": [
        "# Interactive Neuroscope\n",
        "\n",
        "*This is an interactive accompaniment to [neuroscope.io](https://neuroscope.io) and to the [studying learned language features post](https://www.alignmentforum.org/posts/Qup9gorqpd9qKAEav/200-cop-in-mi-studying-learned-features-in-language-models) in [200 Concrete Open Problems in Mechanistic Interpretability](https://neelnanda.io/concrete-open-problems)*\n",
        "\n",
        "There's a surprisingly rich ecosystem of easy ways to create interactive graphics, especially for ML systems. If you're trying to do mechanistic interpretability, the ability to do web dev and to both visualize data and interact with it seems high value! \n",
        "\n",
        "This is a demo of how you can combine HookedTransformer and [Gradio](https://gradio.app/) to create an interactive Neuroscope - a visualization of a neuron's activations on text that will dynamically update as you edit the text. I don't particularly claim that this code is any *good*, but the goal is to illustrate what quickly hacking together a custom visualisation (while knowing fuck all about web dev, like me) can look like! (And as such, I try to explain the basic web dev concepts I use)\n",
        "\n",
        "Note that you'll need to run the code yourself to get the interactive interface, so the cell at the bottom will be blank at first!\n",
        "\n",
        "To emphasise - the point of this notebook is to be a rough proof of concept that just about works, *not* to be the well executed ideal of interactively studying neurons! You are highly encouraged to write your own (and ideally, to [make a pull request](https://github.com/neelnanda-io/TransformerLens/pulls) with improvements!)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cgXdCnavPGL3"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "ZzckbMpPPGL4",
        "outputId": "6da7be63-df51-480b-a0a7-179d7b979261",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Running as a Colab notebook\n"
          ]
        }
      ],
      "source": [
        "import os\n",
        "\n",
        "try:\n",
        "    import google.colab\n",
        "\n",
        "    IN_COLAB = True\n",
        "    print(\"Running as a Colab notebook\")\n",
        "\n",
        "except:\n",
        "    IN_COLAB = False\n",
        "    print(\"Running as a Jupyter notebook - intended for development only!\")\n",
        "    from IPython import get_ipython\n",
        "\n",
        "    ipython = get_ipython()\n",
        "    print(ipython)\n",
        "    # Code to automatically update the HookedTransformer code as its edited without restarting the kernel\n",
        "    ipython.magic(\"load_ext autoreload\")\n",
        "    ipython.magic(\"autoreload 2\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "gxxO_YSGPGL6"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "if IN_COLAB:\n",
        "    os.system(\"pip install git+https://github.com/neelnanda-io/TransformerLens.git\")\n",
        "    os.system(\"pip install gradio\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "ZYc5FQa3PGL7"
      },
      "outputs": [],
      "source": [
        "import gradio as gr\n",
        "from transformer_lens import HookedTransformer\n",
        "from transformer_lens.utils import to_numpy\n",
        "from IPython.display import HTML"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "y8Q64v6mPGL7"
      },
      "source": [
        "## Extracting Model Activations\n",
        "\n",
        "We first write some code using HookedTransformer's cache to extract the neuron activations on a given layer and neuron, for a given text"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "-ZDgyc70PGL8",
        "outputId": "fe356173-56b0-4244-d1e3-c3bc341d9c51",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Using pad_token, but it is not set yet.\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Loaded pretrained model gpt2-small into HookedTransformer\n"
          ]
        }
      ],
      "source": [
        "model_name = \"gpt2-small\"\n",
        "model = HookedTransformer.from_pretrained(model_name)"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "model"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "bN0Zu8NmQH2m",
        "outputId": "38c5e6a9-cf8f-4cda-c260-0bb7eb38f76c"
      },
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "HookedTransformer(\n",
              "  (embed): Embed()\n",
              "  (hook_embed): HookPoint()\n",
              "  (pos_embed): PosEmbed()\n",
              "  (hook_pos_embed): HookPoint()\n",
              "  (blocks): ModuleList(\n",
              "    (0): TransformerBlock(\n",
              "      (ln1): LayerNormPre(\n",
              "        (hook_scale): HookPoint()\n",
              "        (hook_normalized): HookPoint()\n",
              "      )\n",
              "      (ln2): LayerNormPre(\n",
              "        (hook_scale): HookPoint()\n",
              "        (hook_normalized): HookPoint()\n",
              "      )\n",
              "      (attn): Attention(\n",
              "        (hook_k): HookPoint()\n",
              "        (hook_q): HookPoint()\n",
              "        (hook_v): HookPoint()\n",
              "        (hook_z): HookPoint()\n",
              "        (hook_attn_scores): HookPoint()\n",
              "        (hook_pattern): HookPoint()\n",
              "        (hook_result): HookPoint()\n",
              "      )\n",
              "      (mlp): MLP(\n",
              "        (hook_pre): HookPoint()\n",
              "        (hook_post): HookPoint()\n",
              "      )\n",
              "      (hook_attn_out): HookPoint()\n",
              "      (hook_mlp_out): HookPoint()\n",
              "      (hook_resid_pre): HookPoint()\n",
              "      (hook_resid_mid): HookPoint()\n",
              "      (hook_resid_post): HookPoint()\n",
              "    )\n",
              "    (1): TransformerBlock(\n",
              "      (ln1): LayerNormPre(\n",
              "        (hook_scale): HookPoint()\n",
              "        (hook_normalized): HookPoint()\n",
              "      )\n",
              "      (ln2): LayerNormPre(\n",
              "        (hook_scale): HookPoint()\n",
              "        (hook_normalized): HookPoint()\n",
              "      )\n",
              "      (attn): Attention(\n",
              "        (hook_k): HookPoint()\n",
              "        (hook_q): HookPoint()\n",
              "        (hook_v): HookPoint()\n",
              "        (hook_z): HookPoint()\n",
              "        (hook_attn_scores): HookPoint()\n",
              "        (hook_pattern): HookPoint()\n",
              "        (hook_result): HookPoint()\n",
              "      )\n",
              "      (mlp): MLP(\n",
              "        (hook_pre): HookPoint()\n",
              "        (hook_post): HookPoint()\n",
              "      )\n",
              "      (hook_attn_out): HookPoint()\n",
              "      (hook_mlp_out): HookPoint()\n",
              "      (hook_resid_pre): HookPoint()\n",
              "      (hook_resid_mid): HookPoint()\n",
              "      (hook_resid_post): HookPoint()\n",
              "    )\n",
              "    (2): TransformerBlock(\n",
              "      (ln1): LayerNormPre(\n",
              "        (hook_scale): HookPoint()\n",
              "        (hook_normalized): HookPoint()\n",
              "      )\n",
              "      (ln2): LayerNormPre(\n",
              "        (hook_scale): HookPoint()\n",
              "        (hook_normalized): HookPoint()\n",
              "      )\n",
              "      (attn): Attention(\n",
              "        (hook_k): HookPoint()\n",
              "        (hook_q): HookPoint()\n",
              "        (hook_v): HookPoint()\n",
              "        (hook_z): HookPoint()\n",
              "        (hook_attn_scores): HookPoint()\n",
              "        (hook_pattern): HookPoint()\n",
              "        (hook_result): HookPoint()\n",
              "      )\n",
              "      (mlp): MLP(\n",
              "        (hook_pre): HookPoint()\n",
              "        (hook_post): HookPoint()\n",
              "      )\n",
              "      (hook_attn_out): HookPoint()\n",
              "      (hook_mlp_out): HookPoint()\n",
              "      (hook_resid_pre): HookPoint()\n",
              "      (hook_resid_mid): HookPoint()\n",
              "      (hook_resid_post): HookPoint()\n",
              "    )\n",
              "    (3): TransformerBlock(\n",
              "      (ln1): LayerNormPre(\n",
              "        (hook_scale): HookPoint()\n",
              "        (hook_normalized): HookPoint()\n",
              "      )\n",
              "      (ln2): LayerNormPre(\n",
              "        (hook_scale): HookPoint()\n",
              "        (hook_normalized): HookPoint()\n",
              "      )\n",
              "      (attn): Attention(\n",
              "        (hook_k): HookPoint()\n",
              "        (hook_q): HookPoint()\n",
              "        (hook_v): HookPoint()\n",
              "        (hook_z): HookPoint()\n",
              "        (hook_attn_scores): HookPoint()\n",
              "        (hook_pattern): HookPoint()\n",
              "        (hook_result): HookPoint()\n",
              "      )\n",
              "      (mlp): MLP(\n",
              "        (hook_pre): HookPoint()\n",
              "        (hook_post): HookPoint()\n",
              "      )\n",
              "      (hook_attn_out): HookPoint()\n",
              "      (hook_mlp_out): HookPoint()\n",
              "      (hook_resid_pre): HookPoint()\n",
              "      (hook_resid_mid): HookPoint()\n",
              "      (hook_resid_post): HookPoint()\n",
              "    )\n",
              "    (4): TransformerBlock(\n",
              "      (ln1): LayerNormPre(\n",
              "        (hook_scale): HookPoint()\n",
              "        (hook_normalized): HookPoint()\n",
              "      )\n",
              "      (ln2): LayerNormPre(\n",
              "        (hook_scale): HookPoint()\n",
              "        (hook_normalized): HookPoint()\n",
              "      )\n",
              "      (attn): Attention(\n",
              "        (hook_k): HookPoint()\n",
              "        (hook_q): HookPoint()\n",
              "        (hook_v): HookPoint()\n",
              "        (hook_z): HookPoint()\n",
              "        (hook_attn_scores): HookPoint()\n",
              "        (hook_pattern): HookPoint()\n",
              "        (hook_result): HookPoint()\n",
              "      )\n",
              "      (mlp): MLP(\n",
              "        (hook_pre): HookPoint()\n",
              "        (hook_post): HookPoint()\n",
              "      )\n",
              "      (hook_attn_out): HookPoint()\n",
              "      (hook_mlp_out): HookPoint()\n",
              "      (hook_resid_pre): HookPoint()\n",
              "      (hook_resid_mid): HookPoint()\n",
              "      (hook_resid_post): HookPoint()\n",
              "    )\n",
              "    (5): TransformerBlock(\n",
              "      (ln1): LayerNormPre(\n",
              "        (hook_scale): HookPoint()\n",
              "        (hook_normalized): HookPoint()\n",
              "      )\n",
              "      (ln2): LayerNormPre(\n",
              "        (hook_scale): HookPoint()\n",
              "        (hook_normalized): HookPoint()\n",
              "      )\n",
              "      (attn): Attention(\n",
              "        (hook_k): HookPoint()\n",
              "        (hook_q): HookPoint()\n",
              "        (hook_v): HookPoint()\n",
              "        (hook_z): HookPoint()\n",
              "        (hook_attn_scores): HookPoint()\n",
              "        (hook_pattern): HookPoint()\n",
              "        (hook_result): HookPoint()\n",
              "      )\n",
              "      (mlp): MLP(\n",
              "        (hook_pre): HookPoint()\n",
              "        (hook_post): HookPoint()\n",
              "      )\n",
              "      (hook_attn_out): HookPoint()\n",
              "      (hook_mlp_out): HookPoint()\n",
              "      (hook_resid_pre): HookPoint()\n",
              "      (hook_resid_mid): HookPoint()\n",
              "      (hook_resid_post): HookPoint()\n",
              "    )\n",
              "    (6): TransformerBlock(\n",
              "      (ln1): LayerNormPre(\n",
              "        (hook_scale): HookPoint()\n",
              "        (hook_normalized): HookPoint()\n",
              "      )\n",
              "      (ln2): LayerNormPre(\n",
              "        (hook_scale): HookPoint()\n",
              "        (hook_normalized): HookPoint()\n",
              "      )\n",
              "      (attn): Attention(\n",
              "        (hook_k): HookPoint()\n",
              "        (hook_q): HookPoint()\n",
              "        (hook_v): HookPoint()\n",
              "        (hook_z): HookPoint()\n",
              "        (hook_attn_scores): HookPoint()\n",
              "        (hook_pattern): HookPoint()\n",
              "        (hook_result): HookPoint()\n",
              "      )\n",
              "      (mlp): MLP(\n",
              "        (hook_pre): HookPoint()\n",
              "        (hook_post): HookPoint()\n",
              "      )\n",
              "      (hook_attn_out): HookPoint()\n",
              "      (hook_mlp_out): HookPoint()\n",
              "      (hook_resid_pre): HookPoint()\n",
              "      (hook_resid_mid): HookPoint()\n",
              "      (hook_resid_post): HookPoint()\n",
              "    )\n",
              "    (7): TransformerBlock(\n",
              "      (ln1): LayerNormPre(\n",
              "        (hook_scale): HookPoint()\n",
              "        (hook_normalized): HookPoint()\n",
              "      )\n",
              "      (ln2): LayerNormPre(\n",
              "        (hook_scale): HookPoint()\n",
              "        (hook_normalized): HookPoint()\n",
              "      )\n",
              "      (attn): Attention(\n",
              "        (hook_k): HookPoint()\n",
              "        (hook_q): HookPoint()\n",
              "        (hook_v): HookPoint()\n",
              "        (hook_z): HookPoint()\n",
              "        (hook_attn_scores): HookPoint()\n",
              "        (hook_pattern): HookPoint()\n",
              "        (hook_result): HookPoint()\n",
              "      )\n",
              "      (mlp): MLP(\n",
              "        (hook_pre): HookPoint()\n",
              "        (hook_post): HookPoint()\n",
              "      )\n",
              "      (hook_attn_out): HookPoint()\n",
              "      (hook_mlp_out): HookPoint()\n",
              "      (hook_resid_pre): HookPoint()\n",
              "      (hook_resid_mid): HookPoint()\n",
              "      (hook_resid_post): HookPoint()\n",
              "    )\n",
              "    (8): TransformerBlock(\n",
              "      (ln1): LayerNormPre(\n",
              "        (hook_scale): HookPoint()\n",
              "        (hook_normalized): HookPoint()\n",
              "      )\n",
              "      (ln2): LayerNormPre(\n",
              "        (hook_scale): HookPoint()\n",
              "        (hook_normalized): HookPoint()\n",
              "      )\n",
              "      (attn): Attention(\n",
              "        (hook_k): HookPoint()\n",
              "        (hook_q): HookPoint()\n",
              "        (hook_v): HookPoint()\n",
              "        (hook_z): HookPoint()\n",
              "        (hook_attn_scores): HookPoint()\n",
              "        (hook_pattern): HookPoint()\n",
              "        (hook_result): HookPoint()\n",
              "      )\n",
              "      (mlp): MLP(\n",
              "        (hook_pre): HookPoint()\n",
              "        (hook_post): HookPoint()\n",
              "      )\n",
              "      (hook_attn_out): HookPoint()\n",
              "      (hook_mlp_out): HookPoint()\n",
              "      (hook_resid_pre): HookPoint()\n",
              "      (hook_resid_mid): HookPoint()\n",
              "      (hook_resid_post): HookPoint()\n",
              "    )\n",
              "    (9): TransformerBlock(\n",
              "      (ln1): LayerNormPre(\n",
              "        (hook_scale): HookPoint()\n",
              "        (hook_normalized): HookPoint()\n",
              "      )\n",
              "      (ln2): LayerNormPre(\n",
              "        (hook_scale): HookPoint()\n",
              "        (hook_normalized): HookPoint()\n",
              "      )\n",
              "      (attn): Attention(\n",
              "        (hook_k): HookPoint()\n",
              "        (hook_q): HookPoint()\n",
              "        (hook_v): HookPoint()\n",
              "        (hook_z): HookPoint()\n",
              "        (hook_attn_scores): HookPoint()\n",
              "        (hook_pattern): HookPoint()\n",
              "        (hook_result): HookPoint()\n",
              "      )\n",
              "      (mlp): MLP(\n",
              "        (hook_pre): HookPoint()\n",
              "        (hook_post): HookPoint()\n",
              "      )\n",
              "      (hook_attn_out): HookPoint()\n",
              "      (hook_mlp_out): HookPoint()\n",
              "      (hook_resid_pre): HookPoint()\n",
              "      (hook_resid_mid): HookPoint()\n",
              "      (hook_resid_post): HookPoint()\n",
              "    )\n",
              "    (10): TransformerBlock(\n",
              "      (ln1): LayerNormPre(\n",
              "        (hook_scale): HookPoint()\n",
              "        (hook_normalized): HookPoint()\n",
              "      )\n",
              "      (ln2): LayerNormPre(\n",
              "        (hook_scale): HookPoint()\n",
              "        (hook_normalized): HookPoint()\n",
              "      )\n",
              "      (attn): Attention(\n",
              "        (hook_k): HookPoint()\n",
              "        (hook_q): HookPoint()\n",
              "        (hook_v): HookPoint()\n",
              "        (hook_z): HookPoint()\n",
              "        (hook_attn_scores): HookPoint()\n",
              "        (hook_pattern): HookPoint()\n",
              "        (hook_result): HookPoint()\n",
              "      )\n",
              "      (mlp): MLP(\n",
              "        (hook_pre): HookPoint()\n",
              "        (hook_post): HookPoint()\n",
              "      )\n",
              "      (hook_attn_out): HookPoint()\n",
              "      (hook_mlp_out): HookPoint()\n",
              "      (hook_resid_pre): HookPoint()\n",
              "      (hook_resid_mid): HookPoint()\n",
              "      (hook_resid_post): HookPoint()\n",
              "    )\n",
              "    (11): TransformerBlock(\n",
              "      (ln1): LayerNormPre(\n",
              "        (hook_scale): HookPoint()\n",
              "        (hook_normalized): HookPoint()\n",
              "      )\n",
              "      (ln2): LayerNormPre(\n",
              "        (hook_scale): HookPoint()\n",
              "        (hook_normalized): HookPoint()\n",
              "      )\n",
              "      (attn): Attention(\n",
              "        (hook_k): HookPoint()\n",
              "        (hook_q): HookPoint()\n",
              "        (hook_v): HookPoint()\n",
              "        (hook_z): HookPoint()\n",
              "        (hook_attn_scores): HookPoint()\n",
              "        (hook_pattern): HookPoint()\n",
              "        (hook_result): HookPoint()\n",
              "      )\n",
              "      (mlp): MLP(\n",
              "        (hook_pre): HookPoint()\n",
              "        (hook_post): HookPoint()\n",
              "      )\n",
              "      (hook_attn_out): HookPoint()\n",
              "      (hook_mlp_out): HookPoint()\n",
              "      (hook_resid_pre): HookPoint()\n",
              "      (hook_resid_mid): HookPoint()\n",
              "      (hook_resid_post): HookPoint()\n",
              "    )\n",
              "  )\n",
              "  (ln_final): LayerNormPre(\n",
              "    (hook_scale): HookPoint()\n",
              "    (hook_normalized): HookPoint()\n",
              "  )\n",
              "  (unembed): Unembed()\n",
              ")"
            ]
          },
          "metadata": {},
          "execution_count": 5
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# # to be removed later \n",
        "# from transformers import GPT2Tokenizer, GPT2Model\n",
        "# tokenizer = GPT2Tokenizer.from_pretrained('gpt2')\n",
        "# model2 = GPT2Model.from_pretrained('gpt2')\n",
        "# text = \"Replace me by any text you'd like.\"\n",
        "# encoded_input = tokenizer(text, return_tensors='pt')\n",
        "# output = model2(**encoded_input)"
      ],
      "metadata": {
        "id": "TCRUIk5Kmj-1"
      },
      "execution_count": 15,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from transformers.generation.utils import MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING\n",
        "MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING.__dict__"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ytZka3UVm7NG",
        "outputId": "e9f0639f-e832-467f-ce5c-b6baaeabb708"
      },
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'_config_mapping': OrderedDict([('albert', 'AlbertConfig'),\n",
              "              ('audio-spectrogram-transformer', 'ASTConfig'),\n",
              "              ('bart', 'BartConfig'),\n",
              "              ('beit', 'BeitConfig'),\n",
              "              ('bert', 'BertConfig'),\n",
              "              ('bert-generation', 'BertGenerationConfig'),\n",
              "              ('big_bird', 'BigBirdConfig'),\n",
              "              ('bigbird_pegasus', 'BigBirdPegasusConfig'),\n",
              "              ('blenderbot', 'BlenderbotConfig'),\n",
              "              ('blenderbot-small', 'BlenderbotSmallConfig'),\n",
              "              ('bloom', 'BloomConfig'),\n",
              "              ('camembert', 'CamembertConfig'),\n",
              "              ('canine', 'CanineConfig'),\n",
              "              ('chinese_clip', 'ChineseCLIPConfig'),\n",
              "              ('clip', 'CLIPConfig'),\n",
              "              ('clipseg', 'CLIPSegConfig'),\n",
              "              ('codegen', 'CodeGenConfig'),\n",
              "              ('conditional_detr', 'ConditionalDetrConfig'),\n",
              "              ('convbert', 'ConvBertConfig'),\n",
              "              ('convnext', 'ConvNextConfig'),\n",
              "              ('ctrl', 'CTRLConfig'),\n",
              "              ('cvt', 'CvtConfig'),\n",
              "              ('data2vec-audio', 'Data2VecAudioConfig'),\n",
              "              ('data2vec-text', 'Data2VecTextConfig'),\n",
              "              ('data2vec-vision', 'Data2VecVisionConfig'),\n",
              "              ('deberta', 'DebertaConfig'),\n",
              "              ('deberta-v2', 'DebertaV2Config'),\n",
              "              ('decision_transformer', 'DecisionTransformerConfig'),\n",
              "              ('deformable_detr', 'DeformableDetrConfig'),\n",
              "              ('deit', 'DeiTConfig'),\n",
              "              ('detr', 'DetrConfig'),\n",
              "              ('dinat', 'DinatConfig'),\n",
              "              ('distilbert', 'DistilBertConfig'),\n",
              "              ('donut-swin', 'DonutSwinConfig'),\n",
              "              ('dpr', 'DPRConfig'),\n",
              "              ('dpt', 'DPTConfig'),\n",
              "              ('electra', 'ElectraConfig'),\n",
              "              ('encoder-decoder', 'EncoderDecoderConfig'),\n",
              "              ('ernie', 'ErnieConfig'),\n",
              "              ('esm', 'EsmConfig'),\n",
              "              ('flaubert', 'FlaubertConfig'),\n",
              "              ('flava', 'FlavaConfig'),\n",
              "              ('fnet', 'FNetConfig'),\n",
              "              ('fsmt', 'FSMTConfig'),\n",
              "              ('funnel', 'FunnelConfig'),\n",
              "              ('glpn', 'GLPNConfig'),\n",
              "              ('gpt2', 'GPT2Config'),\n",
              "              ('gpt_neo', 'GPTNeoConfig'),\n",
              "              ('gpt_neox', 'GPTNeoXConfig'),\n",
              "              ('gpt_neox_japanese', 'GPTNeoXJapaneseConfig'),\n",
              "              ('gptj', 'GPTJConfig'),\n",
              "              ('groupvit', 'GroupViTConfig'),\n",
              "              ('hubert', 'HubertConfig'),\n",
              "              ('ibert', 'IBertConfig'),\n",
              "              ('imagegpt', 'ImageGPTConfig'),\n",
              "              ('jukebox', 'JukeboxConfig'),\n",
              "              ('layoutlm', 'LayoutLMConfig'),\n",
              "              ('layoutlmv2', 'LayoutLMv2Config'),\n",
              "              ('layoutlmv3', 'LayoutLMv3Config'),\n",
              "              ('led', 'LEDConfig'),\n",
              "              ('levit', 'LevitConfig'),\n",
              "              ('lilt', 'LiltConfig'),\n",
              "              ('longformer', 'LongformerConfig'),\n",
              "              ('longt5', 'LongT5Config'),\n",
              "              ('luke', 'LukeConfig'),\n",
              "              ('lxmert', 'LxmertConfig'),\n",
              "              ('m2m_100', 'M2M100Config'),\n",
              "              ('marian', 'MarianConfig'),\n",
              "              ('markuplm', 'MarkupLMConfig'),\n",
              "              ('maskformer', 'MaskFormerConfig'),\n",
              "              ('maskformer-swin', 'MaskFormerSwinConfig'),\n",
              "              ('mbart', 'MBartConfig'),\n",
              "              ('mctct', 'MCTCTConfig'),\n",
              "              ('megatron-bert', 'MegatronBertConfig'),\n",
              "              ('mobilebert', 'MobileBertConfig'),\n",
              "              ('mobilenet_v1', 'MobileNetV1Config'),\n",
              "              ('mobilenet_v2', 'MobileNetV2Config'),\n",
              "              ('mobilevit', 'MobileViTConfig'),\n",
              "              ('mpnet', 'MPNetConfig'),\n",
              "              ('mt5', 'MT5Config'),\n",
              "              ('mvp', 'MvpConfig'),\n",
              "              ('nat', 'NatConfig'),\n",
              "              ('nezha', 'NezhaConfig'),\n",
              "              ('nystromformer', 'NystromformerConfig'),\n",
              "              ('openai-gpt', 'OpenAIGPTConfig'),\n",
              "              ('opt', 'OPTConfig'),\n",
              "              ('owlvit', 'OwlViTConfig'),\n",
              "              ('pegasus', 'PegasusConfig'),\n",
              "              ('pegasus_x', 'PegasusXConfig'),\n",
              "              ('perceiver', 'PerceiverConfig'),\n",
              "              ('plbart', 'PLBartConfig'),\n",
              "              ('poolformer', 'PoolFormerConfig'),\n",
              "              ('prophetnet', 'ProphetNetConfig'),\n",
              "              ('qdqbert', 'QDQBertConfig'),\n",
              "              ('rag', 'RagConfig'),\n",
              "              ('realm', 'RealmConfig'),\n",
              "              ('reformer', 'ReformerConfig'),\n",
              "              ('regnet', 'RegNetConfig'),\n",
              "              ('rembert', 'RemBertConfig'),\n",
              "              ('resnet', 'ResNetConfig'),\n",
              "              ('retribert', 'RetriBertConfig'),\n",
              "              ('roberta', 'RobertaConfig'),\n",
              "              ('roc_bert', 'RoCBertConfig'),\n",
              "              ('roformer', 'RoFormerConfig'),\n",
              "              ('segformer', 'SegformerConfig'),\n",
              "              ('sew', 'SEWConfig'),\n",
              "              ('sew-d', 'SEWDConfig'),\n",
              "              ('speech-encoder-decoder', 'SpeechEncoderDecoderConfig'),\n",
              "              ('speech_to_text', 'Speech2TextConfig'),\n",
              "              ('speech_to_text_2', 'Speech2Text2Config'),\n",
              "              ('splinter', 'SplinterConfig'),\n",
              "              ('squeezebert', 'SqueezeBertConfig'),\n",
              "              ('swin', 'SwinConfig'),\n",
              "              ('swinv2', 'Swinv2Config'),\n",
              "              ('switch_transformers', 'SwitchTransformersConfig'),\n",
              "              ('t5', 'T5Config'),\n",
              "              ('table-transformer', 'TableTransformerConfig'),\n",
              "              ('tapas', 'TapasConfig'),\n",
              "              ('time_series_transformer', 'TimeSeriesTransformerConfig'),\n",
              "              ('trajectory_transformer', 'TrajectoryTransformerConfig'),\n",
              "              ('transfo-xl', 'TransfoXLConfig'),\n",
              "              ('trocr', 'TrOCRConfig'),\n",
              "              ('unispeech', 'UniSpeechConfig'),\n",
              "              ('unispeech-sat', 'UniSpeechSatConfig'),\n",
              "              ('van', 'VanConfig'),\n",
              "              ('videomae', 'VideoMAEConfig'),\n",
              "              ('vilt', 'ViltConfig'),\n",
              "              ('vision-encoder-decoder', 'VisionEncoderDecoderConfig'),\n",
              "              ('vision-text-dual-encoder', 'VisionTextDualEncoderConfig'),\n",
              "              ('visual_bert', 'VisualBertConfig'),\n",
              "              ('vit', 'ViTConfig'),\n",
              "              ('vit_mae', 'ViTMAEConfig'),\n",
              "              ('vit_msn', 'ViTMSNConfig'),\n",
              "              ('wav2vec2', 'Wav2Vec2Config'),\n",
              "              ('wav2vec2-conformer', 'Wav2Vec2ConformerConfig'),\n",
              "              ('wavlm', 'WavLMConfig'),\n",
              "              ('whisper', 'WhisperConfig'),\n",
              "              ('xclip', 'XCLIPConfig'),\n",
              "              ('xglm', 'XGLMConfig'),\n",
              "              ('xlm', 'XLMConfig'),\n",
              "              ('xlm-prophetnet', 'XLMProphetNetConfig'),\n",
              "              ('xlm-roberta', 'XLMRobertaConfig'),\n",
              "              ('xlm-roberta-xl', 'XLMRobertaXLConfig'),\n",
              "              ('xlnet', 'XLNetConfig'),\n",
              "              ('yolos', 'YolosConfig'),\n",
              "              ('yoso', 'YosoConfig')]),\n",
              " '_reverse_config_mapping': {'AlbertConfig': 'albert',\n",
              "  'ASTConfig': 'audio-spectrogram-transformer',\n",
              "  'BartConfig': 'bart',\n",
              "  'BeitConfig': 'beit',\n",
              "  'BertConfig': 'bert',\n",
              "  'BertGenerationConfig': 'bert-generation',\n",
              "  'BigBirdConfig': 'big_bird',\n",
              "  'BigBirdPegasusConfig': 'bigbird_pegasus',\n",
              "  'BlenderbotConfig': 'blenderbot',\n",
              "  'BlenderbotSmallConfig': 'blenderbot-small',\n",
              "  'BloomConfig': 'bloom',\n",
              "  'CamembertConfig': 'camembert',\n",
              "  'CanineConfig': 'canine',\n",
              "  'ChineseCLIPConfig': 'chinese_clip',\n",
              "  'CLIPConfig': 'clip',\n",
              "  'CLIPSegConfig': 'clipseg',\n",
              "  'CodeGenConfig': 'codegen',\n",
              "  'ConditionalDetrConfig': 'conditional_detr',\n",
              "  'ConvBertConfig': 'convbert',\n",
              "  'ConvNextConfig': 'convnext',\n",
              "  'CTRLConfig': 'ctrl',\n",
              "  'CvtConfig': 'cvt',\n",
              "  'Data2VecAudioConfig': 'data2vec-audio',\n",
              "  'Data2VecTextConfig': 'data2vec-text',\n",
              "  'Data2VecVisionConfig': 'data2vec-vision',\n",
              "  'DebertaConfig': 'deberta',\n",
              "  'DebertaV2Config': 'deberta-v2',\n",
              "  'DecisionTransformerConfig': 'decision_transformer',\n",
              "  'DeformableDetrConfig': 'deformable_detr',\n",
              "  'DeiTConfig': 'deit',\n",
              "  'DetrConfig': 'detr',\n",
              "  'DinatConfig': 'dinat',\n",
              "  'DistilBertConfig': 'distilbert',\n",
              "  'DonutSwinConfig': 'donut-swin',\n",
              "  'DPRConfig': 'dpr',\n",
              "  'DPTConfig': 'dpt',\n",
              "  'ElectraConfig': 'electra',\n",
              "  'EncoderDecoderConfig': 'encoder-decoder',\n",
              "  'ErnieConfig': 'ernie',\n",
              "  'EsmConfig': 'esm',\n",
              "  'FlaubertConfig': 'flaubert',\n",
              "  'FlavaConfig': 'flava',\n",
              "  'FNetConfig': 'fnet',\n",
              "  'FSMTConfig': 'fsmt',\n",
              "  'FunnelConfig': 'funnel',\n",
              "  'GLPNConfig': 'glpn',\n",
              "  'GPT2Config': 'gpt2',\n",
              "  'GPTNeoConfig': 'gpt_neo',\n",
              "  'GPTNeoXConfig': 'gpt_neox',\n",
              "  'GPTNeoXJapaneseConfig': 'gpt_neox_japanese',\n",
              "  'GPTJConfig': 'gptj',\n",
              "  'GroupViTConfig': 'groupvit',\n",
              "  'HubertConfig': 'hubert',\n",
              "  'IBertConfig': 'ibert',\n",
              "  'ImageGPTConfig': 'imagegpt',\n",
              "  'JukeboxConfig': 'jukebox',\n",
              "  'LayoutLMConfig': 'layoutlm',\n",
              "  'LayoutLMv2Config': 'layoutlmv2',\n",
              "  'LayoutLMv3Config': 'layoutlmv3',\n",
              "  'LEDConfig': 'led',\n",
              "  'LevitConfig': 'levit',\n",
              "  'LiltConfig': 'lilt',\n",
              "  'LongformerConfig': 'longformer',\n",
              "  'LongT5Config': 'longt5',\n",
              "  'LukeConfig': 'luke',\n",
              "  'LxmertConfig': 'lxmert',\n",
              "  'M2M100Config': 'm2m_100',\n",
              "  'MarianConfig': 'marian',\n",
              "  'MarkupLMConfig': 'markuplm',\n",
              "  'MaskFormerConfig': 'maskformer',\n",
              "  'MaskFormerSwinConfig': 'maskformer-swin',\n",
              "  'MBartConfig': 'mbart',\n",
              "  'MCTCTConfig': 'mctct',\n",
              "  'MegatronBertConfig': 'megatron-bert',\n",
              "  'MobileBertConfig': 'mobilebert',\n",
              "  'MobileNetV1Config': 'mobilenet_v1',\n",
              "  'MobileNetV2Config': 'mobilenet_v2',\n",
              "  'MobileViTConfig': 'mobilevit',\n",
              "  'MPNetConfig': 'mpnet',\n",
              "  'MT5Config': 'mt5',\n",
              "  'MvpConfig': 'mvp',\n",
              "  'NatConfig': 'nat',\n",
              "  'NezhaConfig': 'nezha',\n",
              "  'NystromformerConfig': 'nystromformer',\n",
              "  'OpenAIGPTConfig': 'openai-gpt',\n",
              "  'OPTConfig': 'opt',\n",
              "  'OwlViTConfig': 'owlvit',\n",
              "  'PegasusConfig': 'pegasus',\n",
              "  'PegasusXConfig': 'pegasus_x',\n",
              "  'PerceiverConfig': 'perceiver',\n",
              "  'PLBartConfig': 'plbart',\n",
              "  'PoolFormerConfig': 'poolformer',\n",
              "  'ProphetNetConfig': 'prophetnet',\n",
              "  'QDQBertConfig': 'qdqbert',\n",
              "  'RagConfig': 'rag',\n",
              "  'RealmConfig': 'realm',\n",
              "  'ReformerConfig': 'reformer',\n",
              "  'RegNetConfig': 'regnet',\n",
              "  'RemBertConfig': 'rembert',\n",
              "  'ResNetConfig': 'resnet',\n",
              "  'RetriBertConfig': 'retribert',\n",
              "  'RobertaConfig': 'roberta',\n",
              "  'RoCBertConfig': 'roc_bert',\n",
              "  'RoFormerConfig': 'roformer',\n",
              "  'SegformerConfig': 'segformer',\n",
              "  'SEWConfig': 'sew',\n",
              "  'SEWDConfig': 'sew-d',\n",
              "  'SpeechEncoderDecoderConfig': 'speech-encoder-decoder',\n",
              "  'Speech2TextConfig': 'speech_to_text',\n",
              "  'Speech2Text2Config': 'speech_to_text_2',\n",
              "  'SplinterConfig': 'splinter',\n",
              "  'SqueezeBertConfig': 'squeezebert',\n",
              "  'SwinConfig': 'swin',\n",
              "  'Swinv2Config': 'swinv2',\n",
              "  'SwitchTransformersConfig': 'switch_transformers',\n",
              "  'T5Config': 't5',\n",
              "  'TableTransformerConfig': 'table-transformer',\n",
              "  'TapasConfig': 'tapas',\n",
              "  'TimeSeriesTransformerConfig': 'time_series_transformer',\n",
              "  'TrajectoryTransformerConfig': 'trajectory_transformer',\n",
              "  'TransfoXLConfig': 'transfo-xl',\n",
              "  'TrOCRConfig': 'trocr',\n",
              "  'UniSpeechConfig': 'unispeech',\n",
              "  'UniSpeechSatConfig': 'unispeech-sat',\n",
              "  'VanConfig': 'van',\n",
              "  'VideoMAEConfig': 'videomae',\n",
              "  'ViltConfig': 'vilt',\n",
              "  'VisionEncoderDecoderConfig': 'vision-encoder-decoder',\n",
              "  'VisionTextDualEncoderConfig': 'vision-text-dual-encoder',\n",
              "  'VisualBertConfig': 'visual_bert',\n",
              "  'ViTConfig': 'vit',\n",
              "  'ViTMAEConfig': 'vit_mae',\n",
              "  'ViTMSNConfig': 'vit_msn',\n",
              "  'Wav2Vec2Config': 'wav2vec2',\n",
              "  'Wav2Vec2ConformerConfig': 'wav2vec2-conformer',\n",
              "  'WavLMConfig': 'wavlm',\n",
              "  'WhisperConfig': 'whisper',\n",
              "  'XCLIPConfig': 'xclip',\n",
              "  'XGLMConfig': 'xglm',\n",
              "  'XLMConfig': 'xlm',\n",
              "  'XLMProphetNetConfig': 'xlm-prophetnet',\n",
              "  'XLMRobertaConfig': 'xlm-roberta',\n",
              "  'XLMRobertaXLConfig': 'xlm-roberta-xl',\n",
              "  'XLNetConfig': 'xlnet',\n",
              "  'YolosConfig': 'yolos',\n",
              "  'YosoConfig': 'yoso'},\n",
              " '_model_mapping': OrderedDict([('speech-encoder-decoder',\n",
              "               'SpeechEncoderDecoderModel'),\n",
              "              ('speech_to_text', 'Speech2TextForConditionalGeneration'),\n",
              "              ('whisper', 'WhisperForConditionalGeneration')]),\n",
              " '_extra_content': {},\n",
              " '_modules': {}}"
            ]
          },
          "metadata": {},
          "execution_count": 6
        }
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "id": "Ei97c4taPGL9"
      },
      "outputs": [],
      "source": [
        "def get_neuron_acts(text, layer, neuron_index):\n",
        "    # Hacky way to get out state from a single hook - we have a single element list and edit that list within the hook.\n",
        "    cache = {}\n",
        "\n",
        "    def caching_hook(act, hook):\n",
        "        cache[\"activation\"] = act[0, :, neuron_index]\n",
        "\n",
        "    model.run_with_hooks(\n",
        "        text, fwd_hooks=[(f\"blocks.{layer}.mlp.hook_post\", caching_hook)]\n",
        "    )\n",
        "    return to_numpy(cache[\"activation\"])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "H6Zw0gjdPGL-"
      },
      "source": [
        "We can run this function and verify that it gives vaguely sensible outputs"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 22,
      "metadata": {
        "id": "YVg-vSqDPGL_",
        "outputId": "c03cf3b1-3529-4b8d-8772-c7bebdc6d7d1",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 396
        }
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "['<|endoftext|>', 'The', ' following', ' is', ' a', ' list', ' of', ' powers', ' of', ' 10', ':', ' 1', ',', ' 10', ',', ' 100', ',', ' 1000', ',', ' 10000', ',', ' 100', '000', ',', ' 100', '0000', ',', ' 100', '00000']\n"
          ]
        },
        {
          "output_type": "error",
          "ename": "IndexError",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mIndexError\u001b[0m                                Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-22-096f4345f60b>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      3\u001b[0m \u001b[0mdefault_text\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m\"The following is a list of powers of 10: 1, 10, 100, 1000, 10000, 100000, 1000000, 10000000\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto_str_tokens\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdefault_text\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 5\u001b[0;31m \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mget_neuron_acts\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdefault_text\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdefault_layer\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdefault_neuron_index\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
            "\u001b[0;32m<ipython-input-11-9ef20eb9f0c2>\u001b[0m in \u001b[0;36mget_neuron_acts\u001b[0;34m(text, layer, neuron_index)\u001b[0m\n\u001b[1;32m      7\u001b[0m         \u001b[0mcache\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m\"activation\"\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mact\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m:\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mneuron_index\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      8\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 9\u001b[0;31m     model.run_with_hooks(\n\u001b[0m\u001b[1;32m     10\u001b[0m         \u001b[0mtext\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfwd_hooks\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34mf\"blocks.{layer}.mlp.hook_post\"\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcaching_hook\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     11\u001b[0m     )\n",
            "\u001b[0;32m/usr/local/lib/python3.8/dist-packages/transformer_lens/hook_points.py\u001b[0m in \u001b[0;36mrun_with_hooks\u001b[0;34m(self, fwd_hooks, bwd_hooks, reset_hooks_end, clear_contexts, *model_args, **model_kwargs)\u001b[0m\n\u001b[1;32m    187\u001b[0m                     \u001b[0;32mif\u001b[0m \u001b[0mname\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mhook_name\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    188\u001b[0m                         \u001b[0mhp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0madd_hook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mhook\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdir\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m\"bwd\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 189\u001b[0;31m         \u001b[0mout\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0mmodel_args\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mmodel_kwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    190\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mreset_hooks_end\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    191\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mbwd_hooks\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m>\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.8/dist-packages/transformer_lens/HookedTransformer.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, input, return_type, loss_per_token, prepend_bos, stop_at_layer, past_kv_cache)\u001b[0m\n\u001b[1;32m    288\u001b[0m             \u001b[0;31m# Note that each block includes skip connections, so we don't need\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    289\u001b[0m             \u001b[0;31m# residual + block(residual)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 290\u001b[0;31m             residual = block(\n\u001b[0m\u001b[1;32m    291\u001b[0m                 \u001b[0mresidual\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    292\u001b[0m                 \u001b[0mpast_kv_cache_entry\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mpast_kv_cache\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mi\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1192\u001b[0m         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks\n\u001b[1;32m   1193\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1194\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1195\u001b[0m         \u001b[0;31m# Do not call functions when jit is used\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1196\u001b[0m         \u001b[0mfull_backward_hooks\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnon_full_backward_hooks\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.8/dist-packages/transformer_lens/components.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, resid_pre, shortformer_pos_embed, past_kv_cache_entry)\u001b[0m\n\u001b[1;32m    714\u001b[0m             \u001b[0mnormalized_resid_mid\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mln2\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mresid_mid\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    715\u001b[0m             mlp_out = self.hook_mlp_out(\n\u001b[0;32m--> 716\u001b[0;31m                 \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmlp\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnormalized_resid_mid\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    717\u001b[0m             )  # [batch, pos, d_model]\n\u001b[1;32m    718\u001b[0m             resid_post = self.hook_resid_post(\n",
            "\u001b[0;32m/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1192\u001b[0m         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks\n\u001b[1;32m   1193\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1194\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1195\u001b[0m         \u001b[0;31m# Do not call functions when jit is used\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1196\u001b[0m         \u001b[0mfull_backward_hooks\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnon_full_backward_hooks\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.8/dist-packages/transformer_lens/components.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, x)\u001b[0m\n\u001b[1;32m    628\u001b[0m         )  # [batch, pos, d_mlp]\n\u001b[1;32m    629\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcfg\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mact_fn\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mendswith\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"_ln\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 630\u001b[0;31m             \u001b[0mpost_act\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mhook_post\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mact_fn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpre_act\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# [batch, pos, d_mlp]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    631\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    632\u001b[0m             \u001b[0mmid_act\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mhook_mid\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mact_fn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpre_act\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# [batch, pos, d_mlp]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1213\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0m_global_forward_hooks\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_forward_hooks\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1214\u001b[0m             \u001b[0;32mfor\u001b[0m \u001b[0mhook\u001b[0m \u001b[0;32min\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0m_global_forward_hooks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_forward_hooks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1215\u001b[0;31m                 \u001b[0mhook_result\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mhook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1216\u001b[0m                 \u001b[0;32mif\u001b[0m \u001b[0mhook_result\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1217\u001b[0m                     \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mhook_result\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.8/dist-packages/transformer_lens/hook_points.py\u001b[0m in \u001b[0;36mfull_hook\u001b[0;34m(module, module_input, module_output)\u001b[0m\n\u001b[1;32m     46\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     47\u001b[0m             \u001b[0;32mdef\u001b[0m \u001b[0mfull_hook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodule\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmodule_input\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmodule_output\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 48\u001b[0;31m                 \u001b[0;32mreturn\u001b[0m \u001b[0mhook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodule_output\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mhook\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     49\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     50\u001b[0m             \u001b[0mhandle\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mregister_forward_hook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfull_hook\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m<ipython-input-11-9ef20eb9f0c2>\u001b[0m in \u001b[0;36mcaching_hook\u001b[0;34m(act, hook)\u001b[0m\n\u001b[1;32m      5\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      6\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mcaching_hook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mact\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mhook\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 7\u001b[0;31m         \u001b[0mcache\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m\"activation\"\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mact\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m:\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mneuron_index\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      8\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      9\u001b[0m     model.run_with_hooks(\n",
            "\u001b[0;31mIndexError\u001b[0m: index 3072 is out of bounds for dimension 2 with size 3072"
          ]
        }
      ],
      "source": [
        "# default_layer = 9\n",
        "# default_neuron_index = 652\n",
        "# default_text = \"The following is a list of powers of 10: 1, 10, 100, 1000, 10000, 100000, 1000000, 10000000\"\n",
        "# print(model.to_str_tokens(default_text))\n",
        "# print(get_neuron_acts(default_text, default_layer, default_neuron_index))"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "default_layer = 0\n",
        "default_neuron_index = 3071\n",
        "default_text = \"Bill Moyers criticized the corporate media for parroting the\"\n",
        "print(model.to_str_tokens(default_text))\n",
        "print(get_neuron_acts(default_text, default_layer, default_neuron_index))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "98Dzpd_RtSDu",
        "outputId": "d636d2a2-d5b4-4bd6-e447-d0dd84282a56"
      },
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "['<|endoftext|>', 'Bill', ' Moy', 'ers', ' criticized', ' the', ' corporate', ' media', ' for', ' par', 'ro', 'ting', ' the']\n",
            "[ 0.02496125  0.10861734  2.7581763  -0.09472307 -0.13635083 -0.16790147\n",
            " -0.15908068 -0.12547089 -0.06177776 -0.10276087  0.4740647  -0.13276054\n",
            " -0.16233377]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cbGY-mscPGMA"
      },
      "source": [
        "## Visualizing Model Activations\n",
        "\n",
        "We now write some code to visualize the neuron activations on some text - we're going to hack something together which just does some string processing to make an HTML string, with each token element colored according to the intensity neuron activation. We normalize the neuron activations so they all lie in [0, 1]. You can do much better, but this is a useful proof of concept of what \"just hack stuff together\" can look like!\n",
        "\n",
        "I'll be keeping neuron 562 in layer 9 as a running example, as it seems to activate strongly on powers of 10.\n",
        "\n",
        "Note that this visualization is very sensitive to `max_val` and `min_val`! You can tune those to whatever seems reasonable for the distribution of neuron activations you care about - I generally default to `min_val=0` and `max_val` as the max activation across the dataset."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {
        "id": "CNpxK5DmPGMA"
      },
      "outputs": [],
      "source": [
        "# This is some CSS (tells us what style )to give each token a thin gray border, to make it easy to see token separation\n",
        "style_string = \"\"\"<style> \n",
        "    span.token {\n",
        "        border: 1px solid rgb(123, 123, 123)\n",
        "        } \n",
        "    </style>\"\"\"\n",
        "\n",
        "\n",
        "def calculate_color(val, max_val, min_val):\n",
        "    # Hacky code that takes in a value val in range [min_val, max_val], normalizes it to [0, 1] and returns a color which interpolates between slightly off-white and red (0 = white, 1 = red)\n",
        "    # We return a string of the form \"rgb(240, 240, 240)\" which is a color CSS knows\n",
        "    normalized_val = (val - min_val) / max_val\n",
        "    return f\"rgb(240, {240*(1-normalized_val)}, {240*(1-normalized_val)})\"\n",
        "\n",
        "\n",
        "def basic_neuron_vis(text, layer, neuron_index, max_val=None, min_val=None):\n",
        "    \"\"\"\n",
        "    text: The text to visualize\n",
        "    layer: The layer index\n",
        "    neuron_index: The neuron index\n",
        "    max_val: The top end of our activation range, defaults to the maximum activation\n",
        "    min_val: The top end of our activation range, defaults to the minimum activation\n",
        "\n",
        "    Returns a string of HTML that displays the text with each token colored according to its activation\n",
        "\n",
        "    Note: It's useful to be able to input a fixed max_val and min_val, because otherwise the colors will change as you edit the text, which is annoying.\n",
        "    \"\"\"\n",
        "    if layer is None:\n",
        "        return \"Please select a Layer\"\n",
        "    if neuron_index is None:\n",
        "        return \"Please select a Neuron\"\n",
        "    acts = get_neuron_acts(text, layer, neuron_index)\n",
        "    print('acts ', acts , '\\n')\n",
        "    act_max = acts.max()\n",
        "    act_min = acts.min()\n",
        "    print('act_max :', act_max,'\\n act_mix : ', act_min )\n",
        "    # Defaults to the max and min of the activations\n",
        "    if max_val is None:\n",
        "        max_val = act_max\n",
        "    if min_val is None:\n",
        "        min_val = act_min\n",
        "    # We want to make a list of HTML strings to concatenate into our final HTML string\n",
        "    # We first add the style to make each token element have a nice border\n",
        "    htmls = [style_string]\n",
        "    # We then add some text to tell us what layer and neuron we're looking at \n",
        "    # - we're just dealing with strings and can use f-strings as normal\n",
        "    # h4 means \"small heading\"\n",
        "    htmls.append(f\"<h4>Layer: <b>{layer}</b>. Neuron Index: <b>{neuron_index}</b></h4>\")\n",
        "    # We then add a line telling us the limits of our range\n",
        "    htmls.append(\n",
        "        f\"<h4>Max Range: <b>{max_val:.4f}</b>. Min Range: <b>{min_val:.4f}</b></h4>\"\n",
        "    )\n",
        "    # If we added a custom range, print a line telling us the range of our activations too.\n",
        "    if act_max != max_val or act_min != min_val:\n",
        "        htmls.append(\n",
        "            f\"<h4>Custom Range Set. Max Act: <b>{act_max:.4f}</b>. Min Act: <b>{act_min:.4f}</b></h4>\"\n",
        "        )\n",
        "    # Convert the text to a list of tokens\n",
        "    str_tokens = model.to_str_tokens(text)\n",
        "    print('str_tokens', str_tokens)\n",
        "    for tok, act in zip(str_tokens, acts):\n",
        "        # A span is an HTML element that lets us style a part of a string (and remains on the same line by default)\n",
        "        # We set the background color of the span to be the color we calculated from the activation\n",
        "        # We set the contents of the span to be the token\n",
        "        htmls.append(\n",
        "            f\"<span class='token' style='background-color:{calculate_color(act, max_val, min_val)}' >{tok}</span>\"\n",
        "        )\n",
        "\n",
        "    return \"\".join(htmls)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {
        "id": "4j2Q39r2PGMB",
        "outputId": "11f6fb80-0c1d-4980-bb59-84649a2f2075",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 420
        }
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "acts  [ 0.02496125  0.10861734  2.7581763  -0.09472307 -0.13635083 -0.16790147\n",
            " -0.15908068 -0.12547089 -0.06177776 -0.10276087  0.4740647  -0.13276054\n",
            " -0.16233377] \n",
            "\n",
            "act_max : 2.7581763 \n",
            " act_mix :  -0.16790147\n",
            "str_tokens ['<|endoftext|>', 'Bill', ' Moy', 'ers', ' criticized', ' the', ' corporate', ' media', ' for', ' par', 'ro', 'ting', ' the']\n",
            "Displayed HTML\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              "<style> \n",
              "    span.token {\n",
              "        border: 1px solid rgb(123, 123, 123)\n",
              "        } \n",
              "    </style><h4>Layer: <b>0</b>. Neuron Index: <b>3071</b></h4><h4>Max Range: <b>4.0000</b>. Min Range: <b>0.0000</b></h4><h4>Custom Range Set. Max Act: <b>2.7582</b>. Min Act: <b>-0.1679</b></h4><span class='token' style='background-color:rgb(240, 238.50232522934675, 238.50232522934675)' ><|endoftext|></span><span class='token' style='background-color:rgb(240, 233.48295986652374, 233.48295986652374)' >Bill</span><span class='token' style='background-color:rgb(240, 74.50942039489746, 74.50942039489746)' > Moy</span><span class='token' style='background-color:rgb(240, 245.683384090662, 245.683384090662)' >ers</span><span class='token' style='background-color:rgb(240, 248.18104952573776, 248.18104952573776)' > criticized</span><span class='token' style='background-color:rgb(240, 250.0740882754326, 250.0740882754326)' > the</span><span class='token' style='background-color:rgb(240, 249.54484105110168, 249.54484105110168)' > corporate</span><span class='token' style='background-color:rgb(240, 247.5282534956932, 247.5282534956932)' > media</span><span class='token' style='background-color:rgb(240, 243.70666556060314, 243.70666556060314)' > for</span><span class='token' style='background-color:rgb(240, 246.16565197706223, 246.16565197706223)' > par</span><span class='token' style='background-color:rgb(240, 211.55611753463745, 211.55611753463745)' >ro</span><span class='token' style='background-color:rgb(240, 247.96563237905502, 247.96563237905502)' >ting</span><span class='token' style='background-color:rgb(240, 249.7400262951851, 249.7400262951851)' > the</span>"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "HTML String - it's just raw HTML code!\n",
            "<style> \n",
            "    span.token {\n",
            "        border: 1px solid rgb(123, 123, 123)\n",
            "        } \n",
            "    </style><h4>Layer: <b>0</b>. Neuron Index: <b>3071</b></h4><h4>Max Range: <b>4.0000</b>. Min Range: <b>0.0000</b></h4><h4>Custom Range Set. Max Act: <b>2.7582</b>. Min Act: <b>-0.1679</b></h4><span class='token' style='background-color:rgb(240, 238.50232522934675, 238.50232522934675)' ><|endoftext|></span><span class='token' style='background-color:rgb(240, 233.48295986652374, 233.48295986652374)' >Bill</span><span class='token' style='background-color:rgb(240, 74.50942039489746, 74.50942039489746)' > Moy</span><span class='token' style='background-color:rgb(240, 245.683384090662, 245.683384090662)' >ers</span><span class='token' style='background-color:rgb(240, 248.18104952573776, 248.18104952573776)' > criticized</span><span class='token' style='background-color:rgb(240, 250.0740882754326, 250.0740882754326)' > the</span><span class='token' style='background-color:rgb(240, 249.54484105110168, 249.54484105110168)' > corporate</span><span class='token' style='background-color:rgb(240, 247.5282534956932, 247.5282534956932)' > media</span><span class='token' style='background-color:rgb(240, 243.70666556060314, 243.70666556060314)' > for</span><span class='token' style='background-color:rgb(240, 246.16565197706223, 246.16565197706223)' > par</span><span class='token' style='background-color:rgb(240, 211.55611753463745, 211.55611753463745)' >ro</span><span class='token' style='background-color:rgb(240, 247.96563237905502, 247.96563237905502)' >ting</span><span class='token' style='background-color:rgb(240, 249.7400262951851, 249.7400262951851)' > the</span>\n"
          ]
        }
      ],
      "source": [
        "# The function outputs a string of HTML\n",
        "default_max_val = 4.0\n",
        "default_min_val = 0.0\n",
        "default_html_string = basic_neuron_vis(\n",
        "    default_text,\n",
        "    default_layer,\n",
        "    default_neuron_index,\n",
        "    max_val=default_max_val,\n",
        "    min_val=default_min_val,\n",
        ")\n",
        "\n",
        "# IPython lets us display HTML\n",
        "print(\"Displayed HTML\")\n",
        "display(HTML(default_html_string))\n",
        "\n",
        "# We can also print the string directly\n",
        "print(\"HTML String - it's just raw HTML code!\")\n",
        "print(default_html_string)"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "print()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "39LHLldx56ag",
        "outputId": "5ba90f55-068e-45c0-f78d-02ad67f1ff8e"
      },
      "execution_count": 15,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BW161buKPGMC"
      },
      "source": [
        "## Create Interactive UI\n",
        "\n",
        "We now put all these together to create an interactive visualization in Gradio! \n",
        "\n",
        "The internal format is that there's a bunch of elements - Textboxes, Numbers, etc which the user can interact with and which return strings and numbers. And we can also define output elements that just display things - in this case, one which takes in an arbitrary HTML string. We call `input.change(update_function, inputs, output)` - this says \"if that input element changes, run the update function on the value of each of the elements in `inputs` and set the value of `output` to the output of the function\". As a bonus, this gives us live interactivity!\n",
        "\n",
        "This is also more complex than a typical Gradio intro example - I wanted to use custom HTML to display the nice colours, which made things much messier! Normally you could just make `out` into another Textbox and pass it a string."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "metadata": {
        "id": "LV5dM9mGPGMC"
      },
      "outputs": [],
      "source": [
        "# The `with gr.Blocks() as demo:` syntax just creates a variable called demo containing all these components\n",
        "with gr.Blocks() as demo:\n",
        "    gr.HTML(value=f\"Hacky Interactive Neuroscope for {model_name}\")\n",
        "    # The input elements\n",
        "    with gr.Row():\n",
        "        with gr.Column():\n",
        "            text = gr.Textbox(label=\"Text\", value=default_text)\n",
        "            # Precision=0 makes it an int, otherwise it's a float\n",
        "            # Value sets the initial default value\n",
        "            layer = gr.Number(label=\"Layer\", value=default_layer, precision=0)\n",
        "            neuron_index = gr.Number(\n",
        "                label=\"Neuron Index\", value=default_neuron_index, precision=0\n",
        "            )\n",
        "            # If empty, these two map to None\n",
        "            max_val = gr.Number(label=\"Max Value\", value=default_max_val)\n",
        "            min_val = gr.Number(label=\"Min Value\", value=default_min_val)\n",
        "            inputs = [text, layer, neuron_index, max_val, min_val]\n",
        "        with gr.Column():\n",
        "            # The output element\n",
        "            out = gr.HTML(label=\"Neuron Acts\", value=default_html_string)\n",
        "    for inp in inputs:\n",
        "        inp.change(basic_neuron_vis, inputs, out)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IVOqJlhcPGMD"
      },
      "source": [
        "We can now launch our demo element, and we're done! The setting share=True even gives you a public link to the demo (though it just redirects to the backend run by this notebook, and will go away once you turn the notebook off!) Sharing makes it much slower, and can be turned off if you aren't in a colab.\n",
        "\n",
        "**Exercise:** Explore where this neuron does and does not activate. Is it just powers of ten? Just comma separated numbers? Numbers in any particular sequence?"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 18,
      "metadata": {
        "id": "OhjuVBj0PGME",
        "outputId": "c85f1872-b22c-4a5e-e4eb-6a866f07209f",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        }
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Rerunning server... use `close()` to stop if you need to change `launch()` parameters.\n",
            "----\n",
            "Colab notebook detected. To show errors in colab notebook, set debug=True in launch()\n",
            "Running on public URL: https://752d6628-8f03-4110.gradio.live\n",
            "\n",
            "This share link expires in 72 hours. For free permanent hosting and GPU upgrades (NEW!), check out Spaces: https://huggingface.co/spaces\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              "<div><iframe src=\"https://752d6628-8f03-4110.gradio.live\" width=\"100%\" height=\"1000\" allow=\"autoplay; camera; microphone; clipboard-read; clipboard-write;\" frameborder=\"0\" allowfullscreen></iframe></div>"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": []
          },
          "metadata": {},
          "execution_count": 18
        }
      ],
      "source": [
        "demo.launch(share=True, height=1000)"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3.7.13 ('base')",
      "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.7.13"
    },
    "orig_nbformat": 4,
    "vscode": {
      "interpreter": {
        "hash": "d4d1e4263499bec80672ea0156c357c1ee493ec2b1c70f0acce89fc37c4a6abe"
      }
    },
    "colab": {
      "provenance": []
    },
    "accelerator": "GPU",
    "gpuClass": "standard"
  },
  "nbformat": 4,
  "nbformat_minor": 0
}