Experiment scripts

@experiment turns a function main(config: Config) into a command-line script. The fields of the Config dataclass become --kebab-case options, with their types, defaults and required-ness taken from the fields themselves.

A script

import dataclasses
from dataclasses import dataclass, field

from transformer_lens.model_bridge import TransformerBridge

from entoscope import ConfigFactory, Lens, LensConfig, LensInput, experiment


@dataclass
class Config:
    config_path: str = field(metadata={"help": "a YAML file of lens configs"})
    config_id: str = "logit_lens"
    prompts: list[str] = field(default_factory=lambda: ["The capital of France is"])
    top_k: int | None = field(default=None, metadata={"help": "override the config's top_k"})
    model_name: str = field(default="gpt2", metadata={"choices": ["gpt2", "gpt2-medium"]})
    dry_run: bool = False


@experiment
def main(config: Config):
    """Run one lens config on some prompts."""
    model = TransformerBridge.boot_transformers(config.model_name)
    model.eval()
    lens_config = ConfigFactory[LensConfig](config.config_path).create(config.config_id)
    if config.top_k is not None:
        lens_config = dataclasses.replace(lens_config, top_k=config.top_k)

    lens = Lens(model)
    for i, prompt in enumerate(config.prompts):
        lens_input = LensInput(lens_config, prompt)
        save_dir = f"runs/{config.config_id}/{i}"
        if config.dry_run:
            lens.plan(lens_input, save_dir)
        else:
            lens.apply(lens_input, save_dir)

Saved as lens_script.py, it runs as:

uv run python lens_script.py --config-path examples/lens.yaml --config-id tuned_lens --dry-run
uv run python lens_script.py --config-path examples/lens.yaml --prompts "The Eiffel Tower is in" "Rome is the capital of"

and describes itself:

$ uv run python lens_script.py --help
usage: lens_script.py [-h] --config-path CONFIG_PATH [--config-id CONFIG_ID]
                      [--prompts [PROMPTS ...]] [--top-k TOP_K] [--model-name {gpt2,gpt2-medium}]
                      [--dry-run | --no-dry-run]

Run one lens config on some prompts.

options:
  -h, --help            show this help message and exit
  --config-path CONFIG_PATH
                        a YAML file of lens configs
  --config-id CONFIG_ID
                        (default: logit_lens)
  --prompts [PROMPTS ...]
                        (default: ['The capital of France is'])
  --top-k TOP_K         override the config's top_k (default: None)
  --model-name {gpt2,gpt2-medium}
                        (default: gpt2)
  --dry-run, --no-dry-run
                        (default: False)

From fields to options

Field type	Option
str, int, float, …	--name VALUE, converted with the type
bool	--name and --no-name
X | None	as X, with None when the option is not given
list[X]	--name A B C, any number of values, each converted with X
dict[K, V]	--name k1:v1 k2:v2, each split at its first colon and converted with K and V
tuple[X, Y]	exactly that many values, each converted with its own type
tuple[X, ...]	any number of values, converted with X

A field with a default is optional, and its help ends with (default: ...). Put %(default)s in the help text to place the default yourself. A field without a default is a required option. Anything argparse-specific goes in the field’s metadata, which is passed to add_argument and overrides what the field implies: help, choices, nargs, metavar, a custom type, and so on.

The script’s description is the decorator’s description= argument if given (@experiment(description="...")), else a module-level DESCRIPTION, else main’s docstring.

When main runs

Run as a script, the module is __main__, and the decorator parses the command line and calls main as it decorates it. Everything main relies on must therefore be defined above it in the file: the config class, and any method the script registers (see Custom methods). Imported from another module, nothing runs.

The decorated function can also be called directly. main(Config(...)) uses that config and skips parsing, and main(argv=["--config-path", "configs.yaml"]) parses the given list instead of the command line. A main that takes no parameters gets no parser and simply runs.

A command line argparse cannot read, such as a missing required option or a value that does not convert, exits with status 2 and a usage message.

The scripts on the Examples page are all built this way.