task_manager.mdx

The Tasks Manager

Exporting a model from one framework to some format (also called backend here) involves specifying inputs and outputs information that the export function needs. The way optimum.exporters is structured for each backend is as follows:

• Configuration classes containing the information for each model to perform the export.
• Exporting functions using the proper configuration for the model to export.

The role of the [~optimum.exporters.tasks.TasksManager] is to be the main entry-point to load a model given a name and a task, and to get the proper configuration for a given (architecture, backend) couple. That way, there is a centralized place to register the task -> model class and (architecture, backend) -> configuration mappings. This allows the export functions to use this, and to rely on the various checks it provides.

Task names

The tasks supported might depend on the backend, but here are the mappings between a task name and the auto class for both PyTorch and TensorFlow.

It is possible to know which tasks are supported for a model for a given backend, by doing:

>>> from optimum.exporters.tasks import TasksManager

>>> model_type = "distilbert"
>>> # For instance, for the ONNX export.
>>> backend = "onnx"
>>> distilbert_tasks = list(TasksManager.get_supported_tasks_for_model_type(model_type, backend).keys())

>>> print(distilbert_tasks)
['default', 'fill-mask', 'text-classification', 'multiple-choice', 'token-classification', 'question-answering']

PyTorch

Task	Auto Class
text-generation, text-generation-with-past	AutoModelForCausalLM
feature-extraction, feature-extraction-with-past	AutoModel
fill-mask	AutoModelForMaskedLM
question-answering	AutoModelForQuestionAnswering
text2text-generation, text2text-generation-with-past	AutoModelForSeq2SeqLM
text-classification	AutoModelForSequenceClassification
token-classification	AutoModelForTokenClassification
multiple-choice	AutoModelForMultipleChoice
image-classification	AutoModelForImageClassification
object-detection	AutoModelForObjectDetection
image-segmentation	AutoModelForImageSegmentation
masked-im	AutoModelForMaskedImageModeling
semantic-segmentation	AutoModelForSemanticSegmentation
automatic-speech-recognition	AutoModelForSpeechSeq2Seq

TensorFlow

Task	Auto Class
text-generation, text-generation-with-past	TFAutoModelForCausalLM
default, default-with-past	TFAutoModel
fill-mask	TFAutoModelForMaskedLM
question-answering	TFAutoModelForQuestionAnswering
text2text-generation, text2text-generation-with-past	TFAutoModelForSeq2SeqLM
text-classification	TFAutoModelForSequenceClassification
token-classification	TFAutoModelForTokenClassification
multiple-choice	TFAutoModelForMultipleChoice
semantic-segmentation	TFAutoModelForSemanticSegmentation

Reference

[[autodoc]] exporters.tasks.TasksManager