#
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 machine learning framework the pipeline supports with either `pt_model` or `tf_model` (add both to ensure it works with either frameworks)
- 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, 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",
    )

##
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. The implementation works for PyTorch and TensorFlow models.

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

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

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"
          ],
          "tf": [
            "TFAutoModelForSequenceClassification"
          ],
        }
      },

Call 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 and tf_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 and test_small_model_tf, 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 nad test_large_model_tf, use a realistic model for these pipelines to make sure they return meaningful results. These tests are slow and should be marked as slow.

< > Update on GitHub