### Install Prisma Project Repository

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/docs/UsageGuide.md

Instructions to clone the ViT-Prisma repository and install it as an editable package using pip.

```shell
git clone https://github.com/soniajoseph/ViT-Prisma
cd ViT-Prisma
pip install -e .
```

--------------------------------

### Environment Setup and Installation - Shell

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/demos/1_Load_SAE.ipynb

Commands to set up a Python virtual environment, activate it, and install the Prisma library along with the 'kaggle' package. It also includes steps to download the ImageNet dataset using Kaggle.

```bash
python3 -m venv myenv
source myenv/bin/activate
pip install -e .
pip install kaggle
cd data
kaggle competitions download -c imagenet-object-localization-challenge
```

--------------------------------

### Load Datasets and Create DataLoader with PyTorch

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/demos/Interactive_Attention_Head_Tour.ipynb

This Python snippet demonstrates how to load a dataset using the `datasets` library and create a PyTorch DataLoader for batch processing. It assumes the `datasets` and `torch` libraries are installed. The function takes a dataset name and batch size as input and returns a DataLoader object.

```python
from torch.utils.data import DataLoader
import datasets
from datasets import load_dataset

# Example usage:
dataset_name = "your_dataset_name"
batch_size = 32

# Load the dataset
data = load_dataset(dataset_name)

# Create a DataLoader
dataloader = DataLoader(data["train"], batch_size=batch_size, shuffle=True)
```

--------------------------------

### Load Dataset

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/tests/test_notebooks/Untitled.ipynb

Loads a dataset using the `datasets` library from Hugging Face. This is a common first step for preparing data for model training or evaluation.

```python
from datasets import load_dataset
```

--------------------------------

### Install and Import ViT-Prisma and Dependencies

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/demos/ViT_Prisma_Main_Demo.ipynb

Installs the vit-prisma package and imports necessary libraries such as numpy, torch, plotly, PIL, and torchvision. This sets up the environment for using the ViT-Prisma model and its associated utilities.

```python
import vit_prisma
from vit_prisma.utils.data_utils.imagenet_dict import IMAGENET_DICT
from vit_prisma.utils import prisma_utils

import numpy as np
import torch
from fancy_einsum import einsum
from collections import defaultdict

import plotly.graph_objs as go
import plotly.express as px

import matplotlib.colors as mcolors

from PIL import Image
from torchvision import transforms
import matplotlib.pyplot as plt

from IPython.core.display import display, HTML
```

--------------------------------

### Iterate Through DataLoader

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/tests/test_notebooks/Untitled.ipynb

Demonstrates how to get a single batch of data from the training DataLoader. This is useful for debugging or inspecting the data format before training a model.

```python
image, label = next(iter(trainloader))
```

--------------------------------

### Import Dependencies - Python

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/demos/1_Load_SAE.ipynb

Imports necessary libraries and modules for the Prisma SAE tutorial. This includes utilities for data loading, model operations, and constants for directories and devices. Assumes the 'vit-prisma' library is installed.

```python
import os
from pathlib import Path

import numpy as np
import torch
from torch.utils.data import DataLoader, Subset

from vit_prisma.utils.constants import BASE_DIR, DATA_DIR, MODEL_DIR, DEVICE, MODEL_CHECKPOINTS_DIR
from vit_prisma.utils.tutorial_utils import calculate_clean_accuracy, load_clip_models, plot_image, get_feature_activations
from vit_prisma.utils.data_utils.loader import load_dataset
from vit_prisma.utils.tutorial_utils import plot_act_distribution
from vit_prisma.utils.tutorial_utils import plot_top_imgs_for_features
```

--------------------------------

### Image Preprocessing with torchvision in Python

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/demos/Interactive_Attention_Head_Tour.ipynb

This Python snippet defines a custom transformation class `ConvertTo3Channels` using `torchvision.transforms`. This class ensures that input images are converted to the RGB format, which is a common requirement for many deep learning models. It requires the `torchvision` library to be installed.

```python
from torchvision import transforms

class ConvertTo3Channels:
    def __call__(self, img):
        if img.mode != 'RGB':
            return img.convert('RGB')
        return img

# Example usage with torchvision transforms:
transform = transforms.Compose([
    ConvertTo3Channels(),
    transforms.Resize((224, 224)),
    transforms.ToTensor()
])
```

--------------------------------

### Instantiate Induction Dataset

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/docs/UsageGuide.md

Demonstrates how to instantiate the InductionDataset from the vit_prisma.dataloaders module. This dataset can be used for training, and the framework can handle train/test splitting if a test split is not provided.

```python
from vit_prisma.dataloaders.induction import InductionDataset

train = InductionDataset('train')
```

--------------------------------

### Create PyTorch DataLoader

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/demos/Interactive_Attention_Head_Tour.ipynb

Creates a PyTorch DataLoader for the dataset, specifying batch size, shuffling, and worker initialization. This is used to iterate over the data in batches during training or evaluation.

```python
from torch.utils.data import DataLoader

# Create PyTorch DataLoader
batch_size = 4
train_loader = DataLoader(dataset['valid'], batch_size=batch_size, shuffle=True, num_workers=0, worker_init_fn=worker_init_fn)
```

--------------------------------

### Instantiate Base HookedViT Model

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/docs/UsageGuide.md

Shows how to create an instance of the HookedViT model using a configuration object. The HookedViT class facilitates flexible model instantiation and internal inspection for ViTs.

```python
from vit_prisma.models.base_vit import HookedViT
from vit_prisma.configs.HookedViTConfig import HookedViTConfig

config = HookedViTConfig()

model = HookedViT(config)
```

--------------------------------

### Load and Transform ImageNet Dataset

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/demos/Interactive_Attention_Head_Tour.ipynb

