




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



	'''simple docstring'''


	def __init__(    self :  List[str]						,							A_ :  Any						,							A_ :  Dict=13						,							A_ :  Any=7						,							A_ :  str=True						,							A_ :  List[str]=True						,							A_ :  List[str]=True						,							A_ :  Union[str, Any]=True						,							A_ :  List[str]=99						,							A_ :  Dict=32						,							A_ :  int=5						,							A_ :  Any=4						,							A_ :  int=37						,							A_ :  Optional[int]="gelu"						,							A_ :  Dict=0.1						,							A_ :  Dict=0.1						,							A_ :  Any=512						,							A_ :  int=16						,							A_ :  Dict=2						,							A_ :  Any=0.02						,							A_ :  Any=4						,							)       ->     Union[str, Any]:

							"""simple docstring"""






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

	def     a__			(    self :  Dict  )       ->     Optional[Any]:

							"""simple docstring"""






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

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

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

							lowerCamelCase_							=   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=__lowerCamelCase						,							initializer_range=self.initializer_range						,							)

							return config, input_ids, token_type_ids, attention_mask





	def     a__			(    self :  str  )       ->     Dict:

							"""simple docstring"""






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







@require_flax
class   A(						_a							,				unittest.TestCase       ):



	'''simple docstring'''


	UpperCamelCase						       =						(
	    (
	        FlaxAlbertModel,
	        FlaxAlbertForPreTraining,
	        FlaxAlbertForMaskedLM,
	        FlaxAlbertForMultipleChoice,
	        FlaxAlbertForQuestionAnswering,
	        FlaxAlbertForSequenceClassification,
	        FlaxAlbertForTokenClassification,
	        FlaxAlbertForQuestionAnswering,
	    )
	    if is_flax_available()
	    else ()
	)

	def     a__			(    self :  Union[str, Any]  )       ->     Optional[Any]:

							"""simple docstring"""






							lowerCamelCase_							=   FlaxAlbertModelTester(self  )





	@slow
	def     a__			(    self :  List[str]  )       ->     int:

							"""simple docstring"""






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







@require_flax
class   A(						unittest.TestCase       ):



	'''simple docstring'''


	@slow
	def     a__			(    self :  Union[str, Any]  )       ->     str:

							"""simple docstring"""






							lowerCamelCase_							=   FlaxAlbertModel.from_pretrained('albert-base-v2'  )
							lowerCamelCase_							=   np.array([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]  )
							lowerCamelCase_							=   np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]  )
							lowerCamelCase_							=   model(__lowerCamelCase						,							attention_mask=__lowerCamelCase  )[0]
							lowerCamelCase_							=   (1, 11, 768)
							self.assertEqual(output.shape						,							__lowerCamelCase  )
							lowerCamelCase_							=   np.array(
							    [[[-0.6513, 1.5035, -0.2766], [-0.6515, 1.5046, -0.2780], [-0.6512, 1.5049, -0.2784]]]  )

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









from ..utils import DummyObject, requires_backends







class   A(						metaclass=UpperCamelCase       ):



 '''simple docstring'''


 UpperCamelCase						       =						['''keras_nlp''']

 def __init__(    self :  Optional[int]						,							*A_ :  Any						,							**A_ :  Dict  )       ->     Optional[int]:

       """simple docstring"""






       requires_backends(self						,							['keras_nlp']  )



