Transformers 文本生成策略实战：贪心搜索、采样、束搜索与 custom_generate 自定义生成方法

Transformers 文本生成策略实战：贪心搜索、采样、束搜索与 custom_generate 自定义生成方法

【免费下载链接】transformers🤗 Transformers: the model-definition framework for state-of-the-art machine learning models in text, vision, audio, and multimodal models, for both inference and training.项目地址: https://gitcode.com/GitHub_Trending/tra/transformers

解码策略（decoding strategy）决定了模型如何选择下一个要生成的 token。Transformers 提供了从贪心搜索、多项式采样到束搜索等多种内置解码方法，并在此基础上提供了custom_generate机制，允许你把任意解码逻辑（如“模型不确定时继续思考”、“生成卡住时回滚”、处理特殊 token 的自定义逻辑、使用专用 KV cache 等）打包成一个 Hub 仓库或本地目录，注入到任何模型中。读完本文，你将掌握 Transformers 内置解码策略的选择与调用方式，并能够创建、测试和分发自己的自定义生成方法。

解码策略决定生成质量

解码策略直接影响生成文本的质量：对短输出、低创造性要求的任务，贪心搜索通常足够；对需要多样性的创作类任务，采样类方法更合适；对图像描述、语音识别等“以输入为基准”的任务，束搜索能取得更好的整体概率。官方文档 generation_strategies.md 将策略分为两大类：

• 基础解码方法（Basic decoding methods）：贪心搜索、采样（Sampling）、束搜索（Beam search），是所有文本生成任务的起点；
• 自定义生成方法（Custom generation methods）：通过custom_generate机制扩展的专用行为。

在源码层面，这些方法统一由GenerationMixin.generate入口调度。src/transformers/generation/utils.py 中的GENERATION_MODES_MAPPING维护了生成模式到具体实现函数的映射：

GENERATION_MODES_MAPPING = { GenerationMode.SAMPLE: "_sample", GenerationMode.GREEDY_SEARCH: "_sample", GenerationMode.BEAM_SEARCH: "_beam_search", GenerationMode.BEAM_SAMPLE: "_beam_search", GenerationMode.ASSISTED_GENERATION: "_assisted_decoding", # Deprecated methods GenerationMode.DOLA_GENERATION: "transformers-community/dola", GenerationMode.CONTRASTIVE_SEARCH: "transformers-community/contrastive-search", GenerationMode.GROUP_BEAM_SEARCH: "transformers-community/group-beam-search", GenerationMode.CONSTRAINED_BEAM_SEARCH: "transformers-community/constrained-beam-search", }

从中可以看到两点事实：贪心搜索和采样共用_sample实现（区别仅在do_sample参数）；而 DOLA、对比搜索（contrastive search）、组束搜索（group beam search）、约束束搜索（constrained beam search）等已不在核心实现中，而是被迁移到了custom_generate仓库（由 _get_deprecated_gen_repo 处理并提示用户显式传入custom_generate仓库名，v4.62.0 后将移除兼容逻辑）。这正是自定义生成机制的典型应用案例。

基础解码方法

贪心搜索（Greedy search）

贪心搜索是默认解码策略：每一步都选取概率最高的 token。除非在GenerationConfig中另行指定，该策略默认最多生成 20 个新 token（max_new_tokens默认值为 20）。

贪心搜索适合输出相对较短、不追求创造性的任务；但生成较长序列时容易开始重复自己。

import torch from transformers import AutoModelForCausalLM, AutoTokenizer from accelerate import Accelerator device = Accelerator().device tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf") inputs = tokenizer("Hugging Face is an open-source company", return_tensors="pt").to(device) model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf", dtype=torch.float16).to(device) # explicitly set to default length because Llama2 generation length is 4096 outputs = model.generate(**inputs, max_new_tokens=20) tokenizer.batch_decode(outputs, skip_special_tokens=True) 'Hugging Face is an open-source company that provides a suite of tools and services for building, deploying, and maintaining natural language processing'

采样（Sampling）

