





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







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







       def 						__lowerCamelCase   (				self					:     Dict      )       ->List[str]:
          lowerCamelCase__       :						List[Any]        =							ids_tensor([self.batch_size, self.seq_length]		, self.vocab_size      )

          lowerCamelCase__       :						Union[str, Any]        =							None
          if self.use_input_mask:
             lowerCamelCase__       :						Tuple        =							random_attention_mask([self.batch_size, self.seq_length]      )

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

          lowerCamelCase__       :						List[Any]        =							None
          lowerCamelCase__       :						Optional[Any]        =							None
          lowerCamelCase__       :						Tuple        =							None
          if self.use_labels:
             lowerCamelCase__       :						Optional[Any]        =							ids_tensor([self.batch_size]		, self.type_sequence_label_size      )
             lowerCamelCase__       :						Union[str, Any]        =							ids_tensor([self.batch_size, self.seq_length]		, self.num_labels      )
             lowerCamelCase__       :						str        =							ids_tensor([self.batch_size]		, self.num_choices      )

          lowerCamelCase__       :						Tuple        =							self.get_config()

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







       def 						__lowerCamelCase   (				self					:     int      )       ->Dict:
          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 						__lowerCamelCase   (				self					:     List[str]		, A					:     Tuple		, A					:     Tuple		, A					:     str		, A					:     Any		, A					:     Union[str, Any]		, A					:     Dict		, A					:     List[str]      )       ->List[Any]:
          lowerCamelCase__       :						Tuple        =							MegatronBertModel(config=__UpperCAmelCase      )
          model.to(__UpperCAmelCase      )
          model.eval()
          lowerCamelCase__       :						str        =							model(__UpperCAmelCase		, attention_mask=__UpperCAmelCase		, token_type_ids=__UpperCAmelCase      )
          lowerCamelCase__       :						List[str]        =							model(__UpperCAmelCase		, token_type_ids=__UpperCAmelCase      )
          lowerCamelCase__       :						Union[str, Any]        =							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 						__lowerCamelCase   (				self					:     Dict		, A					:     List[Any]		, A					:     List[Any]		, A					:     Tuple		, A					:     List[Any]		, A					:     List[str]		, A					:     Optional[int]		, A					:     List[str]      )       ->str:
          lowerCamelCase__       :						Union[str, Any]        =							MegatronBertForMaskedLM(config=__UpperCAmelCase      )
          model.to(__UpperCAmelCase      )
          model.eval()
          lowerCamelCase__       :						int        =							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 						__lowerCamelCase   (				self					:     List[Any]		, A					:     str		, A					:     Union[str, Any]		, A					:     Tuple		, A					:     Optional[Any]		, A					:     Dict		, A					:     Dict		, A					:     List[str]      )       ->Optional[int]:
          lowerCamelCase__       :						int        =							MegatronBertForCausalLM(config=__UpperCAmelCase      )
          model.to(__UpperCAmelCase      )
          model.eval()
          lowerCamelCase__       :						Tuple        =							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 						__lowerCamelCase   (				self					:     Optional[Any]		, A					:     str		, A					:     int		, A					:     int		, A					:     int		, A					:     str		, A					:     int		, A					:     Optional[int]      )       ->Dict:
          lowerCamelCase__       :						List[Any]        =							MegatronBertForNextSentencePrediction(config=__UpperCAmelCase      )
          model.to(__UpperCAmelCase      )
          model.eval()
          lowerCamelCase__       :						List[str]        =							model(
              __UpperCAmelCase		, attention_mask=__UpperCAmelCase		, token_type_ids=__UpperCAmelCase		, labels=__UpperCAmelCase		, )
          self.parent.assertEqual(result.logits.shape		, (self.batch_size, 2)      )







       def 						__lowerCamelCase   (				self					:     int		, A					:     Optional[Any]		, A					:     str		, A					:     List[str]		, A					:     List[Any]		, A					:     Optional[int]		, A					:     Optional[int]		, A					:     str      )       ->List[str]:
          lowerCamelCase__       :						str        =							MegatronBertForPreTraining(config=__UpperCAmelCase      )
          model.to(__UpperCAmelCase      )
          model.eval()
          lowerCamelCase__       :						Tuple        =							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 						__lowerCamelCase   (				self					:     Optional[int]		, A					:     int		, A					:     Dict		, A					:     Optional[Any]		, A					:     str		, A					:     Dict		, A					:     Any		, A					:     List[Any]      )       ->Union[str, Any]:
          lowerCamelCase__       :						List[Any]        =							MegatronBertForQuestionAnswering(config=__UpperCAmelCase      )
          model.to(__UpperCAmelCase      )
          model.eval()
          lowerCamelCase__       :						Union[str, Any]        =							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 						__lowerCamelCase   (				self					:     Union[str, Any]		, A					:     Tuple		, A					:     List[str]		, A					:     Any		, A					:     Tuple		, A					:     Any		, A					:     int		, A					:     List[str]      )       ->Optional[Any]:
          lowerCamelCase__       :						Dict        =							self.num_labels
          lowerCamelCase__       :						Any        =							MegatronBertForSequenceClassification(__UpperCAmelCase      )
          model.to(__UpperCAmelCase      )
          model.eval()
          lowerCamelCase__       :						List[Any]        =							model(__UpperCAmelCase		, attention_mask=__UpperCAmelCase		, token_type_ids=__UpperCAmelCase		, labels=__UpperCAmelCase      )
          self.parent.assertEqual(result.logits.shape		, (self.batch_size, self.num_labels)      )







       def 						__lowerCamelCase   (				self					:     List[Any]		, A					:     List[Any]		, A					:     int		, A					:     int		, A					:     Any		, A					:     Union[str, Any]		, A					:     List[Any]		, A					:     Tuple      )       ->Dict:
          lowerCamelCase__       :						Any        =							self.num_labels
          lowerCamelCase__       :						Tuple        =							MegatronBertForTokenClassification(config=__UpperCAmelCase      )
          model.to(__UpperCAmelCase      )
          model.eval()
          lowerCamelCase__       :						Tuple        =							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 						__lowerCamelCase   (				self					:     List[Any]		, A					:     List[str]		, A					:     List[Any]		, A					:     List[str]		, A					:     Optional[int]		, A					:     Optional[int]		, A					:     Optional[Any]		, A					:     str      )       ->Optional[int]:
          lowerCamelCase__       :						int        =							self.num_choices
          lowerCamelCase__       :						Tuple        =							MegatronBertForMultipleChoice(config=__UpperCAmelCase      )
          model.to(__UpperCAmelCase      )
          model.eval()
          lowerCamelCase__       :						Optional[int]        =							input_ids.unsqueeze(1      ).expand(-1		, self.num_choices		, -1      ).contiguous()
          lowerCamelCase__       :						Tuple        =							token_type_ids.unsqueeze(1      ).expand(-1		, self.num_choices		, -1      ).contiguous()
          lowerCamelCase__       :						Union[str, Any]        =							input_mask.unsqueeze(1      ).expand(-1		, self.num_choices		, -1      ).contiguous()
          lowerCamelCase__       :						List[Any]        =							model(
              __UpperCAmelCase		, attention_mask=__UpperCAmelCase		, token_type_ids=__UpperCAmelCase		, labels=__UpperCAmelCase		, )
          self.parent.assertEqual(result.logits.shape		, (self.batch_size, self.num_choices)      )







       def 						__lowerCamelCase   (				self					:     Dict      )       ->Any:
          lowerCamelCase__       :						List[str]        =							self.prepare_config_and_inputs()
          (
              (
              lowerCamelCase__					
          )       ,							(
              lowerCamelCase__					
          )       ,							(
              lowerCamelCase__					
          )       ,							(
              lowerCamelCase__					
          )       ,							(
              lowerCamelCase__					
          )       ,							(
              lowerCamelCase__					
          )       ,							(
              lowerCamelCase__					
          )       ,												
          )       :						str        =							config_and_inputs
          lowerCamelCase__       :						Union[str, Any]        =							{'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
          return config, inputs_dict


@require_torch
class        __SCREAMING_SNAKE_CASE      (			lowerCAmelCase_					,lowerCAmelCase_					,unittest.TestCase ):
       _UpperCAmelCase				:  Union[str, Any]			   =   (
           (
               MegatronBertModel,
               MegatronBertForMaskedLM,
               MegatronBertForCausalLM,
               MegatronBertForMultipleChoice,
               MegatronBertForNextSentencePrediction,
               MegatronBertForPreTraining,
               MegatronBertForQuestionAnswering,
               MegatronBertForSequenceClassification,
               MegatronBertForTokenClassification,
           )
           if is_torch_available()
           else ()
       )
       _UpperCAmelCase				:  int			   =   (
           {
               "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				:  Dict			   =   True
       # test_resize_embeddings = False
       _UpperCAmelCase				:  Tuple			   =   False







       def 						__lowerCamelCase   (				self					:     Union[str, Any]		, A					:     List[str]		, A					:     int		, A					:     List[str]=False      )       ->Union[str, Any]:
          lowerCamelCase__       :						Any        =							super()._prepare_for_class(__UpperCAmelCase		, __UpperCAmelCase		, return_labels=__UpperCAmelCase      )

          if return_labels:
             if model_class in get_values(__UpperCAmelCase      ):
                lowerCamelCase__       :						Any        =							torch.zeros(
                    (self.model_tester.batch_size, self.model_tester.seq_length)		, dtype=torch.long		, device=__UpperCAmelCase      )
                lowerCamelCase__       :						Optional[Any]        =							torch.zeros(
                    self.model_tester.batch_size		, dtype=torch.long		, device=__UpperCAmelCase      )
          return inputs_dict







       def 						__lowerCamelCase   (				self					:     Optional[Any]      )       ->Any:
          lowerCamelCase__       :						Optional[Any]        =							MegatronBertModelTester(self      )
          lowerCamelCase__       :						Union[str, Any]        =							ConfigTester(self		, config_class=__UpperCAmelCase		, hidden_size=3_7      )







       def 						__lowerCamelCase   (				self					:     List[Any]      )       ->Dict:
          self.config_tester.run_common_tests()







       def 						__lowerCamelCase   (				self					:     List[str]      )       ->Tuple:
          lowerCamelCase__       :						Any        =							self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_megatron_bert_model(*__UpperCAmelCase      )







       def 						__lowerCamelCase   (				self					:     Union[str, Any]      )       ->Optional[Any]:
          lowerCamelCase__       :						List[str]        =							self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_megatron_bert_for_masked_lm(*__UpperCAmelCase      )







       def 						__lowerCamelCase   (				self					:     Optional[int]      )       ->Optional[int]:
          lowerCamelCase__       :						Union[str, Any]        =							self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*__UpperCAmelCase      )







       def 						__lowerCamelCase   (				self					:     List[str]      )       ->Optional[int]:
          lowerCamelCase__       :						int        =							self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*__UpperCAmelCase      )







       def 						__lowerCamelCase   (				self					:     Union[str, Any]      )       ->Optional[Any]:
          lowerCamelCase__       :						Any        =							self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_megatron_bert_for_pretraining(*__UpperCAmelCase      )







       def 						__lowerCamelCase   (				self					:     Optional[Any]      )       ->Union[str, Any]:
          lowerCamelCase__       :						Any        =							self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_megatron_bert_for_question_answering(*__UpperCAmelCase      )







       def 						__lowerCamelCase   (				self					:     Tuple      )       ->Optional[Any]:
          lowerCamelCase__       :						Optional[Any]        =							self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*__UpperCAmelCase      )







       def 						__lowerCamelCase   (				self					:     Union[str, Any]      )       ->List[str]:
          lowerCamelCase__       :						Dict        =							self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_megatron_bert_for_token_classification(*__UpperCAmelCase      )



