




'''simple docstring'''



import argparse
import json
from pathlib import Path

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import ViTConfig, ViTForImageClassification, ViTImageProcessor, ViTModel
from transformers.utils import logging


logging.set_verbosity_info()
UpperCAmelCase_							:	Tuple	     =		logging.get_logger(__name__)







def       _UpperCamelCase			(_lowerCamelCase				:  Optional[Any] ,						_lowerCamelCase				:  Union[str, Any]=False						)-> Optional[Any]:




     '''simple docstring'''





     __snake_case            =	[]
     for i in range(config.num_hidden_layers						):
          # encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
          rename_keys.append((f'''blocks.{i}.norm1.weight''', f'''vit.encoder.layer.{i}.layernorm_before.weight''')						)
          rename_keys.append((f'''blocks.{i}.norm1.bias''', f'''vit.encoder.layer.{i}.layernorm_before.bias''')						)
          rename_keys.append((f'''blocks.{i}.attn.proj.weight''', f'''vit.encoder.layer.{i}.attention.output.dense.weight''')						)
          rename_keys.append((f'''blocks.{i}.attn.proj.bias''', f'''vit.encoder.layer.{i}.attention.output.dense.bias''')						)
          rename_keys.append((f'''blocks.{i}.norm2.weight''', f'''vit.encoder.layer.{i}.layernorm_after.weight''')						)
          rename_keys.append((f'''blocks.{i}.norm2.bias''', f'''vit.encoder.layer.{i}.layernorm_after.bias''')						)
          rename_keys.append((f'''blocks.{i}.mlp.fc1.weight''', f'''vit.encoder.layer.{i}.intermediate.dense.weight''')						)
          rename_keys.append((f'''blocks.{i}.mlp.fc1.bias''', f'''vit.encoder.layer.{i}.intermediate.dense.bias''')						)
          rename_keys.append((f'''blocks.{i}.mlp.fc2.weight''', f'''vit.encoder.layer.{i}.output.dense.weight''')						)
          rename_keys.append((f'''blocks.{i}.mlp.fc2.bias''', f'''vit.encoder.layer.{i}.output.dense.bias''')						)

     # projection layer + position embeddings
     rename_keys.extend(
         [
             ('''cls_token''', '''vit.embeddings.cls_token'''),
             ('''patch_embed.proj.weight''', '''vit.embeddings.patch_embeddings.projection.weight'''),
             ('''patch_embed.proj.bias''', '''vit.embeddings.patch_embeddings.projection.bias'''),
             ('''pos_embed''', '''vit.embeddings.position_embeddings'''),
         ]						)

     if base_model:
          # layernorm + pooler
          rename_keys.extend(
              [
                  ('''norm.weight''', '''layernorm.weight'''),
                  ('''norm.bias''', '''layernorm.bias'''),
              ]						)

          # if just the base model, we should remove "vit" from all keys that start with "vit"
          __snake_case            =	[(pair[0], pair[1][4:]) if pair[1].startswith('''vit'''						) else pair for pair in rename_keys]
     else:
          # layernorm + classification head
          rename_keys.extend(
              [
                  ('''norm.weight''', '''vit.layernorm.weight'''),
                  ('''norm.bias''', '''vit.layernorm.bias'''),
                  ('''head.weight''', '''classifier.weight'''),
                  ('''head.bias''', '''classifier.bias'''),
              ]						)

     return rename_keys







def       _UpperCamelCase			(_lowerCamelCase				:  Optional[Any] ,						_lowerCamelCase				:  Dict ,						_lowerCamelCase				:  Tuple=False						)-> Optional[int]:




     '''simple docstring'''





     for i in range(config.num_hidden_layers						):
          if base_model:
               __snake_case            =	""""""
          else:
               __snake_case            =	"""vit."""
          # read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
          __snake_case            =	state_dict.pop(f'''blocks.{i}.attn.qkv.weight'''						)
          __snake_case            =	state_dict.pop(f'''blocks.{i}.attn.qkv.bias'''						)
          # next, add query, keys and values (in that order) to the state dict
          __snake_case            =	in_proj_weight[
              : config.hidden_size, :
          ]
          __snake_case            =	in_proj_bias[: config.hidden_size]
          __snake_case            =	in_proj_weight[
              config.hidden_size : config.hidden_size * 2, :
          ]
          __snake_case            =	in_proj_bias[
              config.hidden_size : config.hidden_size * 2
          ]
          __snake_case            =	in_proj_weight[
              -config.hidden_size :, :
          ]
          __snake_case            =	in_proj_bias[-config.hidden_size :]







def       _UpperCamelCase			(_lowerCamelCase				:  Tuple						)-> List[Any]:




     '''simple docstring'''





     __snake_case            =	["""head.weight""", """head.bias"""]
     for k in ignore_keys:
          state_dict.pop(__snake_case ,						__snake_case						)







