



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		:		List[str]      ,				__lowercase		:		Optional[int]      ,				__lowercase		:		int=13      ,				__lowercase		:		Any=7      ,				__lowercase		:		Optional[Any]=True      ,				__lowercase		:		List[Any]=True      ,				__lowercase		:		List[str]=True      ,				__lowercase		:		Dict=True      ,				__lowercase		:		str=99      ,				__lowercase		:		List[str]=32      ,				__lowercase		:		Any=5      ,				__lowercase		:		List[Any]=4      ,				__lowercase		:		str=37      ,				__lowercase		:		List[Any]="gelu"      ,				__lowercase		:		Any=0.1      ,				__lowercase		:		Tuple=0.1      ,				__lowercase		:		Tuple=512      ,				__lowercase		:		Any=16      ,				__lowercase		:		str=2      ,				__lowercase		:		Union[str, Any]=0.02      ,				__lowercase		:		str=4      ,				):

           '''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 						UpperCamelCase_					(   self		:		Union[str, 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 						UpperCamelCase_					(   self		:		int	):

           '''simple docstring'''
           __a       =		self.prepare_config_and_inputs()
           __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    SCREAMING_SNAKE_CASE			(							__lowercase			,      unittest.TestCase					):
       __lowerCamelCase	:	Tuple    =(
           (
               FlaxAlbertModel,
               FlaxAlbertForPreTraining,
               FlaxAlbertForMaskedLM,
               FlaxAlbertForMultipleChoice,
               FlaxAlbertForQuestionAnswering,
               FlaxAlbertForSequenceClassification,
               FlaxAlbertForTokenClassification,
               FlaxAlbertForQuestionAnswering,
           )
           if is_flax_available()
           else ()
       )





       def 						UpperCamelCase_					(   self		:		Optional[int]	):

           '''simple docstring'''
           __a       =		FlaxAlbertModelTester(self	)







       @slow
       def 						UpperCamelCase_					(   self		:		int	):

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





       @slow
       def 						UpperCamelCase_					(   self		:		Any	):

           '''simple docstring'''
           __a       =		FlaxAlbertModel.from_pretrained("""albert-base-v2"""	)
           __a       =		np.array([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 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 typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_flax_available, is_torch_available


lowerCamelCase__			= {"""configuration_speech_encoder_decoder""": ["""SpeechEncoderDecoderConfig"""]}

try:
						if not is_torch_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						lowerCamelCase__			= ["""SpeechEncoderDecoderModel"""]

try:
						if not is_flax_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						lowerCamelCase__			= ["""FlaxSpeechEncoderDecoderModel"""]

if TYPE_CHECKING:
						from .configuration_speech_encoder_decoder import SpeechEncoderDecoderConfig

						try:
												if not is_torch_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_speech_encoder_decoder import SpeechEncoderDecoderModel

						try:
												if not is_flax_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_flax_speech_encoder_decoder import FlaxSpeechEncoderDecoderModel

else:
						import sys

						lowerCamelCase__			= _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)
