prompt_utils¶

prompt_utils ¶

Prompts with target tokens to track, on the standard values.

get_first_tokens turns words into the token ids a model would predict for them, Prompt pairs a prompt with named sets of those, and run_prompts runs many prompts and returns each target's probability mass per prompt.

TokenizationError ¶

Bases: Exception

A word could not be tokenized as a standalone first token.

Prompt `dataclass` ¶

Prompt(prompt: str, target_tokens: dict[str, list[int]], target_strings: dict[str, str | list[str]] | None = None)

A prompt with named sets of target tokens to track in the next-token distribution.

ATTRIBUTE	DESCRIPTION
`prompt`	The text. TYPE: `str`
`target_tokens`	Target name -> the token ids that count for it. TYPE: `dict[str, list[int]]`
`target_strings`	What those came from, when built with `from_strings`. TYPE: `dict[str, str \| list[str]] \| None`

prompt `instance-attribute` ¶

prompt: str

target_tokens `instance-attribute` ¶

target_tokens: dict[str, list[int]]

target_strings `class-attribute` `instance-attribute` ¶

target_strings: dict[str, str | list[str]] | None = None

from_strings `classmethod` ¶

from_strings(prompt: str, target_strings: dict[str, str | list[str]] | list[str] | str, model_or_tokenizer: StandardizedTransformer | PreTrainedTokenizerBase) -> 'Prompt'

Build from words: a string or list is one target named "target", a dict names them.

has_no_collisions ¶

has_no_collisions(ignore_targets: str | list[str] | None = None) -> bool

Whether no token id belongs to two targets (ignore_targets left out).

get_target_probs ¶

get_target_probs(probs: Tensor, layer: int | None = None) -> dict[str, Tensor]

Each target's probability mass from probs of shape [batch, layers, vocab_size]; one layer if given.

run ¶

run(model: StandardizedTransformer, get_probs: Callable) -> dict[str, Tensor]

get_probs(model, prompt) ([batch, layers, vocab_size]) reduced to each target's mass.

get_first_tokens ¶

get_first_tokens(words: str | list[str], model_or_tokenizer: StandardizedTransformer | PreTrainedTokenizerBase, use_hacky_implementation: bool = False) -> list[int]

The first token of word and of " word" for each word, deduplicated.

A model's StandardizedTransformer.add_prefix_false_tokenizer is used when a model is given, so that "word" and " word" tokenize differently; with a tokenizer that adds a prefix space the two collide, and the hacky implementation (tokenize "🍐word" and drop the pear) is used instead, or forced with use_hacky_implementation.

next_token_probs_unsqueeze ¶

next_token_probs_unsqueeze(model: StandardizedTransformer, prompt: str | list[str], remote: bool = False, **_) -> Tensor

compute_next_token_probs with a layer axis of one, [batch, 1, vocab_size]: the default get_probs.

run_prompts ¶

run_prompts(model: StandardizedTransformer, prompts: list[Prompt], batch_size: int = 32, get_probs_func: Callable | None = None, func_kwargs: dict | None = None, remote: bool = False, tqdm=None) -> dict[str, Tensor]

Run prompts in batches; each target's probability mass per prompt, [num_prompts, layers].

All prompts must name the same targets. get_probs_func(model, batch, remote=..., **func_kwargs) returns [batch, layers, vocab_size]; the default is the next-token distribution with one layer. tqdm is a progress-bar factory to wrap the batch loop with, or None.