modeling_auto.py 44.6 KB
Newer Older
thomwolf's avatar
thomwolf committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
# coding=utf-8
# Copyright 2018 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.
Sylvain Gugger's avatar
Sylvain Gugger committed
15
""" Auto Model class."""
thomwolf's avatar
thomwolf committed
16

17
import warnings
Julien Chaumond's avatar
Julien Chaumond committed
18
from collections import OrderedDict
thomwolf's avatar
thomwolf committed
19

Sylvain Gugger's avatar
Sylvain Gugger committed
20
from ...utils import logging
21
22
from .auto_factory import _BaseAutoModelClass, _LazyAutoMapping, auto_class_update
from .configuration_auto import CONFIG_MAPPING_NAMES
thomwolf's avatar
thomwolf committed
23

thomwolf's avatar
thomwolf committed
24

Lysandre Debut's avatar
Lysandre Debut committed
25
logger = logging.get_logger(__name__)
thomwolf's avatar
thomwolf committed
26
27


28
MODEL_MAPPING_NAMES = OrderedDict(
Julien Chaumond's avatar
Julien Chaumond committed
29
    [
30
        # Base model mapping
31
32
        ("albert", "AlbertModel"),
        ("bart", "BartModel"),
33
        ("beit", "BeitModel"),
34
35
36
37
38
39
        ("bert", "BertModel"),
        ("bert-generation", "BertGenerationEncoder"),
        ("big_bird", "BigBirdModel"),
        ("bigbird_pegasus", "BigBirdPegasusModel"),
        ("blenderbot", "BlenderbotModel"),
        ("blenderbot-small", "BlenderbotSmallModel"),
Younes Belkada's avatar
Younes Belkada committed
40
        ("bloom", "BloomModel"),
41
        ("camembert", "CamembertModel"),
42
43
        ("canine", "CanineModel"),
        ("clip", "CLIPModel"),
rooa's avatar
rooa committed
44
        ("codegen", "CodeGenModel"),
45
        ("conditional_detr", "ConditionalDetrModel"),
46
47
48
        ("convbert", "ConvBertModel"),
        ("convnext", "ConvNextModel"),
        ("ctrl", "CTRLModel"),
NielsRogge's avatar
NielsRogge committed
49
        ("cvt", "CvtModel"),
50
51
52
53
54
55
56
        ("data2vec-audio", "Data2VecAudioModel"),
        ("data2vec-text", "Data2VecTextModel"),
        ("data2vec-vision", "Data2VecVisionModel"),
        ("deberta", "DebertaModel"),
        ("deberta-v2", "DebertaV2Model"),
        ("decision_transformer", "DecisionTransformerModel"),
        ("decision_transformer_gpt2", "DecisionTransformerGPT2Model"),
NielsRogge's avatar
NielsRogge committed
57
        ("deformable_detr", "DeformableDetrModel"),
58
59
        ("deit", "DeiTModel"),
        ("detr", "DetrModel"),
60
        ("distilbert", "DistilBertModel"),
NielsRogge's avatar
NielsRogge committed
61
        ("donut-swin", "DonutSwinModel"),
62
63
64
        ("dpr", "DPRQuestionEncoder"),
        ("dpt", "DPTModel"),
        ("electra", "ElectraModel"),
65
        ("ernie", "ErnieModel"),
66
        ("esm", "EsmModel"),
67
68
69
70
71
72
73
        ("flaubert", "FlaubertModel"),
        ("flava", "FlavaModel"),
        ("fnet", "FNetModel"),
        ("fsmt", "FSMTModel"),
        ("funnel", ("FunnelModel", "FunnelBaseModel")),
        ("glpn", "GLPNModel"),
        ("gpt2", "GPT2Model"),
74
        ("gpt_neo", "GPTNeoModel"),
75
        ("gpt_neox", "GPTNeoXModel"),
76
        ("gpt_neox_japanese", "GPTNeoXJapaneseModel"),
77
        ("gptj", "GPTJModel"),
78
        ("groupvit", "GroupViTModel"),
79
        ("hubert", "HubertModel"),
80
81
82
83
        ("ibert", "IBertModel"),
        ("imagegpt", "ImageGPTModel"),
        ("layoutlm", "LayoutLMModel"),
        ("layoutlmv2", "LayoutLMv2Model"),
NielsRogge's avatar
NielsRogge committed
84
        ("layoutlmv3", "LayoutLMv3Model"),
85
        ("led", "LEDModel"),
86
        ("levit", "LevitModel"),
NielsRogge's avatar
NielsRogge committed
87
        ("lilt", "LiltModel"),
88
        ("longformer", "LongformerModel"),
Daniel Stancl's avatar
Daniel Stancl committed
89
        ("longt5", "LongT5Model"),
90
91
92
        ("luke", "LukeModel"),
        ("lxmert", "LxmertModel"),
        ("m2m_100", "M2M100Model"),
93
        ("marian", "MarianModel"),
NielsRogge's avatar
NielsRogge committed
94
        ("markuplm", "MarkupLMModel"),
95
        ("maskformer", "MaskFormerModel"),
96
        ("mbart", "MBartModel"),
Chan Woo Kim's avatar
Chan Woo Kim committed
97
        ("mctct", "MCTCTModel"),
98
99
        ("megatron-bert", "MegatronBertModel"),
        ("mobilebert", "MobileBertModel"),
100
        ("mobilevit", "MobileViTModel"),
101
102
        ("mpnet", "MPNetModel"),
        ("mt5", "MT5Model"),
StevenTang1998's avatar
StevenTang1998 committed
103
        ("mvp", "MvpModel"),
104
        ("nezha", "NezhaModel"),
Lysandre Debut's avatar
Lysandre Debut committed
105
        ("nllb", "M2M100Model"),
106
107
        ("nystromformer", "NystromformerModel"),
        ("openai-gpt", "OpenAIGPTModel"),
Younes Belkada's avatar
Younes Belkada committed
108
        ("opt", "OPTModel"),
109
        ("owlvit", "OwlViTModel"),
110
        ("pegasus", "PegasusModel"),
Jason Phang's avatar
Jason Phang committed
111
        ("pegasus_x", "PegasusXModel"),
112
113
114
115
116
117
118
119
120
121
        ("perceiver", "PerceiverModel"),
        ("plbart", "PLBartModel"),
        ("poolformer", "PoolFormerModel"),
        ("prophetnet", "ProphetNetModel"),
        ("qdqbert", "QDQBertModel"),
        ("reformer", "ReformerModel"),
        ("regnet", "RegNetModel"),
        ("rembert", "RemBertModel"),
        ("resnet", "ResNetModel"),
        ("retribert", "RetriBertModel"),
122
        ("roberta", "RobertaModel"),
123
124
125
126
127
128
        ("roformer", "RoFormerModel"),
        ("segformer", "SegformerModel"),
        ("sew", "SEWModel"),
        ("sew-d", "SEWDModel"),
        ("speech_to_text", "Speech2TextModel"),
        ("splinter", "SplinterModel"),
129
        ("squeezebert", "SqueezeBertModel"),
130
        ("swin", "SwinModel"),
131
        ("swinv2", "Swinv2Model"),
132
        ("t5", "T5Model"),
133
        ("table-transformer", "TableTransformerModel"),
134
        ("tapas", "TapasModel"),
135
        ("time_series_transformer", "TimeSeriesTransformerModel"),
Carl's avatar
Carl committed
136
        ("trajectory_transformer", "TrajectoryTransformerModel"),
137
        ("transfo-xl", "TransfoXLModel"),
138
139
140
        ("unispeech", "UniSpeechModel"),
        ("unispeech-sat", "UniSpeechSatModel"),
        ("van", "VanModel"),
NielsRogge's avatar
NielsRogge committed
141
        ("videomae", "VideoMAEModel"),
142
143
144
145
146
        ("vilt", "ViltModel"),
        ("vision-text-dual-encoder", "VisionTextDualEncoderModel"),
        ("visual_bert", "VisualBertModel"),
        ("vit", "ViTModel"),
        ("vit_mae", "ViTMAEModel"),
147
        ("vit_msn", "ViTMSNModel"),
148
        ("wav2vec2", "Wav2Vec2Model"),
149
        ("wav2vec2-conformer", "Wav2Vec2ConformerModel"),
150
        ("wavlm", "WavLMModel"),
151
        ("whisper", "WhisperModel"),
NielsRogge's avatar
NielsRogge committed
152
        ("xclip", "XCLIPModel"),
153
        ("xglm", "XGLMModel"),
154
155
        ("xlm", "XLMModel"),
        ("xlm-prophetnet", "XLMProphetNetModel"),
156
157
158
159
160
        ("xlm-roberta", "XLMRobertaModel"),
        ("xlm-roberta-xl", "XLMRobertaXLModel"),
        ("xlnet", "XLNetModel"),
        ("yolos", "YolosModel"),
        ("yoso", "YosoModel"),
Julien Chaumond's avatar
Julien Chaumond committed
161
162
163
    ]
)

