


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    :




 def __init__(						self   ,    _lowerCamelCase   ,    _lowerCamelCase=13   ,    _lowerCamelCase=7   ,    _lowerCamelCase=True   ,    _lowerCamelCase=True   ,    _lowerCamelCase=True   ,    _lowerCamelCase=True   ,    _lowerCamelCase=99   ,    _lowerCamelCase=64   ,    _lowerCamelCase=32   ,    _lowerCamelCase=5   ,    _lowerCamelCase=4   ,    _lowerCamelCase=37   ,    _lowerCamelCase="gelu"   ,    _lowerCamelCase=0.1   ,    _lowerCamelCase=0.1   ,    _lowerCamelCase=512   ,    _lowerCamelCase=16   ,    _lowerCamelCase=2   ,    _lowerCamelCase=0.02   ,    _lowerCamelCase=3   ,    _lowerCamelCase=4   ,    _lowerCamelCase=None   ,    ):
       a		:Optional[int]									=			parent
       a		:Union[str, Any]									=			batch_size
       a		:Dict									=			seq_length
       a		:Any									=			is_training
       a		:Optional[Any]									=			use_input_mask
       a		:Union[str, Any]									=			use_token_type_ids
       a		:str									=			use_labels
       a		:Any									=			vocab_size
       a		:List[str]									=			hidden_size
       a		:str									=			embedding_size
       a		:Union[str, Any]									=			num_hidden_layers
       a		:str									=			num_attention_heads
       a		:Optional[int]									=			intermediate_size
       a		:Any									=			hidden_act
       a		:Optional[int]									=			hidden_dropout_prob
       a		:Tuple									=			attention_probs_dropout_prob
       a		:Tuple									=			max_position_embeddings
       a		:List[str]									=			type_vocab_size
       a		:Optional[Any]									=			type_sequence_label_size
       a		:int									=			initializer_range
       a		:Optional[int]									=			num_labels
       a		:Any									=			num_choices
       a		:List[str]									=			scope




 def 			SCREAMING_SNAKE_CASE__							(						self       ):
       a		:Optional[Any]									=			ids_tensor([self.batch_size, self.seq_length]   ,    self.vocab_size       )

       a		:Dict									=			None
       if self.use_input_mask:
             a		:List[Any]									=			random_attention_mask([self.batch_size, self.seq_length]       )

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

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

       a		:Optional[Any]									=			self.get_config()

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




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




 def 			SCREAMING_SNAKE_CASE__							(						self   ,    _lowerCamelCase   ,    _lowerCamelCase   ,    _lowerCamelCase   ,    _lowerCamelCase   ,    _lowerCamelCase   ,    _lowerCamelCase   ,    _lowerCamelCase       ):
       a		:Any									=			MegatronBertModel(config=_lowerCamelCase       )
       model.to(_lowerCamelCase       )
       model.eval()
       a		:Optional[Any]									=			model(_lowerCamelCase   ,    attention_mask=_lowerCamelCase   ,    token_type_ids=_lowerCamelCase       )
       a		:int									=			model(_lowerCamelCase   ,    token_type_ids=_lowerCamelCase       )
       a		:Optional[Any]									=			model(_lowerCamelCase       )

       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 			SCREAMING_SNAKE_CASE__							(						self   ,    _lowerCamelCase   ,    _lowerCamelCase   ,    _lowerCamelCase   ,    _lowerCamelCase   ,    _lowerCamelCase   ,    _lowerCamelCase   ,    _lowerCamelCase       ):
       a		:int									=			MegatronBertForMaskedLM(config=_lowerCamelCase       )
       model.to(_lowerCamelCase       )
       model.eval()
       a		:Any									=			model(_lowerCamelCase   ,    attention_mask=_lowerCamelCase   ,    token_type_ids=_lowerCamelCase   ,    labels=_lowerCamelCase       )
       self.parent.assertEqual(result.logits.shape   ,    (self.batch_size, self.seq_length, self.vocab_size)       )




 def 			SCREAMING_SNAKE_CASE__							(						self   ,    _lowerCamelCase   ,    _lowerCamelCase   ,    _lowerCamelCase   ,    _lowerCamelCase   ,    _lowerCamelCase   ,    _lowerCamelCase   ,    _lowerCamelCase       ):
       a		:List[Any]									=			MegatronBertForCausalLM(config=_lowerCamelCase       )
       model.to(_lowerCamelCase       )
       model.eval()
       a		:int									=			model(_lowerCamelCase   ,    attention_mask=_lowerCamelCase   ,    token_type_ids=_lowerCamelCase   ,    labels=_lowerCamelCase       )
       self.parent.assertEqual(result.logits.shape   ,    (self.batch_size, self.seq_length, self.vocab_size)       )




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




 def 			SCREAMING_SNAKE_CASE__							(						self   ,    _lowerCamelCase   ,    _lowerCamelCase   ,    _lowerCamelCase   ,    _lowerCamelCase   ,    _lowerCamelCase   ,    _lowerCamelCase   ,    _lowerCamelCase       ):
       a		:List[str]									=			MegatronBertForPreTraining(config=_lowerCamelCase       )
       model.to(_lowerCamelCase       )
       model.eval()
       a		:str									=			model(
           _lowerCamelCase   ,    attention_mask=_lowerCamelCase   ,    token_type_ids=_lowerCamelCase   ,    labels=_lowerCamelCase   ,    next_sentence_label=_lowerCamelCase   ,    )
       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 			SCREAMING_SNAKE_CASE__							(						self   ,    _lowerCamelCase   ,    _lowerCamelCase   ,    _lowerCamelCase   ,    _lowerCamelCase   ,    _lowerCamelCase   ,    _lowerCamelCase   ,    _lowerCamelCase       ):
       a		:Tuple									=			MegatronBertForQuestionAnswering(config=_lowerCamelCase       )
       model.to(_lowerCamelCase       )
       model.eval()
       a		:Any									=			model(
           _lowerCamelCase   ,    attention_mask=_lowerCamelCase   ,    token_type_ids=_lowerCamelCase   ,    start_positions=_lowerCamelCase   ,    end_positions=_lowerCamelCase   ,    )
       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 			SCREAMING_SNAKE_CASE__							(						self   ,    _lowerCamelCase   ,    _lowerCamelCase   ,    _lowerCamelCase   ,    _lowerCamelCase   ,    _lowerCamelCase   ,    _lowerCamelCase   ,    _lowerCamelCase       ):
       a		:List[Any]									=			self.num_labels
       a		:Any									=			MegatronBertForSequenceClassification(_lowerCamelCase       )
       model.to(_lowerCamelCase       )
       model.eval()
       a		:Tuple									=			model(_lowerCamelCase   ,    attention_mask=_lowerCamelCase   ,    token_type_ids=_lowerCamelCase   ,    labels=_lowerCamelCase       )
       self.parent.assertEqual(result.logits.shape   ,    (self.batch_size, self.num_labels)       )




 def 			SCREAMING_SNAKE_CASE__							(						self   ,    _lowerCamelCase   ,    _lowerCamelCase   ,    _lowerCamelCase   ,    _lowerCamelCase   ,    _lowerCamelCase   ,    _lowerCamelCase   ,    _lowerCamelCase       ):
       a		:int									=			self.num_labels
       a		:Tuple									=			MegatronBertForTokenClassification(config=_lowerCamelCase       )
       model.to(_lowerCamelCase       )
       model.eval()
       a		:int									=			model(_lowerCamelCase   ,    attention_mask=_lowerCamelCase   ,    token_type_ids=_lowerCamelCase   ,    labels=_lowerCamelCase       )
       self.parent.assertEqual(result.logits.shape   ,    (self.batch_size, self.seq_length, self.num_labels)       )




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




 def 			SCREAMING_SNAKE_CASE__							(						self       ):
       a		:Dict									=			self.prepare_config_and_inputs()
       (
           (
           a
       )				,							(
           a
       )				,							(
           a
       )				,							(
           a
       )				,							(
           a
       )				,							(
           a
       )				,							(
           a
       )				,							
       )		:List[Any]									=			config_and_inputs
       a		:Tuple									=			{'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
       return config, inputs_dict







@require_torch
class 						_snake_case    ( _snake_case		,       _snake_case		,       unittest.TestCase				):
 SCREAMING_SNAKE_CASE__ =			(
     (
         MegatronBertModel,
         MegatronBertForMaskedLM,
         MegatronBertForCausalLM,
         MegatronBertForMultipleChoice,
         MegatronBertForNextSentencePrediction,
         MegatronBertForPreTraining,
         MegatronBertForQuestionAnswering,
         MegatronBertForSequenceClassification,
         MegatronBertForTokenClassification,
     )
     if is_torch_available()
     else ()
 )
 SCREAMING_SNAKE_CASE__ =			(
     {
         '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__ =			True
 # test_resize_embeddings = False
 SCREAMING_SNAKE_CASE__ =			False




 def 			SCREAMING_SNAKE_CASE__							(						self   ,    _lowerCamelCase   ,    _lowerCamelCase   ,    _lowerCamelCase=False       ):
       a		:List[Any]									=			super()._prepare_for_class(_lowerCamelCase   ,    _lowerCamelCase   ,    return_labels=_lowerCamelCase       )

       if return_labels:
             if model_class in get_values(_lowerCamelCase       ):
                   a		:Optional[int]									=			torch.zeros(
                       (self.model_tester.batch_size, self.model_tester.seq_length)   ,    dtype=torch.long   ,    device=_lowerCamelCase       )
                   a		:Union[str, Any]									=			torch.zeros(
                       self.model_tester.batch_size   ,    dtype=torch.long   ,    device=_lowerCamelCase       )
       return inputs_dict




 def 			SCREAMING_SNAKE_CASE__							(						self       ):
       a		:str									=			MegatronBertModelTester(self       )
       a		:Optional[int]									=			ConfigTester(self   ,    config_class=_lowerCamelCase   ,    hidden_size=37       )




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




 def 			SCREAMING_SNAKE_CASE__							(						self       ):
       a		:Union[str, Any]									=			self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_model(*_lowerCamelCase       )




 def 			SCREAMING_SNAKE_CASE__							(						self       ):
       a		:Tuple									=			self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_masked_lm(*_lowerCamelCase       )




 def 			SCREAMING_SNAKE_CASE__							(						self       ):
       a		:str									=			self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*_lowerCamelCase       )




 def 			SCREAMING_SNAKE_CASE__							(						self       ):
       a		:Dict									=			self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*_lowerCamelCase       )




 def 			SCREAMING_SNAKE_CASE__							(						self       ):
       a		:Any									=			self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_pretraining(*_lowerCamelCase       )




 def 			SCREAMING_SNAKE_CASE__							(						self       ):
       a		:List[Any]									=			self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_question_answering(*_lowerCamelCase       )




 def 			SCREAMING_SNAKE_CASE__							(						self       ):
       a		:Optional[Any]									=			self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*_lowerCamelCase       )




 def 			SCREAMING_SNAKE_CASE__							(						self       ):
       a		:List[str]									=			self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_token_classification(*_lowerCamelCase       )







