Skip to content

lophius ⚓︎

A workbench for language model research.

Lophius turns a notebook into an interactive workbench for inspecting language models. It loads models and prompts, generates completions while capturing the model's internal signals, and displays all of it in rich, interactive viewers.

The entry points are init, which prepares a notebook for displaying viewers, load, which loads one or several models, and load_prompts, which loads a collection of prompts.

Example
# Run in the first cell.
import lophius
lophius.init()

# Run in the second cell.
model = lophius.load("Qwen/Qwen3.5-4B")
model

# Run in the third cell.
prompts = lophius.load_prompts(["Why is the sky blue?"])
prompts

# Run in the fourth cell.
output = model(prompts)
output

# Run in the fifth cell.
model.chat

Modules:

  • models

    Loading, inspecting, and generating with language models.

  • prompts

    Collections of text prompts to generate for.

  • outputs

    Generation results and the model internals captured along with them.

Functions:

  • init

    Initialize Lophius for notebook use.

  • load

    Load one or several models.

  • load_prompts

    Load prompts from a file, a Hugging Face dataset, or a list.

init ⚓︎

init() -> None

Initialize Lophius for notebook use.

Call this once, before displaying any Lophius viewer (e.g. a Model or Prompts) in a notebook.

load ⚓︎

load(models: str | PathLike[str], /, model_kwargs: dict[str, Any] | None = None, tokenizer_kwargs: dict[str, Any] | None = None) -> Model
load(models: list[ModelSpec], /, model_kwargs: dict[str, Any] | None = None, tokenizer_kwargs: dict[str, Any] | None = None) -> Models
load(models: str | PathLike[str] | list[ModelSpec], /, model_kwargs: dict[str, Any] | None = None, tokenizer_kwargs: dict[str, Any] | None = None) -> Model | Models

Load one or several models.

This is the main entry point for loading models. It dispatches on its first argument, mirroring the two constructors it forwards to:

  • Given a single pretrained_model_name_or_path (a string or path), it loads and returns a single Model.
  • Given a list of model specifications, it loads and returns a Models collection.

In both cases the optional model_kwargs and tokenizer_kwargs are forwarded unchanged to the corresponding constructor.

Parameters:

  • models (str | PathLike[str] | list[ModelSpec]) –

    A single pretrained_model_name_or_path, or a list of model specifications to load as a collection.

  • model_kwargs (dict[str, Any] | None, default: None ) –

    Optional keyword arguments forwarded to the model loader.

  • tokenizer_kwargs (dict[str, Any] | None, default: None ) –

    Optional keyword arguments forwarded to the tokenizer loader.

Returns:

load_prompts ⚓︎

load_prompts(source: str | PathLike[str] | list[str] | list[tuple[str, bool]], /, split: str = _DEFAULT_SPLIT, column: str = _DEFAULT_COLUMN) -> Prompts

Load prompts from a file, a Hugging Face dataset, or a list.

This is the main entry point for loading prompts and forwards directly to Prompts. A string or path naming an existing local file is loaded one nonempty line at a time; any other string or path is treated as a Hugging Face dataset identifier. A list is used directly and may contain either prompt strings or (prompt, active) tuples.

Parameters:

  • source (str | PathLike[str] | list[str] | list[tuple[str, bool]]) –

    An explicit list of prompts, a path to a local text file, or a Hugging Face dataset identifier.

  • split (str, default: _DEFAULT_SPLIT ) –

    The dataset split to load. Ignored for files and lists.

  • column (str, default: _DEFAULT_COLUMN ) –

    The dataset column containing the prompts. Ignored for files and lists.

Returns:

models ⚓︎

Loading, inspecting, and generating with language models.

This module defines Model, a wrapper around a Hugging Face causal or multimodal language model and its tokenizer that exposes the model's architecture, size, and configuration through a single flat set of properties, and generates completions that capture the model's internal signals. Models is the corresponding collection, which applies the same operations to several models at once so their results can be compared.

Classes:

  • Model

    A thin, introspection-friendly wrapper around a Hugging Face model.

  • Models

    A wrapper around multiple Model objects.

Model ⚓︎

Model(pretrained_model_name_or_path: str | PathLike[str], model_kwargs: dict[str, Any] | None = None, tokenizer_kwargs: dict[str, Any] | None = None)

A thin, introspection-friendly wrapper around a Hugging Face model.

This class loads a pretrained causal or multimodal language model together with its tokenizer and exposes a collection of convenience properties that summarize the model's architecture, size, and configuration. These properties are intended to surface the information a researcher most commonly wants when working with a language model, without having to dig through the underlying transformers objects.

Attributes:

  • model (PreTrainedModel) –

    The loaded Hugging Face model.

  • tokenizer (PreTrainedTokenizerBase) –

    The tokenizer associated with the model.

  • source (str) –

    The identifier or path the model was loaded from.

  • model_kwargs (dict[str, Any]) –

    The keyword arguments passed to the model loader, reapplied on every reload.

  • tokenizer_kwargs (dict[str, Any]) –

    The keyword arguments passed to the tokenizer loader, reapplied on every reload_tokenizer.

