Pipelines¶

The pipelines are a great and easy way to use models for inference. These pipelines are objects that abstract most of the complex code from the library, offering a simple API dedicated to several tasks, including Named Entity Recognition, Masked Language Modeling, Sentiment Analysis, Feature Extraction and Question Answering.

There are two categories of pipeline abstractions to be aware about:

The pipeline() which is the most powerful object encapsulating all other pipelines
The other task-specific pipelines, such as TokenClassificationPipeline or QuestionAnsweringPipeline

The pipeline abstraction¶

The pipeline abstraction is a wrapper around all the other available pipelines. It is instantiated as any other pipeline but requires an additional argument which is the task.

transformers.pipeline(task: str, model: Optional = None, config: Optional[Union[str, transformers.configuration_utils.PretrainedConfig]] = None, tokenizer: Optional[Union[str, transformers.tokenization_utils.PreTrainedTokenizer]] = None, framework: Optional[str] = None, **kwargs) → transformers.pipelines.Pipeline[source]¶

Utility factory method to build a pipeline.

Pipeline are made of:

A Tokenizer instance in charge of mapping raw textual input to token

A Model instance

Some (optional) post processing for enhancing model’s output

Parameters

task (str) –
The task defining which pipeline will be returned. Currently accepted tasks are:
- ”feature-extraction”: will return a FeatureExtractionPipeline
- ”sentiment-analysis”: will return a TextClassificationPipeline
- ”ner”: will return a TokenClassificationPipeline
- ”question-answering”: will return a QuestionAnsweringPipeline
- ”fill-mask”: will return a FillMaskPipeline
- ”summarization”: will return a SummarizationPipeline
- ”translation_xx_to_yy”: will return a TranslationPipeline
- ”text-generation”: will return a TextGenerationPipeline
model (str or PreTrainedModel or TFPreTrainedModel, optional, defaults to None) –
The model that will be used by the pipeline to make predictions. This can be None, a model identifier or an actual pre-trained model inheriting from PreTrainedModel for PyTorch and TFPreTrainedModel for TensorFlow.

If None, the default for this pipeline will be loaded.
config (str or PretrainedConfig, optional, defaults to None) –
The configuration that will be used by the pipeline to instantiate the model. This can be None, a model identifier or an actual pre-trained model configuration inheriting from PretrainedConfig.

If None, the default for this pipeline will be loaded.
tokenizer (str or PreTrainedTokenizer, optional, defaults to None) –
The tokenizer that will be used by the pipeline to encode data for the model. This can be None, a model identifier or an actual pre-trained tokenizer inheriting from PreTrainedTokenizer.

If None, the default for this pipeline will be loaded.
framework (str, optional, defaults to None) –
The framework to use, either “pt” for PyTorch or “tf” for TensorFlow. The specified framework must be installed.

If no framework is specified, will default to the one currently installed. If no framework is specified and both frameworks are installed, will default to PyTorch.

Returns

Class inheriting from Pipeline, according to the task.

Return type

Pipeline

Examples:

from transformers import pipeline, AutoModelForTokenClassification, AutoTokenizer

# Sentiment analysis pipeline
pipeline('sentiment-analysis')

# Question answering pipeline, specifying the checkpoint identifier
pipeline('question-answering', model='distilbert-base-cased-distilled-squad', tokenizer='bert-base-cased')

# Named entity recognition pipeline, passing in a specific model and tokenizer
model = AutoModelForTokenClassification.from_pretrained("dbmdz/bert-large-cased-finetuned-conll03-english")
tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
pipeline('ner', model=model, tokenizer=tokenizer)

The task specific pipelines¶

Parent class: Pipeline¶

class transformers.Pipeline(model: Union[PreTrainedModel, TFPreTrainedModel], tokenizer: transformers.tokenization_utils.PreTrainedTokenizer, modelcard: Optional[transformers.modelcard.ModelCard] = None, framework: Optional[str] = None, task: str = '', args_parser: transformers.pipelines.ArgumentHandler = None, device: int = - 1, binary_output: bool = False)[source]¶

The Pipeline class is the class from which all pipelines inherit. Refer to this class for methods shared across different pipelines.

Base class implementing pipelined operations. Pipeline workflow is defined as a sequence of the following operations:

Input -> Tokenization -> Model Inference -> Post-Processing (Task dependent) -> Output

Pipeline supports running on CPU or GPU through the device argument. Users can specify device argument as an integer, -1 meaning “CPU”, >= 0 referring the CUDA device ordinal.

Some pipeline, like for instance FeatureExtractionPipeline (‘feature-extraction’) outputs large tensor object as nested-lists. In order to avoid dumping such large structure as textual data we provide the binary_output constructor argument. If set to True, the output will be stored in the pickle format.

