'''simple docstring'''





from __future__ import annotations

import unittest

from transformers import RoFormerConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
    import tensorflow as tf

    from transformers import (
        TFRoFormerForCausalLM,
        TFRoFormerForMaskedLM,
        TFRoFormerForMultipleChoice,
        TFRoFormerForQuestionAnswering,
        TFRoFormerForSequenceClassification,
        TFRoFormerForTokenClassification,
        TFRoFormerModel,
    )
    from transformers.models.roformer.modeling_tf_roformer import (
        TFRoFormerSelfAttention,
        TFRoFormerSinusoidalPositionalEmbedding,
    )




class 	__A				:



    def __init__(self   :   Union[str, Any]       ,      __a   :   Dict       ,      __a   :   Dict=13       ,      __a   :   List[str]=7       ,      __a   :   Optional[Any]=True       ,      __a   :   Any=True       ,      __a   :   List[Any]=True       ,      __a   :   Any=True       ,      __a   :   Any=99       ,      __a   :   Any=32       ,      __a   :   Union[str, Any]=2       ,      __a   :   Optional[int]=4       ,      __a   :   Optional[int]=37       ,      __a   :   Union[str, Any]="gelu"       ,      __a   :   Dict=0.1       ,      __a   :   str=0.1       ,      __a   :   Any=512       ,      __a   :   Tuple=16       ,      __a   :   Optional[int]=2       ,      __a   :   List[str]=0.02       ,      __a   :   Any=3       ,      __a   :   str=4       ,      __a   :   int=None       ,      ):
         UpperCAmelCase_	   =       parent
         UpperCAmelCase_	   =       13
         UpperCAmelCase_	   =       7
         UpperCAmelCase_	   =       True
         UpperCAmelCase_	   =       True
         UpperCAmelCase_	   =       True
         UpperCAmelCase_	   =       True
         UpperCAmelCase_	   =       99
         UpperCAmelCase_	   =       32
         UpperCAmelCase_	   =       2
         UpperCAmelCase_	   =       4
         UpperCAmelCase_	   =       37
         UpperCAmelCase_	   =       "gelu"
         UpperCAmelCase_	   =       0.1
         UpperCAmelCase_	   =       0.1
         UpperCAmelCase_	   =       512
         UpperCAmelCase_	   =       16
         UpperCAmelCase_	   =       2
         UpperCAmelCase_	   =       0.02
         UpperCAmelCase_	   =       3
         UpperCAmelCase_	   =       4
         UpperCAmelCase_	   =       None



    def  _lowercase						(self   :   Any			):
         UpperCAmelCase_	   =       ids_tensor([self.batch_size, self.seq_length]       ,      self.vocab_size			)

         UpperCAmelCase_	   =       None
         if self.use_input_mask:
              UpperCAmelCase_	   =       random_attention_mask([self.batch_size, self.seq_length]			)

         UpperCAmelCase_	   =       None
         if self.use_token_type_ids:
              UpperCAmelCase_	   =       ids_tensor([self.batch_size, self.seq_length]       ,      self.type_vocab_size			)

         UpperCAmelCase_	   =       None
         UpperCAmelCase_	   =       None
         UpperCAmelCase_	   =       None
         if self.use_labels:
              UpperCAmelCase_	   =       ids_tensor([self.batch_size]       ,      self.type_sequence_label_size			)
              UpperCAmelCase_	   =       ids_tensor([self.batch_size, self.seq_length]       ,      self.num_labels			)
              UpperCAmelCase_	   =       ids_tensor([self.batch_size]       ,      self.num_choices			)

         UpperCAmelCase_	   =       RoFormerConfig(
             vocab_size=self.vocab_size       ,      hidden_size=self.hidden_size       ,      num_hidden_layers=self.num_hidden_layers       ,      num_attention_heads=self.num_attention_heads       ,      intermediate_size=self.intermediate_size       ,      hidden_act=self.hidden_act       ,      hidden_dropout_prob=self.hidden_dropout_prob       ,      attention_probs_dropout_prob=self.attention_probs_dropout_prob       ,      max_position_embeddings=self.max_position_embeddings       ,      type_vocab_size=self.type_vocab_size       ,      initializer_range=self.initializer_range       ,      return_dict=__a       ,      )

         return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels



    def  _lowercase						(self   :   int       ,      __a   :   Tuple       ,      __a   :   Dict       ,      __a   :   Union[str, Any]       ,      __a   :   List[str]       ,      __a   :   Union[str, Any]       ,      __a   :   List[str]       ,      __a   :   Dict			):
         UpperCAmelCase_	   =       TFRoFormerModel(config=__a			)
         UpperCAmelCase_	   =       {"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}

         UpperCAmelCase_	   =       [input_ids, input_mask]
         UpperCAmelCase_	   =       model(__a			)

         UpperCAmelCase_	   =       model(__a			)

         self.parent.assertEqual(result.last_hidden_state.shape       ,      (self.batch_size, self.seq_length, self.hidden_size)			)



    def  _lowercase						(self   :   List[Any]       ,      __a   :   Optional[Any]       ,      __a   :   Optional[Any]       ,      __a   :   Tuple       ,      __a   :   Optional[Any]       ,      __a   :   Tuple       ,      __a   :   Any       ,      __a   :   str			):
         UpperCAmelCase_	   =       True
         UpperCAmelCase_	   =       TFRoFormerForCausalLM(config=__a			)
         UpperCAmelCase_	   =       {
             "input_ids": input_ids,
             "attention_mask": input_mask,
             "token_type_ids": token_type_ids,
         }
         UpperCAmelCase_	   =       model(__a			)["logits"]
         self.parent.assertListEqual(
             list(prediction_scores.numpy().shape			)       ,      [self.batch_size, self.seq_length, self.vocab_size]			)



    def  _lowercase						(self   :   Dict       ,      __a   :   int       ,      __a   :   Optional[int]       ,      __a   :   Union[str, Any]       ,      __a   :   Any       ,      __a   :   Optional[int]       ,      __a   :   Union[str, Any]       ,      __a   :   List[Any]			):
         UpperCAmelCase_	   =       TFRoFormerForMaskedLM(config=__a			)
         UpperCAmelCase_	   =       {
             "input_ids": input_ids,
             "attention_mask": input_mask,
             "token_type_ids": token_type_ids,
         }
         UpperCAmelCase_	   =       model(__a			)
         self.parent.assertEqual(result.logits.shape       ,      (self.batch_size, self.seq_length, self.vocab_size)			)



    def  _lowercase						(self   :   int       ,      __a   :   List[Any]       ,      __a   :   Any       ,      __a   :   Dict       ,      __a   :   List[Any]       ,      __a   :   Any       ,      __a   :   List[str]       ,      __a   :   int			):
         UpperCAmelCase_	   =       self.num_labels
         UpperCAmelCase_	   =       TFRoFormerForSequenceClassification(config=__a			)
         UpperCAmelCase_	   =       {
             "input_ids": input_ids,
             "attention_mask": input_mask,
             "token_type_ids": token_type_ids,
         }

         UpperCAmelCase_	   =       model(__a			)
         self.parent.assertEqual(result.logits.shape       ,      (self.batch_size, self.num_labels)			)



    def  _lowercase						(self   :   Optional[Any]       ,      __a   :   str       ,      __a   :   Any       ,      __a   :   Optional[Any]       ,      __a   :   Tuple       ,      __a   :   Tuple       ,      __a   :   Any       ,      __a   :   Any			):
         UpperCAmelCase_	   =       self.num_choices
         UpperCAmelCase_	   =       TFRoFormerForMultipleChoice(config=__a			)
         UpperCAmelCase_	   =       tf.tile(tf.expand_dims(__a       ,      1			)       ,      (1, self.num_choices, 1)			)
         UpperCAmelCase_	   =       tf.tile(tf.expand_dims(__a       ,      1			)       ,      (1, self.num_choices, 1)			)
         UpperCAmelCase_	   =       tf.tile(tf.expand_dims(__a       ,      1			)       ,      (1, self.num_choices, 1)			)
         UpperCAmelCase_	   =       {
             "input_ids": multiple_choice_inputs_ids,
             "attention_mask": multiple_choice_input_mask,
             "token_type_ids": multiple_choice_token_type_ids,
         }
         UpperCAmelCase_	   =       model(__a			)
         self.parent.assertEqual(result.logits.shape       ,      (self.batch_size, self.num_choices)			)



    def  _lowercase						(self   :   Union[str, Any]       ,      __a   :   Any       ,      __a   :   int       ,      __a   :   Any       ,      __a   :   Optional[int]       ,      __a   :   Tuple       ,      __a   :   Any       ,      __a   :   Any			):
         UpperCAmelCase_	   =       self.num_labels
         UpperCAmelCase_	   =       TFRoFormerForTokenClassification(config=__a			)
         UpperCAmelCase_	   =       {
             "input_ids": input_ids,
             "attention_mask": input_mask,
             "token_type_ids": token_type_ids,
         }
         UpperCAmelCase_	   =       model(__a			)
         self.parent.assertEqual(result.logits.shape       ,      (self.batch_size, self.seq_length, self.num_labels)			)



    def  _lowercase						(self   :   Optional[Any]       ,      __a   :   Optional[int]       ,      __a   :   Any       ,      __a   :   Dict       ,      __a   :   Optional[int]       ,      __a   :   Optional[Any]       ,      __a   :   Any       ,      __a   :   Dict			):
         UpperCAmelCase_	   =       TFRoFormerForQuestionAnswering(config=__a			)
         UpperCAmelCase_	   =       {
             "input_ids": input_ids,
             "attention_mask": input_mask,
             "token_type_ids": token_type_ids,
         }

         UpperCAmelCase_	   =       model(__a			)
         self.parent.assertEqual(result.start_logits.shape       ,      (self.batch_size, self.seq_length)			)
         self.parent.assertEqual(result.end_logits.shape       ,      (self.batch_size, self.seq_length)			)






    def  _lowercase						(self   :   Union[str, Any]			):
         UpperCAmelCase_	   =       self.prepare_config_and_inputs()
         (
             (
             UpperCAmelCase_			
         )   ,			(
             UpperCAmelCase_			
         )   ,			(
             UpperCAmelCase_			
         )   ,			(
             UpperCAmelCase_			
         )   ,			(
             UpperCAmelCase_			
         )   ,			(
             UpperCAmelCase_			
         )   ,			(
             UpperCAmelCase_			
         )   ,						
         )	   =       config_and_inputs
         UpperCAmelCase_	   =       {"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
         return config, inputs_dict




@require_tf
class 	__A				(			UpperCamelCase__		,     UpperCamelCase__		,     unittest.TestCase     ):
    a__	:						Union[str, Any]	       =  (
        (
            TFRoFormerModel,
            TFRoFormerForCausalLM,
            TFRoFormerForMaskedLM,
            TFRoFormerForQuestionAnswering,
            TFRoFormerForSequenceClassification,
            TFRoFormerForTokenClassification,
            TFRoFormerForMultipleChoice,
        )
        if is_tf_available()
        else ()
    )
    a__	:						List[Any]	       =  (
        {
            """feature-extraction""": TFRoFormerModel,
            """fill-mask""": TFRoFormerForMaskedLM,
            """question-answering""": TFRoFormerForQuestionAnswering,
            """text-classification""": TFRoFormerForSequenceClassification,
            """text-generation""": TFRoFormerForCausalLM,
            """token-classification""": TFRoFormerForTokenClassification,
            """zero-shot""": TFRoFormerForSequenceClassification,
        }
        if is_tf_available()
        else {}
    )

    a__	:						List[str]	       =  False
    a__	:						Optional[Any]	       =  False



    def  _lowercase						(self   :   int       ,      __a   :   Any       ,      __a   :   str       ,      __a   :   Union[str, Any]       ,      __a   :   Optional[Any]       ,      __a   :   Optional[int]			):
         if pipeline_test_casse_name == "TextGenerationPipelineTests":
              return True

         return False



    def  _lowercase						(self   :   List[str]			):
         UpperCAmelCase_	   =       TFRoFormerModelTester(self			)
         UpperCAmelCase_	   =       ConfigTester(self       ,      config_class=__a       ,      hidden_size=37			)



    def  _lowercase						(self   :   Any			):
         self.config_tester.run_common_tests()



    def  _lowercase						(self   :   Tuple			):
         UpperCAmelCase_	   =       self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*__a			)



    def  _lowercase						(self   :   Dict			):
         UpperCAmelCase_	   =       self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_masked_lm(*__a			)



    def  _lowercase						(self   :   Tuple			):
         UpperCAmelCase_	   =       self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_lm_head(*__a			)



    def  _lowercase						(self   :   int			):
         UpperCAmelCase_	   =       self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_multiple_choice(*__a			)



    def  _lowercase						(self   :   str			):
         UpperCAmelCase_	   =       self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_question_answering(*__a			)



    def  _lowercase						(self   :   Tuple			):
         UpperCAmelCase_	   =       self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_sequence_classification(*__a			)



    def  _lowercase						(self   :   Tuple			):
         UpperCAmelCase_	   =       self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_token_classification(*__a			)






    @slow
    def  _lowercase						(self   :   int			):
         UpperCAmelCase_	   =       TFRoFormerModel.from_pretrained("junnyu/roformer_chinese_base"			)
         self.assertIsNotNone(__a			)




@require_tf
class 	__A				(			unittest.TestCase     ):



    @slow
    def  _lowercase						(self   :   Optional[Any]			):
         UpperCAmelCase_	   =       TFRoFormerForMaskedLM.from_pretrained("junnyu/roformer_chinese_base"			)
         UpperCAmelCase_	   =       tf.constant([[0, 1, 2, 3, 4, 5]]			)
         UpperCAmelCase_	   =       model(__a			)[0]

         # TODO Replace vocab size
         UpperCAmelCase_	   =       50000

         UpperCAmelCase_	   =       [1, 6, vocab_size]
         self.assertEqual(output.shape       ,      __a			)

         print(output[:, :3, :3]			)

         # TODO Replace values below with what was printed above.
         UpperCAmelCase_	   =       tf.constant(
             [
                 [
                     [-0.12_05_33_41, -1.0_26_49_01, 0.29_22_19_46],
                     [-1.5_13_37_83, 0.19_74_33, 0.15_19_06_07],
                     [-5.0_13_54_03, -3.90_02_56, -0.84_03_87_64],
                 ]
             ]			)
         tf.debugging.assert_near(output[:, :3, :3]       ,      __a       ,      atol=1E-4			)




@require_tf
class 	__A				(			unittest.TestCase     ):
    a__	:						Union[str, Any]	       =  1e-4



    def  _lowercase						(self   :   str			):
         UpperCAmelCase_	   =       tf.constant([[4, 10]]			)
         UpperCAmelCase_	   =       TFRoFormerSinusoidalPositionalEmbedding(num_positions=6       ,      embedding_dim=6			)

         UpperCAmelCase_	   =       emba(input_ids.shape			)
         UpperCAmelCase_	   =       tf.constant(
             [[0.00_00, 0.00_00, 0.00_00, 1.00_00, 1.00_00, 1.00_00], [0.84_15, 0.04_64, 0.00_22, 0.54_03, 0.99_89, 1.00_00]]			)

         tf.debugging.assert_near(__a       ,      __a       ,      atol=self.tolerance			)






    def  _lowercase						(self   :   Optional[Any]			):
         UpperCAmelCase_	   =       tf.constant(
             [
                 [0.00_00, 0.00_00, 0.00_00, 0.00_00, 0.00_00],
                 [0.84_15, 0.82_19, 0.80_20, 0.78_19, 0.76_17],
                 [0.90_93, 0.93_64, 0.95_81, 0.97_49, 0.98_70],
             ]			)
         UpperCAmelCase_	   =       TFRoFormerSinusoidalPositionalEmbedding(num_positions=512       ,      embedding_dim=512			)
         emba([2, 16, 512]			)
         UpperCAmelCase_	   =       emba.weight[:3, :5]

         tf.debugging.assert_near(__a       ,      __a       ,      atol=self.tolerance			)




@require_tf
class 	__A				(			unittest.TestCase     ):
    a__	:						Optional[int]	       =  1e-4



    def  _lowercase						(self   :   Optional[int]			):
         # 2,12,16,64
         UpperCAmelCase_	   =       tf.reshape(tf.range(2 * 12 * 16 * 64       ,      dtype=tf.floataa			)       ,      shape=(2, 12, 16, 64)			) / 100

         UpperCAmelCase_	   =       -tf.reshape(tf.range(2 * 12 * 16 * 64       ,      dtype=tf.floataa			)       ,      shape=(2, 12, 16, 64)			) / 100

         UpperCAmelCase_	   =       TFRoFormerSinusoidalPositionalEmbedding(num_positions=32       ,      embedding_dim=64			)
         UpperCAmelCase_	   =       embed_positions([2, 16, 768]			)[None, None, :, :]

         UpperCAmelCase_   ,			UpperCAmelCase_	   =       TFRoFormerSelfAttention.apply_rotary_position_embeddings(
             __a       ,      __a       ,      __a			)

         UpperCAmelCase_	   =       tf.constant(
             [
                 [0.00_00, 0.01_00, 0.02_00, 0.03_00, 0.04_00, 0.05_00, 0.06_00, 0.07_00],
                 [-0.20_12, 0.88_97, 0.02_63, 0.94_01, 0.20_74, 0.94_63, 0.34_81, 0.93_43],
                 [-1.70_57, 0.62_71, -1.21_45, 1.38_97, -0.63_03, 1.76_47, -0.11_73, 1.89_85],
                 [-2.17_31, -1.63_97, -2.73_58, 0.28_54, -2.18_40, 1.71_83, -1.30_18, 2.48_71],
                 [0.27_17, -3.61_73, -2.92_06, -2.19_88, -3.66_38, 0.38_58, -2.91_55, 2.29_80],
                 [3.98_59, -2.15_80, -0.79_84, -4.49_04, -4.11_81, -2.02_52, -4.47_82, 1.12_53],
             ]			)
         UpperCAmelCase_	   =       tf.constant(
             [
                 [0.00_00, -0.01_00, -0.02_00, -0.03_00, -0.04_00, -0.05_00, -0.06_00, -0.07_00],
                 [0.20_12, -0.88_97, -0.02_63, -0.94_01, -0.20_74, -0.94_63, -0.34_81, -0.93_43],
                 [1.70_57, -0.62_71, 1.21_45, -1.38_97, 0.63_03, -1.76_47, 0.11_73, -1.89_85],
                 [2.17_31, 1.63_97, 2.73_58, -0.28_54, 2.18_40, -1.71_83, 1.30_18, -2.48_71],
                 [-0.27_17, 3.61_73, 2.92_06, 2.19_88, 3.66_38, -0.38_58, 2.91_55, -2.29_80],
                 [-3.98_59, 2.15_80, 0.79_84, 4.49_04, 4.11_81, 2.02_52, 4.47_82, -1.12_53],
             ]			)

         tf.debugging.assert_near(query_layer[0, 0, :6, :8]       ,      __a       ,      atol=self.tolerance			)
         tf.debugging.assert_near(key_layer[0, 0, :6, :8]       ,      __a       ,      atol=self.tolerance			)





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


SCREAMING_SNAKE_CASE_:					Any  =logging.get_logger(__name__)


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

SCREAMING_SNAKE_CASE_:					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'),
    ]
)

