



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   __lowerCAmelCase  :
					"""simple docstring"""






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






								"""simple docstring"""


								A__	  =     parent
								A__	  =     batch_size
								A__	  =     seq_length
								A__	  =     is_training
								A__	  =     use_input_mask
								A__	  =     use_token_type_ids
								A__	  =     use_labels
								A__	  =     vocab_size
								A__	  =     hidden_size
								A__	  =     embedding_size
								A__	  =     num_hidden_layers
								A__	  =     num_attention_heads
								A__	  =     intermediate_size
								A__	  =     hidden_act
								A__	  =     hidden_dropout_prob
								A__	  =     attention_probs_dropout_prob
								A__	  =     max_position_embeddings
								A__	  =     type_vocab_size
								A__	  =     type_sequence_label_size
								A__	  =     initializer_range
								A__	  =     num_labels
								A__	  =     num_choices
								A__	  =     scope



					def      _a				(  self       :							int       ):






								"""simple docstring"""


								A__	  =     ids_tensor([self.batch_size, self.seq_length]					,       self.vocab_size       )

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

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

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

								A__	  =     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=_snake_case					,       initializer_range=self.initializer_range					,       )



					def      _a				(  self       :							Dict					,       _snake_case       :							Any					,       _snake_case       :							Optional[Any]					,       _snake_case       :							List[str]					,       _snake_case       :							int					,       _snake_case       :							Optional[Any]					,       _snake_case       :							str					,       _snake_case       :							Union[str, Any]       ):






								"""simple docstring"""


								A__	  =     MegatronBertModel(config=_snake_case       )
								model.to(_snake_case       )
								model.eval()
								A__	  =     model(_snake_case					,       attention_mask=_snake_case					,       token_type_ids=_snake_case       )
								A__	  =     model(_snake_case					,       token_type_ids=_snake_case       )
								A__	  =     model(_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      _a				(  self       :							Tuple					,       _snake_case       :							List[str]					,       _snake_case       :							Any					,       _snake_case       :							str					,       _snake_case       :							Optional[int]					,       _snake_case       :							str					,       _snake_case       :							List[Any]					,       _snake_case       :							Optional[int]       ):






								"""simple docstring"""


								A__	  =     MegatronBertForMaskedLM(config=_snake_case       )
								model.to(_snake_case       )
								model.eval()
								A__	  =     model(_snake_case					,       attention_mask=_snake_case					,       token_type_ids=_snake_case					,       labels=_snake_case       )
								self.parent.assertEqual(result.logits.shape					,       (self.batch_size, self.seq_length, self.vocab_size)       )



					def      _a				(  self       :							str					,       _snake_case       :							Optional[Any]					,       _snake_case       :							List[str]					,       _snake_case       :							str					,       _snake_case       :							Any					,       _snake_case       :							Any					,       _snake_case       :							Optional[int]					,       _snake_case       :							Dict       ):






								"""simple docstring"""


								A__	  =     MegatronBertForCausalLM(config=_snake_case       )
								model.to(_snake_case       )
								model.eval()
								A__	  =     model(_snake_case					,       attention_mask=_snake_case					,       token_type_ids=_snake_case					,       labels=_snake_case       )
								self.parent.assertEqual(result.logits.shape					,       (self.batch_size, self.seq_length, self.vocab_size)       )



					def      _a				(  self       :							Union[str, Any]					,       _snake_case       :							Any					,       _snake_case       :							List[str]					,       _snake_case       :							List[str]					,       _snake_case       :							List[Any]					,       _snake_case       :							Optional[int]					,       _snake_case       :							List[Any]					,       _snake_case       :							List[Any]       ):






								"""simple docstring"""


								A__	  =     MegatronBertForNextSentencePrediction(config=_snake_case       )
								model.to(_snake_case       )
								model.eval()
								A__	  =     model(
								    _snake_case					,       attention_mask=_snake_case					,       token_type_ids=_snake_case					,       labels=_snake_case					,       )
								self.parent.assertEqual(result.logits.shape					,       (self.batch_size, 2)       )



					def      _a				(  self       :							Dict					,       _snake_case       :							Optional[Any]					,       _snake_case       :							Optional[Any]					,       _snake_case       :							List[str]					,       _snake_case       :							Tuple					,       _snake_case       :							str					,       _snake_case       :							str					,       _snake_case       :							Dict       ):






								"""simple docstring"""


								A__	  =     MegatronBertForPreTraining(config=_snake_case       )
								model.to(_snake_case       )
								model.eval()
								A__	  =     model(
								    _snake_case					,       attention_mask=_snake_case					,       token_type_ids=_snake_case					,       labels=_snake_case					,       next_sentence_label=_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      _a				(  self       :							List[Any]					,       _snake_case       :							Union[str, Any]					,       _snake_case       :							Optional[int]					,       _snake_case       :							Union[str, Any]					,       _snake_case       :							Tuple					,       _snake_case       :							int					,       _snake_case       :							List[Any]					,       _snake_case       :							str       ):






								"""simple docstring"""


								A__	  =     MegatronBertForQuestionAnswering(config=_snake_case       )
								model.to(_snake_case       )
								model.eval()
								A__	  =     model(
								    _snake_case					,       attention_mask=_snake_case					,       token_type_ids=_snake_case					,       start_positions=_snake_case					,       end_positions=_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      _a				(  self       :							Optional[int]					,       _snake_case       :							str					,       _snake_case       :							Any					,       _snake_case       :							int					,       _snake_case       :							Dict					,       _snake_case       :							int					,       _snake_case       :							List[Any]					,       _snake_case       :							Tuple       ):






								"""simple docstring"""


								A__	  =     self.num_labels
								A__	  =     MegatronBertForSequenceClassification(_snake_case       )
								model.to(_snake_case       )
								model.eval()
								A__	  =     model(_snake_case					,       attention_mask=_snake_case					,       token_type_ids=_snake_case					,       labels=_snake_case       )
								self.parent.assertEqual(result.logits.shape					,       (self.batch_size, self.num_labels)       )



					def      _a				(  self       :							Union[str, Any]					,       _snake_case       :							Tuple					,       _snake_case       :							Union[str, Any]					,       _snake_case       :							List[str]					,       _snake_case       :							str					,       _snake_case       :							Optional[Any]					,       _snake_case       :							Any					,       _snake_case       :							Tuple       ):






								"""simple docstring"""


								A__	  =     self.num_labels
								A__	  =     MegatronBertForTokenClassification(config=_snake_case       )
								model.to(_snake_case       )
								model.eval()
								A__	  =     model(_snake_case					,       attention_mask=_snake_case					,       token_type_ids=_snake_case					,       labels=_snake_case       )
								self.parent.assertEqual(result.logits.shape					,       (self.batch_size, self.seq_length, self.num_labels)       )



					def      _a				(  self       :							str					,       _snake_case       :							Dict					,       _snake_case       :							Any					,       _snake_case       :							Dict					,       _snake_case       :							int					,       _snake_case       :							Dict					,       _snake_case       :							List[Any]					,       _snake_case       :							Optional[int]       ):






								"""simple docstring"""


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



					def      _a				(  self       :							Any       ):






								"""simple docstring"""


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







@require_torch
class   __lowerCAmelCase  (		UpperCAmelCase_				,  UpperCAmelCase_				,  unittest.TestCase   ):
					"""simple docstring"""






					A__					:		str      							=     (
					    (
					        MegatronBertModel,
					        MegatronBertForMaskedLM,
					        MegatronBertForCausalLM,
					        MegatronBertForMultipleChoice,
					        MegatronBertForNextSentencePrediction,
					        MegatronBertForPreTraining,
					        MegatronBertForQuestionAnswering,
					        MegatronBertForSequenceClassification,
					        MegatronBertForTokenClassification,
					    )
					    if is_torch_available()
					    else ()
					)
					A__					:		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 {}
					)
					A__					:		Dict      							=     True
					# test_resize_embeddings = False
					A__					:		List[str]      							=     False



					def      _a				(  self       :							Dict					,       _snake_case       :							Tuple					,       _snake_case       :							str					,       _snake_case       :							Optional[Any]=False       ):






								"""simple docstring"""


								A__	  =     super()._prepare_for_class(_snake_case					,       _snake_case					,       return_labels=_snake_case       )

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



					def      _a				(  self       :							str       ):






								"""simple docstring"""


								A__	  =     MegatronBertModelTester(self       )
								A__	  =     ConfigTester(self					,       config_class=_snake_case					,       hidden_size=37       )



					def      _a				(  self       :							Optional[int]       ):






								"""simple docstring"""


								self.config_tester.run_common_tests()



					def      _a				(  self       :							Any       ):






								"""simple docstring"""


								A__	  =     self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_megatron_bert_model(*_snake_case       )



					def      _a				(  self       :							Dict       ):






								"""simple docstring"""


								A__	  =     self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_megatron_bert_for_masked_lm(*_snake_case       )



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






								"""simple docstring"""


								A__	  =     self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*_snake_case       )



					def      _a				(  self       :							Tuple       ):






								"""simple docstring"""


								A__	  =     self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*_snake_case       )



					def      _a				(  self       :							Dict       ):






								"""simple docstring"""


								A__	  =     self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_megatron_bert_for_pretraining(*_snake_case       )



					def      _a				(  self       :							int       ):






								"""simple docstring"""


								A__	  =     self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_megatron_bert_for_question_answering(*_snake_case       )



					def      _a				(  self       :							Tuple       ):






								"""simple docstring"""


								A__	  =     self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*_snake_case       )



					def      _a				(  self       :							Optional[Any]       ):






								"""simple docstring"""


								A__	  =     self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_megatron_bert_for_token_classification(*_snake_case       )







