


"""simple docstring"""
from __future__ import annotations

import unittest

from transformers import BlenderbotSmallConfig, BlenderbotSmallTokenizer, is_tf_available
from transformers.testing_utils import require_tf, require_tokenizers, slow
from transformers.utils import cached_property

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


if is_tf_available():
 import tensorflow as tf

 from transformers import TFAutoModelForSeqaSeqLM, TFBlenderbotSmallForConditionalGeneration, TFBlenderbotSmallModel




@require_tf
class 	lowercase__				:
   __UpperCAmelCase							       =					BlenderbotSmallConfig
   __UpperCAmelCase							       =					{}
   __UpperCAmelCase							       =					'''gelu'''





   def __init__(							self      ,   SCREAMING_SNAKE_CASE      ,   SCREAMING_SNAKE_CASE=13      ,   SCREAMING_SNAKE_CASE=7      ,   SCREAMING_SNAKE_CASE=True      ,   SCREAMING_SNAKE_CASE=False      ,   SCREAMING_SNAKE_CASE=99      ,   SCREAMING_SNAKE_CASE=32      ,   SCREAMING_SNAKE_CASE=2      ,   SCREAMING_SNAKE_CASE=4      ,   SCREAMING_SNAKE_CASE=37      ,   SCREAMING_SNAKE_CASE=0.1      ,   SCREAMING_SNAKE_CASE=0.1      ,   SCREAMING_SNAKE_CASE=20      ,   SCREAMING_SNAKE_CASE=2      ,   SCREAMING_SNAKE_CASE=1      ,   SCREAMING_SNAKE_CASE=0      ,   )  ->  Union[str, Any]:
      _lowerCamelCase	:					List[Any]   =		parent
      _lowerCamelCase	:					Dict   =		batch_size
      _lowerCamelCase	:					List[str]   =		seq_length
      _lowerCamelCase	:					Union[str, Any]   =		is_training
      _lowerCamelCase	:					Optional[int]   =		use_labels
      _lowerCamelCase	:					List[str]   =		vocab_size
      _lowerCamelCase	:					int   =		hidden_size
      _lowerCamelCase	:					Optional[int]   =		num_hidden_layers
      _lowerCamelCase	:					Dict   =		num_attention_heads
      _lowerCamelCase	:					Tuple   =		intermediate_size

      _lowerCamelCase	:					Tuple   =		hidden_dropout_prob
      _lowerCamelCase	:					List[str]   =		attention_probs_dropout_prob
      _lowerCamelCase	:					Any   =		max_position_embeddings
      _lowerCamelCase	:					Optional[int]   =		eos_token_id
      _lowerCamelCase	:					List[str]   =		pad_token_id
      _lowerCamelCase	:					Dict   =		bos_token_id





   def 		UpperCamelCase_  (							self)  ->  List[str]:
      _lowerCamelCase	:					int   =		ids_tensor([self.batch_size, self.seq_length - 1]      ,   self.vocab_size)
      _lowerCamelCase	:					Optional[Any]   =		tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size)      ,   1)
      _lowerCamelCase	:					str   =		tf.concat([input_ids, eos_tensor]      ,   axis=1)

      _lowerCamelCase	:					Union[str, Any]   =		ids_tensor([self.batch_size, self.seq_length]      ,   self.vocab_size)

      _lowerCamelCase	:					Optional[Any]   =		self.config_cls(
          vocab_size=self.vocab_size      ,   d_model=self.hidden_size      ,   encoder_layers=self.num_hidden_layers      ,   decoder_layers=self.num_hidden_layers      ,   encoder_attention_heads=self.num_attention_heads      ,   decoder_attention_heads=self.num_attention_heads      ,   encoder_ffn_dim=self.intermediate_size      ,   decoder_ffn_dim=self.intermediate_size      ,   dropout=self.hidden_dropout_prob      ,   attention_dropout=self.attention_probs_dropout_prob      ,   max_position_embeddings=self.max_position_embeddings      ,   eos_token_ids=[2]      ,   bos_token_id=self.bos_token_id      ,   pad_token_id=self.pad_token_id      ,   decoder_start_token_id=self.pad_token_id      ,   **self.config_updates      ,   )
      _lowerCamelCase	:					Tuple   =		prepare_blenderbot_small_inputs_dict(SCREAMING_SNAKE_CASE      ,   SCREAMING_SNAKE_CASE      ,   SCREAMING_SNAKE_CASE)
      return config, inputs_dict





   def 		UpperCamelCase_  (							self      ,   SCREAMING_SNAKE_CASE      ,   SCREAMING_SNAKE_CASE)  ->  Union[str, Any]:
      _lowerCamelCase	:					List[str]   =		TFBlenderbotSmallModel(config=SCREAMING_SNAKE_CASE).get_decoder()
      _lowerCamelCase	:					int   =		inputs_dict["""input_ids"""]

      _lowerCamelCase	:					List[str]   =		input_ids[:1, :]
      _lowerCamelCase	:					Any   =		inputs_dict["""attention_mask"""][:1, :]
      _lowerCamelCase	:					Any   =		inputs_dict["""head_mask"""]
      _lowerCamelCase	:					Union[str, Any]   =		1

      # first forward pass
      _lowerCamelCase	:					int   =		model(SCREAMING_SNAKE_CASE      ,   attention_mask=SCREAMING_SNAKE_CASE      ,   head_mask=SCREAMING_SNAKE_CASE      ,   use_cache=SCREAMING_SNAKE_CASE)

      _lowerCamelCase   ,    _lowerCamelCase	:					int   =		outputs.to_tuple()

      # create hypothetical next token and extent to next_input_ids
      _lowerCamelCase	:					Tuple   =		ids_tensor((self.batch_size, 3)      ,   config.vocab_size)
      _lowerCamelCase	:					List[Any]   =		tf.cast(ids_tensor((self.batch_size, 3)      ,   2)      ,   tf.inta)

      # append to next input_ids and
      _lowerCamelCase	:					List[str]   =		tf.concat([input_ids, next_tokens]      ,   axis=-1)
      _lowerCamelCase	:					Optional[Any]   =		tf.concat([attention_mask, next_attn_mask]      ,   axis=-1)

      _lowerCamelCase	:					List[Any]   =		model(SCREAMING_SNAKE_CASE      ,   attention_mask=SCREAMING_SNAKE_CASE)[0]
      _lowerCamelCase	:					Optional[Any]   =		model(SCREAMING_SNAKE_CASE      ,   attention_mask=SCREAMING_SNAKE_CASE      ,   past_key_values=SCREAMING_SNAKE_CASE)[0]

      self.parent.assertEqual(next_tokens.shape[1]      ,   output_from_past.shape[1])

      # select random slice
      _lowerCamelCase	:					Union[str, Any]   =		int(ids_tensor((1,)      ,   output_from_past.shape[-1]))
      _lowerCamelCase	:					Tuple   =		output_from_no_past[:, -3:, random_slice_idx]
      _lowerCamelCase	:					Optional[int]   =		output_from_past[:, :, random_slice_idx]

      # test that outputs are equal for slice
      tf.debugging.assert_near(SCREAMING_SNAKE_CASE      ,   SCREAMING_SNAKE_CASE      ,   rtol=1e-3)







