Stopping Rules¶

StoppingRules is a generic output control that exposes stopping criteria as constructor arguments. Without it, a substring, token, or budget stop requires writing a criteria class. With StoppingRules, each stop is specified as a configuration.

StoppingRules is a step-level control rather than a decoding driver. get_stopping_criteria returns fresh criteria anchored to each generation's prompt, meaning two generations with different prompt lengths each stop relative to their own prompt. Since it contributes no logits processors, it composes with a logits processor (such as ValueGuidance or ContrastiveGuidance) and with a decoding driver in the same pipeline.

This notebook runs each stop against one instruction model and shows its effect as a contrast, placing the halted generation beside the un-halted baseline along with the generated token counts.

Method parameters¶

parameter	type	description
stop_texts	list[str]	Substrings that halt a row when they appear in its continuation
stop_token_ids	list[int]	Token ids that halt a row when generated
budget	int \| None	Maximum new tokens before a row halts

At least one of the three must be set. A substring stop decodes each row's continuation every step, which is the cost of a text-level stop. A token-id or budget stop is an integer comparison.

Setup¶

If running this from a Google Colab notebook, uncomment the clone cell below. It is not necessary when running from a virtual environment where the package is already installed.

In [1]:

# !git clone https://github.com/generative-computing/steerability.git
# %cd steerability

# !git clone https://github.com/generative-computing/steerability.git # %cd steerability

In [2]:

import sys
!{sys.executable} -m pip install -q tabulate

import sys !{sys.executable} -m pip install -q tabulate

In [3]:

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

from steerability.algorithms.core.steering_pipeline import SteeringPipeline
from steerability.algorithms.output_control.stopping_rules.control import StoppingRules
from steerability.algorithms.output_control.value_guidance.control import ValueGuidance

from IPython.display import display, HTML
display(HTML("<style>:root { --jp-notebook-max-width: 100% !important; }</style>"))

from tabulate import tabulate
import textwrap

def wrap(text, width=60):
    return '\n'.join(textwrap.wrap(text, width=width))

import torch from transformers import AutoModelForCausalLM, AutoTokenizer from steerability.algorithms.core.steering_pipeline import SteeringPipeline from steerability.algorithms.output_control.stopping_rules.control import StoppingRules from steerability.algorithms.output_control.value_guidance.control import ValueGuidance from IPython.display import display, HTML display(HTML("<style>:root { --jp-notebook-max-width: 100% !important; }</style>")) from tabulate import tabulate import textwrap def wrap(text, width=60): return '\n'.join(textwrap.wrap(text, width=width))

We use Qwen/Qwen2.5-1.5B-Instruct throughout and load it once. Each stop below builds a fresh SteeringPipeline over this shared model, passing the model and tokenizer at construction. Since a pipeline's steer() is one-shot, each configuration gets its own pipeline object.

In [4]:

MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct"

model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, device_map="auto", dtype=torch.bfloat16)
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
device = model.device

gen_params = {
    "max_new_tokens": 60,
    "do_sample": False,
    "repetition_penalty": 1.1,
    "pad_token_id": tokenizer.eos_token_id,
}

MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct" model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, device_map="auto", dtype=torch.bfloat16) tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME) device = model.device gen_params = { "max_new_tokens": 60, "do_sample": False, "repetition_penalty": 1.1, "pad_token_id": tokenizer.eos_token_id, }

Stop on a substring¶

A substring stop halts a row as soon as its continuation contains the given text. We ask the model to list items one per line and stop at the first blank line ("\n\n"), cutting the generation to a single block. The contrast below runs the same prompt with and without the stop, and reports the generated token count for each.

The token count comes from return_output=True, which returns an Output whose output_ids holds the generated tokens (the prompt excluded). output_ids.size(1) is therefore the number of new tokens. The finish_reason on that Output reports "stop" for a substring or token stop (the stop rules are part of the generation parameters and the pipeline classifies them directly) and "length" when the token budget is exhausted. Decoded text is truncated at the first stop-string occurrence, while output_ids keeps the tokens as generated.

In [5]:

substring_prompt = "List a few programming languages, then explain in a paragraph why one of them is popular."

baseline_pipeline = SteeringPipeline(controls=[], model=model, tokenizer=tokenizer)
baseline_pipeline.steer()

stopped_pipeline = SteeringPipeline(controls=[StoppingRules(stop_texts=["\n\n"])], model=model, tokenizer=tokenizer)
stopped_pipeline.steer()

messages = [[{"role": "user", "content": substring_prompt}]]
baseline_out = baseline_pipeline.generate(messages=messages, return_output=True, **gen_params)[0]
stopped_out = stopped_pipeline.generate(messages=messages, return_output=True, **gen_params)[0]

table = [
    ["no stop", baseline_out.output_ids.size(1), wrap(baseline_out.decode(tokenizer)[0], 70)],
    ["stop at \\n\\n", stopped_out.output_ids.size(1), wrap(stopped_out.decode(tokenizer)[0], 70)],
]
print(f"Prompt: {substring_prompt}")
print(tabulate(table, headers=["config", "new tokens", "completion"], tablefmt="grid", maxcolwidths=[14, 10, 70]))

substring_prompt = "List a few programming languages, then explain in a paragraph why one of them is popular." baseline_pipeline = SteeringPipeline(controls=[], model=model, tokenizer=tokenizer) baseline_pipeline.steer() stopped_pipeline = SteeringPipeline(controls=[StoppingRules(stop_texts=["\n\n"])], model=model, tokenizer=tokenizer) stopped_pipeline.steer() messages = [[{"role": "user", "content": substring_prompt}]] baseline_out = baseline_pipeline.generate(messages=messages, return_output=True, **gen_params)[0] stopped_out = stopped_pipeline.generate(messages=messages, return_output=True, **gen_params)[0] table = [ ["no stop", baseline_out.output_ids.size(1), wrap(baseline_out.decode(tokenizer)[0], 70)], ["stop at \\n\\n", stopped_out.output_ids.size(1), wrap(stopped_out.decode(tokenizer)[0], 70)], ] print(f"Prompt: {substring_prompt}") print(tabulate(table, headers=["config", "new tokens", "completion"], tablefmt="grid", maxcolwidths=[14, 10, 70]))

Prompt: List a few programming languages, then explain in a paragraph why one of them is popular.
+--------------+--------------+------------------------------------------------------------------------+
| config       |   new tokens | completion                                                             |
+==============+==============+========================================================================+
| no stop      |           60 | Sure! Here's a list of some popular programming languages:  1. Python: |
|              |              | Known for its simplicity and readability, Python is widely used for    |
|              |              | web development, data analysis, artificial intelligence, and           |
|              |              | scientific computing.  2. JavaScript: Essential for front-end web      |
|              |              | development, JavaScript powers interactive elements on websites like   |
|              |              | buttons                                                                |
+--------------+--------------+------------------------------------------------------------------------+
| stop at \n\n |           12 | Sure! Here's a list of some popular programming languages:             |
+--------------+--------------+------------------------------------------------------------------------+

Stop on a token id or a budget¶

A token-id stop halts on a specific token, and a budget stop halts after a fixed number of new tokens. Both are configuration rather than code. Below, the token-id stop ends the generation at the first period (the "." token), cutting the output to one sentence, and the budget stop caps the generation at sixteen new tokens against an un-capped baseline. The token counts are tabulated for each configuration.

In [6]:

period_id = tokenizer.encode(".")[-1]
budget_prompt = "Describe a walk on the beach at sunset."

token_pipeline = SteeringPipeline(
    controls=[StoppingRules(stop_token_ids=[period_id])], model=model, tokenizer=tokenizer,
)
token_pipeline.steer()

budget_pipeline = SteeringPipeline(controls=[StoppingRules(budget=16)], model=model, tokenizer=tokenizer)
budget_pipeline.steer()

messages = [[{"role": "user", "content": budget_prompt}]]
uncapped = baseline_pipeline.generate(messages=messages, return_output=True, **gen_params)[0]
token_stopped = token_pipeline.generate(messages=messages, return_output=True, **gen_params)[0]
budget_stopped = budget_pipeline.generate(messages=messages, return_output=True, **gen_params)[0]

table = [
    ["no stop", uncapped.output_ids.size(1), wrap(uncapped.decode(tokenizer)[0], 68)],
    [f"stop on '.' (id {period_id})", token_stopped.output_ids.size(1), wrap(token_stopped.decode(tokenizer)[0], 68)],
    ["budget = 16", budget_stopped.output_ids.size(1), wrap(budget_stopped.decode(tokenizer)[0], 68)],
]
print(f"Prompt: {budget_prompt}")
print(tabulate(table, headers=["config", "new tokens", "completion"], tablefmt="grid", maxcolwidths=[24, 10, 68]))

period_id = tokenizer.encode(".")[-1] budget_prompt = "Describe a walk on the beach at sunset." token_pipeline = SteeringPipeline( controls=[StoppingRules(stop_token_ids=[period_id])], model=model, tokenizer=tokenizer, ) token_pipeline.steer() budget_pipeline = SteeringPipeline(controls=[StoppingRules(budget=16)], model=model, tokenizer=tokenizer) budget_pipeline.steer() messages = [[{"role": "user", "content": budget_prompt}]] uncapped = baseline_pipeline.generate(messages=messages, return_output=True, **gen_params)[0] token_stopped = token_pipeline.generate(messages=messages, return_output=True, **gen_params)[0] budget_stopped = budget_pipeline.generate(messages=messages, return_output=True, **gen_params)[0] table = [ ["no stop", uncapped.output_ids.size(1), wrap(uncapped.decode(tokenizer)[0], 68)], [f"stop on '.' (id {period_id})", token_stopped.output_ids.size(1), wrap(token_stopped.decode(tokenizer)[0], 68)], ["budget = 16", budget_stopped.output_ids.size(1), wrap(budget_stopped.decode(tokenizer)[0], 68)], ] print(f"Prompt: {budget_prompt}") print(tabulate(table, headers=["config", "new tokens", "completion"], tablefmt="grid", maxcolwidths=[24, 10, 68]))

