codingutf8 2018 the google flax team s and 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 orglicenseslicense2 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 auto model class from collections import ordereddict from utils import logging from autofactory import baseautomodelclass lazyautomapping autoclassupdate from configurationauto import configmappingnames logger logging getloggername flaxmodelmappingnames ordereddict base model mapping albert flaxalbertmodel bart flaxbartmodel beit flaxbeitmodel bert flaxbertmodel bigbird flaxbigbirdmodel blenderbot flaxblenderbotmodel blenderbotsmall flaxblenderbotsmallmodel bloom flaxbloommodel clip flaxclipmodel distilbert flaxdistilbertmodel electra flaxelectramodel gptsw3 flaxgpt2model gpt2 flaxgpt2model gptneo flaxgptneomodel gptj flaxgptjmodel longt5 flaxlongt5model marian flaxmarianmodel mbart flaxmbartmodel mt5 flaxmt5model opt flaxoptmodel pegasus flaxpegasusmodel regnet flaxregnetmodel resnet flaxresnetmodel roberta flaxrobertamodel robertaprelayernorm flaxrobertaprelayernormmodel roformer flaxroformermodel t5 flaxt5model visiontextdualencoder flaxvisiontextdualencodermodel vit flaxvitmodel wav2vec2 flaxwav2vec2model whisper flaxwhispermodel xglm flaxxglmmodel xlmroberta flaxxlmrobertamodel flaxmodelforpretrainingmappingnames ordereddict model for pretraining mapping albert flaxalbertforpretraining bart flaxbartforconditionalgeneration bert flaxbertforpretraining bigbird flaxbigbirdforpretraining electra flaxelectraforpretraining longt5 flaxlongt5forconditionalgeneration mbart flaxmbartforconditionalgeneration mt5 flaxmt5forconditionalgeneration roberta flaxrobertaformaskedlm robertaprelayernorm flaxrobertaprelayernormformaskedlm roformer flaxroformerformaskedlm t5 flaxt5forconditionalgeneration wav2vec2 flaxwav2vec2forpretraining whisper flaxwhisperforconditionalgeneration xlmroberta flaxxlmrobertaformaskedlm flaxmodelformaskedlmmappingnames ordereddict model for masked lm mapping albert flaxalbertformaskedlm bart flaxbartforconditionalgeneration bert flaxbertformaskedlm bigbird flaxbigbirdformaskedlm distilbert flaxdistilbertformaskedlm electra flaxelectraformaskedlm mbart flaxmbartforconditionalgeneration roberta flaxrobertaformaskedlm robertaprelayernorm flaxrobertaprelayernormformaskedlm roformer flaxroformerformaskedlm xlmroberta flaxxlmrobertaformaskedlm flaxmodelforseqtoseqcausallmmappingnames ordereddict model for seq2seq causal lm mapping bart flaxbartforconditionalgeneration blenderbot flaxblenderbotforconditionalgeneration blenderbotsmall flaxblenderbotsmallforconditionalgeneration encoderdecoder flaxencoderdecodermodel longt5 flaxlongt5forconditionalgeneration marian flaxmarianmtmodel mbart flaxmbartforconditionalgeneration mt5 flaxmt5forconditionalgeneration pegasus flaxpegasusforconditionalgeneration t5 flaxt5forconditionalgeneration flaxmodelforimageclassificationmappingnames ordereddict model for imageclasssification beit flaxbeitforimageclassification regnet flaxregnetforimageclassification resnet flaxresnetforimageclassification vit flaxvitforimageclassification flaxmodelforvision2seqmappingnames ordereddict visionencoderdecoder flaxvisionencoderdecodermodel flaxmodelforcausallmmappingnames ordereddict model for causal lm mapping bart flaxbartforcausallm bert flaxbertforcausallm bigbird flaxbigbirdforcausallm bloom flaxbloomforcausallm electra flaxelectraforcausallm gptsw3 flaxgpt2lmheadmodel gpt2 flaxgpt2lmheadmodel gptneo flaxgptneoforcausallm gptj flaxgptjforcausallm opt flaxoptforcausallm roberta flaxrobertaforcausallm robertaprelayernorm flaxrobertaprelayernormforcausallm xglm flaxxglmforcausallm xlmroberta flaxxlmrobertaforcausallm flaxmodelforsequenceclassificationmappingnames ordereddict model for sequence classification mapping albert flaxalbertforsequenceclassification bart flaxbartforsequenceclassification bert flaxbertforsequenceclassification bigbird flaxbigbirdforsequenceclassification distilbert flaxdistilbertforsequenceclassification electra flaxelectraforsequenceclassification mbart flaxmbartforsequenceclassification roberta flaxrobertaforsequenceclassification robertaprelayernorm flaxrobertaprelayernormforsequenceclassification roformer flaxroformerforsequenceclassification xlmroberta flaxxlmrobertaforsequenceclassification flaxmodelforquestionansweringmappingnames ordereddict model for question answering mapping albert flaxalbertforquestionanswering bart flaxbartforquestionanswering bert flaxbertforquestionanswering bigbird flaxbigbirdforquestionanswering distilbert flaxdistilbertforquestionanswering electra flaxelectraforquestionanswering mbart flaxmbartforquestionanswering roberta flaxrobertaforquestionanswering robertaprelayernorm flaxrobertaprelayernormforquestionanswering roformer flaxroformerforquestionanswering xlmroberta flaxxlmrobertaforquestionanswering flaxmodelfortokenclassificationmappingnames ordereddict model for token classification mapping albert flaxalbertfortokenclassification bert flaxbertfortokenclassification bigbird flaxbigbirdfortokenclassification distilbert flaxdistilbertfortokenclassification electra flaxelectrafortokenclassification roberta flaxrobertafortokenclassification robertaprelayernorm flaxrobertaprelayernormfortokenclassification roformer flaxroformerfortokenclassification xlmroberta flaxxlmrobertafortokenclassification flaxmodelformultiplechoicemappingnames ordereddict model for multiple choice mapping albert flaxalbertformultiplechoice bert flaxbertformultiplechoice bigbird flaxbigbirdformultiplechoice distilbert flaxdistilbertformultiplechoice electra flaxelectraformultiplechoice roberta flaxrobertaformultiplechoice robertaprelayernorm flaxrobertaprelayernormformultiplechoice roformer flaxroformerformultiplechoice xlmroberta flaxxlmrobertaformultiplechoice flaxmodelfornextsentencepredictionmappingnames ordereddict bert flaxbertfornextsentenceprediction flaxmodelforspeechseq2seqmappingnames ordereddict speechencoderdecoder flaxspeechencoderdecodermodel whisper flaxwhisperforconditionalgeneration flaxmodelforaudioclassificationmappingnames ordereddict whisper flaxwhisperforaudioclassification flaxmodelmapping lazyautomappingconfigmappingnames flaxmodelmappingnames flaxmodelforpretrainingmapping lazyautomappingconfigmappingnames flaxmodelforpretrainingmappingnames flaxmodelformaskedlmmapping lazyautomappingconfigmappingnames flaxmodelformaskedlmmappingnames flaxmodelforseqtoseqcausallmmapping lazyautomapping configmappingnames flaxmodelforseqtoseqcausallmmappingnames flaxmodelforimageclassificationmapping lazyautomapping configmappingnames flaxmodelforimageclassificationmappingnames flaxmodelforvision2seqmapping lazyautomappingconfigmappingnames flaxmodelforvision2seqmappingnames flaxmodelforcausallmmapping lazyautomappingconfigmappingnames flaxmodelforcausallmmappingnames flaxmodelforsequenceclassificationmapping lazyautomapping configmappingnames flaxmodelforsequenceclassificationmappingnames flaxmodelforquestionansweringmapping lazyautomapping configmappingnames flaxmodelforquestionansweringmappingnames flaxmodelfortokenclassificationmapping lazyautomapping configmappingnames flaxmodelfortokenclassificationmappingnames flaxmodelformultiplechoicemapping lazyautomapping configmappingnames flaxmodelformultiplechoicemappingnames flaxmodelfornextsentencepredictionmapping lazyautomapping configmappingnames flaxmodelfornextsentencepredictionmappingnames flaxmodelforspeechseq2seqmapping lazyautomapping configmappingnames flaxmodelforspeechseq2seqmappingnames flaxmodelforaudioclassificationmapping lazyautomapping configmappingnames flaxmodelforaudioclassificationmappingnames class flaxautomodelbaseautomodelclass modelmapping flaxmodelmapping flaxautomodel autoclassupdateflaxautomodel class flaxautomodelforpretrainingbaseautomodelclass modelmapping flaxmodelforpretrainingmapping flaxautomodelforpretraining autoclassupdateflaxautomodelforpretraining headdocpretraining class flaxautomodelforcausallmbaseautomodelclass modelmapping flaxmodelforcausallmmapping flaxautomodelforcausallm autoclassupdateflaxautomodelforcausallm headdoccausal language modeling class flaxautomodelformaskedlmbaseautomodelclass modelmapping flaxmodelformaskedlmmapping flaxautomodelformaskedlm autoclassupdateflaxautomodelformaskedlm headdocmasked language modeling class flaxautomodelforseq2seqlmbaseautomodelclass modelmapping flaxmodelforseqtoseqcausallmmapping flaxautomodelforseq2seqlm autoclassupdate flaxautomodelforseq2seqlm headdocsequencetosequence language modeling checkpointforexamplet5base class flaxautomodelforsequenceclassificationbaseautomodelclass modelmapping flaxmodelforsequenceclassificationmapping flaxautomodelforsequenceclassification autoclassupdate flaxautomodelforsequenceclassification headdocsequence classification class flaxautomodelforquestionansweringbaseautomodelclass modelmapping flaxmodelforquestionansweringmapping flaxautomodelforquestionanswering autoclassupdateflaxautomodelforquestionanswering headdocquestion answering class flaxautomodelfortokenclassificationbaseautomodelclass modelmapping flaxmodelfortokenclassificationmapping flaxautomodelfortokenclassification autoclassupdate flaxautomodelfortokenclassification headdoctoken classification class flaxautomodelformultiplechoicebaseautomodelclass modelmapping flaxmodelformultiplechoicemapping flaxautomodelformultiplechoice autoclassupdateflaxautomodelformultiplechoice headdocmultiple choice class flaxautomodelfornextsentencepredictionbaseautomodelclass modelmapping flaxmodelfornextsentencepredictionmapping flaxautomodelfornextsentenceprediction autoclassupdate flaxautomodelfornextsentenceprediction headdocnext sentence prediction class flaxautomodelforimageclassificationbaseautomodelclass modelmapping flaxmodelforimageclassificationmapping flaxautomodelforimageclassification autoclassupdate flaxautomodelforimageclassification headdocimage classification class flaxautomodelforvision2seqbaseautomodelclass modelmapping flaxmodelforvision2seqmapping flaxautomodelforvision2seq autoclassupdateflaxautomodelforvision2seq headdocvisiontotext modeling class flaxautomodelforspeechseq2seqbaseautomodelclass modelmapping flaxmodelforspeechseq2seqmapping flaxautomodelforspeechseq2seq autoclassupdate flaxautomodelforspeechseq2seq headdocsequencetosequence speechtotext modeling coding utf 8 2018 the google flax team s and 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 auto model class base model mapping model for pre training mapping model for masked lm mapping model for seq2seq causal lm mapping model for image classsification model for causal lm mapping model for sequence classification mapping model for question answering mapping model for token classification mapping model for multiple choice mapping
from collections import OrderedDict

