from __future__ import annotations # see PEP 749, no longer needed when 3.13 reaches EOL
from typing import (
Protocol,
runtime_checkable,
TypeAlias,
NamedTuple,
TYPE_CHECKING,
Callable,
Any,
Sequence,
Iterable,
_ProtocolMeta,
get_args,
Mapping,
TypedDict,
TypeVar,
)
from pathlib import Path
from types import UnionType
from enum import auto, Enum, EnumMeta
from dataclasses import dataclass
from os import PathLike
import inspect
import torch
from torch import Tensor
from torch.optim import Optimizer
from typing_extensions import NotRequired, Required
from jsonargparse import Namespace
from transformers import BatchEncoding, PreTrainedTokenizerBase
# Heavy, optional runtime deps (sae_lens, transformer_lens) are only required for
# type-checking and should not be imported at top-level to avoid slowing down package import.
if TYPE_CHECKING:
from sae_lens.config import HfDataset
from interpretune.analysis.backends import ModelBackend
from interpretune.config import ITDataModuleConfig, ITConfig
else:
# Provide light-weight stand-ins for typing at runtime without importing heavy packages.
HfDataset = None # type: ignore
ITDataModuleConfig = None # type: ignore
ITConfig = None # type: ignore
################################################################################
# Interpretune helper types
################################################################################
StrOrPath: TypeAlias = str | Path
[docs]
class GraphComponentPayload(TypedDict):
"""Primitive graph payload used for Arrow-native graph serialization."""
input_string: str
input_tokens: Tensor
active_features: Tensor
adjacency_matrix: Tensor
selected_features: Tensor
activation_values: Tensor
logit_target_ids: Tensor
logit_target_tokens: list[str]
logit_probabilities: Tensor
graph_cfg: Mapping[str, Any]
scan: str | list[str] | None
vocab_size: int
[docs]
@runtime_checkable
class GraphComponentFactoryProtocol(Protocol):
"""Protocol for graph-like objects that can be hydrated from primitive components."""
@classmethod
def from_graph_components(cls, components: GraphComponentPayload) -> Any: ...
################################################################################
# Interpretune Enhanced Enums
################################################################################
[docs]
class AutoStrEnum(Enum):
@staticmethod
def _generate_next_value_(name, _start, _count, _last_values) -> str: # type: ignore
return name
# NOTE [Interpretability Adapters]:
[docs]
class Adapter(AutoStrEnum):
# CORE: The provided module and datamodule will be prepared for use with core PyTorch. The default
# trainer, a custom trainer or no trainer all can be used in combination with any supported and specified
# adapter.
core = auto()
# LIGHTNING: The provided module and datamodule will be prepared for use with the Lightning trainer and any
# supported and specified adapter.
lightning = auto()
# TRANSFORMER_LENS: The provided module and datamodule will be prepared for use with the TransformerLens adapter in
# in combination with any supported and specified adapter.
transformer_lens = auto()
# SAE_LENS: The provided module and datamodule will be prepared for use with the SAELens adapter in
# in combination with any supported and specified adapter.
sae_lens = auto()
# CIRCUIT_TRACER: The provided module and datamodule will be prepared for use with the Circuit Tracer adapter in
# in combination with any supported and specified adapter.
circuit_tracer = auto()
# NNSIGHT: The provided module and datamodule will be prepared for use with the NNsight adapter in
# combination with any supported and specified adapter. NNsight wraps HF models directly.
nnsight = auto()
def __lt__(self, other: "Adapter") -> bool:
return self.value < other.value
# DerivedEnumMeta is a custom metaclass that adds enum members from an input set.
[docs]
class SetDerivedEnum(AutoStrEnum, metaclass=DerivedEnumMeta): ...
################################################################################
# Core Enums
################################################################################
# TODO: consider switching these to a data structure that natively allows for more flexible DRY composition
# (currently preferring a custom enum for IDE autocompletion etc.)
CORE_PHASES = frozenset(["train", "validation", "test", "predict"])
EXT_PHASES = frozenset(["analysis"])
ALL_PHASES = CORE_PHASES.union(EXT_PHASES)
[docs]
class CorePhases(SetDerivedEnum):
# The _input_set class attribute is used by DerivedEnumMeta to generate enum members.
_input_set = CORE_PHASES
[docs]
class CoreSteps(SetDerivedEnum):
_input_set = CORE_PHASES
_transform = lambda x: "training_step" if x == "train" else f"{x}_step"
[docs]
class AllPhases(SetDerivedEnum):
_input_set = ALL_PHASES
[docs]
class AllSteps(SetDerivedEnum):
_input_set = ALL_PHASES
_transform = lambda x: "training_step" if x == "train" else f"{x}_step"
################################################################################
# Framework Compatibility helper types
# originally inspired by https://bit.ly/lightning_types definitions
################################################################################
_DictKey = TypeVar("_DictKey")
[docs]
@runtime_checkable
class Steppable(Protocol):
"""To structurally type ``optimizer.step()``"""
# Inferred from `torch.optim.optimizer.pyi`
def step(self, closure: Callable[[], float] | None = ...) -> float | None: ...
[docs]
@runtime_checkable
class Optimizable(Steppable, Protocol):
"""To structurally type ``optimizer``"""
param_groups: list[dict[Any, Any]]
defaults: dict[Any, Any]
state: dict[Any, Any]
def zero_grad(self) -> None: ...
def state_dict(self) -> dict[str, dict[Any, Any]]: ...
def load_state_dict(self, state_dict: dict[str, dict[Any, Any]]) -> None: ...
@runtime_checkable
class _Stateful(Protocol[_DictKey]):
"""This class is used to detect if an object is stateful using `isinstance(obj, _Stateful)`."""
def state_dict(self) -> dict[_DictKey, Any]: ...
def load_state_dict(self, state_dict: dict[_DictKey, Any]) -> None: ...
[docs]
@runtime_checkable
class LRScheduler(_Stateful[str], Protocol):
optimizer: Optimizer
base_lrs: list[float]
def __init__(self, optimizer: Optimizer, *args: Any, **kwargs: Any) -> None: ...
def step(self, epoch: int | None = None) -> None: ...
# Inferred from `torch.optim.lr_scheduler.pyi`
# Missing attributes were added to improve typing
[docs]
@runtime_checkable
class ReduceLROnPlateau(_Stateful[str], Protocol):
in_cooldown: bool
optimizer: Optimizer
def __init__(
self,
optimizer: Optimizer,
mode: str = ...,
factor: float = ...,
patience: int = ...,
verbose: bool = ...,
threshold: float = ...,
threshold_mode: str = ...,
cooldown: int = ...,
min_lr: float = ...,
eps: float = ...,
) -> None: ...
def step(self, metrics: float | int | Tensor, epoch: int | None = None) -> None: ...
STEP_OUTPUT = Tensor | Mapping[str, Any] | None
LRSchedulerTypeUnion = torch.optim.lr_scheduler.LRScheduler | torch.optim.lr_scheduler.ReduceLROnPlateau
# Protocol-level union covering the scheduler Protocol and the ReduceLROnPlateau Protocol
LRSchedulerProtocolUnion: TypeAlias = LRScheduler | ReduceLROnPlateau
[docs]
@dataclass
class LRSchedulerConfig:
scheduler: torch.optim.lr_scheduler.LRScheduler | torch.optim.lr_scheduler.ReduceLROnPlateau
# no custom name
name: str | None = None
# after epoch is over
interval: str = "epoch"
# every epoch/batch
frequency: int = 1
# most often not ReduceLROnPlateau scheduler
reduce_on_plateau: bool = False
# value to monitor for ReduceLROnPlateau
monitor: str | None = None
# enforce that the monitor exists for ReduceLROnPlateau
strict: bool = True
[docs]
class LRSchedulerConfigType(TypedDict, total=False):
scheduler: Required[LRSchedulerTypeUnion]
name: str | None
interval: str
frequency: int
reduce_on_plateau: bool
monitor: str | None
scrict: bool
[docs]
class OptimizerLRSchedulerConfig(TypedDict):
optimizer: Optimizer
lr_scheduler: NotRequired[LRSchedulerTypeUnion | LRSchedulerConfigType]
OptimizerLRScheduler = (
Optimizer
| Sequence[Optimizer]
| tuple[Sequence[Optimizer], Sequence[LRSchedulerTypeUnion | LRSchedulerConfig]]
| OptimizerLRSchedulerConfig
| None
)
ArgsType = list[str] | dict[str, Any] | Namespace | None
AnyDataClass = TypeVar("AnyDataClass")
################################################################################
# Core Protocols
################################################################################
[docs]
@runtime_checkable
class DataPrepable(Protocol):
"""Minimum requirement for an Interpretunable DataModule is to have a prepare_data method and a valid
datamodule config."""
def prepare_data(self, target_model: torch.nn.Module | None = None) -> None: ...
[docs]
@runtime_checkable
class SaveHyperparametersProtocol(Protocol):
"""Simple protocol indicating an object exposes a save_hyperparameters method.
Intentionally framework-agnostic, but see Lightning `HyperparametersMixin.save_hyperparameters` for an example
implementation. The method is expected to return None.
"""
def save_hyperparameters(self, *args: Any, **kwargs: Any) -> None: ...
[docs]
@runtime_checkable
class TrainLoadable(DataPrepable, Protocol):
def train_dataloader(self) -> torch.utils.data.DataLoader: ...
[docs]
@runtime_checkable
class ValLoadable(DataPrepable, Protocol):
def val_dataloader(self) -> torch.utils.data.DataLoader: ...
[docs]
@runtime_checkable
class TestLoadable(DataPrepable, Protocol):
def test_dataloader(self) -> torch.utils.data.DataLoader: ...
[docs]
@runtime_checkable
class PredictLoadable(DataPrepable, Protocol):
def predict_dataloader(self) -> torch.utils.data.DataLoader: ...
[docs]
@runtime_checkable
class DataModuleInvariants(Protocol):
itdm_cfg: ITDataModuleConfig
def setup(self, *args, **kwargs) -> None: ...
[docs]
@runtime_checkable
class TrainSteppable(Protocol):
def training_step(self, *args, **kwargs) -> STEP_OUTPUT: ...
[docs]
@runtime_checkable
class ValidationSteppable(Protocol):
def validation_step(self, *args, **kwargs) -> STEP_OUTPUT: ...
[docs]
@runtime_checkable
class TestSteppable(Protocol):
def test_step(self, *args, **kwargs) -> STEP_OUTPUT: ...
[docs]
@runtime_checkable
class PredictSteppable(Protocol):
def predict_step(self, *args, **kwargs) -> STEP_OUTPUT: ...
[docs]
@runtime_checkable
class ModuleInvariants(Protocol):
it_cfg: ITConfig
@property
def analysis_backend(self) -> Any | None: ...
def setup(self, *args, **kwargs) -> None: ...
def configure_optimizers(self) -> OptimizerLRScheduler | None: ...
# N.B. runtime protocol validation will check for attribute presence but not validate signatures etc. With this
# protocol-based approach we're providing rudimentary functional checks while erroring on the side of flexibility
ModuleSteppable: TypeAlias = TrainSteppable | ValidationSteppable | TestSteppable | PredictSteppable
DataModuleInitable: TypeAlias = TrainLoadable | ValLoadable | TestLoadable | PredictLoadable
def gen_protocol_variants(
supported_sub_protocols: UnionType, base_protocols: _ProtocolMeta | tuple[_ProtocolMeta]
) -> UnionType:
protocol_components = []
if not isinstance(base_protocols, tuple):
base_protocols = (base_protocols,)
for sub_proto in get_args(supported_sub_protocols):
supported_cls = _ProtocolMeta(f"Built{sub_proto.__name__}", (*base_protocols, sub_proto, Protocol), {}) # type: ignore[arg-type]
supported_cls._is_runtime_protocol = True
frame = inspect.currentframe()
if frame and frame.f_back:
module = inspect.getmodule(frame.f_back)
if module:
supported_cls.__module__ = module.__name__
protocol_components.append(supported_cls)
gen_type_alias = " | ".join([f"protocol_components[{i}]" for i in range(len(protocol_components))])
return eval(gen_type_alias)
# We generate valid datamodule/module protocol variants by composing their respective base protocols with the set of
# valid subprotocols over which `any` semantics apply.
ITDataModuleProtocol: TypeAlias = gen_protocol_variants(DataModuleInitable, DataModuleInvariants) # type: ignore
ITModuleProtocol: TypeAlias = gen_protocol_variants(ModuleSteppable, ModuleInvariants) # type: ignore
# TODO: ensure protocol variants are explicitly documented and possibly add a section describing the approach to
# supported protocol variant generation. Also add an issue tracker for this approach to solicit ideas for a
# cleaner/more pythonic approach. As Python structural subtyping features are still evolving, if a cleaner and
# more pythonic approach isn't available now, one will hopefully be available in the near future.
InterpretunableType: TypeAlias = ITDataModuleProtocol | ITModuleProtocol
[docs]
class InterpretunableTuple(NamedTuple):
datamodule: ITDataModuleProtocol | None = None
module: ITModuleProtocol | None = None
################################################################################
# Base Module / Generative-step Protocols
################################################################################
[docs]
@runtime_checkable
class GenerativeStepProtocol(Protocol):
"""Protocol describing the surface provided by GenerativeStepMixin used by debug utilities.
This is intentionally small: it only includes the methods/properties the debug generation code relies on.
"""
def it_generate(self, batch: Any, **kwargs: Any) -> Any: ...
@property
def generation_cfg(self) -> Any: ...
@property
def gen_sig_keys(self) -> list: ...
def map_gen_inputs(self, batch: Any) -> dict[str, Any]: ...
def map_gen_kwargs(self, kwargs: dict) -> dict[str, Any]: ...
[docs]
@runtime_checkable
class ITModuleBase(
ModuleInvariants,
Protocol,
):
"""Concrete Protocol combining the two core module protocols.
NOTE: As `ITModuleProtocol` is a union of dynamically created Protocol variants it can't be used as a direct base
class in a class statement without metaclass conflict (it's a typealias of a union so it has a metaclass
of types.UnionType and our derived class's metaclass (_ProtocolMeta) isn't a subclass of types.UnionType).
`ITModuleBase` provides a stable Protocol class that composes the structural requirements of our module and can be
used safely as a base for further composite Protocols.
"""
...
[docs]
@runtime_checkable
class ITModuleGenDebuggable(ITModuleBase, GenerativeStepProtocol, Protocol):
"""Composite protocol for ITModuleGenDebuggable objects.
Requires both module-like surface and generative-step helpers.
Use this as an annotation for attributes that must behave like a given module composition
that also supports generation helper methods (e.g., `it_generate`).
"""
...
################################################################################
# Analysis Protocols
################################################################################
NamesFilter = Callable[[str], bool] | Sequence[str] | str | None
[docs]
class LatentModelFqn(NamedTuple):
release: str
sae_id: str
[docs]
class AnalysisOpProtocol(Protocol):
"""Protocol defining required interface for analysis operations."""
name: str
description: str
output_schema: dict
input_schema: dict | None
def save_batch(
self,
analysis_batch: BaseAnalysisBatchProtocol,
batch: BatchEncoding,
tokenizer: PreTrainedTokenizerBase | None = None,
save_prompts: bool = False,
save_tokens: bool = False,
decode_kwargs: dict | None = None,
) -> BaseAnalysisBatchProtocol: ...
[docs]
class LatentDictProtocol(Protocol):
"""Protocol for latent model analysis dictionary operations."""
@property
def shapes(self) -> dict[str, torch.Size | list[torch.Size]]: ...
def batch_join(
self, across_saes: bool = False, join_fn: Callable = torch.cat
) -> "LatentDictProtocol" | list[torch.Tensor]: ...
def apply_op_by_latent_model(self, operation: Callable | str, *args, **kwargs) -> "LatentDictProtocol": ...
[docs]
class AnalysisStoreProtocol(Protocol):
"""Protocol verifying core analysis store functionality."""
dataset: HfDataset | StrOrPath | PathLike | None
streaming: bool
cache_dir: str | None
# save_dir on AnalysisStore is exposed as a property (getter returns Path, setter accepts str/PathLike).
# Model it as a property in the protocol so static checkers consider the descriptor instead of a plain attribute.
@property
def save_dir(self) -> Path: ...
@save_dir.setter
def save_dir(self, path: StrOrPath) -> None: ...
stack_batches: bool
split: str
# op_output_dataset_path may be a string path, a pathlib.Path, or None at runtime
op_output_dataset_path: str | Path | None
def by_latent_model(self, field_name: str, stack_latents: bool = True) -> LatentDictProtocol: ...
# __getattr__ may return dataset columns (lists/tensors), callables, or other attributes — accept Any
def __getattr__(self, name: str) -> Any: ...
def reset_dataset(self) -> None: ...
[docs]
class AnalysisCfgProtocol(Protocol):
"""Protocol verifying core analysis configuration functionality."""
output_store: AnalysisStoreProtocol
input_store: AnalysisStoreProtocol
op: AnalysisOpProtocol
fwd_hooks: list[tuple]
bwd_hooks: list[tuple]
cache_dict: dict
names_filter: NamesFilter | None
# Save configuration fields
save_prompts: bool
save_tokens: bool
decode_kwargs: dict
def check_add_default_hooks(
self, op: AnalysisOpProtocol, names_filter: str | Callable | None, cache_dict: dict | None
) -> tuple[list[tuple], list[tuple]]: ...
[docs]
class LatentAnalysisProtocol(Protocol):
"""Protocol for latent analysis components requiring a subset of SAELensAnalysisMixin methods."""
def construct_names_filter(
self, target_layers: list[int], sae_hook_match_fn: Callable[[str, list[int] | None], bool]
) -> NamesFilter: ...
[docs]
class LatentAnalysisModuleProtocol(ITModuleBase, LatentAnalysisProtocol, Protocol):
"""Protocol requiring both the ITModuleProtocol surface and LatentAnalysisProtocol methods.
Inheriting both interfaces expresses the intersection of the two required structural surfaces so static checkers
(like pyright) will require both sets of attributes/methods.
"""
model_backend: ModelBackend
[docs]
class RunnerCfgProtocol(Protocol):
"""Protocol representing the base Session/Runner configuration object."""
it_session: Any | None
module: Any | None
datamodule: Any | None
[docs]
class AnalysisRunnerCfgProtocol(RunnerCfgProtocol, Protocol):
"""Protocol representing the analysis runner configuration used by adapters.
Matches the public surface of `AnalysisRunnerCfg` (see `config/runner.py`).
"""
analysis_cfgs: Any
limit_analysis_batches: int
cache_dir: Any
op_output_dataset_path: Any
latent_analysis_targets: Any
artifact_cfg: Any
ignore_manual: bool
# The processed, canonicalized list of AnalysisCfg that the runner exposes
_processed_analysis_cfgs: list[Any]
[docs]
class AnalysisRunnerProtocol(Protocol):
"""Protocol verifying presence of analysis_run_cfg attribute.
This minimal structural protocol is used by adapters to access runner-provided analysis configuration without
importing the runner implementation.
"""
analysis_run_cfg: AnalysisRunnerCfgProtocol
[docs]
class ActivationCacheProtocol(Protocol):
"""Core activation cache protocol."""
cache_dict: dict[str, torch.Tensor]
has_batch_dim: bool
has_embed: bool
has_pos_embed: bool
def __getitem__(self, key: str | tuple) -> torch.Tensor: ...
def stack_activation(
self, activation_name: str, layer: int = -1, sublayer_type: str | None = None
) -> torch.Tensor: ...
def items(self) -> Iterable[tuple[str, torch.Tensor]]: ...
[docs]
class BaseAnalysisBatchProtocol(Protocol):
"""Base protocol defining methods all analysis batches should implement.
Subclasses should define which dataset columns will have attribute-based access enabled for associated AnalysisStore
objects.
"""
def update(self, **kwargs) -> None: ...
def to_cpu(self) -> None: ...
[docs]
class DefaultAnalysisBatchProtocol(BaseAnalysisBatchProtocol):
"""Default analysis batch protocol defining which dataset columns should have attribute-based access enabled
for AnalysisStore objects. Subclasses can extend this protocol (or the base one) to add additional attributes
or change existing attributes as needed.
Attributes:
logit_diffs (torch.Tensor | dict[str, dict[int, torch.Tensor]] | None):
Per batch logit differences with shape [batch_size]
answer_logits (torch.Tensor | dict[str, dict[int, torch.Tensor]] | None):
Model output logits with shape [batch_size, 1, num_classes]
loss (torch.Tensor | dict[str, dict[int, torch.Tensor]] | None):
Loss values with shape [batch_size]
label_ids (torch.Tensor | None):
Input labels translated to token ids with shape [batch_size] (if labels provided & translation is needed)
orig_labels (torch.Tensor | None):
Ground truth unmodified labels with shape [batch_size]
preds (torch.Tensor | dict[str, dict[int, torch.Tensor]] | None):
Model predictions with shape [batch_size]
cache (ActivationCacheProtocol | None):
Forward pass activation cache
grad_cache (ActivationCacheProtocol | None):
Backward pass gradient cache
answer_indices (torch.Tensor | None):
Indices of answers with shape [batch_size]
alive_latents (dict[str, list[int]] | None):
Active latent indices per latent model hook
correct_activations (dict[str, torch.Tensor] | None):
Latent model activations after corrections with shape [batch_size, d_sae] for each latent model
attribution_values (dict[str, torch.Tensor] | None):
Attribution values per latent model hook
tokens (torch.Tensor | None):
Input token IDs
prompts (list[str] | None):
Text prompts
"""
logit_diffs: torch.Tensor | dict[str, dict[int, torch.Tensor]] | None
answer_logits: torch.Tensor | dict[str, dict[int, torch.Tensor]] | None
loss: torch.Tensor | dict[str, dict[int, torch.Tensor]] | None
preds: torch.Tensor | dict[str, dict[int, torch.Tensor]] | None
label_ids: torch.Tensor | None
orig_labels: torch.Tensor | None
cache: ActivationCacheProtocol | None
grad_cache: ActivationCacheProtocol | None
answer_indices: torch.Tensor | None
alive_latents: dict[str, list[int]] | None
correct_activations: dict[str, torch.Tensor] | None
attribution_values: dict[str, torch.Tensor] | None
tokens: torch.Tensor | None
prompts: list[str] | None
[docs]
class CircuitAnalysisBatchProtocol(DefaultAnalysisBatchProtocol):
"""Circuit analysis batch protocol defining additional attributes for circuit tracer operations.
Extends the default protocol with circuit tracing specific attributes.
Attributes:
attribution_graphs (list | None):
Generated attribution graphs for each prompt in the batch
graph_metadata (list[dict] | None):
Metadata for each generated graph including parameters used
graph_paths (list[str] | None):
File paths where graphs are saved (if saved)
circuit_prompts (list[str] | None):
Prompts used for circuit attribution (may differ from input prompts)
"""
attribution_graphs: list | None
graph_metadata: list[dict] | None
graph_paths: list[str] | None
circuit_prompts: list[str] | None
input_string: str | None
adjacency_matrix: torch.Tensor | None
active_features: torch.Tensor | None
selected_features: torch.Tensor | None
activation_values: torch.Tensor | None
logit_target_ids: torch.Tensor | None
logit_target_tokens: list[str] | None
logit_probabilities: torch.Tensor | None
input_tokens: torch.Tensor | None
graph_cfg_json: str | None
graph_scan_json: str | None
graph_vocab_size: int | None
intervention_config: str | None
intervention_specs_json: str | None
feature_intervention_dict: Any | None
feature_intervention_dict_json: str | None
intervention_layers: list[int] | torch.Tensor | None
intervention_positions: list[int] | torch.Tensor | None
intervention_feature_ids: list[int] | torch.Tensor | None
intervention_values: list[float] | torch.Tensor | None
intervention_base_values: list[float] | torch.Tensor | None
intervention_score_values: list[float] | torch.Tensor | None
intervention_scale_factors: list[float] | torch.Tensor | None
pre_intervention_logits: torch.Tensor | None
post_intervention_logits: torch.Tensor | None
logit_diff: torch.Tensor | float | None
