





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






			def __init__(self					,							__magic_name__					,							__magic_name__=13					,							__magic_name__=7					,							__magic_name__=True					,							__magic_name__=True					,							__magic_name__=True					,							__magic_name__=True					,							__magic_name__=99					,							__magic_name__=32					,							__magic_name__=5					,							__magic_name__=4					,							__magic_name__=37					,							__magic_name__="gelu"					,							__magic_name__=0.1					,							__magic_name__=0.1					,							__magic_name__=512					,							__magic_name__=16					,							__magic_name__=2					,							__magic_name__=0.02					,							__magic_name__=4					,							)     ->   Union[str, Any]:


							'''simple docstring'''



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






			def 			lowerCamelCase  (self    )     ->   Dict:


							'''simple docstring'''



							snake_case_						:    List[str]				=  ids_tensor([self.batch_size, self.seq_length]					,							self.vocab_size    )

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

							snake_case_						:    Dict				=  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_=__magic_name__					,							)

							return config, input_ids, attention_mask






			def 			lowerCamelCase  (self    )     ->   str:


							'''simple docstring'''



							snake_case_						:    List[str]				=  self.prepare_config_and_inputs()
							snake_case_			,							snake_case_			,							snake_case_						:    Optional[int]				=  config_and_inputs
							snake_case_						:    Tuple				=  {'''input_ids''': input_ids, '''attention_mask''': attention_mask}
							return config, inputs_dict




@require_flax
class 					__lowerCAmelCase       (			_a,     unittest.TestCase     ):
			lowerCamelCase_     :     Optional[Any]       			=      (
			    (
			        FlaxDistilBertModel,
			        FlaxDistilBertForMaskedLM,
			        FlaxDistilBertForMultipleChoice,
			        FlaxDistilBertForQuestionAnswering,
			        FlaxDistilBertForSequenceClassification,
			        FlaxDistilBertForTokenClassification,
			        FlaxDistilBertForQuestionAnswering,
			    )
			    if is_flax_available()
			    else ()
			)






			def 			lowerCamelCase  (self    )     ->   Any:


							'''simple docstring'''



							snake_case_						:    List[str]				=  FlaxDistilBertModelTester(self    )






			@slow
			def 			lowerCamelCase  (self    )     ->   Optional[Any]:


							'''simple docstring'''



							for model_class_name in self.all_model_classes:
											snake_case_						:    Any				=  model_class_name.from_pretrained('''distilbert-base-uncased'''    )
											snake_case_						:    int				=  model(np.ones((1, 1)    )    )
											self.assertIsNotNone(__magic_name__    )




@require_flax
class 					__lowerCAmelCase       (			unittest.TestCase     ):






			@slow
			def 			lowerCamelCase  (self    )     ->   Union[str, Any]:


							'''simple docstring'''



							snake_case_						:    Optional[int]				=  FlaxDistilBertModel.from_pretrained('''distilbert-base-uncased'''    )
							snake_case_						:    Optional[Any]				=  np.array([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]    )
							snake_case_						:    Optional[Any]				=  np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]    )
							snake_case_						:    Optional[Any]				=  model(__magic_name__					,							attention_mask=__magic_name__    )[0]
							snake_case_						:    Optional[int]				=  (1, 11, 768)
							self.assertEqual(output.shape					,							__magic_name__    )
							snake_case_						:    Union[str, Any]				=  np.array([[[-0.1_639, 0.3_299, 0.1_648], [-0.1_746, 0.3_289, 0.1_710], [-0.1_884, 0.3_357, 0.1_810]]]    )

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







from ...configuration_utils import PretrainedConfig
from ...utils import logging


lowerCAmelCase_    =				logging.get_logger(__name__)

lowerCAmelCase_    =				{
    # See all MEGATRON_BERT models at https://huggingface.co/models?filter=bert
}



class 					__lowerCAmelCase       (			_a     ):
			lowerCamelCase_     :     Any       			=      '''megatron-bert'''






			def __init__(self					,							__magic_name__=2_9056					,							__magic_name__=1024					,							__magic_name__=24					,							__magic_name__=16					,							__magic_name__=4096					,							__magic_name__="gelu"					,							__magic_name__=0.1					,							__magic_name__=0.1					,							__magic_name__=512					,							__magic_name__=2					,							__magic_name__=0.02					,							__magic_name__=1e-12					,							__magic_name__=0					,							__magic_name__="absolute"					,							__magic_name__=True					,							**__magic_name__					,							)     ->   Any:


							'''simple docstring'''



							super().__init__(pad_token_id=__magic_name__					,							**__magic_name__    )

							snake_case_						:    Union[str, Any]				=  vocab_size
							snake_case_						:    Dict				=  hidden_size
							snake_case_						:    Dict				=  num_hidden_layers
							snake_case_						:    Optional[int]				=  num_attention_heads
							snake_case_						:    int				=  hidden_act
							snake_case_						:    List[str]				=  intermediate_size
							snake_case_						:    Dict				=  hidden_dropout_prob
							snake_case_						:    str				=  attention_probs_dropout_prob
							snake_case_						:    str				=  max_position_embeddings
							snake_case_						:    Any				=  type_vocab_size
							snake_case_						:    int				=  initializer_range
							snake_case_						:    int				=  layer_norm_eps
							snake_case_						:    List[str]				=  position_embedding_type
							snake_case_						:    Dict				=  use_cache