def       _UpperCamelCase			(_lowerCamelCase				:  Optional[int] ,						_lowerCamelCase				:  Dict ,						_lowerCamelCase				:  Union[str, Any]						)-> Union[str, Any]:




     '''simple docstring'''





     __snake_case            =	dct.pop(__snake_case						)
     __snake_case            =	val







def       _UpperCamelCase			()-> List[str]:




     '''simple docstring'''





     __snake_case            =	"""http://images.cocodataset.org/val2017/000000039769.jpg"""
     __snake_case            =	Image.open(requests.get(__snake_case ,						stream=__snake_case						).raw						)
     return im







@torch.no_grad()
def       _UpperCamelCase			(_lowerCamelCase				:  int ,						_lowerCamelCase				:  str ,						_lowerCamelCase				:  Any=True						)-> Tuple:




     '''simple docstring'''





     __snake_case            =	ViTConfig()
     # patch_size
     if model_name[-1] == "8":
          __snake_case            =	8
     # set labels if required
     if not base_model:
          __snake_case            =	10_00
          __snake_case            =	"""huggingface/label-files"""
          __snake_case            =	"""imagenet-1k-id2label.json"""
          __snake_case            =	json.load(open(hf_hub_download(__snake_case ,						__snake_case ,						repo_type='''dataset'''						) ,						'''r'''						)						)
          __snake_case            =	{int(__snake_case						): v for k, v in idalabel.items()}
          __snake_case            =	idalabel
          __snake_case            =	{v: k for k, v in idalabel.items()}
     # size of the architecture
     if model_name in ["dino_vits8", "dino_vits16"]:
          __snake_case            =	3_84
          __snake_case            =	15_36
          __snake_case            =	12
          __snake_case            =	6

     # load original model from torch hub
     __snake_case            =	torch.hub.load('''facebookresearch/dino:main''' ,						__snake_case						)
     original_model.eval()

     # load state_dict of original model, remove and rename some keys
     __snake_case            =	original_model.state_dict()
     if base_model:
          remove_classification_head_(__snake_case						)
     __snake_case            =	create_rename_keys(__snake_case ,						base_model=__snake_case						)
     for src, dest in rename_keys:
          rename_key(__snake_case ,						__snake_case ,						__snake_case						)
     read_in_q_k_v(__snake_case ,						__snake_case ,						__snake_case						)

     # load HuggingFace model
     if base_model:
          __snake_case            =	ViTModel(__snake_case ,						add_pooling_layer=__snake_case						).eval()
     else:
          __snake_case            =	ViTForImageClassification(__snake_case						).eval()
     model.load_state_dict(__snake_case						)

     # Check outputs on an image, prepared by ViTImageProcessor
     __snake_case            =	ViTImageProcessor()
     __snake_case            =	image_processor(images=prepare_img() ,						return_tensors='''pt'''						)
     __snake_case            =	encoding["""pixel_values"""]
     __snake_case            =	model(__snake_case						)

     if base_model:
          __snake_case            =	original_model(__snake_case						)
          assert torch.allclose(__snake_case ,						outputs.last_hidden_state[:, 0, :] ,						atol=1E-1						)
     else:
          __snake_case            =	original_model(__snake_case						)
          assert logits.shape == outputs.logits.shape
          assert torch.allclose(__snake_case ,						outputs.logits ,						atol=1E-3						)

     Path(__snake_case						).mkdir(exist_ok=__snake_case						)
     print(f'''Saving model {model_name} to {pytorch_dump_folder_path}'''						)
     model.save_pretrained(__snake_case						)
     print(f'''Saving image processor to {pytorch_dump_folder_path}'''						)
     image_processor.save_pretrained(__snake_case						)


if __name__ == "__main__":
     UpperCAmelCase_							:	Dict	     =		argparse.ArgumentParser()
     # Required parameters
     parser.add_argument(
         '''--model_name''',
         default='''dino_vitb16''',
         type=str,
         help='''Name of the model trained with DINO you\'d like to convert.''',
     )
     parser.add_argument(
         '''--pytorch_dump_folder_path''', default=None, type=str, help='''Path to the output PyTorch model directory.'''
     )
     parser.add_argument(
         '''--base_model''',
         action='''store_true''',
         help='''Whether to only convert the base model (no projection head weights).''',
     )

     parser.set_defaults(base_model=True)
     UpperCAmelCase_							:	Optional[int]	     =		parser.parse_args()
     convert_vit_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.base_model)










"""simple docstring"""
from collections import OrderedDict

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


UpperCAmelCase				=						logging.get_logger(__name__)


