+
+## Features
+
+- Customizable dashboards with various plots and data representations for SAE features
+- Support for any SAE in the SAELens library
+- Neuronpedia integration for hosting and comprehensive neuron analysis (note: this requires a Neuronpedia account and is currently only used internally)
+- Ability to handle large datasets and models efficiently
+
+## Installation
+
+Install SAEDashboard using pip:
+```bash
+pip install sae-dashboard
+```
+
+
+## Quick Start
+
+Here's a basic example of how to use SAEDashboard with SaeVisRunner:
+
+```python
+from sae_lens import SAE
+from transformer_lens import HookedTransformer
+from sae_dashboard.sae_vis_data import SaeVisConfig
+from sae_dashboard.sae_vis_runner import SaeVisRunner
+
+# Load model and SAE
+model = HookedTransformer.from_pretrained("gpt2-small", device="cuda", dtype="bfloat16")
+sae, _, _ = SAE.from_pretrained(
+ release="gpt2-small-res-jb",
+ sae_id="blocks.6.hook_resid_pre",
+ device="cuda"
+)
+sae.fold_W_dec_norm()
+
+# Configure visualization
+config = SaeVisConfig(
+ hook_point=sae.cfg.hook_name,
+ features=list(range(256)),
+ minibatch_size_features=64,
+ minibatch_size_tokens=256,
+ device="cuda",
+ dtype="bfloat16"
+)
+
+# Generate data
+data = SaeVisRunner(config).run(encoder=sae, model=model, tokens=your_token_dataset)
+
+# Save feature-centric visualization
+from sae_dashboard.data_writing_fns import save_feature_centric_vis
+save_feature_centric_vis(sae_vis_data=data, filename="feature_dashboard.html")
+```
+
+For a more detailed tutorial, check out our [demo notebook](https://colab.research.google.com/drive/1oqDS35zibmL1IUQrk_OSTxdhcGrSS6yO?usp=drive_link).
+
+## Advanced Usage: Neuronpedia Runner
+
+For internal use or advanced analysis, SAEDashboard provides a Neuronpedia runner that generates data compatible with Neuronpedia. Here's a basic example:
+
+```python
+from sae_dashboard.neuronpedia.neuronpedia_runner_config import NeuronpediaRunnerConfig
+from sae_dashboard.neuronpedia.neuronpedia_runner import NeuronpediaRunner
+
+config = NeuronpediaRunnerConfig(
+ sae_set="your_sae_set",
+ sae_path="path/to/sae",
+ np_set_name="your_neuronpedia_set_name",
+ huggingface_dataset_path="dataset/path",
+ n_prompts_total=1000,
+ n_features_at_a_time=64
+)
+
+runner = NeuronpediaRunner(config)
+runner.run()
+```
+
+For more options and detailed configuration, refer to the `NeuronpediaRunnerConfig` class in the code.
+
+## Cross-Layer Transcoder (CLT) Support
+
+SAEDashboard now supports visualization of Cross-Layer Transcoders (CLTs), which are a variant of SAEs that process activations across transformer layers. To use CLT visualization:
+
+### Required Files
+
+When using a CLT model, you'll need these files in your CLT model directory:
+
+1. **Model weights**: A `.safetensors` or `.pt` file containing the CLT weights
+2. **Configuration**: A `cfg.json` file with the CLT configuration, including:
+ - `num_features`: Number of features in the CLT
+ - `num_layers`: Number of transformer layers
+ - `d_model`: Model dimension
+ - `activation_fn`: Activation function (e.g., "jumprelu", "relu")
+ - `normalization_method`: How inputs are normalized (e.g., "mean_std", "none")
+ - `tl_input_template`: TransformerLens hook template (e.g., "blocks.{}.ln2.hook_normalized"). Note that this will usually differ from the hook name in the model's cfg.json, which is based on NNsight/transformers. You will need to find the corresponding TransformerLens hook name.
+3. **Normalization statistics** (if `normalization_method` is "mean_std"): A `norm_stats.json` file containing the mean and standard deviation for each layer's inputs, generated from the dataset when activations were generated (or afterwards). The file should have this structure:
+ ```json
+ {
+ "0": {
+ "inputs": {
+ "mean": [0.1, -0.2, ...], // Array of d_model values
+ "std": [1.0, 0.9, ...] // Array of d_model values
+ }
+ },
+ "1": {
+ "inputs": {
+ "mean": [...],
+ "std": [...]
+ }
+ },
+ // ... entries for each layer
+ }
+ ```
+
+### Example Usage
+
+```python
+from sae_dashboard.neuronpedia.neuronpedia_runner_config import NeuronpediaRunnerConfig
+from sae_dashboard.neuronpedia.neuronpedia_runner import NeuronpediaRunner
+
+config = NeuronpediaRunnerConfig(
+ sae_set="your_clt_set",
+ sae_path="/path/to/clt/model/directory", # Directory containing the files above
+ model_id="gpt2", # Base model the CLT was trained on
+ outputs_dir="clt_outputs",
+ huggingface_dataset_path="your/dataset",
+ use_clt=True, # Enable CLT mode
+ clt_layer_idx=5, # Which layer to visualize (0-indexed)
+ clt_weights_filename="model.safetensors", # Optional: specify exact weights file
+ n_prompts_total=1000,
+ n_features_at_a_time=64
+)
+
+runner = NeuronpediaRunner(config)
+runner.run()
+```
+
+### Notes on CLT Support
+
+- CLTs must be loaded from local files (HuggingFace Hub loading not yet supported)
+- The `--use-clt` flag is mutually exclusive with `--use-transcoder` and `--use-skip-transcoder`
+- JumpReLU activation functions with learned thresholds are supported
+- The visualization will show features for the specified layer only
+
+## Configuration Options
+
+SAEDashboard offers a wide range of configuration options for both SaeVisRunner and NeuronpediaRunner. Key options include:
+
+- `hook_point`: The layer to analyze in the model
+- `features`: List of feature indices to visualize
+- `minibatch_size_features`: Number of features to process in each batch
+- `minibatch_size_tokens`: Number of tokens to process in each forward pass
+- `device`: Computation device (e.g., "cuda", "cpu")
+- `dtype`: Data type for computations
+- `sparsity_threshold`: Threshold for feature sparsity (Neuronpedia runner)
+- `n_prompts_total`: Total number of prompts to analyze
+- `use_wandb`: Enable logging with Weights & Biases
+
+Refer to `SaeVisConfig` and `NeuronpediaRunnerConfig` for full lists of options.
+
+## Contributing
+
+This project uses [Poetry](https://python-poetry.org/) for dependency management. After cloning the repo, install dependencies with `poetry lock && poetry install`.
+
+We welcome contributions to SAEDashboard! Please follow these steps:
+
+1. Fork the repository
+2. Create a new branch for your feature
+3. Implement your changes
+4. Run tests and checks:
+ - Use `make format` to format your code
+ - Use `make check-ci` to run all checks and tests
+5. Submit a pull request
+
+Ensure your code passes all checks, including:
+- Black and Flake8 for formatting and linting
+- Pyright for type-checking
+- Pytest for tests
+
+## Citing This Work
+
+To cite SAEDashboard in your research, please use the following BibTeX entry:
+
+```bibtex
+@misc{sae_dashboard,
+ title = {{SAE Dashboard}},
+ author = {Decode Research},
+ howpublished = {\url{https://github.com/jbloomAus/sae-dashboard}},
+ year = {2024}
+}
+```
+
+## License
+
+SAE Dashboard is licensed under the MIT License. See the [LICENSE](LICENSE) file for details.
+
+## Acknowledgment and Citation
+
+This project is based on the work by Callum McDougall. If you use SAEDashboard in your research, please cite the original SAEVis project as well:
+
+```bibtex
+@misc{sae_vis,
+ title = {{SAE Visualizer}},
+ author = {Callum McDougall},
+ howpublished = {\url{https://github.com/callummcdougall/sae_vis}},
+ year = {2024}
+}
+```
+
+## Contact
+
+For questions or support, please [open an issue](https://github.com/your-username/sae-dashboard/issues) on our GitHub repository.
+
diff --git a/SAEDashboard/docker/docker-entrypoint.sh b/SAEDashboard/docker/docker-entrypoint.sh
new file mode 100644
index 0000000000000000000000000000000000000000..6911523002c332b35fe01b28036c6ca2b5446f88
--- /dev/null
+++ b/SAEDashboard/docker/docker-entrypoint.sh
@@ -0,0 +1,11 @@
+#!/bin/sh
+
+set -e
+
+# activate our virtual environment here
+. /opt/pysetup/.venv/bin/activate
+
+# You can put other setup logic here
+
+# Evaluating passed command:
+exec "$@"
\ No newline at end of file
diff --git a/SAEDashboard/docker/docker-hub.yaml b/SAEDashboard/docker/docker-hub.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..a3d65f47dc2f6571cae2ba1e1ff1155481049d77
--- /dev/null
+++ b/SAEDashboard/docker/docker-hub.yaml
@@ -0,0 +1,57 @@
+name: Build and Push Docker Image to Docker Hub
+
+on:
+ push:
+ branches: [ "main" ]
+ pull_request:
+ branches: [ "main" ]
+
+env:
+ REGISTRY: docker.io
+ IMAGE_NAME: decoderesearch/saedashboard-cuda
+
+jobs:
+
+ build:
+
+ runs-on: ubuntu-latest
+
+ steps:
+ - uses: actions/checkout@v3
+ - name: Build the Docker image
+ run: docker build --target development -t ${{ env.IMAGE_NAME }} --file Dockerfile .
+ # test:
+ # runs-on: ubuntu-latest
+ # steps:
+ # - uses: actions/checkout@v2
+ # - name: Test the Docker image
+ # run: docker-compose up -d
+ push_to_registry:
+ name: Push Docker image to Docker Hub
+ runs-on: ubuntu-latest
+ steps:
+ - name: Check out the repo
+ uses: actions/checkout@v3
+
+ - name: Set up Docker Buildx
+ uses: docker/setup-buildx-action@v2
+
+ - name: Log in to Docker Hub
+ uses: docker/login-action@v3
+ with:
+ username: ${{ secrets.DOCKERHUB_USERNAME }}
+ password: ${{ secrets.DOCKERHUB_PASSWORD }}
+
+ - name: Extract metadata (tags, labels) for Docker
+ id: meta
+ uses: docker/metadata-action@v5
+ with:
+ images: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}
+
+ - name: Build and push Docker image
+ uses: docker/build-push-action@v2
+ with:
+ context: "{{defaultContext}}"
+ push: true
+ tags: ${{ steps.meta.outputs.tags }}
+ labels: ${{ steps.meta.outputs.labels }}
\ No newline at end of file
diff --git a/SAEDashboard/neuronpedia_vector_pipeline_demo.ipynb b/SAEDashboard/neuronpedia_vector_pipeline_demo.ipynb
new file mode 100644
index 0000000000000000000000000000000000000000..b31dc881d940d5498c8cf84b818e547ed16c8626
--- /dev/null
+++ b/SAEDashboard/neuronpedia_vector_pipeline_demo.ipynb
@@ -0,0 +1,282 @@
+{
+ "cells": [
+ {
+ "cell_type": "code",
+ "execution_count": 1,
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Number of vectors: 1\n",
+ "Vector dimension: 768\n",
+ "Vector names: ['sentiment_vector']\n"
+ ]
+ }
+ ],
+ "source": [
+ "# Example usage\n",
+ "import json\n",
+ "import torch\n",
+ "from pathlib import Path\n",
+ "from sae_dashboard.neuronpedia.vector_set import VectorSet\n",
+ "\n",
+ "\n",
+ "# Load vector from file. Note that the vectors should be stored in this format, as a list of lists of floats:\n",
+ "# {\n",
+ "# \"vectors\": [\n",
+ "# [vector_1],\n",
+ "# [vector_2],\n",
+ "# ...\n",
+ "# ]\n",
+ "# }\n",
+ "json_path = Path(\"test_vectors/logistic_direction.json\")\n",
+ "\n",
+ "# Load the vector into a VectorSet\n",
+ "vector_set = VectorSet.from_json(\n",
+ " json_path=json_path,\n",
+ " d_model=768, # Example dimension for GPT-2 Small\n",
+ " hook_point=\"blocks.7.hook_resid_pre\",\n",
+ " hook_layer=7,\n",
+ " model_name=\"gpt2\",\n",
+ " names=[\"sentiment_vector\"], # Optional custom name\n",
+ ")\n",
+ "\n",
+ "# Now you can use the vector set\n",
+ "print(f\"Number of vectors: {vector_set.vectors.shape[0]}\")\n",
+ "print(f\"Vector dimension: {vector_set.vectors.shape[1]}\")\n",
+ "print(f\"Vector names: {vector_set.names}\")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 3,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# You can also save and load the vector set as a VectorSet object as opposed to a simple list of lists of floats\n",
+ "vector_set.save(Path(\"test_vectors/logistic_direction_vector_set.json\"))\n",
+ "vector_set = VectorSet.load(Path(\"test_vectors/logistic_direction_vector_set.json\"))"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 4,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from sae_dashboard.neuronpedia.neuronpedia_vector_runner import (\n",
+ " NeuronpediaVectorRunner,\n",
+ " NeuronpediaVectorRunnerConfig,\n",
+ ")\n",
+ "\n",
+ "cfg = NeuronpediaVectorRunnerConfig(\n",
+ " outputs_dir=\"test_outputs/\",\n",
+ " huggingface_dataset_path=\"monology/pile-uncopyrighted\",\n",
+ " vector_dtype=\"float32\",\n",
+ " model_dtype=\"float32\",\n",
+ " # Small test settings\n",
+ " n_prompts_total=16384,\n",
+ " n_tokens_in_prompt=128, # Shorter sequences\n",
+ " n_prompts_in_forward_pass=256,\n",
+ " n_vectors_at_a_time=1,\n",
+ " use_wandb=False, # Disable wandb for testing\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 5,
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Device Count: 1\n",
+ "Using specified vector dtype: float32\n",
+ "SAE Device: mps\n",
+ "Model Device: mps\n",
+ "Model Num Devices: 1\n",
+ "Activation Store Device: mps\n",
+ "Dataset Path: monology/pile-uncopyrighted\n",
+ "Forward Pass size: 128\n",
+ "Total number of tokens: 2097152\n",
+ "Total number of contexts (prompts): 16384\n",
+ "Vector DType: float32\n",
+ "Model DType: float32\n"
+ ]
+ },
+ {
+ "name": "stderr",
+ "output_type": "stream",
+ "text": [
+ "/Users/curttigges/miniconda3/envs/sae-d/lib/python3.12/site-packages/transformers/tokenization_utils_base.py:1601: FutureWarning: `clean_up_tokenization_spaces` was not set. It will be set to `True` by default. This behavior will be depracted in transformers v4.45, and will be then set to `False` by default. For more details check this issue: https://github.com/huggingface/transformers/issues/31884\n",
+ " warnings.warn(\n"
+ ]
+ },
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Loaded pretrained model gpt2 into HookedTransformer\n"
+ ]
+ },
+ {
+ "data": {
+ "application/vnd.jupyter.widget-view+json": {
+ "model_id": "f1a49eee02cd482e9de6deaa88e4afde",
+ "version_major": 2,
+ "version_minor": 0
+ },
+ "text/plain": [
+ "Resolving data files: 0%| | 0/30 [00:00 1024). Running this sequence through the model will result in indexing errors\n",
+ "100%|██████████| 2048/2048 [00:18<00:00, 108.67it/s]\n",
+ "0it [00:00, ?it/s]"
+ ]
+ },
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "========== Running Batch #0 ==========\n"
+ ]
+ },
+ {
+ "data": {
+ "application/vnd.jupyter.widget-view+json": {
+ "model_id": "418cf3ccf15d4ae597e06d24e4c89b11",
+ "version_major": 2,
+ "version_minor": 0
+ },
+ "text/plain": [
+ "Forward passes to cache data for vis: 0%| | 0/60 [00:00┏━━━━━━┳━━━━━━┳━━━━━━━┓\n",
+ "┃ Task ┃ Time ┃ Pct % ┃\n",
+ "┡━━━━━━╇━━━━━━╇━━━━━━━┩\n",
+ "└──────┴──────┴───────┘\n",
+ "\n"
+ ],
+ "text/plain": [
+ "┏━━━━━━┳━━━━━━┳━━━━━━━┓\n",
+ "┃\u001b[1m \u001b[0m\u001b[1mTask\u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1mTime\u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1mPct %\u001b[0m\u001b[1m \u001b[0m┃\n",
+ "┡━━━━━━╇━━━━━━╇━━━━━━━┩\n",
+ "└──────┴──────┴───────┘\n"
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ },
+ {
+ "name": "stderr",
+ "output_type": "stream",
+ "text": [
+ "1it [00:02, 2.65s/it]"
+ ]
+ },
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Output written to test_outputs/gpt2_blocks.7.hook_resid_pre/batch-0.json\n"
+ ]
+ },
+ {
+ "name": "stderr",
+ "output_type": "stream",
+ "text": [
+ "\n"
+ ]
+ }
+ ],
+ "source": [
+ "runner = NeuronpediaVectorRunner(vector_set, cfg)\n",
+ "runner.run()"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": []
+ }
+ ],
+ "metadata": {
+ "kernelspec": {
+ "display_name": "sae-d",
+ "language": "python",
+ "name": "python3"
+ },
+ "language_info": {
+ "codemirror_mode": {
+ "name": "ipython",
+ "version": 3
+ },
+ "file_extension": ".py",
+ "mimetype": "text/x-python",
+ "name": "python",
+ "nbconvert_exporter": "python",
+ "pygments_lexer": "ipython3",
+ "version": "3.12.4"
+ }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
diff --git a/SAEDashboard/notebooks/experiment_gemma_2_9b_dashboard_generation_np.py b/SAEDashboard/notebooks/experiment_gemma_2_9b_dashboard_generation_np.py
new file mode 100644
index 0000000000000000000000000000000000000000..444fca15e531192038dc2bb805ebe9ebcf15ec42
--- /dev/null
+++ b/SAEDashboard/notebooks/experiment_gemma_2_9b_dashboard_generation_np.py
@@ -0,0 +1,52 @@
+# I'm running this in an A100 with 90GB of GPU Ram.
+# I'm using TransformerLens 2.2 which I manually installed from source.
+# I'm a few edits to fix bfloat16 errors (but I've since made PR's so latest SAE Lens / SAE dashboard should be fine here).
+import os
+
+from sae_dashboard.neuronpedia.neuronpedia_runner import (
+ NeuronpediaRunner,
+ NeuronpediaRunnerConfig,
+)
+
+# GET WEIGHTS FROM WANDB
+# import wandb
+# run = wandb.init()
+# artifact = run.use_artifact('jbloom/gemma-2-9b_test/sae_gemma-2-9b_blocks.24.hook_resid_post_114688:v7', type='model')
+# artifact_dir = artifact.download()
+
+
+# Get Sparsity from Wandb (and manually move it accross)
+# import wandb
+# run = wandb.init()
+# artifact = run.use_artifact('jbloom/gemma-2-9b_test/sae_gemma-2-9b_blocks.24.hook_resid_post_114688_log_feature_sparsity:v7', type='log_feature_sparsity')
+# artifact_dir = artifact.download()
+
+NP_OUTPUT_FOLDER = "neuronpedia_outputs/gemma-2-9b-test"
+SAE_SET = "res-jb-test"
+SAE_PATH = "artifacts/sae_gemma-2-9b_blocks.24.hook_resid_post_114688:v7"
+print(SAE_PATH)
+
+# delete output files if present
+os.system(f"rm -rf {NP_OUTPUT_FOLDER}")
+cfg = NeuronpediaRunnerConfig(
+ sae_set=SAE_SET,
+ sae_path=SAE_PATH,
+ outputs_dir=NP_OUTPUT_FOLDER,
+ sparsity_threshold=-6,
+ n_prompts_total=4096,
+ huggingface_dataset_path="monology/pile-uncopyrighted",
+ n_features_at_a_time=1024,
+ n_tokens_in_prompt=128,
+ start_batch=0,
+ end_batch=8,
+ use_wandb=True,
+ sae_device="cuda",
+ model_device="cuda",
+ model_n_devices=1,
+ activation_store_device="cuda",
+ model_dtype="bfloat16",
+ sae_dtype="float32",
+)
+
+runner = NeuronpediaRunner(cfg)
+runner.run()
diff --git a/SAEDashboard/notebooks/sae_dashboard_demo_gemma_2_9b.ipynb b/SAEDashboard/notebooks/sae_dashboard_demo_gemma_2_9b.ipynb
new file mode 100644
index 0000000000000000000000000000000000000000..8ddc3b0d4b02c3f1446fb39fb26096c56e969196
--- /dev/null
+++ b/SAEDashboard/notebooks/sae_dashboard_demo_gemma_2_9b.ipynb
@@ -0,0 +1,618 @@
+{
+ "cells": [
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "# Demo Notebook"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "Steps:\n",
+ "1. Download SAE with SAE Lens.\n",
+ "2. Create a dataset consistent with that SAE. \n",
+ "3. Fold the SAE decoder norm weights so that feature activations are \"correct\".\n",
+ "4. Estimate the activation normalization constant if needed, and fold it into the SAE weights.\n",
+ "5. Run the SAE generator for the features you want."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "# Set Up"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Download Gemma-2-9b weights\n",
+ "\n",
+ "import wandb\n",
+ "\n",
+ "run = wandb.init()\n",
+ "artifact = run.use_artifact(\n",
+ " \"jbloom/gemma-2-9b_test/sae_gemma-2-9b_blocks.24.hook_resid_post_114688:v7\",\n",
+ " type=\"model\",\n",
+ ")\n",
+ "artifact_dir = artifact.download()"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "import wandb\n",
+ "\n",
+ "run = wandb.init()\n",
+ "artifact = run.use_artifact(\n",
+ " \"jbloom/gemma-2-9b_test/sae_gemma-2-9b_blocks.24.hook_resid_post_114688_log_feature_sparsity:v7\",\n",
+ " type=\"log_feature_sparsity\",\n",
+ ")\n",
+ "artifact_dir = artifact.download()"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "import torch\n",
+ "import matplotlib.pyplot as plt\n",
+ "from safetensors.torch import load_file\n",
+ "\n",
+ "# Assume we have a PyTorch tensor\n",
+ "feature_sparsity = load_file(\n",
+ " \"artifacts/sae_gemma-2-9b_blocks.24.hook_resid_post_114688:v7/sparsity.safetensors\"\n",
+ ")[\"sparsity\"]\n",
+ "\n",
+ "# Convert the tensor to a numpy array\n",
+ "data = feature_sparsity.numpy()\n",
+ "\n",
+ "# Create the histogram\n",
+ "plt.hist(data, bins=30, edgecolor=\"black\")\n",
+ "\n",
+ "# Add labels and title\n",
+ "plt.xlabel(\"Value\")\n",
+ "plt.ylabel(\"Frequency\")\n",
+ "plt.title(\"Histogram of PyTorch Tensor\")\n",
+ "\n",
+ "# Show the plot\n",
+ "plt.show()"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "import torch\n",
+ "from transformer_lens import HookedTransformer\n",
+ "from sae_lens import ActivationsStore, SAE\n",
+ "from importlib import reload\n",
+ "import sae_dashboard\n",
+ "\n",
+ "torch.set_grad_enabled(False)\n",
+ "\n",
+ "reload(sae_dashboard)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "MODEL = \"gemma-2-9b\"\n",
+ "\n",
+ "if torch.backends.mps.is_available():\n",
+ " device = \"mps\"\n",
+ "else:\n",
+ " device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
+ "\n",
+ "print(f\"Device: {device}\")\n",
+ "\n",
+ "model = HookedTransformer.from_pretrained(MODEL, device=device, dtype=\"bfloat16\")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "sae = SAE.load_from_pretrained(\n",
+ " \"artifacts/sae_gemma-2-9b_blocks.24.hook_resid_post_114688:v7\"\n",
+ ")\n",
+ "sae.fold_W_dec_norm()"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# _, cache = model.run_with_cache(\"Wasssssup\", names_filter = sae.cfg.hook_name)\n",
+ "# sae_in = cache[sae.cfg.hook_name]\n",
+ "# print(sae_in.shape)\n",
+ "sae_in = torch.rand((1, 4, 3584)).to(sae.device)\n",
+ "sae_out = sae(sae_in)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# # the cfg dict is returned alongside the SAE since it may contain useful information for analysing the SAE (eg: instantiating an activation store)\n",
+ "# # Note that this is not the same as the SAEs config dict, rather it is whatever was in the HF repo, from which we can extract the SAE config dict\n",
+ "# # We also return the feature sparsities which are stored in HF for convenience.\n",
+ "# sae, cfg_dict, sparsity = SAE.from_pretrained(\n",
+ "# release = \"mistral-7b-res-wg\", # see other options in sae_lens/pretrained_saes.yaml\n",
+ "# sae_id = \"blocks.8.hook_resid_pre\", # won't always be a hook point\n",
+ "# device = \"cuda:3\",\n",
+ "# )\n",
+ "# # fold w_dec norm so feature activations are accurate\n",
+ "#\n",
+ "activations_store = ActivationsStore.from_sae(\n",
+ " model=model,\n",
+ " sae=sae,\n",
+ " streaming=True,\n",
+ " store_batch_size_prompts=8,\n",
+ " n_batches_in_buffer=8,\n",
+ " device=\"cpu\",\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "sae.encode_fn"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from sae_lens import run_evals\n",
+ "\n",
+ "eval_metrics = run_evals(\n",
+ " sae=sae,\n",
+ " activation_store=activations_store,\n",
+ " model=model,\n",
+ " n_eval_batches=3,\n",
+ " eval_batch_size_prompts=8,\n",
+ ")\n",
+ "\n",
+ "# CE Loss score should be high for residual stream SAEs\n",
+ "print(eval_metrics[\"metrics/CE_loss_score\"])\n",
+ "\n",
+ "# ce loss without SAE should be fairly low < 3.5 suggesting the Model is being run correctly\n",
+ "print(eval_metrics[\"metrics/ce_loss_without_sae\"])\n",
+ "\n",
+ "# ce loss with SAE shouldn't be massively higher\n",
+ "print(eval_metrics[\"metrics/ce_loss_with_sae\"])"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from tqdm import tqdm\n",
+ "\n",
+ "\n",
+ "from sae_dashboard.utils_fns import get_tokens\n",
+ "\n",
+ "# 1000 prompts is plenty for a demo.\n",
+ "token_dataset = get_tokens(activations_store, 4096)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# torch.save(token_dataset, \"to\")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# torch.save(token_dataset, \"token_dataset.pt\")\n",
+ "token_dataset = torch.load(\"token_dataset.pt\")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "import os\n",
+ "\n",
+ "os.rmdir(\"demo_activations_cache\")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "import torch\n",
+ "\n",
+ "\n",
+ "def select_indices_in_range(tensor, min_val, max_val, num_samples=None):\n",
+ " \"\"\"\n",
+ " Select indices of a tensor where values fall within a specified range.\n",
+ "\n",
+ " Args:\n",
+ " tensor (torch.Tensor): Input tensor with values between -10 and 0.\n",
+ " min_val (float): Minimum value of the range (inclusive).\n",
+ " max_val (float): Maximum value of the range (inclusive).\n",
+ " num_samples (int, optional): Number of indices to randomly select. If None, return all indices.\n",
+ "\n",
+ " Returns:\n",
+ " torch.Tensor: Tensor of selected indices.\n",
+ " \"\"\"\n",
+ " # Ensure the input range is valid\n",
+ " if not (-10 <= min_val <= max_val <= 0):\n",
+ " raise ValueError(\n",
+ " \"Range must be within -10 to 0, and min_val must be <= max_val\"\n",
+ " )\n",
+ "\n",
+ " # Find indices where values are within the specified range\n",
+ " mask = (tensor >= min_val) & (tensor <= max_val)\n",
+ " indices = mask.nonzero().squeeze()\n",
+ "\n",
+ " # If num_samples is specified and less than the total number of valid indices,\n",
+ " # randomly select that many indices\n",
+ " if num_samples is not None and num_samples < indices.numel():\n",
+ " perm = torch.randperm(indices.numel())\n",
+ " indices = indices[perm[:num_samples]]\n",
+ "\n",
+ " return indices\n",
+ "\n",
+ "\n",
+ "n_features = 4096\n",
+ "feature_idxs = select_indices_in_range(feature_sparsity, -4, -2, 4096)\n",
+ "feature_sparsity[feature_idxs.tolist()]"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": []
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from importlib import reload\n",
+ "import sys\n",
+ "\n",
+ "\n",
+ "def reload_user_modules(module_names):\n",
+ " \"\"\"Reload specified user modules.\"\"\"\n",
+ " for name in module_names:\n",
+ " if name in sys.modules:\n",
+ " reload(sys.modules[name])\n",
+ "\n",
+ "\n",
+ "# List of your module names\n",
+ "user_modules = [\n",
+ " \"sae_dashboard\",\n",
+ " \"sae_dashboard.sae_vis_runner\",\n",
+ " \"sae_dashboard.data_parsing_fns\",\n",
+ " \"sae_dashboard.feature_data_generator\",\n",
+ "]\n",
+ "\n",
+ "# Reload modules\n",
+ "reload_user_modules(user_modules)\n",
+ "\n",
+ "# Re-import after reload\n",
+ "from sae_dashboard.feature_data_generator import FeatureDataGenerator"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from pathlib import Path\n",
+ "\n",
+ "test_feature_idx_gpt = feature_idxs.tolist()\n",
+ "\n",
+ "feature_vis_config_gpt = sae_vis_runner.SaeVisConfig(\n",
+ " hook_point=sae.cfg.hook_name,\n",
+ " features=test_feature_idx_gpt,\n",
+ " minibatch_size_features=16,\n",
+ " minibatch_size_tokens=4096, # this is really prompt with the number of tokens determined by the sequence length\n",
+ " verbose=True,\n",
+ " device=\"cuda\",\n",
+ " cache_dir=Path(\n",
+ " \"demo_activations_cache\"\n",
+ " ), # this will enable us to skip running the model for subsequent features.\n",
+ " dtype=\"bfloat16\",\n",
+ ")\n",
+ "\n",
+ "runner = sae_vis_runner.SaeVisRunner(feature_vis_config_gpt)\n",
+ "\n",
+ "data = runner.run(\n",
+ " encoder=sae, # type: ignore\n",
+ " model=model,\n",
+ " tokens=token_dataset[:1024],\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": []
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from sae_dashboard.data_writing_fns import save_feature_centric_vis\n",
+ "\n",
+ "filename = f\"demo_feature_dashboards.html\"\n",
+ "save_feature_centric_vis(sae_vis_data=data, filename=filename)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": []
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "# Quick Profiling experiment"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "def mock_feature_acts_subset_for_now(sae: SAE):\n",
+ "\n",
+ " @torch.no_grad()\n",
+ " def sae_lens_get_feature_acts_subset(x: torch.Tensor, feature_idx): # type: ignore\n",
+ " \"\"\"\n",
+ " Get a subset of the feature activations for a dataset.\n",
+ " \"\"\"\n",
+ " original_device = x.device\n",
+ " feature_activations = sae.encode_fn(x.to(device=sae.device, dtype=sae.dtype))\n",
+ " return feature_activations[..., feature_idx].to(original_device)\n",
+ "\n",
+ " sae.get_feature_acts_subset = sae_lens_get_feature_acts_subset # type: ignore\n",
+ "\n",
+ " return sae\n",
+ "\n",
+ "\n",
+ "sae = mock_feature_acts_subset_for_now(sae)\n",
+ "feature_idxs = list(range(128))\n",
+ "sae_in = torch.rand((1, 4, 3584)).to(sae.device)\n",
+ "sae.get_feature_acts_subset(sae_in, feature_idxs)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "for k, v in sae.named_parameters():\n",
+ " print(k, v.shape)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from torch import nn\n",
+ "from typing import List\n",
+ "\n",
+ "\n",
+ "class FeatureMaskingContext:\n",
+ " def __init__(self, sae: SAE, feature_idxs: List):\n",
+ " self.sae = sae\n",
+ " self.feature_idxs = feature_idxs\n",
+ " self.original_weight = {}\n",
+ "\n",
+ " def __enter__(self):\n",
+ "\n",
+ " ## W_dec\n",
+ " self.original_weight[\"W_dec\"] = getattr(self.sae, \"W_dec\").data.clone()\n",
+ " # mask the weight\n",
+ " masked_weight = sae.W_dec[self.feature_idxs]\n",
+ " # set the weight\n",
+ " setattr(self.sae, \"W_dec\", nn.Parameter(masked_weight))\n",
+ "\n",
+ " ## W_enc\n",
+ " # clone the weight.\n",
+ " self.original_weight[\"W_enc\"] = getattr(self.sae, \"W_enc\").data.clone()\n",
+ " # mask the weight\n",
+ " masked_weight = sae.W_enc[:, self.feature_idxs]\n",
+ " # set the weight\n",
+ " setattr(self.sae, \"W_enc\", nn.Parameter(masked_weight))\n",
+ "\n",
+ " if self.sae.cfg.architecture == \"standard\":\n",
+ "\n",
+ " ## b_enc\n",
+ " self.original_weight[\"b_enc\"] = getattr(self.sae, \"b_enc\").data.clone()\n",
+ " # mask the weight\n",
+ " masked_weight = sae.b_enc[self.feature_idxs]\n",
+ " # set the weight\n",
+ " setattr(self.sae, \"b_enc\", nn.Parameter(masked_weight))\n",
+ "\n",
+ " elif self.sae.cfg.architecture == \"gated\":\n",
+ "\n",
+ " ## b_gate\n",
+ " self.original_weight[\"b_gate\"] = getattr(self.sae, \"b_gate\").data.clone()\n",
+ " # mask the weight\n",
+ " masked_weight = sae.b_gate[self.feature_idxs]\n",
+ " # set the weight\n",
+ " setattr(self.sae, \"b_gate\", nn.Parameter(masked_weight))\n",
+ "\n",
+ " ## r_mag\n",
+ " self.original_weight[\"r_mag\"] = getattr(self.sae, \"r_mag\").data.clone()\n",
+ " # mask the weight\n",
+ " masked_weight = sae.r_mag[self.feature_idxs]\n",
+ " # set the weight\n",
+ " setattr(self.sae, \"r_mag\", nn.Parameter(masked_weight))\n",
+ "\n",
+ " ## b_mag\n",
+ " self.original_weight[\"b_mag\"] = getattr(self.sae, \"b_mag\").data.clone()\n",
+ " # mask the weight\n",
+ " masked_weight = sae.b_mag[self.feature_idxs]\n",
+ " # set the weight\n",
+ " setattr(self.sae, \"b_mag\", nn.Parameter(masked_weight))\n",
+ " else:\n",
+ " raise (ValueError(\"Invalid architecture\"))\n",
+ "\n",
+ " return self\n",
+ "\n",
+ " def __exit__(self, exc_type, exc_value, traceback):\n",
+ "\n",
+ " # set everything back to normal\n",
+ " for key, value in self.original_weight.items():\n",
+ " setattr(self.sae, key, nn.Parameter(value))"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "import gc\n",
+ "import torch\n",
+ "\n",
+ "gc.collect()\n",
+ "torch.cuda.empty_cache()\n",
+ "torch.set_grad_enabled(False)\n",
+ "\n",
+ "\n",
+ "def my_function(sae_in):\n",
+ " # Your PyTorch code here\n",
+ " feature_idxs = list(range(2048))\n",
+ " with FeatureMaskingContext(sae, feature_idxs):\n",
+ " features = sae(sae_in)\n",
+ " print(features.mean())\n",
+ "\n",
+ "\n",
+ "tokens = token_dataset[:64]\n",
+ "_, cache = model.run_with_cache(\n",
+ " tokens, stop_at_layer=sae.cfg.hook_layer + 1, names_filter=sae.cfg.hook_name\n",
+ ")\n",
+ "sae_in = cache[sae.cfg.hook_name]"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "tokens.shape"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "sae.W_dec.shape"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "%load_ext memray"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "%%memray_flamegraph --trace-python-allocators --leaks\n",
+ "my_function(sae_in)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": []
+ }
+ ],
+ "metadata": {
+ "kernelspec": {
+ "display_name": ".venv",
+ "language": "python",
+ "name": "python3"
+ },
+ "language_info": {
+ "codemirror_mode": {
+ "name": "ipython",
+ "version": 3
+ },
+ "file_extension": ".py",
+ "mimetype": "text/x-python",
+ "name": "python",
+ "nbconvert_exporter": "python",
+ "pygments_lexer": "ipython3",
+ "version": "3.11.7"
+ }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
diff --git a/SAEDashboard/pyproject.toml b/SAEDashboard/pyproject.toml
new file mode 100644
index 0000000000000000000000000000000000000000..79cb896017f186fb5dfabaaf671c6d6bead797c1
--- /dev/null
+++ b/SAEDashboard/pyproject.toml
@@ -0,0 +1,70 @@
+[tool.poetry]
+name = "sae-dashboard"
+version = "0.7.3"
+description = "Open-source SAE visualizer, based on Anthropic's published visualizer. Forked / Detached from sae_vis."
+authors = ["Callum McDougall ${tok} (${tokID})${tok}${tok}... element, because then we have to use
+ the literal characters.
+ """
+ for k, v in HTML_ALL_REVERSED.items():
+ str_tok = str_tok.replace(k, v)
+
+ return str_tok
+
+
+@overload
+def to_str_tokens(
+ decode_fn: Callable[[int | list[int]], str | list[str]],
+ tokens: int,
+) -> str: ...
+
+
+@overload
+def to_str_tokens(
+ decode_fn: Callable[[int | list[int]], str | list[str]],
+ tokens: list[int],
+) -> list[str]: ...
+
+
+def to_str_tokens(
+ decode_fn: Callable[[int | list[int]], str | list[str]],
+ tokens: int | list[int] | torch.Tensor,
+) -> str | Any:
+ """
+ Helper function which converts tokens to their string representations, but (if tokens is a tensor) keeps
+ them in the same shape as the original tensor (i.e. nested lists).
+ """
+ # Deal with the int case separately
+ if isinstance(tokens, int):
+ return decode_fn(tokens)
+
+ # If the tokens are a (possibly nested) list, turn them into a tensor
+ if isinstance(tokens, list):
+ tokens = torch.tensor(tokens)
+
+ # Get flattened list of tokens
+ str_tokens = [decode_fn(t) for t in tokens.flatten().tolist()]
+
+ # Reshape
+ return np.reshape(str_tokens, tokens.shape).tolist()
+
+
+class TopK:
+ """
+ This function implements a version of torch.topk over the last dimension. It offers the following:
+
+ (1) Nicer type signatures (the default obj returned by torck.topk isn't well typed)
+ (2) Helper functions for indexing & other standard tensor operations like .ndim, .shape, etc.
+ (3) An efficient topk calculation, which doesn't bother applying it to the zero elements of a tensor.
+ """
+
+ values: Arr
+ indices: Arr
+
+ def __init__(
+ self,
+ tensor: Float[Tensor, "... d"],
+ k: int,
+ largest: bool = True,
+ tensor_mask: Bool[Tensor, "..."] | None = None,
+ ):
+ self.k = k
+ self.largest = largest
+ self.values, self.indices = self.topk(tensor, tensor_mask)
+
+ def __getitem__(self, item: int) -> "TopK":
+ new_topk = TopK.__new__(TopK)
+ new_topk.k = self.k
+ new_topk.largest = self.largest
+ new_topk.values = self.values[item]
+ new_topk.indices = self.indices[item]
+ return new_topk
+
+ def __len__(self) -> int:
+ return len(self.values)
+
+ @property
+ def ndim(self) -> int:
+ return self.values.ndim
+
+ @property
+ def shape(self) -> tuple[int]:
+ return tuple(self.values.shape) # type: ignore
+
+ def numel(self) -> int:
+ return self.values.size
+
+ def topk( # type: ignore
+ self,
+ tensor: Float[Tensor, "... d"],
+ tensor_mask: Bool[Tensor, "..."] | None = None,
+ ) -> tuple[Arr, Arr]:
+ """
+ This is an efficient version of `torch.topk(..., dim=-1)`. It saves time by only doing the topk calculation over
+ the bits of `tensor` where `tensor_mask=True`. This is useful when `tensor` is very sparse, e.g. it has shape
+ (batch, seq, d_vocab) and its elements are zero if the corresponding token has feature activation zero. In this
+ case, we don't want to waste time taking topk over a tensor of zeros.
+ """
+ # If no tensor mask is provided, then we just return the topk values and indices
+ if tensor_mask is None or not tensor_mask.any():
+ k = min(self.k, tensor.shape[-1])
+ topk = tensor.topk(k=k, largest=self.largest)
+ return utils.to_numpy(topk.values), utils.to_numpy(topk.indices)
+
+ # Get the topk of the tensor, but only computed over the values of the tensor which are nontrivial
+ assert (
+ tensor_mask.shape == tensor.shape[:-1]
+ ), "Error: unexpected shape for tensor mask."
+ tensor_nontrivial_values = tensor[tensor_mask] # shape [rows d]
+ k = min(self.k, tensor_nontrivial_values.shape[-1])
+ k = self.k
+ topk = tensor_nontrivial_values.topk(
+ k=k, largest=self.largest
+ ) # shape [rows k]
+
+ # Get an array of indices and values (with unimportant elements) which we'll index into using the topk object
+ topk_shape = (*tensor_mask.shape, k)
+ topk_indices = torch.zeros(
+ topk_shape, device=tensor.device, dtype=torch.long
+ ) # ).long() # shape [... k]
+ topk_indices[tensor_mask] = topk.indices
+ topk_values = torch.zeros(
+ topk_shape, device=tensor.device, dtype=tensor.dtype
+ ) # shape [... k]
+ topk_values[tensor_mask] = topk.values
+
+ return utils.to_numpy(topk_values), utils.to_numpy(topk_indices)
+
+
+def merge_lists(*lists: Iterable[T]) -> list[T]:
+ """
+ Merges a bunch of lists into a single list.
+ """
+ return [item for sublist in lists for item in sublist]
+
+
+def extract_and_remove_scripts(html_content: str) -> tuple[str, str]:
+ """
+ Extracts JavaScript from script tags in the HTML content, and returns it as a single string,
+ along with the original content with the script tags removed.
+ """
+ # Pattern to find tags and capture content inside
+ pattern = r"