Skip to content

GitLab

Unverified Commit ddb1a47e authored by Sylvain Gugger's avatar Sylvain Gugger Committed by GitHub
Browse files

Automatically sort auto mappings (#17250) 

* Automatically sort auto mappings

* Better class extraction

* Some auto class magic

* Adapt test and underlying behavior

* Remove re-used config

* Quality
parent 2f611f85
Hide whitespace changes
Inline Side-by-side
Showing with 1096 additions and 991 deletions
.circleci/config.yml
View file @ ddb1a47e
...@@ -857,6 +857,7 @@ jobs: ...@@ -857,6 +857,7 @@ jobs:
- run: black --check --preview examples tests src utils - run: black --check --preview examples tests src utils
- run: isort --check-only examples tests src utils - run: isort --check-only examples tests src utils
- run: python utils/custom_init_isort.py --check_only - run: python utils/custom_init_isort.py --check_only
- run: python utils/sort_auto_mappings.py --check_only
- run: flake8 examples tests src utils - run: flake8 examples tests src utils
- run: doc-builder style src/transformers docs/source --max_len 119 --check_only --path_to_docs docs/source - run: doc-builder style src/transformers docs/source --max_len 119 --check_only --path_to_docs docs/source
......
Makefile
View file @ ddb1a47e
...@@ -48,6 +48,7 @@ quality: ...@@ -48,6 +48,7 @@ quality:
black --check --preview $(check_dirs) black --check --preview $(check_dirs)
isort --check-only $(check_dirs) isort --check-only $(check_dirs)
python utils/custom_init_isort.py --check_only python utils/custom_init_isort.py --check_only
python utils/sort_auto_mappings.py --check_only
flake8 $(check_dirs) flake8 $(check_dirs)
doc-builder style src/transformers docs/source --max_len 119 --check_only --path_to_docs docs/source doc-builder style src/transformers docs/source --max_len 119 --check_only --path_to_docs docs/source
...@@ -55,6 +56,7 @@ quality: ...@@ -55,6 +56,7 @@ quality:
extra_style_checks: extra_style_checks:
python utils/custom_init_isort.py python utils/custom_init_isort.py
python utils/sort_auto_mappings.py
doc-builder style src/transformers docs/source --max_len 119 --path_to_docs docs/source doc-builder style src/transformers docs/source --max_len 119 --path_to_docs docs/source
# this target runs checks on all files and potentially modifies some of them # this target runs checks on all files and potentially modifies some of them
......
docs/source/en/index.mdx
View file @ ddb1a47e
...@@ -259,7 +259,6 @@ Flax), PyTorch, and/or TensorFlow. ...@@ -259,7 +259,6 @@ Flax), PyTorch, and/or TensorFlow.
| Swin | ❌ | ❌ | ✅ | ❌ | ❌ | | Swin | ❌ | ❌ | ✅ | ❌ | ❌ |
| T5 | ✅ | ✅ | ✅ | ✅ | ✅ | | T5 | ✅ | ✅ | ✅ | ✅ | ✅ |
| TAPAS | ✅ | ❌ | ✅ | ✅ | ❌ | | TAPAS | ✅ | ❌ | ✅ | ✅ | ❌ |
| TAPEX | ✅ | ✅ | ✅ | ✅ | ✅ |
| Transformer-XL | ✅ | ❌ | ✅ | ✅ | ❌ | | Transformer-XL | ✅ | ❌ | ✅ | ✅ | ❌ |
| TrOCR | ❌ | ❌ | ✅ | ❌ | ❌ | | TrOCR | ❌ | ❌ | ✅ | ❌ | ❌ |
| UniSpeech | ❌ | ❌ | ✅ | ❌ | ❌ | | UniSpeech | ❌ | ❌ | ✅ | ❌ | ❌ |
......
docs/source/en/serialization.mdx
View file @ ddb1a47e
...@@ -74,7 +74,6 @@ Ready-made configurations include the following architectures: ...@@ -74,7 +74,6 @@ Ready-made configurations include the following architectures:
- RoBERTa - RoBERTa
- RoFormer - RoFormer
- T5 - T5
- TAPEX
- ViT - ViT
- XLM-RoBERTa - XLM-RoBERTa
- XLM-RoBERTa-XL - XLM-RoBERTa-XL
......
src/transformers/models/auto/auto_factory.py
View file @ ddb1a47e
...@@ -560,10 +560,17 @@ class _LazyAutoMapping(OrderedDict): ...@@ -560,10 +560,17 @@ class _LazyAutoMapping(OrderedDict):
if key in self._extra_content: if key in self._extra_content:
return self._extra_content[key] return self._extra_content[key]
model_type = self._reverse_config_mapping[key.__name__] model_type = self._reverse_config_mapping[key.__name__]
if model_type not in self._model_mapping: if model_type in self._model_mapping:
raise KeyError(key) model_name = self._model_mapping[model_type]
model_name = self._model_mapping[model_type] return self._load_attr_from_module(model_type, model_name)
return self._load_attr_from_module(model_type, model_name)
# Maybe there was several model types associated with this config.
model_types = [k for k, v in self._config_mapping.items() if v == key.__name__]
for mtype in model_types:
if mtype in self._model_mapping:
model_name = self._model_mapping[mtype]
return self._load_attr_from_module(mtype, model_name)
raise KeyError(key)
def _load_attr_from_module(self, model_type, attr): def _load_attr_from_module(self, model_type, attr):
module_name = model_type_to_module_name(model_type) module_name = model_type_to_module_name(model_type)
......
src/transformers/models/auto/configuration_auto.py
View file @ ddb1a47e
...@@ -29,339 +29,338 @@ logger = logging.get_logger(__name__) ...@@ -29,339 +29,338 @@ logger = logging.get_logger(__name__)
CONFIG_MAPPING_NAMES = OrderedDict( CONFIG_MAPPING_NAMES = OrderedDict(
[ [
# Add configs here # Add configs here
("yolos", "YolosConfig"), ("albert", "AlbertConfig"),
("tapex", "BartConfig"), ("bart", "BartConfig"),
("dpt", "DPTConfig"),
("decision_transformer", "DecisionTransformerConfig"),
("glpn", "GLPNConfig"),
("maskformer", "MaskFormerConfig"),
("decision_transformer", "DecisionTransformerConfig"),
("poolformer", "PoolFormerConfig"),
("convnext", "ConvNextConfig"),
("van", "VanConfig"),
("resnet", "ResNetConfig"),
("regnet", "RegNetConfig"),
("yoso", "YosoConfig"),
("swin", "SwinConfig"),
("vilt", "ViltConfig"),
("vit_mae", "ViTMAEConfig"),
("realm", "RealmConfig"),
("nystromformer", "NystromformerConfig"),
("xglm", "XGLMConfig"),
("imagegpt", "ImageGPTConfig"),
("qdqbert", "QDQBertConfig"),
("vision-encoder-decoder", "VisionEncoderDecoderConfig"),
("trocr", "TrOCRConfig"),
("fnet", "FNetConfig"),
("segformer", "SegformerConfig"),
("vision-text-dual-encoder", "VisionTextDualEncoderConfig"),
("perceiver", "PerceiverConfig"),
("gptj", "GPTJConfig"),
("layoutlmv2", "LayoutLMv2Config"),
("plbart", "PLBartConfig"),
("beit", "BeitConfig"), ("beit", "BeitConfig"),
("data2vec-vision", "Data2VecVisionConfig"), ("bert", "BertConfig"),
("rembert", "RemBertConfig"), ("bert-generation", "BertGenerationConfig"),
("visual_bert", "VisualBertConfig"), ("big_bird", "BigBirdConfig"),
("bigbird_pegasus", "BigBirdPegasusConfig"),
("blenderbot", "BlenderbotConfig"),
("blenderbot-small", "BlenderbotSmallConfig"),
("camembert", "CamembertConfig"),
("canine", "CanineConfig"), ("canine", "CanineConfig"),
("roformer", "RoFormerConfig"),
("clip", "CLIPConfig"), ("clip", "CLIPConfig"),
("flava", "FlavaConfig"), ("convbert", "ConvBertConfig"),
("bigbird_pegasus", "BigBirdPegasusConfig"), ("convnext", "ConvNextConfig"),
("ctrl", "CTRLConfig"),
("data2vec-audio", "Data2VecAudioConfig"),
("data2vec-text", "Data2VecTextConfig"),
("data2vec-vision", "Data2VecVisionConfig"),
("deberta", "DebertaConfig"),
("deberta-v2", "DebertaV2Config"),
("decision_transformer", "DecisionTransformerConfig"),
("decision_transformer", "DecisionTransformerConfig"),
("deit", "DeiTConfig"), ("deit", "DeiTConfig"),
("luke", "LukeConfig"),
("detr", "DetrConfig"), ("detr", "DetrConfig"),
("distilbert", "DistilBertConfig"),
("dpr", "DPRConfig"),
("dpt", "DPTConfig"),
("electra", "ElectraConfig"),
("encoder-decoder", "EncoderDecoderConfig"),
("flaubert", "FlaubertConfig"),
("flava", "FlavaConfig"),
("fnet", "FNetConfig"),
("fsmt", "FSMTConfig"),
("funnel", "FunnelConfig"),
("glpn", "GLPNConfig"),
("gpt2", "GPT2Config"),
("gpt_neo", "GPTNeoConfig"), ("gpt_neo", "GPTNeoConfig"),
("big_bird", "BigBirdConfig"), ("gptj", "GPTJConfig"),
("speech_to_text_2", "Speech2Text2Config"), ("hubert", "HubertConfig"),
("speech_to_text", "Speech2TextConfig"),
("vit", "ViTConfig"),
("wav2vec2", "Wav2Vec2Config"),
("m2m_100", "M2M100Config"),
("convbert", "ConvBertConfig"),
("led", "LEDConfig"),
("blenderbot-small", "BlenderbotSmallConfig"),
("retribert", "RetriBertConfig"),
("ibert", "IBertConfig"), ("ibert", "IBertConfig"),
("mt5", "MT5Config"), ("imagegpt", "ImageGPTConfig"),
("t5", "T5Config"), ("layoutlm", "LayoutLMConfig"),
("mobilebert", "MobileBertConfig"), ("layoutlmv2", "LayoutLMv2Config"),
("distilbert", "DistilBertConfig"), ("led", "LEDConfig"),
("albert", "AlbertConfig"), ("longformer", "LongformerConfig"),
("bert-generation", "BertGenerationConfig"), ("luke", "LukeConfig"),
("camembert", "CamembertConfig"), ("lxmert", "LxmertConfig"),
("xlm-roberta-xl", "XLMRobertaXLConfig"), ("m2m_100", "M2M100Config"),
("xlm-roberta", "XLMRobertaConfig"),
("pegasus", "PegasusConfig"),
("marian", "MarianConfig"), ("marian", "MarianConfig"),
("maskformer", "MaskFormerConfig"),
("mbart", "MBartConfig"), ("mbart", "MBartConfig"),
("megatron-bert", "MegatronBertConfig"), ("megatron-bert", "MegatronBertConfig"),
("mobilebert", "MobileBertConfig"),
("mpnet", "MPNetConfig"), ("mpnet", "MPNetConfig"),
("bart", "BartConfig"), ("mt5", "MT5Config"),
("nystromformer", "NystromformerConfig"),
("openai-gpt", "OpenAIGPTConfig"),
("opt", "OPTConfig"), ("opt", "OPTConfig"),
("blenderbot", "BlenderbotConfig"), ("pegasus", "PegasusConfig"),
("perceiver", "PerceiverConfig"),
("plbart", "PLBartConfig"),
("poolformer", "PoolFormerConfig"),
("prophetnet", "ProphetNetConfig"),
("qdqbert", "QDQBertConfig"),
("rag", "RagConfig"),
("realm", "RealmConfig"),
("reformer", "ReformerConfig"), ("reformer", "ReformerConfig"),
("longformer", "LongformerConfig"), ("regnet", "RegNetConfig"),
("rembert", "RemBertConfig"),
("resnet", "ResNetConfig"),
("retribert", "RetriBertConfig"),
("roberta", "RobertaConfig"), ("roberta", "RobertaConfig"),
("deberta-v2", "DebertaV2Config"), ("roformer", "RoFormerConfig"),
("deberta", "DebertaConfig"), ("segformer", "SegformerConfig"),
("flaubert", "FlaubertConfig"), ("sew", "SEWConfig"),
("fsmt", "FSMTConfig"), ("sew-d", "SEWDConfig"),
("squeezebert", "SqueezeBertConfig"),
("hubert", "HubertConfig"),
("bert", "BertConfig"),
("openai-gpt", "OpenAIGPTConfig"),
("gpt2", "GPT2Config"),
("transfo-xl", "TransfoXLConfig"),
("xlnet", "XLNetConfig"),
("xlm-prophetnet", "XLMProphetNetConfig"),
("prophetnet", "ProphetNetConfig"),
("xlm", "XLMConfig"),
("ctrl", "CTRLConfig"),
("electra", "ElectraConfig"),
("speech-encoder-decoder", "SpeechEncoderDecoderConfig"), ("speech-encoder-decoder", "SpeechEncoderDecoderConfig"),
("encoder-decoder", "EncoderDecoderConfig"), ("speech_to_text", "Speech2TextConfig"),
("funnel", "FunnelConfig"), ("speech_to_text_2", "Speech2Text2Config"),
("lxmert", "LxmertConfig"),
("dpr", "DPRConfig"),
("layoutlm", "LayoutLMConfig"),
("rag", "RagConfig"),
("tapas", "TapasConfig"),
("splinter", "SplinterConfig"), ("splinter", "SplinterConfig"),
("sew-d", "SEWDConfig"), ("squeezebert", "SqueezeBertConfig"),
("sew", "SEWConfig"), ("swin", "SwinConfig"),
("unispeech-sat", "UniSpeechSatConfig"), ("t5", "T5Config"),
("tapas", "TapasConfig"),
("transfo-xl", "TransfoXLConfig"),
("trocr", "TrOCRConfig"),
("unispeech", "UniSpeechConfig"), ("unispeech", "UniSpeechConfig"),
("unispeech-sat", "UniSpeechSatConfig"),
("van", "VanConfig"),
("vilt", "ViltConfig"),
("vision-encoder-decoder", "VisionEncoderDecoderConfig"),
("vision-text-dual-encoder", "VisionTextDualEncoderConfig"),
("visual_bert", "VisualBertConfig"),
("vit", "ViTConfig"),
("vit_mae", "ViTMAEConfig"),
("wav2vec2", "Wav2Vec2Config"),
("wavlm", "WavLMConfig"), ("wavlm", "WavLMConfig"),
("data2vec-audio", "Data2VecAudioConfig"), ("xglm", "XGLMConfig"),
("data2vec-text", "Data2VecTextConfig"), ("xlm", "XLMConfig"),
("xlm-prophetnet", "XLMProphetNetConfig"),
("xlm-roberta", "XLMRobertaConfig"),
("xlm-roberta-xl", "XLMRobertaXLConfig"),
("xlnet", "XLNetConfig"),
("yolos", "YolosConfig"),
("yoso", "YosoConfig"),
] ]
) )
CONFIG_ARCHIVE_MAP_MAPPING_NAMES = OrderedDict( CONFIG_ARCHIVE_MAP_MAPPING_NAMES = OrderedDict(
[ [
# Add archive maps here) # Add archive maps here)
("yolos", "YOLOS_PRETRAINED_CONFIG_ARCHIVE_MAP"), ("albert", "ALBERT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("dpt", "DPT_PRETRAINED_CONFIG_ARCHIVE_MAP"), ("bart", "BART_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("glpn", "GLPN_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("maskformer", "MASKFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("poolformer", "POOLFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("convnext", "CONVNEXT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("van", "VAN_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("resnet", "RESNET_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("regnet", "REGNET_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("yoso", "YOSO_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("swin", "SWIN_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("vilt", "VILT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("vit_mae", "VIT_MAE_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("realm", "REALM_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("nystromformer", "NYSTROMFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("xglm", "XGLM_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("imagegpt", "IMAGEGPT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("qdqbert", "QDQBERT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("fnet", "FNET_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("pegasus", "PEGASUS_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("segformer", "SEGFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("perceiver", "PERCEIVER_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("gptj", "GPTJ_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("layoutlmv2", "LAYOUTLMV2_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("plbart", "PLBART_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("beit", "BEIT_PRETRAINED_CONFIG_ARCHIVE_MAP"), ("beit", "BEIT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("data2vec-vision", "DATA2VEC_VISION_PRETRAINED_CONFIG_ARCHIVE_MAP"), ("bert", "BERT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("rembert", "REMBERT_PRETRAINED_CONFIG_ARCHIVE_MAP"), ("big_bird", "BIG_BIRD_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("visual_bert", "VISUAL_BERT_PRETRAINED_CONFIG_ARCHIVE_MAP"), ("bigbird_pegasus", "BIGBIRD_PEGASUS_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("blenderbot", "BLENDERBOT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("blenderbot-small", "BLENDERBOT_SMALL_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("camembert", "CAMEMBERT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("canine", "CANINE_PRETRAINED_CONFIG_ARCHIVE_MAP"), ("canine", "CANINE_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("roformer", "ROFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("clip", "CLIP_PRETRAINED_CONFIG_ARCHIVE_MAP"), ("clip", "CLIP_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("flava", "FLAVA_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("bigbird_pegasus", "BIGBIRD_PEGASUS_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("deit", "DEIT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("luke", "LUKE_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("detr", "DETR_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("gpt_neo", "GPT_NEO_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("big_bird", "BIG_BIRD_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("megatron-bert", "MEGATRON_BERT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("speech_to_text", "SPEECH_TO_TEXT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("speech_to_text_2", "SPEECH_TO_TEXT_2_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("vit", "VIT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("wav2vec2", "WAV_2_VEC_2_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("m2m_100", "M2M_100_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("convbert", "CONVBERT_PRETRAINED_CONFIG_ARCHIVE_MAP"), ("convbert", "CONVBERT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("led", "LED_PRETRAINED_CONFIG_ARCHIVE_MAP"), ("convnext", "CONVNEXT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("blenderbot-small", "BLENDERBOT_SMALL_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("bert", "BERT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("bart", "BART_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("opt", "OPT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("blenderbot", "BLENDERBOT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("mbart", "MBART_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("openai-gpt", "OPENAI_GPT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("transfo-xl", "TRANSFO_XL_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("gpt2", "GPT2_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("ctrl", "CTRL_PRETRAINED_CONFIG_ARCHIVE_MAP"), ("ctrl", "CTRL_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("xlnet", "XLNET_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("xlm", "XLM_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("roberta", "ROBERTA_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("data2vec-text", "DATA2VEC_TEXT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("data2vec-audio", "DATA2VEC_AUDIO_PRETRAINED_CONFIG_ARCHIVE_MAP"), ("data2vec-audio", "DATA2VEC_AUDIO_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("data2vec-text", "DATA2VEC_TEXT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("data2vec-vision", "DATA2VEC_VISION_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("deberta", "DEBERTA_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("deberta-v2", "DEBERTA_V2_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("deit", "DEIT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("detr", "DETR_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("distilbert", "DISTILBERT_PRETRAINED_CONFIG_ARCHIVE_MAP"), ("distilbert", "DISTILBERT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("albert", "ALBERT_PRETRAINED_CONFIG_ARCHIVE_MAP"), ("dpr", "DPR_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("camembert", "CAMEMBERT_PRETRAINED_CONFIG_ARCHIVE_MAP"), ("dpt", "DPT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("t5", "T5_PRETRAINED_CONFIG_ARCHIVE_MAP"), ("electra", "ELECTRA_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("xlm-roberta", "XLM_ROBERTA_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("flaubert", "FLAUBERT_PRETRAINED_CONFIG_ARCHIVE_MAP"), ("flaubert", "FLAUBERT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("flava", "FLAVA_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("fnet", "FNET_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("fsmt", "FSMT_PRETRAINED_CONFIG_ARCHIVE_MAP"), ("fsmt", "FSMT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("electra", "ELECTRA_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("longformer", "LONGFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("retribert", "RETRIBERT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("funnel", "FUNNEL_PRETRAINED_CONFIG_ARCHIVE_MAP"), ("funnel", "FUNNEL_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("lxmert", "LXMERT_PRETRAINED_CONFIG_ARCHIVE_MAP"), ("glpn", "GLPN_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("gpt2", "GPT2_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("gpt_neo", "GPT_NEO_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("gptj", "GPTJ_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("hubert", "HUBERT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("ibert", "IBERT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("imagegpt", "IMAGEGPT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("layoutlm", "LAYOUTLM_PRETRAINED_CONFIG_ARCHIVE_MAP"), ("layoutlm", "LAYOUTLM_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("dpr", "DPR_PRETRAINED_CONFIG_ARCHIVE_MAP"), ("layoutlmv2", "LAYOUTLMV2_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("deberta", "DEBERTA_PRETRAINED_CONFIG_ARCHIVE_MAP"), ("led", "LED_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("deberta-v2", "DEBERTA_V2_PRETRAINED_CONFIG_ARCHIVE_MAP"), ("longformer", "LONGFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("squeezebert", "SQUEEZEBERT_PRETRAINED_CONFIG_ARCHIVE_MAP"), ("luke", "LUKE_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("xlm-prophetnet", "XLM_PROPHETNET_PRETRAINED_CONFIG_ARCHIVE_MAP"), ("lxmert", "LXMERT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("prophetnet", "PROPHETNET_PRETRAINED_CONFIG_ARCHIVE_MAP"), ("m2m_100", "M2M_100_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("maskformer", "MASKFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("mbart", "MBART_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("megatron-bert", "MEGATRON_BERT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("mpnet", "MPNET_PRETRAINED_CONFIG_ARCHIVE_MAP"), ("mpnet", "MPNET_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("tapas", "TAPAS_PRETRAINED_CONFIG_ARCHIVE_MAP"), ("nystromformer", "NYSTROMFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("ibert", "IBERT_PRETRAINED_CONFIG_ARCHIVE_MAP"), ("openai-gpt", "OPENAI_GPT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("hubert", "HUBERT_PRETRAINED_CONFIG_ARCHIVE_MAP"), ("opt", "OPT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("splinter", "SPLINTER_PRETRAINED_CONFIG_ARCHIVE_MAP"), ("pegasus", "PEGASUS_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("sew-d", "SEW_D_PRETRAINED_CONFIG_ARCHIVE_MAP"), ("perceiver", "PERCEIVER_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("plbart", "PLBART_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("poolformer", "POOLFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("prophetnet", "PROPHETNET_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("qdqbert", "QDQBERT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("realm", "REALM_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("regnet", "REGNET_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("rembert", "REMBERT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("resnet", "RESNET_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("retribert", "RETRIBERT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("roberta", "ROBERTA_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("roformer", "ROFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("segformer", "SEGFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("sew", "SEW_PRETRAINED_CONFIG_ARCHIVE_MAP"), ("sew", "SEW_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("unispeech-sat", "UNISPEECH_SAT_PRETRAINED_CONFIG_ARCHIVE_MAP"), ("sew-d", "SEW_D_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("speech_to_text", "SPEECH_TO_TEXT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("speech_to_text_2", "SPEECH_TO_TEXT_2_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("splinter", "SPLINTER_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("squeezebert", "SQUEEZEBERT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("swin", "SWIN_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("t5", "T5_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("tapas", "TAPAS_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("transfo-xl", "TRANSFO_XL_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("unispeech", "UNISPEECH_PRETRAINED_CONFIG_ARCHIVE_MAP"), ("unispeech", "UNISPEECH_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("unispeech-sat", "UNISPEECH_SAT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("van", "VAN_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("vilt", "VILT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("visual_bert", "VISUAL_BERT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("vit", "VIT_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("vit_mae", "VIT_MAE_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("wav2vec2", "WAV_2_VEC_2_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("xglm", "XGLM_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("xlm", "XLM_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("xlm-prophetnet", "XLM_PROPHETNET_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("xlm-roberta", "XLM_ROBERTA_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("xlnet", "XLNET_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("yolos", "YOLOS_PRETRAINED_CONFIG_ARCHIVE_MAP"),
("yoso", "YOSO_PRETRAINED_CONFIG_ARCHIVE_MAP"),
] ]
) )
MODEL_NAMES_MAPPING = OrderedDict( MODEL_NAMES_MAPPING = OrderedDict(
[ [
# Add full (and cased) model names here # Add full (and cased) model names here
("yolos", "YOLOS"), ("albert", "ALBERT"),
("tapex", "TAPEX"), ("bart", "BART"),
("dpt", "DPT"), ("barthez", "BARThez"),
("decision_transformer", "Decision Transformer"), ("bartpho", "BARTpho"),
("glpn", "GLPN"),
("maskformer", "MaskFormer"),
("poolformer", "PoolFormer"),
("convnext", "ConvNext"),
("van", "VAN"),
("resnet", "ResNet"),
("regnet", "RegNet"),
("yoso", "YOSO"),
("swin", "Swin"),
("vilt", "ViLT"),
("vit_mae", "ViTMAE"),
("realm", "Realm"),
("nystromformer", "Nystromformer"),
("xglm", "XGLM"),
("imagegpt", "ImageGPT"),
("qdqbert", "QDQBert"),
("vision-encoder-decoder", "Vision Encoder decoder"),
("trocr", "TrOCR"),
("fnet", "FNet"),
("segformer", "SegFormer"),
("vision-text-dual-encoder", "VisionTextDualEncoder"),
("perceiver", "Perceiver"),
("gptj", "GPT-J"),
("beit", "BEiT"), ("beit", "BEiT"),
("plbart", "PLBart"), ("bert", "BERT"),
("rembert", "RemBERT"), ("bert-generation", "Bert Generation"),
("layoutlmv2", "LayoutLMv2"), ("bert-japanese", "BertJapanese"),
("visual_bert", "VisualBert"), ("bertweet", "Bertweet"),
("big_bird", "BigBird"),
("bigbird_pegasus", "BigBirdPegasus"),
("blenderbot", "Blenderbot"),
("blenderbot-small", "BlenderbotSmall"),
("bort", "BORT"),
("byt5", "ByT5"),
("camembert", "CamemBERT"),
("canine", "Canine"), ("canine", "Canine"),
("roformer", "RoFormer"),
("clip", "CLIP"), ("clip", "CLIP"),
("flava", "Flava"), ("convbert", "ConvBERT"),
("bigbird_pegasus", "BigBirdPegasus"), ("convnext", "ConvNext"),
("cpm", "CPM"),
("ctrl", "CTRL"),
("data2vec-audio", "Data2VecAudio"),
("data2vec-text", "Data2VecText"),
("data2vec-vision", "Data2VecVision"),
("deberta", "DeBERTa"),
("deberta-v2", "DeBERTa-v2"),
("decision_transformer", "Decision Transformer"),
("deit", "DeiT"), ("deit", "DeiT"),
("luke", "LUKE"),
("detr", "DETR"), ("detr", "DETR"),
("dialogpt", "DialoGPT"),
("distilbert", "DistilBERT"),
("dit", "DiT"),
("dpr", "DPR"),
("dpt", "DPT"),
("electra", "ELECTRA"),
("encoder-decoder", "Encoder decoder"),
("flaubert", "FlauBERT"),
("flava", "Flava"),
("fnet", "FNet"),
("fsmt", "FairSeq Machine-Translation"),
("funnel", "Funnel Transformer"),
("glpn", "GLPN"),
("gpt2", "OpenAI GPT-2"),
("gpt_neo", "GPT Neo"), ("gpt_neo", "GPT Neo"),
("big_bird", "BigBird"), ("gptj", "GPT-J"),
("speech_to_text_2", "Speech2Text2"), ("herbert", "HerBERT"),
("speech_to_text", "Speech2Text"), ("hubert", "Hubert"),
("vit", "ViT"),
("wav2vec2", "Wav2Vec2"),
("m2m_100", "M2M100"),
("convbert", "ConvBERT"),
("led", "LED"),
("blenderbot-small", "BlenderbotSmall"),
("retribert", "RetriBERT"),
("ibert", "I-BERT"), ("ibert", "I-BERT"),
("t5", "T5"), ("imagegpt", "ImageGPT"),
("mobilebert", "MobileBERT"), ("layoutlm", "LayoutLM"),
("distilbert", "DistilBERT"), ("layoutlmv2", "LayoutLMv2"),
("albert", "ALBERT"), ("layoutxlm", "LayoutXLM"),
("bert-generation", "Bert Generation"), ("led", "LED"),
("camembert", "CamemBERT"), ("longformer", "Longformer"),
("xlm-roberta", "XLM-RoBERTa"), ("luke", "LUKE"),
("xlm-roberta-xl", "XLM-RoBERTa-XL"), ("lxmert", "LXMERT"),
("pegasus", "Pegasus"), ("m2m_100", "M2M100"),
("blenderbot", "Blenderbot"),
("marian", "Marian"), ("marian", "Marian"),
("maskformer", "MaskFormer"),
("mbart", "mBART"), ("mbart", "mBART"),
("mbart50", "mBART-50"),
("megatron-bert", "MegatronBert"), ("megatron-bert", "MegatronBert"),
("bart", "BART"), ("megatron_gpt2", "MegatronGPT2"),
("mluke", "mLUKE"),
("mobilebert", "MobileBERT"),
("mpnet", "MPNet"),
("mt5", "mT5"),
("nystromformer", "Nystromformer"),
("openai-gpt", "OpenAI GPT"),
("opt", "OPT"), ("opt", "OPT"),
("pegasus", "Pegasus"),
("perceiver", "Perceiver"),
("phobert", "PhoBERT"),
("plbart", "PLBart"),
("poolformer", "PoolFormer"),
("prophetnet", "ProphetNet"),
("qdqbert", "QDQBert"),
("rag", "RAG"),
("realm", "Realm"),
("reformer", "Reformer"), ("reformer", "Reformer"),
("longformer", "Longformer"), ("regnet", "RegNet"),
("rembert", "RemBERT"),
("resnet", "ResNet"),
("retribert", "RetriBERT"),
("roberta", "RoBERTa"), ("roberta", "RoBERTa"),
("flaubert", "FlauBERT"), ("roformer", "RoFormer"),
("fsmt", "FairSeq Machine-Translation"), ("segformer", "SegFormer"),
("squeezebert", "SqueezeBERT"), ("sew", "SEW"),
("bert", "BERT"), ("sew-d", "SEW-D"),
("openai-gpt", "OpenAI GPT"),
("gpt2", "OpenAI GPT-2"),
("transfo-xl", "Transformer-XL"),
("xlnet", "XLNet"),
("xlm", "XLM"),
("ctrl", "CTRL"),
("electra", "ELECTRA"),
("encoder-decoder", "Encoder decoder"),
("speech-encoder-decoder", "Speech Encoder decoder"), ("speech-encoder-decoder", "Speech Encoder decoder"),
("vision-encoder-decoder", "Vision Encoder decoder"), ("speech_to_text", "Speech2Text"),
("funnel", "Funnel Transformer"), ("speech_to_text_2", "Speech2Text2"),
("lxmert", "LXMERT"),
("deberta-v2", "DeBERTa-v2"),
("deberta", "DeBERTa"),
("layoutlm", "LayoutLM"),
("dpr", "DPR"),
("rag", "RAG"),
("xlm-prophetnet", "XLMProphetNet"),
("prophetnet", "ProphetNet"),
("mt5", "mT5"),
("mpnet", "MPNet"),
("tapas", "TAPAS"),
("hubert", "Hubert"),
("barthez", "BARThez"),
("phobert", "PhoBERT"),
("bartpho", "BARTpho"),
("cpm", "CPM"),
("bertweet", "Bertweet"),
("bert-japanese", "BertJapanese"),
("byt5", "ByT5"),
("mbart50", "mBART-50"),
("splinter", "Splinter"), ("splinter", "Splinter"),
("sew-d", "SEW-D"), ("squeezebert", "SqueezeBERT"),
("sew", "SEW"), ("swin", "Swin"),
("unispeech-sat", "UniSpeechSat"), ("t5", "T5"),
("t5v1.1", "T5v1.1"),
("tapas", "TAPAS"),
("tapex", "TAPEX"),
("transfo-xl", "Transformer-XL"),
("trocr", "TrOCR"),
("unispeech", "UniSpeech"), ("unispeech", "UniSpeech"),
("unispeech-sat", "UniSpeechSat"),
("van", "VAN"),
("vilt", "ViLT"),
("vision-encoder-decoder", "Vision Encoder decoder"),
("vision-encoder-decoder", "Vision Encoder decoder"),
("vision-text-dual-encoder", "VisionTextDualEncoder"),
("visual_bert", "VisualBert"),
("vit", "ViT"),
("vit_mae", "ViTMAE"),
("wav2vec2", "Wav2Vec2"),
("wav2vec2_phoneme", "Wav2Vec2Phoneme"),
("wavlm", "WavLM"), ("wavlm", "WavLM"),
("bort", "BORT"), ("xglm", "XGLM"),
("dialogpt", "DialoGPT"), ("xlm", "XLM"),
("xlm-prophetnet", "XLMProphetNet"),
("xlm-roberta", "XLM-RoBERTa"),
("xlm-roberta-xl", "XLM-RoBERTa-XL"),
("xlnet", "XLNet"),
("xls_r", "XLS-R"), ("xls_r", "XLS-R"),
("t5v1.1", "T5v1.1"),
("herbert", "HerBERT"),
("wav2vec2_phoneme", "Wav2Vec2Phoneme"),
("megatron_gpt2", "MegatronGPT2"),
("xlsr_wav2vec2", "XLSR-Wav2Vec2"), ("xlsr_wav2vec2", "XLSR-Wav2Vec2"),
("mluke", "mLUKE"), ("yolos", "YOLOS"),
("layoutxlm", "LayoutXLM"), ("yoso", "YOSO"),
("data2vec-audio", "Data2VecAudio"),
("data2vec-text", "Data2VecText"),
("data2vec-vision", "Data2VecVision"),
("dit", "DiT"),
] ]
) )
...@@ -703,9 +702,10 @@ class AutoConfig: ...@@ -703,9 +702,10 @@ class AutoConfig:
return config_class.from_dict(config_dict, **kwargs) return config_class.from_dict(config_dict, **kwargs)
else: else:
# Fallback: use pattern matching on the string. # Fallback: use pattern matching on the string.
for pattern, config_class in CONFIG_MAPPING.items(): # We go from longer names to shorter names to catch roberta before bert (for instance)
for pattern in sorted(CONFIG_MAPPING.keys(), key=len, reverse=True):
if pattern in str(pretrained_model_name_or_path): if pattern in str(pretrained_model_name_or_path):
return config_class.from_dict(config_dict, **kwargs) return CONFIG_MAPPING[pattern].from_dict(config_dict, **kwargs)
raise ValueError( raise ValueError(
f"Unrecognized model in {pretrained_model_name_or_path}. " f"Unrecognized model in {pretrained_model_name_or_path}. "
......
src/transformers/models/auto/feature_extraction_auto.py
View file @ ddb1a47e
...@@ -38,30 +38,30 @@ logger = logging.get_logger(__name__) ...@@ -38,30 +38,30 @@ logger = logging.get_logger(__name__)
FEATURE_EXTRACTOR_MAPPING_NAMES = OrderedDict( FEATURE_EXTRACTOR_MAPPING_NAMES = OrderedDict(
[ [
("beit", "BeitFeatureExtractor"), ("beit", "BeitFeatureExtractor"),
("detr", "DetrFeatureExtractor"),
("deit", "DeiTFeatureExtractor"),
("hubert", "Wav2Vec2FeatureExtractor"),
("speech_to_text", "Speech2TextFeatureExtractor"),
("vit", "ViTFeatureExtractor"),
("wav2vec2", "Wav2Vec2FeatureExtractor"),
("detr", "DetrFeatureExtractor"),
("layoutlmv2", "LayoutLMv2FeatureExtractor"),
("clip", "CLIPFeatureExtractor"), ("clip", "CLIPFeatureExtractor"),
("flava", "FlavaFeatureExtractor"),
("perceiver", "PerceiverFeatureExtractor"),
("swin", "ViTFeatureExtractor"),
("vit_mae", "ViTFeatureExtractor"),
("segformer", "SegformerFeatureExtractor"),
("convnext", "ConvNextFeatureExtractor"), ("convnext", "ConvNextFeatureExtractor"),
("van", "ConvNextFeatureExtractor"),
("resnet", "ConvNextFeatureExtractor"),
("regnet", "ConvNextFeatureExtractor"),
("poolformer", "PoolFormerFeatureExtractor"),
("maskformer", "MaskFormerFeatureExtractor"),
("data2vec-audio", "Wav2Vec2FeatureExtractor"), ("data2vec-audio", "Wav2Vec2FeatureExtractor"),
("data2vec-vision", "BeitFeatureExtractor"), ("data2vec-vision", "BeitFeatureExtractor"),
("deit", "DeiTFeatureExtractor"),
("detr", "DetrFeatureExtractor"),
("detr", "DetrFeatureExtractor"),
("dpt", "DPTFeatureExtractor"), ("dpt", "DPTFeatureExtractor"),
("flava", "FlavaFeatureExtractor"),
("glpn", "GLPNFeatureExtractor"), ("glpn", "GLPNFeatureExtractor"),
("hubert", "Wav2Vec2FeatureExtractor"),
("layoutlmv2", "LayoutLMv2FeatureExtractor"),
("maskformer", "MaskFormerFeatureExtractor"),
("perceiver", "PerceiverFeatureExtractor"),
("poolformer", "PoolFormerFeatureExtractor"),
("regnet", "ConvNextFeatureExtractor"),
("resnet", "ConvNextFeatureExtractor"),
("segformer", "SegformerFeatureExtractor"),
("speech_to_text", "Speech2TextFeatureExtractor"),
("swin", "ViTFeatureExtractor"),
("van", "ConvNextFeatureExtractor"),
("vit", "ViTFeatureExtractor"),
("vit_mae", "ViTFeatureExtractor"),
("wav2vec2", "Wav2Vec2FeatureExtractor"),
("yolos", "YolosFeatureExtractor"), ("yolos", "YolosFeatureExtractor"),
] ]
) )
...@@ -75,8 +75,10 @@ def feature_extractor_class_from_name(class_name: str): ...@@ -75,8 +75,10 @@ def feature_extractor_class_from_name(class_name: str):
module_name = model_type_to_module_name(module_name) module_name = model_type_to_module_name(module_name)
module = importlib.import_module(f".{module_name}", "transformers.models") module = importlib.import_module(f".{module_name}", "transformers.models")
return getattr(module, class_name) try:
break return getattr(module, class_name)
except AttributeError:
continue
for config, extractor in FEATURE_EXTRACTOR_MAPPING._extra_content.items(): for config, extractor in FEATURE_EXTRACTOR_MAPPING._extra_content.items():
if getattr(extractor, "__name__", None) == class_name: if getattr(extractor, "__name__", None) == class_name:
......
src/transformers/models/auto/modeling_auto.py
View file @ ddb1a47e
...@@ -28,257 +28,257 @@ logger = logging.get_logger(__name__) ...@@ -28,257 +28,257 @@ logger = logging.get_logger(__name__)
MODEL_MAPPING_NAMES = OrderedDict( MODEL_MAPPING_NAMES = OrderedDict(
[ [
# Base model mapping # Base model mapping
("yolos", "YolosModel"), ("albert", "AlbertModel"),
("dpt", "DPTModel"), ("bart", "BartModel"),
("decision_transformer", "DecisionTransformerModel"),
("glpn", "GLPNModel"),
("maskformer", "MaskFormerModel"),
("decision_transformer", "DecisionTransformerModel"),
("decision_transformer_gpt2", "DecisionTransformerGPT2Model"),
("poolformer", "PoolFormerModel"),
("convnext", "ConvNextModel"),
("van", "VanModel"),
("resnet", "ResNetModel"),
("regnet", "RegNetModel"),
("yoso", "YosoModel"),
("swin", "SwinModel"),
("vilt", "ViltModel"),
("vit_mae", "ViTMAEModel"),
("nystromformer", "NystromformerModel"),
("xglm", "XGLMModel"),
("imagegpt", "ImageGPTModel"),
("qdqbert", "QDQBertModel"),
("fnet", "FNetModel"),
("segformer", "SegformerModel"),
("vision-text-dual-encoder", "VisionTextDualEncoderModel"),
("perceiver", "PerceiverModel"),
("gptj", "GPTJModel"),
("layoutlmv2", "LayoutLMv2Model"),
("plbart", "PLBartModel"),
("beit", "BeitModel"), ("beit", "BeitModel"),
("data2vec-vision", "Data2VecVisionModel"), ("bert", "BertModel"),
("rembert", "RemBertModel"), ("bert-generation", "BertGenerationEncoder"),
("visual_bert", "VisualBertModel"), ("big_bird", "BigBirdModel"),
("bigbird_pegasus", "BigBirdPegasusModel"),
("blenderbot", "BlenderbotModel"),
("blenderbot-small", "BlenderbotSmallModel"),
("camembert", "CamembertModel"),
("canine", "CanineModel"), ("canine", "CanineModel"),
("roformer", "RoFormerModel"),
("clip", "CLIPModel"), ("clip", "CLIPModel"),
("flava", "FlavaModel"), ("convbert", "ConvBertModel"),
("bigbird_pegasus", "BigBirdPegasusModel"), ("convnext", "ConvNextModel"),
("ctrl", "CTRLModel"),
("data2vec-audio", "Data2VecAudioModel"),
("data2vec-text", "Data2VecTextModel"),
("data2vec-vision", "Data2VecVisionModel"),
("deberta", "DebertaModel"),
("deberta-v2", "DebertaV2Model"),
("decision_transformer", "DecisionTransformerModel"),
("decision_transformer", "DecisionTransformerModel"),
("decision_transformer_gpt2", "DecisionTransformerGPT2Model"),
("deit", "DeiTModel"), ("deit", "DeiTModel"),
("luke", "LukeModel"),
("detr", "DetrModel"), ("detr", "DetrModel"),
("distilbert", "DistilBertModel"),
("dpr", "DPRQuestionEncoder"),
("dpt", "DPTModel"),
("electra", "ElectraModel"),
("flaubert", "FlaubertModel"),
("flava", "FlavaModel"),
("fnet", "FNetModel"),
("fsmt", "FSMTModel"),
("funnel", ("FunnelModel", "FunnelBaseModel")),
("glpn", "GLPNModel"),
("gpt2", "GPT2Model"),
("gpt_neo", "GPTNeoModel"), ("gpt_neo", "GPTNeoModel"),
("big_bird", "BigBirdModel"), ("gptj", "GPTJModel"),
("speech_to_text", "Speech2TextModel"),
("vit", "ViTModel"),
("wav2vec2", "Wav2Vec2Model"),
("unispeech-sat", "UniSpeechSatModel"),
("wavlm", "WavLMModel"),
("unispeech", "UniSpeechModel"),
("hubert", "HubertModel"), ("hubert", "HubertModel"),
("m2m_100", "M2M100Model"), ("ibert", "IBertModel"),
("convbert", "ConvBertModel"), ("imagegpt", "ImageGPTModel"),
("layoutlm", "LayoutLMModel"),
("layoutlmv2", "LayoutLMv2Model"),
("led", "LEDModel"), ("led", "LEDModel"),
("blenderbot-small", "BlenderbotSmallModel"), ("longformer", "LongformerModel"),
("retribert", "RetriBertModel"), ("luke", "LukeModel"),
("mt5", "MT5Model"), ("lxmert", "LxmertModel"),
("t5", "T5Model"), ("m2m_100", "M2M100Model"),
("pegasus", "PegasusModel"),
("marian", "MarianModel"), ("marian", "MarianModel"),
("maskformer", "MaskFormerModel"),
("mbart", "MBartModel"), ("mbart", "MBartModel"),
("blenderbot", "BlenderbotModel"), ("megatron-bert", "MegatronBertModel"),
("distilbert", "DistilBertModel"), ("mobilebert", "MobileBertModel"),
("albert", "AlbertModel"), ("mpnet", "MPNetModel"),
("camembert", "CamembertModel"), ("mt5", "MT5Model"),
("xlm-roberta-xl", "XLMRobertaXLModel"), ("nystromformer", "NystromformerModel"),
("xlm-roberta", "XLMRobertaModel"), ("openai-gpt", "OpenAIGPTModel"),
("bart", "BartModel"),
("opt", "OPTModel"), ("opt", "OPTModel"),
("longformer", "LongformerModel"), ("pegasus", "PegasusModel"),
("perceiver", "PerceiverModel"),
("plbart", "PLBartModel"),
("poolformer", "PoolFormerModel"),
("prophetnet", "ProphetNetModel"),
("qdqbert", "QDQBertModel"),
("reformer", "ReformerModel"),
("regnet", "RegNetModel"),
("rembert", "RemBertModel"),
("resnet", "ResNetModel"),
("retribert", "RetriBertModel"),
("roberta", "RobertaModel"), ("roberta", "RobertaModel"),
("data2vec-text", "Data2VecTextModel"), ("roformer", "RoFormerModel"),
("data2vec-audio", "Data2VecAudioModel"), ("segformer", "SegformerModel"),
("layoutlm", "LayoutLMModel"), ("sew", "SEWModel"),
("sew-d", "SEWDModel"),
("speech_to_text", "Speech2TextModel"),
("splinter", "SplinterModel"),
("squeezebert", "SqueezeBertModel"), ("squeezebert", "SqueezeBertModel"),
("bert", "BertModel"), ("swin", "SwinModel"),
("openai-gpt", "OpenAIGPTModel"), ("t5", "T5Model"),
("gpt2", "GPT2Model"), ("tapas", "TapasModel"),
("megatron-bert", "MegatronBertModel"),
("mobilebert", "MobileBertModel"),
("transfo-xl", "TransfoXLModel"), ("transfo-xl", "TransfoXLModel"),
("xlnet", "XLNetModel"), ("unispeech", "UniSpeechModel"),
("flaubert", "FlaubertModel"), ("unispeech-sat", "UniSpeechSatModel"),
("fsmt", "FSMTModel"), ("van", "VanModel"),
("vilt", "ViltModel"),
("vision-text-dual-encoder", "VisionTextDualEncoderModel"),
("visual_bert", "VisualBertModel"),
("vit", "ViTModel"),
("vit_mae", "ViTMAEModel"),
("wav2vec2", "Wav2Vec2Model"),
("wavlm", "WavLMModel"),
("xglm", "XGLMModel"),
("xlm", "XLMModel"), ("xlm", "XLMModel"),
("ctrl", "CTRLModel"),
("electra", "ElectraModel"),
("reformer", "ReformerModel"),
("funnel", ("FunnelModel", "FunnelBaseModel")),
("lxmert", "LxmertModel"),
("bert-generation", "BertGenerationEncoder"),
("deberta", "DebertaModel"),
("deberta-v2", "DebertaV2Model"),
("dpr", "DPRQuestionEncoder"),
("xlm-prophetnet", "XLMProphetNetModel"), ("xlm-prophetnet", "XLMProphetNetModel"),
("prophetnet", "ProphetNetModel"), ("xlm-roberta", "XLMRobertaModel"),
("mpnet", "MPNetModel"), ("xlm-roberta-xl", "XLMRobertaXLModel"),
("tapas", "TapasModel"), ("xlnet", "XLNetModel"),
("ibert", "IBertModel"), ("yolos", "YolosModel"),
("splinter", "SplinterModel"), ("yoso", "YosoModel"),
("sew", "SEWModel"),
("sew-d", "SEWDModel"),
] ]
) )
MODEL_FOR_PRETRAINING_MAPPING_NAMES = OrderedDict( MODEL_FOR_PRETRAINING_MAPPING_NAMES = OrderedDict(
[ [
# Model for pre-training mapping # Model for pre-training mapping
("flava", "FlavaForPreTraining"),
("vit_mae", "ViTMAEForPreTraining"),
("fnet", "FNetForPreTraining"),
("visual_bert", "VisualBertForPreTraining"),
("layoutlm", "LayoutLMForMaskedLM"),
("retribert", "RetriBertModel"),
("t5", "T5ForConditionalGeneration"),
("distilbert", "DistilBertForMaskedLM"),
("albert", "AlbertForPreTraining"), ("albert", "AlbertForPreTraining"),
("camembert", "CamembertForMaskedLM"),
("xlm-roberta-xl", "XLMRobertaXLForMaskedLM"),
("xlm-roberta", "XLMRobertaForMaskedLM"),
("bart", "BartForConditionalGeneration"), ("bart", "BartForConditionalGeneration"),
("fsmt", "FSMTForConditionalGeneration"),
("longformer", "LongformerForMaskedLM"),
("roberta", "RobertaForMaskedLM"),
("data2vec-text", "Data2VecTextForMaskedLM"),
("squeezebert", "SqueezeBertForMaskedLM"),
("bert", "BertForPreTraining"), ("bert", "BertForPreTraining"),
("big_bird", "BigBirdForPreTraining"), ("big_bird", "BigBirdForPreTraining"),
("openai-gpt", "OpenAIGPTLMHeadModel"), ("camembert", "CamembertForMaskedLM"),
("gpt2", "GPT2LMHeadModel"),
("megatron-bert", "MegatronBertForPreTraining"),
("mobilebert", "MobileBertForPreTraining"),
("transfo-xl", "TransfoXLLMHeadModel"),
("xlnet", "XLNetLMHeadModel"),
("flaubert", "FlaubertWithLMHeadModel"),
("xlm", "XLMWithLMHeadModel"),
("ctrl", "CTRLLMHeadModel"), ("ctrl", "CTRLLMHeadModel"),
("data2vec-text", "Data2VecTextForMaskedLM"),
("deberta", "DebertaForMaskedLM"),
("deberta-v2", "DebertaV2ForMaskedLM"),
("distilbert", "DistilBertForMaskedLM"),
("electra", "ElectraForPreTraining"), ("electra", "ElectraForPreTraining"),
("lxmert", "LxmertForPreTraining"), ("flaubert", "FlaubertWithLMHeadModel"),
("flava", "FlavaForPreTraining"),
("fnet", "FNetForPreTraining"),
("fsmt", "FSMTForConditionalGeneration"),
("funnel", "FunnelForPreTraining"), ("funnel", "FunnelForPreTraining"),
("gpt2", "GPT2LMHeadModel"),
("ibert", "IBertForMaskedLM"),
("layoutlm", "LayoutLMForMaskedLM"),
("longformer", "LongformerForMaskedLM"),
("lxmert", "LxmertForPreTraining"),
("megatron-bert", "MegatronBertForPreTraining"),
("mobilebert", "MobileBertForPreTraining"),
("mpnet", "MPNetForMaskedLM"), ("mpnet", "MPNetForMaskedLM"),
("openai-gpt", "OpenAIGPTLMHeadModel"),
("retribert", "RetriBertModel"),
("roberta", "RobertaForMaskedLM"),
("squeezebert", "SqueezeBertForMaskedLM"),
("t5", "T5ForConditionalGeneration"),
("tapas", "TapasForMaskedLM"), ("tapas", "TapasForMaskedLM"),
("ibert", "IBertForMaskedLM"), ("transfo-xl", "TransfoXLLMHeadModel"),
("deberta", "DebertaForMaskedLM"),
("deberta-v2", "DebertaV2ForMaskedLM"),
("wav2vec2", "Wav2Vec2ForPreTraining"),
("unispeech-sat", "UniSpeechSatForPreTraining"),
("unispeech", "UniSpeechForPreTraining"), ("unispeech", "UniSpeechForPreTraining"),
("unispeech-sat", "UniSpeechSatForPreTraining"),
("visual_bert", "VisualBertForPreTraining"),
("vit_mae", "ViTMAEForPreTraining"),
("wav2vec2", "Wav2Vec2ForPreTraining"),
("xlm", "XLMWithLMHeadModel"),
("xlm-roberta", "XLMRobertaForMaskedLM"),
("xlm-roberta-xl", "XLMRobertaXLForMaskedLM"),
("xlnet", "XLNetLMHeadModel"),
] ]
) )
MODEL_WITH_LM_HEAD_MAPPING_NAMES = OrderedDict( MODEL_WITH_LM_HEAD_MAPPING_NAMES = OrderedDict(
[ [
# Model with LM heads mapping # Model with LM heads mapping
("yoso", "YosoForMaskedLM"),
("nystromformer", "NystromformerForMaskedLM"),
("plbart", "PLBartForConditionalGeneration"),
("qdqbert", "QDQBertForMaskedLM"),
("fnet", "FNetForMaskedLM"),
("gptj", "GPTJForCausalLM"),
("rembert", "RemBertForMaskedLM"),
("roformer", "RoFormerForMaskedLM"),
("bigbird_pegasus", "BigBirdPegasusForConditionalGeneration"),
("gpt_neo", "GPTNeoForCausalLM"),
("big_bird", "BigBirdForMaskedLM"),
("speech_to_text", "Speech2TextForConditionalGeneration"),
("wav2vec2", "Wav2Vec2ForMaskedLM"),
("m2m_100", "M2M100ForConditionalGeneration"),
("convbert", "ConvBertForMaskedLM"),
("led", "LEDForConditionalGeneration"),
("blenderbot-small", "BlenderbotSmallForConditionalGeneration"),
("layoutlm", "LayoutLMForMaskedLM"),
("t5", "T5ForConditionalGeneration"),
("distilbert", "DistilBertForMaskedLM"),
("albert", "AlbertForMaskedLM"), ("albert", "AlbertForMaskedLM"),
("camembert", "CamembertForMaskedLM"),
("xlm-roberta-xl", "XLMRobertaXLForMaskedLM"),
("xlm-roberta", "XLMRobertaForMaskedLM"),
("marian", "MarianMTModel"),
("fsmt", "FSMTForConditionalGeneration"),
("bart", "BartForConditionalGeneration"), ("bart", "BartForConditionalGeneration"),
("longformer", "LongformerForMaskedLM"),
("roberta", "RobertaForMaskedLM"),
("data2vec-text", "Data2VecTextForMaskedLM"),
("squeezebert", "SqueezeBertForMaskedLM"),
("bert", "BertForMaskedLM"), ("bert", "BertForMaskedLM"),
("openai-gpt", "OpenAIGPTLMHeadModel"), ("big_bird", "BigBirdForMaskedLM"),
("gpt2", "GPT2LMHeadModel"), ("bigbird_pegasus", "BigBirdPegasusForConditionalGeneration"),
("megatron-bert", "MegatronBertForCausalLM"), ("blenderbot-small", "BlenderbotSmallForConditionalGeneration"),
("mobilebert", "MobileBertForMaskedLM"), ("camembert", "CamembertForMaskedLM"),
("transfo-xl", "TransfoXLLMHeadModel"), ("convbert", "ConvBertForMaskedLM"),
("xlnet", "XLNetLMHeadModel"),
("flaubert", "FlaubertWithLMHeadModel"),
("xlm", "XLMWithLMHeadModel"),
("ctrl", "CTRLLMHeadModel"), ("ctrl", "CTRLLMHeadModel"),
("data2vec-text", "Data2VecTextForMaskedLM"),
("deberta", "DebertaForMaskedLM"),
("deberta-v2", "DebertaV2ForMaskedLM"),
("distilbert", "DistilBertForMaskedLM"),
("electra", "ElectraForMaskedLM"), ("electra", "ElectraForMaskedLM"),
("encoder-decoder", "EncoderDecoderModel"), ("encoder-decoder", "EncoderDecoderModel"),
("reformer", "ReformerModelWithLMHead"), ("flaubert", "FlaubertWithLMHeadModel"),
("fnet", "FNetForMaskedLM"),
("fsmt", "FSMTForConditionalGeneration"),
("funnel", "FunnelForMaskedLM"), ("funnel", "FunnelForMaskedLM"),
("gpt2", "GPT2LMHeadModel"),
("gpt_neo", "GPTNeoForCausalLM"),
("gptj", "GPTJForCausalLM"),
("ibert", "IBertForMaskedLM"),
("layoutlm", "LayoutLMForMaskedLM"),
("led", "LEDForConditionalGeneration"),
("longformer", "LongformerForMaskedLM"),
("m2m_100", "M2M100ForConditionalGeneration"),
("marian", "MarianMTModel"),
("megatron-bert", "MegatronBertForCausalLM"),
("mobilebert", "MobileBertForMaskedLM"),
("mpnet", "MPNetForMaskedLM"), ("mpnet", "MPNetForMaskedLM"),
("nystromformer", "NystromformerForMaskedLM"),
("openai-gpt", "OpenAIGPTLMHeadModel"),
("plbart", "PLBartForConditionalGeneration"),
("qdqbert", "QDQBertForMaskedLM"),
("reformer", "ReformerModelWithLMHead"),
("rembert", "RemBertForMaskedLM"),
("roberta", "RobertaForMaskedLM"),
("roformer", "RoFormerForMaskedLM"),
("speech_to_text", "Speech2TextForConditionalGeneration"),
("squeezebert", "SqueezeBertForMaskedLM"),
("t5", "T5ForConditionalGeneration"),
("tapas", "TapasForMaskedLM"), ("tapas", "TapasForMaskedLM"),
("deberta", "DebertaForMaskedLM"), ("transfo-xl", "TransfoXLLMHeadModel"),
("deberta-v2", "DebertaV2ForMaskedLM"), ("wav2vec2", "Wav2Vec2ForMaskedLM"),
("ibert", "IBertForMaskedLM"), ("xlm", "XLMWithLMHeadModel"),
("xlm-roberta", "XLMRobertaForMaskedLM"),
("xlm-roberta-xl", "XLMRobertaXLForMaskedLM"),
("xlnet", "XLNetLMHeadModel"),
("yoso", "YosoForMaskedLM"),
] ]
) )
MODEL_FOR_CAUSAL_LM_MAPPING_NAMES = OrderedDict( MODEL_FOR_CAUSAL_LM_MAPPING_NAMES = OrderedDict(
[ [
# Model for Causal LM mapping # Model for Causal LM mapping
("xglm", "XGLMForCausalLM"), ("bart", "BartForCausalLM"),
("bert", "BertLMHeadModel"),
("bert-generation", "BertGenerationDecoder"),
("big_bird", "BigBirdForCausalLM"),
("bigbird_pegasus", "BigBirdPegasusForCausalLM"),
("blenderbot", "BlenderbotForCausalLM"),
("blenderbot-small", "BlenderbotSmallForCausalLM"),
("camembert", "CamembertForCausalLM"),
("ctrl", "CTRLLMHeadModel"),
("data2vec-text", "Data2VecTextForCausalLM"),
("electra", "ElectraForCausalLM"),
("gpt2", "GPT2LMHeadModel"),
("gpt_neo", "GPTNeoForCausalLM"),
("gptj", "GPTJForCausalLM"),
("marian", "MarianForCausalLM"),
("mbart", "MBartForCausalLM"),
("megatron-bert", "MegatronBertForCausalLM"),
("openai-gpt", "OpenAIGPTLMHeadModel"),
("opt", "OPTForCausalLM"),
("pegasus", "PegasusForCausalLM"),
("plbart", "PLBartForCausalLM"), ("plbart", "PLBartForCausalLM"),
("prophetnet", "ProphetNetForCausalLM"),
("qdqbert", "QDQBertLMHeadModel"), ("qdqbert", "QDQBertLMHeadModel"),
("trocr", "TrOCRForCausalLM"), ("reformer", "ReformerModelWithLMHead"),
("gptj", "GPTJForCausalLM"),
("rembert", "RemBertForCausalLM"), ("rembert", "RemBertForCausalLM"),
("roformer", "RoFormerForCausalLM"),
("bigbird_pegasus", "BigBirdPegasusForCausalLM"),
("gpt_neo", "GPTNeoForCausalLM"),
("big_bird", "BigBirdForCausalLM"),
("camembert", "CamembertForCausalLM"),
("xlm-roberta-xl", "XLMRobertaXLForCausalLM"),
("xlm-roberta", "XLMRobertaForCausalLM"),
("roberta", "RobertaForCausalLM"), ("roberta", "RobertaForCausalLM"),
("bert", "BertLMHeadModel"), ("roformer", "RoFormerForCausalLM"),
("openai-gpt", "OpenAIGPTLMHeadModel"), ("speech_to_text_2", "Speech2Text2ForCausalLM"),
("gpt2", "GPT2LMHeadModel"),
("transfo-xl", "TransfoXLLMHeadModel"), ("transfo-xl", "TransfoXLLMHeadModel"),
("xlnet", "XLNetLMHeadModel"), ("trocr", "TrOCRForCausalLM"),
("xglm", "XGLMForCausalLM"),
("xlm", "XLMWithLMHeadModel"), ("xlm", "XLMWithLMHeadModel"),
("electra", "ElectraForCausalLM"),
("ctrl", "CTRLLMHeadModel"),
("reformer", "ReformerModelWithLMHead"),
("bert-generation", "BertGenerationDecoder"),
("xlm-prophetnet", "XLMProphetNetForCausalLM"), ("xlm-prophetnet", "XLMProphetNetForCausalLM"),
("prophetnet", "ProphetNetForCausalLM"), ("xlm-roberta", "XLMRobertaForCausalLM"),
("bart", "BartForCausalLM"), ("xlm-roberta-xl", "XLMRobertaXLForCausalLM"),
("opt", "OPTForCausalLM"), ("xlnet", "XLNetLMHeadModel"),
("mbart", "MBartForCausalLM"),
("pegasus", "PegasusForCausalLM"),
("marian", "MarianForCausalLM"),
("blenderbot", "BlenderbotForCausalLM"),
("blenderbot-small", "BlenderbotSmallForCausalLM"),
("megatron-bert", "MegatronBertForCausalLM"),
("speech_to_text_2", "Speech2Text2ForCausalLM"),
("data2vec-text", "Data2VecTextForCausalLM"),
] ]
) )
MODEL_FOR_MASKED_IMAGE_MODELING_MAPPING_NAMES = OrderedDict( MODEL_FOR_MASKED_IMAGE_MODELING_MAPPING_NAMES = OrderedDict(
[ [
("vit", "ViTForMaskedImageModeling"),
("deit", "DeiTForMaskedImageModeling"), ("deit", "DeiTForMaskedImageModeling"),
("swin", "SwinForMaskedImageModeling"), ("swin", "SwinForMaskedImageModeling"),
("vit", "ViTForMaskedImageModeling"),
] ]
) )
...@@ -293,11 +293,10 @@ MODEL_FOR_CAUSAL_IMAGE_MODELING_MAPPING_NAMES = OrderedDict( ...@@ -293,11 +293,10 @@ MODEL_FOR_CAUSAL_IMAGE_MODELING_MAPPING_NAMES = OrderedDict(
MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING_NAMES = OrderedDict( MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING_NAMES = OrderedDict(
[ [
# Model for Image Classification mapping # Model for Image Classification mapping
("vit", "ViTForImageClassification"),
("deit", ("DeiTForImageClassification", "DeiTForImageClassificationWithTeacher")),
("beit", "BeitForImageClassification"), ("beit", "BeitForImageClassification"),
("convnext", "ConvNextForImageClassification"),
("data2vec-vision", "Data2VecVisionForImageClassification"), ("data2vec-vision", "Data2VecVisionForImageClassification"),
("segformer", "SegformerForImageClassification"), ("deit", ("DeiTForImageClassification", "DeiTForImageClassificationWithTeacher")),
("imagegpt", "ImageGPTForImageClassification"), ("imagegpt", "ImageGPTForImageClassification"),
( (
"perceiver", "perceiver",
...@@ -307,12 +306,13 @@ MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING_NAMES = OrderedDict( ...@@ -307,12 +306,13 @@ MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING_NAMES = OrderedDict(
"PerceiverForImageClassificationConvProcessing", "PerceiverForImageClassificationConvProcessing",
), ),
), ),
("poolformer", "PoolFormerForImageClassification"),
("regnet", "RegNetForImageClassification"),
("resnet", "ResNetForImageClassification"),
("segformer", "SegformerForImageClassification"),
("swin", "SwinForImageClassification"), ("swin", "SwinForImageClassification"),
("convnext", "ConvNextForImageClassification"),
("van", "VanForImageClassification"), ("van", "VanForImageClassification"),
("resnet", "ResNetForImageClassification"), ("vit", "ViTForImageClassification"),
("regnet", "RegNetForImageClassification"),
("poolformer", "PoolFormerForImageClassification"),
] ]
) )
...@@ -329,8 +329,8 @@ MODEL_FOR_SEMANTIC_SEGMENTATION_MAPPING_NAMES = OrderedDict( ...@@ -329,8 +329,8 @@ MODEL_FOR_SEMANTIC_SEGMENTATION_MAPPING_NAMES = OrderedDict(
# Model for Semantic Segmentation mapping # Model for Semantic Segmentation mapping
("beit", "BeitForSemanticSegmentation"), ("beit", "BeitForSemanticSegmentation"),
("data2vec-vision", "Data2VecVisionForSemanticSegmentation"), ("data2vec-vision", "Data2VecVisionForSemanticSegmentation"),
("segformer", "SegformerForSemanticSegmentation"),
("dpt", "DPTForSemanticSegmentation"), ("dpt", "DPTForSemanticSegmentation"),
("segformer", "SegformerForSemanticSegmentation"),
] ]
) )
...@@ -350,72 +350,71 @@ MODEL_FOR_VISION_2_SEQ_MAPPING_NAMES = OrderedDict( ...@@ -350,72 +350,71 @@ MODEL_FOR_VISION_2_SEQ_MAPPING_NAMES = OrderedDict(
MODEL_FOR_MASKED_LM_MAPPING_NAMES = OrderedDict( MODEL_FOR_MASKED_LM_MAPPING_NAMES = OrderedDict(
[ [
# Model for Masked LM mapping # Model for Masked LM mapping
("yoso", "YosoForMaskedLM"),
("nystromformer", "NystromformerForMaskedLM"),
("perceiver", "PerceiverForMaskedLM"),
("qdqbert", "QDQBertForMaskedLM"),
("fnet", "FNetForMaskedLM"),
("rembert", "RemBertForMaskedLM"),
("roformer", "RoFormerForMaskedLM"),
("big_bird", "BigBirdForMaskedLM"),
("wav2vec2", "Wav2Vec2ForMaskedLM"),
("convbert", "ConvBertForMaskedLM"),
("layoutlm", "LayoutLMForMaskedLM"),
("distilbert", "DistilBertForMaskedLM"),
("albert", "AlbertForMaskedLM"), ("albert", "AlbertForMaskedLM"),
("bart", "BartForConditionalGeneration"), ("bart", "BartForConditionalGeneration"),
("mbart", "MBartForConditionalGeneration"), ("bert", "BertForMaskedLM"),
("big_bird", "BigBirdForMaskedLM"),
("camembert", "CamembertForMaskedLM"), ("camembert", "CamembertForMaskedLM"),
("xlm-roberta-xl", "XLMRobertaXLForMaskedLM"), ("convbert", "ConvBertForMaskedLM"),
("xlm-roberta", "XLMRobertaForMaskedLM"),
("longformer", "LongformerForMaskedLM"),
("roberta", "RobertaForMaskedLM"),
("data2vec-text", "Data2VecTextForMaskedLM"), ("data2vec-text", "Data2VecTextForMaskedLM"),
("squeezebert", "SqueezeBertForMaskedLM"), ("deberta", "DebertaForMaskedLM"),
("bert", "BertForMaskedLM"), ("deberta-v2", "DebertaV2ForMaskedLM"),
("megatron-bert", "MegatronBertForMaskedLM"), ("distilbert", "DistilBertForMaskedLM"),
("mobilebert", "MobileBertForMaskedLM"),
("flaubert", "FlaubertWithLMHeadModel"),
("xlm", "XLMWithLMHeadModel"),
("electra", "ElectraForMaskedLM"), ("electra", "ElectraForMaskedLM"),
("reformer", "ReformerForMaskedLM"), ("flaubert", "FlaubertWithLMHeadModel"),
("fnet", "FNetForMaskedLM"),
("funnel", "FunnelForMaskedLM"), ("funnel", "FunnelForMaskedLM"),
("ibert", "IBertForMaskedLM"),
("layoutlm", "LayoutLMForMaskedLM"),
("longformer", "LongformerForMaskedLM"),
("mbart", "MBartForConditionalGeneration"),
("megatron-bert", "MegatronBertForMaskedLM"),
("mobilebert", "MobileBertForMaskedLM"),
("mpnet", "MPNetForMaskedLM"), ("mpnet", "MPNetForMaskedLM"),
("nystromformer", "NystromformerForMaskedLM"),
("perceiver", "PerceiverForMaskedLM"),
("qdqbert", "QDQBertForMaskedLM"),
("reformer", "ReformerForMaskedLM"),
("rembert", "RemBertForMaskedLM"),
("roberta", "RobertaForMaskedLM"),
("roformer", "RoFormerForMaskedLM"),
("squeezebert", "SqueezeBertForMaskedLM"),
("tapas", "TapasForMaskedLM"), ("tapas", "TapasForMaskedLM"),
("deberta", "DebertaForMaskedLM"), ("wav2vec2", "Wav2Vec2ForMaskedLM"),
("deberta-v2", "DebertaV2ForMaskedLM"), ("xlm", "XLMWithLMHeadModel"),
("ibert", "IBertForMaskedLM"), ("xlm-roberta", "XLMRobertaForMaskedLM"),
("xlm-roberta-xl", "XLMRobertaXLForMaskedLM"),
("yoso", "YosoForMaskedLM"),
] ]
) )
MODEL_FOR_OBJECT_DETECTION_MAPPING_NAMES = OrderedDict( MODEL_FOR_OBJECT_DETECTION_MAPPING_NAMES = OrderedDict(
[ [
# Model for Object Detection mapping # Model for Object Detection mapping
("yolos", "YolosForObjectDetection"),
("detr", "DetrForObjectDetection"), ("detr", "DetrForObjectDetection"),
("yolos", "YolosForObjectDetection"),
] ]
) )
MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING_NAMES = OrderedDict( MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING_NAMES = OrderedDict(
[ [
# Model for Seq2Seq Causal LM mapping # Model for Seq2Seq Causal LM mapping
("tapex", "BartForConditionalGeneration"), ("bart", "BartForConditionalGeneration"),
("plbart", "PLBartForConditionalGeneration"),
("bigbird_pegasus", "BigBirdPegasusForConditionalGeneration"), ("bigbird_pegasus", "BigBirdPegasusForConditionalGeneration"),
("m2m_100", "M2M100ForConditionalGeneration"), ("blenderbot", "BlenderbotForConditionalGeneration"),
("led", "LEDForConditionalGeneration"),
("blenderbot-small", "BlenderbotSmallForConditionalGeneration"), ("blenderbot-small", "BlenderbotSmallForConditionalGeneration"),
("mt5", "MT5ForConditionalGeneration"), ("encoder-decoder", "EncoderDecoderModel"),
("t5", "T5ForConditionalGeneration"), ("fsmt", "FSMTForConditionalGeneration"),
("pegasus", "PegasusForConditionalGeneration"), ("led", "LEDForConditionalGeneration"),
("m2m_100", "M2M100ForConditionalGeneration"),
("marian", "MarianMTModel"), ("marian", "MarianMTModel"),
("mbart", "MBartForConditionalGeneration"), ("mbart", "MBartForConditionalGeneration"),
("blenderbot", "BlenderbotForConditionalGeneration"), ("mt5", "MT5ForConditionalGeneration"),
("bart", "BartForConditionalGeneration"), ("pegasus", "PegasusForConditionalGeneration"),
("fsmt", "FSMTForConditionalGeneration"), ("plbart", "PLBartForConditionalGeneration"),
("encoder-decoder", "EncoderDecoderModel"),
("xlm-prophetnet", "XLMProphetNetForConditionalGeneration"),
("prophetnet", "ProphetNetForConditionalGeneration"), ("prophetnet", "ProphetNetForConditionalGeneration"),
("t5", "T5ForConditionalGeneration"),
("xlm-prophetnet", "XLMProphetNetForConditionalGeneration"),
] ]
) )
...@@ -429,98 +428,97 @@ MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING_NAMES = OrderedDict( ...@@ -429,98 +428,97 @@ MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING_NAMES = OrderedDict(
MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING_NAMES = OrderedDict( MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING_NAMES = OrderedDict(
[ [
# Model for Sequence Classification mapping # Model for Sequence Classification mapping
("tapex", "BartForSequenceClassification"),
("yoso", "YosoForSequenceClassification"),
("nystromformer", "NystromformerForSequenceClassification"),
("plbart", "PLBartForSequenceClassification"),
("perceiver", "PerceiverForSequenceClassification"),
("qdqbert", "QDQBertForSequenceClassification"),
("fnet", "FNetForSequenceClassification"),
("gptj", "GPTJForSequenceClassification"),
("layoutlmv2", "LayoutLMv2ForSequenceClassification"),
("rembert", "RemBertForSequenceClassification"),
("canine", "CanineForSequenceClassification"),
("roformer", "RoFormerForSequenceClassification"),
("bigbird_pegasus", "BigBirdPegasusForSequenceClassification"),
("big_bird", "BigBirdForSequenceClassification"),
("convbert", "ConvBertForSequenceClassification"),
("led", "LEDForSequenceClassification"),
("distilbert", "DistilBertForSequenceClassification"),
("albert", "AlbertForSequenceClassification"), ("albert", "AlbertForSequenceClassification"),
("camembert", "CamembertForSequenceClassification"),
("xlm-roberta-xl", "XLMRobertaXLForSequenceClassification"),
("xlm-roberta", "XLMRobertaForSequenceClassification"),
("mbart", "MBartForSequenceClassification"),
("bart", "BartForSequenceClassification"), ("bart", "BartForSequenceClassification"),
("longformer", "LongformerForSequenceClassification"),
("roberta", "RobertaForSequenceClassification"),
("data2vec-text", "Data2VecTextForSequenceClassification"),
("squeezebert", "SqueezeBertForSequenceClassification"),
("layoutlm", "LayoutLMForSequenceClassification"),
("bert", "BertForSequenceClassification"), ("bert", "BertForSequenceClassification"),
("xlnet", "XLNetForSequenceClassification"), ("big_bird", "BigBirdForSequenceClassification"),
("megatron-bert", "MegatronBertForSequenceClassification"), ("bigbird_pegasus", "BigBirdPegasusForSequenceClassification"),
("mobilebert", "MobileBertForSequenceClassification"), ("camembert", "CamembertForSequenceClassification"),
("flaubert", "FlaubertForSequenceClassification"), ("canine", "CanineForSequenceClassification"),
("xlm", "XLMForSequenceClassification"), ("convbert", "ConvBertForSequenceClassification"),
("electra", "ElectraForSequenceClassification"), ("ctrl", "CTRLForSequenceClassification"),
("funnel", "FunnelForSequenceClassification"), ("data2vec-text", "Data2VecTextForSequenceClassification"),
("deberta", "DebertaForSequenceClassification"), ("deberta", "DebertaForSequenceClassification"),
("deberta-v2", "DebertaV2ForSequenceClassification"), ("deberta-v2", "DebertaV2ForSequenceClassification"),
("distilbert", "DistilBertForSequenceClassification"),
("electra", "ElectraForSequenceClassification"),
("flaubert", "FlaubertForSequenceClassification"),
("fnet", "FNetForSequenceClassification"),
("funnel", "FunnelForSequenceClassification"),
("gpt2", "GPT2ForSequenceClassification"), ("gpt2", "GPT2ForSequenceClassification"),
("gpt_neo", "GPTNeoForSequenceClassification"), ("gpt_neo", "GPTNeoForSequenceClassification"),
("gptj", "GPTJForSequenceClassification"),
("ibert", "IBertForSequenceClassification"),
("layoutlm", "LayoutLMForSequenceClassification"),
("layoutlmv2", "LayoutLMv2ForSequenceClassification"),
("led", "LEDForSequenceClassification"),
("longformer", "LongformerForSequenceClassification"),
("mbart", "MBartForSequenceClassification"),
("megatron-bert", "MegatronBertForSequenceClassification"),
("mobilebert", "MobileBertForSequenceClassification"),
("mpnet", "MPNetForSequenceClassification"),
("nystromformer", "NystromformerForSequenceClassification"),
("openai-gpt", "OpenAIGPTForSequenceClassification"), ("openai-gpt", "OpenAIGPTForSequenceClassification"),
("perceiver", "PerceiverForSequenceClassification"),
("plbart", "PLBartForSequenceClassification"),
("qdqbert", "QDQBertForSequenceClassification"),
("reformer", "ReformerForSequenceClassification"), ("reformer", "ReformerForSequenceClassification"),
("ctrl", "CTRLForSequenceClassification"), ("rembert", "RemBertForSequenceClassification"),
("transfo-xl", "TransfoXLForSequenceClassification"), ("roberta", "RobertaForSequenceClassification"),
("mpnet", "MPNetForSequenceClassification"), ("roformer", "RoFormerForSequenceClassification"),
("squeezebert", "SqueezeBertForSequenceClassification"),
("tapas", "TapasForSequenceClassification"), ("tapas", "TapasForSequenceClassification"),
("ibert", "IBertForSequenceClassification"), ("transfo-xl", "TransfoXLForSequenceClassification"),
("xlm", "XLMForSequenceClassification"),
("xlm-roberta", "XLMRobertaForSequenceClassification"),
("xlm-roberta-xl", "XLMRobertaXLForSequenceClassification"),
("xlnet", "XLNetForSequenceClassification"),
("yoso", "YosoForSequenceClassification"),
] ]
) )
MODEL_FOR_QUESTION_ANSWERING_MAPPING_NAMES = OrderedDict( MODEL_FOR_QUESTION_ANSWERING_MAPPING_NAMES = OrderedDict(
[ [
# Model for Question Answering mapping # Model for Question Answering mapping
("yoso", "YosoForQuestionAnswering"), ("albert", "AlbertForQuestionAnswering"),
("nystromformer", "NystromformerForQuestionAnswering"), ("bart", "BartForQuestionAnswering"),
("qdqbert", "QDQBertForQuestionAnswering"), ("bert", "BertForQuestionAnswering"),
("big_bird", "BigBirdForQuestionAnswering"),
("bigbird_pegasus", "BigBirdPegasusForQuestionAnswering"),
("camembert", "CamembertForQuestionAnswering"),
("canine", "CanineForQuestionAnswering"),
("convbert", "ConvBertForQuestionAnswering"),
("data2vec-text", "Data2VecTextForQuestionAnswering"),
("deberta", "DebertaForQuestionAnswering"),
("deberta-v2", "DebertaV2ForQuestionAnswering"),
("distilbert", "DistilBertForQuestionAnswering"),
("electra", "ElectraForQuestionAnswering"),
("flaubert", "FlaubertForQuestionAnsweringSimple"),
("fnet", "FNetForQuestionAnswering"), ("fnet", "FNetForQuestionAnswering"),
("funnel", "FunnelForQuestionAnswering"),
("gptj", "GPTJForQuestionAnswering"), ("gptj", "GPTJForQuestionAnswering"),
("ibert", "IBertForQuestionAnswering"),
("layoutlmv2", "LayoutLMv2ForQuestionAnswering"), ("layoutlmv2", "LayoutLMv2ForQuestionAnswering"),
("rembert", "RemBertForQuestionAnswering"),
("canine", "CanineForQuestionAnswering"),
("roformer", "RoFormerForQuestionAnswering"),
("bigbird_pegasus", "BigBirdPegasusForQuestionAnswering"),
("big_bird", "BigBirdForQuestionAnswering"),
("convbert", "ConvBertForQuestionAnswering"),
("led", "LEDForQuestionAnswering"), ("led", "LEDForQuestionAnswering"),
("distilbert", "DistilBertForQuestionAnswering"),
("albert", "AlbertForQuestionAnswering"),
("camembert", "CamembertForQuestionAnswering"),
("bart", "BartForQuestionAnswering"),
("mbart", "MBartForQuestionAnswering"),
("longformer", "LongformerForQuestionAnswering"), ("longformer", "LongformerForQuestionAnswering"),
("xlm-roberta-xl", "XLMRobertaXLForQuestionAnswering"), ("lxmert", "LxmertForQuestionAnswering"),
("xlm-roberta", "XLMRobertaForQuestionAnswering"), ("mbart", "MBartForQuestionAnswering"),
("roberta", "RobertaForQuestionAnswering"),
("squeezebert", "SqueezeBertForQuestionAnswering"),
("bert", "BertForQuestionAnswering"),
("xlnet", "XLNetForQuestionAnsweringSimple"),
("flaubert", "FlaubertForQuestionAnsweringSimple"),
("megatron-bert", "MegatronBertForQuestionAnswering"), ("megatron-bert", "MegatronBertForQuestionAnswering"),
("mobilebert", "MobileBertForQuestionAnswering"), ("mobilebert", "MobileBertForQuestionAnswering"),
("xlm", "XLMForQuestionAnsweringSimple"),
("electra", "ElectraForQuestionAnswering"),
("reformer", "ReformerForQuestionAnswering"),
("funnel", "FunnelForQuestionAnswering"),
("lxmert", "LxmertForQuestionAnswering"),
("mpnet", "MPNetForQuestionAnswering"), ("mpnet", "MPNetForQuestionAnswering"),
("deberta", "DebertaForQuestionAnswering"), ("nystromformer", "NystromformerForQuestionAnswering"),
("deberta-v2", "DebertaV2ForQuestionAnswering"), ("qdqbert", "QDQBertForQuestionAnswering"),
("ibert", "IBertForQuestionAnswering"), ("reformer", "ReformerForQuestionAnswering"),
("rembert", "RemBertForQuestionAnswering"),
("roberta", "RobertaForQuestionAnswering"),
("roformer", "RoFormerForQuestionAnswering"),
("splinter", "SplinterForQuestionAnswering"), ("splinter", "SplinterForQuestionAnswering"),
("data2vec-text", "Data2VecTextForQuestionAnswering"), ("squeezebert", "SqueezeBertForQuestionAnswering"),
("xlm", "XLMForQuestionAnsweringSimple"),
("xlm-roberta", "XLMRobertaForQuestionAnswering"),
("xlm-roberta-xl", "XLMRobertaXLForQuestionAnswering"),
("xlnet", "XLNetForQuestionAnsweringSimple"),
("yoso", "YosoForQuestionAnswering"),
] ]
) )
...@@ -534,132 +532,132 @@ MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING_NAMES = OrderedDict( ...@@ -534,132 +532,132 @@ MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING_NAMES = OrderedDict(
MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING_NAMES = OrderedDict( MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING_NAMES = OrderedDict(
[ [
# Model for Token Classification mapping # Model for Token Classification mapping
("yoso", "YosoForTokenClassification"), ("albert", "AlbertForTokenClassification"),
("nystromformer", "NystromformerForTokenClassification"), ("bert", "BertForTokenClassification"),
("qdqbert", "QDQBertForTokenClassification"),
("fnet", "FNetForTokenClassification"),
("layoutlmv2", "LayoutLMv2ForTokenClassification"),
("rembert", "RemBertForTokenClassification"),
("canine", "CanineForTokenClassification"),
("roformer", "RoFormerForTokenClassification"),
("big_bird", "BigBirdForTokenClassification"), ("big_bird", "BigBirdForTokenClassification"),
("camembert", "CamembertForTokenClassification"),
("canine", "CanineForTokenClassification"),
("convbert", "ConvBertForTokenClassification"), ("convbert", "ConvBertForTokenClassification"),
("layoutlm", "LayoutLMForTokenClassification"), ("data2vec-text", "Data2VecTextForTokenClassification"),
("deberta", "DebertaForTokenClassification"),
("deberta-v2", "DebertaV2ForTokenClassification"),
("distilbert", "DistilBertForTokenClassification"), ("distilbert", "DistilBertForTokenClassification"),
("camembert", "CamembertForTokenClassification"), ("electra", "ElectraForTokenClassification"),
("flaubert", "FlaubertForTokenClassification"), ("flaubert", "FlaubertForTokenClassification"),
("xlm", "XLMForTokenClassification"), ("fnet", "FNetForTokenClassification"),
("xlm-roberta-xl", "XLMRobertaXLForTokenClassification"), ("funnel", "FunnelForTokenClassification"),
("xlm-roberta", "XLMRobertaForTokenClassification"), ("gpt2", "GPT2ForTokenClassification"),
("ibert", "IBertForTokenClassification"),
("layoutlm", "LayoutLMForTokenClassification"),
("layoutlmv2", "LayoutLMv2ForTokenClassification"),
("longformer", "LongformerForTokenClassification"), ("longformer", "LongformerForTokenClassification"),
("roberta", "RobertaForTokenClassification"),
("squeezebert", "SqueezeBertForTokenClassification"),
("bert", "BertForTokenClassification"),
("megatron-bert", "MegatronBertForTokenClassification"), ("megatron-bert", "MegatronBertForTokenClassification"),
("mobilebert", "MobileBertForTokenClassification"), ("mobilebert", "MobileBertForTokenClassification"),
("xlnet", "XLNetForTokenClassification"),
("albert", "AlbertForTokenClassification"),
("electra", "ElectraForTokenClassification"),
("funnel", "FunnelForTokenClassification"),
("mpnet", "MPNetForTokenClassification"), ("mpnet", "MPNetForTokenClassification"),
("deberta", "DebertaForTokenClassification"), ("nystromformer", "NystromformerForTokenClassification"),
("deberta-v2", "DebertaV2ForTokenClassification"), ("qdqbert", "QDQBertForTokenClassification"),
("gpt2", "GPT2ForTokenClassification"), ("rembert", "RemBertForTokenClassification"),
("ibert", "IBertForTokenClassification"), ("roberta", "RobertaForTokenClassification"),
("data2vec-text", "Data2VecTextForTokenClassification"), ("roformer", "RoFormerForTokenClassification"),
("squeezebert", "SqueezeBertForTokenClassification"),
("xlm", "XLMForTokenClassification"),
("xlm-roberta", "XLMRobertaForTokenClassification"),
("xlm-roberta-xl", "XLMRobertaXLForTokenClassification"),
("xlnet", "XLNetForTokenClassification"),
("yoso", "YosoForTokenClassification"),
] ]
) )
MODEL_FOR_MULTIPLE_CHOICE_MAPPING_NAMES = OrderedDict( MODEL_FOR_MULTIPLE_CHOICE_MAPPING_NAMES = OrderedDict(
[ [
# Model for Multiple Choice mapping # Model for Multiple Choice mapping
("yoso", "YosoForMultipleChoice"), ("albert", "AlbertForMultipleChoice"),
("nystromformer", "NystromformerForMultipleChoice"), ("bert", "BertForMultipleChoice"),
("qdqbert", "QDQBertForMultipleChoice"),
("fnet", "FNetForMultipleChoice"),
("rembert", "RemBertForMultipleChoice"),
("canine", "CanineForMultipleChoice"),
("roformer", "RoFormerForMultipleChoice"),
("big_bird", "BigBirdForMultipleChoice"), ("big_bird", "BigBirdForMultipleChoice"),
("convbert", "ConvBertForMultipleChoice"),
("camembert", "CamembertForMultipleChoice"), ("camembert", "CamembertForMultipleChoice"),
("electra", "ElectraForMultipleChoice"), ("canine", "CanineForMultipleChoice"),
("xlm-roberta-xl", "XLMRobertaXLForMultipleChoice"), ("convbert", "ConvBertForMultipleChoice"),
("xlm-roberta", "XLMRobertaForMultipleChoice"),
("longformer", "LongformerForMultipleChoice"),
("roberta", "RobertaForMultipleChoice"),
("data2vec-text", "Data2VecTextForMultipleChoice"), ("data2vec-text", "Data2VecTextForMultipleChoice"),
("squeezebert", "SqueezeBertForMultipleChoice"), ("deberta-v2", "DebertaV2ForMultipleChoice"),
("bert", "BertForMultipleChoice"),
("distilbert", "DistilBertForMultipleChoice"), ("distilbert", "DistilBertForMultipleChoice"),
("megatron-bert", "MegatronBertForMultipleChoice"), ("electra", "ElectraForMultipleChoice"),
("mobilebert", "MobileBertForMultipleChoice"),
("xlnet", "XLNetForMultipleChoice"),
("albert", "AlbertForMultipleChoice"),
("xlm", "XLMForMultipleChoice"),
("flaubert", "FlaubertForMultipleChoice"), ("flaubert", "FlaubertForMultipleChoice"),
("fnet", "FNetForMultipleChoice"),
("funnel", "FunnelForMultipleChoice"), ("funnel", "FunnelForMultipleChoice"),
("mpnet", "MPNetForMultipleChoice"),
("ibert", "IBertForMultipleChoice"), ("ibert", "IBertForMultipleChoice"),
("deberta-v2", "DebertaV2ForMultipleChoice"), ("longformer", "LongformerForMultipleChoice"),
("megatron-bert", "MegatronBertForMultipleChoice"),
("mobilebert", "MobileBertForMultipleChoice"),
("mpnet", "MPNetForMultipleChoice"),
("nystromformer", "NystromformerForMultipleChoice"),
("qdqbert", "QDQBertForMultipleChoice"),
("rembert", "RemBertForMultipleChoice"),
("roberta", "RobertaForMultipleChoice"),
("roformer", "RoFormerForMultipleChoice"),
("squeezebert", "SqueezeBertForMultipleChoice"),
("xlm", "XLMForMultipleChoice"),
("xlm-roberta", "XLMRobertaForMultipleChoice"),
("xlm-roberta-xl", "XLMRobertaXLForMultipleChoice"),
("xlnet", "XLNetForMultipleChoice"),
("yoso", "YosoForMultipleChoice"),
] ]
) )
MODEL_FOR_NEXT_SENTENCE_PREDICTION_MAPPING_NAMES = OrderedDict( MODEL_FOR_NEXT_SENTENCE_PREDICTION_MAPPING_NAMES = OrderedDict(
[ [
("qdqbert", "QDQBertForNextSentencePrediction"),
("bert", "BertForNextSentencePrediction"), ("bert", "BertForNextSentencePrediction"),
("fnet", "FNetForNextSentencePrediction"), ("fnet", "FNetForNextSentencePrediction"),
("megatron-bert", "MegatronBertForNextSentencePrediction"), ("megatron-bert", "MegatronBertForNextSentencePrediction"),
("mobilebert", "MobileBertForNextSentencePrediction"), ("mobilebert", "MobileBertForNextSentencePrediction"),
("qdqbert", "QDQBertForNextSentencePrediction"),
] ]
) )
MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING_NAMES = OrderedDict( MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING_NAMES = OrderedDict(
[ [
# Model for Audio Classification mapping # Model for Audio Classification mapping
("wav2vec2", "Wav2Vec2ForSequenceClassification"), ("data2vec-audio", "Data2VecAudioForSequenceClassification"),
("unispeech-sat", "UniSpeechSatForSequenceClassification"),
("unispeech", "UniSpeechForSequenceClassification"),
("hubert", "HubertForSequenceClassification"), ("hubert", "HubertForSequenceClassification"),
("sew", "SEWForSequenceClassification"), ("sew", "SEWForSequenceClassification"),
("sew-d", "SEWDForSequenceClassification"), ("sew-d", "SEWDForSequenceClassification"),
("unispeech", "UniSpeechForSequenceClassification"),
("unispeech-sat", "UniSpeechSatForSequenceClassification"),
("wav2vec2", "Wav2Vec2ForSequenceClassification"),
("wavlm", "WavLMForSequenceClassification"), ("wavlm", "WavLMForSequenceClassification"),
("data2vec-audio", "Data2VecAudioForSequenceClassification"),
] ]
) )
MODEL_FOR_CTC_MAPPING_NAMES = OrderedDict( MODEL_FOR_CTC_MAPPING_NAMES = OrderedDict(
[ [
# Model for Connectionist temporal classification (CTC) mapping # Model for Connectionist temporal classification (CTC) mapping
("wav2vec2", "Wav2Vec2ForCTC"), ("data2vec-audio", "Data2VecAudioForCTC"),
("unispeech-sat", "UniSpeechSatForCTC"),
("unispeech", "UniSpeechForCTC"),
("hubert", "HubertForCTC"), ("hubert", "HubertForCTC"),
("sew", "SEWForCTC"), ("sew", "SEWForCTC"),
("sew-d", "SEWDForCTC"), ("sew-d", "SEWDForCTC"),
("unispeech", "UniSpeechForCTC"),
("unispeech-sat", "UniSpeechSatForCTC"),
("wav2vec2", "Wav2Vec2ForCTC"),
("wavlm", "WavLMForCTC"), ("wavlm", "WavLMForCTC"),
("data2vec-audio", "Data2VecAudioForCTC"),
] ]
) )
MODEL_FOR_AUDIO_FRAME_CLASSIFICATION_MAPPING_NAMES = OrderedDict( MODEL_FOR_AUDIO_FRAME_CLASSIFICATION_MAPPING_NAMES = OrderedDict(
[ [
# Model for Audio Classification mapping # Model for Audio Classification mapping
("wav2vec2", "Wav2Vec2ForAudioFrameClassification"), ("data2vec-audio", "Data2VecAudioForAudioFrameClassification"),
("unispeech-sat", "UniSpeechSatForAudioFrameClassification"), ("unispeech-sat", "UniSpeechSatForAudioFrameClassification"),
("wav2vec2", "Wav2Vec2ForAudioFrameClassification"),
("wavlm", "WavLMForAudioFrameClassification"), ("wavlm", "WavLMForAudioFrameClassification"),
("data2vec-audio", "Data2VecAudioForAudioFrameClassification"),
] ]
) )
MODEL_FOR_AUDIO_XVECTOR_MAPPING_NAMES = OrderedDict( MODEL_FOR_AUDIO_XVECTOR_MAPPING_NAMES = OrderedDict(
[ [
# Model for Audio Classification mapping # Model for Audio Classification mapping
("wav2vec2", "Wav2Vec2ForXVector"), ("data2vec-audio", "Data2VecAudioForXVector"),
("unispeech-sat", "UniSpeechSatForXVector"), ("unispeech-sat", "UniSpeechSatForXVector"),
("wav2vec2", "Wav2Vec2ForXVector"),
("wavlm", "WavLMForXVector"), ("wavlm", "WavLMForXVector"),
("data2vec-audio", "Data2VecAudioForXVector"),
] ]
) )
......
src/transformers/models/auto/modeling_flax_auto.py
View file @ ddb1a47e
...@@ -28,31 +28,31 @@ logger = logging.get_logger(__name__) ...@@ -28,31 +28,31 @@ logger = logging.get_logger(__name__)
FLAX_MODEL_MAPPING_NAMES = OrderedDict( FLAX_MODEL_MAPPING_NAMES = OrderedDict(
[ [
# Base model mapping # Base model mapping
("xglm", "FlaxXGLMModel"),
("blenderbot-small", "FlaxBlenderbotSmallModel"),
("pegasus", "FlaxPegasusModel"),
("vision-text-dual-encoder", "FlaxVisionTextDualEncoderModel"),
("distilbert", "FlaxDistilBertModel"),
("albert", "FlaxAlbertModel"), ("albert", "FlaxAlbertModel"),
("roberta", "FlaxRobertaModel"), ("bart", "FlaxBartModel"),
("xlm-roberta", "FlaxXLMRobertaModel"),
("bert", "FlaxBertModel"),
("beit", "FlaxBeitModel"), ("beit", "FlaxBeitModel"),
("bert", "FlaxBertModel"),
("big_bird", "FlaxBigBirdModel"), ("big_bird", "FlaxBigBirdModel"),
("bart", "FlaxBartModel"), ("blenderbot", "FlaxBlenderbotModel"),
("blenderbot-small", "FlaxBlenderbotSmallModel"),
("clip", "FlaxCLIPModel"),
("distilbert", "FlaxDistilBertModel"),
("electra", "FlaxElectraModel"),
("gpt2", "FlaxGPT2Model"), ("gpt2", "FlaxGPT2Model"),
("gpt_neo", "FlaxGPTNeoModel"), ("gpt_neo", "FlaxGPTNeoModel"),
("gptj", "FlaxGPTJModel"), ("gptj", "FlaxGPTJModel"),
("electra", "FlaxElectraModel"), ("marian", "FlaxMarianModel"),
("clip", "FlaxCLIPModel"),
("vit", "FlaxViTModel"),
("mbart", "FlaxMBartModel"), ("mbart", "FlaxMBartModel"),
("t5", "FlaxT5Model"),
("mt5", "FlaxMT5Model"), ("mt5", "FlaxMT5Model"),
("wav2vec2", "FlaxWav2Vec2Model"), ("pegasus", "FlaxPegasusModel"),
("marian", "FlaxMarianModel"), ("roberta", "FlaxRobertaModel"),
("blenderbot", "FlaxBlenderbotModel"),
("roformer", "FlaxRoFormerModel"), ("roformer", "FlaxRoFormerModel"),
("t5", "FlaxT5Model"),
("vision-text-dual-encoder", "FlaxVisionTextDualEncoderModel"),
("vit", "FlaxViTModel"),
("wav2vec2", "FlaxWav2Vec2Model"),
("xglm", "FlaxXGLMModel"),
("xlm-roberta", "FlaxXLMRobertaModel"),
] ]
) )
...@@ -60,56 +60,56 @@ FLAX_MODEL_FOR_PRETRAINING_MAPPING_NAMES = OrderedDict( ...@@ -60,56 +60,56 @@ FLAX_MODEL_FOR_PRETRAINING_MAPPING_NAMES = OrderedDict(
[ [
# Model for pre-training mapping # Model for pre-training mapping
("albert", "FlaxAlbertForPreTraining"), ("albert", "FlaxAlbertForPreTraining"),
("roberta", "FlaxRobertaForMaskedLM"), ("bart", "FlaxBartForConditionalGeneration"),
("xlm-roberta", "FlaxXLMRobertaForMaskedLM"),
("bert", "FlaxBertForPreTraining"), ("bert", "FlaxBertForPreTraining"),
("big_bird", "FlaxBigBirdForPreTraining"), ("big_bird", "FlaxBigBirdForPreTraining"),
("bart", "FlaxBartForConditionalGeneration"),
("electra", "FlaxElectraForPreTraining"), ("electra", "FlaxElectraForPreTraining"),
("mbart", "FlaxMBartForConditionalGeneration"), ("mbart", "FlaxMBartForConditionalGeneration"),
("t5", "FlaxT5ForConditionalGeneration"),
("mt5", "FlaxMT5ForConditionalGeneration"), ("mt5", "FlaxMT5ForConditionalGeneration"),
("wav2vec2", "FlaxWav2Vec2ForPreTraining"), ("roberta", "FlaxRobertaForMaskedLM"),
("roformer", "FlaxRoFormerForMaskedLM"), ("roformer", "FlaxRoFormerForMaskedLM"),
("t5", "FlaxT5ForConditionalGeneration"),
("wav2vec2", "FlaxWav2Vec2ForPreTraining"),
("xlm-roberta", "FlaxXLMRobertaForMaskedLM"),
] ]
) )
FLAX_MODEL_FOR_MASKED_LM_MAPPING_NAMES = OrderedDict( FLAX_MODEL_FOR_MASKED_LM_MAPPING_NAMES = OrderedDict(
[ [
# Model for Masked LM mapping # Model for Masked LM mapping
("distilbert", "FlaxDistilBertForMaskedLM"),
("albert", "FlaxAlbertForMaskedLM"), ("albert", "FlaxAlbertForMaskedLM"),
("roberta", "FlaxRobertaForMaskedLM"), ("bart", "FlaxBartForConditionalGeneration"),
("xlm-roberta", "FlaxXLMRobertaForMaskedLM"),
("bert", "FlaxBertForMaskedLM"), ("bert", "FlaxBertForMaskedLM"),
("big_bird", "FlaxBigBirdForMaskedLM"), ("big_bird", "FlaxBigBirdForMaskedLM"),
("bart", "FlaxBartForConditionalGeneration"), ("distilbert", "FlaxDistilBertForMaskedLM"),
("electra", "FlaxElectraForMaskedLM"), ("electra", "FlaxElectraForMaskedLM"),
("mbart", "FlaxMBartForConditionalGeneration"), ("mbart", "FlaxMBartForConditionalGeneration"),
("roberta", "FlaxRobertaForMaskedLM"),
("roformer", "FlaxRoFormerForMaskedLM"), ("roformer", "FlaxRoFormerForMaskedLM"),
("xlm-roberta", "FlaxXLMRobertaForMaskedLM"),
] ]
) )
FLAX_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING_NAMES = OrderedDict( FLAX_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING_NAMES = OrderedDict(
[ [
# Model for Seq2Seq Causal LM mapping # Model for Seq2Seq Causal LM mapping
("blenderbot-small", "FlaxBlenderbotSmallForConditionalGeneration"),
("pegasus", "FlaxPegasusForConditionalGeneration"),
("bart", "FlaxBartForConditionalGeneration"), ("bart", "FlaxBartForConditionalGeneration"),
("blenderbot", "FlaxBlenderbotForConditionalGeneration"),
("blenderbot-small", "FlaxBlenderbotSmallForConditionalGeneration"),
("encoder-decoder", "FlaxEncoderDecoderModel"),
("marian", "FlaxMarianMTModel"),
("mbart", "FlaxMBartForConditionalGeneration"), ("mbart", "FlaxMBartForConditionalGeneration"),
("t5", "FlaxT5ForConditionalGeneration"),
("mt5", "FlaxMT5ForConditionalGeneration"), ("mt5", "FlaxMT5ForConditionalGeneration"),
("marian", "FlaxMarianMTModel"), ("pegasus", "FlaxPegasusForConditionalGeneration"),
("encoder-decoder", "FlaxEncoderDecoderModel"), ("t5", "FlaxT5ForConditionalGeneration"),
("blenderbot", "FlaxBlenderbotForConditionalGeneration"),
] ]
) )
FLAX_MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING_NAMES = OrderedDict( FLAX_MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING_NAMES = OrderedDict(
[ [
# Model for Image-classsification # Model for Image-classsification
("vit", "FlaxViTForImageClassification"),
("beit", "FlaxBeitForImageClassification"), ("beit", "FlaxBeitForImageClassification"),
("vit", "FlaxViTForImageClassification"),
] ]
) )
...@@ -122,75 +122,75 @@ FLAX_MODEL_FOR_VISION_2_SEQ_MAPPING_NAMES = OrderedDict( ...@@ -122,75 +122,75 @@ FLAX_MODEL_FOR_VISION_2_SEQ_MAPPING_NAMES = OrderedDict(
FLAX_MODEL_FOR_CAUSAL_LM_MAPPING_NAMES = OrderedDict( FLAX_MODEL_FOR_CAUSAL_LM_MAPPING_NAMES = OrderedDict(
[ [
# Model for Causal LM mapping # Model for Causal LM mapping
("gpt2", "FlaxGPT2LMHeadModel"),
("gpt_neo", "FlaxGPTNeoForCausalLM"),
("gptj", "FlaxGPTJForCausalLM"),
("xglm", "FlaxXGLMForCausalLM"),
("bart", "FlaxBartForCausalLM"), ("bart", "FlaxBartForCausalLM"),
("bert", "FlaxBertForCausalLM"), ("bert", "FlaxBertForCausalLM"),
("roberta", "FlaxRobertaForCausalLM"),
("big_bird", "FlaxBigBirdForCausalLM"), ("big_bird", "FlaxBigBirdForCausalLM"),
("electra", "FlaxElectraForCausalLM"), ("electra", "FlaxElectraForCausalLM"),
("gpt2", "FlaxGPT2LMHeadModel"),
("gpt_neo", "FlaxGPTNeoForCausalLM"),
("gptj", "FlaxGPTJForCausalLM"),
("roberta", "FlaxRobertaForCausalLM"),
("xglm", "FlaxXGLMForCausalLM"),
] ]
) )
FLAX_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING_NAMES = OrderedDict( FLAX_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING_NAMES = OrderedDict(
[ [
# Model for Sequence Classification mapping # Model for Sequence Classification mapping
("distilbert", "FlaxDistilBertForSequenceClassification"),
("albert", "FlaxAlbertForSequenceClassification"), ("albert", "FlaxAlbertForSequenceClassification"),
("roberta", "FlaxRobertaForSequenceClassification"), ("bart", "FlaxBartForSequenceClassification"),
("xlm-roberta", "FlaxXLMRobertaForSequenceClassification"),
("bert", "FlaxBertForSequenceClassification"), ("bert", "FlaxBertForSequenceClassification"),
("big_bird", "FlaxBigBirdForSequenceClassification"), ("big_bird", "FlaxBigBirdForSequenceClassification"),
("bart", "FlaxBartForSequenceClassification"), ("distilbert", "FlaxDistilBertForSequenceClassification"),
("electra", "FlaxElectraForSequenceClassification"), ("electra", "FlaxElectraForSequenceClassification"),
("mbart", "FlaxMBartForSequenceClassification"), ("mbart", "FlaxMBartForSequenceClassification"),
("roberta", "FlaxRobertaForSequenceClassification"),
("roformer", "FlaxRoFormerForSequenceClassification"), ("roformer", "FlaxRoFormerForSequenceClassification"),
("xlm-roberta", "FlaxXLMRobertaForSequenceClassification"),
] ]
) )
FLAX_MODEL_FOR_QUESTION_ANSWERING_MAPPING_NAMES = OrderedDict( FLAX_MODEL_FOR_QUESTION_ANSWERING_MAPPING_NAMES = OrderedDict(
[ [
# Model for Question Answering mapping # Model for Question Answering mapping
("distilbert", "FlaxDistilBertForQuestionAnswering"),
("albert", "FlaxAlbertForQuestionAnswering"), ("albert", "FlaxAlbertForQuestionAnswering"),
("roberta", "FlaxRobertaForQuestionAnswering"), ("bart", "FlaxBartForQuestionAnswering"),
("xlm-roberta", "FlaxXLMRobertaForQuestionAnswering"),
("bert", "FlaxBertForQuestionAnswering"), ("bert", "FlaxBertForQuestionAnswering"),
("big_bird", "FlaxBigBirdForQuestionAnswering"), ("big_bird", "FlaxBigBirdForQuestionAnswering"),
("bart", "FlaxBartForQuestionAnswering"), ("distilbert", "FlaxDistilBertForQuestionAnswering"),
("electra", "FlaxElectraForQuestionAnswering"), ("electra", "FlaxElectraForQuestionAnswering"),
("mbart", "FlaxMBartForQuestionAnswering"), ("mbart", "FlaxMBartForQuestionAnswering"),
("roberta", "FlaxRobertaForQuestionAnswering"),
("roformer", "FlaxRoFormerForQuestionAnswering"), ("roformer", "FlaxRoFormerForQuestionAnswering"),
("xlm-roberta", "FlaxXLMRobertaForQuestionAnswering"),
] ]
) )
FLAX_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING_NAMES = OrderedDict( FLAX_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING_NAMES = OrderedDict(
[ [
# Model for Token Classification mapping # Model for Token Classification mapping
("distilbert", "FlaxDistilBertForTokenClassification"),
("albert", "FlaxAlbertForTokenClassification"), ("albert", "FlaxAlbertForTokenClassification"),
("roberta", "FlaxRobertaForTokenClassification"),
("xlm-roberta", "FlaxXLMRobertaForTokenClassification"),
("bert", "FlaxBertForTokenClassification"), ("bert", "FlaxBertForTokenClassification"),
("big_bird", "FlaxBigBirdForTokenClassification"), ("big_bird", "FlaxBigBirdForTokenClassification"),
("distilbert", "FlaxDistilBertForTokenClassification"),
("electra", "FlaxElectraForTokenClassification"), ("electra", "FlaxElectraForTokenClassification"),
("roberta", "FlaxRobertaForTokenClassification"),
("roformer", "FlaxRoFormerForTokenClassification"), ("roformer", "FlaxRoFormerForTokenClassification"),
("xlm-roberta", "FlaxXLMRobertaForTokenClassification"),
] ]
) )
FLAX_MODEL_FOR_MULTIPLE_CHOICE_MAPPING_NAMES = OrderedDict( FLAX_MODEL_FOR_MULTIPLE_CHOICE_MAPPING_NAMES = OrderedDict(
[ [
# Model for Multiple Choice mapping # Model for Multiple Choice mapping
("distilbert", "FlaxDistilBertForMultipleChoice"),
("albert", "FlaxAlbertForMultipleChoice"), ("albert", "FlaxAlbertForMultipleChoice"),
("roberta", "FlaxRobertaForMultipleChoice"),
("xlm-roberta", "FlaxXLMRobertaForMultipleChoice"),
("bert", "FlaxBertForMultipleChoice"), ("bert", "FlaxBertForMultipleChoice"),
("big_bird", "FlaxBigBirdForMultipleChoice"), ("big_bird", "FlaxBigBirdForMultipleChoice"),
("distilbert", "FlaxDistilBertForMultipleChoice"),
("electra", "FlaxElectraForMultipleChoice"), ("electra", "FlaxElectraForMultipleChoice"),
("roberta", "FlaxRobertaForMultipleChoice"),
("roformer", "FlaxRoFormerForMultipleChoice"), ("roformer", "FlaxRoFormerForMultipleChoice"),
("xlm-roberta", "FlaxXLMRobertaForMultipleChoice"),
] ]
) )
......
src/transformers/models/auto/modeling_tf_auto.py
View file @ ddb1a47e
...@@ -29,142 +29,142 @@ logger = logging.get_logger(__name__) ...@@ -29,142 +29,142 @@ logger = logging.get_logger(__name__)
TF_MODEL_MAPPING_NAMES = OrderedDict( TF_MODEL_MAPPING_NAMES = OrderedDict(
[ [
# Base model mapping # Base model mapping
("speech_to_text", "TFSpeech2TextModel"), ("albert", "TFAlbertModel"),
("bart", "TFBartModel"),
("bert", "TFBertModel"),
("blenderbot", "TFBlenderbotModel"),
("blenderbot-small", "TFBlenderbotSmallModel"),
("camembert", "TFCamembertModel"),
("clip", "TFCLIPModel"), ("clip", "TFCLIPModel"),
("deberta-v2", "TFDebertaV2Model"),
("deberta", "TFDebertaModel"),
("rembert", "TFRemBertModel"),
("roformer", "TFRoFormerModel"),
("convbert", "TFConvBertModel"), ("convbert", "TFConvBertModel"),
("convnext", "TFConvNextModel"), ("convnext", "TFConvNextModel"),
("ctrl", "TFCTRLModel"),
("data2vec-vision", "TFData2VecVisionModel"), ("data2vec-vision", "TFData2VecVisionModel"),
("led", "TFLEDModel"), ("deberta", "TFDebertaModel"),
("lxmert", "TFLxmertModel"), ("deberta-v2", "TFDebertaV2Model"),
("mt5", "TFMT5Model"),
("t5", "TFT5Model"),
("distilbert", "TFDistilBertModel"), ("distilbert", "TFDistilBertModel"),
("albert", "TFAlbertModel"), ("dpr", "TFDPRQuestionEncoder"),
("bart", "TFBartModel"), ("electra", "TFElectraModel"),
("camembert", "TFCamembertModel"), ("flaubert", "TFFlaubertModel"),
("xlm-roberta", "TFXLMRobertaModel"), ("funnel", ("TFFunnelModel", "TFFunnelBaseModel")),
("longformer", "TFLongformerModel"),
("roberta", "TFRobertaModel"),
("layoutlm", "TFLayoutLMModel"),
("bert", "TFBertModel"),
("openai-gpt", "TFOpenAIGPTModel"),
("gpt2", "TFGPT2Model"), ("gpt2", "TFGPT2Model"),
("gptj", "TFGPTJModel"), ("gptj", "TFGPTJModel"),
("hubert", "TFHubertModel"),
("layoutlm", "TFLayoutLMModel"),
("led", "TFLEDModel"),
("longformer", "TFLongformerModel"),
("lxmert", "TFLxmertModel"),
("marian", "TFMarianModel"),
("mbart", "TFMBartModel"),
("mobilebert", "TFMobileBertModel"), ("mobilebert", "TFMobileBertModel"),
("transfo-xl", "TFTransfoXLModel"),
("xlnet", "TFXLNetModel"),
("flaubert", "TFFlaubertModel"),
("xlm", "TFXLMModel"),
("ctrl", "TFCTRLModel"),
("electra", "TFElectraModel"),
("funnel", ("TFFunnelModel", "TFFunnelBaseModel")),
("dpr", "TFDPRQuestionEncoder"),
("mpnet", "TFMPNetModel"), ("mpnet", "TFMPNetModel"),
("tapas", "TFTapasModel"), ("mt5", "TFMT5Model"),
("mbart", "TFMBartModel"), ("openai-gpt", "TFOpenAIGPTModel"),
("marian", "TFMarianModel"),
("pegasus", "TFPegasusModel"), ("pegasus", "TFPegasusModel"),
("blenderbot", "TFBlenderbotModel"), ("rembert", "TFRemBertModel"),
("blenderbot-small", "TFBlenderbotSmallModel"), ("roberta", "TFRobertaModel"),
("roformer", "TFRoFormerModel"),
("speech_to_text", "TFSpeech2TextModel"),
("t5", "TFT5Model"),
("tapas", "TFTapasModel"),
("transfo-xl", "TFTransfoXLModel"),
("vit", "TFViTModel"), ("vit", "TFViTModel"),
("vit_mae", "TFViTMAEModel"), ("vit_mae", "TFViTMAEModel"),
("wav2vec2", "TFWav2Vec2Model"), ("wav2vec2", "TFWav2Vec2Model"),
("hubert", "TFHubertModel"), ("xlm", "TFXLMModel"),
("xlm-roberta", "TFXLMRobertaModel"),
("xlnet", "TFXLNetModel"),
] ]
) )
TF_MODEL_FOR_PRETRAINING_MAPPING_NAMES = OrderedDict( TF_MODEL_FOR_PRETRAINING_MAPPING_NAMES = OrderedDict(
[ [
# Model for pre-training mapping # Model for pre-training mapping
("lxmert", "TFLxmertForPreTraining"),
("t5", "TFT5ForConditionalGeneration"),
("distilbert", "TFDistilBertForMaskedLM"),
("albert", "TFAlbertForPreTraining"), ("albert", "TFAlbertForPreTraining"),
("bart", "TFBartForConditionalGeneration"), ("bart", "TFBartForConditionalGeneration"),
("camembert", "TFCamembertForMaskedLM"),
("xlm-roberta", "TFXLMRobertaForMaskedLM"),
("roberta", "TFRobertaForMaskedLM"),
("layoutlm", "TFLayoutLMForMaskedLM"),
("bert", "TFBertForPreTraining"), ("bert", "TFBertForPreTraining"),
("openai-gpt", "TFOpenAIGPTLMHeadModel"), ("camembert", "TFCamembertForMaskedLM"),
("gpt2", "TFGPT2LMHeadModel"),
("mobilebert", "TFMobileBertForPreTraining"),
("transfo-xl", "TFTransfoXLLMHeadModel"),
("xlnet", "TFXLNetLMHeadModel"),
("flaubert", "TFFlaubertWithLMHeadModel"),
("xlm", "TFXLMWithLMHeadModel"),
("ctrl", "TFCTRLLMHeadModel"), ("ctrl", "TFCTRLLMHeadModel"),
("distilbert", "TFDistilBertForMaskedLM"),
("electra", "TFElectraForPreTraining"), ("electra", "TFElectraForPreTraining"),
("tapas", "TFTapasForMaskedLM"), ("flaubert", "TFFlaubertWithLMHeadModel"),
("funnel", "TFFunnelForPreTraining"), ("funnel", "TFFunnelForPreTraining"),
("gpt2", "TFGPT2LMHeadModel"),
("layoutlm", "TFLayoutLMForMaskedLM"),
("lxmert", "TFLxmertForPreTraining"),
("mobilebert", "TFMobileBertForPreTraining"),
("mpnet", "TFMPNetForMaskedLM"), ("mpnet", "TFMPNetForMaskedLM"),
("openai-gpt", "TFOpenAIGPTLMHeadModel"),
("roberta", "TFRobertaForMaskedLM"),
("t5", "TFT5ForConditionalGeneration"),
("tapas", "TFTapasForMaskedLM"),
("transfo-xl", "TFTransfoXLLMHeadModel"),
("vit_mae", "TFViTMAEForPreTraining"), ("vit_mae", "TFViTMAEForPreTraining"),
("xlm", "TFXLMWithLMHeadModel"),
("xlm-roberta", "TFXLMRobertaForMaskedLM"),
("xlnet", "TFXLNetLMHeadModel"),
] ]
) )
TF_MODEL_WITH_LM_HEAD_MAPPING_NAMES = OrderedDict( TF_MODEL_WITH_LM_HEAD_MAPPING_NAMES = OrderedDict(
[ [
# Model with LM heads mapping # Model with LM heads mapping
("speech_to_text", "TFSpeech2TextForConditionalGeneration"),
("rembert", "TFRemBertForMaskedLM"),
("roformer", "TFRoFormerForMaskedLM"),
("convbert", "TFConvBertForMaskedLM"),
("led", "TFLEDForConditionalGeneration"),
("t5", "TFT5ForConditionalGeneration"),
("distilbert", "TFDistilBertForMaskedLM"),
("albert", "TFAlbertForMaskedLM"), ("albert", "TFAlbertForMaskedLM"),
("marian", "TFMarianMTModel"),
("bart", "TFBartForConditionalGeneration"), ("bart", "TFBartForConditionalGeneration"),
("camembert", "TFCamembertForMaskedLM"),
("xlm-roberta", "TFXLMRobertaForMaskedLM"),
("longformer", "TFLongformerForMaskedLM"),
("roberta", "TFRobertaForMaskedLM"),
("layoutlm", "TFLayoutLMForMaskedLM"),
("bert", "TFBertForMaskedLM"), ("bert", "TFBertForMaskedLM"),
("openai-gpt", "TFOpenAIGPTLMHeadModel"), ("camembert", "TFCamembertForMaskedLM"),
("convbert", "TFConvBertForMaskedLM"),
("ctrl", "TFCTRLLMHeadModel"),
("distilbert", "TFDistilBertForMaskedLM"),
("electra", "TFElectraForMaskedLM"),
("flaubert", "TFFlaubertWithLMHeadModel"),
("funnel", "TFFunnelForMaskedLM"),
("gpt2", "TFGPT2LMHeadModel"), ("gpt2", "TFGPT2LMHeadModel"),
("gptj", "TFGPTJForCausalLM"), ("gptj", "TFGPTJForCausalLM"),
("layoutlm", "TFLayoutLMForMaskedLM"),
("led", "TFLEDForConditionalGeneration"),
("longformer", "TFLongformerForMaskedLM"),
("marian", "TFMarianMTModel"),
("mobilebert", "TFMobileBertForMaskedLM"), ("mobilebert", "TFMobileBertForMaskedLM"),
("mpnet", "TFMPNetForMaskedLM"),
("openai-gpt", "TFOpenAIGPTLMHeadModel"),
("rembert", "TFRemBertForMaskedLM"),
("roberta", "TFRobertaForMaskedLM"),
("roformer", "TFRoFormerForMaskedLM"),
("speech_to_text", "TFSpeech2TextForConditionalGeneration"),
("t5", "TFT5ForConditionalGeneration"),
("tapas", "TFTapasForMaskedLM"),
("transfo-xl", "TFTransfoXLLMHeadModel"), ("transfo-xl", "TFTransfoXLLMHeadModel"),
("xlnet", "TFXLNetLMHeadModel"),
("flaubert", "TFFlaubertWithLMHeadModel"),
("xlm", "TFXLMWithLMHeadModel"), ("xlm", "TFXLMWithLMHeadModel"),
("ctrl", "TFCTRLLMHeadModel"), ("xlm-roberta", "TFXLMRobertaForMaskedLM"),
("electra", "TFElectraForMaskedLM"), ("xlnet", "TFXLNetLMHeadModel"),
("tapas", "TFTapasForMaskedLM"),
("funnel", "TFFunnelForMaskedLM"),
("mpnet", "TFMPNetForMaskedLM"),
] ]
) )
TF_MODEL_FOR_CAUSAL_LM_MAPPING_NAMES = OrderedDict( TF_MODEL_FOR_CAUSAL_LM_MAPPING_NAMES = OrderedDict(
[ [
# Model for Causal LM mapping # Model for Causal LM mapping
("camembert", "TFCamembertForCausalLM"),
("rembert", "TFRemBertForCausalLM"),
("roformer", "TFRoFormerForCausalLM"),
("roberta", "TFRobertaForCausalLM"),
("bert", "TFBertLMHeadModel"), ("bert", "TFBertLMHeadModel"),
("openai-gpt", "TFOpenAIGPTLMHeadModel"), ("camembert", "TFCamembertForCausalLM"),
("ctrl", "TFCTRLLMHeadModel"),
("gpt2", "TFGPT2LMHeadModel"), ("gpt2", "TFGPT2LMHeadModel"),
("gptj", "TFGPTJForCausalLM"), ("gptj", "TFGPTJForCausalLM"),
("openai-gpt", "TFOpenAIGPTLMHeadModel"),
("rembert", "TFRemBertForCausalLM"),
("roberta", "TFRobertaForCausalLM"),
("roformer", "TFRoFormerForCausalLM"),
("transfo-xl", "TFTransfoXLLMHeadModel"), ("transfo-xl", "TFTransfoXLLMHeadModel"),
("xlnet", "TFXLNetLMHeadModel"),
("xlm", "TFXLMWithLMHeadModel"), ("xlm", "TFXLMWithLMHeadModel"),
("ctrl", "TFCTRLLMHeadModel"), ("xlnet", "TFXLNetLMHeadModel"),
] ]
) )
TF_MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING_NAMES = OrderedDict( TF_MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING_NAMES = OrderedDict(
[ [
# Model for Image-classsification # Model for Image-classsification
("vit", "TFViTForImageClassification"),
("convnext", "TFConvNextForImageClassification"), ("convnext", "TFConvNextForImageClassification"),
("data2vec-vision", "TFData2VecVisionForImageClassification"), ("data2vec-vision", "TFData2VecVisionForImageClassification"),
("vit", "TFViTForImageClassification"),
] ]
) )
...@@ -177,42 +177,42 @@ TF_MODEL_FOR_VISION_2_SEQ_MAPPING_NAMES = OrderedDict( ...@@ -177,42 +177,42 @@ TF_MODEL_FOR_VISION_2_SEQ_MAPPING_NAMES = OrderedDict(
TF_MODEL_FOR_MASKED_LM_MAPPING_NAMES = OrderedDict( TF_MODEL_FOR_MASKED_LM_MAPPING_NAMES = OrderedDict(
[ [
# Model for Masked LM mapping # Model for Masked LM mapping
("deberta-v2", "TFDebertaV2ForMaskedLM"),
("deberta", "TFDebertaForMaskedLM"),
("rembert", "TFRemBertForMaskedLM"),
("roformer", "TFRoFormerForMaskedLM"),
("convbert", "TFConvBertForMaskedLM"),
("distilbert", "TFDistilBertForMaskedLM"),
("albert", "TFAlbertForMaskedLM"), ("albert", "TFAlbertForMaskedLM"),
("camembert", "TFCamembertForMaskedLM"),
("xlm-roberta", "TFXLMRobertaForMaskedLM"),
("longformer", "TFLongformerForMaskedLM"),
("roberta", "TFRobertaForMaskedLM"),
("layoutlm", "TFLayoutLMForMaskedLM"),
("bert", "TFBertForMaskedLM"), ("bert", "TFBertForMaskedLM"),
("mobilebert", "TFMobileBertForMaskedLM"), ("camembert", "TFCamembertForMaskedLM"),
("flaubert", "TFFlaubertWithLMHeadModel"), ("convbert", "TFConvBertForMaskedLM"),
("xlm", "TFXLMWithLMHeadModel"), ("deberta", "TFDebertaForMaskedLM"),
("deberta-v2", "TFDebertaV2ForMaskedLM"),
("distilbert", "TFDistilBertForMaskedLM"),
("electra", "TFElectraForMaskedLM"), ("electra", "TFElectraForMaskedLM"),
("tapas", "TFTapasForMaskedLM"), ("flaubert", "TFFlaubertWithLMHeadModel"),
("funnel", "TFFunnelForMaskedLM"), ("funnel", "TFFunnelForMaskedLM"),
("layoutlm", "TFLayoutLMForMaskedLM"),
("longformer", "TFLongformerForMaskedLM"),
("mobilebert", "TFMobileBertForMaskedLM"),
("mpnet", "TFMPNetForMaskedLM"), ("mpnet", "TFMPNetForMaskedLM"),
("rembert", "TFRemBertForMaskedLM"),
("roberta", "TFRobertaForMaskedLM"),
("roformer", "TFRoFormerForMaskedLM"),
("tapas", "TFTapasForMaskedLM"),
("xlm", "TFXLMWithLMHeadModel"),
("xlm-roberta", "TFXLMRobertaForMaskedLM"),
] ]
) )
TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING_NAMES = OrderedDict( TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING_NAMES = OrderedDict(
[ [
# Model for Seq2Seq Causal LM mapping # Model for Seq2Seq Causal LM mapping
("bart", "TFBartForConditionalGeneration"),
("blenderbot", "TFBlenderbotForConditionalGeneration"),
("blenderbot-small", "TFBlenderbotSmallForConditionalGeneration"),
("encoder-decoder", "TFEncoderDecoderModel"),
("led", "TFLEDForConditionalGeneration"), ("led", "TFLEDForConditionalGeneration"),
("mt5", "TFMT5ForConditionalGeneration"),
("t5", "TFT5ForConditionalGeneration"),
("marian", "TFMarianMTModel"), ("marian", "TFMarianMTModel"),
("mbart", "TFMBartForConditionalGeneration"), ("mbart", "TFMBartForConditionalGeneration"),
("mt5", "TFMT5ForConditionalGeneration"),
("pegasus", "TFPegasusForConditionalGeneration"), ("pegasus", "TFPegasusForConditionalGeneration"),
("blenderbot", "TFBlenderbotForConditionalGeneration"), ("t5", "TFT5ForConditionalGeneration"),
("blenderbot-small", "TFBlenderbotSmallForConditionalGeneration"),
("bart", "TFBartForConditionalGeneration"),
("encoder-decoder", "TFEncoderDecoderModel"),
] ]
) )
...@@ -225,58 +225,58 @@ TF_MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING_NAMES = OrderedDict( ...@@ -225,58 +225,58 @@ TF_MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING_NAMES = OrderedDict(
TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING_NAMES = OrderedDict( TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING_NAMES = OrderedDict(
[ [
# Model for Sequence Classification mapping # Model for Sequence Classification mapping
("deberta-v2", "TFDebertaV2ForSequenceClassification"),
("deberta", "TFDebertaForSequenceClassification"),
("rembert", "TFRemBertForSequenceClassification"),
("roformer", "TFRoFormerForSequenceClassification"),
("convbert", "TFConvBertForSequenceClassification"),
("distilbert", "TFDistilBertForSequenceClassification"),
("albert", "TFAlbertForSequenceClassification"), ("albert", "TFAlbertForSequenceClassification"),
("camembert", "TFCamembertForSequenceClassification"),
("xlm-roberta", "TFXLMRobertaForSequenceClassification"),
("longformer", "TFLongformerForSequenceClassification"),
("roberta", "TFRobertaForSequenceClassification"),
("layoutlm", "TFLayoutLMForSequenceClassification"),
("bert", "TFBertForSequenceClassification"), ("bert", "TFBertForSequenceClassification"),
("xlnet", "TFXLNetForSequenceClassification"), ("camembert", "TFCamembertForSequenceClassification"),
("mobilebert", "TFMobileBertForSequenceClassification"), ("convbert", "TFConvBertForSequenceClassification"),
("flaubert", "TFFlaubertForSequenceClassification"), ("ctrl", "TFCTRLForSequenceClassification"),
("xlm", "TFXLMForSequenceClassification"), ("deberta", "TFDebertaForSequenceClassification"),
("deberta-v2", "TFDebertaV2ForSequenceClassification"),
("distilbert", "TFDistilBertForSequenceClassification"),
("electra", "TFElectraForSequenceClassification"), ("electra", "TFElectraForSequenceClassification"),
("tapas", "TFTapasForSequenceClassification"), ("flaubert", "TFFlaubertForSequenceClassification"),
("funnel", "TFFunnelForSequenceClassification"), ("funnel", "TFFunnelForSequenceClassification"),
("gpt2", "TFGPT2ForSequenceClassification"), ("gpt2", "TFGPT2ForSequenceClassification"),
("gptj", "TFGPTJForSequenceClassification"), ("gptj", "TFGPTJForSequenceClassification"),
("layoutlm", "TFLayoutLMForSequenceClassification"),
("longformer", "TFLongformerForSequenceClassification"),
("mobilebert", "TFMobileBertForSequenceClassification"),
("mpnet", "TFMPNetForSequenceClassification"), ("mpnet", "TFMPNetForSequenceClassification"),
("openai-gpt", "TFOpenAIGPTForSequenceClassification"), ("openai-gpt", "TFOpenAIGPTForSequenceClassification"),
("rembert", "TFRemBertForSequenceClassification"),
("roberta", "TFRobertaForSequenceClassification"),
("roformer", "TFRoFormerForSequenceClassification"),
("tapas", "TFTapasForSequenceClassification"),
("transfo-xl", "TFTransfoXLForSequenceClassification"), ("transfo-xl", "TFTransfoXLForSequenceClassification"),
("ctrl", "TFCTRLForSequenceClassification"), ("xlm", "TFXLMForSequenceClassification"),
("xlm-roberta", "TFXLMRobertaForSequenceClassification"),
("xlnet", "TFXLNetForSequenceClassification"),
] ]
) )
TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING_NAMES = OrderedDict( TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING_NAMES = OrderedDict(
[ [
# Model for Question Answering mapping # Model for Question Answering mapping
("deberta-v2", "TFDebertaV2ForQuestionAnswering"),
("deberta", "TFDebertaForQuestionAnswering"),
("rembert", "TFRemBertForQuestionAnswering"),
("roformer", "TFRoFormerForQuestionAnswering"),
("convbert", "TFConvBertForQuestionAnswering"),
("distilbert", "TFDistilBertForQuestionAnswering"),
("albert", "TFAlbertForQuestionAnswering"), ("albert", "TFAlbertForQuestionAnswering"),
("camembert", "TFCamembertForQuestionAnswering"),
("xlm-roberta", "TFXLMRobertaForQuestionAnswering"),
("longformer", "TFLongformerForQuestionAnswering"),
("roberta", "TFRobertaForQuestionAnswering"),
("bert", "TFBertForQuestionAnswering"), ("bert", "TFBertForQuestionAnswering"),
("xlnet", "TFXLNetForQuestionAnsweringSimple"), ("camembert", "TFCamembertForQuestionAnswering"),
("mobilebert", "TFMobileBertForQuestionAnswering"), ("convbert", "TFConvBertForQuestionAnswering"),
("flaubert", "TFFlaubertForQuestionAnsweringSimple"), ("deberta", "TFDebertaForQuestionAnswering"),
("xlm", "TFXLMForQuestionAnsweringSimple"), ("deberta-v2", "TFDebertaV2ForQuestionAnswering"),
("distilbert", "TFDistilBertForQuestionAnswering"),
("electra", "TFElectraForQuestionAnswering"), ("electra", "TFElectraForQuestionAnswering"),
("flaubert", "TFFlaubertForQuestionAnsweringSimple"),
("funnel", "TFFunnelForQuestionAnswering"), ("funnel", "TFFunnelForQuestionAnswering"),
("gptj", "TFGPTJForQuestionAnswering"), ("gptj", "TFGPTJForQuestionAnswering"),
("longformer", "TFLongformerForQuestionAnswering"),
("mobilebert", "TFMobileBertForQuestionAnswering"),
("mpnet", "TFMPNetForQuestionAnswering"), ("mpnet", "TFMPNetForQuestionAnswering"),
("rembert", "TFRemBertForQuestionAnswering"),
("roberta", "TFRobertaForQuestionAnswering"),
("roformer", "TFRoFormerForQuestionAnswering"),
("xlm", "TFXLMForQuestionAnsweringSimple"),
("xlm-roberta", "TFXLMRobertaForQuestionAnswering"),
("xlnet", "TFXLNetForQuestionAnsweringSimple"),
] ]
) )
...@@ -291,49 +291,49 @@ TF_MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING_NAMES = OrderedDict( ...@@ -291,49 +291,49 @@ TF_MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING_NAMES = OrderedDict(
TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING_NAMES = OrderedDict( TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING_NAMES = OrderedDict(
[ [
# Model for Token Classification mapping # Model for Token Classification mapping
("deberta-v2", "TFDebertaV2ForTokenClassification"),
("deberta", "TFDebertaForTokenClassification"),
("rembert", "TFRemBertForTokenClassification"),
("roformer", "TFRoFormerForTokenClassification"),
("convbert", "TFConvBertForTokenClassification"),
("distilbert", "TFDistilBertForTokenClassification"),
("albert", "TFAlbertForTokenClassification"), ("albert", "TFAlbertForTokenClassification"),
("bert", "TFBertForTokenClassification"),
("camembert", "TFCamembertForTokenClassification"), ("camembert", "TFCamembertForTokenClassification"),
("convbert", "TFConvBertForTokenClassification"),
("deberta", "TFDebertaForTokenClassification"),
("deberta-v2", "TFDebertaV2ForTokenClassification"),
("distilbert", "TFDistilBertForTokenClassification"),
("electra", "TFElectraForTokenClassification"),
("flaubert", "TFFlaubertForTokenClassification"), ("flaubert", "TFFlaubertForTokenClassification"),
("xlm", "TFXLMForTokenClassification"), ("funnel", "TFFunnelForTokenClassification"),
("xlm-roberta", "TFXLMRobertaForTokenClassification"),
("longformer", "TFLongformerForTokenClassification"),
("roberta", "TFRobertaForTokenClassification"),
("layoutlm", "TFLayoutLMForTokenClassification"), ("layoutlm", "TFLayoutLMForTokenClassification"),
("bert", "TFBertForTokenClassification"), ("longformer", "TFLongformerForTokenClassification"),
("mobilebert", "TFMobileBertForTokenClassification"), ("mobilebert", "TFMobileBertForTokenClassification"),
("xlnet", "TFXLNetForTokenClassification"),
("electra", "TFElectraForTokenClassification"),
("funnel", "TFFunnelForTokenClassification"),
("mpnet", "TFMPNetForTokenClassification"), ("mpnet", "TFMPNetForTokenClassification"),
("rembert", "TFRemBertForTokenClassification"),
("roberta", "TFRobertaForTokenClassification"),
("roformer", "TFRoFormerForTokenClassification"),
("xlm", "TFXLMForTokenClassification"),
("xlm-roberta", "TFXLMRobertaForTokenClassification"),
("xlnet", "TFXLNetForTokenClassification"),
] ]
) )
TF_MODEL_FOR_MULTIPLE_CHOICE_MAPPING_NAMES = OrderedDict( TF_MODEL_FOR_MULTIPLE_CHOICE_MAPPING_NAMES = OrderedDict(
[ [
# Model for Multiple Choice mapping # Model for Multiple Choice mapping
("rembert", "TFRemBertForMultipleChoice"), ("albert", "TFAlbertForMultipleChoice"),
("roformer", "TFRoFormerForMultipleChoice"),
("convbert", "TFConvBertForMultipleChoice"),
("camembert", "TFCamembertForMultipleChoice"),
("xlm", "TFXLMForMultipleChoice"),
("xlm-roberta", "TFXLMRobertaForMultipleChoice"),
("longformer", "TFLongformerForMultipleChoice"),
("roberta", "TFRobertaForMultipleChoice"),
("bert", "TFBertForMultipleChoice"), ("bert", "TFBertForMultipleChoice"),
("camembert", "TFCamembertForMultipleChoice"),
("convbert", "TFConvBertForMultipleChoice"),
("distilbert", "TFDistilBertForMultipleChoice"), ("distilbert", "TFDistilBertForMultipleChoice"),
("mobilebert", "TFMobileBertForMultipleChoice"),
("xlnet", "TFXLNetForMultipleChoice"),
("flaubert", "TFFlaubertForMultipleChoice"),
("albert", "TFAlbertForMultipleChoice"),
("electra", "TFElectraForMultipleChoice"), ("electra", "TFElectraForMultipleChoice"),
("flaubert", "TFFlaubertForMultipleChoice"),
("funnel", "TFFunnelForMultipleChoice"), ("funnel", "TFFunnelForMultipleChoice"),
("longformer", "TFLongformerForMultipleChoice"),
("mobilebert", "TFMobileBertForMultipleChoice"),
("mpnet", "TFMPNetForMultipleChoice"), ("mpnet", "TFMPNetForMultipleChoice"),
("rembert", "TFRemBertForMultipleChoice"),
("roberta", "TFRobertaForMultipleChoice"),
("roformer", "TFRoFormerForMultipleChoice"),
("xlm", "TFXLMForMultipleChoice"),
("xlm-roberta", "TFXLMRobertaForMultipleChoice"),
("xlnet", "TFXLNetForMultipleChoice"),
] ]
) )
......
src/transformers/models/auto/processing_auto.py
View file @ ddb1a47e
...@@ -41,17 +41,17 @@ PROCESSOR_MAPPING_NAMES = OrderedDict( ...@@ -41,17 +41,17 @@ PROCESSOR_MAPPING_NAMES = OrderedDict(
("flava", "FLAVAProcessor"), ("flava", "FLAVAProcessor"),
("layoutlmv2", "LayoutLMv2Processor"), ("layoutlmv2", "LayoutLMv2Processor"),
("layoutxlm", "LayoutXLMProcessor"), ("layoutxlm", "LayoutXLMProcessor"),
("sew", "Wav2Vec2Processor"),
("sew-d", "Wav2Vec2Processor"),
("speech_to_text", "Speech2TextProcessor"), ("speech_to_text", "Speech2TextProcessor"),
("speech_to_text_2", "Speech2Text2Processor"), ("speech_to_text_2", "Speech2Text2Processor"),
("trocr", "TrOCRProcessor"), ("trocr", "TrOCRProcessor"),
("wav2vec2", "Wav2Vec2Processor"),
("wav2vec2_with_lm", "Wav2Vec2ProcessorWithLM"),
("vision-text-dual-encoder", "VisionTextDualEncoderProcessor"),
("unispeech", "Wav2Vec2Processor"), ("unispeech", "Wav2Vec2Processor"),
("unispeech-sat", "Wav2Vec2Processor"), ("unispeech-sat", "Wav2Vec2Processor"),
("sew", "Wav2Vec2Processor"),
("sew-d", "Wav2Vec2Processor"),
("vilt", "ViltProcessor"), ("vilt", "ViltProcessor"),
("vision-text-dual-encoder", "VisionTextDualEncoderProcessor"),
("wav2vec2", "Wav2Vec2Processor"),
("wav2vec2_with_lm", "Wav2Vec2ProcessorWithLM"),
("wavlm", "Wav2Vec2Processor"), ("wavlm", "Wav2Vec2Processor"),
] ]
) )
...@@ -65,7 +65,10 @@ def processor_class_from_name(class_name: str): ...@@ -65,7 +65,10 @@ def processor_class_from_name(class_name: str):
module_name = model_type_to_module_name(module_name) module_name = model_type_to_module_name(module_name)
module = importlib.import_module(f".{module_name}", "transformers.models") module = importlib.import_module(f".{module_name}", "transformers.models")
return getattr(module, class_name) try:
return getattr(module, class_name)
except AttributeError:
continue
for processor in PROCESSOR_MAPPING._extra_content.values(): for processor in PROCESSOR_MAPPING._extra_content.values():
if getattr(processor, "__name__", None) == class_name: if getattr(processor, "__name__", None) == class_name:
......
src/transformers/models/auto/tokenization_auto.py
View file @ ddb1a47e
...@@ -46,34 +46,37 @@ if TYPE_CHECKING: ...@@ -46,34 +46,37 @@ if TYPE_CHECKING:
else: else:
TOKENIZER_MAPPING_NAMES = OrderedDict( TOKENIZER_MAPPING_NAMES = OrderedDict(
[ [
("plbart", ("PLBartTokenizer" if is_sentencepiece_available() else None, None)),
("realm", ("RealmTokenizer", "RealmTokenizerFast" if is_tokenizers_available() else None)),
("fnet", ("FNetTokenizer", "FNetTokenizerFast" if is_tokenizers_available() else None)),
("retribert", ("RetriBertTokenizer", "RetriBertTokenizerFast" if is_tokenizers_available() else None)),
("roformer", ("RoFormerTokenizer", "RoFormerTokenizerFast" if is_tokenizers_available() else None)),
( (
"t5", "albert",
( (
"T5Tokenizer" if is_sentencepiece_available() else None, "AlbertTokenizer" if is_sentencepiece_available() else None,
"T5TokenizerFast" if is_tokenizers_available() else None, "AlbertTokenizerFast" if is_tokenizers_available() else None,
), ),
), ),
("bart", ("BartTokenizer", "BartTokenizerFast")),
( (
"mt5", "barthez",
( (
"MT5Tokenizer" if is_sentencepiece_available() else None, "BarthezTokenizer" if is_sentencepiece_available() else None,
"MT5TokenizerFast" if is_tokenizers_available() else None, "BarthezTokenizerFast" if is_tokenizers_available() else None,
), ),
), ),
("mobilebert", ("MobileBertTokenizer", "MobileBertTokenizerFast" if is_tokenizers_available() else None)), ("bartpho", ("BartphoTokenizer", None)),
("distilbert", ("DistilBertTokenizer", "DistilBertTokenizerFast" if is_tokenizers_available() else None)), ("bert", ("BertTokenizer", "BertTokenizerFast" if is_tokenizers_available() else None)),
("bert-generation", ("BertGenerationTokenizer" if is_sentencepiece_available() else None, None)),
("bert-japanese", ("BertJapaneseTokenizer", None)),
("bertweet", ("BertweetTokenizer", None)),
( (
"albert", "big_bird",
( (
"AlbertTokenizer" if is_sentencepiece_available() else None, "BigBirdTokenizer" if is_sentencepiece_available() else None,
"AlbertTokenizerFast" if is_tokenizers_available() else None, "BigBirdTokenizerFast" if is_tokenizers_available() else None,
), ),
), ),
("bigbird_pegasus", ("PegasusTokenizer", "PegasusTokenizerFast" if is_tokenizers_available() else None)),
("blenderbot", ("BlenderbotTokenizer", "BlenderbotTokenizerFast")),
("blenderbot-small", ("BlenderbotSmallTokenizer", None)),
("byt5", ("ByT5Tokenizer", None)),
( (
"camembert", "camembert",
( (
...@@ -81,136 +84,119 @@ else: ...@@ -81,136 +84,119 @@ else:
"CamembertTokenizerFast" if is_tokenizers_available() else None, "CamembertTokenizerFast" if is_tokenizers_available() else None,
), ),
), ),
("canine", ("CanineTokenizer", None)),
( (
"pegasus", "clip",
( (
"PegasusTokenizer" if is_sentencepiece_available() else None, "CLIPTokenizer",
"PegasusTokenizerFast" if is_tokenizers_available() else None, "CLIPTokenizerFast" if is_tokenizers_available() else None,
), ),
), ),
("convbert", ("ConvBertTokenizer", "ConvBertTokenizerFast" if is_tokenizers_available() else None)),
( (
"mbart", "cpm",
( (
"MBartTokenizer" if is_sentencepiece_available() else None, "CpmTokenizer" if is_sentencepiece_available() else None,
"MBartTokenizerFast" if is_tokenizers_available() else None, "CpmTokenizerFast" if is_tokenizers_available() else None,
), ),
), ),
("ctrl", ("CTRLTokenizer", None)),
("data2vec-text", ("RobertaTokenizer", "RobertaTokenizerFast" if is_tokenizers_available() else None)),
("deberta", ("DebertaTokenizer", "DebertaTokenizerFast" if is_tokenizers_available() else None)),
( (
"xlm-roberta", "deberta-v2",
( (
"XLMRobertaTokenizer" if is_sentencepiece_available() else None, "DebertaV2Tokenizer" if is_sentencepiece_available() else None,
"XLMRobertaTokenizerFast" if is_tokenizers_available() else None, "DebertaV2TokenizerFast" if is_tokenizers_available() else None,
), ),
), ),
("marian", ("MarianTokenizer" if is_sentencepiece_available() else None, None)), ("distilbert", ("DistilBertTokenizer", "DistilBertTokenizerFast" if is_tokenizers_available() else None)),
("blenderbot-small", ("BlenderbotSmallTokenizer", None)),
("blenderbot", ("BlenderbotTokenizer", "BlenderbotTokenizerFast")),
("tapex", ("TapexTokenizer", None)),
("bart", ("BartTokenizer", "BartTokenizerFast")),
("longformer", ("LongformerTokenizer", "LongformerTokenizerFast" if is_tokenizers_available() else None)),
("roberta", ("RobertaTokenizer", "RobertaTokenizerFast" if is_tokenizers_available() else None)),
( (
"reformer", "dpr",
( (
"ReformerTokenizer" if is_sentencepiece_available() else None, "DPRQuestionEncoderTokenizer",
"ReformerTokenizerFast" if is_tokenizers_available() else None, "DPRQuestionEncoderTokenizerFast" if is_tokenizers_available() else None,
), ),
), ),
("electra", ("ElectraTokenizer", "ElectraTokenizerFast" if is_tokenizers_available() else None)), ("electra", ("ElectraTokenizer", "ElectraTokenizerFast" if is_tokenizers_available() else None)),
("flaubert", ("FlaubertTokenizer", None)),
("fnet", ("FNetTokenizer", "FNetTokenizerFast" if is_tokenizers_available() else None)),
("fsmt", ("FSMTTokenizer", None)),
("funnel", ("FunnelTokenizer", "FunnelTokenizerFast" if is_tokenizers_available() else None)), ("funnel", ("FunnelTokenizer", "FunnelTokenizerFast" if is_tokenizers_available() else None)),
("lxmert", ("LxmertTokenizer", "LxmertTokenizerFast" if is_tokenizers_available() else None)), ("gpt2", ("GPT2Tokenizer", "GPT2TokenizerFast" if is_tokenizers_available() else None)),
("gpt_neo", ("GPT2Tokenizer", "GPT2TokenizerFast" if is_tokenizers_available() else None)),
("gptj", ("GPT2Tokenizer", "GPT2TokenizerFast" if is_tokenizers_available() else None)),
("herbert", ("HerbertTokenizer", "HerbertTokenizerFast" if is_tokenizers_available() else None)),
("hubert", ("Wav2Vec2CTCTokenizer", None)),
("ibert", ("RobertaTokenizer", "RobertaTokenizerFast" if is_tokenizers_available() else None)),
("layoutlm", ("LayoutLMTokenizer", "LayoutLMTokenizerFast" if is_tokenizers_available() else None)), ("layoutlm", ("LayoutLMTokenizer", "LayoutLMTokenizerFast" if is_tokenizers_available() else None)),
("layoutlmv2", ("LayoutLMv2Tokenizer", "LayoutLMv2TokenizerFast" if is_tokenizers_available() else None)), ("layoutlmv2", ("LayoutLMv2Tokenizer", "LayoutLMv2TokenizerFast" if is_tokenizers_available() else None)),
("layoutxlm", ("LayoutXLMTokenizer", "LayoutXLMTokenizerFast" if is_tokenizers_available() else None)), ("layoutxlm", ("LayoutXLMTokenizer", "LayoutXLMTokenizerFast" if is_tokenizers_available() else None)),
("led", ("LEDTokenizer", "LEDTokenizerFast" if is_tokenizers_available() else None)),
("longformer", ("LongformerTokenizer", "LongformerTokenizerFast" if is_tokenizers_available() else None)),
("luke", ("LukeTokenizer", None)),
("lxmert", ("LxmertTokenizer", "LxmertTokenizerFast" if is_tokenizers_available() else None)),
("m2m_100", ("M2M100Tokenizer" if is_sentencepiece_available() else None, None)),
("marian", ("MarianTokenizer" if is_sentencepiece_available() else None, None)),
( (
"dpr", "mbart",
( (
"DPRQuestionEncoderTokenizer", "MBartTokenizer" if is_sentencepiece_available() else None,
"DPRQuestionEncoderTokenizerFast" if is_tokenizers_available() else None, "MBartTokenizerFast" if is_tokenizers_available() else None,
), ),
), ),
( (
"squeezebert", "mbart50",
("SqueezeBertTokenizer", "SqueezeBertTokenizerFast" if is_tokenizers_available() else None),
),
("bert", ("BertTokenizer", "BertTokenizerFast" if is_tokenizers_available() else None)),
("openai-gpt", ("OpenAIGPTTokenizer", "OpenAIGPTTokenizerFast" if is_tokenizers_available() else None)),
("gpt2", ("GPT2Tokenizer", "GPT2TokenizerFast" if is_tokenizers_available() else None)),
("gptj", ("GPT2Tokenizer", "GPT2TokenizerFast" if is_tokenizers_available() else None)),
("opt", ("GPT2Tokenizer", None)),
("transfo-xl", ("TransfoXLTokenizer", None)),
(
"xlnet",
( (
"XLNetTokenizer" if is_sentencepiece_available() else None, "MBart50Tokenizer" if is_sentencepiece_available() else None,
"XLNetTokenizerFast" if is_tokenizers_available() else None, "MBart50TokenizerFast" if is_tokenizers_available() else None,
), ),
), ),
("flaubert", ("FlaubertTokenizer", None)), ("megatron-bert", ("BertTokenizer", "BertTokenizerFast" if is_tokenizers_available() else None)),
("xlm", ("XLMTokenizer", None)), ("mluke", ("MLukeTokenizer" if is_sentencepiece_available() else None, None)),
("ctrl", ("CTRLTokenizer", None)), ("mobilebert", ("MobileBertTokenizer", "MobileBertTokenizerFast" if is_tokenizers_available() else None)),
("fsmt", ("FSMTTokenizer", None)), ("mpnet", ("MPNetTokenizer", "MPNetTokenizerFast" if is_tokenizers_available() else None)),
("bert-generation", ("BertGenerationTokenizer" if is_sentencepiece_available() else None, None)),
("deberta", ("DebertaTokenizer", "DebertaTokenizerFast" if is_tokenizers_available() else None)),
( (
"deberta-v2", "mt5",
( (
"DebertaV2Tokenizer" if is_sentencepiece_available() else None, "MT5Tokenizer" if is_sentencepiece_available() else None,
"DebertaV2TokenizerFast" if is_tokenizers_available() else None, "MT5TokenizerFast" if is_tokenizers_available() else None,
), ),
), ),
("rag", ("RagTokenizer", None)),
("xlm-prophetnet", ("XLMProphetNetTokenizer" if is_sentencepiece_available() else None, None)),
("speech_to_text", ("Speech2TextTokenizer" if is_sentencepiece_available() else None, None)),
("speech_to_text_2", ("Speech2Text2Tokenizer", None)),
("m2m_100", ("M2M100Tokenizer" if is_sentencepiece_available() else None, None)),
("prophetnet", ("ProphetNetTokenizer", None)),
("mpnet", ("MPNetTokenizer", "MPNetTokenizerFast" if is_tokenizers_available() else None)),
("tapas", ("TapasTokenizer", None)),
("led", ("LEDTokenizer", "LEDTokenizerFast" if is_tokenizers_available() else None)),
("convbert", ("ConvBertTokenizer", "ConvBertTokenizerFast" if is_tokenizers_available() else None)),
( (
"big_bird", "nystromformer",
( (
"BigBirdTokenizer" if is_sentencepiece_available() else None, "AlbertTokenizer" if is_sentencepiece_available() else None,
"BigBirdTokenizerFast" if is_tokenizers_available() else None, "AlbertTokenizerFast" if is_tokenizers_available() else None,
), ),
), ),
("ibert", ("RobertaTokenizer", "RobertaTokenizerFast" if is_tokenizers_available() else None)), ("openai-gpt", ("OpenAIGPTTokenizer", "OpenAIGPTTokenizerFast" if is_tokenizers_available() else None)),
("qdqbert", ("BertTokenizer", "BertTokenizerFast" if is_tokenizers_available() else None)), ("opt", ("GPT2Tokenizer", None)),
("wav2vec2", ("Wav2Vec2CTCTokenizer", None)),
("hubert", ("Wav2Vec2CTCTokenizer", None)),
("gpt_neo", ("GPT2Tokenizer", "GPT2TokenizerFast" if is_tokenizers_available() else None)),
("luke", ("LukeTokenizer", None)),
("mluke", ("MLukeTokenizer" if is_sentencepiece_available() else None, None)),
("bigbird_pegasus", ("PegasusTokenizer", "PegasusTokenizerFast" if is_tokenizers_available() else None)),
("canine", ("CanineTokenizer", None)),
("bertweet", ("BertweetTokenizer", None)),
("bert-japanese", ("BertJapaneseTokenizer", None)),
("splinter", ("SplinterTokenizer", "SplinterTokenizerFast")),
("byt5", ("ByT5Tokenizer", None)),
( (
"cpm", "pegasus",
( (
"CpmTokenizer" if is_sentencepiece_available() else None, "PegasusTokenizer" if is_sentencepiece_available() else None,
"CpmTokenizerFast" if is_tokenizers_available() else None, "PegasusTokenizerFast" if is_tokenizers_available() else None,
), ),
), ),
("herbert", ("HerbertTokenizer", "HerbertTokenizerFast" if is_tokenizers_available() else None)),
("phobert", ("PhobertTokenizer", None)),
("bartpho", ("BartphoTokenizer", None)),
( (
"barthez", "perceiver",
( (
"BarthezTokenizer" if is_sentencepiece_available() else None, "PerceiverTokenizer",
"BarthezTokenizerFast" if is_tokenizers_available() else None, None,
), ),
), ),
("phobert", ("PhobertTokenizer", None)),
("plbart", ("PLBartTokenizer" if is_sentencepiece_available() else None, None)),
("prophetnet", ("ProphetNetTokenizer", None)),
("qdqbert", ("BertTokenizer", "BertTokenizerFast" if is_tokenizers_available() else None)),
("rag", ("RagTokenizer", None)),
("realm", ("RealmTokenizer", "RealmTokenizerFast" if is_tokenizers_available() else None)),
( (
"mbart50", "reformer",
( (
"MBart50Tokenizer" if is_sentencepiece_available() else None, "ReformerTokenizer" if is_sentencepiece_available() else None,
"MBart50TokenizerFast" if is_tokenizers_available() else None, "ReformerTokenizerFast" if is_tokenizers_available() else None,
), ),
), ),
( (
...@@ -220,21 +206,29 @@ else: ...@@ -220,21 +206,29 @@ else:
"RemBertTokenizerFast" if is_tokenizers_available() else None, "RemBertTokenizerFast" if is_tokenizers_available() else None,
), ),
), ),
("retribert", ("RetriBertTokenizer", "RetriBertTokenizerFast" if is_tokenizers_available() else None)),
("roberta", ("RobertaTokenizer", "RobertaTokenizerFast" if is_tokenizers_available() else None)),
("roformer", ("RoFormerTokenizer", "RoFormerTokenizerFast" if is_tokenizers_available() else None)),
("speech_to_text", ("Speech2TextTokenizer" if is_sentencepiece_available() else None, None)),
("speech_to_text_2", ("Speech2Text2Tokenizer", None)),
("splinter", ("SplinterTokenizer", "SplinterTokenizerFast")),
( (
"clip", "squeezebert",
( ("SqueezeBertTokenizer", "SqueezeBertTokenizerFast" if is_tokenizers_available() else None),
"CLIPTokenizer",
"CLIPTokenizerFast" if is_tokenizers_available() else None,
),
), ),
("wav2vec2_phoneme", ("Wav2Vec2PhonemeCTCTokenizer", None)),
( (
"perceiver", "t5",
( (
"PerceiverTokenizer", "T5Tokenizer" if is_sentencepiece_available() else None,
None, "T5TokenizerFast" if is_tokenizers_available() else None,
), ),
), ),
("tapas", ("TapasTokenizer", None)),
("tapex", ("TapexTokenizer", None)),
("transfo-xl", ("TransfoXLTokenizer", None)),
("visual_bert", ("BertTokenizer", "BertTokenizerFast" if is_tokenizers_available() else None)),
("wav2vec2", ("Wav2Vec2CTCTokenizer", None)),
("wav2vec2_phoneme", ("Wav2Vec2PhonemeCTCTokenizer", None)),
( (
"xglm", "xglm",
( (
...@@ -242,16 +236,23 @@ else: ...@@ -242,16 +236,23 @@ else:
"XGLMTokenizerFast" if is_tokenizers_available() else None, "XGLMTokenizerFast" if is_tokenizers_available() else None,
), ),
), ),
("visual_bert", ("BertTokenizer", "BertTokenizerFast" if is_tokenizers_available() else None)), ("xlm", ("XLMTokenizer", None)),
("megatron-bert", ("BertTokenizer", "BertTokenizerFast" if is_tokenizers_available() else None)), ("xlm-prophetnet", ("XLMProphetNetTokenizer" if is_sentencepiece_available() else None, None)),
( (
"nystromformer", "xlm-roberta",
( (
"AlbertTokenizer" if is_sentencepiece_available() else None, "XLMRobertaTokenizer" if is_sentencepiece_available() else None,
"AlbertTokenizerFast" if is_tokenizers_available() else None, "XLMRobertaTokenizerFast" if is_tokenizers_available() else None,
), ),
), ),
("xlm-roberta-xl", ("RobertaTokenizer", "RobertaTokenizerFast" if is_tokenizers_available() else None)), ("xlm-roberta-xl", ("RobertaTokenizer", "RobertaTokenizerFast" if is_tokenizers_available() else None)),
(
"xlnet",
(
"XLNetTokenizer" if is_sentencepiece_available() else None,
"XLNetTokenizerFast" if is_tokenizers_available() else None,
),
),
( (
"yoso", "yoso",
( (
...@@ -259,7 +260,6 @@ else: ...@@ -259,7 +260,6 @@ else:
"AlbertTokenizerFast" if is_tokenizers_available() else None, "AlbertTokenizerFast" if is_tokenizers_available() else None,
), ),
), ),
("data2vec-text", ("RobertaTokenizer", "RobertaTokenizerFast" if is_tokenizers_available() else None)),
] ]
) )
...@@ -277,7 +277,10 @@ def tokenizer_class_from_name(class_name: str): ...@@ -277,7 +277,10 @@ def tokenizer_class_from_name(class_name: str):
module_name = model_type_to_module_name(module_name) module_name = model_type_to_module_name(module_name)
module = importlib.import_module(f".{module_name}", "transformers.models") module = importlib.import_module(f".{module_name}", "transformers.models")
return getattr(module, class_name) try:
return getattr(module, class_name)
except AttributeError:
continue
for config, tokenizers in TOKENIZER_MAPPING._extra_content.items(): for config, tokenizers in TOKENIZER_MAPPING._extra_content.items():
for tokenizer in tokenizers: for tokenizer in tokenizers:
......
tests/models/auto/test_configuration_auto.py
View file @ ddb1a47e
...@@ -14,6 +14,8 @@ ...@@ -14,6 +14,8 @@
# limitations under the License. # limitations under the License.
import importlib import importlib
import json
import os
import sys import sys
import tempfile import tempfile
import unittest import unittest
...@@ -56,14 +58,14 @@ class AutoConfigTest(unittest.TestCase): ...@@ -56,14 +58,14 @@ class AutoConfigTest(unittest.TestCase):
self.assertIsInstance(config, RobertaConfig) self.assertIsInstance(config, RobertaConfig)
def test_pattern_matching_fallback(self): def test_pattern_matching_fallback(self):
""" with tempfile.TemporaryDirectory() as tmp_dir:
In cases where config.json doesn't include a model_type, # This model name contains bert and roberta, but roberta ends up being picked.
perform a few safety checks on the config mapping's order. folder = os.path.join(tmp_dir, "fake-roberta")
""" os.makedirs(folder, exist_ok=True)
# no key string should be included in a later key string (typical failure case) with open(os.path.join(folder, "config.json"), "w") as f:
keys = list(CONFIG_MAPPING.keys()) f.write(json.dumps({}))
for i, key in enumerate(keys): config = AutoConfig.from_pretrained(folder)
self.assertFalse(any(key in later_key for later_key in keys[i + 1 :])) self.assertEqual(type(config), RobertaConfig)
def test_new_config_registration(self): def test_new_config_registration(self):
try: try:
......
utils/sort_auto_mappings.py 0 → 100644
View file @ ddb1a47e
# coding=utf-8
# Copyright 2022 The HuggingFace Inc. team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import argparse
import os
import re
PATH_TO_AUTO_MODULE = "src/transformers/models/auto"
# re pattern that matches mapping introductions:
# SUPER_MODEL_MAPPING_NAMES = OrderedDict or SUPER_MODEL_MAPPING = OrderedDict
_re_intro_mapping = re.compile("[A-Z_]+_MAPPING(\s+|_[A-Z_]+\s+)=\s+OrderedDict")
# re pattern that matches identifiers in mappings
_re_identifier = re.compile(r'\s*\(\s*"(\S[^"]+)"')
def sort_auto_mapping(fname, overwrite: bool = False):
with open(fname, "r", encoding="utf-8") as f:
content = f.read()
lines = content.split("\n")
new_lines = []
line_idx = 0
while line_idx < len(lines):
if _re_intro_mapping.search(lines[line_idx]) is not None:
indent = len(re.search(r"^(\s*)\S", lines[line_idx]).groups()[0]) + 8
# Start of a new mapping!
while not lines[line_idx].startswith(" " * indent + "("):
new_lines.append(lines[line_idx])
line_idx += 1
blocks = []
while lines[line_idx].strip() != "]":
# Blocks either fit in one line or not
if lines[line_idx].strip() == "(":
start_idx = line_idx
while not lines[line_idx].startswith(" " * indent + ")"):
line_idx += 1
blocks.append("\n".join(lines[start_idx : line_idx + 1]))
else:
blocks.append(lines[line_idx])
line_idx += 1
# Sort blocks by their identifiers
blocks = sorted(blocks, key=lambda x: _re_identifier.search(x).groups()[0])
new_lines += blocks
else:
new_lines.append(lines[line_idx])
line_idx += 1
if overwrite:
with open(fname, "w", encoding="utf-8") as f:
f.write("\n".join(new_lines))
elif "\n".join(new_lines) != content:
return True
def sort_all_auto_mappings(overwrite: bool = False):
fnames = [os.path.join(PATH_TO_AUTO_MODULE, f) for f in os.listdir(PATH_TO_AUTO_MODULE) if f.endswith(".py")]
diffs = [sort_auto_mapping(fname, overwrite=overwrite) for fname in fnames]
if not overwrite and any(diffs):
failures = [f for f, d in zip(fnames, diffs) if d]
raise ValueError(
f"The following files have auto mappings that need sorting: {', '.join(failures)}. Run `make style` to fix"
" this."
)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--check_only", action="store_true", help="Whether to only check or fix style.")
args = parser.parse_args()
sort_all_auto_mappings(not args.check_only)
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment