





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






   '''simple docstring'''





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

   def 		__UpperCamelCase						(							self       ):
    snake_case__     :					Union[str, Any]    =  ids_tensor([self.batch_size, self.seq_length]	,		self.vocab_size       )

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

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

    snake_case__     :					List[Any]    =  None
    snake_case__     :					str    =  None
    snake_case__     :					Optional[int]    =  None
    if self.use_labels:
     snake_case__     :					List[str]    =  ids_tensor([self.batch_size]	,		self.type_sequence_label_size       )
     snake_case__     :					str    =  ids_tensor([self.batch_size, self.seq_length]	,		self.num_labels       )
     snake_case__     :					Dict    =  ids_tensor([self.batch_size]	,		self.num_choices       )

    snake_case__     :					int    =  self.get_config()

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

   def 		__UpperCamelCase						(							self       ):
    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 		__UpperCamelCase						(							self	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE       ):
    snake_case__     :					List[Any]    =  MegatronBertModel(config=__SCREAMING_SNAKE_CASE       )
    model.to(__SCREAMING_SNAKE_CASE       )
    model.eval()
    snake_case__     :					Optional[int]    =  model(__SCREAMING_SNAKE_CASE	,		attention_mask=__SCREAMING_SNAKE_CASE	,		token_type_ids=__SCREAMING_SNAKE_CASE       )
    snake_case__     :					List[str]    =  model(__SCREAMING_SNAKE_CASE	,		token_type_ids=__SCREAMING_SNAKE_CASE       )
    snake_case__     :					Tuple    =  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 		__UpperCamelCase						(							self	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE       ):
    snake_case__     :					str    =  MegatronBertForMaskedLM(config=__SCREAMING_SNAKE_CASE       )
    model.to(__SCREAMING_SNAKE_CASE       )
    model.eval()
    snake_case__     :					str    =  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 		__UpperCamelCase						(							self	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE       ):
    snake_case__     :					List[str]    =  MegatronBertForCausalLM(config=__SCREAMING_SNAKE_CASE       )
    model.to(__SCREAMING_SNAKE_CASE       )
    model.eval()
    snake_case__     :					Optional[int]    =  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 		__UpperCamelCase						(							self	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE       ):
    snake_case__     :					Optional[int]    =  MegatronBertForNextSentencePrediction(config=__SCREAMING_SNAKE_CASE       )
    model.to(__SCREAMING_SNAKE_CASE       )
    model.eval()
    snake_case__     :					Any    =  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 		__UpperCamelCase						(							self	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE       ):
    snake_case__     :					Dict    =  MegatronBertForPreTraining(config=__SCREAMING_SNAKE_CASE       )
    model.to(__SCREAMING_SNAKE_CASE       )
    model.eval()
    snake_case__     :					Optional[int]    =  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 		__UpperCamelCase						(							self	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE       ):
    snake_case__     :					Optional[Any]    =  MegatronBertForQuestionAnswering(config=__SCREAMING_SNAKE_CASE       )
    model.to(__SCREAMING_SNAKE_CASE       )
    model.eval()
    snake_case__     :					Tuple    =  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 		__UpperCamelCase						(							self	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE       ):
    snake_case__     :					int    =  self.num_labels
    snake_case__     :					Optional[int]    =  MegatronBertForSequenceClassification(__SCREAMING_SNAKE_CASE       )
    model.to(__SCREAMING_SNAKE_CASE       )
    model.eval()
    snake_case__     :					int    =  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 		__UpperCamelCase						(							self	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE       ):
    snake_case__     :					Optional[Any]    =  self.num_labels
    snake_case__     :					str    =  MegatronBertForTokenClassification(config=__SCREAMING_SNAKE_CASE       )
    model.to(__SCREAMING_SNAKE_CASE       )
    model.eval()
    snake_case__     :					Tuple    =  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 		__UpperCamelCase						(							self	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE       ):
    snake_case__     :					Any    =  self.num_choices
    snake_case__     :					Dict    =  MegatronBertForMultipleChoice(config=__SCREAMING_SNAKE_CASE       )
    model.to(__SCREAMING_SNAKE_CASE       )
    model.eval()
    snake_case__     :					Dict    =  input_ids.unsqueeze(1       ).expand(-1	,		self.num_choices	,		-1       ).contiguous()
    snake_case__     :					str    =  token_type_ids.unsqueeze(1       ).expand(-1	,		self.num_choices	,		-1       ).contiguous()
    snake_case__     :					List[str]    =  input_mask.unsqueeze(1       ).expand(-1	,		self.num_choices	,		-1       ).contiguous()
    snake_case__     :					Tuple    =  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 		__UpperCamelCase						(							self       ):
    snake_case__     :					Tuple    =  self.prepare_config_and_inputs()
    (
        (
        snake_case__
    )  ,							(
        snake_case__
    )  ,							(
        snake_case__
    )  ,							(
        snake_case__
    )  ,							(
        snake_case__
    )  ,							(
        snake_case__
    )  ,							(
        snake_case__
    )  ,							
    )     :					Optional[int]    =  config_and_inputs
    snake_case__     :					Union[str, Any]    =  {"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
    return config, inputs_dict



@require_torch
class 					__snake_case  (					__SCREAMING_SNAKE_CASE		,				__SCREAMING_SNAKE_CASE		,				unittest.TestCase							):






   '''simple docstring'''





   lowerCamelCase__    							=      (
       (
           MegatronBertModel,
           MegatronBertForMaskedLM,
           MegatronBertForCausalLM,
           MegatronBertForMultipleChoice,
           MegatronBertForNextSentencePrediction,
           MegatronBertForPreTraining,
           MegatronBertForQuestionAnswering,
           MegatronBertForSequenceClassification,
           MegatronBertForTokenClassification,
       )
       if is_torch_available()
       else ()
   )
   lowerCamelCase__    							=      (
       {
           '''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__    							=      True
   # test_resize_embeddings = False
   lowerCamelCase__    							=      False

   def 		__UpperCamelCase						(							self	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE	,		__SCREAMING_SNAKE_CASE=False       ):
    snake_case__     :					List[Any]    =  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__     :					str    =  torch.zeros(
          (self.model_tester.batch_size, self.model_tester.seq_length)	,		dtype=torch.long	,		device=__SCREAMING_SNAKE_CASE       )
      snake_case__     :					Tuple    =  torch.zeros(
          self.model_tester.batch_size	,		dtype=torch.long	,		device=__SCREAMING_SNAKE_CASE       )
    return inputs_dict

   def 		__UpperCamelCase						(							self       ):
    snake_case__     :					Tuple    =  MegatronBertModelTester(self       )
    snake_case__     :					List[str]    =  ConfigTester(self	,		config_class=__SCREAMING_SNAKE_CASE	,		hidden_size=3_7       )

   def 		__UpperCamelCase						(							self       ):
    self.config_tester.run_common_tests()

   def 		__UpperCamelCase						(							self       ):
    snake_case__     :					Optional[Any]    =  self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_megatron_bert_model(*__SCREAMING_SNAKE_CASE       )

   def 		__UpperCamelCase						(							self       ):
    snake_case__     :					int    =  self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_megatron_bert_for_masked_lm(*__SCREAMING_SNAKE_CASE       )

   def 		__UpperCamelCase						(							self       ):
    snake_case__     :					Union[str, Any]    =  self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*__SCREAMING_SNAKE_CASE       )

   def 		__UpperCamelCase						(							self       ):
    snake_case__     :					int    =  self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*__SCREAMING_SNAKE_CASE       )

   def 		__UpperCamelCase						(							self       ):
    snake_case__     :					List[str]    =  self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_megatron_bert_for_pretraining(*__SCREAMING_SNAKE_CASE       )

   def 		__UpperCamelCase						(							self       ):
    snake_case__     :					Tuple    =  self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_megatron_bert_for_question_answering(*__SCREAMING_SNAKE_CASE       )

   def 		__UpperCamelCase						(							self       ):
    snake_case__     :					Optional[Any]    =  self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*__SCREAMING_SNAKE_CASE       )


   def 		__UpperCamelCase						(							self       ):
    snake_case__     :					Union[str, Any]    =  self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_megatron_bert_for_token_classification(*__SCREAMING_SNAKE_CASE       )



def        UpperCamelCase__		(						__magic_name__  :   str      )				->			str:
 '''simple docstring'''






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


A_		:  str		 =						1e-4





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






   '''simple docstring'''





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

    snake_case__     :					str    =  [-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       ):
      snake_case__     :					List[Any]    =  output[0, ii, jj]
      snake_case__     :					Optional[int]    =  expected[3 * ii + jj]
      snake_case__     :					Dict    =  """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       )







def 	__a     (      A__			:  int       ):
    if not isinstance(A__			,  A__       ):
        raise ValueError("Input must be an integer"       )
    if input_num <= 0:
        raise ValueError("Input must be positive"       )
    return sum(
        divisor for divisor in range(1			,  input_num // 2 + 1       ) if input_num % divisor == 0       )


if __name__ == "__main__":
    import doctest

    doctest.testmod()