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





   """simple docstring"""




   def __init__(	self					:  List[Any]		,  _lowercase					:  List[Any]		,  _lowercase					:  Optional[Any]=13		,  _lowercase					:  Optional[int]=7		,  _lowercase					:  Tuple=True		,  _lowercase					:  Union[str, Any]=True		,  _lowercase					:  List[Any]=True		,  _lowercase					:  str=True		,  _lowercase					:  Tuple=99		,  _lowercase					:  Union[str, Any]=64		,  _lowercase					:  Tuple=32		,  _lowercase					:  Optional[int]=5		,  _lowercase					:  int=4		,  _lowercase					:  Optional[int]=37		,  _lowercase					:  str="gelu"		,  _lowercase					:  List[str]=0.1		,  _lowercase					:  str=0.1		,  _lowercase					:  List[Any]=512		,  _lowercase					:  str=16		,  _lowercase					:  int=2		,  _lowercase					:  List[str]=0.02		,  _lowercase					:  List[Any]=3		,  _lowercase					:  Optional[int]=4		,  _lowercase					:  Tuple=None		,  ):



    """simple docstring"""



    _UpperCamelCase:     Any							= parent
    _UpperCamelCase:     str							= batch_size
    _UpperCamelCase:     Optional[Any]							= seq_length
    _UpperCamelCase:     str							= is_training
    _UpperCamelCase:     Dict							= use_input_mask
    _UpperCamelCase:     int							= use_token_type_ids
    _UpperCamelCase:     List[Any]							= use_labels
    _UpperCamelCase:     Tuple							= vocab_size
    _UpperCamelCase:     Any							= hidden_size
    _UpperCamelCase:     Any							= embedding_size
    _UpperCamelCase:     Dict							= num_hidden_layers
    _UpperCamelCase:     Optional[Any]							= num_attention_heads
    _UpperCamelCase:     Optional[Any]							= intermediate_size
    _UpperCamelCase:     List[Any]							= hidden_act
    _UpperCamelCase:     List[Any]							= hidden_dropout_prob
    _UpperCamelCase:     List[str]							= attention_probs_dropout_prob
    _UpperCamelCase:     str							= max_position_embeddings
    _UpperCamelCase:     int							= type_vocab_size
    _UpperCamelCase:     Dict							= type_sequence_label_size
    _UpperCamelCase:     Optional[Any]							= initializer_range
    _UpperCamelCase:     List[str]							= num_labels
    _UpperCamelCase:     Dict							= num_choices
    _UpperCamelCase:     Optional[Any]							= scope



   def   lowerCAmelCase						(	self					:  int							):



    """simple docstring"""



    _UpperCamelCase:     List[Any]							= ids_tensor([self.batch_size, self.seq_length]		,  self.vocab_size							)

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

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

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

    _UpperCamelCase:     str							= self.get_config()

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



   def   lowerCAmelCase						(	self					:  int							):



    """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   lowerCAmelCase						(	self					:  List[str]		,  _lowercase					:  Optional[int]		,  _lowercase					:  List[Any]		,  _lowercase					:  List[Any]		,  _lowercase					:  Any		,  _lowercase					:  List[Any]		,  _lowercase					:  Any		,  _lowercase					:  Union[str, Any]							):



    """simple docstring"""



    _UpperCamelCase:     str							= MegatronBertModel(config=__A							)
    model.to(__A							)
    model.eval()
    _UpperCamelCase:     Any							= model(__A		,  attention_mask=__A		,  token_type_ids=__A							)
    _UpperCamelCase:     Union[str, Any]							= model(__A		,  token_type_ids=__A							)
    _UpperCamelCase:     Union[str, Any]							= 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   lowerCAmelCase						(	self					:  Any		,  _lowercase					:  Union[str, Any]		,  _lowercase					:  List[str]		,  _lowercase					:  Union[str, Any]		,  _lowercase					:  Any		,  _lowercase					:  Tuple		,  _lowercase					:  int		,  _lowercase					:  int							):



    """simple docstring"""



    _UpperCamelCase:     Any							= MegatronBertForMaskedLM(config=__A							)
    model.to(__A							)
    model.eval()
    _UpperCamelCase:     str							= 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   lowerCAmelCase						(	self					:  str		,  _lowercase					:  Tuple		,  _lowercase					:  List[Any]		,  _lowercase					:  Any		,  _lowercase					:  Optional[Any]		,  _lowercase					:  int		,  _lowercase					:  Union[str, Any]		,  _lowercase					:  Union[str, Any]							):



    """simple docstring"""



    _UpperCamelCase:     Optional[Any]							= MegatronBertForCausalLM(config=__A							)
    model.to(__A							)
    model.eval()
    _UpperCamelCase:     Optional[Any]							= 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   lowerCAmelCase						(	self					:  Optional[int]		,  _lowercase					:  Optional[Any]		,  _lowercase					:  Dict		,  _lowercase					:  List[str]		,  _lowercase					:  List[str]		,  _lowercase					:  List[Any]		,  _lowercase					:  Optional[int]		,  _lowercase					:  str							):



    """simple docstring"""



    _UpperCamelCase:     List[Any]							= MegatronBertForNextSentencePrediction(config=__A							)
    model.to(__A							)
    model.eval()
    _UpperCamelCase:     int							= model(
        __A		,  attention_mask=__A		,  token_type_ids=__A		,  labels=__A		,  )
    self.parent.assertEqual(result.logits.shape		,  (self.batch_size, 2)							)



   def   lowerCAmelCase						(	self					:  Optional[int]		,  _lowercase					:  Optional[Any]		,  _lowercase					:  str		,  _lowercase					:  int		,  _lowercase					:  List[str]		,  _lowercase					:  Optional[Any]		,  _lowercase					:  List[Any]		,  _lowercase					:  List[str]							):



    """simple docstring"""



    _UpperCamelCase:     Any							= MegatronBertForPreTraining(config=__A							)
    model.to(__A							)
    model.eval()
    _UpperCamelCase:     int							= 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   lowerCAmelCase						(	self					:  Dict		,  _lowercase					:  List[Any]		,  _lowercase					:  int		,  _lowercase					:  Any		,  _lowercase					:  Optional[int]		,  _lowercase					:  int		,  _lowercase					:  Union[str, Any]		,  _lowercase					:  int							):



    """simple docstring"""



    _UpperCamelCase:     Dict							= MegatronBertForQuestionAnswering(config=__A							)
    model.to(__A							)
    model.eval()
    _UpperCamelCase:     Any							= 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   lowerCAmelCase						(	self					:  str		,  _lowercase					:  Union[str, Any]		,  _lowercase					:  Any		,  _lowercase					:  Optional[Any]		,  _lowercase					:  Union[str, Any]		,  _lowercase					:  str		,  _lowercase					:  int		,  _lowercase					:  Union[str, Any]							):



    """simple docstring"""



    _UpperCamelCase:     Optional[Any]							= self.num_labels
    _UpperCamelCase:     str							= MegatronBertForSequenceClassification(__A							)
    model.to(__A							)
    model.eval()
    _UpperCamelCase:     List[Any]							= model(__A		,  attention_mask=__A		,  token_type_ids=__A		,  labels=__A							)
    self.parent.assertEqual(result.logits.shape		,  (self.batch_size, self.num_labels)							)



   def   lowerCAmelCase						(	self					:  List[Any]		,  _lowercase					:  Tuple		,  _lowercase					:  Union[str, Any]		,  _lowercase					:  Optional[Any]		,  _lowercase					:  List[Any]		,  _lowercase					:  Union[str, Any]		,  _lowercase					:  List[str]		,  _lowercase					:  List[str]							):



    """simple docstring"""



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



    """simple docstring"""



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



   def   lowerCAmelCase						(	self					:  Dict							):



    """simple docstring"""



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




@require_torch
class  __magic_name__  (					UpperCamelCase__      ,	UpperCamelCase__      ,	unittest.TestCase    ):





   """simple docstring"""




   lowerCAmelCase      :    Optional[Any]  			=   (
       (
           MegatronBertModel,
           MegatronBertForMaskedLM,
           MegatronBertForCausalLM,
           MegatronBertForMultipleChoice,
           MegatronBertForNextSentencePrediction,
           MegatronBertForPreTraining,
           MegatronBertForQuestionAnswering,
           MegatronBertForSequenceClassification,
           MegatronBertForTokenClassification,
       )
       if is_torch_available()
       else ()
   )
   lowerCAmelCase      :    List[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      :    Tuple  			=   True
   # test_resize_embeddings = False
   lowerCAmelCase      :    Optional[Any]  			=   False



   def   lowerCAmelCase						(	self					:  Optional[Any]		,  _lowercase					:  Union[str, Any]		,  _lowercase					:  List[Any]		,  _lowercase					:  Optional[int]=False							):



    """simple docstring"""



    _UpperCamelCase:     Optional[int]							= super()._prepare_for_class(__A		,  __A		,  return_labels=__A							)

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



   def   lowerCAmelCase						(	self					:  Optional[int]							):



    """simple docstring"""



    _UpperCamelCase:     Any							= MegatronBertModelTester(self							)
    _UpperCamelCase:     Optional[int]							= ConfigTester(self		,  config_class=__A		,  hidden_size=37							)



   def   lowerCAmelCase						(	self					:  Optional[Any]							):



    """simple docstring"""



    self.config_tester.run_common_tests()



   def   lowerCAmelCase						(	self					:  Tuple							):



    """simple docstring"""



    _UpperCamelCase:     Any							= self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_megatron_bert_model(*__A							)



   def   lowerCAmelCase						(	self					:  List[Any]							):



    """simple docstring"""



    _UpperCamelCase:     Union[str, Any]							= self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_megatron_bert_for_masked_lm(*__A							)



   def   lowerCAmelCase						(	self					:  Optional[Any]							):



    """simple docstring"""



    _UpperCamelCase:     Optional[int]							= self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*__A							)



   def   lowerCAmelCase						(	self					:  Union[str, Any]							):



    """simple docstring"""



    _UpperCamelCase:     Union[str, Any]							= self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*__A							)



   def   lowerCAmelCase						(	self					:  List[Any]							):



    """simple docstring"""



    _UpperCamelCase:     List[Any]							= self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_megatron_bert_for_pretraining(*__A							)



   def   lowerCAmelCase						(	self					:  Optional[Any]							):



    """simple docstring"""



    _UpperCamelCase:     List[str]							= self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_megatron_bert_for_question_answering(*__A							)



   def   lowerCAmelCase						(	self					:  Optional[int]							):



    """simple docstring"""



    _UpperCamelCase:     List[str]							= self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*__A							)



   def   lowerCAmelCase						(	self					:  Dict							):



    """simple docstring"""



    _UpperCamelCase:     Any							= self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_megatron_bert_for_token_classification(*__A							)




def       lowerCAmelCase_     (		lowercase:							Tuple		)    ->				List[str]:







 '''simple docstring'''
 return torch.tensor(
     UpperCamelCase__  ,					dtype=torch.long  ,					device=UpperCamelCase__  ,					)


UpperCAmelCase_			    =							1E-4





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





   """simple docstring"""




   @slow
   @unittest.skip('''Model is not available.'''							)
   def   lowerCAmelCase						(	self					:  str							):



    """simple docstring"""



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

    _UpperCamelCase:     Optional[Any]							= [-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							):
      _UpperCamelCase:     int							= output[0, ii, jj]
      _UpperCamelCase:     Dict							= expected[3 * ii + jj]
      _UpperCamelCase:     Optional[Any]							= '''ii={} jj={} a={} b={}'''.format(__A		,  __A		,  __A		,  __A							)
      self.assertTrue(math.isclose(__A		,  __A		,  rel_tol=__A		,  abs_tol=__A							)		,  msg=__A							)
def       lowerCAmelCase_     (		lowercase:							int = 10**9		)    ->				int:







 '''simple docstring'''
 _UpperCamelCase:     List[Any]							= 1
 _UpperCamelCase:     Dict							= 2

 _UpperCamelCase:     Tuple							= 0
 _UpperCamelCase:     int							= 0
 _UpperCamelCase:     Dict							= 0
 while perimeter <= max_perimeter:
  perimeters_sum += perimeter

  prev_value += 2 * value
  value += prev_value

  _UpperCamelCase:     str							= 2 * value + 2 if i % 2 == 0 else 2 * value - 2
  i += 1

 return perimeters_sum


if __name__ == "__main__":
  print(f"""{solution() = }""")