采样（又称多项式采样，multinomial sampling）不是选概率最高的 token，而是按照整个词表上的概率分布随机抽取一个 token——只要某个 token 概率非零，就有机会被选中。采样类策略能减少重复、产生更有创造力和多样性的输出。

启用方式：do_sample=True且num_beams=1。

import torch from transformers import AutoModelForCausalLM, AutoTokenizer from accelerate import Accelerator device = Accelerator().device tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf") inputs = tokenizer("Hugging Face is an open-source company", return_tensors="pt").to(device) model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf", dtype=torch.float16).to(device) # explicitly set to 100 because Llama2 generation length is 4096 outputs = model.generate(**inputs, max_new_tokens=50, do_sample=True, num_beams=1) tokenizer.batch_decode(outputs, skip_special_tokens=True) 'Hugging Face is an open-source company 🤗\nWe are open-source and believe that open-source is the best way to build technology. Our mission is to make AI accessible to everyone, and we believe that open-source is the best way to achieve that.'

束搜索（Beam search）

束搜索在每个时间步同时维护多条生成序列（beam），在若干步之后选取整体概率最高的序列。与贪心搜索不同，它具备“向前看”的能力：即使某条序列开头的 token 概率较低，只要整体序列概率更高，也可能被选中。它最适合以输入为基准（input-grounded）的任务，例如图像描述、语音识别。

也可以配合do_sample=True使用束搜索：每步内部进行采样，但束搜索仍会在步骤之间贪心地剪掉低概率序列。

启用方式：设置num_beams参数（必须大于 1，否则等价于贪心搜索）。

import torch from transformers import AutoModelForCausalLM, AutoTokenizer from accelerate import Accelerator device = Accelerator().device tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf") inputs = tokenizer("Hugging Face is an open-source company", return_tensors="pt").to(device) model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf", dtype=torch.float16).to(device) # explicitly set to 100 because Llama2 generation length is 4096 outputs = model.generate(**inputs, max_new_tokens=50, num_beams=2) tokenizer.batch_decode(outputs, skip_special_tokens=True) "['Hugging Face is an open-source company that develops and maintains the Hugging Face platform, which is a collection of tools and libraries for building and deploying natural language processing (NLP) models. Hugging Face was founded in 2018 by Thomas Wolf']"

三种基础策略的参数速查：

自定义生成方法（custom_generate）

当内置方法无法满足需求时——例如希望模型不确定时继续“思考”、生成卡住时回滚、用自定义逻辑处理特殊 token、或使用专用 KV cache——可以用custom_generate机制扩展生成行为。这是对 自定义模型代码 能力的进一步延伸，同样要求设置trust_remote_code=True。

该机制有两种使用形态：

形态一：加载自带自定义生成方法的模型仓库

如果某个模型仓库内置了自定义生成方法（仓库内含custom_generate/目录），加载它时generate会被自动覆盖。从源码看，from_pretrained加载流程中会尝试调用load_custom_generate，成功则用functools.partial替换self.generate（见 GenerationMixin.from_pretrained）：

# 加载自定义生成方法如果 `pretrained_model_name_or_path` 定义了它（并覆盖 `generate`） if hasattr(self, "load_custom_generate") and trust_remote_code: try: custom_generate = self.load_custom_generate( pretrained_model_name_or_path, trust_remote_code=trust_remote_code, **repo_loading_kwargs ) self.generate = functools.partial(custom_generate, model=self) except OSError: # 不存在自定义 generate 函数 pass

示例：transformers-community/custom_generate_example仓库是Qwen/Qwen2.5-0.5B-Instruct的一份副本，但附带了自定义生成代码——直接调用generate就会使用它：

