


import unittest

import numpy as np

from transformers import DistilBertConfig, 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.distilbert.modeling_flax_distilbert import (
      FlaxDistilBertForMaskedLM,
      FlaxDistilBertForMultipleChoice,
      FlaxDistilBertForQuestionAnswering,
      FlaxDistilBertForSequenceClassification,
      FlaxDistilBertForTokenClassification,
      FlaxDistilBertModel,
  )
class 			SCREAMING_SNAKE_CASE__  (    unittest.TestCase      ):
      def __init__(	self    :	Union[str, Any]					,							SCREAMING_SNAKE_CASE__    :	Optional[int]					,							SCREAMING_SNAKE_CASE__    :	Optional[Any]=1_3					,							SCREAMING_SNAKE_CASE__    :	Tuple=7					,							SCREAMING_SNAKE_CASE__    :	Dict=True					,							SCREAMING_SNAKE_CASE__    :	str=True					,							SCREAMING_SNAKE_CASE__    :	Dict=True					,							SCREAMING_SNAKE_CASE__    :	Any=True					,							SCREAMING_SNAKE_CASE__    :	Dict=9_9					,							SCREAMING_SNAKE_CASE__    :	List[str]=3_2					,							SCREAMING_SNAKE_CASE__    :	Any=5					,							SCREAMING_SNAKE_CASE__    :	Optional[int]=4					,							SCREAMING_SNAKE_CASE__    :	Dict=3_7					,							SCREAMING_SNAKE_CASE__    :	Tuple="gelu"					,							SCREAMING_SNAKE_CASE__    :	Dict=0.1					,							SCREAMING_SNAKE_CASE__    :	Dict=0.1					,							SCREAMING_SNAKE_CASE__    :	str=5_1_2					,							SCREAMING_SNAKE_CASE__    :	Optional[Any]=1_6					,							SCREAMING_SNAKE_CASE__    :	Dict=2					,							SCREAMING_SNAKE_CASE__    :	List[str]=0.02					,							SCREAMING_SNAKE_CASE__    :	str=4					,							)		->      Dict:
             a_ :		Optional[int]							=  parent
             a_ :		Dict							=  batch_size
             a_ :		Dict							=  seq_length
             a_ :		Optional[Any]							=  is_training
             a_ :		Any							=  use_attention_mask
             a_ :		int							=  use_token_type_ids
             a_ :		List[Any]							=  use_labels
             a_ :		int							=  vocab_size
             a_ :		Optional[Any]							=  hidden_size
             a_ :		int							=  num_hidden_layers
             a_ :		Union[str, Any]							=  num_attention_heads
             a_ :		List[str]							=  intermediate_size
             a_ :		Optional[int]							=  hidden_act
             a_ :		Optional[Any]							=  hidden_dropout_prob
             a_ :		Union[str, Any]							=  attention_probs_dropout_prob
             a_ :		str							=  max_position_embeddings
             a_ :		Optional[Any]							=  type_vocab_size
             a_ :		Tuple							=  type_sequence_label_size
             a_ :		Optional[Any]							=  initializer_range
             a_ :		Any							=  num_choices
      def 	SCREAMING_SNAKE_CASE							(	self    :	List[Any]					)		->      Optional[int]:
             a_ :		Optional[int]							=  ids_tensor([self.batch_size, self.seq_length]					,							self.vocab_size					)

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

             a_ :		Tuple							=  DistilBertConfig(
                 vocab_size=self.vocab_size					,							dim=self.hidden_size					,							n_layers=self.num_hidden_layers					,							n_heads=self.num_attention_heads					,							hidden_dim=self.intermediate_size					,							hidden_act=self.hidden_act					,							dropout=self.hidden_dropout_prob					,							attention_dropout=self.attention_probs_dropout_prob					,							max_position_embeddings=self.max_position_embeddings					,							initializer_range=self.initializer_range					,							tie_weights_=SCREAMING_SNAKE_CASE__					,							)

             return config, input_ids, attention_mask

      def 	SCREAMING_SNAKE_CASE							(	self    :	Tuple					)		->      Tuple:
             a_ :		List[Any]							=  self.prepare_config_and_inputs()
             a_ :		Optional[Any]							=  config_and_inputs
             a_ :		Dict							=  {'input_ids': input_ids, 'attention_mask': attention_mask}
             return config, inputs_dict
