



'''simple docstring'''





import math
import os
import unittest

from transformers import MegatronBertConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
  import torch

  from transformers import (
      MODEL_FOR_PRETRAINING_MAPPING,
      MegatronBertForCausalLM,
      MegatronBertForMaskedLM,
      MegatronBertForMultipleChoice,
      MegatronBertForNextSentencePrediction,
      MegatronBertForPreTraining,
      MegatronBertForQuestionAnswering,
      MegatronBertForSequenceClassification,
      MegatronBertForTokenClassification,
      MegatronBertModel,
  )






class __A    :



 '''simple docstring'''







 def __init__(self				,     A				,     A=13				,     A=7				,     A=True				,     A=True				,     A=True				,     A=True				,     A=99				,     A=64				,     A=32				,     A=5				,     A=4				,     A=37				,     A="gelu"				,     A=0.1				,     A=0.1				,     A=512				,     A=16				,     A=2				,     A=0.02				,     A=3				,     A=4				,     A=None				,     )					->  List[str]:


      """simple docstring"""



      _a											=						parent
      _a											=						batch_size
      _a											=						seq_length
      _a											=						is_training
      _a											=						use_input_mask
      _a											=						use_token_type_ids
      _a											=						use_labels
      _a											=						vocab_size
      _a											=						hidden_size
      _a											=						embedding_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_labels
      _a											=						num_choices
      _a											=						scope



 def       a__						(self )					->  Tuple:


      """simple docstring"""



      _a											=						ids_tensor([self.batch_size, self.seq_length]				,     self.vocab_size )

      _a											=						None
      if self.use_input_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											=						None
      _a											=						None
      _a											=						None
      if self.use_labels:
           _a											=						ids_tensor([self.batch_size]				,     self.type_sequence_label_size )
           _a											=						ids_tensor([self.batch_size, self.seq_length]				,     self.num_labels )
           _a											=						ids_tensor([self.batch_size]				,     self.num_choices )

      _a											=						self.get_config()

      return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels



 def       a__						(self )					->  List[Any]:


      """simple docstring"""



      return MegatronBertConfig(
          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				,     embedding_size=self.embedding_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				,     )



 def       a__						(self				,     A				,     A				,     A				,     A				,     A				,     A				,     A )					->  Optional[Any]:


      """simple docstring"""



      _a											=						MegatronBertModel(config=A )
      model.to(A )
      model.eval()
      _a											=						model(A				,     attention_mask=A				,     token_type_ids=A )
      _a											=						model(A				,     token_type_ids=A )
      _a											=						model(A )

      self.parent.assertEqual(result.last_hidden_state.shape				,     (self.batch_size, self.seq_length, self.hidden_size) )
      self.parent.assertEqual(result.pooler_output.shape				,     (self.batch_size, self.hidden_size) )



 def       a__						(self				,     A				,     A				,     A				,     A				,     A				,     A				,     A )					->  Tuple:


      """simple docstring"""



      _a											=						MegatronBertForMaskedLM(config=A )
      model.to(A )
      model.eval()
      _a											=						model(A				,     attention_mask=A				,     token_type_ids=A				,     labels=A )
      self.parent.assertEqual(result.logits.shape				,     (self.batch_size, self.seq_length, self.vocab_size) )



 def       a__						(self				,     A				,     A				,     A				,     A				,     A				,     A				,     A )					->  Union[str, Any]:


      """simple docstring"""



      _a											=						MegatronBertForCausalLM(config=A )
      model.to(A )
      model.eval()
      _a											=						model(A				,     attention_mask=A				,     token_type_ids=A				,     labels=A )
      self.parent.assertEqual(result.logits.shape				,     (self.batch_size, self.seq_length, self.vocab_size) )



 def       a__						(self				,     A				,     A				,     A				,     A				,     A				,     A				,     A )					->  Union[str, Any]:


      """simple docstring"""



      _a											=						MegatronBertForNextSentencePrediction(config=A )
      model.to(A )
      model.eval()
      _a											=						model(
          A				,     attention_mask=A				,     token_type_ids=A				,     labels=A				,     )
      self.parent.assertEqual(result.logits.shape				,     (self.batch_size, 2) )



 def       a__						(self				,     A				,     A				,     A				,     A				,     A				,     A				,     A )					->  Optional[Any]:


      """simple docstring"""



      _a											=						MegatronBertForPreTraining(config=A )
      model.to(A )
      model.eval()
      _a											=						model(
          A				,     attention_mask=A				,     token_type_ids=A				,     labels=A				,     next_sentence_label=A				,     )
      self.parent.assertEqual(result.prediction_logits.shape				,     (self.batch_size, self.seq_length, self.vocab_size) )
      self.parent.assertEqual(result.seq_relationship_logits.shape				,     (self.batch_size, 2) )



 def       a__						(self				,     A				,     A				,     A				,     A				,     A				,     A				,     A )					->  Tuple:


      """simple docstring"""



      _a											=						MegatronBertForQuestionAnswering(config=A )
      model.to(A )
      model.eval()
      _a											=						model(
          A				,     attention_mask=A				,     token_type_ids=A				,     start_positions=A				,     end_positions=A				,     )
      self.parent.assertEqual(result.start_logits.shape				,     (self.batch_size, self.seq_length) )
      self.parent.assertEqual(result.end_logits.shape				,     (self.batch_size, self.seq_length) )



 def       a__						(self				,     A				,     A				,     A				,     A				,     A				,     A				,     A )					->  str:


      """simple docstring"""



      _a											=						self.num_labels
      _a											=						MegatronBertForSequenceClassification(A )
      model.to(A )
      model.eval()
      _a											=						model(A				,     attention_mask=A				,     token_type_ids=A				,     labels=A )
      self.parent.assertEqual(result.logits.shape				,     (self.batch_size, self.num_labels) )



 def       a__						(self				,     A				,     A				,     A				,     A				,     A				,     A				,     A )					->  str:


      """simple docstring"""



      _a											=						self.num_labels
      _a											=						MegatronBertForTokenClassification(config=A )
      model.to(A )
      model.eval()
      _a											=						model(A				,     attention_mask=A				,     token_type_ids=A				,     labels=A )
      self.parent.assertEqual(result.logits.shape				,     (self.batch_size, self.seq_length, self.num_labels) )



 def       a__						(self				,     A				,     A				,     A				,     A				,     A				,     A				,     A )					->  Union[str, Any]:


      """simple docstring"""



      _a											=						self.num_choices
      _a											=						MegatronBertForMultipleChoice(config=A )
      model.to(A )
      model.eval()
      _a											=						input_ids.unsqueeze(1 ).expand(-1				,     self.num_choices				,     -1 ).contiguous()
      _a											=						token_type_ids.unsqueeze(1 ).expand(-1				,     self.num_choices				,     -1 ).contiguous()
      _a											=						input_mask.unsqueeze(1 ).expand(-1				,     self.num_choices				,     -1 ).contiguous()
      _a											=						model(
          A				,     attention_mask=A				,     token_type_ids=A				,     labels=A				,     )
      self.parent.assertEqual(result.logits.shape				,     (self.batch_size, self.num_choices) )



 def       a__						(self )					->  Any:


      """simple docstring"""



      _a											=						self.prepare_config_and_inputs()
      (
          (
          _a
      )						,					(
          _a
      )						,					(
          _a
      )						,					(
          _a
      )						,					(
          _a
      )						,					(
          _a
      )						,					(
          _a
      )						,					
      )											=						config_and_inputs
      _a											=						{'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
      return config, inputs_dict






@require_torch
class __A    (						A     ,     A     ,     unittest.TestCase		):



 '''simple docstring'''







 __lowerCamelCase					: Optional[Any]    = (
     (
         MegatronBertModel,
         MegatronBertForMaskedLM,
         MegatronBertForCausalLM,
         MegatronBertForMultipleChoice,
         MegatronBertForNextSentencePrediction,
         MegatronBertForPreTraining,
         MegatronBertForQuestionAnswering,
         MegatronBertForSequenceClassification,
         MegatronBertForTokenClassification,
     )
     if is_torch_available()
     else ()
 )
 __lowerCamelCase					: Optional[Any]    = (
     {
         'feature-extraction': MegatronBertModel,
         'fill-mask': MegatronBertForMaskedLM,
         'question-answering': MegatronBertForQuestionAnswering,
         'text-classification': MegatronBertForSequenceClassification,
         'text-generation': MegatronBertForCausalLM,
         'token-classification': MegatronBertForTokenClassification,
         'zero-shot': MegatronBertForSequenceClassification,
     }
     if is_torch_available()
     else {}
 )
 __lowerCamelCase					: Union[str, Any]    = True
 # test_resize_embeddings = False
 __lowerCamelCase					: List[Any]    = False



 def       a__						(self				,     A				,     A				,     A=False )					->  Dict:


      """simple docstring"""



      _a											=						super()._prepare_for_class(A				,     A				,     return_labels=A )

      if return_labels:
           if model_class in get_values(A ):
                _a											=						torch.zeros(
                    (self.model_tester.batch_size, self.model_tester.seq_length)				,     dtype=torch.long				,     device=A )
                _a											=						torch.zeros(
                    self.model_tester.batch_size				,     dtype=torch.long				,     device=A )
      return inputs_dict



 def       a__						(self )					->  str:


      """simple docstring"""



      _a											=						MegatronBertModelTester(self )
      _a											=						ConfigTester(self				,     config_class=A				,     hidden_size=37 )



 def       a__						(self )					->  str:


      """simple docstring"""



      self.config_tester.run_common_tests()



 def       a__						(self )					->  List[str]:


      """simple docstring"""



      _a											=						self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_megatron_bert_model(*A )



 def       a__						(self )					->  int:


      """simple docstring"""



      _a											=						self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_megatron_bert_for_masked_lm(*A )



 def       a__						(self )					->  int:


      """simple docstring"""



      _a											=						self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*A )



 def       a__						(self )					->  Optional[int]:


      """simple docstring"""



      _a											=						self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*A )



 def       a__						(self )					->  List[Any]:


      """simple docstring"""



      _a											=						self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_megatron_bert_for_pretraining(*A )



 def       a__						(self )					->  Tuple:


      """simple docstring"""



      _a											=						self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_megatron_bert_for_question_answering(*A )



 def       a__						(self )					->  List[Any]:


      """simple docstring"""



      _a											=						self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*A )



 def       a__						(self )					->  str:


      """simple docstring"""



      _a											=						self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_megatron_bert_for_token_classification(*A )