The model is loaded with automatic dtype selection and device placement, unless dtype and/or device_map are provided explicitly via model_kwargs. The tokenizer is configured for decoder-only generation by ensuring a pad token exists and enabling left-padding.

The supplied keyword arguments are remembered and reapplied whenever the model or tokenizer is rebuilt (see reload and reload_tokenizer), so any custom loading behavior is preserved across reloads.

Parameters:

  • pretrained_model_name_or_path (str | PathLike[str]) –

    A model identifier on the Hugging Face Hub or a path to a local directory containing the model.

  • model_kwargs (dict[str, Any] | None, default: None ) –

    Optional keyword arguments forwarded to the model's from_pretrained. Any keys given here override the defaults; in particular, the automatic dtype/device_map selection is only applied for keys not present here.

  • tokenizer_kwargs (dict[str, Any] | None, default: None ) –

    Optional keyword arguments forwarded to the tokenizer's from_pretrained.

Raises:

  • ValueError

    If no tokenizer could be loaded for the model, or no token usable for padding could be found.

Methods:

  • __call__

    Alias for generate, so a model can be called directly.

  • __repr__

    Return a concise representation identifying the wrapped model.

  • generate

    Generate completions for one or many prompts, capturing internals.

  • reload

    Rebuild the model from its source using the current configuration.

  • reload_tokenizer

    Rebuild the tokenizer from its source using the current config.

attention_classes property ⚓︎

attention_classes: dict[type[Module], int]

A count of the attention module classes used by the model.

Attention modules are identified heuristically, by the decoder layers' submodules whose class name ends in an attention keyword. For a homogeneous model this typically contains a single entry; hybrid or mixed-attention models (e.g. alternating full and sliding-window attention) yield one entry per distinct attention class.

Returns:

  • dict[type[Module], int]

    A dictionary mapping each attention module class to the number of instances of that class, sorted by descending count.

attention_implementation property ⚓︎

attention_implementation: str | None

The attention implementation in use (e.g. sdpa or eager).

attention_layer_indices property ⚓︎

attention_layer_indices: list[int]

The indices of the decoder layers that contain an attention module.

A layer is considered an attention layer when any of its submodules is classified as attention (by the same heuristic as attention_classes). For a homogeneous model this is every layer; for a hybrid model that interleaves attention with other block types (e.g. gated-delta-net or Mamba layers) it is the subset that actually attends.

Captured attentions cover only the attention layers, so this is what maps them back to their true layer positions.

Returns:

  • list[int]

    The zero-based indices of the attention-bearing decoder layers, in layer order.

chat property ⚓︎

chat: Viewer

A chat viewer for interactively chatting with the model.

Displaying this in a notebook brings up a chat interface backed by the wrapped model, streaming each reply as it is generated.

Returns:

  • Viewer

    A Panel component rendering an interactive chat with the model.

chat_template property ⚓︎

chat_template: str | None

The tokenizer's chat (prompt) template, if one is defined.

Returns:

  • str | None

    The Jinja chat template string used to format conversations, or None if the tokenizer does not define one.

config property ⚓︎

config: PretrainedConfig

The configuration object of the underlying model.

devices property ⚓︎

devices: list[device]

The distinct devices across which the model's parameters are placed.

dtype property ⚓︎

dtype: dtype

The floating-point dtype most commonly used by the model's parameters.

Only the model's parameters (weights) are considered, so that integer buffers cannot be reported as the model's dtype.

dtypes property ⚓︎

dtypes: dict[dtype, int]

A count of parameters and buffers by their dtype.

All tensor dtypes are counted, including non-floating-point buffers such as integer position-id or rotary-embedding index buffers.

Returns:

  • dict[dtype, int]

    A dictionary mapping each dtype present in the model to the number of tensors (parameters and buffers) of that dtype, sorted by descending count.

eos_token_ids property ⚓︎

eos_token_ids: set[int]

The token ids at which the model stops generating.

A model may declare its stop tokens on the tokenizer, in its generation config, or in its model config, and may declare several rather than one. This collects them from every such source, so it reflects what generate actually stops on.

Returns:

  • set[int]

    The end-of-sequence ids, or an empty set if the model declares none (in which case generation only stops on length).

head_dim property ⚓︎

head_dim: int | None

The dimensionality of each attention head.

Returns the explicitly declared head_dim if present, otherwise derives it from the hidden size and number of attention heads when both are available.

Returns:

  • int | None

    The per-head dimensionality, or None if it cannot be determined.

hidden_size property ⚓︎

hidden_size: int | None

The dimensionality of the language model's hidden states, if declared.

intermediate_size property ⚓︎

intermediate_size: int | list[int] | None

The dimensionality of the feed-forward (MLP) intermediate layer.

A model that sizes each layer independently (e.g. Gemma 3n) declares one size per layer rather than a single one.

is_multimodal property ⚓︎

is_multimodal: bool

Whether the model handles any modality beyond text.

This covers audio as well as vision (image and video) components, so an audio-language model counts as multimodal just as a vision one does. It is equivalent to modalities holding more than just "text".

