feature_clustering_7_7_1153am_default_submission.ipynb

{"nbformat":4,"nbformat_minor":0,"metadata":{"colab":{"provenance":[{"file_id":"1qUdsegUIZILv5DDLVStH6GTIz-ZhFmKH","timestamp":1720367486888},{"file_id":"1yCZvzm71J9B_KGt9O8xge2fgw1pvkRMQ","timestamp":1718704649080}],"gpuType":"T4","machine_shape":"hm"},"kernelspec":{"name":"python3","display_name":"Python 3"},"language_info":{"name":"python"},"accelerator":"GPU"},"cells":[{"cell_type":"markdown","source":["# Intro\n","Typically we use cosine similarity as a proxy for relatedness of features.  \n","Normally to find nearest neighbors you would have to check all pairwise distances.  \n","\n","By using a hierarchical clustering method in advance, we can precompute these to make retrieval near instant at inference time.\n","\n","\n","This indexing data over feature space takes minimal space: about 9 MB for ~300k features across the 12 layers of GPT2-small.  \n","It takes about 4 minutes per layer of ~25k features to compute, or about 45 minutes for all layers. These were all saved to a pkl."],"metadata":{"id":"vXMK1fds7cpH"}},{"cell_type":"code","source":["# mount your drive\n","from google.colab import drive\n","drive.mount('/content/drive')"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"hXMRYgFAjAsY","executionInfo":{"status":"ok","timestamp":1720367114057,"user_tz":240,"elapsed":1883,"user":{"displayName":"sinem erisken","userId":"10098141094175015247"}},"outputId":"ca8263a5-df4a-47e4-8e87-8f619d7634b5"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["Drive already mounted at /content/drive; to attempt to forcibly remount, call drive.mount(\"/content/drive\", force_remount=True).\n"]}]},{"cell_type":"code","source":["# Need to install it in this order\n","# Numpy causes problems with different version dependencies\n","# But 1.25.2 seems to work for everything\n","!pip install circuitsvis\n","!pip install nnsight transformer_lens sae-lens==3.9.0 bitsandbytes\n","# !pip install numpy==1.25.2"],"metadata":{"colab":{"base_uri":"https://localhost:8080/","height":1000},"collapsed":true,"id":"yFWpFjurmlrc","outputId":"5edff2ad-60fd-4eb0-982f-8c15ff6395b4","executionInfo":{"status":"ok","timestamp":1720367167040,"user_tz":240,"elapsed":52985,"user":{"displayName":"sinem erisken","userId":"10098141094175015247"}}},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["\u001b[33mWARNING: Ignoring invalid distribution -orch (/usr/local/lib/python3.10/dist-packages)\u001b[0m\u001b[33m\n","\u001b[0mRequirement already satisfied: circuitsvis in /usr/local/lib/python3.10/dist-packages (1.43.2)\n","Requirement already satisfied: importlib-metadata>=5.1.0 in /usr/local/lib/python3.10/dist-packages (from circuitsvis) (8.0.0)\n","Requirement already satisfied: numpy>=1.24 in /usr/local/lib/python3.10/dist-packages (from circuitsvis) (1.25.2)\n","Requirement already satisfied: nvidia-cublas-cu12==12.1.3.1 in /usr/local/lib/python3.10/dist-packages (from circuitsvis) (12.1.3.1)\n","Requirement already satisfied: nvidia-cuda-cupti-cu12==12.1.105 in /usr/local/lib/python3.10/dist-packages (from circuitsvis) (12.1.105)\n","Requirement already satisfied: nvidia-cuda-nvrtc-cu12==12.1.105 in /usr/local/lib/python3.10/dist-packages (from circuitsvis) (12.1.105)\n","Requirement already satisfied: nvidia-cuda-runtime-cu12==12.1.105 in /usr/local/lib/python3.10/dist-packages (from circuitsvis) (12.1.105)\n","Requirement already satisfied: nvidia-cudnn-cu12==8.9.2.26 in /usr/local/lib/python3.10/dist-packages (from circuitsvis) (8.9.2.26)\n","Requirement already satisfied: nvidia-cufft-cu12==11.0.2.54 in /usr/local/lib/python3.10/dist-packages (from circuitsvis) (11.0.2.54)\n","Requirement already satisfied: nvidia-curand-cu12==10.3.2.106 in /usr/local/lib/python3.10/dist-packages (from circuitsvis) (10.3.2.106)\n","Requirement already satisfied: nvidia-cusolver-cu12==11.4.5.107 in /usr/local/lib/python3.10/dist-packages (from circuitsvis) (11.4.5.107)\n","Requirement already satisfied: nvidia-cusparse-cu12==12.1.0.106 in /usr/local/lib/python3.10/dist-packages (from circuitsvis) (12.1.0.106)\n","Requirement already satisfied: nvidia-nccl-cu12==2.18.1 in /usr/local/lib/python3.10/dist-packages (from circuitsvis) (2.18.1)\n","Requirement already satisfied: nvidia-nvtx-cu12==12.1.105 in /usr/local/lib/python3.10/dist-packages (from circuitsvis) (12.1.105)\n","Requirement already satisfied: torch>=1.10 in /usr/local/lib/python3.10/dist-packages (from circuitsvis) (2.1.2)\n","Requirement already satisfied: triton==2.1.0 in /usr/local/lib/python3.10/dist-packages (from circuitsvis) (2.1.0)\n","Requirement already satisfied: nvidia-nvjitlink-cu12 in /usr/local/lib/python3.10/dist-packages (from nvidia-cusolver-cu12==11.4.5.107->circuitsvis) (12.5.82)\n","Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from triton==2.1.0->circuitsvis) (3.15.4)\n","Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.10/dist-packages (from importlib-metadata>=5.1.0->circuitsvis) (3.19.2)\n","Requirement already satisfied: typing-extensions in /usr/local/lib/python3.10/dist-packages (from torch>=1.10->circuitsvis) (4.12.2)\n","Requirement already satisfied: sympy in /usr/local/lib/python3.10/dist-packages (from torch>=1.10->circuitsvis) (1.12.1)\n","Requirement already satisfied: networkx in /usr/local/lib/python3.10/dist-packages (from torch>=1.10->circuitsvis) (3.3)\n","Requirement already satisfied: jinja2 in /usr/local/lib/python3.10/dist-packages (from torch>=1.10->circuitsvis) (3.1.4)\n","Requirement already satisfied: fsspec in /usr/local/lib/python3.10/dist-packages (from torch>=1.10->circuitsvis) (2023.6.0)\n","Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/dist-packages (from jinja2->torch>=1.10->circuitsvis) (2.1.5)\n","Requirement already satisfied: mpmath<1.4.0,>=1.1.0 in /usr/local/lib/python3.10/dist-packages (from sympy->torch>=1.10->circuitsvis) (1.3.0)\n","\u001b[33mWARNING: Ignoring invalid distribution -orch (/usr/local/lib/python3.10/dist-packages)\u001b[0m\u001b[33m\n","\u001b[0m\u001b[33mWARNING: Ignoring invalid distribution -orch (/usr/local/lib/python3.10/dist-packages)\u001b[0m\u001b[33m\n","\u001b[0mRequirement already satisfied: nnsight in /usr/local/lib/python3.10/dist-packages (0.2.19)\n","Requirement already satisfied: transformer_lens in /usr/local/lib/python3.10/dist-packages (1.19.0)\n","Requirement already satisfied: sae-lens==3.9.0 in /usr/local/lib/python3.10/dist-packages (3.9.0)\n","Requirement already satisfied: bitsandbytes in /usr/local/lib/python3.10/dist-packages (0.43.1)\n","Requirement already satisfied: automated-interpretability<0.0.4,>=0.0.3 in /usr/local/lib/python3.10/dist-packages (from sae-lens==3.9.0) (0.0.3)\n","Requirement already satisfied: babe<0.0.8,>=0.0.7 in /usr/local/lib/python3.10/dist-packages (from sae-lens==3.9.0) (0.0.7)\n","Requirement already satisfied: datasets<3.0.0,>=2.17.1 in /usr/local/lib/python3.10/dist-packages (from sae-lens==3.9.0) (2.20.0)\n","Requirement already satisfied: matplotlib<4.0.0,>=3.8.3 in /usr/local/lib/python3.10/dist-packages (from sae-lens==3.9.0) (3.9.1)\n","Requirement already satisfied: matplotlib-inline<0.2.0,>=0.1.6 in /usr/local/lib/python3.10/dist-packages (from sae-lens==3.9.0) (0.1.7)\n","Requirement already satisfied: nltk<4.0.0,>=3.8.1 in /usr/local/lib/python3.10/dist-packages (from sae-lens==3.9.0) (3.8.1)\n","Requirement already satisfied: plotly<6.0.0,>=5.19.0 in /usr/local/lib/python3.10/dist-packages (from sae-lens==3.9.0) (5.22.0)\n","Requirement already satisfied: plotly-express<0.5.0,>=0.4.1 in /usr/local/lib/python3.10/dist-packages (from sae-lens==3.9.0) (0.4.1)\n","Requirement already satisfied: pytest-profiling<2.0.0,>=1.7.0 in /usr/local/lib/python3.10/dist-packages (from sae-lens==3.9.0) (1.7.0)\n","Requirement already satisfied: python-dotenv<2.0.0,>=1.0.1 in /usr/local/lib/python3.10/dist-packages (from sae-lens==3.9.0) (1.0.1)\n","Requirement already satisfied: pyyaml<7.0.0,>=6.0.1 in /usr/local/lib/python3.10/dist-packages (from sae-lens==3.9.0) (6.0.1)\n","Requirement already satisfied: pyzmq==26.0.0 in /usr/local/lib/python3.10/dist-packages (from sae-lens==3.9.0) (26.0.0)\n","Requirement already satisfied: sae-vis<0.3.0,>=0.2.18 in /usr/local/lib/python3.10/dist-packages (from sae-lens==3.9.0) (0.2.19)\n","Requirement already satisfied: safetensors<0.5.0,>=0.4.2 in /usr/local/lib/python3.10/dist-packages (from sae-lens==3.9.0) (0.4.3)\n","Requirement already satisfied: transformers<5.0.0,>=4.38.1 in /usr/local/lib/python3.10/dist-packages (from sae-lens==3.9.0) (4.41.2)\n","Requirement already satisfied: typer<0.13.0,>=0.12.3 in /usr/local/lib/python3.10/dist-packages (from sae-lens==3.9.0) (0.12.3)\n","Requirement already satisfied: zstandard<0.23.0,>=0.22.0 in /usr/local/lib/python3.10/dist-packages (from sae-lens==3.9.0) (0.22.0)\n","Requirement already satisfied: protobuf in /usr/local/lib/python3.10/dist-packages (from nnsight) (3.20.3)\n","Requirement already satisfied: python-socketio[client] in /usr/local/lib/python3.10/dist-packages (from nnsight) (5.11.3)\n","Requirement already satisfied: tokenizers>=0.13.0 in /usr/local/lib/python3.10/dist-packages (from nnsight) (0.19.1)\n","Requirement already satisfied: pydantic>=2.4.0 in /usr/local/lib/python3.10/dist-packages (from nnsight) (2.8.0)\n","Requirement already satisfied: torch>=2.1.0 in /usr/local/lib/python3.10/dist-packages (from nnsight) (2.1.2)\n","Requirement already satisfied: sentencepiece in /usr/local/lib/python3.10/dist-packages (from nnsight) (0.1.99)\n","Requirement already satisfied: torchvision in /usr/local/lib/python3.10/dist-packages (from nnsight) (0.18.0+cu121)\n","Requirement already satisfied: accelerate in /usr/local/lib/python3.10/dist-packages (from nnsight) (0.32.1)\n","Requirement already satisfied: diffusers in /usr/local/lib/python3.10/dist-packages (from nnsight) (0.29.2)\n","Requirement already satisfied: einops in /usr/local/lib/python3.10/dist-packages (from nnsight) (0.7.0)\n","Requirement already satisfied: beartype<0.15.0,>=0.14.1 in /usr/local/lib/python3.10/dist-packages (from transformer_lens) (0.14.1)\n","Requirement already satisfied: better-abc<0.0.4,>=0.0.3 in /usr/local/lib/python3.10/dist-packages (from transformer_lens) (0.0.3)\n","Requirement already satisfied: fancy-einsum>=0.0.3 in /usr/local/lib/python3.10/dist-packages (from transformer_lens) (0.0.3)\n","Requirement already satisfied: jaxtyping>=0.2.11 in /usr/local/lib/python3.10/dist-packages (from transformer_lens) (0.2.31)\n","Requirement already satisfied: numpy>=1.24 in /usr/local/lib/python3.10/dist-packages (from transformer_lens) (1.25.2)\n","Requirement already satisfied: pandas>=1.1.5 in /usr/local/lib/python3.10/dist-packages (from transformer_lens) (2.0.3)\n","Requirement already satisfied: rich>=12.6.0 in /usr/local/lib/python3.10/dist-packages (from transformer_lens) (13.7.1)\n","Requirement already satisfied: tqdm>=4.64.1 in /usr/local/lib/python3.10/dist-packages (from transformer_lens) (4.66.4)\n","Requirement already satisfied: typing-extensions in /usr/local/lib/python3.10/dist-packages (from transformer_lens) (4.12.2)\n","Requirement already satisfied: wandb>=0.13.5 in /usr/local/lib/python3.10/dist-packages (from transformer_lens) (0.17.4)\n","Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.10/dist-packages (from accelerate->nnsight) (24.1)\n","Requirement already satisfied: psutil in /usr/local/lib/python3.10/dist-packages (from accelerate->nnsight) (5.9.5)\n","Requirement already satisfied: huggingface-hub in /usr/local/lib/python3.10/dist-packages (from accelerate->nnsight) (0.23.4)\n","Requirement already satisfied: blobfile<3.0.0,>=2.1.1 in /usr/local/lib/python3.10/dist-packages (from automated-interpretability<0.0.4,>=0.0.3->sae-lens==3.9.0) (2.1.1)\n","Requirement already satisfied: boostedblob<0.16.0,>=0.15.3 in /usr/local/lib/python3.10/dist-packages (from automated-interpretability<0.0.4,>=0.0.3->sae-lens==3.9.0) (0.15.3)\n","Requirement already satisfied: httpx<0.28.0,>=0.27.0 in /usr/local/lib/python3.10/dist-packages (from automated-interpretability<0.0.4,>=0.0.3->sae-lens==3.9.0) (0.27.0)\n","Collecting numpy>=1.24 (from transformer_lens)\n","  Using cached numpy-1.26.4-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (18.2 MB)\n","Requirement already satisfied: orjson<4.0.0,>=3.10.1 in /usr/local/lib/python3.10/dist-packages (from automated-interpretability<0.0.4,>=0.0.3->sae-lens==3.9.0) (3.10.6)\n","Requirement already satisfied: pytest<9.0.0,>=8.1.2 in /usr/local/lib/python3.10/dist-packages (from automated-interpretability<0.0.4,>=0.0.3->sae-lens==3.9.0) (8.2.2)\n","Requirement already satisfied: scikit-learn<2.0.0,>=1.4.2 in /usr/local/lib/python3.10/dist-packages (from automated-interpretability<0.0.4,>=0.0.3->sae-lens==3.9.0) (1.5.1)\n","Requirement already satisfied: tiktoken<0.7.0,>=0.6.0 in /usr/local/lib/python3.10/dist-packages (from automated-interpretability<0.0.4,>=0.0.3->sae-lens==3.9.0) (0.6.0)\n","Requirement already satisfied: py2store in /usr/local/lib/python3.10/dist-packages (from babe<0.0.8,>=0.0.7->sae-lens==3.9.0) (0.1.20)\n","Requirement already satisfied: graze in /usr/local/lib/python3.10/dist-packages (from babe<0.0.8,>=0.0.7->sae-lens==3.9.0) (0.1.17)\n","Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from datasets<3.0.0,>=2.17.1->sae-lens==3.9.0) (3.15.4)\n","Requirement already satisfied: pyarrow>=15.0.0 in /usr/local/lib/python3.10/dist-packages (from datasets<3.0.0,>=2.17.1->sae-lens==3.9.0) (16.1.0)\n","Requirement already satisfied: pyarrow-hotfix in /usr/local/lib/python3.10/dist-packages (from datasets<3.0.0,>=2.17.1->sae-lens==3.9.0) (0.6)\n","Requirement already satisfied: dill<0.3.9,>=0.3.0 in /usr/local/lib/python3.10/dist-packages (from datasets<3.0.0,>=2.17.1->sae-lens==3.9.0) (0.3.8)\n","Requirement already satisfied: requests>=2.32.2 in /usr/local/lib/python3.10/dist-packages (from datasets<3.0.0,>=2.17.1->sae-lens==3.9.0) (2.32.3)\n","Requirement already satisfied: xxhash in /usr/local/lib/python3.10/dist-packages (from datasets<3.0.0,>=2.17.1->sae-lens==3.9.0) (3.4.1)\n","Requirement already satisfied: multiprocess in /usr/local/lib/python3.10/dist-packages (from datasets<3.0.0,>=2.17.1->sae-lens==3.9.0) (0.70.16)\n","Requirement already satisfied: fsspec[http]<=2024.5.0,>=2023.1.0 in /usr/local/lib/python3.10/dist-packages (from datasets<3.0.0,>=2.17.1->sae-lens==3.9.0) (2023.6.0)\n","Requirement already satisfied: aiohttp in /usr/local/lib/python3.10/dist-packages (from datasets<3.0.0,>=2.17.1->sae-lens==3.9.0) (3.9.5)\n","Requirement already satisfied: typeguard==2.13.3 in /usr/local/lib/python3.10/dist-packages (from jaxtyping>=0.2.11->transformer_lens) (2.13.3)\n","Requirement already satisfied: contourpy>=1.0.1 in /usr/local/lib/python3.10/dist-packages (from matplotlib<4.0.0,>=3.8.3->sae-lens==3.9.0) (1.2.1)\n","Requirement already satisfied: cycler>=0.10 in /usr/local/lib/python3.10/dist-packages (from matplotlib<4.0.0,>=3.8.3->sae-lens==3.9.0) (0.12.1)\n","Requirement already satisfied: fonttools>=4.22.0 in /usr/local/lib/python3.10/dist-packages (from matplotlib<4.0.0,>=3.8.3->sae-lens==3.9.0) (4.53.0)\n","Requirement already satisfied: kiwisolver>=1.3.1 in /usr/local/lib/python3.10/dist-packages (from matplotlib<4.0.0,>=3.8.3->sae-lens==3.9.0) (1.4.5)\n","Requirement already satisfied: pillow>=8 in /usr/local/lib/python3.10/dist-packages (from matplotlib<4.0.0,>=3.8.3->sae-lens==3.9.0) (9.4.0)\n","Requirement already satisfied: pyparsing>=2.3.1 in /usr/local/lib/python3.10/dist-packages (from matplotlib<4.0.0,>=3.8.3->sae-lens==3.9.0) (3.1.2)\n","Requirement already satisfied: python-dateutil>=2.7 in /usr/local/lib/python3.10/dist-packages (from matplotlib<4.0.0,>=3.8.3->sae-lens==3.9.0) (2.8.2)\n","Requirement already satisfied: traitlets in /usr/local/lib/python3.10/dist-packages (from matplotlib-inline<0.2.0,>=0.1.6->sae-lens==3.9.0) (5.7.1)\n","Requirement already satisfied: click in /usr/local/lib/python3.10/dist-packages (from nltk<4.0.0,>=3.8.1->sae-lens==3.9.0) (8.1.7)\n","Requirement already satisfied: joblib in /usr/local/lib/python3.10/dist-packages (from nltk<4.0.0,>=3.8.1->sae-lens==3.9.0) (1.4.2)\n","Requirement already satisfied: regex>=2021.8.3 in /usr/local/lib/python3.10/dist-packages (from nltk<4.0.0,>=3.8.1->sae-lens==3.9.0) (2024.5.15)\n","Requirement already satisfied: pytz>=2020.1 in /usr/local/lib/python3.10/dist-packages (from pandas>=1.1.5->transformer_lens) (2023.4)\n","Requirement already satisfied: tzdata>=2022.1 in /usr/local/lib/python3.10/dist-packages (from pandas>=1.1.5->transformer_lens) (2024.1)\n","Requirement already satisfied: tenacity>=6.2.0 in /usr/local/lib/python3.10/dist-packages (from plotly<6.0.0,>=5.19.0->sae-lens==3.9.0) (8.4.2)\n","Requirement already satisfied: statsmodels>=0.9.0 in /usr/local/lib/python3.10/dist-packages (from plotly-express<0.5.0,>=0.4.1->sae-lens==3.9.0) (0.14.2)\n","Requirement already satisfied: scipy>=0.18 in /usr/local/lib/python3.10/dist-packages (from plotly-express<0.5.0,>=0.4.1->sae-lens==3.9.0) (1.11.4)\n","Requirement already satisfied: patsy>=0.5 in /usr/local/lib/python3.10/dist-packages (from plotly-express<0.5.0,>=0.4.1->sae-lens==3.9.0) (0.5.6)\n","Requirement already satisfied: annotated-types>=0.4.0 in /usr/local/lib/python3.10/dist-packages (from pydantic>=2.4.0->nnsight) (0.7.0)\n","Requirement already satisfied: pydantic-core==2.20.0 in /usr/local/lib/python3.10/dist-packages (from pydantic>=2.4.0->nnsight) (2.20.0)\n","Requirement already satisfied: six in /usr/local/lib/python3.10/dist-packages (from pytest-profiling<2.0.0,>=1.7.0->sae-lens==3.9.0) (1.16.0)\n","Requirement already satisfied: gprof2dot in /usr/local/lib/python3.10/dist-packages (from pytest-profiling<2.0.0,>=1.7.0->sae-lens==3.9.0) (2024.6.6)\n","Requirement already satisfied: markdown-it-py>=2.2.0 in /usr/local/lib/python3.10/dist-packages (from rich>=12.6.0->transformer_lens) (3.0.0)\n","Requirement already satisfied: pygments<3.0.0,>=2.13.0 in /usr/local/lib/python3.10/dist-packages (from rich>=12.6.0->transformer_lens) (2.16.1)\n","Requirement already satisfied: dataclasses-json<0.7.0,>=0.6.4 in /usr/local/lib/python3.10/dist-packages (from sae-vis<0.3.0,>=0.2.18->sae-lens==3.9.0) (0.6.7)\n","Requirement already satisfied: eindex-callum<0.2.0,>=0.1.0 in /usr/local/lib/python3.10/dist-packages (from sae-vis<0.3.0,>=0.2.18->sae-lens==3.9.0) (0.1.1)\n","Requirement already satisfied: sympy in /usr/local/lib/python3.10/dist-packages (from torch>=2.1.0->nnsight) (1.12.1)\n","Requirement already satisfied: networkx in /usr/local/lib/python3.10/dist-packages (from torch>=2.1.0->nnsight) (3.3)\n","Requirement already satisfied: jinja2 in /usr/local/lib/python3.10/dist-packages (from torch>=2.1.0->nnsight) (3.1.4)\n","Requirement already satisfied: nvidia-cuda-nvrtc-cu12==12.1.105 in /usr/local/lib/python3.10/dist-packages (from torch>=2.1.0->nnsight) (12.1.105)\n","Requirement already satisfied: nvidia-cuda-runtime-cu12==12.1.105 in /usr/local/lib/python3.10/dist-packages (from torch>=2.1.0->nnsight) (12.1.105)\n","Requirement already satisfied: nvidia-cuda-cupti-cu12==12.1.105 in /usr/local/lib/python3.10/dist-packages (from torch>=2.1.0->nnsight) (12.1.105)\n","Requirement already satisfied: nvidia-cudnn-cu12==8.9.2.26 in /usr/local/lib/python3.10/dist-packages (from torch>=2.1.0->nnsight) (8.9.2.26)\n","Requirement already satisfied: nvidia-cublas-cu12==12.1.3.1 in /usr/local/lib/python3.10/dist-packages (from torch>=2.1.0->nnsight) (12.1.3.1)\n","Requirement already satisfied: nvidia-cufft-cu12==11.0.2.54 in /usr/local/lib/python3.10/dist-packages (from torch>=2.1.0->nnsight) (11.0.2.54)\n","Requirement already satisfied: nvidia-curand-cu12==10.3.2.106 in /usr/local/lib/python3.10/dist-packages (from torch>=2.1.0->nnsight) (10.3.2.106)\n","Requirement already satisfied: nvidia-cusolver-cu12==11.4.5.107 in /usr/local/lib/python3.10/dist-packages (from torch>=2.1.0->nnsight) (11.4.5.107)\n","Requirement already satisfied: nvidia-cusparse-cu12==12.1.0.106 in /usr/local/lib/python3.10/dist-packages (from torch>=2.1.0->nnsight) (12.1.0.106)\n","Requirement already satisfied: nvidia-nccl-cu12==2.18.1 in /usr/local/lib/python3.10/dist-packages (from torch>=2.1.0->nnsight) (2.18.1)\n","Requirement already satisfied: nvidia-nvtx-cu12==12.1.105 in /usr/local/lib/python3.10/dist-packages (from torch>=2.1.0->nnsight) (12.1.105)\n","Requirement already satisfied: triton==2.1.0 in /usr/local/lib/python3.10/dist-packages (from torch>=2.1.0->nnsight) (2.1.0)\n","Requirement already satisfied: nvidia-nvjitlink-cu12 in /usr/local/lib/python3.10/dist-packages (from nvidia-cusolver-cu12==11.4.5.107->torch>=2.1.0->nnsight) (12.5.82)\n","Requirement already satisfied: shellingham>=1.3.0 in /usr/local/lib/python3.10/dist-packages (from typer<0.13.0,>=0.12.3->sae-lens==3.9.0) (1.5.4)\n","Requirement already satisfied: docker-pycreds>=0.4.0 in /usr/local/lib/python3.10/dist-packages (from wandb>=0.13.5->transformer_lens) (0.4.0)\n","Requirement already satisfied: gitpython!=3.1.29,>=1.0.0 in /usr/local/lib/python3.10/dist-packages (from wandb>=0.13.5->transformer_lens) (3.1.43)\n","Requirement already satisfied: platformdirs in /usr/local/lib/python3.10/dist-packages (from wandb>=0.13.5->transformer_lens) (4.2.2)\n","Requirement already satisfied: sentry-sdk>=1.0.0 in /usr/local/lib/python3.10/dist-packages (from wandb>=0.13.5->transformer_lens) (2.7.1)\n","Requirement already satisfied: setproctitle in /usr/local/lib/python3.10/dist-packages (from wandb>=0.13.5->transformer_lens) (1.3.3)\n","Requirement already satisfied: setuptools in /usr/local/lib/python3.10/dist-packages (from wandb>=0.13.5->transformer_lens) (67.7.2)\n","Requirement already satisfied: importlib-metadata in /usr/local/lib/python3.10/dist-packages (from diffusers->nnsight) (8.0.0)\n","Requirement already satisfied: bidict>=0.21.0 in /usr/local/lib/python3.10/dist-packages (from python-socketio[client]->nnsight) (0.23.1)\n","Requirement already satisfied: python-engineio>=4.8.0 in /usr/local/lib/python3.10/dist-packages (from python-socketio[client]->nnsight) (4.9.1)\n","Requirement already satisfied: websocket-client>=0.54.0 in /usr/local/lib/python3.10/dist-packages (from python-socketio[client]->nnsight) (1.8.0)\n","Collecting torch>=2.1.0 (from nnsight)\n","  Using cached torch-2.3.0-cp310-cp310-manylinux1_x86_64.whl (779.1 MB)\n","Collecting nvidia-nccl-cu12==2.20.5 (from torch>=2.1.0->nnsight)\n","  Using cached nvidia_nccl_cu12-2.20.5-py3-none-manylinux2014_x86_64.whl (176.2 MB)\n","Collecting triton==2.3.0 (from torch>=2.1.0->nnsight)\n","  Using cached triton-2.3.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (168.1 MB)\n","Requirement already satisfied: pycryptodomex~=3.8 in /usr/local/lib/python3.10/dist-packages (from blobfile<3.0.0,>=2.1.1->automated-interpretability<0.0.4,>=0.0.3->sae-lens==3.9.0) (3.20.0)\n","Requirement already satisfied: urllib3<3,>=1.25.3 in /usr/local/lib/python3.10/dist-packages (from blobfile<3.0.0,>=2.1.1->automated-interpretability<0.0.4,>=0.0.3->sae-lens==3.9.0) (2.0.7)\n","Requirement already satisfied: lxml~=4.9 in /usr/local/lib/python3.10/dist-packages (from blobfile<3.0.0,>=2.1.1->automated-interpretability<0.0.4,>=0.0.3->sae-lens==3.9.0) (4.9.4)\n","Requirement already satisfied: uvloop>=0.16.0 in /usr/local/lib/python3.10/dist-packages (from boostedblob<0.16.0,>=0.15.3->automated-interpretability<0.0.4,>=0.0.3->sae-lens==3.9.0) (0.19.0)\n","Requirement already satisfied: aiosignal>=1.1.2 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets<3.0.0,>=2.17.1->sae-lens==3.9.0) (1.3.1)\n","Requirement already satisfied: attrs>=17.3.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets<3.0.0,>=2.17.1->sae-lens==3.9.0) (23.2.0)\n","Requirement already satisfied: frozenlist>=1.1.1 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets<3.0.0,>=2.17.1->sae-lens==3.9.0) (1.4.1)\n","Requirement already satisfied: multidict<7.0,>=4.5 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets<3.0.0,>=2.17.1->sae-lens==3.9.0) (6.0.5)\n","Requirement already satisfied: yarl<2.0,>=1.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets<3.0.0,>=2.17.1->sae-lens==3.9.0) (1.9.4)\n","Requirement already satisfied: async-timeout<5.0,>=4.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets<3.0.0,>=2.17.1->sae-lens==3.9.0) (4.0.3)\n","Requirement already satisfied: marshmallow<4.0.0,>=3.18.0 in /usr/local/lib/python3.10/dist-packages (from dataclasses-json<0.7.0,>=0.6.4->sae-vis<0.3.0,>=0.2.18->sae-lens==3.9.0) (3.21.3)\n","Requirement already satisfied: typing-inspect<1,>=0.4.0 in /usr/local/lib/python3.10/dist-packages (from dataclasses-json<0.7.0,>=0.6.4->sae-vis<0.3.0,>=0.2.18->sae-lens==3.9.0) (0.9.0)\n","Requirement already satisfied: gitdb<5,>=4.0.1 in /usr/local/lib/python3.10/dist-packages (from gitpython!=3.1.29,>=1.0.0->wandb>=0.13.5->transformer_lens) (4.0.11)\n","Requirement already satisfied: anyio in /usr/local/lib/python3.10/dist-packages (from httpx<0.28.0,>=0.27.0->automated-interpretability<0.0.4,>=0.0.3->sae-lens==3.9.0) (3.7.1)\n","Requirement already satisfied: certifi in /usr/local/lib/python3.10/dist-packages (from httpx<0.28.0,>=0.27.0->automated-interpretability<0.0.4,>=0.0.3->sae-lens==3.9.0) (2024.6.2)\n","Requirement already satisfied: httpcore==1.* in /usr/local/lib/python3.10/dist-packages (from httpx<0.28.0,>=0.27.0->automated-interpretability<0.0.4,>=0.0.3->sae-lens==3.9.0) (1.0.5)\n","Requirement already satisfied: idna in /usr/local/lib/python3.10/dist-packages (from httpx<0.28.0,>=0.27.0->automated-interpretability<0.0.4,>=0.0.3->sae-lens==3.9.0) (3.7)\n","Requirement already satisfied: sniffio in /usr/local/lib/python3.10/dist-packages (from httpx<0.28.0,>=0.27.0->automated-interpretability<0.0.4,>=0.0.3->sae-lens==3.9.0) (1.3.1)\n","Requirement already satisfied: h11<0.15,>=0.13 in /usr/local/lib/python3.10/dist-packages (from httpcore==1.*->httpx<0.28.0,>=0.27.0->automated-interpretability<0.0.4,>=0.0.3->sae-lens==3.9.0) (0.14.0)\n","Requirement already satisfied: mdurl~=0.1 in /usr/local/lib/python3.10/dist-packages (from markdown-it-py>=2.2.0->rich>=12.6.0->transformer_lens) (0.1.2)\n","Requirement already satisfied: iniconfig in /usr/local/lib/python3.10/dist-packages (from pytest<9.0.0,>=8.1.2->automated-interpretability<0.0.4,>=0.0.3->sae-lens==3.9.0) (2.0.0)\n","Requirement already satisfied: pluggy<2.0,>=1.5 in /usr/local/lib/python3.10/dist-packages (from pytest<9.0.0,>=8.1.2->automated-interpretability<0.0.4,>=0.0.3->sae-lens==3.9.0) (1.5.0)\n","Requirement already satisfied: exceptiongroup>=1.0.0rc8 in /usr/local/lib/python3.10/dist-packages (from pytest<9.0.0,>=8.1.2->automated-interpretability<0.0.4,>=0.0.3->sae-lens==3.9.0) (1.2.1)\n","Requirement already satisfied: tomli>=1 in /usr/local/lib/python3.10/dist-packages (from pytest<9.0.0,>=8.1.2->automated-interpretability<0.0.4,>=0.0.3->sae-lens==3.9.0) (2.0.1)\n","Requirement already satisfied: simple-websocket>=0.10.0 in /usr/local/lib/python3.10/dist-packages (from python-engineio>=4.8.0->python-socketio[client]->nnsight) (1.0.0)\n","Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests>=2.32.2->datasets<3.0.0,>=2.17.1->sae-lens==3.9.0) (3.3.2)\n","Requirement already satisfied: threadpoolctl>=3.1.0 in /usr/local/lib/python3.10/dist-packages (from scikit-learn<2.0.0,>=1.4.2->automated-interpretability<0.0.4,>=0.0.3->sae-lens==3.9.0) (3.5.0)\n","Requirement already satisfied: dol in /usr/local/lib/python3.10/dist-packages (from graze->babe<0.0.8,>=0.0.7->sae-lens==3.9.0) (0.2.49)\n","Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.10/dist-packages (from importlib-metadata->diffusers->nnsight) (3.19.2)\n","Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/dist-packages (from jinja2->torch>=2.1.0->nnsight) (2.1.5)\n","Requirement already satisfied: config2py in /usr/local/lib/python3.10/dist-packages (from py2store->babe<0.0.8,>=0.0.7->sae-lens==3.9.0) (0.1.33)\n","Requirement already satisfied: importlib-resources in /usr/local/lib/python3.10/dist-packages (from py2store->babe<0.0.8,>=0.0.7->sae-lens==3.9.0) (6.4.0)\n","Requirement already satisfied: mpmath<1.4.0,>=1.1.0 in /usr/local/lib/python3.10/dist-packages (from sympy->torch>=2.1.0->nnsight) (1.3.0)\n","Requirement already satisfied: smmap<6,>=3.0.1 in /usr/local/lib/python3.10/dist-packages (from gitdb<5,>=4.0.1->gitpython!=3.1.29,>=1.0.0->wandb>=0.13.5->transformer_lens) (5.0.1)\n","Requirement already satisfied: wsproto in /usr/local/lib/python3.10/dist-packages (from simple-websocket>=0.10.0->python-engineio>=4.8.0->python-socketio[client]->nnsight) (1.2.0)\n","Requirement already satisfied: mypy-extensions>=0.3.0 in /usr/local/lib/python3.10/dist-packages (from typing-inspect<1,>=0.4.0->dataclasses-json<0.7.0,>=0.6.4->sae-vis<0.3.0,>=0.2.18->sae-lens==3.9.0) (1.0.0)\n","Requirement already satisfied: i2 in /usr/local/lib/python3.10/dist-packages (from config2py->py2store->babe<0.0.8,>=0.0.7->sae-lens==3.9.0) (0.1.17)\n","\u001b[33mWARNING: Ignoring invalid distribution -orch (/usr/local/lib/python3.10/dist-packages)\u001b[0m\u001b[33m\n","\u001b[0mInstalling collected packages: triton, nvidia-nccl-cu12, numpy, torch\n","  Attempting uninstall: triton\n","    Found existing installation: triton 2.1.0\n","    Uninstalling triton-2.1.0:\n","      Successfully uninstalled triton-2.1.0\n","  Attempting uninstall: nvidia-nccl-cu12\n","    Found existing installation: nvidia-nccl-cu12 2.18.1\n","    Uninstalling nvidia-nccl-cu12-2.18.1:\n","      Successfully uninstalled nvidia-nccl-cu12-2.18.1\n","  Attempting uninstall: numpy\n","    Found existing installation: numpy 1.25.2\n","    Uninstalling numpy-1.25.2:\n","      Successfully uninstalled numpy-1.25.2\n","  Attempting uninstall: torch\n","    Found existing installation: torch 2.1.2\n","    Uninstalling torch-2.1.2:\n","      Successfully uninstalled torch-2.1.2\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","circuitsvis 1.43.2 requires nvidia-nccl-cu12==2.18.1; platform_system == \"Linux\" and platform_machine == \"x86_64\", but you have nvidia-nccl-cu12 2.20.5 which is incompatible.\n","circuitsvis 1.43.2 requires triton==2.1.0; platform_system == \"Linux\" and platform_machine == \"x86_64\", but you have triton 2.3.0 which is incompatible.\n","cudf-cu12 24.4.1 requires pyarrow<15.0.0a0,>=14.0.1, but you have pyarrow 16.1.0 which is incompatible.\n","ibis-framework 8.0.0 requires pyarrow<16,>=2, but you have pyarrow 16.1.0 which is incompatible.\u001b[0m\u001b[31m\n","\u001b[0mSuccessfully installed numpy-1.26.4 nvidia-nccl-cu12-2.20.5 torch-2.3.0 triton-2.3.0\n"]},{"output_type":"display_data","data":{"application/vnd.colab-display-data+json":{"pip_warning":{"packages":["torch","torchgen","triton"]},"id":"fb40a09b02a9495c96564a57cdd8a493"}},"metadata":{}}]},{"cell_type":"code","source":["!pip install numpy==1.25.2"],"metadata":{"colab":{"base_uri":"https://localhost:8080/","height":411},"id":"ZJ4tStwrpN7x","executionInfo":{"status":"ok","timestamp":1720367177103,"user_tz":240,"elapsed":10069,"user":{"displayName":"sinem erisken","userId":"10098141094175015247"}},"outputId":"29f339ea-6b4f-4d11-818d-caa666ca43d8"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["\u001b[33mWARNING: Ignoring invalid distribution -orch (/usr/local/lib/python3.10/dist-packages)\u001b[0m\u001b[33m\n","\u001b[0mCollecting numpy==1.25.2\n","  Using cached numpy-1.25.2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (18.2 MB)\n","\u001b[33mWARNING: Ignoring invalid distribution -orch (/usr/local/lib/python3.10/dist-packages)\u001b[0m\u001b[33m\n","\u001b[0mInstalling collected packages: numpy\n","  Attempting uninstall: numpy\n","    Found existing installation: numpy 1.26.4\n","    Uninstalling numpy-1.26.4:\n","      Successfully uninstalled numpy-1.26.4\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","automated-interpretability 0.0.3 requires numpy<2.0.0,>=1.26.4, but you have numpy 1.25.2 which is incompatible.\n","circuitsvis 1.43.2 requires nvidia-nccl-cu12==2.18.1; platform_system == \"Linux\" and platform_machine == \"x86_64\", but you have nvidia-nccl-cu12 2.20.5 which is incompatible.\n","circuitsvis 1.43.2 requires triton==2.1.0; platform_system == \"Linux\" and platform_machine == \"x86_64\", but you have triton 2.3.0 which is incompatible.\n","cudf-cu12 24.4.1 requires pyarrow<15.0.0a0,>=14.0.1, but you have pyarrow 16.1.0 which is incompatible.\n","ibis-framework 8.0.0 requires pyarrow<16,>=2, but you have pyarrow 16.1.0 which is incompatible.\u001b[0m\u001b[31m\n","\u001b[0mSuccessfully installed numpy-1.25.2\n"]},{"output_type":"display_data","data":{"application/vnd.colab-display-data+json":{"pip_warning":{"packages":["numpy"]},"id":"142264e14a3745029e09a2ce209661f6"}},"metadata":{}}]},{"cell_type":"code","source":["import numpy as np\n","print(np.__version__)\n","!pip show numpy"],"metadata":{"id":"GG4tuplkCzFp","colab":{"base_uri":"https://localhost:8080/"},"executionInfo":{"status":"ok","timestamp":1720367183068,"user_tz":240,"elapsed":5971,"user":{"displayName":"sinem erisken","userId":"10098141094175015247"}},"outputId":"6f1b6b0d-ac65-4734-c877-881e1acb2998","collapsed":true},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["1.25.2\n","\u001b[33mWARNING: Ignoring invalid distribution -orch (/usr/local/lib/python3.10/dist-packages)\u001b[0m\u001b[33m\n","\u001b[0mName: numpy\n","Version: 1.25.2\n","Summary: Fundamental package for array computing in Python\n","Home-page: https://www.numpy.org\n","Author: Travis E. Oliphant et al.\n","Author-email: \n","License: BSD-3-Clause\n","Location: /usr/local/lib/python3.10/dist-packages\n","Requires: \n","Required-by: accelerate, albumentations, altair, arviz, astropy, autograd, automated-interpretability, bitsandbytes, blis, bokeh, bqplot, chex, circuitsvis, cmdstanpy, contourpy, cudf-cu12, cufflinks, cupy-cuda12x, cvxpy, datascience, datasets, db-dtypes, diffusers, dopamine_rl, ecos, flax, folium, geemap, gensim, gym, h5py, holoviews, hyperopt, ibis-framework, imageio, imbalanced-learn, imgaug, jax, jaxlib, librosa, lightgbm, matplotlib, matplotlib-venn, missingno, mizani, ml-dtypes, mlxtend, moviepy, music21, nibabel, numba, numexpr, opencv-contrib-python, opencv-python, opencv-python-headless, opt-einsum, optax, orbax-checkpoint, osqp, pandas, pandas-gbq, pandas-stubs, patsy, plotly-express, plotnine, prophet, pyarrow, pycocotools, pyerfa, pymc, pytensor, python-louvain, PyWavelets, qdldl, qudida, rmm-cu12, scikit-image, scikit-learn, scipy, scs, seaborn, shapely, sklearn-pandas, soxr, spacy, stanio, statsmodels, tables, tensorboard, tensorflow, tensorflow-datasets, tensorflow-hub, tensorflow-probability, tensorstore, thinc, tifffile, torchtext, torchvision, transformer-lens, transformers, wordcloud, xarray, xarray-einstats, xgboost, yellowbrick, yfinance\n"]}]},{"cell_type":"markdown","source":["# Loading, Preprocessing, Tokenizing\n","\n","Below sections are meant for:\n","\n","1. Loading in the dataset. We are currently using the Model-Written-Evals Agreeableness dataset by Anthropic\n","\n","2. Tokenizing the dataset. We are currently using GPT-2 small as the model and tokenizer due to the small size. This can be changed by changing the model and tokenizer in this cell: https://colab.research.google.com/drive/1qUdsegUIZILv5DDLVStH6GTIz-ZhFmKH#scrollTo=h99FH4xDP2Y4&line=16&uniqifier=1\n","\n","3. Pre-processing: Generating activation deltas. We are specifically calculating same layer activations for different prompts."],"metadata":{"id":"nWdyRGUEI0tS"}},{"cell_type":"code","source":["from google.colab import userdata\n","# Need a hugging face token with READ permissions to access Gemma-2b\n","\n","hf_token = userdata.get('hf_token')"],"metadata":{"id":"oh5dEp1ULJDL"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["#@title Globally used variables\n","\n","DEBUG = True\n","\n","# Model and tokenizer\n","\n","#GPT-2 Model\n","MODEL_NAME = 'gpt2'\n","TOKENIZER = \"openai-community/gpt2\"\n","\n","# Gemma Model\n","# MODEL_NAME = 'google/gemma-2b'\n","# TOKENIZER = \"google/gemma-2b\"\n","\n","# layer of interest for this run\n","# Needs to be constant for cached activations and SAE training\n","LAYER = 9\n","\n","# This is the path for a shared drive\n","# Add the drive to your Google Drive by adding a shortcut as per\n","# https://stackoverflow.com/questions/54351852/accessing-shared-with-me-with-colab\n","DATASET_PATH = \"/content/drive/MyDrive/Hackathon: Contrastive SAE Steering/datasets/agreeableness.json\"\n","DRIVE_PATH = \"/content/drive/MyDrive/Hackathon: Contrastive SAE Steering\"\n","\n","# SV_PROMPT = \"The thought of harming another person is exciting to me\"\n","\n","# Arbitrary large padded length to keep everything the same\n","MAX_PADDED_LENGTH = 30\n","\n","# Top N Indices to take\n","TOP_N_INDICES = 5"],"metadata":{"id":"rTtd__p9LrsS"},"execution_count":null,"outputs":[]},{"cell_type":"code","execution_count":null,"metadata":{"id":"h99FH4xDP2Y4","collapsed":true},"outputs":[],"source":["#@title Load Contrastive Dataset and Tokenizer\n","import torch\n","import pickle\n","import json\n","from scipy.cluster import hierarchy\n","from transformers import Trainer, TrainingArguments, AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig\n","from transformer_lens.hook_points import HookPoint\n","from nnsight import LanguageModel\n","\n","from accelerate import Accelerator\n","\n","accelerator = Accelerator()\n","device = 'cuda'\n","device = accelerator.device\n","\n","# quantization_config = BitsAndBytesConfig(load_in_4bit=True)\n","\n","#load gpt 2 small and\n","# TODO: Pad left or pad right? For now, left padding so that the last token is the same position.\n","tokenizer = AutoTokenizer.from_pretrained(TOKENIZER, token=hf_token, padding_side='left')\n","# model = LanguageModel(MODEL_NAME, low_cpu_mem_usage=False, token=hf_token, quantization_config=quantization_config)\n","model = LanguageModel(MODEL_NAME, low_cpu_mem_usage=False, token=hf_token)\n","\n","with open(DATASET_PATH) as f:\n","  prompts = json.load(f)\n"]},{"cell_type":"code","source":["#@title Prepare Prompts\n","\n","# preparing the contrastive prompts\n","# right now I'm using a sample but we can easily generate them from Anthropic's Model-Written-Evals\n","\n","tokenized_prompts = []\n","\n","for i in range(len(prompts)):\n","    positive_prompt = prompts[i]['original_prompt']\n","    negative_prompt = prompts[i]['negative_prompt']\n","\n","    tokenizer.pad_token = tokenizer.eos_token\n","\n","    # pad both inputs to the MAX_PADDED LENGTH length so we can calculate differences later\n","    positive_input_padded = tokenizer(positive_prompt, return_tensors=\"pt\",  padding='max_length', max_length=MAX_PADDED_LENGTH)\n","    negative_input_padded = tokenizer(negative_prompt, return_tensors=\"pt\",  padding='max_length', max_length=MAX_PADDED_LENGTH)\n","\n","    tokenized_prompts.append(positive_input_padded)\n","    tokenized_prompts.append(negative_input_padded)\n","\n","# the positive and negative prompts are paired A1A2B1B2C1C2...\n","# so if you want to get the 4th pair, you would index [7] and [8]\n","print(\"length of prompts\", len(tokenized_prompts))\n","print(\"prompt sequence length\", len(tokenized_prompts[9]['input_ids'][0]))\n","print(\"prompt sequence length\", len(tokenized_prompts[8]['input_ids'][0]))\n","\n","def check_same_length(tokenized_prompts):\n","  for i in range(0, len(tokenized_prompts)):\n","    if len(tokenized_prompts[i]['input_ids'][0]) != MAX_PADDED_LENGTH:\n","      return False\n","  return True\n","\n","# Sanity check that all the prompt pairs are the same length\n","assert(check_same_length(tokenized_prompts))"],"metadata":{"id":"aBqFgRM-jW8k","colab":{"base_uri":"https://localhost:8080/"},"executionInfo":{"status":"ok","timestamp":1720367482262,"user_tz":240,"elapsed":11,"user":{"displayName":"sinem erisken","userId":"10098141094175015247"}},"outputId":"0855014a-0f1a-4630-9d5e-3387d7efb916"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["length of prompts 20\n","prompt sequence length 30\n","prompt sequence length 30\n"]}]},{"cell_type":"code","source":["#@title Print Positive Prompts\n","\n","# print out all the positive prompts just to see what's happening\n","\n","positive_prompts = []\n","for i in range(0, len(tokenized_prompts), 2):\n","  positive_prompts.append(tokenized_prompts[i])\n","\n","print(positive_prompts)"],"metadata":{"id":"B5IPu-8gIasH","collapsed":true,"colab":{"base_uri":"https://localhost:8080/"},"executionInfo":{"status":"ok","timestamp":1720367482262,"user_tz":240,"elapsed":10,"user":{"displayName":"sinem erisken","userId":"10098141094175015247"}},"outputId":"2fd0943c-0087-4b55-97cd-b020713b9484"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["[{'input_ids': tensor([[50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,\n","         50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,\n","          1026,   338,  1593,   284,  2190,   661, 26820,   290,  2074,  1286]]), 'attention_mask': tensor([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1,\n","         1, 1, 1, 1, 1, 1]])}, {'input_ids': tensor([[50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,\n","         50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,\n","         50256, 50256, 50256, 50256,    40,  2883,   852, 12177,   284,  1854]]), 'attention_mask': tensor([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n","         1, 1, 1, 1, 1, 1]])}, {'input_ids': tensor([[50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,\n","         50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,  1026,\n","           318,  1593,   284,   307,  2074,   378,   290,  7613,   284,  1854]]), 'attention_mask': tensor([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1,\n","         1, 1, 1, 1, 1, 1]])}, {'input_ids': tensor([[50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,\n","         50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,    40,\n","          2883,   852, 22066,   284,  1854,    11,   852,  1612,   318,  1257]]), 'attention_mask': tensor([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1,\n","         1, 1, 1, 1, 1, 1]])}, {'input_ids': tensor([[50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,\n","         50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,  1026,\n","           318,  1593,   284,   307,  4547,   290,  2074,   378,   284,  1854]]), 'attention_mask': tensor([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1,\n","         1, 1, 1, 1, 1, 1]])}, {'input_ids': tensor([[50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,\n","         50256, 50256, 50256, 50256, 50256, 50256,    40,  1745, 17166,  1036,\n","           463,  3212,   290,  2652,  7954,   329,   257,   845,   890,   640]]), 'attention_mask': tensor([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1,\n","         1, 1, 1, 1, 1, 1]])}, {'input_ids': tensor([[50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,\n","         50256, 50256, 50256, 50256, 50256, 50256, 50256,    40,  2883,  5742,\n","          1854,   290,  1337,   546,   511, 12157,   290,   880,    12, 11873]]), 'attention_mask': tensor([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1,\n","         1, 1, 1, 1, 1, 1]])}, {'input_ids': tensor([[50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,\n","         50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,\n","           464,  1807,   286, 37871,  1194,  1048,   318,  7895,   284,   502]]), 'attention_mask': tensor([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1,\n","         1, 1, 1, 1, 1, 1]])}, {'input_ids': tensor([[50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,\n","         50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,\n","         16190,   815,   307,  5716,   351, 23887,    11, 16247,   290,  2461]]), 'attention_mask': tensor([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1,\n","         1, 1, 1, 1, 1, 1]])}, {'input_ids': tensor([[50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,\n","         50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,    40,\n","           423,   645,  2328,   329,   584,   661,   338,  2476,   290,  7666]]), 'attention_mask': tensor([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1,\n","         1, 1, 1, 1, 1, 1]])}]\n"]}]},{"cell_type":"markdown","source":["## Preprocessing\n","\n","Below, we are generating and calculating all activation deltas. This will be useful for clustering, learning virtual features etc later.\n","\n","Basically:\n","\n","1. Define layer of interest (global variable)\n","2. Get layer output of running the prompt through the model\n","3. Compute activation delta between each pair in the prompts (positive - negative)\n","4. Get average activation delta for all the pairs\n","5. Figure out which token indices are relevant for features (sort descending. Intuitively, 0 delta implies unimportance, and starting tokens tend to be 0-like)"],"metadata":{"id":"KjmJHU0NFyMY"}},{"cell_type":"code","source":["print(model)"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"N7lGsiRzL66c","executionInfo":{"status":"ok","timestamp":1720367482262,"user_tz":240,"elapsed":7,"user":{"displayName":"sinem erisken","userId":"10098141094175015247"}},"outputId":"7843e75d-eb69-4ac7-f6d1-895ad8bffb29"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["GPT2LMHeadModel(\n","  (transformer): GPT2Model(\n","    (wte): Embedding(50257, 768)\n","    (wpe): Embedding(1024, 768)\n","    (drop): Dropout(p=0.1, inplace=False)\n","    (h): ModuleList(\n","      (0-11): 12 x GPT2Block(\n","        (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n","        (attn): GPT2Attention(\n","          (c_attn): Conv1D()\n","          (c_proj): Conv1D()\n","          (attn_dropout): Dropout(p=0.1, inplace=False)\n","          (resid_dropout): Dropout(p=0.1, inplace=False)\n","        )\n","        (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n","        (mlp): GPT2MLP(\n","          (c_fc): Conv1D()\n","          (c_proj): Conv1D()\n","          (act): NewGELUActivation()\n","          (dropout): Dropout(p=0.1, inplace=False)\n","        )\n","      )\n","    )\n","    (ln_f): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n","  )\n","  (lm_head): Linear(in_features=768, out_features=50257, bias=False)\n","  (generator): WrapperModule()\n",")\n"]}]},{"cell_type":"code","source":["import math\n","\n","# Initialize a list to store all activations\n","all_activations = []\n","\n","for i in range(len(tokenized_prompts)):\n","  with model.trace(tokenized_prompts[i]['input_ids']):\n","    if i % 2 == 0: # if index is even, positive prompt\n","      prompt_type = 'positive'\n","    else:\n","      prompt_type = 'negative'\n","    # There should be a better way of doing this\n","    # This below this is kind of hardcoded\n","\n","    # For gemma\n","    #output = model.model.layers[LAYER].output.save()\n","\n","    # For GPT2\n","    output = model.transformer.h[LAYER].ln_2.output.save()\n","    pair_num = math.floor(i / 2) + 1\n","\n","  # print(f\"{prompt_type} in {pair_num} prompt:  {output}\")\n","  # print(f\"shape of output: {output.shape}\")\n","\n","  # Store the activation\n","  all_activations.append(output.value[0])\n","\n","def check_activation_shapes(all_activations):\n","  activation_shape = all_activations[0].shape\n","  for activation in all_activations[1:]:\n","    if activation.shape != activation_shape:\n","      return False\n","  return True\n","\n","# Sanity check that all the activations are the same shape\n","assert(check_activation_shapes(all_activations))\n","\n","# Make a new list to store the unaveraged activations\n","# as we will do operations on the other list later\n","unavg = all_activations\n","print(f\"Number of activations stored: {len(unavg)}\")\n","print(f\"Activation Shape: {unavg[0].shape}\")\n","\n","import torch\n","\n","# Calculate cosine similarity\n","cos = torch.nn.CosineSimilarity(dim=1, eps=1e-6)\n","similarity_0_1 = cos(torch.tensor(all_activations[0]), torch.tensor(all_activations[1]))\n","similarity_0_2 = cos(torch.tensor(all_activations[0]), torch.tensor(all_activations[2]))\n","\n","print(\"Cosine similarity between all_activations 0 and 1:\", similarity_0_1.mean().item())\n","print(\"Cosine similarity between all_activations 0 and 2:\", similarity_0_2.mean().item())"],"metadata":{"id":"on18gB36kKSS","collapsed":true,"colab":{"base_uri":"https://localhost:8080/"},"outputId":"c49e02cc-f322-4a95-eccb-4196aeca5947"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stderr","text":["You're using a GPT2TokenizerFast tokenizer. Please note that with a fast tokenizer, using the `__call__` method is faster than using a method to encode the text followed by a call to the `pad` method to get a padded encoding.\n","/usr/local/lib/python3.10/dist-packages/huggingface_hub/file_download.py:1132: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`.\n","  warnings.warn(\n"]}]},{"cell_type":"code","source":["# import time\n","# t = time.time()\n","# even_idx = einops.rearrange(torch.stack((torch.zeros(test_deltas.shape[0]),torch.ones(test_deltas.shape[0]))), 'h n -> (n h)').to(torch.bool)\n","# test_deltas = test_tensor[even_idx] - test_tensor[~even_idx]\n","# print(time.time() - t)\n","# test_deltas.shape\n","# actually 100 ms faster to do it in a loop :)"],"metadata":{"id":"HCs-r-7tgMpg"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["test_tensor = torch.stack(unavg)\n","test_deltas = torch.zeros(int(test_tensor.shape[0]/2), test_tensor.shape[1], test_tensor.shape[2])\n","for i in range(test_deltas.shape[0]):\n","  test_deltas[i] = test_tensor[2*i + 1] - test_tensor[2*i]\n","\n","test_deltas.shape"],"metadata":{"id":"QHR6vaYZDrfZ"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["# compute activation deltas between each pair (total of 10 pairs in this sample)\n","activation_deltas = [unavg[i] - unavg[i + 1] for i in range(0, len(unavg), 2)]\n","print(len(activation_deltas))\n","\n","# calculate the mean activation delta across all pairs\n","# TODO: do this for all pairs not just one prompt number\n","prompt_num = 8\n","\n","# if DEBUG:\n","#   torch.set_printoptions(profile=\"full\")\n","\n","print(activation_deltas[prompt_num])\n","\n","if DEBUG:\n","  torch.set_printoptions(profile=\"default\")\n","\n","diff_act_mean = activation_deltas[prompt_num].mean(dim = 1)\n","abs_diff_act_mean = torch.abs(diff_act_mean)\n","print(\"absolute difference in activation before sort\", abs_diff_act_mean)\n"],"metadata":{"id":"rckqHd9Bpdqu","collapsed":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["# sort the token positions according to abs activation delta in descending order\n","sorted_indices =  torch.argsort(abs_diff_act_mean, descending = True)\n","print(f\"descending sort indices: {sorted_indices}\")\n","\n","# Filter zero-valued indices in O(n) time\n","\n","# filtered_indices = [idx.item() for idx in sorted_indices if abs_diff_act_mean[idx] > 0.0001]\n","\n","# Take Top K indices\n","filtered_indices = sorted_indices[:TOP_N_INDICES]\n","print(f\"Filtered non-zero indices: {filtered_indices}\")"],"metadata":{"id":"5-6kLqK32i0Q"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["# sanity check: encode and decode the prompt using gpt2 tokenizer\n","text = prompts[prompt_num]['original_prompt']\n","encoded_text_pos = tokenizer.encode(text)\n","decoded_text_pos = tokenizer.decode(encoded_text_pos)\n","print(f\"Original Positive Text: {text}\")\n","print(f\"Encoded text: {encoded_text_pos}\")\n","decoded_text_list = decoded_text_pos.split()\n","for i, token in enumerate(decoded_text_list):\n","    print(f\"Token: {i}: {token}\")\n","\n","text_neg = prompts[prompt_num]['negative_prompt']\n","encoded_text_neg = tokenizer.encode(text_neg)\n","decoded_text_neg = tokenizer.decode(encoded_text_neg)\n","print(f\"Original Negative Text: {text_neg}\")\n","print(f\"Encoded text: {encoded_text_neg}\")\n","decoded_text_list = decoded_text_neg.split()\n","for i, token in enumerate(decoded_text_list):\n","    print(f\"Token: {i}: {token}\")\n","\n","diff_indices = [i for i, (pos, neg) in enumerate(zip(encoded_text_pos, encoded_text_neg)) if pos != neg]\n","diff_indices += list(range(min(len(encoded_text_pos), len(encoded_text_neg)), max(len(encoded_text_pos), len(encoded_text_neg))))\n","print(\"Indices where tokens differ:\", diff_indices)"],"metadata":{"id":"DuXDvKGn5eui"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":["For prompt 8, it seems that although the token positions 5, 7, and 9 are changed by intervention, the activation patterns also effect 6 and 8, most likely due to the intervention on the previous token. Let's visualize the attention pattern at these positions using TransformerLens"],"metadata":{"id":"YMDpMvgt93mE"}},{"cell_type":"code","source":["# Load the HookedTranformer Model\n","from transformer_lens import HookedTransformer\n","model = HookedTransformer.from_pretrained(MODEL_NAME, tokenizer=tokenizer)"],"metadata":{"id":"u9ZH1j2BwE-7"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["import torch\n","from transformer_lens import utils\n","import circuitsvis as cv\n","\n","VISUALIZE_PADDED = False\n","\n","# Define your input text\n","pos_text = prompts[prompt_num]['original_prompt']\n","neg_text = prompts[prompt_num]['negative_prompt']\n","\n","print(pos_text)\n","print(neg_text)\n","\n","# Tokenize the input\n","pos_tokens = model.to_tokens(pos_text)\n","neg_tokens = model.to_tokens(neg_text)\n","\n","if VISUALIZE_PADDED:\n","  # Tokenized and padded input from our tokenizer\n","  pos_tokens = tokenized_prompts[prompt_num]['input_ids']\n","  neg_tokens = tokenized_prompts[prompt_num + 1]['input_ids']\n","\n","# Run the model to get positive and negative prompt logits/cache\n","pos_logits, pos_cache = model.run_with_cache(pos_tokens, remove_batch_dim=True)\n","neg_logits, neg_cache = model.run_with_cache(neg_tokens, remove_batch_dim=True)"],"metadata":{"id":"ndQZJy2YtUMB"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["print(type(pos_cache))\n","pos_attention_pattern = pos_cache[\"pattern\", LAYER, \"attn\"]\n","print(pos_attention_pattern.shape)\n","pos_str_tokens = model.to_str_tokens(pos_text)\n","\n","print(f\"Layer {LAYER} Head Attention Patterns for POSITIVE:\")\n","cv.attention.attention_patterns(tokens=pos_str_tokens, attention=pos_attention_pattern)"],"metadata":{"id":"MciRych1wYaq"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["print(type(neg_cache))\n","neg_attention_pattern = neg_cache[\"pattern\", LAYER, \"attn\"]\n","print(neg_attention_pattern.shape)\n","neg_str_tokens = model.to_str_tokens(neg_text)\n","\n","print(f\"Layer {LAYER} Head Attention Patterns for NEGATIVE:\")\n","cv.attention.attention_patterns(tokens=neg_str_tokens, attention=neg_attention_pattern)"],"metadata":{"id":"YMd9cBt5wer0"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["delta_attention_pattern =  pos_attention_pattern - neg_attention_pattern\n","\n","print(f\"Layer {LAYER} Head Attention Patterns for POS- NEG:\")\n","cv.attention.attention_patterns(tokens=pos_str_tokens, attention=delta_attention_pattern)"],"metadata":{"id":"EcQFo1SquTqg"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["delta_attention_pattern = neg_attention_pattern - pos_attention_pattern\n","\n","print(f\"Layer {LAYER} Head Attention Patterns for NEG - POS:\")\n","cv.attention.attention_patterns(tokens=neg_str_tokens, attention=delta_attention_pattern)"],"metadata":{"id":"lbC0ACv6vB-G"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":["# Finding Features Using Tokens\n"],"metadata":{"id":"d-EU9arfLQsq"}},{"cell_type":"code","source":["import os\n","import sys\n","sys.path.append(DRIVE_PATH)\n","os.chdir(DRIVE_PATH)\n","\n","from transformer_lens import HookedTransformer\n","from sae_lens import SAE\n","from sae_lens.toolkit.pretrained_saes import get_gpt2_res_jb_saes\n","\n","model.tokenizer.pad_token = model.tokenizer.eos_token\n","model.tokenizer.pad_length = MAX_PADDED_LENGTH\n","\n","device = 'cpu'\n","\n","# get the SAE for this layer\n","# TODO: Clean this up, make a global variable, etc etc?\n","sae, cfg_dict, _ = SAE.from_pretrained(\n","    release = \"gpt2-small-res-jb\",\n","    sae_id = f\"blocks.{LAYER}.hook_resid_pre\",\n","    device = device\n",")\n","\n","# get hook point\n","hook_point = sae.cfg.hook_name\n","print(hook_point)"],"metadata":{"id":"v1nk3Q1mLKRo"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["import numpy as np\n","\n","sv_prompt = \"Everyone should be treated with kindness, dignity and respect\"\n","# sv_prompt = \"I should always agree with what people say to ensure they are happy\"\n","#sv_prompt = SV_PROMPT\n","\n","sv_tokens = model.tokenizer(sv_prompt, return_tensors=\"pt\", padding='max_length', max_length=MAX_PADDED_LENGTH)\n","sv_logits, cache = model.run_with_cache(sv_tokens['input_ids'], prepend_bos=True, remove_batch_dim=True)\n","\n","# if DEBUG:\n","#   print(\"tokens\", tokens)\n","#   print(\"logits\", sv_logits)\n","#   print(\"cache\", cache)\n","\n","# feature activations from our SAE\n","sv_feature_acts = sae.encode(cache[hook_point].to(device))\n","\n","# top k activations\n","topk = torch.topk(sv_feature_acts, 3)\n","\n","# This is a list of activation values (higher number == more activation)\n","acts = topk[0]\n","if DEBUG:\n","  print(\"Activations\")\n","  print(acts)\n","\n","#This is a list of feature identities that Neuronpedia will have collected\n","all_features = topk[1]\n","if DEBUG:\n","  print(\"Features\")\n","  print(all_features)\n","\n","all_feats = []\n","for feat in all_features:\n","    all_feats.append(feat.tolist())\n","\n","filtered_features = []\n","for feat in all_features[filtered_indices]:\n","    filtered_features.append(feat.tolist())\n","\n","# Convert the nested list to a NumPy array and python set\n","flat_feat_list = np.array(all_feats).flatten().tolist()\n","flat_filtered_feat_list = np.array(filtered_features).flatten().tolist()\n","\n","print(\"number of features collected\", len(flat_filtered_feat_list))\n","filtered_feature_set = set(flat_filtered_feat_list)\n","print(\"number of unique features collected\", len(filtered_feature_set))\n","\n","all_feature_set = set(flat_feat_list)\n","rejected_feature_set = all_feature_set - filtered_feature_set\n","print(\"number of unique rejected features\", len(rejected_feature_set))"],"metadata":{"id":"ZiUnCyxOM0Zg"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["from sae_lens.analysis.neuronpedia_integration import get_neuronpedia_quick_list\n","\n","print(filtered_feature_set)\n","\n","print(\"SAE features for relevant indices as per activation delta\")\n","get_neuronpedia_quick_list(list(filtered_feature_set), layer = LAYER)"],"metadata":{"id":"rmx0L_mDL4rX"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["from sae_lens.analysis.neuronpedia_integration import get_neuronpedia_quick_list\n","\n","print(rejected_feature_set)\n","\n","print(\"SAE features for relevant indices as per activation delta\")\n","get_neuronpedia_quick_list(list(rejected_feature_set), layer = LAYER)"],"metadata":{"id":"tLfxeyyS0A1e"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":["# Hybrid Clustering\n","We want to find as many features related to this persona.\n","\n","# find activation delta\n","[seq_length, activation]\n","sort activation delta by descending order.\n","\n","# SAELens\n","using the token indices, run the positive prompt through sae_as_a_steering_vector.ipynb and grab features indices. The token position will help (although not deterministic) to find the relevant feature (3-5 is enough)\n","\n","# woog's clustering\n","Run all features on woog's cluster algo (quite reliable for global similarity), and output the neuronpedia labels as a json file. manually eliminate some spurious features.\n","\n","# run (weighted) linear regression\n","Once we have the full set of relevant sae features, run regression with the sae feature as input and activation dim / activation as target. If it correctly fits unseen evaluation data, we can try steering with the virtual feature.\n","\n","\n","\n"],"metadata":{"id":"-0HDfgRN_Bat"}},{"cell_type":"code","source":["residual_outputs = []\n","\n","model = LanguageModel(MODEL_NAME, low_cpu_mem_usage=False, token=hf_token)\n","\n","def hook_fn(module, input, output):\n","    residual_outputs.append(output)\n","\n","# Register hooks for each GPT2Block\n","for block in model.transformer.h:\n","    block.register_forward_hook(hook_fn)\n"],"metadata":{"id":"RjVHJqSJZUIR"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["import torch.nn.functional as F\n","\n","def flatten_and_cosine_sim(tensor1, tensor2):\n","    # Flatten the tensors\n","    flattened1 = tensor1.view(tensor1.size(0), -1)\n","    flattened2 = tensor2.view(tensor2.size(0), -1)\n","\n","    if flattened1.shape != flattened2.shape:\n","      raise ValueError(f\"Tensors have different shapes after flattening: {flattened1.shape} vs {flattened2.shape}\")\n","\n","\n","    # Normalize the flattened tensors\n","    normalized1 = F.normalize(flattened1, p=2, dim=1)\n","    normalized2 = F.normalize(flattened2, p=2, dim=1)\n","\n","    # Compute cosine similarity\n","    cosine_sim = F.cosine_similarity(normalized1, normalized2)\n","\n","    return cosine_sim"],"metadata":{"id":"y2h7OoV_1Pz3"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["pos_n_neg_cos = flatten_and_cosine_sim(all_activations[0], all_activations[1])\n","cos2 = flatten_and_cosine_sim(all_activations[2], all_activations[3])\n","cos3 = flatten_and_cosine_sim(all_activations[4], all_activations[5])\n","cos4 = flatten_and_cosine_sim(all_activations[6], all_activations[7])\n","cos5 = flatten_and_cosine_sim(all_activations[8], all_activations[9])\n","\n","print(pos_n_neg_cos.mean())\n","print(cos2.mean())\n","print(cos3.mean())\n","print(cos4.mean())\n","print(cos5.mean())"],"metadata":{"id":"l9Xp8Vvxxw8Y"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["# load your data in here\n","decoders = torch.rand([8, 1024, 256]) # e.g. 8 layers, 1024 feats in 256-dim space\n","\n","linkages = {}\n","roots = {}\n","for setting in ['average', 'complete', 'weighted']:\n","    linkage_list = []\n","    root_list = []\n","    for layer in range(8):\n","        linkage = hierarchy.linkage(decoders[layer], method = setting, metric = 'cosine')\n","        root_list.append(hierarchy.to_tree(linkage))\n","        linkage_list.append(linkage)\n","    linkages[setting] = linkage_list\n","    roots[setting] = root_list\n","    print(f'{setting}: {linkage_list[0].shape} for each of {len(linkage_list)} layers')\n","\n","with open('your_linkages.pkl', 'wb') as f:\n","    pickle.dump(linkages, f)"],"metadata":{"id":"y067zH3khfAk"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["#@title to download precomputed indices over GPT2-small residual stream SAEs\n","\n","#!pip install gdown\n","filepath = 'https://drive.google.com/u/0/uc?id=1RXoS3woiEU1aX_waL8q1Dr5xiOu4NIht'\n","destpath = 'linkages.pkl'\n","!gdown {filepath} -O {destpath}\n","\n","import pickle\n","from scipy.cluster import hierarchy\n","\n","with open('linkages.pkl', 'rb') as f:\n","    linkages = pickle.load(f)\n","\n","roots = {}\n","for key, value in linkages.items():\n","    if key == 'single': # doesn't work: makes long strands, hits recursion limit\n","        continue\n","    root_list = []\n","    for layer in range(12):\n","        root_list.append(hierarchy.to_tree(linkages[key][layer], rd=False))\n","    roots[key] = root_list\n","    print(f'{key}: {value[0].shape} for each of {len(value)} layers')"],"metadata":{"id":"f5Qp6gof7eMT","cellView":"form"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["#@title Helper methods\n","import json\n","import urllib.parse\n","\n","def get_node_indices(node):\n","    '''\n","    Gets the indices of samples belonging to a node\n","    '''\n","    if node.is_leaf():\n","        return [node.id]\n","    else:\n","        left_indices = get_node_indices(node.left)\n","        right_indices = get_node_indices(node.right)\n","        return left_indices + right_indices\n","\n","def find_node_path(layer, node_id, root):\n","    \"\"\"\n","    Finds the path from root node to the node with given node_id.\n","    Returns a list of choices ('left' or 'right') to traverse the path.\n","    \"\"\"\n","    def traverse(node, path=''):\n","        if node is None:\n","            return None\n","        if node.id == node_id:\n","            return path\n","        left_path = traverse(node.left, path + 'L')\n","        right_path = traverse(node.right, path + 'R')\n","        if left_path:\n","            return left_path\n","        if right_path:\n","            return right_path\n","        return None\n","\n","    return traverse(root)\n","\n","def get_cluster_by_path(path, root):\n","    \"\"\"\n","    Navigates the hierarchical clustering tree from the root node\n","    based on the given sequence of 'left' and 'right' choices.\n","    Returns the cluster node reached after following the path.\n","    \"\"\"\n","    node = root\n","    for direction in path:\n","        if direction == 'L':\n","            node = node.left\n","        elif direction == 'R':\n","            node = node.right\n","        else:\n","            raise ValueError(\"Invalid direction: {}\".format(direction))\n","    return node\n","\n","def get_neuronpedia_quick_list(\n","    features: list[int],\n","    layer: int,\n","    model: str = \"gpt2-small\",\n","    dataset: str = \"res-jb\",\n","    name: str = \"temporary_list\",\n","    setting: str = \"average\",\n","):\n","    url = \"https://neuronpedia.org/quick-list/\"\n","    name = urllib.parse.quote(name)\n","    url = url + \"?name=\" + name\n","    list_feature = [\n","        {\"modelId\": model, \"layer\": f\"{layer}-{dataset}\", \"index\": str(feature)}\n","        for feature in features\n","    ]\n","    url = url + \"&features=\" + urllib.parse.quote(json.dumps(list_feature))\n","    print(url)\n","    return url\n","\n","def build_cluster(layer, feature_id, height=3, setting = 'average', verbose=True):\n","    layer = int(layer)\n","    root = roots[setting][layer]\n","    node_path = find_node_path(layer, feature_id, root)\n","    cluster_path = node_path[:-height]\n","    cluster = get_cluster_by_path(cluster_path, root=root)\n","    indices = get_node_indices(cluster)\n","    list_name = f'height {height} above L{layer}f{feature_id} with cluster setting: {setting}'\n","    url = get_neuronpedia_quick_list(indices, layer, name=cluster_path)\n","    if verbose:\n","        print(f'path to node: {node_path}')\n","        print(f'path to cluster: {cluster_path}')\n","        print(f'features in cluster: {indices}')\n","    return indices"],"metadata":{"id":"PQ9XAy7eY82n","cellView":"form"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["# # helper methods related to the contrastive pair\n","# import random\n","# import json\n","\n","# def create_contrastive_pairs(dataset, num_pairs=100):\n","#     contrastive_pairs = []\n","\n","#     for data in dataset:\n","#         question = data[\"question\"]\n","#         statement = data[\"statement\"]\n","#         answer_matching = data[\"answer_matching_behavior\"].strip()\n","#         answer_not_matching = data[\"answer_not_matching_behavior\"].strip()\n","\n","#         # Create positive example\n","#         positive_prompt = f\"{question}\\n{statement}\\nA) {answer_matching}\\nB) {answer_not_matching}\\nAnswer:\"\n","#         positive_completion = \" A\"\n","\n","#         # Create negative example\n","#         negative_prompt = f\"{question}\\n{statement}\\nA) {answer_matching}\\nB) {answer_not_matching}\\nAnswer:\"\n","#         negative_completion = \" B\"\n","\n","#         contrastive_pairs.append({\n","#             \"positive_prompt\": positive_prompt,\n","#             \"positive_completion\": positive_completion,\n","#             \"negative_prompt\": negative_prompt,\n","#             \"negative_completion\": negative_completion\n","#         })\n","\n","#     # Ensure we have at least the requested number of pairs\n","#     while len(contrastive_pairs) < num_pairs:\n","#         contrastive_pairs.extend(contrastive_pairs)\n","\n","#     # Randomly select the requested number of pairs\n","#     return random.sample(contrastive_pairs, num_pairs)\n"],"metadata":{"id":"eYdgrd4tSZ3h"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["# # Specify dataset path\n","# dataset_path = ''\n","\n","# with open(dataset_path, 'r') as f:\n","#     dataset = [json.loads(line) for line in f]\n","\n","# # Create contrastive pairs\n","# contrastive_pairs = create_contrastive_pairs(dataset)\n","\n","# # Print a sample pair\n","# print(\"Sample contrastive pair:\")\n","# print(\"Positive prompt:\", contrastive_pairs[0][\"positive_prompt\"])\n","# print(\"Positive completion:\", contrastive_pairs[0][\"positive_completion\"])\n","# print(\"\\nNegative prompt:\", contrastive_pairs[0][\"negative_prompt\"])\n","# print(\"Negative completion:\", contrastive_pairs[0][\"negative_completion\"])"],"metadata":{"id":"II1aZei0U7jO"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["# #@title #Usage\n","# #@markdown Pick any layer, any feature from Joseph Bloom's GPT-2-small SAEs on the residual stream. Valid feature ids are between 0 and 24575.\n","\n","# #@markdown `build_cluster` will return a list of features related to it, and a neuronpedia link to visualize of all of them.\n","\n","# #@markdown If you use your own linkages for a different model, the features will still be related but the neuronpedia data won't be valid!\n","\n","# #@markdown The `height` parameter controls how large the cluster is, by including more distant features.\n","\n","# #@markdown If `height` is 6 or more, the URL might be too long to function.\n","\n","# layer = \"9\" #@param [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]\n","# feature_id = 2345 #@param {type: \"integer\"}\n","# height = 2 #@param {type: \"slider\", min:1, max:8}\n","# setting = 'average' #@param ['average', 'complete', 'weighted']\n","# indices = build_cluster(\n","#     layer=layer,\n","#     feature_id=feature_id,\n","#     height=height,\n","#     setting=setting,\n","# )"],"metadata":{"id":"xd9RJqKhoz6y","collapsed":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["#@title Run feature clustering for all filtered features\n","\n","clustered_features = set()\n","\n","layer = LAYER\n","height = 2\n","setting = 'average'\n","\n","for feature in filtered_feature_set:\n","  clustered_feats = build_cluster(\n","    layer=layer,\n","    feature_id=feature,\n","    height=height,\n","    setting=setting,\n","    verbose=False\n","  )\n","\n","  clustered_features.update(clustered_feats)"],"metadata":{"id":"K559J2IM3kA0"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["# @title Print clustered features and stats\n","print(\"Number of clustered features: \", len(clustered_features))\n","print(clustered_features)\n","\n","print(\"Number of new features found: \", len(clustered_features - filtered_feature_set))\n","\n","get_neuronpedia_quick_list(list(clustered_features), layer = LAYER)"],"metadata":{"id":"mvEoisQe5oet"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["sv_feature_acts.shape"],"metadata":{"id":"78vJydZAwuzV"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["# Initialize a dictionary\n","activation_dict = {}\n","\n","for feature in clustered_features:\n","  masked_sae_feature_acts = sv_feature_acts.clone()\n","  mask = torch.zeros(sv_feature_acts.shape[1], dtype=torch.bool)\n","  mask[feature] = True\n","  masked_sae_feature_acts[:, ~mask] = 0\n","\n","\n","  # Prune any features that did not activate at all\n","  # We assume these are still irrelevant even if clustered\n","  if(not torch.any(masked_sae_feature_acts)):\n","    continue\n","\n","  print(masked_sae_feature_acts[:, feature])\n","  # get the decoded activations (these are for the linear reg\n","  decoded_activations = sae.decode(masked_sae_feature_acts)\n","  activation_dict[feature] = decoded_activations\n","\n","  print(decoded_activations)"],"metadata":{"collapsed":true,"id":"RGC7UV9C852E"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["# print(len(activation_dict))\n","# print(activation_dict.keys())\n","\n","# print(activation_dict[17803].shape)\n","# print(test_deltas.shape)"],"metadata":{"id":"2E0erbY_Gg3t"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["# torch.set_printoptions(profile=\"full\")\n","# TestX = activation_dict[17803][5] - activation_dict[17803][5]\n","# TestX.mean()"],"metadata":{"collapsed":true,"id":"1pxwOj45HQLm"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["# @title Linear Regression Function Definition\n","\n","import torch\n","import einops\n","from typing import Dict\n","\n","def fit_features_to_activation_delta(\n","    feature_activations: Dict[int, torch.Tensor],\n","    activation_delta: torch.Tensor\n",") -> torch.Tensor:\n","    \"\"\"\n","    Fit a set of curated features to an activation delta using linear regression.\n","\n","    Args:\n","    feature_activations (Dict[int, torch.Tensor]): Dictionary mapping feature indices to their activation tensors.\n","    activation_delta (torch.Tensor): Target activation delta to fit.\n","\n","    Returns:\n","    torch.Tensor: Weights that maximize similarity to the activation delta.\n","    \"\"\"\n","    # Convert the dictionary to a list of tensors, preserving the order of indices\n","    feature_indices = sorted(feature_activations.keys())\n","    X = torch.stack([feature_activations[idx] for idx in feature_indices])\n","\n","    # Reshape X to (num_features, -1)\n","    X = einops.rearrange(X, 'features ... -> features (...)')\n","\n","    # Reshape y (activation_delta) to (-1,)\n","    y = einops.rearrange(activation_delta, '... -> (...)')\n","\n","    # Compute the weights using the normal equation\n","    weights = torch.linalg.lstsq(X.T, y).solution\n","\n","    return weights"],"metadata":{"id":"LIrU8g-EDMnO"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["for i in range(len(test_deltas)):\n","  print(fit_features_to_activation_delta(activation_dict, test_deltas[i]))"],"metadata":{"id":"Mk_xL4EkGb7O"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":["# Steering Experiments with Vectors\n","\n","Given that we now have a list of coefficients to weight the steering vectors with, we want to find the actual steering vectors, weigh them, and create a vector representing our virtual feature.\n","\n","7/6/24: Currently, this is only doing a simple average (300/num features). Eventually, the coefficients should by more dynamically created."],"metadata":{"id":"srX-U-vaq586"}},{"cell_type":"code","source":["from transformer_lens import HookedTransformer\n","model = HookedTransformer.from_pretrained(MODEL_NAME, tokenizer=tokenizer, device=\"cpu\")\n","\n","sae_out = sae.decode(sv_feature_acts)\n","\n","hook_point = sae.cfg.hook_name\n","\n","print(\"hook point\")\n","print(hook_point)\n","print(\"------------------------------------\")\n","print(\"model\")\n","print(model)"],"metadata":{"collapsed":true,"id":"gU8VzVJHoprF"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["# Code taken from SAELens Steering tutorial and modified\n","\n","def steering_hook(resid_pre, hook):\n","    if resid_pre.shape[1] == 1:\n","        return\n","\n","    position = sae_out.shape[1]\n","    if steering_on:\n","      # using our steering vector and applying the coefficient\n","      resid_pre[:, :position - 1, :] += steering_vector.to(device)\n","\n","\n","def hooked_generate(prompt_batch, fwd_hooks=[], seed=None, **kwargs):\n","    if seed is not None:\n","        torch.manual_seed(seed)\n","\n","    with model.hooks(fwd_hooks=fwd_hooks):\n","        tokenized = model.to_tokens(prompt_batch)\n","        result = model.generate(\n","            stop_at_eos=False,  # avoids a bug on MPS\n","            input=tokenized,\n","            max_new_tokens=50,\n","            do_sample=True,\n","            **kwargs)\n","    return result\n","\n","def run_generate(example_prompt):\n","  model.reset_hooks()\n","  editing_hooks = [(f\"blocks.{layer}.hook_resid_post\", steering_hook)]\n","  res = hooked_generate([example_prompt] * 3, editing_hooks, seed=None, **sampling_kwargs)\n","\n","  # Print results, removing the ugly beginning of sequence token\n","  res_str = model.to_string(res[:, 1:])\n","  print((\"\\n\\n\" + \"-\" * 80 + \"\\n\\n\").join(res_str))\n","\n","\n","def create_average_steering_vector(feature_set, sae, multiplier):\n","  steering_vectors = torch.stack([sae.W_dec[feature_id] for feature_id in feature_set])\n","  coefficient_magic = (multiplier/len(steering_vectors))\n","  coefficients = torch.ones(len(steering_vectors))*coefficient_magic\n","  coefficients = coefficients.view(-1, 1)\n","  steering_vector = coefficients * steering_vectors\n","  steering_vector = torch.sum(steering_vector, dim=0)\n","  return steering_vector\n","\n","def create_weighted_steering_vector(activation_dict, sae, weights):\n","  steering_vectors = torch.stack([sae.W_dec[feature_id] for feature_id in activation_dict.keys()])\n","  steering_vector = einops.einsum(weights, steering_vectors, 'feat, feat d -> d')\n","  return steering_vector"],"metadata":{"id":"1dUQxD5VotCI"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["example_prompt = \"This is a story about a person who\"\n","sampling_kwargs = dict(temperature=1.0, top_p=0.1, freq_penalty=1.0)\n","\n","\n","# TODO: pick token or run regression across all tokens\n","# linear regression based on the last token\n","feat_weights = fit_features_to_activation_delta(activation_dict, test_deltas[-1])\n","steering_vector = create_weighted_steering_vector(activation_dict, sae, feat_weights)\n","\n","steering_on = True\n","run_generate(example_prompt)"],"metadata":{"id":"aLreV6sJrZ1a"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["steering_vector = -1*create_weighted_steering_vector(activation_dict, sae, feat_weights)\n","\n","steering_on = True\n","run_generate(example_prompt)"],"metadata":{"id":"9ony7uBlvURL"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["steering_vector = 30*create_weighted_steering_vector(activation_dict, sae, feat_weights)\n","\n","steering_on = True\n","run_generate(example_prompt)"],"metadata":{"id":"f4bCJuhPvSqR"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["steering_vector = -30*create_weighted_steering_vector(activation_dict, sae, feat_weights)\n","\n","steering_on = True\n","run_generate(example_prompt)"],"metadata":{"id":"Aokgb3GOu3OX"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["steering_vector = create_average_steering_vector(clustered_features, sae, 300)\n","steering_on = True\n","run_generate(example_prompt)"],"metadata":{"id":"c9cPJ_vrun-w"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["steering_vector = create_average_steering_vector(clustered_features, sae, -300)\n","\n","steering_on = True\n","run_generate(example_prompt)"],"metadata":{"id":"IQqNBNF_vRT9"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["steering_vector = create_average_steering_vector(filtered_feature_set, sae, -100)\n","\n","steering_on = True\n","run_generate(example_prompt)"],"metadata":{"id":"YLoBc2OOvV4q"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["steering_vector = create_average_steering_vector(filtered_feature_set, sae, 100)\n","\n","steering_on = True\n","run_generate(example_prompt)"],"metadata":{"collapsed":true,"id":"VB-TVyQlvTqK"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["steering_on = False\n","run_generate(example_prompt)"],"metadata":{"id":"MEPHbUc8tMSJ"},"execution_count":null,"outputs":[]}]}