


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

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







      """simple docstring"""







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

    def 					_UpperCAmelCase	(				self			:				int):







      """simple docstring"""







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

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

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

      lowercase_								= 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=_A ,  initializer_range=self.initializer_range ,  )

      return config, input_ids, token_type_ids, attention_mask




    def 					_UpperCAmelCase	(				self			:				Union[str, Any]):







      """simple docstring"""







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







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

    def 					_UpperCAmelCase	(				self			:				str):







      """simple docstring"""







      lowercase_								= FlaxAlbertModelTester(self)




    @slow
    def 					_UpperCAmelCase	(				self			:				Dict):







      """simple docstring"""







      for model_class_name in self.all_model_classes:
        lowercase_								= model_class_name.from_pretrained("""albert-base-v2""")
        lowercase_								= model(np.ones((1, 1)))
        self.assertIsNotNone(_A)







@require_flax
class    SCREAMING_SNAKE_CASE__ (							unittest.TestCase   ):

    @slow
    def 					_UpperCAmelCase	(				self			:				Any):







      """simple docstring"""







      lowercase_								= FlaxAlbertModel.from_pretrained("""albert-base-v2""")
      lowercase_								= 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]])
      lowercase_								= np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]])
      lowercase_								= model(_A ,  attention_mask=_A)[0]
      lowercase_								= (1, 1_1, 7_6_8)
      self.assertEqual(output.shape ,  _A)
      lowercase_								= np.array(
          [[[-0.6_513, 1.5_035, -0.2_766], [-0.6_515, 1.5_046, -0.2_780], [-0.6_512, 1.5_049, -0.2_784]]])

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






"""simple docstring"""
from ..utils import DummyObject, requires_backends


class    SCREAMING_SNAKE_CASE__ (							metaclass=__UpperCAmelCase   ):
    lowercase__      =   ["keras_nlp"]

    def __init__(				self			:				Dict ,  *lowerCAmelCase_			:				List[str] ,  **lowerCAmelCase_			:				Tuple):







      """simple docstring"""







      requires_backends(self ,  ["""keras_nlp"""])