language_model property ⚓︎

language_model: PreTrainedModel

The language (text) model within the overall model.

For multimodal models the decoder stack lives under a dedicated language sub-model, separate from the vision tower or audio encoder. This property resolves to that language model, falling back to the base model for text-only models, so that text-centric statistics are computed over the correct component.

Returns:

  • PreTrainedModel

    The PreTrainedModel representing the language component.

layer_class_sequence property ⚓︎

layer_class_sequence: list[type[Module]]

The decoder layer classes in their actual order within the model.

This preserves the interleaving pattern of hybrid models, so that e.g. an alternating attention/Mamba stack can be read off directly.

Returns:

  • list[type[Module]]

    A list with one class per layer, in layer order.

layer_classes property ⚓︎

layer_classes: dict[type[Module], int]

A count of the decoder layer classes used by the model.

For a homogeneous model this contains a single entry. For hybrid models that interleave different block types (e.g. attention and state-space or Mamba layers) it contains one entry per distinct layer class, revealing the architecture's composition.

Returns:

  • dict[type[Module], int]

    A dictionary mapping each decoder layer class to the number of layers of that class, sorted by descending count.

layers property ⚓︎

layers: ModuleList

The list of transformer decoder layers.

This property transparently handles both multimodal models, where the language layers are nested under a language sub-model, and text-only models, where they live directly under the base model.

Returns:

Raises:

max_position_embeddings property ⚓︎

max_position_embeddings: int | None

The maximum context length supported by the model, if declared.

modalities property ⚓︎

modalities: list[str]

The modalities the model handles, in a stable order.

Every model Lophius loads is a language model, so "text" is always present and comes first; any non-text modality the model declares follows it. A text-only model therefore yields ["text"], and is_multimodal is exactly the question of whether there is more than that.

Returns:

  • list[str]

    A list beginning with "text", followed by "vision" and/or "audio" where declared.

model_class property ⚓︎

model_class: type[PreTrainedModel]

The class of the underlying model (e.g. LlamaForCausalLM).

model_type property ⚓︎

model_type: str | None

The model type string declared by the configuration (e.g. llama).

module_class_counts property ⚓︎

module_class_counts: dict[type[Module], int]

A count of how often each module class appears in the model.

Returns:

  • dict[type[Module], int]

    A dictionary mapping module classes to the number of module instances of that class, sorted by descending count.

modules property ⚓︎

modules: dict[str, Module]

A mapping from qualified module name to module for every submodule.

Returns:

  • dict[str, Module]

    A dictionary whose keys are dot-separated module paths and whose values are the corresponding Module instances. The root module is included under the empty-string key.

norm_classes property ⚓︎

norm_classes: list[type[Module]]

The distinct normalization layer classes used by the model.

Normalization modules are identified heuristically, by the submodules whose class name ends in a normalization keyword.

Returns:

  • list[type[Module]]

    A list of unique normalization module classes, sorted by class name.

num_attention_heads property ⚓︎

num_attention_heads: int | None

The number of attention heads per layer, if declared.

num_key_value_heads property ⚓︎

num_key_value_heads: int | None

The number of key/value heads per layer (for grouped-query attention).

num_layers property ⚓︎

num_layers: int

The number of transformer decoder layers in the model.

num_modules property ⚓︎

num_modules: int

The total number of submodules in the model, including the root.

num_parameters property ⚓︎

num_parameters: int

The total number of parameters in the model.

num_trainable_parameters property ⚓︎

num_trainable_parameters: int

The number of parameters that require gradients.

source_url property ⚓︎

source_url: str | None

The Hugging Face Hub URL the model was loaded from, if applicable.

A source is considered a Hub repository identifier when it has the namespace/name form and does not point at an existing local path.

Returns:

  • str | None

    The URL of the model's page on the Hugging Face Hub, or None if the model was loaded from a local path.

text_config property ⚓︎

text_config: PretrainedConfig

The configuration of the language (text) component of the model.

For multimodal models the top-level configuration is a composite that nests the language model's configuration (which holds the attributes a text-centric analysis cares about, such as num_attention_heads) under a sub-config, alongside e.g. a vision_config. This property resolves to that language sub-config, falling back to the top-level configuration for text-only models.

Returns:

  • PretrainedConfig

    The configuration object describing the language model.

tokenizer_class property ⚓︎

tokenizer_class: type[PreTrainedTokenizerBase]

The class of the tokenizer (e.g. LlamaTokenizerFast).

tokenizer_config property ⚓︎

tokenizer_config: dict[str, Any]

The tokenizer's construction keyword arguments.

This is the tokenizer-side analogue of config: the keyword arguments the tokenizer was instantiated with, which capture options such as model_max_length, clean_up_tokenization_spaces, or special-token settings. It is the configuration that reload_tokenizer re-applies.

Every value is JSON-serializable, so the result can be rendered and edited in a JSON view.

Returns:

  • dict[str, Any]

    A dictionary of the tokenizer's construction keyword arguments.

vocab_size property ⚓︎

vocab_size: int

