

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






     '''simple docstring'''

     def __init__(						self	,    _a	,    _a=13	,    _a=7	,    _a=True	,    _a=True	,    _a=True	,    _a=True	,    _a=99	,    _a=64	,    _a=32	,    _a=5	,    _a=4	,    _a=37	,    _a="gelu"	,    _a=0.1	,    _a=0.1	,    _a=512	,    _a=16	,    _a=2	,    _a=0.02	,    _a=3	,    _a=4	,    _a=None	,    )  ->      Tuple:


           """simple docstring"""


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

     def 					_a			(						self  )  ->      str:


           """simple docstring"""


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

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

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

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

           SCREAMING_SNAKE_CASE__   :      List[str]      		=	self.get_config()

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

     def 					_a			(						self  )  ->      str:


           """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=UpperCAmelCase_	,    initializer_range=self.initializer_range	,    )

     def 					_a			(						self	,    _a	,    _a	,    _a	,    _a	,    _a	,    _a	,    _a  )  ->      Optional[Any]:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      Dict      		=	MegatronBertModel(config=UpperCAmelCase_  )
           model.to(UpperCAmelCase_  )
           model.eval()
           SCREAMING_SNAKE_CASE__   :      Optional[int]      		=	model(UpperCAmelCase_	,    attention_mask=UpperCAmelCase_	,    token_type_ids=UpperCAmelCase_  )
           SCREAMING_SNAKE_CASE__   :      Any      		=	model(UpperCAmelCase_	,    token_type_ids=UpperCAmelCase_  )
           SCREAMING_SNAKE_CASE__   :      Dict      		=	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 					_a			(						self	,    _a	,    _a	,    _a	,    _a	,    _a	,    _a	,    _a  )  ->      int:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      List[str]      		=	MegatronBertForMaskedLM(config=UpperCAmelCase_  )
           model.to(UpperCAmelCase_  )
           model.eval()
           SCREAMING_SNAKE_CASE__   :      Optional[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 					_a			(						self	,    _a	,    _a	,    _a	,    _a	,    _a	,    _a	,    _a  )  ->      Optional[Any]:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      str      		=	MegatronBertForCausalLM(config=UpperCAmelCase_  )
           model.to(UpperCAmelCase_  )
           model.eval()
           SCREAMING_SNAKE_CASE__   :      List[Any]      		=	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 					_a			(						self	,    _a	,    _a	,    _a	,    _a	,    _a	,    _a	,    _a  )  ->      Optional[int]:


           """simple docstring"""


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

     def 					_a			(						self	,    _a	,    _a	,    _a	,    _a	,    _a	,    _a	,    _a  )  ->      int:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      Dict      		=	MegatronBertForPreTraining(config=UpperCAmelCase_  )
           model.to(UpperCAmelCase_  )
           model.eval()
           SCREAMING_SNAKE_CASE__   :      Optional[Any]      		=	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 					_a			(						self	,    _a	,    _a	,    _a	,    _a	,    _a	,    _a	,    _a  )  ->      Optional[int]:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      str      		=	MegatronBertForQuestionAnswering(config=UpperCAmelCase_  )
           model.to(UpperCAmelCase_  )
           model.eval()
           SCREAMING_SNAKE_CASE__   :      List[str]      		=	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 					_a			(						self	,    _a	,    _a	,    _a	,    _a	,    _a	,    _a	,    _a  )  ->      List[Any]:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      Optional[Any]      		=	self.num_labels
           SCREAMING_SNAKE_CASE__   :      Optional[int]      		=	MegatronBertForSequenceClassification(UpperCAmelCase_  )
           model.to(UpperCAmelCase_  )
           model.eval()
           SCREAMING_SNAKE_CASE__   :      Optional[int]      		=	model(UpperCAmelCase_	,    attention_mask=UpperCAmelCase_	,    token_type_ids=UpperCAmelCase_	,    labels=UpperCAmelCase_  )
           self.parent.assertEqual(result.logits.shape	,    (self.batch_size, self.num_labels)  )

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


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      List[str]      		=	self.num_labels
           SCREAMING_SNAKE_CASE__   :      Any      		=	MegatronBertForTokenClassification(config=UpperCAmelCase_  )
           model.to(UpperCAmelCase_  )
           model.eval()
           SCREAMING_SNAKE_CASE__   :      List[Any]      		=	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 					_a			(						self	,    _a	,    _a	,    _a	,    _a	,    _a	,    _a	,    _a  )  ->      Any:


           """simple docstring"""


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





     def 					_a			(						self  )  ->      str:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      Dict      		=	self.prepare_config_and_inputs()
           (
               SCREAMING_SNAKE_CASE__	
           )   :      Dict      		=	config_and_inputs
           SCREAMING_SNAKE_CASE__   :      Dict      		=	{'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):






     '''simple docstring'''
     _SCREAMING_SNAKE_CASE						:Optional[Any]       = (
         (
             MegatronBertModel,
             MegatronBertForMaskedLM,
             MegatronBertForCausalLM,
             MegatronBertForMultipleChoice,
             MegatronBertForNextSentencePrediction,
             MegatronBertForPreTraining,
             MegatronBertForQuestionAnswering,
             MegatronBertForSequenceClassification,
             MegatronBertForTokenClassification,
         )
         if is_torch_available()
         else ()
     )
     _SCREAMING_SNAKE_CASE						:Dict       = (
         {
             """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						:int       = True
     # test_resize_embeddings = False
     _SCREAMING_SNAKE_CASE						:Any       = False

     def 					_a			(						self	,    _a	,    _a	,    _a=False  )  ->      List[str]:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      int      		=	super()._prepare_for_class(UpperCAmelCase_	,    UpperCAmelCase_	,    return_labels=UpperCAmelCase_  )

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

     def 					_a			(						self  )  ->      str:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      List[str]      		=	MegatronBertModelTester(self  )
           SCREAMING_SNAKE_CASE__   :      Optional[Any]      		=	ConfigTester(self	,    config_class=UpperCAmelCase_	,    hidden_size=37  )

     def 					_a			(						self  )  ->      int:


           """simple docstring"""


           self.config_tester.run_common_tests()

     def 					_a			(						self  )  ->      int:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      int      		=	self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_megatron_bert_model(*UpperCAmelCase_  )

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


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      Optional[Any]      		=	self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_megatron_bert_for_masked_lm(*UpperCAmelCase_  )

     def 					_a			(						self  )  ->      int:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      str      		=	self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*UpperCAmelCase_  )

     def 					_a			(						self  )  ->      str:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      Optional[int]      		=	self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*UpperCAmelCase_  )

     def 					_a			(						self  )  ->      Optional[Any]:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      Optional[Any]      		=	self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_megatron_bert_for_pretraining(*UpperCAmelCase_  )

     def 					_a			(						self  )  ->      List[Any]:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      str      		=	self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_megatron_bert_for_question_answering(*UpperCAmelCase_  )

     def 					_a			(						self  )  ->      Dict:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      List[str]      		=	self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*UpperCAmelCase_  )





     def 					_a			(						self  )  ->      str:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      str      		=	self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_megatron_bert_for_token_classification(*UpperCAmelCase_  )