SCREAMING_SNAKE_CASE_:					Any  =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'),
    ]
)

SCREAMING_SNAKE_CASE_:					int  =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'),
    ]
)

SCREAMING_SNAKE_CASE_:					str  =OrderedDict(
    [
        # Model for Image-classsification
        ('beit', 'FlaxBeitForImageClassification'),
        ('regnet', 'FlaxRegNetForImageClassification'),
        ('resnet', 'FlaxResNetForImageClassification'),
        ('vit', 'FlaxViTForImageClassification'),
    ]
)

SCREAMING_SNAKE_CASE_:					str  =OrderedDict(
    [
        ('vision-encoder-decoder', 'FlaxVisionEncoderDecoderModel'),
    ]
)

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

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

SCREAMING_SNAKE_CASE_:					Any  =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'),
    ]
)

SCREAMING_SNAKE_CASE_:					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'),
    ]
)

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

SCREAMING_SNAKE_CASE_:					int  =OrderedDict(
    [
        ('bert', 'FlaxBertForNextSentencePrediction'),
    ]
)

SCREAMING_SNAKE_CASE_:					List[Any]  =OrderedDict(
    [
        ('speech-encoder-decoder', 'FlaxSpeechEncoderDecoderModel'),
        ('whisper', 'FlaxWhisperForConditionalGeneration'),
    ]
)

