GPT2-small-Token-Visualizer.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "03ba0f87-6665-4db3-8484-5ed061f39151",
   "metadata": {},
   "source": [
    "This is a tutorial (code copied and updated from https://github.com/neelnanda-io/EasyTransformer/blob/main/Hacky-Interactive-Lexoscope.ipynb) that explores how to interactively visualize a Transformer language model. We use @neelnanda-io's `Easy-Transformer` library."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "6704c22e-edfb-4eeb-9add-36599db978de",
   "metadata": {},
   "outputs": [],
   "source": [
    "%reload_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "5cb18bb4-0f4e-4182-a908-4c1feb6d1699",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "10303ce5-1ff2-418a-8265-45025b98576e",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: gradio in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (3.9)\n",
      "Requirement already satisfied: websockets in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from gradio) (10.4)\n",
      "Requirement already satisfied: fsspec in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from gradio) (2022.10.0)\n",
      "Requirement already satisfied: pydantic in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from gradio) (1.8.2)\n",
      "Requirement already satisfied: aiohttp in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from gradio) (3.8.3)\n",
      "Requirement already satisfied: matplotlib in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from gradio) (3.5.2)\n",
      "Requirement already satisfied: orjson in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from gradio) (3.8.1)\n",
      "Requirement already satisfied: pycryptodome in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from gradio) (3.15.0)\n",
      "Requirement already satisfied: python-multipart in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from gradio) (0.0.5)\n",
      "Requirement already satisfied: pyyaml in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from gradio) (6.0)\n",
      "Requirement already satisfied: pillow in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from gradio) (9.2.0)\n",
      "Requirement already satisfied: h11<0.13,>=0.11 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from gradio) (0.12.0)\n",
      "Requirement already satisfied: jinja2 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from gradio) (3.1.2)\n",
      "Requirement already satisfied: requests in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from gradio) (2.25.1)\n",
      "Requirement already satisfied: markdown-it-py[linkify,plugins] in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from gradio) (2.1.0)\n",
      "Requirement already satisfied: pandas in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from gradio) (1.4.2)\n",
      "Requirement already satisfied: pydub in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from gradio) (0.25.1)\n",
      "Requirement already satisfied: uvicorn in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from gradio) (0.19.0)\n",
      "Requirement already satisfied: httpx in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from gradio) (0.23.0)\n",
      "Requirement already satisfied: numpy in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from gradio) (1.23.3)\n",
      "Requirement already satisfied: paramiko in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from gradio) (2.12.0)\n",
      "Requirement already satisfied: fastapi in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from gradio) (0.86.0)\n",
      "Requirement already satisfied: ffmpy in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from gradio) (0.3.0)\n",
      "Requirement already satisfied: async-timeout<5.0,>=4.0.0a3 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from aiohttp->gradio) (4.0.2)\n",
      "Requirement already satisfied: multidict<7.0,>=4.5 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from aiohttp->gradio) (6.0.2)\n",
      "Requirement already satisfied: yarl<2.0,>=1.0 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from aiohttp->gradio) (1.8.1)\n",
      "Requirement already satisfied: attrs>=17.3.0 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from aiohttp->gradio) (22.1.0)\n",
      "Requirement already satisfied: aiosignal>=1.1.2 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from aiohttp->gradio) (1.2.0)\n",
      "Requirement already satisfied: charset-normalizer<3.0,>=2.0 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from aiohttp->gradio) (2.1.1)\n",
      "Requirement already satisfied: frozenlist>=1.1.1 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from aiohttp->gradio) (1.3.1)\n",
      "Requirement already satisfied: starlette==0.20.4 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from fastapi->gradio) (0.20.4)\n",
      "Requirement already satisfied: anyio<5,>=3.4.0 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from starlette==0.20.4->fastapi->gradio) (3.6.2)\n",
      "Requirement already satisfied: typing-extensions>=3.7.4.3 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from pydantic->gradio) (4.4.0)\n",
      "Requirement already satisfied: httpcore<0.16.0,>=0.15.0 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from httpx->gradio) (0.15.0)\n",
      "Requirement already satisfied: certifi in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from httpx->gradio) (2022.9.24)\n",
      "Requirement already satisfied: sniffio in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from httpx->gradio) (1.3.0)\n",
      "Requirement already satisfied: rfc3986[idna2008]<2,>=1.3 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from httpx->gradio) (1.5.0)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from jinja2->gradio) (2.1.1)\n",
      "Requirement already satisfied: mdurl~=0.1 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from markdown-it-py[linkify,plugins]->gradio) (0.1.2)\n",
      "Requirement already satisfied: linkify-it-py~=1.0 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from markdown-it-py[linkify,plugins]->gradio) (1.0.3)\n",
      "Requirement already satisfied: mdit-py-plugins in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from markdown-it-py[linkify,plugins]->gradio) (0.3.1)\n",
      "Requirement already satisfied: pyparsing>=2.2.1 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from matplotlib->gradio) (3.0.9)\n",
      "Requirement already satisfied: fonttools>=4.22.0 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from matplotlib->gradio) (4.33.3)\n",
      "Requirement already satisfied: python-dateutil>=2.7 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from matplotlib->gradio) (2.8.2)\n",
      "Requirement already satisfied: kiwisolver>=1.0.1 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from matplotlib->gradio) (1.4.2)\n",
      "Requirement already satisfied: cycler>=0.10 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from matplotlib->gradio) (0.11.0)\n",
      "Requirement already satisfied: packaging>=20.0 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from matplotlib->gradio) (21.3)\n",
      "Requirement already satisfied: pytz>=2020.1 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from pandas->gradio) (2022.5)\n",
      "Requirement already satisfied: cryptography>=2.5 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from paramiko->gradio) (37.0.2)\n",
      "Requirement already satisfied: six in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from paramiko->gradio) (1.16.0)\n",
      "Requirement already satisfied: pynacl>=1.0.1 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from paramiko->gradio) (1.5.0)\n",
      "Requirement already satisfied: bcrypt>=3.1.3 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from paramiko->gradio) (4.0.1)\n",
      "Requirement already satisfied: idna<3,>=2.5 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from requests->gradio) (2.10)\n",
      "Requirement already satisfied: chardet<5,>=3.0.2 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from requests->gradio) (4.0.0)\n",
      "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from requests->gradio) (1.26.11)\n",
      "Requirement already satisfied: click>=7.0 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from uvicorn->gradio) (8.1.3)\n",
      "Requirement already satisfied: cffi>=1.12 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from cryptography>=2.5->paramiko->gradio) (1.15.1)\n",
      "Requirement already satisfied: uc-micro-py in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from linkify-it-py~=1.0->markdown-it-py[linkify,plugins]->gradio) (1.0.1)\n",
      "Requirement already satisfied: pycparser in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from cffi>=1.12->cryptography>=2.5->paramiko->gradio) (2.21)\n",
      "Collecting git+https://github.com/neelnanda-io/Easy-Transformer.git\n",
      "  Cloning https://github.com/neelnanda-io/Easy-Transformer.git to /tmp/pip-req-build-yjykeksm\n",
      "  Running command git clone --filter=blob:none --quiet https://github.com/neelnanda-io/Easy-Transformer.git /tmp/pip-req-build-yjykeksm\n",
      "  Resolved https://github.com/neelnanda-io/Easy-Transformer.git to commit ee0fc75942365dec64d476aa72219a0fca628059\n",
      "  Preparing metadata (setup.py) ... \u001b[?25ldone\n",
      "\u001b[?25hCollecting einops\n",
      "  Downloading einops-0.5.0-py3-none-any.whl (36 kB)\n",
      "Requirement already satisfied: numpy in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from easy-transformer==0.1.0) (1.23.3)\n",
      "Requirement already satisfied: torch in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from easy-transformer==0.1.0) (1.12.1)\n",
      "Requirement already satisfied: datasets in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from easy-transformer==0.1.0) (2.6.1)\n",
      "Requirement already satisfied: transformers in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from easy-transformer==0.1.0) (4.23.1)\n",
      "Requirement already satisfied: tqdm in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from easy-transformer==0.1.0) (4.64.1)\n",
      "Requirement already satisfied: pandas in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from easy-transformer==0.1.0) (1.4.2)\n",
      "Requirement already satisfied: wandb in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from easy-transformer==0.1.0) (0.13.4)\n",
      "Collecting fancy_einsum\n",
      "  Downloading fancy_einsum-0.0.3-py3-none-any.whl (6.2 kB)\n",
      "Collecting torchtyping\n",
      "  Downloading torchtyping-0.1.4-py3-none-any.whl (17 kB)\n",
      "Requirement already satisfied: responses<0.19 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from datasets->easy-transformer==0.1.0) (0.18.0)\n",
      "Requirement already satisfied: pyarrow>=6.0.0 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from datasets->easy-transformer==0.1.0) (9.0.0)\n",
      "Requirement already satisfied: packaging in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from datasets->easy-transformer==0.1.0) (21.3)\n",
      "Requirement already satisfied: fsspec[http]>=2021.11.1 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from datasets->easy-transformer==0.1.0) (2022.10.0)\n",
      "Requirement already satisfied: requests>=2.19.0 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from datasets->easy-transformer==0.1.0) (2.25.1)\n",
      "Requirement already satisfied: multiprocess in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from datasets->easy-transformer==0.1.0) (0.70.13)\n",
      "Requirement already satisfied: pyyaml>=5.1 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from datasets->easy-transformer==0.1.0) (6.0)\n",
      "Requirement already satisfied: xxhash in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from datasets->easy-transformer==0.1.0) (3.1.0)\n",
      "Requirement already satisfied: aiohttp in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from datasets->easy-transformer==0.1.0) (3.8.3)\n",
      "Requirement already satisfied: dill<0.3.6 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from datasets->easy-transformer==0.1.0) (0.3.5.1)\n",
      "Requirement already satisfied: huggingface-hub<1.0.0,>=0.2.0 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from datasets->easy-transformer==0.1.0) (0.10.1)\n",
      "Requirement already satisfied: python-dateutil>=2.8.1 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from pandas->easy-transformer==0.1.0) (2.8.2)\n",
      "Requirement already satisfied: pytz>=2020.1 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from pandas->easy-transformer==0.1.0) (2022.5)\n",
      "Requirement already satisfied: typing_extensions in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from torch->easy-transformer==0.1.0) (4.4.0)\n",
      "Collecting typeguard>=2.11.1\n",
      "  Downloading typeguard-2.13.3-py3-none-any.whl (17 kB)\n",
      "Requirement already satisfied: regex!=2019.12.17 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from transformers->easy-transformer==0.1.0) (2022.9.13)\n",
      "Requirement already satisfied: tokenizers!=0.11.3,<0.14,>=0.11.1 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from transformers->easy-transformer==0.1.0) (0.12.1)\n",
      "Requirement already satisfied: filelock in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from transformers->easy-transformer==0.1.0) (3.8.0)\n",
      "Requirement already satisfied: six>=1.13.0 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from wandb->easy-transformer==0.1.0) (1.16.0)\n",
      "Requirement already satisfied: GitPython>=1.0.0 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from wandb->easy-transformer==0.1.0) (3.1.29)\n",
      "Requirement already satisfied: setproctitle in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from wandb->easy-transformer==0.1.0) (1.3.2)\n",
      "Requirement already satisfied: docker-pycreds>=0.4.0 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from wandb->easy-transformer==0.1.0) (0.4.0)\n",
      "Requirement already satisfied: psutil>=5.0.0 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from wandb->easy-transformer==0.1.0) (5.9.3)\n",
      "Requirement already satisfied: shortuuid>=0.5.0 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from wandb->easy-transformer==0.1.0) (1.0.9)\n",
      "Requirement already satisfied: pathtools in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from wandb->easy-transformer==0.1.0) (0.1.2)\n",
      "Requirement already satisfied: setuptools in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from wandb->easy-transformer==0.1.0) (63.4.1)\n",
      "Requirement already satisfied: Click!=8.0.0,>=7.0 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from wandb->easy-transformer==0.1.0) (8.1.3)\n",
      "Requirement already satisfied: sentry-sdk>=1.0.0 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from wandb->easy-transformer==0.1.0) (1.10.1)\n",
      "Requirement already satisfied: promise<3,>=2.0 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from wandb->easy-transformer==0.1.0) (2.3)\n",
      "Requirement already satisfied: protobuf!=4.0.*,!=4.21.0,<5,>=3.12.0 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from wandb->easy-transformer==0.1.0) (3.19.6)\n",
      "Requirement already satisfied: yarl<2.0,>=1.0 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from aiohttp->datasets->easy-transformer==0.1.0) (1.8.1)\n",
      "Requirement already satisfied: attrs>=17.3.0 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from aiohttp->datasets->easy-transformer==0.1.0) (22.1.0)\n",
      "Requirement already satisfied: frozenlist>=1.1.1 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from aiohttp->datasets->easy-transformer==0.1.0) (1.3.1)\n",
      "Requirement already satisfied: aiosignal>=1.1.2 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from aiohttp->datasets->easy-transformer==0.1.0) (1.2.0)\n",
      "Requirement already satisfied: multidict<7.0,>=4.5 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from aiohttp->datasets->easy-transformer==0.1.0) (6.0.2)\n",
      "Requirement already satisfied: charset-normalizer<3.0,>=2.0 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from aiohttp->datasets->easy-transformer==0.1.0) (2.1.1)\n",
      "Requirement already satisfied: async-timeout<5.0,>=4.0.0a3 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from aiohttp->datasets->easy-transformer==0.1.0) (4.0.2)\n",
      "Requirement already satisfied: gitdb<5,>=4.0.1 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from GitPython>=1.0.0->wandb->easy-transformer==0.1.0) (4.0.9)\n",
      "Requirement already satisfied: pyparsing!=3.0.5,>=2.0.2 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from packaging->datasets->easy-transformer==0.1.0) (3.0.9)\n",
      "Requirement already satisfied: idna<3,>=2.5 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from requests>=2.19.0->datasets->easy-transformer==0.1.0) (2.10)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from requests>=2.19.0->datasets->easy-transformer==0.1.0) (2022.9.24)\n",
      "Requirement already satisfied: chardet<5,>=3.0.2 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from requests>=2.19.0->datasets->easy-transformer==0.1.0) (4.0.0)\n",
      "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from requests>=2.19.0->datasets->easy-transformer==0.1.0) (1.26.11)\n",
      "Requirement already satisfied: smmap<6,>=3.0.1 in /home/muhia/miniconda3/envs/fastai/lib/python3.10/site-packages (from gitdb<5,>=4.0.1->GitPython>=1.0.0->wandb->easy-transformer==0.1.0) (5.0.0)\n",
      "Building wheels for collected packages: easy-transformer\n",
      "  Building wheel for easy-transformer (setup.py) ... \u001b[?25ldone\n",
      "\u001b[?25h  Created wheel for easy-transformer: filename=easy_transformer-0.1.0-py3-none-any.whl size=64508 sha256=9bd6e773390d61a4fcd7bcf113dc5779f75ffb82efd5583aa47516a4a974d005\n",
      "  Stored in directory: /tmp/pip-ephem-wheel-cache-pqoy_c92/wheels/4e/73/bc/ef945055d6b29743c2a3cf4f890eb6c28423dccb1703f7ab06\n",
      "Successfully built easy-transformer\n",
      "Installing collected packages: typeguard, fancy_einsum, einops, torchtyping, easy-transformer\n",
      "Successfully installed easy-transformer-0.1.0 einops-0.5.0 fancy_einsum-0.0.3 torchtyping-0.1.4 typeguard-2.13.3\n"
     ]
    }
   ],
   "source": [
    "!pip install gradio plotly\n",
    "!pip install \"git+https://github.com/neelnanda-io/Easy-Transformer.git\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "74871411-3d11-4edc-b3e1-5152bce97303",
   "metadata": {},
   "outputs": [],
   "source": [
    "import gradio as gr\n",
    "from easy_transformer import EasyTransformer\n",
    "from easy_transformer.utils import to_numpy\n",
    "from IPython.display import HTML"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e8345214-7740-44fc-a239-7f639a2b9094",
   "metadata": {},
   "source": [
    "### GPT2-small (117M)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31227246-07e5-4a2d-b79c-8c2d673a9325",
   "metadata": {},
   "source": [
    "We use `gpt2-small` in this example. This is a language model with 117M parameters. It embeds some input tokens, contextualizes them, then predicts the next word, all while computing loss against a known target. It was trained on the WebText dataset, which emphasizes quality by filtering human-curated documents that have a karma score > 3 on Reddit, scraping them and preprocessing them using `Dragnet` and `https://github.com/codelucas/newspaper` content extractors."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c640469-3551-4f4e-a8a1-d3a506b9d997",
   "metadata": {},
   "source": [
    "### Extracting model activations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "57223ce0-f86a-4904-8780-235df108b918",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading model: gpt2-small\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1b064ccaf55047d387439feda81a5bc5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading:   0%|          | 0.00/665 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cd267794a05748d5a47012ef7d674597",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading:   0%|          | 0.00/548M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c0230b4e44f442eaa1ef1620ac991cc0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading:   0%|          | 0.00/1.04M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "67214433acef4a56bd3708b200e13651",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading:   0%|          | 0.00/456k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "da8c4ff4620a46a69bc0666ce3d71790",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading:   0%|          | 0.00/1.36M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using pad_token, but it is not set yet.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Moving model to device:  cuda\n",
      "Finished loading pretrained model gpt2-small into EasyTransformer!\n"
     ]
    }
   ],
   "source": [
    "model_name = \"gpt2-small\"\n",
    "model = EasyTransformer.from_pretrained(model_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 125,
   "id": "fedbd5f9-7bf8-47dd-811a-e24bf8576158",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_activations(inp, layer, neuron_idx):\n",
    "    \"\"\"Use a hook to extract a particular activation (neuron) from a layer, this could be a callback in a different library.\n",
    "    \"\"\"\n",
    "    cache = {}\n",
    "    \n",
    "    def caching_hook(act, hook):\n",
    "        cache[\"activation\"] = act[0, :, neuron_idx]\n",
    "    \n",
    "    model.run_with_hooks(\n",
    "        inp, fwd_hooks=[(f\"blocks.{layer}.mlp.hook_post\", caching_hook)]\n",
    "    )\n",
    "    return to_numpy(cache[\"activation\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "id": "5728d89a-e7b6-4b66-8cc7-ea6770271ad0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Testing the outputs of the function above"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 90,
   "id": "41e735e1-833c-45af-b042-8c7c1d47ca18",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['<|endoftext|>', 'The', ' following', ' is', ' a', ' list', ' of', ' powers', ' of', ' 10', ':', ' 1', ',', ' 10', ',', ' 100', ',', ' 1000', ',', ' 10000', ',', ' 100', '000', ',', ' 100', '0000', ',', ' 100', '00000']\n",
      "[-0.08643487 -0.14071974 -0.10398163 -0.12390741 -0.04058979 -0.110649\n",
      " -0.0518984  -0.11276124 -0.06905475 -0.11189392 -0.03059199 -0.10336907\n",
      " -0.04322352  1.5935557  -0.14205763  2.5116613  -0.13316426  2.51967\n",
      " -0.11360867  3.0765228  -0.1163746   0.53939086  2.3499641  -0.1495217\n",
      " -0.16476327  1.9449065  -0.1369017  -0.0880252   2.1848845 ]\n"
     ]
    }
   ],
   "source": [
    "default_layer = 9\n",
    "default_neuron_idx = 652\n",
    "default_inp = \"The following is a list of powers of 10: 1, 10, 100, 1000, 10000, 100000, 1000000, 10000000\"\n",
    "print(model.to_str_tokens(default_inp))\n",
    "print(get_activations(default_inp, default_layer, default_neuron_idx))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8312d15b-26c2-488e-b36b-2f729cd65557",
   "metadata": {},
   "source": [
    "### Visualizing Model Activations"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eec6b6ec-8e8b-4757-a281-e35ff8fd2331",
   "metadata": {},
   "source": [
    "Here we will implement a simple interface that visualizes the neurons of a GPT language model on some text.\n",
    "Each token element is coloured according to the intensity of the selected neuron's activation. We normalize the activations to be between [0,1]\n",
    "\n",
    "Question: How does one tell if neuron 562 in layer 9 activates strongly on powers of 10?\n",
    "\n",
    "This visualization is sensitive to `max_val` and `min_val`. When doing research on neuron activations, these values can be tuned to whatever seems reasonable for the distribution of neurons in question. The defaults come from Nanda, min_val=0, max_val is the max activation across the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 91,
   "id": "65cc5309-caf1-424e-958f-e850b2773a11",
   "metadata": {},
   "outputs": [],
   "source": [
    "# We need css to give each token a thin gray border, to make it easy to see token separation\n",
    "style = \"\"\"\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 92,
   "id": "52d9958c-4fcb-40de-9f65-8836650d1d19",
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_color(val, max_val, min_val):\n",
    "    \"\"\"Normalize val to [0,1] and return a colour which interpolates between slightly off-white and red (0 = white, 1 = red)\n",
    "    Returns a string of the form \"rgb(240,240,240) which CSS recognizes\"\n",
    "    \"\"\"\n",
    "    normalized_val = (val - abs(min_val)) / max_val\n",
    "    st = f\"rgb(240, {240*(1-normalized_val)}, {240*(1-normalized_val)})\"\n",
    "    # print(st)\n",
    "    return st\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 106,
   "id": "32f53105-a6d3-4a60-8977-5459e1e66502",
   "metadata": {},
   "outputs": [],
   "source": [
    "def neuron_viz(inp, layer, neuron_idx, max_val=None, min_val=None):\n",
    "    \"\"\"\n",
    "    inp: Input text to be visualised\n",
    "    layer: layer index into the model\n",
    "    neuron_idx: neuron index\n",
    "    max_val: Top range for our activation range, defaults to max activation across the dataset\n",
    "    min_val: Lower end of our activation range, defaults to 0\n",
    "    \n",
    "    Returns: a string of HTML that displays the text with each token colored according to its activation\n",
    "    It is useful to be able to input a layer's fixed max_val and min_val, otherwise the colors change as you edit the text, which is annoying.\n",
    "    \"\"\"\n",
    "    if layer is None:\n",
    "        return \"Please select a layer\"\n",
    "    \n",
    "    if neuron_idx is None:\n",
    "        return \"Please select a neuron\"\n",
    "    \n",
    "    acts = get_activations(inp, layer, neuron_idx)\n",
    "    act_max = acts.max()\n",
    "    act_min = acts.min()\n",
    "    \n",
    "    # Override defaults if not set\n",
    "    if max_val is None:\n",
    "        max_val = act_max\n",
    "    if min_val is None:\n",
    "        min_val = act_min\n",
    "    \n",
    "    # Make a list of HTML markup to concat into our final HTML\n",
    "    # First add style which adds a border to each token element\n",
    "    markup = [style]\n",
    "    \n",
    "    # Add text to tell us which layer and neuron we're looking at.\n",
    "    markup.append(f\"\"\"\n",
    "<section>\n",
    "Layer: <b>{layer}</b>.\n",
    "\n",
    "Neuron Index: <b>{neuron_idx}</b> </section>\"\"\")\n",
    "    # Add a line telling us the limits of our range\n",
    "    markup.append(f\"\"\"\n",
    "<section>\n",
    "Max Range: <b>{max_val:.4f}</b>\n",
    "    \n",
    "Min Range: <b>{min_val:.4f}</b> </section>\"\"\")\n",
    "    \n",
    "    # If we added a custom range, print a line telling us the range of our activations too\n",
    "    if act_max != max_val or act_min != min_val:\n",
    "        markup.append(f\"\"\"\n",
    "<section>Custom range: \n",
    "\n",
    "Max Act: <b>{act_max:.4f}</b>\n",
    "        \n",
    "Min Act: <b>{act_min:.4f}</b>\n",
    "\n",
    "</section>\"\"\")\n",
    "    \n",
    "    # Tokenize input text\n",
    "    toks = model.to_str_tokens(inp)\n",
    "    for tok, act in zip(toks, acts):\n",
    "        # set bgcolor to the colour calculated from the activation\n",
    "        # Set the span\n",
    "        markup.append(f\"\"\"<span class=\"token\" style=\"background-color:{calculate_color(act, act_max, act_min)}\">{tok}</span>\"\"\")\n",
    "    return \"\".join(markup)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 107,
   "id": "cac4d747-a41e-4ed3-9bf5-cd2036af20dc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "<section>\n",
       "Layer: <b>9</b>.\n",
       "\n",
       "Neuron Index: <b>652</b> </section>\n",
       "<section>\n",
       "Max Range: <b>4.0000</b>\n",
       "    \n",
       "Min Range: <b>0.0000</b> </section>\n",
       "<section>Custom range: \n",
       "\n",
       "Max Act: <b>3.0765</b>\n",
       "        \n",
       "Min Act: <b>-0.1648</b>\n",
       "\n",
       "</section><span class=\"token\" style=\"background-color:rgb(240, 259.596004486084, 259.596004486084)\"><|endoftext|></span><span class=\"token\" style=\"background-color:rgb(240, 263.83077442646027, 263.83077442646027)\">The</span><span class=\"token\" style=\"background-color:rgb(240, 260.9648287296295, 260.9648287296295)\"> following</span><span class=\"token\" style=\"background-color:rgb(240, 262.51924216747284, 262.51924216747284)\"> is</span><span class=\"token\" style=\"background-color:rgb(240, 256.019623875618, 256.019623875618)\"> a</span><span class=\"token\" style=\"background-color:rgb(240, 261.4849519729614, 261.4849519729614)\"> list</span><span class=\"token\" style=\"background-color:rgb(240, 256.9018095731735, 256.9018095731735)\"> of</span><span class=\"token\" style=\"background-color:rgb(240, 261.649729013443, 261.649729013443)\"> powers</span><span class=\"token\" style=\"background-color:rgb(240, 258.2401782274246, 258.2401782274246)\"> of</span><span class=\"token\" style=\"background-color:rgb(240, 261.58206939697266, 261.58206939697266)\"> 10</span><span class=\"token\" style=\"background-color:rgb(240, 255.23969292640686, 255.23969292640686)\">:</span><span class=\"token\" style=\"background-color:rgb(240, 260.9170424938202, 260.9170424938202)\"> 1</span><span class=\"token\" style=\"background-color:rgb(240, 256.2250792980194, 256.2250792980194)\">,</span><span class=\"token\" style=\"background-color:rgb(240, 128.53968858718872, 128.53968858718872)\"> 10</span><span class=\"token\" style=\"background-color:rgb(240, 263.93514454364777, 263.93514454364777)\">,</span><span class=\"token\" style=\"background-color:rgb(240, 56.918134689331055, 56.918134689331055)\"> 100</span><span class=\"token\" style=\"background-color:rgb(240, 263.2413697242737, 263.2413697242737)\">,</span><span class=\"token\" style=\"background-color:rgb(240, 56.293373107910156, 56.293373107910156)\"> 1000</span><span class=\"token\" style=\"background-color:rgb(240, 261.71583473682404, 261.71583473682404)\">,</span><span class=\"token\" style=\"background-color:rgb(240, 12.853202819824219, 12.853202819824219)\"> 10000</span><span class=\"token\" style=\"background-color:rgb(240, 261.9316077232361, 261.9316077232361)\">,</span><span class=\"token\" style=\"background-color:rgb(240, 210.77524602413177, 210.77524602413177)\"> 100</span><span class=\"token\" style=\"background-color:rgb(240, 69.53215599060059, 69.53215599060059)\">000</span><span class=\"token\" style=\"background-color:rgb(240, 264.5174181461334, 264.5174181461334)\">,</span><span class=\"token\" style=\"background-color:rgb(240, 265.70641458034515, 265.70641458034515)\"> 100</span><span class=\"token\" style=\"background-color:rgb(240, 101.13076686859131, 101.13076686859131)\">0000</span><span class=\"token\" style=\"background-color:rgb(240, 263.5329294204712, 263.5329294204712)\">,</span><span class=\"token\" style=\"background-color:rgb(240, 259.7200673818588, 259.7200673818588)\"> 100</span><span class=\"token\" style=\"background-color:rgb(240, 82.41003513336182, 82.41003513336182)\">00000</span>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "default_max = 4.0\n",
    "default_min = 0.0\n",
    "default_markup = neuron_viz(\n",
    "    default_inp,\n",
    "    default_layer,\n",
    "    default_neuron_idx,\n",
    "    max_val=default_max,\n",
    "    min_val=default_min,\n",
    ")\n",
    "\n",
    "display(HTML(default_markup))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45242779-cb41-4aaf-9cc6-105be832c5c9",
   "metadata": {},
   "source": [
    "### Interactive UI: Building a gradio app"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ee612df-5d09-46d6-90ee-7ff777318bdd",
   "metadata": {},
   "source": [
    "We introduce a few elements: TextBoxes, Numbers etc that  return strings and numbers. We also define elements that display things. We call `input.change(update_function, inputs, output)`.\n",
    "This says that \"if that input element changes, run the update function on the value of each of the `inputs` and set the value of `output`  to the output of the function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 113,
   "id": "3b5db838-b426-4783-8b2b-e522f276e6a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "with gr.Blocks() as app:\n",
    "    gr.HTML(value=f\"Interactive Token Visualizer for {model_name}\")\n",
    "    \n",
    "    #input elements\n",
    "    with gr.Row():\n",
    "        with gr.Column():\n",
    "            text = gr.Textbox(label=\"Input Text\", value=default_inp)\n",
    "            # Precision=0 makes it an int, otherwise its a float\n",
    "            # Value sets the initial default value\n",
    "            layer = gr.Number(label=\"Layer\", value=default_layer, precision=0)\n",
    "            neuron_idx = gr.Number(label=\"Neuron Index\", value=default_neuron_idx, precision=0)\n",
    "            \n",
    "            # If empty, these two map to None\n",
    "            max_val = gr.Number(label=\"Max Value\", value=default_max)\n",
    "            min_val = gr.Number(label=\"Min Value\", value=default_min)\n",
    "            inputs = [text, layer, neuron_idx, max_val, min_val]\n",
    "        with gr.Column():\n",
    "            # output Element\n",
    "            out = gr.HTML(label=\"Neuron Activations\", value=default_markup)\n",
    "    for inp in inputs:\n",
    "        inp.change(neuron_viz, inputs, out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 114,
   "id": "aecc4d37-9c4a-49b5-abe4-1f189b3edcb7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running on local URL:  http://127.0.0.1:7863\n",
      "Running on public URL: https://6870c2a1d159526b.gradio.app\n",
      "\n",
      "This share link expires in 72 hours. For free permanent hosting and GPU upgrades (NEW!), check out Spaces: https://huggingface.co/spaces\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div><iframe src=\"https://6870c2a1d159526b.gradio.app\" width=\"100%\" height=\"1000\" allow=\"autoplay; camera; microphone; clipboard-read; clipboard-write;\" frameborder=\"0\" allowfullscreen></iframe></div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "(<gradio.routes.App at 0x7fe4d134bfa0>,\n",
       " 'http://127.0.0.1:7863/',\n",
       " 'https://6870c2a1d159526b.gradio.app')"
      ]
     },
     "execution_count": 114,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "app.launch(share=True, height=1000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7a815a4-de39-44f4-9cd4-11c64310abc6",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}