transformer_lens.model_protocol module¶

Structural model types for the model-agnostic interpretability utilities.

The interpretability utilities (patching, head_detector, ActivationCache) only need cfg plus the run_with_* / tokenization surface and a few weight-processing helpers. Typing their model parameter as this Protocol accepts TransformerBridge and any other structural match (e.g. RemoteBridge).

Members are typed loosely on purpose: implementations differ in signature detail but agree in use.

class transformer_lens.model_protocol.TrainableTransformerLensModel(*args, **kwargs)¶

Bases: Protocol

Exactly the surface the tools.training loop touches: callable with return_type="loss" plus the standard torch parameter/mode/device methods. Deliberately standalone (not extending TransformerLensModel): beartype validates protocols via getattr_static, and demanding the full TL surface would spuriously reject plain nn.Module models that train fine through this loop.

parameters(*args: Any, **kwargs: Any) → Any¶

state_dict(*args: Any, **kwargs: Any) → Any¶

to(*args: Any, **kwargs: Any) → Any¶

train(*args: Any, **kwargs: Any) → Any¶

class transformer_lens.model_protocol.TransformerLensModel(*args, **kwargs)¶

Bases: Protocol

Minimal structural interface the interpretability utilities rely on.

property cfg: TransformerLensConfig¶

run_with_cache(*args: Any, **kwargs: Any) → Any¶

run_with_hooks(*args: Any, **kwargs: Any) → Any¶

to_tokens(*args: Any, **kwargs: Any) → Any¶

class transformer_lens.model_protocol.TransformerLensModelWithWeights(*args, **kwargs)¶

Bases: TransformerLensModel, Protocol

Adds the weight-processing surface that ActivationCache’s advanced helpers (LayerNorm folding, residual-direction projection) reach for. The bridge builds them from its adapter.

accumulated_bias(*args: Any, **kwargs: Any) → Any¶

property blocks: Any¶

property ln_final: Any¶

to_single_token(*args: Any, **kwargs: Any) → Any¶

tokens_to_residual_directions(*args: Any, **kwargs: Any) → Any¶