





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 						__magic_name__    :


 def __init__(						self		: Union[str, Any]			,						lowerCAmelCase__		: str			,						lowerCAmelCase__		: List[Any]=1_3			,						lowerCAmelCase__		: Dict=7			,						lowerCAmelCase__		: Union[str, Any]=True			,						lowerCAmelCase__		: str=True			,						lowerCAmelCase__		: List[str]=True			,						lowerCAmelCase__		: List[str]=True			,						lowerCAmelCase__		: Optional[int]=9_9			,						lowerCAmelCase__		: List[Any]=6_4			,						lowerCAmelCase__		: Dict=3_2			,						lowerCAmelCase__		: List[Any]=5			,						lowerCAmelCase__		: Any=4			,						lowerCAmelCase__		: List[str]=3_7			,						lowerCAmelCase__		: Tuple="gelu"			,						lowerCAmelCase__		: Union[str, Any]=0.1			,						lowerCAmelCase__		: Optional[int]=0.1			,						lowerCAmelCase__		: Optional[Any]=5_1_2			,						lowerCAmelCase__		: List[str]=1_6			,						lowerCAmelCase__		: Optional[int]=2			,						lowerCAmelCase__		: Optional[int]=0.02			,						lowerCAmelCase__		: List[str]=3			,						lowerCAmelCase__		: int=4			,						lowerCAmelCase__		: List[Any]=None			,						)    ->      Optional[int]:
   UpperCAmelCase					      =  parent
   UpperCAmelCase					      =  batch_size
   UpperCAmelCase					      =  seq_length
   UpperCAmelCase					      =  is_training
   UpperCAmelCase					      =  use_input_mask
   UpperCAmelCase					      =  use_token_type_ids
   UpperCAmelCase					      =  use_labels
   UpperCAmelCase					      =  vocab_size
   UpperCAmelCase					      =  hidden_size
   UpperCAmelCase					      =  embedding_size
   UpperCAmelCase					      =  num_hidden_layers
   UpperCAmelCase					      =  num_attention_heads
   UpperCAmelCase					      =  intermediate_size
   UpperCAmelCase					      =  hidden_act
   UpperCAmelCase					      =  hidden_dropout_prob
   UpperCAmelCase					      =  attention_probs_dropout_prob
   UpperCAmelCase					      =  max_position_embeddings
   UpperCAmelCase					      =  type_vocab_size
   UpperCAmelCase					      =  type_sequence_label_size
   UpperCAmelCase					      =  initializer_range
   UpperCAmelCase					      =  num_labels
   UpperCAmelCase					      =  num_choices
   UpperCAmelCase					      =  scope


 def _UpperCamelCase	(						self		: Optional[int]    )    ->      Tuple:
   UpperCAmelCase					      =  ids_tensor([self.batch_size, self.seq_length]			,						self.vocab_size    )

   UpperCAmelCase					      =  None
   if self.use_input_mask:
     UpperCAmelCase					      =  random_attention_mask([self.batch_size, self.seq_length]    )

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

   UpperCAmelCase					      =  None
   UpperCAmelCase					      =  None
   UpperCAmelCase					      =  None
   if self.use_labels:
     UpperCAmelCase					      =  ids_tensor([self.batch_size]			,						self.type_sequence_label_size    )
     UpperCAmelCase					      =  ids_tensor([self.batch_size, self.seq_length]			,						self.num_labels    )
     UpperCAmelCase					      =  ids_tensor([self.batch_size]			,						self.num_choices    )

   UpperCAmelCase					      =  self.get_config()

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


 def _UpperCamelCase	(						self		: List[str]    )    ->      List[str]:
   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=UpperCamelCase_			,						initializer_range=self.initializer_range			,						)


 def _UpperCamelCase	(						self		: Optional[Any]			,						lowerCAmelCase__		: List[str]			,						lowerCAmelCase__		: List[str]			,						lowerCAmelCase__		: List[Any]			,						lowerCAmelCase__		: int			,						lowerCAmelCase__		: Optional[int]			,						lowerCAmelCase__		: Any			,						lowerCAmelCase__		: str    )    ->      Any:
   UpperCAmelCase					      =  MegatronBertModel(config=UpperCamelCase_    )
   model.to(UpperCamelCase_    )
   model.eval()
   UpperCAmelCase					      =  model(UpperCamelCase_			,						attention_mask=UpperCamelCase_			,						token_type_ids=UpperCamelCase_    )
   UpperCAmelCase					      =  model(UpperCamelCase_			,						token_type_ids=UpperCamelCase_    )
   UpperCAmelCase					      =  model(UpperCamelCase_    )

   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 _UpperCamelCase	(						self		: int			,						lowerCAmelCase__		: Any			,						lowerCAmelCase__		: Dict			,						lowerCAmelCase__		: Union[str, Any]			,						lowerCAmelCase__		: List[Any]			,						lowerCAmelCase__		: List[Any]			,						lowerCAmelCase__		: List[str]			,						lowerCAmelCase__		: str    )    ->      List[str]:
   UpperCAmelCase					      =  MegatronBertForMaskedLM(config=UpperCamelCase_    )
   model.to(UpperCamelCase_    )
   model.eval()
   UpperCAmelCase					      =  model(UpperCamelCase_			,						attention_mask=UpperCamelCase_			,						token_type_ids=UpperCamelCase_			,						labels=UpperCamelCase_    )
   self.parent.assertEqual(result.logits.shape			,						(self.batch_size, self.seq_length, self.vocab_size)    )


 def _UpperCamelCase	(						self		: Any			,						lowerCAmelCase__		: List[Any]			,						lowerCAmelCase__		: int			,						lowerCAmelCase__		: str			,						lowerCAmelCase__		: Dict			,						lowerCAmelCase__		: List[Any]			,						lowerCAmelCase__		: Union[str, Any]			,						lowerCAmelCase__		: Optional[int]    )    ->      Optional[int]:
   UpperCAmelCase					      =  MegatronBertForCausalLM(config=UpperCamelCase_    )
   model.to(UpperCamelCase_    )
   model.eval()
   UpperCAmelCase					      =  model(UpperCamelCase_			,						attention_mask=UpperCamelCase_			,						token_type_ids=UpperCamelCase_			,						labels=UpperCamelCase_    )
   self.parent.assertEqual(result.logits.shape			,						(self.batch_size, self.seq_length, self.vocab_size)    )


 def _UpperCamelCase	(						self		: Tuple			,						lowerCAmelCase__		: int			,						lowerCAmelCase__		: List[Any]			,						lowerCAmelCase__		: Optional[Any]			,						lowerCAmelCase__		: Any			,						lowerCAmelCase__		: Optional[int]			,						lowerCAmelCase__		: Optional[Any]			,						lowerCAmelCase__		: str    )    ->      Tuple:
   UpperCAmelCase					      =  MegatronBertForNextSentencePrediction(config=UpperCamelCase_    )
   model.to(UpperCamelCase_    )
   model.eval()
   UpperCAmelCase					      =  model(
       UpperCamelCase_			,						attention_mask=UpperCamelCase_			,						token_type_ids=UpperCamelCase_			,						labels=UpperCamelCase_			,						)
   self.parent.assertEqual(result.logits.shape			,						(self.batch_size, 2)    )


 def _UpperCamelCase	(						self		: Optional[int]			,						lowerCAmelCase__		: Union[str, Any]			,						lowerCAmelCase__		: Tuple			,						lowerCAmelCase__		: Dict			,						lowerCAmelCase__		: Tuple			,						lowerCAmelCase__		: Dict			,						lowerCAmelCase__		: List[str]			,						lowerCAmelCase__		: Optional[int]    )    ->      List[Any]:
   UpperCAmelCase					      =  MegatronBertForPreTraining(config=UpperCamelCase_    )
   model.to(UpperCamelCase_    )
   model.eval()
   UpperCAmelCase					      =  model(
       UpperCamelCase_			,						attention_mask=UpperCamelCase_			,						token_type_ids=UpperCamelCase_			,						labels=UpperCamelCase_			,						next_sentence_label=UpperCamelCase_			,						)
   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 _UpperCamelCase	(						self		: List[Any]			,						lowerCAmelCase__		: Any			,						lowerCAmelCase__		: Dict			,						lowerCAmelCase__		: Any			,						lowerCAmelCase__		: Union[str, Any]			,						lowerCAmelCase__		: Optional[int]			,						lowerCAmelCase__		: Optional[int]			,						lowerCAmelCase__		: Optional[int]    )    ->      Optional[int]:
   UpperCAmelCase					      =  MegatronBertForQuestionAnswering(config=UpperCamelCase_    )
   model.to(UpperCamelCase_    )
   model.eval()
   UpperCAmelCase					      =  model(
       UpperCamelCase_			,						attention_mask=UpperCamelCase_			,						token_type_ids=UpperCamelCase_			,						start_positions=UpperCamelCase_			,						end_positions=UpperCamelCase_			,						)
   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 _UpperCamelCase	(						self		: str			,						lowerCAmelCase__		: str			,						lowerCAmelCase__		: Optional[int]			,						lowerCAmelCase__		: Any			,						lowerCAmelCase__		: Optional[int]			,						lowerCAmelCase__		: Union[str, Any]			,						lowerCAmelCase__		: Optional[Any]			,						lowerCAmelCase__		: Union[str, Any]    )    ->      Dict:
   UpperCAmelCase					      =  self.num_labels
   UpperCAmelCase					      =  MegatronBertForSequenceClassification(UpperCamelCase_    )
   model.to(UpperCamelCase_    )
   model.eval()
   UpperCAmelCase					      =  model(UpperCamelCase_			,						attention_mask=UpperCamelCase_			,						token_type_ids=UpperCamelCase_			,						labels=UpperCamelCase_    )
   self.parent.assertEqual(result.logits.shape			,						(self.batch_size, self.num_labels)    )


 def _UpperCamelCase	(						self		: Optional[Any]			,						lowerCAmelCase__		: str			,						lowerCAmelCase__		: List[str]			,						lowerCAmelCase__		: Union[str, Any]			,						lowerCAmelCase__		: Optional[Any]			,						lowerCAmelCase__		: Dict			,						lowerCAmelCase__		: List[Any]			,						lowerCAmelCase__		: Optional[Any]    )    ->      int:
   UpperCAmelCase					      =  self.num_labels
   UpperCAmelCase					      =  MegatronBertForTokenClassification(config=UpperCamelCase_    )
   model.to(UpperCamelCase_    )
   model.eval()
   UpperCAmelCase					      =  model(UpperCamelCase_			,						attention_mask=UpperCamelCase_			,						token_type_ids=UpperCamelCase_			,						labels=UpperCamelCase_    )
   self.parent.assertEqual(result.logits.shape			,						(self.batch_size, self.seq_length, self.num_labels)    )


 def _UpperCamelCase	(						self		: Optional[Any]			,						lowerCAmelCase__		: List[Any]			,						lowerCAmelCase__		: Any			,						lowerCAmelCase__		: Any			,						lowerCAmelCase__		: Union[str, Any]			,						lowerCAmelCase__		: Optional[int]			,						lowerCAmelCase__		: Dict			,						lowerCAmelCase__		: Dict    )    ->      List[str]:
   UpperCAmelCase					      =  self.num_choices
   UpperCAmelCase					      =  MegatronBertForMultipleChoice(config=UpperCamelCase_    )
   model.to(UpperCamelCase_    )
   model.eval()
   UpperCAmelCase					      =  input_ids.unsqueeze(1    ).expand(-1			,						self.num_choices			,						-1    ).contiguous()
   UpperCAmelCase					      =  token_type_ids.unsqueeze(1    ).expand(-1			,						self.num_choices			,						-1    ).contiguous()
   UpperCAmelCase					      =  input_mask.unsqueeze(1    ).expand(-1			,						self.num_choices			,						-1    ).contiguous()
   UpperCAmelCase					      =  model(
       UpperCamelCase_			,						attention_mask=UpperCamelCase_			,						token_type_ids=UpperCamelCase_			,						labels=UpperCamelCase_			,						)
   self.parent.assertEqual(result.logits.shape			,						(self.batch_size, self.num_choices)    )


 def _UpperCamelCase	(						self		: Tuple    )    ->      Union[str, Any]:
   UpperCAmelCase					      =  self.prepare_config_and_inputs()
   (
       (
       UpperCAmelCase     
   )      ,	(
       UpperCAmelCase     
   )      ,	(
       UpperCAmelCase     
   )      ,	(
       UpperCAmelCase     
   )      ,	(
       UpperCAmelCase     
   )      ,	(
       UpperCAmelCase     
   )      ,	(
       UpperCAmelCase     
   )      ,	     
   )					      =  config_and_inputs
   UpperCAmelCase					      =  {"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
   return config, inputs_dict






@require_torch
class 						__magic_name__    (       __lowerCamelCase	,						__lowerCamelCase	,						unittest.TestCase    ):
 UpperCAmelCase							    =						(
     (
         MegatronBertModel,
         MegatronBertForMaskedLM,
         MegatronBertForCausalLM,
         MegatronBertForMultipleChoice,
         MegatronBertForNextSentencePrediction,
         MegatronBertForPreTraining,
         MegatronBertForQuestionAnswering,
         MegatronBertForSequenceClassification,
         MegatronBertForTokenClassification,
     )
     if is_torch_available()
     else ()
 )
 UpperCAmelCase							    =						(
     {
         '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 {}
 )
 UpperCAmelCase							    =						True
 # test_resize_embeddings = False
 UpperCAmelCase							    =						False


 def _UpperCamelCase	(						self		: Union[str, Any]			,						lowerCAmelCase__		: List[Any]			,						lowerCAmelCase__		: str			,						lowerCAmelCase__		: Tuple=False    )    ->      Optional[int]:
   UpperCAmelCase					      =  super()._prepare_for_class(UpperCamelCase_			,						UpperCamelCase_			,						return_labels=UpperCamelCase_    )

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


 def _UpperCamelCase	(						self		: Dict    )    ->      List[str]:
   UpperCAmelCase					      =  MegatronBertModelTester(self    )
   UpperCAmelCase					      =  ConfigTester(self			,						config_class=UpperCamelCase_			,						hidden_size=3_7    )


 def _UpperCamelCase	(						self		: int    )    ->      Optional[int]:
   self.config_tester.run_common_tests()


 def _UpperCamelCase	(						self		: Optional[Any]    )    ->      Union[str, Any]:
   UpperCAmelCase					      =  self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_megatron_bert_model(*UpperCamelCase_    )


 def _UpperCamelCase	(						self		: Union[str, Any]    )    ->      str:
   UpperCAmelCase					      =  self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_megatron_bert_for_masked_lm(*UpperCamelCase_    )


 def _UpperCamelCase	(						self		: Union[str, Any]    )    ->      Optional[Any]:
   UpperCAmelCase					      =  self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*UpperCamelCase_    )


 def _UpperCamelCase	(						self		: List[Any]    )    ->      Optional[Any]:
   UpperCAmelCase					      =  self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*UpperCamelCase_    )


 def _UpperCamelCase	(						self		: Optional[Any]    )    ->      int:
   UpperCAmelCase					      =  self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_megatron_bert_for_pretraining(*UpperCamelCase_    )


 def _UpperCamelCase	(						self		: List[Any]    )    ->      List[Any]:
   UpperCAmelCase					      =  self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_megatron_bert_for_question_answering(*UpperCamelCase_    )


 def _UpperCamelCase	(						self		: str    )    ->      Optional[Any]:
   UpperCAmelCase					      =  self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*UpperCamelCase_    )


 def _UpperCamelCase	(						self		: List[Any]    )    ->      Optional[int]:
   UpperCAmelCase					      =  self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_megatron_bert_for_token_classification(*UpperCamelCase_    )