from ...utils import logging
from .auto_factory import _BaseAutoModelClass, _LazyAutoMapping, auto_class_update
from .configuration_auto import CONFIG_MAPPING_NAMES


logger = logging.get_logger(__name__)


FLAX_MODEL_MAPPING_NAMES = OrderedDict(
    [
        
        ("albert", "FlaxAlbertModel"),
        ("bart", "FlaxBartModel"),
        ("beit", "FlaxBeitModel"),
        ("bert", "FlaxBertModel"),
        ("big_bird", "FlaxBigBirdModel"),
        ("blenderbot", "FlaxBlenderbotModel"),
        ("blenderbot-small", "FlaxBlenderbotSmallModel"),
        ("bloom", "FlaxBloomModel"),
        ("clip", "FlaxCLIPModel"),
        ("distilbert", "FlaxDistilBertModel"),
        ("electra", "FlaxElectraModel"),
        ("gpt-sw3", "FlaxGPT2Model"),
        ("gpt2", "FlaxGPT2Model"),
        ("gpt_neo", "FlaxGPTNeoModel"),
        ("gptj", "FlaxGPTJModel"),
        ("longt5", "FlaxLongT5Model"),
        ("marian", "FlaxMarianModel"),
        ("mbart", "FlaxMBartModel"),
        ("mt5", "FlaxMT5Model"),
        ("opt", "FlaxOPTModel"),
        ("pegasus", "FlaxPegasusModel"),
        ("regnet", "FlaxRegNetModel"),
        ("resnet", "FlaxResNetModel"),
        ("roberta", "FlaxRobertaModel"),
        ("roberta-prelayernorm", "FlaxRobertaPreLayerNormModel"),
        ("roformer", "FlaxRoFormerModel"),
        ("t5", "FlaxT5Model"),
        ("vision-text-dual-encoder", "FlaxVisionTextDualEncoderModel"),
        ("vit", "FlaxViTModel"),
        ("wav2vec2", "FlaxWav2Vec2Model"),
        ("whisper", "FlaxWhisperModel"),
        ("xglm", "FlaxXGLMModel"),
        ("xlm-roberta", "FlaxXLMRobertaModel"),
    ]
)

