


'''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


_lowercase    :    Any			 =  logging.get_logger(__name__)


_lowercase    :    Dict			 =  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"""),
    ]
)

_lowercase    :    List[Any]			 =  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"""),
    ]
)

_lowercase    :    List[str]			 =  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"""),
    ]
)

_lowercase    :    Tuple			 =  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"""),
    ]
)

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

_lowercase    :    Optional[Any]			 =  OrderedDict(
    [
        ("""vision-encoder-decoder""", """FlaxVisionEncoderDecoderModel"""),
    ]
)

_lowercase    :    str			 =  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"""),
    ]
)

_lowercase    :    Dict			 =  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"""),
    ]
)

_lowercase    :    Optional[int]			 =  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"""),
    ]
)

_lowercase    :    Optional[Any]			 =  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"""),
    ]
)

_lowercase    :    Optional[int]			 =  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"""),
    ]
)

_lowercase    :    Dict			 =  OrderedDict(
    [
        ("""bert""", """FlaxBertForNextSentencePrediction"""),
    ]
)

_lowercase    :    List[str]			 =  OrderedDict(
    [
        ("""speech-encoder-decoder""", """FlaxSpeechEncoderDecoderModel"""),
        ("""whisper""", """FlaxWhisperForConditionalGeneration"""),
    ]
)

_lowercase    :    Optional[Any]			 =  OrderedDict(
    [
        ("""whisper""", """FlaxWhisperForAudioClassification"""),
    ]
)

_lowercase    :    Optional[int]			 =  _LazyAutoMapping(CONFIG_MAPPING_NAMES, FLAX_MODEL_MAPPING_NAMES)
_lowercase    :    str			 =  _LazyAutoMapping(CONFIG_MAPPING_NAMES, FLAX_MODEL_FOR_PRETRAINING_MAPPING_NAMES)
_lowercase    :    Optional[Any]			 =  _LazyAutoMapping(CONFIG_MAPPING_NAMES, FLAX_MODEL_FOR_MASKED_LM_MAPPING_NAMES)
_lowercase    :    Optional[int]			 =  _LazyAutoMapping(
    CONFIG_MAPPING_NAMES, FLAX_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING_NAMES
)
_lowercase    :    List[Any]			 =  _LazyAutoMapping(
    CONFIG_MAPPING_NAMES, FLAX_MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING_NAMES
)
_lowercase    :    Dict			 =  _LazyAutoMapping(CONFIG_MAPPING_NAMES, FLAX_MODEL_FOR_VISION_2_SEQ_MAPPING_NAMES)
_lowercase    :    Tuple			 =  _LazyAutoMapping(CONFIG_MAPPING_NAMES, FLAX_MODEL_FOR_CAUSAL_LM_MAPPING_NAMES)
_lowercase    :    int			 =  _LazyAutoMapping(
    CONFIG_MAPPING_NAMES, FLAX_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING_NAMES
)
_lowercase    :    Optional[Any]			 =  _LazyAutoMapping(
    CONFIG_MAPPING_NAMES, FLAX_MODEL_FOR_QUESTION_ANSWERING_MAPPING_NAMES
)
_lowercase    :    List[Any]			 =  _LazyAutoMapping(
    CONFIG_MAPPING_NAMES, FLAX_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING_NAMES
)
_lowercase    :    int			 =  _LazyAutoMapping(
    CONFIG_MAPPING_NAMES, FLAX_MODEL_FOR_MULTIPLE_CHOICE_MAPPING_NAMES
)
_lowercase    :    List[str]			 =  _LazyAutoMapping(
    CONFIG_MAPPING_NAMES, FLAX_MODEL_FOR_NEXT_SENTENCE_PREDICTION_MAPPING_NAMES
)
_lowercase    :    str			 =  _LazyAutoMapping(
    CONFIG_MAPPING_NAMES, FLAX_MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING_NAMES
)
_lowercase    :    str			 =  _LazyAutoMapping(
    CONFIG_MAPPING_NAMES, FLAX_MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING_NAMES
)






class  UpperCamelCase__(     _BaseAutoModelClass ):
				__magic_name__       :   Tuple							= FLAX_MODEL_MAPPING







_lowercase    :    Optional[int]			 =  auto_class_update(FlaxAutoModel)






class  UpperCamelCase__(     _BaseAutoModelClass ):
				__magic_name__       :   Dict							= FLAX_MODEL_FOR_PRETRAINING_MAPPING







_lowercase    :    Dict			 =  auto_class_update(FlaxAutoModelForPreTraining, head_doc="""pretraining""")






class  UpperCamelCase__(     _BaseAutoModelClass ):
				__magic_name__       :   List[Any]							= FLAX_MODEL_FOR_CAUSAL_LM_MAPPING







_lowercase    :    Optional[int]			 =  auto_class_update(FlaxAutoModelForCausalLM, head_doc="""causal language modeling""")






class  UpperCamelCase__(     _BaseAutoModelClass ):
				__magic_name__       :   str							= FLAX_MODEL_FOR_MASKED_LM_MAPPING







_lowercase    :    Tuple			 =  auto_class_update(FlaxAutoModelForMaskedLM, head_doc="""masked language modeling""")






class  UpperCamelCase__(     _BaseAutoModelClass ):
				__magic_name__       :   Tuple							= FLAX_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING







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






class  UpperCamelCase__(     _BaseAutoModelClass ):
				__magic_name__       :   Tuple							= FLAX_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING







_lowercase    :    List[Any]			 =  auto_class_update(
    FlaxAutoModelForSequenceClassification, head_doc="""sequence classification"""
)






class  UpperCamelCase__(     _BaseAutoModelClass ):
				__magic_name__       :   List[Any]							= FLAX_MODEL_FOR_QUESTION_ANSWERING_MAPPING







_lowercase    :    Tuple			 =  auto_class_update(FlaxAutoModelForQuestionAnswering, head_doc="""question answering""")






class  UpperCamelCase__(     _BaseAutoModelClass ):
				__magic_name__       :   Dict							= FLAX_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING







_lowercase    :    str			 =  auto_class_update(
    FlaxAutoModelForTokenClassification, head_doc="""token classification"""
)






class  UpperCamelCase__(     _BaseAutoModelClass ):
				__magic_name__       :   List[Any]							= FLAX_MODEL_FOR_MULTIPLE_CHOICE_MAPPING







_lowercase    :    Union[str, Any]			 =  auto_class_update(FlaxAutoModelForMultipleChoice, head_doc="""multiple choice""")






class  UpperCamelCase__(     _BaseAutoModelClass ):
				__magic_name__       :   Optional[Any]							= FLAX_MODEL_FOR_NEXT_SENTENCE_PREDICTION_MAPPING







_lowercase    :    Optional[int]			 =  auto_class_update(
    FlaxAutoModelForNextSentencePrediction, head_doc="""next sentence prediction"""
)






class  UpperCamelCase__(     _BaseAutoModelClass ):
				__magic_name__       :   Any							= FLAX_MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING







_lowercase    :    List[Any]			 =  auto_class_update(
    FlaxAutoModelForImageClassification, head_doc="""image classification"""
)






class  UpperCamelCase__(     _BaseAutoModelClass ):
				__magic_name__       :   Tuple							= FLAX_MODEL_FOR_VISION_2_SEQ_MAPPING







_lowercase    :    List[Any]			 =  auto_class_update(FlaxAutoModelForVisionaSeq, head_doc="""vision-to-text modeling""")






class  UpperCamelCase__(     _BaseAutoModelClass ):
				__magic_name__       :   str							= FLAX_MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING







_lowercase    :    List[Any]			 =  auto_class_update(
    FlaxAutoModelForSpeechSeqaSeq, head_doc="""sequence-to-sequence speech-to-text modeling"""
)









'''simple docstring'''



from argparse import ArgumentParser

from . import BaseTransformersCLICommand


def    lowerCamelCase__   ( A     :    List[Any]  ):




						'''simple docstring'''






						return DownloadCommand(args.model   ,     args.cache_dir   ,     args.force   ,     args.trust_remote_code  )






class  UpperCamelCase__(     lowerCAmelCase ):




				@staticmethod
				def   a__(	lowerCAmelCase    :				ArgumentParser   )->  Optional[Any]:

										"""simple docstring"""




										UpperCAmelCase         =			parser.add_parser('''download'''   )
										download_parser.add_argument(
										    '''--cache-dir'''   ,					type=lowerCAmelCase   ,					default=lowerCAmelCase   ,					help='''Path to location to store the models'''   )
										download_parser.add_argument(
										    '''--force'''   ,					action='''store_true'''   ,					help='''Force the model to be download even if already in cache-dir'''   )
										download_parser.add_argument(
										    '''--trust-remote-code'''   ,					action='''store_true'''   ,					help='''Whether or not to allow for custom models defined on the Hub in their own modeling files. Use only if you\'ve reviewed the code as it will execute on your local machine'''   ,					)
										download_parser.add_argument('''model'''   ,					type=lowerCAmelCase   ,					help='''Name of the model to download'''   )
										download_parser.set_defaults(func=lowerCAmelCase   )




				def __init__(	self    :				Dict   ,					lowerCAmelCase    :				str   ,					lowerCAmelCase    :				str   ,					lowerCAmelCase    :				bool   ,					lowerCAmelCase    :				bool   )->  Any:

										"""simple docstring"""




										UpperCAmelCase         =			model
										UpperCAmelCase         =			cache
										UpperCAmelCase         =			force
										UpperCAmelCase         =			trust_remote_code





				def   a__(	self    :				int   )->  Optional[Any]:

										"""simple docstring"""




										from ..models.auto import AutoModel, AutoTokenizer

										AutoModel.from_pretrained(
										    self._model   ,					cache_dir=self._cache   ,					force_download=self._force   ,					trust_remote_code=self._trust_remote_code   )
										AutoTokenizer.from_pretrained(
										    self._model   ,					cache_dir=self._cache   ,					force_download=self._force   ,					trust_remote_code=self._trust_remote_code   )





