modeling_auto.py 44 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
133
        ("t5", "T5Model"),
        ("tapas", "TapasModel"),
134
        ("time_series_transformer", "TimeSeriesTransformerModel"),
Carl's avatar
Carl committed
135
        ("trajectory_transformer", "TrajectoryTransformerModel"),
136
        ("transfo-xl", "TransfoXLModel"),
137
138
139
        ("unispeech", "UniSpeechModel"),
        ("unispeech-sat", "UniSpeechSatModel"),
        ("van", "VanModel"),
NielsRogge's avatar
NielsRogge committed
140
        ("videomae", "VideoMAEModel"),
141
142
143
144
145
        ("vilt", "ViltModel"),
        ("vision-text-dual-encoder", "VisionTextDualEncoderModel"),
        ("visual_bert", "VisualBertModel"),
        ("vit", "ViTModel"),
        ("vit_mae", "ViTMAEModel"),
146
        ("vit_msn", "ViTMSNModel"),
147
        ("wav2vec2", "Wav2Vec2Model"),
148
        ("wav2vec2-conformer", "Wav2Vec2ConformerModel"),
149
        ("wavlm", "WavLMModel"),
150
        ("whisper", "WhisperModel"),
NielsRogge's avatar
NielsRogge committed
151
        ("xclip", "XCLIPModel"),
152
        ("xglm", "XGLMModel"),
153
154
        ("xlm", "XLMModel"),
        ("xlm-prophetnet", "XLMProphetNetModel"),
155
156
157
158
159
        ("xlm-roberta", "XLMRobertaModel"),
        ("xlm-roberta-xl", "XLMRobertaXLModel"),
        ("xlnet", "XLNetModel"),
        ("yolos", "YolosModel"),
        ("yoso", "YosoModel"),
Julien Chaumond's avatar
Julien Chaumond committed
160
161
162
    ]
)

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