164
MODEL_FOR_PRETRAINING_MAPPING_NAMES = OrderedDict(
thomwolf's avatar
thomwolf committed
165
    [
166
        # Model for pre-training mapping
167
168
169
170
        ("albert", "AlbertForPreTraining"),
        ("bart", "BartForConditionalGeneration"),
        ("bert", "BertForPreTraining"),
        ("big_bird", "BigBirdForPreTraining"),
Younes Belkada's avatar
Younes Belkada committed
171
        ("bloom", "BloomForCausalLM"),
172
        ("camembert", "CamembertForMaskedLM"),
173
        ("ctrl", "CTRLLMHeadModel"),
174
175
176
177
        ("data2vec-text", "Data2VecTextForMaskedLM"),
        ("deberta", "DebertaForMaskedLM"),
        ("deberta-v2", "DebertaV2ForMaskedLM"),
        ("distilbert", "DistilBertForMaskedLM"),
178
        ("electra", "ElectraForPreTraining"),
179
        ("ernie", "ErnieForPreTraining"),
180
181
182
183
        ("flaubert", "FlaubertWithLMHeadModel"),
        ("flava", "FlavaForPreTraining"),
        ("fnet", "FNetForPreTraining"),
        ("fsmt", "FSMTForConditionalGeneration"),
184
        ("funnel", "FunnelForPreTraining"),
185
186
187
188
        ("gpt2", "GPT2LMHeadModel"),
        ("ibert", "IBertForMaskedLM"),
        ("layoutlm", "LayoutLMForMaskedLM"),
        ("longformer", "LongformerForMaskedLM"),
189
        ("luke", "LukeForMaskedLM"),
190
191
192
        ("lxmert", "LxmertForPreTraining"),
        ("megatron-bert", "MegatronBertForPreTraining"),
        ("mobilebert", "MobileBertForPreTraining"),
193
        ("mpnet", "MPNetForMaskedLM"),
StevenTang1998's avatar
StevenTang1998 committed
194
        ("mvp", "MvpForConditionalGeneration"),
195
        ("nezha", "NezhaForPreTraining"),
196
197
198
        ("openai-gpt", "OpenAIGPTLMHeadModel"),
        ("retribert", "RetriBertModel"),
        ("roberta", "RobertaForMaskedLM"),
199
        ("splinter", "SplinterForPreTraining"),
200
201
        ("squeezebert", "SqueezeBertForMaskedLM"),
        ("t5", "T5ForConditionalGeneration"),
202
        ("tapas", "TapasForMaskedLM"),
203
        ("transfo-xl", "TransfoXLLMHeadModel"),
204
        ("unispeech", "UniSpeechForPreTraining"),
205
        ("unispeech-sat", "UniSpeechSatForPreTraining"),
NielsRogge's avatar
NielsRogge committed
206
        ("videomae", "VideoMAEForPreTraining"),
207
208
209
        ("visual_bert", "VisualBertForPreTraining"),
        ("vit_mae", "ViTMAEForPreTraining"),
        ("wav2vec2", "Wav2Vec2ForPreTraining"),
210
        ("wav2vec2-conformer", "Wav2Vec2ConformerForPreTraining"),
211
212
213
214
        ("xlm", "XLMWithLMHeadModel"),
        ("xlm-roberta", "XLMRobertaForMaskedLM"),
        ("xlm-roberta-xl", "XLMRobertaXLForMaskedLM"),
        ("xlnet", "XLNetLMHeadModel"),
thomwolf's avatar
thomwolf committed
215
216
217
    ]
)

218
MODEL_WITH_LM_HEAD_MAPPING_NAMES = OrderedDict(
219
    [
220
        # Model with LM heads mapping
221
222
223
        ("albert", "AlbertForMaskedLM"),
        ("bart", "BartForConditionalGeneration"),
        ("bert", "BertForMaskedLM"),
224
225
226
        ("big_bird", "BigBirdForMaskedLM"),
        ("bigbird_pegasus", "BigBirdPegasusForConditionalGeneration"),
        ("blenderbot-small", "BlenderbotSmallForConditionalGeneration"),
Younes Belkada's avatar
Younes Belkada committed
227
        ("bloom", "BloomForCausalLM"),
228
        ("camembert", "CamembertForMaskedLM"),
rooa's avatar
rooa committed
229
        ("codegen", "CodeGenForCausalLM"),
230
        ("convbert", "ConvBertForMaskedLM"),
231
        ("ctrl", "CTRLLMHeadModel"),
232
233
234
235
        ("data2vec-text", "Data2VecTextForMaskedLM"),
        ("deberta", "DebertaForMaskedLM"),
        ("deberta-v2", "DebertaV2ForMaskedLM"),
        ("distilbert", "DistilBertForMaskedLM"),
236
237
        ("electra", "ElectraForMaskedLM"),
        ("encoder-decoder", "EncoderDecoderModel"),
238
        ("ernie", "ErnieForMaskedLM"),
239
        ("esm", "EsmForMaskedLM"),
240
241
242
        ("flaubert", "FlaubertWithLMHeadModel"),
        ("fnet", "FNetForMaskedLM"),
        ("fsmt", "FSMTForConditionalGeneration"),
243
        ("funnel", "FunnelForMaskedLM"),
244
245
        ("gpt2", "GPT2LMHeadModel"),
        ("gpt_neo", "GPTNeoForCausalLM"),
246
        ("gpt_neox", "GPTNeoXForCausalLM"),
247
        ("gpt_neox_japanese", "GPTNeoXJapaneseForCausalLM"),
248
249
250
251
252
        ("gptj", "GPTJForCausalLM"),
        ("ibert", "IBertForMaskedLM"),
        ("layoutlm", "LayoutLMForMaskedLM"),
        ("led", "LEDForConditionalGeneration"),
        ("longformer", "LongformerForMaskedLM"),
Daniel Stancl's avatar
Daniel Stancl committed
253
        ("longt5", "LongT5ForConditionalGeneration"),
254
        ("luke", "LukeForMaskedLM"),
255
256
257
258
        ("m2m_100", "M2M100ForConditionalGeneration"),
        ("marian", "MarianMTModel"),
        ("megatron-bert", "MegatronBertForCausalLM"),
        ("mobilebert", "MobileBertForMaskedLM"),
259
        ("mpnet", "MPNetForMaskedLM"),
StevenTang1998's avatar
StevenTang1998 committed
260
        ("mvp", "MvpForConditionalGeneration"),
261
        ("nezha", "NezhaForMaskedLM"),
Lysandre Debut's avatar
Lysandre Debut committed
262
        ("nllb", "M2M100ForConditionalGeneration"),
263
264
        ("nystromformer", "NystromformerForMaskedLM"),
        ("openai-gpt", "OpenAIGPTLMHeadModel"),
Jason Phang's avatar
Jason Phang committed
265
        ("pegasus_x", "PegasusXForConditionalGeneration"),
266
267
268
269
270
271
272
273
274
        ("plbart", "PLBartForConditionalGeneration"),
        ("qdqbert", "QDQBertForMaskedLM"),
        ("reformer", "ReformerModelWithLMHead"),
        ("rembert", "RemBertForMaskedLM"),
        ("roberta", "RobertaForMaskedLM"),
        ("roformer", "RoFormerForMaskedLM"),
        ("speech_to_text", "Speech2TextForConditionalGeneration"),
        ("squeezebert", "SqueezeBertForMaskedLM"),
        ("t5", "T5ForConditionalGeneration"),
275
        ("tapas", "TapasForMaskedLM"),
276
277
        ("transfo-xl", "TransfoXLLMHeadModel"),
        ("wav2vec2", "Wav2Vec2ForMaskedLM"),
278
        ("whisper", "WhisperForConditionalGeneration"),
279
280
281
282
283
        ("xlm", "XLMWithLMHeadModel"),
        ("xlm-roberta", "XLMRobertaForMaskedLM"),
        ("xlm-roberta-xl", "XLMRobertaXLForMaskedLM"),
        ("xlnet", "XLNetLMHeadModel"),
        ("yoso", "YosoForMaskedLM"),
284
285
286
    ]
)