The size of the tokenizer's vocabulary, including added tokens.

__call__ ⚓︎

__call__(prompts: str | list[str] | Prompts | dict[str, str | list[str] | Prompts], system_prompt: str | None = None, return_logits: bool = True, return_hidden_states: bool = False, return_attentions: bool = False, batch_size: int | None = None, generation_kwargs: dict[str, Any] | None = None, show_progress: bool = True, top_logits: int = 100) -> Output

Alias for generate, so a model can be called directly.

Calling model(prompts, ...) is exactly equivalent to model.generate(prompts, ...); see generate for the full description of the arguments and the returned Output.

__repr__ ⚓︎

__repr__() -> str

Return a concise representation identifying the wrapped model.

generate ⚓︎

generate(prompts: str | list[str] | Prompts | dict[str, str | list[str] | Prompts], system_prompt: str | None = None, return_logits: bool = True, return_hidden_states: bool = False, return_attentions: bool = False, batch_size: int | None = None, generation_kwargs: dict[str, Any] | None = None, show_progress: bool = True, top_logits: int = 100) -> Output

Generate completions for one or many prompts, capturing internals.

Every prompt is wrapped as a single user turn, formatted with the chat template (optionally behind a shared system_prompt), and completed by the model. The result is a rich Output that pairs each prompt with its completion and, when requested, the internal signals produced along the way (logits, hidden states, and attentions).

Prompts may be given as a single string, a list of strings, a Prompts (only its active prompts are used), or a dict mapping group labels to any of those. The group labels are retained in the Output, which makes it possible to tell prompt groups apart when analyzing the results.

Generation is deterministic: a fixed seed is applied regardless of the chosen sampling parameters, so a run can be reproduced exactly (for a given batch size).

Parameters:

  • prompts (str | list[str] | Prompts | dict[str, str | list[str] | Prompts]) –

    The prompts to generate for, in any of the forms above.

  • system_prompt (str | None, default: None ) –

    An optional system prompt prepended to every prompt.

  • return_logits (bool, default: True ) –

    Whether to capture per-step logits (truncated to the top top_logits tokens), over the whole completion. On by default: the truncation makes this cheap.

  • return_hidden_states (bool, default: False ) –

    Whether to capture the first-token hidden states.

  • return_attentions (bool, default: False ) –

    Whether to capture the prompt's attention weights. This forces the eager attention implementation, which is slower and more memory-hungry, since the fused kernels do not expose their weights.

  • batch_size (int | None, default: None ) –

    An explicit batch size. When None (the default) a usable batch size is discovered automatically, backing off whenever a batch runs out of memory.

  • generation_kwargs (dict[str, Any] | None, default: None ) –

    Optional keyword arguments forwarded to generate (e.g. max_new_tokens or temperature). A default max_new_tokens is applied only when not overridden.

  • show_progress (bool, default: True ) –

    Whether to display a progress bar over the prompts.

  • top_logits (int, default: 100 ) –

    The number of highest-logit tokens to keep per generated position when capturing logits. Truncating the full vocabulary to this many keeps the captured logits small enough that capturing them by default is practical.

Returns:

Raises:

  • ValueError

    If no prompts are given, if batch_size is not positive, if top_logits is not positive, if the tokenizer does not left-pad, or if the tokenizer has no chat template.

reload ⚓︎

reload() -> None

Rebuild the model from its source using the current configuration.

Many configuration options are consumed only once, when the model is constructed, to determine the shapes of its weight tensors and other structural properties. Editing the live config object therefore has no effect on an already-instantiated model. This method re-instantiates the model from source, applying the (possibly edited) configuration, so that such construction-time options take effect.

Weights whose shapes no longer match the edited configuration are freshly initialized rather than causing a load error, so that dimensional changes are applied instead of rejected.

The tokenizer is left untouched, as it is unaffected by the model configuration.

reload_tokenizer ⚓︎

reload_tokenizer() -> None

Rebuild the tokenizer from its source using the current config.

Like model options, most tokenizer options are consumed only when the tokenizer is constructed, so editing the live tokenizer_config has no effect on an already-instantiated tokenizer. This method re-instantiates the tokenizer from source, applying the (possibly edited) construction keyword arguments, so that such options take effect. The rebuilt tokenizer is configured for decoder-only generation exactly as on initial load.

Raises:

  • ValueError

    If no tokenizer could be loaded, or no token usable for padding could be found.

Models ⚓︎

Models(models: list[ModelSpec], model_kwargs: dict[str, Any] | None = None, tokenizer_kwargs: dict[str, Any] | None = None)

A wrapper around multiple Model objects.

A Models object holds a collection of models and exists to perform actions on all of them at once. The models it wraps can be specified in a mix of forms for convenience (see below): already-loaded Model instances are used as-is, while strings and dictionaries are loaded into Model instances on construction.

Attributes:

  • models (list[Model]) –

    The wrapped models.

