任务管理器

将模型从一个框架导出为某种格式（在这里也称为后端）涉及到指定导出函数所需的输入和输出信息。optimum.exporters 为每个后端构建的方式如下：

包含每个模型执行导出所需信息的配置类。
使用适用于要导出的模型的正确配置的导出函数。

TasksManager 的作用是作为主入口点，根据名称和任务加载模型，并为给定的（架构、后端）组合获取正确的配置。这样，就有一个中心位置来注册 task -> model class 和 (architecture, backend) -> configuration 映射。这允许导出函数使用它，并依赖它提供的各种检查。

任务名称

支持的任务可能取决于后端，但以下是任务名称与 PyTorch 和 TensorFlow 的自动类之间的映射。

可以通过执行以下操作来了解给定后端的模型支持哪些任务：

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

任务	自动类
`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

任务	自动类
`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`

参考

class optimum.exporters.TasksManager

< source >

( )

处理 task name -> model class 和 architecture -> configuration 映射。

create_register

< source >

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

参数

backend (str) — 注册函数将处理的后端名称。
overwrite_existing (bool, 默认为 False) — 注册函数是否允许覆盖已存在的配置。

返回值

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

一个装饰器，接受模型类型和支持的任务。

为指定的后端创建注册函数。

示例

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

参数

model_name_or_path (Union[str, Path]) — 可以是 Hugging Face Hub 上模型仓库的模型 ID，也可以是包含模型的本地目录的路径。
subfolder (str, 可选，默认为 "") — 如果模型文件位于模型目录/Hugging Face Hub 上的仓库的子文件夹中，您可以在此处指定子文件夹名称。
revision (Optional[str], 默认为 None) — Revision 是要使用的特定模型版本。它可以是分支名称、标签名称或提交 ID。
cache_dir (Optional[str], 可选) — 缓存下载的预训练模型权重的目录路径，如果不想使用标准缓存。
token (Optional[Union[bool,str]], 默认为 None) — 用作远程文件 HTTP Bearer 授权的令牌。如果为 True，将使用运行 huggingface-cli login 时生成的令牌（存储在 huggingface_hub.constants.HF_TOKEN_PATH 中）。

返回值

str

用于导出的框架。

确定用于导出的框架。

优先级顺序如下：

用户通过 framework 输入。
如果提供本地检查点，则使用与检查点相同的框架。
如果是模型仓库，尝试从缓存（如果可用）或 Hub 推断框架。
如果无法推断，则使用环境中可用的框架，优先考虑 PyTorch。

get_all_tasks

< source >

( ) → List

返回值

List

所有可能的任务。

检索所有可能的任务。

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.Optional[typing.Dict[str, typing.Any]] = None library_name: typing.Optional[str] = None ) → ExportConfigConstructor

参数

exporter (str) — 要使用的导出器。
model (Optional[Union[PreTrainedModel, TFPreTrainedModel]], 默认为 None) — 模型的实例。
task (str, 默认为 "feature-extraction") — 要检索配置的任务。
model_type (Optional[str], 默认为 None) — 要检索配置的模型类型。
model_name (Optional[str], 默认为 None) — 模型对象的名称属性，仅用于异常消息。
exporter_config_kwargs (Optional[Dict[str, Any]], defaults to None) — 构建配置构造器时将传递给导出器配置类别的参数。
library_name (Optional[str], defaults to None) — 模型的库名称。可以是 “transformers”、“timm”、“diffusers”、“sentence_transformers” 中的任何一个。

返回值

ExportConfigConstructor

针对所请求后端的 ExportConfig 构造器。

获取模型（或模型类型）和任务组合的 ExportConfigConstructor。

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

参数

task (str) — 所需的任务。
framework (str, defaults to "pt") — 用于导出的框架。
model_type (Optional[str], defaults to None) — 要检索模型类别的模型类型。某些架构需要加载自定义类别，并且无法从自动类别加载。
model_class_name (Optional[str], defaults to None) — 模型类别名称，允许覆盖将为任务检测到的默认类别。此参数对于例如 “automatic-speech-recognition” 非常有用，它可能映射到 AutoModelForSpeechSeq2Seq 或 AutoModelForCTC。
library (str, defaults to transformers) — 模型的库名称。可以是 “transformers”、“timm”、“diffusers”、“sentence_transformers” 中的任何一个。

尝试从任务名称检索 AutoModel 类别。

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 )

参数

task (str) — 所需的任务。
model_name_or_path (Union[str, Path]) — 可以是 Hugging Face Hub 上模型仓库的模型 ID，也可以是包含模型的本地目录路径。
subfolder (str, defaults to "") — 如果模型文件位于模型目录/Hugging Face Hub 上的仓库的子文件夹中，您可以在此处指定子文件夹名称。
revision (Optional[str], optional) — Revision 是要使用的特定模型版本。它可以是分支名称、标签名称或提交 ID。
cache_dir (Optional[str], optional) — 缓存下载的预训练模型权重的目录路径（如果应避免使用标准缓存）。
token (Optional[Union[bool,str]], defaults to None) — 用作远程文件的 HTTP Bearer 授权的令牌。如果为 True，将使用运行 huggingface-cli login 时生成的令牌（存储在 huggingface_hub.constants.HF_TOKEN_PATH 中）。
framework (Optional[str], optional) — 用于导出的框架。请参阅 TasksManager.determine_framework 以了解在未提供框架时应遵循的优先级。
torch_dtype (Optional[torch.dtype], defaults to None) — 用于加载模型的数据类型。仅限 PyTorch 参数。
device (Optional[torch.device], defaults to None) — 用于初始化模型的设备。仅限 PyTorch 参数。对于 PyTorch，默认为 “cpu”。
library_name (Optional[str], defaults to None) — 模型的库名称。可以是 “transformers”、“timm”、“diffusers”、“sentence_transformers” 中的任何一个。请参阅 TasksManager.infer_library_from_model 以了解在未提供库名称时应遵循的优先级。
model_kwargs (Dict[str, Any], optional) — 要传递给模型 .from_pretrained() 方法的关键字参数。
library_name (Optional[str], defaults to None) — 模型的库名称。可以是 “transformers”、“timm”、“diffusers”、“sentence_transformers” 中的任何一个。请参阅 TasksManager.infer_library_from_model 以了解在未提供库名称时应遵循的优先级。

从模型名称和要启用的任务中检索模型。

get_supported_model_type_for_task

< source >

( task: str exporter: str )

返回给定任务的导出器支持的架构列表。 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

参数

model_type (str) — 要检索支持任务的模型类型。
exporter (str) — 导出器的名称。
model_name (Optional[str], defaults to None) — 模型对象的名称属性，仅用于异常消息。
library_name (Optional[str], defaults to None) — 模型的库名称。可以是 “transformers”、“timm”、“diffusers”、“sentence_transformers” 中的任何一个。

返回值

TaskNameToExportConfigDict

将每个任务映射到相应的 ExportConfig 构造函数的字典。

从模型类型中检索 task -> exporter backend config constructors 映射。

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

参数

model (Union[str, PreTrainedModel, TFPreTrainedModel, DiffusionPipeline, Type]) — 要从中推断任务的模型。这可以是 HuggingFace Hub 上的 repo 名称、模型实例或模型类。
subfolder (str, defaults to "") — 如果模型文件位于 Hugging Face Hub 上的模型目录/repo 的子文件夹中，您可以在此处指定子文件夹名称。
revision (Optional[str], optional, defaults to None) — Revision 是要使用的特定模型版本。它可以是分支名称、标签名称或提交 ID。
cache_dir (Optional[str], optional) — 缓存下载的预训练模型权重的目录路径，如果不想使用标准缓存。
token (Optional[Union[bool,str]], defaults to None) — 用作远程文件的 HTTP Bearer 授权的令牌。如果为 True，将使用运行 huggingface-cli login 时生成的令牌（存储在 huggingface_hub.constants.HF_TOKEN_PATH 中）。

返回值

str

从模型仓库、模型实例或模型类自动检测到的库名称。

从模型仓库、模型实例或模型类推断库。

infer_task_from_model

< source >

参数

model (Union[str, PreTrainedModel, TFPreTrainedModel, DiffusionPipeline, Type]) — 要从中推断任务的模型。这可以是 HuggingFace Hub 上的 repo 名称、模型实例或模型类。
subfolder (str, optional, defaults to "") — 如果模型文件位于 Hugging Face Hub 上的模型目录/repo 的子文件夹中，您可以在此处指定子文件夹名称。
revision (Optional[str], defaults to None) — Revision 是要使用的特定模型版本。它可以是分支名称、标签名称或提交 ID。
cache_dir (Optional[str], optional) — 缓存下载的预训练模型权重的目录路径，如果不想使用标准缓存。
token (Optional[Union[bool,str]], defaults to None) — 用作远程文件的 HTTP Bearer 授权的令牌。如果为 True，将使用运行 huggingface-cli login 时生成的令牌（存储在 huggingface_hub.constants.HF_TOKEN_PATH 中）。
library_name (Optional[str], defaults to None) — 模型的库名称。可以是 “transformers”、“timm”、“diffusers”、“sentence_transformers” 中的任何一个。有关在未提供任何库名称时应遵循的优先级，请参阅 TasksManager.infer_library_from_model。

返回值

str

从 HF hub repo、模型实例或模型类自动检测的任务名称。

从模型仓库、模型实例或模型类推断任务。

standardize_model_attributes

< source >

( model: typing.Union[ForwardRef('PreTrainedModel'), ForwardRef('TFPreTrainedModel'), ForwardRef('DiffusionPipeline')] library_name: typing.Optional[str] = None )

参数

model (Union[PreTrainedModel, TFPreTrainedModel, DiffusionPipeline]) — 模型的实例。

更新模型以进行导出。此功能适用于对来自不同库的模型进行所需更改，以遵循 transformers 样式。

< > Update on GitHub

Optimum

任务管理器

任务名称

PyTorch

TensorFlow

参考

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