

import unittest

from parameterized import parameterized

from transformers import LlamaConfig, is_torch_available, set_seed
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
 import torch

 from transformers import LlamaForCausalLM, LlamaForSequenceClassification, LlamaModel, LlamaTokenizer



class    _a						:


     '''simple docstring'''






     def __init__(			self   ,	A__   ,	A__=13   ,	A__=7   ,	A__=True   ,	A__=True   ,	A__=False   ,	A__=True   ,	A__=99   ,	A__=32   ,	A__=5   ,	A__=4   ,	A__=37   ,	A__="gelu"   ,	A__=0.1   ,	A__=0.1   ,	A__=512   ,	A__=16   ,	A__=2   ,	A__=0.0_2   ,	A__=3   ,	A__=4   ,	A__=None   ,	):
      A__							:   Tuple					       =		parent
      A__							:   Tuple					       =		batch_size
      A__							:   int					       =		seq_length
      A__							:   Optional[int]					       =		is_training
      A__							:   Optional[Any]					       =		use_input_mask
      A__							:   List[str]					       =		use_token_type_ids
      A__							:   List[str]					       =		use_labels
      A__							:   Optional[Any]					       =		vocab_size
      A__							:   List[Any]					       =		hidden_size
      A__							:   List[Any]					       =		num_hidden_layers
      A__							:   List[str]					       =		num_attention_heads
      A__							:   str					       =		intermediate_size
      A__							:   Any					       =		hidden_act
      A__							:   List[str]					       =		hidden_dropout_prob
      A__							:   Optional[Any]					       =		attention_probs_dropout_prob
      A__							:   Any					       =		max_position_embeddings
      A__							:   Optional[int]					       =		type_vocab_size
      A__							:   Union[str, Any]					       =		type_sequence_label_size
      A__							:   Dict					       =		initializer_range
      A__							:   List[Any]					       =		num_labels
      A__							:   List[str]					       =		num_choices
      A__							:   List[Any]					       =		scope






     def   __A      (			self      ):
      A__							:   Tuple					       =		ids_tensor([self.batch_size, self.seq_length]   ,	self.vocab_size      )

      A__							:   Any					       =		None
      if self.use_input_mask:
       A__							:   Optional[Any]					       =		random_attention_mask([self.batch_size, self.seq_length]      )

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

      A__							:   List[Any]					       =		None
      A__							:   int					       =		None
      A__							:   int					       =		None
      if self.use_labels:
       A__							:   Dict					       =		ids_tensor([self.batch_size]   ,	self.type_sequence_label_size      )
       A__							:   Optional[int]					       =		ids_tensor([self.batch_size, self.seq_length]   ,	self.num_labels      )
       A__							:   List[str]					       =		ids_tensor([self.batch_size]   ,	self.num_choices      )

      A__							:   str					       =		self.get_config()

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






     def   __A      (			self      ):
      return LlamaConfig(
          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   ,	is_decoder=__A   ,	initializer_range=self.initializer_range   ,	)






     def   __A      (			self   ,	A__   ,	A__   ,	A__   ,	A__   ,	A__   ,	A__   ,	A__      ):
      A__							:   Tuple					       =		LlamaModel(config=__A      )
      model.to(__A      )
      model.eval()
      A__							:   Union[str, Any]					       =		model(__A   ,	attention_mask=__A      )
      A__							:   Union[str, Any]					       =		model(__A      )
      self.parent.assertEqual(result.last_hidden_state.shape   ,	(self.batch_size, self.seq_length, self.hidden_size)      )






     def   __A      (			self   ,	A__   ,	A__   ,	A__   ,	A__   ,	A__   ,	A__   ,	A__   ,	A__   ,	A__   ,	):
      A__							:   Optional[int]					       =		True
      A__							:   Dict					       =		LlamaModel(__A      )
      model.to(__A      )
      model.eval()
      A__							:   Dict					       =		model(
          __A   ,	attention_mask=__A   ,	encoder_hidden_states=__A   ,	encoder_attention_mask=__A   ,	)
      A__							:   Any					       =		model(
          __A   ,	attention_mask=__A   ,	encoder_hidden_states=__A   ,	)
      A__							:   Optional[Any]					       =		model(__A   ,	attention_mask=__A      )
      self.parent.assertEqual(result.last_hidden_state.shape   ,	(self.batch_size, self.seq_length, self.hidden_size)      )






     def   __A      (			self   ,	A__   ,	A__   ,	A__   ,	A__   ,	A__   ,	A__   ,	A__   ,	A__   ,	A__   ,	):
      A__							:   List[str]					       =		LlamaForCausalLM(config=__A      )
      model.to(__A      )
      model.eval()
      A__							:   Dict					       =		model(__A   ,	attention_mask=__A   ,	labels=__A      )
      self.parent.assertEqual(result.logits.shape   ,	(self.batch_size, self.seq_length, self.vocab_size)      )






     def   __A      (			self   ,	A__   ,	A__   ,	A__   ,	A__   ,	A__   ,	A__   ,	A__   ,	A__   ,	A__   ,	):
      A__							:   Dict					       =		True
      A__							:   List[str]					       =		True
      A__							:   Optional[Any]					       =		LlamaForCausalLM(config=__A      )
      model.to(__A      )
      model.eval()

      # first forward pass
      A__							:   Any					       =		model(
          __A   ,	attention_mask=__A   ,	encoder_hidden_states=__A   ,	encoder_attention_mask=__A   ,	use_cache=__A   ,	)
      A__							:   Tuple					       =		outputs.past_key_values

      # create hypothetical multiple next token and extent to next_input_ids
      A__							:   str					       =		ids_tensor((self.batch_size, 3)   ,	config.vocab_size      )
      A__							:   int					       =		ids_tensor((self.batch_size, 3)   ,	vocab_size=2      )

      # append to next input_ids and
      A__							:   Optional[int]					       =		torch.cat([input_ids, next_tokens]   ,	dim=-1      )
      A__							:   Union[str, Any]					       =		torch.cat([input_mask, next_mask]   ,	dim=-1      )

      A__							:   Optional[Any]					       =		model(
          __A   ,	attention_mask=__A   ,	encoder_hidden_states=__A   ,	encoder_attention_mask=__A   ,	output_hidden_states=__A   ,	)["""hidden_states"""][0]
      A__							:   Union[str, Any]					       =		model(
          __A   ,	attention_mask=__A   ,	encoder_hidden_states=__A   ,	encoder_attention_mask=__A   ,	past_key_values=__A   ,	output_hidden_states=__A   ,	)["""hidden_states"""][0]

      # select random slice
      A__							:   List[Any]					       =		ids_tensor((1,)   ,	output_from_past.shape[-1]      ).item()
      A__							:   Dict					       =		output_from_no_past[:, -3:, random_slice_idx].detach()
      A__							:   Tuple					       =		output_from_past[:, :, random_slice_idx].detach()

      self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]      )

      # test that outputs are equal for slice
      self.parent.assertTrue(torch.allclose(__A   ,	__A   ,	atol=1e-3      )      )






     def   __A      (			self      ):
      A__							:   Optional[int]					       =		self.prepare_config_and_inputs()
      (
          (
          A__
      )       ,       (
          A__
      )       ,       (
          A__
      )       ,       (
          A__
      )       ,       (
          A__
      )       ,       (
          A__
      )       ,       (
          A__
      )       ,       
      )							:   int					       =		config_and_inputs
      A__							:   str					       =		{"""input_ids""": input_ids, """attention_mask""": input_mask}
      return config, inputs_dict



