sae-lens 6.7.0__py3-none-any.whl → 6.9.0__py3-none-any.whl

sae_lens/__init__.py

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

sae_lens/analysis/neuronpedia_integration.py

@@ -1,14 +1,10 @@
-import asyncio
 import json
-import os
 import urllib.parse
 import webbrowser
-from datetime import datetime
-from typing import Any, TypeVar
+from typing import Any
 
 import requests
 from dotenv import load_dotenv
-from tenacity import retry, stop_after_attempt, wait_random_exponential
 
 from sae_lens import SAE, logger
 
@@ -126,47 +122,6 @@ class NeuronpediaFeature:
         return any(max(activation.act_values) > 0 for activation in self.activations)
 
 
-T = TypeVar("T")
-
-
-@retry(wait=wait_random_exponential(min=1, max=500), stop=stop_after_attempt(10))
-def sleep_identity(x: T) -> T:
-    """Dummy function for retrying."""
-    return x
-
-
-@retry(wait=wait_random_exponential(min=1, max=500), stop=stop_after_attempt(10))
-async def simulate_and_score(  # 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(
-                _simulate_and_score_sequence(
-                    simulator,
-                    activation_record,
-                )
-            )
-            for activation_record in activation_records
-        ]
-    )
-    return aggregate_scored_sequence_simulations(scored_sequence_simulations)
-
-
 def make_neuronpedia_list_with_features(
     api_key: str,
     list_name: str,
@@ -206,305 +161,3 @@ def test_key(api_key: str):
     response = requests.post(url, json=body)
     if response.status_code != 200:
         raise Exception("Neuronpedia API key is not valid.")
-
-
-async def autointerp_neuronpedia_features(  # noqa: C901
-    features: list[NeuronpediaFeature],
-    openai_api_key: str | None = None,
-    autointerp_retry_attempts: int = 3,
-    autointerp_score_max_concurrent: int = 20,
-    neuronpedia_api_key: str | None = None,
-    skip_neuronpedia_api_key_test: bool = False,
-    do_score: bool = True,
-    output_dir: str = "neuronpedia_outputs/autointerp",
-    num_activations_to_use: int = 20,
-    max_explanation_activation_records: int = 20,
-    upload_to_neuronpedia: bool = True,
-    autointerp_explainer_model_name: str = "gpt-4-1106-preview",
-    autointerp_scorer_model_name: str | None = "gpt-3.5-turbo",
-    save_to_disk: bool = True,
-):
-    """
-    Autointerp Neuronpedia features.
-
-    Args:
-        features: List of NeuronpediaFeature objects.
-        openai_api_key: OpenAI API key.
-        autointerp_retry_attempts: Number of retry attempts for autointerp.
-        autointerp_score_max_concurrent: Maximum number of concurrent requests for autointerp scoring.
-        neuronpedia_api_key: Neuronpedia API key.
-        do_score: Whether to score the features.
-        output_dir: Output directory for saving the results.
-        num_activations_to_use: Number of activations to use.
-        max_explanation_activation_records: Maximum number of activation records for explanation.
-        upload_to_neuronpedia: Whether to upload the results to Neuronpedia.
-        autointerp_explainer_model_name: Model name for autointerp explainer.
-        autointerp_scorer_model_name: Model name for autointerp scorer.
-
-    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:
-        if openai_api_key is None:
-            raise Exception(
-                "You need to provide an OpenAI API key either in environment variable OPENAI_API_KEY or as an argument."
-            )
-        os.environ["OPENAI_API_KEY"] = openai_api_key
-
-    if autointerp_explainer_model_name not in HARMONY_V4_MODELS:
-        raise Exception(
-            f"Invalid explainer model name: {autointerp_explainer_model_name}. Must be one of: {HARMONY_V4_MODELS}"
-        )
-
-    if do_score and autointerp_scorer_model_name not in HARMONY_V4_MODELS:
-        raise Exception(
-            f"Invalid scorer model name: {autointerp_scorer_model_name}. Must be one of: {HARMONY_V4_MODELS}"
-        )
-
-    if upload_to_neuronpedia:
-        if neuronpedia_api_key is None:
-            raise Exception(
-                "You need to provide a Neuronpedia API key to upload the results to Neuronpedia."
-            )
-        if not skip_neuronpedia_api_key_test:
-            test_key(neuronpedia_api_key)
-
-    logger.info("\n\n=== Step 1) Fetching features from Neuronpedia")
-    for feature in features:
-        feature_data = get_neuronpedia_feature(
-            feature=feature.feature,
-            layer=feature.layer,
-            model=feature.modelId,
-            dataset=feature.dataset,
-        )
-
-        if "modelId" not in feature_data:
-            raise Exception(
-                f"Feature {feature.feature} in layer {feature.layer} of model {feature.modelId} and dataset {feature.dataset} does not exist."
-            )
-
-        if "activations" not in feature_data or len(feature_data["activations"]) == 0:
-            raise Exception(
-                f"Feature {feature.feature} in layer {feature.layer} of model {feature.modelId} and dataset {feature.dataset} does not have activations."
-            )
-
-        activations = feature_data["activations"]
-        activations_to_add = []
-        for activation in activations:
-            if len(activations_to_add) < num_activations_to_use:
-                activations_to_add.append(
-                    NeuronpediaActivation(
-                        id=activation["id"],
-                        tokens=activation["tokens"],
-                        act_values=activation["values"],
-                    )
-                )
-        feature.activations = activations_to_add
-
-        if not feature.has_activating_text():
-            raise Exception(
-                f"Feature {feature.modelId}@{feature.layer}-{feature.dataset}:{feature.feature} appears dead - it does not have activating text."
-            )
-
-    for iteration_num, feature in enumerate(features):
-        start_time = datetime.now()
-
-        logger.info(
-            f"\n========== Feature {feature.modelId}@{feature.layer}-{feature.dataset}:{feature.feature} ({iteration_num + 1} of {len(features)} Features) =========="
-        )
-        logger.info(
-            f"\n=== Step 2) Explaining feature {feature.modelId}@{feature.layer}-{feature.dataset}:{feature.feature}"
-        )
-
-        if feature.activations is None:
-            feature.activations = []
-        activation_records = [
-            ActivationRecord(
-                tokens=activation.tokens,  # type: ignore
-                activations=activation.act_values,  # type: ignore
-            )  # type: ignore
-            for activation in feature.activations
-        ]
-
-        activation_records_explaining = activation_records[
-            :max_explanation_activation_records
-        ]
-
-        explainer = TokenActivationPairExplainer(
-            model_name=autointerp_explainer_model_name,
-            prompt_format=PromptFormat.HARMONY_V4,
-            context_size=ContextSize.SIXTEEN_K,
-            max_concurrent=1,
-        )
-
-        explanations = []
-        for _ in range(autointerp_retry_attempts):
-            try:
-                explanations = await explainer.generate_explanations(
-                    all_activation_records=activation_records_explaining,
-                    max_activation=calculate_max_activation(
-                        activation_records_explaining
-                    ),
-                    num_samples=1,
-                )
-            except Exception as e:
-                logger.error(f"ERROR, RETRYING: {e}")
-            else:
-                break
-        else:
-            logger.error(
-                f"ERROR: Failed to explain feature {feature.modelId}@{feature.layer}-{feature.dataset}:{feature.feature}"
-            )
-
-        if len(explanations) != 1:
-            raise ValueError(
-                f"Expected exactly one explanation but got {len(explanations)}. This may indicate an issue with the explainer's response."
-            )
-        explanation = explanations[0].rstrip(".")
-        logger.info(
-            f"===== {autointerp_explainer_model_name}'s explanation: {explanation}"
-        )
-        feature.autointerp_explanation = explanation
-
-        scored_simulation = None
-        if do_score and autointerp_scorer_model_name:
-            logger.info(
-                f"\n=== Step 3) Scoring feature {feature.modelId}@{feature.layer}-{feature.dataset}:{feature.feature}"
-            )
-            logger.info("=== This can take up to 30 seconds.")
-
-            temp_activation_records = [
-                ActivationRecord(
-                    tokens=[  # type: ignore
-                        token.replace("<|endoftext|>", "<|not_endoftext|>")
-                        .replace(" 55", "_55")
-                        .encode("ascii", errors="backslashreplace")
-                        .decode("ascii")
-                        for token in activation_record.tokens  # type: ignore
-                    ],
-                    activations=activation_record.activations,  # type: ignore
-                )  # type: ignore
-                for activation_record in activation_records
-            ]
-
-            score = None
-            scored_simulation = None
-            for _ in range(autointerp_retry_attempts):
-                try:
-                    simulator = UncalibratedNeuronSimulator(
-                        LogprobFreeExplanationTokenSimulator(
-                            autointerp_scorer_model_name,
-                            explanation,
-                            json_mode=True,
-                            max_concurrent=autointerp_score_max_concurrent,
-                            few_shot_example_set=FewShotExampleSet.JL_FINE_TUNED,
-                            prompt_format=PromptFormat.HARMONY_V4,
-                        )
-                    )
-                    scored_simulation = await simulate_and_score(
-                        simulator, temp_activation_records
-                    )
-                    score = scored_simulation.get_preferred_score()
-                except Exception as e:
-                    logger.error(f"ERROR, RETRYING: {e}")
-                else:
-                    break
-
-            if (
-                score is None
-                or scored_simulation is None
-                or len(scored_simulation.scored_sequence_simulations)
-                != num_activations_to_use
-            ):
-                logger.error(
-                    f"ERROR: Failed to score feature {feature.modelId}@{feature.layer}-{feature.dataset}:{feature.feature}. Skipping it."
-                )
-                continue
-            feature.autointerp_explanation_score = score
-            logger.info(
-                f"===== {autointerp_scorer_model_name}'s score: {(score * 100):.0f}"
-            )
-
-        else:
-            logger.info("=== Step 3) Skipping scoring as instructed.")
-
-        feature_data = {
-            "modelId": feature.modelId,
-            "layer": f"{feature.layer}-{feature.dataset}",
-            "index": feature.feature,
-            "explanation": feature.autointerp_explanation,
-            "explanationScore": feature.autointerp_explanation_score,
-            "explanationModel": autointerp_explainer_model_name,
-        }
-        if do_score and autointerp_scorer_model_name and scored_simulation:
-            feature_data["activations"] = feature.activations
-            feature_data["simulationModel"] = autointerp_scorer_model_name
-            feature_data["simulationActivations"] = (
-                scored_simulation.scored_sequence_simulations
-            )  # type: ignore
-            feature_data["simulationScore"] = feature.autointerp_explanation_score
-        feature_data_str = json.dumps(feature_data, default=vars)
-
-        if save_to_disk:
-            output_file = f"{output_dir}/{feature.modelId}-{feature.layer}-{feature.dataset}_feature-{feature.feature}_time-{datetime.now().strftime('%Y%m%d-%H%M%S')}.jsonl"
-            os.makedirs(output_dir, exist_ok=True)
-            logger.info(f"\n=== Step 4) Saving feature to {output_file}")
-            with open(output_file, "a") as f:
-                f.write(feature_data_str)
-                f.write("\n")
-        else:
-            logger.info("\n=== Step 4) Skipping saving to disk.")
-
-        if upload_to_neuronpedia:
-            logger.info("\n=== Step 5) Uploading feature to Neuronpedia")
-            upload_data = json.dumps(
-                {
-                    "feature": feature_data,
-                },
-                default=vars,
-            )
-            upload_data_json = json.loads(upload_data, parse_constant=NanAndInfReplacer)
-            url = f"{NEURONPEDIA_DOMAIN}/api/explanation/new"
-            response = requests.post(
-                url, json=upload_data_json, headers={"x-api-key": neuronpedia_api_key}
-            )
-            if response.status_code != 200:
-                logger.error(
-                    f"ERROR: Couldn't upload explanation to Neuronpedia: {response.text}"
-                )
-            else:
-                logger.info(
-                    f"===== Uploaded to Neuronpedia: {NEURONPEDIA_DOMAIN}/{feature.modelId}/{feature.layer}-{feature.dataset}/{feature.feature}"
-                )
-
-        end_time = datetime.now()
-        logger.info(f"\n========== Time Spent for Feature: {end_time - start_time}\n")
-
-    logger.info("\n\n========== Generation and Upload Complete ==========\n\n")

sae_lens/pretrained_saes.yaml

@@ -14744,4 +14744,88 @@ mwhanna-qwen3-0.6b-transcoders-lowl0:
     neuronpedia: qwen3-0.6b/26-transcoder-hp-lowl0
   - id: layer_27
     path: layer_27.safetensors
-    neuronpedia: qwen3-0.6b/27-transcoder-hp-lowl0
+    neuronpedia: qwen3-0.6b/27-transcoder-hp-lowl0
+
+mntss-gemma-2-2b-2.5m-clt-as-per-layer:
+  conversion_func: mntss_clt_layer_transcoder
+  model: gemma-2-2b
+  repo_id: mntss/clt-gemma-2-2b-2.5M
+  saes:
+  - id: layer_0
+    path: 0
+    neuronpedia: gemma-2-2b/0-clt-hp
+  - id: layer_1
+    path: 1
+    neuronpedia: gemma-2-2b/1-clt-hp
+  - id: layer_2
+    path: 2
+    neuronpedia: gemma-2-2b/2-clt-hp
+  - id: layer_3
+    path: 3
+    neuronpedia: gemma-2-2b/3-clt-hp
+  - id: layer_4
+    path: 4
+    neuronpedia: gemma-2-2b/4-clt-hp
+  - id: layer_5
+    path: 5
+    neuronpedia: gemma-2-2b/5-clt-hp
+  - id: layer_6
+    path: 6
+    neuronpedia: gemma-2-2b/6-clt-hp
+  - id: layer_7
+    path: 7
+    neuronpedia: gemma-2-2b/7-clt-hp
+  - id: layer_8
+    path: 8
+    neuronpedia: gemma-2-2b/8-clt-hp
+  - id: layer_9
+    path: 9
+    neuronpedia: gemma-2-2b/9-clt-hp
+  - id: layer_10
+    path: 10
+    neuronpedia: gemma-2-2b/10-clt-hp
+  - id: layer_11
+    path: 11
+    neuronpedia: gemma-2-2b/11-clt-hp
+  - id: layer_12
+    path: 12
+    neuronpedia: gemma-2-2b/12-clt-hp
+  - id: layer_13
+    path: 13
+    neuronpedia: gemma-2-2b/13-clt-hp
+  - id: layer_14
+    path: 14
+    neuronpedia: gemma-2-2b/14-clt-hp
+  - id: layer_15
+    path: 15
+    neuronpedia: gemma-2-2b/15-clt-hp
+  - id: layer_16
+    path: 16
+    neuronpedia: gemma-2-2b/16-clt-hp
+  - id: layer_17
+    path: 17
+    neuronpedia: gemma-2-2b/17-clt-hp
+  - id: layer_18
+    path: 18
+    neuronpedia: gemma-2-2b/18-clt-hp
+  - id: layer_19
+    path: 19
+    neuronpedia: gemma-2-2b/19-clt-hp
+  - id: layer_20
+    path: 20
+    neuronpedia: gemma-2-2b/20-clt-hp
+  - id: layer_21
+    path: 21
+    neuronpedia: gemma-2-2b/21-clt-hp
+  - id: layer_22
+    path: 22
+    neuronpedia: gemma-2-2b/22-clt-hp
+  - id: layer_23
+    path: 23
+    neuronpedia: gemma-2-2b/23-clt-hp
+  - id: layer_24
+    path: 24
+    neuronpedia: gemma-2-2b/24-clt-hp
+  - id: layer_25
+    path: 25
+    neuronpedia: gemma-2-2b/25-clt-hp

{sae_lens-6.7.0.dist-info → sae_lens-6.9.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: sae-lens
-Version: 6.7.0
+Version: 6.9.0
 Summary: Training and Analyzing Sparse Autoencoders (SAEs)
 License: MIT
 Keywords: deep-learning,sparse-autoencoders,mechanistic-interpretability,PyTorch
@@ -14,8 +14,6 @@ Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Provides-Extra: mamba
-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.7.0.dist-info → sae_lens-6.9.0.dist-info}/RECORD

@@ -1,7 +1,7 @@
-sae_lens/__init__.py,sha256=JT234I0wGYcCdf_zNYt-22puv4Tjpz9tuftPnr6SXqM,3588
+sae_lens/__init__.py,sha256=TQtVpMEI_qo11x4o5w2cfOF9pqHeXXtyMCypr3lSRCI,3588
 sae_lens/analysis/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 sae_lens/analysis/hooked_sae_transformer.py,sha256=vRu6JseH1lZaEeILD5bEkQEQ1wYHHDcxD-f2olKmE9Y,14275
-sae_lens/analysis/neuronpedia_integration.py,sha256=Fj4gVyaXMGBUxoK0vPeTwGVFr4n40fmfPrRENo4WzPs,19324
+sae_lens/analysis/neuronpedia_integration.py,sha256=Gx1W7hUBEuMoasNcnOnZ1wmqbXDd1pSZ1nqKEya1HQc,4962
 sae_lens/cache_activations_runner.py,sha256=cNeAtp2JQ_vKbeddZVM-tcPLYyyfTWL8NDna5KQpkLI,12583
 sae_lens/config.py,sha256=IrjbsKBbaZoFXYrsPJ5xBwIqi9uZJIIFXjV_uoErJaE,28176
 sae_lens/constants.py,sha256=CSjmiZ-bhjQeVLyRvWxAjBokCgkfM8mnvd7-vxLIWTY,639
@@ -12,7 +12,7 @@ sae_lens/loading/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,
 sae_lens/loading/pretrained_sae_loaders.py,sha256=IVtgxWN0w96ZORnWPYW2ndYWey7e5GpzlpedWF3NJ8k,46818
 sae_lens/loading/pretrained_saes_directory.py,sha256=4Vn-Jex6SveD7EbxcSOBv8cx1gkPfUMLU1QOP-ww1ZE,3752
 sae_lens/pretokenize_runner.py,sha256=w0f6SfZLAxbp5eAAKnet8RqUB_DKofZ9RGsoJwFnYbA,7058
-sae_lens/pretrained_saes.yaml,sha256=O_FwoOe7fU9_WLEOnMk1IWXRxD4nwzf1tCfbof1r0D0,598578
+sae_lens/pretrained_saes.yaml,sha256=d6FYfWTdVAPlOCM55C1ICS6lF9nWPPVNwjlXCa9p7NU,600468
 sae_lens/registry.py,sha256=nhy7BPSudSATqW4lo9H_k3Na7sfGHmAf9v-3wpnLL_o,1490
 sae_lens/saes/__init__.py,sha256=jVwazK8Q6dW5J6_zFXPoNAuBvSxgziQ8eMOjGM3t-X8,1475
 sae_lens/saes/batchtopk_sae.py,sha256=CyaFG2hMyyDaEaXXrAMJC8wQDW1JoddTKF5mvxxBQKY,3395
@@ -33,7 +33,7 @@ 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.7.0.dist-info/LICENSE,sha256=DW6e-hDosiu4CfW0-imI57sV1I5f9UEslpviNQcOAKs,1069
-sae_lens-6.7.0.dist-info/METADATA,sha256=K5tIuWNs4ho7aEsNfimzbInpxi-eX_hhRneAGw22xVE,5356
-sae_lens-6.7.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-sae_lens-6.7.0.dist-info/RECORD,,
+sae_lens-6.9.0.dist-info/LICENSE,sha256=DW6e-hDosiu4CfW0-imI57sV1I5f9UEslpviNQcOAKs,1069
+sae_lens-6.9.0.dist-info/METADATA,sha256=xqfc_6msZd2ZGEwnLFVBDAK2HdFrvllKwa78bx6iYhs,5244
+sae_lens-6.9.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+sae_lens-6.9.0.dist-info/RECORD,,