Merge pull request #182 from OpenMOSS/dev

Dev

Author: dest1n1s

docs/index.md

@@ -24,32 +24,37 @@ This library provides:
     To add our library as a project dependency, run:
 
     ```bash
-    uv add lm-saes
+    uv add lm-saes==2.0.0b16
     ```
 
     We also support [Ascend NPU](https://github.com/Ascend/pytorch) as an accelerator backend. To add our library as a project dependency with NPU dependency constraints, run:
 
     ```bash
-    uv add lm-saes[npu]
+    uv add lm-saes[npu]==2.0.0b16
     ```
 
 === "Pip"
 
     Of course, you can also directly use [pip](https://pypi.org/project/pip/) to install our library. To install our library with pip, run:
 
     ```bash
-    pip install lm-saes
+    pip install lm-saes==2.0.0b16
     ```
 
     We also support [Ascend NPU](https://github.com/Ascend/pytorch) as an accelerator backend. To install our library with NPU dependency constraints, run:
 
     ```bash
-    pip install lm-saes[npu]
+    pip install lm-saes[npu]==2.0.0b16
     ```
 
 ### Load a trained Sparse Autoencoder from HuggingFace
 
-WIP
+Load any Sparse Autoencoder or other sparse dictionaries in `Language-Model-SAEs` or SAELens format.
+
+```python
+# Load Gemma Scope 2 SAE
+sae = AbstractSparseAutoEncoder.from_pretrained("gemma-scope-2-1b-pt-res-all:layer_12_width_16k_l0_small")
+```
 
 ### Training a Sparse Autoencoder
 
@@ -121,11 +126,42 @@ train_sae(settings)
 
 ### Analyze a trained Sparse Autoencoder
 
-WIP
+Requires setting up [MongoDB](https://www.mongodb.com/). See [analyze-saes](analyze-saes.md) for details.
+
+```python
+settings = AnalyzeSAESettings(
+    sae=PretrainedSAE(pretrained_name_or_path="path/to/sae", device="cuda"),
+    sae_name="pythia-160m-sae",
+    activation_factory=ActivationFactoryConfig(
+        sources=[ActivationFactoryDatasetSource(name="SlimPajama-3B")],
+        target=ActivationFactoryTarget.ACTIVATIONS_2D,
+        hook_points=["blocks.6.hook_resid_post"],
+        batch_size=16,
+        context_size=2048,
+    ),
+    model=LanguageModelConfig(model_name="EleutherAI/pythia-160m", device="cuda"),
+    model_name="pythia-160m",
+    datasets={"SlimPajama-3B": DatasetConfig(dataset_name_or_path="Hzfinfdu/SlimPajama-3B")},
+    analyzer=FeatureAnalyzerConfig(total_analyzing_tokens=100_000_000),
+    mongo=MongoDBConfig(),
+    device_type="cuda",
+)
+
+analyze_sae(settings)
+```
 
 ### Convert trained Sparse Autoencoder to SAELens format
 
-WIP
+Requires `sae_lens` package available. Supports ReLU, JumpReLU, and TopK SAEs.
+
+```python
+from lm_saes import SparseAutoEncoder
+
+sae = SparseAutoEncoder.from_pretrained("path/to/sae")
+sae_saelens = sae.to_saelens(model_name="pythia-160m")
+```
+
+You can use the `sae_saelens` with any tools compatible to SAELens.
 
 ## Citation
 

examples/load_saelens_model.py

@@ -0,0 +1,35 @@
+"""
+Load a Sparse Autoencoder from HuggingFace (using SAELens format) and use it with lm-saes.
+
+Requires: uv add "lm-saes[sae_lens]"
+"""
+
+import torch
+from transformer_lens import HookedTransformer
+
+from lm_saes.abstract_sae import AbstractSparseAutoEncoder
+
+# Load Gemma Scope 2 SAE from HuggingFace
+sae = AbstractSparseAutoEncoder.from_pretrained("gemma-scope-2-1b-pt-res-all:layer_12_width_16k_l0_small").to("cpu")
+
+print(f"Loaded SAE: {sae.cfg}")
+
+# Load Gemma 3 with TransformerLens
+model = HookedTransformer.from_pretrained("google/gemma-3-1b-pt")
+model.to("cpu")
+model.eval()
+
+prompt = "The capital of France is"
+tokens = model.to_tokens(prompt)
+_, cache = model.run_with_cache(tokens, names_filter=[sae.cfg.hook_point_in])
+activations = cache[sae.cfg.hook_point_in]
+
+with torch.no_grad():
+    feature_acts = sae.encode(activations)
+    reconstructed = sae.decode(feature_acts)
+
+l0 = (feature_acts > 0).sum(dim=-1).float().mean()
+mse = (activations.to(sae.cfg.dtype) - reconstructed).pow(2).mean()
+print(f"Prompt: {prompt}")
+print(f"Average L0: {l0.item():.1f}")
+print(f"Reconstruction MSE: {mse.item():.6f}")

pyproject.toml

@@ -114,7 +114,7 @@ npu = ["torch-npu==2.9.0rc1"]
 
 triton = ["triton"]
 
-sae_lens = ["sae-lens>=6.22.3"]
+sae_lens = ["sae-lens>=6.37.5"]
 
 [[tool.uv.index]]
 name = "torch-cpu"
@@ -206,4 +206,8 @@ update_changelog_on_bump = true
 version_files = [
     "README.md:pip install lm-saes==",
     "README.md:uv add lm-saes==",
+    "docs/index.md:uv add lm-saes==",
+    "docs/index.md:uv add lm-saes[npu]==",
+    "docs/index.md:pip install lm-saes==",
+    "docs/index.md:pip install lm-saes[npu]==",
 ]

src/lm_saes/utils/auto.py

@@ -49,6 +49,26 @@ def auto_infer_pretrained_sae_type(pretrained_name_or_path: str) -> PretrainedSA
     elif repo_exists(repo_id=repo_id):
         return PretrainedSAEType.HUGGINGFACE
 
+    likely_saelens = "/" not in repo_id
+    if likely_saelens:
+        if importlib.util.find_spec("sae_lens") is None:
+            raise ValueError(
+                f"Pretrained name or path {pretrained_name_or_path} is likely in SAELens format, but SAELens is not installed."
+            )
+        else:
+            from sae_lens.loading.pretrained_saes_directory import get_pretrained_saes_directory
+
+            lookups = get_pretrained_saes_directory()
+            if lookups.get(repo_id) is None:
+                raise ValueError(
+                    f"Pretrained name or path {pretrained_name_or_path} is likely in SAELens format, but {repo_id} is not a valid SAELens release. If you are sure this is a valid SAELens release, try upgrading SAELens to the latest version."
+                )
+            if lookups[repo_id].saes_map.get(name) is None:
+                raise ValueError(
+                    f"Pretrained name or path {pretrained_name_or_path} is likely in SAELens format, but {name} is not a valid ID in release {repo_id}. If you are sure this is a valid ID, try upgrading SAELens to the latest version."
+                )
+            return PretrainedSAEType.SAELENS
+
     raise ValueError(
         f"Pretrained name or path {pretrained_name_or_path} is not found on disk, nor on HuggingFace, nor in SAELens."
     )