Prompt: Describe a walk on the beach at sunset.
+---------------------+--------------+----------------------------------------------------------------------+
| config              |   new tokens | completion                                                           |
+=====================+==============+======================================================================+
| no stop             |           60 | Walking on the beach at sunset is a serene and beautiful experience  |
|                     |              | that can be both calming and exhilarating. The golden hour of the    |
|                     |              | day when the sun begins to set creates an enchanting atmosphere with |
|                     |              | its warm hues of orange, pink, and purple lighting up the sky.  As   |
|                     |              | you stroll along the sandy                                           |
+---------------------+--------------+----------------------------------------------------------------------+
| stop on '.' (id 13) |           21 | Walking on the beach at sunset is a serene and beautiful experience  |
|                     |              | that can be both calming and exhilarating.                           |
+---------------------+--------------+----------------------------------------------------------------------+
| budget = 16         |           16 | Walking on the beach at sunset is a serene and beautiful experience  |
|                     |              | that can be both                                                     |
+---------------------+--------------+----------------------------------------------------------------------+

Per-generation anchoring¶

get_stopping_criteria builds fresh criteria for each generation, anchored at that call's prompt length. A substring stop therefore measures the continuation from the end of that call's prompt, and two generations whose prompts have different lengths each stop relative to their own prompt. We show this with two sequential single-prompt calls, a short prompt and a long one, under the same substring stop. Each halts at its own first blank line and each reports its own continuation and token count.

We run the two prompts as separate calls rather than as one batch, since the substring criterion anchors on the tokenized batch's common length. That length is exact only when it is a single prompt's true length (batch size one) or when the batch is left-padded and every real prompt ends at the common length. A right-padded multi-prompt batch would misalign the anchor. Per-generation anchoring is therefore demonstrated one prompt at a time.

Note that adapted_input_ids on the returned Output is the prompt as fed to the model, and its length is the anchor position.

In [7]:

short_prompt = "Name three fruits, one per line."
long_prompt = (
    "You are compiling a short reference sheet for a cooking class. "
    "Name three fruits that are common in desserts, one per line."
)

anchor_pipeline = SteeringPipeline(controls=[StoppingRules(stop_texts=["\n\n"])], model=model, tokenizer=tokenizer)
anchor_pipeline.steer()

rows = []
for label, prompt in [("short prompt", short_prompt), ("long prompt", long_prompt)]:
    out = anchor_pipeline.generate(messages=[{"role": "user", "content": prompt}], return_output=True, **gen_params)
    prompt_len = out.adapted_input_ids.size(1)
    rows.append([label, prompt_len, out.output_ids.size(1), wrap(out.decode(tokenizer)[0], 60)])

print(tabulate(rows, headers=["call", "prompt tokens", "new tokens", "continuation"], tablefmt="grid", maxcolwidths=[14, 14, 10, 60]))

short_prompt = "Name three fruits, one per line." long_prompt = ( "You are compiling a short reference sheet for a cooking class. " "Name three fruits that are common in desserts, one per line." ) anchor_pipeline = SteeringPipeline(controls=[StoppingRules(stop_texts=["\n\n"])], model=model, tokenizer=tokenizer) anchor_pipeline.steer() rows = [] for label, prompt in [("short prompt", short_prompt), ("long prompt", long_prompt)]: out = anchor_pipeline.generate(messages=[{"role": "user", "content": prompt}], return_output=True, **gen_params) prompt_len = out.adapted_input_ids.size(1) rows.append([label, prompt_len, out.output_ids.size(1), wrap(out.decode(tokenizer)[0], 60)]) print(tabulate(rows, headers=["call", "prompt tokens", "new tokens", "continuation"], tablefmt="grid", maxcolwidths=[14, 14, 10, 60]))

+--------------+-----------------+--------------+-------------------------------------------------------------+
| call         |   prompt tokens |   new tokens | continuation                                                |
+==============+=================+==============+=============================================================+
| short prompt |              37 |            7 | Apple Banana Orange                                         |
+--------------+-----------------+--------------+-------------------------------------------------------------+
| long prompt  |              54 |           43 | 1. Bananas - Often used in banana bread and smoothies. 2.   |
|              |                 |              | Strawberries - Popular in strawberry shortcake and pies. 3. |
|              |                 |              | Apples - Common in apple pie and other autumn-themed        |
|              |                 |              | desserts.                                                   |
+--------------+-----------------+--------------+-------------------------------------------------------------+

