新しいパイプラインの追加

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 は、preprocess、_forward、postprocess に直接渡される追加のキーワード引数を格納した 3 つの辞書を返します。ユーザーが追加のパラメータなしにパイプラインを呼び出した場合は、何も追加しないでください。これにより、関数定義内のデフォルト引数が保持され、常により自然な状態が維持されます。

例えば、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でモデルが期待通りに動作する特定のリビジョン(ブランチまたはコミットハッシュ)から取得されるべきデフォルトモデル
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 で定義されたモデルで実行される必要があります。これは、将来の新しいモデルとの互換性をテストする上で重要です。

また、run_pipeline_test関数全体でANYが使用されていることに気づくでしょう。モデルはランダムであるため、実際の値を確認することはできません。ANYを使用することで、テストはパイプライン型の出力に一致させることができます。

最後に、以下の4つのテストも実装する必要があります。

test_small_model_ptでは、これらのパイプラインが正しい出力を返すことを確認するために、小さなモデルを使用します。結果が意味をなす必要はありません。各パイプラインは同じ結果を返す必要があります。
test_large_model_ptでは、これらのパイプラインが意味のある結果を返すことを確認するために、現実的なモデルを使用します。これらのテストは低速であるため、slowとしてマークする必要があります。

Update on GitHub