def 							__lowerCamelCase		(					UpperCAmelCase_		:						int		):




      """simple docstring"""





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


snake_case							:       List[str]										=	1E-4







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




 @slow
 @unittest.skip('''Model is not available.'''       )
 def 			SCREAMING_SNAKE_CASE__							(						self       ):
       a		:Union[str, Any]									=			'''nvidia/megatron-bert-uncased-345m'''
       if "MYDIR" in os.environ:
             a		:Tuple									=			os.path.join(os.environ['''MYDIR''']   ,    _lowerCamelCase       )
       a		:List[str]									=			MegatronBertModel.from_pretrained(_lowerCamelCase       )
       model.to(_lowerCamelCase       )
       model.half()
       a		:Tuple									=			_long_tensor([[101, 7110, 1005, 1056, 2023, 1_1333, 1_7413, 1029, 102]]       )
       with torch.no_grad():
             a		:int									=			model(_lowerCamelCase       )[0]
       a		:Any									=			torch.Size((1, 9, 1024)       )
       self.assertEqual(output.shape   ,    _lowerCamelCase       )

       a		:Union[str, 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       ):
                   a		:str									=			output[0, ii, jj]
                   a		:List[str]									=			expected[3 * ii + jj]
                   a		:str									=			'''ii={} jj={} a={} b={}'''.format(_lowerCamelCase   ,    _lowerCamelCase   ,    _lowerCamelCase   ,    _lowerCamelCase       )
                   self.assertTrue(math.isclose(_lowerCamelCase   ,    _lowerCamelCase   ,    rel_tol=_lowerCamelCase   ,    abs_tol=_lowerCamelCase       )   ,    msg=_lowerCamelCase       )