SCREAMING_SNAKE_CASE_:					Optional[int]  =OrderedDict(
    [
        ('whisper', 'FlaxWhisperForAudioClassification'),
    ]
)

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




class 	__A				(			_BaseAutoModelClass     ):
    a__	:						int	       =  FLAX_MODEL_MAPPING




SCREAMING_SNAKE_CASE_:					Dict  =auto_class_update(FlaxAutoModel)




class 	__A				(			_BaseAutoModelClass     ):
    a__	:						str	       =  FLAX_MODEL_FOR_PRETRAINING_MAPPING




SCREAMING_SNAKE_CASE_:					Dict  =auto_class_update(FlaxAutoModelForPreTraining, head_doc='pretraining')




class 	__A				(			_BaseAutoModelClass     ):
    a__	:						Optional[Any]	       =  FLAX_MODEL_FOR_CAUSAL_LM_MAPPING




SCREAMING_SNAKE_CASE_:					Tuple  =auto_class_update(FlaxAutoModelForCausalLM, head_doc='causal language modeling')




class 	__A				(			_BaseAutoModelClass     ):
    a__	:						Optional[Any]	       =  FLAX_MODEL_FOR_MASKED_LM_MAPPING




SCREAMING_SNAKE_CASE_:					Optional[Any]  =auto_class_update(FlaxAutoModelForMaskedLM, head_doc='masked language modeling')




