sae-lens 6.8.0__tar.gz → 6.9.1__tar.gz

{sae_lens-6.8.0 → sae_lens-6.9.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: sae-lens
-Version: 6.8.0
+Version: 6.9.1
 Summary: Training and Analyzing Sparse Autoencoders (SAEs)
 License: MIT
 Keywords: deep-learning,sparse-autoencoders,mechanistic-interpretability,PyTorch

{sae_lens-6.8.0 → sae_lens-6.9.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "sae-lens"
-version = "6.8.0"
+version = "6.9.1"
 description = "Training and Analyzing Sparse Autoencoders (SAEs)"
 authors = ["Joseph Bloom"]
 readme = "README.md"

{sae_lens-6.8.0 → sae_lens-6.9.1}/sae_lens/__init__.py

@@ -1,5 +1,5 @@
 # ruff: noqa: E402
-__version__ = "6.8.0"
+__version__ = "6.9.1"
 
 import logging
 

{sae_lens-6.8.0 → sae_lens-6.9.1}/sae_lens/loading/pretrained_sae_loaders.py

@@ -1,12 +1,14 @@
 import json
 import re
+import warnings
 from pathlib import Path
 from typing import Any, Protocol
 
 import numpy as np
+import requests
 import torch
 import yaml
-from huggingface_hub import hf_hub_download
+from huggingface_hub import hf_hub_download, hf_hub_url
 from huggingface_hub.utils import EntryNotFoundError
 from packaging.version import Version
 from safetensors import safe_open
@@ -1330,6 +1332,48 @@ def mwhanna_transcoder_huggingface_loader(
     return cfg_dict, state_dict, None
 
 
+def get_safetensors_tensor_shapes(url: str) -> dict[str, list[int]]:
+    """
+    Get tensor shapes from a safetensors file using HTTP range requests
+    without downloading the entire file.
+
+    Args:
+        url: Direct URL to the safetensors file
+
+    Returns:
+        Dictionary mapping tensor names to their shapes
+    """
+    # Check if server supports range requests
+    response = requests.head(url, timeout=10)
+    response.raise_for_status()
+
+    accept_ranges = response.headers.get("Accept-Ranges", "")
+    if "bytes" not in accept_ranges:
+        raise ValueError("Server does not support range requests")
+
+    # Fetch first 8 bytes to get metadata size
+    headers = {"Range": "bytes=0-7"}
+    response = requests.get(url, headers=headers, timeout=10)
+    if response.status_code != 206:
+        raise ValueError("Failed to fetch initial bytes for metadata size")
+
+    meta_size = int.from_bytes(response.content, byteorder="little")
+
+    # Fetch the metadata header
+    headers = {"Range": f"bytes=8-{8 + meta_size - 1}"}
+    response = requests.get(url, headers=headers, timeout=10)
+    if response.status_code != 206:
+        raise ValueError("Failed to fetch metadata header")
+
+    metadata_json = response.content.decode("utf-8").strip()
+    metadata = json.loads(metadata_json)
+
+    # Extract tensor shapes, excluding the __metadata__ key
+    return {
+        name: info["shape"] for name, info in metadata.items() if name != "__metadata__"
+    }
+
+
 def mntss_clt_layer_huggingface_loader(
     repo_id: str,
     folder_name: str,
@@ -1341,11 +1385,20 @@ def mntss_clt_layer_huggingface_loader(
     Load a MNTSS CLT layer as a single layer transcoder.
     The assumption is that the `folder_name` is the layer to load as an int
     """
-    base_config_path = hf_hub_download(
-        repo_id, "config.yaml", force_download=force_download
+
+    # warn that this sums the decoders together, so should only be used to find feature activations, not for reconstruction
+    warnings.warn(
+        "This loads the CLT layer as a single layer transcoder by summing all decoders together. This should only be used to find feature activations, not for reconstruction",
+        UserWarning,
+    )
+
+    cfg_dict = get_mntss_clt_layer_config_from_hf(
+        repo_id,
+        folder_name,
+        device,
+        force_download,
+        cfg_overrides,
     )
-    with open(base_config_path) as f:
-        cfg_info: dict[str, Any] = yaml.safe_load(f)
 
     # We need to actually load the weights, since the config is missing most information
     encoder_path = hf_hub_download(
@@ -1370,11 +1423,39 @@ def mntss_clt_layer_huggingface_loader(
             "W_dec": decoder_state_dict[f"W_dec_{folder_name}"].sum(dim=1),  # type: ignore
         }
 
-    cfg_dict = {
+    return cfg_dict, state_dict, None
+
+
+def get_mntss_clt_layer_config_from_hf(
+    repo_id: str,
+    folder_name: str,
+    device: str,
+    force_download: bool = False,  # noqa: ARG001
+    cfg_overrides: dict[str, Any] | None = None,
+) -> dict[str, Any]:
+    """
+    Load a MNTSS CLT layer as a single layer transcoder.
+    The assumption is that the `folder_name` is the layer to load as an int
+    """
+    base_config_path = hf_hub_download(
+        repo_id, "config.yaml", force_download=force_download
+    )
+    with open(base_config_path) as f:
+        cfg_info: dict[str, Any] = yaml.safe_load(f)
+
+    # Get tensor shapes without downloading full files using HTTP range requests
+    encoder_url = hf_hub_url(repo_id, f"W_enc_{folder_name}.safetensors")
+    encoder_shapes = get_safetensors_tensor_shapes(encoder_url)
+
+    # Extract shapes for the required tensors
+    b_dec_shape = encoder_shapes[f"b_dec_{folder_name}"]
+    b_enc_shape = encoder_shapes[f"b_enc_{folder_name}"]
+
+    return {
         "architecture": "transcoder",
-        "d_in": state_dict["b_dec"].shape[0],
-        "d_out": state_dict["b_dec"].shape[0],
-        "d_sae": state_dict["b_enc"].shape[0],
+        "d_in": b_dec_shape[0],
+        "d_out": b_dec_shape[0],
+        "d_sae": b_enc_shape[0],
         "dtype": "float32",
         "device": device if device is not None else "cpu",
         "activation_fn": "relu",
@@ -1387,8 +1468,6 @@ def mntss_clt_layer_huggingface_loader(
         **(cfg_overrides or {}),
     }
 
-    return cfg_dict, state_dict, None
-
 
 NAMED_PRETRAINED_SAE_LOADERS: dict[str, PretrainedSaeHuggingfaceLoader] = {
     "sae_lens": sae_lens_huggingface_loader,
@@ -1416,4 +1495,5 @@ NAMED_PRETRAINED_SAE_CONFIG_GETTERS: dict[str, PretrainedSaeConfigHuggingfaceLoa
     "sparsify": get_sparsify_config_from_hf,
     "gemma_2_transcoder": get_gemma_2_transcoder_config_from_hf,
     "mwhanna_transcoder": get_mwhanna_transcoder_config_from_hf,
+    "mntss_clt_layer_transcoder": get_mntss_clt_layer_config_from_hf,
 }

{sae_lens-6.8.0 → sae_lens-6.9.1}/sae_lens/pretrained_saes.yaml

@@ -14744,4 +14744,88 @@ mwhanna-qwen3-0.6b-transcoders-lowl0:
     neuronpedia: qwen3-0.6b/26-transcoder-hp-lowl0
   - id: layer_27
     path: layer_27.safetensors
-    neuronpedia: qwen3-0.6b/27-transcoder-hp-lowl0
+    neuronpedia: qwen3-0.6b/27-transcoder-hp-lowl0
+
+mntss-gemma-2-2b-2.5m-clt-as-per-layer:
+  conversion_func: mntss_clt_layer_transcoder
+  model: gemma-2-2b
+  repo_id: mntss/clt-gemma-2-2b-2.5M
+  saes:
+  - id: layer_0
+    path: 0
+    neuronpedia: gemma-2-2b/0-clt-hp
+  - id: layer_1
+    path: 1
+    neuronpedia: gemma-2-2b/1-clt-hp
+  - id: layer_2
+    path: 2
+    neuronpedia: gemma-2-2b/2-clt-hp
+  - id: layer_3
+    path: 3
+    neuronpedia: gemma-2-2b/3-clt-hp
+  - id: layer_4
+    path: 4
+    neuronpedia: gemma-2-2b/4-clt-hp
+  - id: layer_5
+    path: 5
+    neuronpedia: gemma-2-2b/5-clt-hp
+  - id: layer_6
+    path: 6
+    neuronpedia: gemma-2-2b/6-clt-hp
+  - id: layer_7
+    path: 7
+    neuronpedia: gemma-2-2b/7-clt-hp
+  - id: layer_8
+    path: 8
+    neuronpedia: gemma-2-2b/8-clt-hp
+  - id: layer_9
+    path: 9
+    neuronpedia: gemma-2-2b/9-clt-hp
+  - id: layer_10
+    path: 10
+    neuronpedia: gemma-2-2b/10-clt-hp
+  - id: layer_11
+    path: 11
+    neuronpedia: gemma-2-2b/11-clt-hp
+  - id: layer_12
+    path: 12
+    neuronpedia: gemma-2-2b/12-clt-hp
+  - id: layer_13
+    path: 13
+    neuronpedia: gemma-2-2b/13-clt-hp
+  - id: layer_14
+    path: 14
+    neuronpedia: gemma-2-2b/14-clt-hp
+  - id: layer_15
+    path: 15
+    neuronpedia: gemma-2-2b/15-clt-hp
+  - id: layer_16
+    path: 16
+    neuronpedia: gemma-2-2b/16-clt-hp
+  - id: layer_17
+    path: 17
+    neuronpedia: gemma-2-2b/17-clt-hp
+  - id: layer_18
+    path: 18
+    neuronpedia: gemma-2-2b/18-clt-hp
+  - id: layer_19
+    path: 19
+    neuronpedia: gemma-2-2b/19-clt-hp
+  - id: layer_20
+    path: 20
+    neuronpedia: gemma-2-2b/20-clt-hp
+  - id: layer_21
+    path: 21
+    neuronpedia: gemma-2-2b/21-clt-hp
+  - id: layer_22
+    path: 22
+    neuronpedia: gemma-2-2b/22-clt-hp
+  - id: layer_23
+    path: 23
+    neuronpedia: gemma-2-2b/23-clt-hp
+  - id: layer_24
+    path: 24
+    neuronpedia: gemma-2-2b/24-clt-hp
+  - id: layer_25
+    path: 25
+    neuronpedia: gemma-2-2b/25-clt-hp