217
MODEL_WITH_LM_HEAD_MAPPING_NAMES = OrderedDict(
218
    [
219
        # Model with LM heads mapping
220
221
222
        ("albert", "AlbertForMaskedLM"),
        ("bart", "BartForConditionalGeneration"),
        ("bert", "BertForMaskedLM"),
223
224
225
        ("big_bird", "BigBirdForMaskedLM"),
        ("bigbird_pegasus", "BigBirdPegasusForConditionalGeneration"),
        ("blenderbot-small", "BlenderbotSmallForConditionalGeneration"),
Younes Belkada's avatar
Younes Belkada committed
226
        ("bloom", "BloomForCausalLM"),
227
        ("camembert", "CamembertForMaskedLM"),
rooa's avatar
rooa committed
228
        ("codegen", "CodeGenForCausalLM"),
229
        ("convbert", "ConvBertForMaskedLM"),
230
        ("ctrl", "CTRLLMHeadModel"),
231
232
233
234
        ("data2vec-text", "Data2VecTextForMaskedLM"),
        ("deberta", "DebertaForMaskedLM"),
        ("deberta-v2", "DebertaV2ForMaskedLM"),
        ("distilbert", "DistilBertForMaskedLM"),
235
236
        ("electra", "ElectraForMaskedLM"),
        ("encoder-decoder", "EncoderDecoderModel"),
237
        ("ernie", "ErnieForMaskedLM"),
238
        ("esm", "EsmForMaskedLM"),
239
240
241
        ("flaubert", "FlaubertWithLMHeadModel"),
        ("fnet", "FNetForMaskedLM"),
        ("fsmt", "FSMTForConditionalGeneration"),
242
        ("funnel", "FunnelForMaskedLM"),
243
244
        ("gpt2", "GPT2LMHeadModel"),
        ("gpt_neo", "GPTNeoForCausalLM"),
245
        ("gpt_neox", "GPTNeoXForCausalLM"),
246
        ("gpt_neox_japanese", "GPTNeoXJapaneseForCausalLM"),
247
248
249
250
251
        ("gptj", "GPTJForCausalLM"),
        ("ibert", "IBertForMaskedLM"),
        ("layoutlm", "LayoutLMForMaskedLM"),
        ("led", "LEDForConditionalGeneration"),
        ("longformer", "LongformerForMaskedLM"),
Daniel Stancl's avatar
Daniel Stancl committed
252
        ("longt5", "LongT5ForConditionalGeneration"),
253
        ("luke", "LukeForMaskedLM"),
254
255
        ("m2m_100", "M2M100ForConditionalGeneration"),
        ("marian", "MarianMTModel"),
NielsRogge's avatar
NielsRogge committed
256
        ("markuplm", "MarkupLMForMaskedLM"),
257
258
        ("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
        ("yolos", "YolosForObjectDetection"),
NielsRogge's avatar
NielsRogge committed
473
474
475
    ]
)

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

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

509
510
511
512
MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING_NAMES = OrderedDict(
    [
        ("speech-encoder-decoder", "SpeechEncoderDecoderModel"),
        ("speech_to_text", "Speech2TextForConditionalGeneration"),
513
        ("whisper", "WhisperForConditionalGeneration"),
514
515
516
    ]
)

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

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

633
MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING_NAMES = OrderedDict(
634
635
    [
        # Model for Table Question Answering mapping
636
        ("tapas", "TapasForQuestionAnswering"),
637
638
639
    ]
)

640
641
642
643
644
645
MODEL_FOR_VISUAL_QUESTION_ANSWERING_MAPPING_NAMES = OrderedDict(
    [
        ("vilt", "ViltForQuestionAnswering"),
    ]
)

646
647
648
649
650
651
652
653
MODEL_FOR_DOCUMENT_QUESTION_ANSWERING_MAPPING_NAMES = OrderedDict(
    [
        ("layoutlm", "LayoutLMForQuestionAnswering"),
        ("layoutlmv2", "LayoutLMv2ForQuestionAnswering"),
        ("layoutlmv3", "LayoutLMv3ForQuestionAnswering"),
    ]
)

654
MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING_NAMES = OrderedDict(
Julien Chaumond's avatar
Julien Chaumond committed
655
    [
656
        # Model for Token Classification mapping
657
658
        ("albert", "AlbertForTokenClassification"),
        ("bert", "BertForTokenClassification"),
659
        ("big_bird", "BigBirdForTokenClassification"),
660
        ("bloom", "BloomForTokenClassification"),
661
662
        ("camembert", "CamembertForTokenClassification"),
        ("canine", "CanineForTokenClassification"),
663
        ("convbert", "ConvBertForTokenClassification"),
664
665
666
        ("data2vec-text", "Data2VecTextForTokenClassification"),
        ("deberta", "DebertaForTokenClassification"),
        ("deberta-v2", "DebertaV2ForTokenClassification"),
667
        ("distilbert", "DistilBertForTokenClassification"),
668
        ("electra", "ElectraForTokenClassification"),
669
        ("ernie", "ErnieForTokenClassification"),
670
        ("esm", "EsmForTokenClassification"),
671
        ("flaubert", "FlaubertForTokenClassification"),
672
673
674
675
676
677
        ("fnet", "FNetForTokenClassification"),
        ("funnel", "FunnelForTokenClassification"),
        ("gpt2", "GPT2ForTokenClassification"),
        ("ibert", "IBertForTokenClassification"),
        ("layoutlm", "LayoutLMForTokenClassification"),
        ("layoutlmv2", "LayoutLMv2ForTokenClassification"),
NielsRogge's avatar
NielsRogge committed
678
        ("layoutlmv3", "LayoutLMv3ForTokenClassification"),
NielsRogge's avatar
NielsRogge committed
679
        ("lilt", "LiltForTokenClassification"),
680
        ("longformer", "LongformerForTokenClassification"),
681
        ("luke", "LukeForTokenClassification"),
NielsRogge's avatar
NielsRogge committed
682
        ("markuplm", "MarkupLMForTokenClassification"),
683
684
685
        ("megatron-bert", "MegatronBertForTokenClassification"),
        ("mobilebert", "MobileBertForTokenClassification"),
        ("mpnet", "MPNetForTokenClassification"),
686
        ("nezha", "NezhaForTokenClassification"),
687
688
689
690
691
692
693
694
695
696
697
        ("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
698
699
700
    ]
)

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

739
MODEL_FOR_NEXT_SENTENCE_PREDICTION_MAPPING_NAMES = OrderedDict(
740
    [
741
        ("bert", "BertForNextSentencePrediction"),
742
        ("ernie", "ErnieForNextSentencePrediction"),
Gunjan Chhablani's avatar
Gunjan Chhablani committed
743
        ("fnet", "FNetForNextSentencePrediction"),
744
745
        ("megatron-bert", "MegatronBertForNextSentencePrediction"),
        ("mobilebert", "MobileBertForNextSentencePrediction"),
746
        ("nezha", "NezhaForNextSentencePrediction"),
747
        ("qdqbert", "QDQBertForNextSentencePrediction"),
748
749
750
    ]
)

751
752
753
MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING_NAMES = OrderedDict(
    [
        # Model for Audio Classification mapping
754
        ("data2vec-audio", "Data2VecAudioForSequenceClassification"),
755
        ("hubert", "HubertForSequenceClassification"),
756
757
        ("sew", "SEWForSequenceClassification"),
        ("sew-d", "SEWDForSequenceClassification"),
758
759
760
        ("unispeech", "UniSpeechForSequenceClassification"),
        ("unispeech-sat", "UniSpeechSatForSequenceClassification"),
        ("wav2vec2", "Wav2Vec2ForSequenceClassification"),
761
        ("wav2vec2-conformer", "Wav2Vec2ConformerForSequenceClassification"),
Patrick von Platen's avatar
Patrick von Platen committed
762
        ("wavlm", "WavLMForSequenceClassification"),
763
764
765
    ]
)

766
767
768
MODEL_FOR_CTC_MAPPING_NAMES = OrderedDict(
    [
        # Model for Connectionist temporal classification (CTC) mapping
769
        ("data2vec-audio", "Data2VecAudioForCTC"),
770
        ("hubert", "HubertForCTC"),
Chan Woo Kim's avatar
Chan Woo Kim committed
771
        ("mctct", "MCTCTForCTC"),
772
773
        ("sew", "SEWForCTC"),
        ("sew-d", "SEWDForCTC"),
774
775
776
        ("unispeech", "UniSpeechForCTC"),
        ("unispeech-sat", "UniSpeechSatForCTC"),
        ("wav2vec2", "Wav2Vec2ForCTC"),
777
        ("wav2vec2-conformer", "Wav2Vec2ConformerForCTC"),
Patrick von Platen's avatar
Patrick von Platen committed
778
        ("wavlm", "WavLMForCTC"),
779
780
781
    ]
)

782
783
784
MODEL_FOR_AUDIO_FRAME_CLASSIFICATION_MAPPING_NAMES = OrderedDict(
    [
        # Model for Audio Classification mapping
785
        ("data2vec-audio", "Data2VecAudioForAudioFrameClassification"),
786
        ("unispeech-sat", "UniSpeechSatForAudioFrameClassification"),
787
        ("wav2vec2", "Wav2Vec2ForAudioFrameClassification"),
788
        ("wav2vec2-conformer", "Wav2Vec2ConformerForAudioFrameClassification"),
789
        ("wavlm", "WavLMForAudioFrameClassification"),
790
791
792
793
794
795
    ]
)

MODEL_FOR_AUDIO_XVECTOR_MAPPING_NAMES = OrderedDict(
    [
        # Model for Audio Classification mapping
796
        ("data2vec-audio", "Data2VecAudioForXVector"),
797
        ("unispeech-sat", "UniSpeechSatForXVector"),
798
        ("wav2vec2", "Wav2Vec2ForXVector"),
799
        ("wav2vec2-conformer", "Wav2Vec2ConformerForXVector"),
800
        ("wavlm", "WavLMForXVector"),
801
802
803
    ]
)

804
805
806
807
808
809
810
_MODEL_FOR_ZERO_SHOT_IMAGE_CLASSIFICATION_MAPPING_NAMES = OrderedDict(
    [
        # Model for Zero Shot Image Classification mapping
        ("clip", "CLIPModel"),
    ]
)

811
812
813
814
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
815
816
817
MODEL_FOR_CAUSAL_IMAGE_MODELING_MAPPING = _LazyAutoMapping(
    CONFIG_MAPPING_NAMES, MODEL_FOR_CAUSAL_IMAGE_MODELING_MAPPING_NAMES
)
818
819
820
MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING = _LazyAutoMapping(
    CONFIG_MAPPING_NAMES, MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING_NAMES
)
821
822
823
MODEL_FOR_IMAGE_SEGMENTATION_MAPPING = _LazyAutoMapping(
    CONFIG_MAPPING_NAMES, MODEL_FOR_IMAGE_SEGMENTATION_MAPPING_NAMES
)
824
825
826
MODEL_FOR_SEMANTIC_SEGMENTATION_MAPPING = _LazyAutoMapping(
    CONFIG_MAPPING_NAMES, MODEL_FOR_SEMANTIC_SEGMENTATION_MAPPING_NAMES
)
827
828
829
MODEL_FOR_INSTANCE_SEGMENTATION_MAPPING = _LazyAutoMapping(
    CONFIG_MAPPING_NAMES, MODEL_FOR_INSTANCE_SEGMENTATION_MAPPING_NAMES
)
NielsRogge's avatar
NielsRogge committed
830
831
832
MODEL_FOR_VIDEO_CLASSIFICATION_MAPPING = _LazyAutoMapping(
    CONFIG_MAPPING_NAMES, MODEL_FOR_VIDEO_CLASSIFICATION_MAPPING_NAMES
)
833
MODEL_FOR_VISION_2_SEQ_MAPPING = _LazyAutoMapping(CONFIG_MAPPING_NAMES, MODEL_FOR_VISION_2_SEQ_MAPPING_NAMES)
834
835
836
MODEL_FOR_VISUAL_QUESTION_ANSWERING_MAPPING = _LazyAutoMapping(
    CONFIG_MAPPING_NAMES, MODEL_FOR_VISUAL_QUESTION_ANSWERING_MAPPING_NAMES
)
837
838
839
MODEL_FOR_DOCUMENT_QUESTION_ANSWERING_MAPPING = _LazyAutoMapping(
    CONFIG_MAPPING_NAMES, MODEL_FOR_DOCUMENT_QUESTION_ANSWERING_MAPPING_NAMES
)
840
MODEL_FOR_MASKED_LM_MAPPING = _LazyAutoMapping(CONFIG_MAPPING_NAMES, MODEL_FOR_MASKED_LM_MAPPING_NAMES)
NielsRogge's avatar
NielsRogge committed
841
842
843
MODEL_FOR_MASKED_IMAGE_MODELING_MAPPING = _LazyAutoMapping(
    CONFIG_MAPPING_NAMES, MODEL_FOR_MASKED_IMAGE_MODELING_MAPPING_NAMES
)
844
MODEL_FOR_OBJECT_DETECTION_MAPPING = _LazyAutoMapping(CONFIG_MAPPING_NAMES, MODEL_FOR_OBJECT_DETECTION_MAPPING_NAMES)
845
846
847
MODEL_FOR_ZERO_SHOT_OBJECT_DETECTION_MAPPING = _LazyAutoMapping(
    CONFIG_MAPPING_NAMES, MODEL_FOR_ZERO_SHOT_OBJECT_DETECTION_MAPPING_NAMES
)
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
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
)
867
868
869
MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING = _LazyAutoMapping(
    CONFIG_MAPPING_NAMES, MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING_NAMES
)
870
871
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)
872
873
874
875
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)
876

