sae-probes 0.1.1__tar.gz

sae_probes-0.1.1/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 sae-probes
+
+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_probes-0.1.1/PKG-INFO

@@ -0,0 +1,123 @@
+Metadata-Version: 2.3
+Name: sae-probes
+Version: 0.1.1
+Summary: Sparse probing benchmark for Sparse Autoencoders derived from the paper "Are Sparse Autoencoders Useful? A Case Study in Sparse Probing"
+Author: Subhash Kantamneni
+Requires-Python: >=3.10,<4.0
+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
+Requires-Dist: hf-transfer (>=0.1.9,<0.2.0)
+Requires-Dist: pandas (>=2.2.3,<3.0.0)
+Requires-Dist: sae-bench (>=0.4.0,<0.5.0)
+Requires-Dist: sae-lens (>=5.0.0,<7.0.0)
+Requires-Dist: scikit-learn (>=1.6.1,<2.0.0)
+Requires-Dist: transformer-lens (>=2.15.4,<3.0.0)
+Requires-Dist: xgboost (>=3.0.2,<4.0.0)
+Description-Content-Type: text/markdown
+
+# SAE Probes Benchmark
+
+This repository contains the code for the paper [_Are Sparse Autoencoders Useful? A Case Study in Sparse Probing_](https://arxiv.org/pdf/2502.16681), but has been reformatted into a Python package that will work with any SAE that can be loaded in [SAELens](https://github.com/jbloomAus/SAELens). This makes it easy to use the sparse probing tasks from the paper as a standalone SAE benchmark.
+
+# Installation
+
+```
+pip install sae-probes
+```
+
+## Running evaluations
+
+You can run benchmarks directly; any missing model activations are generated on demand. If you don't pass a `model_cache_path`, a temporary directory is used and cleaned up when the function completes. To persist activations across runs (recommended for repeated experiments), provide a `model_cache_path`.
+
+## Training Probes
+
+Probes can be trained directly on the model activations (baselines) or on SAE activations. In both cases, the following test data-balance settings are available: `"normal"`, `"scarcity"`, and `"imbalance"`. For more details about these settings, see the original paper. For the most standard sparse-probing benchmark, use the `normal` setting.
+
+### SAE Probes
+
+The most standard use of this library is as a sparse probing benchmark for SAEs using the `normal` setting. This is demonstrated below:
+
+```python
+from sae_probes import run_sae_evals
+from sae_lens import SAE
+
+# run the benchmark on a Gemma Scope SAE
+release = "gemma-scope-2b-pt-res-canonical"
+sae_id = "layer_12/width_16k/canonical"
+sae = SAE.from_pretrained(release, sae_id)
+
+run_sae_evals(
+  sae=sae,
+  model_name="gemma-2-2b",
+  hook_name="blocks.12.hook_resid_post",
+  reg_type="l1",
+  setting="normal",
+  results_path="/results/output/path",
+  # model_cache_path is optional; if omitted, a temp dir is used and cleared after
+  model_cache_path="/path/to/saved/activations",
+  ks=[1, 16],
+)
+```
+
+The sparse probing results for each dataset will be saved to `results_path` as a JSON file per dataset.
+
+### Baseline Probes
+
+You can now run baseline probes using a unified API that matches the SAE evaluation interface:
+
+```python
+from sae_probes import run_baseline_evals
+
+# Run baseline probes with consistent API
+run_baseline_evals(
+  model_name="gemma-2-2b",
+  hook_name="blocks.12.hook_resid_post",
+  setting="normal",  # or "scarcity", "imbalance"
+  results_path="/results/output/path",
+  # model_cache_path is optional; if omitted, a temp dir is used and cleared after
+  model_cache_path="/path/to/saved/activations",
+)
+```
+
+#### Output Format
+
+Both SAE and baseline probes now save results as **JSON files** with consistent structure:
+
+- **SAE results**: `sae_probes_{model_name}/{setting}_setting/{dataset}_{hook_name}_{reg_type}.json`
+- **Baseline results**: `baseline_results_{model_name}/{setting}_setting/{dataset}_{hook_name}_{method}.json`
+
+Each JSON file contains a list with metrics and metadata for easy comparison between SAE and baseline approaches.
+
+#### Optional: Pre-generating model activations
+
+Pre-generating can speed up repeated runs and lets you inspect the saved tensors. It's optional because benchmarks will auto-generate missing activations on their first run if missing.
+
+```python
+from sae_probes import generate_dataset_activations
+
+generate_dataset_activations(
+  model_name="gemma-2-2b", # the TransformerLens name of the model
+  hook_names=["blocks.12.hook_resid_post"], # Any TLens hook names
+  batch_size=64,
+  device="cuda",
+  model_cache_path="/path/to/save/activations",
+)
+```
+
+If you skip pre-generation, the benchmarks will create any missing activations automatically. Passing a `model_cache_path` persists them; if omitted, activations will be written to a temporary directory that is deleted after the run.
+
+## Citation
+
+If you use this code in your research, please cite:
+
+```
+@inproceedings{kantamnenisparse,
+  title={Are Sparse Autoencoders Useful? A Case Study in Sparse Probing},
+  author={Kantamneni, Subhash and Engels, Joshua and Rajamanoharan, Senthooran and Tegmark, Max and Nanda, Neel},
+  booktitle={Forty-second International Conference on Machine Learning}
+}
+```
+

sae_probes-0.1.1/README.md

@@ -0,0 +1,102 @@
+# SAE Probes Benchmark
+
+This repository contains the code for the paper [_Are Sparse Autoencoders Useful? A Case Study in Sparse Probing_](https://arxiv.org/pdf/2502.16681), but has been reformatted into a Python package that will work with any SAE that can be loaded in [SAELens](https://github.com/jbloomAus/SAELens). This makes it easy to use the sparse probing tasks from the paper as a standalone SAE benchmark.
+
+# Installation
+
+```
+pip install sae-probes
+```
+
+## Running evaluations
+
+You can run benchmarks directly; any missing model activations are generated on demand. If you don't pass a `model_cache_path`, a temporary directory is used and cleaned up when the function completes. To persist activations across runs (recommended for repeated experiments), provide a `model_cache_path`.
+
+## Training Probes
+
+Probes can be trained directly on the model activations (baselines) or on SAE activations. In both cases, the following test data-balance settings are available: `"normal"`, `"scarcity"`, and `"imbalance"`. For more details about these settings, see the original paper. For the most standard sparse-probing benchmark, use the `normal` setting.
+
+### SAE Probes
+
+The most standard use of this library is as a sparse probing benchmark for SAEs using the `normal` setting. This is demonstrated below:
+
+```python
+from sae_probes import run_sae_evals
+from sae_lens import SAE
+
+# run the benchmark on a Gemma Scope SAE
+release = "gemma-scope-2b-pt-res-canonical"
+sae_id = "layer_12/width_16k/canonical"
+sae = SAE.from_pretrained(release, sae_id)
+
+run_sae_evals(
+  sae=sae,
+  model_name="gemma-2-2b",
+  hook_name="blocks.12.hook_resid_post",
+  reg_type="l1",
+  setting="normal",
+  results_path="/results/output/path",
+  # model_cache_path is optional; if omitted, a temp dir is used and cleared after
+  model_cache_path="/path/to/saved/activations",
+  ks=[1, 16],
+)
+```
+
+The sparse probing results for each dataset will be saved to `results_path` as a JSON file per dataset.
+
+### Baseline Probes
+
+You can now run baseline probes using a unified API that matches the SAE evaluation interface:
+
+```python
+from sae_probes import run_baseline_evals
+
+# Run baseline probes with consistent API
+run_baseline_evals(
+  model_name="gemma-2-2b",
+  hook_name="blocks.12.hook_resid_post",
+  setting="normal",  # or "scarcity", "imbalance"
+  results_path="/results/output/path",
+  # model_cache_path is optional; if omitted, a temp dir is used and cleared after
+  model_cache_path="/path/to/saved/activations",
+)
+```
+
+#### Output Format
+
+Both SAE and baseline probes now save results as **JSON files** with consistent structure:
+
+- **SAE results**: `sae_probes_{model_name}/{setting}_setting/{dataset}_{hook_name}_{reg_type}.json`
+- **Baseline results**: `baseline_results_{model_name}/{setting}_setting/{dataset}_{hook_name}_{method}.json`
+
+Each JSON file contains a list with metrics and metadata for easy comparison between SAE and baseline approaches.
+
+#### Optional: Pre-generating model activations
+
+Pre-generating can speed up repeated runs and lets you inspect the saved tensors. It's optional because benchmarks will auto-generate missing activations on their first run if missing.
+
+```python
+from sae_probes import generate_dataset_activations
+
+generate_dataset_activations(
+  model_name="gemma-2-2b", # the TransformerLens name of the model
+  hook_names=["blocks.12.hook_resid_post"], # Any TLens hook names
+  batch_size=64,
+  device="cuda",
+  model_cache_path="/path/to/save/activations",
+)
+```
+
+If you skip pre-generation, the benchmarks will create any missing activations automatically. Passing a `model_cache_path` persists them; if omitted, activations will be written to a temporary directory that is deleted after the run.
+
+## Citation
+
+If you use this code in your research, please cite:
+
+```
+@inproceedings{kantamnenisparse,
+  title={Are Sparse Autoencoders Useful? A Case Study in Sparse Probing},
+  author={Kantamneni, Subhash and Engels, Joshua and Rajamanoharan, Senthooran and Tegmark, Max and Nanda, Neel},
+  booktitle={Forty-second International Conference on Machine Learning}
+}
+```

sae_probes-0.1.1/pyproject.toml

@@ -0,0 +1,55 @@
+[tool.poetry]
+name = "sae-probes"
+version = "0.1.1"
+description = "Sparse probing benchmark for Sparse Autoencoders derived from the paper \"Are Sparse Autoencoders Useful? A Case Study in Sparse Probing\""
+authors = ["Subhash Kantamneni", "Joshua Engels"]
+readme = "README.md"
+packages = [{ include = "sae_probes" }]
+
+[tool.poetry.dependencies]
+python = ">=3.10,<4.0"
+transformer-lens = "^2.15.4"
+sae-lens = ">=5.0.0,<7.0.0"
+scikit-learn = "^1.6.1"
+pandas = "^2.2.3"
+sae-bench = "^0.4.0"
+xgboost = "^3.0.2"
+hf-transfer = "^0.1.9"
+
+
+[tool.poetry.group.dev.dependencies]
+pytest = "^8.3.5"
+ruff = "^0.11.11"
+pyright = "^1.1.401"
+
+
+[tool.isort]
+profile = "black"
+
+[tool.ruff]
+exclude = ["dist", "docs", "*.ipynb", "tests/_comparison"]
+
+[tool.ruff.lint]
+ignore = ["E203", "E501", "E731"]
+extend-select = ["UP", "TID", "I", "F", "E", "ARG"]
+
+[tool.ruff.lint.flake8-tidy-imports.banned-api]
+"typing.Union".msg = "Use `|` instead"
+"typing.Optional".msg = "Use `| None` instead"
+"typing.Dict".msg = "Use `dict` instead"
+"typing.Tuple".msg = "Use `tuple` instead"
+"typing.List".msg = "Use `list` instead"
+
+[tool.pyright]
+typeCheckingMode = "basic"
+ignore = ["tests/_comparison"]
+
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"
+
+[tool.semantic_release]
+version_variables = ["sae_probes/__init__.py:__version__"]
+version_toml = ["pyproject.toml:tool.poetry.version"]
+build_command = "pip install poetry && poetry build"
+branches = { main = { match = "main" }, alpha = { match = "alpha", prerelease = true }, beta = { match = "beta", prerelease = true } }

sae_probes-0.1.1/sae_probes/__init__.py

@@ -0,0 +1,21 @@
+__version__ = "0.1.1"
+
+from .generate_model_activations import generate_dataset_activations
+from .run_baselines import (
+    run_all_baseline_class_imbalance,
+    run_all_baseline_corrupt,
+    run_all_baseline_normal,
+    run_all_baseline_scarcity,
+    run_baseline_evals,
+)
+from .run_sae_evals import run_sae_evals
+
+__all__ = [
+    "generate_dataset_activations",
+    "run_sae_evals",
+    "run_baseline_evals",
+    "run_all_baseline_class_imbalance",
+    "run_all_baseline_corrupt",
+    "run_all_baseline_normal",
+    "run_all_baseline_scarcity",
+]

sae_probes-0.1.1/sae_probes/constants.py

@@ -0,0 +1,8 @@
+from pathlib import Path
+from typing import Literal
+
+DATA_PATH = Path(__file__).parent / "data"
+DEFAULT_RESULTS_PATH = Path("results")
+
+RegType = Literal["l1", "l2"]
+Setting = Literal["normal", "scarcity", "imbalance"]