添加新管道

通过继承并实现几个方法来定制您的管道。将代码分享到Hub社区，并将管道注册到Transformers，以便所有人都能快速轻松地使用它。

本指南将引导您完成向Transformers添加新管道的过程。

设计选择

至少，您只需要为管道提供适合任务的输入。这也是您设计管道的起点。

确定管道可以接受的输入类型。它可以是字符串、原始字节、字典等等。尽可能将输入保持为纯 Python，因为它具有更好的兼容性。接下来，确定管道应该返回的输出。同样，将输出保持为 Python 是最简单和最佳的选择，因为它更容易处理。

保持输入和输出简单，理想情况下是 JSON 可序列化的，这使得用户无需学习新的对象类型即可更轻松地运行您的管道。支持多种不同的输入类型以获得更高的易用性也很常见。例如，接受文件名、URL 或原始字节的音频文件，为用户提供了提供音频数据时更大的灵活性。

创建管道

确定输入和输出后，您可以开始实现管道。您的管道应该继承自基础管道类，并包含 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` 允许用户将额外参数传递给管道。这可以在初始化期间或调用管道时进行。`_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` 或 `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",
)

共享您的管道

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

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

上传到 Hub

将您的管道代码以 Python 文件的形式添加到 Hub。

例如，一个用于句子对分类的自定义管道可能如下面的代码所示。该实现适用于 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 函数将管道详细信息（任务类型、管道类、支持的后端）注册到模型的 `config.json` 文件中。

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

调用 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 和 tf_model_mapping 中定义的模型上运行。这对于测试与新模型的未来兼容性非常重要。

您还会注意到在 run_pipeline_test 函数中使用了 `ANY`。模型是随机的，所以您无法检查实际值。使用 `ANY` 允许测试匹配管道类型的输出。

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

test_small_model_pt 和 test_small_model_tf，使用小型模型进行这些管道测试，以确保它们返回正确的输出。结果不一定有意义。每个管道都应该返回相同的结果。
test_large_model_pt 和 test_large_model_tf，使用真实模型进行这些管道测试，以确保它们返回有意义的结果。这些测试速度较慢，应标记为慢速。

< > 在 GitHub 上更新