FLAX_MODEL_FOR_PRETRAINING_MAPPING_NAMES = OrderedDict(
    [
        
        ("albert", "FlaxAlbertForPreTraining"),
        ("bart", "FlaxBartForConditionalGeneration"),
        ("bert", "FlaxBertForPreTraining"),
        ("big_bird", "FlaxBigBirdForPreTraining"),
        ("electra", "FlaxElectraForPreTraining"),
        ("longt5", "FlaxLongT5ForConditionalGeneration"),
        ("mbart", "FlaxMBartForConditionalGeneration"),
        ("mt5", "FlaxMT5ForConditionalGeneration"),
        ("roberta", "FlaxRobertaForMaskedLM"),
        ("roberta-prelayernorm", "FlaxRobertaPreLayerNormForMaskedLM"),
        ("roformer", "FlaxRoFormerForMaskedLM"),
        ("t5", "FlaxT5ForConditionalGeneration"),
        ("wav2vec2", "FlaxWav2Vec2ForPreTraining"),
        ("whisper", "FlaxWhisperForConditionalGeneration"),
        ("xlm-roberta", "FlaxXLMRobertaForMaskedLM"),
    ]
)

FLAX_MODEL_FOR_MASKED_LM_MAPPING_NAMES = OrderedDict(
    [
        
        ("albert", "FlaxAlbertForMaskedLM"),
        ("bart", "FlaxBartForConditionalGeneration"),
        ("bert", "FlaxBertForMaskedLM"),
        ("big_bird", "FlaxBigBirdForMaskedLM"),
        ("distilbert", "FlaxDistilBertForMaskedLM"),
        ("electra", "FlaxElectraForMaskedLM"),
        ("mbart", "FlaxMBartForConditionalGeneration"),
        ("roberta", "FlaxRobertaForMaskedLM"),
        ("roberta-prelayernorm", "FlaxRobertaPreLayerNormForMaskedLM"),
        ("roformer", "FlaxRoFormerForMaskedLM"),
        ("xlm-roberta", "FlaxXLMRobertaForMaskedLM"),
    ]
)

