Using the Interpretune Neuronpedia Extension#
This tutorial demonstrates Interpretune’s integration with Circuit Tracer and Neuronpedia using Gemma-3-1B-IT, guiding the user through generating and sharing a basic model attribution graph. Future tutorials will demonstrate the use of Interpretune’s analytical abstractions to generate, transform and share more advanced model attribution graphs in the context of world model analysis, research or AI agent/application interpretability maximization.
Notebook Parameters#
This cell contains parameters that can be modified for different test configurations using papermill.
# Parameters - These will be injected by papermill during parameterized test runs
use_baseline_salient_logits = True # logits computation mode: True->salient logits, False->specific logits
use_baseline_transcoder_arch = True # transcoder architecture: True->SingleLayerTranscoder, False->CrossLayerTranscoder
core_log_dir = None # Directory to save analysis logs (if None, a temp directory will be created)
# Core imports
import interpretune as it # registered analysis ops will be available as it.<op> when analysis is imported
# this example assumes the user has either installed `interpretune` with the `examples` extra
# (or that circuit_tracer and neuronpedia are installed separately in addition to base interpretune dependencies)
# Import circuit tracer and required modules
from transformer_lens import ActivationCache # noqa: F401
from datetime import datetime
from it_examples import _ACTIVE_PATCHES # noqa: F401 # TODO: add note about this unless patched in SL before release
from it_examples.example_module_registry import MODULE_EXAMPLE_REGISTRY # TODO: move to hub once implemented
from it_examples.utils.example_helpers import required_os_env
from it_examples.example_prompt_configs import GemmaPromptConfig
from interpretune import ITSessionConfig, ITSession
from interpretune.base.call import it_init
import os
env_path: str | None = None # set to '/full/path/to/.env' to override
# Note: We need to set your API keys via environment variables:
# export NEURONPEDIA_API_KEY="your-production-key"
# export DEV_NEURONPEDIA_API_KEY="your-dev-key" # if using dev API
use_localhost = os.environ.get("USE_LOCALHOST", "false").lower() == "true"
# Note there is a planned request for neuronpedia to qualify USE_LOCALHOST as NP_USE_LOCALHOST
os_env_reqs = [("DEV_NEURONPEDIA_API_KEY", "non-empty")] if use_localhost else [("NEURONPEDIA_API_KEY", "non-empty")]
assert required_os_env(env_path=env_path, env_reqs=os_env_reqs)
# Load our demo config (this will be done from the hub once that is available)
base_itdm_cfg, base_it_cfg, dm_cls, m_cls = MODULE_EXAMPLE_REGISTRY.get("gemma3.rte_demo.circuit_tracer_w_neuronpedia")
# Optionally override base_it_cfg.core_log_dir with the notebook parameter if provided
if core_log_dir:
base_it_cfg.core_log_dir = core_log_dir
# configure our session with our desired adapter composition, core and circuit_tracer in this case
session_cfg = ITSessionConfig(
adapter_ctx=(it.Adapter.core, it.Adapter.nnsight, it.Adapter.circuit_tracer),
datamodule_cfg=base_itdm_cfg,
module_cfg=base_it_cfg,
datamodule_cls=dm_cls,
module_cls=m_cls,
)
# start our session
it_session = ITSession(session_cfg)
[INFO] interpretune.extensions.neuronpedia: Neuronpedia package available
INFO:interpretune.extensions.neuronpedia:Neuronpedia package available
[INFO] interpretune.utils.logging: Loading ReplacementModel with backend: nnsight
INFO:interpretune.utils.logging:Loading ReplacementModel with backend: nnsight
[INFO] interpretune.utils.logging: NNsight ReplacementModel initialized for Circuit Tracer
INFO:interpretune.utils.logging:NNsight ReplacementModel initialized for Circuit Tracer
[INFO] interpretune.utils.logging: Attempted to clean a key that was not present, continuing without cleaning that key: 'Gemma3TextConfig' object has no attribute 'quantization_config'
INFO:interpretune.utils.logging:Attempted to clean a key that was not present, continuing without cleaning that key: 'Gemma3TextConfig' object has no attribute 'quantization_config'
[INFO] interpretune.utils.logging: Attempted to clean a key that was not present, continuing without cleaning that key: 'Gemma3TextConfig' object has no attribute '_pre_quantization_dtype'
INFO:interpretune.utils.logging:Attempted to clean a key that was not present, continuing without cleaning that key: 'Gemma3TextConfig' object has no attribute '_pre_quantization_dtype'
# manual init for now
it_init(**it_session)
print("\nIT Session initialized successfully!")
[INFO] interpretune.utils.logging: Preparing data: InterpretunableDataModule
INFO:interpretune.utils.logging:Preparing data: InterpretunableDataModule
[INFO] interpretune.utils.logging: The following columns don't have a corresponding argument in `NNSightReplacementModel.forward` and have been ignored: hypothesis, idx, label, premise, sequences. If hypothesis, idx, label, premise, sequences are not expected by `NNSightReplacementModel.forward`, you can safely ignore this message.
INFO:interpretune.utils.logging:The following columns don't have a corresponding argument in `NNSightReplacementModel.forward` and have been ignored: hypothesis, idx, label, premise, sequences. If hypothesis, idx, label, premise, sequences are not expected by `NNSightReplacementModel.forward`, you can safely ignore this message.
[INFO] interpretune.utils.logging: The following columns don't have a corresponding argument in `NNSightReplacementModel.forward` and have been ignored: hypothesis, idx, label, premise, sequences. If hypothesis, idx, label, premise, sequences are not expected by `NNSightReplacementModel.forward`, you can safely ignore this message.
INFO:interpretune.utils.logging:The following columns don't have a corresponding argument in `NNSightReplacementModel.forward` and have been ignored: hypothesis, idx, label, premise, sequences. If hypothesis, idx, label, premise, sequences are not expected by `NNSightReplacementModel.forward`, you can safely ignore this message.
[INFO] interpretune.utils.logging: The following columns don't have a corresponding argument in `NNSightReplacementModel.forward` and have been ignored: hypothesis, idx, label, premise, sequences. If hypothesis, idx, label, premise, sequences are not expected by `NNSightReplacementModel.forward`, you can safely ignore this message.
INFO:interpretune.utils.logging:The following columns don't have a corresponding argument in `NNSightReplacementModel.forward` and have been ignored: hypothesis, idx, label, premise, sequences. If hypothesis, idx, label, premise, sequences are not expected by `NNSightReplacementModel.forward`, you can safely ignore this message.
[INFO] interpretune.utils.logging: Setting up datamodule: InterpretunableDataModule
INFO:interpretune.utils.logging:Setting up datamodule: InterpretunableDataModule
[INFO] interpretune.utils.logging: Setting up model: InterpretunableModule
INFO:interpretune.utils.logging:Setting up model: InterpretunableModule
[INFO] interpretune.utils.logging: initializing optimizers and schedulers: InterpretunableModule
INFO:interpretune.utils.logging:initializing optimizers and schedulers: InterpretunableModule
[INFO] interpretune.utils.logging: Input gradient requirements handled by circuit tracer internally.
INFO:interpretune.utils.logging:Input gradient requirements handled by circuit tracer internally.
IT Session initialized successfully!
Example: Using Circuit Tracer to generate a custom graph and upload to Neuronpedia#
# Generate and upload graph in one step
raw_prompt = "The capital of the state containing Dallas is"
ct_module = it_session.module
# instruction-tuned replacement models assert chat-formatted inputs
# (circuit-tracer's ensure_tokenized requires a <bos><start_of_turn>user prefix)
prompt = GemmaPromptConfig().apply_chat_template_fn(
it_session.datamodule.tokenizer, raw_prompt, tokenize=False, add_generation_prompt=True
)
# we use the datetime to ensure the slug is unique
slug = f"it-generated-compute-specific-logits-demo-graph-{datetime.now().strftime('%Y%m%d-%H%M%S')}"
# This method does everything:
# - Generates attribution graph
# - Transforms it for Neuronpedia compatibility
# - Uploads to Neuronpedia
# - Returns both the graph and the Neuronpedia metadata
# if you want to set custom metadata for Neuronpedia graph, you can do so here.
# in this demo, since the circuit-tracer generated graph does not currently have cantor-paired feature nodes, we
# need to avoid setting neuronpedia_source_set in the feature_details metadata
# temporary hack: https://bit.ly/non_cantor_edge_case
graph, local_graph_path, neuronpedia_metadata = ct_module.generate_graph(
prompt=prompt,
slug=slug,
upload_to_np=False, # Upload the produced graph to Neuronpedia if True, or just return the graph if False
)
print(f"Generated and saved graph locally to {local_graph_path}")
if neuronpedia_metadata and hasattr(neuronpedia_metadata, "url"):
print("Graph uploaded successfully.")
print(f"View your graph at: {neuronpedia_metadata.url}")
Generated and saved graph locally to /tmp/20260730_155614_default/graph_data/it-generated-compute-specific-logits-demo-graph-20260730-155638.json
Phase 0: Precomputing activations and vectors
Precomputation completed in 1.47s
Found 5883 active features
Phase 1: Running forward pass
/mnt/cache/speediedan/.venvs/it_latest/lib/python3.13/site-packages/circuit_tracer/transcoder/single_layer_transcoder.py:139: UserWarning: Attempting to run cuBLAS, but there was no current CUDA context! Attempting to set the primary context... (Triggered internally at /__w/pytorch/pytorch/aten/src/ATen/cuda/CublasHandlePool.cpp:408.)
return input_acts @ self.W_skip.T
Forward pass completed in 0.74s
Phase 2: Building input vectors
Using 2 specified logit targets with cumulative probability 0.0000
Will include 5883 of 5883 feature nodes
Input vectors built in 0.43s
Phase 3: Computing logit attributions
1 logit attribution(s) completed in 0.10s
Phase 4: Computing feature attributions
Feature influence computation: 0%| | 0/5883 [00:00<?, ?it/s]
Feature influence computation: 26%|██▌ | 1536/5883 [00:00<00:00, 12809.66it/s]
Feature influence computation: 52%|█████▏ | 3072/5883 [00:00<00:00, 9821.86it/s]
Feature influence computation: 74%|███████▍ | 4352/5883 [00:00<00:00, 7483.00it/s]
Feature influence computation: 91%|█████████▏| 5376/5883 [00:00<00:00, 6064.93it/s]
Feature influence computation: 100%|██████████| 5883/5883 [00:00<00:00, 6538.99it/s]
Feature attributions completed in 0.90s
Attribution completed in 4.59s
[INFO] interpretune.utils.logging: [NeuronpediaIntegration] Transforming graph: /tmp/20260730_155614_default/graph_data/it-generated-compute-specific-logits-demo-graph-20260730-155638.json
INFO:interpretune.utils.logging:[NeuronpediaIntegration] Transforming graph: /tmp/20260730_155614_default/graph_data/it-generated-compute-specific-logits-demo-graph-20260730-155638.json
[INFO] interpretune.utils.logging: [NeuronpediaIntegration] Field 'clerps' changed: before=None, after=[]
INFO:interpretune.utils.logging:[NeuronpediaIntegration] Field 'clerps' changed: before=None, after=[]
[INFO] interpretune.utils.logging: [NeuronpediaIntegration] Changes written to /tmp/20260730_155614_default/it-generated-compute-specific-logits-demo-graph-20260730-155638_qparam_changes_20260730_155646.log
INFO:interpretune.utils.logging:[NeuronpediaIntegration] Changes written to /tmp/20260730_155614_default/it-generated-compute-specific-logits-demo-graph-20260730-155638_qparam_changes_20260730_155646.log
[INFO] interpretune.utils.logging: [NeuronpediaIntegration] Transformed qParams for Neuronpedia
INFO:interpretune.utils.logging:[NeuronpediaIntegration] Transformed qParams for Neuronpedia
[INFO] interpretune.utils.logging: [NeuronpediaIntegration] Transformed graph saved to: /tmp/20260730_155614_default/graph_data/it-generated-compute-specific-logits-demo-graph-20260730-155638.json
INFO:interpretune.utils.logging:[NeuronpediaIntegration] Transformed graph saved to: /tmp/20260730_155614_default/graph_data/it-generated-compute-specific-logits-demo-graph-20260730-155638.json
[INFO] interpretune.utils.logging: Neuronpedia upload not requested. Set upload_to_np to `True` to automatically upload.
INFO:interpretune.utils.logging:Neuronpedia upload not requested. Set upload_to_np to `True` to automatically upload.
# Alternative: Just transform an existing graph
# transformed_graph, np_metadata = ct_module.neuronpedia.transform_graph_for_np(
# graph_path=local_graph_path,
# # slug="it-generated-dallas-austin-demo-graph-20250725-121824-retransformed",
# upload_to_np=False # Set to True to upload the transformed graph
# )