def 				_snake_case      (   __snake_case							:      str						, __snake_case							:      Tuple						, __snake_case							:      Optional[Any]						, __snake_case							:      Optional[int]=None						, __snake_case							:      Any=None						, __snake_case							:      Optional[int]=None						, __snake_case							:      str=None						, __snake_case							:      str=None						, ):




   """simple docstring"""







   if attention_mask is None:
      _lowerCamelCase	:					Any   =		tf.cast(tf.math.not_equal(__snake_case						, config.pad_token_id   )						, tf.inta   )
   if decoder_attention_mask is None:
      _lowerCamelCase	:					Any   =		tf.concat(
          [
              tf.ones(decoder_input_ids[:, :1].shape						, dtype=tf.inta   ),
              tf.cast(tf.math.not_equal(decoder_input_ids[:, 1:]						, config.pad_token_id   )						, tf.inta   ),
          ]						, axis=-1						, )
   if head_mask is None:
      _lowerCamelCase	:					str   =		tf.ones((config.encoder_layers, config.encoder_attention_heads)   )
   if decoder_head_mask is None:
      _lowerCamelCase	:					Union[str, Any]   =		tf.ones((config.decoder_layers, config.decoder_attention_heads)   )
   if cross_attn_head_mask is None:
      _lowerCamelCase	:					int   =		tf.ones((config.decoder_layers, config.decoder_attention_heads)   )
   return {
       "input_ids": input_ids,
       "decoder_input_ids": decoder_input_ids,
       "attention_mask": attention_mask,
       "decoder_attention_mask": decoder_attention_mask,
       "head_mask": head_mask,
       "decoder_head_mask": decoder_head_mask,
       "cross_attn_head_mask": cross_attn_head_mask,
   }