Loads a subset of the ImageNet dataset from Huggingface and applies a transformation function. This prepares the data for use with PyTorch DataLoaders.

```python
from datasets import load_dataset

# We'll use a subset of imagenet from Huggingface
dataset = load_dataset("zh-plus/tiny-imagenet")


# Apply transformations to the dataset
dataset.set_transform(transform_batch)
```

--------------------------------

### Run Prisma HookedViT with Cache

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/demos/Interactive_Attention_Head_Tour.ipynb

Executes the Prisma HookedViT model with a batch of images and caches the attention patterns. This allows for later analysis of how the model attends to different parts of the input.

```python
data = next(iter(train_loader))
images, labels = data['image'], data['label']
output, cache = model.run_with_cache(images)
```

--------------------------------

### Import PyTorch and torchvision Libraries

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/tests/test_notebooks/Untitled.ipynb

Imports necessary libraries from PyTorch and torchvision for deep learning tasks. These include torch for general tensor operations and datasets/transforms for data handling.

```python
import torch
from torchvision import datasets, transforms
import numpy as np
```

--------------------------------

### Import Configuration for HookedViT

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/docs/UsageGuide.md

Demonstrates how to import the HookedViTConfig class from the vit_prisma.configs module. This configuration object holds hyperparameters for models and training procedures.

```python
from vit_prisma.configs.HookedViTConfig import HookedViTConfig
```

--------------------------------

### Download Image with wget

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/demos/1_Load_SAE.ipynb

This snippet demonstrates how to download an image file from a given URL using the `wget` command-line utility. It saves the image with a specified local filename. Ensure `wget` is installed and accessible in your environment.

```bash
!wget -O normal_image.jpg https://raw.githubusercontent.com/Prisma-Multimodal/ViT-Prisma/dev/src/vit_prisma/sample_images/n01818515_39.JPEG
```

--------------------------------

### Install vit-prisma Python Library (0.1.4)

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/demos/Interactive_Attention_Head_Tour.ipynb

This command installs the vit-prisma Python library with a specific version (0.1.4). Ensure you have pip and Python installed. This installation includes various dependencies required for the library's functionality.

```bash
# Install the Prisma repo library (update version number or clone from source for latest functionality)

!pip install vit_prisma==0.1.4

```

--------------------------------

### Example Usage of Image Generation Function

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/src/vit_prisma/sae/evals/mustache_decoder.ipynb

Demonstrates how to call the generate_from_embedding function with a pre-computed embedding and the initialized pipeline to produce a final image.

```python
generated_image = generate_from_embedding(your_embedding, pipe)
```

--------------------------------

### CIFAR-10 Dataset Transformations and Loading

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/tests/test_notebooks/Untitled.ipynb

Defines transformations for the CIFAR-10 dataset, including converting images to tensors and normalizing them. It then downloads and loads the training and test sets. The 'download=False' argument assumes the dataset is already present.

```python
# Define transformations for the dataset
transform = transforms.Compose([
    transforms.ToTensor(),  # Convert images to PyTorch tensors
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))  # Normalize the dataset
])

# Download and load the training set
trainset = datasets.CIFAR10(root='../../data/cifar10', train=True, download=False, transform=transform)

# Download and load the test set
testset = datasets.CIFAR10(root='../../data/cifar10', train=False, download=False, transform=transform)
```

--------------------------------

### Train Model with Default Configuration

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/docs/UsageGuide.md

Initiates the training process using the trainer module. It requires the model function, configuration object, and training dataset. Assumes Wandb is set up for experiment tracking.

```python
from vit_prisma.models.base_vit import HookedViT
from vit_prisma.configs.HookedViTConfig import HookedViTConfig
from vit_prisma.training import trainer
from vit_prisma.dataloaders.induction import InductionDataset

train_dataset = InductionDataset('train')

config = HookedViTConfig()

model_function = HookedViT

trainer.train(model_function, config, train_dataset)
```

--------------------------------

### Load Pretrained Vision Transformer Model

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/tests/test_notebooks/Untitled.ipynb

Loads a pretrained Vision Transformer model from Hugging Face Hub and moves it to the specified device (e.g., CUDA). It utilizes the `HookedViT` class from the `vit_prisma.models.base_vit` module.

```python
from vit_prisma.configs import HookedViTConfig
from vit_prisma.models.base_vit import HookedViT

TOLERANCE = 1e-5

model_name = "vit_base_patch16_224"
device = "cuda"

hooked_model = HookedViT.from_pretrained(model_name)
hooked_model.to(device)
```

--------------------------------

### Download Imagenette Dataset

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/tests/test_notebooks/Untitled.ipynb

Utilizes the `untar_data` function to download and extract the Imagenette dataset. This is a smaller subset of ImageNet, suitable for faster experimentation. The URL is specified using `URLs.IMAGENETTE_160`.

```python
path = untar_data(URLs.IMAGENETTE_160)
```

--------------------------------

### Load Pre-trained Vision Transformer Model

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/demos/Interactive_Attention_Head_Tour.ipynb

Loads a pre-trained Vision Transformer model from Hugging Face using the `vit_prisma` library. It configures various model parameters for optimal performance and compatibility, such as centering weights and folding layer normalization.

```python
# We'll use a vanilla vision transformer
import vit_prisma
from vit_prisma.models.base_vit import HookedViT
model = HookedViT.from_pretrained("vit_base_patch32_224",
                                        center_writing_weights=True,
                                        center_unembed=True,
                                        fold_ln=True,
                                        refactor_factored_attn_matrices=True,
                                    )

```

--------------------------------

### Run Model with Cache and Extract Residual Stream

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/tests/test_notebooks/Untitled.ipynb

Executes the hooked model with input data and stores intermediate activations in a cache. It then extracts the final residual stream from the cache for further analysis.

