dew.sampling.decoding

Logit transforms and stopping criteria for the decode loop.

Three concepts extend decoding. A LogitsTransform is a pure callable from the step state and [rows, vocab] logits to new logits. A Stopping is a pure callable from the step state and the tokens just drawn to a per-row finished flag; criteria combine with OR after every committed token. A Strategy in dew.sampling.strategies owns the device loop.

Transforms and criteria read StepState, which carries the token history and nothing about the model: no parameters, no cache. Every built-in here is a pytree, so a configuration holding arrays crosses jax.jit as data instead of entering a compilation cache key. A plain function works as well, and jax.tree_util.Partial carries array configuration for one.

The numerical reference is Transformers 5.16.1 generation/logits_process.py and generation/stopping_criteria.py, with two differences that follow from Dew’s decode loop. Each row reads its own unpadded history rather than the batch’s padded width, and frequency and presence penalties follow vLLM’s formula (model_executor/layers/utils.py), which Transformers does not implement.

Name	Summary
FILTER	The score a removed token keeps, as logits_process.py’s filter value.
StepState	What a transform or a criterion sees at one decode step.
LogitsTransform	A pure [rows, vocab] score rewrite, applied before the draw.
Stopping	A pure per-row finish test over the tokens a step just drew.
Greedy	The argmax as a distribution: zero on the best token, -inf elsewhere.
Temperature	logits / value, as TemperatureLogitsWarper.
TopK	Keep the k highest scores, as TopKLogitsWarper.
TopP	Nucleus filtering, as TopPLogitsWarper.
MinP	Relative filtering at p times the top probability, as MinPLogitsWarper.
Typical	Locally typical filtering, as TypicalLogitsWarper.
EpsilonCutoff	Remove tokens below an absolute probability, as EpsilonLogitsWarper.
EtaCutoff	Entropy-scaled cutoff, as EtaLogitsWarper.
TopH	Entropy-budget filtering, as TopHLogitsWarper.
Renormalize	Replace scores by their log softmax, as LogitNormalization.
RemoveInvalidValues	Map NaN to zero and infinities to the float range, as InfNanRemoveLogitsProcessor.
RepetitionPenalty	Divide positive scores of seen tokens and multiply negative ones.
PromptRepetitionPenalty	Raise the scores of prompt tokens, as EncoderRepetitionPenaltyLogitsProcessor.
FrequencyPenalty	Subtract penalty times each token’s count among the drawn tokens.
PresencePenalty	Subtract penalty from every token already drawn.
NoRepeatNGram	Ban tokens that would repeat an n-gram of the row’s own history.
PromptNoRepeatNGram	Ban tokens that would repeat an n-gram of the prompt.
SequenceBias	Add a bias to the token that would complete each biased sequence.
sequence_bias	A SequenceBias table from (token ids, bias) pairs.
bad_words	A -inf SequenceBias over forbidden sequences, as NoBadWordsLogitsProcessor.
SuppressTokens	Remove a fixed set of tokens, as SuppressTokensLogitsProcessor.
BeginSuppressTokens	Remove tokens at one generated position, as SuppressTokensAtBeginLogitsProcessor.
ForcedBOS	Force one token as the first of the whole sequence, as ForcedBOSTokenLogitsProcessor.
ForcedEOS	Force EOS one step before the end, as ForcedEOSTokenLogitsProcessor.
MinLength	Suppress EOS until the whole sequence reaches length, as MinLengthLogitsProcessor.
MinNewTokens	Suppress EOS until count tokens are drawn, as MinNewTokensLengthLogitsProcessor.
ExponentialDecayLengthPenalty	Grow the EOS score after start drawn tokens, as ExponentialDecayLengthPenalty.
EndOfSequence	Finish a row that drew one of the EOS ids, as EosTokenCriteria.
MaxNewTokens	Finish a row once it has drawn count tokens.
MaxLength	Finish a row once prompt and generated tokens reach length, as MaxLengthCriteria.
StopStrings	Finish a row whose text ends with one of the compiled stop strings.
Vocabulary	The tokenizer surface stop_strings reads once, on the host.
PRINTABLE	The bytes GPT-2’s byte-level alphabet maps to themselves.
byte_alphabet	GPT-2’s byte-to-unicode table, inverted.
PieceDecoder	A tokenizers decoder, which pickles as the JSON that configures it; matching_mode reads the decoder kinds that JSON names.
Backend	The Rust tokenizer a fast Transformers tokenizer wraps: its decoder says how pieces spell bytes, and is None on a tokenizer without one.
Fast	A fast Transformers tokenizer, which carries its Rust backend; a slow one has no backend and its pieces are read through their text.
Referencing	A processor that holds the source’s own processor or tokenizer as reference, as dew.interop.pretrained.Processor does.
Tokenizing	A processor that holds its tokenizer: a Transformers processor, a run’s RunProcessor, or dew.data.HFTokenizer over the hub one.
matching_mode	Whether a tokenizer’s pieces are bytes, and in which spelling.
vocabulary_pieces	What each vocabulary entry contributes to the text, and its id.
stop_strings	Compile a tokenizer’s vocabulary against strings into a StopStrings.
as_pytree	value in a form jax.jit accepts as data.
components	values as a tuple of pytrees jax.jit accepts as data.
chain	The transforms as one callable, applied in order.
criterion	The criteria as one callable, combined with OR.