Each entry in models may be:

  • A Model instance, which is used as-is. Such an entry is already loaded and carries its own keyword arguments, so the shared model_kwargs/tokenizer_kwargs below do not apply to it.
  • A pretrained_model_name_or_path (a string or path), which is loaded into a Model using the shared keyword arguments.
  • A dictionary with a pretrained_model_name_or_path key and optional model_kwargs and tokenizer_kwargs keys. The per-entry keyword arguments are merged on top of the shared ones, so an entry can add to or override the shared defaults.

Parameters:

  • models (list[ModelSpec]) –

    The models to wrap, in any combination of the forms above.

  • model_kwargs (dict[str, Any] | None, default: None ) –

    Optional keyword arguments applied to every entry that is not already a Model. For dictionary entries, these are overridden by the entry's own model_kwargs.

  • tokenizer_kwargs (dict[str, Any] | None, default: None ) –

    Optional keyword arguments applied to every entry that is not already a Model. For dictionary entries, these are overridden by the entry's own tokenizer_kwargs.

Raises:

  • TypeError

    If an entry is not a Model, string, path, or dictionary, or if a dictionary entry's pretrained_model_name_or_path is not a string or path or its model_kwargs/tokenizer_kwargs are not dictionaries.

  • ValueError

    If models is empty, or if a dictionary entry lacks a pretrained_model_name_or_path key.

Methods:

  • __call__

    Alias for generate, so a collection can be called directly.

  • __getitem__

    Return the model(s) selected by an index, slice, or source.

  • __iter__

    Iterate over the wrapped models, in order.

  • __len__

    Return the number of wrapped models.

  • __repr__

    Return a concise representation listing the wrapped models.

  • generate

    Generate for the same prompts with every wrapped model.

chat property ⚓︎

chat: Viewer

A chat viewer for interactively chatting with all models at once.

Displaying this in a notebook brings up a chat interface backed by every wrapped model. Each user message is answered by all models in turn, and the user selects which reply to continue the conversation from.

Returns:

  • Viewer

    A Panel component rendering an interactive multi-model chat.

__call__ ⚓︎

__call__(prompts: str | list[str] | Prompts | dict[str, str | list[str] | Prompts], system_prompt: str | None = None, return_logits: bool = True, return_hidden_states: bool = False, return_attentions: bool = False, batch_size: int | None = None, generation_kwargs: dict[str, Any] | None = None, show_progress: bool = True, top_logits: int = 100) -> Outputs

Alias for generate, so a collection can be called directly.

Calling models(prompts, ...) is exactly equivalent to models.generate(prompts, ...); see generate for the full description of the arguments and the returned Outputs.

__getitem__ ⚓︎

__getitem__(index: int) -> Model
__getitem__(index: slice) -> Models
__getitem__(index: str) -> Model | Models
__getitem__(index: int | slice | str) -> Model | Models

Return the model(s) selected by an index, slice, or source.

Indexing with a single integer returns the corresponding Model. Indexing with a slice returns a new Models holding the selected models, which are reused as-is (not reloaded). Indexing with a string selects models by their source: a single match returns that Model, while multiple matches return a new Models holding all of them (in order). In every case the selected models are reused as-is.

Parameters:

  • index (int | slice | str) –

    An integer index, a slice, or a source string.

Returns:

  • Model | Models

    The Model at an integer index; a new Models for a slice; and, for a source string, the matching Model (one match) or a Models of all matches (several).

Raises:

  • ValueError

    If index is a slice that selects no models, since a Models collection must contain at least one model.

  • KeyError

    If index is a source string that matches no model.

__iter__ ⚓︎

__iter__() -> Iterator[Model]

Iterate over the wrapped models, in order.

__len__ ⚓︎

__len__() -> int

Return the number of wrapped models.

__repr__ ⚓︎

__repr__() -> str

Return a concise representation listing the wrapped models.

generate ⚓︎

generate(prompts: str | list[str] | Prompts | dict[str, str | list[str] | Prompts], system_prompt: str | None = None, return_logits: bool = True, return_hidden_states: bool = False, return_attentions: bool = False, batch_size: int | None = None, generation_kwargs: dict[str, Any] | None = None, show_progress: bool = True, top_logits: int = 100) -> Outputs

Generate for the same prompts with every wrapped model.

Each model runs Model.generate on the same prompts and settings, so the resulting Outputs collects one Output per model, in models order, all produced with identical generation parameters, ready to be compared side by side. Every argument is forwarded unchanged to each model's Model.generate; see it for their full description and the shape of each Output.

Parameters:

  • prompts (str | list[str] | Prompts | dict[str, str | list[str] | Prompts]) –

    The prompts to generate for, in any of the forms accepted by Model.generate.

  • system_prompt (str | None, default: None ) –

    An optional system prompt prepended to every prompt.

  • return_logits (bool, default: True ) –

    Whether to capture per-step logits.

  • return_hidden_states (bool, default: False ) –

    Whether to capture first-token hidden states.

  • return_attentions (bool, default: False ) –

    Whether to capture the prompt's attentions.

  • batch_size (int | None, default: None ) –

    An explicit batch size, or None to discover one automatically (independently per model).

  • generation_kwargs (dict[str, Any] | None, default: None ) –

    Optional keyword arguments forwarded to generate.

  • show_progress (bool, default: True ) –

    Whether to display a progress bar over the prompts.

  • top_logits (int, default: 100 ) –

    The number of highest-logit tokens to keep per generated position when capturing logits.

