owt_comma_ablation.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "README.md\t\t    miniconda.sh  quick_start_pytorch.ipynb   wandb\n",
      "eliciting-latent-sentiment  miniconda3\t  quick_start_pytorch_images\n"
     ]
    }
   ],
   "source": [
    "!ls"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/notebooks/eliciting-latent-sentiment\n"
     ]
    }
   ],
   "source": [
    "%cd eliciting-latent-sentiment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collecting fancy_einsum==0.0.3\n",
      "  Downloading fancy_einsum-0.0.3-py3-none-any.whl (6.2 kB)\n",
      "Installing collected packages: fancy_einsum\n",
      "Successfully installed fancy_einsum-0.0.3\n",
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
      "\u001b[0mCollecting transformer_lens\n",
      "  Downloading transformer_lens-1.6.0-py3-none-any.whl (105 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m106.0/106.0 kB\u001b[0m \u001b[31m16.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hCollecting transformers>=4.25.1\n",
      "  Downloading transformers-4.32.1-py3-none-any.whl (7.5 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m7.5/7.5 MB\u001b[0m \u001b[31m3.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m00:01\u001b[0mm\n",
      "\u001b[?25hCollecting wandb>=0.13.5\n",
      "  Downloading wandb-0.15.9-py3-none-any.whl (2.1 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2.1/2.1 MB\u001b[0m \u001b[31m40.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0ma \u001b[36m0:00:01\u001b[0m\n",
      "\u001b[?25hCollecting beartype<0.15.0,>=0.14.1\n",
      "  Downloading beartype-0.14.1-py3-none-any.whl (739 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m739.7/739.7 kB\u001b[0m \u001b[31m44.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hRequirement already satisfied: rich>=12.6.0 in /usr/local/lib/python3.9/dist-packages (from transformer_lens) (13.2.0)\n",
      "Collecting einops>=0.6.0\n",
      "  Downloading einops-0.6.1-py3-none-any.whl (42 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m42.2/42.2 kB\u001b[0m \u001b[31m17.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hRequirement already satisfied: tqdm>=4.64.1 in /usr/local/lib/python3.9/dist-packages (from transformer_lens) (4.64.1)\n",
      "Requirement already satisfied: pandas>=1.1.5 in /usr/local/lib/python3.9/dist-packages (from transformer_lens) (1.5.0)\n",
      "Requirement already satisfied: numpy>=1.21 in /usr/local/lib/python3.9/dist-packages (from transformer_lens) (1.23.4)\n",
      "Collecting jaxtyping>=0.2.11\n",
      "  Downloading jaxtyping-0.2.21-py3-none-any.whl (25 kB)\n",
      "Collecting datasets>=2.7.1\n",
      "  Downloading datasets-2.14.4-py3-none-any.whl (519 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m519.3/519.3 kB\u001b[0m \u001b[31m23.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hRequirement already satisfied: fancy-einsum>=0.0.3 in /usr/local/lib/python3.9/dist-packages (from transformer_lens) (0.0.3)\n",
      "Requirement already satisfied: torch>=1.10 in /usr/local/lib/python3.9/dist-packages (from transformer_lens) (1.12.1+cu116)\n",
      "Requirement already satisfied: multiprocess in /usr/local/lib/python3.9/dist-packages (from datasets>=2.7.1->transformer_lens) (0.70.13)\n",
      "Requirement already satisfied: pyyaml>=5.1 in /usr/local/lib/python3.9/dist-packages (from datasets>=2.7.1->transformer_lens) (5.4.1)\n",
      "Requirement already satisfied: fsspec[http]>=2021.11.1 in /usr/local/lib/python3.9/dist-packages (from datasets>=2.7.1->transformer_lens) (2023.1.0)\n",
      "Requirement already satisfied: dill<0.3.8,>=0.3.0 in /usr/local/lib/python3.9/dist-packages (from datasets>=2.7.1->transformer_lens) (0.3.5.1)\n",
      "Collecting huggingface-hub<1.0.0,>=0.14.0\n",
      "  Downloading huggingface_hub-0.16.4-py3-none-any.whl (268 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m268.8/268.8 kB\u001b[0m \u001b[31m5.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0ma \u001b[36m0:00:01\u001b[0m\n",
      "\u001b[?25hRequirement already satisfied: requests>=2.19.0 in /usr/local/lib/python3.9/dist-packages (from datasets>=2.7.1->transformer_lens) (2.28.2)\n",
      "Requirement already satisfied: pyarrow>=8.0.0 in /usr/local/lib/python3.9/dist-packages (from datasets>=2.7.1->transformer_lens) (10.0.1)\n",
      "Requirement already satisfied: packaging in /usr/local/lib/python3.9/dist-packages (from datasets>=2.7.1->transformer_lens) (23.0)\n",
      "Requirement already satisfied: aiohttp in /usr/local/lib/python3.9/dist-packages (from datasets>=2.7.1->transformer_lens) (3.8.3)\n",
      "Requirement already satisfied: xxhash in /usr/local/lib/python3.9/dist-packages (from datasets>=2.7.1->transformer_lens) (3.2.0)\n",
      "Collecting typeguard>=2.13.3\n",
      "  Downloading typeguard-4.1.3-py3-none-any.whl (33 kB)\n",
      "Requirement already satisfied: typing-extensions>=3.7.4.1 in /usr/local/lib/python3.9/dist-packages (from jaxtyping>=0.2.11->transformer_lens) (4.4.0)\n",
      "Requirement already satisfied: pytz>=2020.1 in /usr/local/lib/python3.9/dist-packages (from pandas>=1.1.5->transformer_lens) (2022.7.1)\n",
      "Requirement already satisfied: python-dateutil>=2.8.1 in /usr/local/lib/python3.9/dist-packages (from pandas>=1.1.5->transformer_lens) (2.8.2)\n",
      "Requirement already satisfied: pygments<3.0.0,>=2.6.0 in /usr/local/lib/python3.9/dist-packages (from rich>=12.6.0->transformer_lens) (2.14.0)\n",
      "Requirement already satisfied: markdown-it-py<3.0.0,>=2.1.0 in /usr/local/lib/python3.9/dist-packages (from rich>=12.6.0->transformer_lens) (2.1.0)\n",
      "Collecting safetensors>=0.3.1\n",
      "  Downloading safetensors-0.3.3-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (1.3 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.3/1.3 MB\u001b[0m \u001b[31m3.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0ma \u001b[36m0:00:01\u001b[0mm\n",
      "\u001b[?25hRequirement already satisfied: tokenizers!=0.11.3,<0.14,>=0.11.1 in /usr/local/lib/python3.9/dist-packages (from transformers>=4.25.1->transformer_lens) (0.12.1)\n",
      "Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.9/dist-packages (from transformers>=4.25.1->transformer_lens) (2022.10.31)\n",
      "Requirement already satisfied: filelock in /usr/local/lib/python3.9/dist-packages (from transformers>=4.25.1->transformer_lens) (3.9.0)\n",
      "Requirement already satisfied: Click!=8.0.0,>=7.1 in /usr/local/lib/python3.9/dist-packages (from wandb>=0.13.5->transformer_lens) (8.1.3)\n",
      "Collecting appdirs>=1.4.3\n",
      "  Downloading appdirs-1.4.4-py2.py3-none-any.whl (9.6 kB)\n",
      "Requirement already satisfied: docker-pycreds>=0.4.0 in /usr/local/lib/python3.9/dist-packages (from wandb>=0.13.5->transformer_lens) (0.4.0)\n",
      "Requirement already satisfied: setuptools in /usr/local/lib/python3.9/dist-packages (from wandb>=0.13.5->transformer_lens) (66.1.1)\n",
      "Requirement already satisfied: GitPython!=3.1.29,>=1.0.0 in /usr/local/lib/python3.9/dist-packages (from wandb>=0.13.5->transformer_lens) (3.1.30)\n",
      "Requirement already satisfied: protobuf!=4.21.0,<5,>=3.15.0 in /usr/local/lib/python3.9/dist-packages (from wandb>=0.13.5->transformer_lens) (3.19.6)\n",
      "Requirement already satisfied: sentry-sdk>=1.0.0 in /usr/local/lib/python3.9/dist-packages (from wandb>=0.13.5->transformer_lens) (1.14.0)\n",
      "Requirement already satisfied: psutil>=5.0.0 in /usr/local/lib/python3.9/dist-packages (from wandb>=0.13.5->transformer_lens) (5.9.4)\n",
      "Requirement already satisfied: pathtools in /usr/local/lib/python3.9/dist-packages (from wandb>=0.13.5->transformer_lens) (0.1.2)\n",
      "Requirement already satisfied: setproctitle in /usr/local/lib/python3.9/dist-packages (from wandb>=0.13.5->transformer_lens) (1.3.2)\n",
      "Requirement already satisfied: six>=1.4.0 in /usr/lib/python3/dist-packages (from docker-pycreds>=0.4.0->wandb>=0.13.5->transformer_lens) (1.14.0)\n",
      "Requirement already satisfied: multidict<7.0,>=4.5 in /usr/local/lib/python3.9/dist-packages (from aiohttp->datasets>=2.7.1->transformer_lens) (6.0.4)\n",
      "Requirement already satisfied: aiosignal>=1.1.2 in /usr/local/lib/python3.9/dist-packages (from aiohttp->datasets>=2.7.1->transformer_lens) (1.3.1)\n",
      "Requirement already satisfied: frozenlist>=1.1.1 in /usr/local/lib/python3.9/dist-packages (from aiohttp->datasets>=2.7.1->transformer_lens) (1.3.3)\n",
      "Requirement already satisfied: attrs>=17.3.0 in /usr/local/lib/python3.9/dist-packages (from aiohttp->datasets>=2.7.1->transformer_lens) (18.2.0)\n",
      "Requirement already satisfied: async-timeout<5.0,>=4.0.0a3 in /usr/local/lib/python3.9/dist-packages (from aiohttp->datasets>=2.7.1->transformer_lens) (4.0.2)\n",
      "Requirement already satisfied: charset-normalizer<3.0,>=2.0 in /usr/local/lib/python3.9/dist-packages (from aiohttp->datasets>=2.7.1->transformer_lens) (2.1.1)\n",
      "Requirement already satisfied: yarl<2.0,>=1.0 in /usr/local/lib/python3.9/dist-packages (from aiohttp->datasets>=2.7.1->transformer_lens) (1.8.2)\n",
      "Requirement already satisfied: gitdb<5,>=4.0.1 in /usr/local/lib/python3.9/dist-packages (from GitPython!=3.1.29,>=1.0.0->wandb>=0.13.5->transformer_lens) (4.0.10)\n",
      "Requirement already satisfied: mdurl~=0.1 in /usr/local/lib/python3.9/dist-packages (from markdown-it-py<3.0.0,>=2.1.0->rich>=12.6.0->transformer_lens) (0.1.2)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /usr/lib/python3/dist-packages (from requests>=2.19.0->datasets>=2.7.1->transformer_lens) (2019.11.28)\n",
      "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /usr/local/lib/python3.9/dist-packages (from requests>=2.19.0->datasets>=2.7.1->transformer_lens) (1.26.14)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /usr/lib/python3/dist-packages (from requests>=2.19.0->datasets>=2.7.1->transformer_lens) (2.8)\n",
      "Requirement already satisfied: importlib-metadata>=3.6 in /usr/local/lib/python3.9/dist-packages (from typeguard>=2.13.3->jaxtyping>=0.2.11->transformer_lens) (6.0.0)\n",
      "Collecting typing-extensions>=3.7.4.1\n",
      "  Downloading typing_extensions-4.7.1-py3-none-any.whl (33 kB)\n",
      "Requirement already satisfied: smmap<6,>=3.0.1 in /usr/local/lib/python3.9/dist-packages (from gitdb<5,>=4.0.1->GitPython!=3.1.29,>=1.0.0->wandb>=0.13.5->transformer_lens) (5.0.0)\n",
      "Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.9/dist-packages (from importlib-metadata>=3.6->typeguard>=2.13.3->jaxtyping>=0.2.11->transformer_lens) (3.11.0)\n",
      "Installing collected packages: safetensors, appdirs, typing-extensions, einops, beartype, typeguard, huggingface-hub, wandb, transformers, jaxtyping, datasets, transformer_lens\n",
      "  Attempting uninstall: typing-extensions\n",
      "    Found existing installation: typing_extensions 4.4.0\n",
      "    Uninstalling typing_extensions-4.4.0:\n",
      "      Successfully uninstalled typing_extensions-4.4.0\n",
      "  Attempting uninstall: huggingface-hub\n",
      "    Found existing installation: huggingface-hub 0.12.0\n",
      "    Uninstalling huggingface-hub-0.12.0:\n",
      "      Successfully uninstalled huggingface-hub-0.12.0\n",
      "  Attempting uninstall: wandb\n",
      "    Found existing installation: wandb 0.13.4\n",
      "    Uninstalling wandb-0.13.4:\n",
      "      Successfully uninstalled wandb-0.13.4\n",
      "  Attempting uninstall: transformers\n",
      "    Found existing installation: transformers 4.21.3\n",
      "    Uninstalling transformers-4.21.3:\n",
      "      Successfully uninstalled transformers-4.21.3\n",
      "  Attempting uninstall: datasets\n",
      "    Found existing installation: datasets 2.4.0\n",
      "    Uninstalling datasets-2.4.0:\n",
      "      Successfully uninstalled datasets-2.4.0\n",
      "Successfully installed appdirs-1.4.4 beartype-0.14.1 datasets-2.14.4 einops-0.6.1 huggingface-hub-0.16.4 jaxtyping-0.2.21 safetensors-0.3.3 transformer_lens-1.6.0 transformers-4.32.1 typeguard-4.1.3 typing-extensions-4.7.1 wandb-0.15.9\n",
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
      "\u001b[0mCollecting jaxtyping==0.2.13\n",
      "  Downloading jaxtyping-0.2.13-py3-none-any.whl (19 kB)\n",
      "Requirement already satisfied: typing-extensions>=3.7.4.1 in /usr/local/lib/python3.9/dist-packages (from jaxtyping==0.2.13) (4.7.1)\n",
      "Requirement already satisfied: typeguard>=2.13.3 in /usr/local/lib/python3.9/dist-packages (from jaxtyping==0.2.13) (4.1.3)\n",
      "Requirement already satisfied: numpy>=1.20.0 in /usr/local/lib/python3.9/dist-packages (from jaxtyping==0.2.13) (1.23.4)\n",
      "Requirement already satisfied: importlib-metadata>=3.6 in /usr/local/lib/python3.9/dist-packages (from typeguard>=2.13.3->jaxtyping==0.2.13) (6.0.0)\n",
      "Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.9/dist-packages (from importlib-metadata>=3.6->typeguard>=2.13.3->jaxtyping==0.2.13) (3.11.0)\n",
      "Installing collected packages: jaxtyping\n",
      "  Attempting uninstall: jaxtyping\n",
      "    Found existing installation: jaxtyping 0.2.21\n",
      "    Uninstalling jaxtyping-0.2.21:\n",
      "      Successfully uninstalled jaxtyping-0.2.21\n",
      "Successfully installed jaxtyping-0.2.13\n",
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
      "\u001b[0mRequirement already satisfied: einops in /usr/local/lib/python3.9/dist-packages (0.6.1)\n",
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
      "\u001b[0mCollecting protobuf==3.20.*\n",
      "  Downloading protobuf-3.20.3-cp39-cp39-manylinux_2_5_x86_64.manylinux1_x86_64.whl (1.0 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.0/1.0 MB\u001b[0m \u001b[31m80.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hInstalling collected packages: protobuf\n",
      "  Attempting uninstall: protobuf\n",
      "    Found existing installation: protobuf 3.19.6\n",
      "    Uninstalling protobuf-3.19.6:\n",
      "      Successfully uninstalled protobuf-3.19.6\n",
      "\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
      "tensorflow 2.9.2 requires protobuf<3.20,>=3.9.2, but you have protobuf 3.20.3 which is incompatible.\n",
      "tensorboard 2.9.1 requires protobuf<3.20,>=3.9.2, but you have protobuf 3.20.3 which is incompatible.\u001b[0m\u001b[31m\n",
      "\u001b[0mSuccessfully installed protobuf-3.20.3\n",
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
      "\u001b[0mCollecting plotly\n",
      "  Downloading plotly-5.16.1-py2.py3-none-any.whl (15.6 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m15.6/15.6 MB\u001b[0m \u001b[31m84.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m00:01\u001b[0m\n",
      "\u001b[?25hCollecting tenacity>=6.2.0\n",
      "  Downloading tenacity-8.2.3-py3-none-any.whl (24 kB)\n",
      "Requirement already satisfied: packaging in /usr/local/lib/python3.9/dist-packages (from plotly) (23.0)\n",
      "Installing collected packages: tenacity, plotly\n",
      "Successfully installed plotly-5.16.1 tenacity-8.2.3\n",
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
      "\u001b[0mCollecting torchtyping\n",
      "  Downloading torchtyping-0.1.4-py3-none-any.whl (17 kB)\n",
      "Requirement already satisfied: torch>=1.7.0 in /usr/local/lib/python3.9/dist-packages (from torchtyping) (1.12.1+cu116)\n",
      "Requirement already satisfied: typeguard>=2.11.1 in /usr/local/lib/python3.9/dist-packages (from torchtyping) (4.1.3)\n",
      "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.9/dist-packages (from torch>=1.7.0->torchtyping) (4.7.1)\n",
      "Requirement already satisfied: importlib-metadata>=3.6 in /usr/local/lib/python3.9/dist-packages (from typeguard>=2.11.1->torchtyping) (6.0.0)\n",
      "Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.9/dist-packages (from importlib-metadata>=3.6->typeguard>=2.11.1->torchtyping) (3.11.0)\n",
      "Installing collected packages: torchtyping\n",
      "Successfully installed torchtyping-0.1.4\n",
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
      "\u001b[0mCollecting git+https://github.com/neelnanda-io/neel-plotly.git\n",
      "  Cloning https://github.com/neelnanda-io/neel-plotly.git to /tmp/pip-req-build-el7an8z5\n",
      "  Running command git clone --filter=blob:none --quiet https://github.com/neelnanda-io/neel-plotly.git /tmp/pip-req-build-el7an8z5\n",
      "  Resolved https://github.com/neelnanda-io/neel-plotly.git to commit 6dc24b26f8dec991908479d7445dae496b3430b7\n",
      "  Preparing metadata (setup.py) ... \u001b[?25ldone\n",
      "\u001b[?25hRequirement already satisfied: einops in /usr/local/lib/python3.9/dist-packages (from neel-plotly==0.0.0) (0.6.1)\n",
      "Requirement already satisfied: numpy in /usr/local/lib/python3.9/dist-packages (from neel-plotly==0.0.0) (1.23.4)\n",
      "Requirement already satisfied: torch in /usr/local/lib/python3.9/dist-packages (from neel-plotly==0.0.0) (1.12.1+cu116)\n",
      "Requirement already satisfied: plotly in /usr/local/lib/python3.9/dist-packages (from neel-plotly==0.0.0) (5.16.1)\n",
      "Requirement already satisfied: tqdm in /usr/local/lib/python3.9/dist-packages (from neel-plotly==0.0.0) (4.64.1)\n",
      "Requirement already satisfied: pandas in /usr/local/lib/python3.9/dist-packages (from neel-plotly==0.0.0) (1.5.0)\n",
      "Requirement already satisfied: python-dateutil>=2.8.1 in /usr/local/lib/python3.9/dist-packages (from pandas->neel-plotly==0.0.0) (2.8.2)\n",
      "Requirement already satisfied: pytz>=2020.1 in /usr/local/lib/python3.9/dist-packages (from pandas->neel-plotly==0.0.0) (2022.7.1)\n",
      "Requirement already satisfied: packaging in /usr/local/lib/python3.9/dist-packages (from plotly->neel-plotly==0.0.0) (23.0)\n",
      "Requirement already satisfied: tenacity>=6.2.0 in /usr/local/lib/python3.9/dist-packages (from plotly->neel-plotly==0.0.0) (8.2.3)\n",
      "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.9/dist-packages (from torch->neel-plotly==0.0.0) (4.7.1)\n",
      "Requirement already satisfied: six>=1.5 in /usr/lib/python3/dist-packages (from python-dateutil>=2.8.1->pandas->neel-plotly==0.0.0) (1.14.0)\n",
      "Building wheels for collected packages: neel-plotly\n",
      "  Building wheel for neel-plotly (setup.py) ... \u001b[?25ldone\n",
      "\u001b[?25h  Created wheel for neel-plotly: filename=neel_plotly-0.0.0-py3-none-any.whl size=10186 sha256=afa51279c0624f470bb7f5c73e532b069ba7531f6dd0da9d003fd01f22030a8e\n",
      "  Stored in directory: /tmp/pip-ephem-wheel-cache-31poqqir/wheels/e1/3c/c0/b5897c402b85e7fc329feb205ad5948b518f0423d891a79f7f\n",
      "Successfully built neel-plotly\n",
      "Installing collected packages: neel-plotly\n",
      "Successfully installed neel-plotly-0.0.0\n",
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
      "\u001b[0mCollecting circuitsvis\n",
      "  Downloading circuitsvis-1.41.0-py3-none-any.whl (1.8 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.8/1.8 MB\u001b[0m \u001b[31m67.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hCollecting importlib-metadata<6.0.0,>=5.1.0\n",
      "  Downloading importlib_metadata-5.2.0-py3-none-any.whl (21 kB)\n",
      "Requirement already satisfied: numpy<2.0,>=1.21 in /usr/local/lib/python3.9/dist-packages (from circuitsvis) (1.23.4)\n",
      "Requirement already satisfied: torch>=1.10 in /usr/local/lib/python3.9/dist-packages (from circuitsvis) (1.12.1+cu116)\n",
      "Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.9/dist-packages (from importlib-metadata<6.0.0,>=5.1.0->circuitsvis) (3.11.0)\n",
      "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.9/dist-packages (from torch>=1.10->circuitsvis) (4.7.1)\n",
      "Installing collected packages: importlib-metadata, circuitsvis\n",
      "  Attempting uninstall: importlib-metadata\n",
      "    Found existing installation: importlib-metadata 6.0.0\n",
      "    Uninstalling importlib-metadata-6.0.0:\n",
      "      Successfully uninstalled importlib-metadata-6.0.0\n",
      "Successfully installed circuitsvis-1.41.0 importlib-metadata-5.2.0\n",
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "#!pip install git+https://github.com/callummcdougall/CircuitsVis.git#subdirectory=python\n",
    "!pip install fancy_einsum==0.0.3\n",
    "!pip install transformer_lens\n",
    "!pip install jaxtyping==0.2.13\n",
    "!pip install einops\n",
    "!pip install protobuf==3.20.*\n",
    "!pip install plotly\n",
    "!pip install torchtyping\n",
    "!pip install git+https://github.com/neelnanda-io/neel-p\n",
    "!pip install circuitsvis\n",
    "# !curl -fsSL https://deb.nodesource.com/setup_16.x | sudo -E bash -; sudo apt-get install -y nodejs\n",
    "# %pip install git+https://github.com/neelnanda-io/PySvelte.git\n",
    "# %pip install typeguard==2.13.3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython import get_ipython\n",
    "ipython = get_ipython()\n",
    "ipython.run_line_magic(\"load_ext\", \"autoreload\")\n",
    "ipython.run_line_magic(\"autoreload\", \"2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import einops\n",
    "from functools import partial\n",
    "import torch\n",
    "import datasets\n",
    "from torch import Tensor\n",
    "from torch.utils.data import DataLoader\n",
    "from datasets import load_dataset\n",
    "from jaxtyping import Float, Int, Bool\n",
    "from typing import Dict, Iterable, List, Tuple, Union\n",
    "from transformer_lens import HookedTransformer\n",
    "from transformer_lens.utils import get_dataset, tokenize_and_concatenate, get_act_name, test_prompt\n",
    "from transformer_lens.hook_points import HookPoint\n",
    "from tqdm.notebook import tqdm\n",
    "import pandas as pd\n",
    "from circuitsvis.activations import text_neuron_activations\n",
    "from utils.store import load_array, save_html, save_array, is_file, get_model_name, clean_label, save_text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def find_positions(tensor, token_ids=[11, 13]):\n",
    "    positions = []\n",
    "    for batch_item in tensor:\n",
    "        token_positions = {token_id: [] for token_id in token_ids}\n",
    "        for position, token in enumerate(batch_item):\n",
    "            if token.item() in token_ids:\n",
    "                token_positions[token.item()].append(position)\n",
    "        positions.append([token_positions[token_id] for token_id in token_ids])\n",
    "    return positions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def zero_attention_pos_hook(\n",
    "    pattern: Float[Tensor, \"batch head seq_Q seq_K\"], hook: HookPoint,\n",
    "    pos_by_batch: List[List[int]], layer: int = 0, head_idx: int = 0,\n",
    ") -> Float[Tensor, \"batch head seq_Q seq_K\"]:\n",
    "    \"\"\"Zero-ablates an attention pattern tensor at a particular position\"\"\"\n",
    "    assert 'pattern' in hook.name\n",
    "\n",
    "    batch_size = pattern.shape[0]\n",
    "    assert len(pos_by_batch) == batch_size\n",
    "\n",
    "    for i in range(batch_size):\n",
    "        for p in pos_by_batch[i]:\n",
    "            pattern[i, head_idx, p, p] = 0\n",
    "            \n",
    "    return pattern"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def names_filter(name: str):\n",
    "    \"\"\"Filter for the names of the activations we want to keep to study the resid stream.\"\"\"\n",
    "    return name.endswith('resid_post') or name == get_act_name('resid_pre', 0)\n",
    "\n",
    "def get_layerwise_token_mean_activations(model: HookedTransformer, data_loader: DataLoader, token_id: int = 13) -> Float[Tensor, \"layer d_model\"]:\n",
    "    \"\"\"Get the mean value of a token across layers\"\"\"\n",
    "    num_layers = model.cfg.n_layers\n",
    "    d_model = model.cfg.d_model\n",
    "    \n",
    "    activation_sums = torch.stack([torch.zeros(d_model) for _ in range(num_layers)]).to(device)\n",
    "    comma_counts = [0] * num_layers\n",
    "\n",
    "    print(activation_sums.shape)\n",
    "\n",
    "    token_mean_values = torch.zeros((num_layers, d_model))\n",
    "    for _, batch_value in tqdm(enumerate(data_loader), total=100):\n",
    "        \n",
    "        batch_tokens = batch_value['tokens'].to(device)\n",
    "\n",
    "        # get positions of all 11 and 13 token ids in batch\n",
    "        punct_pos = find_positions(batch_tokens, token_ids=[13])\n",
    "\n",
    "        _, cache = model.run_with_cache(\n",
    "            batch_tokens, \n",
    "            names_filter=names_filter\n",
    "        )\n",
    "\n",
    "        \n",
    "        for i in range(batch_tokens.shape[0]):\n",
    "            for p in punct_pos[i][0]:\n",
    "                for layer in range(num_layers):\n",
    "                    activation_sums[layer] += cache[f\"blocks.{layer}.hook_resid_post\"][i, p, :]\n",
    "                    comma_counts[layer] += 1\n",
    "\n",
    "    for layer in range(num_layers):\n",
    "        token_mean_values[layer] = activation_sums[layer] / comma_counts[layer]\n",
    "\n",
    "    return token_mean_values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_zeroed_attn_modified_loss(model: HookedTransformer, data_loader: DataLoader) -> float:\n",
    "    total_loss = 0\n",
    "    loss_list = []\n",
    "    for _, batch_value in tqdm(enumerate(data_loader), total=len(data_loader)):\n",
    "        batch_tokens = batch_value['tokens'].to(device)\n",
    "\n",
    "        # get positions of all 11 and 13 token ids in batch\n",
    "        punct_pos = find_positions(batch_tokens, token_ids=[13])\n",
    "\n",
    "        # get the loss for each token in the batch\n",
    "        initial_loss = model(batch_tokens, return_type=\"loss\", prepend_bos=False, loss_per_token=True)\n",
    "        \n",
    "        # add hooks for the activations of the 11 and 13 tokens\n",
    "        for layer, head in heads_to_ablate:\n",
    "            ablate_punct = partial(zero_attention_pos_hook, pos_by_batch=punct_pos, layer=layer, head_idx=head)\n",
    "            model.blocks[layer].attn.hook_pattern.add_hook(ablate_punct)\n",
    "\n",
    "        # get the loss for each token when run with hooks\n",
    "        hooked_loss = model(batch_tokens, return_type=\"loss\", prepend_bos=False, loss_per_token=True)\n",
    "\n",
    "        # compute the percent difference between the two losses\n",
    "        loss_diff = (hooked_loss - initial_loss) / initial_loss\n",
    "\n",
    "        loss_list.append(loss_diff)\n",
    "\n",
    "    model.reset_hooks()\n",
    "    return loss_list, batch_tokens"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 128,
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils.ablation import ablate_resid_with_precalc_mean\n",
    "\n",
    "def compute_mean_ablation_modified_loss(model: HookedTransformer, data_loader: DataLoader, cached_means, target_token_ids) -> float:\n",
    "    total_loss = 0\n",
    "    loss_diff_list = []\n",
    "    orig_loss_list = []\n",
    "    for _, batch_value in tqdm(enumerate(data_loader), total=len(data_loader)):\n",
    "        if isinstance(batch_value['tokens'], list):\n",
    "            batch_tokens = torch.stack(batch_value['tokens']).to(device)\n",
    "        else:\n",
    "            batch_tokens = batch_value['tokens'].to(device)\n",
    "\n",
    "        batch_tokens = einops.rearrange(batch_tokens, 'seq batch -> batch seq')\n",
    "        punct_pos = batch_value['positions']\n",
    "        print(f\"punct_pos: {punct_pos}\")\n",
    "\n",
    "        # get the loss for each token in the batch\n",
    "        initial_loss = model(batch_tokens, return_type=\"loss\", prepend_bos=False, loss_per_token=True)\n",
    "        print(f\"initial loss shape: {initial_loss.shape}\")\n",
    "        orig_loss_list.append(initial_loss)\n",
    "        \n",
    "        # add hooks for the activations of the 11 and 13 tokens\n",
    "        for layer, head in heads_to_ablate:\n",
    "            mean_ablate_comma = partial(ablate_resid_with_precalc_mean_no_batch, cached_means=cached_means, pos_by_batch=punct_pos, layer=layer)\n",
    "            model.blocks[layer].hook_resid_post.add_hook(mean_ablate_comma)\n",
    "\n",
    "        # get the loss for each token when run with hooks\n",
    "        print(f\"batch tokens shape: {batch_tokens.shape}\")\n",
    "        \n",
    "        hooked_loss = model(batch_tokens, return_type=\"loss\", prepend_bos=False, loss_per_token=True)\n",
    "        print(f\"hooked loss shape: {hooked_loss.shape}\")\n",
    "\n",
    "        # compute the difference between the two losses\n",
    "        loss_diff = hooked_loss - initial_loss\n",
    "        \n",
    "        # set all positions right after punct_pos to zero\n",
    "        for p in punct_pos:\n",
    "            print(f\"zeroing {p}\")\n",
    "            loss_diff[0, p] = 0\n",
    "\n",
    "        loss_diff_list.append(loss_diff)\n",
    "\n",
    "    model.reset_hooks()\n",
    "    return loss_diff_list, orig_loss_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 123,
   "metadata": {},
   "outputs": [],
   "source": [
    "def ablate_resid_with_precalc_mean_no_batch(\n",
    "    component: Float[Tensor, \"batch ...\"],\n",
    "    hook: HookPoint,\n",
    "    cached_means: Float[Tensor, \"layer ...\"],\n",
    "    pos_by_batch: List[Tensor],\n",
    "    layer: int = 0,\n",
    ") -> Float[Tensor, \"batch ...\"]:\n",
    "    \"\"\"\n",
    "    Mean-ablates a batch tensor\n",
    "\n",
    "    :param component: the tensor to compute the mean over the batch dim of\n",
    "    :return: the mean over the cache component of the tensor\n",
    "    \"\"\"\n",
    "    assert 'resid' in hook.name\n",
    "\n",
    "    #print(f\"batch size: {batch_size} pos_by_batch: {len(pos_by_batch)}\")\n",
    "\n",
    "    for p in pos_by_batch:\n",
    "        component[:, p] = cached_means[layer]\n",
    "            \n",
    "    return component"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "metadata": {},
   "outputs": [],
   "source": [
    "def convert_to_tensors(dataset, column_name='tokens'):\n",
    "    token_buffer = []\n",
    "    final_batches = []\n",
    "    \n",
    "    for batch in dataset:\n",
    "        trimmed_batch = batch[column_name] #[batch[column_name][0]] + [token for token in batch[column_name] if token != 0]\n",
    "        final_batches.append(trimmed_batch)\n",
    "    \n",
    "    # Convert list of batches to tensors\n",
    "    final_batches = [torch.tensor(batch, dtype=torch.long) for batch in final_batches]\n",
    "    # Create a new dataset with specified features\n",
    "    features = Features({\"tokens\": Sequence(Value(\"int64\"))})\n",
    "    final_dataset = Dataset.from_dict({\"tokens\": final_batches}, features=features)\n",
    "\n",
    "    final_dataset.set_format(type=\"torch\", columns=[\"tokens\"])\n",
    "    \n",
    "    return final_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_neuroscope(\n",
    "    tokens: Int[Tensor, \"batch pos\"], centred: bool = False, activations: Float[Tensor, \"pos layer 1\"] = None,\n",
    "    verbose=False,\n",
    "):\n",
    "    \n",
    "    str_tokens = model.to_str_tokens(tokens, prepend_bos=False)\n",
    "\n",
    "    if verbose:\n",
    "        print(f\"Tokens shape: {tokens.shape}\")\n",
    "  \n",
    "    if centred:\n",
    "        if verbose:\n",
    "            print(\"Centering activations\")\n",
    "        layer_means = einops.reduce(activations, \"pos layer 1 -> 1 layer 1\", reduction=\"mean\")\n",
    "        layer_means = einops.repeat(layer_means, \"1 layer 1 -> pos layer 1\", pos=activations.shape[0])\n",
    "        activations -= layer_means\n",
    "    elif verbose:\n",
    "        print(\"Activations already centered\")\n",
    "    assert (\n",
    "        activations.ndim == 3\n",
    "    ), f\"activations must be of shape [tokens x layers x neurons], found {activations.shape}\"\n",
    "    assert len(str_tokens) == activations.shape[0], (\n",
    "        f\"tokens and activations must have the same length, found tokens={len(str_tokens)} and acts={activations.shape[0]}, \"\n",
    "        f\"tokens={str_tokens}, \"\n",
    "        f\"activations={activations.shape}\"\n",
    "\n",
    "    )\n",
    "    return text_neuron_activations(\n",
    "        tokens=str_tokens, \n",
    "        activations=activations,\n",
    "        first_dimension_name=\"Layer (resid_pre)\",\n",
    "        second_dimension_name=\"Model\",\n",
    "        second_dimension_labels=[\"pythia-2.8b\"],\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Comma Ablation on Natural Text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using pad_token, but it is not set yet.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded pretrained model EleutherAI/pythia-2.8b into HookedTransformer\n"
     ]
    }
   ],
   "source": [
    "torch.set_grad_enabled(False)\n",
    "device = \"cuda\"\n",
    "MODEL_NAME = \"EleutherAI/pythia-2.8b\"\n",
    "model = HookedTransformer.from_pretrained(\n",
    "    MODEL_NAME,\n",
    "    center_unembed=True,\n",
    "    center_writing_weights=True,\n",
    "    fold_ln=True,\n",
    "    refactor_factored_attn_matrices=False,\n",
    "    device=device,\n",
    ")\n",
    "model.name = MODEL_NAME"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['text', 'label'],\n",
       "        num_rows: 7864\n",
       "    })\n",
       "    test: Dataset({\n",
       "        features: ['text', 'label'],\n",
       "        num_rows: 2058\n",
       "    })\n",
       "    dev: Dataset({\n",
       "        features: ['text', 'label'],\n",
       "        num_rows: 1007\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from datasets import load_from_disk\n",
    "sst_data = load_from_disk(\"sst2\")\n",
    "sst_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "from datasets import Dataset, Features, Sequence, Value\n",
    "from transformers import AutoTokenizer\n",
    "import torch\n",
    "\n",
    "def tokenize_and_concatenate2(dataset, tokenizer, max_length=1024, column_name='text', add_bos_token=True):\n",
    "    token_buffer = []\n",
    "    final_batches = []\n",
    "    \n",
    "    for batch in dataset:\n",
    "        text = batch[column_name]\n",
    "        if add_bos_token:\n",
    "            text = tokenizer.bos_token + text\n",
    "        tokenized_text = tokenizer(text, add_special_tokens=False)['input_ids']\n",
    "        eos_token_id = tokenizer.eos_token_id\n",
    "        tokenized_text.append(eos_token_id)\n",
    "        token_buffer.extend(tokenized_text)\n",
    "        \n",
    "        while len(token_buffer) >= max_length:\n",
    "            final_batch = token_buffer[:max_length]\n",
    "            token_buffer = token_buffer[max_length:]\n",
    "            final_batches.append(final_batch)\n",
    "    \n",
    "    # Handle any remaining tokens\n",
    "    if len(token_buffer) > 0:\n",
    "        final_batches.append(token_buffer)\n",
    "    \n",
    "    # Convert list of batches to tensors\n",
    "    final_batches = [torch.tensor(batch) for batch in final_batches]\n",
    "    \n",
    "    # Create a new dataset with specified features\n",
    "    features = Features({\"tokens\": Sequence(Value(\"int64\"))})\n",
    "    final_dataset = Dataset.from_dict({\"tokens\": final_batches}, features=features)\n",
    "\n",
    "    final_dataset.set_format(type=\"torch\", columns=[\"tokens\"])\n",
    "    \n",
    "    return final_dataset\n",
    "\n",
    "# # Example usage\n",
    "# tokenizer = AutoTokenizer.from_pretrained(\"gpt2\")  # Make sure the tokenizer has bos_token_id and eos_token_id\n",
    "# text_dataset = Dataset.from_dict({\"text\": [\"This is a sample text.\", \"Another sample text.\"]})  # Example dataset\n",
    "# tokenized_dataset = tokenize_and_concatenate2(text_dataset, tokenizer, max_length=1024)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"The Rock is destined to be the 21st Century's new ''Conan'' and that he's going to make a splash even greater than Arnold Schwarzenegger, Jean-Claud Van Damme or Steven Segal.\""
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sst_data['train'][0]['text']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 96,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3e3711aceeda4e888e63d7e4e01d810d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/7864 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1f360faa1ae1430998946a544a786db9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/7864 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0bd714397c7148a49071d041203c2e1e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Filter:   0%|          | 0/7864 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from datasets import load_dataset\n",
    "from torch.utils.data import DataLoader\n",
    "from transformers import AutoTokenizer\n",
    "\n",
    "\n",
    "# Define the batch size\n",
    "BATCH_SIZE = 1\n",
    "\n",
    "# Load a tokenizer (you'll need to specify the appropriate model)\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"EleutherAI/pythia-2.8b\")\n",
    "# set padding token\n",
    "tokenizer.pad_token = tokenizer.eos_token\n",
    "\n",
    "#dataset = text_dataset.map(lambda x: tokenize_and_concatenate(x, tokenizer))\n",
    "\n",
    "def tokenize_function(examples):\n",
    "    return tokenizer(examples[\"text\"], padding=\"max_length\", truncation=True, max_length=32)\n",
    "\n",
    "\n",
    "def find_dataset_positions(example, token_id=13):\n",
    "    positions = []\n",
    "    has_token = False\n",
    "    for position, token in enumerate(example['tokens']):\n",
    "        if token==token_id:\n",
    "            has_token = True\n",
    "            positions.append(position)\n",
    "    return {'positions': positions, 'has_token': has_token}\n",
    "\n",
    "\n",
    "dataset = sst_data['train'].map(tokenize_function, batched=False)\n",
    "dataset = dataset.rename_column(\"input_ids\", \"tokens\")\n",
    "dataset = dataset.map(find_dataset_positions, batched=False)\n",
    "dataset = dataset.filter(lambda example: example['has_token'])\n",
    "\n",
    "# Create a PyTorch DataLoader\n",
    "train_data_loader = DataLoader(\n",
    "    dataset, batch_size=BATCH_SIZE, shuffle=False, drop_last=True\n",
    ")\n",
    "\n",
    "# Now, train_data_loader is ready to be used for training or analysis\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 97,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Dataset({\n",
       "    features: ['text', 'label', 'tokens', 'attention_mask', 'positions', 'has_token'],\n",
       "    num_rows: 4119\n",
       "})"
      ]
     },
     "execution_count": 97,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 98,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Singer\\\\/composer Bryan Adams contributes a slew of songs--a few potential hits, a few more simply intrusive to the story--but the whole package certainly captures the intended, er, spirit of the piece.'"
      ]
     },
     "execution_count": 98,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset[0]['text']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 99,
   "metadata": {},
   "outputs": [],
   "source": [
    "# comma_mean_values = get_layerwise_token_mean_activations(model, data_loader, token_id=13)\n",
    "# save_array(comma_mean_values, 'comma_mean_values.npy', model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 100,
   "metadata": {},
   "outputs": [],
   "source": [
    "# period_mean_values = get_layerwise_token_mean_activations(model, train_data_loader, token_id=15)\n",
    "# save_array(period_mean_values, 'period_mean_values.npy', model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 101,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load the files\n",
    "comma_mean_values = torch.from_numpy(load_array('comma_mean_values.npy', model)).to(device)\n",
    "period_mean_values = torch.from_numpy(load_array('period_mean_values.npy', model)).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 139,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "100"
      ]
     },
     "execution_count": 139,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import random\n",
    "\n",
    "# Generate random indices\n",
    "n = 100  # Size of the random subset\n",
    "total_size = len(dataset)\n",
    "random_indices = random.sample(range(total_size), n)\n",
    "\n",
    "# Get a random subset of the dataset of a given size\n",
    "subset_dataset = dataset.select(random_indices)\n",
    "\n",
    "# Create a new dataloader from the subset, converting the data to tensors\n",
    "subset_data_loader = DataLoader(\n",
    "    subset_dataset, batch_size=1, shuffle=False, drop_last=True\n",
    ")\n",
    "len(subset_data_loader)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 140,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "462cef5931e54b47a83b65b60af59b96",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/100 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "punct_pos: [tensor([14]), tensor([19])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([14])\n",
      "zeroing tensor([19])\n",
      "punct_pos: [tensor([3]), tensor([9]), tensor([19]), tensor([23])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([3])\n",
      "zeroing tensor([9])\n",
      "zeroing tensor([19])\n",
      "zeroing tensor([23])\n",
      "punct_pos: [tensor([12]), tensor([17])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([12])\n",
      "zeroing tensor([17])\n",
      "punct_pos: [tensor([7])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([7])\n",
      "punct_pos: [tensor([5])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([5])\n",
      "punct_pos: [tensor([5])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([5])\n",
      "punct_pos: [tensor([9])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([9])\n",
      "punct_pos: [tensor([9]), tensor([15])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([9])\n",
      "zeroing tensor([15])\n",
      "punct_pos: [tensor([14])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([14])\n",
      "punct_pos: [tensor([14]), tensor([17]), tensor([19])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([14])\n",
      "zeroing tensor([17])\n",
      "zeroing tensor([19])\n",
      "punct_pos: [tensor([4])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([4])\n",
      "punct_pos: [tensor([9])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([9])\n",
      "punct_pos: [tensor([15])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([15])\n",
      "punct_pos: [tensor([19])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([19])\n",
      "punct_pos: [tensor([13]), tensor([17])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([13])\n",
      "zeroing tensor([17])\n",
      "punct_pos: [tensor([4])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([4])\n",
      "punct_pos: [tensor([3]), tensor([20])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([3])\n",
      "zeroing tensor([20])\n",
      "punct_pos: [tensor([4])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([4])\n",
      "punct_pos: [tensor([5]), tensor([11]), tensor([20]), tensor([24])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([5])\n",
      "zeroing tensor([11])\n",
      "zeroing tensor([20])\n",
      "zeroing tensor([24])\n",
      "punct_pos: [tensor([3])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([3])\n",
      "punct_pos: [tensor([11]), tensor([17])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([11])\n",
      "zeroing tensor([17])\n",
      "punct_pos: [tensor([3])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([3])\n",
      "punct_pos: [tensor([19]), tensor([24])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([19])\n",
      "zeroing tensor([24])\n",
      "punct_pos: [tensor([17])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([17])\n",
      "punct_pos: [tensor([2])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([2])\n",
      "punct_pos: [tensor([18]), tensor([21])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([18])\n",
      "zeroing tensor([21])\n",
      "punct_pos: [tensor([17])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([17])\n",
      "punct_pos: [tensor([16])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([16])\n",
      "punct_pos: [tensor([13])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([13])\n",
      "punct_pos: [tensor([7]), tensor([16])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([7])\n",
      "zeroing tensor([16])\n",
      "punct_pos: [tensor([12])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([12])\n",
      "punct_pos: [tensor([11]), tensor([14])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([11])\n",
      "zeroing tensor([14])\n",
      "punct_pos: [tensor([12])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([12])\n",
      "punct_pos: [tensor([10]), tensor([27])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([10])\n",
      "zeroing tensor([27])\n",
      "punct_pos: [tensor([2]), tensor([16])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([2])\n",
      "zeroing tensor([16])\n",
      "punct_pos: [tensor([6])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([6])\n",
      "punct_pos: [tensor([5]), tensor([21])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([5])\n",
      "zeroing tensor([21])\n",
      "punct_pos: [tensor([6])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([6])\n",
      "punct_pos: [tensor([9]), tensor([13]), tensor([27])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([9])\n",
      "zeroing tensor([13])\n",
      "zeroing tensor([27])\n",
      "punct_pos: [tensor([9]), tensor([26])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([9])\n",
      "zeroing tensor([26])\n",
      "punct_pos: [tensor([7])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([7])\n",
      "punct_pos: [tensor([4])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([4])\n",
      "punct_pos: [tensor([14])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([14])\n",
      "punct_pos: [tensor([9])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([9])\n",
      "punct_pos: [tensor([17])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([17])\n",
      "punct_pos: [tensor([12]), tensor([16]), tensor([21]), tensor([24])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([12])\n",
      "zeroing tensor([16])\n",
      "zeroing tensor([21])\n",
      "zeroing tensor([24])\n",
      "punct_pos: [tensor([6]), tensor([16])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([6])\n",
      "zeroing tensor([16])\n",
      "punct_pos: [tensor([9]), tensor([20])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([9])\n",
      "zeroing tensor([20])\n",
      "punct_pos: [tensor([2]), tensor([13]), tensor([29])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([2])\n",
      "zeroing tensor([13])\n",
      "zeroing tensor([29])\n",
      "punct_pos: [tensor([14])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([14])\n",
      "punct_pos: [tensor([10])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([10])\n",
      "punct_pos: [tensor([14])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([14])\n",
      "punct_pos: [tensor([10])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([10])\n",
      "punct_pos: [tensor([5]), tensor([8])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([5])\n",
      "zeroing tensor([8])\n",
      "punct_pos: [tensor([8])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([8])\n",
      "punct_pos: [tensor([20])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([20])\n",
      "punct_pos: [tensor([14])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([14])\n",
      "punct_pos: [tensor([11]), tensor([19])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([11])\n",
      "zeroing tensor([19])\n",
      "punct_pos: [tensor([8])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([8])\n",
      "punct_pos: [tensor([21]), tensor([23])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([21])\n",
      "zeroing tensor([23])\n",
      "punct_pos: [tensor([1]), tensor([5]), tensor([23])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([1])\n",
      "zeroing tensor([5])\n",
      "zeroing tensor([23])\n",
      "punct_pos: [tensor([1])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([1])\n",
      "punct_pos: [tensor([5])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([5])\n",
      "punct_pos: [tensor([20])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([20])\n",
      "punct_pos: [tensor([14])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([14])\n",
      "punct_pos: [tensor([10]), tensor([18])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([10])\n",
      "zeroing tensor([18])\n",
      "punct_pos: [tensor([4])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([4])\n",
      "punct_pos: [tensor([9])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([9])\n",
      "punct_pos: [tensor([4])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([4])\n",
      "punct_pos: [tensor([4]), tensor([7]), tensor([11]), tensor([14])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([4])\n",
      "zeroing tensor([7])\n",
      "zeroing tensor([11])\n",
      "zeroing tensor([14])\n",
      "punct_pos: [tensor([12])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([12])\n",
      "punct_pos: [tensor([2]), tensor([5])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([2])\n",
      "zeroing tensor([5])\n",
      "punct_pos: [tensor([3]), tensor([6])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([3])\n",
      "zeroing tensor([6])\n",
      "punct_pos: [tensor([13])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([13])\n",
      "punct_pos: [tensor([3]), tensor([7]), tensor([12])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([3])\n",
      "zeroing tensor([7])\n",
      "zeroing tensor([12])\n",
      "punct_pos: [tensor([9]), tensor([19]), tensor([27])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([9])\n",
      "zeroing tensor([19])\n",
      "zeroing tensor([27])\n",
      "punct_pos: [tensor([25])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([25])\n",
      "punct_pos: [tensor([2]), tensor([9])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([2])\n",
      "zeroing tensor([9])\n",
      "punct_pos: [tensor([4])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([4])\n",
      "punct_pos: [tensor([9]), tensor([19])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([9])\n",
      "zeroing tensor([19])\n",
      "punct_pos: [tensor([1])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([1])\n",
      "punct_pos: [tensor([17])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([17])\n",
      "punct_pos: [tensor([18])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([18])\n",
      "punct_pos: [tensor([16])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([16])\n",
      "punct_pos: [tensor([4]), tensor([13]), tensor([17])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([4])\n",
      "zeroing tensor([13])\n",
      "zeroing tensor([17])\n",
      "punct_pos: [tensor([14]), tensor([19])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([14])\n",
      "zeroing tensor([19])\n",
      "punct_pos: [tensor([12])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([12])\n",
      "punct_pos: [tensor([20])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([20])\n",
      "punct_pos: [tensor([2])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([2])\n",
      "punct_pos: [tensor([18])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([18])\n",
      "punct_pos: [tensor([12])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([12])\n",
      "punct_pos: [tensor([5])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([5])\n",
      "punct_pos: [tensor([19]), tensor([25]), tensor([29])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([19])\n",
      "zeroing tensor([25])\n",
      "zeroing tensor([29])\n",
      "punct_pos: [tensor([13])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([13])\n",
      "punct_pos: [tensor([12])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([12])\n",
      "punct_pos: [tensor([6]), tensor([19])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([6])\n",
      "zeroing tensor([19])\n",
      "punct_pos: [tensor([25])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([25])\n",
      "punct_pos: [tensor([7])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([7])\n",
      "punct_pos: [tensor([6]), tensor([10])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([6])\n",
      "zeroing tensor([10])\n",
      "punct_pos: [tensor([5])]\n",
      "initial loss shape: torch.Size([1, 31])\n",
      "batch tokens shape: torch.Size([1, 32])\n",
      "hooked loss shape: torch.Size([1, 31])\n",
      "zeroing tensor([5])\n"
     ]
    }
   ],
   "source": [
    "heads_to_ablate = [(layer, head) for layer in range(model.cfg.n_layers) for head in range(model.cfg.n_heads)]\n",
    "loss_change_by_token, orig_loss = compute_mean_ablation_modified_loss(model, subset_data_loader, comma_mean_values, target_token_ids=[13])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 141,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(len(loss_change_by_token)):\n",
    "    # add one column of zeros to the loss change tensor\n",
    "    loss_change_by_token[i] = torch.cat([torch.zeros(loss_change_by_token[i].shape[0], 1).to(device), loss_change_by_token[i]], dim=1)\n",
    "\n",
    "loss_change_by_token = torch.stack(loss_change_by_token).cpu()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 142,
   "metadata": {},
   "outputs": [],
   "source": [
    "subset_dataset_tokens = convert_to_tensors(subset_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 143,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "100"
      ]
     },
     "execution_count": 143,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "subset_data_loader_tkns = DataLoader(\n",
    "    subset_dataset_tokens, batch_size=1, shuffle=False, drop_last=True\n",
    ")\n",
    "len(subset_data_loader_tkns)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 144,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([100, 32, 1])"
      ]
     },
     "execution_count": 144,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "loss_change_by_token_by_row = einops.rearrange(loss_change_by_token, \"batch item token -> (batch item) token\")\n",
    "loss_change_by_token_by_row = loss_change_by_token_by_row.unsqueeze(2)\n",
    "loss_change_by_token_by_row.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 145,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'data/pythia-2.8b/loss_change_by_token_by_row_sst.npy'"
      ]
     },
     "execution_count": 145,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "save_array(loss_change_by_token_by_row, 'loss_change_by_token_by_row_sst.npy', model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 146,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(8.0657)"
      ]
     },
     "execution_count": 146,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "loss_change_by_token_by_row.sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 153,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 0.0000],\n",
       "        [ 0.0000],\n",
       "        [ 0.0000],\n",
       "        [ 0.0000],\n",
       "        [ 0.0000],\n",
       "        [ 0.4284],\n",
       "        [ 0.4954],\n",
       "        [ 0.3777],\n",
       "        [ 0.7570],\n",
       "        [ 0.6130],\n",
       "        [ 0.0000],\n",
       "        [ 0.4558],\n",
       "        [ 0.1512],\n",
       "        [-0.5029],\n",
       "        [-0.4370],\n",
       "        [ 0.0000],\n",
       "        [ 0.5531],\n",
       "        [-0.4186],\n",
       "        [ 0.0285],\n",
       "        [-0.0923],\n",
       "        [ 0.0000],\n",
       "        [ 0.0777],\n",
       "        [ 0.0282],\n",
       "        [ 0.0304],\n",
       "        [ 0.0000],\n",
       "        [ 0.3828],\n",
       "        [-0.1881],\n",
       "        [-0.0471],\n",
       "        [-0.1300],\n",
       "        [-0.0969],\n",
       "        [ 0.0932],\n",
       "        [ 0.0637]])"
      ]
     },
     "execution_count": 153,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "loss_change_by_token_by_row[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 150,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Top 40 most positive examples:\n",
      "Example:  Yas, Activation: 2.0176, Batch: 76, Pos: 31\n",
      "Example: ., Activation: 1.3744, Batch: 33, Pos: 29\n",
      "Example:  II, Activation: 1.2075, Batch: 7, Pos: 28\n",
      "Example: ,, Activation: 1.1415, Batch: 14, Pos: 17\n",
      "Example:  weight, Activation: 0.8663, Batch: 18, Pos: 17\n",
      "Example:  roles, Activation: 0.7570, Batch: 1, Pos: 8\n",
      "Example:  echoes, Activation: 0.7151, Batch: 2, Pos: 14\n",
      "Example:  buy, Activation: 0.6960, Batch: 35, Pos: 9\n",
      "Example: 'm, Activation: 0.6842, Batch: 48, Pos: 31\n",
      "Example: .'', Activation: 0.6383, Batch: 10, Pos: 12\n",
      "Example: ., Activation: 0.6022, Batch: 25, Pos: 26\n",
      "Example: Min, Activation: 0.5999, Batch: 4, Pos: 7\n",
      "Example: <|endoftext|>, Activation: 0.5333, Batch: 0, Pos: 25\n",
      "Example:  as, Activation: 0.4545, Batch: 39, Pos: 29\n",
      "Example:  it, Activation: 0.3246, Batch: 53, Pos: 23\n",
      "Example:  string, Activation: 0.2914, Batch: 3, Pos: 12\n",
      "Example:  the, Activation: 0.2676, Batch: 27, Pos: 26\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div id=\"circuits-vis-def83de9-eb91\" style=\"margin: 15px 0;\"/>\n",
       "    <script crossorigin type=\"module\">\n",
       "    import { render, TextNeuronActivations } from \"https://unpkg.com/circuitsvis@1.41.0/dist/cdn/esm.js\";\n",
       "    render(\n",
       "      \"circuits-vis-def83de9-eb91\",\n",
       "      TextNeuronActivations,\n",
       "      {\"tokens\": [\"<|endoftext|>\", \"Ts\", \"ai\", \" has\", \" managed\", \" to\", \" create\", \" an\", \" under\", \"played\", \" mel\", \"od\", \"r\", \"ama\", \" about\", \" family\", \" dynamics\", \" and\", \" dysfunction\", \" that\", \" har\", \"ks\", \" back\", \" to\", \" the\", \" spare\", \",\", \" un\", \"checked\", \" heart\", \"ache\", \" of\", \" Yas\", \"\\n\", \"What\", \" sets\", \" it\", \" apart\", \" is\", \" the\", \" vision\", \" that\", \" Tay\", \"mor\", \",\", \" the\", \" avant\", \" gard\", \"e\", \" director\", \" of\", \" Broadway\", \"'s\", \" The\", \" Lion\", \" King\", \" and\", \" the\", \" film\", \" T\", \"itus\", \",\", \" brings\", \".\", \"<|endoftext|>\", \"<|endoftext|>\", \"\\n\", \"The\", \" problem\", \" with\", \" the\", \" film\", \" is\", \" whether\", \" these\", \" ambitions\", \",\", \" la\", \"ud\", \"able\", \" in\", \" themselves\", \",\", \" justify\", \" a\", \" theatrical\", \" simulation\", \" of\", \" the\", \" death\", \" camp\", \" of\", \" Aus\", \"ch\", \"witz\", \" II\", \"-\", \"Bir\", \"ken\", \"\\n\", \"Like\", \" a\", \" medium\", \"-\", \"grade\", \" network\", \" sit\", \"com\", \"--\", \"mostly\", \" in\", \"off\", \"ensive\", \",\", \" fit\", \"fully\", \" amusing\", \",\", \" but\", \" ultimately\", \" so\", \" weight\", \"less\", \" that\", \" a\", \" decent\", \" draft\", \" in\", \" the\", \" auditor\", \"ium\", \" might\", \"\\n\", \"Although\", \" sensitive\", \" to\", \" a\", \" fault\", \",\", \" it\", \"'s\", \" often\", \" over\", \"written\", \",\", \" with\", \" a\", \" sur\", \"feit\", \" of\", \" weight\", \"y\", \" revelations\", \",\", \" flow\", \"ery\", \" dialogue\", \",\", \" and\", \" nost\", \"algia\", \" for\", \" the\", \" past\", \" and\", \"\\n\", \"The\", \" two\", \" leads\", \",\", \" nearly\", \" perfect\", \" in\", \" their\", \" roles\", \",\", \" bring\", \" a\", \" heart\", \" and\", \" reality\", \" that\", \" buoy\", \" the\", \" film\", \",\", \" and\", \" at\", \" times\", \",\", \" elevate\", \" it\", \" to\", \" a\", \" superior\", \" crime\", \" movie\", \".\", \"\\n\", \"...\", \" a\", \" low\", \" rate\", \" Annie\", \" featuring\", \" some\", \" kid\", \" who\", \" can\", \"'t\", \" act\", \",\", \" only\", \" echoes\", \" of\", \" Jordan\", \",\", \" and\", \" weird\", \"o\", \" actor\", \" Cris\", \"pin\", \" Gl\", \"over\", \" screw\", \"ing\", \" things\", \" up\", \" old\", \" school\", \"\\n\", \"If\", \" I\", \" want\", \" a\", \" real\", \" movie\", \",\", \" I\", \"'ll\", \" buy\", \" the\", \" C\", \"riterion\", \" DVD\", \".\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"\\n\", \"A\", \" boring\", \",\", \" pret\", \"ent\", \"ious\", \" m\", \"uddle\", \" that\", \" uses\", \" a\", \" sens\", \"ational\", \",\", \" real\", \"-\", \"life\", \" 19\", \"th\", \"-\", \"Cent\", \"ury\", \" crime\", \" as\", \" a\", \" metaphor\", \" for\", \"--\", \"well\", \",\", \" I\", \"'m\", \"\\n\", \"If\", \" you\", \"'re\", \" not\", \",\", \" you\", \"'ll\", \" still\", \" have\", \" a\", \" good\", \" time\", \".''\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"\\n\", \"The\", \" aud\", \"acity\", \" to\", \" view\", \" one\", \" of\", \" Shakespeare\", \"'s\", \" better\", \" known\", \" traged\", \"ies\", \" as\", \" a\", \" dark\", \" comedy\", \" is\", \",\", \" by\", \" itself\", \",\", \" des\", \"erving\", \" of\", \" discussion\", \".\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"\\n\", \"As\", \" a\", \" science\", \" fiction\", \" movie\", \",\", \" ''\", \"Min\", \"ority\", \" Report\", \"''\", \" ast\", \"ounds\", \".\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"\\n\", \"N\", \"ett\", \"el\", \"beck\", \" has\", \" crafted\", \" an\", \" engaging\", \" fantasy\", \" of\", \" flav\", \"ours\", \" and\", \" emotions\", \",\", \" one\", \" part\", \" romance\", \" novel\", \",\", \" one\", \" part\", \" recipe\", \" book\", \".\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"\\n\", \"I\", \" stopped\", \" thinking\", \" about\", \" how\", \" good\", \" it\", \" all\", \" was\", \",\", \" and\", \" started\", \" doing\", \" nothing\", \" but\", \" reacting\", \" to\", \" it\", \" -\", \" feeling\", \" a\", \" part\", \" of\", \" its\", \" grand\", \" locations\", \",\", \" thinking\", \" urgently\", \" as\", \" the\", \" protagon\", \"\\n\", \"The\", \" appeal\", \" of\", \" the\", \" vulgar\", \",\", \" sex\", \"ist\", \",\", \" racist\", \" humour\", \" went\", \" over\", \" my\", \" head\", \" or\", \"--\", \"consider\", \"ing\", \" just\", \" how\", \" low\", \" brow\", \" it\", \" is\", \"--\", \"perhaps\", \" it\", \" sn\", \"uck\", \" under\", \" my\", \"\\n\", \"While\", \" the\", \" film\", \" is\", \" not\", \" entirely\", \" successful\", \",\", \" it\", \" still\", \" manages\", \" to\", \" string\", \" together\", \" enough\", \" charming\", \" moments\", \" to\", \" work\", \".\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"\\n\", \"Mr\", \".\", \" W\", \"oll\", \"ter\", \" and\", \" Ms\", \".\", \" S\", \"eld\", \"hal\", \" give\", \" strong\", \" and\", \" convincing\", \" performances\", \",\", \" but\", \" neither\", \" reaches\", \" into\", \" the\", \" deepest\", \" recess\", \"es\", \" of\", \" the\", \" character\", \" to\", \" un\", \"earth\", \" the\", \"\\n\"], \"activations\": [[[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[-0.19274234771728516]], [[0.0]], [[2.0175647735595703]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[-0.20722222328186035]], [[0.0]], [[0.0]], [[-0.2772040367126465]], [[0.893836259841919]], [[0.0]], [[1.3744421005249023]], [[0.05347633361816406]], [[0.021844863891601562]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.48612093925476074]], [[0.0]], [[0.38088035583496094]], [[0.0]], [[0.4289255142211914]], [[0.01771068572998047]], [[-0.27133989334106445]], [[0.0]], [[-0.10186576843261719]], [[-0.06450269371271133]], [[-0.0004067039117217064]], [[1.2074971199035645]], [[0.23783433437347412]], [[0.049881041049957275]], [[-0.0023308824747800827]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[1.1414662599563599]], [[0.0]], [[-0.9171280860900879]], [[0.0]], [[0.05997943878173828]], [[-0.2600143551826477]], [[-0.42669153213500977]], [[0.0]], [[-0.12856006622314453]], [[0.08962059020996094]], [[-0.13900971412658691]], [[-0.0025559663772583008]], [[0.025457382202148438]], [[-0.008208153769373894]], [[-0.21365785598754883]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.8663253784179688]], [[0.0]], [[0.3511333465576172]], [[0.0]], [[0.0]], [[-0.34648358821868896]], [[-0.35437679290771484]], [[0.0]], [[0.0]], [[-0.011941909790039062]], [[0.2753291130065918]], [[-0.1596698760986328]], [[-0.0003930330276489258]], [[0.3423302173614502]], [[-0.06328916549682617]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.4283900260925293]], [[0.4953716993331909]], [[0.37772059440612793]], [[0.7569708824157715]], [[0.6129941344261169]], [[0.0]], [[0.45584189891815186]], [[0.15118169784545898]], [[-0.5029170513153076]], [[-0.436953067779541]], [[0.0]], [[0.553135871887207]], [[-0.4185950756072998]], [[0.02854180335998535]], [[-0.09227514266967773]], [[0.0]], [[0.07766056060791016]], [[0.028180718421936035]], [[0.030414223670959473]], [[0.0]], [[0.3828401565551758]], [[-0.18807339668273926]], [[-0.04705369472503662]], [[-0.12997913360595703]], [[-0.09692716598510742]], [[0.09320187568664551]], [[0.06368634104728699]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.715144157409668]], [[0.0]], [[0.2103872299194336]], [[0.20577597618103027]], [[0.0]], [[-0.010138511657714844]], [[0.0]], [[0.2643156051635742]], [[-0.5618381500244141]], [[-0.13305020332336426]], [[0.0]], [[-0.2560608386993408]], [[0.24546241760253906]], [[-0.13008368015289307]], [[-0.4933905601501465]], [[-0.2784861922264099]], [[0.19512557983398438]], [[-0.2857980728149414]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.696044921875]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.005031585693359375]], [[0.0]], [[0.0]], [[0.008228302001953125]], [[0.008072853088378906]], [[0.0]], [[0.009116172790527344]], [[0.006091117858886719]], [[0.0050411224365234375]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.684215784072876]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[-0.4266018867492676]], [[0.0]], [[-0.1007765531539917]], [[0.0]], [[0.08285665512084961]], [[0.6382961273193359]], [[0.0]], [[0.009791374206542969]], [[0.0]], [[0.0]], [[0.0025320053100585938]], [[0.0]], [[0.0001888275146484375]], [[0.0]], [[0.00039768218994140625]], [[0.0001239776611328125]], [[0.000713348388671875]], [[0.0]], [[0.0014286041259765625]], [[0.0013103485107421875]], [[0.00066375732421875]], [[-0.0006532669067382812]], [[-0.0006895065307617188]], [[0.00017070770263671875]], [[2.47955322265625e-05]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[-0.3078150749206543]], [[0.0]], [[0.0]], [[0.6022284030914307]], [[-0.16152381896972656]], [[0.035714149475097656]], [[0.020914077758789062]], [[0.014864921569824219]], [[0.014507293701171875]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.599884033203125]], [[0.0]], [[-0.7714405059814453]], [[0.0]], [[0.44145870208740234]], [[-0.5480990409851074]], [[0.0]], [[0.3767681121826172]], [[0.0]], [[-0.002948760986328125]], [[0.00031948089599609375]], [[0.0]], [[0.0003490447998046875]], [[0.0]], [[-0.00018215179443359375]], [[0.0011615753173828125]], [[0.0018367767333984375]], [[0.0]], [[0.0010967254638671875]], [[-0.0009307861328125]], [[-0.0007410049438476562]], [[0.0005178451538085938]], [[0.0008134841918945312]], [[0.0011816024780273438]], [[0.0012922286987304688]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[-1.33705472946167]], [[-0.06672334671020508]], [[-1.0656061172485352]], [[-0.07907816767692566]], [[0.0]], [[0.033959269523620605]], [[-0.27968692779541016]], [[0.054088741540908813]], [[-0.484147310256958]], [[0.533322811126709]], [[0.07065200805664062]], [[0.05722332000732422]], [[0.05430793762207031]], [[0.05743598937988281]], [[0.060379981994628906]], [[0.061964988708496094]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.4545450210571289]], [[0.17342853546142578]], [[-0.2978391647338867]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.3246316909790039]], [[0.0]], [[0.0]], [[-0.3036928176879883]], [[0.0]], [[0.0]], [[0.12363910675048828]], [[0.0]], [[-0.01715373992919922]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[-0.07805109024047852]], [[0.0]], [[0.02230234444141388]], [[0.29138946533203125]], [[0.0]], [[-0.11905670166015625]], [[0.0]], [[0.09117794036865234]], [[-0.07906484603881836]], [[0.0]], [[-0.06593608856201172]], [[0.0]], [[0.009602546691894531]], [[0.0069599151611328125]], [[0.004940032958984375]], [[0.0]], [[0.003753662109375]], [[0.0038690567016601562]], [[0.0034427642822265625]], [[0.0030078887939453125]], [[0.00290679931640625]], [[0.0030012130737304688]], [[0.0031442642211914062]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.005962371826171875]], [[0.0]], [[0.0]], [[0.2676198482513428]], [[0.15470027923583984]], [[-0.14868831634521484]], [[0.10778427124023438]], [[-0.2531794309616089]], [[-0.015482425689697266]], [[0.0]]], \"firstDimensionName\": \"Layer (resid_pre)\", \"secondDimensionName\": \"Model\", \"secondDimensionLabels\": [\"pythia-2.8b\"]}\n",
       "    )\n",
       "    </script>"
      ],
      "text/plain": [
       "<circuitsvis.utils.render.RenderedHTML at 0x7fc2a86f4c10>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Top 40 most negative examples:\n",
      "Example:  part, Activation: -1.3371, Batch: 0, Pos: 16\n",
      "Example:  ultimately, Activation: -0.9171, Batch: 14, Pos: 19\n",
      "Example:  Report, Activation: -0.7714, Batch: 4, Pos: 9\n",
      "Example:  Cris, Activation: -0.5618, Batch: 2, Pos: 22\n",
      "Example:  and, Activation: -0.5029, Batch: 1, Pos: 13\n",
      "Example:  still, Activation: -0.4266, Batch: 10, Pos: 7\n",
      "Example:  seem, Activation: -0.3844, Batch: 24, Pos: 7\n",
      "Example:  dialogue, Activation: -0.3544, Batch: 18, Pos: 23\n",
      "Example: erving, Activation: -0.3078, Batch: 25, Pos: 23\n",
      "Example: perhaps, Activation: -0.3037, Batch: 53, Pos: 26\n",
      "Example:  protagon, Activation: -0.2978, Batch: 39, Pos: 31\n",
      "Example: itus, Activation: -0.2772, Batch: 33, Pos: 26\n",
      "Example:  camp, Activation: -0.2713, Batch: 7, Pos: 23\n",
      "Example: earth, Activation: -0.2532, Batch: 27, Pos: 30\n",
      "Example: ache, Activation: -0.1927, Batch: 76, Pos: 29\n",
      "Example:  the, Activation: -0.1866, Batch: 16, Pos: 23\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div id=\"circuits-vis-f88bb3a6-34ea\" style=\"margin: 15px 0;\"/>\n",
       "    <script crossorigin type=\"module\">\n",
       "    import { render, TextNeuronActivations } from \"https://unpkg.com/circuitsvis@1.41.0/dist/cdn/esm.js\";\n",
       "    render(\n",
       "      \"circuits-vis-f88bb3a6-34ea\",\n",
       "      TextNeuronActivations,\n",
       "      {\"tokens\": [\"<|endoftext|>\", \"N\", \"ett\", \"el\", \"beck\", \" has\", \" crafted\", \" an\", \" engaging\", \" fantasy\", \" of\", \" flav\", \"ours\", \" and\", \" emotions\", \",\", \" one\", \" part\", \" romance\", \" novel\", \",\", \" one\", \" part\", \" recipe\", \" book\", \".\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"\\n\", \"Like\", \" a\", \" medium\", \"-\", \"grade\", \" network\", \" sit\", \"com\", \"--\", \"mostly\", \" in\", \"off\", \"ensive\", \",\", \" fit\", \"fully\", \" amusing\", \",\", \" but\", \" ultimately\", \" so\", \" weight\", \"less\", \" that\", \" a\", \" decent\", \" draft\", \" in\", \" the\", \" auditor\", \"ium\", \" might\", \"\\n\", \"As\", \" a\", \" science\", \" fiction\", \" movie\", \",\", \" ''\", \"Min\", \"ority\", \" Report\", \"''\", \" ast\", \"ounds\", \".\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"\\n\", \"...\", \" a\", \" low\", \" rate\", \" Annie\", \" featuring\", \" some\", \" kid\", \" who\", \" can\", \"'t\", \" act\", \",\", \" only\", \" echoes\", \" of\", \" Jordan\", \",\", \" and\", \" weird\", \"o\", \" actor\", \" Cris\", \"pin\", \" Gl\", \"over\", \" screw\", \"ing\", \" things\", \" up\", \" old\", \" school\", \"\\n\", \"The\", \" two\", \" leads\", \",\", \" nearly\", \" perfect\", \" in\", \" their\", \" roles\", \",\", \" bring\", \" a\", \" heart\", \" and\", \" reality\", \" that\", \" buoy\", \" the\", \" film\", \",\", \" and\", \" at\", \" times\", \",\", \" elevate\", \" it\", \" to\", \" a\", \" superior\", \" crime\", \" movie\", \".\", \"\\n\", \"If\", \" you\", \"'re\", \" not\", \",\", \" you\", \"'ll\", \" still\", \" have\", \" a\", \" good\", \" time\", \".''\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"\\n\", \"In\", \" fact\", \",\", \" it\", \" doesn\", \"'t\", \" even\", \" seem\", \" like\", \" she\", \" tried\", \".\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"\\n\", \"Although\", \" sensitive\", \" to\", \" a\", \" fault\", \",\", \" it\", \"'s\", \" often\", \" over\", \"written\", \",\", \" with\", \" a\", \" sur\", \"feit\", \" of\", \" weight\", \"y\", \" revelations\", \",\", \" flow\", \"ery\", \" dialogue\", \",\", \" and\", \" nost\", \"algia\", \" for\", \" the\", \" past\", \" and\", \"\\n\", \"The\", \" aud\", \"acity\", \" to\", \" view\", \" one\", \" of\", \" Shakespeare\", \"'s\", \" better\", \" known\", \" traged\", \"ies\", \" as\", \" a\", \" dark\", \" comedy\", \" is\", \",\", \" by\", \" itself\", \",\", \" des\", \"erving\", \" of\", \" discussion\", \".\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"<|endoftext|>\", \"\\n\", \"The\", \" appeal\", \" of\", \" the\", \" vulgar\", \",\", \" sex\", \"ist\", \",\", \" racist\", \" humour\", \" went\", \" over\", \" my\", \" head\", \" or\", \"--\", \"consider\", \"ing\", \" just\", \" how\", \" low\", \" brow\", \" it\", \" is\", \"--\", \"perhaps\", \" it\", \" sn\", \"uck\", \" under\", \" my\", \"\\n\", \"I\", \" stopped\", \" thinking\", \" about\", \" how\", \" good\", \" it\", \" all\", \" was\", \",\", \" and\", \" started\", \" doing\", \" nothing\", \" but\", \" reacting\", \" to\", \" it\", \" -\", \" feeling\", \" a\", \" part\", \" of\", \" its\", \" grand\", \" locations\", \",\", \" thinking\", \" urgently\", \" as\", \" the\", \" protagon\", \"\\n\", \"What\", \" sets\", \" it\", \" apart\", \" is\", \" the\", \" vision\", \" that\", \" Tay\", \"mor\", \",\", \" the\", \" avant\", \" gard\", \"e\", \" director\", \" of\", \" Broadway\", \"'s\", \" The\", \" Lion\", \" King\", \" and\", \" the\", \" film\", \" T\", \"itus\", \",\", \" brings\", \".\", \"<|endoftext|>\", \"<|endoftext|>\", \"\\n\", \"The\", \" problem\", \" with\", \" the\", \" film\", \" is\", \" whether\", \" these\", \" ambitions\", \",\", \" la\", \"ud\", \"able\", \" in\", \" themselves\", \",\", \" justify\", \" a\", \" theatrical\", \" simulation\", \" of\", \" the\", \" death\", \" camp\", \" of\", \" Aus\", \"ch\", \"witz\", \" II\", \"-\", \"Bir\", \"ken\", \"\\n\", \"Mr\", \".\", \" W\", \"oll\", \"ter\", \" and\", \" Ms\", \".\", \" S\", \"eld\", \"hal\", \" give\", \" strong\", \" and\", \" convincing\", \" performances\", \",\", \" but\", \" neither\", \" reaches\", \" into\", \" the\", \" deepest\", \" recess\", \"es\", \" of\", \" the\", \" character\", \" to\", \" un\", \"earth\", \" the\", \"\\n\", \"Ts\", \"ai\", \" has\", \" managed\", \" to\", \" create\", \" an\", \" under\", \"played\", \" mel\", \"od\", \"r\", \"ama\", \" about\", \" family\", \" dynamics\", \" and\", \" dysfunction\", \" that\", \" har\", \"ks\", \" back\", \" to\", \" the\", \" spare\", \",\", \" un\", \"checked\", \" heart\", \"ache\", \" of\", \" Yas\", \"\\n\", \"Like\", \" Sh\", \"rek\", \",\", \" Spirit\", \"'s\", \" visual\", \" imagination\", \" reminds\", \" you\", \" of\", \" why\", \" animation\", \" is\", \" such\", \" a\", \" perfect\", \" medium\", \" for\", \" children\", \",\", \" because\", \" of\", \" the\", \" way\", \" it\", \" allows\", \" the\", \" mind\", \" to\", \" enter\", \" and\", \"\\n\"], \"activations\": [[[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[-1.33705472946167]], [[-0.06672334671020508]], [[-1.0656061172485352]], [[-0.07907816767692566]], [[0.0]], [[0.033959269523620605]], [[-0.27968692779541016]], [[0.054088741540908813]], [[-0.484147310256958]], [[0.533322811126709]], [[0.07065200805664062]], [[0.05722332000732422]], [[0.05430793762207031]], [[0.05743598937988281]], [[0.060379981994628906]], [[0.061964988708496094]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[1.1414662599563599]], [[0.0]], [[-0.9171280860900879]], [[0.0]], [[0.05997943878173828]], [[-0.2600143551826477]], [[-0.42669153213500977]], [[0.0]], [[-0.12856006622314453]], [[0.08962059020996094]], [[-0.13900971412658691]], [[-0.0025559663772583008]], [[0.025457382202148438]], [[-0.008208153769373894]], [[-0.21365785598754883]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.599884033203125]], [[0.0]], [[-0.7714405059814453]], [[0.0]], [[0.44145870208740234]], [[-0.5480990409851074]], [[0.0]], [[0.3767681121826172]], [[0.0]], [[-0.002948760986328125]], [[0.00031948089599609375]], [[0.0]], [[0.0003490447998046875]], [[0.0]], [[-0.00018215179443359375]], [[0.0011615753173828125]], [[0.0018367767333984375]], [[0.0]], [[0.0010967254638671875]], [[-0.0009307861328125]], [[-0.0007410049438476562]], [[0.0005178451538085938]], [[0.0008134841918945312]], [[0.0011816024780273438]], [[0.0012922286987304688]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.715144157409668]], [[0.0]], [[0.2103872299194336]], [[0.20577597618103027]], [[0.0]], [[-0.010138511657714844]], [[0.0]], [[0.2643156051635742]], [[-0.5618381500244141]], [[-0.13305020332336426]], [[0.0]], [[-0.2560608386993408]], [[0.24546241760253906]], [[-0.13008368015289307]], [[-0.4933905601501465]], [[-0.2784861922264099]], [[0.19512557983398438]], [[-0.2857980728149414]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.4283900260925293]], [[0.4953716993331909]], [[0.37772059440612793]], [[0.7569708824157715]], [[0.6129941344261169]], [[0.0]], [[0.45584189891815186]], [[0.15118169784545898]], [[-0.5029170513153076]], [[-0.436953067779541]], [[0.0]], [[0.553135871887207]], [[-0.4185950756072998]], [[0.02854180335998535]], [[-0.09227514266967773]], [[0.0]], [[0.07766056060791016]], [[0.028180718421936035]], [[0.030414223670959473]], [[0.0]], [[0.3828401565551758]], [[-0.18807339668273926]], [[-0.04705369472503662]], [[-0.12997913360595703]], [[-0.09692716598510742]], [[0.09320187568664551]], [[0.06368634104728699]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[-0.4266018867492676]], [[0.0]], [[-0.1007765531539917]], [[0.0]], [[0.08285665512084961]], [[0.6382961273193359]], [[0.0]], [[0.009791374206542969]], [[0.0]], [[0.0]], [[0.0025320053100585938]], [[0.0]], [[0.0001888275146484375]], [[0.0]], [[0.00039768218994140625]], [[0.0001239776611328125]], [[0.000713348388671875]], [[0.0]], [[0.0014286041259765625]], [[0.0013103485107421875]], [[0.00066375732421875]], [[-0.0006532669067382812]], [[-0.0006895065307617188]], [[0.00017070770263671875]], [[2.47955322265625e-05]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[-0.3843812942504883]], [[0.0]], [[-0.13643217086791992]], [[0.0]], [[-0.08007335662841797]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[-0.007172584533691406]], [[0.0]], [[-0.00307464599609375]], [[0.0]], [[0.0]], [[-0.0005655288696289062]], [[0.0013132095336914062]], [[0.0]], [[0.0]], [[0.0018339157104492188]], [[0.001819610595703125]], [[0.0017557144165039062]], [[0.0016222000122070312]], [[0.0013027191162109375]], [[0.0008392333984375]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.8663253784179688]], [[0.0]], [[0.3511333465576172]], [[0.0]], [[0.0]], [[-0.34648358821868896]], [[-0.35437679290771484]], [[0.0]], [[0.0]], [[-0.011941909790039062]], [[0.2753291130065918]], [[-0.1596698760986328]], [[-0.0003930330276489258]], [[0.3423302173614502]], [[-0.06328916549682617]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[-0.3078150749206543]], [[0.0]], [[0.0]], [[0.6022284030914307]], [[-0.16152381896972656]], [[0.035714149475097656]], [[0.020914077758789062]], [[0.014864921569824219]], [[0.014507293701171875]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.3246316909790039]], [[0.0]], [[0.0]], [[-0.3036928176879883]], [[0.0]], [[0.0]], [[0.12363910675048828]], [[0.0]], [[-0.01715373992919922]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.4545450210571289]], [[0.17342853546142578]], [[-0.2978391647338867]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[-0.20722222328186035]], [[0.0]], [[0.0]], [[-0.2772040367126465]], [[0.893836259841919]], [[0.0]], [[1.3744421005249023]], [[0.05347633361816406]], [[0.021844863891601562]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.48612093925476074]], [[0.0]], [[0.38088035583496094]], [[0.0]], [[0.4289255142211914]], [[0.01771068572998047]], [[-0.27133989334106445]], [[0.0]], [[-0.10186576843261719]], [[-0.06450269371271133]], [[-0.0004067039117217064]], [[1.2074971199035645]], [[0.23783433437347412]], [[0.049881041049957275]], [[-0.0023308824747800827]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.005962371826171875]], [[0.0]], [[0.0]], [[0.2676198482513428]], [[0.15470027923583984]], [[-0.14868831634521484]], [[0.10778427124023438]], [[-0.2531794309616089]], [[-0.015482425689697266]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[-0.19274234771728516]], [[0.0]], [[2.0175647735595703]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[0.0]], [[-0.037001729011535645]], [[-0.18659627437591553]], [[0.0]], [[-0.1013951301574707]], [[-0.027297019958496094]], [[-0.0771942138671875]], [[0.0826864242553711]], [[-0.07775256037712097]], [[0.013264656066894531]], [[0.0632627010345459]], [[0.0]]], \"firstDimensionName\": \"Layer (resid_pre)\", \"secondDimensionName\": \"Model\", \"secondDimensionLabels\": [\"pythia-2.8b\"]}\n",
       "    )\n",
       "    </script>"
      ],
      "text/plain": [
       "<circuitsvis.utils.render.RenderedHTML at 0x7fc2a80cdd00>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from utils.neuroscope import plot_topk\n",
    "loss_change_by_token = torch.from_numpy(load_array('loss_change_by_token_by_row_sst.npy', model))\n",
    "plot_topk(\n",
    "    activations=loss_change_by_token_by_row, \n",
    "    dataloader=subset_data_loader_tkns, \n",
    "    window_size=64, \n",
    "    model=model, \n",
    "    k=40, \n",
    "    centred=False, \n",
    "    #exclusions=[\" '\", \" ,\", \",\", \".\",\" .\",\" Fig\", \"'t\", \" Pinterest\", \" Kampf\", \"m\", \"uk\", \" Kamp\", \"com\", \"edu\", \"S\", \"youtube\", \"twitter\", \"0\", \"js\", \"py\", \" Protein\", \" Fiber\", \" Carbohydrates\", \" Sugar\", \" Grant\", \" Pub\", \",\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Comma Ablation for Classification"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "import random\n",
    "\n",
    "from transformers import AutoTokenizer, AutoModelForSequenceClassification, AutoModelForCausalLM\n",
    "from transformers import TrainingArguments, Trainer\n",
    "from transformer_lens import HookedTransformer\n",
    "from datasets import load_dataset, Dataset, DatasetDict\n",
    "from tqdm.notebook import tqdm\n",
    "from utils.store import load_pickle, load_array\n",
    "from utils.ablation import ablate_resid_with_precalc_mean"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "aa26894959d64ac296e1cb2a60ee4322",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading builder script:   0%|          | 0.00/4.31k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "feaf5af2c04446a19ad811c5d51bc8a6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading metadata:   0%|          | 0.00/2.17k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c91197d5949b4f31875956f0771215fe",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading readme:   0%|          | 0.00/7.59k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "842429bbf5f74813b225c25895b01693",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/84.1M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5ace358243414877942fbf540a8f7ca9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Generating train split:   0%|          | 0/25000 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "dd9c8aaf413142a8af17d4571b68f1e3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Generating test split:   0%|          | 0/25000 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "689d325ea6dd4dac97d840fcf80c6fc3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Generating unsupervised split:   0%|          | 0/50000 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "{'text': \"Terrible movie. Nuff Said.<br /><br />These Lines are Just Filler. The movie was bad. Why I have to expand on that I don't know. This is already a waste of my time. I just wanted to warn others. Avoid this movie. The acting sucks and the writing is just moronic. Bad in every way. The only nice thing about the movie are Deniz Akkaya's breasts. Even that was ruined though by a terrible and unneeded rape scene. The movie is a poorly contrived and totally unbelievable piece of garbage.<br /><br />OK now I am just going to rag on IMDb for this stupid rule of 10 lines of text minimum. First I waste my time watching this offal. Then feeling compelled to warn others I create an account with IMDb only to discover that I have to write a friggen essay on the film just to express how bad I think it is. Totally unnecessary.\",\n",
       " 'label': 0}"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset = load_dataset(\"imdb\")\n",
    "dataset[\"train\"][100]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### GPT-2 Classifier"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a1d270f58324445a8d73448eed87a48a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)lve/main/config.json:   0%|          | 0.00/665 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c9ea50eedf2e478e8ccd45d4ccb09dd7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)olve/main/vocab.json:   0%|          | 0.00/1.04M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ad29e22fbc7e4f69a7d82af85339ba5e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)olve/main/merges.txt:   0%|          | 0.00/456k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "97368b04fa5e41508c3114c3942669a4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading (…)/main/tokenizer.json:   0%|          | 0.00/1.36M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8e000b40262f44d791f6868822079f45",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/25000 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8f30f553b1444566aba7d5279bbc97c1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/25000 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "81d933534b0b487d9955dcb66d4077d7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/50000 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "MODEL_NAME = \"gpt2\"\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)\n",
    "\n",
    "tokenizer.pad_token = tokenizer.eos_token\n",
    "\n",
    "def tokenize_function(examples):\n",
    "    return tokenizer(examples[\"text\"], padding=\"max_length\", truncation=True, max_length=512)\n",
    "\n",
    "\n",
    "tokenized_datasets = dataset.map(tokenize_function, batched=True)\n",
    "\n",
    "small_train_dataset = tokenized_datasets[\"train\"].shuffle(seed=42).select(range(8000))\n",
    "small_eval_dataset = tokenized_datasets[\"test\"].shuffle(seed=42).select(range(2000))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using pad_token, but it is not set yet.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded pretrained model gpt2 into HookedTransformer\n"
     ]
    }
   ],
   "source": [
    "model = AutoModelForCausalLM.from_pretrained(\"./gpt2_imdb_classifier\")\n",
    "class_layer_weights = load_pickle(\"gpt2_imdb_classifier_classification_head_weights\", 'gpt2')\n",
    "\n",
    "model = HookedTransformer.from_pretrained(\n",
    "    \"gpt2\",\n",
    "    hf_model=model,\n",
    "    center_unembed=True,\n",
    "    center_writing_weights=True,\n",
    "    fold_ln=True,\n",
    "    refactor_factored_attn_matrices=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_classification_prediction(eval_dataset, dataset_idx, verbose=False):\n",
    "\n",
    "    logits, cache = model.run_with_cache(small_eval_dataset[dataset_idx]['text'])\n",
    "    last_token_act = cache['ln_final.hook_normalized'][0, -1, :]\n",
    "    res = torch.softmax(torch.tensor(class_layer_weights['score.weight']) @ last_token_act.cpu(), dim=-1)\n",
    "    if verbose:\n",
    "        print(f\"Sentence: {small_eval_dataset[dataset_idx]['text']}\")\n",
    "        print(f\"Prediction: {res.argmax()} Label: {small_eval_dataset[dataset_idx]['label']}\")\n",
    "\n",
    "    return res.argmax(), small_eval_dataset[dataset_idx]['label'], res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.89"
      ]
     },
     "execution_count": 47,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def get_accuracy(eval_dataset, n=100):\n",
    "    correct = 0\n",
    "    for idx in range(min(len(eval_dataset), n)):\n",
    "        pred, label, _ = get_classification_prediction(eval_dataset, idx)\n",
    "        if pred == label:\n",
    "            correct += 1\n",
    "    return correct / n\n",
    "\n",
    "get_accuracy(small_eval_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_mean_ablation_modified_accuracy(model: HookedTransformer, dataset: Dataset, target_token_ids) -> float:\n",
    "\n",
    "    for _, item in tqdm(enumerate(dataset), total=len(dataset)):\n",
    "        batch_tokens = model.to_tokens(item['text'], prepend_bos=False)\n",
    "        print(batch_tokens.shape)\n",
    "\n",
    "    #     # get positions of all 13 and 15 token ids in batch\n",
    "    #     punct_pos = find_positions(batch_tokens, token_ids=target_token_ids)\n",
    "\n",
    "    #     # get the loss for each token in the batch\n",
    "    #     initial_loss = model(batch_tokens, return_type=\"loss\", prepend_bos=False, loss_per_token=True)\n",
    "    #     orig_loss_list.append(initial_loss)\n",
    "        \n",
    "    #     # add hooks for the activations of the 13 and 15 tokens\n",
    "    #     for layer, head in heads_to_ablate:\n",
    "    #         mean_ablate_comma = partial(ablate_resid_with_precalc_mean, cached_means=cached_means, pos_by_batch=punct_pos, layer=layer)\n",
    "    #         model.blocks[layer].hook_resid_post.add_hook(mean_ablate_comma)\n",
    "\n",
    "    #     # get the loss for each token when run with hooks\n",
    "    #     hooked_loss = model.run_with_cache(batch_tokens, return_type=\"loss\", prepend_bos=False, loss_per_token=True)\n",
    "\n",
    "    #     # compute the percent difference between the two losses\n",
    "    #     loss_diff = hooked_loss - initial_loss\n",
    "    #     loss_diff_list.append(loss_diff)\n",
    "\n",
    "    # model.reset_hooks()\n",
    "    # return loss_diff_list, orig_loss_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [
    {
     "ename": "FileNotFoundError",
     "evalue": "[Errno 2] No such file or directory: 'data/gpt2-small/comma_mean_values.npy'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mFileNotFoundError\u001b[0m                         Traceback (most recent call last)",
      "\u001b[1;32m/home/curttigges/proj/eliciting-latent-sentiment/owt_comma_ablation.ipynb Cell 42\u001b[0m in \u001b[0;36m2\n\u001b[1;32m      <a href='vscode-notebook-cell:/home/curttigges/proj/eliciting-latent-sentiment/owt_comma_ablation.ipynb#Y252sZmlsZQ%3D%3D?line=0'>1</a>\u001b[0m \u001b[39m# load the files\u001b[39;00m\n\u001b[0;32m----> <a href='vscode-notebook-cell:/home/curttigges/proj/eliciting-latent-sentiment/owt_comma_ablation.ipynb#Y252sZmlsZQ%3D%3D?line=1'>2</a>\u001b[0m comma_mean_values \u001b[39m=\u001b[39m torch\u001b[39m.\u001b[39mfrom_numpy(load_array(\u001b[39m'\u001b[39;49m\u001b[39mcomma_mean_values.npy\u001b[39;49m\u001b[39m'\u001b[39;49m, model))\u001b[39m.\u001b[39mto(device)\n\u001b[1;32m      <a href='vscode-notebook-cell:/home/curttigges/proj/eliciting-latent-sentiment/owt_comma_ablation.ipynb#Y252sZmlsZQ%3D%3D?line=2'>3</a>\u001b[0m period_mean_values \u001b[39m=\u001b[39m torch\u001b[39m.\u001b[39mfrom_numpy(load_array(\u001b[39m'\u001b[39m\u001b[39mperiod_mean_values.npy\u001b[39m\u001b[39m'\u001b[39m, model))\u001b[39m.\u001b[39mto(device)\n",
      "File \u001b[0;32m/notebooks/eliciting-latent-sentiment/utils/store.py:186\u001b[0m, in \u001b[0;36mload_array\u001b[0;34m(label, model)\u001b[0m\n\u001b[1;32m    184\u001b[0m model_path \u001b[39m=\u001b[39m os\u001b[39m.\u001b[39mpath\u001b[39m.\u001b[39mjoin(\u001b[39m'\u001b[39m\u001b[39mdata\u001b[39m\u001b[39m'\u001b[39m, model)\n\u001b[1;32m    185\u001b[0m path \u001b[39m=\u001b[39m os\u001b[39m.\u001b[39mpath\u001b[39m.\u001b[39mjoin(model_path, label \u001b[39m+\u001b[39m \u001b[39m'\u001b[39m\u001b[39m.npy\u001b[39m\u001b[39m'\u001b[39m)\n\u001b[0;32m--> 186\u001b[0m \u001b[39mwith\u001b[39;00m \u001b[39mopen\u001b[39;49m(path, \u001b[39m'\u001b[39;49m\u001b[39mrb\u001b[39;49m\u001b[39m'\u001b[39;49m) \u001b[39mas\u001b[39;00m f:\n\u001b[1;32m    187\u001b[0m     array \u001b[39m=\u001b[39m np\u001b[39m.\u001b[39mload(f)\n\u001b[1;32m    188\u001b[0m \u001b[39mreturn\u001b[39;00m array\n",
      "\u001b[0;31mFileNotFoundError\u001b[0m: [Errno 2] No such file or directory: 'data/gpt2-small/comma_mean_values.npy'"
     ]
    }
   ],
   "source": [
    "# load the files\n",
    "comma_mean_values = torch.from_numpy(load_array('comma_mean_values.npy', model)).to(device)\n",
    "period_mean_values = torch.from_numpy(load_array('period_mean_values.npy', model)).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ea6ec2bc38df4f9d87fcdfd05723f17b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/2000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 413])\n",
      "torch.Size([1, 230])\n",
      "torch.Size([1, 157])\n",
      "torch.Size([1, 139])\n",
      "torch.Size([1, 108])\n",
      "torch.Size([1, 171])\n",
      "torch.Size([1, 147])\n",
      "torch.Size([1, 164])\n",
      "torch.Size([1, 177])\n",
      "torch.Size([1, 158])\n",
      "torch.Size([1, 173])\n",
      "torch.Size([1, 154])\n",
      "torch.Size([1, 274])\n",
      "torch.Size([1, 340])\n",
      "torch.Size([1, 606])\n",
      "torch.Size([1, 232])\n",
      "torch.Size([1, 217])\n",
      "torch.Size([1, 344])\n",
      "torch.Size([1, 193])\n",
      "torch.Size([1, 229])\n",
      "torch.Size([1, 175])\n",
      "torch.Size([1, 581])\n",
      "torch.Size([1, 634])\n",
      "torch.Size([1, 206])\n",
      "torch.Size([1, 161])\n",
      "torch.Size([1, 296])\n",
      "torch.Size([1, 225])\n",
      "torch.Size([1, 163])\n",
      "torch.Size([1, 245])\n",
      "torch.Size([1, 166])\n",
      "torch.Size([1, 438])\n",
      "torch.Size([1, 153])\n",
      "torch.Size([1, 114])\n",
      "torch.Size([1, 94])\n",
      "torch.Size([1, 209])\n",
      "torch.Size([1, 75])\n",
      "torch.Size([1, 156])\n",
      "torch.Size([1, 385])\n",
      "torch.Size([1, 209])\n",
      "torch.Size([1, 666])\n",
      "torch.Size([1, 67])\n",
      "torch.Size([1, 137])\n",
      "torch.Size([1, 423])\n",
      "torch.Size([1, 151])\n",
      "torch.Size([1, 142])\n",
      "torch.Size([1, 208])\n",
      "torch.Size([1, 270])\n",
      "torch.Size([1, 257])\n",
      "torch.Size([1, 349])\n",
      "torch.Size([1, 166])\n",
      "torch.Size([1, 503])\n",
      "torch.Size([1, 207])\n",
      "torch.Size([1, 142])\n",
      "torch.Size([1, 252])\n",
      "torch.Size([1, 146])\n",
      "torch.Size([1, 158])\n",
      "torch.Size([1, 165])\n",
      "torch.Size([1, 181])\n",
      "torch.Size([1, 211])\n",
      "torch.Size([1, 207])\n",
      "torch.Size([1, 446])\n",
      "torch.Size([1, 144])\n",
      "torch.Size([1, 500])\n",
      "torch.Size([1, 123])\n",
      "torch.Size([1, 144])\n",
      "torch.Size([1, 253])\n",
      "torch.Size([1, 256])\n",
      "torch.Size([1, 57])\n",
      "torch.Size([1, 200])\n",
      "torch.Size([1, 165])\n",
      "torch.Size([1, 371])\n",
      "torch.Size([1, 349])\n",
      "torch.Size([1, 191])\n",
      "torch.Size([1, 72])\n",
      "torch.Size([1, 258])\n",
      "torch.Size([1, 69])\n",
      "torch.Size([1, 333])\n",
      "torch.Size([1, 168])\n",
      "torch.Size([1, 295])\n",
      "torch.Size([1, 119])\n",
      "torch.Size([1, 335])\n",
      "torch.Size([1, 362])\n",
      "torch.Size([1, 191])\n",
      "torch.Size([1, 222])\n",
      "torch.Size([1, 170])\n",
      "torch.Size([1, 149])\n",
      "torch.Size([1, 312])\n",
      "torch.Size([1, 156])\n",
      "torch.Size([1, 962])\n",
      "torch.Size([1, 201])\n",
      "torch.Size([1, 385])\n",
      "torch.Size([1, 241])\n",
      "torch.Size([1, 408])\n",
      "torch.Size([1, 136])\n",
      "torch.Size([1, 418])\n",
      "torch.Size([1, 199])\n",
      "torch.Size([1, 245])\n",
      "torch.Size([1, 70])\n",
      "torch.Size([1, 381])\n",
      "torch.Size([1, 196])\n",
      "torch.Size([1, 154])\n",
      "torch.Size([1, 151])\n",
      "torch.Size([1, 226])\n",
      "torch.Size([1, 74])\n",
      "torch.Size([1, 79])\n",
      "torch.Size([1, 188])\n",
      "torch.Size([1, 227])\n",
      "torch.Size([1, 82])\n",
      "torch.Size([1, 58])\n",
      "torch.Size([1, 154])\n",
      "torch.Size([1, 203])\n",
      "torch.Size([1, 358])\n",
      "torch.Size([1, 353])\n",
      "torch.Size([1, 467])\n",
      "torch.Size([1, 147])\n",
      "torch.Size([1, 359])\n",
      "torch.Size([1, 161])\n",
      "torch.Size([1, 1024])\n",
      "torch.Size([1, 185])\n",
      "torch.Size([1, 521])\n",
      "torch.Size([1, 500])\n",
      "torch.Size([1, 134])\n",
      "torch.Size([1, 375])\n",
      "torch.Size([1, 66])\n",
      "torch.Size([1, 164])\n",
      "torch.Size([1, 148])\n",
      "torch.Size([1, 149])\n",
      "torch.Size([1, 797])\n",
      "torch.Size([1, 250])\n",
      "torch.Size([1, 400])\n",
      "torch.Size([1, 183])\n",
      "torch.Size([1, 149])\n",
      "torch.Size([1, 396])\n",
      "torch.Size([1, 367])\n",
      "torch.Size([1, 583])\n",
      "torch.Size([1, 134])\n",
      "torch.Size([1, 1024])\n",
      "torch.Size([1, 578])\n",
      "torch.Size([1, 55])\n",
      "torch.Size([1, 77])\n",
      "torch.Size([1, 82])\n",
      "torch.Size([1, 680])\n",
      "torch.Size([1, 229])\n",
      "torch.Size([1, 166])\n",
      "torch.Size([1, 222])\n",
      "torch.Size([1, 227])\n",
      "torch.Size([1, 132])\n",
      "torch.Size([1, 311])\n",
      "torch.Size([1, 467])\n",
      "torch.Size([1, 143])\n",
      "torch.Size([1, 179])\n",
      "torch.Size([1, 234])\n",
      "torch.Size([1, 117])\n",
      "torch.Size([1, 72])\n",
      "torch.Size([1, 98])\n",
      "torch.Size([1, 464])\n",
      "torch.Size([1, 54])\n",
      "torch.Size([1, 66])\n",
      "torch.Size([1, 441])\n",
      "torch.Size([1, 63])\n",
      "torch.Size([1, 153])\n",
      "torch.Size([1, 174])\n",
      "torch.Size([1, 66])\n",
      "torch.Size([1, 213])\n",
      "torch.Size([1, 291])\n",
      "torch.Size([1, 659])\n",
      "torch.Size([1, 140])\n",
      "torch.Size([1, 180])\n",
      "torch.Size([1, 291])\n",
      "torch.Size([1, 317])\n",
      "torch.Size([1, 189])\n",
      "torch.Size([1, 285])\n",
      "torch.Size([1, 141])\n",
      "torch.Size([1, 245])\n",
      "torch.Size([1, 420])\n",
      "torch.Size([1, 619])\n",
      "torch.Size([1, 215])\n",
      "torch.Size([1, 384])\n",
      "torch.Size([1, 139])\n",
      "torch.Size([1, 165])\n",
      "torch.Size([1, 98])\n",
      "torch.Size([1, 203])\n",
      "torch.Size([1, 147])\n",
      "torch.Size([1, 125])\n",
      "torch.Size([1, 166])\n",
      "torch.Size([1, 173])\n",
      "torch.Size([1, 301])\n",
      "torch.Size([1, 154])\n",
      "torch.Size([1, 180])\n",
      "torch.Size([1, 255])\n",
      "torch.Size([1, 220])\n",
      "torch.Size([1, 185])\n",
      "torch.Size([1, 413])\n",
      "torch.Size([1, 179])\n",
      "torch.Size([1, 239])\n",
      "torch.Size([1, 189])\n",
      "torch.Size([1, 520])\n",
      "torch.Size([1, 119])\n",
      "torch.Size([1, 391])\n",
      "torch.Size([1, 166])\n",
      "torch.Size([1, 134])\n",
      "torch.Size([1, 169])\n",
      "torch.Size([1, 261])\n",
      "torch.Size([1, 238])\n",
      "torch.Size([1, 324])\n",
      "torch.Size([1, 360])\n",
      "torch.Size([1, 460])\n",
      "torch.Size([1, 130])\n",
      "torch.Size([1, 61])\n",
      "torch.Size([1, 143])\n",
      "torch.Size([1, 157])\n",
      "torch.Size([1, 165])\n",
      "torch.Size([1, 562])\n",
      "torch.Size([1, 152])\n",
      "torch.Size([1, 290])\n",
      "torch.Size([1, 387])\n",
      "torch.Size([1, 61])\n",
      "torch.Size([1, 165])\n",
      "torch.Size([1, 206])\n",
      "torch.Size([1, 467])\n",
      "torch.Size([1, 122])\n",
      "torch.Size([1, 155])\n",
      "torch.Size([1, 571])\n",
      "torch.Size([1, 174])\n",
      "torch.Size([1, 182])\n",
      "torch.Size([1, 139])\n",
      "torch.Size([1, 483])\n",
      "torch.Size([1, 84])\n",
      "torch.Size([1, 292])\n",
      "torch.Size([1, 132])\n",
      "torch.Size([1, 67])\n",
      "torch.Size([1, 293])\n",
      "torch.Size([1, 263])\n",
      "torch.Size([1, 96])\n",
      "torch.Size([1, 207])\n",
      "torch.Size([1, 525])\n",
      "torch.Size([1, 391])\n",
      "torch.Size([1, 230])\n",
      "torch.Size([1, 119])\n",
      "torch.Size([1, 367])\n",
      "torch.Size([1, 235])\n",
      "torch.Size([1, 400])\n",
      "torch.Size([1, 341])\n",
      "torch.Size([1, 63])\n",
      "torch.Size([1, 128])\n",
      "torch.Size([1, 433])\n",
      "torch.Size([1, 139])\n",
      "torch.Size([1, 141])\n",
      "torch.Size([1, 214])\n",
      "torch.Size([1, 308])\n",
      "torch.Size([1, 86])\n",
      "torch.Size([1, 140])\n",
      "torch.Size([1, 424])\n",
      "torch.Size([1, 62])\n",
      "torch.Size([1, 75])\n",
      "torch.Size([1, 197])\n",
      "torch.Size([1, 121])\n",
      "torch.Size([1, 151])\n",
      "torch.Size([1, 247])\n",
      "torch.Size([1, 395])\n",
      "torch.Size([1, 319])\n",
      "torch.Size([1, 206])\n",
      "torch.Size([1, 156])\n",
      "torch.Size([1, 147])\n",
      "torch.Size([1, 156])\n",
      "torch.Size([1, 132])\n",
      "torch.Size([1, 551])\n",
      "torch.Size([1, 533])\n",
      "torch.Size([1, 156])\n",
      "torch.Size([1, 276])\n",
      "torch.Size([1, 228])\n",
      "torch.Size([1, 255])\n",
      "torch.Size([1, 48])\n",
      "torch.Size([1, 142])\n",
      "torch.Size([1, 173])\n",
      "torch.Size([1, 167])\n",
      "torch.Size([1, 227])\n",
      "torch.Size([1, 212])\n",
      "torch.Size([1, 468])\n",
      "torch.Size([1, 330])\n",
      "torch.Size([1, 74])\n",
      "torch.Size([1, 369])\n",
      "torch.Size([1, 89])\n",
      "torch.Size([1, 482])\n",
      "torch.Size([1, 754])\n",
      "torch.Size([1, 255])\n",
      "torch.Size([1, 163])\n",
      "torch.Size([1, 154])\n",
      "torch.Size([1, 189])\n",
      "torch.Size([1, 143])\n",
      "torch.Size([1, 794])\n",
      "torch.Size([1, 318])\n",
      "torch.Size([1, 559])\n",
      "torch.Size([1, 377])\n",
      "torch.Size([1, 64])\n",
      "torch.Size([1, 733])\n",
      "torch.Size([1, 140])\n",
      "torch.Size([1, 540])\n",
      "torch.Size([1, 160])\n",
      "torch.Size([1, 297])\n",
      "torch.Size([1, 188])\n",
      "torch.Size([1, 400])\n",
      "torch.Size([1, 154])\n",
      "torch.Size([1, 288])\n",
      "torch.Size([1, 87])\n",
      "torch.Size([1, 143])\n",
      "torch.Size([1, 78])\n",
      "torch.Size([1, 195])\n",
      "torch.Size([1, 321])\n",
      "torch.Size([1, 289])\n",
      "torch.Size([1, 529])\n",
      "torch.Size([1, 760])\n",
      "torch.Size([1, 324])\n",
      "torch.Size([1, 166])\n",
      "torch.Size([1, 177])\n",
      "torch.Size([1, 867])\n",
      "torch.Size([1, 182])\n",
      "torch.Size([1, 581])\n",
      "torch.Size([1, 204])\n",
      "torch.Size([1, 343])\n",
      "torch.Size([1, 388])\n",
      "torch.Size([1, 152])\n",
      "torch.Size([1, 169])\n",
      "torch.Size([1, 213])\n",
      "torch.Size([1, 260])\n",
      "torch.Size([1, 185])\n",
      "torch.Size([1, 222])\n",
      "torch.Size([1, 229])\n",
      "torch.Size([1, 181])\n",
      "torch.Size([1, 223])\n",
      "torch.Size([1, 193])\n",
      "torch.Size([1, 259])\n",
      "torch.Size([1, 222])\n",
      "torch.Size([1, 101])\n",
      "torch.Size([1, 149])\n",
      "torch.Size([1, 111])\n",
      "torch.Size([1, 189])\n",
      "torch.Size([1, 233])\n",
      "torch.Size([1, 160])\n",
      "torch.Size([1, 149])\n",
      "torch.Size([1, 172])\n",
      "torch.Size([1, 52])\n",
      "torch.Size([1, 129])\n",
      "torch.Size([1, 208])\n",
      "torch.Size([1, 159])\n",
      "torch.Size([1, 89])\n",
      "torch.Size([1, 149])\n",
      "torch.Size([1, 699])\n",
      "torch.Size([1, 154])\n",
      "torch.Size([1, 151])\n",
      "torch.Size([1, 234])\n",
      "torch.Size([1, 840])\n",
      "torch.Size([1, 239])\n",
      "torch.Size([1, 159])\n",
      "torch.Size([1, 164])\n",
      "torch.Size([1, 529])\n",
      "torch.Size([1, 71])\n",
      "torch.Size([1, 317])\n",
      "torch.Size([1, 154])\n",
      "torch.Size([1, 457])\n",
      "torch.Size([1, 202])\n",
      "torch.Size([1, 184])\n",
      "torch.Size([1, 355])\n",
      "torch.Size([1, 221])\n",
      "torch.Size([1, 181])\n",
      "torch.Size([1, 249])\n",
      "torch.Size([1, 638])\n",
      "torch.Size([1, 538])\n",
      "torch.Size([1, 523])\n",
      "torch.Size([1, 88])\n",
      "torch.Size([1, 193])\n",
      "torch.Size([1, 791])\n",
      "torch.Size([1, 226])\n",
      "torch.Size([1, 829])\n",
      "torch.Size([1, 85])\n",
      "torch.Size([1, 161])\n",
      "torch.Size([1, 104])\n",
      "torch.Size([1, 288])\n",
      "torch.Size([1, 431])\n",
      "torch.Size([1, 145])\n",
      "torch.Size([1, 260])\n",
      "torch.Size([1, 210])\n",
      "torch.Size([1, 264])\n",
      "torch.Size([1, 129])\n",
      "torch.Size([1, 77])\n",
      "torch.Size([1, 193])\n",
      "torch.Size([1, 149])\n",
      "torch.Size([1, 189])\n",
      "torch.Size([1, 79])\n",
      "torch.Size([1, 80])\n",
      "torch.Size([1, 591])\n",
      "torch.Size([1, 334])\n",
      "torch.Size([1, 536])\n",
      "torch.Size([1, 166])\n",
      "torch.Size([1, 572])\n",
      "torch.Size([1, 167])\n",
      "torch.Size([1, 148])\n",
      "torch.Size([1, 226])\n",
      "torch.Size([1, 304])\n",
      "torch.Size([1, 88])\n",
      "torch.Size([1, 344])\n",
      "torch.Size([1, 161])\n",
      "torch.Size([1, 355])\n",
      "torch.Size([1, 221])\n",
      "torch.Size([1, 142])\n",
      "torch.Size([1, 370])\n",
      "torch.Size([1, 190])\n",
      "torch.Size([1, 196])\n",
      "torch.Size([1, 266])\n",
      "torch.Size([1, 184])\n",
      "torch.Size([1, 94])\n",
      "torch.Size([1, 122])\n",
      "torch.Size([1, 214])\n",
      "torch.Size([1, 213])\n",
      "torch.Size([1, 134])\n",
      "torch.Size([1, 307])\n",
      "torch.Size([1, 236])\n",
      "torch.Size([1, 143])\n",
      "torch.Size([1, 207])\n",
      "torch.Size([1, 129])\n",
      "torch.Size([1, 260])\n",
      "torch.Size([1, 284])\n",
      "torch.Size([1, 64])\n",
      "torch.Size([1, 282])\n",
      "torch.Size([1, 1024])\n",
      "torch.Size([1, 63])\n",
      "torch.Size([1, 311])\n",
      "torch.Size([1, 166])\n",
      "torch.Size([1, 189])\n",
      "torch.Size([1, 302])\n",
      "torch.Size([1, 171])\n",
      "torch.Size([1, 366])\n",
      "torch.Size([1, 528])\n",
      "torch.Size([1, 81])\n",
      "torch.Size([1, 142])\n",
      "torch.Size([1, 199])\n",
      "torch.Size([1, 893])\n",
      "torch.Size([1, 317])\n",
      "torch.Size([1, 42])\n",
      "torch.Size([1, 97])\n",
      "torch.Size([1, 160])\n",
      "torch.Size([1, 159])\n",
      "torch.Size([1, 110])\n",
      "torch.Size([1, 469])\n",
      "torch.Size([1, 252])\n",
      "torch.Size([1, 158])\n",
      "torch.Size([1, 193])\n",
      "torch.Size([1, 161])\n",
      "torch.Size([1, 165])\n",
      "torch.Size([1, 283])\n",
      "torch.Size([1, 283])\n",
      "torch.Size([1, 212])\n",
      "torch.Size([1, 67])\n",
      "torch.Size([1, 114])\n",
      "torch.Size([1, 310])\n",
      "torch.Size([1, 171])\n",
      "torch.Size([1, 149])\n",
      "torch.Size([1, 87])\n",
      "torch.Size([1, 57])\n",
      "torch.Size([1, 297])\n",
      "torch.Size([1, 241])\n",
      "torch.Size([1, 146])\n",
      "torch.Size([1, 1012])\n",
      "torch.Size([1, 42])\n",
      "torch.Size([1, 409])\n",
      "torch.Size([1, 220])\n",
      "torch.Size([1, 306])\n",
      "torch.Size([1, 132])\n",
      "torch.Size([1, 264])\n",
      "torch.Size([1, 124])\n",
      "torch.Size([1, 190])\n",
      "torch.Size([1, 305])\n",
      "torch.Size([1, 223])\n",
      "torch.Size([1, 249])\n",
      "torch.Size([1, 137])\n",
      "torch.Size([1, 211])\n",
      "torch.Size([1, 136])\n",
      "torch.Size([1, 183])\n",
      "torch.Size([1, 160])\n",
      "torch.Size([1, 250])\n",
      "torch.Size([1, 185])\n",
      "torch.Size([1, 144])\n",
      "torch.Size([1, 277])\n",
      "torch.Size([1, 171])\n",
      "torch.Size([1, 104])\n",
      "torch.Size([1, 141])\n",
      "torch.Size([1, 160])\n",
      "torch.Size([1, 139])\n",
      "torch.Size([1, 213])\n",
      "torch.Size([1, 1024])\n",
      "torch.Size([1, 197])\n",
      "torch.Size([1, 177])\n",
      "torch.Size([1, 273])\n",
      "torch.Size([1, 192])\n",
      "torch.Size([1, 357])\n",
      "torch.Size([1, 88])\n",
      "torch.Size([1, 154])\n",
      "torch.Size([1, 158])\n",
      "torch.Size([1, 254])\n",
      "torch.Size([1, 225])\n",
      "torch.Size([1, 366])\n",
      "torch.Size([1, 173])\n",
      "torch.Size([1, 184])\n",
      "torch.Size([1, 218])\n",
      "torch.Size([1, 264])\n",
      "torch.Size([1, 231])\n",
      "torch.Size([1, 170])\n",
      "torch.Size([1, 308])\n",
      "torch.Size([1, 139])\n",
      "torch.Size([1, 593])\n",
      "torch.Size([1, 185])\n",
      "torch.Size([1, 374])\n",
      "torch.Size([1, 143])\n",
      "torch.Size([1, 1024])\n",
      "torch.Size([1, 271])\n",
      "torch.Size([1, 345])\n",
      "torch.Size([1, 481])\n",
      "torch.Size([1, 222])\n",
      "torch.Size([1, 152])\n",
      "torch.Size([1, 219])\n",
      "torch.Size([1, 568])\n",
      "torch.Size([1, 156])\n",
      "torch.Size([1, 485])\n",
      "torch.Size([1, 317])\n",
      "torch.Size([1, 250])\n",
      "torch.Size([1, 150])\n",
      "torch.Size([1, 112])\n",
      "torch.Size([1, 286])\n",
      "torch.Size([1, 233])\n",
      "torch.Size([1, 138])\n",
      "torch.Size([1, 170])\n",
      "torch.Size([1, 286])\n",
      "torch.Size([1, 232])\n",
      "torch.Size([1, 148])\n",
      "torch.Size([1, 150])\n",
      "torch.Size([1, 275])\n",
      "torch.Size([1, 143])\n",
      "torch.Size([1, 55])\n",
      "torch.Size([1, 1024])\n",
      "torch.Size([1, 408])\n",
      "torch.Size([1, 147])\n",
      "torch.Size([1, 466])\n",
      "torch.Size([1, 684])\n",
      "torch.Size([1, 289])\n",
      "torch.Size([1, 256])\n",
      "torch.Size([1, 707])\n",
      "torch.Size([1, 609])\n",
      "torch.Size([1, 206])\n",
      "torch.Size([1, 442])\n",
      "torch.Size([1, 210])\n",
      "torch.Size([1, 196])\n",
      "torch.Size([1, 351])\n",
      "torch.Size([1, 459])\n",
      "torch.Size([1, 188])\n",
      "torch.Size([1, 197])\n",
      "torch.Size([1, 886])\n",
      "torch.Size([1, 217])\n",
      "torch.Size([1, 688])\n",
      "torch.Size([1, 1017])\n",
      "torch.Size([1, 431])\n",
      "torch.Size([1, 331])\n",
      "torch.Size([1, 335])\n",
      "torch.Size([1, 154])\n",
      "torch.Size([1, 235])\n",
      "torch.Size([1, 212])\n",
      "torch.Size([1, 150])\n",
      "torch.Size([1, 132])\n",
      "torch.Size([1, 166])\n",
      "torch.Size([1, 60])\n",
      "torch.Size([1, 82])\n",
      "torch.Size([1, 54])\n",
      "torch.Size([1, 516])\n",
      "torch.Size([1, 189])\n",
      "torch.Size([1, 254])\n",
      "torch.Size([1, 144])\n",
      "torch.Size([1, 426])\n",
      "torch.Size([1, 163])\n",
      "torch.Size([1, 361])\n",
      "torch.Size([1, 251])\n",
      "torch.Size([1, 101])\n",
      "torch.Size([1, 266])\n",
      "torch.Size([1, 222])\n",
      "torch.Size([1, 499])\n",
      "torch.Size([1, 277])\n",
      "torch.Size([1, 210])\n",
      "torch.Size([1, 148])\n",
      "torch.Size([1, 103])\n",
      "torch.Size([1, 193])\n",
      "torch.Size([1, 353])\n",
      "torch.Size([1, 718])\n",
      "torch.Size([1, 392])\n",
      "torch.Size([1, 151])\n",
      "torch.Size([1, 205])\n",
      "torch.Size([1, 302])\n",
      "torch.Size([1, 213])\n",
      "torch.Size([1, 180])\n",
      "torch.Size([1, 325])\n",
      "torch.Size([1, 234])\n",
      "torch.Size([1, 483])\n",
      "torch.Size([1, 252])\n",
      "torch.Size([1, 185])\n",
      "torch.Size([1, 184])\n",
      "torch.Size([1, 174])\n",
      "torch.Size([1, 423])\n",
      "torch.Size([1, 128])\n",
      "torch.Size([1, 188])\n",
      "torch.Size([1, 143])\n",
      "torch.Size([1, 79])\n",
      "torch.Size([1, 454])\n",
      "torch.Size([1, 265])\n",
      "torch.Size([1, 208])\n",
      "torch.Size([1, 464])\n",
      "torch.Size([1, 550])\n",
      "torch.Size([1, 155])\n",
      "torch.Size([1, 270])\n",
      "torch.Size([1, 659])\n",
      "torch.Size([1, 616])\n",
      "torch.Size([1, 528])\n",
      "torch.Size([1, 216])\n",
      "torch.Size([1, 463])\n",
      "torch.Size([1, 571])\n",
      "torch.Size([1, 248])\n",
      "torch.Size([1, 157])\n",
      "torch.Size([1, 119])\n",
      "torch.Size([1, 220])\n",
      "torch.Size([1, 659])\n",
      "torch.Size([1, 586])\n",
      "torch.Size([1, 282])\n",
      "torch.Size([1, 249])\n",
      "torch.Size([1, 237])\n",
      "torch.Size([1, 386])\n",
      "torch.Size([1, 425])\n",
      "torch.Size([1, 389])\n",
      "torch.Size([1, 315])\n",
      "torch.Size([1, 210])\n",
      "torch.Size([1, 131])\n",
      "torch.Size([1, 180])\n",
      "torch.Size([1, 144])\n",
      "torch.Size([1, 330])\n",
      "torch.Size([1, 292])\n",
      "torch.Size([1, 157])\n",
      "torch.Size([1, 150])\n",
      "torch.Size([1, 170])\n",
      "torch.Size([1, 316])\n",
      "torch.Size([1, 51])\n",
      "torch.Size([1, 217])\n",
      "torch.Size([1, 175])\n",
      "torch.Size([1, 195])\n",
      "torch.Size([1, 147])\n",
      "torch.Size([1, 166])\n",
      "torch.Size([1, 156])\n",
      "torch.Size([1, 735])\n",
      "torch.Size([1, 514])\n",
      "torch.Size([1, 471])\n",
      "torch.Size([1, 214])\n",
      "torch.Size([1, 682])\n",
      "torch.Size([1, 144])\n",
      "torch.Size([1, 173])\n",
      "torch.Size([1, 238])\n",
      "torch.Size([1, 370])\n",
      "torch.Size([1, 165])\n",
      "torch.Size([1, 121])\n",
      "torch.Size([1, 169])\n",
      "torch.Size([1, 205])\n",
      "torch.Size([1, 194])\n",
      "torch.Size([1, 482])\n",
      "torch.Size([1, 170])\n",
      "torch.Size([1, 591])\n",
      "torch.Size([1, 543])\n",
      "torch.Size([1, 223])\n",
      "torch.Size([1, 264])\n",
      "torch.Size([1, 286])\n",
      "torch.Size([1, 744])\n",
      "torch.Size([1, 420])\n",
      "torch.Size([1, 231])\n",
      "torch.Size([1, 102])\n",
      "torch.Size([1, 165])\n",
      "torch.Size([1, 96])\n",
      "torch.Size([1, 211])\n",
      "torch.Size([1, 177])\n",
      "torch.Size([1, 213])\n",
      "torch.Size([1, 518])\n",
      "torch.Size([1, 390])\n",
      "torch.Size([1, 281])\n",
      "torch.Size([1, 103])\n",
      "torch.Size([1, 331])\n",
      "torch.Size([1, 244])\n",
      "torch.Size([1, 423])\n",
      "torch.Size([1, 256])\n",
      "torch.Size([1, 73])\n",
      "torch.Size([1, 302])\n",
      "torch.Size([1, 582])\n",
      "torch.Size([1, 242])\n",
      "torch.Size([1, 279])\n",
      "torch.Size([1, 318])\n",
      "torch.Size([1, 276])\n",
      "torch.Size([1, 496])\n",
      "torch.Size([1, 410])\n",
      "torch.Size([1, 171])\n",
      "torch.Size([1, 227])\n",
      "torch.Size([1, 156])\n",
      "torch.Size([1, 387])\n",
      "torch.Size([1, 195])\n",
      "torch.Size([1, 863])\n",
      "torch.Size([1, 421])\n",
      "torch.Size([1, 508])\n",
      "torch.Size([1, 131])\n",
      "torch.Size([1, 746])\n",
      "torch.Size([1, 161])\n",
      "torch.Size([1, 487])\n",
      "torch.Size([1, 323])\n",
      "torch.Size([1, 1008])\n",
      "torch.Size([1, 231])\n",
      "torch.Size([1, 88])\n",
      "torch.Size([1, 446])\n",
      "torch.Size([1, 104])\n",
      "torch.Size([1, 342])\n",
      "torch.Size([1, 832])\n",
      "torch.Size([1, 198])\n",
      "torch.Size([1, 351])\n",
      "torch.Size([1, 791])\n",
      "torch.Size([1, 241])\n",
      "torch.Size([1, 179])\n",
      "torch.Size([1, 358])\n",
      "torch.Size([1, 154])\n",
      "torch.Size([1, 163])\n",
      "torch.Size([1, 151])\n",
      "torch.Size([1, 496])\n",
      "torch.Size([1, 182])\n",
      "torch.Size([1, 731])\n",
      "torch.Size([1, 156])\n",
      "torch.Size([1, 212])\n",
      "torch.Size([1, 381])\n",
      "torch.Size([1, 355])\n",
      "torch.Size([1, 159])\n",
      "torch.Size([1, 425])\n",
      "torch.Size([1, 902])\n",
      "torch.Size([1, 275])\n",
      "torch.Size([1, 167])\n",
      "torch.Size([1, 148])\n",
      "torch.Size([1, 187])\n",
      "torch.Size([1, 221])\n",
      "torch.Size([1, 265])\n",
      "torch.Size([1, 63])\n",
      "torch.Size([1, 229])\n",
      "torch.Size([1, 213])\n",
      "torch.Size([1, 62])\n",
      "torch.Size([1, 219])\n",
      "torch.Size([1, 78])\n",
      "torch.Size([1, 271])\n",
      "torch.Size([1, 77])\n",
      "torch.Size([1, 697])\n",
      "torch.Size([1, 138])\n",
      "torch.Size([1, 229])\n",
      "torch.Size([1, 124])\n",
      "torch.Size([1, 943])\n",
      "torch.Size([1, 70])\n",
      "torch.Size([1, 92])\n",
      "torch.Size([1, 99])\n",
      "torch.Size([1, 88])\n",
      "torch.Size([1, 172])\n",
      "torch.Size([1, 232])\n",
      "torch.Size([1, 227])\n",
      "torch.Size([1, 696])\n",
      "torch.Size([1, 334])\n",
      "torch.Size([1, 457])\n",
      "torch.Size([1, 267])\n",
      "torch.Size([1, 205])\n",
      "torch.Size([1, 164])\n",
      "torch.Size([1, 270])\n",
      "torch.Size([1, 512])\n",
      "torch.Size([1, 337])\n",
      "torch.Size([1, 253])\n",
      "torch.Size([1, 403])\n",
      "torch.Size([1, 459])\n",
      "torch.Size([1, 232])\n",
      "torch.Size([1, 316])\n",
      "torch.Size([1, 442])\n",
      "torch.Size([1, 88])\n",
      "torch.Size([1, 294])\n",
      "torch.Size([1, 497])\n",
      "torch.Size([1, 202])\n",
      "torch.Size([1, 281])\n",
      "torch.Size([1, 828])\n",
      "torch.Size([1, 451])\n",
      "torch.Size([1, 353])\n",
      "torch.Size([1, 143])\n",
      "torch.Size([1, 277])\n",
      "torch.Size([1, 401])\n",
      "torch.Size([1, 118])\n",
      "torch.Size([1, 451])\n",
      "torch.Size([1, 276])\n",
      "torch.Size([1, 318])\n",
      "torch.Size([1, 223])\n",
      "torch.Size([1, 150])\n",
      "torch.Size([1, 189])\n",
      "torch.Size([1, 642])\n",
      "torch.Size([1, 282])\n",
      "torch.Size([1, 88])\n",
      "torch.Size([1, 206])\n",
      "torch.Size([1, 109])\n",
      "torch.Size([1, 391])\n",
      "torch.Size([1, 160])\n",
      "torch.Size([1, 777])\n",
      "torch.Size([1, 155])\n",
      "torch.Size([1, 193])\n",
      "torch.Size([1, 373])\n",
      "torch.Size([1, 188])\n",
      "torch.Size([1, 151])\n",
      "torch.Size([1, 164])\n",
      "torch.Size([1, 74])\n",
      "torch.Size([1, 157])\n",
      "torch.Size([1, 107])\n",
      "torch.Size([1, 166])\n",
      "torch.Size([1, 313])\n",
      "torch.Size([1, 208])\n",
      "torch.Size([1, 174])\n",
      "torch.Size([1, 162])\n",
      "torch.Size([1, 136])\n",
      "torch.Size([1, 295])\n",
      "torch.Size([1, 402])\n",
      "torch.Size([1, 225])\n",
      "torch.Size([1, 443])\n",
      "torch.Size([1, 132])\n",
      "torch.Size([1, 452])\n",
      "torch.Size([1, 86])\n",
      "torch.Size([1, 221])\n",
      "torch.Size([1, 320])\n",
      "torch.Size([1, 187])\n",
      "torch.Size([1, 318])\n",
      "torch.Size([1, 135])\n",
      "torch.Size([1, 244])\n",
      "torch.Size([1, 172])\n",
      "torch.Size([1, 149])\n",
      "torch.Size([1, 482])\n",
      "torch.Size([1, 296])\n",
      "torch.Size([1, 399])\n",
      "torch.Size([1, 228])\n",
      "torch.Size([1, 166])\n",
      "torch.Size([1, 196])\n",
      "torch.Size([1, 332])\n",
      "torch.Size([1, 116])\n",
      "torch.Size([1, 187])\n",
      "torch.Size([1, 138])\n",
      "torch.Size([1, 225])\n",
      "torch.Size([1, 73])\n",
      "torch.Size([1, 164])\n",
      "torch.Size([1, 197])\n",
      "torch.Size([1, 170])\n",
      "torch.Size([1, 420])\n",
      "torch.Size([1, 273])\n",
      "torch.Size([1, 338])\n",
      "torch.Size([1, 307])\n",
      "torch.Size([1, 125])\n",
      "torch.Size([1, 283])\n",
      "torch.Size([1, 54])\n",
      "torch.Size([1, 660])\n",
      "torch.Size([1, 255])\n",
      "torch.Size([1, 226])\n",
      "torch.Size([1, 103])\n",
      "torch.Size([1, 210])\n",
      "torch.Size([1, 249])\n",
      "torch.Size([1, 909])\n",
      "torch.Size([1, 185])\n",
      "torch.Size([1, 203])\n",
      "torch.Size([1, 185])\n",
      "torch.Size([1, 152])\n",
      "torch.Size([1, 132])\n",
      "torch.Size([1, 182])\n",
      "torch.Size([1, 80])\n",
      "torch.Size([1, 52])\n",
      "torch.Size([1, 207])\n",
      "torch.Size([1, 584])\n",
      "torch.Size([1, 181])\n",
      "torch.Size([1, 492])\n",
      "torch.Size([1, 882])\n",
      "torch.Size([1, 150])\n",
      "torch.Size([1, 380])\n",
      "torch.Size([1, 158])\n",
      "torch.Size([1, 159])\n",
      "torch.Size([1, 169])\n",
      "torch.Size([1, 154])\n",
      "torch.Size([1, 33])\n",
      "torch.Size([1, 176])\n",
      "torch.Size([1, 705])\n",
      "torch.Size([1, 273])\n",
      "torch.Size([1, 146])\n",
      "torch.Size([1, 373])\n",
      "torch.Size([1, 69])\n",
      "torch.Size([1, 214])\n",
      "torch.Size([1, 655])\n",
      "torch.Size([1, 910])\n",
      "torch.Size([1, 201])\n",
      "torch.Size([1, 209])\n",
      "torch.Size([1, 130])\n",
      "torch.Size([1, 181])\n",
      "torch.Size([1, 567])\n",
      "torch.Size([1, 175])\n",
      "torch.Size([1, 677])\n",
      "torch.Size([1, 177])\n",
      "torch.Size([1, 325])\n",
      "torch.Size([1, 205])\n",
      "torch.Size([1, 371])\n",
      "torch.Size([1, 151])\n",
      "torch.Size([1, 339])\n",
      "torch.Size([1, 221])\n",
      "torch.Size([1, 294])\n",
      "torch.Size([1, 90])\n",
      "torch.Size([1, 582])\n",
      "torch.Size([1, 291])\n",
      "torch.Size([1, 1024])\n",
      "torch.Size([1, 83])\n",
      "torch.Size([1, 131])\n",
      "torch.Size([1, 76])\n",
      "torch.Size([1, 200])\n",
      "torch.Size([1, 131])\n",
      "torch.Size([1, 163])\n",
      "torch.Size([1, 254])\n",
      "torch.Size([1, 172])\n",
      "torch.Size([1, 504])\n",
      "torch.Size([1, 235])\n",
      "torch.Size([1, 197])\n",
      "torch.Size([1, 279])\n",
      "torch.Size([1, 196])\n",
      "torch.Size([1, 206])\n",
      "torch.Size([1, 161])\n",
      "torch.Size([1, 149])\n",
      "torch.Size([1, 124])\n",
      "torch.Size([1, 904])\n",
      "torch.Size([1, 398])\n",
      "torch.Size([1, 50])\n",
      "torch.Size([1, 175])\n",
      "torch.Size([1, 193])\n",
      "torch.Size([1, 420])\n",
      "torch.Size([1, 273])\n",
      "torch.Size([1, 196])\n",
      "torch.Size([1, 696])\n",
      "torch.Size([1, 286])\n",
      "torch.Size([1, 174])\n",
      "torch.Size([1, 159])\n",
      "torch.Size([1, 364])\n",
      "torch.Size([1, 244])\n",
      "torch.Size([1, 277])\n",
      "torch.Size([1, 190])\n",
      "torch.Size([1, 933])\n",
      "torch.Size([1, 53])\n",
      "torch.Size([1, 229])\n",
      "torch.Size([1, 201])\n",
      "torch.Size([1, 252])\n",
      "torch.Size([1, 186])\n",
      "torch.Size([1, 177])\n",
      "torch.Size([1, 108])\n",
      "torch.Size([1, 209])\n",
      "torch.Size([1, 66])\n",
      "torch.Size([1, 144])\n",
      "torch.Size([1, 238])\n",
      "torch.Size([1, 226])\n",
      "torch.Size([1, 333])\n",
      "torch.Size([1, 101])\n",
      "torch.Size([1, 285])\n",
      "torch.Size([1, 206])\n",
      "torch.Size([1, 503])\n",
      "torch.Size([1, 221])\n",
      "torch.Size([1, 128])\n",
      "torch.Size([1, 752])\n",
      "torch.Size([1, 776])\n",
      "torch.Size([1, 202])\n",
      "torch.Size([1, 834])\n",
      "torch.Size([1, 245])\n",
      "torch.Size([1, 180])\n",
      "torch.Size([1, 834])\n",
      "torch.Size([1, 751])\n",
      "torch.Size([1, 212])\n",
      "torch.Size([1, 275])\n",
      "torch.Size([1, 192])\n",
      "torch.Size([1, 202])\n",
      "torch.Size([1, 512])\n",
      "torch.Size([1, 631])\n",
      "torch.Size([1, 137])\n",
      "torch.Size([1, 363])\n",
      "torch.Size([1, 184])\n",
      "torch.Size([1, 503])\n",
      "torch.Size([1, 288])\n",
      "torch.Size([1, 143])\n",
      "torch.Size([1, 148])\n",
      "torch.Size([1, 258])\n",
      "torch.Size([1, 254])\n",
      "torch.Size([1, 163])\n",
      "torch.Size([1, 263])\n",
      "torch.Size([1, 94])\n",
      "torch.Size([1, 156])\n",
      "torch.Size([1, 725])\n",
      "torch.Size([1, 255])\n",
      "torch.Size([1, 239])\n",
      "torch.Size([1, 205])\n",
      "torch.Size([1, 146])\n",
      "torch.Size([1, 186])\n",
      "torch.Size([1, 574])\n",
      "torch.Size([1, 215])\n",
      "torch.Size([1, 211])\n",
      "torch.Size([1, 234])\n",
      "torch.Size([1, 142])\n",
      "torch.Size([1, 227])\n",
      "torch.Size([1, 303])\n",
      "torch.Size([1, 64])\n",
      "torch.Size([1, 90])\n",
      "torch.Size([1, 176])\n",
      "torch.Size([1, 221])\n",
      "torch.Size([1, 300])\n",
      "torch.Size([1, 72])\n",
      "torch.Size([1, 220])\n",
      "torch.Size([1, 186])\n",
      "torch.Size([1, 485])\n",
      "torch.Size([1, 153])\n",
      "torch.Size([1, 226])\n",
      "torch.Size([1, 493])\n",
      "torch.Size([1, 165])\n",
      "torch.Size([1, 229])\n",
      "torch.Size([1, 297])\n",
      "torch.Size([1, 146])\n",
      "torch.Size([1, 329])\n",
      "torch.Size([1, 156])\n",
      "torch.Size([1, 635])\n",
      "torch.Size([1, 77])\n",
      "torch.Size([1, 54])\n",
      "torch.Size([1, 197])\n",
      "torch.Size([1, 174])\n",
      "torch.Size([1, 308])\n",
      "torch.Size([1, 346])\n",
      "torch.Size([1, 237])\n",
      "torch.Size([1, 53])\n",
      "torch.Size([1, 65])\n",
      "torch.Size([1, 137])\n",
      "torch.Size([1, 283])\n",
      "torch.Size([1, 169])\n",
      "torch.Size([1, 773])\n",
      "torch.Size([1, 527])\n",
      "torch.Size([1, 464])\n",
      "torch.Size([1, 355])\n",
      "torch.Size([1, 229])\n",
      "torch.Size([1, 158])\n",
      "torch.Size([1, 211])\n",
      "torch.Size([1, 255])\n",
      "torch.Size([1, 255])\n",
      "torch.Size([1, 622])\n",
      "torch.Size([1, 299])\n",
      "torch.Size([1, 177])\n",
      "torch.Size([1, 204])\n",
      "torch.Size([1, 262])\n",
      "torch.Size([1, 666])\n",
      "torch.Size([1, 225])\n",
      "torch.Size([1, 237])\n",
      "torch.Size([1, 191])\n",
      "torch.Size([1, 263])\n",
      "torch.Size([1, 363])\n",
      "torch.Size([1, 614])\n",
      "torch.Size([1, 466])\n",
      "torch.Size([1, 144])\n",
      "torch.Size([1, 888])\n",
      "torch.Size([1, 143])\n",
      "torch.Size([1, 142])\n",
      "torch.Size([1, 289])\n",
      "torch.Size([1, 234])\n",
      "torch.Size([1, 222])\n",
      "torch.Size([1, 573])\n",
      "torch.Size([1, 367])\n",
      "torch.Size([1, 158])\n",
      "torch.Size([1, 303])\n",
      "torch.Size([1, 58])\n",
      "torch.Size([1, 185])\n",
      "torch.Size([1, 985])\n",
      "torch.Size([1, 218])\n",
      "torch.Size([1, 144])\n",
      "torch.Size([1, 574])\n",
      "torch.Size([1, 249])\n",
      "torch.Size([1, 297])\n",
      "torch.Size([1, 333])\n",
      "torch.Size([1, 341])\n",
      "torch.Size([1, 191])\n",
      "torch.Size([1, 156])\n",
      "torch.Size([1, 439])\n",
      "torch.Size([1, 95])\n",
      "torch.Size([1, 96])\n",
      "torch.Size([1, 285])\n",
      "torch.Size([1, 194])\n",
      "torch.Size([1, 759])\n",
      "torch.Size([1, 67])\n",
      "torch.Size([1, 594])\n",
      "torch.Size([1, 133])\n",
      "torch.Size([1, 182])\n",
      "torch.Size([1, 274])\n",
      "torch.Size([1, 907])\n",
      "torch.Size([1, 84])\n",
      "torch.Size([1, 147])\n",
      "torch.Size([1, 539])\n",
      "torch.Size([1, 516])\n",
      "torch.Size([1, 339])\n",
      "torch.Size([1, 174])\n",
      "torch.Size([1, 398])\n",
      "torch.Size([1, 409])\n",
      "torch.Size([1, 536])\n",
      "torch.Size([1, 197])\n",
      "torch.Size([1, 212])\n",
      "torch.Size([1, 212])\n",
      "torch.Size([1, 365])\n",
      "torch.Size([1, 207])\n",
      "torch.Size([1, 287])\n",
      "torch.Size([1, 105])\n",
      "torch.Size([1, 169])\n",
      "torch.Size([1, 206])\n",
      "torch.Size([1, 166])\n",
      "torch.Size([1, 142])\n",
      "torch.Size([1, 234])\n",
      "torch.Size([1, 137])\n",
      "torch.Size([1, 515])\n",
      "torch.Size([1, 229])\n",
      "torch.Size([1, 155])\n",
      "torch.Size([1, 1024])\n",
      "torch.Size([1, 157])\n",
      "torch.Size([1, 162])\n",
      "torch.Size([1, 346])\n",
      "torch.Size([1, 172])\n",
      "torch.Size([1, 44])\n",
      "torch.Size([1, 171])\n",
      "torch.Size([1, 230])\n",
      "torch.Size([1, 225])\n",
      "torch.Size([1, 94])\n",
      "torch.Size([1, 901])\n",
      "torch.Size([1, 105])\n",
      "torch.Size([1, 653])\n",
      "torch.Size([1, 208])\n",
      "torch.Size([1, 332])\n",
      "torch.Size([1, 211])\n",
      "torch.Size([1, 399])\n",
      "torch.Size([1, 166])\n",
      "torch.Size([1, 211])\n",
      "torch.Size([1, 147])\n",
      "torch.Size([1, 145])\n",
      "torch.Size([1, 229])\n",
      "torch.Size([1, 348])\n",
      "torch.Size([1, 150])\n",
      "torch.Size([1, 859])\n",
      "torch.Size([1, 221])\n",
      "torch.Size([1, 336])\n",
      "torch.Size([1, 465])\n",
      "torch.Size([1, 294])\n",
      "torch.Size([1, 328])\n",
      "torch.Size([1, 296])\n",
      "torch.Size([1, 568])\n",
      "torch.Size([1, 252])\n",
      "torch.Size([1, 766])\n",
      "torch.Size([1, 234])\n",
      "torch.Size([1, 211])\n",
      "torch.Size([1, 245])\n",
      "torch.Size([1, 226])\n",
      "torch.Size([1, 281])\n",
      "torch.Size([1, 192])\n",
      "torch.Size([1, 75])\n",
      "torch.Size([1, 343])\n",
      "torch.Size([1, 644])\n",
      "torch.Size([1, 86])\n",
      "torch.Size([1, 852])\n",
      "torch.Size([1, 266])\n",
      "torch.Size([1, 433])\n",
      "torch.Size([1, 337])\n",
      "torch.Size([1, 198])\n",
      "torch.Size([1, 153])\n",
      "torch.Size([1, 147])\n",
      "torch.Size([1, 71])\n",
      "torch.Size([1, 245])\n",
      "torch.Size([1, 141])\n",
      "torch.Size([1, 829])\n",
      "torch.Size([1, 249])\n",
      "torch.Size([1, 208])\n",
      "torch.Size([1, 645])\n",
      "torch.Size([1, 686])\n",
      "torch.Size([1, 316])\n",
      "torch.Size([1, 156])\n",
      "torch.Size([1, 260])\n",
      "torch.Size([1, 194])\n",
      "torch.Size([1, 100])\n",
      "torch.Size([1, 174])\n",
      "torch.Size([1, 355])\n",
      "torch.Size([1, 213])\n",
      "torch.Size([1, 189])\n",
      "torch.Size([1, 82])\n",
      "torch.Size([1, 588])\n",
      "torch.Size([1, 416])\n",
      "torch.Size([1, 402])\n",
      "torch.Size([1, 273])\n",
      "torch.Size([1, 128])\n",
      "torch.Size([1, 488])\n",
      "torch.Size([1, 517])\n",
      "torch.Size([1, 366])\n",
      "torch.Size([1, 261])\n",
      "torch.Size([1, 228])\n",
      "torch.Size([1, 323])\n",
      "torch.Size([1, 361])\n",
      "torch.Size([1, 446])\n",
      "torch.Size([1, 248])\n",
      "torch.Size([1, 969])\n",
      "torch.Size([1, 323])\n",
      "torch.Size([1, 236])\n",
      "torch.Size([1, 426])\n",
      "torch.Size([1, 78])\n",
      "torch.Size([1, 330])\n",
      "torch.Size([1, 95])\n",
      "torch.Size([1, 234])\n",
      "torch.Size([1, 148])\n",
      "torch.Size([1, 274])\n",
      "torch.Size([1, 742])\n",
      "torch.Size([1, 168])\n",
      "torch.Size([1, 633])\n",
      "torch.Size([1, 202])\n",
      "torch.Size([1, 311])\n",
      "torch.Size([1, 180])\n",
      "torch.Size([1, 258])\n",
      "torch.Size([1, 84])\n",
      "torch.Size([1, 228])\n",
      "torch.Size([1, 718])\n",
      "torch.Size([1, 210])\n",
      "torch.Size([1, 375])\n",
      "torch.Size([1, 642])\n",
      "torch.Size([1, 148])\n",
      "torch.Size([1, 341])\n",
      "torch.Size([1, 318])\n",
      "torch.Size([1, 209])\n",
      "torch.Size([1, 283])\n",
      "torch.Size([1, 165])\n",
      "torch.Size([1, 244])\n",
      "torch.Size([1, 149])\n",
      "torch.Size([1, 221])\n",
      "torch.Size([1, 194])\n",
      "torch.Size([1, 336])\n",
      "torch.Size([1, 717])\n",
      "torch.Size([1, 181])\n",
      "torch.Size([1, 202])\n",
      "torch.Size([1, 133])\n",
      "torch.Size([1, 241])\n",
      "torch.Size([1, 142])\n",
      "torch.Size([1, 248])\n",
      "torch.Size([1, 136])\n",
      "torch.Size([1, 226])\n",
      "torch.Size([1, 213])\n",
      "torch.Size([1, 780])\n",
      "torch.Size([1, 165])\n",
      "torch.Size([1, 295])\n",
      "torch.Size([1, 224])\n",
      "torch.Size([1, 323])\n",
      "torch.Size([1, 132])\n",
      "torch.Size([1, 213])\n",
      "torch.Size([1, 387])\n",
      "torch.Size([1, 167])\n",
      "torch.Size([1, 219])\n",
      "torch.Size([1, 184])\n",
      "torch.Size([1, 275])\n",
      "torch.Size([1, 362])\n",
      "torch.Size([1, 354])\n",
      "torch.Size([1, 397])\n",
      "torch.Size([1, 193])\n",
      "torch.Size([1, 427])\n",
      "torch.Size([1, 308])\n",
      "torch.Size([1, 165])\n",
      "torch.Size([1, 202])\n",
      "torch.Size([1, 188])\n",
      "torch.Size([1, 107])\n",
      "torch.Size([1, 143])\n",
      "torch.Size([1, 253])\n",
      "torch.Size([1, 309])\n",
      "torch.Size([1, 106])\n",
      "torch.Size([1, 106])\n",
      "torch.Size([1, 290])\n",
      "torch.Size([1, 159])\n",
      "torch.Size([1, 192])\n",
      "torch.Size([1, 160])\n",
      "torch.Size([1, 291])\n",
      "torch.Size([1, 397])\n",
      "torch.Size([1, 130])\n",
      "torch.Size([1, 168])\n",
      "torch.Size([1, 136])\n",
      "torch.Size([1, 70])\n",
      "torch.Size([1, 385])\n",
      "torch.Size([1, 120])\n",
      "torch.Size([1, 933])\n",
      "torch.Size([1, 101])\n",
      "torch.Size([1, 319])\n",
      "torch.Size([1, 637])\n",
      "torch.Size([1, 309])\n",
      "torch.Size([1, 588])\n",
      "torch.Size([1, 152])\n",
      "torch.Size([1, 66])\n",
      "torch.Size([1, 149])\n",
      "torch.Size([1, 214])\n",
      "torch.Size([1, 183])\n",
      "torch.Size([1, 192])\n",
      "torch.Size([1, 218])\n",
      "torch.Size([1, 107])\n",
      "torch.Size([1, 211])\n",
      "torch.Size([1, 241])\n",
      "torch.Size([1, 310])\n",
      "torch.Size([1, 158])\n",
      "torch.Size([1, 173])\n",
      "torch.Size([1, 288])\n",
      "torch.Size([1, 216])\n",
      "torch.Size([1, 393])\n",
      "torch.Size([1, 150])\n",
      "torch.Size([1, 177])\n",
      "torch.Size([1, 430])\n",
      "torch.Size([1, 92])\n",
      "torch.Size([1, 150])\n",
      "torch.Size([1, 199])\n",
      "torch.Size([1, 457])\n",
      "torch.Size([1, 95])\n",
      "torch.Size([1, 99])\n",
      "torch.Size([1, 165])\n",
      "torch.Size([1, 370])\n",
      "torch.Size([1, 201])\n",
      "torch.Size([1, 160])\n",
      "torch.Size([1, 508])\n",
      "torch.Size([1, 277])\n",
      "torch.Size([1, 122])\n",
      "torch.Size([1, 284])\n",
      "torch.Size([1, 192])\n",
      "torch.Size([1, 206])\n",
      "torch.Size([1, 266])\n",
      "torch.Size([1, 147])\n",
      "torch.Size([1, 77])\n",
      "torch.Size([1, 175])\n",
      "torch.Size([1, 306])\n",
      "torch.Size([1, 215])\n",
      "torch.Size([1, 574])\n",
      "torch.Size([1, 222])\n",
      "torch.Size([1, 153])\n",
      "torch.Size([1, 108])\n",
      "torch.Size([1, 304])\n",
      "torch.Size([1, 747])\n",
      "torch.Size([1, 45])\n",
      "torch.Size([1, 202])\n",
      "torch.Size([1, 185])\n",
      "torch.Size([1, 160])\n",
      "torch.Size([1, 133])\n",
      "torch.Size([1, 330])\n",
      "torch.Size([1, 250])\n",
      "torch.Size([1, 158])\n",
      "torch.Size([1, 1024])\n",
      "torch.Size([1, 200])\n",
      "torch.Size([1, 44])\n",
      "torch.Size([1, 412])\n",
      "torch.Size([1, 594])\n",
      "torch.Size([1, 873])\n",
      "torch.Size([1, 388])\n",
      "torch.Size([1, 220])\n",
      "torch.Size([1, 157])\n",
      "torch.Size([1, 199])\n",
      "torch.Size([1, 250])\n",
      "torch.Size([1, 296])\n",
      "torch.Size([1, 204])\n",
      "torch.Size([1, 98])\n",
      "torch.Size([1, 174])\n",
      "torch.Size([1, 189])\n",
      "torch.Size([1, 145])\n",
      "torch.Size([1, 1024])\n",
      "torch.Size([1, 158])\n",
      "torch.Size([1, 260])\n",
      "torch.Size([1, 37])\n",
      "torch.Size([1, 188])\n",
      "torch.Size([1, 729])\n",
      "torch.Size([1, 175])\n",
      "torch.Size([1, 126])\n",
      "torch.Size([1, 557])\n",
      "torch.Size([1, 56])\n",
      "torch.Size([1, 124])\n",
      "torch.Size([1, 160])\n",
      "torch.Size([1, 154])\n",
      "torch.Size([1, 210])\n",
      "torch.Size([1, 203])\n",
      "torch.Size([1, 137])\n",
      "torch.Size([1, 189])\n",
      "torch.Size([1, 140])\n",
      "torch.Size([1, 208])\n",
      "torch.Size([1, 291])\n",
      "torch.Size([1, 227])\n",
      "torch.Size([1, 448])\n",
      "torch.Size([1, 278])\n",
      "torch.Size([1, 201])\n",
      "torch.Size([1, 99])\n",
      "torch.Size([1, 328])\n",
      "torch.Size([1, 253])\n",
      "torch.Size([1, 149])\n",
      "torch.Size([1, 255])\n",
      "torch.Size([1, 222])\n",
      "torch.Size([1, 148])\n",
      "torch.Size([1, 275])\n",
      "torch.Size([1, 422])\n",
      "torch.Size([1, 328])\n",
      "torch.Size([1, 217])\n",
      "torch.Size([1, 320])\n",
      "torch.Size([1, 724])\n",
      "torch.Size([1, 453])\n",
      "torch.Size([1, 66])\n",
      "torch.Size([1, 242])\n",
      "torch.Size([1, 151])\n",
      "torch.Size([1, 157])\n",
      "torch.Size([1, 188])\n",
      "torch.Size([1, 605])\n",
      "torch.Size([1, 206])\n",
      "torch.Size([1, 811])\n",
      "torch.Size([1, 414])\n",
      "torch.Size([1, 332])\n",
      "torch.Size([1, 553])\n",
      "torch.Size([1, 178])\n",
      "torch.Size([1, 152])\n",
      "torch.Size([1, 86])\n",
      "torch.Size([1, 148])\n",
      "torch.Size([1, 166])\n",
      "torch.Size([1, 201])\n",
      "torch.Size([1, 298])\n",
      "torch.Size([1, 817])\n",
      "torch.Size([1, 249])\n",
      "torch.Size([1, 728])\n",
      "torch.Size([1, 421])\n",
      "torch.Size([1, 415])\n",
      "torch.Size([1, 76])\n",
      "torch.Size([1, 226])\n",
      "torch.Size([1, 347])\n",
      "torch.Size([1, 76])\n",
      "torch.Size([1, 69])\n",
      "torch.Size([1, 403])\n",
      "torch.Size([1, 378])\n",
      "torch.Size([1, 235])\n",
      "torch.Size([1, 146])\n",
      "torch.Size([1, 221])\n",
      "torch.Size([1, 254])\n",
      "torch.Size([1, 154])\n",
      "torch.Size([1, 233])\n",
      "torch.Size([1, 102])\n",
      "torch.Size([1, 290])\n",
      "torch.Size([1, 102])\n",
      "torch.Size([1, 618])\n",
      "torch.Size([1, 73])\n",
      "torch.Size([1, 406])\n",
      "torch.Size([1, 136])\n",
      "torch.Size([1, 73])\n",
      "torch.Size([1, 150])\n",
      "torch.Size([1, 178])\n",
      "torch.Size([1, 539])\n",
      "torch.Size([1, 180])\n",
      "torch.Size([1, 347])\n",
      "torch.Size([1, 140])\n",
      "torch.Size([1, 248])\n",
      "torch.Size([1, 122])\n",
      "torch.Size([1, 374])\n",
      "torch.Size([1, 163])\n",
      "torch.Size([1, 892])\n",
      "torch.Size([1, 257])\n",
      "torch.Size([1, 544])\n",
      "torch.Size([1, 269])\n",
      "torch.Size([1, 138])\n",
      "torch.Size([1, 140])\n",
      "torch.Size([1, 173])\n",
      "torch.Size([1, 209])\n",
      "torch.Size([1, 235])\n",
      "torch.Size([1, 695])\n",
      "torch.Size([1, 496])\n",
      "torch.Size([1, 692])\n",
      "torch.Size([1, 161])\n",
      "torch.Size([1, 718])\n",
      "torch.Size([1, 344])\n",
      "torch.Size([1, 1024])\n",
      "torch.Size([1, 202])\n",
      "torch.Size([1, 394])\n",
      "torch.Size([1, 179])\n",
      "torch.Size([1, 147])\n",
      "torch.Size([1, 486])\n",
      "torch.Size([1, 363])\n",
      "torch.Size([1, 227])\n",
      "torch.Size([1, 94])\n",
      "torch.Size([1, 309])\n",
      "torch.Size([1, 1024])\n",
      "torch.Size([1, 1024])\n",
      "torch.Size([1, 216])\n",
      "torch.Size([1, 173])\n",
      "torch.Size([1, 136])\n",
      "torch.Size([1, 172])\n",
      "torch.Size([1, 185])\n",
      "torch.Size([1, 325])\n",
      "torch.Size([1, 214])\n",
      "torch.Size([1, 1007])\n",
      "torch.Size([1, 159])\n",
      "torch.Size([1, 88])\n",
      "torch.Size([1, 108])\n",
      "torch.Size([1, 493])\n",
      "torch.Size([1, 743])\n",
      "torch.Size([1, 119])\n",
      "torch.Size([1, 211])\n",
      "torch.Size([1, 252])\n",
      "torch.Size([1, 208])\n",
      "torch.Size([1, 404])\n",
      "torch.Size([1, 158])\n",
      "torch.Size([1, 1024])\n",
      "torch.Size([1, 234])\n",
      "torch.Size([1, 404])\n",
      "torch.Size([1, 312])\n",
      "torch.Size([1, 147])\n",
      "torch.Size([1, 152])\n",
      "torch.Size([1, 152])\n",
      "torch.Size([1, 243])\n",
      "torch.Size([1, 417])\n",
      "torch.Size([1, 193])\n",
      "torch.Size([1, 338])\n",
      "torch.Size([1, 143])\n",
      "torch.Size([1, 146])\n",
      "torch.Size([1, 175])\n",
      "torch.Size([1, 495])\n",
      "torch.Size([1, 477])\n",
      "torch.Size([1, 180])\n",
      "torch.Size([1, 246])\n",
      "torch.Size([1, 238])\n",
      "torch.Size([1, 267])\n",
      "torch.Size([1, 129])\n",
      "torch.Size([1, 252])\n",
      "torch.Size([1, 416])\n",
      "torch.Size([1, 235])\n",
      "torch.Size([1, 222])\n",
      "torch.Size([1, 164])\n",
      "torch.Size([1, 239])\n",
      "torch.Size([1, 99])\n",
      "torch.Size([1, 150])\n",
      "torch.Size([1, 161])\n",
      "torch.Size([1, 168])\n",
      "torch.Size([1, 451])\n",
      "torch.Size([1, 387])\n",
      "torch.Size([1, 88])\n",
      "torch.Size([1, 61])\n",
      "torch.Size([1, 634])\n",
      "torch.Size([1, 229])\n",
      "torch.Size([1, 359])\n",
      "torch.Size([1, 606])\n",
      "torch.Size([1, 178])\n",
      "torch.Size([1, 300])\n",
      "torch.Size([1, 168])\n",
      "torch.Size([1, 438])\n",
      "torch.Size([1, 275])\n",
      "torch.Size([1, 254])\n",
      "torch.Size([1, 180])\n",
      "torch.Size([1, 556])\n",
      "torch.Size([1, 88])\n",
      "torch.Size([1, 355])\n",
      "torch.Size([1, 122])\n",
      "torch.Size([1, 229])\n",
      "torch.Size([1, 151])\n",
      "torch.Size([1, 192])\n",
      "torch.Size([1, 205])\n",
      "torch.Size([1, 167])\n",
      "torch.Size([1, 222])\n",
      "torch.Size([1, 496])\n",
      "torch.Size([1, 698])\n",
      "torch.Size([1, 75])\n",
      "torch.Size([1, 205])\n",
      "torch.Size([1, 185])\n",
      "torch.Size([1, 176])\n",
      "torch.Size([1, 159])\n",
      "torch.Size([1, 391])\n",
      "torch.Size([1, 1015])\n",
      "torch.Size([1, 252])\n",
      "torch.Size([1, 1024])\n",
      "torch.Size([1, 227])\n",
      "torch.Size([1, 203])\n",
      "torch.Size([1, 672])\n",
      "torch.Size([1, 187])\n",
      "torch.Size([1, 245])\n",
      "torch.Size([1, 139])\n",
      "torch.Size([1, 270])\n",
      "torch.Size([1, 180])\n",
      "torch.Size([1, 173])\n",
      "torch.Size([1, 46])\n",
      "torch.Size([1, 156])\n",
      "torch.Size([1, 131])\n",
      "torch.Size([1, 128])\n",
      "torch.Size([1, 107])\n",
      "torch.Size([1, 295])\n",
      "torch.Size([1, 147])\n",
      "torch.Size([1, 86])\n",
      "torch.Size([1, 441])\n",
      "torch.Size([1, 196])\n",
      "torch.Size([1, 127])\n",
      "torch.Size([1, 246])\n",
      "torch.Size([1, 223])\n",
      "torch.Size([1, 337])\n",
      "torch.Size([1, 147])\n",
      "torch.Size([1, 248])\n",
      "torch.Size([1, 304])\n",
      "torch.Size([1, 80])\n",
      "torch.Size([1, 150])\n",
      "torch.Size([1, 315])\n",
      "torch.Size([1, 172])\n",
      "torch.Size([1, 375])\n",
      "torch.Size([1, 1024])\n",
      "torch.Size([1, 226])\n",
      "torch.Size([1, 228])\n",
      "torch.Size([1, 199])\n",
      "torch.Size([1, 156])\n",
      "torch.Size([1, 198])\n",
      "torch.Size([1, 146])\n",
      "torch.Size([1, 167])\n",
      "torch.Size([1, 157])\n",
      "torch.Size([1, 306])\n",
      "torch.Size([1, 165])\n",
      "torch.Size([1, 110])\n",
      "torch.Size([1, 566])\n",
      "torch.Size([1, 243])\n",
      "torch.Size([1, 63])\n",
      "torch.Size([1, 174])\n",
      "torch.Size([1, 437])\n",
      "torch.Size([1, 150])\n",
      "torch.Size([1, 154])\n",
      "torch.Size([1, 402])\n",
      "torch.Size([1, 164])\n",
      "torch.Size([1, 138])\n",
      "torch.Size([1, 278])\n",
      "torch.Size([1, 410])\n",
      "torch.Size([1, 160])\n",
      "torch.Size([1, 188])\n",
      "torch.Size([1, 169])\n",
      "torch.Size([1, 309])\n",
      "torch.Size([1, 47])\n",
      "torch.Size([1, 619])\n",
      "torch.Size([1, 261])\n",
      "torch.Size([1, 155])\n",
      "torch.Size([1, 50])\n",
      "torch.Size([1, 94])\n",
      "torch.Size([1, 253])\n",
      "torch.Size([1, 224])\n",
      "torch.Size([1, 203])\n",
      "torch.Size([1, 33])\n",
      "torch.Size([1, 469])\n",
      "torch.Size([1, 191])\n",
      "torch.Size([1, 516])\n",
      "torch.Size([1, 283])\n",
      "torch.Size([1, 176])\n",
      "torch.Size([1, 394])\n",
      "torch.Size([1, 901])\n",
      "torch.Size([1, 471])\n",
      "torch.Size([1, 78])\n",
      "torch.Size([1, 142])\n",
      "torch.Size([1, 108])\n",
      "torch.Size([1, 453])\n",
      "torch.Size([1, 195])\n",
      "torch.Size([1, 157])\n",
      "torch.Size([1, 139])\n",
      "torch.Size([1, 621])\n",
      "torch.Size([1, 205])\n",
      "torch.Size([1, 1024])\n",
      "torch.Size([1, 424])\n",
      "torch.Size([1, 613])\n",
      "torch.Size([1, 360])\n",
      "torch.Size([1, 545])\n",
      "torch.Size([1, 309])\n",
      "torch.Size([1, 243])\n",
      "torch.Size([1, 249])\n",
      "torch.Size([1, 375])\n",
      "torch.Size([1, 144])\n",
      "torch.Size([1, 439])\n",
      "torch.Size([1, 317])\n",
      "torch.Size([1, 84])\n",
      "torch.Size([1, 467])\n",
      "torch.Size([1, 192])\n",
      "torch.Size([1, 320])\n",
      "torch.Size([1, 290])\n",
      "torch.Size([1, 304])\n",
      "torch.Size([1, 174])\n",
      "torch.Size([1, 122])\n",
      "torch.Size([1, 95])\n",
      "torch.Size([1, 144])\n",
      "torch.Size([1, 166])\n",
      "torch.Size([1, 315])\n",
      "torch.Size([1, 628])\n",
      "torch.Size([1, 532])\n",
      "torch.Size([1, 237])\n",
      "torch.Size([1, 159])\n",
      "torch.Size([1, 504])\n",
      "torch.Size([1, 244])\n",
      "torch.Size([1, 476])\n",
      "torch.Size([1, 280])\n",
      "torch.Size([1, 144])\n",
      "torch.Size([1, 117])\n",
      "torch.Size([1, 295])\n",
      "torch.Size([1, 405])\n",
      "torch.Size([1, 171])\n",
      "torch.Size([1, 489])\n",
      "torch.Size([1, 145])\n",
      "torch.Size([1, 407])\n",
      "torch.Size([1, 158])\n",
      "torch.Size([1, 196])\n",
      "torch.Size([1, 268])\n",
      "torch.Size([1, 106])\n",
      "torch.Size([1, 235])\n",
      "torch.Size([1, 260])\n",
      "torch.Size([1, 121])\n",
      "torch.Size([1, 287])\n",
      "torch.Size([1, 1024])\n",
      "torch.Size([1, 470])\n",
      "torch.Size([1, 218])\n",
      "torch.Size([1, 431])\n",
      "torch.Size([1, 148])\n",
      "torch.Size([1, 213])\n",
      "torch.Size([1, 238])\n",
      "torch.Size([1, 193])\n",
      "torch.Size([1, 249])\n",
      "torch.Size([1, 252])\n",
      "torch.Size([1, 152])\n",
      "torch.Size([1, 141])\n",
      "torch.Size([1, 154])\n",
      "torch.Size([1, 181])\n",
      "torch.Size([1, 785])\n",
      "torch.Size([1, 439])\n",
      "torch.Size([1, 463])\n",
      "torch.Size([1, 316])\n",
      "torch.Size([1, 219])\n",
      "torch.Size([1, 184])\n",
      "torch.Size([1, 188])\n",
      "torch.Size([1, 135])\n",
      "torch.Size([1, 624])\n",
      "torch.Size([1, 291])\n",
      "torch.Size([1, 164])\n",
      "torch.Size([1, 797])\n",
      "torch.Size([1, 257])\n",
      "torch.Size([1, 64])\n",
      "torch.Size([1, 628])\n",
      "torch.Size([1, 194])\n",
      "torch.Size([1, 77])\n",
      "torch.Size([1, 478])\n",
      "torch.Size([1, 486])\n",
      "torch.Size([1, 94])\n",
      "torch.Size([1, 79])\n",
      "torch.Size([1, 419])\n",
      "torch.Size([1, 149])\n",
      "torch.Size([1, 123])\n",
      "torch.Size([1, 326])\n",
      "torch.Size([1, 328])\n",
      "torch.Size([1, 169])\n",
      "torch.Size([1, 175])\n",
      "torch.Size([1, 499])\n",
      "torch.Size([1, 317])\n",
      "torch.Size([1, 238])\n",
      "torch.Size([1, 276])\n",
      "torch.Size([1, 214])\n",
      "torch.Size([1, 280])\n",
      "torch.Size([1, 156])\n",
      "torch.Size([1, 281])\n",
      "torch.Size([1, 168])\n",
      "torch.Size([1, 899])\n",
      "torch.Size([1, 207])\n",
      "torch.Size([1, 198])\n",
      "torch.Size([1, 893])\n",
      "torch.Size([1, 265])\n",
      "torch.Size([1, 220])\n",
      "torch.Size([1, 240])\n",
      "torch.Size([1, 173])\n",
      "torch.Size([1, 474])\n",
      "torch.Size([1, 335])\n",
      "torch.Size([1, 163])\n",
      "torch.Size([1, 134])\n",
      "torch.Size([1, 570])\n",
      "torch.Size([1, 223])\n",
      "torch.Size([1, 183])\n",
      "torch.Size([1, 285])\n",
      "torch.Size([1, 144])\n",
      "torch.Size([1, 524])\n",
      "torch.Size([1, 304])\n",
      "torch.Size([1, 201])\n",
      "torch.Size([1, 54])\n",
      "torch.Size([1, 185])\n",
      "torch.Size([1, 238])\n",
      "torch.Size([1, 277])\n",
      "torch.Size([1, 88])\n",
      "torch.Size([1, 525])\n",
      "torch.Size([1, 217])\n",
      "torch.Size([1, 395])\n",
      "torch.Size([1, 139])\n",
      "torch.Size([1, 124])\n",
      "torch.Size([1, 543])\n",
      "torch.Size([1, 153])\n",
      "torch.Size([1, 71])\n",
      "torch.Size([1, 405])\n",
      "torch.Size([1, 221])\n",
      "torch.Size([1, 106])\n",
      "torch.Size([1, 1024])\n",
      "torch.Size([1, 205])\n",
      "torch.Size([1, 122])\n",
      "torch.Size([1, 191])\n",
      "torch.Size([1, 174])\n",
      "torch.Size([1, 73])\n",
      "torch.Size([1, 184])\n",
      "torch.Size([1, 406])\n",
      "torch.Size([1, 278])\n",
      "torch.Size([1, 97])\n",
      "torch.Size([1, 526])\n",
      "torch.Size([1, 135])\n",
      "torch.Size([1, 609])\n",
      "torch.Size([1, 308])\n",
      "torch.Size([1, 312])\n",
      "torch.Size([1, 166])\n",
      "torch.Size([1, 204])\n",
      "torch.Size([1, 331])\n",
      "torch.Size([1, 522])\n",
      "torch.Size([1, 647])\n",
      "torch.Size([1, 418])\n",
      "torch.Size([1, 292])\n",
      "torch.Size([1, 321])\n",
      "torch.Size([1, 888])\n",
      "torch.Size([1, 191])\n",
      "torch.Size([1, 201])\n",
      "torch.Size([1, 335])\n",
      "torch.Size([1, 1024])\n",
      "torch.Size([1, 456])\n",
      "torch.Size([1, 143])\n",
      "torch.Size([1, 176])\n",
      "torch.Size([1, 152])\n",
      "torch.Size([1, 227])\n",
      "torch.Size([1, 369])\n",
      "torch.Size([1, 206])\n",
      "torch.Size([1, 625])\n",
      "torch.Size([1, 165])\n",
      "torch.Size([1, 468])\n",
      "torch.Size([1, 272])\n",
      "torch.Size([1, 469])\n",
      "torch.Size([1, 185])\n",
      "torch.Size([1, 157])\n",
      "torch.Size([1, 57])\n",
      "torch.Size([1, 286])\n",
      "torch.Size([1, 213])\n",
      "torch.Size([1, 726])\n",
      "torch.Size([1, 730])\n",
      "torch.Size([1, 197])\n",
      "torch.Size([1, 655])\n",
      "torch.Size([1, 164])\n",
      "torch.Size([1, 277])\n",
      "torch.Size([1, 281])\n",
      "torch.Size([1, 51])\n",
      "torch.Size([1, 483])\n",
      "torch.Size([1, 250])\n",
      "torch.Size([1, 374])\n",
      "torch.Size([1, 156])\n",
      "torch.Size([1, 139])\n",
      "torch.Size([1, 91])\n",
      "torch.Size([1, 211])\n",
      "torch.Size([1, 307])\n",
      "torch.Size([1, 174])\n",
      "torch.Size([1, 179])\n",
      "torch.Size([1, 426])\n",
      "torch.Size([1, 199])\n",
      "torch.Size([1, 211])\n",
      "torch.Size([1, 134])\n",
      "torch.Size([1, 278])\n",
      "torch.Size([1, 178])\n",
      "torch.Size([1, 157])\n",
      "torch.Size([1, 600])\n",
      "torch.Size([1, 199])\n",
      "torch.Size([1, 137])\n",
      "torch.Size([1, 131])\n",
      "torch.Size([1, 292])\n",
      "torch.Size([1, 432])\n",
      "torch.Size([1, 147])\n",
      "torch.Size([1, 188])\n",
      "torch.Size([1, 317])\n",
      "torch.Size([1, 143])\n",
      "torch.Size([1, 504])\n",
      "torch.Size([1, 333])\n",
      "torch.Size([1, 382])\n",
      "torch.Size([1, 68])\n",
      "torch.Size([1, 697])\n",
      "torch.Size([1, 88])\n",
      "torch.Size([1, 160])\n",
      "torch.Size([1, 128])\n",
      "torch.Size([1, 142])\n",
      "torch.Size([1, 226])\n",
      "torch.Size([1, 165])\n",
      "torch.Size([1, 352])\n",
      "torch.Size([1, 274])\n",
      "torch.Size([1, 150])\n",
      "torch.Size([1, 136])\n",
      "torch.Size([1, 313])\n",
      "torch.Size([1, 91])\n",
      "torch.Size([1, 76])\n",
      "torch.Size([1, 59])\n",
      "torch.Size([1, 183])\n",
      "torch.Size([1, 268])\n",
      "torch.Size([1, 251])\n",
      "torch.Size([1, 984])\n",
      "torch.Size([1, 216])\n",
      "torch.Size([1, 429])\n",
      "torch.Size([1, 172])\n",
      "torch.Size([1, 432])\n",
      "torch.Size([1, 385])\n",
      "torch.Size([1, 181])\n",
      "torch.Size([1, 205])\n",
      "torch.Size([1, 62])\n",
      "torch.Size([1, 620])\n",
      "torch.Size([1, 348])\n",
      "torch.Size([1, 147])\n",
      "torch.Size([1, 332])\n",
      "torch.Size([1, 638])\n",
      "torch.Size([1, 231])\n",
      "torch.Size([1, 163])\n",
      "torch.Size([1, 227])\n",
      "torch.Size([1, 938])\n",
      "torch.Size([1, 159])\n",
      "torch.Size([1, 149])\n",
      "torch.Size([1, 132])\n",
      "torch.Size([1, 75])\n",
      "torch.Size([1, 106])\n",
      "torch.Size([1, 214])\n",
      "torch.Size([1, 66])\n",
      "torch.Size([1, 221])\n",
      "torch.Size([1, 506])\n",
      "torch.Size([1, 136])\n",
      "torch.Size([1, 137])\n",
      "torch.Size([1, 181])\n",
      "torch.Size([1, 138])\n",
      "torch.Size([1, 221])\n",
      "torch.Size([1, 132])\n",
      "torch.Size([1, 156])\n",
      "torch.Size([1, 835])\n",
      "torch.Size([1, 99])\n",
      "torch.Size([1, 379])\n",
      "torch.Size([1, 183])\n",
      "torch.Size([1, 348])\n",
      "torch.Size([1, 395])\n",
      "torch.Size([1, 96])\n",
      "torch.Size([1, 47])\n",
      "torch.Size([1, 359])\n",
      "torch.Size([1, 169])\n",
      "torch.Size([1, 355])\n",
      "torch.Size([1, 458])\n",
      "torch.Size([1, 212])\n",
      "torch.Size([1, 236])\n",
      "torch.Size([1, 50])\n",
      "torch.Size([1, 248])\n",
      "torch.Size([1, 178])\n",
      "torch.Size([1, 160])\n",
      "torch.Size([1, 145])\n",
      "torch.Size([1, 170])\n",
      "torch.Size([1, 399])\n",
      "torch.Size([1, 180])\n",
      "torch.Size([1, 216])\n",
      "torch.Size([1, 183])\n",
      "torch.Size([1, 94])\n",
      "torch.Size([1, 1024])\n",
      "torch.Size([1, 283])\n",
      "torch.Size([1, 130])\n",
      "torch.Size([1, 380])\n",
      "torch.Size([1, 320])\n",
      "torch.Size([1, 216])\n",
      "torch.Size([1, 185])\n",
      "torch.Size([1, 206])\n",
      "torch.Size([1, 448])\n",
      "torch.Size([1, 227])\n",
      "torch.Size([1, 233])\n",
      "torch.Size([1, 383])\n",
      "torch.Size([1, 382])\n",
      "torch.Size([1, 178])\n",
      "torch.Size([1, 196])\n",
      "torch.Size([1, 201])\n",
      "torch.Size([1, 235])\n",
      "torch.Size([1, 92])\n",
      "torch.Size([1, 227])\n",
      "torch.Size([1, 188])\n",
      "torch.Size([1, 261])\n",
      "torch.Size([1, 164])\n",
      "torch.Size([1, 91])\n",
      "torch.Size([1, 150])\n",
      "torch.Size([1, 155])\n",
      "torch.Size([1, 408])\n",
      "torch.Size([1, 116])\n",
      "torch.Size([1, 230])\n",
      "torch.Size([1, 220])\n",
      "torch.Size([1, 191])\n",
      "torch.Size([1, 341])\n",
      "torch.Size([1, 147])\n",
      "torch.Size([1, 152])\n",
      "torch.Size([1, 134])\n",
      "torch.Size([1, 168])\n",
      "torch.Size([1, 340])\n",
      "torch.Size([1, 488])\n",
      "torch.Size([1, 301])\n",
      "torch.Size([1, 444])\n",
      "torch.Size([1, 216])\n",
      "torch.Size([1, 127])\n",
      "torch.Size([1, 259])\n",
      "torch.Size([1, 91])\n",
      "torch.Size([1, 94])\n",
      "torch.Size([1, 310])\n",
      "torch.Size([1, 550])\n",
      "torch.Size([1, 96])\n",
      "torch.Size([1, 1024])\n",
      "torch.Size([1, 146])\n",
      "torch.Size([1, 197])\n",
      "torch.Size([1, 215])\n",
      "torch.Size([1, 724])\n"
     ]
    }
   ],
   "source": [
    "compute_mean_ablation_modified_accuracy(model, small_eval_dataset, target_token_ids=[15])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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.9.16"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}