@require_torch
class    _a						(__magic_name__		,       __magic_name__		,       __magic_name__		,       unittest.TestCase    ):


     '''simple docstring'''






     UpperCAmelCase__:       List[str]  							=	(LlamaModel, LlamaForCausalLM, LlamaForSequenceClassification) if is_torch_available() else ()
     UpperCAmelCase__:       Optional[int]  							=	(LlamaForCausalLM,) if is_torch_available() else ()
     UpperCAmelCase__:       List[Any]  							=	(
         {
             '''feature-extraction''': LlamaModel,
             '''text-classification''': LlamaForSequenceClassification,
             '''text-generation''': LlamaForCausalLM,
             '''zero-shot''': LlamaForSequenceClassification,
         }
         if is_torch_available()
         else {}
     )
     UpperCAmelCase__:       Dict  							=	False
     UpperCAmelCase__:       Union[str, Any]  							=	False






     def   __A      (			self      ):
      A__							:   Optional[int]					       =		LlamaModelTester(self      )
      A__							:   Tuple					       =		ConfigTester(self   ,	config_class=__A   ,	hidden_size=37      )






     def   __A      (			self      ):
      self.config_tester.run_common_tests()






     def   __A      (			self      ):
      A__							:   Optional[Any]					       =		self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*__A      )






     def   __A      (			self      ):
      A__							:   Union[str, Any]					       =		self.model_tester.prepare_config_and_inputs()
      for type in ["absolute", "relative_key", "relative_key_query"]:
       A__							:   List[Any]					       =		type
       self.model_tester.create_and_check_model(*__A      )






     def   __A      (			self      ):
      A__       ,       A__							:   Union[str, Any]					       =		self.model_tester.prepare_config_and_inputs_for_common()
      A__							:   List[str]					       =		3
      A__							:   Dict					       =		input_dict["""input_ids"""]
      A__							:   List[str]					       =		input_ids.ne(1      ).to(__A      )
      A__							:   int					       =		ids_tensor([self.model_tester.batch_size]   ,	self.model_tester.type_sequence_label_size      )
      A__							:   Optional[Any]					       =		LlamaForSequenceClassification(__A      )
      model.to(__A      )
      model.eval()
      A__							:   Union[str, Any]					       =		model(__A   ,	attention_mask=__A   ,	labels=__A      )
      self.assertEqual(result.logits.shape   ,	(self.model_tester.batch_size, self.model_tester.num_labels)      )






     def   __A      (			self      ):
      A__       ,       A__							:   Dict					       =		self.model_tester.prepare_config_and_inputs_for_common()
      A__							:   Tuple					       =		3
      A__							:   int					       =		"""single_label_classification"""
      A__							:   Optional[Any]					       =		input_dict["""input_ids"""]
      A__							:   Tuple					       =		input_ids.ne(1      ).to(__A      )
      A__							:   str					       =		ids_tensor([self.model_tester.batch_size]   ,	self.model_tester.type_sequence_label_size      )
      A__							:   Optional[Any]					       =		LlamaForSequenceClassification(__A      )
      model.to(__A      )
      model.eval()
      A__							:   List[Any]					       =		model(__A   ,	attention_mask=__A   ,	labels=__A      )
      self.assertEqual(result.logits.shape   ,	(self.model_tester.batch_size, self.model_tester.num_labels)      )






     def   __A      (			self      ):
      A__       ,       A__							:   int					       =		self.model_tester.prepare_config_and_inputs_for_common()
      A__							:   Any					       =		3
      A__							:   str					       =		"""multi_label_classification"""
      A__							:   Union[str, Any]					       =		input_dict["""input_ids"""]
      A__							:   Any					       =		input_ids.ne(1      ).to(__A      )
      A__							:   Any					       =		ids_tensor(
          [self.model_tester.batch_size, config.num_labels]   ,	self.model_tester.type_sequence_label_size      ).to(torch.float      )
      A__							:   Any					       =		LlamaForSequenceClassification(__A      )
      model.to(__A      )
      model.eval()
      A__							:   Any					       =		model(__A   ,	attention_mask=__A   ,	labels=__A      )
      self.assertEqual(result.logits.shape   ,	(self.model_tester.batch_size, self.model_tester.num_labels)      )






     @unittest.skip("""LLaMA buffers include complex numbers, which breaks this test"""      )
     def   __A      (			self      ):
      pass






     @parameterized.expand([("""linear""",), ("""dynamic""",)]      )
     def   __A      (			self   ,	A__      ):
      A__       ,       A__							:   List[str]					       =		self.model_tester.prepare_config_and_inputs_for_common()
      A__							:   Union[str, Any]					       =		ids_tensor([1, 10]   ,	config.vocab_size      )
      A__							:   Any					       =		ids_tensor([1, int(config.max_position_embeddings * 1.5      )]   ,	config.vocab_size      )

      set_seed(42      )  # Fixed seed at init time so the two models get the same random weights
      A__							:   Optional[int]					       =		LlamaModel(__A      )
      original_model.to(__A      )
      original_model.eval()
      A__							:   List[str]					       =		original_model(__A      ).last_hidden_state
      A__							:   int					       =		original_model(__A      ).last_hidden_state

      set_seed(42      )  # Fixed seed at init time so the two models get the same random weights
      A__							:   str					       =		{"""type""": scaling_type, """factor""": 1_0.0}
      A__							:   int					       =		LlamaModel(__A      )
      scaled_model.to(__A      )
      scaled_model.eval()
      A__							:   List[Any]					       =		scaled_model(__A      ).last_hidden_state
      A__							:   str					       =		scaled_model(__A      ).last_hidden_state

      # Dynamic scaling does not change the RoPE embeddings until it receives an input longer than the original
      # maximum sequence length, so the outputs for the short input should match.
      if scaling_type == "dynamic":
       self.assertTrue(torch.allclose(__A   ,	__A   ,	atol=1e-5      )      )
      else:
       self.assertFalse(torch.allclose(__A   ,	__A   ,	atol=1e-5      )      )

      # The output should be different for long inputs
      self.assertFalse(torch.allclose(__A   ,	__A   ,	atol=1e-5      )      )