Returns:

prompts ⚓︎

Collections of text prompts to generate for.

This module defines Prompts, a collection of prompts loaded from a local file, a Hugging Face dataset, or an explicit list. Each prompt carries an active flag, so a subset of a larger collection can be selected (interactively in a notebook or in code) and handed to Model.generate.

Classes:

  • Prompts

    A collection of text prompts loaded from a file, a dataset, or a list.

Prompts ⚓︎

Prompts(source: str | PathLike[str] | list[str] | list[tuple[str, bool]], split: str = _DEFAULT_SPLIT, column: str = _DEFAULT_COLUMN)

A collection of text prompts loaded from a file, a dataset, or a list.

Prompts can be sourced in one of three ways:

  • From a local text file, in which case each nonempty line (after stripping surrounding whitespace) becomes one prompt.
  • From a Hugging Face dataset, in which case the values of a chosen column across a chosen split become the prompts.
  • From an explicit list, either of plain strings or of (prompt, active) tuples, which is used as-is.

Which source is used is decided by the type and value of source; see below.

Each prompt additionally carries an active flag. All prompts are active by default. Iterating over the collection yields only the active prompts, and the flags are preserved when the collection is sliced.

Attributes:

  • source (str | None) –

    The identifier or path the prompts were loaded from, or None when constructed from an explicit list of strings.

  • split (str | None) –

    The dataset split the prompts were loaded from, or None when not loaded from a dataset.

  • column (str | None) –

    The dataset column the prompts were taken from, or None when not loaded from a dataset.

The source is interpreted as follows:

  • If source is a list, it is used directly as the prompts. Its items may be plain strings or (prompt, active) tuples; plain strings default to active.
  • Otherwise, if source is a path to an existing local file, the prompts are the nonempty lines of that file (each stripped of surrounding whitespace).
  • Otherwise, source is treated as a Hugging Face dataset identifier and the prompts are read from the given column of the given split.

Prompts loaded from a file or dataset, and those given as plain strings, start out active; prompts given as (prompt, active) tuples use the flag from the tuple.

Parameters:

  • source (str | PathLike[str] | list[str] | list[tuple[str, bool]]) –

    An explicit list of prompts (plain strings or (prompt, active) tuples), a path to a local text file, or a Hugging Face dataset identifier on the Hub.

  • split (str, default: _DEFAULT_SPLIT ) –

    The dataset split to load. Only used when loading from a dataset. Defaults to the first 100 rows of the training split.

  • column (str, default: _DEFAULT_COLUMN ) –

    The dataset column to read prompts from. Only used when loading from a dataset. Defaults to "text".

Methods:

  • __getitem__

    Return the prompt at an index, or a sub-collection for a slice.

  • __iter__

    Iterate over the active prompts, in order.

  • __len__

    Return the total number of prompts, regardless of their active flag.

  • __repr__

    Return a concise representation showing the source and prompt count.

  • set_active

    Set whether the prompt at index is active.

  • set_prompt

    Set the text of the prompt at index.

active_prompts property ⚓︎

active_prompts: list[str]

The active prompts, in order.

flagged_prompts property ⚓︎

flagged_prompts: list[tuple[str, bool]]

Each prompt paired with its active flag, in order.

Returns:

inactive_prompts property ⚓︎

inactive_prompts: list[str]

The inactive prompts, in order.

is_from_dataset property ⚓︎

is_from_dataset: bool

Whether the prompts were loaded from a Hugging Face dataset.

is_from_file property ⚓︎

is_from_file: bool

Whether the prompts were loaded from a local file.

num_active_prompts property ⚓︎

num_active_prompts: int

The number of active prompts.

num_prompts property ⚓︎

num_prompts: int

The total number of prompts, regardless of their active flag.

prompts property ⚓︎

prompts: list[str]

All prompts, regardless of their active flag, in order.

source_url property ⚓︎

source_url: str | None

The Hugging Face Hub URL the prompts were loaded from, if applicable.

A source has a Hub URL only when the prompts were loaded from a dataset (i.e. split is set).

Returns:

  • str | None

    The URL of the dataset's page on the Hugging Face Hub, or None if the prompts were loaded from a local file or an explicit list.

__getitem__ ⚓︎

__getitem__(index: int) -> str
__getitem__(index: slice) -> Prompts
__getitem__(index: int | slice) -> str | Prompts

Return the prompt at an index, or a sub-collection for a slice.

Indexing with a single integer returns the corresponding prompt string. Indexing with a slice returns a new Prompts holding the selected prompts, preserving each prompt's active flag.

Parameters:

  • index (int | slice) –

    An integer index or a slice.

Returns:

  • str | Prompts

    The prompt at index for an integer, or a new Prompts for a slice.

__iter__ ⚓︎

__iter__() -> Iterator[str]

Iterate over the active prompts, in order.