877

Sylvain Gugger's avatar
Sylvain Gugger committed
878
879
880
881
882
883
884
885
886
887
888
889
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")
890

thomwolf's avatar
thomwolf committed
891

892
# Private on purpose, the public class will add the deprecation warnings.
Sylvain Gugger's avatar
Sylvain Gugger committed
893
894
class _AutoModelWithLMHead(_BaseAutoModelClass):
    _model_mapping = MODEL_WITH_LM_HEAD_MAPPING
thomwolf's avatar
thomwolf committed
895
896


Sylvain Gugger's avatar
Sylvain Gugger committed
897
_AutoModelWithLMHead = auto_class_update(_AutoModelWithLMHead, head_doc="language modeling")
thomwolf's avatar
thomwolf committed
898
899


Sylvain Gugger's avatar
Sylvain Gugger committed
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
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"
920
)
thomwolf's avatar
thomwolf committed
921

Sylvain Gugger's avatar
Sylvain Gugger committed
922
923
924
925
926
927
928

class AutoModelForSequenceClassification(_BaseAutoModelClass):
    _model_mapping = MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING


AutoModelForSequenceClassification = auto_class_update(
    AutoModelForSequenceClassification, head_doc="sequence classification"
929
)
thomwolf's avatar
thomwolf committed
930

