Upload folder using huggingface_hub

SAEDashboard/.dockerignore

@@ -0,0 +1,184 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# ruff
+.ruff_cache
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+*.pkl
+*.pt
+sae_vis/archive_fns.py
+*__pycache__
+mats_sae_training
+callum_instructions.md
+april-fools
+*large.html
+requirements.txt
+tests/fixtures/cache_benchmark/
+tests/fixtures/cache_unit/
+
+neuronpedia_outputs/
+cached_activations/
+wandb/
+
+
+**.safetensors
+**flamegraph.html
+artifacts/

SAEDashboard/.flake8

@@ -0,0 +1,8 @@
+[flake8]
+extend-ignore = E203, E266, E501, W503, E721, F722, E731, E402, F821
+max-line-length = 79
+max-complexity = 25
+extend-select = E9, F63, F7, F82
+show-source = true
+statistics = true
+exclude = ./wandb/*, ./research/wandb/*, .venv/*

SAEDashboard/.github/workflows/ci.yaml

@@ -0,0 +1,117 @@
+name: "ci"
+
+on:
+  pull_request:
+    branches: ["**"]
+  push:
+    branches: ["**"]
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ["3.10", "3.11", "3.12"]
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Cache Huggingface assets
+        uses: actions/cache@v4
+        with:
+          key: huggingface-${{ runner.os }}-${{ matrix.python-version }}-${{ hashFiles('**/pyproject.toml') }}
+          path: ~/.cache/huggingface
+          restore-keys: |
+            huggingface-${{ runner.os }}-${{ matrix.python-version }}-
+
+      - name: Load cached Poetry installation
+        id: cached-poetry
+        uses: actions/cache@v4
+        with:
+          path: ~/.local
+          key: poetry-${{ runner.os }}-${{ matrix.python-version }}-1  # Incremented to reset cache
+
+      - name: Install Poetry
+        if: steps.cached-poetry.outputs.cache-hit != 'true'
+        uses: snok/install-poetry@v1
+        with:
+          version: 1.5.1  # Specify a version explicitly
+          virtualenvs-create: true
+          virtualenvs-in-project: true
+          installer-parallel: true
+
+      - name: Check Poetry Version
+        run: poetry --version
+
+      - name: Load cached venv
+        id: cached-poetry-dependencies
+        uses: actions/cache@v4
+        with:
+          path: .venv
+          key: venv-${{ runner.os }}-${{ matrix.python-version }}-${{ hashFiles('**/pyproject.toml') }}-1  # Incremented to reset cache
+          restore-keys: |
+            venv-${{ runner.os }}-${{ matrix.python-version }}-
+
+      - name: Install dependencies
+        if: steps.cached-poetry-dependencies.outputs.cache-hit != 'true'
+        run: poetry install --no-interaction
+
+      - name: List installed packages
+        run: poetry run pip list
+
+      - name: Check flake8 installation
+        run: poetry run which flake8
+
+      - name: check linting
+        run: poetry run flake8 .
+
+      - name: check formatting
+        run: poetry run black --check .
+
+      - name: check types
+        run: poetry run pyright .
+
+      - name: test
+        run: poetry run pytest --cov=sae_dashboard --cov-report=term-missing tests/unit
+
+      - name: build
+        run: poetry build
+
+  release:
+    needs: build
+    permissions:
+      contents: write
+      id-token: write
+    if: github.event_name == 'push' && github.ref == 'refs/heads/main' && !contains(github.event.head_commit.message, 'chore(release):')
+    runs-on: ubuntu-latest
+    concurrency: release
+    environment:
+      name: pypi
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.11"
+
+      - name: Semantic Release
+        id: release
+        uses: python-semantic-release/python-semantic-release@v9.8.8
+        with:
+          github_token: ${{ secrets.GITHUB_TOKEN }}
+
+      - name: Publish package distributions to PyPI
+        uses: pypa/gh-action-pypi-publish@release/v1
+        if: steps.release.outputs.released == 'true'
+
+      - name: Publish package distributions to GitHub Releases
+        uses: python-semantic-release/upload-to-gh-release@main
+        if: steps.release.outputs.released == 'true'
+        with:
+          github_token: ${{ secrets.GITHUB_TOKEN }}

SAEDashboard/.gitignore

@@ -0,0 +1,204 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# ruff
+.ruff_cache
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+*.pkl
+*.pt
+sae_vis/archive_fns.py
+*__pycache__
+mats_sae_training
+callum_instructions.md
+april-fools
+*large.html
+requirements.txt
+tests/fixtures/cache_benchmark/
+tests/fixtures/cache_unit/
+
+neuronpedia_outputs/
+cached_activations/
+wandb/
+demo_activations_cache/
+test_activations_cache/
+demo_feature_dashboards.html
+
+**.safetensors
+**flamegraph.html
+artifacts/
+prof/
+
+.vscode/settings.json
+dfa_tests.ipynb
+
+.DS_Store
+
+# Test and temporary directories
+crosslayer-coding/
+SAELens/
+clt_test*/
+test_output/
+test_outputs/
+clt-technical-description.md
+
+ignore_data/
+
+outputs/

SAEDashboard/.vscode/settings.json

@@ -0,0 +1,18 @@
+{
+  "python.testing.pytestArgs": [
+    "tests"
+],
+  "python.testing.unittestEnabled": false,
+  "python.testing.pytestEnabled": true,
+
+  "[python]": {
+    "editor.defaultFormatter": "ms-python.black-formatter",
+    "editor.formatOnSave": true,
+    "editor.codeActionsOnSave": {
+      "source.organizeImports": "explicit"
+    }
+  },
+  "isort.args": ["--profile", "black"],
+"editor.defaultFormatter": "mikoz.black-py",
+"liveServer.settings.port": 5501
+}

SAEDashboard/CHANGELOG.md