287
MODEL_FOR_CAUSAL_LM_MAPPING_NAMES = OrderedDict(
288
    [
289
        # Model for Causal LM mapping
290
291
292
293
294
295
296
        ("bart", "BartForCausalLM"),
        ("bert", "BertLMHeadModel"),
        ("bert-generation", "BertGenerationDecoder"),
        ("big_bird", "BigBirdForCausalLM"),
        ("bigbird_pegasus", "BigBirdPegasusForCausalLM"),
        ("blenderbot", "BlenderbotForCausalLM"),
        ("blenderbot-small", "BlenderbotSmallForCausalLM"),
Younes Belkada's avatar
Younes Belkada committed
297
        ("bloom", "BloomForCausalLM"),
298
        ("camembert", "CamembertForCausalLM"),
rooa's avatar
rooa committed
299
        ("codegen", "CodeGenForCausalLM"),
300
301
302
        ("ctrl", "CTRLLMHeadModel"),
        ("data2vec-text", "Data2VecTextForCausalLM"),
        ("electra", "ElectraForCausalLM"),
303
        ("ernie", "ErnieForCausalLM"),
304
305
        ("gpt2", "GPT2LMHeadModel"),
        ("gpt_neo", "GPTNeoForCausalLM"),
306
        ("gpt_neox", "GPTNeoXForCausalLM"),
307
        ("gpt_neox_japanese", "GPTNeoXJapaneseForCausalLM"),
308
309
310
311
        ("gptj", "GPTJForCausalLM"),
        ("marian", "MarianForCausalLM"),
        ("mbart", "MBartForCausalLM"),
        ("megatron-bert", "MegatronBertForCausalLM"),
StevenTang1998's avatar
StevenTang1998 committed
312
        ("mvp", "MvpForCausalLM"),
313
314
315
        ("openai-gpt", "OpenAIGPTLMHeadModel"),
        ("opt", "OPTForCausalLM"),
        ("pegasus", "PegasusForCausalLM"),
Gunjan Chhablani's avatar
Gunjan Chhablani committed
316
        ("plbart", "PLBartForCausalLM"),
317
        ("prophetnet", "ProphetNetForCausalLM"),
318
        ("qdqbert", "QDQBertLMHeadModel"),
319
        ("reformer", "ReformerModelWithLMHead"),
320
321
        ("rembert", "RemBertForCausalLM"),
        ("roberta", "RobertaForCausalLM"),
322
323
        ("roformer", "RoFormerForCausalLM"),
        ("speech_to_text_2", "Speech2Text2ForCausalLM"),
324
        ("transfo-xl", "TransfoXLLMHeadModel"),
325
326
        ("trocr", "TrOCRForCausalLM"),
        ("xglm", "XGLMForCausalLM"),
327
328
        ("xlm", "XLMWithLMHeadModel"),
        ("xlm-prophetnet", "XLMProphetNetForCausalLM"),
329
330
331
        ("xlm-roberta", "XLMRobertaForCausalLM"),
        ("xlm-roberta-xl", "XLMRobertaXLForCausalLM"),
        ("xlnet", "XLNetLMHeadModel"),
332
333
334
    ]
)

NielsRogge's avatar
NielsRogge committed
335
336
337
338
MODEL_FOR_MASKED_IMAGE_MODELING_MAPPING_NAMES = OrderedDict(
    [
        ("deit", "DeiTForMaskedImageModeling"),
        ("swin", "SwinForMaskedImageModeling"),
339
        ("swinv2", "Swinv2ForMaskedImageModeling"),
340
        ("vit", "ViTForMaskedImageModeling"),
NielsRogge's avatar
NielsRogge committed
341
342
343
344
    ]
)


NielsRogge's avatar
NielsRogge committed
345
346
347
348
349
350
351
MODEL_FOR_CAUSAL_IMAGE_MODELING_MAPPING_NAMES = OrderedDict(
    # Model for Causal Image Modeling mapping
    [
        ("imagegpt", "ImageGPTForCausalImageModeling"),
    ]
)

352
MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING_NAMES = OrderedDict(
353
354
    [
        # Model for Image Classification mapping
355
        ("beit", "BeitForImageClassification"),
356
        ("convnext", "ConvNextForImageClassification"),
NielsRogge's avatar
NielsRogge committed
357
        ("cvt", "CvtForImageClassification"),
358
        ("data2vec-vision", "Data2VecVisionForImageClassification"),
359
        ("deit", ("DeiTForImageClassification", "DeiTForImageClassificationWithTeacher")),
NielsRogge's avatar
NielsRogge committed
360
        ("imagegpt", "ImageGPTForImageClassification"),
361
        ("levit", ("LevitForImageClassification", "LevitForImageClassificationWithTeacher")),
362
        ("mobilevit", "MobileViTForImageClassification"),
NielsRogge's avatar
NielsRogge committed
363
364
365
366
367
368
369
370
        (
            "perceiver",
            (
                "PerceiverForImageClassificationLearned",
                "PerceiverForImageClassificationFourier",
                "PerceiverForImageClassificationConvProcessing",
            ),
        ),
371
372
373
374
        ("poolformer", "PoolFormerForImageClassification"),
        ("regnet", "RegNetForImageClassification"),
        ("resnet", "ResNetForImageClassification"),
        ("segformer", "SegformerForImageClassification"),
novice's avatar
novice committed
375
        ("swin", "SwinForImageClassification"),
376
        ("swinv2", "Swinv2ForImageClassification"),
377
        ("van", "VanForImageClassification"),
378
        ("vit", "ViTForImageClassification"),
379
        ("vit_msn", "ViTMSNForImageClassification"),
380
381
382
    ]
)

383
384
MODEL_FOR_IMAGE_SEGMENTATION_MAPPING_NAMES = OrderedDict(
    [
385
        # Do not add new models here, this class will be deprecated in the future.
386
387
388
389
390
        # Model for Image Segmentation mapping
        ("detr", "DetrForSegmentation"),
    ]
)

391
392
393
394
MODEL_FOR_SEMANTIC_SEGMENTATION_MAPPING_NAMES = OrderedDict(
    [
        # Model for Semantic Segmentation mapping
        ("beit", "BeitForSemanticSegmentation"),
395
        ("data2vec-vision", "Data2VecVisionForSemanticSegmentation"),
NielsRogge's avatar
NielsRogge committed
396
        ("dpt", "DPTForSemanticSegmentation"),
397
        ("mobilevit", "MobileViTForSemanticSegmentation"),
398
        ("segformer", "SegformerForSemanticSegmentation"),
399
400
401
    ]
)

402
403
404
405
406
407
408
MODEL_FOR_INSTANCE_SEGMENTATION_MAPPING_NAMES = OrderedDict(
    [
        # Model for Instance Segmentation mapping
        ("maskformer", "MaskFormerForInstanceSegmentation"),
    ]
)

NielsRogge's avatar
NielsRogge committed
409
410
411
412
413
414
MODEL_FOR_VIDEO_CLASSIFICATION_MAPPING_NAMES = OrderedDict(
    [
        ("videomae", "VideoMAEForVideoClassification"),
    ]
)

415
416
417
418
419
420
MODEL_FOR_VISION_2_SEQ_MAPPING_NAMES = OrderedDict(
    [
        ("vision-encoder-decoder", "VisionEncoderDecoderModel"),
    ]
)