@require_torch
class    _a						(unittest.TestCase    ):


     '''simple docstring'''






     @unittest.skip("""Logits are not exactly the same, once we fix the instabalities somehow, will update!"""      )
     @slow
     def   __A      (			self      ):
      A__							:   Optional[Any]					       =		[1, 306, 4658, 278, 6593, 310, 2834, 338]
      A__							:   str					       =		LlamaForCausalLM.from_pretrained("""meta-llama/Llama-2-7b-hf"""   ,	device_map="""auto"""      )
      A__							:   List[Any]					       =		model(torch.tensor([input_ids]      )      )
      # Expected mean on dim = -1
      A__							:   int					       =		torch.tensor([[-6.6_5_5_0, -4.1_2_2_7, -4.9_8_5_9, -3.2_4_0_6, 0.8_2_6_2, -3.0_0_3_3, 1.2_9_6_4, -3.3_6_9_9]]      )
      torch.testing.assert_close(out.mean(-1      )   ,	__A   ,	atol=1e-2   ,	rtol=1e-2      )
      # slicing logits[0, 0, 0:30]
      # fmt: off
      A__							:   str					       =		torch.tensor([-1_2.8_2_8_1, -7.4_4_5_3, -0.4_6_3_9, -8.0_6_2_5, -7.2_5_0_0, -8.0_0_0_0, -6.4_8_8_3, -7.7_6_9_5, -7.8_4_3_8, -7.0_3_1_2, -6.2_1_8_8, -7.1_3_2_8, -1.8_4_9_6, 1.9_9_6_1, -8.6_2_5_0, -6.7_2_2_7, -1_2.8_2_8_1, -6.9_4_9_2, -7.0_7_4_2, -7.7_8_5_2, -7.5_8_2_0, -7.9_0_6_2, -6.9_3_7_5, -7.9_8_0_5, -8.3_4_3_8, -8.1_5_6_2, -8.0_4_6_9, -7.6_2_5_0, -7.7_4_2_2, -7.3_3_9_8,]      )
      # fmt: on
      torch.testing.assert_close(out[0, 0, :30]   ,	__A   ,	atol=1e-5   ,	rtol=1e-5      )






     @unittest.skip("""Logits are not exactly the same, once we fix the instabalities somehow, will update!"""      )
     @slow
     def   __A      (			self      ):
      A__							:   Any					       =		[1, 306, 4658, 278, 6593, 310, 2834, 338]
      A__							:   List[Any]					       =		LlamaForCausalLM.from_pretrained("""meta-llama/Llama-2-13b-hf"""   ,	device_map="""auto"""      )
      A__							:   List[str]					       =		model(torch.tensor(__A      )      )
      # Expected mean on dim = -1
      A__							:   Optional[Any]					       =		torch.tensor([[-2.0_6_2_2, -1.2_7_9_4, -1.1_6_3_8, -0.9_7_8_8, -1.4_6_0_3, -1.0_2_3_8, -1.7_8_9_3, -1.4_4_1_1]]      )
      torch.testing.assert_close(out.mean(-1      )   ,	__A   ,	atol=1e-2   ,	rtol=1e-2      )
      # slicing logits[0, 0, 0:30]
      # fmt: off
      A__							:   str					       =		torch.tensor([-8.1_4_0_6, -8.0_5_4_7, 2.7_4_6_1, -1.2_3_4_4, -0.1_4_4_8, -1.8_2_6_2, -1.0_0_2_0, -1.8_1_5_4, -1.6_8_9_5, -1.8_5_1_6, -2.3_5_7_4, -0.9_2_7_7, 3.7_5_9_8, 6.5_7_4_2, -1.2_9_9_8, -0.1_1_7_7, -8.1_4_0_6, -2.9_6_8_8, -2.9_1_9_9, -3.1_6_9_9, -3.5_2_5_4, -2.3_5_5_5, -2.7_9_8_8, -3.4_1_4_1, -2.8_2_6_2, -4.5_1_9_5, -3.3_3_7_9, -3.3_1_6_4, -2.7_8_3_2, -3.0_2_7_3]      )
      # fmt: on
      torch.testing.assert_close(out[0, 0, :30]   ,	__A   ,	atol=1e-5   ,	rtol=1e-5      )






     @unittest.skip("""Logits are not exactly the same, once we fix the instabalities somehow, will update!"""      )
     @slow
     def   __A      (			self      ):
      A__							:   int					       =		[1, 306, 4658, 278, 6593, 310, 2834, 338]
      A__							:   Any					       =		LlamaForCausalLM.from_pretrained("""meta-llama/Llama-2-13b-chat-hf"""   ,	device_map="""auto"""      )
      A__							:   Tuple					       =		model(torch.tensor(__A      )      )
      # Expected mean on dim = -1
      A__							:   Tuple					       =		torch.tensor([[-0.8_5_6_2, -1.8_5_2_0, -0.7_5_5_1, -0.4_1_6_2, -1.5_1_6_1, -1.2_0_3_8, -2.4_8_2_3, -2.3_2_5_4]]      )
      torch.testing.assert_close(out.mean(-1      )   ,	__A   ,	atol=1e-2   ,	rtol=1e-2      )
      # slicing logits[0, 0, 0:30]
      # fmt: off
      A__							:   Optional[int]					       =		torch.tensor([-2.2_2_2_7, 4.8_8_2_8, 0.9_0_2_3, -0.4_5_7_8, -0.7_8_7_1, -0.1_0_3_3, -0.6_2_2_1, -0.5_7_8_6, -0.7_8_0_3, -1.0_6_7_4, -1.2_9_2_0, -0.1_5_7_0, 0.8_0_0_8, 2.0_7_2_3, -0.9_4_9_7, 0.2_7_7_1, -2.2_2_2_7, -0.7_6_1_2, -1.4_3_4_6, -1.2_0_6_1, -1.6_4_2_6, -0.3_0_0_0, -0.7_1_3_9, -1.1_9_3_4, -1.8_6_9_1, -1.6_9_7_3, -1.5_9_4_7, -1.2_7_0_5, -0.3_5_2_3, -0.5_5_1_3]      )
      # fmt: on
      torch.testing.assert_close(out.mean(-1      )   ,	__A   ,	atol=1e-2   ,	rtol=1e-2      )






     @unittest.skip(
         """Logits are not exactly the same, once we fix the instabalities somehow, will update! Also it is gonna be a `too_slow` test"""      )
     @slow
     def   __A      (			self      ):
      A__							:   Tuple					       =		[1, 306, 4658, 278, 6593, 310, 2834, 338]
      A__							:   Tuple					       =		LlamaForCausalLM.from_pretrained("""meta-llama/Llama-2-70b-hf"""   ,	device_map="""auto"""      )
      A__							:   Union[str, Any]					       =		model(torch.tensor(__A      )      )

      A__							:   List[str]					       =		torch.tensor(
          [[-4.2_3_2_7, -3.3_3_6_0, -4.6_6_6_5, -4.7_6_3_1, -1.8_1_8_0, -3.4_1_7_0, -1.4_2_1_1, -3.1_8_1_0]]   ,	dtype=torch.floataa      )
      torch.testing.assert_close(out.mean(-1      )   ,	__A   ,	atol=1e-2   ,	rtol=1e-2      )
      # fmt: off
      A__							:   List[Any]					       =		torch.tensor([-9.4_9_2_2, -3.9_5_5_1, 1.7_9_9_8, -5.6_7_5_8, -5.1_0_5_5, -5.8_9_8_4, -4.8_3_2_0, -6.8_0_8_6, -6.5_3_9_1, -5.6_1_7_2, -5.5_8_2_0, -5.5_3_5_2, 1.7_8_8_1, 3.6_2_8_9, -6.5_1_1_7, -3.4_7_8_5, -9.5_0_0_0, -6.0_3_5_2, -6.8_1_2_5, -6.0_1_9_5, -6.6_8_3_6, -5.4_7_2_7, -6.2_8_1_2, -6.0_3_9_1, -7.3_3_9_8, -7.4_2_9_7, -7.4_8_4_4, -6.5_8_2_0, -5.8_7_8_9, -5.5_3_1_2]      )
      # fmt: on
      torch.testing.assert_close(out[0, 0, :30]   ,	__A   ,	atol=1e-5   ,	rtol=1e-5      )






     @unittest.skip("""Model is curently gated"""      )
     @slow
     def   __A      (			self      ):
      A__							:   str					       =		"""Simply put, the theory of relativity states that 1) the laws of physics are the same everywhere in the universe and 2) the passage of time and the length of objects can vary depending on the observer\'s frame of reference.\n\nThe first part of the theory, that the laws of physics are the same everywhere, is known as the \"princi"""
      A__							:   Any					       =		"""Simply put, the theory of relativity states that """
      A__							:   Optional[int]					       =		LlamaTokenizer.from_pretrained("""meta-llama/Llama-2-13b-chat-hf"""      )
      A__							:   int					       =		tokenizer.encode(__A   ,	return_tensors="""pt"""      )
      A__							:   Optional[int]					       =		LlamaForCausalLM.from_pretrained(
          """meta-llama/Llama-2-13b-chat-hf"""   ,	device_map="""sequential"""   ,	use_safetensors=__A      )

      # greedy generation outputs
      A__							:   Optional[int]					       =		model.generate(__A   ,	max_new_tokens=64   ,	top_p=__A   ,	temperature=1   ,	do_sample=__A      )
      A__							:   List[str]					       =		tokenizer.decode(generated_ids[0]   ,	skip_special_tokens=__A      )
      self.assertEqual(__A   ,	__A      )






