



'''simple docstring'''




import unittest

import numpy as np

from transformers import AlbertConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
      import jax.numpy as jnp

      from transformers.models.albert.modeling_flax_albert import (
          FlaxAlbertForMaskedLM,
          FlaxAlbertForMultipleChoice,
          FlaxAlbertForPreTraining,
          FlaxAlbertForQuestionAnswering,
          FlaxAlbertForSequenceClassification,
          FlaxAlbertForTokenClassification,
          FlaxAlbertModel,
      )






class 		_UpperCamelCase						(							unittest.TestCase					):


   '''simple docstring'''

   def __init__(							self  :				List[Any]							,  lowerCAmelCase__  :				str							,  lowerCAmelCase__  :				int=1_3							,  lowerCAmelCase__  :				Union[str, Any]=7							,  lowerCAmelCase__  :				Tuple=True							,  lowerCAmelCase__  :				List[str]=True							,  lowerCAmelCase__  :				List[str]=True							,  lowerCAmelCase__  :				Union[str, Any]=True							,  lowerCAmelCase__  :				List[Any]=9_9							,  lowerCAmelCase__  :				Any=3_2							,  lowerCAmelCase__  :				str=5							,  lowerCAmelCase__  :				Any=4							,  lowerCAmelCase__  :				Tuple=3_7							,  lowerCAmelCase__  :				Dict="gelu"							,  lowerCAmelCase__  :				Union[str, Any]=0.1							,  lowerCAmelCase__  :				List[Any]=0.1							,  lowerCAmelCase__  :				List[Any]=5_1_2							,  lowerCAmelCase__  :				Any=1_6							,  lowerCAmelCase__  :				Any=2							,  lowerCAmelCase__  :				Tuple=0.02							,  lowerCAmelCase__  :				Optional[int]=4							,  ):





     """simple docstring"""






     __SCREAMING_SNAKE_CASE :	Optional[int]					 =	parent
     __SCREAMING_SNAKE_CASE :	str					 =	batch_size
     __SCREAMING_SNAKE_CASE :	str					 =	seq_length
     __SCREAMING_SNAKE_CASE :	str					 =	is_training
     __SCREAMING_SNAKE_CASE :	Optional[int]					 =	use_attention_mask
     __SCREAMING_SNAKE_CASE :	Dict					 =	use_token_type_ids
     __SCREAMING_SNAKE_CASE :	Any					 =	use_labels
     __SCREAMING_SNAKE_CASE :	int					 =	vocab_size
     __SCREAMING_SNAKE_CASE :	str					 =	hidden_size
     __SCREAMING_SNAKE_CASE :	List[str]					 =	num_hidden_layers
     __SCREAMING_SNAKE_CASE :	int					 =	num_attention_heads
     __SCREAMING_SNAKE_CASE :	Optional[int]					 =	intermediate_size
     __SCREAMING_SNAKE_CASE :	Tuple					 =	hidden_act
     __SCREAMING_SNAKE_CASE :	List[str]					 =	hidden_dropout_prob
     __SCREAMING_SNAKE_CASE :	Optional[int]					 =	attention_probs_dropout_prob
     __SCREAMING_SNAKE_CASE :	Optional[int]					 =	max_position_embeddings
     __SCREAMING_SNAKE_CASE :	Dict					 =	type_vocab_size
     __SCREAMING_SNAKE_CASE :	Optional[int]					 =	type_sequence_label_size
     __SCREAMING_SNAKE_CASE :	List[str]					 =	initializer_range
     __SCREAMING_SNAKE_CASE :	Union[str, Any]					 =	num_choices

   def UpperCamelCase__      (							self  :				Any		):





     """simple docstring"""






     __SCREAMING_SNAKE_CASE :	Any					 =	ids_tensor([self.batch_size, self.seq_length]							,  self.vocab_size		)

     __SCREAMING_SNAKE_CASE :	List[str]					 =	None
     if self.use_attention_mask:
       __SCREAMING_SNAKE_CASE :	Dict					 =	random_attention_mask([self.batch_size, self.seq_length]		)

     __SCREAMING_SNAKE_CASE :	List[Any]					 =	None
     if self.use_token_type_ids:
       __SCREAMING_SNAKE_CASE :	Optional[Any]					 =	ids_tensor([self.batch_size, self.seq_length]							,  self.type_vocab_size		)

     __SCREAMING_SNAKE_CASE :	str					 =	AlbertConfig(
         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=UpperCamelCase__							,  initializer_range=self.initializer_range							,  )

     return config, input_ids, token_type_ids, attention_mask



   def UpperCamelCase__      (							self  :				Dict		):





     """simple docstring"""






     __SCREAMING_SNAKE_CASE :	str					 =	self.prepare_config_and_inputs()
     __SCREAMING_SNAKE_CASE :	Tuple					 =	config_and_inputs
     __SCREAMING_SNAKE_CASE :	Any					 =	{'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': attention_mask}
     return config, inputs_dict