from transformers import AutoModelForCausalLM, AutoTokenizer # `transformers-community/custom_generate_example` 是 `Qwen/Qwen2.5-0.5B-Instruct` 的副本， # 但带有自定义生成代码 -> 调用 `generate` 即使用自定义生成方法！ tokenizer = AutoTokenizer.from_pretrained("transformers-community/custom_generate_example") model = AutoModelForCausalLM.from_pretrained( "transformers-community/custom_generate_example", device_map="auto", trust_remote_code=True ) inputs = tokenizer(["The quick brown"], return_tensors="pt").to(model.device) # 自定义生成方法是一个最简贪心解码实现，运行时还会打印一条自定义消息 gen_out = model.generate(**inputs) # 此时应能看到它的自定义消息："✨ using a custom generation method ✨" print(tokenizer.batch_decode(gen_out, skip_special_tokens=True)) 'The quick brown fox jumps over a lazy dog, and the dog is a type of animal. Is'

形态二：任意模型通过custom_generate参数注入方法

自定义生成方法还有一个关键特性：它可以从任何模型加载。~GenerationMixin.generate提供了custom_generate参数，任何人都能创建并分享可作用于任意 Transformers 模型的自定义生成方法，用户无需安装额外的 Python 包：

from transformers import AutoModelForCausalLM, AutoTokenizer tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct") model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct", device_map="auto") inputs = tokenizer(["The quick brown"], return_tensors="pt").to(model.device) # `custom_generate` 用 `transformers-community/custom_generate_example` 中定义的 # 自定义生成方法替换原有的 `generate` gen_out = model.generate(**inputs, custom_generate="transformers-community/custom_generate_example", trust_remote_code=True) print(tokenizer.batch_decode(gen_out, skip_special_tokens=True)[0]) 'The quick brown fox jumps over a lazy dog, and the dog is a type of animal. Is'

从源码看，generate在处理任何输入准备之前（“0.a”步骤）就会拦截字符串形式的custom_generate：收集除self、kwargs、trust_remote_code、custom_generate之外的全部参数，交给load_custom_generate加载出的函数，并把model=self一并转发（见 generate 的 custom_generate 分支）：

if custom_generate is not None and isinstance(custom_generate, str): global_keys_to_exclude = {"self", "kwargs", "global_keys_to_exclude", "trust_remote_code", "custom_generate"} generate_arguments = {key: value for key, value in locals().items() if key not in global_keys_to_exclude} generate_arguments.update(kwargs) custom_generate_function = self.load_custom_generate( custom_generate, trust_remote_code=trust_remote_code, **kwargs ) return custom_generate_function(model=self, **generate_arguments)

也就是说，你的自定义generate收到的参数与原生generate完全一致，只是把self换成了model，并且可以访问GenerationMixin定义的所有属性和方法。

使用自定义方法前，应阅读该仓库的README.md，确认是否有新的输入参数或输出类型差异；如果没有，可认为其行为与基础generate一致。以transformers-community/custom_generate_example为例，其 README 声明了一个额外参数left_padding（在 prompt 前添加若干 pad token）：

gen_out = model.generate( **inputs, custom_generate="transformers-community/custom_generate_example", trust_remote_code=True, left_padding=5 ) print(tokenizer.batch_decode(gen_out)[0]) ' The quick brown fox jumps over the lazy dog.\n\nThe sentence "The quick'

依赖检查：requirements 缺失时的报错

如果自定义方法固定了当前环境不满足的 Python 依赖，load_custom_generate会先执行check_python_requirements校验custom_generate/requirements.txt（见 load_custom_generate 实现），不满足则抛出异常。例如transformers-community/custom_generate_bad_requirements仓库定义了不可能满足的依赖，运行会得到类似报错：

ImportError: Missing requirements in your local environment for `transformers-community/custom_generate_bad_requirements`: foo (installed: None) bar==0.0.0 (installed: None) torch>=99.0 (installed: 2.6.0)

按提示更新环境依赖即可消除该错误。相关行为在测试中有覆盖，见 tests/generation/test_utils.py 中的test_custom_generate_bad_requirements等用例。

创建自定义生成方法

创建一个新生成方法，需要建立一个新的模型仓库并推送以下文件：

1. 你设计生成方法所用的模型；
2. custom_generate/generate.py——自定义生成方法的全部逻辑；
3. custom_generate/requirements.txt——可选的额外 Python 依赖及版本锁定；
4. README.md——添加custom_generate标签，并记录新方法的所有新参数与输出类型差异。

