新增新管道

透過繼承並實現幾個方法來定製您的管道。將程式碼分享到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 上更新