Parameters

model (PreTrainedModel or TFPreTrainedModel) – The model that will be used by the pipeline to make predictions. This needs to be a model inheriting from PreTrainedModel for PyTorch and TFPreTrainedModel for TensorFlow.
tokenizer (PreTrainedTokenizer) – The tokenizer that will be used by the pipeline to encode data for the model. This object inherits from PreTrainedTokenizer.
modelcard (str or ModelCard, optional, defaults to None) – Model card attributed to the model for this pipeline.
framework (str, optional, defaults to None) –
The framework to use, either “pt” for PyTorch or “tf” for TensorFlow. The specified framework must be installed.

If no framework is specified, will default to the one currently installed. If no framework is specified and both frameworks are installed, will default to PyTorch.
args_parser (ArgumentHandler, optional, defaults to None) – Reference to the object in charge of parsing supplied pipeline parameters.
device (int, optional, defaults to -1) – Device ordinal for CPU/GPU supports. Setting this to -1 will leverage CPU, >=0 will run the model on the associated CUDA device id.
binary_output (bool, optional, defaults to False) – Flag indicating if the output the pipeline should happen in a binary format (i.e. pickle) or as raw text.

Returns

Pipeline returns list or dictionary depending on:

Whether the user supplied multiple samples

Whether the pipeline exposes multiple fields in the output object

Return type

List or Dict

predict(X)[source]¶: Scikit / Keras interface to transformers’ pipelines. This method will forward to __call__().

save_pretrained(save_directory)[source]¶: Save the pipeline’s model and tokenizer to the specified save_directory

transform(X)[source]¶: Scikit / Keras interface to transformers’ pipelines. This method will forward to __call__().

TokenClassificationPipeline¶

class transformers.TokenClassificationPipeline(model: Union[PreTrainedModel, TFPreTrainedModel], tokenizer: transformers.tokenization_utils.PreTrainedTokenizer, modelcard: Optional[transformers.modelcard.ModelCard] = None, framework: Optional[str] = None, args_parser: transformers.pipelines.ArgumentHandler = None, device: int = - 1, binary_output: bool = False, ignore_labels=['O'], task: str = '', grouped_entities: bool = False)[source]¶

Named Entity Recognition pipeline using ModelForTokenClassification head. See the named entity recognition usage examples for more information.

This token recognition pipeline can currently be loaded from the pipeline() method using the following task identifier(s):

“ner”, for predicting the classes of tokens in a sequence: person, organisation, location or miscellaneous.

The models that this pipeline can use are models that have been fine-tuned on a token classification task. See the up-to-date list of available models on huggingface.co/models.

Parameters

model (PreTrainedModel or TFPreTrainedModel) – The model that will be used by the pipeline to make predictions. This needs to be a model inheriting from PreTrainedModel for PyTorch and TFPreTrainedModel for TensorFlow.
tokenizer (PreTrainedTokenizer) – The tokenizer that will be used by the pipeline to encode data for the model. This object inherits from PreTrainedTokenizer.
modelcard (str or ModelCard, optional, defaults to None) – Model card attributed to the model for this pipeline.
framework (str, optional, defaults to None) –
The framework to use, either “pt” for PyTorch or “tf” for TensorFlow. The specified framework must be installed.

If no framework is specified, will default to the one currently installed. If no framework is specified and both frameworks are installed, will default to PyTorch.
args_parser (ArgumentHandler, optional, defaults to None) – Reference to the object in charge of parsing supplied pipeline parameters.
device (int, optional, defaults to -1) – Device ordinal for CPU/GPU supports. Setting this to -1 will leverage CPU, >=0 will run the model on the associated CUDA device id.

NerPipeline¶

This class is an alias of the TokenClassificationPipeline defined above. Please refer to that pipeline for documentation and usage examples.

FillMaskPipeline¶

class transformers.FillMaskPipeline(model: Union[PreTrainedModel, TFPreTrainedModel], tokenizer: transformers.tokenization_utils.PreTrainedTokenizer, modelcard: Optional[transformers.modelcard.ModelCard] = None, framework: Optional[str] = None, args_parser: transformers.pipelines.ArgumentHandler = None, device: int = - 1, topk=5, task: str = '')[source]¶

Masked language modeling prediction pipeline using ModelWithLMHead head. See the masked language modeling usage examples for more information.

This mask filling pipeline can currently be loaded from the pipeline() method using the following task identifier(s):

“fill-mask”, for predicting masked tokens in a sequence.

The models that this pipeline can use are models that have been trained with a masked language modeling objective, which includes the bi-directional models in the library. See the up-to-date list of available models on huggingface.co/models.