421
MODEL_FOR_MASKED_LM_MAPPING_NAMES = OrderedDict(
422
    [
423
        # Model for Masked LM mapping
424
425
        ("albert", "AlbertForMaskedLM"),
        ("bart", "BartForConditionalGeneration"),
426
427
        ("bert", "BertForMaskedLM"),
        ("big_bird", "BigBirdForMaskedLM"),
428
        ("camembert", "CamembertForMaskedLM"),
429
        ("convbert", "ConvBertForMaskedLM"),
430
        ("data2vec-text", "Data2VecTextForMaskedLM"),
431
432
433
        ("deberta", "DebertaForMaskedLM"),
        ("deberta-v2", "DebertaV2ForMaskedLM"),
        ("distilbert", "DistilBertForMaskedLM"),
434
        ("electra", "ElectraForMaskedLM"),
435
        ("ernie", "ErnieForMaskedLM"),
436
437
        ("flaubert", "FlaubertWithLMHeadModel"),
        ("fnet", "FNetForMaskedLM"),
438
        ("funnel", "FunnelForMaskedLM"),
439
440
441
        ("ibert", "IBertForMaskedLM"),
        ("layoutlm", "LayoutLMForMaskedLM"),
        ("longformer", "LongformerForMaskedLM"),
Ryokan RI's avatar
Ryokan RI committed
442
        ("luke", "LukeForMaskedLM"),
443
444
445
        ("mbart", "MBartForConditionalGeneration"),
        ("megatron-bert", "MegatronBertForMaskedLM"),
        ("mobilebert", "MobileBertForMaskedLM"),
446
        ("mpnet", "MPNetForMaskedLM"),
StevenTang1998's avatar
StevenTang1998 committed
447
        ("mvp", "MvpForConditionalGeneration"),
448
        ("nezha", "NezhaForMaskedLM"),
449
450
451
452
453
454
455
456
        ("nystromformer", "NystromformerForMaskedLM"),
        ("perceiver", "PerceiverForMaskedLM"),
        ("qdqbert", "QDQBertForMaskedLM"),
        ("reformer", "ReformerForMaskedLM"),
        ("rembert", "RemBertForMaskedLM"),
        ("roberta", "RobertaForMaskedLM"),
        ("roformer", "RoFormerForMaskedLM"),
        ("squeezebert", "SqueezeBertForMaskedLM"),
457
        ("tapas", "TapasForMaskedLM"),
458
459
460
461
462
        ("wav2vec2", "Wav2Vec2ForMaskedLM"),
        ("xlm", "XLMWithLMHeadModel"),
        ("xlm-roberta", "XLMRobertaForMaskedLM"),
        ("xlm-roberta-xl", "XLMRobertaXLForMaskedLM"),
        ("yoso", "YosoForMaskedLM"),
463
464
465
    ]
)

466
MODEL_FOR_OBJECT_DETECTION_MAPPING_NAMES = OrderedDict(
NielsRogge's avatar
NielsRogge committed
467
468
    [
        # Model for Object Detection mapping
469
        ("conditional_detr", "ConditionalDetrForObjectDetection"),
NielsRogge's avatar
NielsRogge committed
470
        ("deformable_detr", "DeformableDetrForObjectDetection"),
471
        ("detr", "DetrForObjectDetection"),
472
        ("table-transformer", "TableTransformerForObjectDetection"),
473
        ("yolos", "YolosForObjectDetection"),
NielsRogge's avatar
NielsRogge committed
474
475
476
    ]
)

477
478
479
480
481
482
483
MODEL_FOR_ZERO_SHOT_OBJECT_DETECTION_MAPPING_NAMES = OrderedDict(
    [
        # Model for Zero Shot Object Detection mapping
        ("owlvit", "OwlViTForObjectDetection")
    ]
)

484
485
486
487
488
489
490
MODEL_FOR_DEPTH_ESTIMATION_MAPPING_NAMES = OrderedDict(
    [
        # Model for depth estimation mapping
        ("dpt", "DPTForDepthEstimation"),
        ("glpn", "GLPNForDepthEstimation"),
    ]
)
491
MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING_NAMES = OrderedDict(
492
    [
493
        # Model for Seq2Seq Causal LM mapping
494
        ("bart", "BartForConditionalGeneration"),
495
        ("bigbird_pegasus", "BigBirdPegasusForConditionalGeneration"),
496
        ("blenderbot", "BlenderbotForConditionalGeneration"),
497
        ("blenderbot-small", "BlenderbotSmallForConditionalGeneration"),
498
499
500
        ("encoder-decoder", "EncoderDecoderModel"),
        ("fsmt", "FSMTForConditionalGeneration"),
        ("led", "LEDForConditionalGeneration"),
Daniel Stancl's avatar
Daniel Stancl committed
501
        ("longt5", "LongT5ForConditionalGeneration"),
502
        ("m2m_100", "M2M100ForConditionalGeneration"),
503
504
        ("marian", "MarianMTModel"),
        ("mbart", "MBartForConditionalGeneration"),
505
        ("mt5", "MT5ForConditionalGeneration"),
StevenTang1998's avatar
StevenTang1998 committed
506
        ("mvp", "MvpForConditionalGeneration"),
Lysandre Debut's avatar
Lysandre Debut committed
507
        ("nllb", "M2M100ForConditionalGeneration"),
508
        ("pegasus", "PegasusForConditionalGeneration"),
Jason Phang's avatar
Jason Phang committed
509
        ("pegasus_x", "PegasusXForConditionalGeneration"),
510
        ("plbart", "PLBartForConditionalGeneration"),
511
        ("prophetnet", "ProphetNetForConditionalGeneration"),
512
513
        ("t5", "T5ForConditionalGeneration"),
        ("xlm-prophetnet", "XLMProphetNetForConditionalGeneration"),
514
515
516
    ]
)

517
518
519
520
MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING_NAMES = OrderedDict(
    [
        ("speech-encoder-decoder", "SpeechEncoderDecoderModel"),
        ("speech_to_text", "Speech2TextForConditionalGeneration"),
521
        ("whisper", "WhisperForConditionalGeneration"),
522
523
524
    ]
)

525
MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING_NAMES = OrderedDict(
526
    [
527
        # Model for Sequence Classification mapping
528
529
530
        ("albert", "AlbertForSequenceClassification"),
        ("bart", "BartForSequenceClassification"),
        ("bert", "BertForSequenceClassification"),
531
532
        ("big_bird", "BigBirdForSequenceClassification"),
        ("bigbird_pegasus", "BigBirdPegasusForSequenceClassification"),
533
        ("bloom", "BloomForSequenceClassification"),
534
535
536
537
538
        ("camembert", "CamembertForSequenceClassification"),
        ("canine", "CanineForSequenceClassification"),
        ("convbert", "ConvBertForSequenceClassification"),
        ("ctrl", "CTRLForSequenceClassification"),
        ("data2vec-text", "Data2VecTextForSequenceClassification"),
539
540
        ("deberta", "DebertaForSequenceClassification"),
        ("deberta-v2", "DebertaV2ForSequenceClassification"),
541
542
        ("distilbert", "DistilBertForSequenceClassification"),
        ("electra", "ElectraForSequenceClassification"),
543
        ("ernie", "ErnieForSequenceClassification"),
544
        ("esm", "EsmForSequenceClassification"),
545
546
547
        ("flaubert", "FlaubertForSequenceClassification"),
        ("fnet", "FNetForSequenceClassification"),
        ("funnel", "FunnelForSequenceClassification"),
548
549
        ("gpt2", "GPT2ForSequenceClassification"),
        ("gpt_neo", "GPTNeoForSequenceClassification"),
550
551
552
553
        ("gptj", "GPTJForSequenceClassification"),
        ("ibert", "IBertForSequenceClassification"),
        ("layoutlm", "LayoutLMForSequenceClassification"),
        ("layoutlmv2", "LayoutLMv2ForSequenceClassification"),
NielsRogge's avatar
NielsRogge committed
554
        ("layoutlmv3", "LayoutLMv3ForSequenceClassification"),
555
        ("led", "LEDForSequenceClassification"),
NielsRogge's avatar
NielsRogge committed
556
        ("lilt", "LiltForSequenceClassification"),
557
        ("longformer", "LongformerForSequenceClassification"),
558
        ("luke", "LukeForSequenceClassification"),
NielsRogge's avatar
NielsRogge committed
559
        ("markuplm", "MarkupLMForSequenceClassification"),
560
561
562
563
        ("mbart", "MBartForSequenceClassification"),
        ("megatron-bert", "MegatronBertForSequenceClassification"),
        ("mobilebert", "MobileBertForSequenceClassification"),
        ("mpnet", "MPNetForSequenceClassification"),
StevenTang1998's avatar
StevenTang1998 committed
564
        ("mvp", "MvpForSequenceClassification"),
565
        ("nezha", "NezhaForSequenceClassification"),
566
        ("nystromformer", "NystromformerForSequenceClassification"),
567
        ("openai-gpt", "OpenAIGPTForSequenceClassification"),
568
        ("opt", "OPTForSequenceClassification"),
569
570
571
        ("perceiver", "PerceiverForSequenceClassification"),
        ("plbart", "PLBartForSequenceClassification"),
        ("qdqbert", "QDQBertForSequenceClassification"),
572
        ("reformer", "ReformerForSequenceClassification"),
573
574
575
576
        ("rembert", "RemBertForSequenceClassification"),
        ("roberta", "RobertaForSequenceClassification"),
        ("roformer", "RoFormerForSequenceClassification"),
        ("squeezebert", "SqueezeBertForSequenceClassification"),
577
        ("tapas", "TapasForSequenceClassification"),
578
579
580
581
582
583
        ("transfo-xl", "TransfoXLForSequenceClassification"),
        ("xlm", "XLMForSequenceClassification"),
        ("xlm-roberta", "XLMRobertaForSequenceClassification"),
        ("xlm-roberta-xl", "XLMRobertaXLForSequenceClassification"),
        ("xlnet", "XLNetForSequenceClassification"),
        ("yoso", "YosoForSequenceClassification"),
584
585
586
    ]
)