FILTER

attribute source

FILTER = -jnp.inf

The score a removed token keeps, as logits_process.py’s filter value.

StepState

dataclass source

class StepState(prompt_width: int = struct.field(pytree_node=False, default=0))

What a transform or a criterion sees at one decode step.

tokens is the fixed-capacity buffer of the prompt followed by the draw slots, [rows, prompt_width + budget], and valid marks the slots that hold a real token. A row’s history is therefore its own, whatever padding the prompt batch needed. step counts the tokens the row has committed, active marks the rows still generating, and keys holds one PRNG key per row.

width: int

Slots in the buffer, prompt plus budget.

StepState.total

def total() -> jax.Array

Real tokens each row holds, prompt and generated together.

StepState.history

def history() -> tuple[jax.Array, jax.Array]

Each row’s real tokens left aligned, and how many there are.

Prompts pad wherever their batch needed it, so a transform that reads order (n-grams, biased sequences, stop strings) needs the row’s own tokens without holes. Padding slots hold -1, which no token id equals.

StepState.prompt_history

def prompt_history() -> tuple[jax.Array, jax.Array]

The prompt region’s real tokens left aligned, and how many.

StepState.generated

def generated() -> tuple[jax.Array, jax.Array]

The drawn region’s real tokens left aligned, and how many.

StepState.commit

def commit(tokens: jax.Array, drawn: jax.Array) -> StepState

The state after drawn rows appended tokens at their next slot.

LogitsTransform

class source

class LogitsTransform(Protocol)

A pure [rows, vocab] score rewrite, applied before the draw.

Stopping

class source

class Stopping(Protocol)

A pure per-row finish test over the tokens a step just drew.

Greedy

dataclass source

class Greedy()

The argmax as a distribution: zero on the best token, -inf elsewhere.

Sampling(temperature=0) compiles to this, so a zero-temperature draw stays the deterministic argmax and its behaviour log probability stays exactly zero while running through the same categorical draw as any other policy. Transforms placed before it still shape the argmax, which is what greedy search does with a processor list.

A row that arrives without a distribution leaves without one. A point mass over an all-removed row, or over a NaN or +inf the model or an earlier transform produced, would turn an undefined draw into a confident token, so those rows pass through and the draw refuses them.

Temperature

dataclass source

class Temperature(value: float = struct.field(pytree_node=False, default=1.0))

logits / value, as TemperatureLogitsWarper.

TopK

dataclass source

class TopK(k: int = struct.field(pytree_node=False, default=1))

Keep the k highest scores, as TopKLogitsWarper.

TopP

dataclass source

class TopP(p: float = struct.field(pytree_node=False, default=1.0))

Nucleus filtering, as TopPLogitsWarper.

The ascending tail holding cumulative mass at most 1 - p is removed and the best token always survives.

MinP

dataclass source

class MinP(p: float = struct.field(pytree_node=False, default=0.0))

Relative filtering at p times the top probability, as MinPLogitsWarper.

Typical

dataclass source

class Typical(mass: float = struct.field(pytree_node=False, default=1.0))

Locally typical filtering, as TypicalLogitsWarper.