```python
image, label = next(iter(testloader))
output, cache = hooked_model.run_with_cache(image.cuda())

final_stream = cache["resid_post", -1]
```

--------------------------------

### Load Pretrained Model and Transforms

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/docs/UsageGuide.md

Loads a pretrained Vision Transformer model (e.g., from OpenCLIP or Timm) and its corresponding transformation pipeline using utility functions from the vit_prisma library.

```python
import vit_prisma
model_name = "open-clip:timm/vit_base_patch32_clip_224.laion2b_e16"
hooked_clip = vit_prisma.load_hooked_model(model_name)
transform = vit_prisma.get_model_transforms(model_name)
```

--------------------------------

### Implement and Use Custom Callback for Training

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/docs/UsageGuide.md

Shows how to define a custom callback class inheriting from PrismaCallback and integrate it into the training process. Callbacks can execute custom logic at the end of epochs or steps.

```python
from vit_prisma.models.base_vit import HookedViT
from vit_prisma.configs.HookedViTConfig import HookedViTConfig
from vit_prisma.training import trainer
from vit_prisma.dataloaders.induction import InductionDataset

from vit_prisma.training.training_utils import PrismaCallback

class DemoCallback(PrismaCallback):
    def on_epoch_end(self, epoch, net, val_loader, wandb_logger):
        # Specify a condition if you don't want the callback to execute after each epoch
        if epoch % 5 == 0: 
            # perform some function with the network, validation set, and log it if required.
            pass

    def on_step_end(self, step, net, val_loader, wandb_logger):
        # It is similar to on_epoch_end but runs after desired number of steps instead of epochs
        pass

train_dataset = InductionDataset('train')

config = HookedViTConfig()

model_function = HookedViT

trainer.train(model_function, config, train_dataset, callbacks=[DemoCallback()])
```

--------------------------------

### Create DataLoader for CIFAR-10 Dataset

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/tests/test_notebooks/Untitled.ipynb

Creates PyTorch DataLoaders for the training and test sets. These loaders facilitate efficient batching and shuffling of data during model training and evaluation. The batch size is set to 64, with shuffling enabled for the training set.

```python
# Create DataLoader for the training and test sets
trainloader = torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=True)
testloader = torch.utils.data.DataLoader(testset, batch_size=64, shuffle=False)
```

--------------------------------

### Interactive Attention Head Visualization (JavaScript)

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/demos/Interactive_Attention_Head_Tour.ipynb

Generates an interactive JavaScript visualization for a specific attention head and batch index. This allows for detailed inspection of attention scores between image patches by hovering over the visualization.

```python
from vit_prisma.visualization.visualize_attention_js import plot_javascript
from IPython.core.display import display, HTML

attn_head_idx = 1
batch_idx = 0
patterns = visualize_attention(attn_head_idx, cache, "Attention Scores", 700, batch_idx = batch_idx, attention_type="attn_scores")
image = images[batch_idx]
html_code = plot_javascript(patterns, [image], ATTN_SCALING=8)
display(HTML(html_code))
```

--------------------------------

### Set Random Seeds for Reproducibility

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/demos/Interactive_Attention_Head_Tour.ipynb

Sets the random seeds for PyTorch, NumPy, and Python's random module to ensure reproducibility of experiments. It also configures CUDA and cuDNN for deterministic behavior. Includes a worker initialization function for data loaders.

```python
import torch
import numpy as np
import random

# Set global seeds
def set_seed(seed_value):
    torch.manual_seed(seed_value)
    torch.cuda.manual_seed(seed_value)
    torch.cuda.manual_seed_all(seed_value)  # if you are using multi-GPU.
    np.random.seed(seed_value)  # Numpy module.
    random.seed(seed_value)  # Python random module.
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True

seed_value = 42
set_seed(seed_value)

# Define a worker init function that sets the seed
def worker_init_fn(worker_id):
    worker_seed = torch.initial_seed() % 2**32
    np.random.seed(worker_seed)
    random.seed(worker_seed)
```

--------------------------------

### Imagenette Dataset Transformations

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/tests/test_notebooks/Untitled.ipynb

Defines transformations for the Imagenette dataset, including resizing images to 224x224, converting them to PyTorch tensors, and normalizing them using ImageNet's mean and standard deviation. This prepares images for models trained on ImageNet.

```python
# Define the transformations. Adjust the sizes as needed.
transform = transforms.Compose([
    transforms.Resize((224, 224)),  # Resize to a square of 256x256
    transforms.ToTensor(),  # Convert images to PyTorch tensors
    transforms.Normalize(mean=[0.485, 0.456, 0.406],  # Normalize using ImageNet mean and std
                         std=[0.229, 0.224, 0.225])
])

# Assuming you've extracted Imagenette to 'path/to/imagenette'
imagenette_dataset = datasets.ImageFolder(root= path, transform=transform)
```

--------------------------------

### Load Legacy Pretrained Models (Hugging Face & Timm)

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/docs/UsageGuide.md

Loads pretrained models from Hugging Face and Timm using the PretrainedModel class. This method allows specifying the model name and configuration, with an option to indicate if it's a Timm model.

```python
from vit_prisma.models.pretrained_model import PretrainedModel
from vit_prisma.configs.HookedViTConfig import HookedViTConfig

config = HookedViTConfig()

hf_model = PretrainedModel('google/vit-base-patch16-224-in21k', config)

timm_model = PretrainedModel('vit_base_patch32_224', config, is_timm=True)
```

--------------------------------

### Visualize Top Activations with Heatmaps

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/src/vit_prisma/sae/evals/Evaluating_Vision_SAE.ipynb

This comprehensive snippet iterates through top activations per feature, retrieves corresponding image data and labels, computes heatmaps using `get_heatmap` and `image_patch_heatmap`, and displays the images with overlaid heatmaps. It utilizes libraries like `tqdm`, `numpy`, `matplotlib.pyplot`, and `torch` for data processing and visualization, creating subplots to show multiple examples.