from collections import OrderedDict

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


A_	:			Union[str, Any]  	=      logging.get_logger(__name__)


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

A_	:			Tuple  	=      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'),
    ]
)

A_	:			Tuple  	=      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'),
    ]
)

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

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

A_	:			Union[str, Any]  	=      OrderedDict(
    [
        ('vision-encoder-decoder', 'FlaxVisionEncoderDecoderModel'),
    ]
)

A_	:			Tuple  	=      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'),
    ]
)

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

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

A_	:			Dict  	=      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'),
    ]
)

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

A_	:			List[str]  	=      OrderedDict(
    [
        ('bert', 'FlaxBertForNextSentencePrediction'),
    ]
)

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

A_	:			Optional[Any]  	=      OrderedDict(
    [
        ('whisper', 'FlaxWhisperForAudioClassification'),
    ]
)

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



class    _a						(_BaseAutoModelClass    ):


					'''simple docstring'''






					UpperCAmelCase__:       str  							=	FLAX_MODEL_MAPPING



A_	:			Any  	=      auto_class_update(FlaxAutoModel)



class    _a						(_BaseAutoModelClass    ):


					'''simple docstring'''






					UpperCAmelCase__:       List[Any]  							=	FLAX_MODEL_FOR_PRETRAINING_MAPPING