def lowerCAmelCase						(__A):



     """simple docstring"""



     return torch.tensor(
         __A						,			dtype=torch.long						,			device=__A						,			)


lowercase_	 =						1e-4






@require_torch
@require_sentencepiece
@require_tokenizers
class __A    (						unittest.TestCase		):



 '''simple docstring'''







 @slow
 @unittest.skip('''Model is not available.''' )
 def       a__						(self )					->  List[Any]:


      """simple docstring"""



      _a											=						'''nvidia/megatron-bert-uncased-345m'''
      if "MYDIR" in os.environ:
           _a											=						os.path.join(os.environ['''MYDIR''']				,     A )
      _a											=						MegatronBertModel.from_pretrained(A )
      model.to(A )
      model.half()
      _a											=						_long_tensor([[101, 7_110, 1_005, 1_056, 2_023, 11_333, 17_413, 1_029, 102]] )
      with torch.no_grad():
           _a											=						model(A )[0]
      _a											=						torch.Size((1, 9, 1_024) )
      self.assertEqual(output.shape				,     A )

      _a											=						[-0.6040, -0.2517, -0.1025, 0.3420, -0.6758, -0.0017, -0.1089, -0.1990, 0.5728]
      for ii in range(3 ):
           for jj in range(3 ):
                _a											=						output[0, ii, jj]
                _a											=						expected[3 * ii + jj]
                _a											=						'''ii={} jj={} a={} b={}'''.format(A				,     A				,     A				,     A )
                self.assertTrue(math.isclose(A				,     A				,     rel_tol=A				,     abs_tol=A )				,     msg=A )








'''simple docstring'''





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


lowercase_	 =						logging.get_logger(__name__)

# TODO: upload to AWS
lowercase_	 =						{
    "yjernite/retribert-base-uncased": (
        "https://huggingface.co/yjernite/retribert-base-uncased/resolve/main/config.json"
    ),
}






class __A    (						A		):



 '''simple docstring'''







 __lowerCamelCase					: Optional[Any]    = 'retribert'



 def __init__(self				,     A=30_522				,     A=768				,     A=8				,     A=12				,     A=3_072				,     A="gelu"				,     A=0.1				,     A=0.1				,     A=512				,     A=2				,     A=0.02				,     A=1E-12				,     A=True				,     A=128				,     A=0				,     **A				,     )					->  Any:


      """simple docstring"""



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

      _a											=						vocab_size
      _a											=						hidden_size
      _a											=						num_hidden_layers
      _a											=						num_attention_heads
      _a											=						hidden_act
      _a											=						intermediate_size
      _a											=						hidden_dropout_prob
      _a											=						attention_probs_dropout_prob
      _a											=						max_position_embeddings
      _a											=						type_vocab_size
      _a											=						initializer_range
      _a											=						layer_norm_eps
      _a											=						share_encoders
      _a											=						projection_dim