```python
for feature_ids, cat, logfreq in tqdm(zip(top_activations_per_feature.keys(), interesting_features_category, interesting_features_values), total=len(interesting_features_category)):
  #  print(f"looking at {feature_ids}, {cat}")
    max_vals, max_inds = top_activations_per_feature[feature_ids]
    images = []
    model_images = []
    gt_labels = []
    for bid, v in zip(max_inds, max_vals):

        image, label, image_ind = val_data_visualize[bid]

        assert image_ind.item() == bid
        images.append(image)

        # model_img, _, _ = imagenet_data[bid]
        model_image, _, _ = val_data[bid]
        model_images.append(model_image)
        gt_labels.append(ind_to_name[str(label)][1])
    
    grid_size = int(np.ceil(np.sqrt(len(images))))
    fig, axs = plt.subplots(int(np.ceil(len(images)/grid_size)), grid_size, figsize=(15, 15))
    name=  f"Category: {cat},  Feature: {feature_ids}"
    fig.suptitle(name)#, y=0.95)
    for ax in axs.flatten():
        ax.axis('off')
    complete_bid = []

    for i, (image_tensor, label, val, bid,model_img) in enumerate(zip(images, gt_labels, max_vals,max_inds,model_images )):
        if bid in complete_bid:
            continue 
        complete_bid.append(bid)


        row = i // grid_size
        col = i % grid_size
        heatmap = get_heatmap(model_img,model,sparse_autoencoder, feature_ids )
        heatmap = image_patch_heatmap(heatmap, pixel_num=224//cfg.patch_size)

        display = image_tensor.numpy().transpose(1, 2, 0)

        has_zero = False
        

        axs[row, col].imshow(display)
        axs[row, col].imshow(heatmap, cmap='viridis', alpha=0.3)  # Overlaying the heatmap
        axs[row, col].set_title(f"{label} {val.item():0.03f} {'class token!' if has_zero else ''}")  
        axs[row, col].axis('off')  

    plt.tight_layout()
```

--------------------------------

### Error: Numpy Not Available

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/tests/test_notebooks/Untitled.ipynb

This snippet illustrates a common runtime error encountered when PyTorch operations depend on NumPy, but NumPy is not available or not installed correctly in the environment. Ensure NumPy is installed (`pip install numpy`).

```python
# This code block is part of an error message and demonstrates a potential issue.
# img = torch.from_numpy(np.array(pic, mode_to_nptype.get(pic.mode, np.uint8), copy=True))
# if pic.mode == "1":
#     img = 255 * img
# RuntimeError: Numpy is not available
```

--------------------------------

### Download Sample Images using Wget

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/demos/ViT_Prisma_Main_Demo.ipynb

Downloads sample JPEG images ('cat_dog.jpeg' and 'cat_crop.jpeg') from a GitHub repository using the `wget` command. These images are intended for use with the ViT-Prisma model.

```shell
# Get images we'll feed into the model
!wget https://github.com/soniajoseph/ViT-Prisma/blob/main/src/vit_prisma/sample_images/cat_dog.jpeg?raw=true -O cat_dog.jpeg --quiet
!wget https://github.com/soniajoseph/ViT-Prisma/blob/main/src/vit_prisma/sample_images/cat_crop.jpeg?raw=true -O crop_cat.png --quiet
```

--------------------------------

### Initialize Activations Store and DataLoader (Python)

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/src/vit_prisma/sae/evals/Evaluating_Vision_SAE.ipynb

Imports the `VisionActivationsStore` class for managing activations and `DataLoader` for batching data. These are essential for the subsequent steps in the SAE evaluation pipeline.

```python
from vit_prisma.sae.training.activations_store import VisionActivationsStore
# import dataloader
from torch.utils.data import DataLoader
```

--------------------------------

### Download Sample Images for ViT-Prisma

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/demos/Emoji_Logit_Lens_Demo.ipynb

This code snippet downloads four sample JPEG images (toilet_cat.jpeg, fruit_bowl.jpeg, child_lion.jpeg, cheetah.jpeg) required for the tutorial. It uses wget to fetch the images from a GitHub repository and saves them with specific filenames. Standard output and error are redirected to /dev/null to keep the output clean.

```python
!wget https://github.com/soniajoseph/ViT-Prisma/blob/main/src/vit_prisma/sample_images/toilet_cat.jpeg?raw=true -O toilet_cat.jpg > /dev/null 2>&1
!wget https://github.com/soniajoseph/ViT-Prisma/blob/main/src/vit_prisma/sample_images/fruit_bowl.jpeg?raw=true -O fruit.jpg > /dev/null 2>&1
!wget https://github.com/soniajoseph/ViT-Prisma/blob/main/src/vit_prisma/sample_images/child_lion.jpeg?raw=true -O child_lion.jpg > /dev/null 2>&1
!wget https://github.com/soniajoseph/ViT-Prisma/blob/main/src/vit_prisma/sample_images/cheetah.jpeg?raw=true -O cheetah.jpg > /dev/null 2>&1
```

--------------------------------

### Convert Tokens to Residual Directions

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/tests/test_notebooks/Untitled.ipynb

Converts the token embeddings to residual directions, which can be useful for understanding how tokens are represented in the model's internal state.

```python
test = hooked_model.tokens_to_residual_directions()
```

--------------------------------

### Get ImageNet Indices for Classes

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/demos/ViT_Prisma_Main_Demo.ipynb

Fetches the ImageNet indices for specified class names, such as 'tabby' and 'border collie'. This is a prerequisite for subsequent logit attribution calculations.

```python
## Accumulated residual
cat_index = imagenet_index_from_word("tabby")
dog_index = imagenet_index_from_word("border collie")

print("Tabby index:", cat_index)
print("Border collie index: ", dog_index)
```