UpperCAmelCase				=						OrderedDict(
    [
        # Base model mapping
        ("""albert""", """FlaxAlbertModel"""),
        ("""bart""", """FlaxBartModel"""),
        ("""beit""", """FlaxBeitModel"""),
        ("""bert""", """FlaxBertModel"""),
        ("""big_bird""", """FlaxBigBirdModel"""),
        ("""blenderbot""", """FlaxBlenderbotModel"""),
        ("""blenderbot-small""", """FlaxBlenderbotSmallModel"""),
        ("""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"""),
    ]
)

UpperCAmelCase				=						OrderedDict(
    [
        # Model for pre-training mapping
        ("""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"""),
    ]
)

UpperCAmelCase				=						OrderedDict(
    [
        # Model for Masked LM mapping
        ("""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"""),
    ]
)

UpperCAmelCase				=						OrderedDict(
    [
        # Model for Seq2Seq Causal LM mapping
        ("""bart""", """FlaxBartForConditionalGeneration"""),
        ("""blenderbot""", """FlaxBlenderbotForConditionalGeneration"""),
        ("""blenderbot-small""", """FlaxBlenderbotSmallForConditionalGeneration"""),
        ("""encoder-decoder""", """FlaxEncoderDecoderModel"""),
        ("""longt5""", """FlaxLongT5ForConditionalGeneration"""),
        ("""marian""", """FlaxMarianMTModel"""),
        ("""mbart""", """FlaxMBartForConditionalGeneration"""),
        ("""mt5""", """FlaxMT5ForConditionalGeneration"""),
        ("""pegasus""", """FlaxPegasusForConditionalGeneration"""),
        ("""t5""", """FlaxT5ForConditionalGeneration"""),
    ]
)

UpperCAmelCase				=						OrderedDict(
    [
        # Model for Image-classsification
        ("""beit""", """FlaxBeitForImageClassification"""),
        ("""regnet""", """FlaxRegNetForImageClassification"""),
        ("""resnet""", """FlaxResNetForImageClassification"""),
        ("""vit""", """FlaxViTForImageClassification"""),
    ]
)

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

UpperCAmelCase				=						OrderedDict(
    [
        # Model for Causal LM mapping
        ("""bart""", """FlaxBartForCausalLM"""),
        ("""bert""", """FlaxBertForCausalLM"""),
        ("""big_bird""", """FlaxBigBirdForCausalLM"""),
        ("""electra""", """FlaxElectraForCausalLM"""),
        ("""gpt-sw3""", """FlaxGPT2LMHeadModel"""),
        ("""gpt2""", """FlaxGPT2LMHeadModel"""),
        ("""gpt_neo""", """FlaxGPTNeoForCausalLM"""),
        ("""gptj""", """FlaxGPTJForCausalLM"""),
        ("""opt""", """FlaxOPTForCausalLM"""),
        ("""roberta""", """FlaxRobertaForCausalLM"""),
        ("""roberta-prelayernorm""", """FlaxRobertaPreLayerNormForCausalLM"""),
        ("""xglm""", """FlaxXGLMForCausalLM"""),
        ("""xlm-roberta""", """FlaxXLMRobertaForCausalLM"""),
    ]
)

UpperCAmelCase				=						OrderedDict(
    [
        # Model for Sequence Classification mapping
        ("""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"""),
    ]
)

UpperCAmelCase				=						OrderedDict(
    [
        # Model for Question Answering mapping
        ("""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"""),
    ]
)

UpperCAmelCase				=						OrderedDict(
    [
        # Model for Token Classification mapping
        ("""albert""", """FlaxAlbertForTokenClassification"""),
        ("""bert""", """FlaxBertForTokenClassification"""),
        ("""big_bird""", """FlaxBigBirdForTokenClassification"""),
        ("""distilbert""", """FlaxDistilBertForTokenClassification"""),
        ("""electra""", """FlaxElectraForTokenClassification"""),
        ("""roberta""", """FlaxRobertaForTokenClassification"""),
        ("""roberta-prelayernorm""", """FlaxRobertaPreLayerNormForTokenClassification"""),
        ("""roformer""", """FlaxRoFormerForTokenClassification"""),
        ("""xlm-roberta""", """FlaxXLMRobertaForTokenClassification"""),
    ]
)

UpperCAmelCase				=						OrderedDict(
    [
        # Model for Multiple Choice mapping
        ("""albert""", """FlaxAlbertForMultipleChoice"""),
        ("""bert""", """FlaxBertForMultipleChoice"""),
        ("""big_bird""", """FlaxBigBirdForMultipleChoice"""),
        ("""distilbert""", """FlaxDistilBertForMultipleChoice"""),
        ("""electra""", """FlaxElectraForMultipleChoice"""),
        ("""roberta""", """FlaxRobertaForMultipleChoice"""),
        ("""roberta-prelayernorm""", """FlaxRobertaPreLayerNormForMultipleChoice"""),
        ("""roformer""", """FlaxRoFormerForMultipleChoice"""),
        ("""xlm-roberta""", """FlaxXLMRobertaForMultipleChoice"""),
    ]
)

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

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

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

