




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






     def __init__(					self   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE=13   ,_SCREAMING_SNAKE_CASE=7   ,_SCREAMING_SNAKE_CASE=True   ,_SCREAMING_SNAKE_CASE=True   ,_SCREAMING_SNAKE_CASE=True   ,_SCREAMING_SNAKE_CASE=True   ,_SCREAMING_SNAKE_CASE=99   ,_SCREAMING_SNAKE_CASE=64   ,_SCREAMING_SNAKE_CASE=32   ,_SCREAMING_SNAKE_CASE=5   ,_SCREAMING_SNAKE_CASE=4   ,_SCREAMING_SNAKE_CASE=37   ,_SCREAMING_SNAKE_CASE="gelu"   ,_SCREAMING_SNAKE_CASE=0.1   ,_SCREAMING_SNAKE_CASE=0.1   ,_SCREAMING_SNAKE_CASE=512   ,_SCREAMING_SNAKE_CASE=16   ,_SCREAMING_SNAKE_CASE=2   ,_SCREAMING_SNAKE_CASE=0.0_2   ,_SCREAMING_SNAKE_CASE=3   ,_SCREAMING_SNAKE_CASE=4   ,_SCREAMING_SNAKE_CASE=None   ,)				->     Union[str, Any]:
          _snake_case				= parent
          _snake_case				= batch_size
          _snake_case				= seq_length
          _snake_case				= is_training
          _snake_case				= use_input_mask
          _snake_case				= use_token_type_ids
          _snake_case				= use_labels
          _snake_case				= vocab_size
          _snake_case				= hidden_size
          _snake_case				= embedding_size
          _snake_case				= num_hidden_layers
          _snake_case				= num_attention_heads
          _snake_case				= intermediate_size
          _snake_case				= hidden_act
          _snake_case				= hidden_dropout_prob
          _snake_case				= attention_probs_dropout_prob
          _snake_case				= max_position_embeddings
          _snake_case				= type_vocab_size
          _snake_case				= type_sequence_label_size
          _snake_case				= initializer_range
          _snake_case				= num_labels
          _snake_case				= num_choices
          _snake_case				= scope






     def 							_lowercase      (					self       )				->     str:
          _snake_case				= ids_tensor([self.batch_size, self.seq_length]   ,self.vocab_size       )

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

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

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

          _snake_case				= self.get_config()

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






     def 							_lowercase      (					self       )				->     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=_SCREAMING_SNAKE_CASE   ,initializer_range=self.initializer_range   ,)






     def 							_lowercase      (					self   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE       )				->     Tuple:
          _snake_case				= MegatronBertModel(config=_SCREAMING_SNAKE_CASE       )
          model.to(_SCREAMING_SNAKE_CASE       )
          model.eval()
          _snake_case				= model(_SCREAMING_SNAKE_CASE   ,attention_mask=_SCREAMING_SNAKE_CASE   ,token_type_ids=_SCREAMING_SNAKE_CASE       )
          _snake_case				= model(_SCREAMING_SNAKE_CASE   ,token_type_ids=_SCREAMING_SNAKE_CASE       )
          _snake_case				= model(_SCREAMING_SNAKE_CASE       )

          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 							_lowercase      (					self   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE       )				->     Tuple:
          _snake_case				= MegatronBertForMaskedLM(config=_SCREAMING_SNAKE_CASE       )
          model.to(_SCREAMING_SNAKE_CASE       )
          model.eval()
          _snake_case				= model(_SCREAMING_SNAKE_CASE   ,attention_mask=_SCREAMING_SNAKE_CASE   ,token_type_ids=_SCREAMING_SNAKE_CASE   ,labels=_SCREAMING_SNAKE_CASE       )
          self.parent.assertEqual(result.logits.shape   ,(self.batch_size, self.seq_length, self.vocab_size)       )






     def 							_lowercase      (					self   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE       )				->     Optional[Any]:
          _snake_case				= MegatronBertForCausalLM(config=_SCREAMING_SNAKE_CASE       )
          model.to(_SCREAMING_SNAKE_CASE       )
          model.eval()
          _snake_case				= model(_SCREAMING_SNAKE_CASE   ,attention_mask=_SCREAMING_SNAKE_CASE   ,token_type_ids=_SCREAMING_SNAKE_CASE   ,labels=_SCREAMING_SNAKE_CASE       )
          self.parent.assertEqual(result.logits.shape   ,(self.batch_size, self.seq_length, self.vocab_size)       )






     def 							_lowercase      (					self   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE       )				->     Tuple:
          _snake_case				= MegatronBertForNextSentencePrediction(config=_SCREAMING_SNAKE_CASE       )
          model.to(_SCREAMING_SNAKE_CASE       )
          model.eval()
          _snake_case				= model(
              _SCREAMING_SNAKE_CASE   ,attention_mask=_SCREAMING_SNAKE_CASE   ,token_type_ids=_SCREAMING_SNAKE_CASE   ,labels=_SCREAMING_SNAKE_CASE   ,)
          self.parent.assertEqual(result.logits.shape   ,(self.batch_size, 2)       )






     def 							_lowercase      (					self   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE       )				->     Tuple:
          _snake_case				= MegatronBertForPreTraining(config=_SCREAMING_SNAKE_CASE       )
          model.to(_SCREAMING_SNAKE_CASE       )
          model.eval()
          _snake_case				= model(
              _SCREAMING_SNAKE_CASE   ,attention_mask=_SCREAMING_SNAKE_CASE   ,token_type_ids=_SCREAMING_SNAKE_CASE   ,labels=_SCREAMING_SNAKE_CASE   ,next_sentence_label=_SCREAMING_SNAKE_CASE   ,)
          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 							_lowercase      (					self   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE       )				->     int:
          _snake_case				= MegatronBertForQuestionAnswering(config=_SCREAMING_SNAKE_CASE       )
          model.to(_SCREAMING_SNAKE_CASE       )
          model.eval()
          _snake_case				= model(
              _SCREAMING_SNAKE_CASE   ,attention_mask=_SCREAMING_SNAKE_CASE   ,token_type_ids=_SCREAMING_SNAKE_CASE   ,start_positions=_SCREAMING_SNAKE_CASE   ,end_positions=_SCREAMING_SNAKE_CASE   ,)
          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 							_lowercase      (					self   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE       )				->     Tuple:
          _snake_case				= self.num_labels
          _snake_case				= MegatronBertForSequenceClassification(_SCREAMING_SNAKE_CASE       )
          model.to(_SCREAMING_SNAKE_CASE       )
          model.eval()
          _snake_case				= model(_SCREAMING_SNAKE_CASE   ,attention_mask=_SCREAMING_SNAKE_CASE   ,token_type_ids=_SCREAMING_SNAKE_CASE   ,labels=_SCREAMING_SNAKE_CASE       )
          self.parent.assertEqual(result.logits.shape   ,(self.batch_size, self.num_labels)       )






     def 							_lowercase      (					self   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE       )				->     Optional[Any]:
          _snake_case				= self.num_labels
          _snake_case				= MegatronBertForTokenClassification(config=_SCREAMING_SNAKE_CASE       )
          model.to(_SCREAMING_SNAKE_CASE       )
          model.eval()
          _snake_case				= model(_SCREAMING_SNAKE_CASE   ,attention_mask=_SCREAMING_SNAKE_CASE   ,token_type_ids=_SCREAMING_SNAKE_CASE   ,labels=_SCREAMING_SNAKE_CASE       )
          self.parent.assertEqual(result.logits.shape   ,(self.batch_size, self.seq_length, self.num_labels)       )






     def 							_lowercase      (					self   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE       )				->     Any:
          _snake_case				= self.num_choices
          _snake_case				= MegatronBertForMultipleChoice(config=_SCREAMING_SNAKE_CASE       )
          model.to(_SCREAMING_SNAKE_CASE       )
          model.eval()
          _snake_case				= input_ids.unsqueeze(1       ).expand(-1   ,self.num_choices   ,-1       ).contiguous()
          _snake_case				= token_type_ids.unsqueeze(1       ).expand(-1   ,self.num_choices   ,-1       ).contiguous()
          _snake_case				= input_mask.unsqueeze(1       ).expand(-1   ,self.num_choices   ,-1       ).contiguous()
          _snake_case				= model(
              _SCREAMING_SNAKE_CASE   ,attention_mask=_SCREAMING_SNAKE_CASE   ,token_type_ids=_SCREAMING_SNAKE_CASE   ,labels=_SCREAMING_SNAKE_CASE   ,)
          self.parent.assertEqual(result.logits.shape   ,(self.batch_size, self.num_choices)       )






     def 							_lowercase      (					self       )				->     Dict:
          _snake_case				= self.prepare_config_and_inputs()
          (
              (
              _snake_case
          )				,   (
              _snake_case
          )				,   (
              _snake_case
          )				,   (
              _snake_case
          )				,   (
              _snake_case
          )				,   (
              _snake_case
          )				,   (
              _snake_case
          )				,   
          )				= config_and_inputs
          _snake_case				= {"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
          return config, inputs_dict







@require_torch
class   _a							(		__lowerCAmelCase   ,   __lowerCAmelCase   ,   unittest.TestCase				):
     SCREAMING_SNAKE_CASE_				:       List[str]       							=		(
         (
             MegatronBertModel,
             MegatronBertForMaskedLM,
             MegatronBertForCausalLM,
             MegatronBertForMultipleChoice,
             MegatronBertForNextSentencePrediction,
             MegatronBertForPreTraining,
             MegatronBertForQuestionAnswering,
             MegatronBertForSequenceClassification,
             MegatronBertForTokenClassification,
         )
         if is_torch_available()
         else ()
     )
     SCREAMING_SNAKE_CASE_				:       Optional[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 {}
     )
     SCREAMING_SNAKE_CASE_				:       List[Any]       							=		True
     # test_resize_embeddings = False
     SCREAMING_SNAKE_CASE_				:       List[str]       							=		False






     def 							_lowercase      (					self   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE=False       )				->     str:
          _snake_case				= super()._prepare_for_class(_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,return_labels=_SCREAMING_SNAKE_CASE       )

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






     def 							_lowercase      (					self       )				->     Dict:
          _snake_case				= MegatronBertModelTester(self       )
          _snake_case				= ConfigTester(self   ,config_class=_SCREAMING_SNAKE_CASE   ,hidden_size=37       )






     def 							_lowercase      (					self       )				->     Optional[Any]:
          self.config_tester.run_common_tests()






     def 							_lowercase      (					self       )				->     Optional[Any]:
          _snake_case				= self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_megatron_bert_model(*_SCREAMING_SNAKE_CASE       )






     def 							_lowercase      (					self       )				->     Optional[int]:
          _snake_case				= self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_megatron_bert_for_masked_lm(*_SCREAMING_SNAKE_CASE       )






     def 							_lowercase      (					self       )				->     List[str]:
          _snake_case				= self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*_SCREAMING_SNAKE_CASE       )






     def 							_lowercase      (					self       )				->     List[str]:
          _snake_case				= self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*_SCREAMING_SNAKE_CASE       )






     def 							_lowercase      (					self       )				->     Optional[Any]:
          _snake_case				= self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_megatron_bert_for_pretraining(*_SCREAMING_SNAKE_CASE       )






     def 							_lowercase      (					self       )				->     str:
          _snake_case				= self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_megatron_bert_for_question_answering(*_SCREAMING_SNAKE_CASE       )






     def 							_lowercase      (					self       )				->     Any:
          _snake_case				= self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*_SCREAMING_SNAKE_CASE       )






     def 							_lowercase      (					self       )				->     int:
          _snake_case				= self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_megatron_bert_for_token_classification(*_SCREAMING_SNAKE_CASE       )