--------------------------------

### Extract and Prepare Attention Patterns

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/demos/Interactive_Attention_Head_Tour.ipynb

Extracts attention patterns from the cache for all heads and layers, then concatenates them for visualization. It also selects a specific batch index for focused analysis.

```python
from vit_prisma.visualization.visualize_attention import plot_attn_heads
import torch

all_attentions = [cache['pattern', head_idx] for head_idx in range(12)]

BATCH_IDX = 0
all_attentions = [attn[BATCH_IDX] for attn in all_attentions]
all_attentions = torch.cat(all_attentions, dim=0)
all_attentions.shape # n_heads x n_patches x n_patches
```

--------------------------------

### Load SAE from Hugging Face Hub

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/demos/1_Load_SAE.ipynb

Downloads SAE weights and configuration from a Hugging Face repository and loads the SparseAutoencoder model. Requires `huggingface_hub` and `vit_prisma` libraries. Takes repository ID, file name, and config name as input, and returns a SparseAutoencoder object.

```python
from huggingface_hub import hf_hub_download, list_repo_files
from vit_prisma.sae import SparseAutoencoder

def load_sae(repo_id, file_name, config_name):
    # Step 1: Download SAE from Hugginface
    sae_path = hf_hub_download(repo_id, file_name) # Download weights
    hf_hub_download(repo_id, config_name) # Download config

    # Step 2: Load the pretrained SAE weights from the downloaded path
    print(f"Loading SAE from {sae_path}...")
    sae = SparseAutoencoder.load_from_pretrained(sae_path) # This now automatically gets config.json and converts into the VisionSAERunnerConfig object
    return sae

repo_id = "Prisma-Multimodal/sparse-autoencoder-clip-b-32-sae-vanilla-x64-layer-10-hook_mlp_out-l1-1e-05" # Change this to your chosen SAE. See /docs for a list of SAEs.
file_name = "weights.pt" # Usually weights.pt but may have slight naming variation. See the original HF repo for the exact file name
config_name = "config.json"
sae = load_sae(repo_id, file_name, config_name)
```

--------------------------------

### Get Timm Model Output Shape in Python

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/tests/test_notebooks/timm_comparison.ipynb

Calculates and returns the shape of the output tensor from the Timm model when passed the same predefined input image.

```python
timm_model(input_image).shape
```

--------------------------------

### Get HookedViT Model Output Shape in Python

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/tests/test_notebooks/timm_comparison.ipynb

Calculates and returns the shape of the output tensor from the HookedViT model when passed a predefined input image.

```python
hooked_model(input_image).shape
```

--------------------------------

### Initialize ViT-Prisma with Data Utilities (Python)

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/demos/Emoji_Logit_Lens_Demo.ipynb

This snippet demonstrates how to import and utilize core components of the ViT-Prisma library, specifically focusing on data utility functions for ImageNet mappings and emoji representations. It also shows the import of Plotly for potential data visualization.

```python
import vit_prisma
from vit_prisma.utils.data_utils.imagenet_emoji import IMAGENET_EMOJI
from vit_prisma.utils.data_utils.imagenet_dict import IMAGENET_DICT

import plotly.express as px
```

--------------------------------

### Stack Head Results for Analysis

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/tests/test_notebooks/Untitled.ipynb

Stacks the results from specific attention heads across a layer and position slice. This is useful for analyzing the contribution of individual heads to the model's output.

```python
per_head_residual, labels = cache.stack_head_results(
    layer=-1, pos_slice=-1, return_labels=True
)
```

--------------------------------

### Inspect SAE Configuration

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/demos/1_Load_SAE.ipynb

Prints the configuration object of a loaded Sparse Autoencoder. This helps in understanding the model's architecture, training parameters, and dataset details. Requires the `pprint` module for formatted output.

```python
from pprint import pprint
sae.cfg
pprint(sae.cfg)
```

--------------------------------

### Verifying Hugging Face Cache Configuration

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/demos/Bonus_SAE_notebooks/Evaluating_Gytis_SAE.ipynb

Prints the configured Hugging Face home directory and transformers cache directory to the console. This confirms that the environment variables for cache locations have been set correctly.

```python
import os
print(os.environ.get("HF_HOME"))
print(os.environ.get("TRANSFORMERS_CACHE"))
```

--------------------------------

### Load Pretrained CLIP Model with HookedViT

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/demos/1_Load_SAE.ipynb

Loads a pretrained CLIP model from Huggingface using the `load_hooked_model` function. This function integrates the pretrained weights into a HookedViT object, enabling easy access to intermediate activations. It requires the `vit_prisma.models.model_loader` module and the model name as input. The loaded model is then moved to the specified device.

```python
# Load model
from vit_prisma.models.model_loader import load_hooked_model

model_name = sae.cfg.model_name
model = load_hooked_model(model_name)
model.to(DEVICE) # Move to device
```

--------------------------------

### Visualize All Attention Heads

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/demos/Interactive_Attention_Head_Tour.ipynb

Plots all attention heads of the ViT model, using log transformation for better visibility of subtle patterns. It helps in understanding the overall attention distribution across the image patches.

```python
from vit_prisma.visualization.visualize_attention import plot_attn_heads

plot_attn_heads(all_attentions, global_min_max=True, figsize=(15,15), global_normalize=False, log_transform=True)
```

--------------------------------

### Initialize and Compare ViT Models in Python

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/tests/test_notebooks/timm_comparison.ipynb

Initializes a ViT model using HookedViT (Prisma) and Timm libraries with the same configuration and compares their outputs on a random input image. It sets up models on a specified device and uses a fixed random seed for reproducible comparisons.