Sylvain Gugger's avatar
Sylvain Gugger committed
931
932
933
934
935
936
937
938
939
940
941
942
943
944

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,
945
946
947
    head_doc="table question answering",
    checkpoint_for_example="google/tapas-base-finetuned-wtq",
)
thomwolf's avatar
thomwolf committed
948
949


950
951
952
953
954
955
956
957
958
959
960
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",
)


961
962
963
964
965
966
967
class AutoModelForDocumentQuestionAnswering(_BaseAutoModelClass):
    _model_mapping = MODEL_FOR_DOCUMENT_QUESTION_ANSWERING_MAPPING


AutoModelForDocumentQuestionAnswering = auto_class_update(
    AutoModelForDocumentQuestionAnswering,
    head_doc="document question answering",
968
    checkpoint_for_example='impira/layoutlm-document-qa", revision="52e01b3',
969
970
971
)


Sylvain Gugger's avatar
Sylvain Gugger committed
972
973
class AutoModelForTokenClassification(_BaseAutoModelClass):
    _model_mapping = MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING
974

975

Sylvain Gugger's avatar
Sylvain Gugger committed
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
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"
992
)
993

994

Sylvain Gugger's avatar
Sylvain Gugger committed
995
996
997
998
999
1000
1001
class AutoModelForImageClassification(_BaseAutoModelClass):
    _model_mapping = MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING


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