@require_flax
class 			SCREAMING_SNAKE_CASE__  (    lowercase__						,   unittest.TestCase      ):
      snake_case__    :					Optional[Any]					       =						(
          (
              FlaxDistilBertModel,
              FlaxDistilBertForMaskedLM,
              FlaxDistilBertForMultipleChoice,
              FlaxDistilBertForQuestionAnswering,
              FlaxDistilBertForSequenceClassification,
              FlaxDistilBertForTokenClassification,
              FlaxDistilBertForQuestionAnswering,
          )
          if is_flax_available()
          else ()
      )
      def 	SCREAMING_SNAKE_CASE							(	self    :	Tuple					)		->      str:
             a_ :		int							=  FlaxDistilBertModelTester(self					)

      @slow
      def 	SCREAMING_SNAKE_CASE							(	self    :	str					)		->      List[str]:
             for model_class_name in self.all_model_classes:
                    a_ :		Optional[Any]							=  model_class_name.from_pretrained('distilbert-base-uncased'					)
                    a_ :		Optional[Any]							=  model(np.ones((1, 1)					)					)
                    self.assertIsNotNone(SCREAMING_SNAKE_CASE__					)
@require_flax
class 			SCREAMING_SNAKE_CASE__  (    unittest.TestCase      ):
      @slow
      def 	SCREAMING_SNAKE_CASE							(	self    :	str					)		->      Any:
             a_ :		Optional[Any]							=  FlaxDistilBertModel.from_pretrained('distilbert-base-uncased'					)
             a_ :		List[Any]							=  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]]					)
             a_ :		List[str]							=  np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]					)
             a_ :		Any							=  model(SCREAMING_SNAKE_CASE__					,							attention_mask=SCREAMING_SNAKE_CASE__					)[0]
             a_ :		int							=  (1, 1_1, 7_6_8)
             self.assertEqual(output.shape					,							SCREAMING_SNAKE_CASE__					)
             a_ :		Dict							=  np.array([[[-0.1639, 0.3299, 0.1648], [-0.1746, 0.3289, 0.1710], [-0.1884, 0.3357, 0.1810]]]					)

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






from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


UpperCAmelCase_				:			Union[str, Any]   =	{
    'configuration_whisper': ['WHISPER_PRETRAINED_CONFIG_ARCHIVE_MAP', 'WhisperConfig', 'WhisperOnnxConfig'],
    'feature_extraction_whisper': ['WhisperFeatureExtractor'],
    'processing_whisper': ['WhisperProcessor'],
    'tokenization_whisper': ['WhisperTokenizer'],
}

try:
   if not is_tokenizers_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   UpperCAmelCase_				:			Optional[Any]   =	['WhisperTokenizerFast']

try:
   if not is_torch_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   UpperCAmelCase_				:			int   =	[
       'WHISPER_PRETRAINED_MODEL_ARCHIVE_LIST',
       'WhisperForConditionalGeneration',
       'WhisperModel',
       'WhisperPreTrainedModel',
       'WhisperForAudioClassification',
   ]

try:
   if not is_tf_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   UpperCAmelCase_				:			int   =	[
       'TF_WHISPER_PRETRAINED_MODEL_ARCHIVE_LIST',
       'TFWhisperForConditionalGeneration',
       'TFWhisperModel',
       'TFWhisperPreTrainedModel',
   ]

try:
   if not is_flax_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   UpperCAmelCase_				:			Any   =	[
       'FlaxWhisperForConditionalGeneration',
       'FlaxWhisperModel',
       'FlaxWhisperPreTrainedModel',
       'FlaxWhisperForAudioClassification',
   ]


if TYPE_CHECKING:
   from .configuration_whisper import WHISPER_PRETRAINED_CONFIG_ARCHIVE_MAP, WhisperConfig, WhisperOnnxConfig
   from .feature_extraction_whisper import WhisperFeatureExtractor
   from .processing_whisper import WhisperProcessor
   from .tokenization_whisper import WhisperTokenizer

   try:
      if not is_tokenizers_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .tokenization_whisper_fast import WhisperTokenizerFast

   try:
      if not is_torch_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_whisper import (
          WHISPER_PRETRAINED_MODEL_ARCHIVE_LIST,
          WhisperForAudioClassification,
          WhisperForConditionalGeneration,
          WhisperModel,
          WhisperPreTrainedModel,
      )

   try:
      if not is_tf_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_tf_whisper import (
          TF_WHISPER_PRETRAINED_MODEL_ARCHIVE_LIST,
          TFWhisperForConditionalGeneration,
          TFWhisperModel,
          TFWhisperPreTrainedModel,
      )

   try:
      if not is_flax_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_flax_whisper import (
          FlaxWhisperForAudioClassification,
          FlaxWhisperForConditionalGeneration,
          FlaxWhisperModel,
          FlaxWhisperPreTrainedModel,
      )

else:
   import sys

   UpperCAmelCase_				:			int   =	_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)