仓库结构如下：

your_repo/ ├── README.md # include the 'custom_generate' tag ├── config.json ├── ... └── custom_generate/ ├── generate.py └── requirements.txt

添加基础模型：起点就是一个普通的模型仓库。应放入你设计该方法时使用的模型，它与生成方法构成一个可独立工作的“模型-生成”对。加载该仓库的模型时，你的自定义方法会覆盖generate；但方法本身仍可按上文方式加载到任意其他 Transformers 模型上。如果只是复制现有模型：

from transformers import AutoModelForCausalLM, AutoTokenizer tokenizer = AutoTokenizer.from_pretrained("source/model_repo") model = AutoModelForCausalLM.from_pretrained("source/model_repo") tokenizer.save_pretrained("your/generation_method", push_to_hub=True) model.save_pretrained("your/generation_method", push_to_hub=True)

generate.py：这是方法的核心。它必须包含一个名为generate的函数，且该函数必须以model作为第一个参数。model就是模型实例，因此你可以访问模型的全部属性和方法，包括GenerationMixin中定义的（如基础generate方法）。

注意：generate.py必须放在名为custom_generate的目录内，不能放在仓库根目录——这些文件路径在框架中是硬编码的（对应get_cached_module_file(..., module_file="custom_generate/generate.py")的调用）。

底层流程是：当基础generate收到custom_generate参数时，先检查 Python 依赖（如有），再定位generate.py中的自定义generate，最后调用它。除用于触发该机制的trust_remote_code和custom_generate两个参数外，其余收到的参数与model会全部转发给你的函数。因此你的generate可以混用原有参数与自定义参数（甚至返回不同输出类型）：

import torch def generate(model, input_ids, generation_config=None, left_padding=None, **kwargs): generation_config = generation_config or model.generation_config # 回落到模型的生成配置 cur_length = input_ids.shape[1] max_length = generation_config.max_length or cur_length + generation_config.max_new_tokens # 自定义参数示例：在 prompt 前添加 `left_padding`（整数个）pad token if left_padding is not None: if not isinstance(left_padding, int) or left_padding < 0: raise ValueError(f"left_padding must be an integer larger than 0, but is {left_padding}") pad_token = kwargs.pop("pad_token", None) or generation_config.pad_token_id or model.config.pad_token_id if pad_token is None: raise ValueError("pad_token is not defined") batch_size = input_ids.shape[0] pad_tensor = torch.full(size=(batch_size, left_padding), fill_value=pad_token).to(input_ids.device) input_ids = torch.cat((pad_tensor, input_ids), dim=1) cur_length = input_ids.shape[1] # 最简贪心解码循环 while cur_length < max_length: logits = model(input_ids).logits next_token_logits = logits[:, -1, :] next_tokens = torch.argmax(next_token_logits, dim=-1) input_ids = torch.cat((input_ids, next_tokens[:, None]), dim=-1) cur_length += 1 return input_ids

推荐实践：

• 可以放心复用原生generate中参数校验与输入准备的逻辑；
• 如果使用了model上的私有方法/属性，应在 requirements 中锁定transformers版本；
• 建议加入模型/输入校验，甚至单独的测试文件，方便用户在自己环境中做健全性检查。

本地开发与相对导入：自定义generate可以相对导入custom_generate目录内的代码，例如存在utils.py时：

from .utils import some_function

只支持与custom_generate同层的相对导入，父目录/兄弟目录导入无效。另外custom_generate参数同样支持本地目录——任何包含custom_generate结构的目录都可以直接传入，这是开发自定义方法时推荐的工作流：

gen_out = model.generate(**inputs, custom_generate="path/to/local/dir", trust_remote_code=True)

警告：加载本地目录同样会执行其中的custom_generate/generate.py，因此与 Hub 仓库一样必须trust_remote_code=True。请只对你自己编写或审查过的代码开启该选项。这一点在 tests/generation/test_utils.py 的本地目录相关测试中得到验证。