def  _lowercase					(    __lowerCAmelCase							)       ->				Optional[int]:
      return torch.tensor(
          lowerCamelCase_  ,					dtype=torch.long  ,					device=lowerCamelCase_  ,					)


a					:Optional[int]    = 1e-4



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






     '''simple docstring'''

     @slow
     @unittest.skip("""Model is not available."""  )
     def 					_a			(						self  )  ->      Tuple:


           """simple docstring"""


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

           SCREAMING_SNAKE_CASE__   :      Any      		=	[-0.6_040, -0.2_517, -0.1_025, 0.3_420, -0.6_758, -0.0_017, -0.1_089, -0.1_990, 0.5_728]
           for ii in range(3  ):
                 for jj in range(3  ):
                       SCREAMING_SNAKE_CASE__   :      Tuple      		=	output[0, ii, jj]
                       SCREAMING_SNAKE_CASE__   :      Optional[Any]      		=	expected[3 * ii + jj]
                       SCREAMING_SNAKE_CASE__   :      Optional[Any]      		=	'ii={} jj={} a={} b={}'.format(UpperCAmelCase_	,    UpperCAmelCase_	,    UpperCAmelCase_	,    UpperCAmelCase_  )
                       self.assertTrue(math.isclose(UpperCAmelCase_	,    UpperCAmelCase_	,    rel_tol=UpperCAmelCase_	,    abs_tol=UpperCAmelCase_  )	,    msg=UpperCAmelCase_  )







"""simple docstring"""


a					:List[str]    = [
    (1_000, "M"),
    (900, "CM"),
    (500, "D"),
    (400, "CD"),
    (100, "C"),
    (90, "XC"),
    (50, "L"),
    (40, "XL"),
    (10, "X"),
    (9, "IX"),
    (5, "V"),
    (4, "IV"),
    (1, "I"),
]




def  _lowercase					(    __lowerCAmelCase							)       ->				int:
      SCREAMING_SNAKE_CASE__   :      Optional[Any]      		=	{"""I""": 1, """V""": 5, """X""": 10, """L""": 50, """C""": 100, """D""": 500, """M""": 1000}
      SCREAMING_SNAKE_CASE__   :      List[Any]      		=	0
      SCREAMING_SNAKE_CASE__   :      List[str]      		=	0
      while place < len(__lowerCAmelCase							):
            if (place + 1 < len(__lowerCAmelCase							)) and (vals[roman[place]] < vals[roman[place + 1]]):
                  total += vals[roman[place + 1]] - vals[roman[place]]
                  place += 2
            else:
                  total += vals[roman[place]]
                  place += 1
      return total




def  _lowercase					(    __lowerCAmelCase							)       ->				str:
      SCREAMING_SNAKE_CASE__   :      Any      		=	[]
      for arabic, roman in ROMAN:
            ((SCREAMING_SNAKE_CASE__)		,		(SCREAMING_SNAKE_CASE__))   :      List[str]      		=	divmod(__lowerCAmelCase  ,					__lowerCAmelCase							)
            result.append(roman * factor							)
            if number == 0:
                  break
      return "".join(__lowerCAmelCase							)


if __name__ == "__main__":
  import doctest

  doctest.testmod()