```python
#currently only vit_base_patch16_224 supported (config loading issue)
TOLERANCE = 1e-5

model_name = "vit_base_patch16_224"
batch_size = 5
channels = 3
height = 224
width = 224
device = "cuda"

hooked_model = HookedViT.from_pretrained(model_name)
hooked_model.to(device)
timm_model = timm.create_model(model_name, pretrained=True)
timm_model.to(device)

with torch.random.fork_rng():
    torch.manual_seed(1)
    input_image = torch.rand((batch_size, channels, height, width)).to(device)

assert torch.allclose(hooked_model(input_image), timm_model(input_image), atol=TOLERANCE), "Model output diverges!"
```

--------------------------------

### Citation for Prisma: An Open Source Toolkit

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/docs/README.md

BibTeX entry for citing the 'Prisma: An Open Source Toolkit for Mechanistic Interpretability in Vision and Video' paper. This entry includes author information, title, year, arXiv details, and URL for academic and research purposes.

```bibtex
@misc{joseph2025prismaopensourcetoolkit,
      title={Prisma: An Open Source Toolkit for Mechanistic Interpretability in Vision and Video},
      author={Sonia Joseph and Praneet Suresh and Lorenz Hufe and Edward Stevinson and Robert Graham and Yash Vadi and Danilo Bzdok and Sebastian Lapuschkin and Lee Sharkey and Blake Aaron Richards},
      year={2025},
      eprint={2504.19475},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2504.19475}
}
```

--------------------------------

### Load and Display Image with ViT-Prisma

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/demos/Bonus_SAE_notebooks/Evaluating_Gytis_SAE.ipynb

This code snippet demonstrates loading an image from a URL using `load_image` and displaying it. The output indicates the image loading process and the resulting PIL Image object.

```python
img = load_image("https://images.featurelab.xyz/example_2.jpg")
img
```

--------------------------------

### Load Hooked Model using Prisma

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/docs/README.md

Demonstrates how to load a vision/video model using Prisma's `load_hooked_model` function. It specifies a model name from OpenCLIP and shows how to move the model to the CUDA device if available. This is crucial for preparing models for interpretability analysis.

```python
from vit_prisma.models.model_loader import load_hooked_model

model_name = "open-clip:laion/CLIP-ViT-B-16-DataComp.L-s1B-b8K"
model = load_hooked_model(model_name)
model.to('cuda') # Move to cuda if available
```

--------------------------------

### Define Image Transformations

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/demos/Interactive_Attention_Head_Tour.ipynb

Defines a series of image transformations using torchvision.transforms, including converting images to 3 channels, resizing, and converting to a PyTorch tensor. This is crucial for preparing image data for a Vision Transformer model.

```python
from torchvision import transforms

# Define your transformations, including the custom ConvertTo3Channels
transform = transforms.Compose([
    ConvertTo3Channels(),  # Ensure all images are 3-channel without turning them grayscale
    transforms.Resize((224, 224)),  # Resize images to a common size.
    transforms.ToTensor(),  # Convert images to tensor.
    # You can include normalization if desired, using correct values for 3-channel images.
])

def transform_batch(examples):
    images = [transform(image) for image in examples['image']]
    labels = torch.tensor(examples['label'])
    return {'image': images, 'label': labels}
```

--------------------------------

### Generalize Attention Head Visualization Across Images

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/demos/Interactive_Attention_Head_Tour.ipynb

Visualizes a specific attention head across multiple images from the batch to check for pattern generality. It utilizes helper functions to retrieve attention patterns and format images for plotting.

```python
# Our plot_javascript function takes in a list of attention heads and list of images
# Let's plot the same attention head for multiple images to validate its generality

# Specify attention head index
attn_head_idx = 1

# Put in list so right format for function
list_of_attn_heads = attn_head_idx
all_patterns = get_attn_across_datapoints(attn_head_idx)

# Put images in right format
list_of_images = [image for image in images] # turn batch tensor into list
```

--------------------------------

### Load and Prepare Image for ViT

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/demos/ViT_Prisma_Main_Demo.ipynb

Loads an image from a file, applies necessary transformations, and displays the processed image. Assumes 'Image' and 'transform' are pre-defined, and 'plot_image' is available for visualization.

```python
image = Image.open('cat_dog.jpeg')
image = transform(image)
plot_image(image)
```

--------------------------------

### Get ImageNet Index to Name Mapping

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/src/vit_prisma/sae/evals/Evaluating_Vision_SAE.ipynb

This code snippet demonstrates how to retrieve a mapping from ImageNet indices to their corresponding class names using the `get_imagenet_index_to_name` function from `vit_prisma.dataloaders.imagenet_dataset`. This mapping is essential for interpreting the labels associated with image data.

```python
from vit_prisma.dataloaders.imagenet_dataset import get_imagenet_index_to_name
ind_to_name = get_imagenet_index_to_name()
```

--------------------------------

### Get Text Embeddings with CLIP

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/src/vit_prisma/sae/evals/Evaluating_Vision_SAE.ipynb

Generates text embeddings for a given list of texts using a specified CLIP model. It handles batching for efficient processing and normalizes the embeddings. Dependencies include PyTorch and Hugging Face Transformers.

```python
import torch
from transformers import CLIPModel, CLIPProcessor


device = 'cuda'

def get_text_embeddings(model_name, original_text, batch_size=32):
    vanilla_model = CLIPModel.from_pretrained(model_name)
    
    processor = CLIPProcessor.from_pretrained(model_name, do_rescale=False)

    # Split the text into batches
    text_batches = [original_text[i:i+batch_size] for i in range(0, len(original_text), batch_size)]

    all_embeddings = []

    for batch in text_batches:
        inputs = processor(text=batch, return_tensors='pt', padding=True, truncation=True, max_length=77)
        # inputs = {k: v.to(cfg.device) for k, v in inputs.items()}

        with torch.no_grad():
            text_embeddings = vanilla_model.get_text_features(**inputs)

        text_embeddings = text_embeddings / text_embeddings.norm(dim=-1, keepdim=True)
        all_embeddings.append(text_embeddings)

    # Concatenate all batches
    final_embeddings = torch.cat(all_embeddings, dim=0)

    return final_embeddings
```