def 	_a						(   UpperCAmelCase					)    ->    int:







   """simple docstring"""
   return torch.tensor(
       UpperCAmelCase      ,  dtype=torch.long      ,  device=UpperCAmelCase      ,  )


_A :   List[Any]  			=  1e-4




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







       @slow
       @unittest.skip('''Model is not available.'''      )
       def 						__lowerCamelCase   (				self					:     Union[str, Any]      )       ->Dict:
          lowerCamelCase__       :						Optional[Any]        =							'''nvidia/megatron-bert-uncased-345m'''
          if "MYDIR" in os.environ:
             lowerCamelCase__       :						int        =							os.path.join(os.environ['''MYDIR''']		, __UpperCAmelCase      )
          lowerCamelCase__       :						Dict        =							MegatronBertModel.from_pretrained(__UpperCAmelCase      )
          model.to(__UpperCAmelCase      )
          model.half()
          lowerCamelCase__       :						Optional[Any]        =							_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():
             lowerCamelCase__       :						Optional[int]        =							model(__UpperCAmelCase      )[0]
          lowerCamelCase__       :						List[Any]        =							torch.Size((1, 9, 1_0_2_4)      )
          self.assertEqual(output.shape		, __UpperCAmelCase      )

          lowerCamelCase__       :						str        =							[-0.60_40, -0.25_17, -0.10_25, 0.34_20, -0.67_58, -0.00_17, -0.10_89, -0.19_90, 0.57_28]
          for ii in range(3      ):
             for jj in range(3      ):
                lowerCamelCase__       :						List[Any]        =							output[0, ii, jj]
                lowerCamelCase__       :						Union[str, Any]        =							expected[3 * ii + jj]
                lowerCamelCase__       :						int        =							'''ii={} jj={} a={} b={}'''.format(__UpperCAmelCase		, __UpperCAmelCase		, __UpperCAmelCase		, __UpperCAmelCase      )
                self.assertTrue(math.isclose(__UpperCAmelCase		, __UpperCAmelCase		, rel_tol=__UpperCAmelCase		, abs_tol=__UpperCAmelCase      )		, msg=__UpperCAmelCase      )










def 	_a						(   UpperCAmelCase					)    ->    int:







   """simple docstring"""
   if not isinstance(UpperCAmelCase      ,  UpperCAmelCase					) or number < 0:
      raise ValueError('''Input must be a non-negative integer'''					)

   lowerCamelCase__       :						List[str]        =							0
   while number:
      # This way we arrive at next set bit (next 1) instead of looping
      # through each bit and checking for 1s hence the
      # loop won't run 32 times it will only run the number of `1` times
      number &= number - 1
      count += 1
   return count


if __name__ == "__main__":
       import doctest

       doctest.testmod()



