transformer_lens.lit.dataset#

LIT Dataset wrapper for TransformerLens.

This module provides LIT-compatible Dataset wrappers for use with TransformerLens models. It includes utilities for loading common datasets and creating custom datasets for model analysis.

Example usage:

>>> from transformer_lens.lit import SimpleTextDataset  
>>>
>>> # Create a dataset from examples
>>> examples = [  
...     {"text": "The capital of France is Paris."},
...     {"text": "Machine learning is a subset of AI."},
... ]
>>> dataset = SimpleTextDataset(examples)  
>>>
>>> # Use with LIT server
>>> from lit_nlp import dev_server  
>>> server = dev_server.Server(models, {"my_data": dataset})  

References

LIT Dataset API: https://pair-code.github.io/lit/documentation/api#datasets
TransformerLens: https://github.com/TransformerLensOrg/TransformerLens

class transformer_lens.lit.dataset.DatasetConfig(max_examples: int | None = None, shuffle: bool = False, seed: int = 42)#

Bases: object

Configuration for LIT datasets.

max_examples: int | None = None#: Maximum number of examples to load.

seed: int = 42#: Random seed for shuffling.

shuffle: bool = False#: Whether to shuffle the examples.

class transformer_lens.lit.dataset.IOIDataset(examples: List[Dict[str, Any]] | None = None, name: str = 'IOI Dataset')#

Bases: object

Indirect Object Identification (IOI) dataset.

This dataset contains examples for the Indirect Object Identification task, commonly used in mechanistic interpretability research.

Each example has the format: “When {name1} and {name2} went to the {place}, {name1} gave a {object} to”

The model should complete with name2 (the indirect object).

Reference:: Wang et al. “Interpretability in the Wild: a Circuit for Indirect Object Identification in GPT-2 small” https://arxiv.org/abs/2211.00593

NAMES = ['Mary', 'John', 'Alice', 'Bob', 'Charlie', 'Diana', 'Emma', 'Frank', 'Grace', 'Henry', 'Ivy', 'Jack']#

OBJECTS = ['book', 'gift', 'letter', 'key', 'phone', 'drink', 'flower', 'card', 'ticket', 'bag']#

PLACES = ['store', 'park', 'beach', 'restaurant', 'library', 'museum', 'cafe', 'market', 'school', 'hospital']#

TEMPLATE = 'When {name1} and {name2} went to the {place}, {name1} gave a {object} to'#

__init__(examples: List[Dict[str, Any]] | None = None, name: str = 'IOI Dataset')#

Initialize the IOI dataset.

Parameters:

examples – Optional pre-defined examples.
name – Dataset name.

__iter__()#: Iterate over examples.

__len__() → int#: Return the number of examples.

add_example(name1: str, name2: str, place: str, obj: str) → None#

Add a single IOI example.

Parameters:

name1 – Subject name (gives the object).
name2 – Indirect object name (receives the object).
place – Location.
obj – Object being given.

description() → str#: Return a description of the dataset.

property examples: List[Dict[str, Any]]#: Return all examples.

classmethod generate(n_examples: int = 100, seed: int = 42, name: str = 'IOI Dataset') → IOIDataset#

Generate random IOI examples.

Parameters:

n_examples – Number of examples to generate.
seed – Random seed for reproducibility.
name – Dataset name.

Returns:

IOIDataset with generated examples.

spec() → Dict[str, Any]#: Return the spec describing the dataset fields.

class transformer_lens.lit.dataset.InductionDataset(examples: List[Dict[str, Any]] | None = None, name: str = 'Induction Dataset')#

Bases: object

Dataset for induction head analysis.

Induction heads are attention heads that perform pattern matching of the form [A][B] … [A] -> [B]. This dataset provides examples designed to trigger induction behavior.

Example pattern: “The cat sat on the mat. The cat sat on the” -> “ mat”

Reference:: Olsson et al. “In-context Learning and Induction Heads” https://arxiv.org/abs/2209.11895

__init__(examples: List[Dict[str, Any]] | None = None, name: str = 'Induction Dataset')#

Initialize the induction dataset.

Parameters:

examples – Optional pre-defined examples.
name – Dataset name.

__iter__()#: Iterate over examples.

__len__() → int#: Return the number of examples.

add_example(pattern: str, repeated_text: str, completion: str) → None#