587
MODEL_FOR_QUESTION_ANSWERING_MAPPING_NAMES = OrderedDict(
588
    [
589
        # Model for Question Answering mapping
590
591
592
593
594
        ("albert", "AlbertForQuestionAnswering"),
        ("bart", "BartForQuestionAnswering"),
        ("bert", "BertForQuestionAnswering"),
        ("big_bird", "BigBirdForQuestionAnswering"),
        ("bigbird_pegasus", "BigBirdPegasusForQuestionAnswering"),
595
        ("bloom", "BloomForQuestionAnswering"),
596
597
598
599
600
601
602
603
        ("camembert", "CamembertForQuestionAnswering"),
        ("canine", "CanineForQuestionAnswering"),
        ("convbert", "ConvBertForQuestionAnswering"),
        ("data2vec-text", "Data2VecTextForQuestionAnswering"),
        ("deberta", "DebertaForQuestionAnswering"),
        ("deberta-v2", "DebertaV2ForQuestionAnswering"),
        ("distilbert", "DistilBertForQuestionAnswering"),
        ("electra", "ElectraForQuestionAnswering"),
604
        ("ernie", "ErnieForQuestionAnswering"),
605
        ("flaubert", "FlaubertForQuestionAnsweringSimple"),
Gunjan Chhablani's avatar
Gunjan Chhablani committed
606
        ("fnet", "FNetForQuestionAnswering"),
607
        ("funnel", "FunnelForQuestionAnswering"),
608
        ("gptj", "GPTJForQuestionAnswering"),
609
        ("ibert", "IBertForQuestionAnswering"),
610
        ("layoutlmv2", "LayoutLMv2ForQuestionAnswering"),
NielsRogge's avatar
NielsRogge committed
611
        ("layoutlmv3", "LayoutLMv3ForQuestionAnswering"),
612
        ("led", "LEDForQuestionAnswering"),
NielsRogge's avatar
NielsRogge committed
613
        ("lilt", "LiltForQuestionAnswering"),
614
        ("longformer", "LongformerForQuestionAnswering"),
615
        ("luke", "LukeForQuestionAnswering"),
616
        ("lxmert", "LxmertForQuestionAnswering"),
NielsRogge's avatar
NielsRogge committed
617
        ("markuplm", "MarkupLMForQuestionAnswering"),
618
        ("mbart", "MBartForQuestionAnswering"),
619
620
621
        ("megatron-bert", "MegatronBertForQuestionAnswering"),
        ("mobilebert", "MobileBertForQuestionAnswering"),
        ("mpnet", "MPNetForQuestionAnswering"),
StevenTang1998's avatar
StevenTang1998 committed
622
        ("mvp", "MvpForQuestionAnswering"),
623
        ("nezha", "NezhaForQuestionAnswering"),
624
        ("nystromformer", "NystromformerForQuestionAnswering"),
625
        ("opt", "OPTForQuestionAnswering"),
626
627
628
629
630
        ("qdqbert", "QDQBertForQuestionAnswering"),
        ("reformer", "ReformerForQuestionAnswering"),
        ("rembert", "RemBertForQuestionAnswering"),
        ("roberta", "RobertaForQuestionAnswering"),
        ("roformer", "RoFormerForQuestionAnswering"),
Ori Ram's avatar
Ori Ram committed
631
        ("splinter", "SplinterForQuestionAnswering"),
632
633
634
635
636
637
        ("squeezebert", "SqueezeBertForQuestionAnswering"),
        ("xlm", "XLMForQuestionAnsweringSimple"),
        ("xlm-roberta", "XLMRobertaForQuestionAnswering"),
        ("xlm-roberta-xl", "XLMRobertaXLForQuestionAnswering"),
        ("xlnet", "XLNetForQuestionAnsweringSimple"),
        ("yoso", "YosoForQuestionAnswering"),
638
639
640
    ]
)

641
MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING_NAMES = OrderedDict(
642
643
    [
        # Model for Table Question Answering mapping
644
        ("tapas", "TapasForQuestionAnswering"),
645
646
647
    ]
)

648
649
650
651
652
653
MODEL_FOR_VISUAL_QUESTION_ANSWERING_MAPPING_NAMES = OrderedDict(
    [
        ("vilt", "ViltForQuestionAnswering"),
    ]
)

654
655
656
657
658
659
660
661
MODEL_FOR_DOCUMENT_QUESTION_ANSWERING_MAPPING_NAMES = OrderedDict(
    [
        ("layoutlm", "LayoutLMForQuestionAnswering"),
        ("layoutlmv2", "LayoutLMv2ForQuestionAnswering"),
        ("layoutlmv3", "LayoutLMv3ForQuestionAnswering"),
    ]
)

662
MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING_NAMES = OrderedDict(
Julien Chaumond's avatar
Julien Chaumond committed
663
    [
664
        # Model for Token Classification mapping
665
666
        ("albert", "AlbertForTokenClassification"),
        ("bert", "BertForTokenClassification"),
667
        ("big_bird", "BigBirdForTokenClassification"),
668
        ("bloom", "BloomForTokenClassification"),
669
670
        ("camembert", "CamembertForTokenClassification"),
        ("canine", "CanineForTokenClassification"),
671
        ("convbert", "ConvBertForTokenClassification"),
672
673
674
        ("data2vec-text", "Data2VecTextForTokenClassification"),
        ("deberta", "DebertaForTokenClassification"),
        ("deberta-v2", "DebertaV2ForTokenClassification"),
675
        ("distilbert", "DistilBertForTokenClassification"),
676
        ("electra", "ElectraForTokenClassification"),
677
        ("ernie", "ErnieForTokenClassification"),
678
        ("esm", "EsmForTokenClassification"),
679
        ("flaubert", "FlaubertForTokenClassification"),
680
681
682
683
684
685
        ("fnet", "FNetForTokenClassification"),
        ("funnel", "FunnelForTokenClassification"),
        ("gpt2", "GPT2ForTokenClassification"),
        ("ibert", "IBertForTokenClassification"),
        ("layoutlm", "LayoutLMForTokenClassification"),
        ("layoutlmv2", "LayoutLMv2ForTokenClassification"),
NielsRogge's avatar
NielsRogge committed
686
        ("layoutlmv3", "LayoutLMv3ForTokenClassification"),
NielsRogge's avatar
NielsRogge committed
687
        ("lilt", "LiltForTokenClassification"),
688
        ("longformer", "LongformerForTokenClassification"),
689
        ("luke", "LukeForTokenClassification"),
NielsRogge's avatar
NielsRogge committed
690
        ("markuplm", "MarkupLMForTokenClassification"),
691
692
693
        ("megatron-bert", "MegatronBertForTokenClassification"),
        ("mobilebert", "MobileBertForTokenClassification"),
        ("mpnet", "MPNetForTokenClassification"),
694
        ("nezha", "NezhaForTokenClassification"),
695
696
697
698
699
700
701
702
703
704
705
        ("nystromformer", "NystromformerForTokenClassification"),
        ("qdqbert", "QDQBertForTokenClassification"),
        ("rembert", "RemBertForTokenClassification"),
        ("roberta", "RobertaForTokenClassification"),
        ("roformer", "RoFormerForTokenClassification"),
        ("squeezebert", "SqueezeBertForTokenClassification"),
        ("xlm", "XLMForTokenClassification"),
        ("xlm-roberta", "XLMRobertaForTokenClassification"),
        ("xlm-roberta-xl", "XLMRobertaXLForTokenClassification"),
        ("xlnet", "XLNetForTokenClassification"),
        ("yoso", "YosoForTokenClassification"),
Julien Chaumond's avatar
Julien Chaumond committed
706
707
708
    ]
)