--------------------------------

### Citation for ViT Prisma: A Mechanistic Interpretability Library

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/docs/README.md

BibTeX entry for citing the 'ViT Prisma: A Mechanistic Interpretability Library for Vision Transformers' project. This entry provides author, title, year, publisher, journal, and a direct URL to the GitHub repository.

```bibtex
@misc{
  joseph2023vit,
  author = {Sonia Joseph},
  title = {ViT Prisma: A Mechanistic Interpretability Library for Vision Transformers},
  year = {2023},
  publisher = {GitHub},
  journal = {GitHub repository},
  howpublished = {\url{https://github.com/soniajoseph/vit-prisma}}
}
```

--------------------------------

### Error: SSL Certificate Verification Failed

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/tests/test_notebooks/Untitled.ipynb

This snippet highlights a common network-related error during data download, specifically an SSL certificate verification failure. This can occur due to system-wide certificate issues or network configurations preventing proper verification.

```python
# This code block is part of an error message and demonstrates a potential issue.
# URLError: <urlopen error [SSL: CERTIFICATE_VERIFY_FAILED] certificate verify failed: unable to get local issuer certificate (_ssl.c:1129)>
```

--------------------------------

### Load ImageNet Dataset with CLIP Transforms

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/demos/1_Load_SAE.ipynb

Loads the ImageNet validation dataset using `torchvision.datasets.ImageFolder`. It applies CLIP-specific transformations for model input and also loads a version with standard transformations for visualization. The function returns a subset dataloader, the transformed validation dataset, and the visualization dataset. It requires `get_clip_val_transforms` from `vit_prisma.transforms` and `torchvision`.

```python
# Put paths here
from vit_prisma.transforms import get_clip_val_transforms
import torchvision

def load_imagenet(imagenet_validation_path):
    torch.set_grad_enabled(False)


    # Load dataset with CLIP transform
    data_transforms = get_clip_val_transforms()
    val_data = torchvision.datasets.ImageFolder(imagenet_validation_path, transform=data_transforms)

    # We'll also load a version of the dataset without the CLIP transform so that we can visualize it beautifully
    viz_transforms = torchvision.transforms.Compose(
                    [
                        torchvision.transforms.Resize((224, 224)),
                        torchvision.transforms.ToTensor(),
                    ]
                )
    viz_data = torchvision.datasets.ImageFolder(imagenet_validation_path, transform=viz_transforms)

    # We only want a subset of validation
            
    subset_dataloader = DataLoader(val_data, batch_size=BATCH_SIZE, shuffle=False, num_workers=4)

    # We only want a subset 
    
    return subset_dataloader, val_data, viz_data
```

--------------------------------

### Initialize Stable UnCLIP Img2Img Pipeline

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/src/vit_prisma/sae/evals/mustache_decoder.ipynb

Initializes the StableUnCLIPImg2ImgPipeline from the diffusers library, specifying the model ID, data type, and cache directory. The pipeline is then moved to the CUDA device.

```python
from diffusers import StableUnCLIPImg2ImgPipeline
import torch

cache_dir = '/network/scratch/s/sonia.joseph/diffusion'

model_id = "stabilityai/stable-diffusion-2-1-unclip"
cache_dir = ...
pipe = StableUnCLIPImg2IMGPipeline.from_pretrained(model_id, torch_dtype=torch.float16, cache_dir=cache_dir)
pipe = pipe.to("cuda")
```

--------------------------------

### Compare Timm and Prisma QKV Activations in Python

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/tests/test_notebooks/timm_comparison.ipynb

Compares Query, Key, and Value (QKV) activations from a Timm ViT model against cached values. It utilizes forward hooks to capture intermediate outputs from the Timm model and then asserts their closeness using torch.allclose.

```python
# First layer
activations = []
def hook_fn(module, input, output):
    activations.append(output)

hook_handle = timm_model.blocks[0].attn.qkv.register_forward_hook(hook_fn)

timm_output = timm_model(image)
hook_handle.remove()

print("timm output", activations[0].shape)
# qkv = activations[0].reshape(-1, 197, 3, 12, 64).permute(2, 0, 3, 1, 4)
# qkv = activations[0].reshape(-1, 197, 3, 12, 64).permute(2, 0, 3, 1, 4)
q, k, v = qkv.unbind(0)

print("timm shape", qkv.shape)
# print("prisma shape", cache['blocks.0.attn.hook_qkv'].shape)

print("prisma q shape", cache['blocks.0.attn.hook_q'].shape)

# assert torch.allclose(qkv, cache['blocks.0.attn.hook_qkv'], atol=1e-6), "Activations differ more than the allowed tolerance"
assert torch.allclose(q, cache['blocks.0.attn.hook_q'], atol=1e-6), "Activations differ more than the allowed tolerance"
assert torch.allclose(k, cache['blocks.0.attn.hook_k'], atol=1e-6), "Activations differ more than the allowed tolerance"
assert torch.allclose(v, cache['blocks.0.attn.hook_v'], atol=1e-6), "Activations differ more than the allowed tolerance"
```

--------------------------------

### Visualize Attention Patterns

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/demos/Interactive_Attention_Head_Tour.ipynb

Visualizes attention scores or patterns from a Vision Transformer model. It allows specifying which attention heads to visualize and extracts the corresponding attention data from the model's cache. The output is a tensor representing the combined attention patterns.

```python
import torch
from typing import List
from jaxtyping import Float

def visualize_attention(
    heads,
    local_cache,
    title,
    max_width,
    batch_idx,
    attention_type = 'attn_scores' # or 'attn_patterns'
) -> str:
    # If a single head is given, convert to a list
    if isinstance(heads, int):
        heads = [heads]

    # Create the plotting data
    labels: List[str] = []
    patterns: List[Float[torch.Tensor, "dest_pos src_pos"]] = []

    for head in heads:
        # Set the label
        layer = head // model.cfg.n_heads
        head_index = head % model.cfg.n_heads
        labels.append(f"L{layer}H{head_index}")

        # Get the attention patterns for the head
        # Attention patterns have shape [batch, head_index, query_pos, key_pos]
        patterns.append(local_cache[attention_type, layer][batch_idx, head_index])

    # Combine the patterns into a single tensor
    patterns = torch.stack(
        patterns, dim=0
    )
    return patterns

def get_attn_across_datapoints(attn_head_idx, attention_type="attn_scores"):

  list_of_attn_heads = [attn_head_idx]

  # Retrieve the activations from each batch idx
  all_patterns = []
  for batch_idx in range(images.shape[0]):
    patterns = visualize_attention(list_of_attn_heads, cache, "Attention Scores", 700, batch_idx = batch_idx, attention_type=attention_type)
    all_patterns.extend(patterns)
  all_patterns = torch.stack(all_patterns, dim=0)
  return all_patterns
```

--------------------------------

### Initialize and Run Vision SAE Trainer

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/demos/2_Train_SAE.ipynb

This Python snippet initializes the VisionSAETrainer with the defined configuration and datasets, then executes the training process by calling the `run()` method. The result, representing the trained SAE, is stored in the `sae` variable.

```python
trainer = VisionSAETrainer(sae_trainer_cfg, model, train_dataset, eval_dataset)
sae = trainer.run()
```

--------------------------------

### Get Class Index and Residual Direction (Python)

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/demos/ViT_Prisma_Main_Demo.ipynb

This snippet retrieves the index for a given word (e.g., 'banana') from the ImageNet index and then computes the corresponding residual direction vector. It's used to isolate the model's focus on a specific class.

```python
misc_index = imagenet_index_from_word("banana")
print("Banana index:", misc_index)
misc_residual_direction = model.tokens_to_residual_directions(misc_index)
```

--------------------------------

### Device and Batch Size Configuration - Python

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/demos/1_Load_SAE.ipynb

Sets the computation device (GPU or CPU) and the batch size for data loading. The 'DEVICE' variable should be set to 'cuda' if a GPU is available, otherwise 'cpu'. 'BATCH_SIZE' controls the number of samples processed at once.

```python
DEVICE = 'cuda' # change to cpu if cpu only paradigm
BATCH_SIZE = 16
```

--------------------------------

### Loading the Prisma HookedViT Model

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/demos/Bonus_SAE_notebooks/Evaluating_Gytis_SAE.ipynb

Loads the pre-trained HookedViT model from the Prisma library using a specified checkpoint ('open-clip:laion/CLIP-ViT-bigG-14-laion2B-39B-b160k'). The model is moved to the designated device (GPU), converted to half-precision, and its output normalization is disabled for compatibility.

```python
# Load the Prisma HookedViT model
from vit_prisma.models.model_loader import load_hooked_model
from pathlib import Path


def load_prisma_model():
    model = load_hooked_model('open-clip:laion/CLIP-ViT-bigG-14-laion2B-39B-b160k')
    model.half().to(DEVICE)
    model.cfg.normalize_output = False
    return model
```

--------------------------------

### Mapping WNIDs to Numeric Labels and Class Names

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/tests/test_notebooks/Untitled.ipynb

Defines dictionaries to map WordNet IDs (wnids) to numeric labels and corresponding class names. It then creates a new mapping from numeric labels directly to class names, which is convenient for model output interpretation.

```python
# Your initial mapping from wnids to numeric labels
wnid_to_numeric = {
    'n01440764': 0,
    'n02102040': 1,
    'n02979186': 2,
    'n03000684': 3,
    'n03028079': 4,
    'n03394916': 5,
    'n03417042': 6,
    'n03425413': 7,
    'n03445777': 8,
    'n03888257': 9
}

# Your mapping from wnids to class names
lbl_dict = {
    'n01440764': 'tench',
    'n02102040': 'English springer',
    'n02979186': 'cassette player',
    'n03000684': 'chain saw',
    'n03028079': 'church',
    'n03394916': 'French horn',
    'n03417042': 'garbage truck',
    'n03425413': 'gas pump',
    'n03445777': 'golf ball',
    'n03888257': 'parachute'
}

# Create a new mapping from numeric labels to class names
numeric_to_class_name = {value: lbl_dict[key] for key, value in wnid_to_numeric.items()}

# Now, numeric_to_class_name maps numbers to class names directly
print(numeric_to_class_name)

# Example: Get the class name for numeric label 0
numeric_label = 0
class_name = numeric_to_class_name[numeric_label]
print(f"Class name for label {numeric_label}")
```

--------------------------------

### Process Image and Generate Patch Logits Dictionary

Source: https://github.com/prisma-multimodal/vit-prisma/blob/main/demos/Emoji_Logit_Lens_Demo.ipynb

Prepares an image for ViT analysis, extracts patch logit directions, and compiles them into a dictionary for visualization. This involves loading the image, transforming it, running it through the model with caching, and then processing the cache to get patch-specific logits.

```python
image = Image.open('fruit.jpg')
image = transform(image)
output, cache = model.run_with_cache(image.unsqueeze(0))
all_answers = model.tokens_to_residual_directions(np.arange(1000))
patch_logit_directions, labels = get_patch_logit_directions(cache, all_answers)
patch_logit_dictionary = get_patch_logit_dictionary(patch_logit_directions, batch_idx=0, rank_label='toilet')
```