Parameters

model (PreTrainedModel or TFPreTrainedModel) – The model that will be used by the pipeline to make predictions. This needs to be a model inheriting from PreTrainedModel for PyTorch and TFPreTrainedModel for TensorFlow.
tokenizer (PreTrainedTokenizer) – The tokenizer that will be used by the pipeline to encode data for the model. This object inherits from PreTrainedTokenizer.
modelcard (str or ModelCard, optional, defaults to None) – Model card attributed to the model for this pipeline.
framework (str, optional, defaults to None) –
The framework to use, either “pt” for PyTorch or “tf” for TensorFlow. The specified framework must be installed.

If no framework is specified, will default to the one currently installed. If no framework is specified and both frameworks are installed, will default to PyTorch.
args_parser (ArgumentHandler, optional, defaults to None) – Reference to the object in charge of parsing supplied pipeline parameters.
device (int, optional, defaults to -1) – Device ordinal for CPU/GPU supports. Setting this to -1 will leverage CPU, >=0 will run the model on the associated CUDA device id.

FeatureExtractionPipeline¶

class transformers.FeatureExtractionPipeline(model: Union[PreTrainedModel, TFPreTrainedModel], tokenizer: transformers.tokenization_utils.PreTrainedTokenizer, modelcard: Optional[transformers.modelcard.ModelCard] = None, framework: Optional[str] = None, args_parser: transformers.pipelines.ArgumentHandler = None, device: int = - 1, task: str = '')[source]¶

Feature extraction pipeline using Model head. This pipeline extracts the hidden states from the base transformer, which can be used as features in downstream tasks.

This feature extraction pipeline can currently be loaded from the pipeline() method using the following task identifier(s):

“feature-extraction”, for extracting features of a sequence.

All models may be used for this pipeline. See a list of all models, including community-contributed models on huggingface.co/models.

Parameters

model (PreTrainedModel or TFPreTrainedModel) – The model that will be used by the pipeline to make predictions. This needs to be a model inheriting from PreTrainedModel for PyTorch and TFPreTrainedModel for TensorFlow.
tokenizer (PreTrainedTokenizer) – The tokenizer that will be used by the pipeline to encode data for the model. This object inherits from PreTrainedTokenizer.
modelcard (str or ModelCard, optional, defaults to None) – Model card attributed to the model for this pipeline.
framework (str, optional, defaults to None) –
The framework to use, either “pt” for PyTorch or “tf” for TensorFlow. The specified framework must be installed.

If no framework is specified, will default to the one currently installed. If no framework is specified and both frameworks are installed, will default to PyTorch.
args_parser (ArgumentHandler, optional, defaults to None) – Reference to the object in charge of parsing supplied pipeline parameters.
device (int, optional, defaults to -1) – Device ordinal for CPU/GPU supports. Setting this to -1 will leverage CPU, >=0 will run the model on the associated CUDA device id.

TextClassificationPipeline¶

class transformers.TextClassificationPipeline(return_all_scores: bool = False, **kwargs)[source]¶

Text classification pipeline using ModelForSequenceClassification head. See the sequence classification usage examples for more information.

This text classification pipeline can currently be loaded from the pipeline() method using the following task identifier(s):

“sentiment-analysis”, for classifying sequences according to positive or negative sentiments.

The models that this pipeline can use are models that have been fine-tuned on a sequence classification task. See the up-to-date list of available models on huggingface.co/models.

Parameters

model (PreTrainedModel or TFPreTrainedModel) – The model that will be used by the pipeline to make predictions. This needs to be a model inheriting from PreTrainedModel for PyTorch and TFPreTrainedModel for TensorFlow.
tokenizer (PreTrainedTokenizer) – The tokenizer that will be used by the pipeline to encode data for the model. This object inherits from PreTrainedTokenizer.
modelcard (str or ModelCard, optional, defaults to None) – Model card attributed to the model for this pipeline.
framework (str, optional, defaults to None) –
The framework to use, either “pt” for PyTorch or “tf” for TensorFlow. The specified framework must be installed.

If no framework is specified, will default to the one currently installed. If no framework is specified and both frameworks are installed, will default to PyTorch.
args_parser (ArgumentHandler, optional, defaults to None) – Reference to the object in charge of parsing supplied pipeline parameters.
device (int, optional, defaults to -1) – Device ordinal for CPU/GPU supports. Setting this to -1 will leverage CPU, >=0 will run the model on the associated CUDA device id.

QuestionAnsweringPipeline¶

