






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






      '''simple docstring'''




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




       """simple docstring"""


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







      def  __lowerCAmelCase   ( self :		Any    )       ->List[Any]:




       """simple docstring"""


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

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

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

       a						    = 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=__lowercase					,   initializer_range=self.initializer_range					,   )

       return config, input_ids, token_type_ids, attention_mask







      def  __lowerCAmelCase   ( self :		List[Any]    )       ->Tuple:




       """simple docstring"""


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






@require_flax
class 						lowercase_   (      UpperCAmelCase__					,    unittest.TestCase  ):






      '''simple docstring'''




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







      def  __lowerCAmelCase   ( self :		str    )       ->int:




       """simple docstring"""


       a						    = FlaxAlbertModelTester(self    )







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




       """simple docstring"""


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






@require_flax
class 						lowercase_   (      unittest.TestCase  ):






      '''simple docstring'''




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




       """simple docstring"""


       a						    = FlaxAlbertModel.from_pretrained('''albert-base-v2'''    )
       a						    = np.array([[0, 345, 232, 328, 740, 140, 1_695, 69, 6_078, 1_588, 2]]    )
       a						    = np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]    )
       a						    = model(__lowercase					,   attention_mask=__lowercase    )[0]
       a						    = (1, 11, 768)
       self.assertEqual(output.shape					,   __lowercase    )
       a						    = 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]					,   __lowercase					,   atol=1e-4    )    )








from abc import ABC, abstractmethod
from argparse import ArgumentParser
class 						UpperCAmelCase  (							UpperCAmelCase__     ):







						'''simple docstring'''





						@staticmethod
						@abstractmethod
						def 			snake_case__		(		__lowercase				:	ArgumentParser		):






							"""simple docstring"""






							raise NotImplementedError()






						@abstractmethod
						def 			snake_case__		(		self				:	Optional[Any]		):






							"""simple docstring"""






							raise NotImplementedError()
