NNsight_Walkthrough.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "i4WnS0_mkkww"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/ndif-team/nnsight/blob/main/NNsight_Walkthrough.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-Scfr942GwO4"
      },
      "source": [
        "<img src=\"https://nnsight.net/_static/images/nnsight_logo.svg\" alt=\"drawing\" width=\"200\"/>\n",
        "\n",
        "# **NNsight**\n",
        "\n",
        "## The API for a transparent science on black-box AI\n",
        "\n",
        "In this era of large-scale deep learning, the most interesting AI models are\n",
        "massive black boxes that are hard to run. Ordinary commercial inference service\n",
        "APIs let us interact with huge models, but they do not let us access model\n",
        "internals.\n",
        "\n",
        "The `nnsight` library is different: it provides full access to all the neural\n",
        "network internals. When used together with a remote service like the\n",
        "[National Deep Inference Fabric](https://thevisible.net/docs/NDIF-proposal.pdf)\n",
        "(NDIF), it makes possible to run complex experiments on huge open  models easily,\n",
        "with fully transparent access.\n",
        "\n",
        "Our team wants to enable entire labs and independent researchers alike, as we\n",
        "believe a large, passionate, and collaborative community will produce the next\n",
        "big insights on a profoundly important field."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "N1OemD2VGyZx"
      },
      "source": [
        "# 1 First, let's start small"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Install nnsight\n",
        "!pip install nnsight\n",
        "!pip install --upgrade transformers torch\n",
        "\n",
        "from IPython.display import clear_output\n",
        "\n",
        "clear_output()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WyLMmhrAKTNM"
      },
      "source": [
        "## Tracing Context\n",
        "\n",
        "To demonstrate the core functionality and syntax of nnsight, we'll define and\n",
        "use a tiny two layer neural network.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bgydw7i3HmIH"
      },
      "source": [
        "Our little model here is composed of two submodules – linear layers 'layer1' and 'layer2'. We specify the sizes of each of these modules and create\n",
        "some complementary example input."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "2pX2Wg8Ceo6N"
      },
      "outputs": [],
      "source": [
        "from collections import OrderedDict\n",
        "import torch\n",
        "\n",
        "input_size = 5\n",
        "hidden_dims = 10\n",
        "output_size = 2\n",
        "\n",
        "net = torch.nn.Sequential(\n",
        "    OrderedDict(\n",
        "        [\n",
        "            (\"layer1\", torch.nn.Linear(input_size, hidden_dims)),\n",
        "            (\"layer2\", torch.nn.Linear(hidden_dims, output_size)),\n",
        "        ]\n",
        "    )\n",
        ").requires_grad_(False)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tIPa2h2pJIwl"
      },
      "source": [
        "The core object of the nnsight package is `NNsight`. This wraps around a given\n",
        "PyTorch model to enable investigation of its internal parameters.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "8H9R_ynTJI5y"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "/opt/anaconda3/envs/nnsight/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
            "  from .autonotebook import tqdm as notebook_tqdm\n"
          ]
        }
      ],
      "source": [
        "import nnsight\n",
        "from nnsight import NNsight\n",
        "\n",
        "tiny_model = NNsight(net)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2NfISlQ_Ilvp"
      },
      "source": [
        "Printing a PyTorch model shows a named hierarchy of modules which is very useful\n",
        "when accessing sub-components directly. NNsight reflect the same hierarchy and can be similarly printed.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "bYtnbJHvlGZV",
        "outputId": "9bec5a3e-647c-4c49-b4ec-e1420ea0c964"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Sequential(\n",
            "  (layer1): Linear(in_features=5, out_features=10, bias=True)\n",
            "  (layer2): Linear(in_features=10, out_features=2, bias=True)\n",
            ")\n"
          ]
        }
      ],
      "source": [
        "print(tiny_model)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "djC5kyJWLuUH"
      },
      "source": [
        "Before we actually get to using the model we just created, let's talk about\n",
        "Python contexts.\n",
        "\n",
        "Python contexts define a scope using the `with` statement and are often used to\n",
        "create some object, or initiate some logic, that you later want to destroy or\n",
        "conclude.\n",
        "\n",
        "The most common application is opening files as in the following example:\n",
        "\n",
        "```python\n",
        "with open('myfile.txt', 'r') as file:\n",
        "  text = file.read()\n",
        "```\n",
        "\n",
        "Python uses the `with` keyword to enter a context-like object. This object\n",
        "defines logic to be run at the start of the `with` block, as well as logic to be\n",
        "run when exiting. When using `with` for a file, entering the context opens the\n",
        "file and exiting the context closes it. Being within the context means we can\n",
        "read from the file.\n",
        "\n",
        "Simple enough! Now we can discuss how `nnsight` uses\n",
        "contexts to enable intuitive access into the internals of a neural network.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iNvuCOeyojcA"
      },
      "source": [
        "The main tool with `nnsight` is a context for tracing.\n",
        "\n",
        "We enter the tracing context by calling `model.trace(<input>)` on an `NNsight`\n",
        "model, which defines how we want to run the model. Inside the context, we will\n",
        "be able to customize how the neural network runs. The model is actually run upon\n",
        "exiting the tracing context.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "qEXQ4auPSL-m"
      },
      "outputs": [],
      "source": [
        "# random input\n",
        "input = torch.rand((1, input_size))\n",
        "\n",
        "with tiny_model.trace(input) as tracer:\n",
        "    pass"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UZQsHinjqicJ"
      },
      "source": [
        "But where's the output? To get that, we'll have to learn how to request it from\n",
        "within the tracing context.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xMfBpYzDPMoB"
      },
      "source": [
        "## Getting\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "B5_aFwFRv0ax"
      },
      "source": [
        "Earlier, when we wrapped our little neural net with the `NNsight` class. This\n",
        "added a couple properties to each module in the model (including the root model\n",
        "itself). The two most important ones are `.input` and `.output`.\n",
        "\n",
        "```python\n",
        "model.input\n",
        "model.output\n",
        "```\n",
        "\n",
        "The names are self explanatory. They correspond to the inputs and outputs of\n",
        "their respective modules during a forward pass of the model. We can use these\n",
        "attributes inside the `with` block.\n",
        "\n",
        "However, it is important to understand that the model is not executed until the\n",
        "end of the tracing context. How can we access inputs and outputs before the\n",
        "model is run? The trick is deferred execution.\n",
        "\n",
        "`.input` and `.output` are Proxies for the eventual inputs and outputs of a\n",
        "module. In other words, when we access `model.output` what we are\n",
        "communicating to `nnsight` is, \"When you compute the output of `model`, please\n",
        "grab it for me and put the value into its corresponding Proxy object. Let's try it:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 329
        },
        "id": "cYe8-r9ptGaG",
        "outputId": "5a1ebb19-be61-4eea-9ec1-141a62ebc313"
      },
      "outputs": [
        {
          "ename": "ValueError",
          "evalue": "Accessing value before it's been set.",
          "output_type": "error",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
            "Cell \u001b[0;32mIn[5], line 5\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m tiny_model\u001b[38;5;241m.\u001b[39mtrace(\u001b[38;5;28minput\u001b[39m) \u001b[38;5;28;01mas\u001b[39;00m tracer:\n\u001b[1;32m      3\u001b[0m     output \u001b[38;5;241m=\u001b[39m tiny_model\u001b[38;5;241m.\u001b[39moutput\n\u001b[0;32m----> 5\u001b[0m \u001b[38;5;28;43mprint\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43moutput\u001b[49m\u001b[43m)\u001b[49m\n",
            "File \u001b[0;32m/opt/anaconda3/envs/nnsight/lib/python3.10/site-packages/nnsight/tracing/graph/proxy.py:70\u001b[0m, in \u001b[0;36mProxy.__str__\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m     66\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__str__\u001b[39m(\u001b[38;5;28mself\u001b[39m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m \u001b[38;5;28mstr\u001b[39m:\n\u001b[1;32m     68\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mnode\u001b[38;5;241m.\u001b[39mattached:\n\u001b[0;32m---> 70\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mstr\u001b[39m(\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mvalue\u001b[49m)\n\u001b[1;32m     72\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mtype\u001b[39m(\u001b[38;5;28mself\u001b[39m)\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m (\u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mnode\u001b[38;5;241m.\u001b[39mtarget\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m)\u001b[39m\u001b[38;5;124m\"\u001b[39m\n",
            "File \u001b[0;32m/opt/anaconda3/envs/nnsight/lib/python3.10/site-packages/nnsight/tracing/graph/proxy.py:64\u001b[0m, in \u001b[0;36mProxy.value\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m     56\u001b[0m \u001b[38;5;129m@property\u001b[39m\n\u001b[1;32m     57\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mvalue\u001b[39m(\u001b[38;5;28mself\u001b[39m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Any:\n\u001b[1;32m     58\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124;03m\"\"\"Property to return the value of this proxy's node.\u001b[39;00m\n\u001b[1;32m     59\u001b[0m \n\u001b[1;32m     60\u001b[0m \u001b[38;5;124;03m    Returns:\u001b[39;00m\n\u001b[1;32m     61\u001b[0m \u001b[38;5;124;03m        Any: The stored value of the proxy, populated during execution of the model.\u001b[39;00m\n\u001b[1;32m     62\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[0;32m---> 64\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mnode\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mvalue\u001b[49m\n",
            "File \u001b[0;32m/opt/anaconda3/envs/nnsight/lib/python3.10/site-packages/nnsight/tracing/graph/node.py:143\u001b[0m, in \u001b[0;36mNode.value\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    133\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"Property to return the value of this node.\u001b[39;00m\n\u001b[1;32m    134\u001b[0m \n\u001b[1;32m    135\u001b[0m \u001b[38;5;124;03mReturns:\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    139\u001b[0m \u001b[38;5;124;03m    ValueError: If the underlying ._value is inspect._empty (therefore never set or was destroyed).\u001b[39;00m\n\u001b[1;32m    140\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    142\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdone:\n\u001b[0;32m--> 143\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mAccessing value before it\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124ms been set.\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m    145\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_value\n",
            "\u001b[0;31mValueError\u001b[0m: Accessing value before it's been set."
          ]
        }
      ],
      "source": [
        "with tiny_model.trace(input) as tracer:\n",
        "\n",
        "    output = tiny_model.output\n",
        "\n",
        "print(output)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4CGa5zSI6kkI"
      },
      "source": [
        "Oh no an error! \"Accessing value before it's been set.\"\n",
        "\n",
        "Why doesn't our `output` have a `value`?\n",
        "\n",
        "Proxy objects will only have their value at the end of a context if we call\n",
        "`.save()` on them. This helps to reduce memory costs. Adding `.save()` fixes the\n",
        "error:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 67,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "y_bXRd5dvsBu",
        "outputId": "feded172-02f5-429b-f9a9-766eed3a787d"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "tensor([[-0.1301, -0.4906]])\n"
          ]
        }
      ],
      "source": [
        "with tiny_model.trace(input) as tracer:\n",
        "\n",
        "    output = tiny_model.output.save()\n",
        "\n",
        "print(output)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "t5C0UZCwvrYn"
      },
      "source": [
        "Success! We now have the model output. We just completed out first\n",
        "intervention using `nnsight`.\n",
        "\n",
        "Each time we access a module's input or output, we create an _intervention_ in\n",
        "the neural network's forward pass. Collectively these requests form the\n",
        "_intervention graph_. We call the process of executing it alongside the model's\n",
        "normal computation graph, _interleaving_.\n",
        "\n",
        "<details>\n",
        "<summary>On Model output</summary>\n",
        "\n",
        "---\n",
        "\n",
        "If we don't need to access anything other than the final model output, we can\n",
        "call the tracing context with `trace=False` and not use it as a context. This could be especially useful for easy remote inference.\n",
        "\n",
        "```python\n",
        "  output = model.trace(<inputs>, trace=False)\n",
        "```\n",
        "\n",
        "---\n",
        "\n",
        "</details>\n",
        "\n",
        "Just like we saved the output of the model as a whole, we can save the output of\n",
        "any of its submodules. We use normal Python attribute syntax. We can discover\n",
        "how to access them by name by printing out the model:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 68,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0WcVUSP-0CJi",
        "outputId": "f40600e1-719d-46c2-f420-e5a9826dd315"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Sequential(\n",
            "  (layer1): Linear(in_features=5, out_features=10, bias=True)\n",
            "  (layer2): Linear(in_features=10, out_features=2, bias=True)\n",
            ")\n"
          ]
        }
      ],
      "source": [
        "print(tiny_model)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MXn9XRQqkkw1"
      },
      "source": [
        "Let's access the output of the first layer (which we've named 'layer1'):"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 69,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "akWr-cNqy-9O",
        "outputId": "13faddb0-fead-4fd0-a500-2bff00c095b3"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "tensor([[ 0.2732,  0.2355, -0.6433,  0.0475,  0.0904,  0.4407, -0.3099,  1.4903,\n",
            "         -0.0748,  0.0088]])\n"
          ]
        }
      ],
      "source": [
        "with tiny_model.trace(input) as tracer:\n",
        "\n",
        "    l1_output = tiny_model.layer1.output.save()\n",
        "\n",
        "print(l1_output)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "85A-aP_03ht6"
      },
      "source": [
        "Let's do the same for the input of layer2. While we're at it, let's also drop\n",
        "the `as tracer`, as we won't be needing the tracer object itself for a few\n",
        "sections:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 70,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "3EHEN38N3nXR",
        "outputId": "74f25b48-89fd-4b15-9dab-cead54f2386b"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "tensor([[ 0.2732,  0.2355, -0.6433,  0.0475,  0.0904,  0.4407, -0.3099,  1.4903,\n",
            "         -0.0748,  0.0088]])\n"
          ]
        }
      ],
      "source": [
        "with tiny_model.trace(input):\n",
        "\n",
        "    l2_input = tiny_model.layer2.input.save()\n",
        "\n",
        "print(l2_input)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jk-U8zi33Gi-"
      },
      "source": [
        "<details>\n",
        "  <summary>On module inputs</summary>\n",
        "\n",
        "---\n",
        "\n",
        "Notice how the value for `l2_input`, is just a single tensor. By default, the `.input` attribute of a module will return the **first** tensor input to the module.\n",
        "\n",
        "We can also access the full input to a module by using the `.inputs` attribute which will return the values in the form of:\n",
        "\n",
        "      tuple(tuple(args), dictionary(kwargs))\n",
        "\n",
        "Where the first index of the tuple is itself a tuple of all positional\n",
        "arguments, and the second index is a dictionary of the keyword arguments.\n",
        "\n",
        "---\n",
        "\n",
        "</details>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d5DE1ZJAkkw1"
      },
      "source": [
        "Until now we were saving the output of the model and its submodules within the `Trace` context to then print it after exiting the context. We will continuing doing this in the rest of the tutorial since it's a good practice to save the computation results for later analysis.\n",
        "\n",
        "However, we can also log the outputs of the model and its submodules within the `Trace` context. This is useful for debugging and understanding the model's behavior while saving memory. Let's see how to do this:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 71,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "oZzS86tzkkw1",
        "outputId": "1c1f3a94-eb4d-48c7-c69f-9241f2fa08f5"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Layer 1 - out:  tensor([[ 0.2732,  0.2355, -0.6433,  0.0475,  0.0904,  0.4407, -0.3099,  1.4903,\n",
            "         -0.0748,  0.0088]])\n"
          ]
        }
      ],
      "source": [
        "with tiny_model.trace(input) as tracer:\n",
        "  tracer.log(\"Layer 1 - out: \", tiny_model.layer1.output)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "p7Xo_PHyPr4p"
      },
      "source": [
        "## Functions, Methods, and Operations\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ehSofWbx5DSx"
      },
      "source": [
        "Now that we can access activations, we also want to do some post-processing on\n",
        "it. Let's find out which dimension of layer1's output has the highest value.\n",
        "\n",
        "We could do this by calling `torch.argmax(...)` after the tracing context or we\n",
        "can just leverage the fact that `nnsight` handles Pytorch functions and methods within\n",
        "the tracing context, by creating a Proxy request for it:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 72,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "c5XCiSZn2p3k",
        "outputId": "c6c5060b-4cd9-4154-8bc6-adf4f5aec5cc"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "tensor(7)\n"
          ]
        }
      ],
      "source": [
        "with tiny_model.trace(input):\n",
        "\n",
        "    # Note we don't need to call .save() on the output,\n",
        "    # as we're only using its value within the tracing context.\n",
        "    l1_output = tiny_model.layer1.output\n",
        "\n",
        "    # We do need to save the argmax tensor however,\n",
        "    # as we're using it outside the tracing context.\n",
        "    l1_amax = torch.argmax(l1_output, dim=1).save()\n",
        "\n",
        "print(l1_amax[0])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BGPUJvWq_bOp"
      },
      "source": [
        "Nice! That worked seamlessly, but hold on, how come we didn't need to call\n",
        "`.value[0]` on the result? In previous sections, we were just being explicit to\n",
        "get an understanding of Proxies and their value. In practice, however, `nnsight`\n",
        "knows that when outside of the tracing context we only care about the actual\n",
        "value, and so printing, indexing, and applying functions all immediately return\n",
        "and reflect the data in `.value`. So for the rest of the tutorial we won't use\n",
        "it.\n",
        "\n",
        "The same principles work for Pytorch methods and all operators as well:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 73,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "FIcOYeEjFuln",
        "outputId": "074fd78f-bc18-48b8-ccef-7db9df756ea6"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "tensor(0.9377)\n"
          ]
        }
      ],
      "source": [
        "with tiny_model.trace(input):\n",
        "\n",
        "    value = (tiny_model.layer1.output.sum() + tiny_model.layer2.output.sum()).save()\n",
        "\n",
        "print(value)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "V0vNOJtJ2oFR"
      },
      "source": [
        "The code block above is saying to `nnsight`, \"Run the model with\n",
        "the given `input`. When the output of `tiny_model.layer1` is computed, take its sum. Then do\n",
        "the same for `tiny_model.layer2`. Now that both of those are computed, add them and make sure\n",
        "not to delete this value as I wish to use it outside of the tracing context.\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kLmzsHaMkkw2"
      },
      "source": [
        "## Custom Functions"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GJFCuZHtkkw2"
      },
      "source": [
        "Everything within the tracing context operates on the intervention graph. Therefore, for `nnsight` to trace a  function it must also be a part of the intervention graph.\n",
        "\n",
        "Out-of-the-box `nnsight` supports PyTorch functions and methods, all operators, as well the `einops` library. We don't need to do anything special to use them. But what do we do if we want to use custom functions? How do we add them to the intervention graph?\n",
        "\n",
        "Enter `nnsight.apply()`. It allows us to add new functions to the intervention graph. Let's see how it works:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 74,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "DF7MPuhckkw2",
        "outputId": "c263468b-affe-4f69-a045-d674e07333d3"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "tensor(1.5584) tensor(1.5584)\n"
          ]
        }
      ],
      "source": [
        "# Take a tensor and return the sum of its elements\n",
        "def tensor_sum(tensor):\n",
        "    flat = tensor.flatten()\n",
        "    total = 0\n",
        "    for element in flat:\n",
        "        total += element.item()\n",
        "\n",
        "    return torch.tensor(total)\n",
        "\n",
        "with tiny_model.trace(input) as tracer:\n",
        "\n",
        "    # Specify the function name and its arguments (in a comma-separated form) to add to the intervention graph\n",
        "    custom_sum = nnsight.apply(tensor_sum, tiny_model.layer1.output).save()\n",
        "    sum = tiny_model.layer1.output.sum()\n",
        "    sum.save()\n",
        "\n",
        "\n",
        "print(custom_sum, sum)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "02UKwk_Lkkw2"
      },
      "source": [
        "`nnsight.apply()` executes the function it wraps and returns its output as a Proxy object. We can then use this Proxy object as we would any other.\n",
        "\n",
        "The applications of `nnsight.apply` are wide: it can be used to wrap any custom function or functions from libraries that `nnsight` does not support out-of-the-box."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "m_5qH5gHPOT_"
      },
      "source": [
        "## Setting\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rgju-b_IOLlq"
      },
      "source": [
        "Getting and analyzing the activations from various points in a model can be\n",
        "really insightful, and a number of ML techniques do exactly that. However, often we not only want to view the computation of a model, but also to influence it.\n",
        "\n",
        "To demonstrate the effect of editing the flow of information through the model,\n",
        "let's set the first dimension of the first layer's output to 0. `NNsight` makes\n",
        "this really easy using the '=' operator:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 75,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "-6y2wzJqOz3a",
        "outputId": "13de53a9-c396-4167-e2f6-950957233c8e"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Before: tensor([[ 0.2732,  0.2355, -0.6433,  0.0475,  0.0904,  0.4407, -0.3099,  1.4903,\n",
            "         -0.0748,  0.0088]])\n",
            "After: tensor([[ 0.0000,  0.2355, -0.6433,  0.0475,  0.0904,  0.4407, -0.3099,  1.4903,\n",
            "         -0.0748,  0.0088]])\n"
          ]
        }
      ],
      "source": [
        "with tiny_model.trace(input):\n",
        "\n",
        "    # Save the output before the edit to compare.\n",
        "    # Notice we apply .clone() before saving as the setting operation is in-place.\n",
        "    l1_output_before = tiny_model.layer1.output.clone().save()\n",
        "\n",
        "    # Access the 0th index of the hidden state dimension and set it to 0.\n",
        "    tiny_model.layer1.output[:, 0] = 0\n",
        "\n",
        "    # Save the output after to see our edit.\n",
        "    l1_output_after = tiny_model.layer1.output.save()\n",
        "\n",
        "print(\"Before:\", l1_output_before)\n",
        "print(\"After:\", l1_output_after)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oZz9SMs3Y_iS"
      },
      "source": [
        "Seems our change was reflected. Now let's do the same for the last dimension:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 76,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 305
        },
        "id": "qwlqvHFcxld2",
        "outputId": "b1ac58fb-dbbd-4375-fcd7-3362762f00fe"
      },
      "outputs": [
        {
          "ename": "NNsightError",
          "evalue": "index 10 is out of bounds for dimension 1 with size 10",
          "output_type": "error",
          "traceback": [
            "Traceback (most recent call last):",
            "  File \"/Users/emmabortz/Documents/Projects/nnsight/src/nnsight/tracing/graph/node.py\", line 297, in execute",
            "    output = self.target(*args, **kwargs)",
            "IndexError: index 10 is out of bounds for dimension 1 with size 10",
            "",
            "During handling of the above exception, another exception occurred:",
            "",
            "Traceback (most recent call last):",
            "  File \"/var/folders/rx/0nl_h2cd54q90chs9hf0n5qr0000gn/T/ipykernel_80599/3404137504.py\", line 8, in <module>",
            "    tiny_model.layer1.output[:, hidden_dims] = 0",
            "",
            "NNsightError: index 10 is out of bounds for dimension 1 with size 10"
          ]
        }
      ],
      "source": [
        "with tiny_model.trace(input):\n",
        "\n",
        "    # Save the output before the edit to compare.\n",
        "    # Notice we apply .clone() before saving as the setting operation is in-place.\n",
        "    l1_output_before = tiny_model.layer1.output.clone().save()\n",
        "\n",
        "    # Access the last index of the hidden state dimension and set it to 0.\n",
        "    tiny_model.layer1.output[:, hidden_dims] = 0\n",
        "\n",
        "    # Save the output after to see our edit.\n",
        "    l1_output_after = tiny_model.layer1.output.save()\n",
        "\n",
        "print(\"Before:\", l1_output_before)\n",
        "print(\"After:\", l1_output_after)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DKaN1oi76djp"
      },
      "source": [
        "Oh no, we are getting an error! Ah of course, we needed to index at `hidden_dims - 1` not `hidden_dims`.\n",
        "\n",
        "If you've been using `nnsight`, you are probably familiar with error messages that can be quite difficult to troubleshoot. In `nnsight 0.4` we've now improved error messaging to be descriptive and line-specific, as you should see in the above example!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Yz-XpUpZ36yd"
      },
      "source": [
        "<details>\n",
        "\n",
        "<summary>\n",
        "Old NNsight error messaging\n",
        "</summary>\n",
        "\n",
        "If you've been using NNsight prior to the NNsight 0.4 release, you will be familiar with the following non-descriptive error messaging. If you choose to turn off NNsight 0.4's new error messaging feature, this is how errors within the tracing context will appear.\n",
        "\n",
        "\n",
        "```\n",
        "---------------------------------------------------------------------------\n",
        "IndexError                                Traceback (most recent call last)\n",
        "/usr/local/lib/python3.11/dist-packages/nnsight/tracing/Node.py in execute(self)\n",
        "    379                 # Call the target to get value.\n",
        "--> 380                 output = self.target(*args, **kwargs)\n",
        "    381\n",
        "\n",
        "IndexError: index 10 is out of bounds for dimension 1 with size 10\n",
        "\n",
        "The above exception was the direct cause of the following exception:\n",
        "\n",
        "IndexError                                Traceback (most recent call last)\n",
        "20 frames\n",
        "<ipython-input-16-5c81de91fb1f> in <cell line: 0>()\n",
        "----> 1 with tiny_model.trace(input):\n",
        "      2\n",
        "      3     # Save the output before the edit to compare.\n",
        "      4     # Notice we apply .clone() before saving as the setting operation is in-place.\n",
        "      5     l1_output_before = tiny_model.layer1.output.clone().save()\n",
        "\n",
        "/usr/local/lib/python3.11/dist-packages/nnsight/contexts/Tracer.py in __exit__(self, exc_type, exc_val, exc_tb)\n",
        "    100\n",
        "    101\n",
        "--> 102         super().__exit__(exc_type, exc_val, exc_tb)\n",
        "    103\n",
        "    104     def invoke(self, *inputs: Any, **kwargs) -> Invoker:\n",
        "\n",
        "/usr/local/lib/python3.11/dist-packages/nnsight/contexts/GraphBasedContext.py in __exit__(self, exc_type, exc_val, exc_tb)\n",
        "    215             raise exc_val\n",
        "    216\n",
        "--> 217         self.backend(self)\n",
        "    218\n",
        "    219     ### BACKENDS ########\n",
        "\n",
        "/usr/local/lib/python3.11/dist-packages/nnsight/contexts/backends/LocalBackend.py in __call__(self, obj)\n",
        "     25     def __call__(self, obj: LocalMixin):\n",
        "     26\n",
        "---> 27         obj.local_backend_execute()\n",
        "\n",
        "/usr/local/lib/python3.11/dist-packages/nnsight/contexts/Tracer.py in local_backend_execute(self)\n",
        "    144         self.graph.execute()\n",
        "    145\n",
        "--> 146         self.model.interleave(\n",
        "    147             self.model._execute,\n",
        "    148             self.graph,\n",
        "\n",
        "/usr/local/lib/python3.11/dist-packages/nnsight/models/NNsightModel.py in interleave(self, fn, intervention_graph, *inputs, **kwargs)\n",
        "    467         module_paths = InterventionProtocol.get_interventions(intervention_graph).keys()\n",
        "    468\n",
        "--> 469         with HookHandler(\n",
        "    470             self._model,\n",
        "    471             list(module_paths),\n",
        "\n",
        "/usr/local/lib/python3.11/dist-packages/nnsight/intervention.py in __exit__(self, exc_type, exc_val, exc_tb)\n",
        "    579\n",
        "    580         if isinstance(exc_val, Exception):\n",
        "--> 581             raise exc_val\n",
        "    582\n",
        "    583\n",
        "\n",
        "/usr/local/lib/python3.11/dist-packages/nnsight/models/NNsightModel.py in interleave(self, fn, intervention_graph, *inputs, **kwargs)\n",
        "    478         ):\n",
        "    479             try:\n",
        "--> 480                 fn(*inputs, **kwargs)\n",
        "    481             except protocols.EarlyStopProtocol.EarlyStopException:\n",
        "    482                 # TODO: Log.\n",
        "\n",
        "/usr/local/lib/python3.11/dist-packages/nnsight/models/NNsightModel.py in _execute(self, *prepared_inputs, **kwargs)\n",
        "    585             pass\n",
        "    586\n",
        "--> 587         return self._model(\n",
        "    588             *prepared_inputs,\n",
        "    589             **kwargs,\n",
        "\n",
        "/usr/local/lib/python3.11/dist-packages/torch/nn/modules/module.py in _wrapped_call_impl(self, *args, **kwargs)\n",
        "   1734             return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]\n",
        "   1735         else:\n",
        "-> 1736             return self._call_impl(*args, **kwargs)\n",
        "   1737\n",
        "   1738     # torchrec tests the code consistency with the following code\n",
        "\n",
        "/usr/local/lib/python3.11/dist-packages/torch/nn/modules/module.py in _call_impl(self, *args, **kwargs)\n",
        "   1842\n",
        "   1843         try:\n",
        "-> 1844             return inner()\n",
        "   1845         except Exception:\n",
        "   1846             # run always called hooks if they have not already been run\n",
        "\n",
        "/usr/local/lib/python3.11/dist-packages/torch/nn/modules/module.py in inner()\n",
        "   1788                 args = bw_hook.setup_input_hook(args)\n",
        "   1789\n",
        "-> 1790             result = forward_call(*args, **kwargs)\n",
        "   1791             if _global_forward_hooks or self._forward_hooks:\n",
        "   1792                 for hook_id, hook in (\n",
        "\n",
        "/usr/local/lib/python3.11/dist-packages/torch/nn/modules/container.py in forward(self, input)\n",
        "    248     def forward(self, input):\n",
        "    249         for module in self:\n",
        "--> 250             input = module(input)\n",
        "    251         return input\n",
        "    252\n",
        "\n",
        "/usr/local/lib/python3.11/dist-packages/torch/nn/modules/module.py in _wrapped_call_impl(self, *args, **kwargs)\n",
        "   1734             return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]\n",
        "   1735         else:\n",
        "-> 1736             return self._call_impl(*args, **kwargs)\n",
        "   1737\n",
        "   1738     # torchrec tests the code consistency with the following code\n",
        "\n",
        "/usr/local/lib/python3.11/dist-packages/torch/nn/modules/module.py in _call_impl(self, *args, **kwargs)\n",
        "   1842\n",
        "   1843         try:\n",
        "-> 1844             return inner()\n",
        "   1845         except Exception:\n",
        "   1846             # run always called hooks if they have not already been run\n",
        "\n",
        "/usr/local/lib/python3.11/dist-packages/torch/nn/modules/module.py in inner()\n",
        "   1801                         hook_result = hook(self, args, kwargs, result)\n",
        "   1802                     else:\n",
        "-> 1803                         hook_result = hook(self, args, result)\n",
        "   1804\n",
        "   1805                     if hook_result is not None:\n",
        "\n",
        "/usr/local/lib/python3.11/dist-packages/nnsight/intervention.py in output_hook(module, input, output, module_path)\n",
        "    564\n",
        "    565                 def output_hook(module, input, output, module_path=module_path):\n",
        "--> 566                     return self.output_hook(output, module_path)\n",
        "    567\n",
        "    568                 self.handles.append(\n",
        "\n",
        "/usr/local/lib/python3.11/dist-packages/nnsight/models/NNsightModel.py in <lambda>(activations, module_path)\n",
        "    473                 activations, module_path, \"input\", intervention_handler\n",
        "    474             ),\n",
        "--> 475             output_hook=lambda activations, module_path: InterventionProtocol.intervene(\n",
        "    476                 activations, module_path, \"output\", intervention_handler\n",
        "    477             ),\n",
        "\n",
        "/usr/local/lib/python3.11/dist-packages/nnsight/intervention.py in intervene(cls, activations, module_path, key, intervention_handler)\n",
        "    454\n",
        "    455                 # Value injection.\n",
        "--> 456                 node.set_value(value)\n",
        "    457\n",
        "    458                 # Check if through the previous value injection, there was a 'swap' intervention.\n",
        "\n",
        "/usr/local/lib/python3.11/dist-packages/nnsight/tracing/Node.py in set_value(self, value)\n",
        "    408\n",
        "    409             if listener.fulfilled() and not self.graph.sequential:\n",
        "--> 410                 listener.execute()\n",
        "    411\n",
        "    412         for dependency in self.arg_dependencies:\n",
        "\n",
        "/usr/local/lib/python3.11/dist-packages/nnsight/tracing/Node.py in execute(self)\n",
        "    385         except Exception as e:\n",
        "    386\n",
        "--> 387             raise type(e)(\n",
        "    388                 f\"Above exception when execution Node: '{self.name}' in Graph: '{self.graph.id}'\"\n",
        "    389             ) from e\n",
        "\n",
        "IndexError: Above exception when execution Node: 'setitem_0' in Graph: '132147685816016'\n",
        "\n",
        "\n",
        "```\n",
        "\n",
        "\n",
        "</details>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Lo7oHC0yzI_p"
      },
      "source": [
        "The error messaging feature can be toggled using `nnsight.CONFIG.APP.DEBUG` which defaults to true.\n",
        "\n",
        "<details>\n",
        "\n",
        "<summary>\n",
        "Toggle Error Messaging \n",
        "</summary>\n",
        "\n",
        "Turn off debugging:\n",
        "```\n",
        "import nnsight\n",
        "\n",
        "nnsight.CONFIG.APP.DEBUG = False\n",
        "nnsight.CONFIG.save()\n",
        "```\n",
        "\n",
        "Turn on debugging:\n",
        "```\n",
        "import nnsight\n",
        "\n",
        "nnsight.CONFIG.APP.DEBUG = True\n",
        "nnsight.CONFIG.save()\n",
        "```\n",
        "</details>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JCNR_Jkpxr8l"
      },
      "source": [
        "Now that we know more about NNsight's error messaging, let's try our setting operation again with the correct indexing and view the shape of the output\n",
        "before leaving the tracing context:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 77,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Wf8ugJKB2mru",
        "outputId": "f076fa36-0271-4185-8583-9e7190e1e8e9"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Layer 1 output shape: InterventionProxy (fetch_attr)\n",
            "Before: tensor([[ 0.2732,  0.2355, -0.6433,  0.0475,  0.0904,  0.4407, -0.3099,  1.4903,\n",
            "         -0.0748,  0.0088]])\n",
            "After: tensor([[ 0.2732,  0.2355, -0.6433,  0.0475,  0.0904,  0.4407, -0.3099,  1.4903,\n",
            "         -0.0748,  0.0000]])\n"
          ]
        }
      ],
      "source": [
        "with tiny_model.trace(input):\n",
        "\n",
        "    # Save the output before the edit to compare.\n",
        "    # Notice we apply .clone() before saving as the setting operation is in-place.\n",
        "    l1_output_before = tiny_model.layer1.output.clone().save()\n",
        "\n",
        "    print(f\"Layer 1 output shape: {tiny_model.layer1.output.shape}\")\n",
        "\n",
        "    # Access the last index of the hidden state dimension and set it to 0.\n",
        "    tiny_model.layer1.output[:, hidden_dims - 1] = 0\n",
        "\n",
        "    # Save the output after to see our edit.\n",
        "    l1_output_after = tiny_model.layer1.output.save()\n",
        "\n",
        "print(\"Before:\", l1_output_before)\n",
        "print(\"After:\", l1_output_after)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6R3wj1BKkkw5"
      },
      "source": [
        "## Scan and Validate\n",
        "Error codes are helpful, but sometimes you may want to quickly troubleshoot your code without actually running it. \n",
        "\n",
        "Enter \"Scanning\" and \"Validating\"! We can enable this features by setting the `scan=True` and `validate=True` flag in the `trace` method."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\"Scanning\" runs \"fake\" inputs throught the model to collect information like shapes and types (i.e., scanning will populate all called `.inputs` and `.outputs`). \n",
        "\n",
        "\"Validating\" attempts to execute the intervention proxies with \"fake\" inputs to check if they work (i.e., executes all interventions in your code with fake tensors).\n",
        "\n",
        "\"Validating\" is dependent on \"Scanning\" to work correctly, so we need to run the scan of the model at least once to debug with validate. Let's try it out on our example above."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 78,
      "metadata": {},
      "outputs": [
        {
          "ename": "IndexError",
          "evalue": "index 10 is out of bounds for dimension 1 with size 10",
          "output_type": "error",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mIndexError\u001b[0m                                Traceback (most recent call last)",
            "Cell \u001b[0;32mIn[78], line 2\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;66;03m# turn on scan and validate\u001b[39;00m\n\u001b[0;32m----> 2\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m tiny_model\u001b[38;5;241m.\u001b[39mtrace(\u001b[38;5;28minput\u001b[39m, scan\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m, validate\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m):\n\u001b[1;32m      4\u001b[0m     l1_output_before \u001b[38;5;241m=\u001b[39m tiny_model\u001b[38;5;241m.\u001b[39mlayer1\u001b[38;5;241m.\u001b[39moutput\u001b[38;5;241m.\u001b[39mclone()\u001b[38;5;241m.\u001b[39msave()\n\u001b[1;32m      6\u001b[0m     \u001b[38;5;66;03m# the error is happening here\u001b[39;00m\n",
            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/intervention/contexts/interleaving.py:96\u001b[0m, in \u001b[0;36mInterleavingTracer.__exit__\u001b[0;34m(self, exc_type, exc_val, exc_tb)\u001b[0m\n\u001b[1;32m     92\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39minvoker\u001b[38;5;241m.\u001b[39m\u001b[38;5;21m__exit__\u001b[39m(\u001b[38;5;28;01mNone\u001b[39;00m, \u001b[38;5;28;01mNone\u001b[39;00m, \u001b[38;5;28;01mNone\u001b[39;00m)\n\u001b[1;32m     94\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_model\u001b[38;5;241m.\u001b[39m_envoy\u001b[38;5;241m.\u001b[39m_reset()\n\u001b[0;32m---> 96\u001b[0m \u001b[38;5;28;43msuper\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[38;5;21;43m__exit__\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43mexc_type\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mexc_val\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mexc_tb\u001b[49m\u001b[43m)\u001b[49m\n",
            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/tracing/contexts/tracer.py:25\u001b[0m, in \u001b[0;36mTracer.__exit__\u001b[0;34m(self, exc_type, exc_val, exc_tb)\u001b[0m\n\u001b[1;32m     21\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mglobals\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m GlobalTracingContext\n\u001b[1;32m     23\u001b[0m GlobalTracingContext\u001b[38;5;241m.\u001b[39mtry_deregister(\u001b[38;5;28mself\u001b[39m)\n\u001b[0;32m---> 25\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43msuper\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[38;5;21;43m__exit__\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43mexc_type\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mexc_val\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mexc_tb\u001b[49m\u001b[43m)\u001b[49m\n",
            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/tracing/contexts/base.py:72\u001b[0m, in \u001b[0;36mContext.__exit__\u001b[0;34m(self, exc_type, exc_val, exc_tb)\u001b[0m\n\u001b[1;32m     69\u001b[0m graph \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mgraph\u001b[38;5;241m.\u001b[39mstack\u001b[38;5;241m.\u001b[39mpop()\n\u001b[1;32m     71\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(exc_val, \u001b[38;5;167;01mBaseException\u001b[39;00m):\n\u001b[0;32m---> 72\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m exc_val\n\u001b[1;32m     74\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39madd(graph\u001b[38;5;241m.\u001b[39mstack[\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m], graph, \u001b[38;5;241m*\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mkwargs)\n\u001b[1;32m     76\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mbackend \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n",
            "Cell \u001b[0;32mIn[78], line 7\u001b[0m\n\u001b[1;32m      4\u001b[0m     l1_output_before \u001b[38;5;241m=\u001b[39m tiny_model\u001b[38;5;241m.\u001b[39mlayer1\u001b[38;5;241m.\u001b[39moutput\u001b[38;5;241m.\u001b[39mclone()\u001b[38;5;241m.\u001b[39msave()\n\u001b[1;32m      6\u001b[0m     \u001b[38;5;66;03m# the error is happening here\u001b[39;00m\n\u001b[0;32m----> 7\u001b[0m     \u001b[43mtiny_model\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mlayer1\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43moutput\u001b[49m\u001b[43m[\u001b[49m\u001b[43m:\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mhidden_dims\u001b[49m\u001b[43m]\u001b[49m \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m0\u001b[39m\n\u001b[1;32m      9\u001b[0m     l1_output_after \u001b[38;5;241m=\u001b[39m tiny_model\u001b[38;5;241m.\u001b[39mlayer1\u001b[38;5;241m.\u001b[39moutput\u001b[38;5;241m.\u001b[39msave()\n\u001b[1;32m     11\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mBefore:\u001b[39m\u001b[38;5;124m\"\u001b[39m, l1_output_before)\n",
            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/tracing/graph/proxy.py:126\u001b[0m, in \u001b[0;36mProxy.__setitem__\u001b[0;34m(self, key, value)\u001b[0m\n\u001b[1;32m    125\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__setitem__\u001b[39m(\u001b[38;5;28mself\u001b[39m, key: Union[Self, Any], value: Union[Self, Any]) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[0;32m--> 126\u001b[0m     \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mnode\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcreate\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    127\u001b[0m \u001b[43m        \u001b[49m\u001b[43moperator\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43msetitem\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    128\u001b[0m \u001b[43m        \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mnode\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    129\u001b[0m \u001b[43m        \u001b[49m\u001b[43mkey\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    130\u001b[0m \u001b[43m        \u001b[49m\u001b[43mvalue\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    131\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n",
            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/tracing/graph/node.py:250\u001b[0m, in \u001b[0;36mNode.create\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    247\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m value\n\u001b[1;32m    249\u001b[0m \u001b[38;5;66;03m# Otherwise just create the Node on the Graph like normal.\u001b[39;00m\n\u001b[0;32m--> 250\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mgraph\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcreate\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    251\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    252\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    253\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n",
            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/tracing/graph/graph.py:131\u001b[0m, in \u001b[0;36mGraph.create\u001b[0;34m(self, target, redirect, *args, **kwargs)\u001b[0m\n\u001b[1;32m    128\u001b[0m \u001b[38;5;66;03m# Redirection.\u001b[39;00m\n\u001b[1;32m    129\u001b[0m graph \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mstack[\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m] \u001b[38;5;28;01mif\u001b[39;00m redirect \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mstack \u001b[38;5;28;01melse\u001b[39;00m \u001b[38;5;28mself\u001b[39m\n\u001b[0;32m--> 131\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mproxy_class(\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mnode_class\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtarget\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mgraph\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mgraph\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m)\n",
            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/intervention/graph/node.py:118\u001b[0m, in \u001b[0;36mValidatingInterventionNode.__init__\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    111\u001b[0m \u001b[38;5;28msuper\u001b[39m()\u001b[38;5;241m.\u001b[39m\u001b[38;5;21m__init__\u001b[39m(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[1;32m    113\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m (\n\u001b[1;32m    114\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mattached\n\u001b[1;32m    115\u001b[0m     \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mfake_value \u001b[38;5;129;01mis\u001b[39;00m inspect\u001b[38;5;241m.\u001b[39m_empty\n\u001b[1;32m    116\u001b[0m     \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m Protocol\u001b[38;5;241m.\u001b[39mis_protocol(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtarget)\n\u001b[1;32m    117\u001b[0m ):\n\u001b[0;32m--> 118\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mfake_value \u001b[38;5;241m=\u001b[39m \u001b[43mvalidate\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtarget\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/intervention/graph/node.py:147\u001b[0m, in \u001b[0;36mvalidate\u001b[0;34m(target, *args, **kwargs)\u001b[0m\n\u001b[1;32m    141\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m FakeTensorMode(\n\u001b[1;32m    142\u001b[0m     allow_non_fake_inputs\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m,\n\u001b[1;32m    143\u001b[0m     shape_env\u001b[38;5;241m=\u001b[39mShapeEnv(assume_static_by_default\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m),\n\u001b[1;32m    144\u001b[0m ) \u001b[38;5;28;01mas\u001b[39;00m fake_mode:\n\u001b[1;32m    145\u001b[0m     \u001b[38;5;28;01mwith\u001b[39;00m FakeCopyMode(fake_mode):\n\u001b[0;32m--> 147\u001b[0m         \u001b[38;5;28;01mwith\u001b[39;00m GlobalTracingContext\u001b[38;5;241m.\u001b[39mexit_global_tracing_context():\n\u001b[1;32m    149\u001b[0m             \u001b[38;5;28;01mif\u001b[39;00m backwards_check(target, \u001b[38;5;241m*\u001b[39margs):\n\u001b[1;32m    150\u001b[0m                 \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n",
            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/tracing/contexts/globals.py:100\u001b[0m, in \u001b[0;36mGlobalTracingContext.GlobalTracingExit.__exit__\u001b[0;34m(self, exc_type, exc_val, traceback)\u001b[0m\n\u001b[1;32m     96\u001b[0m GlobalTracingContext\u001b[38;5;241m.\u001b[39mPATCHER\u001b[38;5;241m.\u001b[39m\u001b[38;5;21m__enter__\u001b[39m()\n\u001b[1;32m     98\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(exc_val, \u001b[38;5;167;01mBaseException\u001b[39;00m):\n\u001b[0;32m--> 100\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m exc_val\n",
            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/intervention/graph/node.py:156\u001b[0m, in \u001b[0;36mvalidate\u001b[0;34m(target, *args, **kwargs)\u001b[0m\n\u001b[1;32m    150\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[1;32m    152\u001b[0m args, kwargs \u001b[38;5;241m=\u001b[39m InterventionNode\u001b[38;5;241m.\u001b[39mprepare_inputs(\n\u001b[1;32m    153\u001b[0m     (args, kwargs), fake\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m\n\u001b[1;32m    154\u001b[0m )\n\u001b[0;32m--> 156\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mtarget\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    157\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    158\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    159\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n",
            "File \u001b[0;32m/opt/anaconda3/envs/nnsight/lib/python3.10/site-packages/torch/_subclasses/fake_tensor.py:2350\u001b[0m, in \u001b[0;36mFakeCopyMode.__torch_function__\u001b[0;34m(self, func, types, args, kwargs)\u001b[0m\n\u001b[1;32m   2348\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m   2349\u001b[0m     \u001b[38;5;28;01mwith\u001b[39;00m torch\u001b[38;5;241m.\u001b[39m_C\u001b[38;5;241m.\u001b[39mDisableTorchFunctionSubclass():\n\u001b[0;32m-> 2350\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mfunc\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
            "\u001b[0;31mIndexError\u001b[0m: index 10 is out of bounds for dimension 1 with size 10"
          ]
        }
      ],
      "source": [
        "# turn on scan and validate\n",
        "with tiny_model.trace(input, scan=True, validate=True):\n",
        "\n",
        "    l1_output_before = tiny_model.layer1.output.clone().save()\n",
        "\n",
        "    # the error is happening here\n",
        "    tiny_model.layer1.output[:, hidden_dims] = 0\n",
        "\n",
        "    l1_output_after = tiny_model.layer1.output.save()\n",
        "\n",
        "print(\"Before:\", l1_output_before)\n",
        "print(\"After:\", l1_output_after)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The operations are never executed using tensors with real values so it doesn't incur any memory costs. Then, when creating proxy requests like the setting one above, `nnsight` also attempts to execute the request on the \"fake\" values we recorded. Hence, it lets us know if our request is feasible before even running the model. [Here](https://nnsight.net/notebooks/features/scan_validate/) is a more detailed example of scan and validate in action!\n",
        "\n",
        "<details>\n",
        "<summary>A word of caution</summary>\n",
        "\n",
        "---\n",
        "\n",
        "Some pytorch operations and related libraries don't work well with fake tensors\n",
        "\n",
        "If you are doing anything in a loop where efficiency is important, you should keep scanning and validating off. It's best to use them only when debugging or when you are unsure if your intervention will work.\n",
        "\n",
        "---\n",
        "\n",
        "</details>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "45E98sIAVc5A"
      },
      "source": [
        "We can also use the `.scan()` method to get the shape of a module without having to fully run the model. If scan  is enabled, our input is run though the model under its own \"fake\" context. This means the input makes its way through all of the model operations, allowing `nnsight` to record the shapes and data types of module inputs and outputs!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 79,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 382
        },
        "id": "5J79HytJkkw5",
        "outputId": "3b11405b-7bd7-45b6-dc16-11c8456b277b"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "10\n"
          ]
        }
      ],
      "source": [
        "with tiny_model.scan(input):\n",
        "\n",
        "    dim = tiny_model.layer1.output.shape[-1]\n",
        "\n",
        "print(dim)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9DH7PeL13-WU"
      },
      "source": [
        "We can also just replace proxy inputs and outputs with tensors of the same shape\n",
        "and type. Let's use the shape information we have at our disposal to add noise\n",
        "to the output, and replace it with this new noised tensor:\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "s016CelFP8sx"
      },
      "source": [
        "## Gradients\n",
        "\n",
        "`NNsight` also lets us apply backpropagation and access gradients with respect to a\n",
        "loss. Like `.input` and `.output` on modules, `nnsight` exposes `.grad` on\n",
        "Proxies themselves (assuming they are proxies of tensors):\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 80,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "L-cYUNNyRDn1",
        "outputId": "c1521deb-34d7-4ea5-c80f-0c5c873252d9"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Layer 1 output gradient: tensor([[-0.1296,  0.4562, -0.1182, -0.3536,  0.0703,  0.1411, -0.3201, -0.4890,\n",
            "          0.0196, -0.0452]])\n",
            "Layer 2 output gradient: tensor([[1., 1.]])\n"
          ]
        }
      ],
      "source": [
        "with tiny_model.trace(input):\n",
        "\n",
        "    # We need to explicitly have the tensor require grad\n",
        "    # as the model we defined earlier turned off requiring grad.\n",
        "    tiny_model.layer1.output.requires_grad = True\n",
        "\n",
        "    # We call .grad on a tensor Proxy to communicate we want to store its gradient.\n",
        "    # We need to call .save() since .grad is its own Proxy.\n",
        "    layer1_output_grad = tiny_model.layer1.output.grad.save()\n",
        "    layer2_output_grad = tiny_model.layer2.output.grad.save()\n",
        "\n",
        "    # Need a loss to propagate through the later modules in order to have a grad.\n",
        "    loss = tiny_model.output.sum()\n",
        "    loss.backward()\n",
        "\n",
        "print(\"Layer 1 output gradient:\", layer1_output_grad)\n",
        "print(\"Layer 2 output gradient:\", layer2_output_grad)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1o0JaaYvWlHG"
      },
      "source": [
        "All of the features we learned previously, also apply to `.grad`. In other\n",
        "words, we can apply operations to and edit the gradients. Let's zero the grad of\n",
        "`layer1` and double the grad of `layer2`.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 81,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "9bFoaJpOWlRb",
        "outputId": "d3c78c79-4b0f-496f-f722-573f6db30e3a"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Layer 1 output gradient: tensor([[0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]])\n",
            "Layer 2 output gradient: tensor([[2., 2.]])\n"
          ]
        }
      ],
      "source": [
        "with tiny_model.trace(input):\n",
        "\n",
        "    # We need to explicitly have the tensor require grad\n",
        "    # as the model we defined earlier turned off requiring grad.\n",
        "    tiny_model.layer1.output.requires_grad = True\n",
        "\n",
        "    tiny_model.layer1.output.grad[:] = 0\n",
        "    tiny_model.layer2.output.grad = tiny_model.layer2.output.grad * 2\n",
        "\n",
        "    layer1_output_grad = tiny_model.layer1.output.grad.save()\n",
        "    layer2_output_grad = tiny_model.layer2.output.grad.save()\n",
        "\n",
        "    # Need a loss to propagate through the later modules in order to have a grad.\n",
        "    loss = tiny_model.output.sum()\n",
        "    loss.backward()\n",
        "\n",
        "print(\"Layer 1 output gradient:\", layer1_output_grad)\n",
        "print(\"Layer 2 output gradient:\", layer2_output_grad)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SJ6dAgzQkkw6"
      },
      "source": [
        "## Early Stopping"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FZLrnIxAkkw6"
      },
      "source": [
        "If we are only interested in a model's intermediate computations, we can halt a forward pass run at any module level, reducing runtime and conserving compute resources. One examples where this could be particularly useful would if we are working with SAEs - we can train an SAE on one layer and then stop the execution."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 82,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "MwAeiZhGkkw6",
        "outputId": "85e6c10d-50dd-4517-b34e-2a9757e29c80"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "L1 - Output:  tensor([[ 0.2732,  0.2355, -0.6433,  0.0475,  0.0904,  0.4407, -0.3099,  1.4903,\n",
            "         -0.0748,  0.0088]])\n"
          ]
        }
      ],
      "source": [
        "with tiny_model.trace(input):\n",
        "   l1_out = tiny_model.layer1.output.save()\n",
        "   tiny_model.layer1.output.stop()\n",
        "\n",
        "# get the output of the first layer and stop tracing\n",
        "print(\"L1 - Output: \", l1_out)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eO_8oZPXkkw6"
      },
      "source": [
        "Interventions within the tracing context do not necessarily execute in the order they are defined. Instead, their execution is tied to the module they are associated with.\n",
        "\n",
        "As a result, if the forward pass is terminated early any interventions linked to modules beyond that point will be skipped, even if they were defined earlier in the context.\n",
        "\n",
        "In the example below, the output of layer 2 _**cannot**_ be accessed since the model's execution was stopped at layer 1."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 83,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 347
        },
        "id": "uPJ6j2Zwkkw6",
        "outputId": "aeff829d-e973-489e-c24e-1d5cdc80c7d0"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "L2 - Output:  "
          ]
        },
        {
          "ename": "ValueError",
          "evalue": "Accessing value before it's been set.",
          "output_type": "error",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
            "Cell \u001b[0;32mIn[83], line 5\u001b[0m\n\u001b[1;32m      2\u001b[0m    l2_out \u001b[38;5;241m=\u001b[39m tiny_model\u001b[38;5;241m.\u001b[39mlayer2\u001b[38;5;241m.\u001b[39moutput\u001b[38;5;241m.\u001b[39msave()\n\u001b[1;32m      3\u001b[0m    tiny_model\u001b[38;5;241m.\u001b[39mlayer1\u001b[38;5;241m.\u001b[39moutput\u001b[38;5;241m.\u001b[39mstop()\n\u001b[0;32m----> 5\u001b[0m \u001b[38;5;28;43mprint\u001b[39;49m\u001b[43m(\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mL2 - Output: \u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43ml2_out\u001b[49m\u001b[43m)\u001b[49m\n",
            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/tracing/graph/proxy.py:70\u001b[0m, in \u001b[0;36mProxy.__str__\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m     66\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__str__\u001b[39m(\u001b[38;5;28mself\u001b[39m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m \u001b[38;5;28mstr\u001b[39m:\n\u001b[1;32m     68\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mnode\u001b[38;5;241m.\u001b[39mattached:\n\u001b[0;32m---> 70\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mstr\u001b[39m(\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mvalue\u001b[49m)\n\u001b[1;32m     72\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mtype\u001b[39m(\u001b[38;5;28mself\u001b[39m)\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m (\u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mnode\u001b[38;5;241m.\u001b[39mtarget\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__name__\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m)\u001b[39m\u001b[38;5;124m\"\u001b[39m\n",
            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/tracing/graph/proxy.py:64\u001b[0m, in \u001b[0;36mProxy.value\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m     56\u001b[0m \u001b[38;5;129m@property\u001b[39m\n\u001b[1;32m     57\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mvalue\u001b[39m(\u001b[38;5;28mself\u001b[39m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Any:\n\u001b[1;32m     58\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124;03m\"\"\"Property to return the value of this proxy's node.\u001b[39;00m\n\u001b[1;32m     59\u001b[0m \n\u001b[1;32m     60\u001b[0m \u001b[38;5;124;03m    Returns:\u001b[39;00m\n\u001b[1;32m     61\u001b[0m \u001b[38;5;124;03m        Any: The stored value of the proxy, populated during execution of the model.\u001b[39;00m\n\u001b[1;32m     62\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[0;32m---> 64\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mnode\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mvalue\u001b[49m\n",
            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/tracing/graph/node.py:143\u001b[0m, in \u001b[0;36mNode.value\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    133\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"Property to return the value of this node.\u001b[39;00m\n\u001b[1;32m    134\u001b[0m \n\u001b[1;32m    135\u001b[0m \u001b[38;5;124;03mReturns:\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    139\u001b[0m \u001b[38;5;124;03m    ValueError: If the underlying ._value is inspect._empty (therefore never set or was destroyed).\u001b[39;00m\n\u001b[1;32m    140\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    142\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdone:\n\u001b[0;32m--> 143\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mAccessing value before it\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124ms been set.\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m    145\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_value\n",
            "\u001b[0;31mValueError\u001b[0m: Accessing value before it's been set."
          ]
        }
      ],
      "source": [
        "with tiny_model.trace(input):\n",
        "   l2_out = tiny_model.layer2.output.save()\n",
        "   tiny_model.layer1.output.stop()\n",
        "\n",
        "print(\"L2 - Output: \", l2_out)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "i1iJOI29kkw6"
      },
      "source": [
        "## Conditional Interventions"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NvC2f8qqkkw6"
      },
      "source": [
        "Interventions can also be made conditional.\n",
        "\n",
        "Inside the tracing context we can specify a new - conditional - context. This context will only execute the interventions within it if the condition is met."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 84,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "F_SIV2CDkkw6",
        "outputId": "2ba0cfb8-dfd8-4403-e2e9-db3de2045ec7"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Random Integer  tensor([-5])  is Odd\n"
          ]
        }
      ],
      "source": [
        "with tiny_model.trace(input) as tracer:\n",
        "\n",
        "  rand_int = torch.randint(low=-10, high=10, size=(1,))\n",
        "\n",
        "  with tracer.cond(rand_int % 2 == 0):\n",
        "    tracer.log(\"Random Integer \", rand_int, \" is Even\")\n",
        "\n",
        "  with tracer.cond(rand_int % 2 == 1):\n",
        "    tracer.log(\"Random Integer \", rand_int, \" is Odd\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CZLQecmrkkw6"
      },
      "source": [
        "Conditional contexts can also be nested, if we want our interventions to depend on more than one condition at a time."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 21,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "qZsFtC5jkkw6",
        "outputId": "4bbad29b-2d9d-4816-ed73-859a829e5ba5"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Rand Int  8  is Positive and Even\n"
          ]
        }
      ],
      "source": [
        "with tiny_model.trace(input) as tracer:\n",
        "\n",
        "  non_rand_int = 8\n",
        "\n",
        "  with tracer.cond(non_rand_int > 0):\n",
        "    with tracer.cond(non_rand_int % 2 == 0):\n",
        "      tracer.log(\"Rand Int \", non_rand_int, \" is Positive and Even\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hl99yP3kkkw6"
      },
      "source": [
        "With `nnsight 0.4` we can now also use Python `if` statements within the tracing context to create a conditional context!\n",
        "\n",
        "*Note: Colab behaves a little strangely with this feature the first time you run it - expect some lagging and warnings*"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 85,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "52a1vwtCfdIS",
        "outputId": "51680d97-0f1f-4ead-ab97-26ec81587fff"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Random Integer  tensor([2])  is Even\n"
          ]
        }
      ],
      "source": [
        "with tiny_model.trace(input) as tracer:\n",
        "\n",
        "  rand_int = torch.randint(low=-10, high=10, size=(1,))\n",
        "\n",
        "  # Since this if statement is inside the tracing context the if will\n",
        "  # create a conditional context and will only execute the intervention\n",
        "  # if this condition is met\n",
        "  if rand_int % 2 == 0:\n",
        "    tracer.log(\"Random Integer \", rand_int, \" is Even\")\n",
        "\n",
        "  if rand_int % 2 == 1:\n",
        "    tracer.log(\"Random Integer \", rand_int, \" is Odd\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tEML88wTf9Lv"
      },
      "source": [
        "`elif` statements should also work as `if` statements within the tracing context:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 86,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "3njmT1szf-Lq",
        "outputId": "b1088864-8e0e-4526-da98-b6275adbabdf"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Random Integer  tensor([-3])  is Odd\n"
          ]
        }
      ],
      "source": [
        "with tiny_model.trace(input) as tracer:\n",
        "\n",
        "  rand_int = torch.randint(low=-10, high=10, size=(1,))\n",
        "\n",
        "  # Since this if statement is inside the tracing context the if will\n",
        "  # create a conditional context and will only execute the intervention\n",
        "  # if this condition is met\n",
        "  if rand_int % 2 == 0:\n",
        "    tracer.log(\"Random Integer \", rand_int, \" is Even\")\n",
        "  elif rand_int % 2 == 1:\n",
        "    tracer.log(\"Random Integer \", rand_int, \" is Odd\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ONzsQcHtxKkf"
      },
      "source": [
        "## Iterative Interventions"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NrwpkfYXxRBO"
      },
      "source": [
        "With the iterator context, you can now run an intervention loop at scale. It iteratively executes and updates a single intervention graph. Use a `.session()` to define the Iterator context and pass in a sequence of items that you want to loop over at each iteration\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 87,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "AYWJVln7xP85",
        "outputId": "3b7a1271-b199-4078-c70e-0fefc4cd5d7b"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "List:  [0, 1, 2]\n"
          ]
        }
      ],
      "source": [
        "with tiny_model.session() as session:\n",
        "\n",
        "  li = nnsight.list() # an NNsight built-in list object\n",
        "  [li.append([num]) for num in range(0, 3)] # adding [0], [1], [2] to the list\n",
        "  li2 = nnsight.list().save()\n",
        "\n",
        "  # You can create nested Iterator contexts\n",
        "  with session.iter(li) as item:\n",
        "    with session.iter(item) as item_2:\n",
        "      li2.append(item_2)\n",
        "\n",
        "print(\"\\nList: \", li2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5Dbnsw6fx8xL"
      },
      "source": [
        "With `nnsight 0.4` we can now also use Python `for` loops within a tracer context at scale.\n",
        "\n",
        "*NOTE: inline for loops (i.e., `[x for x in <Proxy object>`]) are not currently supported.*"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 88,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "14rgY7cZyNUr",
        "outputId": "a028118d-cf04-45d3-ea38-71dd04405a9b"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "List:  [0, 1, 2]\n"
          ]
        }
      ],
      "source": [
        "# New: Using Python for loops for iterative interventions\n",
        "with tiny_model.session() as session:\n",
        "\n",
        "    li = nnsight.list()\n",
        "    [li.append([num]) for num in range(0, 3)]\n",
        "    li2 = nnsight.list().save()\n",
        "\n",
        "    # Using regular for loops\n",
        "    for item in li:\n",
        "        for item_2 in item: # for loops can be nested!\n",
        "            li2.append(item_2)\n",
        "\n",
        "print(\"\\nList: \", li2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gvQ1nZgYQDG3"
      },
      "source": [
        "# 2️ Bigger"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "miQgUY4NAmDQ"
      },
      "source": [
        "Now that we have the basics of `nnsight` under our belt, we can scale our model\n",
        "up and combine the techniques we've learned into more interesting experiments.\n",
        "\n",
        "The `NNsight` class is very bare bones. It wraps a pre-defined model and does no\n",
        "pre-processing on the inputs we enter. It's designed to be extended with more\n",
        "complex and powerful types of models, and we're excited to see what can be done\n",
        "to leverage its features!\n",
        "\n",
        "However, if you'd like to load a Language Model from HuggingFace with its tokenizer, the`LanguageModel` subclass greatly simplifies this process."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5TJDblHiQpp1"
      },
      "source": [
        "## LanguageModel\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9l9mZOY5HFH2"
      },
      "source": [
        "`LanguageModel` is a subclass of `NNsight`. While we could define and create a\n",
        "model to pass in directly, `LanguageModel` includes special support for\n",
        "Huggingface language models, including automatically loading models from a\n",
        "Huggingface ID, and loading the model together with the appropriate tokenizer.\n",
        "\n",
        "Here is how we can use `LanguageModel` to load `GPT-2`:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 713,
          "referenced_widgets": [
            "26dcd7b75feb41a6b86d4260ecd64d06",
            "39c9f1effe664e899125748e97fad198",
            "383b893224684aee916e455c0d3d8c06",
            "af24a12bd05b44148b5d6f91ac2445a0",
            "d00750501f2548d0b1c69f4753aec877",
            "98d28dddf75a45a2b1e05131b0b865a1",
            "6e81d6f31f2446c999df5511e8afd8d0",
            "68b205f60f9d4599a7d1f78583d0b93c",
            "69e8c149f676457d95029f21e9a2956a",
            "ba90058201ee48dcae9411ded457e85f",
            "346ecc4734014cb4931bea53d5a32bf1",
            "79a221c9228d459aa17402c9fef1ad1d",
            "9c13a3d768eb47c8b26a7c6bfa8ec957",
            "e97c052113f54dddbb619ed228f6b4ef",
            "41fc21ee95c14d29943a56f2b63f4989",
            "5250cf68037a43e389493733c5b9d4ac",
            "e24033c5ba7b4e4fa4f4375d85b03bad",
            "4b798a784ae748b1a0aef025ead976fc",
            "19c0e121e52040b18393aee887730644",
            "2702f30b5d3b4b0b8184fec58863750b",
            "42fd0ea3be3c4b3f98ca7818efee3669",
            "184b5c75820b4c42ae861669ac1a47e5",
            "d8feaceb6d854cb2b2ff19a0309d9a7b",
            "fd84ffef58604ce3a16ef8f9868072f3",
            "64b713188f054535a3063c079bceead0",
            "9cea52b7835b40628fc0a65199c5e82f",
            "87a5c035bc6345aea6441b3db4b94321",
            "7f4f63d94f8942fcb3e4b286d5848f0c",
            "a310703500c5481aa37b698c659e1ca6",
            "7afff81060f0481aa95a56daa570e1b8",
            "a038f9037ed540f2ba43f3f85920b17e",
            "ec4f48d5a9ef4afa8fa44078dfb236c6",
            "405308ec76074c7ea319e1c1358baf39",
            "7479281e090d4eeba4c21a9c15b02165",
            "f4b25b3aac414499b653404ea3c066da",
            "0383a2798a0c4fcfbd1f068817770fd6",
            "385291e96814453b8f805cbe8a5e594d",
            "f7a41e267f79401681bc4693c52e55fa",
            "caff23247fb9432db8d5860a8132e1b3",
            "f47609606aa14f6fa2407aa26edae8e3",
            "cef2fe4eb51146afabe820140a3b0df3",
            "52281c8594ad465397e6079385c558a1",
            "e5302eb0481c4e529fc68353eed4b79f",
            "0c46bce8cc3e4393898b025ba6b4e543",
            "14008810744549208f147a132b419d20",
            "cb4be83440e44e1c9e664ac1874b830d",
            "f54c77317ca04861a958c8f982b25653",
            "db24a855fb5846e28c98dcfe9afff259",
            "647c861fcaa44148a90da061b1d5c99c",
            "e4e0b1de76ff40a88b71e0b5a9f7bc39",
            "504038c4efd3439daa958333f108b5fb",
            "ded8c188917b424dbf264ce494fe0532",
            "76fe13ab697c46cc9fcbe46112fb45b1",
            "a1c25910a31a4b5a9c48f1a7fd8047c5",
            "5d877979f2db4dbfa7581d2fe4381cf8"
          ]
        },
        "id": "1OD2z7d3HQJU",
        "outputId": "3a8f135a-ac39-47fe-da78-6680d42572ed"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "GPT2LMHeadModel(\n",
            "  (transformer): GPT2Model(\n",
            "    (wte): Embedding(50257, 768)\n",
            "    (wpe): Embedding(1024, 768)\n",
            "    (drop): Dropout(p=0.1, inplace=False)\n",
            "    (h): ModuleList(\n",
            "      (0-11): 12 x GPT2Block(\n",
            "        (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
            "        (attn): GPT2SdpaAttention(\n",
            "          (c_attn): Conv1D()\n",
            "          (c_proj): Conv1D()\n",
            "          (attn_dropout): Dropout(p=0.1, inplace=False)\n",
            "          (resid_dropout): Dropout(p=0.1, inplace=False)\n",
            "        )\n",
            "        (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
            "        (mlp): GPT2MLP(\n",
            "          (c_fc): Conv1D()\n",
            "          (c_proj): Conv1D()\n",
            "          (act): NewGELUActivation()\n",
            "          (dropout): Dropout(p=0.1, inplace=False)\n",
            "        )\n",
            "      )\n",
            "    )\n",
            "    (ln_f): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
            "  )\n",
            "  (lm_head): Linear(in_features=768, out_features=50257, bias=False)\n",
            "  (generator): Generator(\n",
            "    (streamer): Streamer()\n",
            "  )\n",
            ")\n"
          ]
        }
      ],
      "source": [
        "from nnsight import LanguageModel\n",
        "\n",
        "llm = LanguageModel(\"openai-community/gpt2\", device_map=\"auto\")\n",
        "\n",
        "print(llm)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Hw2VaAr_ezkj"
      },
      "source": [
        "<details>\n",
        "<summary>On Model Initialization</summary>\n",
        "\n",
        "---\n",
        "\n",
        "A few important things to note:\n",
        "\n",
        "Keyword arguments passed to the initialization of `LanguageModel` is forwarded\n",
        "to HuggingFace specific loading logic. In this case, `device_map` specifies\n",
        "which devices to use and its value `auto` indicates to evenly distribute it to\n",
        "all available GPUs (and CPU if no GPUs available). Other arguments can be found\n",
        "here:\n",
        "https://huggingface.co/docs/transformers/model_doc/auto#transformers.AutoModelForCausalLM\n",
        "\n",
        "When we initialize `LanguageModel`, we aren't yet loading the parameters of the\n",
        "model into memory. We are actually loading a 'meta' version of the model which\n",
        "doesn't take up any memory, but still allows us to view and trace actions on it.\n",
        "After exiting the first tracing context, the model is then fully loaded into\n",
        "memory. To load into memory on initialization, you can pass `dispatch=True` into\n",
        "`LanguageModel` like\n",
        "`LanguageModel('openai-community/gpt2', device_map=\"auto\", dispatch=True)`.\n",
        "\n",
        "---\n",
        "\n",
        "</details>\n",
        "\n",
        "Let's now apply some of the features that we used on the small model to `GPT-2`. Unlike `NNsight`, `LanguageModel` does define logic to pre-process\n",
        "inputs upon entering the tracing context. This makes interacting with the model\n",
        "simpler (i.e., you can send prompts to the model without having to directly access the tokenizer).\n",
        "\n",
        "In the following example, we ablate the value coming from the last layer's MLP\n",
        "module and decode the logits to see what token the model predicts without\n",
        "influence from that particular module:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 90,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 152,
          "referenced_widgets": [
            "2fedb8b529c34049932f07b577b515d4",
            "bd1a11654c2d47a09610e4cc23b8bc32",
            "26acae3d321743d2ad8efbf8727dbd05",
            "f42e44122a2d496cb7b2ea85c2858a71",
            "c1586aeef6814447ad46c1f8ddaae625",
            "5e091b15cc9748469676b165533f0c44",
            "f143977aeaf24798b7744681069762d4",
            "a4b5c225f5ac4719ade8e6aafcff78d2",
            "ee27d18a555c46e78c92b282bfa6dc9d",
            "a3154c60744d4a1a87b1c86b89a142b9",
            "de023d4526524f48a34883a287aa8a4f",
            "a15895ad1cb541ca89fa9def63779104",
            "ee60494714c2485f802298dce918edfe",
            "2a0f2484053f41298aceacea5c1cc195",
            "57aebe107749473390af2532806920ce",
            "ff44e896c55048c5a7d92a56aa9eedfc",
            "7fdb623e5b2d4eb1b6cee27da2314062",
            "95832c5076ab4c80810d3c443e061156",
            "08fb0620499c44eb9d160a0f21263451",
            "287a8fe3ddf54d36b7603171ec7a2688",
            "55c0bb35417e454d842e12bd2dfe62df",
            "d5821d69e1fa4673bb39ab70a7f60466"
          ]
        },
        "id": "mLSapaMCgLNU",
        "outputId": "107769db-2d2e-4b52-ed59-31555ebcbfcb"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Token IDs: tensor([[ 262,   12,  417, 8765,   11,  257,  262, 3504,  338, 3576]],\n",
            "       device='mps:0')\n",
            "Prediction:  London\n"
          ]
        }
      ],
      "source": [
        "with llm.trace(\"The Eiffel Tower is in the city of\"):\n",
        "\n",
        "    # Access the last layer using h[-1] as it's a ModuleList\n",
        "    # Access the first index of .output as that's where the hidden states are.\n",
        "    llm.transformer.h[-1].mlp.output[0][:] = 0\n",
        "\n",
        "    # Logits come out of model.lm_head and we apply argmax to get the predicted token ids.\n",
        "    token_ids = llm.lm_head.output.argmax(dim=-1).save()\n",
        "\n",
        "print(\"\\nToken IDs:\", token_ids)\n",
        "\n",
        "# Apply the tokenizer to decode the ids into words after the tracing context.\n",
        "print(\"Prediction:\", llm.tokenizer.decode(token_ids[0][-1]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "X2-HiCQhkv-B"
      },
      "source": [
        "We just ran a little intervention on a much more complex model with many more\n",
        "parameters! However, we're missing an important piece of information: what the\n",
        "prediction would have looked like without our ablation.\n",
        "\n",
        "We could just run two tracing contexts and compare the outputs. However, this would require two forward passes through the model. `NNsight` can do\n",
        "better than that with batching.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vLjHnRyRPXjp"
      },
      "source": [
        "<a name=\"batching-id\"></a>\n",
        "\n",
        "## Batching\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_Gc-Zr6NmEOn"
      },
      "source": [
        "Batching is a way to process multiple inputs in one forward pass. To better understand how batching works, we're going to bring back the `Tracer` object that we dropped before.\n",
        "\n",
        "When we call `.trace(...)`, it's actually creating two different contexts behind the scenes. The first one is the tracing context that we've discussed previously, and the second one is the invoker context. The invoker context defines the values of the `.input` and `.output` Proxies.\n",
        "\n",
        "If we call `.trace(...)` with some input, the input is passed on to the invoker. As there is only one input, only one invoker context is created.\n",
        "\n",
        "If we call `.trace()` without an input, then we can call `tracer.invoke(input1)` to manually create the invoker context with an input, `input1`. We can also repeatedly call `tracer.invoke(...)` to create the invoker context for additional inputs. Every subsequent time we call\n",
        "`.invoke(...)`, interventions within its context will only refer to the input in that particular invoke statement.\n",
        "\n",
        "When exiting the tracing context, the inputs from all of the invokers will be batched together, and they will be executed in one forward pass! To test this out, let's do the same ablation experiment, but also add a 'control' output for comparison:\n",
        "\n",
        "<details>\n",
        "<summary>More on the invoker context</summary>\n",
        "\n",
        "---\n",
        "\n",
        "Note that when injecting data to only the relevant invoker interventions, `nnsight` tries, but can't guarantee, to narrow the data into the right\n",
        "batch indices. Thus, there are cases\n",
        "where all invokes will get all of the data. Specifically, if the input or output data is stored\n",
        "as an object that is not an arbitrary collection of tensors, it will be broadcasted to all invokes.\n",
        "\n",
        "Just like `.trace(...)` created a `Tracer` object, `.invoke(...)` creates an `Invoker` object. For `LanguageModel` models, the `Invoker` prepares the input by running a tokenizer on it.\n",
        "`Invoker` stores pre-processed inputs at `invoker.inputs`, which can be accessed to see information about our inputs.\n",
        "In a case where we pass a single input to `.trace(...)` directly, we can still access the invoker\n",
        "object at `tracer.invoker` without having to call `tracer.invoke(...)`.\n",
        "\n",
        "Keyword arguments given to `.invoke(..)` make their way to the input pre-processing.  \n",
        "`LanguageModel` has keyword arguments `max_length` and `truncation` used for tokenization which can be\n",
        "passed to the invoker. If we want to pass keyword arguments to the invoker for a single-input `.trace(...)`, we can pass `invoker_args` as a dictionary of invoker keyword arguments.\n",
        "\n",
        "Here is an example to demonstrate everything we've described:\n",
        "\n",
        "**This snippet**\n",
        "\n",
        "```\n",
        "with llm.trace(\"hello\", invoker_args={\"max_length\":10}) as tracer:\n",
        "  invoker = tracer.invoker\n",
        "\n",
        "```\n",
        "  **does the same as**\n",
        "  \n",
        "\n",
        "```\n",
        "with llm.trace() as tracer:\n",
        "  with tracer.invoke(\"hello\", max_length=10) as invoker:\n",
        "    invoker = invoker\n",
        "```\n",
        "\n",
        "---\n",
        "\n",
        "</details>\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 91,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Kdcq4oCNmEua",
        "outputId": "33702cb9-46df-4c60-f108-f7ad3c1cd446"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "You're using a GPT2TokenizerFast tokenizer. Please note that with a fast tokenizer, using the `__call__` method is faster than using a method to encode the text followed by a call to the `pad` method to get a padded encoding.\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Original token IDs: tensor([[ 198,   12,  417, 8765,  318,  257,  262, 3504, 7372, 6342]],\n",
            "       device='mps:0')\n",
            "Modified token IDs: tensor([[ 262,   12,  417, 8765,   11,  257,  262, 3504,  338, 3576]],\n",
            "       device='mps:0')\n",
            "Original prediction:  Paris\n",
            "Modified prediction:  London\n"
          ]
        }
      ],
      "source": [
        "with llm.trace() as tracer:\n",
        "\n",
        "    with tracer.invoke(\"The Eiffel Tower is in the city of\"):\n",
        "\n",
        "        # Ablate the last MLP for only this batch.\n",
        "        llm.transformer.h[-1].mlp.output[0][:] = 0\n",
        "\n",
        "        # Get the output for only the intervened on batch.\n",
        "        token_ids_intervention = llm.lm_head.output.argmax(dim=-1).save()\n",
        "\n",
        "    with tracer.invoke(\"The Eiffel Tower is in the city of\"):\n",
        "\n",
        "        # Get the output for only the original batch.\n",
        "        token_ids_original = llm.lm_head.output.argmax(dim=-1).save()\n",
        "\n",
        "\n",
        "print(\"Original token IDs:\", token_ids_original)\n",
        "print(\"Modified token IDs:\", token_ids_intervention)\n",
        "\n",
        "print(\"Original prediction:\", llm.tokenizer.decode(token_ids_original[0][-1]))\n",
        "print(\"Modified prediction:\", llm.tokenizer.decode(token_ids_intervention[0][-1]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8BsybgsK2Bbr"
      },
      "source": [
        "Based on our control results, our ablation did end up affecting what the model predicted. That's pretty neat!\n",
        "\n",
        "Another cool thing with multiple invokes is that Proxies can interact between them.\n",
        "\n",
        "Here, we transfer the token embeddings from a real prompt into another placeholder prompt. Therefore the latter prompt produces the output of the former prompt:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 92,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 519
        },
        "id": "syKTI_KhpvCY",
        "outputId": "35afac9a-a34a-4e60-edc8-4c1b404e4389"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "original prediction shape torch.Size([])\n",
            "Original prediction:  _\n",
            "modified prediction shape torch.Size([])\n",
            "Modified prediction:  Paris\n"
          ]
        }
      ],
      "source": [
        "with llm.trace() as tracer:\n",
        "\n",
        "    with tracer.invoke(\"The Eiffel Tower is in the city of\"):\n",
        "        embeddings = llm.transformer.wte.output\n",
        "\n",
        "    with tracer.invoke(\"_ _ _ _ _ _ _ _ _ _\"):\n",
        "        llm.transformer.wte.output = embeddings\n",
        "        token_ids_intervention = llm.lm_head.output.argmax(dim=-1).save()\n",
        "\n",
        "    with tracer.invoke(\"_ _ _ _ _ _ _ _ _ _\"):\n",
        "      token_ids_original = llm.lm_head.output.argmax(dim=-1).save()\n",
        "\n",
        "print(\"original prediction shape\", token_ids_original[0][-1].shape)\n",
        "print(\"Original prediction:\", llm.tokenizer.decode(token_ids_original[0][-1]))\n",
        "\n",
        "print(\"modified prediction shape\", token_ids_intervention[0][-1].shape)\n",
        "print(\"Modified prediction:\", llm.tokenizer.decode(token_ids_intervention[0][-1]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_bXvUGGqbwv1"
      },
      "source": [
        "## Multiple Token Generation"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vvWA-CWqQtah"
      },
      "source": [
        "### .next()\n",
        "\n",
        "Some HuggingFace models define methods to generate multiple outputs at a time.\n",
        "`LanguageModel` wraps that functionality to provide the same tracing features by\n",
        "using `.generate(...)` instead of `.trace(...)`. This calls the underlying\n",
        "model's `.generate` method. It passes the output through a `.generator`\n",
        "module that we've added onto the model, allowing us to get the generate output\n",
        "at `.generator.output`.\n",
        "\n",
        "In a case like this, the underlying model is called more than once; the modules\n",
        "of said model produce more than one output. Which iteration should a given\n",
        "`module.output` refer to? That's where `Module.next()` comes in!\n",
        "\n",
        "Each module has a call index associated with it and `.next()` simply increments\n",
        "that attribute. At the time of execution, data is injected into the intervention\n",
        "graph only at the iteration that matches the call index.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 93,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Yy5Z9NE1GkaN",
        "outputId": "3827231c-86b1-482e-bd2f-ff17191692e5"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "Setting `pad_token_id` to `eos_token_id`:None for open-end generation.\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "torch.Size([1, 10, 768])\n",
            "torch.Size([1, 1, 768])\n",
            "tensor([[ 464,  412,  733,  417, 8765,  318,  287,  262, 1748,  286, 6342,   11,\n",
            "          290]], device='mps:0')\n"
          ]
        }
      ],
      "source": [
        "with llm.generate('The Eiffel Tower is in the city of', max_new_tokens=3) as tracer:\n",
        "\n",
        "    hidden_states1 = llm.transformer.h[-1].output[0].save()\n",
        "\n",
        "    # use module.next() to access the next intervention\n",
        "    hidden_states2 = llm.transformer.h[-1].next().output[0].save()\n",
        "\n",
        "    # saving the output allows you to save the hidden state across the initial prompt\n",
        "    out = llm.generator.output.save()\n",
        "\n",
        "print(hidden_states1.shape)\n",
        "print(hidden_states2.shape)\n",
        "print(out)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oHy6FqyNRBzm"
      },
      "source": [
        "### New! using .all()\n",
        "\n",
        "With `nnsight 0.4` you can now use `.all()` to recursively apply interventions to a model. Calling `.all()` on a module within a model will recursively apply its `.input` and `.output` across all iterations. Previously, we'd need to loop across each new generated token, saving the intervention for every generated token and calling `.next()` to move forward."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 94,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "C1C1KmclTJDH",
        "outputId": "97f4b5ba-3b94-4541-df98-4025b273ba97"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "Setting `pad_token_id` to `eos_token_id`:None for open-end generation.\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Hidden state length:  3\n"
          ]
        }
      ],
      "source": [
        "# Old approach:\n",
        "prompt = 'The Eiffel Tower is in the city of'\n",
        "layers = llm.transformer.h\n",
        "n_new_tokens = 3\n",
        "hidden_states = []\n",
        "with llm.generate(prompt, max_new_tokens=n_new_tokens) as tracer:\n",
        "    for i in range(n_new_tokens):\n",
        "        # Apply intervention - set first layer output to zero\n",
        "        layers[0].output[0][:] = 0\n",
        "\n",
        "        # Append desired hidden state post-intervention\n",
        "        hidden_states.append(layers[-1].output.save())\n",
        "\n",
        "        # Move to next generated token\n",
        "        layers[0].next()\n",
        "\n",
        "print(\"Hidden state length: \",len(hidden_states))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tuBLYSx_XKXK"
      },
      "source": [
        "We can use also `.all()` to streamline the multiple token generation process. We simply call `.all` on the module where we are applying the intervention (in this case GPT-2's layers), apply our intervention, and append our hidden states (stored in an `nnsight.list()` object).\n",
        "<br> <br>\n",
        "\n",
        "Let's test this out for the multiple token generation case:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 95,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "IgxXlvxDX1zF",
        "outputId": "9588942b-19b8-4d8b-90a3-9b8ef463b53f"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "Setting `pad_token_id` to `eos_token_id`:None for open-end generation.\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Hidden state length:  3\n"
          ]
        }
      ],
      "source": [
        "# using .all():\n",
        "prompt = 'The Eiffel Tower is in the city of'\n",
        "layers = llm.transformer.h\n",
        "n_new_tokens = 3\n",
        "with llm.generate(prompt, max_new_tokens=n_new_tokens) as tracer:\n",
        "    hidden_states = nnsight.list().save() # Initialize & .save() nnsight list\n",
        "\n",
        "    # Call .all() to apply intervention to each new token\n",
        "    layers.all()\n",
        "\n",
        "    # Apply intervention - set first layer output to zero\n",
        "    layers[0].output[0][:] = 0\n",
        "\n",
        "    # Append desired hidden state post-intervention\n",
        "    hidden_states.append(layers[-1].output) # no need to call .save\n",
        "    # Don't need to loop or call .next()!\n",
        "\n",
        "print(\"Hidden state length: \",len(hidden_states))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Z3dMEb6HX-Wp"
      },
      "source": [
        "Easy! Note that because `.all()` is recursive, it will only work to append outputs called on children of the module that `.all()` was called on. See example below for more information. TL;DR: apply `.all()` on the highest-level accessed module if interventions and outputs have different hierarchies within model structure."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3J_IeO96jVkZ"
      },
      "source": [
        "<details>\n",
        "<summary>Recursive properties of .all()</summary>\n",
        "\n",
        "`.all()` recursively acts on model components. In the below code example, only the first token generation is saved, because `.all()` applied to `layers`, while the saved variable `hidden_states` is produced from `model.lm_head`, which is not a child of `layers`.\n",
        "\n",
        "```\n",
        "prompt = 'The Eiffel Tower is in the city of'\n",
        "layers = model.transformer.h\n",
        "n_new_tokens = 3\n",
        "with model.generate(prompt, max_new_tokens=n_new_tokens) as tracer:\n",
        "    hidden_states = nnsight.list().save() # Initialize & .save() nnsight list\n",
        "\n",
        "    # Call .all() on layers\n",
        "    layers.all()\n",
        "\n",
        "    # Apply same intervention - set first layer output to zero\n",
        "    layers[0].output[0][:] = 0\n",
        "\n",
        "    # Append desired hidden state post-intervention\n",
        "    hidden_states.append(model.lm_head.output) # no need to call .save, it's already initialized\n",
        "\n",
        "print(\"Hidden state length: \",len(hidden_states)) # length is 1, meaning it only saved the first token generation\n",
        "```\n",
        "\n",
        "If you want to apply an intervention during multiple token generation while saving the state of a model component that isn't a child of that module, you can instead apply `.all()` to the full model:\n",
        "\n",
        "```\n",
        "prompt = 'The Eiffel Tower is in the city of'\n",
        "layers = model.transformer.h\n",
        "n_new_tokens = 3\n",
        "with model.generate(prompt, max_new_tokens=n_new_tokens) as tracer:\n",
        "    hidden_states = nnsight.list().save() # Initialize & .save() nnsight list\n",
        "\n",
        "    # Call .all() on model\n",
        "    model.all()\n",
        "\n",
        "    # Apply same intervention - set first layer output to zero\n",
        "    layers[0].output[0][:] = 0\n",
        "\n",
        "    # Append desired hidden state post-intervention\n",
        "    hidden_states.append(model.lm_head.output) # no need to call .save\n",
        "\n",
        "print(\"Hidden state length: \",len(hidden_states)) # length is 3, as expected!\n",
        "```\n",
        "\n",
        "</details>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ty-luCwCkkw7"
      },
      "source": [
        "## Model Editing"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hgrZVQLEkkw8"
      },
      "source": [
        "NNsight's model editing feature allows you to create persistently modified versions of a model with a use of `.edit()`. Unlike interventions in a tracing context, which are temporary, the **Editor** context enables you to make lasting changes to a model instance.\n",
        "\n",
        "This feature is useful for:\n",
        "* Creating modified model variants without altering the original\n",
        "* Applying changes that persist across multiple forward passes\n",
        "* Comparing interventions between original and edited models\n",
        "\n",
        "Let's explore how to use the **Editor** context to make a simple persistent change to a model:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 96,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "-_oAprpwkkw8",
        "outputId": "388b31e3-9585-4898-98be-7a1a043fa49c"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Original Prediction:   Rome\n",
            "Modified Prediction:   Paris\n"
          ]
        }
      ],
      "source": [
        "# we take the hidden states with the expected output \"Paris\"\n",
        "with llm.trace(\"The Eiffel Tower is located in the city of\") as tracer:\n",
        "    hs11 = llm.transformer.h[11].output[0][:, -1, :].save()\n",
        "\n",
        "# the edited model will now always predict \"Paris\" as the next token\n",
        "with llm.edit() as llm_edited:\n",
        "    llm.transformer.h[11].output[0][:, -1, :] = hs11\n",
        "\n",
        "# we demonstrate this by comparing the output of an unmodified model...\n",
        "with llm.trace(\"Vatican is located in the city of\") as tracer:\n",
        "    original_tokens = llm.lm_head.output.argmax(dim=-1).save()\n",
        "\n",
        "# ...with the output of the edited model\n",
        "with llm_edited.trace(\"Vatican is located in the city of\") as tracer:\n",
        "    modified_tokens = llm.lm_head.output.argmax(dim=-1).save()\n",
        "\n",
        "\n",
        "print(\"\\nOriginal Prediction: \", llm.tokenizer.decode(original_tokens[0][-1]))\n",
        "print(\"Modified Prediction: \", llm.tokenizer.decode(modified_tokens[0][-1]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "S8J6cVw0kkw8"
      },
      "source": [
        "Edits defined within an **Editor** context create a new, modified version of the model by default, preserving the original. This allows for safe experimentation with model changes. If you wish to modify the original model directly, you can set `inplace=True` when calling `.edit()`.\n",
        "\n",
        "Use this option cautiously, as in-place edits alter the base model for all the consequent model calls."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 97,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "b1atxBp8kkw8",
        "outputId": "d176ce44-a861-4407-fbf3-5c4411bd4343"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Modified In-place:   Paris\n"
          ]
        }
      ],
      "source": [
        "# we use the hidden state we saved above (hs11)\n",
        "with llm.edit(inplace=True) as llm_edited:\n",
        "    llm.transformer.h[11].output[0][:, -1, :] = hs11\n",
        "\n",
        "# we demonstrate this by comparing the output of an unmodified model...\n",
        "with llm.trace(\"Vatican is located in the city of\") as tracer:\n",
        "    modified_tokens = llm.lm_head.output.argmax(dim=-1).save()\n",
        "\n",
        "print(\"Modified In-place: \", llm.tokenizer.decode(modified_tokens[0][-1]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nVeZO96jkkw8"
      },
      "source": [
        "If you've made in-place edits to your model and need to revert these changes, you can apply `.clear_edits()`. This method removes all edits applied to the model, effectively restoring it to its original state."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 98,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "_SHK4JzPkkw8",
        "outputId": "c6da4d26-e5b8-415e-aa01-071fbdbd4802"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Edits cleared:   Rome\n"
          ]
        }
      ],
      "source": [
        "llm.clear_edits()\n",
        "\n",
        "with llm.trace(\"Vatican is located in the city of\"):\n",
        "    modified_tokens = llm.lm_head.output.argmax(dim=-1).save()\n",
        "\n",
        "print(\"Edits cleared: \", llm.tokenizer.decode(modified_tokens[0][-1]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rpkX-LwBQZHo"
      },
      "source": [
        "# 3 I thought you said huge models?"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uCZR65VmILEb"
      },
      "source": [
        "`NNsight` is only one part of our project to democratize access to AI internals. The other half is the National Deep Inference Fabric, or `NDIF`. `NDIF` hosts large models for shared access using `NNsight`, so you don't have to worry about any of the headaches of hosting large models yourself!\n",
        "\n",
        "The interaction between `NDIF` and `NNsight` is fairly straightforward. The\n",
        "**intervention graph** we create via the tracing context can be encoded into a\n",
        "custom json format and sent via an http request to the `NDIF` servers. `NDIF`\n",
        "then decodes the **intervention graph** and **interleaves** it alongside the\n",
        "specified model.\n",
        "\n",
        "To see which models are currently being hosted, check out the following status\n",
        "page: https://nnsight.net/status/\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "65Ks1LUvQaER"
      },
      "source": [
        "## Remote execution\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Aa-ZuVOFKS5k"
      },
      "source": [
        "In its current state, `NDIF` requires you to receive an API key. Therefore, to\n",
        "run the rest of this walkthrough, you need one of your own. To get one, simply\n",
        "register at https://#.ndif.us.\n",
        "\n",
        "With a valid API key, you then can configure `nnsight` as follows:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 100,
      "metadata": {
        "id": "ax1NWoS9MJeZ"
      },
      "outputs": [],
      "source": [
        "from nnsight import CONFIG\n",
        "\n",
        "CONFIG.set_default_api_key(\"YOUR_API_KEY\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mvXLOh65MXmF"
      },
      "source": [
        "If you're running in a local IDE, this only needs to be run once as it will save the API key as the default in a\n",
        ".config file along with your `nnsight` installation. You can also add your API key to Google Colab secrets.\n",
        "\n",
        "To amp things up a few levels, let's demonstrate using `nnsight`'s tracing\n",
        "context with `Llama-3.1-8b`!\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 101,
      "metadata": {
        "id": "Q1N04sJPZnJt"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "# Llama 3.1 8b is a gated model, so you need to apply for access on HuggingFace and include your token.\n",
        "os.environ['HF_TOKEN'] = \"YOUR_HUGGING_FACE_TOKEN\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "oqVzjNoyNGc7",
        "outputId": "abf3f296-c41b-475a-946b-6e0b046333e8"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "2025-01-31 15:47:57,111 69e5f67e-d0b4-4014-ba43-07c4d2b952bb - RECEIVED: Your job has been received and is waiting approval.\n",
            "2025-01-31 15:47:57,317 69e5f67e-d0b4-4014-ba43-07c4d2b952bb - APPROVED: Your job was approved and is waiting to be run.\n",
            "2025-01-31 15:47:57,519 69e5f67e-d0b4-4014-ba43-07c4d2b952bb - RUNNING: Your job has started running.\n",
            "2025-01-31 15:47:58,468 69e5f67e-d0b4-4014-ba43-07c4d2b952bb - COMPLETED: Your job has been completed.\n",
            "Downloading result:   0%|          | 0.00/4.37M [00:00<?, ?B/s]huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
            "To disable this warning, you can either:\n",
            "\t- Avoid using `tokenizers` before the fork if possible\n",
            "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n",
            "Downloading result: 100%|██████████| 4.37M/4.37M [00:00<00:00, 16.5MB/s]"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "(tensor([[[ 1.7734,  2.6875,  0.8047,  ..., -1.8594,  2.2344,  3.2656],\n",
            "         [ 0.0312, -0.0352, -2.8750,  ..., -0.8906, -0.0547,  1.6172],\n",
            "         [ 1.3594, -2.0156,  1.7031,  ..., -1.7031, -0.7422,  1.4375],\n",
            "         ...,\n",
            "         [ 1.0000,  0.3203, -0.2656,  ..., -0.0723, -0.2559,  0.2090],\n",
            "         [ 0.4707, -0.3496,  0.2422,  ...,  0.7344, -0.0078,  0.1133],\n",
            "         [-0.0566, -0.3496,  0.4746,  ...,  0.9844,  0.6797, -0.8750]]],\n",
            "       dtype=torch.bfloat16),)\n",
            "tensor([[[ 6.3438,  8.3750, 12.8125,  ..., -4.3750, -4.3750, -4.3750],\n",
            "         [-2.4375, -1.7266, -2.0156,  ..., -9.1250, -9.1250, -9.1250],\n",
            "         [ 9.6875,  4.5625,  5.8750,  ..., -3.3906, -3.3906, -3.3906],\n",
            "         ...,\n",
            "         [ 2.3281,  1.0703, -0.3203,  ..., -7.1562, -7.1562, -7.1562],\n",
            "         [11.1875,  6.0312,  4.9062,  ..., -3.5156, -3.5156, -3.5156],\n",
            "         [ 8.0000,  5.2500,  4.3750,  ..., -3.9844, -3.9844, -3.9844]]],\n",
            "       dtype=torch.bfloat16)\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\n"
          ]
        }
      ],
      "source": [
        "from nnsight import LanguageModel\n",
        "\n",
        "# We'll never actually load the parameters locally, so no need to specify a device_map.\n",
        "\n",
        "llama = LanguageModel(\"meta-llama/Meta-Llama-3.1-8B\")\n",
        "# All we need to specify using NDIF vs executing locally is remote=True.\n",
        "with llama.trace(\"The Eiffel Tower is in the city of\", remote=True) as runner:\n",
        "\n",
        "    hidden_states = llama.model.layers[-1].output.save()\n",
        "\n",
        "    output = llama.output.save()\n",
        "\n",
        "print(hidden_states)\n",
        "\n",
        "print(output[\"logits\"])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IceVpnZcZ9F0"
      },
      "source": [
        "It really is as simple as `remote=True`. All of the techniques we went through\n",
        "in earlier sections work just the same when running locally or remotely."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lfRDZox5kkw8"
      },
      "source": [
        "## Sessions"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "33XLOvtukkw8"
      },
      "source": [
        "NDIF uses a queue to handle concurrent requests from multiple users. To optimize the execution of our experiments we can use the `session` context to efficiently package multiple interventions together as one single request to the server.\n",
        "\n",
        "This offers the following benefits:\n",
        "1.   All interventions within a session will be executed one after another without additional wait in the NDIF queue\n",
        "2.   All intermediate outputs for each intervention are stored on the server and can be accessed by other interventions in the same session without moving the data back and forth between NDIF and the local machine\n",
        "\n",
        "Let's take a look:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "KA6GROiTR3gf",
        "outputId": "dfefe423-9caf-48ce-cfe5-83c2f808a6b2"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "2025-01-31 15:48:17,378 d82e4471-0bb0-4ecc-8bb3-6e8fe35f0c03 - RECEIVED: Your job has been received and is waiting approval.\n",
            "2025-01-31 15:48:17,558 d82e4471-0bb0-4ecc-8bb3-6e8fe35f0c03 - APPROVED: Your job was approved and is waiting to be run.\n",
            "2025-01-31 15:48:17,755 d82e4471-0bb0-4ecc-8bb3-6e8fe35f0c03 - RUNNING: Your job has started running.\n",
            "2025-01-31 15:48:18,205 d82e4471-0bb0-4ecc-8bb3-6e8fe35f0c03 - COMPLETED: Your job has been completed.\n",
            "Downloading result: 100%|██████████| 1.62k/1.62k [00:00<00:00, 6.57MB/s]"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "T1 - Original Prediction:   Paris\n",
            "T2 - Modified Prediction:  ://\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\n"
          ]
        }
      ],
      "source": [
        "with llama.session(remote=True) as session:\n",
        "\n",
        "  with llama.trace(\"The Eiffel Tower is in the city of\") as t1:\n",
        "    # capture the hidden state from layer 32 at the last token\n",
        "    hs_31 = llama.model.layers[31].output[0][:, -1, :] # no .save()\n",
        "    t1_tokens_out = llama.lm_head.output.argmax(dim=-1).save()\n",
        "\n",
        "  with llama.trace(\"Buckingham Palace is in the city of\") as t2:\n",
        "    llama.model.layers[1].output[0][:, -1, :] = hs_31[:]\n",
        "    t2_tokens_out = llama.lm_head.output.argmax(dim=-1).save()\n",
        "\n",
        "print(\"\\nT1 - Original Prediction: \", llama.tokenizer.decode(t1_tokens_out[0][-1]))\n",
        "print(\"T2 - Modified Prediction: \", llama.tokenizer.decode(t2_tokens_out[0][-1]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "End2DR7qkkw8"
      },
      "source": [
        "In the example above, we are interested in replacing the hidden state of a later layer with an earlier one. Since we are using a `session`, we don't have to save the hidden state from Tracer 1 to reference it in Tracer 2.\n",
        "\n",
        "It is important to note that all the traces defined within the `session` context are executed sequentially, strictly following the order of definition (i.e. `t2` being executed after `t1` and `t3` after `t2` etc.)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "M-Nq8q0Qkkw8"
      },
      "source": [
        "The `session` context object has its own methods to log values and be terminated early."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 104,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 311
        },
        "id": "Q3c6M5Jwkkw8",
        "outputId": "0d937244-d022-4745-f3ea-f10b305bc3e3"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "2025-01-31 15:17:51,636 3d5144b4-e0b6-4106-912f-5abec8e78e7f - RECEIVED: Your job has been received and is waiting approval.\n",
            "2025-01-31 15:17:51,804 3d5144b4-e0b6-4106-912f-5abec8e78e7f - APPROVED: Your job was approved and is waiting to be run.\n",
            "2025-01-31 15:17:52,023 3d5144b4-e0b6-4106-912f-5abec8e78e7f - RUNNING: Your job has started running.\n",
            "2025-01-31 15:17:52,028 3d5144b4-e0b6-4106-912f-5abec8e78e7f - LOG: -- Early Stop --\n",
            "2025-01-31 15:17:52,364 3d5144b4-e0b6-4106-912f-5abec8e78e7f - COMPLETED: Your job has been completed.\n",
            "Downloading result: 100%|██████████| 928/928 [00:00<00:00, 7.25MB/s]\n"
          ]
        }
      ],
      "source": [
        "with llama.session(remote=True) as session:\n",
        "  session.log(\"-- Early Stop --\")\n",
        "  nnsight.stop"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5Z-tTpQekkw8"
      },
      "source": [
        "In addition to the benefits mentioned above, the `session` context also enables interesting experiments not possible with other `nnsight` tools — since every trace is run on its own model, it means that within one session we can run interventions between different models — for example, we could swap activations between base and instruct versions of the Llama model and compare their outputs. And `session` can also be used to run similar experiments entirely locally!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "H6hxSZ8zTFvT"
      },
      "source": [
        "## Streaming"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4bDUYaluTRCM"
      },
      "source": [
        "Streaming enables users apply functions and datasets locally during remote model execution. This allows users to stream results for immediate consumption (i.e., seeing tokens as they are generated) or applying non-whitelisted functions such as model tokenizers, large local datasets, and more!\n",
        "\n",
        "*   `nnsight.local()` context sends values immediately to user's local machine from server\n",
        "*   Intervention graph is executed locally on downstream nodes\n",
        "*   Exiting local context uploads data back to server\n",
        "*   `@nnsight.trace` function decorator enables custom functions to be added to intervention graph when using `nnsight.local()`"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fBDttzi7TT26"
      },
      "source": [
        "### `nnsight.local()`"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NRsfBs5MTW6f"
      },
      "source": [
        "You may sometimes want to locally access and manipulate values during remote execution. Using `.local()` on a proxy, you can send remote content to your local machine and apply local functions. The intervention graph is then executed locally on downstream nodes (until you send execution back to the remote server by exiting the `.local()` context).\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gN8yWbc4TXZc"
      },
      "source": [
        "There are a few use cases for streaming with `.local()`, including live chat generation and applying large datasets or non-whitelisted local functions to the intervention graph."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fFxwkDC6TZ8C"
      },
      "source": [
        "Now let's explore how streaming works. We'll start by grabbing some hidden states of the model and printing their value using `tracer.log()`. Without calling `nnsight.local()`, these operations will all occur remotely."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 120,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "QKMwwmoUTcCI",
        "outputId": "f1400168-1f44-43c4-91c1-14deaeedb4d0"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "2025-01-31 15:21:29,379 94c37060-d3f6-4f21-b55f-0e72f03d08da - RECEIVED: Your job has been received and is waiting approval.\n",
            "2025-01-31 15:21:29,592 94c37060-d3f6-4f21-b55f-0e72f03d08da - APPROVED: Your job was approved and is waiting to be run.\n",
            "2025-01-31 15:21:29,786 94c37060-d3f6-4f21-b55f-0e72f03d08da - RUNNING: Your job has started running.\n",
            "2025-01-31 15:21:29,822 94c37060-d3f6-4f21-b55f-0e72f03d08da - LOG: tensor(1.7656, device='cuda:0')\n",
            "2025-01-31 15:21:30,161 94c37060-d3f6-4f21-b55f-0e72f03d08da - COMPLETED: Your job has been completed.\n",
            "Downloading result: 100%|██████████| 514k/514k [00:00<00:00, 4.80MB/s]\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "tensor([[[ 6.3438,  8.3750, 12.8125,  ..., -4.3750, -4.3750, -4.3750],\n",
            "         [10.2500,  2.1094,  2.8281,  ..., -8.2500, -8.2500, -8.2500]]],\n",
            "       dtype=torch.bfloat16)\n"
          ]
        }
      ],
      "source": [
        "# This will give you a remote LOG response because it's coming from the remote server\n",
        "with llama.trace(\"hello\", remote=True) as tracer:\n",
        "\n",
        "    hs = llama.model.layers[-1].output[0]\n",
        "\n",
        "    tracer.log(hs[0,0,0])\n",
        "\n",
        "    out =  llama.lm_head.output.save()\n",
        "\n",
        "print(out)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1IjPUn8oTouV"
      },
      "source": [
        "Now, let's try the same operation using the `nnsight.local()` context. This will send the operations to get and print the hidden states to your local machine, changing how the logging message is formatted (local formatting instead of remote)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 121,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "x9byLVl_Tqy-",
        "outputId": "e1514fdd-61b5-4cbe-c577-1bdab9605adc"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "2025-01-31 15:21:32,674 0414b4df-a160-485b-89dd-4e70a05a43e2 - RECEIVED: Your job has been received and is waiting approval.\n",
            "2025-01-31 15:21:32,830 0414b4df-a160-485b-89dd-4e70a05a43e2 - APPROVED: Your job was approved and is waiting to be run.\n",
            "2025-01-31 15:21:33,092 0414b4df-a160-485b-89dd-4e70a05a43e2 - RUNNING: Your job has started running.\n",
            "2025-01-31 15:21:33,352 0414b4df-a160-485b-89dd-4e70a05a43e2 - COMPLETED: Your job has been completed.\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "tensor(1.7656, dtype=torch.bfloat16)\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "Downloading result: 100%|██████████| 514k/514k [00:00<00:00, 4.28MB/s]"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "tensor([[[ 6.3438,  8.3750, 12.8125,  ..., -4.3750, -4.3750, -4.3750],\n",
            "         [10.2500,  2.1094,  2.8281,  ..., -8.2500, -8.2500, -8.2500]]],\n",
            "       dtype=torch.bfloat16)\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\n"
          ]
        }
      ],
      "source": [
        "# This will print locally because it's already local\n",
        "with llama.trace(\"hello\", remote=True) as tracer:\n",
        "\n",
        "    with nnsight.local():\n",
        "        hs = llama.model.layers[-1].output[0]\n",
        "        tracer.log(hs[0,0,0])\n",
        "\n",
        "    out =  llama.lm_head.output.save()\n",
        "\n",
        "print(out)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tNbmIZSCTwo2"
      },
      "source": [
        "### `@nnsight.trace` function decorator"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xiQOxe9_TzqI"
      },
      "source": [
        "We can also use function decorators to create custom functions to be used during `.local` calls. This is a handy way to enable live streaming of a chat or to train probing classifiers on model hidden states."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "X7A3IjhLT2TE"
      },
      "source": [
        "Let's try out `@nnsight.trace` and `nnsight.local()` to access a custom function during remote execution."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 122,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "P8ghjV-7TzPj",
        "outputId": "6e5a2efb-8edd-40ff-a4de-84c535bbbd3d"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "2025-01-31 15:21:35,853 3eeee112-3074-4b3e-adea-001a85ce554d - RECEIVED: Your job has been received and is waiting approval.\n",
            "2025-01-31 15:21:36,035 3eeee112-3074-4b3e-adea-001a85ce554d - APPROVED: Your job was approved and is waiting to be run.\n",
            "2025-01-31 15:21:36,258 3eeee112-3074-4b3e-adea-001a85ce554d - RUNNING: Your job has started running.\n",
            "2025-01-31 15:21:36,653 3eeee112-3074-4b3e-adea-001a85ce554d - COMPLETED: Your job has been completed.\n",
            "Downloading result: 100%|██████████| 258k/258k [00:00<00:00, 2.91MB/s]\n"
          ]
        }
      ],
      "source": [
        "# first, let's define our function\n",
        "@nnsight.trace # decorator that enables this function to be added to the intervention graph\n",
        "def my_local_fn(value):\n",
        "    return value * 0\n",
        "\n",
        "# We use a local function to ablate some hidden states\n",
        "# This downloads the data for the .local context, and then uploads it back to set the value.\n",
        "with llama.generate(\"hello\", remote=True) as tracer:\n",
        "\n",
        "    hs = llama.model.layers[-1].output[0]\n",
        "\n",
        "    with nnsight.local():\n",
        "\n",
        "        hs = my_local_fn(hs)\n",
        "\n",
        "    llama.model.layers[-1].output[0][:] = hs\n",
        "\n",
        "    out =  llama.lm_head.output.save()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_7d5AdU6T72j"
      },
      "source": [
        "Note that without calling `.local`, the remote API does not know about `my_local_fn` and will throw a whitelist error. A whitelist error occurs because you are being allowed access to the function."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 123,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 472
        },
        "id": "n4OIIqQhT98b",
        "outputId": "54e6ad3d-a882-4892-aac0-7bdba52cac40"
      },
      "outputs": [
        {
          "ename": "FunctionWhitelistError",
          "evalue": "Function with name `__main__.my_local_fn` not in function whitelist.",
          "output_type": "error",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mFunctionWhitelistError\u001b[0m                    Traceback (most recent call last)",
            "Cell \u001b[0;32mIn[123], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m llama\u001b[38;5;241m.\u001b[39mtrace(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mhello\u001b[39m\u001b[38;5;124m\"\u001b[39m, remote\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m) \u001b[38;5;28;01mas\u001b[39;00m tracer:\n\u001b[1;32m      3\u001b[0m     hs \u001b[38;5;241m=\u001b[39m llama\u001b[38;5;241m.\u001b[39mmodel\u001b[38;5;241m.\u001b[39mlayers[\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m]\u001b[38;5;241m.\u001b[39moutput[\u001b[38;5;241m0\u001b[39m]\n\u001b[1;32m      5\u001b[0m     hs \u001b[38;5;241m=\u001b[39m my_local_fn(hs) \u001b[38;5;66;03m# no .local - will cause an error\u001b[39;00m\n",
            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/intervention/contexts/interleaving.py:96\u001b[0m, in \u001b[0;36mInterleavingTracer.__exit__\u001b[0;34m(self, exc_type, exc_val, exc_tb)\u001b[0m\n\u001b[1;32m     92\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39minvoker\u001b[38;5;241m.\u001b[39m\u001b[38;5;21m__exit__\u001b[39m(\u001b[38;5;28;01mNone\u001b[39;00m, \u001b[38;5;28;01mNone\u001b[39;00m, \u001b[38;5;28;01mNone\u001b[39;00m)\n\u001b[1;32m     94\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_model\u001b[38;5;241m.\u001b[39m_envoy\u001b[38;5;241m.\u001b[39m_reset()\n\u001b[0;32m---> 96\u001b[0m \u001b[38;5;28;43msuper\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[38;5;21;43m__exit__\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43mexc_type\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mexc_val\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mexc_tb\u001b[49m\u001b[43m)\u001b[49m\n",
            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/tracing/contexts/tracer.py:25\u001b[0m, in \u001b[0;36mTracer.__exit__\u001b[0;34m(self, exc_type, exc_val, exc_tb)\u001b[0m\n\u001b[1;32m     21\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mglobals\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m GlobalTracingContext\n\u001b[1;32m     23\u001b[0m GlobalTracingContext\u001b[38;5;241m.\u001b[39mtry_deregister(\u001b[38;5;28mself\u001b[39m)\n\u001b[0;32m---> 25\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43msuper\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[38;5;21;43m__exit__\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43mexc_type\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mexc_val\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mexc_tb\u001b[49m\u001b[43m)\u001b[49m\n",
            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/tracing/contexts/base.py:82\u001b[0m, in \u001b[0;36mContext.__exit__\u001b[0;34m(self, exc_type, exc_val, exc_tb)\u001b[0m\n\u001b[1;32m     78\u001b[0m graph \u001b[38;5;241m=\u001b[39m graph\u001b[38;5;241m.\u001b[39mstack\u001b[38;5;241m.\u001b[39mpop()\n\u001b[1;32m     80\u001b[0m graph\u001b[38;5;241m.\u001b[39malive \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mFalse\u001b[39;00m\n\u001b[0;32m---> 82\u001b[0m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbackend\u001b[49m\u001b[43m(\u001b[49m\u001b[43mgraph\u001b[49m\u001b[43m)\u001b[49m\n",
            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/intervention/backends/remote.py:77\u001b[0m, in \u001b[0;36mRemoteBackend.__call__\u001b[0;34m(self, graph)\u001b[0m\n\u001b[1;32m     72\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__call__\u001b[39m(\u001b[38;5;28mself\u001b[39m, graph: Graph):\n\u001b[1;32m     74\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mblocking:\n\u001b[1;32m     75\u001b[0m \n\u001b[1;32m     76\u001b[0m         \u001b[38;5;66;03m# Do blocking request.\u001b[39;00m\n\u001b[0;32m---> 77\u001b[0m         result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mblocking_request\u001b[49m\u001b[43m(\u001b[49m\u001b[43mgraph\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     79\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m     80\u001b[0m \n\u001b[1;32m     81\u001b[0m         \u001b[38;5;66;03m# Otherwise we are getting the status / result of the existing job.\u001b[39;00m\n\u001b[1;32m     82\u001b[0m         result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mnon_blocking_request(graph)\n",
            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/intervention/backends/remote.py:289\u001b[0m, in \u001b[0;36mRemoteBackend.blocking_request\u001b[0;34m(self, graph)\u001b[0m\n\u001b[1;32m    280\u001b[0m sio\u001b[38;5;241m.\u001b[39mconnect(\n\u001b[1;32m    281\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mws_address,\n\u001b[1;32m    282\u001b[0m     socketio_path\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m/ws/socket.io\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m    283\u001b[0m     transports\u001b[38;5;241m=\u001b[39m[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mwebsocket\u001b[39m\u001b[38;5;124m\"\u001b[39m],\n\u001b[1;32m    284\u001b[0m     wait_timeout\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m10\u001b[39m,\n\u001b[1;32m    285\u001b[0m )\n\u001b[1;32m    287\u001b[0m remote_graph \u001b[38;5;241m=\u001b[39m preprocess(graph)\n\u001b[0;32m--> 289\u001b[0m data, headers \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrequest\u001b[49m\u001b[43m(\u001b[49m\u001b[43mremote_graph\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    291\u001b[0m headers[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msession_id\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m=\u001b[39m sio\u001b[38;5;241m.\u001b[39msid\n\u001b[1;32m    293\u001b[0m \u001b[38;5;66;03m# Submit request via\u001b[39;00m\n",
            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/intervention/backends/remote.py:60\u001b[0m, in \u001b[0;36mRemoteBackend.request\u001b[0;34m(self, graph)\u001b[0m\n\u001b[1;32m     58\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mrequest\u001b[39m(\u001b[38;5;28mself\u001b[39m, graph: Graph) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Tuple[\u001b[38;5;28mbytes\u001b[39m, Dict[\u001b[38;5;28mstr\u001b[39m, \u001b[38;5;28mstr\u001b[39m]]:\n\u001b[0;32m---> 60\u001b[0m     data \u001b[38;5;241m=\u001b[39m \u001b[43mRequestModel\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mserialize\u001b[49m\u001b[43m(\u001b[49m\u001b[43mgraph\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mformat\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mzlib\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     62\u001b[0m     headers \u001b[38;5;241m=\u001b[39m {\n\u001b[1;32m     63\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mmodel_key\u001b[39m\u001b[38;5;124m\"\u001b[39m: \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mmodel_key,\n\u001b[1;32m     64\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mformat\u001b[39m\u001b[38;5;124m\"\u001b[39m: \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mformat,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     67\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msent-timestamp\u001b[39m\u001b[38;5;124m\"\u001b[39m: \u001b[38;5;28mstr\u001b[39m(time\u001b[38;5;241m.\u001b[39mtime()),\n\u001b[1;32m     68\u001b[0m     }\n\u001b[1;32m     70\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m data, headers\n",
            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/schema/request.py:43\u001b[0m, in \u001b[0;36mRequestModel.serialize\u001b[0;34m(graph, format, _zlib)\u001b[0m\n\u001b[1;32m     38\u001b[0m \u001b[38;5;129m@staticmethod\u001b[39m\n\u001b[1;32m     39\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mserialize\u001b[39m(graph: Graph, \u001b[38;5;28mformat\u001b[39m:\u001b[38;5;28mstr\u001b[39m, _zlib:\u001b[38;5;28mbool\u001b[39m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m \u001b[38;5;28mbytes\u001b[39m:\n\u001b[1;32m     41\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mformat\u001b[39m \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mjson\u001b[39m\u001b[38;5;124m\"\u001b[39m:\n\u001b[0;32m---> 43\u001b[0m         data \u001b[38;5;241m=\u001b[39m \u001b[43mRequestModel\u001b[49m\u001b[43m(\u001b[49m\u001b[43mgraph\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mgraph\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     45\u001b[0m         json \u001b[38;5;241m=\u001b[39m data\u001b[38;5;241m.\u001b[39mmodel_dump(mode\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mjson\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m     47\u001b[0m         data \u001b[38;5;241m=\u001b[39m msgspec\u001b[38;5;241m.\u001b[39mjson\u001b[38;5;241m.\u001b[39mencode(json)\n",
            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/schema/request.py:30\u001b[0m, in \u001b[0;36mRequestModel.__init__\u001b[0;34m(self, memo, *args, **kwargs)\u001b[0m\n\u001b[1;32m     28\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__init__\u001b[39m(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;241m*\u001b[39margs, memo: Dict \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs):\n\u001b[0;32m---> 30\u001b[0m     \u001b[38;5;28;43msuper\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[38;5;21;43m__init__\u001b[39;49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmemo\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mmemo\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01mor\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;28;43mdict\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     32\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m memo \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m     34\u001b[0m         \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mmemo \u001b[38;5;241m=\u001b[39m {\u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mMEMO}\n",
            "    \u001b[0;31m[... skipping hidden 1 frame]\u001b[0m\n",
            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/schema/format/types.py:276\u001b[0m, in \u001b[0;36mGraphModel.to_model\u001b[0;34m(value)\u001b[0m\n\u001b[1;32m    273\u001b[0m \u001b[38;5;129m@staticmethod\u001b[39m\n\u001b[1;32m    274\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mto_model\u001b[39m(value: Graph) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Self:\n\u001b[0;32m--> 276\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mGraphModel\u001b[49m\u001b[43m(\u001b[49m\u001b[43mgraph\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mvalue\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mnodes\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mvalue\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mnodes\u001b[49m\u001b[43m)\u001b[49m\n",
            "    \u001b[0;31m[... skipping hidden 1 frame]\u001b[0m\n",
            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/schema/format/types.py:77\u001b[0m, in \u001b[0;36mmemoized.<locals>.inner\u001b[0;34m(value)\u001b[0m\n\u001b[1;32m     75\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21minner\u001b[39m(value):\n\u001b[0;32m---> 77\u001b[0m     model \u001b[38;5;241m=\u001b[39m \u001b[43mfn\u001b[49m\u001b[43m(\u001b[49m\u001b[43mvalue\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     79\u001b[0m     _id \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mid\u001b[39m(value)\n\u001b[1;32m     81\u001b[0m     MEMO[_id] \u001b[38;5;241m=\u001b[39m model\n",
            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/schema/format/types.py:101\u001b[0m, in \u001b[0;36mNodeModel.to_model\u001b[0;34m(value)\u001b[0m\n\u001b[1;32m     97\u001b[0m \u001b[38;5;129m@staticmethod\u001b[39m\n\u001b[1;32m     98\u001b[0m \u001b[38;5;129m@memoized\u001b[39m\n\u001b[1;32m     99\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mto_model\u001b[39m(value: Node) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Self:\n\u001b[0;32m--> 101\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mNodeModel\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtarget\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mvalue\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtarget\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43margs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mvalue\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mkwargs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mvalue\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
            "    \u001b[0;31m[... skipping hidden 1 frame]\u001b[0m\n",
            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/schema/format/types.py:244\u001b[0m, in \u001b[0;36mFunctionModel.to_model\u001b[0;34m(value)\u001b[0m\n\u001b[1;32m    239\u001b[0m \u001b[38;5;129m@staticmethod\u001b[39m\n\u001b[1;32m    240\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mto_model\u001b[39m(value:FUNCTION):\n\u001b[1;32m    242\u001b[0m     model \u001b[38;5;241m=\u001b[39m FunctionModel(function_name\u001b[38;5;241m=\u001b[39mget_function_name(value))\n\u001b[0;32m--> 244\u001b[0m     \u001b[43mFunctionModel\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcheck_function_whitelist\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmodel\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfunction_name\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    246\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m model\n",
            "File \u001b[0;32m~/Documents/Projects/nnsight/src/nnsight/schema/format/types.py:251\u001b[0m, in \u001b[0;36mFunctionModel.check_function_whitelist\u001b[0;34m(cls, qualname)\u001b[0m\n\u001b[1;32m    248\u001b[0m \u001b[38;5;129m@classmethod\u001b[39m\n\u001b[1;32m    249\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mcheck_function_whitelist\u001b[39m(\u001b[38;5;28mcls\u001b[39m, qualname: \u001b[38;5;28mstr\u001b[39m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m \u001b[38;5;28mstr\u001b[39m:\n\u001b[1;32m    250\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m qualname \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;129;01min\u001b[39;00m FUNCTIONS_WHITELIST:\n\u001b[0;32m--> 251\u001b[0m         \u001b[38;5;28;01mraise\u001b[39;00m FunctionWhitelistError(\n\u001b[1;32m    252\u001b[0m             \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mFunction with name `\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mqualname\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m` not in function whitelist.\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    253\u001b[0m         )\n\u001b[1;32m    255\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m qualname\n",
            "\u001b[0;31mFunctionWhitelistError\u001b[0m: Function with name `__main__.my_local_fn` not in function whitelist."
          ]
        }
      ],
      "source": [
        "with llama.trace(\"hello\", remote=True) as tracer:\n",
        "\n",
        "    hs = llama.model.layers[-1].output[0]\n",
        "\n",
        "    hs = my_local_fn(hs) # no .local - will cause an error\n",
        "\n",
        "    llama.model.layers[-1].output[0][:] = hs * 2\n",
        "\n",
        "    out =  llama.lm_head.output.save()\n",
        "\n",
        "print(out)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "v_bdhze-UBec"
      },
      "source": [
        "### Example: Live-streaming remote chat\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hguvYgvVUE4b"
      },
      "source": [
        "Now that we can access data within the tracing context on our local computer, we can apply non-whitelisted functions, such as the model's tokenizer, within our tracing context.\n",
        "\n",
        "Let's build a decoding function that will decode tokens into words and print the result."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 124,
      "metadata": {
        "id": "8OAXweC5UHBJ"
      },
      "outputs": [],
      "source": [
        "@nnsight.trace\n",
        "def my_decoding_function(tokens, model, max_length=80, state=None):\n",
        "    # Initialize state if not provided\n",
        "    if state is None:\n",
        "        state = {'current_line': '', 'current_line_length': 0}\n",
        "\n",
        "    token = tokens[-1] # only use last token\n",
        "\n",
        "    # Decode the token\n",
        "    decoded_token = llama.tokenizer.decode(token).encode(\"unicode_escape\").decode()\n",
        "\n",
        "    if decoded_token == '\\\\n':  # Handle explicit newline tokens\n",
        "        # Print the current line and reset state\n",
        "        print('',flush=True)\n",
        "        state['current_line'] = ''\n",
        "        state['current_line_length'] = 0\n",
        "    else:\n",
        "        # Check if adding the token would exceed the max length\n",
        "        if state['current_line_length'] + len(decoded_token) > max_length:\n",
        "            print('',flush=True)\n",
        "            state['current_line'] = decoded_token  # Start a new line with the current token\n",
        "            state['current_line_length'] = len(decoded_token)\n",
        "            print(state['current_line'], flush=True, end=\"\")  # Print the current line\n",
        "        else:\n",
        "            # Add a space if the line isn't empty and append the token\n",
        "            if state['current_line']:\n",
        "                state['current_line'] += decoded_token\n",
        "            else:\n",
        "                state['current_line'] = decoded_token\n",
        "            state['current_line_length'] += len(decoded_token)\n",
        "            print(state['current_line'], flush=True, end=\"\")  # Print the current line\n",
        "\n",
        "    return state"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Uf-Ir1YgUKJt"
      },
      "source": [
        "Now we can decode and print our model outputs throughout token generation by accessing our decoding function through `nnsight.local()`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 125,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "aT56JuDzULyJ",
        "outputId": "bfa7593d-01c2-46e0-8c67-d8dbc1c7aa93"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Prompt:  A press release is an official statement delivered to members of the news media for the purpose of \n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "2025-01-31 15:21:48,944 94a118bf-3d5d-4ee8-8cf6-2778e95b745f - RECEIVED: Your job has been received and is waiting approval.\n",
            "2025-01-31 15:21:49,121 94a118bf-3d5d-4ee8-8cf6-2778e95b745f - APPROVED: Your job was approved and is waiting to be run.\n",
            "2025-01-31 15:21:49,363 94a118bf-3d5d-4ee8-8cf6-2778e95b745f - RUNNING: Your job has started running.\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            " providing information, an official statement, or making an announcement.A press release is also written or recorded communication directed at members of the news media for the purpose of announcing something ostensibly newsworthy. Typically, they are mailed, faxed, or e"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "2025-01-31 15:21:52,277 94a118bf-3d5d-4ee8-8cf6-2778e95b745f - COMPLETED: Your job has been completed.\n",
            "Downloading result: 100%|██████████| 258k/258k [00:00<00:00, 3.03MB/s]\n"
          ]
        }
      ],
      "source": [
        "import torch\n",
        "\n",
        "nnsight.CONFIG.APP.REMOTE_LOGGING = False\n",
        "\n",
        "prompt = \"A press release is an official statement delivered to members of the news media for the purpose of\"\n",
        "# prompt = \"Your favorite board game is\"\n",
        "\n",
        "print(\"Prompt: \",prompt,'\\n', end =\"\")\n",
        "\n",
        "# Initialize the state for decoding\n",
        "state = {'current_line': '', 'current_line_length': 0}\n",
        "\n",
        "with llama.generate(prompt, remote=True, max_new_tokens = 50) as generator:\n",
        "    # Call .all() to apply to each new token\n",
        "    llama.all()\n",
        "\n",
        "    all_tokens = nnsight.list().save()\n",
        "\n",
        "    # Access model output\n",
        "    out = llama.lm_head.output.save()\n",
        "\n",
        "    # Apply softmax to obtain probabilities and save the result\n",
        "    probs = torch.nn.functional.softmax(out, dim=-1)\n",
        "    max_probs = torch.max(probs, dim=-1)\n",
        "    tokens = max_probs.indices.cpu().tolist()\n",
        "    all_tokens.append(tokens[0]).save()\n",
        "\n",
        "    with nnsight.local():\n",
        "        state = my_decoding_function(tokens[0], llama, max_length=20, state=state)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DUCWX1BDkkw9"
      },
      "source": [
        "## Looping across sessions"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mWifpUjfkkw9"
      },
      "source": [
        "We mention earlier that the `session` context enables multi-tracing execution. But how can we optimize a process that would require running an intervention graph in a loop? If we create a simple `for` loop with a **Tracer context** inside, this will result in creating a new intervention graph at each iteration, which is not scalable.\n",
        "\n",
        "We solve this problem the `nnsight` way via the **Iterator context**: an intervention loop that iteratively executes and updates a single intervention graph.\n",
        "\n",
        "Use a `session` to define the **Iterator context** and pass in a sequence of items that you want to loop over at each iteration:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 126,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "to-hAyuKkkw9",
        "outputId": "ab21cb42-adba-4630-de5e-639cb6af3a35"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "2025-01-31 15:22:03,714 de3f072d-7ec7-4975-8ebe-4307ba8c7258 - RECEIVED: Your job has been received and is waiting approval.\n",
            "2025-01-31 15:22:03,869 de3f072d-7ec7-4975-8ebe-4307ba8c7258 - APPROVED: Your job was approved and is waiting to be run.\n",
            "2025-01-31 15:22:04,120 de3f072d-7ec7-4975-8ebe-4307ba8c7258 - RUNNING: Your job has started running.\n",
            "2025-01-31 15:22:05,086 de3f072d-7ec7-4975-8ebe-4307ba8c7258 - COMPLETED: Your job has been completed.\n",
            "Downloading result: 100%|██████████| 928/928 [00:00<00:00, 1.82MB/s]\n"
          ]
        }
      ],
      "source": [
        "with llama.session(remote=True) as session:\n",
        "\n",
        "  with session.iter([0, 1, 2]) as item:\n",
        "    # define intervention body here ...\n",
        "\n",
        "    with llama.trace(\"_\"):\n",
        "      # define interventions here ...\n",
        "      pass\n",
        "\n",
        "    with llama.trace(\"_\"):\n",
        "      # define interventions here ...\n",
        "      pass"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "L14XEU8jkkw9"
      },
      "source": [
        "The `Iterator` context extends all the `nnsight` graph-based functionalities, but also closely mimics the conventional `for` loop statement in Python, which allows it to support all kind of iterative operations with a use of `as item` syntax:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 127,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Rd-PcFJakkw9",
        "outputId": "a5b20d0c-5f0d-4065-a874-68bfe9ec696e"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "2025-01-31 15:22:07,235 570729dd-1d6d-4cd6-9876-9a096378119c - RECEIVED: Your job has been received and is waiting approval.\n",
            "2025-01-31 15:22:07,440 570729dd-1d6d-4cd6-9876-9a096378119c - APPROVED: Your job was approved and is waiting to be run.\n",
            "2025-01-31 15:22:07,674 570729dd-1d6d-4cd6-9876-9a096378119c - RUNNING: Your job has started running.\n",
            "2025-01-31 15:22:10,270 570729dd-1d6d-4cd6-9876-9a096378119c - COMPLETED: Your job has been completed.\n",
            "Downloading result: 100%|██████████| 928/928 [00:00<00:00, 1.47MB/s]"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "List:  [0, 1, 2]\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\n"
          ]
        }
      ],
      "source": [
        "with llama.session(remote=True) as session:\n",
        "\n",
        "  li = nnsight.list()\n",
        "  [li.append([num]) for num in range(0, 3)] # adding [0], [1], [2] to the list\n",
        "  li2 = nnsight.list().save()\n",
        "\n",
        "  # You can create nested Iterator contexts\n",
        "  with session.iter(li) as item:\n",
        "    with session.iter(item) as item_2:\n",
        "      li2.append(item_2)\n",
        "\n",
        "print(\"\\nList: \", li2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CZTVwF7Fkkw9"
      },
      "source": [
        "Notice how we used the `nnsight.list()` method to create a list of lists to loop over. This type of method is what we call an **NNsight Built-in**. It is a special type of methods that serve as a wrapper around `nnsight.apply()` to provide a more user-friendly interface for adding common datatypes to the Intervention Graph.\n",
        "\n",
        "<details>\n",
        "<summary>A full list of NNsight Built-ins</summary>\n",
        "\n",
        "\n",
        "`nnsight.bool()` creates a traceable Boolean\n",
        "\n",
        "`nnsight.bytes()` creates a traceable Bytes\n",
        "\n",
        "`nnsight.int()` creates a traceable Integer\n",
        "\n",
        "`nnsight.float()` creates a traceable Float\n",
        "\n",
        "`nnsight.str()` creates a traceable String\n",
        "\n",
        "`nnsight.comples()` creates a traceable Complex number\n",
        "\n",
        "`nnsight.bytearray()` creates a traceable Bytearray\n",
        "\n",
        "`nnsight.tuple()` creates a traceable Tuple\n",
        "\n",
        "`nnsight.list()` creates a traceable List\n",
        "\n",
        "`nnsight.set()` creates a traceable Set\n",
        "\n",
        "`nnsight.dict()` creates a traceable Dictionary\n",
        "\n",
        "</details>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pDF_eijqkkw9"
      },
      "source": [
        "We can also expose the `iterator` context object via a `return_context` flag. You can then use it to `exit` out of the Iteration loop early and log the intermediate outputs within the loop:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 128,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 365
        },
        "id": "IhnG03yckkw9",
        "outputId": "fe9ef3ca-e1b3-474d-b15f-f08d5dcd7e90"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "2025-01-31 15:22:12,730 cfa97488-c1e7-4ca7-b47d-dd856905edc9 - RECEIVED: Your job has been received and is waiting approval.\n",
            "2025-01-31 15:22:12,900 cfa97488-c1e7-4ca7-b47d-dd856905edc9 - APPROVED: Your job was approved and is waiting to be run.\n",
            "2025-01-31 15:22:13,237 cfa97488-c1e7-4ca7-b47d-dd856905edc9 - RUNNING: Your job has started running.\n",
            "2025-01-31 15:22:13,240 cfa97488-c1e7-4ca7-b47d-dd856905edc9 - LOG: 0\n",
            "2025-01-31 15:22:13,242 cfa97488-c1e7-4ca7-b47d-dd856905edc9 - LOG: 2\n",
            "2025-01-31 15:22:13,242 cfa97488-c1e7-4ca7-b47d-dd856905edc9 - LOG: 1\n",
            "2025-01-31 15:22:13,759 cfa97488-c1e7-4ca7-b47d-dd856905edc9 - COMPLETED: Your job has been completed.\n",
            "Downloading result: 100%|██████████| 928/928 [00:00<00:00, 4.74MB/s]\n"
          ]
        }
      ],
      "source": [
        "with llama.session(remote=True) as session:\n",
        "\n",
        "  # with session.iter([0, 1, 2, 3], return_context=True) as (item, iterator):\n",
        "  with session.iter([0, 1, 2, 3]) as item:\n",
        "\n",
        "      nnsight.log(item)\n",
        "\n",
        "      with nnsight.cond(item == 2):\n",
        "        nnsight.stop()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VvoZz5OGkkw9"
      },
      "source": [
        "The **Iterator** context is a niece piece of functionality that allows you to define a bunch of basic code operations that can now be \"traceable\" by `nnsight`.\n",
        "\n",
        "But in what kind of experimental scenario would someone need iterators?\n",
        "\n",
        "In the next section, we delve into a powerful use case of the `Iterator` context and see how it enables it!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1OcolbxAkkw9"
      },
      "source": [
        "## Training a LoRA"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jzyt5d6Dkkw9"
      },
      "source": [
        "Here is an example of a task that uses everything we have covered in the last section - remote execution, **Session** context and iterative interventions. Using session and iterator contexts, we're going apply a very simple fine-tuning approach called low-rank adaptation (LoRA).\n",
        "\n",
        "Let's try training a LoRA that, when applied, makes our model always predict \"Paris\" no matter what."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "id": "e9YlPQ-lkkw9"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import nnsight\n",
        "# from nnsight.envoy import Envoy # this moved in 0.4\n",
        "from nnsight import Envoy\n",
        "\n",
        "# We will define a LORA class.\n",
        "# The LORA class call method operations are simply traced like you would normally do in a .trace.\n",
        "class LORA(nn.Module):\n",
        "    def __init__(self, module: Envoy, dim: int, r: int) -> None:\n",
        "        \"\"\"Init.\n",
        "\n",
        "        Args:\n",
        "            module (Envoy): Which model Module we are adding the LORA to.\n",
        "            dim (int): Dimension of the layer we are adding to (This could potentially be auto populated if the user scanned first so we know the shape)\n",
        "            r (int): Inner dimension of the LORA\n",
        "        \"\"\"\n",
        "        super(LORA, self).__init__()\n",
        "        self.r = r\n",
        "        self.module = module\n",
        "        self.WA = torch.nn.Parameter(torch.randn(dim, self.r), requires_grad=True).save()\n",
        "        self.WB = torch.nn.Parameter(torch.zeros(self.r, dim), requires_grad=True).save()\n",
        "\n",
        "    # The Call method defines how to actually apply the LORA.\n",
        "    def __call__(self, alpha: float = 1.0):\n",
        "        \"\"\"Call.\n",
        "\n",
        "        Args:\n",
        "            alpha (float, optional): How much to apply the LORA. Can be altered after training for inference. Defaults to 1.0.\n",
        "        \"\"\"\n",
        "\n",
        "        # We apply WA to the first positional arg (the hidden states)\n",
        "        A_x = torch.matmul(self.module.input[0][0], self.WA)\n",
        "        BA_x = torch.matmul(A_x, self.WB)\n",
        "\n",
        "        # LORA is additive\n",
        "        h = BA_x + self.module.output\n",
        "\n",
        "        # Replace the output with our new one * alpha\n",
        "        # Could also have been self.module.output[:] = h * alpha, for in-place\n",
        "        self.module.output = h * alpha\n",
        "\n",
        "    def parameters(self):\n",
        "        # Some way to get all the parameters.\n",
        "        return [self.WA, self.WB]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_Ub6ApqAkkw9"
      },
      "source": [
        "Let's define all the variables to use in LoRA training."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "id": "yToy3qp1kkw9"
      },
      "outputs": [],
      "source": [
        "# We need the token id of the correct answer.\n",
        "answer = \" Paris\"\n",
        "answer_token = llama.tokenizer.encode(answer)[1]\n",
        "# Inner LORA dimension\n",
        "lora_dim = 4\n",
        "# Module to train LORA on\n",
        "module = llama.model.layers[-1].mlp"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NNriyaBikkw9"
      },
      "source": [
        "We can use the `.scan()` method to get the shape of the module without having to fully run the model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "id": "emv134jkkkw9"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "4096\n"
          ]
        }
      ],
      "source": [
        "with llama.scan(\" \"):\n",
        "    dim = module.output.shape[-1]\n",
        "\n",
        "print(dim)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "63aWV_ZTkkw9"
      },
      "source": [
        "It's time to run the LORA training loop! We using the **Session** and the **Iterator** contexts to achieve this."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {
        "id": "Q2SYDaawkkw9"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "2025-01-31 15:48:38,456 810d88f8-f1fc-43fd-b16d-77e11841eecd - RECEIVED: Your job has been received and is waiting approval.\n",
            "2025-01-31 15:48:38,634 810d88f8-f1fc-43fd-b16d-77e11841eecd - APPROVED: Your job was approved and is waiting to be run.\n",
            "2025-01-31 15:48:38,859 810d88f8-f1fc-43fd-b16d-77e11841eecd - RUNNING: Your job has started running.\n",
            "2025-01-31 15:48:39,402 810d88f8-f1fc-43fd-b16d-77e11841eecd - LOG: Parameter containing:\n",
            "tensor([[ 0.7539,  0.0732, -0.4023, -0.0615],\n",
            "        [-0.5312,  0.5391,  0.5195, -1.1328],\n",
            "        [-1.5781, -0.2480,  0.6953,  0.3535],\n",
            "        ...,\n",
            "        [ 0.1670,  0.5391, -0.5703, -0.6289],\n",
            "        [-0.0762, -1.3438,  0.8320,  1.2656],\n",
            "        [ 0.2500,  1.2188,  0.2891, -1.2578]], requires_grad=True)\n",
            "2025-01-31 15:48:39,410 810d88f8-f1fc-43fd-b16d-77e11841eecd - LOG: Parameter containing:\n",
            "tensor([[ 0.7305,  0.0713, -0.3906, -0.0598],\n",
            "        [-0.5156,  0.5234,  0.5039, -1.1016],\n",
            "        [-1.5312, -0.2402,  0.6758,  0.3438],\n",
            "        ...,\n",
            "        [ 0.1621,  0.5234, -0.5547, -0.6094],\n",
            "        [-0.0737, -1.3047,  0.8086,  1.2266],\n",
            "        [ 0.2422,  1.1797,  0.2812, -1.2188]], requires_grad=True)\n",
            "2025-01-31 15:48:39,448 810d88f8-f1fc-43fd-b16d-77e11841eecd - LOG: Parameter containing:\n",
            "tensor([[ 0.7070,  0.0688, -0.3789, -0.0586],\n",
            "        [-0.5000,  0.5078,  0.4883, -1.0703],\n",
            "        [-1.4844, -0.2334,  0.6562,  0.3340],\n",
            "        ...,\n",
            "        [ 0.1572,  0.5078, -0.5391, -0.5898],\n",
            "        [-0.0718, -1.2656,  0.7852,  1.1875],\n",
            "        [ 0.2344,  1.1406,  0.2734, -1.1797]], requires_grad=True)\n",
            "2025-01-31 15:48:39,485 810d88f8-f1fc-43fd-b16d-77e11841eecd - LOG: Parameter containing:\n",
            "tensor([[ 0.6836,  0.0645, -0.3633, -0.0598],\n",
            "        [-0.4883,  0.4902,  0.4746, -1.0391],\n",
            "        [-1.4375, -0.2285,  0.6406,  0.3203],\n",
            "        ...,\n",
            "        [ 0.1543,  0.4941, -0.5234, -0.5703],\n",
            "        [-0.0718, -1.2266,  0.7617,  1.1484],\n",
            "        [ 0.2227,  1.1094,  0.2695, -1.1484]], requires_grad=True)\n",
            "2025-01-31 15:48:39,521 810d88f8-f1fc-43fd-b16d-77e11841eecd - LOG: Parameter containing:\n",
            "tensor([[-0.9648, -1.5703,  0.6094, -1.6875],\n",
            "        [-2.0938, -1.1562,  1.3906, -2.6406],\n",
            "        [-3.0156, -1.8516,  1.5938, -1.3047],\n",
            "        ...,\n",
            "        [ 1.7578,  2.0781, -1.1406,  1.0625],\n",
            "        [-1.6875, -2.7969,  1.4688, -0.5117],\n",
            "        [-1.4219, -0.5508,  1.3516, -2.7344]], requires_grad=True)\n",
            "2025-01-31 15:48:39,558 810d88f8-f1fc-43fd-b16d-77e11841eecd - LOG: Parameter containing:\n",
            "tensor([[-2.7031e+00, -3.6719e-01, -9.6484e-01, -1.0938e+00],\n",
            "        [-3.7969e+00,  2.1606e-02, -1.9238e-01, -2.0156e+00],\n",
            "        [-4.6875e+00, -6.4453e-01,  3.3875e-03, -7.1484e-01],\n",
            "        ...,\n",
            "        [ 3.4688e+00,  8.6719e-01,  4.1602e-01,  4.8047e-01],\n",
            "        [-3.4062e+00, -1.5703e+00, -1.0645e-01,  5.7129e-02],\n",
            "        [-3.1406e+00,  6.1719e-01, -2.3633e-01, -2.0938e+00]],\n",
            "       requires_grad=True)\n",
            "2025-01-31 15:48:39,630 810d88f8-f1fc-43fd-b16d-77e11841eecd - LOG: Parameter containing:\n",
            "tensor([[-4.7812,  1.3125, -2.5625,  0.5117],\n",
            "        [-5.8125,  1.6953, -1.8203, -0.3789],\n",
            "        [-6.6875,  1.0469, -1.6250,  0.8789],\n",
            "        ...,\n",
            "        [ 5.5000, -0.8320,  2.0312, -1.1016],\n",
            "        [-5.4375,  0.1494, -1.7266,  1.6406],\n",
            "        [-5.1875,  2.2656, -1.8594, -0.4629]], requires_grad=True)\n",
            "2025-01-31 15:48:39,697 810d88f8-f1fc-43fd-b16d-77e11841eecd - LOG: Parameter containing:\n",
            "tensor([[-3.2812,  2.9062, -4.0938,  2.1094],\n",
            "        [-4.3125,  3.2812, -3.3750,  1.2422],\n",
            "        [-5.1562,  2.6406, -3.1719,  2.4531],\n",
            "        ...,\n",
            "        [ 4.0000, -2.4531,  3.5625, -2.6875],\n",
            "        [-3.9375,  1.7812, -3.2656,  3.2031],\n",
            "        [-3.6875,  3.8438, -3.4062,  1.1719]], requires_grad=True)\n",
            "2025-01-31 15:48:39,764 810d88f8-f1fc-43fd-b16d-77e11841eecd - LOG: Parameter containing:\n",
            "tensor([[-2.0000,  4.2500, -5.3750,  3.4688],\n",
            "        [-3.0156,  4.6250, -4.6875,  2.6406],\n",
            "        [-3.8125,  4.0000, -4.5000,  3.7969],\n",
            "        ...,\n",
            "        [ 2.6875, -3.8438,  4.8750, -4.0312],\n",
            "        [-2.6250,  3.1719, -4.5938,  4.5312],\n",
            "        [-2.3906,  5.1875, -4.7188,  2.5625]], requires_grad=True)\n",
            "2025-01-31 15:48:39,830 810d88f8-f1fc-43fd-b16d-77e11841eecd - LOG: Parameter containing:\n",
            "tensor([[-0.8867,  5.4062, -6.4688,  4.6250],\n",
            "        [-1.8750,  5.7812, -5.8125,  3.8438],\n",
            "        [-2.6406,  5.1562, -5.6562,  4.9375],\n",
            "        ...,\n",
            "        [ 1.5469, -5.0312,  6.0000, -5.1875],\n",
            "        [-1.4922,  4.3750, -5.7188,  5.6875],\n",
            "        [-1.2500,  6.3438, -5.8125,  3.7656]], requires_grad=True)\n",
            "2025-01-31 15:48:40,224 810d88f8-f1fc-43fd-b16d-77e11841eecd - COMPLETED: Your job has been completed.\n",
            "Downloading result: 100%|██████████| 67.1k/67.1k [00:00<00:00, 1.65MB/s]\n"
          ]
        }
      ],
      "source": [
        "from torch.utils.data import DataLoader\n",
        "\n",
        "# The LORA object itself isn't transmitted to the server. Only the forward / call method.\n",
        "# The parameters are created remotely and never sent only retrieved\n",
        "with llama.session(remote=True) as session:\n",
        "\n",
        "    # Create dataset of 100 pairs of a blank prompt and the \" Paris \" id\n",
        "    dataset = [[\"_\", answer_token]] * 100\n",
        "\n",
        "    # Create a dataloader from it.\n",
        "    dataloader = DataLoader(dataset, batch_size=10)\n",
        "\n",
        "    # Create our LORA on the last mlp\n",
        "    lora = LORA(module, dim, lora_dim)\n",
        "\n",
        "    # Create an optimizer. Use the parameters from LORA\n",
        "    optimizer = torch.optim.AdamW(lora.parameters(), lr=3)\n",
        "\n",
        "    # Iterate over dataloader using .iter.\n",
        "    with session.iter(dataloader) as batch:\n",
        "\n",
        "        prompt = batch[0]\n",
        "        correct_token = batch[1]\n",
        "\n",
        "        # Run .trace with prompt\n",
        "        with llama.trace(prompt) as tracer:\n",
        "\n",
        "            # Apply LORA to intervention graph just by calling it with .trace\n",
        "            lora()\n",
        "\n",
        "            # Get logits\n",
        "            logits = llama.lm_head.output\n",
        "\n",
        "            # Do cross entropy on last predicted token and correct_token\n",
        "            loss = torch.nn.functional.cross_entropy(logits[:, -1], batch[1])\n",
        "            # Call backward\n",
        "            loss.backward()\n",
        "\n",
        "        # Call methods on optimizer. Graphs that arent from .trace (so in this case session and iterator both have their own graph) are executed sequentially.\n",
        "        # The Graph of Iterator here will be:\n",
        "        # 1.) Index batch at 0 for prompt\n",
        "        # 2.) Index batch at 1 for correct_token\n",
        "        # 3.) Execute the .trace using the prompt\n",
        "        # 4.) Call .step() on optimizer\n",
        "        optimizer.step()\n",
        "        # 5.) Call .zero_grad() in optimizer\n",
        "        optimizer.zero_grad()\n",
        "        # 6.) Print out the lora WA weights to show they are indeed changing\n",
        "        nnsight.log(lora.WA)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2FtFv2Nykkw9"
      },
      "source": [
        "Now `WA` and `WB` are optimized! So we generate with the LoRA just by calling `lora()` in the `.generate` and save the output to then de-tokenize it."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {
        "id": "K5_F4ktlkkw9"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "2025-01-31 15:48:45,111 c9fcc827-5c78-4796-87cf-0f246c75a359 - RECEIVED: Your job has been received and is waiting approval.\n",
            "2025-01-31 15:48:45,281 c9fcc827-5c78-4796-87cf-0f246c75a359 - APPROVED: Your job was approved and is waiting to be run.\n",
            "2025-01-31 15:48:45,535 c9fcc827-5c78-4796-87cf-0f246c75a359 - RUNNING: Your job has started running.\n",
            "2025-01-31 15:48:46,013 c9fcc827-5c78-4796-87cf-0f246c75a359 - COMPLETED: Your job has been completed.\n",
            "Downloading result: 100%|██████████| 1.24k/1.24k [00:00<00:00, 4.30MB/s]\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "['<|begin_of_text|>Hello Paris']\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "2025-01-31 15:48:46,446 9f088820-f00d-492d-9861-c0f0a7c4d784 - RECEIVED: Your job has been received and is waiting approval.\n",
            "2025-01-31 15:48:46,650 9f088820-f00d-492d-9861-c0f0a7c4d784 - APPROVED: Your job was approved and is waiting to be run.\n",
            "2025-01-31 15:48:46,862 9f088820-f00d-492d-9861-c0f0a7c4d784 - RUNNING: Your job has started running.\n",
            "2025-01-31 15:48:47,340 9f088820-f00d-492d-9861-c0f0a7c4d784 - COMPLETED: Your job has been completed.\n",
            "Downloading result: 100%|██████████| 1.24k/1.24k [00:00<00:00, 4.61MB/s]"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "['<|begin_of_text|>Hello,']\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\n"
          ]
        }
      ],
      "source": [
        "# With lora. Should produce \"Hello Paris\"\n",
        "with llama.generate(\"Hello\", remote=True) as generator:\n",
        "\n",
        "    lora()\n",
        "\n",
        "    out = llama.generator.output.save()\n",
        "\n",
        "print(llama.tokenizer.batch_decode(out.value))\n",
        "\n",
        "# Then without. Should produce \"Hello,\"\n",
        "with llama.generate(\"Hello\", remote=True) as generator:\n",
        "\n",
        "    out = llama.generator.output.save()\n",
        "\n",
        "print(llama.tokenizer.batch_decode(out.value))\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IZfiTNGgu-HI"
      },
      "source": [
        "# Next Steps\n",
        "Check out [nnsight.net/tutorials](https://nnsight.net/tutorials) for more walkthroughs implementating classic interpretability techniques using `nnsight`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "P3zRm-7VRRov"
      },
      "source": [
        "# Getting Involved!\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BnCc9xgjvxEP"
      },
      "source": [
        "Note that both `nnsight` and `NDIF` are in active development, so changes may be made and errors may arise during use. If you’re interested in following updates to `nnsight`, contributing, giving feedback, or finding collaborators, please join the [NDIF discord](https://discord.gg/6uFJmCSwW7). We’d love to hear about your work using nnsight!\n",
        "\n",
        "You can also follow us on [LinkedIn](https://www.linkedin.com/company/national-deep-inference-fabric/), Bluesky: [@ndif-team.bsky.social](https://bsky.app/profile/ndif-team.bsky.social), and X: [@ndif_team](https://x.com/ndif_team).\n",
        "\n",
        "💟\n"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "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.13"
    },
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "0383a2798a0c4fcfbd1f068817770fd6": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "FloatProgressModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_cef2fe4eb51146afabe820140a3b0df3",
            "max": 456318,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_52281c8594ad465397e6079385c558a1",
            "value": 456318
          }
        },
        "08fb0620499c44eb9d160a0f21263451": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "0c46bce8cc3e4393898b025ba6b4e543": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "14008810744549208f147a132b419d20": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HBoxModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_cb4be83440e44e1c9e664ac1874b830d",
              "IPY_MODEL_f54c77317ca04861a958c8f982b25653",
              "IPY_MODEL_db24a855fb5846e28c98dcfe9afff259"
            ],
            "layout": "IPY_MODEL_647c861fcaa44148a90da061b1d5c99c"
          }
        },
        "184b5c75820b4c42ae861669ac1a47e5": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "19c0e121e52040b18393aee887730644": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "26acae3d321743d2ad8efbf8727dbd05": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "FloatProgressModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_a4b5c225f5ac4719ade8e6aafcff78d2",
            "max": 548105171,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_ee27d18a555c46e78c92b282bfa6dc9d",
            "value": 548105171
          }
        },
        "26dcd7b75feb41a6b86d4260ecd64d06": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HBoxModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_39c9f1effe664e899125748e97fad198",
              "IPY_MODEL_383b893224684aee916e455c0d3d8c06",
              "IPY_MODEL_af24a12bd05b44148b5d6f91ac2445a0"
            ],
            "layout": "IPY_MODEL_d00750501f2548d0b1c69f4753aec877"
          }
        },
        "2702f30b5d3b4b0b8184fec58863750b": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "ProgressStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "287a8fe3ddf54d36b7603171ec7a2688": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "ProgressStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "2a0f2484053f41298aceacea5c1cc195": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "FloatProgressModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_08fb0620499c44eb9d160a0f21263451",
            "max": 124,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_287a8fe3ddf54d36b7603171ec7a2688",
            "value": 124
          }
        },
        "2fedb8b529c34049932f07b577b515d4": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HBoxModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_bd1a11654c2d47a09610e4cc23b8bc32",
              "IPY_MODEL_26acae3d321743d2ad8efbf8727dbd05",
              "IPY_MODEL_f42e44122a2d496cb7b2ea85c2858a71"
            ],
            "layout": "IPY_MODEL_c1586aeef6814447ad46c1f8ddaae625"
          }
        },
        "346ecc4734014cb4931bea53d5a32bf1": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "383b893224684aee916e455c0d3d8c06": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "FloatProgressModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_68b205f60f9d4599a7d1f78583d0b93c",
            "max": 665,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_69e8c149f676457d95029f21e9a2956a",
            "value": 665
          }
        },
        "385291e96814453b8f805cbe8a5e594d": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_e5302eb0481c4e529fc68353eed4b79f",
            "placeholder": "​",
            "style": "IPY_MODEL_0c46bce8cc3e4393898b025ba6b4e543",
            "value": " 456k/456k [00:00&lt;00:00, 3.39MB/s]"
          }
        },
        "39c9f1effe664e899125748e97fad198": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_98d28dddf75a45a2b1e05131b0b865a1",
            "placeholder": "​",
            "style": "IPY_MODEL_6e81d6f31f2446c999df5511e8afd8d0",
            "value": "config.json: 100%"
          }
        },
        "405308ec76074c7ea319e1c1358baf39": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "41fc21ee95c14d29943a56f2b63f4989": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_42fd0ea3be3c4b3f98ca7818efee3669",
            "placeholder": "​",
            "style": "IPY_MODEL_184b5c75820b4c42ae861669ac1a47e5",
            "value": " 26.0/26.0 [00:00&lt;00:00, 2.05kB/s]"
          }
        },
        "42fd0ea3be3c4b3f98ca7818efee3669": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "4b798a784ae748b1a0aef025ead976fc": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "504038c4efd3439daa958333f108b5fb": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "52281c8594ad465397e6079385c558a1": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "ProgressStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "5250cf68037a43e389493733c5b9d4ac": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "55c0bb35417e454d842e12bd2dfe62df": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "57aebe107749473390af2532806920ce": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_55c0bb35417e454d842e12bd2dfe62df",
            "placeholder": "​",
            "style": "IPY_MODEL_d5821d69e1fa4673bb39ab70a7f60466",
            "value": " 124/124 [00:00&lt;00:00, 8.42kB/s]"
          }
        },
        "5d877979f2db4dbfa7581d2fe4381cf8": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "5e091b15cc9748469676b165533f0c44": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "647c861fcaa44148a90da061b1d5c99c": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "64b713188f054535a3063c079bceead0": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "FloatProgressModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_7afff81060f0481aa95a56daa570e1b8",
            "max": 1042301,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_a038f9037ed540f2ba43f3f85920b17e",
            "value": 1042301
          }
        },
        "68b205f60f9d4599a7d1f78583d0b93c": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "69e8c149f676457d95029f21e9a2956a": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "ProgressStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "6e81d6f31f2446c999df5511e8afd8d0": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "7479281e090d4eeba4c21a9c15b02165": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HBoxModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_f4b25b3aac414499b653404ea3c066da",
              "IPY_MODEL_0383a2798a0c4fcfbd1f068817770fd6",
              "IPY_MODEL_385291e96814453b8f805cbe8a5e594d"
            ],
            "layout": "IPY_MODEL_f7a41e267f79401681bc4693c52e55fa"
          }
        },
        "76fe13ab697c46cc9fcbe46112fb45b1": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "ProgressStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "79a221c9228d459aa17402c9fef1ad1d": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HBoxModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_9c13a3d768eb47c8b26a7c6bfa8ec957",
              "IPY_MODEL_e97c052113f54dddbb619ed228f6b4ef",
              "IPY_MODEL_41fc21ee95c14d29943a56f2b63f4989"
            ],
            "layout": "IPY_MODEL_5250cf68037a43e389493733c5b9d4ac"
          }
        },
        "7afff81060f0481aa95a56daa570e1b8": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "7f4f63d94f8942fcb3e4b286d5848f0c": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "7fdb623e5b2d4eb1b6cee27da2314062": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "87a5c035bc6345aea6441b3db4b94321": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "95832c5076ab4c80810d3c443e061156": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "98d28dddf75a45a2b1e05131b0b865a1": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "9c13a3d768eb47c8b26a7c6bfa8ec957": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_e24033c5ba7b4e4fa4f4375d85b03bad",
            "placeholder": "​",
            "style": "IPY_MODEL_4b798a784ae748b1a0aef025ead976fc",
            "value": "tokenizer_config.json: 100%"
          }
        },
        "9cea52b7835b40628fc0a65199c5e82f": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_ec4f48d5a9ef4afa8fa44078dfb236c6",
            "placeholder": "​",
            "style": "IPY_MODEL_405308ec76074c7ea319e1c1358baf39",
            "value": " 1.04M/1.04M [00:00&lt;00:00, 7.50MB/s]"
          }
        },
        "a038f9037ed540f2ba43f3f85920b17e": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "ProgressStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "a15895ad1cb541ca89fa9def63779104": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HBoxModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_ee60494714c2485f802298dce918edfe",
              "IPY_MODEL_2a0f2484053f41298aceacea5c1cc195",
              "IPY_MODEL_57aebe107749473390af2532806920ce"
            ],
            "layout": "IPY_MODEL_ff44e896c55048c5a7d92a56aa9eedfc"
          }
        },
        "a1c25910a31a4b5a9c48f1a7fd8047c5": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "a310703500c5481aa37b698c659e1ca6": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "a3154c60744d4a1a87b1c86b89a142b9": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "a4b5c225f5ac4719ade8e6aafcff78d2": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "af24a12bd05b44148b5d6f91ac2445a0": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_ba90058201ee48dcae9411ded457e85f",
            "placeholder": "​",
            "style": "IPY_MODEL_346ecc4734014cb4931bea53d5a32bf1",
            "value": " 665/665 [00:00&lt;00:00, 36.5kB/s]"
          }
        },
        "ba90058201ee48dcae9411ded457e85f": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "bd1a11654c2d47a09610e4cc23b8bc32": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_5e091b15cc9748469676b165533f0c44",
            "placeholder": "​",
            "style": "IPY_MODEL_f143977aeaf24798b7744681069762d4",
            "value": "model.safetensors: 100%"
          }
        },
        "c1586aeef6814447ad46c1f8ddaae625": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "caff23247fb9432db8d5860a8132e1b3": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "cb4be83440e44e1c9e664ac1874b830d": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_e4e0b1de76ff40a88b71e0b5a9f7bc39",
            "placeholder": "​",
            "style": "IPY_MODEL_504038c4efd3439daa958333f108b5fb",
            "value": "tokenizer.json: 100%"
          }
        },
        "cef2fe4eb51146afabe820140a3b0df3": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "d00750501f2548d0b1c69f4753aec877": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "d5821d69e1fa4673bb39ab70a7f60466": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "d8feaceb6d854cb2b2ff19a0309d9a7b": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HBoxModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_fd84ffef58604ce3a16ef8f9868072f3",
              "IPY_MODEL_64b713188f054535a3063c079bceead0",
              "IPY_MODEL_9cea52b7835b40628fc0a65199c5e82f"
            ],
            "layout": "IPY_MODEL_87a5c035bc6345aea6441b3db4b94321"
          }
        },
        "db24a855fb5846e28c98dcfe9afff259": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_a1c25910a31a4b5a9c48f1a7fd8047c5",
            "placeholder": "​",
            "style": "IPY_MODEL_5d877979f2db4dbfa7581d2fe4381cf8",
            "value": " 1.36M/1.36M [00:00&lt;00:00, 6.83MB/s]"
          }
        },
        "de023d4526524f48a34883a287aa8a4f": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "ded8c188917b424dbf264ce494fe0532": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "e24033c5ba7b4e4fa4f4375d85b03bad": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "e4e0b1de76ff40a88b71e0b5a9f7bc39": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "e5302eb0481c4e529fc68353eed4b79f": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "e97c052113f54dddbb619ed228f6b4ef": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "FloatProgressModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_19c0e121e52040b18393aee887730644",
            "max": 26,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_2702f30b5d3b4b0b8184fec58863750b",
            "value": 26
          }
        },
        "ec4f48d5a9ef4afa8fa44078dfb236c6": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "ee27d18a555c46e78c92b282bfa6dc9d": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "ProgressStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "ee60494714c2485f802298dce918edfe": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_7fdb623e5b2d4eb1b6cee27da2314062",
            "placeholder": "​",
            "style": "IPY_MODEL_95832c5076ab4c80810d3c443e061156",
            "value": "generation_config.json: 100%"
          }
        },
        "f143977aeaf24798b7744681069762d4": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "f42e44122a2d496cb7b2ea85c2858a71": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_a3154c60744d4a1a87b1c86b89a142b9",
            "placeholder": "​",
            "style": "IPY_MODEL_de023d4526524f48a34883a287aa8a4f",
            "value": " 548M/548M [00:02&lt;00:00, 253MB/s]"
          }
        },
        "f47609606aa14f6fa2407aa26edae8e3": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "f4b25b3aac414499b653404ea3c066da": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_caff23247fb9432db8d5860a8132e1b3",
            "placeholder": "​",
            "style": "IPY_MODEL_f47609606aa14f6fa2407aa26edae8e3",
            "value": "merges.txt: 100%"
          }
        },
        "f54c77317ca04861a958c8f982b25653": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "FloatProgressModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_ded8c188917b424dbf264ce494fe0532",
            "max": 1355256,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_76fe13ab697c46cc9fcbe46112fb45b1",
            "value": 1355256
          }
        },
        "f7a41e267f79401681bc4693c52e55fa": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "fd84ffef58604ce3a16ef8f9868072f3": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "HTMLModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_7f4f63d94f8942fcb3e4b286d5848f0c",
            "placeholder": "​",
            "style": "IPY_MODEL_a310703500c5481aa37b698c659e1ca6",
            "value": "vocab.json: 100%"
          }
        },
        "ff44e896c55048c5a7d92a56aa9eedfc": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        }
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}