__len__ ⚓︎

__len__() -> int

Return the total number of prompts, regardless of their active flag.

__repr__ ⚓︎

__repr__() -> str

Return a concise representation showing the source and prompt count.

set_active ⚓︎

set_active(index: int, active: bool) -> None

Set whether the prompt at index is active.

Parameters:

  • index (int) –

    The position of the prompt whose flag should be changed.

  • active (bool) –

    The new active flag.

set_prompt ⚓︎

set_prompt(index: int, prompt: str) -> None

Set the text of the prompt at index.

Parameters:

  • index (int) –

    The position of the prompt whose text should be changed.

  • prompt (str) –

    The new prompt text.

outputs ⚓︎

Generation results and the model internals captured along with them.

This module defines the containers Model.generate returns: a Generation pairs one prompt with its completion and the internal signals (logits, hidden states, attentions) recorded while producing it, an Output collects the generations of one model together with the settings they were produced with, and an Outputs collects the outputs of several models for the same prompts.

Classes:

  • Generation

    The completion of a single prompt, with any captured internal signals.

  • Output

    The result of generating for one or many prompts with a single model.

  • Outputs

    A wrapper around multiple Output objects.

Generation ⚓︎

Generation(group: str, prompt_text: str, prompt_token_ids: list[int], prompt_token_strings: list[str], generated_text: str, generated_token_ids: list[int], generated_token_strings: list[str], top_token_ids: Tensor | None = None, top_logits: Tensor | None = None, top_probs: Tensor | None = None, entropy: Tensor | None = None, first_token_hidden_states: Tensor | None = None, prompt_attentions: list[Tensor] | None = None, attention_layer_indices: list[int] | None = None)

The completion of a single prompt, with any captured internal signals.

A Generation is the per-prompt unit of an Output: it pairs one prompt with the text the model generated for it and, when the corresponding capture was requested at generation time, the internal signals produced along the way. Every stored tensor lives on the CPU, so a Generation outlives the generation call and can be inspected freely.