A_	:			Union[str, Any]  	=      auto_class_update(FlaxAutoModelForPreTraining, head_doc='pretraining')



class    _a						(_BaseAutoModelClass    ):


					'''simple docstring'''






					UpperCAmelCase__:       Any  							=	FLAX_MODEL_FOR_CAUSAL_LM_MAPPING



A_	:			Union[str, Any]  	=      auto_class_update(FlaxAutoModelForCausalLM, head_doc='causal language modeling')



class    _a						(_BaseAutoModelClass    ):


					'''simple docstring'''






					UpperCAmelCase__:       Optional[Any]  							=	FLAX_MODEL_FOR_MASKED_LM_MAPPING



A_	:			Tuple  	=      auto_class_update(FlaxAutoModelForMaskedLM, head_doc='masked language modeling')



class    _a						(_BaseAutoModelClass    ):


					'''simple docstring'''






					UpperCAmelCase__:       Dict  							=	FLAX_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING



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



class    _a						(_BaseAutoModelClass    ):


					'''simple docstring'''






					UpperCAmelCase__:       Dict  							=	FLAX_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING



A_	:			Tuple  	=      auto_class_update(
    FlaxAutoModelForSequenceClassification, head_doc='sequence classification'
)



class    _a						(_BaseAutoModelClass    ):


					'''simple docstring'''






					UpperCAmelCase__:       Any  							=	FLAX_MODEL_FOR_QUESTION_ANSWERING_MAPPING