class 	__A				(			_BaseAutoModelClass     ):
    a__	:						List[Any]	       =  FLAX_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING




SCREAMING_SNAKE_CASE_:					Optional[Any]  =auto_class_update(
    FlaxAutoModelForSeqaSeqLM, head_doc='sequence-to-sequence language modeling', checkpoint_for_example='t5-base'
)




class 	__A				(			_BaseAutoModelClass     ):
    a__	:						Union[str, Any]	       =  FLAX_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING




SCREAMING_SNAKE_CASE_:					Optional[int]  =auto_class_update(
    FlaxAutoModelForSequenceClassification, head_doc='sequence classification'
)




class 	__A				(			_BaseAutoModelClass     ):
    a__	:						Optional[int]	       =  FLAX_MODEL_FOR_QUESTION_ANSWERING_MAPPING




SCREAMING_SNAKE_CASE_:					List[Any]  =auto_class_update(FlaxAutoModelForQuestionAnswering, head_doc='question answering')




class 	__A				(			_BaseAutoModelClass     ):
    a__	:						Tuple	       =  FLAX_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING




SCREAMING_SNAKE_CASE_:					List[Any]  =auto_class_update(
    FlaxAutoModelForTokenClassification, head_doc='token classification'
)




class 	__A				(			_BaseAutoModelClass     ):
    a__	:						Dict	       =  FLAX_MODEL_FOR_MULTIPLE_CHOICE_MAPPING