FLAX_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING_NAMES = OrderedDict(
    [
        
        ("bart", "FlaxBartForConditionalGeneration"),
        ("blenderbot", "FlaxBlenderbotForConditionalGeneration"),
        ("blenderbot-small", "FlaxBlenderbotSmallForConditionalGeneration"),
        ("encoder-decoder", "FlaxEncoderDecoderModel"),
        ("longt5", "FlaxLongT5ForConditionalGeneration"),
        ("marian", "FlaxMarianMTModel"),
        ("mbart", "FlaxMBartForConditionalGeneration"),
        ("mt5", "FlaxMT5ForConditionalGeneration"),
        ("pegasus", "FlaxPegasusForConditionalGeneration"),
        ("t5", "FlaxT5ForConditionalGeneration"),
    ]
)

FLAX_MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING_NAMES = OrderedDict(
    [
        
        ("beit", "FlaxBeitForImageClassification"),
        ("regnet", "FlaxRegNetForImageClassification"),
        ("resnet", "FlaxResNetForImageClassification"),
        ("vit", "FlaxViTForImageClassification"),
    ]
)

FLAX_MODEL_FOR_VISION_2_SEQ_MAPPING_NAMES = OrderedDict(
    [
        ("vision-encoder-decoder", "FlaxVisionEncoderDecoderModel"),
    ]
)