class transformers.QuestionAnsweringPipeline(model: Union[PreTrainedModel, TFPreTrainedModel], tokenizer: transformers.tokenization_utils.PreTrainedTokenizer, modelcard: Optional[transformers.modelcard.ModelCard] = None, framework: Optional[str] = None, device: int = - 1, task: str = '', **kwargs)[source]¶

Question Answering pipeline using ModelForQuestionAnswering head. See the question answering usage examples for more information.

This question answering can currently be loaded from the pipeline() method using the following task identifier(s):

“question-answering”, for answering questions given a context.

The models that this pipeline can use are models that have been fine-tuned on a question answering task. See the up-to-date list of available models on huggingface.co/models.

Parameters

model (PreTrainedModel or TFPreTrainedModel) – The model that will be used by the pipeline to make predictions. This needs to be a model inheriting from PreTrainedModel for PyTorch and TFPreTrainedModel for TensorFlow.
tokenizer (PreTrainedTokenizer) – The tokenizer that will be used by the pipeline to encode data for the model. This object inherits from PreTrainedTokenizer.
modelcard (str or ModelCard, optional, defaults to None) – Model card attributed to the model for this pipeline.
framework (str, optional, defaults to None) –
The framework to use, either “pt” for PyTorch or “tf” for TensorFlow. The specified framework must be installed.

If no framework is specified, will default to the one currently installed. If no framework is specified and both frameworks are installed, will default to PyTorch.
args_parser (ArgumentHandler, optional, defaults to None) – Reference to the object in charge of parsing supplied pipeline parameters.
device (int, optional, defaults to -1) – Device ordinal for CPU/GPU supports. Setting this to -1 will leverage CPU, >=0 will run the model on the associated CUDA device id.

SummarizationPipeline¶

class transformers.SummarizationPipeline(**kwargs)[source]¶

Summarize news articles and other documents

Usage:

# use bart in pytorch
summarizer = pipeline("summarization")
summarizer("Sam Shleifer writes the best docstring examples in the whole world.", min_length=5, max_length=20)

# use t5 in tf
summarizer = pipeline("summarization", model="t5-base", tokenizer="t5-base", framework="tf")
summarizer("Sam Shleifer writes the best docstring examples in the whole world.", min_length=5, max_length=20)

The models that this pipeline can use are models that have been fine-tuned on a summarization task, which is currently, ‘bart-large-cnn’, ‘t5-small’, ‘t5-base’, ‘t5-large’, ‘t5-3b’, ‘t5-11b’. See the up-to-date list of available models on huggingface.co/models.

Parameters

model (str or PreTrainedModel or TFPreTrainedModel, optional, defaults to None) –
The model that will be used by the pipeline to make predictions. This can be None, a string checkpoint identifier or an actual pre-trained model inheriting from PreTrainedModel for PyTorch and TFPreTrainedModel for TensorFlow.

If None, the default of the pipeline will be loaded.
tokenizer (str or PreTrainedTokenizer, optional, defaults to None) –
The tokenizer that will be used by the pipeline to encode data for the model. This can be None, a string checkpoint identifier or an actual pre-trained tokenizer inheriting from PreTrainedTokenizer.

If None, the default of the pipeline will be loaded.
modelcard (str or ModelCard, optional, defaults to None) – Model card attributed to the model for this pipeline.
framework (str, optional, defaults to None) –
The framework to use, either “pt” for PyTorch or “tf” for TensorFlow. The specified framework must be installed.

If no framework is specified, will default to the one currently installed. If no framework is specified and both frameworks are installed, will default to PyTorch.
args_parser (ArgumentHandler, optional, defaults to None) – Reference to the object in charge of parsing supplied pipeline parameters.
device (int, optional, defaults to -1) – Device ordinal for CPU/GPU supports. Setting this to -1 will leverage CPU, >=0 will run the model on the associated CUDA device id.

TextGenerationPipeline¶

class transformers.TextGenerationPipeline(model: Union[PreTrainedModel, TFPreTrainedModel], tokenizer: transformers.tokenization_utils.PreTrainedTokenizer, modelcard: Optional[transformers.modelcard.ModelCard] = None, framework: Optional[str] = None, task: str = '', args_parser: transformers.pipelines.ArgumentHandler = None, device: int = - 1, binary_output: bool = False)[source]¶

Language generation pipeline using any ModelWithLMHead head. This pipeline predicts the words that will follow a specified text prompt.

This language generation pipeline can currently be loaded from the pipeline() method using the following task identifier(s):

“text-generation”, for generating text from a specified prompt.

The models that this pipeline can use are models that have been trained with an autoregressive language modeling objective, which includes the uni-directional models in the library (e.g. gpt2). See the list of available community models on huggingface.co/models.