



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




      """simple docstring"""




      def __init__(      self				:      Any    ,			SCREAMING_SNAKE_CASE_				:      Any    ,			SCREAMING_SNAKE_CASE_				:      Union[str, Any]=1_3    ,			SCREAMING_SNAKE_CASE_				:      int=7    ,			SCREAMING_SNAKE_CASE_				:      int=True    ,			SCREAMING_SNAKE_CASE_				:      Optional[int]=True    ,			SCREAMING_SNAKE_CASE_				:      Optional[Any]=True    ,			SCREAMING_SNAKE_CASE_				:      Tuple=True    ,			SCREAMING_SNAKE_CASE_				:      str=9_9    ,			SCREAMING_SNAKE_CASE_				:      Optional[Any]=3_2    ,			SCREAMING_SNAKE_CASE_				:      Union[str, Any]=5    ,			SCREAMING_SNAKE_CASE_				:      Optional[Any]=4    ,			SCREAMING_SNAKE_CASE_				:      Optional[int]=3_7    ,			SCREAMING_SNAKE_CASE_				:      str="gelu"    ,			SCREAMING_SNAKE_CASE_				:      List[Any]=0.1    ,			SCREAMING_SNAKE_CASE_				:      int=0.1    ,			SCREAMING_SNAKE_CASE_				:      List[Any]=5_1_2    ,			SCREAMING_SNAKE_CASE_				:      Any=1_6    ,			SCREAMING_SNAKE_CASE_				:      Optional[Any]=2    ,			SCREAMING_SNAKE_CASE_				:      Optional[int]=0.02    ,			SCREAMING_SNAKE_CASE_				:      Any=4    ,			):
        lowerCAmelCase_   :				Optional[int]								= parent
        lowerCAmelCase_   :				Optional[Any]								= batch_size
        lowerCAmelCase_   :				Optional[Any]								= seq_length
        lowerCAmelCase_   :				Any								= is_training
        lowerCAmelCase_   :				List[Any]								= use_attention_mask
        lowerCAmelCase_   :				List[Any]								= use_token_type_ids
        lowerCAmelCase_   :				Any								= use_labels
        lowerCAmelCase_   :				int								= vocab_size
        lowerCAmelCase_   :				Dict								= hidden_size
        lowerCAmelCase_   :				List[Any]								= num_hidden_layers
        lowerCAmelCase_   :				List[str]								= num_attention_heads
        lowerCAmelCase_   :				Optional[Any]								= intermediate_size
        lowerCAmelCase_   :				Dict								= hidden_act
        lowerCAmelCase_   :				List[str]								= hidden_dropout_prob
        lowerCAmelCase_   :				Dict								= attention_probs_dropout_prob
        lowerCAmelCase_   :				List[Any]								= max_position_embeddings
        lowerCAmelCase_   :				Optional[Any]								= type_vocab_size
        lowerCAmelCase_   :				Optional[int]								= type_sequence_label_size
        lowerCAmelCase_   :				int								= initializer_range
        lowerCAmelCase_   :				List[Any]								= num_choices



      def      SCREAMING_SNAKE_CASE__ (      self				:      List[str]       ):
        lowerCAmelCase_   :				Tuple								= ids_tensor([self.batch_size, self.seq_length]    ,			self.vocab_size       )

        lowerCAmelCase_   :				Tuple								= None
        if self.use_attention_mask:
          lowerCAmelCase_   :				Any								= random_attention_mask([self.batch_size, self.seq_length]       )

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

        lowerCAmelCase_   :				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      SCREAMING_SNAKE_CASE__ (      self				:      List[str]       ):
        lowerCAmelCase_   :				Any								= self.prepare_config_and_inputs()
        lowerCAmelCase_   :				Any								= config_and_inputs
        lowerCAmelCase_   :				Union[str, Any]								= {'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': attention_mask}
        return config, inputs_dict







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




      """simple docstring"""




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



      def      SCREAMING_SNAKE_CASE__ (      self				:      Union[str, Any]       ):
        lowerCAmelCase_   :				Dict								= FlaxAlbertModelTester(self       )





      @slow
      def      SCREAMING_SNAKE_CASE__ (      self				:      List[str]       ):
        for model_class_name in self.all_model_classes:
          lowerCAmelCase_   :				Optional[int]								= model_class_name.from_pretrained('albert-base-v2'       )
          lowerCAmelCase_   :				Optional[int]								= model(np.ones((1, 1)       )       )
          self.assertIsNotNone(_a       )







@require_flax
class       UpperCamelCase__   (	unittest.TestCase  ):




      """simple docstring"""




      @slow
      def      SCREAMING_SNAKE_CASE__ (      self				:      Any       ):
        lowerCAmelCase_   :				Dict								= FlaxAlbertModel.from_pretrained('albert-base-v2'       )
        lowerCAmelCase_   :				Optional[int]								= 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]]       )
        lowerCAmelCase_   :				Any								= np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]       )
        lowerCAmelCase_   :				Dict								= model(_a    ,			attention_mask=_a       )[0]
        lowerCAmelCase_   :				Union[str, Any]								= (1, 1_1, 7_6_8)
        self.assertEqual(output.shape    ,			_a       )
        lowerCAmelCase_   :				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]    ,			_a    ,			atol=1E-4       )       )










'''simple docstring'''





from typing import TYPE_CHECKING

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


_UpperCamelCase											=						{'configuration_speech_encoder_decoder': ['SpeechEncoderDecoderConfig']}

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

try:
   if not is_flax_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   _UpperCamelCase											=						['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

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