A_	:			int  	=      auto_class_update(FlaxAutoModelForQuestionAnswering, head_doc='question answering')



class    _a						(_BaseAutoModelClass    ):


					'''simple docstring'''






					UpperCAmelCase__:       Tuple  							=	FLAX_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING



A_	:			Dict  	=      auto_class_update(
    FlaxAutoModelForTokenClassification, head_doc='token classification'
)



class    _a						(_BaseAutoModelClass    ):


					'''simple docstring'''






					UpperCAmelCase__:       List[Any]  							=	FLAX_MODEL_FOR_MULTIPLE_CHOICE_MAPPING



A_	:			Union[str, Any]  	=      auto_class_update(FlaxAutoModelForMultipleChoice, head_doc='multiple choice')



class    _a						(_BaseAutoModelClass    ):


					'''simple docstring'''






					UpperCAmelCase__:       str  							=	FLAX_MODEL_FOR_NEXT_SENTENCE_PREDICTION_MAPPING



A_	:			Optional[Any]  	=      auto_class_update(
    FlaxAutoModelForNextSentencePrediction, head_doc='next sentence prediction'
)



class    _a						(_BaseAutoModelClass    ):


					'''simple docstring'''






					UpperCAmelCase__:       str  							=	FLAX_MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING



A_	:			List[Any]  	=      auto_class_update(
    FlaxAutoModelForImageClassification, head_doc='image classification'
)



class    _a						(_BaseAutoModelClass    ):


					'''simple docstring'''






					UpperCAmelCase__:       Dict  							=	FLAX_MODEL_FOR_VISION_2_SEQ_MAPPING



A_	:			Union[str, Any]  	=      auto_class_update(FlaxAutoModelForVisionaSeq, head_doc='vision-to-text modeling')



class    _a						(_BaseAutoModelClass    ):


					'''simple docstring'''






					UpperCAmelCase__:       int  							=	FLAX_MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING



A_	:			List[str]  	=      auto_class_update(
    FlaxAutoModelForSpeechSeqaSeq, head_doc='sequence-to-sequence speech-to-text modeling'
)