@require_flax
class 		_UpperCamelCase						(							SCREAMING_SNAKE_CASE_					,      unittest.TestCase					):


   '''simple docstring'''

   _A   :	Optional[Any]    					=	(
       (
           FlaxAlbertModel,
           FlaxAlbertForPreTraining,
           FlaxAlbertForMaskedLM,
           FlaxAlbertForMultipleChoice,
           FlaxAlbertForQuestionAnswering,
           FlaxAlbertForSequenceClassification,
           FlaxAlbertForTokenClassification,
           FlaxAlbertForQuestionAnswering,
       )
       if is_flax_available()
       else ()
   )

   def UpperCamelCase__      (							self  :				Any		):





     """simple docstring"""






     __SCREAMING_SNAKE_CASE :	str					 =	FlaxAlbertModelTester(self		)



   @slow
   def UpperCamelCase__      (							self  :				List[str]		):





     """simple docstring"""






     for model_class_name in self.all_model_classes:
       __SCREAMING_SNAKE_CASE :	int					 =	model_class_name.from_pretrained("""albert-base-v2"""		)
       __SCREAMING_SNAKE_CASE :	Optional[Any]					 =	model(np.ones((1, 1)		)		)
       self.assertIsNotNone(UpperCamelCase__		)






@require_flax
class 		_UpperCamelCase						(							unittest.TestCase					):


   '''simple docstring'''

   @slow
   def UpperCamelCase__      (							self  :				Optional[int]		):





     """simple docstring"""






     __SCREAMING_SNAKE_CASE :	Union[str, Any]					 =	FlaxAlbertModel.from_pretrained("""albert-base-v2"""		)
     __SCREAMING_SNAKE_CASE :	Any					 =	np.array([[0, 3_4_5, 2_3_2, 3_2_8, 7_4_0, 1_4_0, 1_6_9_5, 6_9, 6_0_7_8, 1_5_8_8, 2]]		)
     __SCREAMING_SNAKE_CASE :	Dict					 =	np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]		)
     __SCREAMING_SNAKE_CASE :	str					 =	model(UpperCamelCase__							,  attention_mask=UpperCamelCase__		)[0]
     __SCREAMING_SNAKE_CASE :	Dict					 =	(1, 1_1, 7_6_8)
     self.assertEqual(output.shape							,  UpperCamelCase__		)
     __SCREAMING_SNAKE_CASE :	str					 =	np.array(
         [[[-0.65_13, 1.50_35, -0.27_66], [-0.65_15, 1.50_46, -0.27_80], [-0.65_12, 1.50_49, -0.27_84]]]		)

     self.assertTrue(jnp.allclose(output[:, 1:4, 1:4]							,  UpperCamelCase__							,  atol=1E-4		)		)




'''simple docstring'''




from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxSeqaSeqConfigWithPast
from ...utils import logging


UpperCamelCase__ :			Tuple												=				logging.get_logger(__name__)

UpperCamelCase__ :			List[str]												=				{
    '''t5-small''': '''https://huggingface.co/t5-small/resolve/main/config.json''',
    '''t5-base''': '''https://huggingface.co/t5-base/resolve/main/config.json''',
    '''t5-large''': '''https://huggingface.co/t5-large/resolve/main/config.json''',
    '''t5-3b''': '''https://huggingface.co/t5-3b/resolve/main/config.json''',
    '''t5-11b''': '''https://huggingface.co/t5-11b/resolve/main/config.json''',
}






