添加新流水线

通过继承 Pipeline 并实现几个方法，打造属于你自己的 Pipeline。在 Hub 上与社区分享代码，并将该 pipeline 注册到 Transformers 中，以便所有人都能快速轻松地使用它。

本指南将引导你了解向 Transformers 添加新 pipeline 的流程。

设计选择

至少，您只需为 Pipeline 提供适合该任务的输入即可。在设计您的流水线时，也应从这里开始。

确定 Pipeline 可以接受的输入类型。它可以是字符串、原始字节、字典等。在可能的情况下，尽量将输入保持为纯 Python 类型，因为这样兼容性更好。接下来，确定 Pipeline 应返回的输出。同样，将输出保持为 Python 类型是最简单且最佳的选择，因为这样更易于处理。

保持输入和输出简单，并理想情况下使其可 JSON 序列化，能让用户无需学习新的对象类型即可轻松运行您的 Pipeline。为了进一步提升易用性，支持多种不同的输入类型也是常见做法。例如，允许用户通过文件名、URL 或原始字节来提供音频文件，可以让他们在提供音频数据时拥有更大的灵活性。

创建流水线

在确定输入和输出后，您可以开始实现 Pipeline。您的 pipeline 应继承自基类 Pipeline，并包含 4 个方法。

from transformers import Pipeline

class MyPipeline(Pipeline):
    def _sanitize_parameters(self, **kwargs):

    def preprocess(self, inputs, args=2):

    def _forward(self, model_inputs):

    def postprocess(self, model_outputs):

preprocess 接收输入并将其转换为模型所需的适当输入格式。

def preprocess(self, inputs, maybe_arg=2):
    model_input = Tensor(inputs["input_ids"])
    return {"model_input": model_input}

_forward 不应被直接调用。推荐使用 forward，因为它包含保障措施，以确保一切在预期设备上正确运行。与模型相关的任何内容都应放在 _forward 中，其余内容则应放在 preprocess 或 postprocess 中。

def _forward(self, model_inputs):
    outputs = self.model(**model_inputs)
    return outputs

postprocess 根据 _forward 中模型的输出生成最终结果。

def postprocess(self, model_outputs, top_k=5):
    best_class = model_outputs["logits"].softmax(-1)
    return best_class

_sanitize_parameters 允许用户向 Pipeline 传递额外参数。这可以在初始化时进行，也可以在调用 Pipeline 时进行。_sanitize_parameters 返回 3 个包含额外关键字参数的字典，这些参数会直接传递给 preprocess、_forward 和 postprocess。如果用户未通过管道传入额外参数，则不要添加任何内容。这样可以保留函数定义中的默认参数，这种方式通常更自然。

例如，在 postprocess 中添加一个 top_k 参数，以返回最可能的 5 个类别。然后在 _sanitize_parameters 中，检查用户是否传入了 top_k，并将其添加到 postprocess_kwargs 中。

def _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 可以返回最可能的标签。

from transformers import pipeline

pipeline = pipeline("my-task")
# returns 3 most likely labels
pipeline("This is the best meal I've ever had", top_k=3)
# returns 5 most likely labels by default
pipeline("This is the best meal I've ever had")

注册流水线

在 PIPELINE_REGISTRY 中注册你的 pipeline 所支持的新任务。该注册表定义了：

使用 pt_model 支持的 Pytorch 模型类
一个默认模型，该模型应来自一个特定的修订版本（分支或提交哈希），在此修订版本中，模型在 default 下能按预期正常工作
type 类型的预期输入

from transformers.pipelines import PIPELINE_REGISTRY
from transformers import AutoModelForSequenceClassification

PIPELINE_REGISTRY.register_pipeline(
    "new-task",
    pipeline_class=MyPipeline,
    pt_model=AutoModelForSequenceClassification,
    default={"pt": ("user/awesome-model", "branch-name")},
    type="text",
)

分享您的 pipeline

在 Hub 上与社区分享您的 pipeline，或者您可以直接将其添加到 Transformers 中。

将您的流水线代码上传到 Hub 的速度更快，因为它不需要 Transformers 团队进行审核。将流水线添加到 Transformers 中可能会更慢，因为这需要审核，并且您需要添加测试以确保您的 Pipeline 能够正常工作。

上传到 Hub

在 Python 文件中添加您的流水线代码，并将其上传到 Hub。

例如，用于句子对分类的自定义流水线可能如下面的代码所示。

import numpy as np
from transformers import Pipeline

def softmax(outputs):
    maxes = np.max(outputs, axis=-1, keepdims=True)
    shifted_exp = np.exp(outputs - maxes)
    return shifted_exp / shifted_exp.sum(axis=-1, keepdims=True)

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_text=None):
        return self.tokenizer(text, text_pair=second_text, return_tensors=self.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}

将代码保存为名为 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_class=PairClassificationPipeline,
    pt_model=AutoModelForSequenceClassification,
)

register_pipeline 函数会将流水线详情（任务类型、流水线类、支持的后端）注册到模型的 config.json 文件中。

  "custom_pipelines": {
    "pair-classification": {
      "impl": "pair_classification.PairClassificationPipeline",
      "pt": [
        "AutoModelForSequenceClassification"
      ],
    }
  },

调用 push_to_hub() 将 pipeline 推送到 Hub。包含代码的 Python 文件会被复制到 Hub，pipeline 的模型和分词器也会被保存并推送到 Hub。现在，你的 pipeline 应该可以在 Hub 上你的命名空间下找到。

from transformers import pipeline

pipeline = pipeline(task="pair-classification", model="sgugger/finetuned-bert-mrpc")
pipeline.push_to_hub("pair-classification-pipeline")

要使用 pipeline，在加载 pipeline 时添加 trust_remote_code=True。

from transformers import pipeline

pipeline = pipeline(task="pair-classification", trust_remote_code=True)

添加到 Transformers

向 Transformers 添加自定义 pipeline 需要补充测试以确保一切按预期运行，并向 Transformers 团队请求代码审查。

将您的流水线代码作为新模块添加到 pipelines 子模块中，并将其添加到 pipelines/init.py 中定义的任务列表中。

接下来，在 transformers/tests/pipelines 中为该 pipeline 添加一个新测试。你可以参考其他测试，了解如何测试你的 pipeline。

run_pipeline_test 函数应当非常通用，并针对 model_mapping 中定义的模型运行。这对于测试与未来新模型的兼容性至关重要。

您还会注意到 run_pipeline_test 函数中多处使用了 ANY。由于模型是随机的，因此无法检查实际值。使用 ANY 可以让测试匹配流水线类型的输出。

最后，您还应该实现以下 4 个测试。

test_small_small_model_pt，在这些流水线中使用一个小模型，以确保它们返回正确的输出。结果无需有意义。每个流水线应返回相同的结果。
test_large_model_pt，在这些流水线中使用一个真实的模型，以确保它们返回有意义的结果。这些测试运行缓慢，应标记为 slow。

Update on GitHub