EpsilonCutoff

dataclass source

class EpsilonCutoff(epsilon: float = struct.field(pytree_node=False, default=0.0))

Remove tokens below an absolute probability, as EpsilonLogitsWarper.

EtaCutoff

dataclass source

class EtaCutoff(epsilon: float = struct.field(pytree_node=False, default=0.0))

Entropy-scaled cutoff, as EtaLogitsWarper.

TopH

dataclass source

class TopH(
    fraction: float = struct.field(pytree_node=False, default=1.0),
    candidates: int = struct.field(pytree_node=False, default=100),
)

Entropy-budget filtering, as TopHLogitsWarper.

Tokens enter in probability order while the cumulative entropy of the truncated head stays within fraction of its total entropy, and the best token always enters. candidates is the head the reference fixes at 100.

The two entropies are computed the way the reference computes them, and they are not the same expression. The budget is torch.distributions.Categorical.entropy, which clamps the log probabilities to the dtype’s minimum so a removed token contributes nothing. The running sum is the reference’s own -p * log(p), whose removed tokens are NaN, and a NaN ends the selection because every comparison against it is false. Substituting one for the other keeps a token the reference drops.

Renormalize

dataclass source

class Renormalize()

Replace scores by their log softmax, as LogitNormalization.

RemoveInvalidValues

dataclass source

class RemoveInvalidValues()

Map NaN to zero and infinities to the float range, as InfNanRemoveLogitsProcessor.

Nothing else in the chain repairs a broken distribution: an undefined draw raises instead. Ask for this transform to sanitize one.

RepetitionPenalty

dataclass source

class RepetitionPenalty(penalty: float = struct.field(pytree_node=False, default=1.0))

Divide positive scores of seen tokens and multiply negative ones.

The history is the row’s valid prompt and drawn tokens, as RepetitionPenaltyLogitsProcessor reads the whole input_ids.

PromptRepetitionPenalty

dataclass source

class PromptRepetitionPenalty(
    penalty: float = struct.field(pytree_node=False, default=1.0),
)

Raise the scores of prompt tokens, as EncoderRepetitionPenaltyLogitsProcessor.

The reference inverts its argument, so a penalty above one rewards repeating the prompt. A decoder-only prompt is the encoder input here.

FrequencyPenalty

dataclass source

class FrequencyPenalty(penalty: float = struct.field(pytree_node=False, default=0.0))

Subtract penalty times each token’s count among the drawn tokens.

vLLM’s formula, logits -= frequency_penalties * output_bin_counts (model_executor/layers/utils.py), which is also OpenAI’s frequency_penalty. It counts generated tokens, not the prompt.

PresencePenalty

dataclass source

class PresencePenalty(penalty: float = struct.field(pytree_node=False, default=0.0))

Subtract penalty from every token already drawn.

vLLM’s logits -= presence_penalties * output_mask, which is OpenAI’s presence_penalty. It reads generated tokens, not the prompt.

NoRepeatNGram

dataclass source

class NoRepeatNGram(size: int = struct.field(pytree_node=False, default=0))

Ban tokens that would repeat an n-gram of the row’s own history.

The tensorised form of NoRepeatNGramLogitsProcessor: the current suffix is matched against every window, and a matching window bans the token that followed it. The window starting at the suffix itself needs one token more than the row has, so a suffix never bans its own successor.

PromptNoRepeatNGram

dataclass source

class PromptNoRepeatNGram(size: int = struct.field(pytree_node=False, default=0))

Ban tokens that would repeat an n-gram of the prompt.

EncoderNoRepeatNGramLogitsProcessor builds its table from the encoder input and matches it against the decoder’s suffix. The prompt is the encoder input of a decoder-only model.

SequenceBias

dataclass source

class SequenceBias()

Add a bias to the token that would complete each biased sequence.

SequenceBiasLogitsProcessor as a table: sequences is [count, width] right-aligned token ids, lengths their real lengths and bias the value added to the last id when the row’s suffix matches the preceding ones. A sequence longer than the row’s history is skipped, as the reference skips one longer than the context.

sequence_bias

function source

def sequence_bias(entries: Sequence[tuple[Sequence[int], float]]) -> SequenceBias

A SequenceBias table from (token ids, bias) pairs.

bad_words

function source