FLAX_MODEL_FOR_CAUSAL_LM_MAPPING_NAMES = OrderedDict(
    [
        
        ("bart", "FlaxBartForCausalLM"),
        ("bert", "FlaxBertForCausalLM"),
        ("big_bird", "FlaxBigBirdForCausalLM"),
        ("bloom", "FlaxBloomForCausalLM"),
        ("electra", "FlaxElectraForCausalLM"),
        ("gpt-sw3", "FlaxGPT2LMHeadModel"),
        ("gpt2", "FlaxGPT2LMHeadModel"),
        ("gpt_neo", "FlaxGPTNeoForCausalLM"),
        ("gptj", "FlaxGPTJForCausalLM"),
        ("opt", "FlaxOPTForCausalLM"),
        ("roberta", "FlaxRobertaForCausalLM"),
        ("roberta-prelayernorm", "FlaxRobertaPreLayerNormForCausalLM"),
        ("xglm", "FlaxXGLMForCausalLM"),
        ("xlm-roberta", "FlaxXLMRobertaForCausalLM"),
    ]
)

FLAX_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING_NAMES = OrderedDict(
    [
        
        ("albert", "FlaxAlbertForSequenceClassification"),
        ("bart", "FlaxBartForSequenceClassification"),
        ("bert", "FlaxBertForSequenceClassification"),
        ("big_bird", "FlaxBigBirdForSequenceClassification"),
        ("distilbert", "FlaxDistilBertForSequenceClassification"),
        ("electra", "FlaxElectraForSequenceClassification"),
        ("mbart", "FlaxMBartForSequenceClassification"),
        ("roberta", "FlaxRobertaForSequenceClassification"),
        ("roberta-prelayernorm", "FlaxRobertaPreLayerNormForSequenceClassification"),
        ("roformer", "FlaxRoFormerForSequenceClassification"),
        ("xlm-roberta", "FlaxXLMRobertaForSequenceClassification"),
    ]
)

FLAX_MODEL_FOR_QUESTION_ANSWERING_MAPPING_NAMES = OrderedDict(
    [
        
        ("albert", "FlaxAlbertForQuestionAnswering"),
        ("bart", "FlaxBartForQuestionAnswering"),
        ("bert", "FlaxBertForQuestionAnswering"),
        ("big_bird", "FlaxBigBirdForQuestionAnswering"),
        ("distilbert", "FlaxDistilBertForQuestionAnswering"),
        ("electra", "FlaxElectraForQuestionAnswering"),
        ("mbart", "FlaxMBartForQuestionAnswering"),
        ("roberta", "FlaxRobertaForQuestionAnswering"),
        ("roberta-prelayernorm", "FlaxRobertaPreLayerNormForQuestionAnswering"),
        ("roformer", "FlaxRoFormerForQuestionAnswering"),
        ("xlm-roberta", "FlaxXLMRobertaForQuestionAnswering"),
    ]
)

