添加新的pipeline

通过继承 Pipeline 并实现一些方法，使 Pipeline 成为你自己的。在 Hub 上与社区分享代码，并在 Transformers 中注册 pipeline，以便每个人都可以快速轻松地使用它。

本指南将引导你完成向 Transformers 添加新 pipeline 的过程。

设计选择

至少，你只需要为 Pipeline 提供适合任务的输入。这也是你在设计 pipeline 时应该开始的地方。

决定 Pipeline 可以接受哪些输入类型。它可以是字符串、原始字节、字典等等。尽可能将输入保持为纯 Python，因为它更兼容。接下来，决定 Pipeline 应该返回的输出。同样，将输出保持为 Python 是最简单和最佳的选择，因为它更容易使用。

保持输入和输出简单，最好是 JSON 可序列化的，这样用户可以更轻松地运行你的 Pipeline，而无需学习新的对象类型。支持多种不同的输入类型也很常见，以便更易于使用。例如，允许从文件名、URL 或原始字节接受音频文件，使用户在提供音频数据方面具有更大的灵活性。

创建 pipeline

在确定输入和输出后，你可以开始实现 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 的附加关键字参数。如果用户没有使用额外的参数调用 pipeline，则不要添加任何内容。这会将默认参数保留在函数定义中，而这始终更自然。

例如，在 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

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

pipeline 支持的机器学习框架，使用 pt_model 或 tf_model（两者都添加以确保它适用于任一框架）
默认模型，它应该来自特定修订版（分支或提交哈希），其中模型按预期与 default 一起工作
预期的输入，使用 type

from transformers.pipelines import PIPELINE_REGISTRY
from transformers import AutoModelForSequenceClassification, TFAutoModelForSequenceClassification

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

分享你的 pipeline

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

将你的 pipeline 代码上传到 Hub 会更快，因为它不需要 Transformers 团队的审查。将 pipeline 添加到 Transformers 可能会比较慢，因为它需要审查，并且你需要添加测试以确保你的 Pipeline 工作正常。

上传到 Hub

将你的 pipeline 代码添加到 Hub 的 Python 文件中。

例如，用于句子对分类的自定义 pipeline 可能如下面的代码所示。该实现适用于 PyTorch 和 TensorFlow 模型。

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, TFAutoModelForSequenceClassification

PIPELINE_REGISTRY.register_pipeline(
    "pair-classification",
    pipeline_class=PairClassificationPipeline,
    pt_model=AutoModelForSequenceClassification,
    tf_model=TFAutoModelForSequenceClassification,
)

register_pipeline 函数将 pipeline 详细信息（任务类型、pipeline 类、支持的后端）注册到模型的 config.json 文件。

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

调用 push_to_hub() 将 pipeline 推送到 Hub。包含代码的 Python 文件被复制到 Hub，并且 pipelines 模型和 tokenizer 也被保存并推送到 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

将自定义 pipeline 添加到 Transformers 需要添加测试以确保一切按预期工作，并请求 Transformers 团队的审查。

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

接下来，在 transformers/tests/pipelines 中为 pipeline 添加新测试。你可以查看其他测试，了解如何测试你的 pipeline 的示例。

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

你还会注意到 ANY 在整个 run_pipeline_test 函数中被使用。模型是随机的，因此你无法检查实际值。使用 ANY 允许测试匹配 pipeline 类型的输出。

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

test_small_model_pt 和 test_small_model_tf，为这些 pipelines 使用小型模型，以确保它们返回正确的输出。结果不必有意义。每个 pipeline 应返回相同的结果。
test_large_model_pt 和 test_large_model_tf，为这些 pipelines 使用实际模型，以确保它们返回有意义的结果。这些测试速度较慢，应标记为 slow。

< > 更新在 GitHub 上