from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_flax_available, is_torch_available


snake_case							:       Any										=	{
    '''configuration_gpt_neo''': ['''GPT_NEO_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''GPTNeoConfig''', '''GPTNeoOnnxConfig'''],
}

try:
     if not is_torch_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     snake_case							:       List[Any]										=	[
         '''GPT_NEO_PRETRAINED_MODEL_ARCHIVE_LIST''',
         '''GPTNeoForCausalLM''',
         '''GPTNeoForQuestionAnswering''',
         '''GPTNeoForSequenceClassification''',
         '''GPTNeoForTokenClassification''',
         '''GPTNeoModel''',
         '''GPTNeoPreTrainedModel''',
         '''load_tf_weights_in_gpt_neo''',
     ]

try:
     if not is_flax_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     snake_case							:       Any										=	[
         '''FlaxGPTNeoForCausalLM''',
         '''FlaxGPTNeoModel''',
         '''FlaxGPTNeoPreTrainedModel''',
     ]


if TYPE_CHECKING:
     from .configuration_gpt_neo import GPT_NEO_PRETRAINED_CONFIG_ARCHIVE_MAP, GPTNeoConfig, GPTNeoOnnxConfig

     try:
          if not is_torch_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_gpt_neo import (
              GPT_NEO_PRETRAINED_MODEL_ARCHIVE_LIST,
              GPTNeoForCausalLM,
              GPTNeoForQuestionAnswering,
              GPTNeoForSequenceClassification,
              GPTNeoForTokenClassification,
              GPTNeoModel,
              GPTNeoPreTrainedModel,
              load_tf_weights_in_gpt_neo,
          )

     try:
          if not is_flax_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_flax_gpt_neo import FlaxGPTNeoForCausalLM, FlaxGPTNeoModel, FlaxGPTNeoPreTrainedModel


else:
     import sys

     snake_case							:       Optional[int]										=	_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)