FLAX_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING_NAMES = OrderedDict(
    [
        
        ("albert", "FlaxAlbertForTokenClassification"),
        ("bert", "FlaxBertForTokenClassification"),
        ("big_bird", "FlaxBigBirdForTokenClassification"),
        ("distilbert", "FlaxDistilBertForTokenClassification"),
        ("electra", "FlaxElectraForTokenClassification"),
        ("roberta", "FlaxRobertaForTokenClassification"),
        ("roberta-prelayernorm", "FlaxRobertaPreLayerNormForTokenClassification"),
        ("roformer", "FlaxRoFormerForTokenClassification"),
        ("xlm-roberta", "FlaxXLMRobertaForTokenClassification"),
    ]
)

FLAX_MODEL_FOR_MULTIPLE_CHOICE_MAPPING_NAMES = OrderedDict(
    [
        
        ("albert", "FlaxAlbertForMultipleChoice"),
        ("bert", "FlaxBertForMultipleChoice"),
        ("big_bird", "FlaxBigBirdForMultipleChoice"),
        ("distilbert", "FlaxDistilBertForMultipleChoice"),
        ("electra", "FlaxElectraForMultipleChoice"),
        ("roberta", "FlaxRobertaForMultipleChoice"),
        ("roberta-prelayernorm", "FlaxRobertaPreLayerNormForMultipleChoice"),
        ("roformer", "FlaxRoFormerForMultipleChoice"),
        ("xlm-roberta", "FlaxXLMRobertaForMultipleChoice"),
    ]
)

FLAX_MODEL_FOR_NEXT_SENTENCE_PREDICTION_MAPPING_NAMES = OrderedDict(
    [
        ("bert", "FlaxBertForNextSentencePrediction"),
    ]
)

FLAX_MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING_NAMES = OrderedDict(
    [
        ("speech-encoder-decoder", "FlaxSpeechEncoderDecoderModel"),
        ("whisper", "FlaxWhisperForConditionalGeneration"),
    ]
)

FLAX_MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING_NAMES = OrderedDict(
    [
        ("whisper", "FlaxWhisperForAudioClassification"),
    ]
)

FLAX_MODEL_MAPPING = _LazyAutoMapping(CONFIG_MAPPING_NAMES, FLAX_MODEL_MAPPING_NAMES)
FLAX_MODEL_FOR_PRETRAINING_MAPPING = _LazyAutoMapping(CONFIG_MAPPING_NAMES, FLAX_MODEL_FOR_PRETRAINING_MAPPING_NAMES)
FLAX_MODEL_FOR_MASKED_LM_MAPPING = _LazyAutoMapping(CONFIG_MAPPING_NAMES, FLAX_MODEL_FOR_MASKED_LM_MAPPING_NAMES)
FLAX_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING = _LazyAutoMapping(
    CONFIG_MAPPING_NAMES, FLAX_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING_NAMES
)
FLAX_MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING = _LazyAutoMapping(
    CONFIG_MAPPING_NAMES, FLAX_MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING_NAMES
)
FLAX_MODEL_FOR_VISION_2_SEQ_MAPPING = _LazyAutoMapping(CONFIG_MAPPING_NAMES, FLAX_MODEL_FOR_VISION_2_SEQ_MAPPING_NAMES)
FLAX_MODEL_FOR_CAUSAL_LM_MAPPING = _LazyAutoMapping(CONFIG_MAPPING_NAMES, FLAX_MODEL_FOR_CAUSAL_LM_MAPPING_NAMES)
FLAX_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING = _LazyAutoMapping(
    CONFIG_MAPPING_NAMES, FLAX_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING_NAMES
)
FLAX_MODEL_FOR_QUESTION_ANSWERING_MAPPING = _LazyAutoMapping(
    CONFIG_MAPPING_NAMES, FLAX_MODEL_FOR_QUESTION_ANSWERING_MAPPING_NAMES
)
FLAX_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING = _LazyAutoMapping(
    CONFIG_MAPPING_NAMES, FLAX_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING_NAMES
)
FLAX_MODEL_FOR_MULTIPLE_CHOICE_MAPPING = _LazyAutoMapping(
    CONFIG_MAPPING_NAMES, FLAX_MODEL_FOR_MULTIPLE_CHOICE_MAPPING_NAMES
)
FLAX_MODEL_FOR_NEXT_SENTENCE_PREDICTION_MAPPING = _LazyAutoMapping(
    CONFIG_MAPPING_NAMES, FLAX_MODEL_FOR_NEXT_SENTENCE_PREDICTION_MAPPING_NAMES
)
FLAX_MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING = _LazyAutoMapping(
    CONFIG_MAPPING_NAMES, FLAX_MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING_NAMES
)
FLAX_MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING = _LazyAutoMapping(
    CONFIG_MAPPING_NAMES, FLAX_MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING_NAMES
)


