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



      """simple docstring"""







      def __init__(   self		:							Tuple					,							_snake_case		:							str					,							_snake_case		:							Union[str, Any]=13					,							_snake_case		:							Tuple=7					,							_snake_case		:							int=True					,							_snake_case		:							Union[str, Any]=True					,							_snake_case		:							str=True					,							_snake_case		:							Union[str, Any]=True					,							_snake_case		:							List[str]=99					,							_snake_case		:							Dict=32					,							_snake_case		:							int=5					,							_snake_case		:							List[Any]=4					,							_snake_case		:							List[str]=37					,							_snake_case		:							str="gelu"					,							_snake_case		:							Dict=0.1					,							_snake_case		:							Any=0.1					,							_snake_case		:							Any=512					,							_snake_case		:							Optional[int]=16					,							_snake_case		:							Any=2					,							_snake_case		:							Dict=0.02					,							_snake_case		:							int=4					,							):
          __lowercase		:      Tuple						     =							parent
          __lowercase		:      Optional[int]						     =							batch_size
          __lowercase		:      Dict						     =							seq_length
          __lowercase		:      Any						     =							is_training
          __lowercase		:      int						     =							use_attention_mask
          __lowercase		:      str						     =							use_token_type_ids
          __lowercase		:      str						     =							use_labels
          __lowercase		:      str						     =							vocab_size
          __lowercase		:      int						     =							hidden_size
          __lowercase		:      Tuple						     =							num_hidden_layers
          __lowercase		:      Optional[Any]						     =							num_attention_heads
          __lowercase		:      List[Any]						     =							intermediate_size
          __lowercase		:      Dict						     =							hidden_act
          __lowercase		:      Optional[int]						     =							hidden_dropout_prob
          __lowercase		:      Optional[Any]						     =							attention_probs_dropout_prob
          __lowercase		:      Dict						     =							max_position_embeddings
          __lowercase		:      Tuple						     =							type_vocab_size
          __lowercase		:      List[Any]						     =							type_sequence_label_size
          __lowercase		:      Dict						     =							initializer_range
          __lowercase		:      List[str]						     =							num_choices




      def 		snake_case_      (   self		:							str					):
          __lowercase		:      Dict						     =							ids_tensor([self.batch_size, self.seq_length]					,							self.vocab_size					)

          __lowercase		:      Optional[Any]						     =							None
          if self.use_attention_mask:
              __lowercase		:      int						     =							random_attention_mask([self.batch_size, self.seq_length]					)

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

          __lowercase		:      List[Any]						     =							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=__SCREAMING_SNAKE_CASE					,							initializer_range=self.initializer_range					,							)

          return config, input_ids, token_type_ids, attention_mask




      def 		snake_case_      (   self		:							int					):
          __lowercase		:      Any						     =							self.prepare_config_and_inputs()
          __lowercase		:      Union[str, Any]						     =							config_and_inputs
          __lowercase		:      Optional[int]						     =							{'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': attention_mask}
          return config, inputs_dict






@require_flax
class     __lowerCAmelCase						(						SCREAMING_SNAKE_CASE_			,			unittest.TestCase   ):



      """simple docstring"""







      A__						:      List[str]   =     (
          (
              FlaxAlbertModel,
              FlaxAlbertForPreTraining,
              FlaxAlbertForMaskedLM,
              FlaxAlbertForMultipleChoice,
              FlaxAlbertForQuestionAnswering,
              FlaxAlbertForSequenceClassification,
              FlaxAlbertForTokenClassification,
              FlaxAlbertForQuestionAnswering,
          )
          if is_flax_available()
          else ()
      )




      def 		snake_case_      (   self		:							Union[str, Any]					):
          __lowercase		:      int						     =							FlaxAlbertModelTester(self					)




      @slow
      def 		snake_case_      (   self		:							List[Any]					):
          for model_class_name in self.all_model_classes:
              __lowercase		:      Tuple						     =							model_class_name.from_pretrained('''albert-base-v2'''					)
              __lowercase		:      int						     =							model(np.ones((1, 1)					)					)
              self.assertIsNotNone(__SCREAMING_SNAKE_CASE					)






@require_flax
class     __lowerCAmelCase						(						unittest.TestCase   ):



      """simple docstring"""







      @slow
      def 		snake_case_      (   self		:							Tuple					):
          __lowercase		:      Tuple						     =							FlaxAlbertModel.from_pretrained('''albert-base-v2'''					)
          __lowercase		:      str						     =							np.array([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]					)
          __lowercase		:      Union[str, Any]						     =							np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]					)
          __lowercase		:      Tuple						     =							model(__SCREAMING_SNAKE_CASE					,							attention_mask=__SCREAMING_SNAKE_CASE					)[0]
          __lowercase		:      int						     =							(1, 11, 768)
          self.assertEqual(output.shape					,							__SCREAMING_SNAKE_CASE					)
          __lowercase		:      Union[str, Any]						     =							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]					,							__SCREAMING_SNAKE_CASE					,							atol=1E-4					)					)










'''simple docstring'''







from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_tf_available,
    is_torch_available,
    is_vision_available,
)


__lowerCAmelCase									=   {'configuration_deit': ['DEIT_PRETRAINED_CONFIG_ARCHIVE_MAP', 'DeiTConfig', 'DeiTOnnxConfig']}

try:
		if not is_vision_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		__lowerCAmelCase									=   ['DeiTFeatureExtractor']
		__lowerCAmelCase									=   ['DeiTImageProcessor']

try:
		if not is_torch_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		__lowerCAmelCase									=   [
		    'DEIT_PRETRAINED_MODEL_ARCHIVE_LIST',
		    'DeiTForImageClassification',
		    'DeiTForImageClassificationWithTeacher',
		    'DeiTForMaskedImageModeling',
		    'DeiTModel',
		    'DeiTPreTrainedModel',
		]

try:
		if not is_tf_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		__lowerCAmelCase									=   [
		    'TF_DEIT_PRETRAINED_MODEL_ARCHIVE_LIST',
		    'TFDeiTForImageClassification',
		    'TFDeiTForImageClassificationWithTeacher',
		    'TFDeiTForMaskedImageModeling',
		    'TFDeiTModel',
		    'TFDeiTPreTrainedModel',
		]


if TYPE_CHECKING:
		from .configuration_deit import DEIT_PRETRAINED_CONFIG_ARCHIVE_MAP, DeiTConfig, DeiTOnnxConfig

		try:
				if not is_vision_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .feature_extraction_deit import DeiTFeatureExtractor
				from .image_processing_deit import DeiTImageProcessor

		try:
				if not is_torch_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_deit import (
				    DEIT_PRETRAINED_MODEL_ARCHIVE_LIST,
				    DeiTForImageClassification,
				    DeiTForImageClassificationWithTeacher,
				    DeiTForMaskedImageModeling,
				    DeiTModel,
				    DeiTPreTrainedModel,
				)

		try:
				if not is_tf_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_tf_deit import (
				    TF_DEIT_PRETRAINED_MODEL_ARCHIVE_LIST,
				    TFDeiTForImageClassification,
				    TFDeiTForImageClassificationWithTeacher,
				    TFDeiTForMaskedImageModeling,
				    TFDeiTModel,
				    TFDeiTPreTrainedModel,
				)


else:
		import sys

		__lowerCAmelCase									=   _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)




