Hugging Face Transformers generate 怎样用 max_new_tokens 控制输出长度

我在批量生成摘要时最容易踩到的坑，是把 max_length 当成“最多生成多少 token”。它统计的是整条序列：输入提示词和新输出加在一起；max_new_tokens 才是只限制新生成 token 数的参数。只要目标是“无论提示词多长，都给模型最多生成 N 个 token”，就优先设置 max_new_tokens，并避免同时用两个长度上限制造歧义。

官方文档：https://huggingface.co/docs/transformers/main_classes/text_generation

先分清 max_new_tokens 和 max_length 的计数对象

这两个参数都能让生成停止，但解决的不是同一个问题。假设输入已经占用 120 个 token，希望最多再生成 64 个 token：使用 max_new_tokens=64 时，目标总长度大约是 184；使用 max_length=64 时，输入本身已经超过上限，结果可能很短，甚至触发长度边界相关的警告。

把生成上限写进可复用的 generate 调用

实际项目里，我会把长度上限直接写在 generate 调用旁边，让代码审查时能看出资源预算。下面的例子只展示调用关系；模型名称需要换成你实际使用、且与 tokenizer 匹配的模型。

from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "your-model-id"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

inputs = tokenizer(
    ["请用三句话概括这段文本：", "请列出这段文本的两个风险点："],
    return_tensors="pt",
    padding=True,
)

# 这里只限制新生成 token，不让提示词长度改变输出预算。
outputs = model.generate(
    **inputs,
    max_new_tokens=64,
    do_sample=False,
    return_dict_in_generate=True,
)

# decoder-only 模型的序列包含输入；从输入长度后取出真正的新 token。
input_length = inputs["input_ids"].shape[1]
new_tokens = outputs.sequences[:, input_length:]
texts = tokenizer.batch_decode(new_tokens, skip_special_tokens=True)
for text in texts:
    print(text.strip())  # 输出每条样本的生成部分，避免重复打印提示词

如果希望多处调用共享策略，可以构造 GenerationConfig，但仍建议只选择一个主要长度控制参数。max_new_tokens 是上限而不是保证值：模型遇到结束 token、停止字符串或其他停止条件时，生成长度可能明显小于 64。

按模型结构读取真正的新文本

上面的切片适用于 decoder-only 模型，因为它的 sequences 通常是“输入序列 + 续写序列”。批量输入有 padding 时，input_length 取的是这一批张量的序列宽度，因此解码前要让 tokenizer 的 padding 设置与模型要求一致；不要把不同样本各自的原始字符长度当成 token 边界。

encoder-decoder 模型的输出序列通常是 decoder 侧生成结果，不应机械套用 sequences[:, input_length:]。遇到 T5、BART 这类模型，先按模型文档确认输出张量语义，再直接对生成序列调用 batch_decode。这一步能避免“摘要只剩半句”的隐蔽问题。

批量推理时用清单检查长度与资源边界

批量设置 max_new_tokens 后，我会按下面的顺序复查：

• 先算输入：确认最长提示词是否已经接近模型上下文窗口；新 token 上限不能突破模型可接受的总长度。
• 再看停止：输出短于上限通常是正常的，可能因为 eos token、停止条件或模型自然结束，不要为了“凑满 64”盲目放大参数。
• 最后看显存：批大小、输入长度和新 token 上限都会影响 KV cache；一味增大上限会让尾部样本等待更久。

一个实用的取舍是：在线接口用中等的 max_new_tokens 配合请求级超时，离线任务再按最长输入和显存预算提高它。max_length 只有在你明确需要固定总序列长度时才更合适；否则它很容易把“输出预算”与“上下文预算”混成一个数字。

相关问题

为什么设置了 max_new_tokens，输出还是比它短？

它表示最大值，不是固定长度。模型提前生成 eos token，或者命中停止条件时，会在达到上限前结束。

max_new_tokens 和 max_length 能同时设置吗？

技术上可以传入，但两个参数同时参与长度控制时容易出现覆盖或警告。工程上最好明确只保留一个主长度策略。

批量生成为什么不能直接解码整个 sequences？

对 decoder-only 模型，整个序列包含输入提示词。直接解码会把 prompt 一起返回；从输入 token 宽度之后切片，才是新生成内容。