Attributes:

  • group (str) –

    The label of the prompt group this prompt belonged to. The empty string when no prompt groups were used.

  • prompt_text (str) –

    The original prompt text (before chat templating).

  • prompt_token_ids (list[int]) –

    The token ids of the templated prompt.

  • prompt_token_strings (list[str]) –

    The individual token strings of the prompt, one per id in prompt_token_ids.

  • generated_text (str) –

    The generated tokens decoded to text, with special tokens (including the end-of-sequence token) removed.

  • generated_token_ids (list[int]) –

    The token ids the model generated, up to and including the first end-of-sequence token. Unlike generated_text, the end-of-sequence token is retained here, so the model's decision to stop is visible as an explicit token (and its logits are kept).

  • generated_token_strings (list[str]) –

    The individual token strings of the generated tokens, one per id in generated_token_ids.

  • top_token_ids (Tensor | None) –

    For each generated position, the ids of the tokens with the highest logits, of shape (num_generated, k), or None if logits were not captured. The full vocabulary is truncated to the top k at generation time to keep memory small (see Model.generate's top_logits). The chosen token is only present when it falls within the top k (essentially always under greedy decoding); when it does not, it is simply absent.

  • top_logits (Tensor | None) –

    The raw logits of the tokens in top_token_ids, of shape (num_generated, k), or None if not captured.

  • top_probs (Tensor | None) –

    The full-vocabulary softmax probabilities of the tokens in top_token_ids, of shape (num_generated, k), or None if not captured. Recorded directly rather than derived from the truncated logits, so the true probabilities are preserved.

  • entropy (Tensor | None) –

    The full-vocabulary entropy (in nats) of each generated position's probability distribution, of shape (num_generated,), or None if not captured. Computed over the whole distribution before truncation, so it reflects the model's true uncertainty.

  • first_token_hidden_states (Tensor | None) –

    The hidden states of the last prompt position (the representation from which the first token was generated) across all layers, of shape (num_layers + 1, hidden) (the extra layer is the embedding output), or None if hidden states were not captured.

  • prompt_attentions (list[Tensor] | None) –

    The prompt's self-attention weights: a list with one entry per attention layer, each a tensor of shape (num_heads, prompt_len, prompt_len) over the prompt positions. Kept per-layer rather than stacked because layers may differ in head count. Only attention-bearing layers appear, in order; the corresponding true layer numbers are in attention_layer_indices. None if attentions were not captured.

  • attention_layer_indices (list[int] | None) –

    The zero-based decoder-layer index of each entry in prompt_attentions, so a hybrid model's attentions can be labeled with their real layer numbers. None if attentions were not captured.

This is a plain data container; instances are normally created by Model.generate rather than directly. See the class docstring for the meaning of each argument.

Methods:

  • __repr__

    Return a concise representation showing the group and text lengths.

num_generated_tokens property ⚓︎

num_generated_tokens: int

The number of tokens generated for this prompt.

num_prompt_tokens property ⚓︎

num_prompt_tokens: int

The number of prompt tokens.

__repr__ ⚓︎

__repr__() -> str

Return a concise representation showing the group and text lengths.

Output ⚓︎

Output(model: Model, generations: list[Generation], system_prompt: str | None = None, has_logits: bool = False, has_hidden_states: bool = False, has_attentions: bool = False, generation_kwargs: dict[str, Any] | None = None, batch_size: int | None = None, top_logits: int = 100)

The result of generating for one or many prompts with a single model.

An Output bundles the model that produced it, the settings it was produced with, and one Generation per prompt. It holds a reference back to its live Model, while the settings that describe this generation specifically (system prompt, generation kwargs, and which signals were captured) are snapshotted so they remain accurate even if the model is later reloaded.

Attributes:

  • model (Model) –

    The model that produced these generations.

  • generations (list[Generation]) –

    One Generation per prompt, in the order the prompts were given (grouped by prompt group).

  • system_prompt (str | None) –

    The system prompt applied to every prompt, if any.

  • has_logits (bool) –

    Whether logits were captured.

  • has_hidden_states (bool) –

    Whether hidden states were captured.

  • has_attentions (bool) –

    Whether attentions were captured.

  • generation_kwargs (dict[str, Any]) –

    The resolved keyword arguments passed to generate.

  • batch_size (int | None) –

    The batch size used (the discovered size when it was found automatically), or None if generation processed everything in a single batch.

  • top_logits (int) –

    The number of highest-logit tokens kept per generated position when logits were captured.

This is normally created by Model.generate rather than directly. See the class docstring for the meaning of each argument.

Methods:

  • __getitem__

    Return the generation at index.

  • __iter__

    Iterate over the generations, in order.

  • __len__

    Return the number of generations (one per prompt).

  • __repr__

    Return a concise representation showing the model and captured signals.

  • group

    Return the generations belonging to a prompt group.

groups property ⚓︎

groups: list[str]

The distinct prompt-group labels, in first-seen order.

When no prompt groups were used every generation shares the implicit empty label, so this is simply [""].

selected_hidden_states property ⚓︎

selected_hidden_states: Tensor | None

The hidden-state vectors currently selected in the viewer.

When points are selected in the Hidden States tab's scatter plot (by dragging a selection box), this returns their first-token hidden states at the layer currently displayed, stacked into a (num_selected, hidden) tensor, letting a visual selection be pulled straight into code (e.g. assigned to a variable in another notebook cell). The rows are in prompt order, and the layer is whichever the plot currently shows.

Returns:

  • Tensor | None

    A (num_selected, hidden) tensor of the selected layer's hidden states, or None if no points are selected (or hidden states were never captured).

__getitem__ ⚓︎

__getitem__(index: int) -> Generation

Return the generation at index.

__iter__ ⚓︎

__iter__() -> Iterator[Generation]

Iterate over the generations, in order.

__len__ ⚓︎

__len__() -> int

Return the number of generations (one per prompt).

__repr__ ⚓︎

__repr__() -> str

Return a concise representation showing the model and captured signals.

group ⚓︎

group(label: str) -> list[Generation]

Return the generations belonging to a prompt group.

Parameters:

  • label (str) –

    The prompt-group label to select (the empty string selects the implicit group used when no groups were given).

Returns:

  • list[Generation]

    The generations whose group is label, in order.

Outputs ⚓︎

Outputs(outputs: list[Output])

A wrapper around multiple Output objects.

An Outputs object holds a collection of outputs and exists to display and compare them together. It is the multi-model analogue of an Output: Models.generate produces one, collecting the Output of every model for the same prompts and settings so the per-model results can be viewed side by side.

It simply wraps an already-produced list of Output objects.

Attributes:

Parameters:

Raises:

  • ValueError

    If outputs is empty, since an Outputs collection must contain at least one output.

Methods:

  • __getitem__

    Return the output(s) selected by an index, slice, or model source.

  • __iter__

    Iterate over the wrapped outputs, in order.

  • __len__

    Return the number of wrapped outputs.

  • __repr__

    Return a concise representation listing the wrapped outputs.

__getitem__ ⚓︎

__getitem__(index: int) -> Output
__getitem__(index: slice) -> Outputs
__getitem__(index: str) -> Output | Outputs
__getitem__(index: int | slice | str) -> Output | Outputs

Return the output(s) selected by an index, slice, or model source.

Mirroring Models' own indexing: indexing with a single integer returns the corresponding Output; a slice returns a new Outputs holding the selected outputs (reused as-is); and a string selects outputs by their producing model's source, returning that Output for a single match or a new Outputs of all matches (in order) for several.

Parameters:

  • index (int | slice | str) –

    An integer index, a slice, or a model source string.

Returns:

Raises:

  • ValueError

    If index is a slice that selects no outputs, since an Outputs collection must contain at least one output.

  • KeyError

    If index is a source string that matches no output.

__iter__ ⚓︎

__iter__() -> Iterator[Output]

Iterate over the wrapped outputs, in order.

__len__ ⚓︎

__len__() -> int

Return the number of wrapped outputs.

__repr__ ⚓︎

__repr__() -> str

Return a concise representation listing the wrapped outputs.