任務管理器

將模型從一個框架匯出到某種格式（此處也稱為後端）需要指定匯出功能所需的輸入和輸出資訊。`optimum.exporters` 為每個後端構建的方式如下：

包含每個模型執行匯出所需資訊的配置類。
使用適當的模型匯出配置進行匯出。

TasksManager 的作用是作為主要入口點，根據名稱和任務載入模型，並獲取給定（架構、後端）組合的適當配置。這樣，就有一個集中的位置來註冊 `任務 -> 模型類` 和 `（架構、後端）-> 配置` 對映。這允許匯出功能利用此功能，並依賴其提供的各種檢查。

任務名稱

支援的任務可能取決於後端，但這裡是 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`	`自動模型`
`fill-mask`	`AutoModelForMaskedLM`
`問題回答`	`AutoModelForQuestionAnswering`
`text2text-generation`, `text2text-generation-with-past`	`AutoModelForSeq2SeqLM`
`文字分類`	`AutoModelForSequenceClassification`
`詞元分類`	`AutoModelForTokenClassification`
`multiple-choice`	`AutoModelForMultipleChoice`
`image-classification`	`AutoModelForImageClassification`
`object-detection`	`AutoModelForObjectDetection`
`image-segmentation`	`AutoModelForImageSegmentation`
`masked-im`	`AutoModelForMaskedImageModeling`
`semantic-segmentation`	`AutoModelForSemanticSegmentation`
`自動語音識別`	`AutoModelForSpeechSeq2Seq`

TensorFlow

任務	自動類
`text-generation`, `text-generation-with-past`	`TFAutoModelForCausalLM`
`default`, `default-with-past`	`TFAutoModel`
`fill-mask`	`TFAutoModelForMaskedLM`
`問題回答`	`TFAutoModelForQuestionAnswering`
`text2text-generation`, `text2text-generation-with-past`	`TFAutoModelForSeq2SeqLM`
`文字分類`	`TFAutoModelForSequenceClassification`
`詞元分類`	`TFAutoModelForTokenClassification`
`multiple-choice`	`TFAutoModelForMultipleChoice`
`semantic-segmentation`	`TFAutoModelForSemanticSegmentation`

參考

class optimum.exporters.TasksManager

< 源 >

( )

處理 `任務名稱 -> 模型類` 和 `架構 -> 配置` 對映。

create_register

< 源 >

( 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

< 源 >

( 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[str, bool, NoneType] = None ) → str

引數

model_name_or_path (Union[str, Path]) — 可以是 Hugging Face Hub 上模型倉庫的模型 ID，也可以是包含模型的本地目錄路徑。
subfolder (str, 可選, 預設為 "") — 如果模型檔案位於 Hugging Face Hub 上的模型目錄/倉庫的子資料夾中，您可以在此處指定子資料夾名稱。
revision (Optional[str], 可選, 預設為 None) — 版本是要使用的特定模型版本。它可以是分支名稱、標籤名稱或提交 ID。
cache_dir (Optional[str], 可選) — 如果不應使用標準快取，則下載的預訓練模型權重已快取到的目錄路徑。
token (Optional[Union[bool,str]], 預設為 None) — 用作遠端檔案 HTTP 承載授權的令牌。如果為 True，將使用執行 huggingface-cli login 時生成的令牌（儲存在 huggingface_hub.constants.HF_TOKEN_PATH 中）。

字串

用於匯出的框架。

確定用於匯出的框架。

優先順序順序如下：

透過 `framework` 的使用者輸入。
如果提供了本地檢查點，則使用與檢查點相同的框架。
如果模型倉庫可用，則嘗試從快取中推斷框架，否則從 Hub 中推斷。
如果無法推斷，則使用環境中可用的框架，PyTorch 優先。

get_all_tasks

< 源 >

( ) → List

列表

所有可能的任務。

檢索所有可能的任務。

get_exporter_config_constructor

< 源 >

( 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]], 預設為 None) — 在構建配置建構函式時將傳遞給匯出器配置類的引數。
library_name (Optional[str], 預設為 None) — 模型的庫名稱。可以是“transformers”、“timm”、“diffusers”、“sentence_transformers”中的任意一個。

ExportConfigConstructor

請求後端的 `ExporterConfig` 建構函式。

獲取模型（或模型型別）和任務組合的 `ExportConfigConstructor`。

get_model_class_for_task

< 源 >

( 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, 預設為 "pt") — 用於匯出的框架。
model_type (Optional[str], 預設為 None) — 用於檢索模型類的模型型別。某些架構需要載入自定義類，不能從自動類載入。
model_class_name (Optional[str], 預設為 None) — 模型類名稱，允許覆蓋任務的預設檢測類。例如，此引數對於“自動語音識別”很有用，它可能對映到 AutoModelForSpeechSeq2Seq 或 AutoModelForCTC。
library (str, 預設為 transformers) — 模型的庫名稱。可以是“transformers”、“timm”、“diffusers”、“sentence_transformers”中的任意一個。

嘗試從任務名稱中檢索 AutoModel 類。

get_model_from_task

< 源 >

( 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[str, bool, 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, 預設為 "") — 如果模型檔案位於 Hugging Face Hub 上的模型目錄/倉庫的子資料夾中，您可以在此處指定子資料夾名稱。
revision (Optional[str], 可選) — 版本是要使用的特定模型版本。它可以是分支名稱、標籤名稱或提交 ID。
cache_dir (Optional[str], 可選) — 如果不應使用標準快取，則下載的預訓練模型權重已快取到的目錄路徑。
token (Optional[Union[bool,str]], 預設為 None) — 用作遠端檔案 HTTP 承載授權的令牌。如果為 True，將使用執行 huggingface-cli login 時生成的令牌（儲存在 huggingface_hub.constants.HF_TOKEN_PATH 中）。
framework (Optional[str], 可選) — 用於匯出的框架。有關未提供框架時的優先順序，請參見 TasksManager.determine_framework。
torch_dtype (Optional[torch.dtype], 預設為 None) — 載入模型的資料型別。僅限 PyTorch 引數。
device (Optional[torch.device], 預設為 None) — 初始化模型的裝置。僅限 PyTorch 引數。對於 PyTorch，預設為“cpu”。
library_name (Optional[str], 預設為 None) — 模型的庫名稱。可以是“transformers”、“timm”、“diffusers”、“sentence_transformers”中的任意一個。有關未提供庫時的優先順序，請參見 TasksManager.infer_library_from_model。
model_kwargs (Dict[str, Any], 可選) — 傳遞給模型 .from_pretrained() 方法的關鍵字引數。
library_name (Optional[str], 預設為 None) — 模型的庫名稱。可以是“transformers”、“timm”、“diffusers”、“sentence_transformers”中的任意一個。有關未提供庫時的優先順序，請參見 TasksManager.infer_library_from_model。

根據其名稱和要啟用的任務檢索模型。

get_supported_model_type_for_task

< 源 >

( task: str exporter: str )

返回匯出器針對給定任務支援的架構列表。Transformers 特有。

get_supported_tasks_for_model_type

< 源 >

( 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], 預設為 None) — 模型物件的名稱屬性，僅用於異常訊息。
library_name (Optional[str], 預設為 None) — 模型的庫名稱。可以是“transformers”、“timm”、“diffusers”、“sentence_transformers”中的任意一個。

TaskNameToExportConfigDict

將每個任務對映到相應的 `ExporterConfig` 建構函式的字典。

從模型型別檢索 `任務 -> 匯出器後端配置建構函式` 對映。

infer_library_from_model

< 源 >

( 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[str, bool, NoneType] = None ) → str

引數

model (Union[str, PreTrainedModel, TFPreTrainedModel, DiffusionPipeline, Type]) — 要推斷任務的模型。這可以是 HuggingFace Hub 上的倉庫名稱、模型例項或模型類。
subfolder (str, 預設為 "") — 如果模型檔案位於 Hugging Face Hub 上的模型目錄/倉庫的子資料夾中，您可以在此處指定子資料夾名稱。
revision (Optional[str], optional, 預設為 None) — revision 是要使用的特定模型版本。它可以是分支名稱、標籤名稱或提交ID。
cache_dir (Optional[str], optional) — 如果不使用標準快取，則為下載的預訓練模型權重快取目錄的路徑。
token (Optional[Union[bool,str]], 預設為 None) — 用作遠端檔案HTTP bearer授權的token。如果為True，將使用執行huggingface-cli login時生成的token（儲存在huggingface_hub.constants.HF_TOKEN_PATH中）。

字串

根據模型倉庫、模型例項或模型類自動檢測到的庫名稱。

從模型倉庫、模型例項或模型類推斷庫。

infer_task_from_model

< source >

引數

model (Union[str, PreTrainedModel, TFPreTrainedModel, DiffusionPipeline, Type]) — 要推斷任務的模型。這可以是HuggingFace Hub上的倉庫名稱、模型例項或模型類。
subfolder (str, optional, 預設為 "") — 如果模型檔案位於Hugging Face Hub上的模型目錄/倉庫的子資料夾中，可以在此處指定子資料夾名稱。
revision (Optional[str], 預設為 None) — revision 是要使用的特定模型版本。它可以是分支名稱、標籤名稱或提交ID。
cache_dir (Optional[str], optional) — 如果不使用標準快取，則為下載的預訓練模型權重快取目錄的路徑。
token (Optional[Union[bool,str]], 預設為 None) — 用作遠端檔案HTTP bearer授權的token。如果為True，將使用執行huggingface-cli login時生成的token（儲存在huggingface_hub.constants.HF_TOKEN_PATH中）。
library_name (Optional[str], 預設為 None) — 模型的庫名稱。可以是“transformers”、“timm”、“diffusers”、“sentence_transformers”中的任何一個。如果未提供，請參見TasksManager.infer_library_from_model瞭解優先順序。

字串

從HF hub倉庫、模型例項或模型類自動檢測到的任務名稱。

從模型倉庫、模型例項或模型類推斷任務。

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風格。

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