Transformers 自定义 Pipeline 开发指南继承 Pipeline 基类、注册新任务并发布到 Hub【免费下载链接】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本文基于 Transformers 仓库官方文档 add_new_pipeline.md完整讲解如何为 Transformers 添加一条全新的 pipeline从输入/输出设计决策到继承Pipeline基类并实现四个核心方法再到通过PIPELINE_REGISTRY注册任务最后给出发布到 Hubpush_to_hubtrust_remote_code与直接合入 Transformers补充测试用例两条落地路径。读完本文你可以独立完成一条自定义 pipeline 从实现、注册到社区共享的全过程。一、设计决策从输入和输出开始官方文档给出的第一原则是设计 pipeline 时最少的工作量只需要为Pipeline提供一个适配任务的输入——这也是设计工作的起点。具体而言需要决策两件事确定Pipeline能接受哪些输入类型可以是字符串、原始字节、字典等。文档建议尽可能保持输入为纯 Python 类型因为兼容性更好确定Pipeline应返回的输出同样保持为纯 Python 类型最简单也最容易使用。保持输入输出简单、理想情况下是 JSON 可序列化的可以让用户无需学习新的对象类型即可运行你的Pipeline。此外支持多种输入类型也是常见做法以换取更好的易用性。例如允许音频文件以文件名、URL 或原始字节三种方式提供能给用户更大的灵活性。二、创建 Pipeline继承基类并实现 4 个方法输入与输出确定后就可以开始实现Pipeline了。你的 pipeline 应继承基类Pipeline并包含 4 个方法from transformers import Pipeline class MyPipeline(Pipeline): def _sanitize_parameters(self, **kwargs): def preprocess(self, inputs, args2): def _forward(self, model_inputs): def postprocess(self, model_outputs):这 4 个方法在基类 base.py 中全部声明为abstractmethod因此子类必须实现否则实例化会直接失败。下面逐一说明。1. preprocess把输入转成模型可用的张量格式preprocess接收用户输入将其转换为模型的合适输入格式def preprocess(self, inputs, maybe_arg2): model_input Tensor(inputs[input_ids]) return {model_input: model_input}从基类 docstringbase.py可以看到更精确的契约preprocess接收特定 pipeline 的input_返回一个包含_forward正常运行所需全部内容的字典——至少应包含一个张量但可以附带任意其他字段。2. _forward模型相关逻辑的唯一归属地文档强调_forward不应被直接调用forward才是推荐入口因为它包含保障一切在预期设备上正确运行的防护逻辑。凡是与模型相关的事情都放在_forward其他事情则分别归属preprocess或postprocess。def _forward(self, model_inputs): outputs self.model(**model_inputs) return outputs这一设计在源码中有直接印证基类的forward方法base.py在调用_forward之前会进入self.device_placement()上下文、通过get_inference_context()PyTorch 下即torch.no_grad关闭梯度计算并用_ensure_tensor_on_device把输入张量搬到self.device、把输出搬回 CPU。也就是说设备管理、推理模式等热路径外围逻辑由基类统一兜底你的_forward只需专注把输入喂给模型这一件事这也是preprocess/postprocess分离出来的根本原因——让热路径尽可能快。3. postprocess把模型输出整理成用户友好的结果postprocess从_forward的模型输出中生成最终输出def postprocess(self, model_outputs, top_k5): best_class model_outputs[logits].softmax(-1) return best_class基类 docstringbase.py指出postprocess接收_forward的原始输出通常是张量并将其重新组织为更友好的形式通常输出仅包含字符串和数字的列表或字典——这与第一节输入输出保持简单、JSON 可序列化的设计原则首尾呼应。4. _sanitize_parameters让用户能传入额外参数_sanitize_parameters允许用户在初始化或调用Pipeline时传入额外参数。它返回3 个字典分别作为关键字参数直接传给preprocess、_forward和postprocess。文档特别提醒如果用户没有以额外参数调用 pipeline就不要往字典里添加任何东西——这样函数签名中的默认值得以保留行为更自然。基类对这一点的说明是它不应被直接调用而是由__init__和__call__自动调用并完成最终参数解析base.py。例如在postprocess中加一个top_k参数以返回最可能的 5 个类别然后在_sanitize_parameters中检查用户是否传入了top_k并加入postprocess_kwargsdef _sanitize_parameters(self, **kwargs): preprocess_kwargs {} if maybe_arg in kwargs: preprocess_kwargs[maybe_arg] kwargs[maybe_arg] postprocess_kwargs {} if top_k in kwargs: postprocess_kwargs[top_k] kwargs[top_k] return preprocess_kwargs, {}, postprocess_kwargs于是 pipeline 就支持了按用户意愿返回 top-k 最可能标签from transformers import pipeline pipeline pipeline(my-task) # 返回 3 个最可能标签 pipeline(This is the best meal Ive ever had, top_k3) # 默认返回 5 个最可能标签 pipeline(This is the best meal Ive ever had)从源码结构看这条参数链路是完整闭合的Pipeline.__call__先调用self._sanitize_parameters(**kwargs)再把返回的三组参数与初始化时保存的_preprocess_params/_forward_params/_postprocess_params合并base.py单条输入场景下最终由run_single按preprocess→forward→postprocess的顺序串联执行base.py批量输入则走get_iterator构建DataLoader数据流。一个真实实现可参考文本分类 pipeline 的_sanitize_parameterstext_classification.py它接收function_to_apply、top_k并透传 tokenizer 相关 kwargs。三、注册 pipeline写入 PIPELINE_REGISTRYpipeline 写好后需要把你支持的新任务注册进PIPELINE_REGISTRY。注册项定义三样东西用pt_model指定支持的 PyTorch 模型类用default指定默认模型——必须来自一个模型表现符合预期的特定 revision分支或 commit hash用type指定期望的输入类型。from transformers.pipelines import PIPELINE_REGISTRY from transformers import AutoModelForSequenceClassification PIPELINE_REGISTRY.register_pipeline( new-task, pipeline_classMyPipeline, pt_modelAutoModelForSequenceClassification, default{pt: (user/awesome-model, branch-name)}, typetext, )源码层面base.pyPipelineRegistry.register_pipeline的签名为task, pipeline_class, pt_modelNone, defaultNone, typeNone其内部行为可以补充文档未展开的细节若任务名已存在会打印 already registered. Overwriting... 警告并覆盖旧注册项pt_model为None时置为空元组非元组则自动包成单元素元组default若不是{model: ...}形式会被自动包一层{model: default}最终结构统一为{model: (model_id, revision)}注册结果task_impl {impl: pipeline_class, pt: pt_model, ...}会写入self.supported_tasks[task]同时挂到pipeline_class._registered_impl上——后者正是后面push_to_hub能把注册信息写进config.json的关键见下文第五节。对照仓库内置任务表的实现pipelines/init.py可以看到官方任务的注册结构与上述机制完全一致。例如text-classification条目text-classification: { impl: TextClassificationPipeline, pt: (AutoModelForSequenceClassification,) if is_torch_available() else (), default: {model: (distilbert/distilbert-base-uncased-finetuned-sst-2-english, 714eb0f)}, type: text, },注意default中的第二个元素就是一个 commit 短 hash714eb0f印证了文档默认模型应来自特定 revision的要求。所有内置任务集中在SUPPORTED_TASKS字典中并由PIPELINE_REGISTRY PipelineRegistry(supported_tasksSUPPORTED_TASKS, task_aliasesTASK_ALIASES)实例化pipelines/init.py。任务别名如sentiment-analysis - text-classification在TASK_ALIASES中定义PIPELINE_REGISTRY.check_task会先做别名归一化再查任务表未知任务抛出KeyErrorbase.py。四、分享你的 pipeline官方给出两条分享路径上传到 Hub快不需要 Transformers 团队评审或直接把 pipeline 加进 Transformers慢需要评审并补充测试。路径一上传到 Hub把 pipeline 代码放入 Hub 上的一个 Python 文件。文档以句子对分类sentence pair classification为例给出完整自定义 pipelineimport numpy as np from transformers import Pipeline def softmax(outputs): maxes np.max(outputs, axis-1, keepdimsTrue) shifted_exp np.exp(outputs - maxes) return shifted_exp / shifted_exp.sum(axis-1, keepdimsTrue) class PairClassificationPipeline(Pipeline): def _sanitize_parameters(self, **kwargs): preprocess_kwargs {} if second_text in kwargs: preprocess_kwargs[second_text] kwargs[second_text] return preprocess_kwargs, {}, {} def preprocess(self, text, second_textNone): return self.tokenizer(text, text_pairsecond_text, return_tensorsself.framework) def _forward(self, model_inputs): return self.model(**model_inputs) def postprocess(self, model_outputs): logits model_outputs.logits[0].numpy() probabilities softmax(logits) best_class np.argmax(probabilities) label self.model.config.id2label[best_class] score probabilities[best_class].item() logits logits.tolist() return {label: label, score: score, logits: logits}这段代码完整体现了第二节的四个方法preprocess借助self.tokenizer处理text_pairpostprocess输出纯 JSON 友好的字典。把代码保存为pair_classification.py然后导入并注册from pair_classification import PairClassificationPipeline from transformers.pipelines import PIPELINE_REGISTRY from transformers import AutoModelForSequenceClassification PIPELINE_REGISTRY.register_pipeline( pair-classification, pipeline_classPairClassificationPipeline, pt_modelAutoModelForSequenceClassification, )文档称register_pipeline函数会把 pipeline 的详细信息任务类型、pipeline 类、支持的后端注册进模型的config.json。具体机制在源码中清晰可见Pipeline.push_to_hub时若实例带有_registered_impl属性就会把实现类转换成模块名.类名字符串、模型类转换为其类名最终写为self.model.config.custom_pipelinesbase.py。对应的config.json片段形如custom_pipelines: { pair-classification: { impl: pair_classification.PairClassificationPipeline, pt: [ AutoModelForSequenceClassification ], } },接着调用Pipeline.push_to_hub推送 pipeline 到 Hub包含代码的 Python 文件会被复制到 Hubpipeline 的模型与 tokenizer 也会被保存并推送。此时你的 pipeline 就已可用在 Hub 上的命名空间下了from transformers import pipeline pipeline pipeline(taskpair-classification, modelsgugger/finetuned-bert-mrpc) pipeline.push_to_hub(pair-classification-pipeline)使用时加载 pipeline 需加上trust_remote_codeTruefrom transformers import pipeline pipeline pipeline(taskpair-classification, trust_remote_codeTrue)从源码结构看这条远程加载链路的落点是工厂函数pipeline()它会检查config.custom_pipelinespipelines/init.py经clean_custom_task把pt中的类名字符串还原为真正的模型类后通过get_class_from_dynamic_module从仓库动态加载impl指向的类若未设置trust_remote_code则会抛出明确提示需先阅读代码以规避恶意使用的ValueErrorpipelines/init.py。这解释了为何消费端必须显式传trust_remote_codeTrue。路径二直接加入 Transformers把自定义 pipeline 加入 Transformers 本体要求补充测试以确保一切按预期工作并请求 Transformers 团队评审。流程如下添加代码模块把 pipeline 代码作为新模块加入 pipelines 子包并把它加入 pipelines/init.py 中定义的任务列表——在当前的仓库结构中即导入你的 pipeline 类并新增一条SUPPORTED_TASKS条目参考上文text-classification条目的impl/pt/default/type四字段写法添加测试在 tests/pipelines 下为新 pipeline 添加测试文件可参考现有测试文件了解如何测试你的 pipeline。测试写法上有三个文档明确强调、且在仓库测试中可对照实现的要点run_pipeline_test必须足够通用并且要能运行在model_mapping定义的模型上参考 test_pipelines_text_classification.py 中model_mapping MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING的定义以及 test_pipelines_text_classification.py 的run_pipeline_test实现。这一点非常重要因为它保证了未来与新模型的兼容性大量使用ANY占位符。由于测试模型权重是随机的无法断言具体数值使用ANY可以让测试匹配 pipeline 输出的结构类型而非具体值ANY定义于 test_pipelines_common.py例如self.assertEqual(nested_simplify(outputs), [{label: ANY(str), score: ANY(float)}])见 test_pipelines_text_classification.py额外实现 4 个测试文档原文如此列出其中两项为test_small_model_pt使用小模型如hf-internal-testing/tiny-random-distilbert验证 pipeline 返回正确的输出结构。结果不必有语义意义但每个 pipeline 应返回同一结果该测试同时验证了top_k等额外参数路径test_large_model_pt使用真实模型验证 pipeline 返回有意义的结果。这类测试运行缓慢应标记为 slow对应slow装饰器。五、小结自定义 pipeline 的完整生命周期把本文流程串起来一条自定义 pipeline 的落地路径是设计确定输入/输出为尽量简单、JSON 可序列化的纯 Python 类型实现继承Pipeline实现_sanitize_parameters参数路由返回三组 kwargs 字典、preprocess输入→张量字典、_forward只放模型相关逻辑由基类forward统一处理设备与推理上下文、postprocess模型输出→友好结果注册通过PIPELINE_REGISTRY.register_pipeline声明pipeline_class、pt_model、带 revision 的default与type发布轻量路线用push_to_hub把代码与custom_pipelines配置一并推到 Hub消费端以trust_remote_codeTrue加载正式路线则进入 src/transformers/pipelines 与 tests/pipelines 参与仓库评审与 CI 测试。所有结论均可在上述仓库文件中对照验证基类四方法与调用链在 base.py任务注册表与工厂函数在 pipelines/init.py测试范式在 tests/pipelines/test_pipelines_text_classification.py。【免费下载链接】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创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考