@require_tf
class 	lowercase__				(					A_ ,A_ ,unittest.TestCase				):
   __UpperCAmelCase							       =					(
       (TFBlenderbotSmallForConditionalGeneration, TFBlenderbotSmallModel) if is_tf_available() else ()
   )
   __UpperCAmelCase							       =					(TFBlenderbotSmallForConditionalGeneration,) if is_tf_available() else ()
   __UpperCAmelCase							       =					(
       {
           '''conversational''': TFBlenderbotSmallForConditionalGeneration,
           '''feature-extraction''': TFBlenderbotSmallModel,
           '''summarization''': TFBlenderbotSmallForConditionalGeneration,
           '''text2text-generation''': TFBlenderbotSmallForConditionalGeneration,
           '''translation''': TFBlenderbotSmallForConditionalGeneration,
       }
       if is_tf_available()
       else {}
   )
   __UpperCAmelCase							       =					True
   __UpperCAmelCase							       =					False
   __UpperCAmelCase							       =					False





   def 		UpperCamelCase_  (							self)  ->  List[str]:
      _lowerCamelCase	:					Optional[Any]   =		TFBlenderbotSmallModelTester(self)
      _lowerCamelCase	:					Tuple   =		ConfigTester(self      ,   config_class=SCREAMING_SNAKE_CASE)





   def 		UpperCamelCase_  (							self)  ->  List[Any]:
      self.config_tester.run_common_tests()





   def 		UpperCamelCase_  (							self)  ->  Dict:
      _lowerCamelCase	:					str   =		self.model_tester.prepare_config_and_inputs_for_common()
      self.model_tester.check_decoder_model_past_large_inputs(*SCREAMING_SNAKE_CASE)







@require_tokenizers
@require_tf
class 	lowercase__				(					unittest.TestCase				):
   __UpperCAmelCase							       =					[
       '''Social anxiety\nWow, I am never shy. Do you have anxiety?\nYes. I end up sweating and blushing and feel like '''
       '''  i\'m going to throw up.\nand why is that?'''
   ]
   __UpperCAmelCase							       =					'''facebook/blenderbot_small-90M'''





   @cached_property
   def 		UpperCamelCase_  (							self)  ->  Dict:
      # use "old" tokenizer here because of bug when downloading new tokenizer
      return BlenderbotSmallTokenizer.from_pretrained("""facebook/blenderbot-90M""")





   @cached_property
   def 		UpperCamelCase_  (							self)  ->  Any:
      _lowerCamelCase	:					int   =		TFAutoModelForSeqaSeqLM.from_pretrained(self.model_name)
      return model





   @slow
   def 		UpperCamelCase_  (							self)  ->  Optional[Any]:
      _lowerCamelCase	:					int   =		self.tokenizer(self.src_text      ,   return_tensors="""tf""")
      _lowerCamelCase	:					List[Any]   =		self.model.generate(
          model_inputs.input_ids      ,   attention_mask=model_inputs.attention_mask      ,   num_beams=2      ,   use_cache=SCREAMING_SNAKE_CASE      ,   )
      _lowerCamelCase	:					Dict   =		self.tokenizer.batch_decode(generated_ids.numpy()      ,   skip_special_tokens=SCREAMING_SNAKE_CASE)[0]
      assert generated_words in (
          "i don't know. i just feel like i'm going to throw up. it's not fun.",
          "i'm not sure. i just feel like i've been feeling like i have to be in a certain place",
          "i'm not sure. i just feel like i've been in a bad situation.",
      )





"""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"""
)

