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

class optimum.exporters.TasksManager

< source >

( )

Handles the task name -> model class and architecture -> configuration mappings.

create_register

< source >

( backend: str overwrite_existing: bool = False ) → Callable[[str, Tuple[str, ...]], Callable[[Type], Type]]

Parameters

backend (str) — The name of the backend that the register function will handle.
overwrite_existing (bool, defaults to False) — Whether or not the register function is allowed to overwrite an already existing config.

Returns

Callable[[str, Tuple[str, ...]], Callable[[Type], Type]]

A decorator taking the model type and a the supported tasks.

Creates a register function for the specified backend.

Example:

>>> register_for_new_backend = create_register("new-backend")

>>> @register_for_new_backend("bert", "text-classification", "token-classification")
>>> class BertNewBackendConfig(NewBackendConfig):
>>>     pass

determine_framework

< source >

( model_name_or_path: typing.Union[str, pathlib.Path] subfolder: str = '' revision: typing.Optional[str] = None cache_dir: str = '/root/.cache/huggingface/hub' token: typing.Union[bool, str, NoneType] = None ) → str

Parameters

model_name_or_path (Union[str, Path]) — Can be either the model id of a model repo on the Hugging Face Hub, or a path to a local directory containing a model.
subfolder (str, optional, defaults to "") — In case the model files are located inside a subfolder of the model directory / repo on the Hugging Face Hub, you can specify the subfolder name here.
revision (Optional[str], defaults to None) — Revision is the specific model version to use. It can be a branch name, a tag name, or a commit id.
cache_dir (Optional[str], optional) — Path to a directory in which a downloaded pretrained model weights have been cached if the standard cache should not be used.
token (Optional[Union[bool,str]], defaults to None) — The token to use as HTTP bearer authorization for remote files. If True, will use the token generated when running huggingface-cli login (stored in huggingface_hub.constants.HF_TOKEN_PATH).

Returns

str

The framework to use for the export.

Determines the framework to use for the export.

The priority is in the following order:

User input via framework.
If local checkpoint is provided, use the same framework as the checkpoint.
If model repo, try to infer the framework from the cache if available, else from the Hub.
If could not infer, use available framework in environment, with priority given to PyTorch.

get_all_tasks

< source >

( ) → List

Returns

List

all the possible tasks.

Retrieves all the possible tasks.

get_exporter_config_constructor

< source >

( exporter: str model: typing.Union[ForwardRef('PreTrainedModel'), ForwardRef('TFPreTrainedModel'), NoneType] = None task: str = 'feature-extraction' model_type: typing.Optional[str] = None model_name: typing.Optional[str] = None exporter_config_kwargs: typing.Union[typing.Dict[str, typing.Any], NoneType] = None library_name: typing.Optional[str] = None ) → ExportConfigConstructor

Parameters

exporter (str) — The exporter to use.
model (Optional[Union[PreTrainedModel, TFPreTrainedModel]], defaults to None) — The instance of the model.
task (str, defaults to "feature-extraction") — The task to retrieve the config for.
model_type (Optional[str], defaults to None) — The model type to retrieve the config for.
model_name (Optional[str], defaults to None) — The name attribute of the model object, only used for the exception message.
exporter_config_kwargs (Optional[Dict[str, Any]], defaults to None) — Arguments that will be passed to the exporter config class when building the config constructor.
library_name (Optional[str], defaults to None) — The library name of the model. Can be any of “transformers”, “timm”, “diffusers”, “sentence_transformers”.

Returns

ExportConfigConstructor

The ExportConfig constructor for the requested backend.

Gets the ExportConfigConstructor for a model (or alternatively for a model type) and task combination.

get_model_class_for_task

< source >

( task: str framework: str = 'pt' model_type: typing.Optional[str] = None model_class_name: typing.Optional[str] = None library: str = 'transformers' )

Parameters

task (str) — The task required.
framework (str, defaults to "pt") — The framework to use for the export.
model_type (Optional[str], defaults to None) — The model type to retrieve the model class for. Some architectures need a custom class to be loaded, and can not be loaded from auto class.
model_class_name (Optional[str], defaults to None) — A model class name, allowing to override the default class that would be detected for the task. This parameter is useful for example for “automatic-speech-recognition”, that may map to AutoModelForSpeechSeq2Seq or to AutoModelForCTC.
library (str, defaults to transformers) — The library name of the model. Can be any of “transformers”, “timm”, “diffusers”, “sentence_transformers”.

Attempts to retrieve an AutoModel class from a task name.

get_model_from_task

< source >

( task: str model_name_or_path: typing.Union[str, pathlib.Path] subfolder: str = '' revision: typing.Optional[str] = None cache_dir: str = '/root/.cache/huggingface/hub' token: typing.Union[bool, str, NoneType] = None framework: typing.Optional[str] = None torch_dtype: typing.Optional[ForwardRef('torch.dtype')] = None device: typing.Union[ForwardRef('torch.device'), str, NoneType] = None library_name: typing.Optional[str] = None **model_kwargs )

Parameters

