Force the model to write some tokens mid-generation?
Author: ccocks
Pull request: no
Labels: name: Feature request

### Feature request

Here’s an example:

User: Hello make a python function for something
Assistant: Here’s an function for that:

def function():
pass
<codetests> ← This is a line we tuned the model to generate
import pytest
assert foo == bar
</codetests> ← Execute the tests right after this token was predicted
Result: tests **succeeded** ← THIS is the forced tokens, we also tuned the model to generate this
Ok, looks like the function is working…

EDIT:

The LLM is trained to respond with the same block given above, however since LLMs are bad at detecting when they have done a mistake they will lean towards saying succeeded for everything.
However after the inference pass for the token “succeeded” there will be a probablity distribution e.g.

succeeded 0.5
failed 0.3
etc.

So I want to “force” the model to pick failed (or succeeded) even though it is a less likely token. Seems like something very simple, but there is no support.

### Motivation

In case you didn't realize the point already, doing this could be opensource LLMs becoming significantly better for agentic workflows. unlike stopping generation, calling tools, and otherwise creating delays, this works right between inference passes. **Agentic workflows for proprietary LLMs can add up costs FAST.**

### Your contribution

I'm not used to this codebase. It seems very complex. but the feature is very simple. Maybe is someone could give me pointers

Comments:
Asked ChatGPT about it:

Below are some practical “entry points” in 🤗 Transformers that let you overwrite the next-token choice, even if the token you want (e.g. "failed") has lower probability than the model’s top guess.
Everything can be done without changing any C++/CUDA kernels – just a few Python classes and a call to model.generate().

Understand where the decision is made
model.generate() ──► _generate() ├── _greedy_search() / _sample() ├── logits = model(...) ├── logits = logits_processor(logits) ← place to hack ├── logits = logits_warper(logits) └── next_token = sample(argmax) ← after hack └── append next_token, loop …

Key locations in the source tree (v4.39.*):

src/transformers/generation/utils.py – generation loops
src/transformers/generation/logits_process.py – all LogitsProcessor / Warper classes
src/transformers/generation/configuration_utils.py – GenerationConfig helpers

Three out-of-the-box hooks you can already use
A. forced_decoder_ids / ForcedTokensLogitsProcessor

Works when you know in advance that at position N the token must be X.
Example:
gen_cfg = model.generation_config
gen_cfg.forced_decoder_ids = [(offset_from_prompt, tokenizer('failed').input_ids[0])]
model.generate(..., generation_config=gen_cfg)
B. prefix_allowed_tokens_fn

Callable invoked every step; you receive (batch_id, input_ids) and return the set of IDs that are legal next tokens.
If you return only [failed_id] at the step after your asserts, the model is forced to output “failed”.
def constraint(batch_id, input_ids):
    text = tokenizer.decode(input_ids[0])
    if text.endswith("Result: "):
        return [failed_id]          # only one legal token
    return list(range(tokenizer.vocab_size))
model.generate(..., prefix_allowed_tokens_fn=constraint)

C. stop-generate-then-append

Let generation stop after the code block, run pytest, then feed the original prompt + "Result: failed\n" back into generate() with past_key_values to keep it fast.
Costs one extra forward pass but zero code inside Transformers.
If you need total control: write a custom LogitsProcessor
A LogitsProcessor is only ~15 lines of code; it receives the raw logits and must return modified logits of the same shape.

Example: force a single token at a particular time step

from transformers import LogitsProcessor

class ForceTokenAtPosition(LogitsProcessor):
    def __init__(self, pos: int, token_id: int):
        self.pos = pos          # 0-based index *after* prompt
        self.token_id = token_id

    def __call__(self, input_ids, scores):
        cur_len = input_ids.shape[-1]            # length incl. prompt
        if cur_len == self.pos:
            mask = scores.new_full(scores.size(), float("-inf"))
            mask[:, self.token_id] = 0.0         # only token_id allowed
            return mask
        return scores

Usage:

processor = ForceTokenAtPosition(
        pos=prompt_len + code_block_len + 1,     # where “failed” must appear
        token_id=tokenizer.convert_tokens_to_ids("failed")
)
outputs = model.generate(
        input_ids,
        logits_processor=[processor],
        do_sample=False   # or True – doesn’t matter, there’s only one choice
)

Dynamic forcing (decide after pytest)

class DynamicForceNext(LogitsProcessor):
    def __init__(self):
        self.force_next: int | None = None

    def set(self, token_id: int | None):
        self.force_next = token_id

    def __call__(self, input_ids, scores):
        if self.force_next is not None:
            mask = scores.new_full(scores.size(), float("-inf"))
            mask[:, self.force_next] = 0
            self.force_next = None          # reset for later steps
            return mask
        return scores

In your outer loop run dyn_proc.set(failed_id) immediately after executing the tests.

Where to drop the code in the repo
• src/transformers/generation/logits_process.py
– add the new class(es) next to MinLengthLogitsProcessor.
• transformers/generation/init.py
– export the symbol so users can import.
• Optionally extend GenerationConfig (add a flag and construct the processor inside GenerationMixin).

Then unit-test:

tests/generation/test_force_token.py

def test_force_token():
    proc = ForceTokenAtPosition(3, 42)
    scores = torch.randn(1, 50)
    ids    = torch.tensor([[1,2,3]])
    out    = proc(ids, scores.clone())
    assert out[0,42] == 0 and torch.isinf(out).sum() == 49

Performance notes
• A LogitsProcessor runs on the same device as the logits (CUDA tensor) – no CPU-GPU sync if you keep your math inside the tensor world.
• For agentic workflows you usually step-through generation anyway, so the cost is negligible compared to another call to the model or, worse, spawning a sandboxed pytest.

TL;DR
You don’t need a new feature in Transformers – the library already has three hooks (forced_decoder_ids, prefix_allowed_tokens_fn, custom LogitsProcessor) that let you deterministically pick “failed” (or any token) even when its probability is low.
Start by prototyping a tiny LogitsProcessor in your own script; once it works, upstream it to src/transformers/generation/logits_process.py so everyone can use it.

Doesn't make much sense to me, but someone whos familiar may know what to do

It seems like you're asking for structured generation here - cc @gante what do we recommend for people these days? Just using Outlines?

Oh, so we don't need a custom Logits Processor?