Adding a new pipeline

Make [Pipeline] your own by subclassing it and implementing a few methods. Share the code with the community on the Hub and register the pipeline with Transformers so that everyone can quickly and easily use it.

This guide will walk you through the process of adding a new pipeline to Transformers.

Design choices

At a minimum, you only need to provide [Pipeline] with an appropriate input for a task. This is also where you should begin when designing your pipeline.

Decide what input types [Pipeline] can accept. It can be strings, raw bytes, dictionaries, and so on. Try to keep the inputs in pure Python where possible because it's more compatible. Next, decide on the output [Pipeline] should return. Again, keeping the output in Python is the simplest and best option because it's easier to work with.

Keeping the inputs and outputs simple, and ideally JSON-serializable, makes it easier for users to run your [Pipeline] without needing to learn new object types. It's also common to support many different input types for even greater ease of use. For example, making an audio file acceptable from a filename, URL, or raw bytes gives the user more flexibility in how they provide the audio data.

Create a pipeline

With an input and output decided, you can start implementing [Pipeline]. Your pipeline should inherit from the base [Pipeline] class and include 4 methods.

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):

1. preprocess takes the inputs and transforms them into the appropriate input format for the model.

def preprocess(self, inputs, maybe_arg=2):
    model_input = Tensor(inputs["input_ids"])
    return {"model_input": model_input}

2. _forward shouldn't be called directly. forward is the preferred method because it includes safeguards to make sure everything works correctly on the expected device. Anything linked to the model belongs in _forward and everything else belongs in either preprocess or postprocess.

def _forward(self, model_inputs):
    outputs = self.model(**model_inputs)
    return outputs

3. postprocess generates the final output from the models output in _forward.

def postprocess(self, model_outputs, top_k=5):
    best_class = model_outputs["logits"].softmax(-1)
    return best_class

4. _sanitize_parameters lets users pass additional parameters to [Pipeline]. This could be during initialization or when [Pipeline] is called. _sanitize_parameters returns 3 dicts of additional keyword arguments that are passed directly to preprocess, _forward, and postprocess. Don't add anything if a user didn't call the pipeline with extra parameters. This keeps the default arguments in the function definition which is always more natural.

For example, add a top_k parameter in postprocess to return the top 5 most likely classes. Then in _sanitize_parameters, check if the user passed in top_k and add it to 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

Now the pipeline can return the top most likely labels if a user chooses to.

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")

Register a pipeline

Register the new task your pipeline supports in the PIPELINE_REGISTRY. The registry defines:

• The supported Pytorch model class with pt_model
• a default model which should come from a specific revision (branch, or commit hash) where the model works as expected with default
• the expected input with 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",
)

Share your pipeline

Share your pipeline with the community on the Hub or you can add it directly to Transformers.

It's faster to upload your pipeline code to the Hub because it doesn't require a review from the Transformers team. Adding the pipeline to Transformers may be slower because it requires a review and you need to add tests to ensure your [Pipeline] works.

Upload to the Hub

Add your pipeline code to the Hub in a Python file.

For example, a custom pipeline for sentence pair classification might look like the following code below.

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}

Save the code in a file named pair_classification.py, and import and register it as shown below.

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

The register_pipeline function registers the pipeline details (task type, pipeline class, supported backends) to a models config.json file.

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

Call [~Pipeline.push_to_hub] to push the pipeline to the Hub. The Python file containing the code is copied to the Hub, and the pipelines model and tokenizer are also saved and pushed to the Hub. Your pipeline should now be available on the Hub under your namespace.

from transformers import pipeline

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

To use the pipeline, add trust_remote_code=True when loading the pipeline.

from transformers import pipeline

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

Add to Transformers

Adding a custom pipeline to Transformers requires adding tests to make sure everything works as expected, and requesting a review from the Transformers team.

Add your pipeline code as a new module to the pipelines submodule, and add it to the list of tasks defined in pipelines/init.py.

Next, add a new test for the pipeline in transformers/tests/pipelines. You can look at the other tests for examples of how to test your pipeline.

The run_pipeline_test function should be very generic and run on the models defined in model_mapping. This is important for testing future compatibility with new models.

You'll also notice ANY is used throughout the run_pipeline_test function. The models are random, so you can't check the actual values. Using ANY allows the test to match the output of the pipeline type instead.

Finally, you should also implement the following 4 tests.

1. test_small_model_pt, use a small model for these pipelines to make sure they return the correct outputs. The results don't have to make sense. Each pipeline should return the same result.
2. test_large_model_pt, use a realistic model for these pipelines to make sure they return meaningful results. These tests are slow and should be marked as slow.

Source: docs/source/en/add_new_pipeline.md