




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 :  Tuple						,							A_ :  str						,							A_ :  str=13						,							A_ :  Tuple=7						,							A_ :  List[Any]=True						,							A_ :  List[Any]=True						,							A_ :  Optional[Any]=True						,							A_ :  Dict=True						,							A_ :  List[Any]=99						,							A_ :  List[Any]=64						,							A_ :  Tuple=32						,							A_ :  Dict=5						,							A_ :  int=4						,							A_ :  List[str]=37						,							A_ :  Optional[int]="gelu"						,							A_ :  List[Any]=0.1						,							A_ :  int=0.1						,							A_ :  Optional[int]=512						,							A_ :  Tuple=16						,							A_ :  Tuple=2						,							A_ :  List[Any]=0.02						,							A_ :  List[str]=3						,							A_ :  List[str]=4						,							A_ :  List[Any]=None						,							)       ->     Optional[Any]:

       """simple docstring"""






       lowerCamelCase_							=   parent
       lowerCamelCase_							=   batch_size
       lowerCamelCase_							=   seq_length
       lowerCamelCase_							=   is_training
       lowerCamelCase_							=   use_input_mask
       lowerCamelCase_							=   use_token_type_ids
       lowerCamelCase_							=   use_labels
       lowerCamelCase_							=   vocab_size
       lowerCamelCase_							=   hidden_size
       lowerCamelCase_							=   embedding_size
       lowerCamelCase_							=   num_hidden_layers
       lowerCamelCase_							=   num_attention_heads
       lowerCamelCase_							=   intermediate_size
       lowerCamelCase_							=   hidden_act
       lowerCamelCase_							=   hidden_dropout_prob
       lowerCamelCase_							=   attention_probs_dropout_prob
       lowerCamelCase_							=   max_position_embeddings
       lowerCamelCase_							=   type_vocab_size
       lowerCamelCase_							=   type_sequence_label_size
       lowerCamelCase_							=   initializer_range
       lowerCamelCase_							=   num_labels
       lowerCamelCase_							=   num_choices
       lowerCamelCase_							=   scope

 def     a__			(    self :  str  )       ->     Dict:

       """simple docstring"""






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

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

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

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

       lowerCamelCase_							=   self.get_config()

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

 def     a__			(    self :  Dict  )       ->     Tuple:

       """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 :  Union[str, Any]						,							A_ :  Dict						,							A_ :  int						,							A_ :  Dict						,							A_ :  Any						,							A_ :  Dict						,							A_ :  List[Any]						,							A_ :  Dict  )       ->     Union[str, Any]:

       """simple docstring"""






       lowerCamelCase_							=   MegatronBertModel(config=A_  )
       model.to(A_  )
       model.eval()
       lowerCamelCase_							=   model(A_						,							attention_mask=A_						,							token_type_ids=A_  )
       lowerCamelCase_							=   model(A_						,							token_type_ids=A_  )
       lowerCamelCase_							=   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 :  Dict						,							A_ :  Dict						,							A_ :  int						,							A_ :  List[Any]						,							A_ :  List[str]						,							A_ :  List[Any]						,							A_ :  Dict						,							A_ :  List[str]  )       ->     Any:

       """simple docstring"""






       lowerCamelCase_							=   MegatronBertForMaskedLM(config=A_  )
       model.to(A_  )
       model.eval()
       lowerCamelCase_							=   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 :  Any						,							A_ :  Dict						,							A_ :  str						,							A_ :  Dict						,							A_ :  Union[str, Any]						,							A_ :  Optional[Any]						,							A_ :  List[str]						,							A_ :  Optional[int]  )       ->     List[Any]:

       """simple docstring"""






       lowerCamelCase_							=   MegatronBertForCausalLM(config=A_  )
       model.to(A_  )
       model.eval()
       lowerCamelCase_							=   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 :  int						,							A_ :  Any						,							A_ :  int						,							A_ :  Union[str, Any]						,							A_ :  List[str]						,							A_ :  Dict						,							A_ :  List[str]						,							A_ :  Optional[int]  )       ->     Union[str, Any]:

       """simple docstring"""






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

 def     a__			(    self :  int						,							A_ :  Any						,							A_ :  Tuple						,							A_ :  Union[str, Any]						,							A_ :  Any						,							A_ :  List[str]						,							A_ :  Union[str, Any]						,							A_ :  Optional[int]  )       ->     Tuple:

       """simple docstring"""






       lowerCamelCase_							=   MegatronBertForPreTraining(config=A_  )
       model.to(A_  )
       model.eval()
       lowerCamelCase_							=   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 :  Optional[int]						,							A_ :  Optional[Any]						,							A_ :  Dict						,							A_ :  int						,							A_ :  List[str]						,							A_ :  Optional[Any]						,							A_ :  List[str]						,							A_ :  Dict  )       ->     List[str]:

       """simple docstring"""






       lowerCamelCase_							=   MegatronBertForQuestionAnswering(config=A_  )
       model.to(A_  )
       model.eval()
       lowerCamelCase_							=   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 :  Optional[int]						,							A_ :  Tuple						,							A_ :  Dict						,							A_ :  Any						,							A_ :  Dict						,							A_ :  List[Any]						,							A_ :  Tuple						,							A_ :  Union[str, Any]  )       ->     Optional[Any]:

       """simple docstring"""






       lowerCamelCase_							=   self.num_labels
       lowerCamelCase_							=   MegatronBertForSequenceClassification(A_  )
       model.to(A_  )
       model.eval()
       lowerCamelCase_							=   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 :  Optional[Any]						,							A_ :  Optional[Any]						,							A_ :  Union[str, Any]						,							A_ :  Any						,							A_ :  Optional[Any]						,							A_ :  Union[str, Any]						,							A_ :  List[Any]						,							A_ :  Optional[int]  )       ->     Optional[int]:

       """simple docstring"""






       lowerCamelCase_							=   self.num_labels
       lowerCamelCase_							=   MegatronBertForTokenClassification(config=A_  )
       model.to(A_  )
       model.eval()
       lowerCamelCase_							=   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 :  str						,							A_ :  List[str]						,							A_ :  List[str]						,							A_ :  Dict						,							A_ :  str						,							A_ :  Union[str, Any]						,							A_ :  List[Any]						,							A_ :  Optional[Any]  )       ->     Tuple:

       """simple docstring"""






       lowerCamelCase_							=   self.num_choices
       lowerCamelCase_							=   MegatronBertForMultipleChoice(config=A_  )
       model.to(A_  )
       model.eval()
       lowerCamelCase_							=   input_ids.unsqueeze(1  ).expand(-1						,							self.num_choices						,							-1  ).contiguous()
       lowerCamelCase_							=   token_type_ids.unsqueeze(1  ).expand(-1						,							self.num_choices						,							-1  ).contiguous()
       lowerCamelCase_							=   input_mask.unsqueeze(1  ).expand(-1						,							self.num_choices						,							-1  ).contiguous()
       lowerCamelCase_							=   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 :  str  )       ->     str:

       """simple docstring"""






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







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



 '''simple docstring'''


 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     a__			(    self :  str						,							A_ :  Optional[int]						,							A_ :  Optional[int]						,							A_ :  Any=False  )       ->     Optional[int]:

       """simple docstring"""






       lowerCamelCase_							=   super()._prepare_for_class(A_						,							A_						,							return_labels=A_  )

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

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

       """simple docstring"""






       lowerCamelCase_							=   MegatronBertModelTester(self  )
       lowerCamelCase_							=   ConfigTester(self						,							config_class=A_						,							hidden_size=37  )

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

       """simple docstring"""






       self.config_tester.run_common_tests()

 def     a__			(    self :  int  )       ->     Dict:

       """simple docstring"""






       lowerCamelCase_							=   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_model(*A_  )

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

       """simple docstring"""






       lowerCamelCase_							=   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_masked_lm(*A_  )

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

       """simple docstring"""






       lowerCamelCase_							=   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*A_  )

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

       """simple docstring"""






       lowerCamelCase_							=   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*A_  )

 def     a__			(    self :  Dict  )       ->     str:

       """simple docstring"""






       lowerCamelCase_							=   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_pretraining(*A_  )

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

       """simple docstring"""






       lowerCamelCase_							=   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_question_answering(*A_  )

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

       """simple docstring"""






       lowerCamelCase_							=   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*A_  )





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

       """simple docstring"""






       lowerCamelCase_							=   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_token_classification(*A_  )

def       _SCREAMING_SNAKE_CASE      (							lowercase     :	Optional[Any]					):


      '''simple docstring'''



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


lowerCamelCase    :		Optional[Any]			   =       1e-4







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



 '''simple docstring'''


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

       """simple docstring"""






       lowerCamelCase_							=   'nvidia/megatron-bert-uncased-345m'
       if "MYDIR" in os.environ:
             lowerCamelCase_							=   os.path.join(os.environ['MYDIR']						,							A_  )
       lowerCamelCase_							=   MegatronBertModel.from_pretrained(A_  )
       model.to(A_  )
       model.half()
       lowerCamelCase_							=   _long_tensor([[101, 7110, 1005, 1056, 2023, 11333, 17413, 1029, 102]]  )
       with torch.no_grad():
             lowerCamelCase_							=   model(A_  )[0]
       lowerCamelCase_							=   torch.Size((1, 9, 1024)  )
       self.assertEqual(output.shape						,							A_  )

       lowerCamelCase_							=   [-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  ):
                   lowerCamelCase_							=   output[0, ii, jj]
                   lowerCamelCase_							=   expected[3 * ii + jj]
                   lowerCamelCase_							=   'ii={} jj={} a={} b={}'.format(A_						,							A_						,							A_						,							A_  )
                   self.assertTrue(math.isclose(A_						,							A_						,							rel_tol=A_						,							abs_tol=A_  )						,							msg=A_  )









import math

def       _SCREAMING_SNAKE_CASE      (							lowercase     :	int					):


      '''simple docstring'''



      lowerCamelCase_							=   0
      lowerCamelCase_							=   0
      while num > 0:
            lowerCamelCase_							=   num % 8
            lowerCamelCase_							=   octal + (remainder * math.floor(math.pow(10     ,							lowercase					)					))
            counter += 1
            lowerCamelCase_							=   math.floor(num / 8					)  # basically /= 8 without remainder if any
            # This formatting removes trailing '.0' from `octal`.
      return f"""0o{int(lowercase					)}"""

def       _SCREAMING_SNAKE_CASE      (							):


      '''simple docstring'''



      print('\n2 in octal is:'					)
      print(decimal_to_octal(2					)					)  # = 2
      print('\n8 in octal is:'					)
      print(decimal_to_octal(8					)					)  # = 10
      print('\n65 in octal is:'					)
      print(decimal_to_octal(65					)					)  # = 101
      print('\n216 in octal is:'					)
      print(decimal_to_octal(2_16					)					)  # = 330
      print('\n512 in octal is:'					)
      print(decimal_to_octal(5_12					)					)  # = 1000
      print('\n'					)


if __name__ == "__main__":
 main()