Composition with a logits processor¶

Since StoppingRules contributes only stopping criteria, it composes independently of a logits processor in the same controls list. Here we pair a ValueGuidance sentiment control (which shifts the distribution toward positive continuations) with a budget stop, and both effects show in one output, i.e., the text is steered positive and the generation is cut at thirty-two new tokens. The ValueGuidance config here is the FUDGE-style sentiment control from the value-guidance notebook.

In [8]:

SENTIMENT = "distilbert-base-uncased-finetuned-sst-2-english"
compose_prompt = "Write a few sentences about your first day at a new job."

sentiment_value = ValueGuidance(
    value={"kind": "classifier", "model_id": SENTIMENT, "label_index": 1},
    policy="top_k", k=50, beta=4.0,
)

composed_pipeline = SteeringPipeline(
    controls=[sentiment_value, StoppingRules(budget=32)], model=model, tokenizer=tokenizer,
)
composed_pipeline.steer()

messages = [[{"role": "user", "content": compose_prompt}]]
plain = baseline_pipeline.generate(messages=messages, return_output=True, **gen_params)[0]
composed_out = composed_pipeline.generate(messages=messages, return_output=True, **gen_params)[0]

table = [
    ["no control", plain.output_ids.size(1), wrap(plain.decode(tokenizer)[0], 68)],
    ["sentiment + budget=32", composed_out.output_ids.size(1), wrap(composed_out.decode(tokenizer)[0], 68)],
]
print(f"Prompt: {compose_prompt}")
print(tabulate(table, headers=["config", "new tokens", "completion"], tablefmt="grid", maxcolwidths=[22, 10, 68]))

SENTIMENT = "distilbert-base-uncased-finetuned-sst-2-english" compose_prompt = "Write a few sentences about your first day at a new job." sentiment_value = ValueGuidance( value={"kind": "classifier", "model_id": SENTIMENT, "label_index": 1}, policy="top_k", k=50, beta=4.0, ) composed_pipeline = SteeringPipeline( controls=[sentiment_value, StoppingRules(budget=32)], model=model, tokenizer=tokenizer, ) composed_pipeline.steer() messages = [[{"role": "user", "content": compose_prompt}]] plain = baseline_pipeline.generate(messages=messages, return_output=True, **gen_params)[0] composed_out = composed_pipeline.generate(messages=messages, return_output=True, **gen_params)[0] table = [ ["no control", plain.output_ids.size(1), wrap(plain.decode(tokenizer)[0], 68)], ["sentiment + budget=32", composed_out.output_ids.size(1), wrap(composed_out.decode(tokenizer)[0], 68)], ] print(f"Prompt: {compose_prompt}") print(tabulate(table, headers=["config", "new tokens", "completion"], tablefmt="grid", maxcolwidths=[22, 10, 68]))

Prompt: Write a few sentences about your first day at a new job.
+-----------------------+--------------+----------------------------------------------------------------------+
| config                |   new tokens | completion                                                           |
+=======================+==============+======================================================================+
| no control            |           60 | As an AI language model, I don't have personal experiences or        |
|                       |              | emotions like humans do. However, I can tell you that my "first day" |
|                       |              | would be when I was installed and integrated into the system to      |
|                       |              | assist with tasks such as answering questions, generating text, and  |
|                       |              | providing information on various topics.                             |
+-----------------------+--------------+----------------------------------------------------------------------+
| sentiment + budget=32 |           32 | As an AI language model, I don't have personal experiences or        |
|                       |              | emotions like humans do. However, I can provide you with some        |
|                       |              | general information on what it might                                 |
+-----------------------+--------------+----------------------------------------------------------------------+

Semantics¶

Criteria are not applied during compute_logprobs (there is no generation loop to stop). Under a segment or phase driver, the composed criteria apply inside every rollout or phase with the prompt-anchored lengths fixed at composition time. The stop is therefore global and anchored to the original prompt. StopOnSubstring decodes the continuation each step, the cost of a text-level stop.

The phased-decoding notebook shows this global behavior directly. Its phases-and-stops section composes a StoppingRules alongside a PhasedDecoding driver, and the stop fires inside a generated phase, anchored to the original prompt.

Summary¶

Each stop in this notebook was a configuration on StoppingRules, run against one instruction model and shown as a contrast in token counts and text. A substring, token-id, or budget stop halts generation without a criteria class, and each stop's effect is visible in the generated token count. The criteria are rebuilt per generation and anchored at that call's prompt length, meaning different-length prompts each stop relative to their own prompt. Because StoppingRules contributes only criteria, it composes with a logits processor such as the sentiment value here in one pipeline, and each mechanism composes independently.