SCREAMING_SNAKE_CASE_:					Any  =auto_class_update(FlaxAutoModelForMultipleChoice, head_doc='multiple choice')




class 	__A				(			_BaseAutoModelClass     ):
    a__	:						Union[str, Any]	       =  FLAX_MODEL_FOR_NEXT_SENTENCE_PREDICTION_MAPPING




SCREAMING_SNAKE_CASE_:					int  =auto_class_update(
    FlaxAutoModelForNextSentencePrediction, head_doc='next sentence prediction'
)




class 	__A				(			_BaseAutoModelClass     ):
    a__	:						int	       =  FLAX_MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING




SCREAMING_SNAKE_CASE_:					Dict  =auto_class_update(
    FlaxAutoModelForImageClassification, head_doc='image classification'
)




class 	__A				(			_BaseAutoModelClass     ):
    a__	:						Any	       =  FLAX_MODEL_FOR_VISION_2_SEQ_MAPPING




SCREAMING_SNAKE_CASE_:					Optional[int]  =auto_class_update(FlaxAutoModelForVisionaSeq, head_doc='vision-to-text modeling')




class 	__A				(			_BaseAutoModelClass     ):
    a__	:						List[Any]	       =  FLAX_MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING




SCREAMING_SNAKE_CASE_:					Union[str, Any]  =auto_class_update(
    FlaxAutoModelForSpeechSeqaSeq, head_doc='sequence-to-sequence speech-to-text modeling'
)