UpperCAmelCase				=						_LazyAutoMapping(CONFIG_MAPPING_NAMES, FLAX_MODEL_MAPPING_NAMES)
UpperCAmelCase				=						_LazyAutoMapping(CONFIG_MAPPING_NAMES, FLAX_MODEL_FOR_PRETRAINING_MAPPING_NAMES)
UpperCAmelCase				=						_LazyAutoMapping(CONFIG_MAPPING_NAMES, FLAX_MODEL_FOR_MASKED_LM_MAPPING_NAMES)
UpperCAmelCase				=						_LazyAutoMapping(
    CONFIG_MAPPING_NAMES, FLAX_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING_NAMES
)
UpperCAmelCase				=						_LazyAutoMapping(
    CONFIG_MAPPING_NAMES, FLAX_MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING_NAMES
)
UpperCAmelCase				=						_LazyAutoMapping(CONFIG_MAPPING_NAMES, FLAX_MODEL_FOR_VISION_2_SEQ_MAPPING_NAMES)
UpperCAmelCase				=						_LazyAutoMapping(CONFIG_MAPPING_NAMES, FLAX_MODEL_FOR_CAUSAL_LM_MAPPING_NAMES)
UpperCAmelCase				=						_LazyAutoMapping(
    CONFIG_MAPPING_NAMES, FLAX_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING_NAMES
)
UpperCAmelCase				=						_LazyAutoMapping(
    CONFIG_MAPPING_NAMES, FLAX_MODEL_FOR_QUESTION_ANSWERING_MAPPING_NAMES
)
UpperCAmelCase				=						_LazyAutoMapping(
    CONFIG_MAPPING_NAMES, FLAX_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING_NAMES
)
UpperCAmelCase				=						_LazyAutoMapping(
    CONFIG_MAPPING_NAMES, FLAX_MODEL_FOR_MULTIPLE_CHOICE_MAPPING_NAMES
)
UpperCAmelCase				=						_LazyAutoMapping(
    CONFIG_MAPPING_NAMES, FLAX_MODEL_FOR_NEXT_SENTENCE_PREDICTION_MAPPING_NAMES
)
UpperCAmelCase				=						_LazyAutoMapping(
    CONFIG_MAPPING_NAMES, FLAX_MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING_NAMES
)
UpperCAmelCase				=						_LazyAutoMapping(
    CONFIG_MAPPING_NAMES, FLAX_MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING_NAMES
)




class 	lowercase__				(					_BaseAutoModelClass				):
   __UpperCAmelCase							       =					FLAX_MODEL_MAPPING







UpperCAmelCase				=						auto_class_update(FlaxAutoModel)




class 	lowercase__				(					_BaseAutoModelClass				):
   __UpperCAmelCase							       =					FLAX_MODEL_FOR_PRETRAINING_MAPPING







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




class 	lowercase__				(					_BaseAutoModelClass				):
   __UpperCAmelCase							       =					FLAX_MODEL_FOR_CAUSAL_LM_MAPPING







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




class 	lowercase__				(					_BaseAutoModelClass				):
   __UpperCAmelCase							       =					FLAX_MODEL_FOR_MASKED_LM_MAPPING







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




class 	lowercase__				(					_BaseAutoModelClass				):
   __UpperCAmelCase							       =					FLAX_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING







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




class 	lowercase__				(					_BaseAutoModelClass				):
   __UpperCAmelCase							       =					FLAX_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING







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




class 	lowercase__				(					_BaseAutoModelClass				):
   __UpperCAmelCase							       =					FLAX_MODEL_FOR_QUESTION_ANSWERING_MAPPING







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




class 	lowercase__				(					_BaseAutoModelClass				):
   __UpperCAmelCase							       =					FLAX_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING







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




class 	lowercase__				(					_BaseAutoModelClass				):
   __UpperCAmelCase							       =					FLAX_MODEL_FOR_MULTIPLE_CHOICE_MAPPING







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




class 	lowercase__				(					_BaseAutoModelClass				):
   __UpperCAmelCase							       =					FLAX_MODEL_FOR_NEXT_SENTENCE_PREDICTION_MAPPING







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




class 	lowercase__				(					_BaseAutoModelClass				):
   __UpperCAmelCase							       =					FLAX_MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING







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




class 	lowercase__				(					_BaseAutoModelClass				):
   __UpperCAmelCase							       =					FLAX_MODEL_FOR_VISION_2_SEQ_MAPPING







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




class 	lowercase__				(					_BaseAutoModelClass				):
   __UpperCAmelCase							       =					FLAX_MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING







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