def       __a      (  _UpperCamelCase:       Dict       )		->					List[Any]:

     """simple docstring"""




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


UpperCamelCase_		:     List[Any]          =		1E-4





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






     @slow
     @unittest.skip("Model is not available."       )
     def 							_lowercase      (					self       )				->     int:
          _snake_case				= "nvidia/megatron-bert-uncased-345m"
          if "MYDIR" in os.environ:
               _snake_case				= os.path.join(os.environ["MYDIR"]   ,_SCREAMING_SNAKE_CASE       )
          _snake_case				= MegatronBertModel.from_pretrained(_SCREAMING_SNAKE_CASE       )
          model.to(_SCREAMING_SNAKE_CASE       )
          model.half()
          _snake_case				= _long_tensor([[101, 7_110, 1_005, 1_056, 2_023, 11_333, 17_413, 1_029, 102]]       )
          with torch.no_grad():
               _snake_case				= model(_SCREAMING_SNAKE_CASE       )[0]
          _snake_case				= torch.Size((1, 9, 1_024)       )
          self.assertEqual(output.shape   ,_SCREAMING_SNAKE_CASE       )

          _snake_case				= [-0.6_0_4_0, -0.2_5_1_7, -0.1_0_2_5, 0.3_4_2_0, -0.6_7_5_8, -0.0_0_1_7, -0.1_0_8_9, -0.1_9_9_0, 0.5_7_2_8]
          for ii in range(3       ):
               for jj in range(3       ):
                    _snake_case				= output[0, ii, jj]
                    _snake_case				= expected[3 * ii + jj]
                    _snake_case				= "ii={} jj={} a={} b={}".format(_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE       )
                    self.assertTrue(math.isclose(_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,rel_tol=_SCREAMING_SNAKE_CASE   ,abs_tol=_SCREAMING_SNAKE_CASE       )   ,msg=_SCREAMING_SNAKE_CASE       )








'''simple docstring'''





from math import ceil





def       __a      (  _UpperCamelCase:       int = 1_001       )		->					int:

     """simple docstring"""




     _snake_case				= 1

     for i in range(1   ,		int(ceil(n / 2.0       )       )       ):
          _snake_case				= 2 * i + 1
          _snake_case				= 2 * i
          _snake_case				= total + 4 * odd**2 - 6 * even

     return total


if __name__ == "__main__":
   import sys

   if len(sys.argv) == 1:
      print(solution())
   else:
      try:
         UpperCamelCase_		:     int          =		int(sys.argv[1])
         print(solution(n))
      except ValueError:
         print('''Invalid entry - please enter a number''')


