




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    lowercase_		(unittest.TestCase ):



						def __init__(			self	:   List[Any]					,							snake_case__	:   Tuple					,							snake_case__	:   Optional[int]=13					,							snake_case__	:   Any=7					,							snake_case__	:   Optional[Any]=True					,							snake_case__	:   str=True					,							snake_case__	:   List[Any]=True					,							snake_case__	:   Tuple=True					,							snake_case__	:   Optional[int]=99					,							snake_case__	:   List[Any]=32					,							snake_case__	:   Tuple=5					,							snake_case__	:   List[Any]=4					,							snake_case__	:   List[Any]=37					,							snake_case__	:   List[str]="gelu"					,							snake_case__	:   Any=0.1					,							snake_case__	:   Dict=0.1					,							snake_case__	:   Dict=5_12					,							snake_case__	:   Any=16					,							snake_case__	:   Tuple=2					,							snake_case__	:   Any=0.02					,							snake_case__	:   Dict=4					,							):

								"""simple docstring"""

								SCREAMING_SNAKE_CASE_					      =							parent
								SCREAMING_SNAKE_CASE_					      =							batch_size
								SCREAMING_SNAKE_CASE_					      =							seq_length
								SCREAMING_SNAKE_CASE_					      =							is_training
								SCREAMING_SNAKE_CASE_					      =							use_attention_mask
								SCREAMING_SNAKE_CASE_					      =							use_token_type_ids
								SCREAMING_SNAKE_CASE_					      =							use_labels
								SCREAMING_SNAKE_CASE_					      =							vocab_size
								SCREAMING_SNAKE_CASE_					      =							hidden_size
								SCREAMING_SNAKE_CASE_					      =							num_hidden_layers
								SCREAMING_SNAKE_CASE_					      =							num_attention_heads
								SCREAMING_SNAKE_CASE_					      =							intermediate_size
								SCREAMING_SNAKE_CASE_					      =							hidden_act
								SCREAMING_SNAKE_CASE_					      =							hidden_dropout_prob
								SCREAMING_SNAKE_CASE_					      =							attention_probs_dropout_prob
								SCREAMING_SNAKE_CASE_					      =							max_position_embeddings
								SCREAMING_SNAKE_CASE_					      =							type_vocab_size
								SCREAMING_SNAKE_CASE_					      =							type_sequence_label_size
								SCREAMING_SNAKE_CASE_					      =							initializer_range
								SCREAMING_SNAKE_CASE_					      =							num_choices



						def 	__a					(			self	:   Any  ):

								"""simple docstring"""

								SCREAMING_SNAKE_CASE_					      =							ids_tensor([self.batch_size, self.seq_length]					,							self.vocab_size  )

								SCREAMING_SNAKE_CASE_					      =							None
								if self.use_attention_mask:
										SCREAMING_SNAKE_CASE_					      =							random_attention_mask([self.batch_size, self.seq_length]  )

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

								SCREAMING_SNAKE_CASE_					      =							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=snake_case__					,							initializer_range=self.initializer_range					,							)

								return config, input_ids, token_type_ids, attention_mask






						def 	__a					(			self	:   int  ):

								"""simple docstring"""

								SCREAMING_SNAKE_CASE_					      =							self.prepare_config_and_inputs()
								SCREAMING_SNAKE_CASE_      ,				SCREAMING_SNAKE_CASE_      ,				SCREAMING_SNAKE_CASE_      ,				SCREAMING_SNAKE_CASE_					      =							config_and_inputs
								SCREAMING_SNAKE_CASE_					      =							{'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': attention_mask}
								return config, inputs_dict



@require_flax
class    lowercase_		(SCREAMING_SNAKE_CASE__							,	unittest.TestCase ):
						lowerCAmelCase__			      =(
						    (
						        FlaxAlbertModel,
						        FlaxAlbertForPreTraining,
						        FlaxAlbertForMaskedLM,
						        FlaxAlbertForMultipleChoice,
						        FlaxAlbertForQuestionAnswering,
						        FlaxAlbertForSequenceClassification,
						        FlaxAlbertForTokenClassification,
						        FlaxAlbertForQuestionAnswering,
						    )
						    if is_flax_available()
						    else ()
						)



						def 	__a					(			self	:   Any  ):

								"""simple docstring"""

								SCREAMING_SNAKE_CASE_					      =							FlaxAlbertModelTester(self  )






						@slow
						def 	__a					(			self	:   Any  ):

								"""simple docstring"""

								for model_class_name in self.all_model_classes:
										SCREAMING_SNAKE_CASE_					      =							model_class_name.from_pretrained('albert-base-v2'  )
										SCREAMING_SNAKE_CASE_					      =							model(np.ones((1, 1)  )  )
										self.assertIsNotNone(snake_case__  )



@require_flax
class    lowercase_		(unittest.TestCase ):



						@slow
						def 	__a					(			self	:   Dict  ):

								"""simple docstring"""

								SCREAMING_SNAKE_CASE_					      =							FlaxAlbertModel.from_pretrained('albert-base-v2'  )
								SCREAMING_SNAKE_CASE_					      =							np.array([[0, 3_45, 2_32, 3_28, 7_40, 1_40, 16_95, 69, 60_78, 15_88, 2]]  )
								SCREAMING_SNAKE_CASE_					      =							np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]  )
								SCREAMING_SNAKE_CASE_					      =							model(snake_case__					,							attention_mask=snake_case__  )[0]
								SCREAMING_SNAKE_CASE_					      =							(1, 11, 7_68)
								self.assertEqual(output.shape					,							snake_case__  )
								SCREAMING_SNAKE_CASE_					      =							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]					,							snake_case__					,							atol=1e-4  )  )





from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


SCREAMING_SNAKE_CASE:     Any									= {
    '''configuration_pegasus_x''': ['''PEGASUS_X_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''PegasusXConfig'''],
}

try:
						if not is_torch_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						SCREAMING_SNAKE_CASE:     List[Any]									= [
						    '''PEGASUS_X_PRETRAINED_MODEL_ARCHIVE_LIST''',
						    '''PegasusXForConditionalGeneration''',
						    '''PegasusXModel''',
						    '''PegasusXPreTrainedModel''',
						]


if TYPE_CHECKING:
						from .configuration_pegasus_x import PEGASUS_X_PRETRAINED_CONFIG_ARCHIVE_MAP, PegasusXConfig

						try:
												if not is_torch_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_pegasus_x import (
												    PEGASUS_X_PRETRAINED_MODEL_ARCHIVE_LIST,
												    PegasusXForConditionalGeneration,
												    PegasusXModel,
												    PegasusXPreTrainedModel,
												)


else:
						import sys

						SCREAMING_SNAKE_CASE:     List[Any]									= _LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)