Transformers generate 的停止条件与批量输出

在 Hugging Face Transformers 中，generate() 的“什么时候停”和“批量结果怎么拆”是两个经常一起出现的问题。单条提示词时，只看最终文本往往感觉不到差别；一旦同时输入多条提示词，又为每条输入返回多个候选，就会遇到输出行数翻倍、短序列补齐、提示词被重复解码等现象。

本文以 decoder-only 模型为例，集中说明 max_new_tokens、EOS、stop_strings、StoppingCriteria 以及 num_return_sequences 的组合方式。官方文档：https://huggingface.co/docs/transformers/main_classes/text_generation

一、先建立可批处理的最小配方

批量生成前，先解决两个基础问题：没有原生 pad token 的模型需要指定补齐 token；decoder-only 模型建议使用左侧补齐，避免右侧 padding 干扰“最后一个有效 token”的位置。

from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "openai-community/gpt2"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id)

# GPT-2 没有独立的 pad token，批处理时可复用 EOS
if tokenizer.pad_token_id is None:
    tokenizer.pad_token = tokenizer.eos_token

# decoder-only 模型批量生成通常使用左侧补齐
tokenizer.padding_side = "left"

prompts = [
    "Python 中生成器的作用是",
    "Transformer 的注意力机制用于",
]
inputs = tokenizer(prompts, return_tensors="pt", padding=True)

outputs = model.generate(
    **inputs,
    max_new_tokens=48,  # 只限制新增 token 数，不把提示词长度算进去
    do_sample=False,
    eos_token_id=tokenizer.eos_token_id,
    pad_token_id=tokenizer.pad_token_id,
)

# decoder-only 返回值包含输入提示词，按输入张量宽度裁掉前缀
new_tokens = outputs[:, inputs["input_ids"].shape[1]:]
texts = tokenizer.batch_decode(new_tokens, skip_special_tokens=True)

for prompt, text in zip(prompts, texts):
    print(prompt, "=>", text)

max_new_tokens 只约束新增部分，适合提示词长短不一致的批量输入。旧参数 max_length 在 decoder-only 模型中会把提示词长度也算进去，导致长提示词可生成空间更少，因此新代码优先使用 max_new_tokens。

二、四类停止条件如何配合

generate() 可以同时接收多个停止信号。它们不是“只选一个”，而是共同构成安全边界：某一行遇到 EOS、命中停止字符串或自定义条件后可以提前结束；如果一直没有命中，则由新增 token 上限兜底。

图1：generate 的停止信号与批量张量关系，任一停止条件命中后，该行进入结束状态并在批量张量中按需补齐。

• EOS token：模型生成 eos_token_id 后结束当前序列。部分模型可能配置多个 EOS ID。
• max_new_tokens：新增 token 数的硬上限，是防止异常长输出的首选兜底参数。
• stop_strings：生成文本出现指定字符串时停止，适合协议结束标记、分隔符或模板尾标。
• StoppingCriteria：自定义逐行判断，可表达 token 组合、业务状态或其他复杂条件。

max_time 也可限制生成耗时，但它会在当前模型计算过程完成后再返回，不应替代 token 数量上限。生产接口通常同时保留 max_new_tokens 作为确定性保护。

使用 stop_strings

停止字符串可以传单个字符串或字符串列表。调用时需要一并传入 tokenizer，让生成器能够处理字符串与 token 序列之间的对应关系。

outputs = model.generate(
    **inputs,
    max_new_tokens=96,  # 即使结束标记未出现，也会在此处兜底
    stop_strings=["\nEND", "\n用户："],
    tokenizer=tokenizer,  # stop_strings 必须配合 tokenizer
    pad_token_id=tokenizer.pad_token_id,
)

# 根据业务协议决定是否从最终文本中移除结束标记
new_tokens = outputs[:, inputs["input_ids"].shape[1]:]
texts = tokenizer.batch_decode(new_tokens, skip_special_tokens=True)

停止字符串用于判定终止位置，并不等价于“自动清理业务标记”。如果接口返回值不应包含 END 一类尾标，仍需在解码后做明确、可测试的裁剪。

三、为什么批量结果长度看起来不一致

张量要求每一行拥有相同宽度，但不同提示词可能在不同时间结束。先结束的行会使用 pad_token_id 补齐，直到同一批次完成。解码时设置 skip_special_tokens=True，通常可以去掉 padding 与 EOS 等特殊 token。

另一个容易混淆的点是：decoder-only 模型返回的序列通常包含“输入提示词 + 新生成内容”。如果直接对 outputs 做 batch_decode()，提示词会被再次打印。按 inputs["input_ids"].shape[1] 切片，才能取得统一的新增部分。

这里使用的是补齐后的输入宽度，而不是逐条字符串的字符数，也不是每一行非 padding token 的数量。这样才能与输出张量的列位置保持一致。

四、num_return_sequences 如何改变输出行数