@@ -0,0 +1,1263 @@
+# CHANGELOG
+
+## v0.7.3 (2025-10-11)
+
+### Fix
+
+* fix: broken dependencies ([`25ce6e8`](https://github.com/jbloomAus/SAEDashboard/commit/25ce6e8ae2debe232b0eff4ba910fa10fc816480))
+
+### Unknown
+
+* Merge pull request #70 from jbloomAus/fix_deps
+
+fix: broken dependencies ([`352b9b2`](https://github.com/jbloomAus/SAEDashboard/commit/352b9b2148b62c8fabb7adcf6bf0cbacfa345a74))
+
+* update .gitignore ([`6133ca6`](https://github.com/jbloomAus/SAEDashboard/commit/6133ca67b39bf6e033e8b4792f5eba9850668821))
+
+## v0.7.2 (2025-09-01)
+
+### Fix
+
+* fix: use clt pypi library ([`a3197b8`](https://github.com/jbloomAus/SAEDashboard/commit/a3197b8870d43107bba0356af11c64ddc054392d))
+
+## v0.7.1 (2025-08-31)
+
+### Fix
+
+* fix: force build ([`f432ba2`](https://github.com/jbloomAus/SAEDashboard/commit/f432ba2d14edd4c11ac71e947fbb1e97790e753e))
+
+## v0.7.0 (2025-08-31)
+
+### Feature
+
+* feat: Merge pull request #69 from jbloomAus/qwen-transcoder
+
+Transcoder Support + SAELens v6 ([`8f8651e`](https://github.com/jbloomAus/SAEDashboard/commit/8f8651edaf8c20bc8eaff09e05de238a4ce780fb))
+
+### Fix
+
+* fix: relax saelens to not break saelens demo project ([`e38d140`](https://github.com/jbloomAus/SAEDashboard/commit/e38d1408f30f47f9e91646fa8646300750b23fd3))
+
+### Unknown
+
+* Merge branch 'main' into qwen-transcoder ([`ae9dede`](https://github.com/jbloomAus/SAEDashboard/commit/ae9dedead671996ebf4f7de84eb4252b2af85fc2))
+
+* Merge pull request #67 from jbloomAus/relax_saelens
+
+fix: relax saelens to not break saelens demo project ([`fa1691a`](https://github.com/jbloomAus/SAEDashboard/commit/fa1691ab224e684618b2e800b8fda8af741eb81b))
+
+## v0.6.11 (2025-08-05)
+
+### Fix
+
+* fix: fixes tool.semantic_release subtable (#66) ([`eb36157`](https://github.com/jbloomAus/SAEDashboard/commit/eb361571550a4653f7fbcc5a9cc2c98c329aaf41))
+
+* fix: fixes tool.semantic_release subtable (#66) ([`725d76d`](https://github.com/jbloomAus/SAEDashboard/commit/725d76d9ac00e3b295c6b11f4657f4432c925e9e))
+
+### Unknown
+
+* upgrades python-semantic-release (#65) ([`fcdae8b`](https://github.com/jbloomAus/SAEDashboard/commit/fcdae8b0e18bbf2c5184977bdf14cc14280fb6bc))
+
+* fix CI ([`d7dffdd`](https://github.com/jbloomAus/SAEDashboard/commit/d7dffdd79b0165cbdfb360f13622385f884a2158))
+
+* upgrades python-semantic-release (#65) ([`721d683`](https://github.com/jbloomAus/SAEDashboard/commit/721d683b437f378f22c3e713cb4e3f16bdc82e1a))
+
+* converter ([`30ff988`](https://github.com/jbloomAus/SAEDashboard/commit/30ff9881f4ea1adc3305675088f5fb808367d5eb))
+
+* bos override ([`bd04133`](https://github.com/jbloomAus/SAEDashboard/commit/bd04133033ce5b422f62fc12099eb6527d6f8070))
+
+* add prefix tokens to cli ([`4560dd7`](https://github.com/jbloomAus/SAEDashboard/commit/4560dd7301a713b3b45d836df3e43bf25ba52c64))
+
+* add prefix tokens to cli ([`50636fb`](https://github.com/jbloomAus/SAEDashboard/commit/50636fb667384fbf1df8d7809cfe3c5ebc44beab))
+
+* top acts group 20 ([`1e3d3d4`](https://github.com/jbloomAus/SAEDashboard/commit/1e3d3d4228b5f526f9a2e1d90cb51ff74d8e60b8))
+
+* temp updates for qwen transcoder ([`5727ac9`](https://github.com/jbloomAus/SAEDashboard/commit/5727ac944feed9eb78f60f6553feb40c7b6622d8))
+
+* some config fixes ([`51d903a`](https://github.com/jbloomAus/SAEDashboard/commit/51d903a98ebc2402f6d4863f1efc1062420ea3eb))
+
+* olved double normalization ([`5293073`](https://github.com/jbloomAus/SAEDashboard/commit/5293073adef90efc0804c61d9dbe3ab55430f62b))
+
+* updated readme ([`6ea06f2`](https://github.com/jbloomAus/SAEDashboard/commit/6ea06f2371e80191d8a29cca0c5e134943db02d0))
+
+* formatting changes ([`ad7422e`](https://github.com/jbloomAus/SAEDashboard/commit/ad7422e11a5525d8584991e278a3342ebf4ff892))
+
+* Update CLT test script parameters ([`9961859`](https://github.com/jbloomAus/SAEDashboard/commit/996185947217ccd13d6763462d764cbd2977a28e))
+
+* Merge pull request #64 from jbloomAus/clt-support
+
+CLT Support ([`57971a9`](https://github.com/jbloomAus/SAEDashboard/commit/57971a9e85b62cc2a9e7bf04ee8d8c26fac9cecc))
+
+* Add Cross-Layer Transcoder (CLT) support to SAEDashboard
+
+- Add CLTLayerWrapper to provide SAE-compatible interface for CLTs
+- Integrate CLT loading into NeuronpediaRunner with --use-clt flag
+- Add CLT-specific configuration parameters (clt_layer_idx, clt_weights_filename)
+- Support JumpReLU activation with learned thresholds
+- Add normalization statistics loading from norm_stats.json
+- Handle CLT-specific hook naming conventions (tl_input_template)
+- Add comprehensive unit tests for CLT functionality
+- Fix existing unit tests to use StandardSAE/StandardSAEConfig
+
+🤖 Generated with [Claude Code](https://claude.ai/code)
+
+Co-Authored-By: Claude <noreply@anthropic.com> ([`fab4c6c`](https://github.com/jbloomAus/SAEDashboard/commit/fab4c6cc7a026bbf9beed229734a34323f22c158))
+
+* Add CLT (Cross-Layer Transcoder) support
+
+- Add CLTLayerWrapper to wrap CLT models for SAE-compatible interface
+- Add CLT loading logic in neuronpedia_runner with local file support
+- Add conditional logic to skip fold_W_dec_norm() for CLT wrappers
+- Add conditional logic to skip hook_z_reshaping_mode for CLT wrappers
+- Add support for additional hook types (hook_mlp_out, hook_attn_out, etc.)
+- Add CLI arguments for CLT configuration (--use-clt, --clt-layer-idx, etc.)
+- Ensure set_use_hook_mlp_in is called for CLT models ([`0c58760`](https://github.com/jbloomAus/SAEDashboard/commit/0c587608d2da05aa5ed3656afc0ca6b001fbf79b))
+
+* script for CLT dashboard generation ([`31e7154`](https://github.com/jbloomAus/SAEDashboard/commit/31e7154363af98af2d6ec269ee9a78f9723048bb))
+
+* formatting ([`210fdc4`](https://github.com/jbloomAus/SAEDashboard/commit/210fdc4bd3d8e2dc957c2925719a7a7f7bca1de3))
+
+* simplified init function ([`883deb4`](https://github.com/jbloomAus/SAEDashboard/commit/883deb46d75df68af96c2366f59a41dd4a6db964))
+
+* formatted tests ([`676b0f1`](https://github.com/jbloomAus/SAEDashboard/commit/676b0f17db4ad919c951414e29c1eb224efa8cef))
+
+* added tests and formatting ([`70cbad3`](https://github.com/jbloomAus/SAEDashboard/commit/70cbad3074d2c0a86929921d78d20954e1caf6d8))
+
+* Fix compatibility with new SAELens API structure
+
+- Fix hook_layer extraction from hook_name when not in config
+- Remove deprecated unpacking of SAE.from_pretrained() return values
+- Handle prepend_bos in both config and metadata locations
+- Add support for extracting layer number from hook_name pattern
+
+Co-Authored-By: Claude <noreply@anthropic.com> ([`9018991`](https://github.com/jbloomAus/SAEDashboard/commit/9018991a40b62b7f4ece4f91aa8f07b9f2119d9f))
+
+* Fix indentation and hook_name access for transcoders
+
+- Fix indentation errors in neuronpedia_runner.py
+- Fix hook_name access - it's always in metadata for both SAEs and transcoders
+- Add test scripts for transcoder functionality
+- Successfully tested transcoder dashboard generation
+
+Co-Authored-By: Claude <noreply@anthropic.com> ([`31368e4`](https://github.com/jbloomAus/SAEDashboard/commit/31368e4387757063df87e3b04bd512e13a8cc7d5))
+
+* Update .gitignore to exclude test directories and submodules ([`5bf67c5`](https://github.com/jbloomAus/SAEDashboard/commit/5bf67c56e242f496b849957675dd610863485aeb))
+
+* Add transcoder support to SAEDashboard
+
+- Update imports from sae_lens to use new API structure
+- Add support for loading Transcoder and SkipTranscoder
+- Handle differences between SAE and Transcoder configs
+- Add support for normalized hooks in transformer_lens_wrapper
+- Fix architecture handling in FeatureMaskingContext
+- Update ActivationsStore.from_sae() to include dataset parameter
+
+Co-Authored-By: Claude <noreply@anthropic.com> ([`02f78a0`](https://github.com/jbloomAus/SAEDashboard/commit/02f78a0b2a6d08729e60caf06649fca4dfc38ec7))
+
+## v0.6.10 (2025-07-16)
+
+### Fix
+
+* fix: relax SAELens requirement ([`a83147e`](https://github.com/jbloomAus/SAEDashboard/commit/a83147efbf30ef4c4380f306a03468a0c8d41be0))
+
+* fix: Merge pull request #45 from Hzfinfdu/main
+
+fix: reading model_from_pretrained_kwargs from SAELens config with th… ([`0a509fe`](https://github.com/jbloomAus/SAEDashboard/commit/0a509fede04737b8087140fb4fe5f7addc259806))
+
+* fix: reading model_from_pretrained_kwargs from SAELens config with the correct key ([`9938812`](https://github.com/jbloomAus/SAEDashboard/commit/9938812ad209764ceb021eedffb08c0fc5a31c89))
+
+### Unknown
+
+* Merge pull request #60 from jbloomAus/fix-unit-tests
+
+fixes unit tests ([`1a3975d`](https://github.com/jbloomAus/SAEDashboard/commit/1a3975df60198dc169dfcf3354a8e8da5383029f))
+
+* fixes unit tests ([`2a35d5d`](https://github.com/jbloomAus/SAEDashboard/commit/2a35d5dd4ad08e3d3532158030d86dbef93ad309))
+
+* dedupes get_tokens() (#55)
+
+* dedupes get_tokens()
+
+* adds newline ([`faeb6f1`](https://github.com/jbloomAus/SAEDashboard/commit/faeb6f119d35a275a304d39c6e8cc9c7c40d31ce))
+
+* fixes make commands (#57) ([`cb74411`](https://github.com/jbloomAus/SAEDashboard/commit/cb74411039d0c9d0b0883c85434f27601cf940a5))
+
+* deletes print statements in tests (#56) ([`026ba30`](https://github.com/jbloomAus/SAEDashboard/commit/026ba305f4e31f5b47d4f9ada04c7cb0c3aae7f0))
+
+* deletes unused direct_effect_feature_ablation_experiment() (#52) ([`391ff94`](https://github.com/jbloomAus/SAEDashboard/commit/391ff949a997a99b605bd55a706d6fed2892249c))
+
+* removes unused files (#54) ([`5381cc7`](https://github.com/jbloomAus/SAEDashboard/commit/5381cc7118c7655c6c14cdbbd12e1f6c00278fc2))
+
+* Merge pull request #47 from Marlon154/main
+
+Fixing deprecated fn call for SAE Lens ([`61c9bd4`](https://github.com/jbloomAus/SAEDashboard/commit/61c9bd4ad8ccd5d96cb5c89eb961db0e7fbc2ab0))
+
+* Merge branch 'main' into main ([`50b202a`](https://github.com/jbloomAus/SAEDashboard/commit/50b202a0b2fef413dd46b4ce2838bae27c0ac252))
+
+* Merge pull request #35 from chanind/relax-saelens-dep
+
+fix: relax SAELens and einops requirements ([`6c71bbf`](https://github.com/jbloomAus/SAEDashboard/commit/6c71bbfd7b6f1562093f1192616a7a55188631d3))
+
+* fixing type checking ([`42a9845`](https://github.com/jbloomAus/SAEDashboard/commit/42a9845bba856c942f6d70182429cdb49e0ea917))
+
+* Merge branch 'main' into relax-saelens-dep ([`3e6c870`](https://github.com/jbloomAus/SAEDashboard/commit/3e6c8703afd5ce80c29ec1ed0fc729def3f7f8fa))
+
+* also relax einops ([`62614ac`](https://github.com/jbloomAus/SAEDashboard/commit/62614ac27ca50527556cc7c891e589e63a14e9bc))
+
+* fix type checks ([`5a2cca0`](https://github.com/jbloomAus/SAEDashboard/commit/5a2cca0334a0907e7685cbef798cda71cd249ba4))
+
+* Fixing deprecated fn call ([`f1da0e6`](https://github.com/jbloomAus/SAEDashboard/commit/f1da0e6ea7d663e5ff54612d7979d1b1ed9a6b77))
+
+## v0.6.9 (2025-02-25)
+
+### Fix
+
+* fix: Merge pull request #44 from jbloomAus/update_saelens
+
+fix: don't use sparsity ([`f30a19b`](https://github.com/jbloomAus/SAEDashboard/commit/f30a19b9cb42f15302848c31ddf1d14462209a42))
+
+* fix: don't use sparsity ([`d5ba79b`](https://github.com/jbloomAus/SAEDashboard/commit/d5ba79bbf3d51cbc67e276297d18d85add9d33e7))
+
+* fix: update SAELens version and remove unsupported load_sparsity ([`63192ba`](https://github.com/jbloomAus/SAEDashboard/commit/63192ba7d9de7afae9cb67f65db2e79a39b898c6))
+
+### Unknown
+
+* Merge pull request #43 from jbloomAus/update_saelens
+
+fix: update SAELens version and remove unsupported load_sparsity ([`c083723`](https://github.com/jbloomAus/SAEDashboard/commit/c083723237090165725e587f8bdb8f01338394b4))
+
+## v0.6.8 (2025-02-15)
+
+### Fix
+
+* fix: prepended chat template text should not be in activations ([`f3c20ee`](https://github.com/jbloomAus/SAEDashboard/commit/f3c20eec31976db48c1f1d37aabd077e068f66ac))
+
+### Unknown
+
+* Merge pull request #42 from jbloomAus/prepend_text_fix
+
+fix: prepended chat template text should not be in activations ([`eea0b83`](https://github.com/jbloomAus/SAEDashboard/commit/eea0b830e97e791571986cf1ccae1605606ddb4f))
+
+## v0.6.7 (2025-02-13)
+
+### Fix
+
+* fix: force build ([`9b96ac5`](https://github.com/jbloomAus/SAEDashboard/commit/9b96ac57e0b23c1a4cc73fbd9fd855ab0961cce7))
+
+### Unknown
+
+* Merge pull request #41 from jbloomAus/prepend_chat_template
+
+feat: Prepend chat template and activation threshold ([`c7347fa`](https://github.com/jbloomAus/SAEDashboard/commit/c7347faa7d1c800dae398ee8dbded53afced9aa4))
+
+* add example ([`9ab42d6`](https://github.com/jbloomAus/SAEDashboard/commit/9ab42d66a2a3e445b300bf83f677f143da3fecd9))
+
+* proper 'activation threshold' ([`a6d7c1c`](https://github.com/jbloomAus/SAEDashboard/commit/a6d7c1c8ca4a1ec46676298c4661f623bada9049))
+
+* prepend chat template text ([`c8829a1`](https://github.com/jbloomAus/SAEDashboard/commit/c8829a14d6351b5cfd52af927942af5c6897db60))
+
+## v0.6.6 (2025-02-11)
+
+### Fix
+
+* fix: run_settings.json should properly log model_id and layer ([`2e661d9`](https://github.com/jbloomAus/SAEDashboard/commit/2e661d95f30bc28e7d818bbd67de931a334d837f))
+
+### Unknown
+
+* Merge pull request #40 from jbloomAus/run_settings_fix
+
+fix: run_settings.json should properly log model_id and layer ([`f3bde39`](https://github.com/jbloomAus/SAEDashboard/commit/f3bde395843720674d4c60e21bc2453d958ff402))
+
+## v0.6.5 (2025-02-11)
+
+### Fix
+
+* fix: Force Build ([`2e4979c`](https://github.com/jbloomAus/SAEDashboard/commit/2e4979c07ad8bcd2760ee0981ee415d17fef2e5a))
+
+### Unknown
+
+* Merge pull request #39 from jbloomAus/allow_vector_output
+
+feat: allow outputting raw vector in neuronpedia outputs ([`1444786`](https://github.com/jbloomAus/SAEDashboard/commit/14447862418b18d112053a7af8810b049400089a))
+
+* remove debug log ([`6efeb6c`](https://github.com/jbloomAus/SAEDashboard/commit/6efeb6c1976dea8374800e7625d63a14a3b6438d))
+
+* allow outputting vector ([`4c6cb35`](https://github.com/jbloomAus/SAEDashboard/commit/4c6cb35752317db9f22476d26dc1bab7e4d6e511))
+
+* Merge pull request #37 from jbloomAus/feature/vector-dashboards
+
+Feature/vector dashboards ([`64c44a9`](https://github.com/jbloomAus/SAEDashboard/commit/64c44a9c11b2dce26b030d5e7bbf782ef90a2985))
+
+* typing ([`09aeeab`](https://github.com/jbloomAus/SAEDashboard/commit/09aeeabb4c0f45f4bdabb884f683208eb7073142))
+
+* Fixed missing parameter ([`a91d9f5`](https://github.com/jbloomAus/SAEDashboard/commit/a91d9f5dc8c65249c032dc4088aead4364bc42e9))
+
+* Fixed parameterization and formatting ([`1956fbc`](https://github.com/jbloomAus/SAEDashboard/commit/1956fbc0ab6e4200d771ad4b504946ff81707969))
+
+* Renamed demo notebook, some cleanup ([`6a486a5`](https://github.com/jbloomAus/SAEDashboard/commit/6a486a5834377823f380f59cafcbc3debbdcc3ed))
+
+* Working pipeline flow ([`fdb2292`](https://github.com/jbloomAus/SAEDashboard/commit/fdb2292ad84b083246fdf3be2820e0b31168dce2))
+
+* First draft of vector vis pipeline ([`4351ef9`](https://github.com/jbloomAus/SAEDashboard/commit/4351ef938a82d2b5e5a37391c236791cc23b41e5))
+
+* Merge pull request #38 from jbloomAus/feature/hf-model-override
+
+enable passing custom HF model to replace model weights ([`5d98417`](https://github.com/jbloomAus/SAEDashboard/commit/5d98417877c2cfe52bb09ddada0b4b53849b344a))
+
+* enable passing custom HF model to replace model weights ([`b2d6ae5`](https://github.com/jbloomAus/SAEDashboard/commit/b2d6ae5446fb79f4662bcdac6030cb6072b09b60))
+
+* Don't copy to output folder by default ([`4dbde12`](https://github.com/jbloomAus/SAEDashboard/commit/4dbde1214d49eaaf9b591f083f34e57c8c0c1dbd))
+
+* Don't save html file for NP outputs ([`a160bff`](https://github.com/jbloomAus/SAEDashboard/commit/a160bff204b7464d2de00e3f80c255123d11171b))
+
+## v0.6.4 (2024-10-24)
+
+### Fix
+
+* fix: Merge pull request #33 from jbloomAus/fix/topk-selection-purview
+
+Fix/topk selection purview ([`afccd5a`](https://github.com/jbloomAus/SAEDashboard/commit/afccd5aaa00d00672eb1270b258b69f0e51c046a))
+
+### Unknown
+
+* updated formatting/typing ([`fb141ae`](https://github.com/jbloomAus/SAEDashboard/commit/fb141ae991261408d296286bf6777b2ec5f1f319))
+
+* TopK will now select from all latents regardless of feature batch size ([`c1f0e14`](https://github.com/jbloomAus/SAEDashboard/commit/c1f0e14dda7aa3364bfd78ca2b8c04c95b2d14b3))
+
+* Update README.md ([`8235a9e`](https://github.com/jbloomAus/SAEDashboard/commit/8235a9e3adaea50b6b9f26f575e25a254d67a135))
+
+* Merge pull request #32 from jbloomAus/docs/readme-update
+
+docs: updated readme ([`b5e5480`](https://github.com/jbloomAus/SAEDashboard/commit/b5e54808ee05fc75e68d74ec319bf49826b45508))
+
+* Update README.md ([`a1546fd`](https://github.com/jbloomAus/SAEDashboard/commit/a1546fdef32745cdc862a5a2dd0478e57e45320d))
+
+* Removed outdated vis type ([`b0676af`](https://github.com/jbloomAus/SAEDashboard/commit/b0676afcca0845b73a54d983eaa9d72b0e9dff05))
+
+* Update README.md ([`9b8446a`](https://github.com/jbloomAus/SAEDashboard/commit/9b8446aa47f287ba80bf0ac4a39f7c77f0492990))
+
+* Updated format ([`90e4a09`](https://github.com/jbloomAus/SAEDashboard/commit/90e4a09eedd7f428b64e58d5ca2fd1cfa658b0da))
+
+* Updated readme ([`f6819a6`](https://github.com/jbloomAus/SAEDashboard/commit/f6819a6da594673cad65c9ccd3a4f67746de796d))
+
+## v0.6.3 (2024-10-23)
+
+### Fix
+
+* fix: update cached_activations directory to include number of prompts ([`0308cb1`](https://github.com/jbloomAus/SAEDashboard/commit/0308cb146bf2eb9cee26f03d3098511d03022485))
+
+## v0.6.2 (2024-10-23)
+
+### Fix
+
+* fix: lint ([`3fc0e2c`](https://github.com/jbloomAus/SAEDashboard/commit/3fc0e2ccb39ed1d3e31d66ae0aba2b2b367d46aa))
+
+### Unknown
+
+* Merge branch 'main' of https://github.com/jbloomAus/SAEDashboard ([`8f74a96`](https://github.com/jbloomAus/SAEDashboard/commit/8f74a969f48a7e0fd8de17cc983acf3886db95ef))
+
+## v0.6.1 (2024-10-22)
+
+### Unknown
+
+* Fix: divide by zero, cached_activations folder name ([`1792298`](https://github.com/jbloomAus/SAEDashboard/commit/179229805ae6489d86e235240c65d26db64b5cd7))
+
+* Merge branch 'main' of https://github.com/jbloomAus/SAEDashboard ([`508a74d`](https://github.com/jbloomAus/SAEDashboard/commit/508a74df8ff279716501e4179c501b5089a8d706))
+
+## v0.6.0 (2024-10-21)
+
+### Feature
+
+* feat: np sae id suffix ([`448b14e`](https://github.com/jbloomAus/SAEDashboard/commit/448b14e0b3aea8ff854a5365f164b6ce5f419f0d))
+
+### Fix
+
+* fix: update saelens to v4 ([`ef1a330`](https://github.com/jbloomAus/SAEDashboard/commit/ef1a3302d0483eddb247defab5c88816850f7f63))
+
+### Unknown
+
+* Merge pull request #31 from jbloomAus/fix/reduce-mem
+
+fix: added mem cleanup ([`60bd716`](https://github.com/jbloomAus/SAEDashboard/commit/60bd716c7b52bb0eaea0937e097eb77ed78bd33d))
+
+* Fixed formatting ([`f1fab0c`](https://github.com/jbloomAus/SAEDashboard/commit/f1fab0c1fd5be281e2162ab3f54ffc7f4c09a1ce))
+
+* Added cleanup ([`305c46d`](https://github.com/jbloomAus/SAEDashboard/commit/305c46d7a30330bbae6893b83cb6d498c2c975f1))
+
+* Merge pull request #30 from jbloomAus/feat-mask-via-position
+
+feat: prepending/appending tokens for prompt template + feat mask via Position ([`4c60e4c`](https://github.com/jbloomAus/SAEDashboard/commit/4c60e4c834dfb5759ce55dc90d1f88768abfea0d))
+
+* add a few tests ([`96247d5`](https://github.com/jbloomAus/SAEDashboard/commit/96247d5afaf141b8b1279c17fd135240b0d8e869))
+
+* handle prefixes / suffixes and ignored positions ([`bff7fd9`](https://github.com/jbloomAus/SAEDashboard/commit/bff7fd98b09318a1b01d2bc4a06467f8afa156f9))
+
+* simplify masking ([`385b6e1`](https://github.com/jbloomAus/SAEDashboard/commit/385b6e116ecac53ad4df8585f7513c3416707d8b))
+
+* add option for ignoring tokens at particular positions ([`ed3426d`](https://github.com/jbloomAus/SAEDashboard/commit/ed3426de5cb1495c138f770eefa5f941408aa390))
+
+* Merge pull request #29 from jbloomAus/refactor/optimize-dfa-speed
+
+Sped up DFA calculation 60x ([`f992e3c`](https://github.com/jbloomAus/SAEDashboard/commit/f992e3cf116189625b3a92529cf68d6226a1221c))
+
+* Sped up DFA calculation ([`be11cd5`](https://github.com/jbloomAus/SAEDashboard/commit/be11cd5652f0f8a8ae425555666b747b9b99314e))
+
+* Added test to check for decoder weight dist (head dist) ([`f147696`](https://github.com/jbloomAus/SAEDashboard/commit/f1476967af5fee95313264ccaee668605d23b9ad))
+
+* Merge pull request #28 from jbloomAus/feature/np-topk-size-arg
+
+Feature/np topk size arg ([`c5c1365`](https://github.com/jbloomAus/SAEDashboard/commit/c5c136576609991177d3a8924b5bf75a42b66399))
+
+* Simply updated default value for top K ([`5c855fe`](https://github.com/jbloomAus/SAEDashboard/commit/5c855fec0e58a114a537590d1400eaa42dd3610c))
+
+* Testing variable topk sizes ([`79fe14b`](https://github.com/jbloomAus/SAEDashboard/commit/79fe14b840991bd1f8ada8462aeb65d72821c4aa))
+
+* Merge pull request #25 from jbloomAus/fix/dfa-for-gqa
+
+Fix/dfa for gqa ([`85c345f`](https://github.com/jbloomAus/SAEDashboard/commit/85c345f3ad8069a59be8d495242395c50381ab01))
+
+* Fixed formatting ([`48a67c7`](https://github.com/jbloomAus/SAEDashboard/commit/48a67c79247d05745d355e6a4bf380e9df20474e))
+
+* Removed redundant code from rebase ([`a71fb9d`](https://github.com/jbloomAus/SAEDashboard/commit/a71fb9dde6e880b0f4297277d27696c9d524d052))
+
+* fixed rebase ([`57ee280`](https://github.com/jbloomAus/SAEDashboard/commit/57ee28021efd3678bcd9d12d55e048c14a2f2d47))
+
+* Added tests for DFA for GQA ([`3b99e36`](https://github.com/jbloomAus/SAEDashboard/commit/3b99e36c74d2c61617cfed107bee3b0eb3b63294))
+
+* Removed duplicate code ([`7093773`](https://github.com/jbloomAus/SAEDashboard/commit/7093773d079cd235aea99273a1365363a5bf8b6d))
+
+* More rebasing stuff ([`59c6cd8`](https://github.com/jbloomAus/SAEDashboard/commit/59c6cd85ead287b2774aa591463d131840c7f270))
+
+* Fixed formatting ([`ed7d3b1`](https://github.com/jbloomAus/SAEDashboard/commit/ed7d3b16a99e3e3a272e73356cc0509b2c59a292))
+
+* Removed debugging statements ([`6489d1c`](https://github.com/jbloomAus/SAEDashboard/commit/6489d1c5b52ed86cb280c237c08e10238e0d0564))
+
+* more debug prints x3 ([`5ba2b8a`](https://github.com/jbloomAus/SAEDashboard/commit/5ba2b8a69f1881b901131976c7d52f142068dbd2))
+
+* more debug prints x2 ([`e124ff9`](https://github.com/jbloomAus/SAEDashboard/commit/e124ff906ec7b37083af4e4721b9e33902146e47))
+
+* more debug prints ([`e2b0c35`](https://github.com/jbloomAus/SAEDashboard/commit/e2b0c35467e5d405abd3cca664dfd1960dbba0eb))
+
+* temp print statements ([`95df55b`](https://github.com/jbloomAus/SAEDashboard/commit/95df55b29f9250f67c5b986216e587c37f72aa9e))
+
+* Lowered default threshold ([`dc1f31a`](https://github.com/jbloomAus/SAEDashboard/commit/dc1f31a55400231e46feb58a8c100f66472baa1b))
+
+* updated ignore ([`eb0d56a`](https://github.com/jbloomAus/SAEDashboard/commit/eb0d56a9f813b9cf82742093fae00bb0ccfdac45))
+
+* Reduced memory load of GQA DFA ([`05867f1`](https://github.com/jbloomAus/SAEDashboard/commit/05867f1d0c8b5f2a5b76f3ea45ab9c87eaae9c09))
+
+* DFA will now work for models with grouped query attention ([`91a5dd1`](https://github.com/jbloomAus/SAEDashboard/commit/91a5dd17a2e567efa7d8a89d228eb7de47ae6766))
+
+* Added head attr weights functionality for when DFA is use ([`03a615f`](https://github.com/jbloomAus/SAEDashboard/commit/03a615f7c70a6f6e634845dab4051874698fac5b))
+
+* Edited default chunk size ([`7d68f9e`](https://github.com/jbloomAus/SAEDashboard/commit/7d68f9e7131b8c5558e886022625dac267f20aab))
+
+* Fixed formatting ([`4d5f38b`](https://github.com/jbloomAus/SAEDashboard/commit/4d5f38beca15f2ce05c89f83eb3e955c291f9687))
+
+* Removed debugging statements and added device changes ([`76e17c9`](https://github.com/jbloomAus/SAEDashboard/commit/76e17c91a41b5df6047baa5bcfa33d253b029d29))
+
+* more debug prints x3 ([`06535d3`](https://github.com/jbloomAus/SAEDashboard/commit/06535d3df168d92ac79d2f5a14b345c757dfd9de))
+
+* more debug prints x2 ([`26e8297`](https://github.com/jbloomAus/SAEDashboard/commit/26e8297888de066f0097e3b73245eb149bfb327f))
+
+* more debug prints ([`9ded356`](https://github.com/jbloomAus/SAEDashboard/commit/9ded356ea8c3c5dd841bf5a45ea65ae8c67935f5))
+
+* temp print statements ([`024ad57`](https://github.com/jbloomAus/SAEDashboard/commit/024ad578b65b8f3592b42b66dc6a56aeae2a3116))
+
+* Lowered default threshold ([`a3b5977`](https://github.com/jbloomAus/SAEDashboard/commit/a3b5977c0f1bb7a865f7349304a5dd8092f7c2e8))
+
+* updated ignore ([`d5d325a`](https://github.com/jbloomAus/SAEDashboard/commit/d5d325a63b3b26b890c2bab512f2a8473bdc926a))
+
+* Reduced memory load of GQA DFA ([`93eb1a9`](https://github.com/jbloomAus/SAEDashboard/commit/93eb1a9a92320d9f4645b500e22a566135918e3d))
+
+* DFA will now work for models with grouped query attention ([`6594155`](https://github.com/jbloomAus/SAEDashboard/commit/65941559bac03a3e4fb128d5327033e01f19c18d))
+
+* Added head attr weights functionality for when DFA is use ([`9312d90`](https://github.com/jbloomAus/SAEDashboard/commit/9312d901bf17e14400199c86e0284be6c750162a))
+
+* Added tests for DFA for GQA ([`fcfac37`](https://github.com/jbloomAus/SAEDashboard/commit/fcfac37e148461e585f38fddf868ad2a32d908a8))
+
+* Removed duplicate code ([`cc00944`](https://github.com/jbloomAus/SAEDashboard/commit/cc00944855720d5b8139d4267b44c1a230ef5319))
+
+* Fixed formatting ([`50b08b4`](https://github.com/jbloomAus/SAEDashboard/commit/50b08b4eb50734afe0f085274ccaee71ec4017a4))
+
+* Removed debugging statements ([`f7b949b`](https://github.com/jbloomAus/SAEDashboard/commit/f7b949b4af6bc8ca7557bfa5fa2441fbaa0284a0))
+
+* more debug prints x3 ([`53536b0`](https://github.com/jbloomAus/SAEDashboard/commit/53536b03d624783b6b2f95b07b9318139ef0c49e))
+
+* more debug prints x2 ([`6f2c504`](https://github.com/jbloomAus/SAEDashboard/commit/6f2c504a355f9071e766fc7fa3b6aad9890572a8))
+
+* more debug prints ([`e1bef90`](https://github.com/jbloomAus/SAEDashboard/commit/e1bef90d16e8c73c9532b19a08c842757828c7ed))
+
+* temp print statements ([`fd75714`](https://github.com/jbloomAus/SAEDashboard/commit/fd75714ee4631463c1f754d68f83b9ef75eb2285))
+
+* updated ignore ([`c01062f`](https://github.com/jbloomAus/SAEDashboard/commit/c01062faecfaa132d87c56a7ba7add573c6b0f4e))
+
+* Reduced memory load of GQA DFA ([`1ae40e9`](https://github.com/jbloomAus/SAEDashboard/commit/1ae40e9d487af7e8a7b148629588ef87fdd0a6e5))
+
+* DFA will now work for models with grouped query attention ([`c66c90f`](https://github.com/jbloomAus/SAEDashboard/commit/c66c90f5d51961cafd5f13c26a94193ee38f828a))
+
+* Edited default chunk size ([`3c78bdc`](https://github.com/jbloomAus/SAEDashboard/commit/3c78bdcfda12e5873de082a7f1e631a801bd9407))
+
+* Fixed formatting ([`10a36e3`](https://github.com/jbloomAus/SAEDashboard/commit/10a36e3e8da3c7593058d3638ac3b7a32953b1b0))
+
+* Removed debugging statements and added device changes ([`0f51dd9`](https://github.com/jbloomAus/SAEDashboard/commit/0f51dd953cd214244c71e8b9156b90483ceaa2be))
+
+* more debug prints x3 ([`112ef42`](https://github.com/jbloomAus/SAEDashboard/commit/112ef4292b81a64f6168e7527ec583faa9ba20a4))
+
+* more debug prints x2 ([`ef154d6`](https://github.com/jbloomAus/SAEDashboard/commit/ef154d6044bb67d17a2aa225ddf4099ccfc16b55))
+
+* more debug prints ([`1b18d14`](https://github.com/jbloomAus/SAEDashboard/commit/1b18d141dd33e3a99c2abd5a6d195ab5142890d8))
+
+* temp print statements ([`2194d2c`](https://github.com/jbloomAus/SAEDashboard/commit/2194d2cea16856c96ace47ad5ac560f088e769b0))
+
+* Lowered default threshold ([`a49d1e5`](https://github.com/jbloomAus/SAEDashboard/commit/a49d1e5b94c8ef680448f20ded849c7752fb5131))
+
+* updated ignore ([`2067655`](https://github.com/jbloomAus/SAEDashboard/commit/20676554541d29fddd87215a47e8e94891e342ac))
+
+* Reduced memory load of GQA DFA ([`8ec1956`](https://github.com/jbloomAus/SAEDashboard/commit/8ec19566e8898413d349fe3f2e43fbff232ffa62))
+
+* DFA will now work for models with grouped query attention ([`8f3cf55`](https://github.com/jbloomAus/SAEDashboard/commit/8f3cf5532e57abc6e694fb11c5f9c7c2915215c0))
+
+* Added head attr weights functionality for when DFA is use ([`234ea32`](https://github.com/jbloomAus/SAEDashboard/commit/234ea3211ce7dbf84d101c4e8bfe844c3903b16a))
+
+* Merge pull request #27 from jbloomAus/fix/resolve-duplication
+
+Removed sources of duplicate sequences ([`525bffe`](https://github.com/jbloomAus/SAEDashboard/commit/525bffee516a630c4b4f033d3971fad8c6dd5a74))
+
+* Updated location of wandb finish() ([`921da77`](https://github.com/jbloomAus/SAEDashboard/commit/921da77132a560505fa61decf287ca3833f96ec7))
+
+* Added two sets of tests for duplication checks ([`3e95ffd`](https://github.com/jbloomAus/SAEDashboard/commit/3e95ffd1dafd01deb1f7817845ccb6229fb4ae09))
+
+* Restored original random indices function as it seemed ok ([`388719b`](https://github.com/jbloomAus/SAEDashboard/commit/388719bec99b4306e81e0cdb772b9924db210774))
+
+* Removed sources of duplicate sequences ([`853306c`](https://github.com/jbloomAus/SAEDashboard/commit/853306c4e08d9ec95674fdc5c87f807019055d0d))
+
+## v0.5.1 (2024-08-27)
+
+### Fix
+
+* fix: multi-gpu-tlens
+
+fix: handle multiple tlens devices ([`ed1e967`](https://github.com/jbloomAus/SAEDashboard/commit/ed1e967d44b887f4b99d2257934ca920d5c6a508))
+
+* fix: handle multiple tlens devices ([`ba5368f`](https://github.com/jbloomAus/SAEDashboard/commit/ba5368f9999f08332c153816ba5836f8a1eb9ba1))
+
+## v0.5.0 (2024-08-25)
+
+### Feature
+
+* feat: accelerate caching. Torch load / save faster when files are small. 
+
+Refactor/accelerate caching ([`6027d0a`](https://github.com/jbloomAus/SAEDashboard/commit/6027d0a3fc0d70908bad036a9658caa406d9f809))
+
+### Unknown
+
+* Updated formatting ([`c1ea288`](https://github.com/jbloomAus/SAEDashboard/commit/c1ea2882a17e0d1b7b28743a34fca9d0754bd8a7))
+
+* Sped up caching with native torch functions ([`230840a`](https://github.com/jbloomAus/SAEDashboard/commit/230840aea50b8b7055a6aa61961d7ac50855b763))
+
+* Increased cache loading speed ([`83fe5f4`](https://github.com/jbloomAus/SAEDashboard/commit/83fe5f4bdf1252d533f203bc3f53ea9f71880ab8))
+
+## v0.4.0 (2024-08-22)
+
+### Feature
+
+* feat: Refactor json writer and trigger DFA release
+
+JSON writer has been refactored for reusability and readability ([`664f487`](https://github.com/jbloomAus/SAEDashboard/commit/664f4874b585c5510d2d3dd639c5e893023f6332))
+
+### Unknown
+
+* Refactored JSON creation from the neuronpedia runner ([`d6bb24b`](https://github.com/jbloomAus/SAEDashboard/commit/d6bb24b6d773874d8e99be4d84402d559741907b))
+
+* Merge pull request #20 from jbloomAus/feature/dfa
+
+SAEVisRunner DFA Implementation ([`926ea87`](https://github.com/jbloomAus/SAEDashboard/commit/926ea87dd344548489201f68cc92b33662430813))
+
+* Update ci.yaml ([`4b2807d`](https://github.com/jbloomAus/SAEDashboard/commit/4b2807dd865904120d236b355c0ccb1680c2919e))
+
+* Fixed formatting ([`a62cc8f`](https://github.com/jbloomAus/SAEDashboard/commit/a62cc8f1bdd4c6e49b76d2d594e5a6b4b8183a8c))
+
+* Fixed target index ([`ca2668d`](https://github.com/jbloomAus/SAEDashboard/commit/ca2668da03ea4d06cdc9f198988b80e0db844316))
+
+* Corrected DFA indexing ([`d5028ae`](https://github.com/jbloomAus/SAEDashboard/commit/d5028aec875db4c03196726400c3b90b5d9d4d01))
+
+* Adding temporary testing notebook ([`98e4b2f`](https://github.com/jbloomAus/SAEDashboard/commit/98e4b2f93d300ad4e94985d8d2594739a277e0c8))
+
+* Added DFA output to neuronpedia runner ([`68eeff3`](https://github.com/jbloomAus/SAEDashboard/commit/68eeff3172b0c8637a6566c07951c28fd14a1c03))
+
+* Fixed test typehints ([`d358e6f`](https://github.com/jbloomAus/SAEDashboard/commit/d358e6f5cc37304935eed949a0b0b985ba12b94f))
+
+* Fixed formatting ([`5cb19e2`](https://github.com/jbloomAus/SAEDashboard/commit/5cb19e241051503730b6982813a6730556990c92))
+
+* Corrected typehints ([`6173fbd`](https://github.com/jbloomAus/SAEDashboard/commit/6173fbd3824b7cba58e1cf0c7ee239762ee533ce))
+
+* Removed another unused import ([`8be1572`](https://github.com/jbloomAus/SAEDashboard/commit/8be1572370b1adf341e2a650953bf17cd179808d))
+
+* Removed unused imports ([`9071210`](https://github.com/jbloomAus/SAEDashboard/commit/90712105f74b287d77a06c045e8c32fd05f2e668))
+
+* Added support for DFA calculations up to SAE Vis runner ([`4a08ffd`](https://github.com/jbloomAus/SAEDashboard/commit/4a08ffd13a8f29ff16808a20cd663c9d2d369e6a))
+
+* Added activation collection flow for DFA ([`0ebb1f3`](https://github.com/jbloomAus/SAEDashboard/commit/0ebb1f3ca61603662f4f2cc8b1341470bf75b5d1))
+
+* Merge pull request #19 from jbloomAus/fix/remove_precision_reduction
+
+Removed precision reduction option ([`a5f8df1`](https://github.com/jbloomAus/SAEDashboard/commit/a5f8df15ef8619c4d08655e777d379a05b453346))
+
+* Removed float16 option entirely from quantile calc ([`1b6a4a9`](https://github.com/jbloomAus/SAEDashboard/commit/1b6a4a93403ca2e9a869aa73600f37960090f03d))
+
+* Removed precision reduction option ([`cd03ffb`](https://github.com/jbloomAus/SAEDashboard/commit/cd03ffb182e93a42480c01408b47ebae94d4c349))
+
+## v0.3.0 (2024-08-15)
+
+### Feature
+
+* feat: seperate files per dashboard html ([`cd8d050`](https://github.com/jbloomAus/SAEDashboard/commit/cd8d050218ae3c6eeb7a9779072e60b78bfe0b58))
+
+### Unknown
+
+* Merge pull request #17 from jbloomAus/refactor/remove_enc_b
+
+Removed all encoder B code ([`67c9c3f`](https://github.com/jbloomAus/SAEDashboard/commit/67c9c3fdc8bd220938f65c1f97214034cc7528b4))
+
+* Removed all encoder B code ([`5174e2e`](https://github.com/jbloomAus/SAEDashboard/commit/5174e2e161030dc756c148f1740e50c52baf6a91))
+
+* Merge pull request #18 from jbloomAus/feat-seperate-files-per-html-dashboard
+
+feat: seperate files per dashboard html ([`8ff69ba`](https://github.com/jbloomAus/SAEDashboard/commit/8ff69ba207692d4acb8d5fc19d038090067690df))
+
+* Merge pull request #16 from jbloomAus/performance_refactor
+
+Create() will now reduce precision by default ([`fb07b90`](https://github.com/jbloomAus/SAEDashboard/commit/fb07b90eaac395a58f02ba927460dcc2c9e61d1a))
+
+* Removed line ([`d795490`](https://github.com/jbloomAus/SAEDashboard/commit/d795490c1c9d8193c8cf84d0352b9d93c41947fe))
+
+* Removed unnecessary print ([`4544f86`](https://github.com/jbloomAus/SAEDashboard/commit/4544f86472480f0df00344fa84111a7c2a52fcef))
+
+* Precision will now be reduced by default for quantile calc ([`539d222`](https://github.com/jbloomAus/SAEDashboard/commit/539d222ded9e3a0944f5240f3a4cd84497d11a74))
+
+* Merge pull request #15 from jbloomAus/quantile_efficiency
+
+Quantile OOM prevention ([`4a40c37`](https://github.com/jbloomAus/SAEDashboard/commit/4a40c3704aab9363163fef3e2830d42f2fecdc6b))
+
+* Made quantile batch optional and removed sampling code ([`2df51d3`](https://github.com/jbloomAus/SAEDashboard/commit/2df51d353f818a196916a15f2bc56f70480dd853))
+
+* Added device check for test ([`afbb960`](https://github.com/jbloomAus/SAEDashboard/commit/afbb960d3c9376ad512607146826b7d1c1e68d48))
+
+* Added parameter for quantile calculation batching ([`49d0a7a`](https://github.com/jbloomAus/SAEDashboard/commit/49d0a7ab37896a085f80409900e3d0b261b8c9e0))
+
+* Added type annotation ([`c71c4aa`](https://github.com/jbloomAus/SAEDashboard/commit/c71c4aa1c8bc25d85b9a955b482823cbde445a51))
+
+* Removed unused imports ([`ec01bfe`](https://github.com/jbloomAus/SAEDashboard/commit/ec01bfefc2f0f4d880cd5744ff6a2ea71991349b))
+
+* Added float16 version of quantile calculation ([`2f01eb8`](https://github.com/jbloomAus/SAEDashboard/commit/2f01eb8d9f84a20918f19e81c23df86ddc9d7f0c))
+
+* Merge pull request #13 from jbloomAus/hook_z_support
+
+fix: restore hook_z support following regression. ([`ea87559`](https://github.com/jbloomAus/SAEDashboard/commit/ea87559359f9821e352dcab582e23b42fef1cebf))
+
+* format ([`21e3617`](https://github.com/jbloomAus/SAEDashboard/commit/21e3617196ef57944c141563e9263101baf9c7f1))
+
+* make sure hook_z works ([`efaeec0`](https://github.com/jbloomAus/SAEDashboard/commit/efaeec0fdf8c2c43bb13bfd652b812a38ebc0200))
+
+* Merge pull request #12 from jbloomAus/use_sae_lens_loading
+
+Use sae lens loading ([`89bba3e`](https://github.com/jbloomAus/SAEDashboard/commit/89bba3e7a10877782608c50f4b8dd9054f204381))
+
+* add settings.json ([`d8f3034`](https://github.com/jbloomAus/SAEDashboard/commit/d8f3034c0ed7241c35e9761d60a9ee4072403fd0))
+
+* add dtype ([`0d8008a`](https://github.com/jbloomAus/SAEDashboard/commit/0d8008afe93a2a2a5bfc954571c680a529ab883f))
+
+* cli util ([`9da440e`](https://github.com/jbloomAus/SAEDashboard/commit/9da440eb3d50d48a7fdc4d3ee3d26de13a458593))
+
+* wandb logging improvement ([`a077369`](https://github.com/jbloomAus/SAEDashboard/commit/a077369ca43009f4e50c0b1e7176cae398703856))
+
+* add override for np set name ([`8906d10`](https://github.com/jbloomAus/SAEDashboard/commit/8906d103ab8d10bd01b791331dfc5485ac047a4f))
+
+* auto add folder path to output dir ([`35e06ab`](https://github.com/jbloomAus/SAEDashboard/commit/35e06ab89bce257fc15ffaa4918b9598577d6df0))
+
+* update tests ([`50163b0`](https://github.com/jbloomAus/SAEDashboard/commit/50163b04ca29b492b9fb71244aa26798655b663f))
+
+* first step towards sae_lens remote loading ([`415a2d1`](https://github.com/jbloomAus/SAEDashboard/commit/415a2d1e484e9ea2351bf98de221f6a83a805107))
+
+## v0.2.3 (2024-08-06)
+
+### Fix
+
+* fix: neuronpedia uses api_key for uploading features, and update sae_id->sae_set ([`0336a35`](https://github.com/jbloomAus/SAEDashboard/commit/0336a3587f825f0be15af79cc9a0033dda3d4a3f))
+
+### Unknown
+
+* Merge pull request #11 from jbloomAus/ignore_bos_option
+
+Ignore bos option ([`ae34b70`](https://github.com/jbloomAus/SAEDashboard/commit/ae34b70b61993b4cce49a758bf85514410c67bd8))
+
+* change threshold ([`4a0be67`](https://github.com/jbloomAus/SAEDashboard/commit/4a0be67622826f879191ced225c8c075d34bfe56))
+
+* type fix ([`525b6a1`](https://github.com/jbloomAus/SAEDashboard/commit/525b6a10331b9fa0a464ae0c7f01af90ae97d0bb))
+
+* default ignore bos eos pad ([`d2396a7`](https://github.com/jbloomAus/SAEDashboard/commit/d2396a714dd9ea3d59e516aa0fe30a9c9225e22f))
+
+* ignore bos tokens ([`96cf6e9`](https://github.com/jbloomAus/SAEDashboard/commit/96cf6e9427cadf13fa13b55b7d1bc83ae81d9ec0))
+
+* jump relu support in feature masking context ([`a1ba87a`](https://github.com/jbloomAus/SAEDashboard/commit/a1ba87a5c5e03687d7d7b5c5677bd9773fa49517))
+
+* depend on latest sae lens ([`4988207`](https://github.com/jbloomAus/SAEDashboard/commit/4988207abaca24256f52235e474fe5fbb5028c1a))
+
+* Merge pull request #10 from jbloomAus/auth_and_sae_set
+
+fix: neuronpedia uses api_key for uploading features, and update sae_id -> sae_set ([`4684aca`](https://github.com/jbloomAus/SAEDashboard/commit/4684aca54b69dbc913c1122f1a322ed4d808dce0))
+
+* Combine upload-features and upload-dead-stubs ([`faac839`](https://github.com/jbloomAus/SAEDashboard/commit/faac8398fee8582b12c2d1a29df6d4de7e542bed))
+
+* Activation store device should be cuda when available ([`93050b1`](https://github.com/jbloomAus/SAEDashboard/commit/93050b1f5c2b87c8e889fe3449d440016c996762))
+
+* Activation store device should be cuda when available ([`4469066`](https://github.com/jbloomAus/SAEDashboard/commit/4469066af06bb4944832f2e596e36afa09adf160))
+
+* Better support for huggingface dataset path ([`3dc4b78`](https://github.com/jbloomAus/SAEDashboard/commit/3dc4b783a1ced7b938ab45c4d10effedd148a829))
+
+* Docker tweak ([`a1a70cb`](https://github.com/jbloomAus/SAEDashboard/commit/a1a70cb28c726887de9439024b7b1d01082d3932))
+
+## v0.2.2 (2024-07-12)
+
+### Fix
+
+* fix: don't sample too many tokens + other fixes
+
+fix: don't sample too many tokens ([`b2554b0`](https://github.com/jbloomAus/SAEDashboard/commit/b2554b017e75d14b38b343fc6e0c1bcc32be2359))
+
+* fix: don't sample too many tokens ([`0cbb2ed`](https://github.com/jbloomAus/SAEDashboard/commit/0cbb2edb480b83823dc1a98dd7e5978ecdda0d81))
+
+### Unknown
+
+* - Don't force manual overrides for dtype - default to SAE's dtype
+- Add n_prompts_in_forward_pass to neuronpedia.py
+- Add n_prompts_total, n_tokens_in_prompt, and dataset to neuronpedia artifact
+- Remove NPDashboardSettings for now (just save the NPRunnerConfig later)
+- Fix lint error
+- Consolidate minibatch_size_features/tokens to n_feats_at_a_time and n_prompts_in_fwd_pass
+- Update/Fix NP acceptance test ([`b6282c8`](https://github.com/jbloomAus/SAEDashboard/commit/b6282c83e1898e356e271af0926e2271fb23f707))
+
+* Merge pull request #7 from jbloomAus/performance-improvement
+
+feat: performance improvement ([`f98b3dc`](https://github.com/jbloomAus/SAEDashboard/commit/f98b3dcf84c42687dfc92fa38377edd1c3f6fa30))
+
+* delete unused snapshots ([`4210b48`](https://github.com/jbloomAus/SAEDashboard/commit/4210b48608792adc9b841ea92a64050311e66cd6))
+
+* format ([`de57a2d`](https://github.com/jbloomAus/SAEDashboard/commit/de57a2d84564fc0eb7d5e42799c00f73c7007cf8))
+
+* linter ([`4725ffa`](https://github.com/jbloomAus/SAEDashboard/commit/4725ffa2cbe743aa0bb615213f11105b6911f10d))
+
+* hope flaky tests start passing ([`8ac9e8e`](https://github.com/jbloomAus/SAEDashboard/commit/8ac9e8e93127d4ab811019fc62bbe050a9a00e2c))
+
+* np.memmap caching and more explicit hyperparams ([`9a24186`](https://github.com/jbloomAus/SAEDashboard/commit/9a24186cc1c118725c6db7dc3c77feb815cf938f))
+
+* Move docker" ([`27b1a27`](https://github.com/jbloomAus/SAEDashboard/commit/27b1a27118bcccf54576eb1891b936bd92848f3f))
+
+* Add docker to workflow ([`a354fa4`](https://github.com/jbloomAus/SAEDashboard/commit/a354fa47cfb005dd2304b4237f9182e2408daeed))
+
+* Dockerignore file ([`ed9fcf3`](https://github.com/jbloomAus/SAEDashboard/commit/ed9fcf3a634cd57f6517170784d56d86431e1710))
+
+* new versions ([`f64e54d`](https://github.com/jbloomAus/SAEDashboard/commit/f64e54df5c1b643fc3acaff7f4d40d5597edf61a))
+
+* Add tools to docker image ([`2a70f64`](https://github.com/jbloomAus/SAEDashboard/commit/2a70f64cfd4177d807a8345e64699054dd103e8d))
+
+* Fix docker ([`3805f20`](https://github.com/jbloomAus/SAEDashboard/commit/3805f20bff622582d16fd6603bef4b77e6bada9e))
+
+* Fix docker image ([`7f9ff2f`](https://github.com/jbloomAus/SAEDashboard/commit/7f9ff2f9b10ce08264b2153e8191eca32f9ee48a))
+
+* Fix NP simple test, remove check for correlated neurons/features ([`355fad5`](https://github.com/jbloomAus/SAEDashboard/commit/355fad58ab2ab036a33375c02d9006db634702b9))
+
+* Dockerfile, small batching fix ([`4df4c51`](https://github.com/jbloomAus/SAEDashboard/commit/4df4c5138341a1c233c3d0fe1a3d399846e92407))
+
+* set sae_device, activation_store device ([`6d65b22`](https://github.com/jbloomAus/SAEDashboard/commit/6d65b22ef541326cc9558119b40baeb95cc2e47e))
+
+* Fix NP dtype error ([`8bb4d9d`](https://github.com/jbloomAus/SAEDashboard/commit/8bb4d9de0c75ffed5daaba4d5ec563fbbee38f86))
+
+* format ([`f667d92`](https://github.com/jbloomAus/SAEDashboard/commit/f667d92d9359e5c7976e21e821ac0dde8a081da6))
+
+* depend on latest sae_lens ([`4a2a6a0`](https://github.com/jbloomAus/SAEDashboard/commit/4a2a6a0fd70d7b4a3f1f870a510a800b31f57264))
+
+* use a much better method for getting subsets of feature activations ([`7101f13`](https://github.com/jbloomAus/SAEDashboard/commit/7101f13e13b4de5659623433ec359ecf2142daef))
+
+* add to gitignore ([`20180e0`](https://github.com/jbloomAus/SAEDashboard/commit/20180e06a279ef93d6127b467511911db352bce5))
+
+* add isort ([`3ab0fda`](https://github.com/jbloomAus/SAEDashboard/commit/3ab0fdaf75f735ec2eedc904529909111d0db0de))
+
+## v0.2.1 (2024-07-08)
+
+### Fix
+
+* fix: trigger release ([`87bf0b5`](https://github.com/jbloomAus/SAEDashboard/commit/87bf0b5f21f0d1f5397e514090601ec21c718e35))
+
+### Unknown
+
+* Merge pull request #6 from jbloomAus/fix-bfloat16
+
+fix bfloat 16 error ([`2f3c597`](https://github.com/jbloomAus/SAEDashboard/commit/2f3c597c1795357679e92caec3dd7e522c669fdb))
+
+* fix bfloat 16 error ([`63c3c62`](https://github.com/jbloomAus/SAEDashboard/commit/63c3c62f0a03e5656ed78cc0e8f853bea3f0938e))
+
+* Merge pull request #5 from jbloomAus/np-updates
+
+Updates + fixes for Neuronpedia ([`9e6b5c4`](https://github.com/jbloomAus/SAEDashboard/commit/9e6b5c427024b8a468b0d06e4e096c2561c35d5d))
+
+* Fix SAELens compatibility ([`139e1a2`](https://github.com/jbloomAus/SAEDashboard/commit/139e1a2f219d790c6f8faa9be34d9fbc9403dda3))
+
+* Rename file ([`16709ad`](https://github.com/jbloomAus/SAEDashboard/commit/16709add9ee5063b3682be34eef0aea2ddf4eceb))
+
+* Fix type ([`6b20386`](https://github.com/jbloomAus/SAEDashboard/commit/6b2038682ca41423dda3a3597bbe88120b120262))
+
+* Make Neuronpedia outputs an object, and add a real acceptance test ([`a5db256`](https://github.com/jbloomAus/SAEDashboard/commit/a5db2560e5f90a49257124635b3fdbee117ed860))
+
+* Np Runner: Multi-gpu defaults ([`07f7128`](https://github.com/jbloomAus/SAEDashboard/commit/07f71282681ffa801dd15f9265be349cd5745b42))
+
+* Ensure minibatch is on correct device ([`e206546`](https://github.com/jbloomAus/SAEDashboard/commit/e2065462c445df0e0985fb6588d4c01cb39bbef5))
+
+* NP Runner: Automatically use multi-gpu, devices ([`bf280e6`](https://github.com/jbloomAus/SAEDashboard/commit/bf280e685dc4dd2018cd41aa94a29bc853fcee18))
+
+* Allow dtype override ([`a40077d`](https://github.com/jbloomAus/SAEDashboard/commit/a40077dac1fa2ae880fcdabe3227878ef2cfaebe))
+
+* NP-Runner: Remove unnecessary layer of batching. ([`e2ac92b`](https://github.com/jbloomAus/SAEDashboard/commit/e2ac92b036d0192e132c8a8700a5a2f448d1983b))
+
+* NP Runner: Allow skipping sparsity check ([`ef74d2a`](https://github.com/jbloomAus/SAEDashboard/commit/ef74d2aeea2463afe150a5e8824da5a5206cd3d0))
+
+* Merge pull request #2 from jbloomAus/multiple-devices
+
+feat: Multiple devices ([`535e6c9`](https://github.com/jbloomAus/SAEDashboard/commit/535e6c9689d855f82a6ddfd9f169720fe367bde3))
+
+* format ([`7f892ad`](https://github.com/jbloomAus/SAEDashboard/commit/7f892ad0efb42025df0bcf26bdddd6fac4c2d8b1))
+
+* NP runner takes device args seperately ([`8fc31dd`](https://github.com/jbloomAus/SAEDashboard/commit/8fc31dd6ccd59f4f35742a4e15c380673c8cb2a3))
+
+* multi-gpu-support ([`5e24e4e`](https://github.com/jbloomAus/SAEDashboard/commit/5e24e4e6598dd7943f8d677042dcf84bc6f7a0a6))
+
+## v0.2.0 (2024-06-10)
+
+### Feature
+
+* feat: experimental release 2 ([`e264f97`](https://github.com/jbloomAus/SAEDashboard/commit/e264f97d90299f6ade294db8ed03aed9cd7491ee))
+
+## v0.1.0 (2024-06-10)
+
+### Feature
+
+* feat: experimental release ([`d79310a`](https://github.com/jbloomAus/SAEDashboard/commit/d79310a7b6599f7b813e214c9268d736e0cb87f0))
+
+### Unknown
+
+* fix pyproject.toml ([`a27c87d`](https://github.com/jbloomAus/SAEDashboard/commit/a27c87da987f043b470abce3404e305ec3f0d620))
+
+* test deployment ([`288a2d9`](https://github.com/jbloomAus/SAEDashboard/commit/288a2d9bf797a1a2f9947b1ceac5e47edc1684ba))
+
+* refactor np runner and add acceptance test ([`212593c`](https://github.com/jbloomAus/SAEDashboard/commit/212593c33b3aec33078a121738c0a826f705722f))
+
+* Fix: Default context tokens length for neuronpedia runner ([`aefe95c`](https://github.com/jbloomAus/SAEDashboard/commit/aefe95cb1be4139ac45f042abdc78e0feccfb490))
+
+* Allow custom context tokens length for Neuronpedia runner ([`d204cc8`](https://github.com/jbloomAus/SAEDashboard/commit/d204cc8fbb2ef376a1a5e00cd4f1cc5db2afb279))
+
+* Fix: Streaming default true ([`1b91dff`](https://github.com/jbloomAus/SAEDashboard/commit/1b91dff045fdbd8c118c5f209750eca60c260f5f))
+
+* Fix n_devices error for non-cuda ([`70b2dbd`](https://github.com/jbloomAus/SAEDashboard/commit/70b2dbdb2da51f5d78b1c2ce3210865fc259c97b))
+
+* fix import path for ci ([`3bd4687`](https://github.com/jbloomAus/SAEDashboard/commit/3bd468727e2ab0b7d77224b7c0dad88e0727b773))
+
+* make pyright happy, start config ([`b39ae85`](https://github.com/jbloomAus/SAEDashboard/commit/b39ae85d938a0db7c70b7dff9683f68f255dfb67))
+
+* add black ([`236855b`](https://github.com/jbloomAus/SAEDashboard/commit/236855be1ef1464ea85b2afc6aaee963326f9257))
+
+* fix ci ([`12818d7`](https://github.com/jbloomAus/SAEDashboard/commit/12818d7e6cd3e483258598b668805c1a9a048049))
+
+* add pytest cov ([`aae0571`](https://github.com/jbloomAus/SAEDashboard/commit/aae057159639cd247a82fdeda9eddb98612ceec6))
+
+* bring checks in line with sae_lens ([`7cd9679`](https://github.com/jbloomAus/SAEDashboard/commit/7cd9679cc18c64a7c8a0a07a1f12e6fc87543537))
+
+* activation scaling factor ([`333d377`](https://github.com/jbloomAus/SAEDashboard/commit/333d3770d0d1d3c40dfeb3335dcfc46e9b7da717))
+
+* Move Neuronpedia runner to SAEDashboard ([`4e691ea`](https://github.com/jbloomAus/SAEDashboard/commit/4e691eaad919e12b9cae6ff707eaa3cf322ea030))
+
+* fold w_dec norm by default ([`b6c9bc7`](https://github.com/jbloomAus/SAEDashboard/commit/b6c9bc70dc419d1e32bfb5580997369215e15429))
+
+* rename sae_vis to sae_dashboard ([`f0f5341`](https://github.com/jbloomAus/SAEDashboard/commit/f0f5341ffdf31a11884777d6ba8100cd302b9dab))
+
+* rename feature data generator ([`e02ed0a`](https://github.com/jbloomAus/SAEDashboard/commit/e02ed0a18e92c497aea3e137cf43e9f354f8f30f))
+
+* update demo ([`8aa9e52`](https://github.com/jbloomAus/SAEDashboard/commit/8aa9e5272f54d04b741e63aa335bfa1212a2d0f7))
+
+* add demo ([`dd3036f`](https://github.com/jbloomAus/SAEDashboard/commit/dd3036f90e6a4ed459ec21647744d491911900ac))
+
+* delete old demo files ([`3d86202`](https://github.com/jbloomAus/SAEDashboard/commit/3d8620204cf6acb21b5e7f9983c300341345cd88))
+
+* remove unnecessary print statement ([`9d3d937`](https://github.com/jbloomAus/SAEDashboard/commit/9d3d937e74f5575dde68d5a21fb73ce6f826d0d4))
+
+* set sae lens version ([`87a7691`](https://github.com/jbloomAus/SAEDashboard/commit/87a76911ff0f0d46ab421d9b5107aef27216e88b))
+
+* update older readme ([`c5c98e5`](https://github.com/jbloomAus/SAEDashboard/commit/c5c98e53531874efab5bc16235d9c72816fa61d5))
+
+* test ([`923da42`](https://github.com/jbloomAus/SAEDashboard/commit/923da427b56178acd99b988d6d6b51368b5d2359))
+
+* remove sae lens dep ([`2c26d5f`](https://github.com/jbloomAus/SAEDashboard/commit/2c26d5f4c40c41f750971601968577f316e15598))
+
+* Merge branch 'refactor_b' ([`3154d63`](https://github.com/jbloomAus/SAEDashboard/commit/3154d636e1a9f8a30b54c17e62a842bed3f8b2a1))
+
+* pass linting ([`0c079a1`](https://github.com/jbloomAus/SAEDashboard/commit/0c079a105b1b98e0edf2ff1a15593567c81bb103))
+
+* format ([`6f37e2e`](https://github.com/jbloomAus/SAEDashboard/commit/6f37e2eb050a3207a2d3b9defd5d416645215c7c))
+
+* run ci on all branches ([`faa0cc4`](https://github.com/jbloomAus/SAEDashboard/commit/faa0cc4eed4ff35f1e04656a968214c4fefbd573))
+
+* don't use feature ablations ([`dc6e6dc`](https://github.com/jbloomAus/SAEDashboard/commit/dc6e6dc2d2affce331894d8bb61942e103182652))
+
+* mock information in sequences to make normal sequence generation pass ([`c87b82f`](https://github.com/jbloomAus/SAEDashboard/commit/c87b82fdcc5e849d970cdc8bd1e841ec3e3e48ce))
+
+* Remove resid ([`ff83737`](https://github.com/jbloomAus/SAEDashboard/commit/ff837373b65e60d8a9ba7c6e61f78bddc4d170f2))
+
+* adding a test for direct_effect_feature_ablation_experiment ([`a9f3d1b`](https://github.com/jbloomAus/SAEDashboard/commit/a9f3d1b8021d8eeb60cf465934037d07583fa0b2))
+
+* shortcut direct_effect_feature_ablation_experiment if everything is zero ([`2c68ff0`](https://github.com/jbloomAus/SAEDashboard/commit/2c68ff0c8496c58cc0732f3c51905c9c9f405393))
+
+* fixing CI and replacing manual snapshots with syrupy snapshots ([`3b97640`](https://github.com/jbloomAus/SAEDashboard/commit/3b97640803cab3e3915202ac80c43b855c69c1cb))
+
+* more refactor, WIP ([`81657c8`](https://github.com/jbloomAus/SAEDashboard/commit/81657c8c897a81102c0df7b29c49d526e639bb44))
+
+* continue refactor, make data generator ([`eb1ae0f`](https://github.com/jbloomAus/SAEDashboard/commit/eb1ae0fc621407b33481c50b78b041079b08393d))
+
+* add use of safetensors cache for repeated calculations ([`a241c32`](https://github.com/jbloomAus/SAEDashboard/commit/a241c322334340a84c2a252bc0b4a40ed2f19bc9))
+
+* more refactor / benchmarking ([`d65ee87`](https://github.com/jbloomAus/SAEDashboard/commit/d65ee87cd191b2ed279f9f6efabb9e98bb700855))
+
+* only run unit tests ([`5f11ddd`](https://github.com/jbloomAus/SAEDashboard/commit/5f11ddd9bc25f9c9bb7cbeba11224ba12b260ea8))
+
+* fix lint issue ([`24daf17`](https://github.com/jbloomAus/SAEDashboard/commit/24daf17cb92534901681affbcebea314e2cf6580))
+
+* format ([`83e89ed`](https://github.com/jbloomAus/SAEDashboard/commit/83e89ed4860d886ccf19be591bf72d0e029e7344))
+
+* organise tests, make sure only unit tests run on CI ([`21f5fb1`](https://github.com/jbloomAus/SAEDashboard/commit/21f5fb155665329531b16d10673ddd988e7034ea))
+
+*  see if we can do some caching ([`c1dca6f`](https://github.com/jbloomAus/SAEDashboard/commit/c1dca6faa61de0849453acc83ae23baab6cf48be))
+
+* more refactoring ([`b3f0f41`](https://github.com/jbloomAus/SAEDashboard/commit/b3f0f41f36f0eee57a08142880d4b6654309e62c))
+
+* further refactor, possible significant speed up ([`ddd3496`](https://github.com/jbloomAus/SAEDashboard/commit/ddd3496206c0f3e751b596ca51e3544c77ddaf94))
+
+* more refactor ([`a5f6deb`](https://github.com/jbloomAus/SAEDashboard/commit/a5f6deb4263c58803e7af23d767fc5cb17dfd2b2))
+
+* refactoring in progress ([`d210b60`](https://github.com/jbloomAus/SAEDashboard/commit/d210b6056aa5316d2fd917e24ca8a819331a8114))
+
+* use named arguments ([`4a81053`](https://github.com/jbloomAus/SAEDashboard/commit/4a8105355d3b86e460f32cd5c736dde0dbeaa2e3))
+
+* remove create method ([`43b2018`](https://github.com/jbloomAus/SAEDashboard/commit/43b20184ed5ed0c2f08cfd13423f2271fd871274))
+
+* move chunk ([`0f26aa8`](https://github.com/jbloomAus/SAEDashboard/commit/0f26aa85bc9fbe358f4c5f90971d51b86159f095))
+
+* use fixtures ([`7c11dd9`](https://github.com/jbloomAus/SAEDashboard/commit/7c11dd914d467957e5c00b914f302c291924e411))
+
+* refactor to create runner ([`9202c19`](https://github.com/jbloomAus/SAEDashboard/commit/9202c19f4ad6134eb6b68f857c9e4bfd0b911cf8))
+
+* format ([`abd8747`](https://github.com/jbloomAus/SAEDashboard/commit/abd87472b76cfc151abfe2a6e312ea43b29c2250))
+
+* target ci at this branch ([`ea3b2a3`](https://github.com/jbloomAus/SAEDashboard/commit/ea3b2a3181f2eb1ff52d83b2040b586d6fdfef4a))
+
+* comment out release process for now ([`7084b5b`](https://github.com/jbloomAus/SAEDashboard/commit/7084b5ba3325bb559a8377d379bd2f3ba6d68348))
+
+* test generated output ([`7b8b2ab`](https://github.com/jbloomAus/SAEDashboard/commit/7b8b2abd94213d67c378b7746107f6a7c811d93c))
+
+* commit current demo html ([`00a03a0`](https://github.com/jbloomAus/SAEDashboard/commit/00a03a02fbf181caa55704defac25578b4444452))
+
+## v0.0.1 (2024-04-25)
+
+### Chore
+
+* chore: setting up pytest ([`2079d00`](https://github.com/jbloomAus/SAEDashboard/commit/2079d00911d1a00ee19cde478b5cab61ca9c0495))
+
+* chore: setting up semantic-release ([`09075af`](https://github.com/jbloomAus/SAEDashboard/commit/09075afbec279fb89d157f73e9a0ed47ba66d3c8))
+
+### Fix
+
+* fix: remove circular dep with sae lens ([`1dd9f6c`](https://github.com/jbloomAus/SAEDashboard/commit/1dd9f6cd22f879e8d6904ba72f3e52b4344433cd))
+
+### Unknown
+
+* Merge pull request #44 from chanind/pytest-setup
+
+chore: setting up pytest ([`034eefa`](https://github.com/jbloomAus/SAEDashboard/commit/034eefa5a4163e9a560b574e2e255cd06f8f49a1))
+
+* Merge pull request #43 from callummcdougall/move_saelens_dep
+
+Remove dependency on saelens from pyproject, add to demo.ipynb ([`147d87e`](https://github.com/jbloomAus/SAEDashboard/commit/147d87ee9534d30e764851cbe73aadb5783d2515))
+
+* Add missing matplotlib ([`572a3cc`](https://github.com/jbloomAus/SAEDashboard/commit/572a3cc79709a14117bbeafb871a33f0107600d8))
+
+* Remove dependency on saelens from pyproject, add to demo.ipynb ([`1e6f3cf`](https://github.com/jbloomAus/SAEDashboard/commit/1e6f3cf9b2bcfb381a73d9333581c430faa531fd))
+
+* Merge branch 'main' of https://github.com/callummcdougall/sae_vis ([`4e7a24c`](https://github.com/jbloomAus/SAEDashboard/commit/4e7a24c37444f11d718035eede68ac728d949a20))
+
+* Merge pull request #41 from callummcdougall/allow_disable_buffer
+
+oops I forgot to switch back to main before pushing ([`1312cd0`](https://github.com/jbloomAus/SAEDashboard/commit/1312cd09d6e274b1163e79d2ac01f2df54c65157))
+
+* Merge branch 'main' into allow_disable_buffer ([`e7edf5a`](https://github.com/jbloomAus/SAEDashboard/commit/e7edf5a9bae4714bf4983ce6a19a0fe6fdf1f118))
+
+* Merge pull request #40 from chanind/semantic-release-autodeploy
+
+chore: setting up semantic-release for auto-deploy ([`a4d44d1`](https://github.com/jbloomAus/SAEDashboard/commit/a4d44d1a0e86055fb82ef41f51f0adbb7868df3c))
+
+* Merge pull request #38 from chanind/type-checking
+
+Enabling type checking with Pyright ([`f1fd792`](https://github.com/jbloomAus/SAEDashboard/commit/f1fd7926f46f00dca46024377f53aa8f2db98773))
+
+* enabling type checking with Pyright ([`05d14ea`](https://github.com/jbloomAus/SAEDashboard/commit/05d14eafea707d3db81e78b4be87199087cb8e37))
+
+* Merge pull request #39 from callummcdougall/fix_loading_saelens_sae
+
+FIX: SAELens new format has "scaling_factor" key, which causes assert to fail ([`983aee5`](https://github.com/jbloomAus/SAEDashboard/commit/983aee562aea31e90657caf8c6ab6e450e952120))
+
+* Fix Formatting ([`13b8106`](https://github.com/jbloomAus/SAEDashboard/commit/13b81062485f5dce2568e7832bfb2aae218dd4e9))
+
+* Merge branch 'main' into fix_loading_saelens_sae ([`21b0086`](https://github.com/jbloomAus/SAEDashboard/commit/21b0086b8af3603441795e925a15e7cded122acb))
+
+* format ([`8f1506b`](https://github.com/jbloomAus/SAEDashboard/commit/8f1506b6eb7dc0a2d4437d2aa23a0898c46a156d))
+
+* Allow SAELens autoencoder keys to be superset of required keys, instead of exact match ([`6852170`](https://github.com/jbloomAus/SAEDashboard/commit/6852170d55e7d3cf22632c5807cfab219516da98))
+
+* v0.2.17 ([`2bb14da`](https://github.com/jbloomAus/SAEDashboard/commit/2bb14daa88a0af601e13f4e51b50a2b00cd75b48))
+
+* Use main branch of SAELens ([`2b34505`](https://github.com/jbloomAus/SAEDashboard/commit/2b345052bdc92ee9c1255cab0978916307a0a9dc))
+
+* Update version 0.2.16 ([`bf90293`](https://github.com/jbloomAus/SAEDashboard/commit/bf902930844db9b0f8db4fbe8b3610557352660b))
+
+* Merge pull request #36 from callummcdougall/allow_disable_buffer
+
+FEATURE: Allow setting buffer to None, which gives the whole activation sequence ([`f5f9594`](https://github.com/jbloomAus/SAEDashboard/commit/f5f9594fcaf5edb6036a85446e092278004ea200))
+
+* 16 ([`64e7018`](https://github.com/jbloomAus/SAEDashboard/commit/64e701849570d9e172dc065812c9a3e7149a9176))
+
+* version 0.2.16 ([`afca0be`](https://github.com/jbloomAus/SAEDashboard/commit/afca0be8826e0c007b5730fa9fa18454699d16a3))
+
+* Fix version ([`5a43916`](https://github.com/jbloomAus/SAEDashboard/commit/5a43916cbd9836396f051f7a258fdca8664e05e9))
+
+* fix all indices view ([`5f87d52`](https://github.com/jbloomAus/SAEDashboard/commit/5f87d52154d6a8e8c8984836bbe8f85ee25f279d))
+
+* Merge branch 'fix_gpt2_demo' into allow_disable_buffer ([`ea57bfc`](https://github.com/jbloomAus/SAEDashboard/commit/ea57bfc2ee1e23666810982abf32e6e9cbb74193))
+
+* Allow disabling the buffer ([`c1be9f8`](https://github.com/jbloomAus/SAEDashboard/commit/c1be9f8e4b8ee6d8f18c4a1a0445840304440c1d))
+
+* fix conflicts ([`ea3d624`](https://github.com/jbloomAus/SAEDashboard/commit/ea3d624013b9aa7cbd2d6eaa7212a1f7c4ee8e28))
+
+* Merge pull request #35 from callummcdougall/fix_gpt2_demo
+
+Fix usage of SAELens and demo notebook ([`88b5933`](https://github.com/jbloomAus/SAEDashboard/commit/88b59338d3cadbd5c70f0c1117dff00f01a54e6a))
+
+* Import updated SAELens, use correct tokens, fix missing file cfg.json file error. ([`14ba9b0`](https://github.com/jbloomAus/SAEDashboard/commit/14ba9b03d4ce791ba8f4cac553fb82a93c47dfb8))
+
+* Merge pull request #34 from ArthurConmy/patch-1
+
+Update README.md ([`3faac82`](https://github.com/jbloomAus/SAEDashboard/commit/3faac82686f546800492d8aeb5e1d5919cbf1517))
+
+* Update README.md ([`416eca8`](https://github.com/jbloomAus/SAEDashboard/commit/416eca8073c6cb2b120c759330ec47f52ab32d1e))
+
+* Merge pull request #33 from chanind/setup-poetry-and-ruff
+
+Setting up poetry / ruff / github actions ([`287f30f`](https://github.com/jbloomAus/SAEDashboard/commit/287f30f1d8fc39ab583f202c9277e07e5eeeaf62))
+
+* setting up poetry and ruff for linting/formatting ([`0e0eba9`](https://github.com/jbloomAus/SAEDashboard/commit/0e0eba9e4d54c746cddc835ef4f6ddf2bab96844))
+
+* fix feature vis demo gpt ([`821781e`](https://github.com/jbloomAus/SAEDashboard/commit/821781e96b732a5909d8735714482c965891b2ea))
+
+* add scatter plot support ([`6eab28b`](https://github.com/jbloomAus/SAEDashboard/commit/6eab28bef9ef5cd9360fef73e02763301fa1a028))
+
+* update setup ([`8d2ca53`](https://github.com/jbloomAus/SAEDashboard/commit/8d2ca53e8a6bba860fe71368741d06a718adaa27))
+
+* fix setup ([`9cae8f4`](https://github.com/jbloomAus/SAEDashboard/commit/9cae8f461bd780e23eb2d994f56b495ede16201a))
+
+* Merge branch 'main' of https://github.com/callummcdougall/sae_vis ([`ed8f8cb`](https://github.com/jbloomAus/SAEDashboard/commit/ed8f8cb7ad1fba2383dcdd471c33ce4a1b9f32e3))
+
+* Merge pull request #27 from wllgrnt/will-add-eindex-dependency
+
+Update setup.py with eindex dependency ([`8d7ed12`](https://github.com/jbloomAus/SAEDashboard/commit/8d7ed123505ac7ecf93dd310f57888547aead1d7))
+
+* two more deps ([`7f231a8`](https://github.com/jbloomAus/SAEDashboard/commit/7f231a83acfef2494c1866249f57e10c21a1a443))
+
+* Update setup.py with eindex
+
+Without this, 'pip install sae-vis' will cause errors when e.g. you do 'from sae_vis.data_fetching_fns import get_feature_data' ([`a9d7de9`](https://github.com/jbloomAus/SAEDashboard/commit/a9d7de90b492f7305758e15303ba890fb9b503d0))
+
+* fix sae bug ([`247d14b`](https://github.com/jbloomAus/SAEDashboard/commit/247d14b55f209ed9ccf50e5ce091ed66ffbf19d2))
+
+* Merge pull request #32 from hijohnnylin/pin_older_sae_training
+
+Demo notebook errors under "Multi-layer models" vis ([`9ac1dac`](https://github.com/jbloomAus/SAEDashboard/commit/9ac1dac51af32909666977cb5b3794965c70f62f))
+
+* Pin older commit of mats_sae_training ([`8ca7ac1`](https://github.com/jbloomAus/SAEDashboard/commit/8ca7ac14b919fedb91240630ac7072cac40a6d6a))
+
+* update version number ([`72e584b`](https://github.com/jbloomAus/SAEDashboard/commit/72e584b6492ed1ef3989968f6588a17fca758650))
+
+* add gifs to readme ([`1393740`](https://github.com/jbloomAus/SAEDashboard/commit/13937405da31cca70cd1027aaca6c9cc84797ff1))
+
+* test gif ([`4fbafa6`](https://github.com/jbloomAus/SAEDashboard/commit/4fbafa69343dc58dc18d0f78e393b5fcc9e24c0c))
+
+* fix height issue ([`3f272f6`](https://github.com/jbloomAus/SAEDashboard/commit/3f272f61a954effef7bd648cc8117346da3bb971))
+
+* fix pypi ([`7151164`](https://github.com/jbloomAus/SAEDashboard/commit/7151164cc0df8af278617147f07cbfbe3977cfeb))
+
+* update setup ([`8c43478`](https://github.com/jbloomAus/SAEDashboard/commit/8c43478ad2eba8d3d4106fe4239c1229b8720fe6))
+
+* Merge pull request #26 from hijohnnylin/update_html_anomalies
+
+Update and add some HTML_ANOMALIES ([`1874a47`](https://github.com/jbloomAus/SAEDashboard/commit/1874a47a099ce32795bdbb5f98b9167dcca85ff2))
+
+* Update and add some HTML_ANOMALIES ([`c541b7f`](https://github.com/jbloomAus/SAEDashboard/commit/c541b7f06108046ad1e2eb82c89f30f061f4411e))
+
+* 0.2.9 ([`a5c8a6d`](https://github.com/jbloomAus/SAEDashboard/commit/a5c8a6d2008b818db90566cba50211845c753444))
+
+* fix readme ([`5a8a7e3`](https://github.com/jbloomAus/SAEDashboard/commit/5a8a7e3173fc50fdb5ff0e56d7fa83e475af38a3))
+
+* include feature tables ([`7c4c263`](https://github.com/jbloomAus/SAEDashboard/commit/7c4c263a2e069482d341b6265015664792bde817))
+
+* add license ([`fa02a3d`](https://github.com/jbloomAus/SAEDashboard/commit/fa02a3dc93b721322b3902e2ac416ed156bf9d80))
+
+* Merge branch 'main' of https://github.com/callummcdougall/sae_vis ([`ca5efcd`](https://github.com/jbloomAus/SAEDashboard/commit/ca5efcdc81074d3c3002bd997b35e326a44a4a25))
+
+* Merge pull request #24 from chanind/fix-pypi-repo-link
+
+fixing repo URL in setup.py ([`14a0be5`](https://github.com/jbloomAus/SAEDashboard/commit/14a0be54a57b1bc73ac4741611f9c8d1bd229e6f))
+
+* fixing repo URL in setup.py ([`4faeca5`](https://github.com/jbloomAus/SAEDashboard/commit/4faeca5da06c0bb4384e202a91d895a217365d30))
+
+* re-fix html anomalies ([`2fbae4c`](https://github.com/jbloomAus/SAEDashboard/commit/2fbae4c9a7dd663737bae25e73e978d40c59064a))
+
+* fix hook point bug ([`9b573b2`](https://github.com/jbloomAus/SAEDashboard/commit/9b573b27590db1cbd6c8ef08fca7ff8c9d26b340))
+
+* Merge pull request #20 from chanind/fix-final-resid-layer
+
+fixing bug if hook_point == hook_point_resid_final ([`d6882e3`](https://github.com/jbloomAus/SAEDashboard/commit/d6882e3f813ef0d399e07548871f61b1f6a98ac6))
+
+* fixing bug using hook_point_resid_final ([`cfe9b30`](https://github.com/jbloomAus/SAEDashboard/commit/cfe9b3042cfe127d5f7958064ffe817c25a19b56))
+
+* fix indexing speed ([`865ff64`](https://github.com/jbloomAus/SAEDashboard/commit/865ff64329538641cd863dc7668dfc77907fb384))
+
+* enable JSON saving ([`feea47a`](https://github.com/jbloomAus/SAEDashboard/commit/feea47a342d52296b72784ed18ea628848d4c7d4))
+
+* Merge pull request #19 from chanind/support-mlp-and-attn-out
+
+supporting mlp and attn out hooks ([`1c5463b`](https://github.com/jbloomAus/SAEDashboard/commit/1c5463b12f85cd0598b4e2fba5c556b1e9c0fbbe))
+
+* supporting mlp and attn out hooks ([`a100e58`](https://github.com/jbloomAus/SAEDashboard/commit/a100e586498e8cae14df475bc7924cdecaed71ea))
+
+* Merge branch 'main' of https://github.com/callummcdougall/sae_vis ([`083aeba`](https://github.com/jbloomAus/SAEDashboard/commit/083aeba0e4048d9976ec5cbee8df7dc8fd4db4e9))
+
+* Merge pull request #18 from chanind/remove-build-artifacts
+
+removing Python build artifacts and adding to .gitignore ([`b0e0594`](https://github.com/jbloomAus/SAEDashboard/commit/b0e0594590b4472b34052c6eb3ebceb6c9f58a11))
+
+* removing Python build artifacts and adding to .gitignore ([`b6486f5`](https://github.com/jbloomAus/SAEDashboard/commit/b6486f56bea9d4bb7544c36afe70e6f891101b63))
+
+* fix variable naming ([`2507918`](https://github.com/jbloomAus/SAEDashboard/commit/25079186b3f31d2271b1ecdb11f26904af7146d2))
+
+* update readme ([`0ee3608`](https://github.com/jbloomAus/SAEDashboard/commit/0ee3608af396a1a6586dfb809f2f6480bb4f6390))
+
+* update readme ([`f8351f8`](https://github.com/jbloomAus/SAEDashboard/commit/f8351f88e8432ccd4b2206e859daea316304d6c6))
+
+* update version number ([`1e74408`](https://github.com/jbloomAus/SAEDashboard/commit/1e7440883f44a92705299430215f802fea4e1915))
+
+* fix formatting and docstrings ([`b9fe2bb`](https://github.com/jbloomAus/SAEDashboard/commit/b9fe2bbb15a48e4b0415f6f4240d895990d54c9a))
+
+* Merge pull request #17 from jordansauce/sae-agnostic-functions-new
+
+Added SAE class agnostic functions ([`0039c6f`](https://github.com/jbloomAus/SAEDashboard/commit/0039c6f8f99d6e8a1b2ff56aa85f60a3eba3afb0))
+
+* Added sae class agnostic functions
+
+Added parse_feature_data() and parse_prompt_data() ([`e2709d0`](https://github.com/jbloomAus/SAEDashboard/commit/e2709d0b4c55d73d6026f3b9ce534f59ce61f344))
+
+* add to pypi ([`02a5b9a`](https://github.com/jbloomAus/SAEDashboard/commit/02a5b9acd15433cc59d438271b9bd5e12d62b662))
+
+* update notebook images ([`b87ad4d`](https://github.com/jbloomAus/SAEDashboard/commit/b87ad4d256f12c23605b0e7db307ee56913c93ef))
+
+* fix layer parse and custom device ([`14c7ae9`](https://github.com/jbloomAus/SAEDashboard/commit/14c7ae9d0c8b7dad21b953cfc93fe7f34c74e149))
+
+* update dropdown styling ([`83be219`](https://github.com/jbloomAus/SAEDashboard/commit/83be219bfe31b985a26762e06345c574aa0e6fe1))
+
+* add custom prompt vis ([`cabdc5c`](https://github.com/jbloomAus/SAEDashboard/commit/cabdc5cb31f881cddf236490c41332c525d2ee74))
+
+* d3 & multifeature refactor ([`f79a919`](https://github.com/jbloomAus/SAEDashboard/commit/f79a919691862f60a9e30fe0f79fd8e771bc932a))
+
+* remove readme links ([`4bcef48`](https://github.com/jbloomAus/SAEDashboard/commit/4bcef489b644dd3357b1975f3245d534f6f0d2e0))
+
+* add demo html ([`629c713`](https://github.com/jbloomAus/SAEDashboard/commit/629c713345407562dc4ccd9875bf3cfab5480bdd))
+
+* remove demos ([`beedea9`](https://github.com/jbloomAus/SAEDashboard/commit/beedea9667761534a5293015aff9cc17638666a5))
+
+* fix quantile error ([`3a23cfd`](https://github.com/jbloomAus/SAEDashboard/commit/3a23cfd56f21fe0775a1a9957db340d15f75f51a))
+
+* width 425 ([`f25c776`](https://github.com/jbloomAus/SAEDashboard/commit/f25c776d5cb746916d3f2fdf368cbd5448742949))
+
+* fix device bug ([`85dfa49`](https://github.com/jbloomAus/SAEDashboard/commit/85dfa497bc804945911e80607ac31cf3afbdc759))
+
+* dont return vocab dict ([`b4c7138`](https://github.com/jbloomAus/SAEDashboard/commit/b4c713873870acb4035986cc5bff3a4ce1e466c9))
+
+* save as JSON, fix device ([`eba2cff`](https://github.com/jbloomAus/SAEDashboard/commit/eba2cff3eb6215558577a6b4d4f8cc716766b927))
+
+* simple fixed and issues ([`b28a0f7`](https://github.com/jbloomAus/SAEDashboard/commit/b28a0f7c7e936f4bea05528d952dfcd438533cce))
+
+* Merge pull request #8 from lucyfarnik/topk-empty-mask
+
+Topk error handling for empty masks ([`2740c00`](https://github.com/jbloomAus/SAEDashboard/commit/2740c0047e78df7e56d7bcf707c909ac18e71c1f))
+
+* Topk error handling for empty masks ([`1c2627e`](https://github.com/jbloomAus/SAEDashboard/commit/1c2627e237f8f67725fc44e60a190bc141d36fc8))
+
+* viz to vis ([`216d02b`](https://github.com/jbloomAus/SAEDashboard/commit/216d02b550d6fbcb9b37d39c1b272a7dda91aadc))
+
+* update readme links ([`f9b3f95`](https://github.com/jbloomAus/SAEDashboard/commit/f9b3f95e31e7150024be27ec62246f43bf9bcbb8))
+
+* update for TL ([`1941db1`](https://github.com/jbloomAus/SAEDashboard/commit/1941db1e22093d6fc88fb3fcd6f4c7d535d8b3b4))
+
+* Merge pull request #5 from lucyfarnik/transformer-lens-models
+
+Compatibility with TransformerLens models ([`8d59c6c`](https://github.com/jbloomAus/SAEDashboard/commit/8d59c6c5a5f2b98c486e5c74130371ad9254d1c9))
+
+* Merge branch 'main' into transformer-lens-models ([`73057d7`](https://github.com/jbloomAus/SAEDashboard/commit/73057d7e2a3e4e9669fc0556e64190811ac8b52d))
+
+* Merge pull request #4 from lucyfarnik/resid-saes-support
+
+Added support for residual-adjacent SAEs ([`b02e98b`](https://github.com/jbloomAus/SAEDashboard/commit/b02e98b3b852c0613a890f8949d04b5560fb6fd6))
+
+* Added support for residual-adjacent SAEs ([`89aacf1`](https://github.com/jbloomAus/SAEDashboard/commit/89aacf1b22aa81b393b10eca8611c9dbf406c638))
+
+* Merge pull request #7 from lucyfarnik/fix-histogram-div-zero
+
+Fixed division by zero in histogram calculation ([`3aee20e`](https://github.com/jbloomAus/SAEDashboard/commit/3aee20ea7f99cc07e6c5085fddb70cadd8327f4d))
+
+* Fixed division by zero in histogram calculation ([`e986e90`](https://github.com/jbloomAus/SAEDashboard/commit/e986e907cc42790efc93ce75ebf7b28a0278aaa2))
+
+* Merge pull request #6 from lucyfarnik/handling-dead-features
+
+Edge case handling for dead features ([`9e43c30`](https://github.com/jbloomAus/SAEDashboard/commit/9e43c308e58769828234e1505f1c1102ba651dfd))
+
+* Edge case handling for dead features ([`5197aee`](https://github.com/jbloomAus/SAEDashboard/commit/5197aee2c9f92bce7c5fd6d22201152a68c2e6ca))
+
+* add features argument ([`f24ef7e`](https://github.com/jbloomAus/SAEDashboard/commit/f24ef7ebebb3d4fd92e299858dbd5b968b78c69e))
+
+* fix image link ([`22c8734`](https://github.com/jbloomAus/SAEDashboard/commit/22c873434dfa84e3aed5ee0aab0fd25b288428a6))
+
+* Merge pull request #1 from lucyfarnik/read-me-links-fix
+
+Fixed readme links pointing to the old colab ([`86f8e20`](https://github.com/jbloomAus/SAEDashboard/commit/86f8e2012e376b6c498e5e708324f812af6fbc98))
+
+* Fixed readme links pointing to the old colab ([`28ef1cb`](https://github.com/jbloomAus/SAEDashboard/commit/28ef1cbd1b91f6c09c842f48e1f997d189ca04e7))
+
+* Added readme section about models ([`7523e7f`](https://github.com/jbloomAus/SAEDashboard/commit/7523e7f6363e030196496b3c6a3dc70b234c2d9a))
+
+* Compatibility with TransformerLens models ([`ba708e9`](https://github.com/jbloomAus/SAEDashboard/commit/ba708e987be6cc7a09d34ea8fb83de009312684d))
+
+* Added support for MPS ([`196c0a2`](https://github.com/jbloomAus/SAEDashboard/commit/196c0a24d0e8277b327eb2d57662075f9106990b))
+
+* black font ([`d81e74d`](https://github.com/jbloomAus/SAEDashboard/commit/d81e74d575326ef786881fb9182a768f9de2cb70))
+
+* fix html bug ([`265dedd`](https://github.com/jbloomAus/SAEDashboard/commit/265dedd376991230e2041fd37d5b6a0eda048545))
+
+* add jax and dataset deps ([`f1caeaf`](https://github.com/jbloomAus/SAEDashboard/commit/f1caeafc9613e27c7663447cf862301ac11d842d))
+
+* remove TL dependency ([`155991f`](https://github.com/jbloomAus/SAEDashboard/commit/155991fe61d0199d081d344ac44996edce35d118))
+
+* first commit ([`7782eb6`](https://github.com/jbloomAus/SAEDashboard/commit/7782eb6d5058372630c5bbb8693eb540a7bceaf4))

SAEDashboard/Dockerfile

@@ -0,0 +1,45 @@
+# docker build --target development -t decoderesearch/saedashboard-cuda --file Dockerfile .
+# docker run --entrypoint /bin/bash -it decoderesearch/saedashboard-cuda
+
+ARG APP_NAME=sae_dashboard
+ARG APP_PATH=/opt/$APP_NAME
+ARG PYTHON_VERSION=3.12.2
+ARG POETRY_VERSION=1.8.3
+
+FROM pytorch/pytorch:2.2.0-cuda12.1-cudnn8-devel AS staging
+ARG APP_NAME
+ARG APP_PATH
+ARG POETRY_VERSION
+
+ENV \
+    PYTHONDONTWRITEBYTECODE=1 \
+    PYTHONUNBUFFERED=1 \
+    PYTHONFAULTHANDLER=1
+ENV \
+    POETRY_VERSION=$POETRY_VERSION \
+    POETRY_HOME="/opt/poetry" \
+    POETRY_VIRTUALENVS_IN_PROJECT=true \
+    POETRY_NO_INTERACTION=1
+
+RUN apt-get update && apt-get install --no-install-recommends -y curl git-lfs vim && rm -rf /var/lib/apt/lists/*
+
+RUN curl -sSL https://install.python-poetry.org | python
+ENV PATH="$POETRY_HOME/bin:$PATH"
+
+WORKDIR $APP_PATH
+COPY ./pyproject.toml ./README.md ./
+COPY ./$APP_NAME ./$APP_NAME
+
+FROM staging AS development
+ARG APP_NAME
+ARG APP_PATH
+
+WORKDIR $APP_PATH
+RUN poetry lock
+RUN poetry install --no-dev --no-cache
+
+RUN curl -s https://packagecloud.io/install/repositories/github/git-lfs/script.deb.sh | bash
+RUN git lfs install
+
+ENTRYPOINT ["/bin/bash"]
+CMD ["poetry", "shell"]

SAEDashboard/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Decode Research
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

SAEDashboard/Makefile

@@ -0,0 +1,27 @@
+format:
+	poetry run black .
+	poetry run isort .
+
+lint:
+	poetry run flake8 .
+	poetry run black --check .
+	poetry run isort --check-only --diff .
+
+check-type:
+	poetry run pyright .
+
+test:
+	poetry run pytest --cov=sae_dashboard --cov-report=term-missing tests/unit
+
+check-ci:
+	make format
+	make lint
+	make check-type
+	make test
+
+profile-memory-unit:
+	poetry run pytest --memray tests/unit
+
+profile-speed-unit:
+	poetry run py.test tests/unit --profile-svg -k "test_SaeVisData_create_results_look_reasonable[Default]"
+	open prof/combined.svg

SAEDashboard/README.md

@@ -0,0 +1,221 @@
+# SAEDashboard
+
+SAEDashboard is a tool for visualizing and analyzing Sparse Autoencoders (SAEs) in neural networks. This repository is an adaptation and extension of Callum McDougal's [SAEVis](https://github.com/callummcdougall/sae_vis/tree/main), providing enhanced functionality for feature visualization and analysis as well as feature dashboard creation at scale.
+
+## Overview
+
+This codebase was originally designed to replicate Anthropic's sparse autoencoder visualizations, which you can see [here](https://transformer-circuits.pub/2023/monosemantic-features/vis/a1.html). SAEDashboard primarily provides visualizations of features, including their activations, logits, and correlations--similar to what is shown in the Anthropic link. 
+
+<img src="https://raw.githubusercontent.com/callummcdougall/computational-thread-art/master/example_images/misc/feature-vis-video.gif" width="800">
+
+## 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.
+

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 "$@"

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 }}

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<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Warning: Dataset is not tokenized. Pre-tokenizing will improve performance and allows for more control over special tokens. See https://jbloomaus.github.io/SAELens/training_saes/#pretokenizing-datasets for more info.\n",
+      "Tokens don't exist, making them.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  0%|          | 0/2048 [00:00<?, ?it/s]Token indices sequence length is longer than the specified maximum sequence length for this model (3180 > 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<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "2a0ddbda94d0407598edf564b4487407",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Extracting vis data from cached data:   0%|          | 0/1 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/Users/curttigges/Projects/SAEDashboard/sae_dashboard/vector_data_generator.py:205: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
+      "  return torch.load(\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "feature_indices: [0]\n"
+     ]
+    },
+    {
+     "data": {
+      "text/html": [
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┏━━━━━━┳━━━━━━┳━━━━━━━┓\n",
+       "┃<span style=\"font-weight: bold\"> Task </span>┃<span style=\"font-weight: bold\"> Time </span>┃<span style=\"font-weight: bold\"> Pct % </span>┃\n",
+       "┡━━━━━━╇━━━━━━╇━━━━━━━┩\n",
+       "└──────┴──────┴───────┘\n",
+       "</pre>\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
+}

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()

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
+}

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 <cal.s.mcdougall@gmail.com>", "Joseph Bloom, <jbloomaus@gmail.com>"]
+readme = "README.md"
+license = "MIT"
+
+[tool.poetry.dependencies]
+python = "^3.10"
+torch = "^2.0.0"
+einops = ">=0.7.0"
+datasets = "^3.0.0"
+dataclasses-json = "^0.6.4"
+jaxtyping = "^0.2.28"
+transformer-lens = "^2.2.0"
+transformers = "<4.57.0"
+eindex-callum = "^0.1.0"
+rich = "^13.7.1"
+matplotlib = "^3.8.4"
+safetensors = "^0.4.3"
+typer = "^0.12.3"
+sae-lens = "^6.8.0"
+decode-clt = "^0.0.1"
+hf-transfer = "^0.1.9"
+
+[tool.poetry.group.dev.dependencies]
+isort = "^5.13.2"
+ruff = "^0.3.7"
+pytest = "^8.1.1"
+ipykernel = "^6.29.4"
+pyright = "^1.1.359"
+pytest-profiling = "^1.7.0"
+memray = "^1.12.0"
+syrupy = "^4.6.1"
+flake8 = "^7.0.0"
+pytest-cov = "^5.0.0"
+black = "^24.4.2"
+pytest-memray = "^1.7.0"
+
+[tool.poetry.scripts]
+neuronpedia-runner = "sae_dashboard.neuronpedia.neuronpedia_runner:main"
+
+[tool.isort]
+profile = "black"
+src_paths = ["sae_dashboard", "tests"]
+
+[tool.pyright]
+typeCheckingMode = "strict"
+reportMissingTypeStubs = "none"
+reportUnknownMemberType = "none"
+reportUnknownArgumentType = "none"
+reportUnknownVariableType = "none"
+reportUntypedFunctionDecorator = "none"
+reportUnnecessaryIsInstance = "none"
+reportUnnecessaryComparison = "none"
+reportConstantRedefinition = "none"
+reportUnknownLambdaType = "none"
+reportPrivateUsage = "none"
+reportPrivateImportUsage = "none"
+
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"
+
+[tool.semantic_release]
+version_variables = ["sae_dashboard/__init__.py:__version__"]
+version_toml = ["pyproject.toml:tool.poetry.version"]
+build_command = "pip install poetry && poetry build"
+branches = { main = { match = "main" } }

SAEDashboard/sae_dashboard/__init__.py

@@ -0,0 +1,10 @@
+__version__ = "0.7.3"
+
+# from .data_fetching_fns import *
+# from .data_storing_fns import *
+# from .html_fns import *
+# from .transformer_lens_wrapper import *
+# from .utils_fns import *
+
+
+# from autoencoder import AutoEncoder, AutoEncoderConfig

SAEDashboard/sae_dashboard/clt_layer_wrapper.py

@@ -0,0 +1,697 @@
+import json
+
+# Added dataclass, field, asdict
+from dataclasses import asdict, dataclass, field
+from pathlib import Path
+
+# import torch.nn as nn # Unused
+# from torch.distributed import ProcessGroup # Unused
+# from types import SimpleNamespace # Unused import
+from typing import (  # Added Optional, Union and List
+    TYPE_CHECKING,
+    Any,
+    List,
+    Optional,
+    Union,
+)
+
+import torch
+import torch.nn.functional as F
+from clt.models.activations import BatchTopK  # type: ignore
+
+if TYPE_CHECKING:
+    import torch.distributed  # Import for ProcessGroup type hint
+    from clt.models.clt import CrossLayerTranscoder  # type: ignore
+
+
+# Placeholder for dist if torch.distributed is not available or initialized
+class MockDist:
+    def is_initialized(self) -> bool:
+        return False
+
+    def get_world_size(
+        self, group: "Optional[torch.distributed.ProcessGroup]" = None
+    ) -> int:
+        return 1
+
+    def all_gather_into_tensor(
+        self,
+        output_tensor: torch.Tensor,
+        input_tensor: torch.Tensor,
+        group: "Optional[torch.distributed.ProcessGroup]" = None,
+    ) -> None:
+        # In non-distributed setting, just copy input to output (assuming output is sized correctly)
+        if output_tensor.shape[0] == 1 * input_tensor.shape[0]:
+            output_tensor.copy_(input_tensor)
+        else:
+            # This case shouldn't happen if called correctly, but handle defensively
+            raise ValueError(
+                "Output tensor size doesn't match input tensor size in mock all_gather"
+            )
+
+    def all_gather(
+        self,
+        tensor_list: List[torch.Tensor],
+        input_tensor: torch.Tensor,
+        group: "Optional[torch.distributed.ProcessGroup]" = None,
+    ) -> None:
+        """Mock all_gather for a list of tensors."""
+        if self.get_world_size(group) == 1:
+            if len(tensor_list) == 1:
+                tensor_list[0].copy_(input_tensor)
+            else:
+                raise ValueError(
+                    "tensor_list size must be 1 in mock all_gather when world_size is 1"
+                )
+        else:
+            # This mock doesn't support actual gathering for world_size > 1.
+            # It's primarily for the dist.all_gather call in _gather_weight,
+            # which should ideally not proceed if world_size > 1 and dist is MockDist.
+            # However, _gather_weight checks dist.is_initialized() and dist.get_world_size() first.
+            raise NotImplementedError(
+                "MockDist.all_gather not implemented for world_size > 1"
+            )
+
+
+try:
+    import torch.distributed as dist
+
+    if not dist.is_available():
+        dist = MockDist()  # type: ignore
+except ImportError:
+    dist = MockDist()  # type: ignore
+
+
+@dataclass
+class CLTMetadata:
+    """Simple metadata class for CLT wrapper compatibility."""
+
+    hook_name: str
+    hook_layer: int
+    model_name: Optional[str] = None
+    context_size: Optional[int] = None
+    prepend_bos: bool = True
+    hook_head_index: Optional[int] = None
+    seqpos_slice: Optional[slice] = None
+
+
+@dataclass
+class CLTWrapperConfig:
+    """Configuration dataclass for the CLTLayerWrapper."""
+
+    # Fields without defaults first
+    d_sae: int
+    d_in: int
+    hook_name: str
+    hook_layer: int
+    dtype: str
+    device: str
+    # Fields with defaults last
+    architecture: str = "jumprelu"
+    hook_head_index: Optional[int] = None
+    model_name: Optional[str] = None
+    dataset_path: Optional[str] = None
+    context_size: Optional[int] = None
+    prepend_bos: bool = True
+    normalize_activations: bool = False
+    dataset_trust_remote_code: bool = False
+    seqpos_slice: Optional[slice] = None
+    model_from_pretrained_kwargs: dict[str, Any] = field(default_factory=dict)
+    metadata: Optional[CLTMetadata] = None
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert config to dictionary for compatibility with SAE interface."""
+        return asdict(self)
+
+
+class CLTLayerWrapper:
+    """Wraps a single layer of a CrossLayerTranscoder to mimic the SAE interface.
+
+    This allows reusing existing dashboard components that expect an SAE object.
+    It specifically provides access to the encoder and the *same-layer* decoder weights
+    for the specified layer index.
+    """
+
+    cfg: CLTWrapperConfig  # Add type hint for the config attribute
+    threshold: Optional[torch.Tensor] = (
+        None  # For JumpReLU, set by FeatureMaskingContext
+    )
+
+    def __init__(
+        self,
+        clt: "CrossLayerTranscoder",
+        layer_idx: int,
+        clt_model_dir_path: Optional[str] = None,
+    ):
+        self.clt = clt
+        self.layer_idx = layer_idx
+        self.device = clt.device
+        self.dtype = clt.dtype
+        self.hook_z_reshaping_mode = False  # Added to satisfy SAE interface
+
+        # Validate layer index
+        if not (0 <= layer_idx < clt.config.num_layers):
+            raise ValueError(
+                f"Invalid layer_idx {layer_idx} for CLT with {clt.config.num_layers} layers."
+            )
+
+        # --- Create the Wrapper Config ---
+        # Try to get model_name from the underlying clt config if it exists
+        clt_model_name = getattr(clt.config, "model_name", None)
+        clt_dataset_path = getattr(clt.config, "dataset_path", None)
+        clt_context_size = getattr(
+            clt.config, "context_size", 128
+        )  # Default to 128 if not set
+        clt_prepend_bos = getattr(clt.config, "prepend_bos", True)
+        # Use the activation_fn from CLT config for the wrapper's architecture and encode method
+        self.activation_fn = getattr(clt.config, "activation_fn", "jumprelu")
+        clt_model_from_pretrained_kwargs = getattr(
+            clt.config, "model_from_pretrained_kwargs", {}
+        )
+
+        # --- Load CLT-specific normalization stats if applicable ---
+        self.clt_norm_mean: Optional[torch.Tensor] = None
+        self.clt_norm_std: Optional[torch.Tensor] = None
+        wrapper_will_normalize_specifically = False
+        clt_norm_method = getattr(clt.config, "normalization_method", "none")
+
+        if clt_norm_method in ["auto", "estimated_mean_std", "mean_std"]:
+            if clt_model_dir_path:
+                norm_stats_file = Path(clt_model_dir_path) / "norm_stats.json"
+                if norm_stats_file.exists():
+                    try:
+                        with open(norm_stats_file, "r") as f:
+                            stats_data = json.load(f)
+
+                        layer_stats = stats_data.get(str(self.layer_idx), {}).get(
+                            "inputs", {}
+                        )
+                        mean_vals = layer_stats.get("mean")
+                        std_vals = layer_stats.get("std")
+
+                        if mean_vals is not None and std_vals is not None:
+                            self.clt_norm_mean = torch.tensor(
+                                mean_vals, device=self.device, dtype=torch.float32
+                            ).unsqueeze(0)
+                            self.clt_norm_std = (
+                                torch.tensor(
+                                    std_vals, device=self.device, dtype=torch.float32
+                                )
+                                + 1e-6
+                            ).unsqueeze(0)
+                            if torch.any(self.clt_norm_std <= 0):
+                                print(
+                                    f"Warning: Loaded std for layer {self.layer_idx} contains non-positive values after adding epsilon. Disabling specific normalization."
+                                )
+                                self.clt_norm_mean = None
+                                self.clt_norm_std = None
+                            else:
+                                wrapper_will_normalize_specifically = True
+                                print(
+                                    f"CLTLayerWrapper: Loaded norm_stats.json for layer {self.layer_idx}. Wrapper will apply specific normalization."
+                                )
+                        else:
+                            print(
+                                f"Warning: norm_stats.json found, but missing 'mean' or 'std' for layer {self.layer_idx} inputs. Wrapper will not normalize specifically."
+                            )
+                    except Exception as e:
+                        print(
+                            f"Warning: Error loading or parsing norm_stats.json from {norm_stats_file}: {e}. Wrapper will not normalize specifically."
+                        )
+                else:
+                    print(
+                        f"Warning: normalization_method is '{clt_norm_method}' but norm_stats.json not found at {norm_stats_file}. Wrapper will not normalize specifically."
+                    )
+            else:
+                print(
+                    f"Warning: normalization_method is '{clt_norm_method}' but clt_model_dir_path not provided. Wrapper cannot load norm_stats.json and will not normalize specifically."
+                )
+
+        # Determine normalize_activations flag for ActivationsStore based on CLT config and wrapper's capability
+        # This flag in self.cfg controls ActivationsStore. ActivationsStore should only normalize if the wrapper *isn't* doing specific normalization AND the CLT expected some form of normalization.
+        clt_config_indicated_normalization = clt_norm_method != "none"
+        normalize_activations_for_store = clt_config_indicated_normalization and (
+            not wrapper_will_normalize_specifically
+        )
+        if normalize_activations_for_store:
+            print(
+                f"CLTLayerWrapper: Setting normalize_activations=True for ActivationsStore (CLT method: {clt_norm_method}, wrapper specific norm: False)."
+            )
+        elif clt_config_indicated_normalization and wrapper_will_normalize_specifically:
+            print(
+                f"CLTLayerWrapper: Setting normalize_activations=False for ActivationsStore (CLT method: {clt_norm_method}, wrapper specific norm: True)."
+            )
+        else:  # not clt_config_indicated_normalization
+            print(
+                f"CLTLayerWrapper: Setting normalize_activations=False for ActivationsStore (CLT method: {clt_norm_method})."
+            )
+
+        # Initialize self.threshold if activation is jumprelu
+        # This must happen AFTER self.activation_fn, self.device, self.dtype, self.layer_idx, and self.clt are set.
+        if self.activation_fn == "jumprelu":
+            if (
+                hasattr(self.clt, "log_threshold")
+                and self.clt.log_threshold is not None
+            ):
+                if 0 <= self.layer_idx < self.clt.log_threshold.shape[0]:
+                    # The log_threshold from CLT is [num_layers, num_features]
+                    # We need the threshold for the current layer_idx
+                    layer_thresholds = torch.exp(
+                        self.clt.log_threshold[self.layer_idx].clone().detach()
+                    )
+                    self.threshold = layer_thresholds.to(
+                        device=self.device, dtype=self.dtype
+                    )
+                    print(
+                        f"CLTLayerWrapper: Initialized self.threshold for layer {self.layer_idx} from clt.log_threshold."
+                    )
+                else:
+                    print(
+                        f"Warning: CLTLayerWrapper layer_idx {self.layer_idx} is out of bounds for clt.log_threshold "
+                        f"(shape {self.clt.log_threshold.shape}). self.threshold will be None."
+                    )
+                    self.threshold = None
+            else:
+                print(
+                    f"Warning: Underlying CLT model for layer {self.layer_idx} does not have 'log_threshold' or it's None, "
+                    f"but activation_fn is 'jumprelu'. self.threshold will be None."
+                )
+                self.threshold = None
+        # else: self.threshold remains its default None, which is fine for other activation functions.
+
+        # Get the hook name using prioritized templates
+        hook_name_template = getattr(clt.config, "tl_input_template", None)
+        if hook_name_template:
+            hook_name = hook_name_template.format(layer_idx)
+            print(f"Using TL hook name template: {hook_name_template} -> {hook_name}")
+        else:
+            hook_name_template = getattr(clt.config, "mlp_input_template", None)
+            if hook_name_template:
+                hook_name = hook_name_template.format(layer_idx)
+                print(
+                    f"Warning: tl_input_template not found. Using mlp_input_template: {hook_name_template} -> {hook_name}"
+                )
+            else:
+                # Fallback for older configs without any template
+                hook_name = f"blocks.{layer_idx}.hook_mlp_in"
+                print(
+                    f"Warning: Neither tl_input_template nor mlp_input_template found. Falling back to hardcoded: {hook_name}"
+                )
+
+        self.cfg = CLTWrapperConfig(
+            d_sae=clt.config.num_features,  # This is the d_sae of the *entire* CLT layer, not a sub-batch
+            d_in=clt.config.d_model,
+            hook_name=hook_name,
+            hook_layer=layer_idx,
+            hook_head_index=None,
+            dtype=str(self.dtype).replace("torch.", ""),
+            device=str(self.device),
+            architecture=self.activation_fn,  # Use the determined activation_fn
+            model_name=clt_model_name,
+            dataset_path=clt_dataset_path,
+            context_size=clt_context_size,
+            prepend_bos=clt_prepend_bos,
+            normalize_activations=normalize_activations_for_store,
+            dataset_trust_remote_code=False,
+            seqpos_slice=None,
+            model_from_pretrained_kwargs=clt_model_from_pretrained_kwargs,
+            metadata=CLTMetadata(
+                hook_name=hook_name,
+                hook_layer=layer_idx,
+                model_name=clt_model_name,
+                context_size=clt_context_size,
+                prepend_bos=clt_prepend_bos,
+                hook_head_index=None,
+                seqpos_slice=None,
+            ),
+        )
+        # --- End Config Creation ---
+
+        # Extract and potentially gather weights
+        # Ensure weights are detached and cloned to avoid modifying the original CLT
+        # Original W_enc from CLT encoder module is [d_sae_layer, d_model]
+        # We transpose to match sae-lens W_enc convention: [d_model, d_sae_layer]
+        self.W_enc = (
+            self._gather_encoder_weight(clt.encoder_module.encoders[layer_idx].weight)  # type: ignore
+            .t()
+            .contiguous()
+        )
+        # For W_dec, use the decoder from the same layer to itself
+        decoder_key = f"{layer_idx}->{layer_idx}"
+        if decoder_key not in clt.decoder_module.decoders:  # type: ignore
+            raise KeyError(f"Decoder key {decoder_key} not found in CLT decoders.")
+        # Original W_dec from CLT decoder module is [d_model, d_sae_layer]
+        # We transpose to match sae-lens W_dec convention: [d_sae_layer, d_model]
+        self.W_dec = (
+            self._gather_decoder_weight(clt.decoder_module.decoders[decoder_key].weight)  # type: ignore
+            .t()
+            .contiguous()
+        )
+
+        self.b_enc = self._gather_encoder_bias(
+            clt.encoder_module.encoders[layer_idx].bias_param  # type: ignore
+        )
+        # For b_dec, use the bias from the same-layer decoder
+        self.b_dec = self._gather_decoder_bias(
+            clt.decoder_module.decoders[decoder_key].bias_param  # type: ignore
+        )
+
+        # Cache for folded weights if needed
+        self._W_dec_folded = False
+        # Thresholds for JumpReLU will be handled by FeatureMaskingContext if architecture is 'jumprelu'
+        # by setting self.threshold directly on the wrapper instance.
+
+    # --- Façade methods mimicking SAE --- #
+    def encode(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Encodes input using the CLTLayerWrapper's own W_enc and b_enc,
+        respecting masks applied by FeatureMaskingContext.
+        Applies the activation function specified in self.activation_fn.
+        """
+        # x is [..., d_model]
+        # self.W_enc after masking (by FeatureMaskingContext) should be [d_model, N_FEATURES_IN_BATCH]
+        # self.b_enc after masking (by FeatureMaskingContext) should be [N_FEATURES_IN_BATCH]
+
+        original_shape = x.shape
+        if x.ndim > 2:  # Ensure x is [N, d_model] for F.linear
+            # self.cfg.d_in should be d_model
+            x_reshaped = x.reshape(-1, self.cfg.d_in)
+        else:
+            x_reshaped = x
+
+        x_to_process = x_reshaped
+        # Apply CLT-specific normalization if stats were loaded
+        if self.clt_norm_mean is not None and self.clt_norm_std is not None:
+            # Ensure calculation is done in float32 for precision, then cast back
+            x_float32 = x_to_process.to(torch.float32)
+            normalized_x = (x_float32 - self.clt_norm_mean) / self.clt_norm_std
+            x_to_process = normalized_x.to(x.dtype)
+
+        # F.linear(input, weight, bias) expects weight to be [out_features, in_features]
+        # self.W_enc is [d_model, N_FEATURES_IN_BATCH], so its transpose is [N_FEATURES_IN_BATCH, d_model]
+        hidden_pre = F.linear(
+            x_to_process, self.W_enc.T, self.b_enc
+        )  # Output: [N, N_FEATURES_IN_BATCH]
+
+        # Apply activation function
+        if self.activation_fn == "relu":
+            encoded_acts = torch.relu(hidden_pre)
+        elif self.activation_fn == "jumprelu":
+            if not hasattr(self, "threshold") or self.threshold is None:
+                raise AttributeError(
+                    "JumpReLU activation selected, but 'self.threshold' is not available on CLTLayerWrapper. "
+                    "FeatureMaskingContext should set this if architecture is 'jumprelu'."
+                )
+            encoded_acts = torch.where(
+                hidden_pre > self.threshold, hidden_pre, torch.zeros_like(hidden_pre)
+            )
+        elif self.activation_fn == "batchtopk":
+            k_val: float
+            batchtopk_k_abs = getattr(self.clt.config, "batchtopk_k", None)
+            batchtopk_k_frac = getattr(self.clt.config, "batchtopk_frac", None)
+
+            if batchtopk_k_abs is not None:
+                # This k is global. For the current batch of features, we use a per-layer approximation.
+                k_val = float(batchtopk_k_abs) / self.clt.config.num_layers
+                k_val = max(
+                    1.0, k_val
+                )  # Ensure at least 1 feature is kept if k/num_layers is small
+            elif batchtopk_k_frac is not None:
+                k_val = float(
+                    batchtopk_k_frac
+                )  # Fraction applies directly to current N_FEATURES_IN_BATCH
+            else:
+                # Fallback: if neither k nor frac is specified, keep all features currently being processed.
+                # This matches the fallback in CrossLayerTranscoder.encode for its per-layer batchtopk.
+                print(
+                    f"Warning: CLTLayerWrapper using batchtopk, but neither 'batchtopk_k' nor 'batchtopk_frac' defined in CLTConfig. Defaulting to keeping all {hidden_pre.size(-1)} features in the current batch."
+                )
+                k_val = float(hidden_pre.size(-1))
+
+            straight_through_flag = getattr(
+                self.clt.config, "batchtopk_straight_through", False
+            )
+            encoded_acts = BatchTopK.apply(hidden_pre, k_val, straight_through_flag)
+        else:
+            raise ValueError(
+                f"Unsupported activation function in CLTLayerWrapper: {self.activation_fn}"
+            )
+
+        if x.ndim > 2:
+            # Reshape back to original batch/sequence dimensions, with the last dim being N_FEATURES_IN_BATCH
+            encoded_acts = encoded_acts.reshape(*original_shape[:-1], -1)  # type: ignore
+
+        return encoded_acts  # type: ignore
+
+    def turn_off_forward_pass_hook_z_reshaping(self):
+        """Stub method to satisfy SAE interface. CLTWrapper does not use this."""
+        # This mode is not applicable to CLTLayerWrapper, so this method is a no-op.
+        pass
+
+    # Note: CLTLayerWrapper does not have a separate `decode` method façade
+    # because the dashboard primarily uses W_dec directly for analysis (e.g., logits).
+    # The CLT's actual decode logic (summing across layers) isn't needed here.
+
+    def fold_W_dec_norm(self):
+        """Folds the L2 norm of W_dec into W_enc and b_enc.
+
+        Mirrors the logic in sae_lens.SAE.fold_W_dec_norm.
+        Important for ensuring that W_enc activations directly correspond
+        to the output norm when using the wrapped W_dec.
+        """
+        if self._W_dec_folded:
+            print("Warning: W_dec norm already folded.")
+            return
+
+        if self.W_dec is None or self.W_enc is None:
+            print("Warning: Cannot fold W_dec norm, weights not available.")
+            return
+
+        # Detach W_dec before calculating norm to avoid gradient issues
+        # W_dec is [N_FEATURES_IN_BATCH, d_model] (after masking context and init)
+        # Norm should be taken over d_model dim (dim=1)
+
+        # Use W_dec with its original dtype for norm calculation
+        w_dec_for_norm = self.W_dec.detach()
+        w_dec_norms = torch.norm(
+            w_dec_for_norm, dim=1, keepdim=True
+        )  # [N_FEATURES_IN_BATCH, 1]
+
+        w_dec_norms = torch.where(
+            w_dec_norms == 0, torch.ones_like(w_dec_norms), w_dec_norms
+        )
+
+        # self.W_enc is [d_model, N_FEATURES_IN_BATCH]
+        # We want to scale each column of W_enc (each feature's encoder vector)
+        # by the corresponding feature's w_dec_norm.
+        # Ensure dtypes match for multiplication, then cast W_enc back if necessary
+        original_w_enc_dtype = self.W_enc.dtype
+        self.W_enc.data = (self.W_enc.data.to(w_dec_norms.dtype) * w_dec_norms.t()).to(
+            original_w_enc_dtype
+        )
+
+        if self.b_enc is not None:
+            # self.b_enc is [N_FEATURES_IN_BATCH]
+            # w_dec_norms.squeeze() is [N_FEATURES_IN_BATCH]
+            original_b_enc_dtype = self.b_enc.dtype
+            self.b_enc.data = (
+                self.b_enc.data.to(w_dec_norms.dtype) * w_dec_norms.squeeze()
+            ).to(original_b_enc_dtype)
+
+        # Store the norms for potential unfolding or reference
+        self._w_dec_norms_backup = w_dec_norms
+        self._W_dec_folded = True
+        print("Folded W_dec norm into W_enc and b_enc.")
+
+    def unfold_W_dec_norm(self):
+        """Unfolds the L2 norm of W_dec from W_enc and b_enc."""
+        if not self._W_dec_folded or not hasattr(self, "_w_dec_norms_backup"):
+            print("Warning: W_dec norm not folded or backup norms not found.")
+            return
+
+        if self.W_enc is None:
+            print("Warning: Cannot unfold W_dec norm, W_enc not available.")
+            return
+
+        # Retrieve the norms used for folding
+        w_dec_norms = self._w_dec_norms_backup
+        # Avoid division by zero (should have been handled in fold, but double check)
+        w_dec_norms = torch.where(
+            w_dec_norms == 0, torch.ones_like(w_dec_norms), w_dec_norms
+        )
+
+        original_w_enc_dtype = self.W_enc.dtype
+        self.W_enc.data = (self.W_enc.data.to(w_dec_norms.dtype) / w_dec_norms.t()).to(
+            original_w_enc_dtype
+        )
+
+        if self.b_enc is not None:
+            original_b_enc_dtype = self.b_enc.dtype
+            self.b_enc.data = (
+                self.b_enc.data.to(w_dec_norms.dtype) / w_dec_norms.squeeze()
+            ).to(original_b_enc_dtype)
+
+        del self._w_dec_norms_backup
+        self._W_dec_folded = False
+        print("Unfolded W_dec norm from W_enc and b_enc.")
+
+    def to(self, device: Union[str, torch.device]):
+        """Moves the wrapper and underlying components to the specified device."""
+        target_device = torch.device(device)
+
+        # Move the underlying CLT model
+        try:
+            self.clt.to(target_device)
+        except Exception as e:
+            print(
+                f"Warning: Failed to move underlying CLT model to {target_device}: {e}"
+            )
+            # Continue trying to move wrapper components
+
+        # Move the wrapper's stored tensors
+        if self.W_enc is not None:
+            self.W_enc = self.W_enc.to(target_device)
+        if self.W_dec is not None:
+            self.W_dec = self.W_dec.to(target_device)
+        if self.b_enc is not None:
+            self.b_enc = self.b_enc.to(target_device)
+        if self.b_dec is not None:
+            self.b_dec = self.b_dec.to(target_device)
+        if (
+            hasattr(self, "_w_dec_norms_backup")
+            and self._w_dec_norms_backup is not None
+        ):
+            self._w_dec_norms_backup = self._w_dec_norms_backup.to(target_device)
+
+        # Update device attributes
+        self.device = target_device
+        self.cfg.device = str(target_device)
+
+        # Update activation_fn related thresholds if they exist (e.g. for JumpReLU)
+        if hasattr(self, "threshold") and self.threshold is not None:
+            self.threshold = self.threshold.to(target_device)
+
+        if self.clt_norm_mean is not None:  # Added to move norm stats
+            self.clt_norm_mean = self.clt_norm_mean.to(target_device)
+        if self.clt_norm_std is not None:  # Added to move norm stats
+            self.clt_norm_std = self.clt_norm_std.to(target_device)
+
+        print(f"Moved CLTLayerWrapper to {target_device}")
+        return self
+
+    # --- Helper methods for Tensor Parallelism --- #
+
+    def _gather_weight(
+        self,
+        weight_shard: torch.Tensor,
+        gather_dim: int = 0,
+        target_full_dim_size: Optional[int] = None,
+    ) -> torch.Tensor:
+        """Gather a weight tensor shard across TP ranks."""
+        if not dist.is_initialized() or dist.get_world_size() == 1:
+            return weight_shard.clone().detach()
+
+        world_size = dist.get_world_size()
+        # Create a list to hold all gathered tensors
+        tensor_list = [torch.empty_like(weight_shard) for _ in range(world_size)]
+        dist.all_gather(tensor_list, weight_shard)
+
+        # Concatenate along the specified dimension
+        full_weight = torch.cat(tensor_list, dim=gather_dim)
+
+        # Trim padding if necessary
+        if target_full_dim_size is not None:
+            if gather_dim == 0:
+                if full_weight.shape[0] > target_full_dim_size:
+                    full_weight = full_weight[:target_full_dim_size, :]
+            elif gather_dim == 1:
+                if full_weight.shape[1] > target_full_dim_size:
+                    full_weight = full_weight[:, :target_full_dim_size]
+            # Add other gather_dim cases if needed
+
+        return full_weight.detach()
+
+    def _gather_encoder_weight(self, weight_shard: torch.Tensor) -> torch.Tensor:
+        """Gather ColumnParallelLinear weight (sharded along output/feature dim)."""
+        # ColumnParallel weight is [d_sae_local, d_model]
+        # We need to gather along dim 0 to get [d_sae_full_for_layer, d_model]
+        return self._gather_weight(
+            weight_shard,
+            gather_dim=0,
+            target_full_dim_size=self.clt.config.num_features,
+        )
+
+    def _gather_decoder_weight(self, weight_shard: torch.Tensor) -> torch.Tensor:
+        """Gather RowParallelLinear weight (sharded along input/feature dim)."""
+        # RowParallel weight is [d_model, d_sae_local]
+        # We need to gather along dim 1 to get [d_model, d_sae_full_for_layer]
+        return self._gather_weight(
+            weight_shard,
+            gather_dim=1,
+            target_full_dim_size=self.clt.config.num_features,
+        )
+
+    def _gather_bias(
+        self,
+        bias_shard: Optional[torch.Tensor],
+        gather_dim: int = 0,
+        target_full_dim_size: Optional[int] = None,
+    ) -> Optional[torch.Tensor]:
+        """Gather a bias tensor shard across TP ranks."""
+        if bias_shard is None:
+            return None
+        # Biases are typically sharded along the same dimension as the weight's corresponding output dim
+        return self._gather_weight(
+            bias_shard, gather_dim=gather_dim, target_full_dim_size=target_full_dim_size
+        )
+
+    def _gather_encoder_bias(
+        self, bias_shard_candidate: Optional[torch.Tensor]
+    ) -> Optional[torch.Tensor]:
+        """Gather ColumnParallelLinear bias (sharded along output/feature dim).
+
+        Defensively checks if the provided candidate is actually a Tensor.
+        """
+        # Check if the provided object is a Tensor
+        if isinstance(bias_shard_candidate, torch.Tensor):
+            # Encoder bias shape [d_sae_local], gather along dim 0
+            return self._gather_bias(
+                bias_shard_candidate,
+                gather_dim=0,
+                target_full_dim_size=self.clt.config.num_features,
+            )
+        else:
+            # If it's None, bool, or anything else, treat as no bias
+            return None
+
+    def _gather_decoder_bias(
+        self, bias_shard_candidate: Optional[torch.Tensor]
+    ) -> Optional[torch.Tensor]:
+        """Gather RowParallelLinear bias (NOT sharded, but might need broadcast/check).
+
+        Defensively checks if the provided candidate is actually a Tensor.
+        """
+        # Check if the provided object is a Tensor
+        if isinstance(bias_shard_candidate, torch.Tensor):
+            # RowParallelLinear bias is typically not sharded (added after all-reduce)
+            # However, let's check world size and return a clone if TP=1, or verify replication if TP>1
+            if not dist.is_initialized() or dist.get_world_size() == 1:
+                return bias_shard_candidate.clone().detach()
+
+            # In TP > 1, the bias should be identical across ranks. Verify this.
+            world_size = dist.get_world_size()
+            tensor_list = [
+                torch.empty_like(bias_shard_candidate) for _ in range(world_size)
+            ]
+            dist.all_gather(tensor_list, bias_shard_candidate)
+            # Check if all gathered biases are the same
+            for i in range(1, world_size):
+                if not torch.equal(tensor_list[0], tensor_list[i]):
+                    raise RuntimeError(
+                        "RowParallelLinear bias shards are not identical across TP ranks, which is unexpected."
+                    )
+            # Return the bias from rank 0 (or any rank, as they are identical)
+            return tensor_list[0].clone().detach()
+        else:
+            # If it's None, bool, or anything else, treat as no bias
+            return None

SAEDashboard/sae_dashboard/components.py

@@ -0,0 +1,774 @@
+from copy import deepcopy
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Callable, List
+
+import numpy as np
+from dataclasses_json import dataclass_json
+
+from sae_dashboard.components_config import (
+    ActsHistogramConfig,
+    FeatureTablesConfig,
+    LogitsHistogramConfig,
+    LogitsTableConfig,
+    PromptConfig,
+    SequencesConfig,
+)
+from sae_dashboard.html_fns import HTML, bgColorMap, uColorMap
+from sae_dashboard.utils_fns import (
+    HistogramData,
+    max_or_1,
+    to_str_tokens,
+    unprocess_str_tok,
+)
+
+PRECISION = 4
+
+
+@dataclass
+class DecoderWeightsDistribution:
+    n_heads: int
+    allocation_by_head: List[float]
+
+
+@dataclass_json
+@dataclass
+class FeatureTablesData:
+    """
+    This contains all the data necessary to make the left-hand tables in prompt-centric visualization. See diagram
+    in readme:
+
+        https://github.com/callummcdougall/sae_vis#data_storing_fnspy
+
+    Inputs:
+        neuron_alignment...
+            The data for the neuron alignment table (each of its 3 columns). In other words, the data containing which
+            neurons in the transformer the encoder feature is most aligned with.
+
+        correlated_neurons...
+            The data for the correlated neurons table (each of its 3 columns). In other words, the data containing which
+            neurons in the transformer are most correlated with the encoder feature.
+
+        correlated_features...
+            The data for the correlated features table (each of its 3 columns). In other words, the data containing
+            which features in this encoder are most correlated with each other.
+
+        correlated_b_features...
+            The data for the correlated features table (each of its 3 columns). In other words, the data containing
+            which features in encoder-B are most correlated with those in the original encoder. Note, this one might be
+            absent if we're not using a B-encoder.
+    """
+
+    neuron_alignment_indices: list[int] = field(default_factory=list)
+    neuron_alignment_values: list[float] = field(default_factory=list)
+    neuron_alignment_l1: list[float] = field(default_factory=list)
+    correlated_neurons_indices: list[int] = field(default_factory=list)
+    correlated_neurons_pearson: list[float] = field(default_factory=list)
+    correlated_neurons_cossim: list[float] = field(default_factory=list)
+    correlated_features_indices: list[int] = field(default_factory=list)
+    correlated_features_pearson: list[float] = field(default_factory=list)
+    correlated_features_cossim: list[float] = field(default_factory=list)
+    correlated_b_features_indices: list[int] = field(default_factory=list)
+    correlated_b_features_pearson: list[float] = field(default_factory=list)
+    correlated_b_features_cossim: list[float] = field(default_factory=list)
+
+    def _get_html_data(
+        self,
+        cfg: FeatureTablesConfig,
+        decode_fn: Callable[[int | list[int]], str | list[str]],
+        id_suffix: str,
+        column: int | tuple[int, int],
+        component_specific_kwargs: dict[str, Any] = {},
+    ) -> HTML:
+        """
+        Returns the HTML for the left-hand tables, wrapped in a 'grid-column' div.
+
+        Note, we only ever use this obj in the context of the left-hand column of the feature-centric vis, and it's
+        always the same width & height, which is why there's no customization available for this function.
+        """
+        # Read HTML from file, and replace placeholders with real ID values
+        html_str = (
+            Path(__file__).parent / "html" / "feature_tables_template.html"
+        ).read_text()
+        html_str = html_str.replace("FEATURE_TABLES_ID", f"feature-tables-{id_suffix}")
+
+        # Create dictionary storing the data
+        data: dict[str, list[dict[str, str | float]]] = {}
+
+        # Store the neuron alignment data, if it exists
+        if len(self.neuron_alignment_indices) > 0:
+            assert len(self.neuron_alignment_indices) >= cfg.n_rows, "Not enough rows!"
+            data["neuronAlignment"] = [
+                {
+                    "index": index,
+                    "value": f"{value:+.3f}",
+                    "percentageL1": f"{percent_l1:.1%}",
+                }
+                for index, value, percent_l1 in zip(
+                    self.neuron_alignment_indices,
+                    self.neuron_alignment_values,
+                    self.neuron_alignment_l1,
+                )
+            ]
+
+        # Store the other 3, if they exist (they're all in the same format, so we can do it in a for loop)
+        for name, js_name in zip(
+            ["correlated_neurons", "correlated_features", "correlated_b_features"],
+            ["correlatedNeurons", "correlatedFeatures", "correlatedFeaturesB"],
+        ):
+            if len(getattr(self, f"{name}_indices")) > 0:
+                # assert len(getattr(self, f"{name}_indices")) >= cfg.n_rows, "Not enough rows!"
+                data[js_name] = [
+                    {
+                        "index": index,
+                        "value": f"{value:+.3f}",
+                        "percentageL1": f"{percent_L1:+.3f}",
+                    }
+                    for index, value, percent_L1 in zip(
+                        getattr(self, f"{name}_indices")[: cfg.n_rows],
+                        getattr(self, f"{name}_pearson")[: cfg.n_rows],
+                        getattr(self, f"{name}_cossim")[: cfg.n_rows],
+                    )
+                ]
+
+        return HTML(
+            html_data={column: html_str},
+            js_data={"featureTablesData": {id_suffix: data}},
+        )
+
+
+@dataclass_json
+@dataclass
+class ActsHistogramData(HistogramData):
+    def _get_html_data(
+        self,
+        cfg: ActsHistogramConfig,
+        decode_fn: Callable[[int | list[int]], str | list[str]],
+        id_suffix: str,
+        column: int | tuple[int, int],
+        component_specific_kwargs: dict[str, Any] = {},
+    ) -> HTML:
+        """
+        Converts data -> HTML object, for the feature activations histogram (i.e. the histogram over all sampled tokens,
+        showing the distribution of activations for this feature).
+        """
+        # We can't post-hoc change the number of bins, so check this wasn't changed in the config
+        # assert cfg.n_bins == len(self.bar_heights),\
+        #     "Can't post-hoc change `n_bins` in histogram config - you need to regenerate data."
+
+        # Read HTML from file, and replace placeholders with real ID values
+        html_str = (
+            Path(__file__).parent / "html" / "acts_histogram_template.html"
+        ).read_text()
+        html_str = html_str.replace("HISTOGRAM_ACTS_ID", f"histogram-acts-{id_suffix}")
+
+        # Process colors for frequency histogram; it's darker at higher values
+        bar_values_normed = [
+            (0.4 * max(self.bar_values) + 0.6 * v)
+            / max(max(self.bar_values), 1e-6)  # avoid divide by zero
+            for v in self.bar_values
+        ]
+        bar_colors = [bgColorMap(v) for v in bar_values_normed]
+
+        # Next we create the data dict
+        data: dict[str, Any] = {
+            "y": self.bar_heights,
+            "x": self.bar_values,
+            "ticks": self.tick_vals,
+            "colors": bar_colors,
+        }
+        if self.title is not None:
+            data["title"] = self.title
+
+        return HTML(
+            html_data={column: html_str},
+            js_data={"actsHistogramData": {id_suffix: data}},
+        )
+
+
+@dataclass_json
+@dataclass
+class LogitsHistogramData(HistogramData):
+    def _get_html_data(
+        self,
+        cfg: LogitsHistogramConfig,
+        decode_fn: Callable[[int | list[int]], str | list[str]],
+        id_suffix: str,
+        column: int | tuple[int, int],
+        component_specific_kwargs: dict[str, Any] = {},
+    ) -> HTML:
+        """
+        Converts data -> HTML object, for the logits histogram (i.e. the histogram over all tokens in the vocab, showing
+        the distribution of direct logit effect on that token).
+        """
+        # We can't post-hoc change the number of bins, so check this wasn't changed in the config
+        # assert cfg.n_bins == len(self.bar_heights),\
+        #     "Can't post-hoc change `n_bins` in histogram config - you need to regenerate data."
+
+        # Read HTML from file, and replace placeholders with real ID values
+        html_str = (
+            Path(__file__).parent / "html" / "logits_histogram_template.html"
+        ).read_text()
+        html_str = html_str.replace(
+            "HISTOGRAM_LOGITS_ID", f"histogram-logits-{id_suffix}"
+        )
+
+        data: dict[str, Any] = {
+            "y": self.bar_heights,
+            "x": self.bar_values,
+            "ticks": self.tick_vals,
+        }
+        if self.title is not None:
+            data["title"] = self.title
+
+        return HTML(
+            html_data={column: html_str},
+            js_data={"logitsHistogramData": {id_suffix: data}},
+        )
+
+
+@dataclass_json
+@dataclass
+class LogitsTableData:
+    bottom_token_ids: list[int] = field(default_factory=list)
+    bottom_logits: list[float] = field(default_factory=list)
+    top_token_ids: list[int] = field(default_factory=list)
+    top_logits: list[float] = field(default_factory=list)
+
+    def _get_html_data(
+        self,
+        cfg: LogitsTableConfig,
+        decode_fn: Callable[[int | list[int]], str | list[str]],
+        id_suffix: str,
+        column: int | tuple[int, int],
+        component_specific_kwargs: dict[str, Any] = {},
+    ) -> HTML:
+        """
+        Converts data -> HTML object, for the logits table (i.e. the top and bottom affected tokens by this feature).
+        """
+        # Crop the lists to `cfg.n_rows` (first checking the config doesn't ask for more rows than we have)
+        assert cfg.n_rows <= len(self.bottom_logits)
+        bottom_token_ids = self.bottom_token_ids[: cfg.n_rows]
+        bottom_logits = self.bottom_logits[: cfg.n_rows]
+        top_token_ids = self.top_token_ids[: cfg.n_rows]
+        top_logits = self.top_logits[: cfg.n_rows]
+
+        # Get the negative and positive background values (darkest when equals max abs)
+        max_value = max(
+            max(top_logits[: cfg.n_rows]), -min(bottom_logits[: cfg.n_rows])
+        )
+        neg_bg_values = np.absolute(bottom_logits[: cfg.n_rows]) / max_value
+        pos_bg_values = np.absolute(top_logits[: cfg.n_rows]) / max_value
+
+        # Get the string tokens, using the decode function
+        neg_str = to_str_tokens(decode_fn, bottom_token_ids[: cfg.n_rows])
+        pos_str = to_str_tokens(decode_fn, top_token_ids[: cfg.n_rows])
+
+        # Read HTML from file, and replace placeholders with real ID values
+        html_str = (
+            Path(__file__).parent / "html" / "logits_table_template.html"
+        ).read_text()
+        html_str = html_str.replace("LOGITS_TABLE_ID", f"logits-table-{id_suffix}")
+
+        # Create object for storing JS data
+        data: dict[str, list[dict[str, str | float]]] = {
+            "negLogits": [],
+            "posLogits": [],
+        }
+
+        # Get data for the tables of pos/neg logits
+        for i in range(len(neg_str)):
+            data["negLogits"].append(
+                {
+                    "symbol": unprocess_str_tok(neg_str[i]),
+                    "value": round(bottom_logits[i], 2),
+                    "color": f"rgba(255,{int(255*(1-neg_bg_values[i]))},{int(255*(1-neg_bg_values[i]))},0.5)",
+                }
+            )
+            data["posLogits"].append(
+                {
+                    "symbol": unprocess_str_tok(pos_str[i]),
+                    "value": round(top_logits[i], 2),
+                    "color": f"rgba({int(255*(1-pos_bg_values[i]))},{int(255*(1-pos_bg_values[i]))},255,0.5)",
+                }
+            )
+
+        return HTML(
+            html_data={column: html_str},
+            js_data={"logitsTableData": {id_suffix: data}},
+        )
+
+
+@dataclass_json
+@dataclass
+class SequenceData:
+    """
+    This contains all the data necessary to make a sequence of tokens in the vis. See diagram in readme:
+
+        https://github.com/callummcdougall/sae_vis#data_storing_fnspy
+
+    Always-visible data:
+        token_ids:              List of token IDs in the sequence
+        feat_acts:              Sizes of activations on this sequence
+        loss_contribution:   Effect on loss of this feature, for this particular token (neg = helpful)
+
+    Data which is visible on hover:
+        token_logits:       The logits of the particular token in that sequence (used for line on logits histogram)
+        top_token_ids:     List of the top 5 logit-boosted tokens by this feature
+        top_logits:        List of the corresponding 5 changes in logits for those tokens
+        bottom_token_ids:  List of the bottom 5 logit-boosted tokens by this feature
+        bottom_logits:     List of the corresponding 5 changes in logits for those tokens
+    """
+
+    original_index: int = 0
+    qualifying_token_index: int = 0
+    token_ids: list[int] = field(default_factory=list)
+    feat_acts: list[float] = field(default_factory=list)
+    loss_contribution: list[float] = field(default_factory=list)
+
+    token_logits: list[float] = field(default_factory=list)
+    top_token_ids: list[list[int]] = field(default_factory=list)
+    top_logits: list[list[float]] = field(default_factory=list)
+    bottom_token_ids: list[list[int]] = field(default_factory=list)
+    bottom_logits: list[list[float]] = field(default_factory=list)
+
+    def __post_init__(self) -> None:
+        """
+        Filters the logits & token IDs by removing any elements which are zero (this saves space in the eventual
+        JavaScript).
+        """
+        self.seq_len = len(self.token_ids)
+        self.top_logits, self.top_token_ids = self._filter(
+            self.top_logits, self.top_token_ids
+        )
+        self.bottom_logits, self.bottom_token_ids = self._filter(
+            self.bottom_logits, self.bottom_token_ids
+        )
+
+    def _filter(
+        self, float_list: list[list[float]], int_list: list[list[int]]
+    ) -> tuple[list[list[float]], list[list[int]]]:
+        """
+        Filters the list of floats and ints, by removing any elements which are zero. Note - the absolute values of the
+        floats are monotonic non-increasing, so we can assume that all the elements we keep will be the first elements
+        of their respective lists. Also reduces precisions of feature activations & logits.
+        """
+        # Next, filter out zero-elements and reduce precision
+        float_list = [
+            [round(f, PRECISION) for f in floats if abs(f) > 1e-6]
+            for floats in float_list
+        ]
+        int_list = [ints[: len(floats)] for ints, floats in zip(int_list, float_list)]
+        return float_list, int_list
+
+    def _get_html_data(
+        self,
+        cfg: PromptConfig | SequencesConfig,
+        decode_fn: Callable[[int | list[int]], str | list[str]],
+        id_suffix: str,
+        column: int | tuple[int, int],
+        component_specific_kwargs: dict[str, Any] = {},
+    ) -> HTML:
+        """
+        Args:
+
+        Returns:
+            js_data: list[dict[str, Any]]
+                The data for this sequence, in the form of a list of dicts for each token (where the dict stores things
+                like token, feature activations, etc).
+        """
+        assert isinstance(
+            cfg, (PromptConfig, SequencesConfig)
+        ), f"Invalid config type: {type(cfg)}"
+        seq_group_id = component_specific_kwargs.get("seq_group_id", None)
+        max_feat_act = component_specific_kwargs.get("max_feat_act", None)
+        max_loss_contribution = component_specific_kwargs.get(
+            "max_loss_contribution", None
+        )
+        bold_idx = component_specific_kwargs.get("bold_idx", None)
+        permanent_line = component_specific_kwargs.get("permanent_line", False)
+        first_in_group = component_specific_kwargs.get("first_in_group", True)
+        title = component_specific_kwargs.get("title", None)
+        hover_above = component_specific_kwargs.get("hover_above", False)
+
+        # If we didn't supply a sequence group ID, then we assume this sequence is on its own, and give it a unique ID
+        if seq_group_id is None:
+            seq_group_id = f"prompt-{column:03d}"
+
+        # If we didn't specify bold_idx, then set it to be the midpoint
+        if bold_idx is None:
+            bold_idx = self.seq_len // 2
+
+        # If we only have data for the bold token, we pad out everything with zeros or empty lists
+        only_bold = isinstance(cfg, SequencesConfig) and not (cfg.compute_buffer)
+        if only_bold:
+            assert bold_idx != "max", "Don't know how to deal with this case yet."
+            feat_acts = [
+                self.feat_acts[0] if (i == bold_idx) else 0.0
+                for i in range(self.seq_len)
+            ]
+            loss_contribution = [
+                self.loss_contribution[0] if (i == bold_idx) + 1 else 0.0
+                for i in range(self.seq_len)
+            ]
+            pos_ids = [
+                self.top_token_ids[0] if (i == bold_idx) + 1 else []
+                for i in range(self.seq_len)
+            ]
+            neg_ids = [
+                self.bottom_token_ids[0] if (i == bold_idx) + 1 else []
+                for i in range(self.seq_len)
+            ]
+            pos_val = [
+                self.top_logits[0] if (i == bold_idx) + 1 else []
+                for i in range(self.seq_len)
+            ]
+            neg_val = [
+                self.bottom_logits[0] if (i == bold_idx) + 1 else []
+                for i in range(self.seq_len)
+            ]
+        else:
+            feat_acts = deepcopy(self.feat_acts)
+            loss_contribution = deepcopy(self.loss_contribution)
+            pos_ids = deepcopy(self.top_token_ids)
+            neg_ids = deepcopy(self.bottom_token_ids)
+            pos_val = deepcopy(self.top_logits)
+            neg_val = deepcopy(self.bottom_logits)
+
+        # EXPERIMENT: let's just hardcode everything except feature acts to be 0's for now.
+        loss_contribution = [0.0 for _ in range(self.seq_len)]
+        pos_ids = [[] for _ in range(self.seq_len)]
+        neg_ids = [[] for _ in range(self.seq_len)]
+        pos_val = [[] for _ in range(self.seq_len)]
+        neg_val = [[] for _ in range(self.seq_len)]
+        ### END EXPERIMENT
+
+        # Get values for converting into colors later
+        bg_denom = max_feat_act or max_or_1(self.feat_acts)
+        u_denom = max_loss_contribution or max_or_1(self.loss_contribution, abs=True)
+        bg_values = (np.maximum(feat_acts, 0.0) / max(1e-4, bg_denom)).tolist()
+        u_values = (np.array(loss_contribution) / max(1e-4, u_denom)).tolist()
+
+        # If we sent in a prompt rather than this being sliced from a longer sequence, then the pos_ids etc will be shorter
+        # than the token list by 1, so we need to pad it at the first token
+        if isinstance(cfg, PromptConfig):
+            assert (
+                len(pos_ids)
+                == len(neg_ids)
+                == len(pos_val)
+                == len(neg_val)
+                == len(self.token_ids) - 1
+            ), "If this is a single prompt, these lists must be the same length as token_ids or 1 less"
+            pos_ids = [[]] + pos_ids
+            neg_ids = [[]] + neg_ids
+            pos_val = [[]] + pos_val
+            neg_val = [[]] + neg_val
+
+        assert (
+            len(pos_ids)
+            == len(neg_ids)
+            == len(pos_val)
+            == len(neg_val)
+            == len(self.token_ids)
+        ), "If this is part of a sequence group etc are given, they must be the same length as token_ids"
+
+        # Process the tokens to get str toks
+        toks = to_str_tokens(decode_fn, self.token_ids)
+        pos_toks = [to_str_tokens(decode_fn, pos) for pos in pos_ids]
+        neg_toks = [to_str_tokens(decode_fn, neg) for neg in neg_ids]
+
+        # Define the JavaScript object which will be used to populate the HTML string
+        js_data_list = []
+
+        for i in range(len(self.token_ids)):
+            # We might store a bunch of different case-specific data in the JavaScript object for each token. This is
+            # done in the form of a disjoint union over different dictionaries (which can each be empty or not), this
+            # minimizes the size of the overall JavaScript object. See function in `tokens_script.js` for more.
+            kwargs_bold: dict[str, bool] = {}
+            kwargs_hide: dict[str, bool] = {}
+            kwargs_this_token_active: dict[str, Any] = {}
+            kwargs_prev_token_active: dict[str, Any] = {}
+            kwargs_hover_above: dict[str, bool] = {}
+
+            # Get args if this is the bolded token (we make it bold, and maybe add permanent line to histograms)
+            if bold_idx is not None:
+                kwargs_bold["isBold"] = (bold_idx == i) or (
+                    bold_idx == "max" and i == np.argmax(feat_acts).item()
+                )
+                if kwargs_bold["isBold"] and permanent_line:
+                    kwargs_bold["permanentLine"] = True
+
+            # If we only have data for the bold token, we hide all other tokens' hoverdata (and skip other kwargs)
+            if (
+                only_bold
+                and isinstance(bold_idx, int)
+                and (i not in {bold_idx, bold_idx + 1})
+            ):
+                kwargs_hide["hide"] = True
+
+            else:
+                # Get args if we're making the tooltip hover above token (default is below)
+                if hover_above:
+                    kwargs_hover_above["hoverAbove"] = True
+
+                # If feature active on this token, get background color and feature act (for hist line)
+                if abs(feat_acts[i]) > 1e-8:
+                    kwargs_this_token_active = dict(
+                        featAct=round(feat_acts[i], PRECISION),
+                        bgColor=bgColorMap(bg_values[i]),
+                    )
+
+                # If prev token active, get the top/bottom logits table, underline color, and loss effect (for hist line)
+                pos_toks_i, neg_toks_i, pos_val_i, neg_val_i = (
+                    pos_toks[i],
+                    neg_toks[i],
+                    pos_val[i],
+                    neg_val[i],
+                )
+                if len(pos_toks_i) + len(neg_toks_i) > 0:
+                    # Create dictionary
+                    kwargs_prev_token_active = dict(
+                        posToks=pos_toks_i,
+                        negToks=neg_toks_i,
+                        posVal=pos_val_i,
+                        negVal=neg_val_i,
+                        lossEffect=round(loss_contribution[i], PRECISION),
+                        uColor=uColorMap(u_values[i]),
+                    )
+
+            js_data_list.append(
+                dict(
+                    tok=unprocess_str_tok(toks[i]),
+                    tokID=self.token_ids[i],
+                    tokenLogit=round(self.token_logits[i], PRECISION),
+                    **kwargs_bold,
+                    **kwargs_this_token_active,
+                    **kwargs_prev_token_active,
+                    **kwargs_hover_above,
+                )
+            )
+
+        # Create HTML string (empty by default since sequences are added by JavaScript) and JS data
+        html_str = ""
+        js_seq_group_data: dict[str, Any] = {"data": [js_data_list]}
+
+        # Add group-specific stuff if this is the first sequence in the group
+        if first_in_group:
+            # Read HTML from file, replace placeholders with real ID values
+            html_str = (
+                Path(__file__).parent / "html" / "sequences_group_template.html"
+            ).read_text()
+            html_str = html_str.replace("SEQUENCE_GROUP_ID", seq_group_id)
+
+            # Get title of sequence group, and the idSuffix to match up with a histogram
+            js_seq_group_data["idSuffix"] = id_suffix
+            if title is not None:
+                js_seq_group_data["title"] = title
+
+        return HTML(
+            html_data={column: html_str},
+            js_data={"tokenData": {seq_group_id: js_seq_group_data}},
+        )
+
+
+@dataclass_json
+@dataclass
+class SequenceGroupData:
+    """
+    This contains all the data necessary to make a single group of sequences (e.g. a quantile in prompt-centric
+    visualization). See diagram in readme:
+
+        https://github.com/callummcdougall/sae_vis#data_storing_fnspy
+
+    Inputs:
+        title:      The title that this sequence group will have, if any. This is used in `_get_html_data`. The titles
+                    will actually be in the HTML strings, not in the JavaScript data.
+        seq_data:   The data for the sequences in this group.
+    """
+
+    title: str = ""
+    seq_data: list[SequenceData] = field(default_factory=list)
+
+    def __len__(self) -> int:
+        return len(self.seq_data)
+
+    @property
+    def max_feat_act(self) -> float:
+        """Returns maximum value of feature activation over all sequences in this group."""
+        return max_or_1([act for seq in self.seq_data for act in seq.feat_acts])
+
+    @property
+    def max_loss_contribution(self) -> float:
+        """Returns maximum value of loss contribution over all sequences in this group."""
+        return max_or_1(
+            [loss for seq in self.seq_data for loss in seq.loss_contribution], abs=True
+        )
+
+    def _get_html_data(
+        self,
+        cfg: SequencesConfig,
+        decode_fn: Callable[[int | list[int]], str | list[str]],
+        id_suffix: str,
+        column: int | tuple[int, int],
+        component_specific_kwargs: dict[str, Any] = {},
+        # These default values should be correct when we only have one sequence group, because when we call this from
+        # a SequenceMultiGroupData we'll override them)
+    ) -> HTML:
+        """
+        This creates a single group of sequences, i.e. title plus some number of vertically stacked sequences.
+
+        Note, `column` is treated specially here, because the col might overflow (hence colulmn could be a tuple).
+
+        Args (from component-specific kwargs):
+            seq_group_id:   The id of the sequence group div. This will usually be passed as e.g. "seq-group-001".
+            group_size:     Max size of sequences in the group (i.e. we truncate after this many, if argument supplied).
+            max_feat_act:   If supplied, then we use this as the most extreme value (for coloring by feature act).
+
+        Returns:
+            html_obj:       Object containing the HTML and JavaScript data for this seq group.
+        """
+        seq_group_id = component_specific_kwargs.get("seq_group_id", None)
+        group_size = component_specific_kwargs.get("group_size", None)
+        max_feat_act = component_specific_kwargs.get("max_feat_act", self.max_feat_act)
+        max_loss_contribution = component_specific_kwargs.get(
+            "max_loss_contribution", self.max_loss_contribution
+        )
+
+        # Get the data that will go into the div (list of list of dicts, i.e. containing all data for seqs in group). We
+        # start with the title.
+        html_obj = HTML()
+
+        # If seq_group_id is not supplied, then we assume this is the only sequence in the column, and we name the group
+        # after the column
+        if seq_group_id is None:
+            seq_group_id = f"seq-group-{column:03d}"
+
+        # Accumulate the HTML data for each sequence in this group
+        for i, seq in enumerate(self.seq_data[:group_size]):
+            html_obj += seq._get_html_data(
+                cfg=cfg,
+                # pass in a PromptConfig object
+                decode_fn=decode_fn,
+                id_suffix=id_suffix,
+                column=column,
+                component_specific_kwargs=dict(
+                    bold_idx="max" if cfg.buffer is None else cfg.buffer[0],
+                    permanent_line=False,  # in a group, we're never showing a permanent line (only for single seqs)
+                    max_feat_act=max_feat_act,
+                    max_loss_contribution=max_loss_contribution,
+                    seq_group_id=seq_group_id,
+                    first_in_group=(i == 0),
+                    title=self.title,
+                ),
+            )
+
+        return html_obj
+
+
+@dataclass_json
+@dataclass
+class SequenceMultiGroupData:
+    """
+    This contains all the data necessary to make multiple groups of sequences (e.g. the different quantiles in the
+    prompt-centric visualization). See diagram in readme:
+
+        https://github.com/callummcdougall/sae_vis#data_storing_fnspy
+    """
+
+    seq_group_data: list[SequenceGroupData] = field(default_factory=list)
+
+    def __getitem__(self, idx: int) -> SequenceGroupData:
+        return self.seq_group_data[idx]
+
+    @property
+    def max_feat_act(self) -> float:
+        """Returns maximum value of feature activation over all sequences in this group."""
+        return max_or_1([seq_group.max_feat_act for seq_group in self.seq_group_data])
+
+    @property
+    def max_loss_contribution(self) -> float:
+        """Returns maximum value of loss contribution over all sequences in this group."""
+        return max_or_1(
+            [seq_group.max_loss_contribution for seq_group in self.seq_group_data]
+        )
+
+    def _get_html_data(
+        self,
+        cfg: SequencesConfig,
+        decode_fn: Callable[[int | list[int]], str | list[str]],
+        id_suffix: str,
+        column: int | tuple[int, int],
+        component_specific_kwargs: dict[str, Any] = {},
+    ) -> HTML:
+        """
+        Args:
+            decode_fn:                  Mapping from token IDs to string tokens.
+            id_suffix:                  The suffix for the ID of the div containing the sequences.
+            column:                     The index of this column. Note that this will be an int, but we might end up
+                                        turning it into a tuple if we overflow into a new column.
+            component_specific_kwargs:  Contains any specific kwargs that could be used to customize this component.
+
+        Returns:
+            html_obj:  Object containing the HTML and JavaScript data for these multiple seq groups.
+        """
+        assert isinstance(column, int)
+
+        # Get max activation value & max loss contributions, over all sequences in all groups
+        max_feat_act = component_specific_kwargs.get("max_feat_act", self.max_feat_act)
+        max_loss_contribution = component_specific_kwargs.get(
+            "max_loss_contribution", self.max_loss_contribution
+        )
+
+        # Get the correct column indices for the sequence groups, depending on how group_wrap is configured. Note, we
+        # deal with overflowing columns by extending the dictionary, i.e. our column argument isn't just `column`, but
+        # is a tuple of `(column, x)` where `x` is the number of times we've overflowed. For instance, if we have mode
+        # 'stack-none' then our columns are `(column, 0), (column, 1), (column, 1), (column, 1), (column, 2), ...`
+        n_groups = len(self.seq_group_data)
+        n_quantile_groups = n_groups - 1
+        match cfg.stack_mode:
+            case "stack-all":
+                # Here, we stack all groups into 1st column
+                cols = [column for _ in range(n_groups)]
+            case "stack-quantiles":
+                # Here, we give 1st group its own column, and stack all groups into second column
+                cols = [(column, 0)] + [(column, 1) for _ in range(n_quantile_groups)]
+            case "stack-none":
+                # Here, we stack groups into columns as [1, 3, 3, ...]
+                cols = [
+                    (column, 0),
+                    *[(column, 1 + int(i / 3)) for i in range(n_quantile_groups)],
+                ]
+            case _:
+                raise ValueError(
+                    f"Invalid stack_mode: {cfg.stack_mode}. Expected in 'stack-{{all,quantiles,none}}'."
+                )
+
+        # Create the HTML object, and add all the sequence groups to it, possibly across different columns
+        html_obj = HTML()
+        for i, (col, group_size, sequences_group) in enumerate(
+            zip(cols, cfg.group_sizes, self.seq_group_data)
+        ):
+            html_obj += sequences_group._get_html_data(
+                cfg=cfg,
+                decode_fn=decode_fn,
+                id_suffix=id_suffix,
+                column=col,
+                component_specific_kwargs=dict(
+                    group_size=group_size,
+                    max_feat_act=max_feat_act,
+                    max_loss_contribution=max_loss_contribution,
+                    seq_group_id=f"seq-group-{column}-{i}",  # we label our sequence groups with (index, column)
+                ),
+            )
+
+        return html_obj
+
+
+GenericData = (
+    FeatureTablesData
+    | ActsHistogramData
+    | LogitsTableData
+    | LogitsHistogramData
+    | SequenceMultiGroupData
+    | SequenceData
+)

SAEDashboard/sae_dashboard/components_config.py

@@ -0,0 +1,206 @@
+from dataclasses import dataclass
+from typing import Any, Iterator, Literal
+
+SEQUENCES_CONFIG_HELP = dict(
+    buffer="How many tokens to add as context to each sequence, on each side. The tokens chosen for the top acts / \
+quantile groups can't be outside the buffer range. If None, we use the entire sequence as context.",
+    compute_buffer="If False, then we don't compute the loss effect, activations, or any other data for tokens \
+other than the bold tokens in our sequences (saving time).",
+    n_quantiles="Number of quantile groups for the sequences. If zero, we only show top activations, no quantile \
+groups.",
+    top_acts_group_size="Number of sequences in the 'top activating sequences' group.",
+    quantile_group_size="Number of sequences in each of the sequence quantile groups.",
+    top_logits_hoverdata="Number of top/bottom logits to show in the hoverdata for each token.",
+    stack_mode="How to stack the sequence groups.\n  'stack-all' = all groups are stacked in a single column \
+(scrolls vertically if it overflows)\n  'stack-quantiles' = first col contains top acts, second col contains all \
+quantile groups\n  'stack-none' = we stack in a way which ensures no vertical scrolling.",
+    hover_below="Whether the hover information about a token appears below or above the token.",
+)
+
+ACTIVATIONS_HISTOGRAM_CONFIG_HELP = dict(
+    n_bins="Number of bins for the histogram.",
+)
+
+LOGITS_HISTOGRAM_CONFIG_HELP = dict(
+    n_bins="Number of bins for the histogram.",
+)
+
+LOGITS_TABLE_CONFIG_HELP = dict(
+    n_rows="Number of top/bottom logits to show in the table.",
+)
+
+FEATURE_TABLES_CONFIG_HELP = dict(
+    n_rows="Number of rows to show for each feature table.",
+    neuron_alignment_table="Whether to show the neuron alignment table.",
+    correlated_neurons_table="Whether to show the correlated neurons table.",
+    correlated_features_table="Whether to show the (pairwise) correlated features table.",
+    correlated_b_features_table="Whether to show the correlated encoder-B features table.",
+)
+
+
+@dataclass
+class BaseComponentConfig:
+    def data_is_contained_in(self, other: "BaseComponentConfig") -> bool:
+        """
+        This returns False only when the data that was computed based on `other` wouldn't be enough to show the data
+        that was computed based on `self`. For instance, if `self` was a config object with 10 rows, and `other` had
+        just 5 rows, then this would return False. A less obvious example: if `self` was a histogram config with 50 bins
+        then `other` would need to have exactly 50 bins (because we can't change the bins after generating them).
+        """
+        return True
+
+    @property
+    def help_dict(self) -> dict[str, str]:
+        """
+        This is a dictionary which maps the name of each argument to a description of what it does. This is used when
+        printing out the help for a config object, to show what each argument does.
+        """
+        return {}
+
+
+@dataclass
+class PromptConfig(BaseComponentConfig):
+    pass
+
+
+@dataclass
+class SequencesConfig(BaseComponentConfig):
+    buffer: tuple[int, int] | None = (5, 5)
+    compute_buffer: bool = True
+    n_quantiles: int = 10
+    top_acts_group_size: int = 20
+    quantile_group_size: int = 5
+    top_logits_hoverdata: int = 5
+    stack_mode: Literal["stack-all", "stack-quantiles", "stack-none"] = "stack-all"
+    hover_below: bool = True
+
+    def data_is_contained_in(self, other: BaseComponentConfig) -> bool:
+        assert isinstance(other, self.__class__)
+        return all(
+            [
+                self.buffer is None
+                or (
+                    other.buffer is not None and self.buffer[0] <= other.buffer[0]
+                ),  # the buffer needs to be <=
+                self.buffer is None
+                or (other.buffer is not None and self.buffer[1] <= other.buffer[1]),
+                int(self.compute_buffer)
+                <= int(
+                    other.compute_buffer
+                ),  # we can't compute the buffer if we didn't in `other`
+                self.n_quantiles
+                in {
+                    0,
+                    other.n_quantiles,
+                },  # we actually need the quantiles identical (or one to be zero)
+                self.top_acts_group_size
+                <= other.top_acts_group_size,  # group size needs to be <=
+                self.quantile_group_size
+                <= other.quantile_group_size,  # each quantile group needs to be <=
+                self.top_logits_hoverdata
+                <= other.top_logits_hoverdata,  # hoverdata rows need to be <=
+            ]
+        )
+
+    def __post_init__(self):
+        # Get list of group lengths, based on the config params
+        self.group_sizes = [self.top_acts_group_size] + [
+            self.quantile_group_size
+        ] * self.n_quantiles
+
+    @property
+    def help_dict(self) -> dict[str, str]:
+        return SEQUENCES_CONFIG_HELP
+
+
+@dataclass
+class ActsHistogramConfig(BaseComponentConfig):
+    n_bins: int = 50
+
+    def data_is_contained_in(self, other: BaseComponentConfig) -> bool:
+        assert isinstance(other, self.__class__)
+        return self.n_bins == other.n_bins
+
+    @property
+    def help_dict(self) -> dict[str, str]:
+        return ACTIVATIONS_HISTOGRAM_CONFIG_HELP
+
+
+@dataclass
+class LogitsHistogramConfig(BaseComponentConfig):
+    n_bins: int = 50
+
+    def data_is_contained_in(self, other: BaseComponentConfig) -> bool:
+        assert isinstance(other, self.__class__)
+        return self.n_bins == other.n_bins
+
+    @property
+    def help_dict(self) -> dict[str, str]:
+        return LOGITS_HISTOGRAM_CONFIG_HELP
+
+
+@dataclass
+class LogitsTableConfig(BaseComponentConfig):
+    n_rows: int = 10
+
+    def data_is_contained_in(self, other: BaseComponentConfig) -> bool:
+        assert isinstance(other, self.__class__)
+        return self.n_rows <= other.n_rows
+
+    @property
+    def help_dict(self) -> dict[str, str]:
+        return LOGITS_TABLE_CONFIG_HELP
+
+
+@dataclass
+class FeatureTablesConfig(BaseComponentConfig):
+    n_rows: int = 3
+    neuron_alignment_table: bool = True
+    correlated_neurons_table: bool = True
+    correlated_features_table: bool = True
+    correlated_b_features_table: bool = False
+
+    def data_is_contained_in(self, other: BaseComponentConfig) -> bool:
+        assert isinstance(other, self.__class__)
+        return all(
+            [
+                self.n_rows <= other.n_rows,
+                self.neuron_alignment_table <= other.neuron_alignment_table,
+                self.correlated_neurons_table <= other.correlated_neurons_table,
+                self.correlated_features_table <= other.correlated_features_table,
+                self.correlated_b_features_table <= other.correlated_b_features_table,
+            ]
+        )
+
+    @property
+    def help_dict(self) -> dict[str, str]:
+        return FEATURE_TABLES_CONFIG_HELP
+
+
+GenericComponentConfig = (
+    PromptConfig
+    | SequencesConfig
+    | ActsHistogramConfig
+    | LogitsHistogramConfig
+    | LogitsTableConfig
+    | FeatureTablesConfig
+)
+
+
+class Column:
+    def __init__(
+        self,
+        *args: GenericComponentConfig,
+        width: int | None = None,
+    ):
+        self.components = list(args)
+        self.width = width
+
+    def __iter__(self) -> Iterator[Any]:
+        return iter(self.components)
+
+    def __getitem__(self, idx: int) -> Any:
+        return self.components[idx]
+
+    def __len__(self) -> int:
+        return len(self.components)

SAEDashboard/sae_dashboard/css/dropdown.css

@@ -0,0 +1,40 @@
+/* Styling of the dropdowns */
+select {
+    appearance: none;
+    border: 0;
+    flex: 1;
+    padding: 0 1em;
+    background-color: #eee;
+    cursor: pointer;
+}
+.select {
+    box-shadow: 0 5px 5px rgba(0, 0, 0, 0.25);
+    cursor: pointer;
+    display: flex;
+    width: 100px;
+    height: 25px;
+    border-radius: .25em;
+    overflow: hidden;
+    position: relative;
+    margin-right: 15px;
+}
+.select::after {
+    position: absolute;
+    content: '\25BC';
+    font-size: 9px;
+    top: 0;
+    right: 0;
+    padding: 1em;
+    background-color: #ddd;
+    transition: .25s all ease;
+    pointer-events: none;
+}
+.select:hover::after {
+    color: black;
+}
+#dropdown-container {
+    margin-left: 10px;
+    margin-top: 20px;
+    display: flex;
+    flex-wrap: wrap;
+}

SAEDashboard/sae_dashboard/css/general.css

@@ -0,0 +1,53 @@
+/* Styling of the top-level container */
+.grid-container {
+    font-family: 'system-ui';
+    border: 1px solid #e6e6e6;
+    background-color: #fff;
+    margin: 30px 10px;
+    box-shadow: 0 5px 5px rgba(0, 0, 0, 0.25);
+    display: grid;
+    justify-content: start;
+    grid-template-columns: auto;
+    overflow-x: auto;
+    overflow-y: visible;
+    grid-auto-flow: column;
+    white-space: nowrap;
+    padding-bottom: 12px;
+    padding-top: 35px;
+    padding-left: 20px;
+}
+/* Styling each grid column (note, the max-height controls height of grid-container) */
+.grid-column {
+    margin-left: 20px;
+    padding-right: 20px;
+    width: max-content;
+    overflow-y: auto;
+    max-height: 750px;
+}
+/* Styling the scrollbars */
+::-webkit-scrollbar {
+    height: 10px;
+    width: 10px;
+}
+::-webkit-scrollbar-track {
+    background: #f1f1f1;
+}
+::-webkit-scrollbar-thumb {
+    background: #999;
+}
+::-webkit-scrollbar-thumb:hover {
+    background: #555;
+}
+/* Margin at the bottom of each histogram */
+.plotly-hist {
+    margin-bottom: 25px;
+}
+/* Margins below the titles (most subtitles are h4, except for the prompt-centric view which has h2 titles) */
+h4 {
+    margin-top: 0px;
+    margin-bottom: 10px;
+}
+/* Some space below the <hr> line in prompt-centric vis */
+hr {
+    margin-bottom: 35px;
+}

SAEDashboard/sae_dashboard/css/sequences.css

@@ -0,0 +1,61 @@
+/* Default font & appearance for the words in the sequence, before being hovered over */
+code {
+    font-family: Consolas, Menlo, Monaco;
+}
+/* Margin at the bottom of every sequence group, plus handle how overflow works (maybe not necessary) */
+.seq-group {
+    overflow-x: auto;
+    overflow-y: visible;
+    padding-top: 5px;
+    padding-bottom: 10px;
+    margin-bottom: 10px;
+}
+/* Margin between single sequences */
+.seq {
+    margin-bottom: 11px;
+}
+/* Styling for each token in a sequence */
+.token {
+    font-family: Consolas, Menlo, Monaco;
+    font-size: 0.9em;
+    border-top-left-radius: 3px;
+    border-top-right-radius: 3px;
+    padding: 1px;
+    color: black;
+    display: inline;
+    white-space: pre-wrap;
+}
+/* All the messy hovering stuff! */
+.hover-text {
+    position: relative;
+    cursor: pointer;
+    display: inline-block; /* Needed to contain the tooltip */
+    box-sizing: border-box;
+}
+.tooltip {
+    background-color: #fff;
+    color: #333;
+    text-align: center;
+    border-radius: 10px;
+    padding: 5px;
+    box-shadow: 0 4px 8px rgba(0, 0, 0, 0.25);
+    align-items: center;
+    justify-content: center;
+    overflow: hidden;
+    font-family: 'system-ui';
+    font-size: 1.1em;
+    display: none;
+    position: fixed;
+    z-index: 1000;
+}
+.token:hover {
+    border-top: 3px solid black;
+}
+.tooltip-container {
+    position: absolute;
+    pointer-events: none;
+}
+.hover-text:hover + .tooltip-container .tooltip {
+    display: block;
+}
+

SAEDashboard/sae_dashboard/css/tables.css

@@ -0,0 +1,81 @@
+table {
+    border: unset;
+    color: black;
+    border-collapse: collapse;
+    width: -moz-fit-content;
+    width: -webkit-fit-content;
+    width: fit-content;
+    margin-left: auto;
+    margin-right: auto;
+    font-size: 0.8em;
+}
+table.table-left tr {
+    border-bottom: 1px solid #eee;
+    padding: 15px;
+}
+table.table-left td {
+    padding: 3px 4px;
+}
+table.table-left {
+    width: 100%;
+}
+table.table-left td.left-aligned {
+    max-width: 120px;
+    overflow-x: hidden;
+}
+td {
+    border: none;
+    padding: 2px 4px;
+    white-space: nowrap;
+}
+.right-aligned {
+    text-align: right;
+}
+.left-aligned {
+    text-align: left;
+}
+.center-aligned {
+    text-align: center;
+    padding-bottom: 8px;
+}
+table code {
+    background-color: #ddd;
+    padding: 2px;
+    border-radius: 3px;
+}
+.table-container {
+    width: 100%;
+}
+.half-width-container {
+    display: flex;
+}
+.half-width {
+    width: 50%;
+    margin-right: -4px;
+}
+
+/* Feature tables should have space below them, also they should have a min column width */
+div.feature-tables table {
+    margin-bottom: 25px;
+    min-width: 250px;
+}
+/* Configure logits table container (i.e. the thing containing the smaller and larger tables) */
+div.logits-table {
+    min-width: 375px;
+    display: flex;
+    overflow-x: hidden;
+    margin-bottom: 20px;
+}
+/* Code is always bold in this table (this is just the neg/pos string tokens) */
+div.logits-table code {
+    font-weight: bold;
+}
+/* Set width of the tables inside the container (so they can stack horizontally), also put a gap between them */
+div.logits-table > div.positive {
+    width: 47%;
+}
+div.logits-table > div.negative {
+    width: 47%;
+    margin-right: 5%;
+}
+

SAEDashboard/sae_dashboard/data_parsing_fns.py

@@ -0,0 +1,412 @@
+import einops
+import numpy as np
+import torch
+from eindex import eindex
+from jaxtyping import Float, Int
+from sae_lens import SAE
+from torch import Tensor
+from transformer_lens import HookedTransformer, utils
+
+from sae_dashboard.components import LogitsTableData, SequenceData
+from sae_dashboard.sae_vis_data import SaeVisData
+from sae_dashboard.transformer_lens_wrapper import (
+    ActivationConfig,
+    TransformerLensWrapper,
+    to_resid_direction,
+)
+from sae_dashboard.utils_fns import RollingCorrCoef, TopK
+
+Arr = np.ndarray
+
+
+def get_features_table_data(
+    feature_out_dir: Float[Tensor, "feats d_out"],
+    n_rows: int,
+    corrcoef_neurons: RollingCorrCoef | None = None,
+    corrcoef_encoder: RollingCorrCoef | None = None,
+) -> dict[str, list[list[int]] | list[list[float]]]:
+    # ! Calculate all data for the left-hand column visualisations, i.e. the 3 tables
+    # Store kwargs (makes it easier to turn the tables on and off individually)
+    feature_tables_data: dict[str, list[list[int]] | list[list[float]]] = {}
+
+    # Table 1: neuron alignment, based on decoder weights
+    # if layout.feature_tables_cfg.neuron_alignment_table:
+    # Let's just always do this.
+    add_neuron_alignment_data(
+        feature_out_dir=feature_out_dir,
+        feature_tables_data=feature_tables_data,
+        n_rows=n_rows,
+    )
+
+    # Table 2: neurons correlated with this feature, based on their activations
+    if corrcoef_neurons is not None:
+        add_feature_neuron_correlations(
+            corrcoef_neurons=corrcoef_neurons,
+            feature_tables_data=feature_tables_data,
+            n_rows=n_rows,
+        )
+
+    # Table 3: primary encoder features correlated with this feature, based on their activations
+    if corrcoef_encoder is not None:
+        add_intra_encoder_correlations(
+            corrcoef_encoder=corrcoef_encoder,
+            feature_tables_data=feature_tables_data,
+            n_rows=n_rows,
+        )
+
+    return feature_tables_data
+
+
+def add_intra_encoder_correlations(
+    corrcoef_encoder: RollingCorrCoef,
+    feature_tables_data: dict[str, list[list[int]] | list[list[float]]],
+    n_rows: int,
+):
+    enc_indices, enc_pearson, enc_cossim = corrcoef_encoder.topk_pearson(
+        k=n_rows,
+    )
+    feature_tables_data["correlated_features_indices"] = enc_indices
+    feature_tables_data["correlated_features_pearson"] = enc_pearson
+    feature_tables_data["correlated_features_cossim"] = enc_cossim
+
+
+def add_neuron_alignment_data(
+    feature_out_dir: Float[Tensor, "feats d_out"],
+    feature_tables_data: dict[str, list[list[int]] | list[list[float]]],
+    n_rows: int,
+):
+    top3_neurons_aligned = TopK(tensor=feature_out_dir.float(), k=n_rows, largest=True)
+    feature_out_l1_norm = feature_out_dir.abs().sum(dim=-1, keepdim=True)
+    pct_of_l1: Arr = np.absolute(top3_neurons_aligned.values) / utils.to_numpy(
+        feature_out_l1_norm.float()
+    )
+    feature_tables_data["neuron_alignment_indices"] = (
+        top3_neurons_aligned.indices.tolist()
+    )
+    feature_tables_data["neuron_alignment_values"] = (
+        top3_neurons_aligned.values.tolist()
+    )
+    feature_tables_data["neuron_alignment_l1"] = pct_of_l1.tolist()
+
+
+def add_feature_neuron_correlations(
+    corrcoef_neurons: RollingCorrCoef,
+    feature_tables_data: dict[str, list[list[int]] | list[list[float]]],
+    n_rows: int,
+):
+    neuron_indices, neuron_pearson, neuron_cossim = corrcoef_neurons.topk_pearson(
+        k=n_rows,
+    )
+
+    feature_tables_data["correlated_neurons_indices"] = neuron_indices
+    feature_tables_data["correlated_neurons_pearson"] = neuron_pearson
+    feature_tables_data["correlated_neurons_cossim"] = neuron_cossim
+
+
+def get_logits_table_data(
+    logit_vector: Float[Tensor, "d_vocab"], n_rows: int  # noqa: F821
+):
+    # Get logits table data
+    top_logits = TopK(logit_vector.float(), k=n_rows, largest=True)
+    bottom_logits = TopK(logit_vector.float(), k=n_rows, largest=False)
+
+    top_logit_values = top_logits.values.tolist()
+    top_token_ids = top_logits.indices.tolist()
+
+    bottom_logit_values = bottom_logits.values.tolist()
+    bottom_token_ids = bottom_logits.indices.tolist()
+
+    logits_table_data = LogitsTableData(
+        bottom_logits=bottom_logit_values,
+        bottom_token_ids=bottom_token_ids,
+        top_logits=top_logit_values,
+        top_token_ids=top_token_ids,
+    )
+
+    return logits_table_data
+
+
+# @torch.inference_mode()
+# def get_feature_data(
+#     encoder: AutoEncoder,
+#     model: HookedTransformer,
+#     tokens: Int[Tensor, "batch seq"],
+#     cfg: SaeVisConfig,
+# ) -> SaeVisData:
+#     """
+#     This is the main function which users will run to generate the feature visualization data. It batches this
+#     computation over features, in accordance with the arguments in the SaeVisConfig object (we don't want to compute all
+#     the features at once, since might give OOMs).
+
+#     See the `_get_feature_data` function for an explanation of the arguments, as well as a more detailed explanation
+#     of what this function is doing.
+
+#     The return object is the merged SaeVisData objects returned by the `_get_feature_data` function.
+#     """
+#     pass
+
+#     # return sae_vis_data
+
+
+@torch.inference_mode()
+def parse_prompt_data(
+    tokens: Int[Tensor, "batch seq"],
+    str_toks: list[str],
+    sae_vis_data: SaeVisData,
+    feat_acts: Float[Tensor, "seq feats"],
+    feature_resid_dir: Float[Tensor, "feats d_model"],
+    resid_post: Float[Tensor, "seq d_model"],
+    W_U: Float[Tensor, "d_model d_vocab"],
+    feature_idx: list[int] | None = None,
+    num_top_features: int = 10,
+) -> dict[str, tuple[list[int], list[str]]]:
+    """
+    Gets data needed to create the sequences in the prompt-centric vis (displaying dashboards for the most relevant
+    features on a prompt).
+
+    This function exists so that prompt dashboards can be generated without using our AutoEncoder or
+    TransformerLens(Wrapper) classes.
+
+    Args:
+        tokens: Int[Tensor, "batch seq"]
+            The tokens we'll be using to get the feature activations. Note that we might not be using all of them; the
+            number used is determined by `fvp.total_batch_size`.
+
+        str_toks:  list[str]
+            The tokens as a list of strings, so that they can be visualized in HTML.
+
+        sae_vis_data: SaeVisData
+             The object storing all data for each feature. We'll set each `feature_data.prompt_data` to the
+             data we get from `prompt`.
+
+        feat_acts: Float[Tensor, "seq feats"]
+            The activations values of the features across the sequence.
+
+        feature_resid_dir: Float[Tensor, "feats d_model"]
+            The directions that each feature writes to the residual stream.
+
+        resid_post: Float[Tensor, "seq d_model"]
+            The activations of the final layer of the model before the unembed.
+
+        W_U: Float[Tensor, "d_model d_vocab"]
+            The model's unembed weights for the logit lens.
+
+        feature_idx: list[int] or None
+            The features we're actually computing. These might just be a subset of the model's full features.
+
+        num_top_features: int
+            The number of top features to display in this view, for any given metric.
+
+    Returns:
+        scores_dict: dict[str, tuple[list[int], list[str]]]
+            A dictionary mapping keys like "act_quantile|'django' (0)" to a tuple of lists, where the first list is the
+            feature indices, and the second list is the string-formatted values of the scores.
+
+    As well as returning this dictionary, this function will also set `FeatureData.prompt_data` for each feature in
+    `sae_vis_data` (this is necessary for getting the prompts in the prompt-centric vis). Note this design choice could
+    have been done differently (i.e. have this function return a list of the prompt data for each feature). I chose this
+    way because it means the FeatureData._get_html_data_prompt_centric can work fundamentally the same way as
+    FeatureData._get_html_data_feature_centric, rather than treating the prompt data object as a different kind of
+    component in the vis.
+    """
+
+    device = sae_vis_data.cfg.device
+
+    if feature_idx is None:
+        feature_idx = list(sae_vis_data.feature_data_dict.keys())
+    n_feats = len(feature_idx)
+    assert (
+        feature_resid_dir.shape[0] == n_feats
+    ), f"The number of features in feature_resid_dir ({feature_resid_dir.shape[0]}) does not match the number of feature indices ({n_feats})"
+
+    assert (
+        feat_acts.shape[1] == n_feats
+    ), f"The number of features in feat_acts ({feat_acts.shape[1]}) does not match the number of feature indices ({n_feats})"
+
+    feats_loss_contribution = torch.empty(
+        size=(n_feats, tokens.shape[1] - 1), device=device
+    )
+    # Some logit computations which we only need to do once
+    # correct_token_unembeddings = model_wrapped.W_U[:, tokens[0, 1:]] # [d_model seq]
+    orig_logits = (
+        resid_post / resid_post.std(dim=-1, keepdim=True)
+    ) @ W_U  # [seq d_vocab]
+    raw_logits = feature_resid_dir @ W_U  # [feats d_vocab]
+
+    for i, feat in enumerate(feature_idx):
+        # ! Calculate the sequence data for each feature, and store it as FeatureData.prompt_data
+
+        # Get this feature's output vector, using an outer product over the feature activations for all tokens
+        resid_post_feature_effect = einops.einsum(
+            feat_acts[:, i], feature_resid_dir[i], "seq, d_model -> seq d_model"
+        )
+
+        # Ablate the output vector from the residual stream, and get logits post-ablation
+        new_resid_post = resid_post - resid_post_feature_effect
+        new_logits = (new_resid_post / new_resid_post.std(dim=-1, keepdim=True)) @ W_U
+
+        # Get the top5 & bottom5 changes in logits (don't bother with `efficient_topk` cause it's small)
+        contribution_to_logprobs = orig_logits.log_softmax(
+            dim=-1
+        ) - new_logits.log_softmax(dim=-1)
+        top_contribution_to_logits = TopK(contribution_to_logprobs[:-1], k=5)
+        bottom_contribution_to_logits = TopK(
+            contribution_to_logprobs[:-1], k=5, largest=False
+        )
+
+        # Get the change in loss (which is negative of change of logprobs for correct token)
+        loss_contribution = eindex(
+            -contribution_to_logprobs[:-1], tokens[0, 1:], "seq [seq]"
+        )
+        feats_loss_contribution[i, :] = loss_contribution
+
+        # Store the sequence data
+        sae_vis_data.feature_data_dict[feat].prompt_data = SequenceData(
+            token_ids=tokens.squeeze(0).tolist(),
+            feat_acts=[round(f, 4) for f in feat_acts[:, i].tolist()],
+            loss_contribution=[0.0] + loss_contribution.tolist(),
+            token_logits=raw_logits[i, tokens.squeeze(0)].tolist(),
+            top_token_ids=top_contribution_to_logits.indices.tolist(),
+            top_logits=top_contribution_to_logits.values.tolist(),
+            bottom_token_ids=bottom_contribution_to_logits.indices.tolist(),
+            bottom_logits=bottom_contribution_to_logits.values.tolist(),
+        )
+
+    # ! Lastly, return a dictionary mapping each key like 'act_quantile|"django" (0)' to a list of feature indices & scores
+
+    # Get a dict with keys like f"act_quantile|'My' (1)" and values (feature indices list, feature score values list)
+    scores_dict: dict[str, tuple[list[int], list[str]]] = {}
+
+    for seq_pos, seq_key in enumerate([f"{t!r} ({i})" for i, t in enumerate(str_toks)]):
+        # Filter the feature activations, since we only need the ones that are non-zero
+        feat_acts_nonzero_filter = utils.to_numpy(feat_acts[seq_pos] > 0)
+        feat_acts_nonzero_locations = np.nonzero(feat_acts_nonzero_filter)[0].tolist()
+        _feat_acts = feat_acts[seq_pos, feat_acts_nonzero_filter]  # [feats_filtered,]
+        _feature_idx = np.array(feature_idx)[feat_acts_nonzero_filter]
+
+        if feat_acts_nonzero_filter.sum() > 0:
+            k = min(num_top_features, _feat_acts.numel())
+
+            # Get the top features by activation size. This is just applying a TopK function to the feat acts (which
+            # were stored by the code before this). The feat acts are formatted to 3dp.
+            act_size_topk = TopK(_feat_acts, k=k, largest=True)
+            top_features = _feature_idx[act_size_topk.indices].tolist()
+            formatted_scores = [f"{v:.3f}" for v in act_size_topk.values]
+            scores_dict[f"act_size|{seq_key}"] = (top_features, formatted_scores)
+
+            # Get the top features by activation quantile. We do this using the `feature_act_quantiles` object, which
+            # was stored `sae_vis_data`. This quantiles object has a method to return quantiles for a given set of
+            # data, as well as the precision (we make the precision higher for quantiles closer to 100%, because these
+            # are usually the quantiles we're interested in, and it lets us to save space in `feature_act_quantiles`).
+            act_quantile, act_precision = sae_vis_data.feature_stats.get_quantile(
+                _feat_acts, feat_acts_nonzero_locations
+            )
+            act_quantile_topk = TopK(act_quantile, k=k, largest=True)
+            act_formatting = [
+                f".{act_precision[i]-2}%" for i in act_quantile_topk.indices
+            ]
+            top_features = _feature_idx[act_quantile_topk.indices].tolist()
+            formatted_scores = [
+                f"{v:{f}}" for v, f in zip(act_quantile_topk.values, act_formatting)
+            ]
+            scores_dict[f"act_quantile|{seq_key}"] = (top_features, formatted_scores)
+
+        # We don't measure loss effect on the first token
+        if seq_pos == 0:
+            continue
+
+        # Filter the loss effects, since we only need the ones which have non-zero feature acts on the tokens before them
+        prev_feat_acts_nonzero_filter = utils.to_numpy(feat_acts[seq_pos - 1] > 0)
+        _loss_contribution = feats_loss_contribution[
+            prev_feat_acts_nonzero_filter, seq_pos - 1
+        ]  # [feats_filtered,]
+        _feature_idx_prev = np.array(feature_idx)[prev_feat_acts_nonzero_filter]
+
+        if prev_feat_acts_nonzero_filter.sum() > 0:
+            k = min(num_top_features, _loss_contribution.numel())
+
+            # Get the top features by loss effect. This is just applying a TopK function to the loss effects (which were
+            # stored by the code before this). The loss effects are formatted to 3dp. We look for the most negative
+            # values, i.e. the most loss-reducing features.
+            loss_contribution_topk = TopK(_loss_contribution, k=k, largest=False)
+            top_features = _feature_idx_prev[loss_contribution_topk.indices].tolist()
+            formatted_scores = [f"{v:+.3f}" for v in loss_contribution_topk.values]
+            scores_dict[f"loss_effect|{seq_key}"] = (top_features, formatted_scores)
+    return scores_dict
+
+
+@torch.inference_mode()
+def get_prompt_data(
+    sae_vis_data: SaeVisData,
+    prompt: str,
+    num_top_features: int,
+) -> dict[str, tuple[list[int], list[str]]]:
+    """
+    Gets data that will be used to create the sequences in the prompt-centric HTML visualisation, i.e. an object of
+    type SequenceData for each of our features.
+
+    Args:
+        sae_vis_data:     The object storing all data for each feature. We'll set each `feature_data.prompt_data` to the
+                          data we get from `prompt`.
+        prompt:           The prompt we'll be using to get the feature activations.#
+        num_top_features: The number of top features we'll be getting data for.
+
+    Returns:
+        scores_dict:      A dictionary mapping keys like "act_quantile|0" to a tuple of lists, where the first list is
+                          the feature indices, and the second list is the string-formatted values of the scores.
+
+    As well as returning this dictionary, this function will also set `FeatureData.prompt_data` for each feature in
+    `sae_vis_data`. This is because the prompt-centric vis will call `FeatureData._get_html_data_prompt_centric` on each
+    feature data object, so it's useful to have all the data in once place! Even if this will get overwritten next
+    time we call `get_prompt_data` for this same `sae_vis_data` object.
+    """
+
+    # ! Boring setup code
+    feature_idx = list(sae_vis_data.feature_data_dict.keys())
+    encoder = sae_vis_data.encoder
+    assert isinstance(encoder, SAE)
+    model = sae_vis_data.model
+    assert isinstance(model, HookedTransformer)
+    cfg = sae_vis_data.cfg
+    assert isinstance(cfg.hook_point, str), f"{cfg.hook_point=}, expected a string"
+
+    str_toks: list[str] = model.tokenizer.tokenize(prompt)  # type: ignore
+    tokens = model.tokenizer.encode(prompt, return_tensors="pt").to(  # type: ignore
+        sae_vis_data.cfg.device
+    )
+    assert isinstance(tokens, torch.Tensor)
+
+    model_wrapped = TransformerLensWrapper(model, ActivationConfig(cfg.hook_point, []))  # type: ignore
+
+    feature_act_dir = encoder.W_enc[:, feature_idx]  # [d_in feats]
+    feature_out_dir = encoder.W_dec[feature_idx]  # [feats d_in]
+    feature_resid_dir = to_resid_direction(
+        feature_out_dir, model_wrapped
+    )  # [feats d_model]
+    assert (
+        feature_act_dir.T.shape
+        == feature_out_dir.shape
+        == (len(feature_idx), encoder.cfg.d_in)
+    )
+
+    # ! Define hook functions to cache all the info required for feature ablation, then run those hook fns
+    resid_post, act_post = model_wrapped(tokens, return_logits=False)
+    resid_post: Tensor = resid_post.squeeze(0)  # type: ignore
+    feat_acts = encoder.get_feature_acts_subset(act_post, feature_idx).squeeze(  # type: ignore
+        0
+    )  # [seq feats]  # type: ignore
+
+    # ! Use the data we've collected to make the scores_dict and update the sae_vis_data
+    scores_dict = parse_prompt_data(
+        tokens=tokens,
+        str_toks=str_toks,
+        sae_vis_data=sae_vis_data,
+        feat_acts=feat_acts,
+        feature_resid_dir=feature_resid_dir,
+        resid_post=resid_post,
+        W_U=model.W_U,
+        feature_idx=feature_idx,
+        num_top_features=num_top_features,
+    )
+
+    return scores_dict

SAEDashboard/sae_dashboard/data_writing_fns.py

@@ -0,0 +1,210 @@
+import itertools
+from copy import deepcopy
+from pathlib import Path
+
+from tqdm.auto import tqdm
+
+from sae_dashboard.data_parsing_fns import get_prompt_data
+from sae_dashboard.html_fns import HTML
+from sae_dashboard.sae_vis_data import SaeVisData
+from sae_dashboard.utils_fns import get_decode_html_safe_fn
+
+METRIC_TITLES = {
+    "act_size": "Activation Size",
+    "act_quantile": "Activation Quantile",
+    "loss_effect": "Loss Effect",
+}
+
+
+def save_feature_centric_vis(
+    sae_vis_data: SaeVisData,
+    filename: str | Path,
+    feature_idx: int | None = None,
+    include_only: list[int] | None = None,
+    separate_files: bool = False,
+) -> None:
+    """
+    Returns the HTML string for the view which lets you navigate between different features.
+
+    Args:
+        sae_vis_data:   Object containing visualization data.
+        filename:       The HTML filepath we'll save the visualization to. If separate_files is True, this is used as a base name.
+        feature_idx:    This is the default feature index we'll start on. If None, we use the first feature.
+        include_only:   Optional list of specific features to include.
+        separate_files: If True, saves each feature to a separate HTML file.
+    """
+    # Set the default argument for the dropdown (i.e. when the page first loads)
+    first_feature = (
+        next(iter(sae_vis_data.feature_data_dict))
+        if (feature_idx is None)
+        else feature_idx
+    )
+
+    # Get tokenize function (we only need to define it once)
+    assert sae_vis_data.model is not None
+    assert sae_vis_data.model.tokenizer is not None
+    decode_fn = get_decode_html_safe_fn(sae_vis_data.model.tokenizer)
+
+    # Create iterator
+    if include_only is not None:
+        iterator = [(i, sae_vis_data.feature_data_dict[i]) for i in include_only]
+    else:
+        iterator = list(sae_vis_data.feature_data_dict.items())
+    if sae_vis_data.cfg.verbose:
+        iterator = tqdm(iterator, desc="Saving feature-centric vis")
+
+    HTML_OBJ = HTML()  # Initialize HTML object for combined file
+
+    # For each FeatureData object, we get the html_obj for its feature-centric vis
+    for feature, feature_data in iterator:
+        html_obj = feature_data._get_html_data_feature_centric(
+            sae_vis_data.cfg.feature_centric_layout, decode_fn
+        )
+
+        if separate_files:
+            feature_HTML_OBJ = HTML()  # Initialize a new HTML object for each feature
+            feature_HTML_OBJ.js_data[str(feature)] = deepcopy(html_obj.js_data)
+            feature_HTML_OBJ.html_data = deepcopy(html_obj.html_data)
+
+            # Add the aggdata
+            feature_HTML_OBJ.js_data = {
+                "AGGDATA": sae_vis_data.feature_stats.aggdata,
+                "DASHBOARD_DATA": feature_HTML_OBJ.js_data,
+            }
+
+            # Generate filename for this feature
+            feature_filename = Path(filename).with_stem(
+                f"{Path(filename).stem}_feature_{feature}"
+            )
+
+            # Save the HTML for this feature
+            feature_HTML_OBJ.get_html(
+                layout_columns=sae_vis_data.cfg.feature_centric_layout.columns,
+                layout_height=sae_vis_data.cfg.feature_centric_layout.height,
+                filename=feature_filename,
+                first_key=str(feature),
+            )
+        else:
+            # Original behavior: accumulate all features in one HTML object
+            HTML_OBJ.js_data[str(feature)] = deepcopy(html_obj.js_data)
+            if feature == first_feature:
+                HTML_OBJ.html_data = deepcopy(html_obj.html_data)
+
+    if not separate_files:
+        # Add the aggdata
+        HTML_OBJ.js_data = {
+            "AGGDATA": sae_vis_data.feature_stats.aggdata,
+            "DASHBOARD_DATA": HTML_OBJ.js_data,
+        }
+
+        # Save our full HTML
+        HTML_OBJ.get_html(
+            layout_columns=sae_vis_data.cfg.feature_centric_layout.columns,
+            layout_height=sae_vis_data.cfg.feature_centric_layout.height,
+            filename=filename,
+            first_key=str(first_feature),
+        )
+
+
+def save_prompt_centric_vis(
+    sae_vis_data: SaeVisData,
+    prompt: str,
+    filename: str | Path,
+    metric: str | None = None,
+    seq_pos: int | None = None,
+    num_top_features: int = 10,
+):
+    """
+    Returns the HTML string for the view which lets you navigate between different features.
+
+    Args:
+        prompt:     The user-input prompt.
+        model:      Used to get the tokenizer (for converting token IDs to string tokens).
+        filename:   The HTML filepath we'll save the visualization to.
+        metric:     This is the default scoring metric we'll start on. If None, we use 'act_quantile'.
+        seq_pos:    This is the default seq pos we'll start on. If None, we use 0.
+    """
+    # Initialize the object we'll eventually get_html from
+    HTML_OBJ = HTML()
+
+    # Run forward passes on our prompt, and store the data within each FeatureData object as `self.prompt_data` as
+    # well as returning the scores_dict (which maps from score hash to a list of feature indices & formatted scores)
+
+    scores_dict = get_prompt_data(
+        sae_vis_data=sae_vis_data,
+        prompt=prompt,
+        num_top_features=num_top_features,
+    )
+
+    # Get all possible values for dropdowns
+    str_toks = sae_vis_data.model.tokenizer.tokenize(prompt)  # type: ignore
+    str_toks = [
+        t.replace("|", "│") for t in str_toks
+    ]  # vertical line -> pipe (hacky, so key splitting on | works)
+    str_toks_list = [f"{t!r} ({i})" for i, t in enumerate(str_toks)]
+    metric_list = ["act_quantile", "act_size", "loss_effect"]
+
+    # Get default values for dropdowns
+    first_metric = "act_quantile" or metric
+    first_seq_pos = str_toks_list[0 if seq_pos is None else seq_pos]
+    first_key = f"{first_metric}|{first_seq_pos}"
+
+    # Get tokenize function (we only need to define it once)
+    assert sae_vis_data.model is not None
+    assert sae_vis_data.model.tokenizer is not None
+    decode_fn = get_decode_html_safe_fn(sae_vis_data.model.tokenizer)
+
+    # For each (metric, seqpos) object, we merge the prompt-centric views of each of the top features, then we merge
+    # these all together into our HTML_OBJ
+    for _metric, _seq_pos in itertools.product(metric_list, range(len(str_toks))):
+        # Create the key for this given combination of metric & seqpos, and get our top features & scores
+        key = f"{_metric}|{str_toks_list[_seq_pos]}"
+        if key not in scores_dict:
+            continue
+        feature_idx_list, scores_formatted = scores_dict[key]
+
+        # Create HTML object, to store each feature column for all the top features for this particular key
+        html_obj = HTML()
+
+        for i, (feature_idx, score_formatted) in enumerate(
+            zip(feature_idx_list, scores_formatted)
+        ):
+            # Get HTML object at this column (which includes JavaScript to dynamically set the title)
+            html_obj += sae_vis_data.feature_data_dict[
+                feature_idx
+            ]._get_html_data_prompt_centric(
+                layout=sae_vis_data.cfg.prompt_centric_layout,
+                decode_fn=decode_fn,
+                column_idx=i,
+                bold_idx=_seq_pos,
+                title=f"<h3>#{feature_idx}<br>{METRIC_TITLES[_metric]} = {score_formatted}</h3><hr>",
+            )
+
+        # Add the JavaScript (which includes the titles for each column)
+        HTML_OBJ.js_data[key] = deepcopy(html_obj.js_data)
+
+        # Set the HTML data to be the one with the most columns (since different options might have fewer cols)
+        if len(HTML_OBJ.html_data) < len(html_obj.html_data):
+            HTML_OBJ.html_data = deepcopy(html_obj.html_data)
+
+    # Check our first key is in the scores_dict (if not, we should pick a different key)
+    assert first_key in scores_dict, "\n".join(
+        [
+            f"Key {first_key} not found in {scores_dict.keys()=}.",
+            "This means that there are no features with a nontrivial score for this choice of key & metric.",
+        ]
+    )
+
+    # Add the aggdata
+    HTML_OBJ.js_data = {
+        "AGGDATA": sae_vis_data.feature_stats.aggdata,
+        "DASHBOARD_DATA": HTML_OBJ.js_data,
+    }
+
+    # Save our full HTML
+    HTML_OBJ.get_html(
+        layout_columns=sae_vis_data.cfg.prompt_centric_layout.columns,
+        layout_height=sae_vis_data.cfg.prompt_centric_layout.height,
+        filename=filename,
+        first_key=first_key,
+    )

SAEDashboard/sae_dashboard/dfa_calculator.py

@@ -0,0 +1,159 @@
+from typing import Any, Dict, List, Union
+
+import einops
+import numpy as np
+import torch
+from sae_lens import SAE
+from transformer_lens import ActivationCache, HookedTransformer
+
+
+class DFACalculator:
+    """Calculate DFA values for a given layer and set of feature indices."""
+
+    def __init__(self, model: HookedTransformer, sae: SAE[Any]):
+        self.model = model
+        self.sae = sae
+        if (
+            hasattr(model.cfg, "n_key_value_heads")
+            and model.cfg.n_key_value_heads is not None
+            and model.cfg.n_key_value_heads < model.cfg.n_heads
+        ):
+            print("Using GQA")
+            self.use_gqa = True
+        else:
+            self.use_gqa = False
+
+    def calculate(
+        self,
+        activations: Union[Dict[str, torch.Tensor], ActivationCache],
+        layer_num: int,
+        feature_indices: List[int],
+        max_value_indices: torch.Tensor,
+    ) -> Dict[int, Any]:  # type: ignore
+        """Calculate DFA values for a given layer and set of feature indices."""
+        if not feature_indices:
+            return {}
+
+        v = activations[f"blocks.{layer_num}.attn.hook_v"]
+        attn_weights = activations[f"blocks.{layer_num}.attn.hook_pattern"]
+
+        if self.use_gqa:
+            per_src_pos_dfa = self.calculate_gqa_intermediate_tensor(
+                attn_weights, v, feature_indices
+            )
+        else:
+            per_src_pos_dfa = self.calculate_standard_intermediate_tensor(
+                attn_weights, v, feature_indices
+            )
+
+        n_prompts, seq_len, _, n_features = per_src_pos_dfa.shape
+
+        # Use advanced indexing to get per_src_dfa
+        prompt_indices = torch.arange(n_prompts)[:, None, None]
+        src_pos_indices = torch.arange(seq_len)[None, :, None]
+        feature_indices_tensor = torch.arange(n_features)[None, None, :]
+        max_value_indices_expanded = max_value_indices[:, None, :]
+
+        per_src_dfa = per_src_pos_dfa[
+            prompt_indices,
+            max_value_indices_expanded,
+            src_pos_indices,
+            feature_indices_tensor,
+        ]
+
+        max_values, _ = per_src_dfa.max(dim=1)
+
+        # Create a structured numpy array to hold all the data
+        dtype = np.dtype(
+            [
+                ("dfa_values", np.float32, (seq_len,)),
+                ("dfa_target_index", np.int32),
+                ("dfa_max_value", np.float32),
+            ]
+        )
+        results = np.zeros((len(feature_indices), n_prompts), dtype=dtype)
+
+        # Fill the numpy array with data
+        results["dfa_values"] = per_src_dfa.detach().cpu().numpy().transpose(2, 0, 1)
+        results["dfa_target_index"] = max_value_indices.detach().cpu().numpy().T
+        results["dfa_max_value"] = max_values.detach().cpu().numpy().T
+
+        # Create a dictionary mapping feature indices to their respective data
+        final_results = {
+            feat_idx: results[i] for i, feat_idx in enumerate(feature_indices)
+        }
+
+        return final_results
+
+    def calculate_standard_intermediate_tensor(
+        self,
+        attn_weights: torch.Tensor,
+        v: torch.Tensor,
+        feature_indices: List[int],
+    ) -> torch.Tensor:
+        v_cat = einops.rearrange(
+            v, "batch src_pos n_heads d_head -> batch src_pos (n_heads d_head)"
+        )
+
+        attn_weights_bcast = einops.repeat(
+            attn_weights,
+            "batch n_heads dest_pos src_pos -> batch dest_pos src_pos (n_heads d_head)",
+            d_head=self.model.cfg.d_head,
+        )
+
+        decomposed_z_cat = attn_weights_bcast * v_cat.unsqueeze(1)
+
+        W_enc_selected = self.sae.W_enc[:, feature_indices]  # [d_model, num_indices]
+
+        per_src_pos_dfa = einops.einsum(
+            decomposed_z_cat,
+            W_enc_selected,
+            "batch dest_pos src_pos d_model, d_model num_features -> batch dest_pos src_pos num_features",
+        )
+
+        return per_src_pos_dfa
+
+    def calculate_gqa_intermediate_tensor(
+        self, attn_weights: torch.Tensor, v: torch.Tensor, feature_indices: List[int]
+    ) -> torch.Tensor:
+        n_query_heads = attn_weights.shape[1]
+        n_kv_heads = v.shape[2]
+        expansion_factor = n_query_heads // n_kv_heads
+        v = v.repeat_interleave(expansion_factor, dim=2)
+
+        v_cat = einops.rearrange(
+            v, "batch src_pos n_heads d_head -> batch src_pos (n_heads d_head)"
+        )
+
+        W_enc_selected = self.sae.W_enc[:, feature_indices]  # [d_model, num_indices]
+
+        # Initialize the result tensor
+        n_prompts, seq_len, _ = v_cat.shape
+        n_features = len(feature_indices)
+        per_src_pos_dfa = torch.zeros(
+            (n_prompts, seq_len, seq_len, n_features), device=v_cat.device
+        )
+
+        # Process in chunks
+        chunk_size = 16  # Adjust this based on your memory constraints
+        for i in range(0, seq_len, chunk_size):
+            chunk_end = min(i + chunk_size, seq_len)
+
+            # Process a chunk of destination positions
+            attn_weights_chunk = attn_weights[:, :, i:chunk_end, :]
+            attn_weights_bcast_chunk = einops.repeat(
+                attn_weights_chunk,
+                "batch n_heads dest_pos src_pos -> batch dest_pos src_pos (n_heads d_head)",
+                d_head=self.model.cfg.d_head,
+            )
+            decomposed_z_cat_chunk = attn_weights_bcast_chunk * v_cat.unsqueeze(1)
+
+            per_src_pos_dfa_chunk = einops.einsum(
+                decomposed_z_cat_chunk,
+                W_enc_selected,
+                "batch dest_pos src_pos d_model, d_model num_features -> batch dest_pos src_pos num_features",
+            )
+
+            per_src_pos_dfa[:, i:chunk_end, :, :] = per_src_pos_dfa_chunk
+
+        return per_src_pos_dfa

SAEDashboard/sae_dashboard/feature_data.py

@@ -0,0 +1,211 @@
+from dataclasses import dataclass, field
+from typing import Any, Callable, List, Literal, Optional
+
+from sae_dashboard.components import (
+    ActsHistogramData,
+    DecoderWeightsDistribution,
+    FeatureTablesData,
+    GenericData,
+    LogitsHistogramData,
+    LogitsTableData,
+    SequenceData,
+    SequenceMultiGroupData,
+)
+from sae_dashboard.components_config import (
+    ActsHistogramConfig,
+    FeatureTablesConfig,
+    GenericComponentConfig,
+    LogitsHistogramConfig,
+    LogitsTableConfig,
+    PromptConfig,
+    SequencesConfig,
+)
+from sae_dashboard.html_fns import HTML
+from sae_dashboard.layout import SaeVisLayoutConfig
+
+
+@dataclass
+class DFAData:
+    dfaValues: List[List[float]] = field(default_factory=list)
+    dfaTargetIndex: List[int] = field(default_factory=list)
+    dfaMaxValue: float = 0.0
+
+
+@dataclass
+class FeatureData:
+    """
+    This contains all the data necessary to make the feature-centric visualization, for a single feature. See
+    diagram in readme:
+
+        https://github.com/callummcdougall/sae_vis#data_storing_fnspy
+
+    Args:
+        feature_idx:    Index of the feature in question (not used within this class's methods, but used elsewhere).
+        cfg:            Contains layout parameters which are important in the `get_html` function.
+
+        The other args are the 6 possible components we might have in the feature-centric vis, i.e. this is where we
+        store the actual data. Note that one of these arguments is `prompt_data` which is only applicable in the prompt-
+        centric view.
+
+    This is used in both the feature-centric and prompt-centric views. In the feature-centric view, a single one
+    of these objects creates the HTML for a single feature (i.e. a full screen). In the prompt-centric view, a single
+    one of these objects will create one column of the full screen vis.
+    """
+
+    feature_tables_data: FeatureTablesData = field(
+        default_factory=lambda: FeatureTablesData()
+    )
+    acts_histogram_data: ActsHistogramData = field(
+        default_factory=lambda: ActsHistogramData()
+    )
+    logits_table_data: LogitsTableData = field(
+        default_factory=lambda: LogitsTableData()
+    )
+    logits_histogram_data: LogitsHistogramData = field(
+        default_factory=lambda: LogitsHistogramData()
+    )
+    sequence_data: SequenceMultiGroupData = field(
+        default_factory=lambda: SequenceMultiGroupData()
+    )
+    prompt_data: SequenceData = field(default_factory=lambda: SequenceData())
+    dfa_data: Optional[dict[int, dict[str, Any]]] = None
+    decoder_weights_data: Optional[DecoderWeightsDistribution] = None
+
+    def __post_init__(self):
+        if self.dfa_data is None:
+            self.dfa_data = {}
+
+    def get_component_from_config(self, config: GenericComponentConfig) -> GenericData:
+        """
+        Given a config object, returns the corresponding data object stored by this instance. For instance, if the input
+        is an `FeatureTablesConfig` instance, then this function returns `self.feature_tables_data`.
+        """
+        CONFIG_CLASS_MAP = {
+            FeatureTablesConfig.__name__: self.feature_tables_data,
+            ActsHistogramConfig.__name__: self.acts_histogram_data,
+            LogitsTableConfig.__name__: self.logits_table_data,
+            LogitsHistogramConfig.__name__: self.logits_histogram_data,
+            SequencesConfig.__name__: self.sequence_data,
+            PromptConfig.__name__: self.prompt_data,
+            # Add DFA config here if we create a specific config for it
+        }
+        config_class_name = config.__class__.__name__
+        assert (
+            config_class_name in CONFIG_CLASS_MAP
+        ), f"Invalid component config: {config_class_name}"
+        return CONFIG_CLASS_MAP[config_class_name]
+
+    def _get_html_data_feature_centric(
+        self,
+        layout: SaeVisLayoutConfig,
+        decode_fn: Callable[[int | list[int]], str | list[str]],
+    ) -> HTML:
+        """
+        Returns the HTML object for a single feature-centric view. These are assembled together into the full feature-
+        centric view.
+
+        Args:
+            decode_fn:  We use this function to decode the token IDs into string tokens.
+
+        Returns:
+            html_obj.html_data:
+                Contains a dictionary with keys equal to columns, and values equal to the HTML strings. These will be
+                turned into grid-column elements, and concatenated.
+            html_obj.js_data:
+                Contains a dictionary with keys = component names, and values = JavaScript data that will be used by the
+                scripts we'll eventually dump in.
+        """
+        # Create object to store all HTML
+        html_obj = HTML()
+
+        # For every column in this feature-centric layout, we add all the components in that column
+        for column_idx, column_components in layout.columns.items():
+            for component_config in column_components:
+                component = self.get_component_from_config(component_config)
+
+                html_obj += component._get_html_data(
+                    cfg=component_config,
+                    decode_fn=decode_fn,
+                    column=column_idx,
+                    id_suffix="0",  # we only use this if we have >1 set of histograms, i.e. prompt-centric vis
+                )
+
+        return html_obj
+
+    def _get_html_data_prompt_centric(
+        self,
+        layout: SaeVisLayoutConfig,
+        decode_fn: Callable[[int | list[int]], str | list[str]],
+        column_idx: int,
+        bold_idx: int | Literal["max"],
+        title: str,
+    ) -> HTML:
+        """
+        Returns the HTML object for a single column of the prompt-centric view. These are assembled together into a full
+        screen of a prompt-centric view, and then they're further assembled together into the full prompt-centric view.
+
+        Args:
+            decode_fn:  We use this function to decode the token IDs into string tokens.
+            column_idx: This method only gives us a single column (of the prompt-centric vis), so we need to know which
+                        column this is (for the JavaScript data).
+            bold_idx:   Which index should be bolded in the sequence data. If "max", we default to bolding the max-act
+                        token in each sequence.
+            title:      The title for this column, which will be used in the JavaScript data.
+
+        Returns:
+            html_obj.html_data:
+                Contains a dictionary with the single key `str(column_idx)`, representing the single column. This will
+                become a single grid-column element, and will get concatenated with others of these.
+            html_obj.js_data:
+                Contains a dictionary with keys = component names, and values = JavaScript data that will be used by the
+                scripts we'll eventually dump in.
+        """
+        # Create object to store all HTML
+        html_obj = HTML()
+
+        # Verify that we only have a single column
+        assert layout.columns.keys() == {
+            0
+        }, f"prompt_centric_layout should only have 1 column, instead found cols {layout.columns.keys()}"
+        assert (
+            layout.prompt_cfg is not None
+        ), "prompt_centric_cfg should include a PromptConfig, but found None"
+        if layout.seq_cfg is not None:
+            assert (layout.seq_cfg.n_quantiles == 0) or (
+                layout.seq_cfg.stack_mode == "stack-all"
+            ), "prompt_centric_layout should have stack_mode='stack-all' if n_quantiles > 0, so that it fits in 1 col"
+
+        # Get the maximum color over both the prompt and the sequences
+        max_feat_act = max(
+            max(self.prompt_data.feat_acts), self.sequence_data.max_feat_act
+        )
+        max_loss_contribution = max(
+            max(self.prompt_data.loss_contribution),
+            self.sequence_data.max_loss_contribution,
+        )
+
+        # For every component in the single column of this prompt-centric layout, add all the components in that column
+        for component_config in layout.columns[0]:
+            component = self.get_component_from_config(component_config)
+
+            html_obj += component._get_html_data(
+                cfg=component_config,
+                decode_fn=decode_fn,
+                column=column_idx,
+                id_suffix=str(column_idx),
+                component_specific_kwargs=dict(  # only used by SequenceData (the prompt)
+                    bold_idx=bold_idx,
+                    permanent_line=True,
+                    hover_above=True,
+                    max_feat_act=max_feat_act,
+                    max_loss_contribution=max_loss_contribution,
+                ),
+            )
+
+        # Add the title in JavaScript, and the empty title element in HTML
+        html_obj.html_data[column_idx] = (
+            f"<div id='column-{column_idx}-title'></div>\n{html_obj.html_data[column_idx]}"
+        )
+        html_obj.js_data["gridColumnTitlesData"] = {str(column_idx): title}
+
+        return html_obj

SAEDashboard/sae_dashboard/feature_data_generator.py

@@ -0,0 +1,313 @@
+from pathlib import Path
+from typing import Any, Dict, List
+
+import einops
+import numpy as np
+import torch
+from jaxtyping import Float, Int
+from sae_lens import SAE, HookedSAETransformer
+from sae_lens.config import DTYPE_MAP as DTYPES
+from sae_lens.saes.topk_sae import TopK
+from torch import Tensor, nn
+from tqdm.auto import tqdm
+
+from sae_dashboard.dfa_calculator import DFACalculator
+from sae_dashboard.sae_vis_data import SaeVisConfig
+from sae_dashboard.transformer_lens_wrapper import to_resid_direction
+from sae_dashboard.utils_fns import RollingCorrCoef
+
+Arr = np.ndarray
+
+
+class FeatureDataGenerator:
+    def __init__(
+        self,
+        cfg: SaeVisConfig,
+        tokens: Int[Tensor, "batch seq"],
+        model: HookedSAETransformer,
+        encoder: SAE[Any],
+    ):
+        self.cfg = cfg
+        self.model = model
+        self.encoder = encoder
+        self.token_minibatches = self.batch_tokens(tokens)
+        self.dfa_calculator = (
+            DFACalculator(model.model, encoder) if cfg.use_dfa else None  # type: ignore
+        )
+
+        if cfg.use_dfa:
+            assert (
+                "hook_z" in encoder.cfg.hook_name
+            ), f"DFAs are only supported for hook_z, but got {encoder.cfg.hook_name}"
+
+    @torch.inference_mode()
+    def batch_tokens(
+        self, tokens: Int[Tensor, "batch seq"]
+    ) -> list[Int[Tensor, "batch seq"]]:
+        # Get tokens into minibatches, for the fwd pass
+        token_minibatches = (
+            (tokens,)
+            if self.cfg.minibatch_size_tokens is None
+            else tokens.split(self.cfg.minibatch_size_tokens)
+        )
+        token_minibatches = [tok.to(self.cfg.device) for tok in token_minibatches]
+
+        return token_minibatches
+
+    @torch.inference_mode()
+    def get_feature_data(  # type: ignore
+        self,
+        feature_indices: list[int],
+        progress: list[tqdm] | None = None,  # type: ignore
+    ):  # type: ignore
+        # Create lists to store the feature activations & final values of the residual stream
+        all_feat_acts = []
+        all_dfa_results = {feature_idx: {} for feature_idx in feature_indices}
+        total_prompts = 0
+
+        # Create objects to store the data for computing rolling stats
+        corrcoef_neurons = RollingCorrCoef()
+        corrcoef_encoder = RollingCorrCoef(indices=feature_indices, with_self=True)
+
+        # Get encoder & decoder directions
+        feature_out_dir = self.encoder.W_dec[feature_indices]  # [feats d_autoencoder]
+        feature_resid_dir = to_resid_direction(
+            feature_out_dir, self.model  # type: ignore
+        )  # [feats d_model]
+
+        # ! Compute & concatenate together all feature activations & post-activation function values
+        for i, minibatch in enumerate(self.token_minibatches):
+            minibatch.to(self.cfg.device)
+            model_activation_dict = self.get_model_acts(i, minibatch)
+            primary_acts = model_activation_dict[
+                self.model.activation_config.primary_hook_point  # type: ignore
+            ].to(
+                self.encoder.device
+            )  # make sure acts are on the correct device
+
+            # For TopK, compute all activations first, then select features
+            if isinstance(self.encoder.activation_fn, TopK):
+                # Get all features' activations
+                all_features_acts = self.encoder.encode(primary_acts)
+                # Then select only the features we're interested in
+                feature_acts = all_features_acts[:, :, feature_indices].to(
+                    DTYPES[self.cfg.dtype]
+                )
+            else:
+                # For other activation functions, use the masking context
+                with FeatureMaskingContext(self.encoder, feature_indices):
+                    feature_acts = self.encoder.encode(primary_acts).to(
+                        DTYPES[self.cfg.dtype]
+                    )
+
+            self.update_rolling_coefficients(
+                model_acts=primary_acts,
+                feature_acts=feature_acts,
+                corrcoef_neurons=corrcoef_neurons,
+                corrcoef_encoder=corrcoef_encoder,
+            )
+
+            # Add these to the lists (we'll eventually concat)
+            all_feat_acts.append(feature_acts)
+
+            # Calculate DFA
+            if self.cfg.use_dfa and self.dfa_calculator:
+                max_value_indices = torch.argmax(feature_acts, dim=1)
+                batch_dfa_results = self.dfa_calculator.calculate(
+                    model_activation_dict,
+                    self.model.hook_layer,  # type: ignore
+                    feature_indices,
+                    max_value_indices,
+                )
+                for feature_idx, feature_data in batch_dfa_results.items():
+                    for prompt_idx in range(feature_data.shape[0]):
+                        global_prompt_idx = total_prompts + prompt_idx
+                        all_dfa_results[feature_idx][global_prompt_idx] = {
+                            "dfaValues": feature_data[prompt_idx][
+                                "dfa_values"
+                            ].tolist(),
+                            "dfaTargetIndex": int(
+                                feature_data[prompt_idx]["dfa_target_index"]
+                            ),
+                            "dfaMaxValue": float(
+                                feature_data[prompt_idx]["dfa_max_value"]
+                            ),
+                        }
+
+                total_prompts += len(minibatch)
+
+            # Update the 1st progress bar (fwd passes & getting sequence data dominates the runtime of these computations)
+            if progress is not None:
+                progress[0].update(1)
+
+        all_feat_acts = torch.cat(all_feat_acts, dim=0)
+
+        return (
+            all_feat_acts,
+            torch.tensor([]),  # all_resid_post, no longer used
+            feature_resid_dir,
+            feature_out_dir,
+            corrcoef_neurons,
+            corrcoef_encoder,
+            all_dfa_results,
+        )
+
+    @torch.inference_mode()
+    def get_model_acts(
+        self,
+        minibatch_index: int,
+        minibatch_tokens: torch.Tensor,
+        use_cache: bool = True,
+    ) -> Dict[str, torch.Tensor]:
+        """
+        A function that gets the model activations for a given minibatch of tokens.
+        Uses np.memmap for efficient caching.
+        """
+        if self.cfg.cache_dir is not None:
+            cache_path = self.cfg.cache_dir / f"model_activations_{minibatch_index}.pt"
+            if use_cache and cache_path.exists():
+                activation_dict = load_tensor_dict_torch(cache_path, self.cfg.device)
+            else:
+                activation_dict = self.model.forward(
+                    minibatch_tokens.to("cpu"), return_logits=False  # type: ignore
+                )
+                save_tensor_dict_torch(activation_dict, cache_path)
+        else:
+            activation_dict = self.model.forward(
+                minibatch_tokens.to("cpu"), return_logits=False  # type: ignore
+            )
+
+        return activation_dict
+
+    @torch.inference_mode()
+    def update_rolling_coefficients(
+        self,
+        model_acts: Float[Tensor, "batch seq d_in"],
+        feature_acts: Float[Tensor, "batch seq feats"],
+        corrcoef_neurons: RollingCorrCoef | None,
+        corrcoef_encoder: RollingCorrCoef | None,
+    ) -> None:
+        """
+
+        Args:
+            model_acts: Float[Tensor, "batch seq d_in"]
+                The activations of the model, which the SAE was trained on.
+            feature_idx: list[int]
+                The features we're computing the activations for. This will be used to index the encoder's weights.
+            corrcoef_neurons: Optional[RollingCorrCoef]
+                The object storing the minimal data necessary to compute corrcoef between feature activations & neurons.
+            corrcoef_encoder: Optional[RollingCorrCoef]
+                The object storing the minimal data necessary to compute corrcoef between pairwise feature activations.
+        """
+        # Update the CorrCoef object between feature activation & neurons
+        if corrcoef_neurons is not None:
+            corrcoef_neurons.update(
+                einops.rearrange(feature_acts, "batch seq feats -> feats (batch seq)"),
+                einops.rearrange(model_acts, "batch seq d_in -> d_in (batch seq)"),
+            )
+
+        # Update the CorrCoef object between pairwise feature activations
+        if corrcoef_encoder is not None:
+            corrcoef_encoder.update(
+                einops.rearrange(feature_acts, "batch seq feats -> feats (batch seq)"),
+                einops.rearrange(feature_acts, "batch seq feats -> feats (batch seq)"),
+            )
+
+
+def save_tensor_dict_torch(tensor_dict: Dict[str, torch.Tensor], filename: Path):
+    torch.save(tensor_dict, filename)
+
+
+def load_tensor_dict_torch(filename: Path, device: str) -> Dict[str, torch.Tensor]:
+    return torch.load(
+        filename, map_location=torch.device(device)
+    )  # Directly load to GPU
+
+
+class FeatureMaskingContext:
+    def __init__(self, sae: SAE[Any], feature_idxs: List[int]):
+        self.sae = sae
+        self.feature_idxs = feature_idxs
+        self.original_weight = {}
+
+    def __enter__(self):
+        ## W_dec
+        self.original_weight["W_dec"] = getattr(self.sae, "W_dec").data.clone()
+        # mask the weight
+        masked_weight = self.sae.W_dec[self.feature_idxs]
+        # set the weight
+        setattr(self.sae, "W_dec", nn.Parameter(masked_weight))
+
+        ## W_enc
+        # clone the weight.
+        self.original_weight["W_enc"] = getattr(self.sae, "W_enc").data.clone()
+        # mask the weight
+        masked_weight = self.sae.W_enc[:, self.feature_idxs]
+        # set the weight
+        setattr(self.sae, "W_enc", nn.Parameter(masked_weight))
+
+        # Handle architecture as either attribute or method
+        architecture = self.sae.cfg.architecture
+        if callable(architecture):
+            architecture = architecture()
+
+        if architecture in [
+            "standard",
+            "standard_transcoder",
+            "transcoder",
+            "skip_transcoder",
+        ]:
+            ## b_enc
+            self.original_weight["b_enc"] = getattr(self.sae, "b_enc").data.clone()
+            # mask the weight
+            masked_weight = self.sae.b_enc[self.feature_idxs]  # type: ignore
+            # set the weight
+            setattr(self.sae, "b_enc", nn.Parameter(masked_weight))
+
+        elif architecture in ["jumprelu", "jumprelu_transcoder"]:
+            ## b_enc
+            self.original_weight["b_enc"] = getattr(self.sae, "b_enc").data.clone()
+            # mask the weight
+            masked_weight = self.sae.b_enc[self.feature_idxs]  # type: ignore
+            # set the weight
+            setattr(self.sae, "b_enc", nn.Parameter(masked_weight))
+
+            ## threshold
+            self.original_weight["threshold"] = getattr(
+                self.sae, "threshold"
+            ).data.clone()
+            # mask the weight
+            masked_weight = self.sae.threshold[self.feature_idxs]  # type: ignore
+            # set the weight
+            setattr(self.sae, "threshold", nn.Parameter(masked_weight))
+
+        elif architecture in ["gated", "gated_transcoder"]:
+            ## b_gate
+            self.original_weight["b_gate"] = getattr(self.sae, "b_gate").data.clone()
+            # mask the weight
+            masked_weight = self.sae.b_gate[self.feature_idxs]  # type: ignore
+            # set the weight
+            setattr(self.sae, "b_gate", nn.Parameter(masked_weight))
+
+            ## r_mag
+            self.original_weight["r_mag"] = getattr(self.sae, "r_mag").data.clone()
+            # mask the weight
+            masked_weight = self.sae.r_mag[self.feature_idxs]  # type: ignore
+            # set the weight
+            setattr(self.sae, "r_mag", nn.Parameter(masked_weight))
+
+            ## b_mag
+            self.original_weight["b_mag"] = getattr(self.sae, "b_mag").data.clone()
+            # mask the weight
+            masked_weight = self.sae.b_mag[self.feature_idxs]  # type: ignore
+            # set the weight
+            setattr(self.sae, "b_mag", nn.Parameter(masked_weight))
+        else:
+            raise (ValueError("Invalid architecture"))
+
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):  # type: ignore
+        # set everything back to normal
+        for key, value in self.original_weight.items():
+            setattr(self.sae, key, nn.Parameter(value))

SAEDashboard/sae_dashboard/html/acts_histogram_template.html

@@ -0,0 +1,2 @@
+<!-- Activation densities histogram -->
+<div class="plotly-hist" id="HISTOGRAM_ACTS_ID" style="height: 150px; margin-top: 0px;"></div>

SAEDashboard/sae_dashboard/html/feature_tables_template.html

@@ -0,0 +1,2 @@
+<!-- Feature Info Tables -->
+<div class="feature-tables" id="FEATURE_TABLES_ID"></div>

SAEDashboard/sae_dashboard/html/logits_histogram_template.html

@@ -0,0 +1,2 @@
+<!-- Logits histogram -->
+<div class="plotly-hist" id="HISTOGRAM_LOGITS_ID" style="height: 150px; margin-top: 0px;"></div>

SAEDashboard/sae_dashboard/html/logits_table_template.html

@@ -0,0 +1,2 @@
+<!-- Logits table -->
+<div class="logits-table" id="LOGITS_TABLE_ID"></div>

SAEDashboard/sae_dashboard/html/sequences_group_template.html

@@ -0,0 +1,2 @@
+<!-- Sequence group -->
+<div class="seq-group" id="SEQUENCE_GROUP_ID"></div>