709
MODEL_FOR_MULTIPLE_CHOICE_MAPPING_NAMES = OrderedDict(
Julien Chaumond's avatar
Julien Chaumond committed
710
    [
711
        # Model for Multiple Choice mapping
712
713
        ("albert", "AlbertForMultipleChoice"),
        ("bert", "BertForMultipleChoice"),
714
715
        ("big_bird", "BigBirdForMultipleChoice"),
        ("camembert", "CamembertForMultipleChoice"),
716
717
        ("canine", "CanineForMultipleChoice"),
        ("convbert", "ConvBertForMultipleChoice"),
718
        ("data2vec-text", "Data2VecTextForMultipleChoice"),
719
        ("deberta-v2", "DebertaV2ForMultipleChoice"),
720
        ("distilbert", "DistilBertForMultipleChoice"),
721
        ("electra", "ElectraForMultipleChoice"),
722
        ("ernie", "ErnieForMultipleChoice"),
723
        ("flaubert", "FlaubertForMultipleChoice"),
724
        ("fnet", "FNetForMultipleChoice"),
725
726
        ("funnel", "FunnelForMultipleChoice"),
        ("ibert", "IBertForMultipleChoice"),
727
        ("longformer", "LongformerForMultipleChoice"),
728
        ("luke", "LukeForMultipleChoice"),
729
730
731
        ("megatron-bert", "MegatronBertForMultipleChoice"),
        ("mobilebert", "MobileBertForMultipleChoice"),
        ("mpnet", "MPNetForMultipleChoice"),
732
        ("nezha", "NezhaForMultipleChoice"),
733
734
735
736
737
738
739
740
741
742
743
        ("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"),
Julien Chaumond's avatar
Julien Chaumond committed
744
745
746
    ]
)

747
MODEL_FOR_NEXT_SENTENCE_PREDICTION_MAPPING_NAMES = OrderedDict(
748
    [
749
        ("bert", "BertForNextSentencePrediction"),
750
        ("ernie", "ErnieForNextSentencePrediction"),
Gunjan Chhablani's avatar
Gunjan Chhablani committed
751
        ("fnet", "FNetForNextSentencePrediction"),
752
753
        ("megatron-bert", "MegatronBertForNextSentencePrediction"),
        ("mobilebert", "MobileBertForNextSentencePrediction"),
754
        ("nezha", "NezhaForNextSentencePrediction"),
755
        ("qdqbert", "QDQBertForNextSentencePrediction"),
756
757
758
    ]
)

759
760
761
MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING_NAMES = OrderedDict(
    [
        # Model for Audio Classification mapping
762
        ("data2vec-audio", "Data2VecAudioForSequenceClassification"),
763
        ("hubert", "HubertForSequenceClassification"),
764
765
        ("sew", "SEWForSequenceClassification"),
        ("sew-d", "SEWDForSequenceClassification"),
766
767
768
        ("unispeech", "UniSpeechForSequenceClassification"),
        ("unispeech-sat", "UniSpeechSatForSequenceClassification"),
        ("wav2vec2", "Wav2Vec2ForSequenceClassification"),
769
        ("wav2vec2-conformer", "Wav2Vec2ConformerForSequenceClassification"),
Patrick von Platen's avatar
Patrick von Platen committed
770
        ("wavlm", "WavLMForSequenceClassification"),
771
772
773
    ]
)

774
775
776
MODEL_FOR_CTC_MAPPING_NAMES = OrderedDict(
    [
        # Model for Connectionist temporal classification (CTC) mapping
777
        ("data2vec-audio", "Data2VecAudioForCTC"),
778
        ("hubert", "HubertForCTC"),
Chan Woo Kim's avatar
Chan Woo Kim committed
779
        ("mctct", "MCTCTForCTC"),
780
781
        ("sew", "SEWForCTC"),
        ("sew-d", "SEWDForCTC"),
782
783
784
        ("unispeech", "UniSpeechForCTC"),
        ("unispeech-sat", "UniSpeechSatForCTC"),
        ("wav2vec2", "Wav2Vec2ForCTC"),
785
        ("wav2vec2-conformer", "Wav2Vec2ConformerForCTC"),
Patrick von Platen's avatar
Patrick von Platen committed
786
        ("wavlm", "WavLMForCTC"),
787
788
789
    ]
)

790
791
792
MODEL_FOR_AUDIO_FRAME_CLASSIFICATION_MAPPING_NAMES = OrderedDict(
    [
        # Model for Audio Classification mapping
793
        ("data2vec-audio", "Data2VecAudioForAudioFrameClassification"),
794
        ("unispeech-sat", "UniSpeechSatForAudioFrameClassification"),
795
        ("wav2vec2", "Wav2Vec2ForAudioFrameClassification"),
796
        ("wav2vec2-conformer", "Wav2Vec2ConformerForAudioFrameClassification"),
797
        ("wavlm", "WavLMForAudioFrameClassification"),
798
799
800
801
802
803
    ]
)

MODEL_FOR_AUDIO_XVECTOR_MAPPING_NAMES = OrderedDict(
    [
        # Model for Audio Classification mapping
804
        ("data2vec-audio", "Data2VecAudioForXVector"),
805
        ("unispeech-sat", "UniSpeechSatForXVector"),
806
        ("wav2vec2", "Wav2Vec2ForXVector"),
807
        ("wav2vec2-conformer", "Wav2Vec2ConformerForXVector"),
808
        ("wavlm", "WavLMForXVector"),
809
810
811
    ]
)

812
813
814
815
816
817
818
_MODEL_FOR_ZERO_SHOT_IMAGE_CLASSIFICATION_MAPPING_NAMES = OrderedDict(
    [
        # Model for Zero Shot Image Classification mapping
        ("clip", "CLIPModel"),
    ]
)