def      _lowerCAmelCase(	__A			):
  return torch.tensor(
      A__   ,       dtype=torch.long   ,       device=A__   ,       )


lowerCAmelCase__       = 1e-4






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


 @slow
 @unittest.skip("Model is not available."    )
 def _UpperCamelCase	(						self		: Any    )    ->      Dict:
   UpperCAmelCase					      =  "nvidia/megatron-bert-uncased-345m"
   if "MYDIR" in os.environ:
     UpperCAmelCase					      =  os.path.join(os.environ["MYDIR"]			,						UpperCamelCase_    )
   UpperCAmelCase					      =  MegatronBertModel.from_pretrained(UpperCamelCase_    )
   model.to(UpperCamelCase_    )
   model.half()
   UpperCAmelCase					      =  _long_tensor([[1_0_1, 7_1_1_0, 1_0_0_5, 1_0_5_6, 2_0_2_3, 1_1_3_3_3, 1_7_4_1_3, 1_0_2_9, 1_0_2]]    )
   with torch.no_grad():
     UpperCAmelCase					      =  model(UpperCamelCase_    )[0]
   UpperCAmelCase					      =  torch.Size((1, 9, 1_0_2_4)    )
   self.assertEqual(output.shape			,						UpperCamelCase_    )

   UpperCAmelCase					      =  [-0.6_040, -0.2_517, -0.1_025, 0.3_420, -0.6_758, -0.0_017, -0.1_089, -0.1_990, 0.5_728]
   for ii in range(3    ):
     for jj in range(3    ):
       UpperCAmelCase					      =  output[0, ii, jj]
       UpperCAmelCase					      =  expected[3 * ii + jj]
       UpperCAmelCase					      =  "ii={} jj={} a={} b={}".format(UpperCamelCase_			,						UpperCamelCase_			,						UpperCamelCase_			,						UpperCamelCase_    )
       self.assertTrue(math.isclose(UpperCamelCase_			,						UpperCamelCase_			,						rel_tol=UpperCamelCase_			,						abs_tol=UpperCamelCase_    )			,						msg=UpperCamelCase_    )






def      _lowerCAmelCase(	__A			):

  if not isinstance(__A   ,       __A			):
    raise TypeError("only integers accepted as input"			)
  else:
    UpperCAmelCase					      =  str(abs(__A			)			)
    UpperCAmelCase					      =  [list(__A			) for char in range(len(__A			)			)]
    for index in range(len(__A			)			):
      num_transpositions[index].pop(__A			)
    return max(
        int("".join(list(__A			)			)			) for transposition in num_transpositions			)


if __name__ == "__main__":
      __import__("doctest").testmod()