task (str) — The task required.
model_name_or_path (Union[str, Path]) — Can be either the model id of a model repo on the Hugging Face Hub, or a path to a local directory containing a model.
subfolder (str, defaults to "") — In case the model files are located inside a subfolder of the model directory / repo on the Hugging Face Hub, you can specify the subfolder name here.
revision (Optional[str], optional) — Revision is the specific model version to use. It can be a branch name, a tag name, or a commit id.
cache_dir (Optional[str], optional) — Path to a directory in which a downloaded pretrained model weights have been cached if the standard cache should not be used.
token (Optional[Union[bool,str]], defaults to None) — The token to use as HTTP bearer authorization for remote files. If True, will use the token generated when running huggingface-cli login (stored in huggingface_hub.constants.HF_TOKEN_PATH).
framework (Optional[str], optional) — The framework to use for the export. See TasksManager.determine_framework for the priority should none be provided.
torch_dtype (Optional[torch.dtype], defaults to None) — Data type to load the model on. PyTorch-only argument.
device (Optional[torch.device], defaults to None) — Device to initialize the model on. PyTorch-only argument. For PyTorch, defaults to “cpu”.
library_name (Optional[str], defaults to None) — The library name of the model. Can be any of “transformers”, “timm”, “diffusers”, “sentence_transformers”. See TasksManager.infer_library_from_model for the priority should none be provided.
model_kwargs (Dict[str, Any], optional) — Keyword arguments to pass to the model .from_pretrained() method.

Retrieves a model from its name and the task to be enabled.

get_supported_model_type_for_task

< source >

( task: str exporter: str )

Returns the list of supported architectures by the exporter for a given task. Transformers-specific.

get_supported_tasks_for_model_type

< source >

( model_type: str exporter: str model_name: typing.Optional[str] = None library_name: typing.Optional[str] = None ) → TaskNameToExportConfigDict

Parameters

model_type (str) — The model type to retrieve the supported tasks for.
exporter (str) — The name of the exporter.
model_name (Optional[str], defaults to None) — The name attribute of the model object, only used for the exception message.
library_name (Optional[str], defaults to None) — The library name of the model. Can be any of “transformers”, “timm”, “diffusers”, “sentence_transformers”.

Returns

TaskNameToExportConfigDict

The dictionary mapping each task to a corresponding ExportConfig constructor.

Retrieves the task -> exporter backend config constructors map from the model type.

infer_library_from_model

< source >

( model: typing.Union[str, ForwardRef('PreTrainedModel'), ForwardRef('TFPreTrainedModel'), ForwardRef('DiffusionPipeline'), typing.Type] subfolder: str = '' revision: typing.Optional[str] = None cache_dir: str = '/root/.cache/huggingface/hub' token: typing.Union[bool, str, NoneType] = None ) → str

Parameters

model (Union[str, PreTrainedModel, TFPreTrainedModel, DiffusionPipeline, Type]) — The model to infer the task from. This can either be the name of a repo on the HuggingFace Hub, an instance of a model, or a model class.
subfolder (str, defaults to "") — In case the model files are located inside a subfolder of the model directory / repo on the Hugging Face Hub, you can specify the subfolder name here.
revision (Optional[str], optional, defaults to None) — Revision is the specific model version to use. It can be a branch name, a tag name, or a commit id.
cache_dir (Optional[str], optional) — Path to a directory in which a downloaded pretrained model weights have been cached if the standard cache should not be used.
token (Optional[Union[bool,str]], defaults to None) — The token to use as HTTP bearer authorization for remote files. If True, will use the token generated when running huggingface-cli login (stored in huggingface_hub.constants.HF_TOKEN_PATH).

Returns

str

The library name automatically detected from the model repo, model instance, or model class.

Infers the library from the model repo, model instance, or model class.

infer_task_from_model

< source >

Parameters

model (Union[str, PreTrainedModel, TFPreTrainedModel, DiffusionPipeline, Type]) — The model to infer the task from. This can either be the name of a repo on the HuggingFace Hub, an instance of a model, or a model class.
subfolder (str, optional, defaults to "") — In case the model files are located inside a subfolder of the model directory / repo on the Hugging Face Hub, you can specify the subfolder name here.
revision (Optional[str], defaults to None) — Revision is the specific model version to use. It can be a branch name, a tag name, or a commit id.
cache_dir (Optional[str], optional) — Path to a directory in which a downloaded pretrained model weights have been cached if the standard cache should not be used.
token (Optional[Union[bool,str]], defaults to None) — The token to use as HTTP bearer authorization for remote files. If True, will use the token generated when running huggingface-cli login (stored in huggingface_hub.constants.HF_TOKEN_PATH).

Returns

str

The task name automatically detected from the HF hub repo, model instance, or model class.

Infers the task from the model repo, model instance, or model class.

standardize_model_attributes

< source >

( model: typing.Union[ForwardRef('PreTrainedModel'), ForwardRef('TFPreTrainedModel'), ForwardRef('DiffusionPipeline')] )

Parameters

model (Union[PreTrainedModel, TFPreTrainedModel, DiffusionPipeline]) — The instance of the model.

Updates the model for export. This function is suitable to make required changes to the models from different libraries to follow transformers style.

< > Update on GitHub

Optimum

The Tasks Manager

Task names

PyTorch

TensorFlow

Reference

class optimum.exporters.TasksManager

create_register

determine_framework

get_all_tasks

get_exporter_config_constructor

get_model_class_for_task

get_model_from_task

get_supported_model_type_for_task

get_supported_tasks_for_model_type

infer_library_from_model

infer_task_from_model

standardize_model_attributes