Add an induction example.

Parameters:

pattern – The pattern that is repeated.
repeated_text – The text before the second occurrence.
completion – The expected completion.

description() → str#: Return a description of the dataset.

property examples: List[Dict[str, Any]]#: Return all examples.

classmethod generate_simple(n_examples: int = 50, seed: int = 42, name: str = 'Induction Dataset') → InductionDataset#

Generate simple induction examples.

Parameters:

n_examples – Number of examples to generate.
seed – Random seed.
name – Dataset name.

Returns:

InductionDataset with generated examples.

spec() → Dict[str, Any]#: Return the spec describing the dataset fields.

class transformer_lens.lit.dataset.PromptCompletionDataset(examples: List[Dict[str, Any]] | None = None, name: str = 'PromptCompletionDataset')#

Bases: object

Dataset with prompt-completion pairs for generation analysis.

This dataset type is useful for analyzing model generation behavior, where each example has a prompt and an expected completion.

Example

>>> dataset = PromptCompletionDataset([  
...     {"prompt": "The capital of France is", "completion": " Paris"},
...     {"prompt": "2 + 2 =", "completion": " 4"},
... ])

COMPLETION_FIELD = 'completion'#

FULL_TEXT_FIELD = 'text'#

PROMPT_FIELD = 'prompt'#

__init__(examples: List[Dict[str, Any]] | None = None, name: str = 'PromptCompletionDataset')#

Initialize the dataset.

Parameters:

examples – List of example dictionaries with prompt/completion.
name – Name for the dataset.

__iter__()#: Iterate over examples.

__len__() → int#: Return the number of examples.

description() → str#: Return a description of the dataset.

property examples: List[Dict[str, Any]]#: Return all examples.

classmethod from_pairs(pairs: Sequence[tuple], name: str = 'PromptCompletionDataset') → PromptCompletionDataset#

Create a dataset from (prompt, completion) tuples.

Parameters:

pairs – Sequence of (prompt, completion) tuples.
name – Dataset name.

Returns:

PromptCompletionDataset instance.

Example

>>> dataset = PromptCompletionDataset.from_pairs([  
...     ("Hello, my name is", " Alice"),
...     ("The weather today is", " sunny"),
... ])

spec() → Dict[str, Any]#: Return the spec describing the dataset fields.

class transformer_lens.lit.dataset.SimpleTextDataset(examples: List[Dict[str, Any]] | None = None, name: str = 'SimpleTextDataset')#

Bases: object

Simple text dataset for use with HookedTransformerLIT.

This is a basic dataset class that holds text examples for analysis with LIT. Each example is a dictionary with at least a “text” field.

Example

>>> dataset = SimpleTextDataset([  
...     {"text": "Hello world"},
...     {"text": "How are you?"},
... ])
>>> len(dataset.examples)  
2

__init__(examples: List[Dict[str, Any]] | None = None, name: str = 'SimpleTextDataset')#

Initialize the dataset.

Parameters:

examples – List of example dictionaries with “text” field.
name – Name for the dataset (shown in LIT UI).

__iter__()#: Iterate over examples.

__len__() → int#: Return the number of examples.

description() → str#: Return a description of the dataset.

property examples: List[Dict[str, Any]]#: Return all examples in the dataset.

classmethod from_file(filepath: str | Path, name: str | None = None, max_examples: int | None = None) → SimpleTextDataset#

Load a dataset from a text file.

Each line in the file becomes one example.

Parameters:

filepath – Path to the text file.
name – Optional dataset name (defaults to filename).
max_examples – Maximum number of examples to load.

Returns:

SimpleTextDataset instance.

classmethod from_strings(texts: Sequence[str], name: str = 'TextDataset') → SimpleTextDataset#

Create a dataset from a list of strings.

Parameters:

texts – Sequence of text strings.
name – Dataset name.

Returns:

SimpleTextDataset instance.

Example

>>> dataset = SimpleTextDataset.from_strings([  
...     "First example",
...     "Second example",
... ])

spec() → Dict[str, Any]#

Return the spec describing the dataset fields.

This tells LIT what fields each example contains and their types.

Returns:: Dictionary mapping field names to LIT type specs.

transformer_lens.lit.dataset.wrap_for_lit(dataset: Any) → Any#: Placeholder when LIT is not available.