transformer_lens.lit.model module

LIT Model wrapper for TransformerLens models.

This module provides a LIT-compatible wrapper around TransformerLens models, enabling the use of Google’s Learning Interpretability Tool (LIT) for model visualization and analysis.

The wrapper exposes: - Token predictions (logits, top-k tokens) - Per-layer embeddings (residual stream) - Attention patterns (all layers/heads) - Token gradients for salience maps - Loss computation

Example usage:
>>> from transformer_lens import TransformerBridge  
>>> from transformer_lens.lit import TransformerLensLIT  
>>>
>>> # Load model
>>> model = TransformerBridge.boot_transformers("gpt2")  
>>>
>>> # Create LIT wrapper
>>> lit_model = TransformerLensLIT(model)  
>>>
>>> # Run prediction
>>> inputs = [{"text": "Hello, world!"}]  
>>> outputs = list(lit_model.predict(inputs))

References

transformer_lens.lit.model.HookedTransformerLIT

alias of TransformerLensLIT

transformer_lens.lit.model.HookedTransformerLITConfig

alias of TransformerLensLITConfig

class transformer_lens.lit.model.TransformerLensLIT(model: Any, config: TransformerLensLITConfig | None = None)

Bases: object

LIT Model wrapper for TransformerLens models.

This wrapper implements the LIT Model API, enabling the use of LIT’s visualization and analysis tools with TransformerLens models.

The wrapper provides: - Token predictions with top-k probabilities - Per-layer embeddings for embedding projector - Attention patterns for attention visualization - Token gradients for salience maps

Example

>>> model = TransformerBridge.boot_transformers("gpt2")  
>>> lit_model = TransformerLensLIT(model)  
>>> lit_model.input_spec()  
{'text': TextSegment(), ...}
__init__(model: Any, config: TransformerLensLITConfig | None = None)

Initialize the LIT wrapper.

Parameters:

  • model – A TransformerLens model (e.g. TransformerBridge).

  • config – Optional configuration. Uses defaults if not provided.

Raises:

  • ImportError – If lit-nlp is not installed.

  • TypeError – If model does not implement the TransformerLens model interface.

description() → str

Return a human-readable description of the model.

Returns:

Model description string.

classmethod from_pretrained(model_name: str, config: TransformerLensLITConfig | None = None, **model_kwargs) → TransformerLensLIT

Create a LIT wrapper from a pretrained model name.

Convenience method that boots a TransformerBridge (with compatibility mode enabled for the legacy hook/weight surface) and wraps it for LIT.

Parameters:

  • model_name – Name of the pretrained model (e.g., “gpt2”).

  • config – Optional wrapper configuration.

  • **model_kwargs – Additional arguments for TransformerBridge.boot_transformers.

Returns:

TransformerLensLIT wrapper instance.

Example

>>> lit_model = TransformerLensLIT.from_pretrained("gpt2")
get_embedding_table() → tuple

Return the token embedding table.

Required by LIT for certain generators like HotFlip.

Returns:

Tuple of (vocab_list, embedding_matrix) where vocab_list is a list of token strings and embedding_matrix is [vocab, d_model].

classmethod init_spec() → Dict[str, Any]

Return spec for model initialization in LIT UI.

This allows loading new models through the LIT interface.

Returns:

Specification for initialization parameters.

input_spec() → Dict[str, Any]

Return spec describing the model inputs.

Defines the expected input format for the model. LIT uses this to validate inputs and generate appropriate UI controls.

Returns:

Dictionary mapping field names to LIT type specs.

max_minibatch_size() → int

Return the maximum batch size for prediction.

Returns:

Maximum batch size.

output_spec() → Dict[str, Any]

Return spec describing the model outputs.

Defines all the outputs that the model produces. LIT uses this to determine which visualizations to show.

Returns:

Dictionary mapping field names to LIT type specs.

predict(inputs: Iterable[Dict[str, Any]]) → Iterator[Dict[str, Any]]

Run prediction on a sequence of inputs.

This is the main entry point for LIT to get model outputs.

Parameters:

inputs – Iterable of input dictionaries, each with fields matching input_spec().

Yields:

Output dictionaries for each input, with fields matching output_spec().

property supports_concurrent_predictions: bool

Whether this model supports concurrent predictions.

Returns False as PyTorch models typically aren’t thread-safe.

class transformer_lens.lit.model.TransformerLensLITConfig(max_seq_length: int = 512, batch_size: int = 8, top_k: int = 10, compute_gradients: bool = True, output_attention: bool = True, output_embeddings: bool = True, output_all_layers: bool = False, embedding_layers: List[int] | None = None, prepend_bos: bool = True, device: str | None = None)

Bases: object

Configuration for the TransformerLensLIT wrapper.

batch_size: int = 8
compute_gradients: bool = True
device: str | None = None
embedding_layers: List[int] | None = None
max_seq_length: int = 512
output_all_layers: bool = False
output_attention: bool = True
output_embeddings: bool = True
prepend_bos: bool = True
top_k: int = 10