sae-lens 6.6.0__tar.gz → 6.6.2__tar.gz

{sae_lens-6.6.0 → sae_lens-6.6.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: sae-lens
-Version: 6.6.0
+Version: 6.6.2
 Summary: Training and Analyzing Sparse Autoencoders (SAEs)
 License: MIT
 Keywords: deep-learning,sparse-autoencoders,mechanistic-interpretability,PyTorch
@@ -16,24 +16,19 @@ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Provides-Extra: mamba
 Requires-Dist: automated-interpretability (>=0.0.5,<1.0.0)
 Requires-Dist: babe (>=0.0.7,<0.0.8)
-Requires-Dist: datasets (>=3.1.0,<4.0.0)
+Requires-Dist: datasets (>=3.1.0)
 Requires-Dist: mamba-lens (>=0.0.4,<0.0.5) ; extra == "mamba"
-Requires-Dist: matplotlib (>=3.8.3,<4.0.0)
-Requires-Dist: matplotlib-inline (>=0.1.6,<0.2.0)
 Requires-Dist: nltk (>=3.8.1,<4.0.0)
-Requires-Dist: plotly (>=5.19.0,<6.0.0)
-Requires-Dist: plotly-express (>=0.4.1,<0.5.0)
-Requires-Dist: pytest-profiling (>=1.7.0,<2.0.0)
-Requires-Dist: python-dotenv (>=1.0.1,<2.0.0)
+Requires-Dist: plotly (>=5.19.0)
+Requires-Dist: plotly-express (>=0.4.1)
+Requires-Dist: python-dotenv (>=1.0.1)
 Requires-Dist: pyyaml (>=6.0.1,<7.0.0)
-Requires-Dist: pyzmq (==26.0.0)
-Requires-Dist: safetensors (>=0.4.2,<0.5.0)
+Requires-Dist: safetensors (>=0.4.2,<1.0.0)
 Requires-Dist: simple-parsing (>=0.1.6,<0.2.0)
+Requires-Dist: tenacity (>=9.0.0)
 Requires-Dist: transformer-lens (>=2.16.1,<3.0.0)
 Requires-Dist: transformers (>=4.38.1,<5.0.0)
-Requires-Dist: typer (>=0.12.3,<0.13.0)
 Requires-Dist: typing-extensions (>=4.10.0,<5.0.0)
-Requires-Dist: zstandard (>=0.22.0,<0.23.0)
 Project-URL: Homepage, https://jbloomaus.github.io/SAELens
 Project-URL: Repository, https://github.com/jbloomAus/SAELens
 Description-Content-Type: text/markdown

{sae_lens-6.6.0 → sae_lens-6.6.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "sae-lens"
-version = "6.6.0"
+version = "6.6.2"
 description = "Training and Analyzing Sparse Autoencoders (SAEs)"
 authors = ["Joseph Bloom"]
 readme = "README.md"
@@ -21,24 +21,19 @@ classifiers = ["Topic :: Scientific/Engineering :: Artificial Intelligence"]
 python = "^3.10"
 transformer-lens = "^2.16.1"
 transformers = "^4.38.1"
-plotly = "^5.19.0"
-plotly-express = "^0.4.1"
-matplotlib = "^3.8.3"
-matplotlib-inline = "^0.1.6"
-datasets = "^3.1.0"
+plotly = ">=5.19.0"
+plotly-express = ">=0.4.1"
+datasets = ">=3.1.0"
 babe = "^0.0.7"
 nltk = "^3.8.1"
-safetensors = "^0.4.2"
-typer = "^0.12.3"
+safetensors = ">=0.4.2,<1.0.0"
 mamba-lens = { version = "^0.0.4", optional = true }
-pyzmq = "26.0.0"
 automated-interpretability = ">=0.0.5,<1.0.0"
-python-dotenv = "^1.0.1"
+python-dotenv = ">=1.0.1"
 pyyaml = "^6.0.1"
-pytest-profiling = "^1.7.0"
-zstandard = "^0.22.0"
 typing-extensions = "^4.10.0"
 simple-parsing = "^0.1.6"
+tenacity = ">=9.0.0"
 
 [tool.poetry.group.dev.dependencies]
 pytest = "^8.0.2"

{sae_lens-6.6.0 → sae_lens-6.6.2}/sae_lens/__init__.py

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

{sae_lens-6.6.0 → sae_lens-6.6.2}/sae_lens/cache_activations_runner.py

@@ -82,7 +82,7 @@ class CacheActivationsRunner:
             )
             for hook_name in [self.cfg.hook_name]
         }
-        features_dict["token_ids"] = Sequence(
+        features_dict["token_ids"] = Sequence(  # type: ignore
             Value(dtype="int32"), length=self.context_size
         )
         self.features = Features(features_dict)

{sae_lens-6.6.0 → sae_lens-6.6.2}/sae_lens/evals.py

@@ -459,14 +459,16 @@ def get_sparsity_and_variance_metrics(
             original_act = cache[hook_name]
 
         # normalise if necessary (necessary in training only, otherwise we should fold the scaling in)
-        original_act = activation_scaler.scale(original_act)
+        original_act_scaled = activation_scaler.scale(original_act)
 
         # send the (maybe normalised) activations into the SAE
-        sae_feature_activations = sae.encode(original_act.to(sae.device))
-        sae_out = sae.decode(sae_feature_activations).to(original_act.device)
+        sae_feature_activations = sae.encode(original_act_scaled.to(sae.device))
+        sae_out_scaled = sae.decode(sae_feature_activations).to(
+            original_act_scaled.device
+        )
         del cache
 
-        sae_out = activation_scaler.unscale(sae_out)
+        sae_out = activation_scaler.unscale(sae_out_scaled)
 
         flattened_sae_input = einops.rearrange(original_act, "b ctx d -> (b ctx) d")
         flattened_sae_feature_acts = einops.rearrange(

{sae_lens-6.6.0 → sae_lens-6.6.2}/sae_lens/loading/pretrained_sae_loaders.py

@@ -1245,9 +1245,15 @@ def get_mwhanna_transcoder_config_from_hf(
     # Extract layer from folder name
     layer = int(folder_name.replace(".safetensors", "").split("_")[-1])
 
-    wandb_config_path = hf_hub_download(
-        repo_id, "wanb-config.yaml", force_download=force_download
-    )
+    try:
+        # mwhanna transcoders sometimes have a typo in the config file name, so check for both
+        wandb_config_path = hf_hub_download(
+            repo_id, "wanb-config.yaml", force_download=force_download
+        )
+    except EntryNotFoundError:
+        wandb_config_path = hf_hub_download(
+            repo_id, "wandb-config.yaml", force_download=force_download
+        )
     base_config_path = hf_hub_download(
         repo_id, "config.yaml", force_download=force_download
     )

{sae_lens-6.6.0 → sae_lens-6.6.2}/sae_lens/training/activations_store.py

@@ -289,7 +289,7 @@ class ActivationsStore:
                 "Dataset must have a 'tokens', 'input_ids', 'text', or 'problem' column."
             )
         if self.is_dataset_tokenized:
-            ds_context_size = len(dataset_sample[self.tokens_column])
+            ds_context_size = len(dataset_sample[self.tokens_column])  # type: ignore
             if ds_context_size < self.context_size:
                 raise ValueError(
                     f"""pretokenized dataset has context_size {ds_context_size}, but the provided context_size is {self.context_size}.