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

sae_lens/training/upload_saes_to_huggingface.py

@@ -2,23 +2,24 @@ import io
 from pathlib import Path
 from tempfile import TemporaryDirectory
 from textwrap import dedent
-from typing import Iterable
+from typing import Any, Iterable
 
 from huggingface_hub import HfApi, create_repo, get_hf_file_metadata, hf_hub_url
 from huggingface_hub.utils import EntryNotFoundError, RepositoryNotFoundError
 from tqdm.autonotebook import tqdm
 
 from sae_lens import logger
-from sae_lens.config import (
+from sae_lens.constants import (
+    RUNNER_CFG_FILENAME,
     SAE_CFG_FILENAME,
     SAE_WEIGHTS_FILENAME,
     SPARSITY_FILENAME,
 )
-from sae_lens.sae import SAE
+from sae_lens.saes.sae import SAE
 
 
 def upload_saes_to_huggingface(
-    saes_dict: dict[str, SAE | Path | str],
+    saes_dict: dict[str, SAE[Any] | Path | str],
     hf_repo_id: str,
     hf_revision: str = "main",
     show_progress: bool = True,
@@ -87,7 +88,7 @@ def _create_default_readme(repo_id: str, sae_ids: Iterable[str]) -> str:
         ```python
         from sae_lens import SAE
 
-        sae, cfg_dict, sparsity = SAE.from_pretrained("{repo_id}", "<sae_id>")
+        sae = SAE.from_pretrained("{repo_id}", "<sae_id>")
         ```
         """
     )
@@ -119,11 +120,16 @@ def _upload_sae(api: HfApi, sae_path: Path, repo_id: str, sae_id: str, revision:
         revision=revision,
         repo_type="model",
         commit_message=f"Upload SAE {sae_id}",
-        allow_patterns=[SAE_CFG_FILENAME, SAE_WEIGHTS_FILENAME, SPARSITY_FILENAME],
+        allow_patterns=[
+            SAE_CFG_FILENAME,
+            SAE_WEIGHTS_FILENAME,
+            SPARSITY_FILENAME,
+            RUNNER_CFG_FILENAME,
+        ],
     )
 
 
-def _build_sae_path(sae_ref: SAE | Path | str, tmp_dir: str) -> Path:
+def _build_sae_path(sae_ref: SAE[Any] | Path | str, tmp_dir: str) -> Path:
     if isinstance(sae_ref, SAE):
         sae_ref.save_model(tmp_dir)
         return Path(tmp_dir)

sae_lens/util.py

@@ -0,0 +1,47 @@
+import re
+from dataclasses import asdict, fields, is_dataclass
+from typing import Sequence, TypeVar
+
+K = TypeVar("K")
+V = TypeVar("V")
+
+
+def filter_valid_dataclass_fields(
+    source: dict[str, V] | object,
+    destination: object | type,
+    whitelist_fields: Sequence[str] | None = None,
+) -> dict[str, V]:
+    """Filter a source dict or dataclass instance to only include fields that are present in the destination dataclass."""
+
+    if not is_dataclass(destination):
+        raise ValueError(f"{destination} is not a dataclass")
+
+    if is_dataclass(source) and not isinstance(source, type):
+        source_dict = asdict(source)
+    elif isinstance(source, dict):
+        source_dict = source
+    else:
+        raise ValueError(f"{source} is not a dict or dataclass")
+
+    valid_field_names = {field.name for field in fields(destination)}
+    if whitelist_fields is not None:
+        valid_field_names = valid_field_names.union(whitelist_fields)
+    return {key: val for key, val in source_dict.items() if key in valid_field_names}
+
+
+def extract_stop_at_layer_from_tlens_hook_name(hook_name: str) -> int | None:
+    """Extract the stop_at layer from a HookedTransformer hook name.
+
+    Returns None if the hook name is not a valid HookedTransformer hook name.
+    """
+    layer = extract_layer_from_tlens_hook_name(hook_name)
+    return None if layer is None else layer + 1
+
+
+def extract_layer_from_tlens_hook_name(hook_name: str) -> int | None:
+    """Extract the layer from a HookedTransformer hook name.
+
+    Returns None if the hook name is not a valid HookedTransformer hook name.
+    """
+    hook_match = re.search(r"\.(\d+)\.", hook_name)
+    return None if hook_match is None else int(hook_match.group(1))

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

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

sae_lens-6.0.0.dist-info/RECORD

@@ -0,0 +1,37 @@
+sae_lens/__init__.py,sha256=kAAUdxuMPOM3qLV1vK245fOEruEdsKd_GgAEKJsvxgY,2856
+sae_lens/analysis/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+sae_lens/analysis/hooked_sae_transformer.py,sha256=Eyg1Y2hVIHNuiiLOCTgzstOuW6iA-7hPHqaGR8y_vMs,13809
+sae_lens/analysis/neuronpedia_integration.py,sha256=MrENqc81Mc2SMbxGjbwHzpkGUCAFKSf0i4EdaUF2Oj4,18707
+sae_lens/cache_activations_runner.py,sha256=L5hhuU2-zPQr2S3L64GMKKLeMQfqXxwDl8NbuOtrybI,12567
+sae_lens/config.py,sha256=9Lg4HkQvj1t9QZJdmC071lyJMc_iqNQknosT7zOYfwM,27278
+sae_lens/constants.py,sha256=CSjmiZ-bhjQeVLyRvWxAjBokCgkfM8mnvd7-vxLIWTY,639
+sae_lens/evals.py,sha256=kQyrzczKaVD9rHwfFa_DxL_gMXDxsoIVHmsFIPIU2bY,38696
+sae_lens/llm_sae_training_runner.py,sha256=58XbDylw2fPOD7C-ZfSAjeNqJLXB05uHGTuiYVVbXXY,13354
+sae_lens/load_model.py,sha256=dBB_9gO6kWyQ4sXHq7qB8T3YUlXm3PGwYcpR4UVW4QY,8633
+sae_lens/loading/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+sae_lens/loading/pretrained_sae_loaders.py,sha256=5XEU4uFFeGCePwqDwhlE7CrFGRSI0U9Cu-UQVa33Y1E,36432
+sae_lens/loading/pretrained_saes_directory.py,sha256=4Vn-Jex6SveD7EbxcSOBv8cx1gkPfUMLU1QOP-ww1ZE,3752
+sae_lens/pretokenize_runner.py,sha256=0nHQq3s_d80VS8iVK4-e6y_orAYVO8c4RrLGtIDfK_E,6885
+sae_lens/pretrained_saes.yaml,sha256=nhHW1auhyi4GHYrjUnHQqbNVhI5cMJv-HThzbzU1xG0,574145
+sae_lens/registry.py,sha256=nhy7BPSudSATqW4lo9H_k3Na7sfGHmAf9v-3wpnLL_o,1490
+sae_lens/saes/__init__.py,sha256=v6mfeDzyGYtT6x5SszAQtkldTXwPE-V_iwOlrT_pDwQ,1008
+sae_lens/saes/gated_sae.py,sha256=0zd66bH04nsaGk3bxHk10hsZofa2GrFbMo15LOsuqgU,9233
+sae_lens/saes/jumprelu_sae.py,sha256=iwmPQJ4XpIxzgosty680u8Zj7x1uVZhM75kPOT3obi0,12060
+sae_lens/saes/sae.py,sha256=ZEXEXFVtrtFrzuOV3nyweTBleNCV4EDGh1ImaF32uqg,39618
+sae_lens/saes/standard_sae.py,sha256=PfkGLsw_6La3PXHOQL0u7qQsaZsXCJqYCeCcRDj5n64,6274
+sae_lens/saes/topk_sae.py,sha256=kmry1FE1H06OvCfn84V-j2JfWGKcU5b2urwAq_Oq5j4,9893
+sae_lens/tokenization_and_batching.py,sha256=oUAscjy_LPOrOb8_Ty6eLAcZ0B3HB_wiWjWktgolhG0,4314
+sae_lens/training/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+sae_lens/training/activation_scaler.py,sha256=seEE-2Qd2JMHxqgnsNWPt-DGtYGZxWPnOwCGuVNSOtI,1719
+sae_lens/training/activations_store.py,sha256=z8erbiB6ODbsqlu-bwEWbyj4XZvgsVgjCRBuQovqp2Q,32612
+sae_lens/training/mixing_buffer.py,sha256=vDpYG5ZE70szDvBsRKcNHEES3h_WTKJ16qDYk5jPOVA,2015
+sae_lens/training/optim.py,sha256=KXdOym-Ly3f2aFbndRc0JEH0Wa7u1BE5ljxGN3YtouQ,6836
+sae_lens/training/sae_trainer.py,sha256=9K0VudwSTJp9OlCVzaU_ngZ0WlYNrN6-ozTCCAxR9_k,15421
+sae_lens/training/types.py,sha256=qSjmGzXf3MLalygG0psnVjmhX_mpLmL47MQtZfe7qxg,81
+sae_lens/training/upload_saes_to_huggingface.py,sha256=r_WzI1zLtGZ5TzAxuG3xa_8T09j3zXJrWd_vzPsPGkQ,4469
+sae_lens/tutorial/tsea.py,sha256=fd1am_XXsf2KMbByDapJo-2qlxduKaa62Z2qcQZ3QKU,18145
+sae_lens/util.py,sha256=mCwLAilGMVo8Scm7CIsCafU7GsfmBvCcjwmloI4Ly7Y,1718
+sae_lens-6.0.0.dist-info/LICENSE,sha256=DW6e-hDosiu4CfW0-imI57sV1I5f9UEslpviNQcOAKs,1069
+sae_lens-6.0.0.dist-info/METADATA,sha256=LsDx7IEn0Mll0pUSKkBw6jlAXM9NwOVwseomhTb0AiE,5323
+sae_lens-6.0.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+sae_lens-6.0.0.dist-info/RECORD,,