添加新流水线

通过继承 Pipeline 并实现几个方法，你可以定制自己的流水线。将代码在 Hub 上与社区分享，并将该流水线注册到 Transformers 中，以便每个人都能快速便捷地使用它。

本指南将引导你完成在 Transformers 中添加新流水线的整个过程。

设计选择

至少，你只需要为 Pipeline 提供一个适用于该任务的适当输入。这也是设计流水线时应该开始的地方。

决定 Pipeline 可以接受哪些输入类型。可以是字符串、原始字节、字典等等。尽可能将输入保持在纯 Python 类型中，因为这样兼容性更好。接着，决定 Pipeline 应该返回什么样的输出。同样，将输出保持在 Python 类型是最简单且最佳的选择，因为这样更易于处理。

保持输入和输出简单，且理想情况下可被 JSON 序列化，这样用户无需学习新的对象类型即可更轻松地运行你的 Pipeline。为了进一步提高易用性，支持多种不同的输入类型也很常见。例如，使音频文件可以通过文件名、URL 或原始字节提供，能给用户在提供音频数据时带来更大的灵活性。

创建流水线

在确定了输入和输出后，你可以开始实现 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

现在，如果用户选择这样做，流水线可以返回概率最高的标签。

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 中注册你的流水线所支持的新任务。该注册表定义了：

通过 pt_model 定义支持的 Pytorch 模型类
通过 default 定义默认模型，该模型应来自一个特定的版本（分支或 commit hash），在该版本中模型可以按预期工作
通过 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",
)

分享你的流水线

在 Hub 上与社区分享你的流水线，或者你可以将其直接添加到 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() 将流水线推送到 Hub。包含代码的 Python 文件将被复制到 Hub，同时流水线的模型和分词器也会被保存并推送到 Hub。你的流水线现在应该可以在你的命名空间下的 Hub 中使用了。

from transformers import pipeline

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

要使用该流水线，在加载流水线时请添加 trust_remote_code=True。

from transformers import pipeline

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

添加到 Transformers

将自定义流水线添加到 Transformers 需要添加测试以确保一切按预期工作，并请求 Transformers 团队进行审核。

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

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

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

您还会注意到 ANY 被贯穿用于 run_pipeline_test 函数中。由于模型是随机的，因此您无法检查实际数值。使用 ANY 可以让测试改为匹配 pipeline 类型的输出。

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

test_small_model_pt，为这些 pipeline 使用一个小模型，以确保它们返回正确的输出。结果不需要有实际意义。每个 pipeline 应该返回相同的结果。
test_large_model_pt，为这些 pipeline 使用一个真实的模型，以确保它们返回有意义的结果。这些测试运行较慢，应被标记为 slow。

在 GitHub 上更新