sae-lens 5.10.7__py3-none-any.whl → 5.11.0__py3-none-any.whl

sae_lens/__init__.py

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

sae_lens/config.py

@@ -33,6 +33,7 @@ HfDataset = DatasetDict | Dataset | IterableDatasetDict | IterableDataset
 
 SPARSITY_FILENAME = "sparsity.safetensors"
 SAE_WEIGHTS_FILENAME = "sae_weights.safetensors"
+SPARSIFY_WEIGHTS_FILENAME = "sae.safetensors"
 SAE_CFG_FILENAME = "cfg.json"
 
 

sae_lens/load_model.py

@@ -159,7 +159,7 @@ class HookedProxyLM(HookedRootModule):
 
         # We don't want to prepend bos but the tokenizer does it automatically, so we remove it manually
         if hasattr(self.tokenizer, "add_bos_token") and self.tokenizer.add_bos_token:  # type: ignore
-            tokens = get_tokens_with_bos_removed(self.tokenizer, tokens)
+            tokens = get_tokens_with_bos_removed(self.tokenizer, tokens)  # type: ignore
         return tokens  # type: ignore
 
 

sae_lens/toolkit/pretrained_sae_loaders.py

@@ -15,6 +15,7 @@ from sae_lens.config import (
     DTYPE_MAP,
     SAE_CFG_FILENAME,
     SAE_WEIGHTS_FILENAME,
+    SPARSIFY_WEIGHTS_FILENAME,
     SPARSITY_FILENAME,
 )
 from sae_lens.toolkit.pretrained_saes_directory import (
@@ -898,6 +899,146 @@ def llama_scope_r1_distill_sae_huggingface_loader(
     return cfg_dict, state_dict, log_sparsity
 
 
+def get_sparsify_config_from_hf(
+    repo_id: str,
+    folder_name: str,
+    device: str,
+    force_download: bool = False,
+    cfg_overrides: dict[str, Any] | None = None,
+) -> dict[str, Any]:
+    cfg_filename = f"{folder_name}/{SAE_CFG_FILENAME}"
+    cfg_path = hf_hub_download(
+        repo_id,
+        filename=cfg_filename,
+        force_download=force_download,
+    )
+    sae_path = Path(cfg_path).parent
+    return get_sparsify_config_from_disk(
+        sae_path, device=device, cfg_overrides=cfg_overrides
+    )
+
+
+def get_sparsify_config_from_disk(
+    path: str | Path,
+    device: str | None = None,
+    cfg_overrides: dict[str, Any] | None = None,
+) -> dict[str, Any]:
+    path = Path(path)
+
+    with open(path / SAE_CFG_FILENAME) as f:
+        old_cfg_dict = json.load(f)
+
+    config_path = path.parent / "config.json"
+    if config_path.exists():
+        with open(config_path) as f:
+            config_dict = json.load(f)
+    else:
+        config_dict = {}
+
+    folder_name = path.name
+    if folder_name == "embed_tokens":
+        hook_name, layer = "hook_embed", 0
+    else:
+        match = re.search(r"layers[._](\d+)", folder_name)
+        if match is None:
+            raise ValueError(f"Unrecognized Sparsify folder: {folder_name}")
+        layer = int(match.group(1))
+        hook_name = f"blocks.{layer}.hook_resid_post"
+
+    cfg_dict: dict[str, Any] = {
+        "architecture": "standard",
+        "d_in": old_cfg_dict["d_in"],
+        "d_sae": old_cfg_dict["d_in"] * old_cfg_dict["expansion_factor"],
+        "dtype": "bfloat16",
+        "device": device or "cpu",
+        "model_name": config_dict.get("model", path.parts[-2]),
+        "hook_name": hook_name,
+        "hook_layer": layer,
+        "hook_head_index": None,
+        "activation_fn_str": "topk",
+        "activation_fn_kwargs": {
+            "k": old_cfg_dict["k"],
+            "signed": old_cfg_dict.get("signed", False),
+        },
+        "apply_b_dec_to_input": not old_cfg_dict.get("normalize_decoder", False),
+        "dataset_path": config_dict.get(
+            "dataset", "togethercomputer/RedPajama-Data-1T-Sample"
+        ),
+        "context_size": config_dict.get("ctx_len", 2048),
+        "finetuning_scaling_factor": False,
+        "sae_lens_training_version": None,
+        "prepend_bos": True,
+        "dataset_trust_remote_code": True,
+        "normalize_activations": "none",
+        "neuronpedia_id": None,
+    }
+
+    if cfg_overrides:
+        cfg_dict.update(cfg_overrides)
+
+    return cfg_dict
+
+
+def sparsify_huggingface_loader(
+    repo_id: str,
+    folder_name: str,
+    device: str = "cpu",
+    force_download: bool = False,
+    cfg_overrides: dict[str, Any] | None = None,
+) -> tuple[dict[str, Any], dict[str, torch.Tensor], None]:
+    weights_filename = f"{folder_name}/{SPARSIFY_WEIGHTS_FILENAME}"
+    sae_path = hf_hub_download(
+        repo_id,
+        filename=weights_filename,
+        force_download=force_download,
+    )
+    cfg_dict, state_dict = sparsify_disk_loader(
+        Path(sae_path).parent, device=device, cfg_overrides=cfg_overrides
+    )
+    return cfg_dict, state_dict, None
+
+
+def sparsify_disk_loader(
+    path: str | Path,
+    device: str = "cpu",
+    cfg_overrides: dict[str, Any] | None = None,
+) -> tuple[dict[str, Any], dict[str, torch.Tensor]]:
+    cfg_dict = get_sparsify_config_from_disk(path, device, cfg_overrides)
+
+    weight_path = Path(path) / SPARSIFY_WEIGHTS_FILENAME
+    state_dict_loaded = load_file(weight_path, device=device)
+
+    dtype = DTYPE_MAP[cfg_dict["dtype"]]
+
+    W_enc = (
+        state_dict_loaded["W_enc"]
+        if "W_enc" in state_dict_loaded
+        else state_dict_loaded["encoder.weight"].T
+    ).to(dtype)
+
+    if "W_dec" in state_dict_loaded:
+        W_dec = state_dict_loaded["W_dec"].T.to(dtype)
+    else:
+        W_dec = state_dict_loaded["decoder.weight"].T.to(dtype)
+
+    if "b_enc" in state_dict_loaded:
+        b_enc = state_dict_loaded["b_enc"].to(dtype)
+    elif "encoder.bias" in state_dict_loaded:
+        b_enc = state_dict_loaded["encoder.bias"].to(dtype)
+    else:
+        b_enc = torch.zeros(cfg_dict["d_sae"], dtype=dtype, device=device)
+
+    if "b_dec" in state_dict_loaded:
+        b_dec = state_dict_loaded["b_dec"].to(dtype)
+    elif "decoder.bias" in state_dict_loaded:
+        b_dec = state_dict_loaded["decoder.bias"].to(dtype)
+    else:
+        b_dec = torch.zeros(cfg_dict["d_in"], dtype=dtype, device=device)
+
+    state_dict = {"W_enc": W_enc, "b_enc": b_enc, "W_dec": W_dec, "b_dec": b_dec}
+    return cfg_dict, state_dict
+
+
 NAMED_PRETRAINED_SAE_LOADERS: dict[str, PretrainedSaeHuggingfaceLoader] = {
     "sae_lens": sae_lens_huggingface_loader,
     "connor_rob_hook_z": connor_rob_hook_z_huggingface_loader,
@@ -906,6 +1047,7 @@ NAMED_PRETRAINED_SAE_LOADERS: dict[str, PretrainedSaeHuggingfaceLoader] = {
     "llama_scope_r1_distill": llama_scope_r1_distill_sae_huggingface_loader,
     "dictionary_learning_1": dictionary_learning_sae_huggingface_loader_1,
     "deepseek_r1": deepseek_r1_sae_huggingface_loader,
+    "sparsify": sparsify_huggingface_loader,
 }
 
 
@@ -917,4 +1059,5 @@ NAMED_PRETRAINED_SAE_CONFIG_GETTERS: dict[str, PretrainedSaeConfigHuggingfaceLoa
     "llama_scope_r1_distill": get_llama_scope_r1_distill_config_from_hf,
     "dictionary_learning_1": get_dictionary_learning_config_1_from_hf,
     "deepseek_r1": get_deepseek_r1_config_from_hf,
+    "sparsify": get_sparsify_config_from_hf,
 }

{sae_lens-5.10.7.dist-info → sae_lens-5.11.0.dist-info}/METADATA

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

{sae_lens-5.10.7.dist-info → sae_lens-5.11.0.dist-info}/RECORD

@@ -1,18 +1,18 @@
-sae_lens/__init__.py,sha256=NL-ObVNCrl6VcgcCwKLmQfl6PXcHW-7hZq98tLMICHs,1307
+sae_lens/__init__.py,sha256=A2ttZHoobEQm6YKaCqWrztd6LIDGmXlOvuyfp1aGb_E,1307
 sae_lens/analysis/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 sae_lens/analysis/hooked_sae_transformer.py,sha256=fkSsvWcTM_d7M3rRM6N2oFpeSvGhj_ENZtqmfWOzZTQ,13717
 sae_lens/analysis/neuronpedia_integration.py,sha256=dFiKRWfuT5iUfTPBPmZydSaNG3VwqZ1asuNbbQv_NCM,18488
 sae_lens/cache_activations_runner.py,sha256=dGK5EHJMHAKDAFyr25fy1COSm-61q-q6kpWENHFMaKk,12561
-sae_lens/config.py,sha256=6SP10H4U91u6UDeN1F9Lb4p1lpTv7ZKKL29-WPpyRr8,32988
+sae_lens/config.py,sha256=y3dgA_lNSpwi_n442dtrQ6RxfFKbnvUKjb7Qe1ZNoA4,33034
 sae_lens/evals.py,sha256=7cuLlT0ZTAhZ7eQbsZEFT-M3oixmaXSCBJtjh9hGnVQ,38527
-sae_lens/load_model.py,sha256=tE70sXsyyyGYW7o506O3eiw1MXyyW6DCQojLG49hWYI,6771
+sae_lens/load_model.py,sha256=TRxyUpudPCwGzSccQiHxww9OtLiwBBRurvi-HUnfdKg,6787
 sae_lens/pretokenize_runner.py,sha256=0nHQq3s_d80VS8iVK4-e6y_orAYVO8c4RrLGtIDfK_E,6885
 sae_lens/pretrained_saes.yaml,sha256=nhHW1auhyi4GHYrjUnHQqbNVhI5cMJv-HThzbzU1xG0,574145
 sae_lens/sae.py,sha256=8DFVKG72Ml_hVm49YIHJ0zAS6Pbd7O_7wDkQV5kyhxk,27965
 sae_lens/sae_training_runner.py,sha256=tduPN8BGtMatua0bNY-tXGGGxhedMu6F_O9ugDOdRmQ,9004
 sae_lens/tokenization_and_batching.py,sha256=oUAscjy_LPOrOb8_Ty6eLAcZ0B3HB_wiWjWktgolhG0,4314
 sae_lens/toolkit/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-sae_lens/toolkit/pretrained_sae_loaders.py,sha256=ANMLti9n8ipf6cBJiPFs6Cln0ug41EkoA0EhyH4AtVY,30194
+sae_lens/toolkit/pretrained_sae_loaders.py,sha256=vzDdDjy7EnNpmOcEwwXyd0AzzX4Up2Gdhs8wlogww8M,34840
 sae_lens/toolkit/pretrained_saes_directory.py,sha256=4Vn-Jex6SveD7EbxcSOBv8cx1gkPfUMLU1QOP-ww1ZE,3752
 sae_lens/training/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 sae_lens/training/activations_store.py,sha256=1H9sb1mAVrlHU5WVci1FFcCVNj1BydezI4zu7_Usm3s,35985
@@ -22,7 +22,7 @@ sae_lens/training/sae_trainer.py,sha256=xenSV0xw06y1_qLhw82_966DmWOp2nydqlrVgJA6
 sae_lens/training/training_sae.py,sha256=0A4x74qUfinLhwaK9RSoWZ7POrGc8kIU5EgBOp4UJtE,27998
 sae_lens/training/upload_saes_to_huggingface.py,sha256=P1K3nxv-IM7JptfLHj5Agiis7A_adn-g_tiq1d8PdaU,4361
 sae_lens/tutorial/tsea.py,sha256=fd1am_XXsf2KMbByDapJo-2qlxduKaa62Z2qcQZ3QKU,18145
-sae_lens-5.10.7.dist-info/LICENSE,sha256=DW6e-hDosiu4CfW0-imI57sV1I5f9UEslpviNQcOAKs,1069
-sae_lens-5.10.7.dist-info/METADATA,sha256=kRJB-7-34JB888NVsjO6ee75QvPZB5Fyqb_gB6zPECM,5324
-sae_lens-5.10.7.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-sae_lens-5.10.7.dist-info/RECORD,,
+sae_lens-5.11.0.dist-info/LICENSE,sha256=DW6e-hDosiu4CfW0-imI57sV1I5f9UEslpviNQcOAKs,1069
+sae_lens-5.11.0.dist-info/METADATA,sha256=LPXNsAtjYMtDveFW2aBFrLCjWmmnN7hiFt3WIDOxpBU,5324
+sae_lens-5.11.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+sae_lens-5.11.0.dist-info/RECORD,,