1002
1003
1004
1005
1006
1007
1008
class AutoModelForImageSegmentation(_BaseAutoModelClass):
    _model_mapping = MODEL_FOR_IMAGE_SEGMENTATION_MAPPING


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


1009
1010
1011
1012
1013
1014
1015
1016
1017
class AutoModelForSemanticSegmentation(_BaseAutoModelClass):
    _model_mapping = MODEL_FOR_SEMANTIC_SEGMENTATION_MAPPING


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


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


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


1027
1028
1029
1030
1031
1032
1033
class AutoModelForObjectDetection(_BaseAutoModelClass):
    _model_mapping = MODEL_FOR_OBJECT_DETECTION_MAPPING


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


1034
1035
1036
1037
1038
1039
1040
1041
1042
class AutoModelForZeroShotObjectDetection(_BaseAutoModelClass):
    _model_mapping = MODEL_FOR_ZERO_SHOT_OBJECT_DETECTION_MAPPING


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


NielsRogge's avatar
NielsRogge committed
1043
1044
1045
1046
1047
1048
1049
class AutoModelForVideoClassification(_BaseAutoModelClass):
    _model_mapping = MODEL_FOR_VIDEO_CLASSIFICATION_MAPPING


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


1050
1051
1052
1053
1054
1055
1056
class AutoModelForVision2Seq(_BaseAutoModelClass):
    _model_mapping = MODEL_FOR_VISION_2_SEQ_MAPPING


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


1057
1058
1059
1060
1061
1062
1063
class AutoModelForAudioClassification(_BaseAutoModelClass):
    _model_mapping = MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING


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


1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
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
1076
    AutoModelForSpeechSeq2Seq, head_doc="sequence-to-sequence speech-to-text modeling"
1077
1078
1079
)


1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
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
1096
1097
1098
1099
1100
1101
1102
class AutoModelForMaskedImageModeling(_BaseAutoModelClass):
    _model_mapping = MODEL_FOR_MASKED_IMAGE_MODELING_MAPPING


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


1103
class AutoModelWithLMHead(_AutoModelWithLMHead):
1104
1105
    @classmethod
    def from_config(cls, config):
1106
        warnings.warn(
1107
1108
1109
            "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.",
1110
1111
            FutureWarning,
        )
1112
        return super().from_config(config)
1113
1114
1115

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
1116
        warnings.warn(
1117
1118
1119
            "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.",
1120
1121
            FutureWarning,
        )
1122
        return super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)