sae-lens 6.6.5__py3-none-any.whl → 6.8.0__py3-none-any.whl

sae_lens/__init__.py

@@ -1,5 +1,5 @@
 # ruff: noqa: E402
-__version__ = "6.6.5"
+__version__ = "6.8.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/loading/pretrained_sae_loaders.py

@@ -1252,11 +1252,11 @@ def get_mwhanna_transcoder_config_from_hf(
     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
+            repo_id, "wandb-config.yaml", force_download=force_download
         )
     except EntryNotFoundError:
         wandb_config_path = hf_hub_download(
-            repo_id, "wandb-config.yaml", force_download=force_download
+            repo_id, "wanb-config.yaml", force_download=force_download
         )
     try:
         base_config_path = hf_hub_download(
@@ -1330,6 +1330,66 @@ def mwhanna_transcoder_huggingface_loader(
     return cfg_dict, state_dict, None
 
 
+def mntss_clt_layer_huggingface_loader(
+    repo_id: str,
+    folder_name: str,
+    device: str = "cpu",
+    force_download: bool = False,  # noqa: ARG001
+    cfg_overrides: dict[str, Any] | None = None,
+) -> tuple[dict[str, Any], dict[str, torch.Tensor], torch.Tensor | None]:
+    """
+    Load a MNTSS CLT layer as a single layer transcoder.
+    The assumption is that the `folder_name` is the layer to load as an int
+    """
+    base_config_path = hf_hub_download(
+        repo_id, "config.yaml", force_download=force_download
+    )
+    with open(base_config_path) as f:
+        cfg_info: dict[str, Any] = yaml.safe_load(f)
+
+    # We need to actually load the weights, since the config is missing most information
+    encoder_path = hf_hub_download(
+        repo_id,
+        f"W_enc_{folder_name}.safetensors",
+        force_download=force_download,
+    )
+    decoder_path = hf_hub_download(
+        repo_id,
+        f"W_dec_{folder_name}.safetensors",
+        force_download=force_download,
+    )
+
+    encoder_state_dict = load_file(encoder_path, device=device)
+    decoder_state_dict = load_file(decoder_path, device=device)
+
+    with torch.no_grad():
+        state_dict = {
+            "W_enc": encoder_state_dict[f"W_enc_{folder_name}"].T,  # type: ignore
+            "b_enc": encoder_state_dict[f"b_enc_{folder_name}"],  # type: ignore
+            "b_dec": encoder_state_dict[f"b_dec_{folder_name}"],  # type: ignore
+            "W_dec": decoder_state_dict[f"W_dec_{folder_name}"].sum(dim=1),  # type: ignore
+        }
+
+    cfg_dict = {
+        "architecture": "transcoder",
+        "d_in": state_dict["b_dec"].shape[0],
+        "d_out": state_dict["b_dec"].shape[0],
+        "d_sae": state_dict["b_enc"].shape[0],
+        "dtype": "float32",
+        "device": device if device is not None else "cpu",
+        "activation_fn": "relu",
+        "normalize_activations": "none",
+        "model_name": cfg_info["model_name"],
+        "hook_name": f"blocks.{folder_name}.{cfg_info['feature_input_hook']}",
+        "hook_name_out": f"blocks.{folder_name}.{cfg_info['feature_output_hook']}",
+        "apply_b_dec_to_input": False,
+        "model_from_pretrained_kwargs": {"fold_ln": False},
+        **(cfg_overrides or {}),
+    }
+
+    return cfg_dict, state_dict, None
+
+
 NAMED_PRETRAINED_SAE_LOADERS: dict[str, PretrainedSaeHuggingfaceLoader] = {
     "sae_lens": sae_lens_huggingface_loader,
     "connor_rob_hook_z": connor_rob_hook_z_huggingface_loader,
@@ -1341,6 +1401,7 @@ NAMED_PRETRAINED_SAE_LOADERS: dict[str, PretrainedSaeHuggingfaceLoader] = {
     "sparsify": sparsify_huggingface_loader,
     "gemma_2_transcoder": gemma_2_transcoder_huggingface_loader,
     "mwhanna_transcoder": mwhanna_transcoder_huggingface_loader,
+    "mntss_clt_layer_transcoder": mntss_clt_layer_huggingface_loader,
 }
 
 

{sae_lens-6.6.5.dist-info → sae_lens-6.8.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: sae-lens
-Version: 6.6.5
+Version: 6.8.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.6.5.dist-info → sae_lens-6.8.0.dist-info}/RECORD

@@ -1,7 +1,7 @@
-sae_lens/__init__.py,sha256=gvg9photJRtatuXa9YF-uDv1tYiHwHTMh29X1GNQd6Y,3588
+sae_lens/__init__.py,sha256=NKyGkQ37vjuyhYLswmZd6xGFcukOw7lKHFATjgFidD0,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
@@ -9,7 +9,7 @@ sae_lens/evals.py,sha256=4hanbyG8qZLItWqft94F4ZjUoytPVB7fw5s0P4Oi0VE,39504
 sae_lens/llm_sae_training_runner.py,sha256=exxNX_OEhdiUrlgmBP9bjX9DOf0HUcNQGO4unKeDjKM,13713
 sae_lens/load_model.py,sha256=C8AMykctj6H7tz_xRwB06-EXj6TfW64PtSJZR5Jxn1Y,8649
 sae_lens/loading/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-sae_lens/loading/pretrained_sae_loaders.py,sha256=tLeHArWFpu8CI6vXH1ZxFkhmsrhO2UsZyi7DzVzqAUs,44477
+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
@@ -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.6.5.dist-info/LICENSE,sha256=DW6e-hDosiu4CfW0-imI57sV1I5f9UEslpviNQcOAKs,1069
-sae_lens-6.6.5.dist-info/METADATA,sha256=U5oP3RYgIE2EnHA2mwRImUcoyVBhYYwiRU199LM_R7c,5356
-sae_lens-6.6.5.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-sae_lens-6.6.5.dist-info/RECORD,,
+sae_lens-6.8.0.dist-info/LICENSE,sha256=DW6e-hDosiu4CfW0-imI57sV1I5f9UEslpviNQcOAKs,1069
+sae_lens-6.8.0.dist-info/METADATA,sha256=DkYerQAH4AseU876uNcejPkZu4I3M2EscDlgG1NWwc0,5244
+sae_lens-6.8.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+sae_lens-6.8.0.dist-info/RECORD,,