sae-lens 6.6.3__tar.gz → 6.6.4__tar.gz

{sae_lens-6.6.3 → sae_lens-6.6.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: sae-lens
-Version: 6.6.3
+Version: 6.6.4
 Summary: Training and Analyzing Sparse Autoencoders (SAEs)
 License: MIT
 Keywords: deep-learning,sparse-autoencoders,mechanistic-interpretability,PyTorch
@@ -14,7 +14,8 @@ Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Provides-Extra: mamba
-Requires-Dist: automated-interpretability (>=0.0.5,<1.0.0)
+Provides-Extra: neuronpedia
+Requires-Dist: automated-interpretability (>=0.0.5,<1.0.0) ; extra == "neuronpedia"
 Requires-Dist: babe (>=0.0.7,<0.0.8)
 Requires-Dist: datasets (>=3.1.0)
 Requires-Dist: mamba-lens (>=0.0.4,<0.0.5) ; extra == "mamba"

{sae_lens-6.6.3 → sae_lens-6.6.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "sae-lens"
-version = "6.6.3"
+version = "6.6.4"
 description = "Training and Analyzing Sparse Autoencoders (SAEs)"
 authors = ["Joseph Bloom"]
 readme = "README.md"
@@ -28,12 +28,12 @@ babe = "^0.0.7"
 nltk = "^3.8.1"
 safetensors = ">=0.4.2,<1.0.0"
 mamba-lens = { version = "^0.0.4", optional = true }
-automated-interpretability = ">=0.0.5,<1.0.0"
 python-dotenv = ">=1.0.1"
 pyyaml = "^6.0.1"
 typing-extensions = "^4.10.0"
 simple-parsing = "^0.1.6"
 tenacity = ">=9.0.0"
+automated-interpretability = { version = ">=0.0.5,<1.0.0", optional = true }
 
 [tool.poetry.group.dev.dependencies]
 pytest = "^8.0.2"
@@ -56,9 +56,11 @@ ruff = "^0.7.4"
 eai-sparsify = "^1.1.1"
 mike = "^2.0.0"
 trio = "^0.30.0"
+automated-interpretability = "^0.0.13"
 
 [tool.poetry.extras]
 mamba = ["mamba-lens"]
+neuronpedia = ["automated-interpretability"]
 
 [tool.ruff.lint]
 exclude = ["*.ipynb"]

{sae_lens-6.6.3 → sae_lens-6.6.4}/sae_lens/__init__.py

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

{sae_lens-6.6.3 → sae_lens-6.6.4}/sae_lens/analysis/neuronpedia_integration.py

@@ -8,27 +8,6 @@ from typing import Any, TypeVar
 
 import requests
 from dotenv import load_dotenv
-from neuron_explainer.activations.activation_records import calculate_max_activation
-from neuron_explainer.activations.activations import ActivationRecord
-from neuron_explainer.explanations.calibrated_simulator import (
-    UncalibratedNeuronSimulator,
-)
-from neuron_explainer.explanations.explainer import (
-    HARMONY_V4_MODELS,
-    ContextSize,
-    TokenActivationPairExplainer,
-)
-from neuron_explainer.explanations.explanations import ScoredSimulation
-from neuron_explainer.explanations.few_shot_examples import FewShotExampleSet
-from neuron_explainer.explanations.prompt_builder import PromptFormat
-from neuron_explainer.explanations.scoring import (
-    _simulate_and_score_sequence,
-    aggregate_scored_sequence_simulations,
-)
-from neuron_explainer.explanations.simulator import (
-    LogprobFreeExplanationTokenSimulator,
-    NeuronSimulator,
-)
 from tenacity import retry, stop_after_attempt, wait_random_exponential
 
 from sae_lens import SAE, logger
@@ -158,10 +137,22 @@ def sleep_identity(x: T) -> T:
 
 @retry(wait=wait_random_exponential(min=1, max=500), stop=stop_after_attempt(10))
 async def simulate_and_score(  # type: ignore
-    simulator: NeuronSimulator,
-    activation_records: list[ActivationRecord],  # type: ignore
-) -> ScoredSimulation:  # type: ignore
+    simulator: Any,
+    activation_records: list[Any],
+) -> Any:
     """Score an explanation of a neuron by how well it predicts activations on the given text sequences."""
+    try:
+        from neuron_explainer.explanations.scoring import (
+            _simulate_and_score_sequence,
+            aggregate_scored_sequence_simulations,
+        )
+    except ImportError as e:
+        raise ImportError(
+            "The neuron_explainer package is required to use this function. "
+            "Please install SAELens with the neuronpedia optional dependencies: "
+            "pip install sae-lens[neuronpedia]"
+        ) from e
+
     scored_sequence_simulations = await asyncio.gather(
         *[
             sleep_identity(
@@ -253,6 +244,31 @@ async def autointerp_neuronpedia_features(  # noqa: C901
     Returns:
         None
     """
+    try:
+        from neuron_explainer.activations.activation_records import (
+            calculate_max_activation,
+        )
+        from neuron_explainer.activations.activations import ActivationRecord
+        from neuron_explainer.explanations.calibrated_simulator import (
+            UncalibratedNeuronSimulator,
+        )
+        from neuron_explainer.explanations.explainer import (
+            HARMONY_V4_MODELS,
+            ContextSize,
+            TokenActivationPairExplainer,
+        )
+        from neuron_explainer.explanations.few_shot_examples import FewShotExampleSet
+        from neuron_explainer.explanations.prompt_builder import PromptFormat
+        from neuron_explainer.explanations.simulator import (
+            LogprobFreeExplanationTokenSimulator,
+        )
+    except ImportError as e:
+        raise ImportError(
+            "The automated-interpretability package is required to use autointerp functionality. "
+            "Please install SAELens with the neuronpedia optional dependencies: "
+            "pip install sae-lens[neuronpedia]"
+        ) from e
+
     logger.info("\n\n")
 
     if os.getenv("OPENAI_API_KEY") is None: