





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





      def __init__(					self			,				A_			,				A_=13			,				A_=7			,				A_=True			,				A_=True			,				A_=True			,				A_=True			,				A_=99			,				A_=32			,				A_=5			,				A_=4			,				A_=37			,				A_="gelu"			,				A_=0.1			,				A_=0.1			,				A_=512			,				A_=16			,				A_=2			,				A_=0.02			,				A_=4			,				):
       '''simple docstring'''




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





      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       UpperCamelCase      :   Tuple			      =  ids_tensor([self.batch_size, self.seq_length]			,				self.vocab_size  )

       UpperCamelCase      :   Optional[int]			      =  None
       if self.use_attention_mask:
        UpperCamelCase      :   List[str]			      =  random_attention_mask([self.batch_size, self.seq_length]  )

       UpperCamelCase      :   Union[str, Any]			      =  None
       if self.use_token_type_ids:
        UpperCamelCase      :   List[Any]			      =  ids_tensor([self.batch_size, self.seq_length]			,				self.type_vocab_size  )

       UpperCamelCase      :   List[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=A_			,				initializer_range=self.initializer_range			,				)

       return config, input_ids, token_type_ids, attention_mask





      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       UpperCamelCase      :   List[str]			      =  self.prepare_config_and_inputs()
       UpperCamelCase      ,       UpperCamelCase      ,       UpperCamelCase      ,       UpperCamelCase      :   Optional[Any]			      =  config_and_inputs
       UpperCamelCase      :   Optional[Any]			      =  {"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": attention_mask}
       return config, inputs_dict



@require_flax
class    A__  (			__snake_case			,  unittest.TestCase     ):
      _UpperCAmelCase      :int		       =  (
          (
              FlaxAlbertModel,
              FlaxAlbertForPreTraining,
              FlaxAlbertForMaskedLM,
              FlaxAlbertForMultipleChoice,
              FlaxAlbertForQuestionAnswering,
              FlaxAlbertForSequenceClassification,
              FlaxAlbertForTokenClassification,
              FlaxAlbertForQuestionAnswering,
          )
          if is_flax_available()
          else ()
      )





      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       UpperCamelCase      :   Tuple			      =  FlaxAlbertModelTester(self  )





      @slow
      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       for model_class_name in self.all_model_classes:
        UpperCamelCase      :   Dict			      =  model_class_name.from_pretrained("albert-base-v2"  )
        UpperCamelCase      :   Any			      =  model(np.ones((1, 1)  )  )
        self.assertIsNotNone(A_  )



@require_flax
class    A__  (			unittest.TestCase     ):





      @slow
      def 			__UpperCamelCase(					self  ):
       '''simple docstring'''




       UpperCamelCase      :   int			      =  FlaxAlbertModel.from_pretrained("albert-base-v2"  )
       UpperCamelCase      :   Optional[Any]			      =  np.array([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]  )
       UpperCamelCase      :   Dict			      =  np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]  )
       UpperCamelCase      :   Optional[int]			      =  model(A_			,				attention_mask=A_  )[0]
       UpperCamelCase      :   Optional[Any]			      =  (1, 11, 768)
       self.assertEqual(output.shape			,				A_  )
       UpperCamelCase      :   Dict			      =  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]			,				A_			,				atol=1e-4  )  )







from typing import TYPE_CHECKING

from ...utils import _LazyModule


__lowerCamelCase			:	str					    =	{"""tokenization_byt5""": ["""ByT5Tokenizer"""]}


if TYPE_CHECKING:
    from .tokenization_byta import ByTaTokenizer
else:
    import sys

    __lowerCamelCase			:	List[str]					    =	_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)