def A    (							__UpperCamelCase       )    ->   Any:
			return torch.tensor(
			    __UpperCamelCase			,						dtype=torch.long			,						device=__UpperCamelCase			,						)


SCREAMING_SNAKE_CASE__   			=   1e-4






@require_torch
@require_sentencepiece
@require_tokenizers
class   __lowerCAmelCase  (		unittest.TestCase   ):
					"""simple docstring"""






					@slow
					@unittest.skip('Model is not available.'       )
					def      _a				(  self       :							int       ):






								"""simple docstring"""


								A__	  =     'nvidia/megatron-bert-uncased-345m'
								if "MYDIR" in os.environ:
											A__	  =     os.path.join(os.environ['MYDIR']					,       _snake_case       )
								A__	  =     MegatronBertModel.from_pretrained(_snake_case       )
								model.to(_snake_case       )
								model.half()
								A__	  =     _long_tensor([[1_01, 71_10, 10_05, 10_56, 20_23, 1_13_33, 1_74_13, 10_29, 1_02]]       )
								with torch.no_grad():
											A__	  =     model(_snake_case       )[0]
								A__	  =     torch.Size((1, 9, 10_24)       )
								self.assertEqual(output.shape					,       _snake_case       )

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









SCREAMING_SNAKE_CASE__   			=   '''
# Transformers installation
! pip install transformers datasets
# To install from source instead of the last release, comment the command above and uncomment the following one.
# ! pip install git+https://github.com/huggingface/transformers.git
'''

SCREAMING_SNAKE_CASE__   			=   [{'''type''': '''code''', '''content''': INSTALL_CONTENT}]
SCREAMING_SNAKE_CASE__   			=   {
    '''{processor_class}''': '''FakeProcessorClass''',
    '''{model_class}''': '''FakeModelClass''',
    '''{object_class}''': '''FakeObjectClass''',
}




