micm_nlp.training.logits_processors

Generation-time constraints on the decoder’s output.

ConstrainedPrefixLogitsProcessor restricts free-form generation to a closed set of allowed strings. At each step it keeps only the candidate sequences that still match what has been generated so far and masks every token that no survivor could produce next; when nothing matches it forces EOS. This turns a generative model into a classifier over allowed_texts without touching the model itself.

Wired up through custom_training_args.generation_whitelist.

Classes

ConstrainedPrefixLogitsProcessor

Restrict generation to a closed set of strings, one step at a time.

Module Contents

class micm_nlp.training.logits_processors.ConstrainedPrefixLogitsProcessor(tokenizer, allowed_texts)

Bases: transformers.LogitsProcessor

Restrict generation to a closed set of strings, one step at a time.

Each allowed string is tokenized once, up front, prefixed with the pad token that opens a decoder sequence. At every step the processor keeps the candidates whose prefix still matches what has been generated, and passes through the scores of only those tokens any survivor could produce next; everything else is masked to -inf. A candidate that has been generated in full contributes EOS, and if nothing matches at all, EOS is forced.

The decoder prefix length is inferred on the first call rather than configured, so this works whatever the model prepends.

Parameters:

• tokenizer – tokenizer used to encode allowed_texts.

• allowed_texts – the strings generation is restricted to.

allowed_token_seqs

tokenizer