class 		_UpperCamelCase						(							lowerCamelCase__					):


			'''simple docstring'''

			_A   :	Union[str, Any]    					=	'''t5'''
			_A   :	Optional[Any]    					=	['''past_key_values''']
			_A   :	Any    					=	{'''hidden_size''': '''d_model''', '''num_attention_heads''': '''num_heads''', '''num_hidden_layers''': '''num_layers'''}

			def __init__(							self  :				Union[str, Any]							,  lowerCAmelCase__  :				List[str]=3_2_1_2_8							,  lowerCAmelCase__  :				Dict=5_1_2							,  lowerCAmelCase__  :				str=6_4							,  lowerCAmelCase__  :				Tuple=2_0_4_8							,  lowerCAmelCase__  :				Optional[Any]=6							,  lowerCAmelCase__  :				Any=None							,  lowerCAmelCase__  :				Any=8							,  lowerCAmelCase__  :				Dict=3_2							,  lowerCAmelCase__  :				Dict=1_2_8							,  lowerCAmelCase__  :				Any=0.1							,  lowerCAmelCase__  :				str=1E-6							,  lowerCAmelCase__  :				str=1.0							,  lowerCAmelCase__  :				Optional[Any]="relu"							,  lowerCAmelCase__  :				List[str]=True							,  lowerCAmelCase__  :				Union[str, Any]=True							,  lowerCAmelCase__  :				Optional[Any]=0							,  lowerCAmelCase__  :				List[Any]=1							,  **lowerCAmelCase__  :				Dict							,  ):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE :	List[str]					 =	vocab_size
					__SCREAMING_SNAKE_CASE :	str					 =	d_model
					__SCREAMING_SNAKE_CASE :	str					 =	d_kv
					__SCREAMING_SNAKE_CASE :	Optional[Any]					 =	d_ff
					__SCREAMING_SNAKE_CASE :	Optional[Any]					 =	num_layers
					__SCREAMING_SNAKE_CASE :	List[str]					 =	(
					    num_decoder_layers if num_decoder_layers is not None else self.num_layers
					)  # default = symmetry
					__SCREAMING_SNAKE_CASE :	Optional[int]					 =	num_heads
					__SCREAMING_SNAKE_CASE :	Any					 =	relative_attention_num_buckets
					__SCREAMING_SNAKE_CASE :	int					 =	relative_attention_max_distance
					__SCREAMING_SNAKE_CASE :	Optional[int]					 =	dropout_rate
					__SCREAMING_SNAKE_CASE :	Optional[int]					 =	layer_norm_epsilon
					__SCREAMING_SNAKE_CASE :	Optional[int]					 =	initializer_factor
					__SCREAMING_SNAKE_CASE :	Any					 =	feed_forward_proj
					__SCREAMING_SNAKE_CASE :	str					 =	use_cache

					__SCREAMING_SNAKE_CASE :	List[str]					 =	self.feed_forward_proj.split("""-"""		)
					__SCREAMING_SNAKE_CASE :	Tuple					 =	act_info[-1]
					__SCREAMING_SNAKE_CASE :	int					 =	act_info[0] == """gated"""

					if len(lowerCAmelCase__		) > 1 and act_info[0] != "gated" or len(lowerCAmelCase__		) > 2:
							raise ValueError(
							    F"`feed_forward_proj`: {feed_forward_proj} is not a valid activation function of the dense layer."
							    """Please make sure `feed_forward_proj` is of the format `gated-{ACT_FN}` or `{ACT_FN}`, e.g. """
							    """'gated-gelu' or 'relu'"""		)

					# for backwards compatibility
					if feed_forward_proj == "gated-gelu":
							__SCREAMING_SNAKE_CASE :	List[Any]					 =	"""gelu_new"""

					super().__init__(
					    pad_token_id=lowerCAmelCase__							,  eos_token_id=lowerCAmelCase__							,  is_encoder_decoder=lowerCAmelCase__							,  **lowerCAmelCase__							,  )






class 		_UpperCamelCase						(							lowerCamelCase__					):


			'''simple docstring'''

			@property
			def UpperCamelCase__      (							self  :				Tuple		):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE :	List[Any]					 =	{
					    """input_ids""": {0: """batch""", 1: """encoder_sequence"""},
					    """attention_mask""": {0: """batch""", 1: """encoder_sequence"""},
					}
					if self.use_past:
							__SCREAMING_SNAKE_CASE :	Optional[Any]					 =	"""past_encoder_sequence + sequence"""
							__SCREAMING_SNAKE_CASE :	Optional[Any]					 =	{0: """batch"""}
							__SCREAMING_SNAKE_CASE :	Optional[int]					 =	{0: """batch""", 1: """past_decoder_sequence + sequence"""}
					else:
							__SCREAMING_SNAKE_CASE :	List[Any]					 =	{0: """batch""", 1: """decoder_sequence"""}
							__SCREAMING_SNAKE_CASE :	Optional[Any]					 =	{0: """batch""", 1: """decoder_sequence"""}

					if self.use_past:
							self.fill_with_past_key_values_(lowerCAmelCase__							,  direction="""inputs"""		)

					return common_inputs



			@property
			def UpperCamelCase__      (							self  :				List[str]		):





					"""simple docstring"""






					return 1_3