当 num_return_sequences=2 时，每条输入生成两个候选。若批次中有两条提示词，最终输出张量就有四行。返回值不会自动嵌套成二维 Python 列表，而是按输入顺序展平。

图2：批量多候选结果按输入顺序展平，每条提示词连续占用 num_return_sequences 行。

n = 2
outputs = model.generate(
    **inputs,
    max_new_tokens=48,
    do_sample=True,       # 多候选通常配合采样，避免候选完全相同
    top_p=0.9,
    temperature=0.8,
    num_return_sequences=n,
    pad_token_id=tokenizer.pad_token_id,
)

# 输出行数为 len(prompts) * n，提示词前缀宽度仍与批量输入一致
new_tokens = outputs[:, inputs["input_ids"].shape[1]:]
flat_texts = tokenizer.batch_decode(new_tokens, skip_special_tokens=True)

# 每条输入连续占用 n 行，将展平结果重新分组
grouped_texts = [
    flat_texts[i * n:(i + 1) * n]
    for i in range(len(prompts))
]

for prompt, candidates in zip(prompts, grouped_texts):
    print(f"输入：{prompt}")
    for index, candidate in enumerate(candidates, start=1):
        print(f"候选 {index}：{candidate}")

如果使用 beam search，num_return_sequences 不能大于 num_beams。如果使用采样，多候选结果依赖随机性；需要复现实验时应固定随机种子，但线上服务通常还要考虑并发环境下的随机数管理。

五、需要更多元数据时返回 ModelOutput

默认情况下，generate() 直接返回生成序列张量。若还要查看每一步分数、beam 信息或其他生成元数据，可启用结构化返回。

result = model.generate(
    **inputs,
    max_new_tokens=32,
    return_dict_in_generate=True,  # 返回 ModelOutput，而不是单个张量
    output_scores=True,            # 同时保留每一步的分数
    pad_token_id=tokenizer.pad_token_id,
)

sequences = result.sequences
new_tokens = sequences[:, inputs["input_ids"].shape[1]:]
texts = tokenizer.batch_decode(new_tokens, skip_special_tokens=True)

自定义停止条件如果依赖生成分数，也需要开启 return_dict_in_generate=True 与 output_scores=True。仅按 token ID 判断时则不必额外保留分数。

六、自定义逐行停止条件

内置参数无法表达业务规则时，可以继承 StoppingCriteria。下面示例在最后一个 token 等于指定 ID 时返回逐行布尔结果，形状为 (batch_size, 1)。

from transformers import StoppingCriteria, StoppingCriteriaList

class StopOnToken(StoppingCriteria):
    def __init__(self, token_id: int):
        self.token_id = token_id

    def __call__(self, input_ids, scores, **kwargs):
        # 每一行独立判断，返回形状必须是 (batch_size, 1)
        return input_ids[:, -1:].eq(self.token_id)

stop_token_id = tokenizer.convert_tokens_to_ids(".")
criteria = StoppingCriteriaList([StopOnToken(stop_token_id)])

outputs = model.generate(
    **inputs,
    max_new_tokens=64,  # 自定义条件之外仍保留长度兜底
    stopping_criteria=criteria,
    pad_token_id=tokenizer.pad_token_id,
)

实际项目中不要直接假设某个字符串只对应一个 token。复杂结束词可能被 tokenizer 拆成多个 token，应先编码结束词并按 token 序列匹配，或者优先使用 stop_strings。

七、常见问题排查清单

• 长提示词几乎没有输出：检查是否还在使用包含提示词长度的 max_length，改用 max_new_tokens。
• 批量结果重复了输入：decoder-only 输出包含提示词，先按输入张量宽度裁剪再解码。
• 批量生成警告或行为异常：确认已配置 pad_token_id，并为 decoder-only 模型使用左侧补齐。
• 停止字符串没有生效：确认 stop_strings 调用同时传入了 tokenizer。
• 两条输入得到四条结果：检查 num_return_sequences，总行数等于批次大小乘以候选数。
• 候选归属错乱：不要用 zip 直接配对展平结果，按每条输入连续占用的候选数切片。
• 提前结束后仍看到相同宽度：这是批量张量 padding 的结果，不代表模型为每一行生成了相同数量的有效 token。

八、推荐的生产组合

大多数批量文本生成接口可以采用一组稳定默认值：用 max_new_tokens 提供长度兜底，保留模型的 EOS 配置；确有协议尾标时增加 stop_strings 并传入 tokenizer；只有内置条件无法表达业务规则时才实现 StoppingCriteria。输出阶段统一裁掉 decoder-only 提示词前缀，再按 num_return_sequences 分组。

把“停止条件”和“结果整理”拆成两个独立步骤后，批量生成会更容易测试：前者验证每行何时结束，后者验证输出行数、候选归属与纯新增文本。这样即使模型、采样参数或批次大小变化，接口返回结构也能保持稳定。