sae-lens 4.0.4__tar.gz

sae_lens-4.0.4/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 Joseph Bloom
+
+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.

sae_lens-4.0.4/PKG-INFO

@@ -0,0 +1,95 @@
+Metadata-Version: 2.1
+Name: sae-lens
+Version: 4.0.4
+Summary: Training and Analyzing Sparse Autoencoders (SAEs)
+Home-page: https://jbloomaus.github.io/SAELens
+License: MIT
+Keywords: deep-learning,sparse-autoencoders,mechanistic-interpretability,PyTorch
+Author: Joseph Bloom
+Requires-Python: >=3.10,<4.0
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Provides-Extra: mamba
+Requires-Dist: automated-interpretability (>=0.0.5,<1.0.0)
+Requires-Dist: babe (>=0.0.7,<0.0.8)
+Requires-Dist: datasets (>=2.17.1,<3.0.0)
+Requires-Dist: mamba-lens (>=0.0.4,<0.0.5) ; extra == "mamba"
+Requires-Dist: matplotlib (>=3.8.3,<4.0.0)
+Requires-Dist: matplotlib-inline (>=0.1.6,<0.2.0)
+Requires-Dist: nltk (>=3.8.1,<4.0.0)
+Requires-Dist: plotly (>=5.19.0,<6.0.0)
+Requires-Dist: plotly-express (>=0.4.1,<0.5.0)
+Requires-Dist: pytest-profiling (>=1.7.0,<2.0.0)
+Requires-Dist: python-dotenv (>=1.0.1,<2.0.0)
+Requires-Dist: pyyaml (>=6.0.1,<7.0.0)
+Requires-Dist: pyzmq (==26.0.0)
+Requires-Dist: safetensors (>=0.4.2,<0.5.0)
+Requires-Dist: transformer-lens (>=2.0.0,<3.0.0)
+Requires-Dist: transformers (>=4.38.1,<5.0.0)
+Requires-Dist: typer (>=0.12.3,<0.13.0)
+Requires-Dist: typing-extensions (>=4.10.0,<5.0.0)
+Requires-Dist: zstandard (>=0.22.0,<0.23.0)
+Project-URL: Repository, https://github.com/jbloomAus/SAELens
+Description-Content-Type: text/markdown
+
+<img width="1308" alt="Screenshot 2024-03-21 at 3 08 28 pm" src="https://github.com/jbloomAus/mats_sae_training/assets/69127271/209012ec-a779-4036-b4be-7b7739ea87f6">
+
+# SAE Lens 
+[![PyPI](https://img.shields.io/pypi/v/sae-lens?color=blue)](https://pypi.org/project/sae-lens/)
+[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
+[![build](https://github.com/jbloomAus/SAELens/actions/workflows/build.yml/badge.svg)](https://github.com/jbloomAus/SAELens/actions/workflows/build.yml)
+[![Deploy Docs](https://github.com/jbloomAus/SAELens/actions/workflows/deploy_docs.yml/badge.svg)](https://github.com/jbloomAus/SAELens/actions/workflows/deploy_docs.yml)
+[![codecov](https://codecov.io/gh/jbloomAus/SAELens/graph/badge.svg?token=N83NGH8CGE)](https://codecov.io/gh/jbloomAus/SAELens)
+
+SAELens exists to help researchers:
+- Train sparse autoencoders.
+- Analyse sparse autoencoders / research mechanistic interpretability. 
+- Generate insights which make it easier to create safe and aligned AI systems.
+
+Please refer to the [documentation](https://jbloomaus.github.io/SAELens/) for information on how to:
+- Download and Analyse pre-trained sparse autoencoders. 
+- Train your own sparse autoencoders.
+- Generate feature dashboards with the [SAE-Vis Library](https://github.com/callummcdougall/sae_vis/tree/main).
+
+SAE Lens is the result of many contributors working collectively to improve humanity's understanding of neural networks, many of whom are motivated by a desire to [safeguard humanity from risks posed by artificial intelligence](https://80000hours.org/problem-profiles/artificial-intelligence/).
+
+This library is maintained by [Joseph Bloom](https://www.jbloomaus.com/) and [David Chanin](https://github.com/chanind).
+
+## Loading Pre-trained SAEs. 
+
+Pre-trained SAEs for various models can be imported via SAE Lens. See this [page](https://jbloomaus.github.io/SAELens/sae_table/) in the readme for a list of all SAEs.
+## Tutorials
+
+- [SAE Lens + Neuronpedia](tutorials/tutorial_2_0.ipynb)[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://githubtocolab.com/jbloomAus/SAELens/blob/main/tutorials/tutorial_2_0.ipynb)
+- [Loading and Analysing Pre-Trained Sparse Autoencoders](tutorials/basic_loading_and_analysing.ipynb)
+ [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://githubtocolab.com/jbloomAus/SAELens/blob/main/tutorials/basic_loading_and_analysing.ipynb)
+ - [Understanding SAE Features with the Logit Lens](tutorials/logits_lens_with_features.ipynb)
+ [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://githubtocolab.com/jbloomAus/SAELens/blob/main/tutorials/logits_lens_with_features.ipynb)
+  - [Training a Sparse Autoencoder](tutorials/training_a_sparse_autoencoder.ipynb)
+ [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://githubtocolab.com/jbloomAus/SAELens/blob/main/tutorials/training_a_sparse_autoencoder.ipynb)
+
+
+## Join the Slack!
+
+Feel free to join the [Open Source Mechanistic Interpretability Slack](https://join.slack.com/t/opensourcemechanistic/shared_invite/zt-2k0id7mv8-CsIgPLmmHd03RPJmLUcapw) for support!
+
+
+## Citation
+
+Please cite the package as follows:
+
+```
+@misc{bloom2024saetrainingcodebase,
+   title = {SAELens},
+   author = {Joseph Bloom, Curt Tigges and David Chanin},
+   year = {2024},
+   howpublished = {\url{https://github.com/jbloomAus/SAELens}},
+}
+```
+
+

sae_lens-4.0.4/README.md

@@ -0,0 +1,55 @@
+<img width="1308" alt="Screenshot 2024-03-21 at 3 08 28 pm" src="https://github.com/jbloomAus/mats_sae_training/assets/69127271/209012ec-a779-4036-b4be-7b7739ea87f6">
+
+# SAE Lens 
+[![PyPI](https://img.shields.io/pypi/v/sae-lens?color=blue)](https://pypi.org/project/sae-lens/)
+[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
+[![build](https://github.com/jbloomAus/SAELens/actions/workflows/build.yml/badge.svg)](https://github.com/jbloomAus/SAELens/actions/workflows/build.yml)
+[![Deploy Docs](https://github.com/jbloomAus/SAELens/actions/workflows/deploy_docs.yml/badge.svg)](https://github.com/jbloomAus/SAELens/actions/workflows/deploy_docs.yml)
+[![codecov](https://codecov.io/gh/jbloomAus/SAELens/graph/badge.svg?token=N83NGH8CGE)](https://codecov.io/gh/jbloomAus/SAELens)
+
+SAELens exists to help researchers:
+- Train sparse autoencoders.
+- Analyse sparse autoencoders / research mechanistic interpretability. 
+- Generate insights which make it easier to create safe and aligned AI systems.
+
+Please refer to the [documentation](https://jbloomaus.github.io/SAELens/) for information on how to:
+- Download and Analyse pre-trained sparse autoencoders. 
+- Train your own sparse autoencoders.
+- Generate feature dashboards with the [SAE-Vis Library](https://github.com/callummcdougall/sae_vis/tree/main).
+
+SAE Lens is the result of many contributors working collectively to improve humanity's understanding of neural networks, many of whom are motivated by a desire to [safeguard humanity from risks posed by artificial intelligence](https://80000hours.org/problem-profiles/artificial-intelligence/).
+
+This library is maintained by [Joseph Bloom](https://www.jbloomaus.com/) and [David Chanin](https://github.com/chanind).
+
+## Loading Pre-trained SAEs. 
+
+Pre-trained SAEs for various models can be imported via SAE Lens. See this [page](https://jbloomaus.github.io/SAELens/sae_table/) in the readme for a list of all SAEs.
+## Tutorials
+
+- [SAE Lens + Neuronpedia](tutorials/tutorial_2_0.ipynb)[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://githubtocolab.com/jbloomAus/SAELens/blob/main/tutorials/tutorial_2_0.ipynb)
+- [Loading and Analysing Pre-Trained Sparse Autoencoders](tutorials/basic_loading_and_analysing.ipynb)
+ [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://githubtocolab.com/jbloomAus/SAELens/blob/main/tutorials/basic_loading_and_analysing.ipynb)
+ - [Understanding SAE Features with the Logit Lens](tutorials/logits_lens_with_features.ipynb)
+ [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://githubtocolab.com/jbloomAus/SAELens/blob/main/tutorials/logits_lens_with_features.ipynb)
+  - [Training a Sparse Autoencoder](tutorials/training_a_sparse_autoencoder.ipynb)
+ [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://githubtocolab.com/jbloomAus/SAELens/blob/main/tutorials/training_a_sparse_autoencoder.ipynb)
+
+
+## Join the Slack!
+
+Feel free to join the [Open Source Mechanistic Interpretability Slack](https://join.slack.com/t/opensourcemechanistic/shared_invite/zt-2k0id7mv8-CsIgPLmmHd03RPJmLUcapw) for support!
+
+
+## Citation
+
+Please cite the package as follows:
+
+```
+@misc{bloom2024saetrainingcodebase,
+   title = {SAELens},
+   author = {Joseph Bloom, Curt Tigges and David Chanin},
+   year = {2024},
+   howpublished = {\url{https://github.com/jbloomAus/SAELens}},
+}
+```
+

sae_lens-4.0.4/pyproject.toml

@@ -0,0 +1,103 @@
+[tool.poetry]
+name = "sae-lens"
+version = "4.0.4"
+description = "Training and Analyzing Sparse Autoencoders (SAEs)"
+authors = ["Joseph Bloom"]
+readme = "README.md"
+packages = [{include = "sae_lens"}]
+include = ["pretrained_saes.yaml"]
+repository = "https://github.com/jbloomAus/SAELens"
+homepage = "https://jbloomaus.github.io/SAELens"
+license = "MIT"
+keywords = [
+    "deep-learning",
+    "sparse-autoencoders",
+    "mechanistic-interpretability",
+    "PyTorch",
+]
+classifiers = ["Topic :: Scientific/Engineering :: Artificial Intelligence"]
+
+[tool.poetry.dependencies]
+python = "^3.10"
+transformer-lens = "^2.0.0"
+transformers = "^4.38.1"
+plotly = "^5.19.0"
+plotly-express = "^0.4.1"
+matplotlib = "^3.8.3"
+matplotlib-inline = "^0.1.6"
+datasets = "^2.17.1"
+babe = "^0.0.7"
+nltk = "^3.8.1"
+safetensors = "^0.4.2"
+typer = "^0.12.3"
+mamba-lens = { version = "^0.0.4", optional = true }
+pyzmq = "26.0.0"
+automated-interpretability = ">=0.0.5,<1.0.0"
+python-dotenv = "^1.0.1"
+pyyaml = "^6.0.1"
+pytest-profiling = "^1.7.0"
+zstandard = "^0.22.0"
+typing-extensions = "^4.10.0"
+
+
+[tool.poetry.group.dev.dependencies]
+black = { version = "24.4.0", extras = ["jupyter"] }
+pytest = "^8.0.2"
+pytest-cov = "^4.1.0"
+pre-commit = "^3.6.2"
+flake8 = "7.0.0"
+isort = "5.13.2"
+pyright = "1.1.365"
+mamba-lens = "^0.0.4"
+ansible-lint = { version = "^24.2.3", markers = "platform_system != 'Windows'" }
+botocore = "^1.34.101"
+boto3 = "^1.34.101"
+docstr-coverage = "^2.3.2"
+mkdocs = "^1.6.1"
+mkdocs-material = "^9.5.34"
+mkdocs-autorefs = "^1.1.0"
+mkdocs-section-index = "^0.3.9"
+mkdocstrings = "^0.25.2"
+mkdocstrings-python = "^1.10.9"
+tabulate = "^0.9.0"
+
+[tool.poetry.extras]
+mamba = ["mamba-lens"]
+
+
+[tool.isort]
+profile = "black"
+src_paths = ["sae_lens", "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"
+reportDeprecated = "none"
+reportPrivateImportUsage = "none"
+ignore = [
+    "**/wandb/**"
+]
+
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"
+
+
+[tool.semantic_release]
+version_variables = [
+    "sae_lens/__init__.py:__version__",
+]
+version_toml = [
+    "pyproject.toml:tool.poetry.version",
+]
+branch = "main"
+build_command = "pip install poetry && poetry build"

sae_lens-4.0.4/sae_lens/__init__.py

@@ -0,0 +1,35 @@
+__version__ = "4.0.4"
+
+
+from .analysis.hooked_sae_transformer import HookedSAETransformer
+from .cache_activations_runner import CacheActivationsRunner
+from .config import (
+    CacheActivationsRunnerConfig,
+    LanguageModelSAERunnerConfig,
+    PretokenizeRunnerConfig,
+)
+from .evals import run_evals
+from .pretokenize_runner import PretokenizeRunner, pretokenize_runner
+from .sae import SAE, SAEConfig
+from .sae_training_runner import SAETrainingRunner
+from .training.activations_store import ActivationsStore
+from .training.training_sae import TrainingSAE, TrainingSAEConfig
+from .training.upload_saes_to_huggingface import upload_saes_to_huggingface
+
+__all__ = [
+    "SAE",
+    "SAEConfig",
+    "TrainingSAE",
+    "TrainingSAEConfig",
+    "HookedSAETransformer",
+    "ActivationsStore",
+    "LanguageModelSAERunnerConfig",
+    "SAETrainingRunner",
+    "CacheActivationsRunnerConfig",
+    "CacheActivationsRunner",
+    "PretokenizeRunnerConfig",
+    "PretokenizeRunner",
+    "pretokenize_runner",
+    "run_evals",
+    "upload_saes_to_huggingface",
+]

sae_lens-4.0.4/sae_lens/analysis/feature_statistics.py

@@ -0,0 +1,98 @@
+import pandas as pd
+import torch
+from tqdm import tqdm
+from transformer_lens import HookedTransformer
+
+from sae_lens.sae import SAE
+
+
+@torch.no_grad()
+def get_feature_property_df(sae: SAE, feature_sparsity: torch.Tensor):
+    """
+    feature_property_df = get_feature_property_df(sae, log_feature_density.cpu())
+    """
+
+    W_dec_normalized = (
+        sae.W_dec.cpu()
+    )  # / sparse_autoencoder.W_dec.cpu().norm(dim=-1, keepdim=True)
+    W_enc_normalized = (sae.W_enc.cpu() / sae.W_enc.cpu().norm(dim=-1, keepdim=True)).T
+
+    d_e_projection = (W_dec_normalized * W_enc_normalized).sum(-1)
+    b_dec_projection = sae.b_dec.cpu() @ W_dec_normalized.T
+
+    temp_df = pd.DataFrame(
+        {
+            "log_feature_sparsity": feature_sparsity + 1e-10,
+            "d_e_projection": d_e_projection,
+            # "d_e_projection_normalized": d_e_projection_normalized,
+            "b_enc": sae.b_enc.detach().cpu(),
+            "b_dec_projection": b_dec_projection,
+            "feature": list(range(sae.cfg.d_sae)),  # type: ignore
+            "dead_neuron": (feature_sparsity < -9).cpu(),
+        }
+    )
+
+    return temp_df
+
+
+@torch.no_grad()
+def get_stats_df(projection: torch.Tensor):
+    """
+    Returns a dataframe with the mean, std, skewness and kurtosis of the projection
+    """
+    mean = projection.mean(dim=1, keepdim=True)
+    diffs = projection - mean
+    var = (diffs**2).mean(dim=1, keepdim=True)
+    std = torch.pow(var, 0.5)
+    zscores = diffs / std
+    skews = torch.mean(torch.pow(zscores, 3.0), dim=1)
+    kurtosis = torch.mean(torch.pow(zscores, 4.0), dim=1)
+
+    stats_df = pd.DataFrame(
+        {
+            "feature": range(len(skews)),
+            "mean": mean.numpy().squeeze(),
+            "std": std.numpy().squeeze(),
+            "skewness": skews.numpy(),
+            "kurtosis": kurtosis.numpy(),
+        }
+    )
+
+    return stats_df
+
+
+@torch.no_grad()
+def get_all_stats_dfs(
+    gpt2_small_sparse_autoencoders: dict[str, SAE],  # [hook_point, sae]
+    gpt2_small_sae_sparsities: dict[str, torch.Tensor],  # [hook_point, sae]
+    model: HookedTransformer,
+    cosine_sim: bool = False,
+):
+    stats_dfs = []
+    pbar = tqdm(gpt2_small_sparse_autoencoders.keys())
+    for key in pbar:
+        layer = int(key.split(".")[1])
+        sparse_autoencoder = gpt2_small_sparse_autoencoders[key]
+        pbar.set_description(f"Processing layer {sparse_autoencoder.cfg.hook_name}")
+        W_U_stats_df_dec, _ = get_W_U_W_dec_stats_df(
+            sparse_autoencoder.W_dec.cpu(), model, cosine_sim
+        )
+        log_feature_sparsity = gpt2_small_sae_sparsities[key].detach().cpu()
+        W_U_stats_df_dec["log_feature_sparsity"] = log_feature_sparsity
+        W_U_stats_df_dec["layer"] = layer + (1 if "post" in key else 0)
+        stats_dfs.append(W_U_stats_df_dec)
+
+    W_U_stats_df_dec_all_layers = pd.concat(stats_dfs, axis=0)
+    return W_U_stats_df_dec_all_layers
+
+
+@torch.no_grad()
+def get_W_U_W_dec_stats_df(
+    W_dec: torch.Tensor, model: HookedTransformer, cosine_sim: bool = False
+) -> tuple[pd.DataFrame, torch.Tensor]:
+    W_U = model.W_U.detach().cpu()
+    if cosine_sim:
+        W_U = W_U / W_U.norm(dim=0, keepdim=True)
+    dec_projection_onto_W_U = W_dec @ W_U
+    W_U_stats_df = get_stats_df(dec_projection_onto_W_U)
+    return W_U_stats_df, dec_projection_onto_W_U