requirements.txt：可在custom_generate目录内提供requirements.txt指定额外 Python 依赖。这些依赖在运行时被检查，缺失时会抛出异常，提示用户更新环境（即前文的ImportError行为）。

README.md：模型仓库根目录的README.md通常描述模型，但既然该仓库的核心是自定义生成方法，强烈建议把重心转向方法本身的说明，并记录相对于原生generate的输入/输出差异——用户可以聚焦“新在哪里”，通用实现细节仍依赖 Transformers 文档。为便于发现，建议在 README 顶部添加custom_generate标签：

--- library_name: transformers tags: - custom_generate --- (your markdown content here)

README 推荐实践：

• 记录相对于原生generate的输入/输出差异；
• 提供自包含示例，方便快速实验；
• 说明软性约束，例如该方法只在某类模型家族上效果良好。

复用generate的输入准备：传入可调用对象

如果你想新增一个解码循环，但想保留generate里已有的输入准备逻辑（batch 扩展、attention mask、logits processors、stopping criteria 等），可以给custom_generate传一个可调用对象（Callable）：generate会执行完整的标准准备流程，然后调用你提供的可调用对象来运行解码循环，从而只覆盖解码循环本身。此时generate会先执行完整的输入准备，再调用可调用对象，并自动比对可调用对象的签名以提取新增参数（见 _extract_generation_mode_kwargs）。

def custom_loop(model, input_ids, attention_mask, logits_processor, stopping_criteria, generation_config, **model_kwargs): next_tokens = input_ids while input_ids.shape[1] < stopping_criteria[0].max_length: logits = model(next_tokens, attention_mask=attention_mask, **model_kwargs).logits next_token_logits = logits_processor(input_ids, logits[:, -1, :]) next_tokens = torch.argmax(next_token_logits, dim=-1)[:, None] input_ids = torch.cat((input_ids, next_tokens), dim=-1) attention_mask = torch.cat((attention_mask, torch.ones_like(next_tokens)), dim=-1) return input_ids output = model.generate( **inputs, custom_generate=custom_loop, max_new_tokens=10, )

提示：如果发布custom_generate仓库，你的generate实现内部同样可以定义一个可调用对象并传给model.generate()，这样既能自定义解码循环，又能享受 Transformers 内建的输入准备逻辑。

如何发现自定义生成方法

在模型库中搜索custom_generate标签即可找到全部自定义生成方法。除标签外，官方还维护了两个精选集合：社区贡献的方法集合，以及包含“此前属于 transformers 内置、现迁移为 custom_generate 仓库的参考实现”的教程集合。如前文GENERATION_MODES_MAPPING所示，DOLA、contrastive search、group beam search、constrained beam search 均已以这种形式迁移。

策略选择与验证小结

• 默认/短输出、要确定性：贪心搜索（默认max_new_tokens=20，注意显式设置上限，很多模型默认生成长度远大于 20，如 Llama2 为 4096）；
• 要多样性/创造性：do_sample=True, num_beams=1；
• 输入基准任务（描述、ASR）：num_beams>1；
• 需要特殊解码逻辑：优先评估是否可只覆盖解码循环（传 Callable，复用输入准备）；确需完整替换时用 Hub 仓库/本地目录 +trust_remote_code=True；
• 发布前：核对 README 是否说明了新参数、输出差异与适用模型范围，依赖锁定是否写入custom_generate/requirements.txt。

实现与测试证据集中在 src/transformers/generation/utils.py（generate入口、load_custom_generate、模式映射与废弃策略迁移逻辑）和 tests/generation/test_utils.py（custom_generate的参数注入、模型仓库覆盖、依赖检查、trust_remote_code强制要求、本地目录与 Callable 等用例）。深入理解常见解码策略的数学细节，可参考官方博客 “How to generate text: using different decoding methods for language generation with Transformers”。

【免费下载链接】transformers🤗 Transformers: the model-definition framework for state-of-the-art machine learning models in text, vision, audio, and multimodal models, for both inference and training.项目地址: https://gitcode.com/GitHub_Trending/tra/transformers