def bad_words(
    ids: Sequence[Sequence[int]],
    eos_id: int | Sequence[int] | None = None,
) -> SequenceBias

A -inf SequenceBias over forbidden sequences, as NoBadWordsLogitsProcessor.

Single-token sequences that name an EOS id are dropped, as the reference drops them, so banning bad words cannot ban termination.

SuppressTokens

dataclass source

class SuppressTokens()

Remove a fixed set of tokens, as SuppressTokensLogitsProcessor.

BeginSuppressTokens

dataclass source

class BeginSuppressTokens(offset: int = struct.field(pytree_node=False, default=0))

Remove tokens at one generated position, as SuppressTokensAtBeginLogitsProcessor.

The reference suppresses where the sequence still has its prompt width, so offset is the generated index the suppression applies at, zero for the first drawn token and one where a forced BOS occupies that slot.

ForcedBOS

dataclass source

class ForcedBOS(token: int = struct.field(pytree_node=False, default=0))

Force one token as the first of the whole sequence, as ForcedBOSTokenLogitsProcessor.

ForcedEOS

dataclass source

class ForcedEOS(
    eos: jax.Array = struct.field(default_factory=lambda: jnp.zeros((0,), jnp.int32)),
    max_length: int | None = struct.field(pytree_node=False, default=None),
)

Force EOS one step before the end, as ForcedEOSTokenLogitsProcessor.

eos may name several ids, all of which the forced step allows, as the reference allows every id its tensor holds. max_length counts prompt and generated tokens together; left as None the end is the request’s own, the prompt width plus the token budget, so a caller that changes the budget per call forces at the new end rather than at a length the task was built with.

MinLength

dataclass source

class MinLength(
    length: int = struct.field(pytree_node=False, default=0),
    eos: jax.Array = struct.field(default_factory=lambda: jnp.zeros((0,), jnp.int32)),
)

Suppress EOS until the whole sequence reaches length, as MinLengthLogitsProcessor.

MinNewTokens

dataclass source

class MinNewTokens(
    count: int = struct.field(pytree_node=False, default=0),
    eos: jax.Array = struct.field(default_factory=lambda: jnp.zeros((0,), jnp.int32)),
)

Suppress EOS until count tokens are drawn, as MinNewTokensLengthLogitsProcessor.

ExponentialDecayLengthPenalty

dataclass source

class ExponentialDecayLengthPenalty(
    start: int = struct.field(pytree_node=False, default=0),
    factor: float = struct.field(pytree_node=False, default=1.0),
    eos: jax.Array = struct.field(default_factory=lambda: jnp.zeros((0,), jnp.int32)),
)

Grow the EOS score after start drawn tokens, as ExponentialDecayLengthPenalty.

The reference measures from start_index + prompt_width, which is the generated count used here, and adds |score| * (factor ** index - 1) so a negative score also rises. A removed EOS (-inf, a grammar’s mask) stays removed.

EndOfSequence

dataclass source

class EndOfSequence()

Finish a row that drew one of the EOS ids, as EosTokenCriteria.

MaxNewTokens

dataclass source

class MaxNewTokens(count: int = struct.field(pytree_node=False, default=0))

Finish a row once it has drawn count tokens.

MaxLength

dataclass source

class MaxLength(length: int = struct.field(pytree_node=False, default=0))

Finish a row once prompt and generated tokens reach length, as MaxLengthCriteria.

StopStrings

dataclass source

class StopStrings(
    positions: int = struct.field(pytree_node=False, default=1),
    ends: int = struct.field(pytree_node=False, default=1),
    span: int = struct.field(pytree_node=False, default=1),
)

Finish a row whose text ends with one of the compiled stop strings.

The tables come from stop_strings, which reads the tokenizer once. The device check is StopStringCriteria’s: walk the row’s tokens backwards, require the last token to overlap the end of a stop string, and keep matching earlier tokens against the positions where they can sit. A match counts only when the string touches the final token, so a string produced earlier does not stop the row later.

Vocabulary

class source

class Vocabulary(Protocol)

The tokenizer surface stop_strings reads once, on the host.

These are a Transformers tokenizer’s own public vocabulary methods plus the piece-name lookup its slow and fast classes both expose. Nothing here runs generation code; the tables are built from token strings.