class FlaxAutoModel(_BaseAutoModelClass):
    _model_mapping = FLAX_MODEL_MAPPING


FlaxAutoModel = auto_class_update(FlaxAutoModel)


class FlaxAutoModelForPreTraining(_BaseAutoModelClass):
    _model_mapping = FLAX_MODEL_FOR_PRETRAINING_MAPPING


FlaxAutoModelForPreTraining = auto_class_update(FlaxAutoModelForPreTraining, head_doc="pretraining")


class FlaxAutoModelForCausalLM(_BaseAutoModelClass):
    _model_mapping = FLAX_MODEL_FOR_CAUSAL_LM_MAPPING


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


class FlaxAutoModelForMaskedLM(_BaseAutoModelClass):
    _model_mapping = FLAX_MODEL_FOR_MASKED_LM_MAPPING


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


class FlaxAutoModelForSeq2SeqLM(_BaseAutoModelClass):
    _model_mapping = FLAX_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING


FlaxAutoModelForSeq2SeqLM = auto_class_update(
    FlaxAutoModelForSeq2SeqLM, head_doc="sequence-to-sequence language modeling", checkpoint_for_example="t5-base"
)


class FlaxAutoModelForSequenceClassification(_BaseAutoModelClass):
    _model_mapping = FLAX_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING


FlaxAutoModelForSequenceClassification = auto_class_update(
    FlaxAutoModelForSequenceClassification, head_doc="sequence classification"
)


class FlaxAutoModelForQuestionAnswering(_BaseAutoModelClass):
    _model_mapping = FLAX_MODEL_FOR_QUESTION_ANSWERING_MAPPING


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


class FlaxAutoModelForTokenClassification(_BaseAutoModelClass):
    _model_mapping = FLAX_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING


FlaxAutoModelForTokenClassification = auto_class_update(
    FlaxAutoModelForTokenClassification, head_doc="token classification"
)


class FlaxAutoModelForMultipleChoice(_BaseAutoModelClass):
    _model_mapping = FLAX_MODEL_FOR_MULTIPLE_CHOICE_MAPPING


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


class FlaxAutoModelForNextSentencePrediction(_BaseAutoModelClass):
    _model_mapping = FLAX_MODEL_FOR_NEXT_SENTENCE_PREDICTION_MAPPING


FlaxAutoModelForNextSentencePrediction = auto_class_update(
    FlaxAutoModelForNextSentencePrediction, head_doc="next sentence prediction"
)


class FlaxAutoModelForImageClassification(_BaseAutoModelClass):
    _model_mapping = FLAX_MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING


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


class FlaxAutoModelForVision2Seq(_BaseAutoModelClass):
    _model_mapping = FLAX_MODEL_FOR_VISION_2_SEQ_MAPPING


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


class FlaxAutoModelForSpeechSeq2Seq(_BaseAutoModelClass):
    _model_mapping = FLAX_MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING


FlaxAutoModelForSpeechSeq2Seq = auto_class_update(
    FlaxAutoModelForSpeechSeq2Seq, head_doc="sequence-to-sequence speech-to-text modeling"
)