sae-lens 6.5.0__tar.gz → 6.5.2__tar.gz

{sae_lens-6.5.0 → sae_lens-6.5.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: sae-lens
-Version: 6.5.0
+Version: 6.5.2
 Summary: Training and Analyzing Sparse Autoencoders (SAEs)
 License: MIT
 Keywords: deep-learning,sparse-autoencoders,mechanistic-interpretability,PyTorch
@@ -16,7 +16,7 @@ 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 (>=2.17.1,<3.0.0)
+Requires-Dist: datasets (>=3.1.0,<4.0.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)

{sae_lens-6.5.0 → sae_lens-6.5.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "sae-lens"
-version = "6.5.0"
+version = "6.5.2"
 description = "Training and Analyzing Sparse Autoencoders (SAEs)"
 authors = ["Joseph Bloom"]
 readme = "README.md"
@@ -25,7 +25,7 @@ plotly = "^5.19.0"
 plotly-express = "^0.4.1"
 matplotlib = "^3.8.3"
 matplotlib-inline = "^0.1.6"
-datasets = "^2.17.1"
+datasets = "^3.1.0"
 babe = "^0.0.7"
 nltk = "^3.8.1"
 safetensors = "^0.4.2"

{sae_lens-6.5.0 → sae_lens-6.5.2}/sae_lens/__init__.py

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

{sae_lens-6.5.0 → sae_lens-6.5.2}/sae_lens/evals.py

@@ -776,6 +776,7 @@ def multiple_evals(
     n_eval_sparsity_variance_batches: int,
     eval_batch_size_prompts: int = 8,
     datasets: list[str] = ["Skylion007/openwebtext", "lighteval/MATH"],
+    dataset_trust_remote_code: bool = False,
     ctx_lens: list[int] = [128],
     output_dir: str = "eval_results",
     verbose: bool = False,
@@ -822,7 +823,11 @@ def multiple_evals(
         for ctx_len in ctx_lens:
             for dataset in datasets:
                 activation_store = ActivationsStore.from_sae(
-                    current_model, sae, context_size=ctx_len, dataset=dataset
+                    current_model,
+                    sae,
+                    context_size=ctx_len,
+                    dataset=dataset,
+                    dataset_trust_remote_code=dataset_trust_remote_code,
                 )
                 activation_store.shuffle_input_dataset(seed=42)
 
@@ -882,6 +887,7 @@ def run_evaluations(args: argparse.Namespace) -> list[dict[str, Any]]:
         n_eval_sparsity_variance_batches=args.n_eval_sparsity_variance_batches,
         eval_batch_size_prompts=args.batch_size_prompts,
         datasets=args.datasets,
+        dataset_trust_remote_code=args.dataset_trust_remote_code,
         ctx_lens=args.ctx_lens,
         output_dir=args.output_dir,
         verbose=args.verbose,
@@ -1004,6 +1010,11 @@ def process_args(args: list[str]) -> argparse.Namespace:
         default=["Skylion007/openwebtext"],
         help="Datasets to evaluate on, such as 'Skylion007/openwebtext' or 'lighteval/MATH'.",
     )
+    arg_parser.add_argument(
+        "--dataset_trust_remote_code",
+        action="store_true",
+        help="Allow execution of remote code when loading datasets for evaluation.",
+    )
     arg_parser.add_argument(
         "--ctx_lens",
         nargs="+",

{sae_lens-6.5.0 → sae_lens-6.5.2}/sae_lens/load_model.py

@@ -52,7 +52,7 @@ def load_model(
     if model_class_name == "AutoModelForCausalLM":
         hf_model = AutoModelForCausalLM.from_pretrained(
             model_name, **model_from_pretrained_kwargs
-        ).to(device)
+        ).to(device)  # type: ignore
         tokenizer = AutoTokenizer.from_pretrained(model_name)
         return HookedProxyLM(hf_model, tokenizer)