819
820
821
822
MODEL_MAPPING = _LazyAutoMapping(CONFIG_MAPPING_NAMES, MODEL_MAPPING_NAMES)
MODEL_FOR_PRETRAINING_MAPPING = _LazyAutoMapping(CONFIG_MAPPING_NAMES, MODEL_FOR_PRETRAINING_MAPPING_NAMES)
MODEL_WITH_LM_HEAD_MAPPING = _LazyAutoMapping(CONFIG_MAPPING_NAMES, MODEL_WITH_LM_HEAD_MAPPING_NAMES)
MODEL_FOR_CAUSAL_LM_MAPPING = _LazyAutoMapping(CONFIG_MAPPING_NAMES, MODEL_FOR_CAUSAL_LM_MAPPING_NAMES)
NielsRogge's avatar
NielsRogge committed
823
824
825
MODEL_FOR_CAUSAL_IMAGE_MODELING_MAPPING = _LazyAutoMapping(
    CONFIG_MAPPING_NAMES, MODEL_FOR_CAUSAL_IMAGE_MODELING_MAPPING_NAMES
)
826
827
828
MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING = _LazyAutoMapping(
    CONFIG_MAPPING_NAMES, MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING_NAMES
)
829
830
831
MODEL_FOR_IMAGE_SEGMENTATION_MAPPING = _LazyAutoMapping(
    CONFIG_MAPPING_NAMES, MODEL_FOR_IMAGE_SEGMENTATION_MAPPING_NAMES
)
832
833
834
MODEL_FOR_SEMANTIC_SEGMENTATION_MAPPING = _LazyAutoMapping(
    CONFIG_MAPPING_NAMES, MODEL_FOR_SEMANTIC_SEGMENTATION_MAPPING_NAMES
)
835
836
837
MODEL_FOR_INSTANCE_SEGMENTATION_MAPPING = _LazyAutoMapping(
    CONFIG_MAPPING_NAMES, MODEL_FOR_INSTANCE_SEGMENTATION_MAPPING_NAMES
)
NielsRogge's avatar
NielsRogge committed
838
839
840
MODEL_FOR_VIDEO_CLASSIFICATION_MAPPING = _LazyAutoMapping(
    CONFIG_MAPPING_NAMES, MODEL_FOR_VIDEO_CLASSIFICATION_MAPPING_NAMES
)
841
MODEL_FOR_VISION_2_SEQ_MAPPING = _LazyAutoMapping(CONFIG_MAPPING_NAMES, MODEL_FOR_VISION_2_SEQ_MAPPING_NAMES)
842
843
844
MODEL_FOR_VISUAL_QUESTION_ANSWERING_MAPPING = _LazyAutoMapping(
    CONFIG_MAPPING_NAMES, MODEL_FOR_VISUAL_QUESTION_ANSWERING_MAPPING_NAMES
)
845
846
847
MODEL_FOR_DOCUMENT_QUESTION_ANSWERING_MAPPING = _LazyAutoMapping(
    CONFIG_MAPPING_NAMES, MODEL_FOR_DOCUMENT_QUESTION_ANSWERING_MAPPING_NAMES
)
848
MODEL_FOR_MASKED_LM_MAPPING = _LazyAutoMapping(CONFIG_MAPPING_NAMES, MODEL_FOR_MASKED_LM_MAPPING_NAMES)
NielsRogge's avatar
NielsRogge committed
849
850
851
MODEL_FOR_MASKED_IMAGE_MODELING_MAPPING = _LazyAutoMapping(
    CONFIG_MAPPING_NAMES, MODEL_FOR_MASKED_IMAGE_MODELING_MAPPING_NAMES
)
852
MODEL_FOR_OBJECT_DETECTION_MAPPING = _LazyAutoMapping(CONFIG_MAPPING_NAMES, MODEL_FOR_OBJECT_DETECTION_MAPPING_NAMES)
853
854
855
MODEL_FOR_ZERO_SHOT_OBJECT_DETECTION_MAPPING = _LazyAutoMapping(
    CONFIG_MAPPING_NAMES, MODEL_FOR_ZERO_SHOT_OBJECT_DETECTION_MAPPING_NAMES
)
856
MODEL_FOR_DEPTH_ESTIMATION_MAPPING = _LazyAutoMapping(CONFIG_MAPPING_NAMES, MODEL_FOR_DEPTH_ESTIMATION_MAPPING_NAMES)
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING = _LazyAutoMapping(
    CONFIG_MAPPING_NAMES, MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING_NAMES
)
MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING = _LazyAutoMapping(
    CONFIG_MAPPING_NAMES, MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING_NAMES
)
MODEL_FOR_QUESTION_ANSWERING_MAPPING = _LazyAutoMapping(
    CONFIG_MAPPING_NAMES, MODEL_FOR_QUESTION_ANSWERING_MAPPING_NAMES
)
MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING = _LazyAutoMapping(
    CONFIG_MAPPING_NAMES, MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING_NAMES
)
MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING = _LazyAutoMapping(
    CONFIG_MAPPING_NAMES, MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING_NAMES
)
MODEL_FOR_MULTIPLE_CHOICE_MAPPING = _LazyAutoMapping(CONFIG_MAPPING_NAMES, MODEL_FOR_MULTIPLE_CHOICE_MAPPING_NAMES)
MODEL_FOR_NEXT_SENTENCE_PREDICTION_MAPPING = _LazyAutoMapping(
    CONFIG_MAPPING_NAMES, MODEL_FOR_NEXT_SENTENCE_PREDICTION_MAPPING_NAMES
)
876
877
878
MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING = _LazyAutoMapping(
    CONFIG_MAPPING_NAMES, MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING_NAMES
)
879
880
MODEL_FOR_CTC_MAPPING = _LazyAutoMapping(CONFIG_MAPPING_NAMES, MODEL_FOR_CTC_MAPPING_NAMES)
MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING = _LazyAutoMapping(CONFIG_MAPPING_NAMES, MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING_NAMES)
881
882
883
884
MODEL_FOR_AUDIO_FRAME_CLASSIFICATION_MAPPING = _LazyAutoMapping(
    CONFIG_MAPPING_NAMES, MODEL_FOR_AUDIO_FRAME_CLASSIFICATION_MAPPING_NAMES
)
MODEL_FOR_AUDIO_XVECTOR_MAPPING = _LazyAutoMapping(CONFIG_MAPPING_NAMES, MODEL_FOR_AUDIO_XVECTOR_MAPPING_NAMES)
885

886

Sylvain Gugger's avatar
Sylvain Gugger committed
887
888
889
890
891
892
893
894
895
896
897
898
class AutoModel(_BaseAutoModelClass):
    _model_mapping = MODEL_MAPPING


AutoModel = auto_class_update(AutoModel)


class AutoModelForPreTraining(_BaseAutoModelClass):
    _model_mapping = MODEL_FOR_PRETRAINING_MAPPING


AutoModelForPreTraining = auto_class_update(AutoModelForPreTraining, head_doc="pretraining")
899

thomwolf's avatar
thomwolf committed
900

901
# Private on purpose, the public class will add the deprecation warnings.
Sylvain Gugger's avatar
Sylvain Gugger committed
902
903
class _AutoModelWithLMHead(_BaseAutoModelClass):
    _model_mapping = MODEL_WITH_LM_HEAD_MAPPING
thomwolf's avatar
thomwolf committed
904
905


Sylvain Gugger's avatar
Sylvain Gugger committed
906
_AutoModelWithLMHead = auto_class_update(_AutoModelWithLMHead, head_doc="language modeling")
thomwolf's avatar
thomwolf committed
907
908


Sylvain Gugger's avatar
Sylvain Gugger committed
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
class AutoModelForCausalLM(_BaseAutoModelClass):
    _model_mapping = MODEL_FOR_CAUSAL_LM_MAPPING


AutoModelForCausalLM = auto_class_update(AutoModelForCausalLM, head_doc="causal language modeling")


class AutoModelForMaskedLM(_BaseAutoModelClass):
    _model_mapping = MODEL_FOR_MASKED_LM_MAPPING


AutoModelForMaskedLM = auto_class_update(AutoModelForMaskedLM, head_doc="masked language modeling")


class AutoModelForSeq2SeqLM(_BaseAutoModelClass):
    _model_mapping = MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING


AutoModelForSeq2SeqLM = auto_class_update(
    AutoModelForSeq2SeqLM, head_doc="sequence-to-sequence language modeling", checkpoint_for_example="t5-base"
929
)
thomwolf's avatar
thomwolf committed
930

Sylvain Gugger's avatar
Sylvain Gugger committed
931
932
933
934
935
936
937

class AutoModelForSequenceClassification(_BaseAutoModelClass):
    _model_mapping = MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING


AutoModelForSequenceClassification = auto_class_update(
    AutoModelForSequenceClassification, head_doc="sequence classification"
938
)
thomwolf's avatar
thomwolf committed
939

Sylvain Gugger's avatar
Sylvain Gugger committed
940
941
942
943
944
945
946
947
948
949
950
951
952
953

class AutoModelForQuestionAnswering(_BaseAutoModelClass):
    _model_mapping = MODEL_FOR_QUESTION_ANSWERING_MAPPING


AutoModelForQuestionAnswering = auto_class_update(AutoModelForQuestionAnswering, head_doc="question answering")


class AutoModelForTableQuestionAnswering(_BaseAutoModelClass):
    _model_mapping = MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING


AutoModelForTableQuestionAnswering = auto_class_update(
    AutoModelForTableQuestionAnswering,
954
955
956
    head_doc="table question answering",
    checkpoint_for_example="google/tapas-base-finetuned-wtq",
)
thomwolf's avatar
thomwolf committed
957
958


959
960
961
962
963
964
965
966
967
968
969
class AutoModelForVisualQuestionAnswering(_BaseAutoModelClass):
    _model_mapping = MODEL_FOR_VISUAL_QUESTION_ANSWERING_MAPPING


AutoModelForVisualQuestionAnswering = auto_class_update(
    AutoModelForVisualQuestionAnswering,
    head_doc="visual question answering",
    checkpoint_for_example="dandelin/vilt-b32-finetuned-vqa",
)


970
971
972
973
974
975
976
class AutoModelForDocumentQuestionAnswering(_BaseAutoModelClass):
    _model_mapping = MODEL_FOR_DOCUMENT_QUESTION_ANSWERING_MAPPING


AutoModelForDocumentQuestionAnswering = auto_class_update(
    AutoModelForDocumentQuestionAnswering,
    head_doc="document question answering",
977
    checkpoint_for_example='impira/layoutlm-document-qa", revision="52e01b3',
978
979
980
)


Sylvain Gugger's avatar
Sylvain Gugger committed
981
982
class AutoModelForTokenClassification(_BaseAutoModelClass):
    _model_mapping = MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING
983

984

Sylvain Gugger's avatar
Sylvain Gugger committed
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
AutoModelForTokenClassification = auto_class_update(AutoModelForTokenClassification, head_doc="token classification")


class AutoModelForMultipleChoice(_BaseAutoModelClass):
    _model_mapping = MODEL_FOR_MULTIPLE_CHOICE_MAPPING


AutoModelForMultipleChoice = auto_class_update(AutoModelForMultipleChoice, head_doc="multiple choice")


class AutoModelForNextSentencePrediction(_BaseAutoModelClass):
    _model_mapping = MODEL_FOR_NEXT_SENTENCE_PREDICTION_MAPPING


AutoModelForNextSentencePrediction = auto_class_update(
    AutoModelForNextSentencePrediction, head_doc="next sentence prediction"
1001
)
1002

1003

Sylvain Gugger's avatar
Sylvain Gugger committed
1004
1005
1006
1007
1008
1009
1010
class AutoModelForImageClassification(_BaseAutoModelClass):
    _model_mapping = MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING


AutoModelForImageClassification = auto_class_update(AutoModelForImageClassification, head_doc="image classification")


1011
1012
1013
1014
1015
1016
1017
class AutoModelForImageSegmentation(_BaseAutoModelClass):
    _model_mapping = MODEL_FOR_IMAGE_SEGMENTATION_MAPPING


AutoModelForImageSegmentation = auto_class_update(AutoModelForImageSegmentation, head_doc="image segmentation")


1018
1019
1020
1021
1022
1023
1024
1025
1026
class AutoModelForSemanticSegmentation(_BaseAutoModelClass):
    _model_mapping = MODEL_FOR_SEMANTIC_SEGMENTATION_MAPPING


AutoModelForSemanticSegmentation = auto_class_update(
    AutoModelForSemanticSegmentation, head_doc="semantic segmentation"
)


1027
1028
1029
1030
1031
1032
1033
1034
1035
class AutoModelForInstanceSegmentation(_BaseAutoModelClass):
    _model_mapping = MODEL_FOR_INSTANCE_SEGMENTATION_MAPPING


AutoModelForInstanceSegmentation = auto_class_update(
    AutoModelForInstanceSegmentation, head_doc="instance segmentation"
)


1036
1037
1038
1039
1040
1041
1042
class AutoModelForObjectDetection(_BaseAutoModelClass):
    _model_mapping = MODEL_FOR_OBJECT_DETECTION_MAPPING


AutoModelForObjectDetection = auto_class_update(AutoModelForObjectDetection, head_doc="object detection")


1043
1044
1045
1046
1047
1048
1049
1050
1051
class AutoModelForZeroShotObjectDetection(_BaseAutoModelClass):
    _model_mapping = MODEL_FOR_ZERO_SHOT_OBJECT_DETECTION_MAPPING


AutoModelForZeroShotObjectDetection = auto_class_update(
    AutoModelForZeroShotObjectDetection, head_doc="zero-shot object detection"
)


1052
1053
1054
1055
1056
1057
1058
class AutoModelForDepthEstimation(_BaseAutoModelClass):
    _model_mapping = MODEL_FOR_DEPTH_ESTIMATION_MAPPING


AutoModelForDepthEstimation = auto_class_update(AutoModelForDepthEstimation, head_doc="depth estimation")


NielsRogge's avatar
NielsRogge committed
1059
1060
1061
1062
1063
1064
1065
class AutoModelForVideoClassification(_BaseAutoModelClass):
    _model_mapping = MODEL_FOR_VIDEO_CLASSIFICATION_MAPPING


AutoModelForVideoClassification = auto_class_update(AutoModelForVideoClassification, head_doc="video classification")


1066
1067
1068
1069
1070
1071
1072
class AutoModelForVision2Seq(_BaseAutoModelClass):
    _model_mapping = MODEL_FOR_VISION_2_SEQ_MAPPING


AutoModelForVision2Seq = auto_class_update(AutoModelForVision2Seq, head_doc="vision-to-text modeling")


1073
1074
1075
1076
1077
1078
1079
class AutoModelForAudioClassification(_BaseAutoModelClass):
    _model_mapping = MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING


AutoModelForAudioClassification = auto_class_update(AutoModelForAudioClassification, head_doc="audio classification")


1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
class AutoModelForCTC(_BaseAutoModelClass):
    _model_mapping = MODEL_FOR_CTC_MAPPING


AutoModelForCTC = auto_class_update(AutoModelForCTC, head_doc="connectionist temporal classification")


class AutoModelForSpeechSeq2Seq(_BaseAutoModelClass):
    _model_mapping = MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING


AutoModelForSpeechSeq2Seq = auto_class_update(
Joao Gante's avatar
Joao Gante committed
1092
    AutoModelForSpeechSeq2Seq, head_doc="sequence-to-sequence speech-to-text modeling"
1093
1094
1095
)


1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
class AutoModelForAudioFrameClassification(_BaseAutoModelClass):
    _model_mapping = MODEL_FOR_AUDIO_FRAME_CLASSIFICATION_MAPPING


AutoModelForAudioFrameClassification = auto_class_update(
    AutoModelForAudioFrameClassification, head_doc="audio frame (token) classification"
)


class AutoModelForAudioXVector(_BaseAutoModelClass):
    _model_mapping = MODEL_FOR_AUDIO_XVECTOR_MAPPING


AutoModelForAudioXVector = auto_class_update(AutoModelForAudioXVector, head_doc="audio retrieval via x-vector")


NielsRogge's avatar
NielsRogge committed
1112
1113
1114
1115
1116
1117
1118
class AutoModelForMaskedImageModeling(_BaseAutoModelClass):
    _model_mapping = MODEL_FOR_MASKED_IMAGE_MODELING_MAPPING


AutoModelForMaskedImageModeling = auto_class_update(AutoModelForMaskedImageModeling, head_doc="masked image modeling")


1119
class AutoModelWithLMHead(_AutoModelWithLMHead):
1120
1121
    @classmethod
    def from_config(cls, config):
1122
        warnings.warn(
1123
1124
1125
            "The class `AutoModelWithLMHead` is deprecated and will be removed in a future version. Please use "
            "`AutoModelForCausalLM` for causal language models, `AutoModelForMaskedLM` for masked language models and "
            "`AutoModelForSeq2SeqLM` for encoder-decoder models.",
1126
1127
            FutureWarning,
        )
1128
        return super().from_config(config)
1129
1130
1131

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
1132
        warnings.warn(
1133
1134
1135
            "The class `AutoModelWithLMHead` is deprecated and will be removed in a future version. Please use "
            "`AutoModelForCausalLM` for causal language models, `AutoModelForMaskedLM` for masked language models and "
            "`AutoModelForSeq2SeqLM` for encoder-decoder models.",
1136
1137
            FutureWarning,
        )
1138
        return super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)