Vocabulary.get_vocab

def get_vocab() -> dict[str, int]

Vocabulary.convert_tokens_to_string

def convert_tokens_to_string(tokens: list[str]) -> str

PRINTABLE

attribute source

PRINTABLE = list(range(ord('!'), ord('~') + 1)) + list(range(ord('¡'), ord('¬') + 1)) + list(range(ord('®'), ord('ÿ') + 1))

The bytes GPT-2’s byte-level alphabet maps to themselves.

byte_alphabet

function source

def byte_alphabet() -> dict[str, int]

GPT-2’s byte-to-unicode table, inverted.

A byte-level tokenizer stores each byte of a piece as one of these characters, so reading a piece back byte by byte is the only way to keep a code point that two tokens split between them.

PieceDecoder

class source

class PieceDecoder(Protocol)

A tokenizers decoder, which pickles as the JSON that configures it; matching_mode reads the decoder kinds that JSON names.

Backend

class source

class Backend(Protocol)

The Rust tokenizer a fast Transformers tokenizer wraps: its decoder says how pieces spell bytes, and is None on a tokenizer without one.

Fast

class source

class Fast(Protocol)

A fast Transformers tokenizer, which carries its Rust backend; a slow one has no backend and its pieces are read through their text.

Referencing

class source

class Referencing(Protocol)

A processor that holds the source’s own processor or tokenizer as reference, as dew.interop.pretrained.Processor does.

Tokenizing

class source

class Tokenizing(Protocol)

A processor that holds its tokenizer: a Transformers processor, a run’s RunProcessor, or dew.data.HFTokenizer over the hub one.

matching_mode

function source

def matching_mode(tokenizer: Vocabulary) -> str | None

Whether a tokenizer’s pieces are bytes, and in which spelling.

StopStringCriteria._get_stop_string_matching_mode: a byte-level decoder stores pieces in GPT-2’s alphabet and a byte-fallback one spells unknown bytes <0xNN>. Either way the match runs over bytes, so a stop string is encoded to UTF-8 and a piece that is half a code point still counts.

vocabulary_pieces

function source

def vocabulary_pieces(
    tokenizer: Vocabulary,
    mode: str | None,
    prefix: str = 'abcdef',
) -> tuple[list[str | bytes], list[int]]

What each vocabulary entry contributes to the text, and its id.

StopStringCriteria.clean_tokenizer_vocab: a byte-mode piece is read through its byte spelling, and anything else through convert_tokens_to_string behind an ordinary prefix, because a decoder adds or removes a leading space depending on what came before. The prefix is tokenized once and its text is cut off the front of every piece.

stop_strings

function source

def stop_strings(
    tokenizer: Vocabulary | Referencing | Tokenizing,
    strings: str | Sequence[str],
    vocab_size: int | None = None,
) -> StopStrings

Compile a tokenizer’s vocabulary against strings into a StopStrings.

The tables record, for every token, where its piece can sit inside a stop string and how many of the string’s trailing units its start can cover. This runs once on the host; the criterion never decodes.

A byte-level or byte-fallback vocabulary matches over UTF-8 bytes, so a stop string whose code point two tokens split still ends a row. vocab_size sizes the table for the model rather than the tokenizer when a checkpoint pads its head.

as_pytree

function source

def as_pytree(value: LogitsTransform) -> LogitsTransform

value in a form jax.jit accepts as data.

Validated scalar policies lower to partials whose numerical arguments are dynamic leaves. Other built-ins retain their registered pytrees. A plain function is wrapped in a Partial that keeps the function static and its bound arguments as data; strategies use that same callable rule.

components

function source

def components(
    values: LogitsTransform | Sequence[LogitsTransform],
    where: str,
) -> tuple[LogitsTransform, ...]

values as a tuple of pytrees jax.jit accepts as data.

Takes one transform or criterion or a sequence of either. where names the argument in the refusal a non-callable earns.

chain

function source

def chain(
    transforms: Sequence[LogitsTransform],
) -> Callable[[StepState, jax.Array], jax.Array]

The transforms as one callable, applied in order.

criterion

function source

def criterion(
    stopping: Sequence[Stopping],
) -> Callable[[StepState, jax.Array], jax.Array]

The criteria as one callable, combined with OR.