

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







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

						'''simple docstring'''





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







	def 	_SCREAMING_SNAKE_CASE							(   self: Dict      ):

						'''simple docstring'''





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

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

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

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

						__magic_name__							=    self.get_config()

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







	def 	_SCREAMING_SNAKE_CASE							(   self: Dict      ):

						'''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 	_SCREAMING_SNAKE_CASE							(   self: Any  , __UpperCamelCase: Dict  , __UpperCamelCase: Union[str, Any]  , __UpperCamelCase: List[str]  , __UpperCamelCase: Tuple  , __UpperCamelCase: Optional[int]  , __UpperCamelCase: Any  , __UpperCamelCase: Optional[int]      ):

						'''simple docstring'''





						__magic_name__							=    MegatronBertModel(config=__UpperCamelCase      )
						model.to(__UpperCamelCase      )
						model.eval()
						__magic_name__							=    model(__UpperCamelCase  , attention_mask=__UpperCamelCase  , token_type_ids=__UpperCamelCase      )
						__magic_name__							=    model(__UpperCamelCase  , token_type_ids=__UpperCamelCase      )
						__magic_name__							=    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 	_SCREAMING_SNAKE_CASE							(   self: Optional[int]  , __UpperCamelCase: List[Any]  , __UpperCamelCase: Dict  , __UpperCamelCase: Dict  , __UpperCamelCase: List[str]  , __UpperCamelCase: int  , __UpperCamelCase: Union[str, Any]  , __UpperCamelCase: Dict      ):

						'''simple docstring'''





						__magic_name__							=    MegatronBertForMaskedLM(config=__UpperCamelCase      )
						model.to(__UpperCamelCase      )
						model.eval()
						__magic_name__							=    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 	_SCREAMING_SNAKE_CASE							(   self: List[str]  , __UpperCamelCase: Any  , __UpperCamelCase: Tuple  , __UpperCamelCase: Dict  , __UpperCamelCase: int  , __UpperCamelCase: Union[str, Any]  , __UpperCamelCase: str  , __UpperCamelCase: List[str]      ):

						'''simple docstring'''





						__magic_name__							=    MegatronBertForCausalLM(config=__UpperCamelCase      )
						model.to(__UpperCamelCase      )
						model.eval()
						__magic_name__							=    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 	_SCREAMING_SNAKE_CASE							(   self: Dict  , __UpperCamelCase: Union[str, Any]  , __UpperCamelCase: List[Any]  , __UpperCamelCase: Dict  , __UpperCamelCase: List[Any]  , __UpperCamelCase: str  , __UpperCamelCase: Dict  , __UpperCamelCase: Optional[Any]      ):

						'''simple docstring'''





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







	def 	_SCREAMING_SNAKE_CASE							(   self: List[Any]  , __UpperCamelCase: List[Any]  , __UpperCamelCase: Any  , __UpperCamelCase: int  , __UpperCamelCase: int  , __UpperCamelCase: Union[str, Any]  , __UpperCamelCase: Any  , __UpperCamelCase: List[str]      ):

						'''simple docstring'''





						__magic_name__							=    MegatronBertForPreTraining(config=__UpperCamelCase      )
						model.to(__UpperCamelCase      )
						model.eval()
						__magic_name__							=    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 	_SCREAMING_SNAKE_CASE							(   self: Optional[Any]  , __UpperCamelCase: List[Any]  , __UpperCamelCase: List[Any]  , __UpperCamelCase: Any  , __UpperCamelCase: List[str]  , __UpperCamelCase: Optional[int]  , __UpperCamelCase: Optional[Any]  , __UpperCamelCase: List[str]      ):

						'''simple docstring'''





						__magic_name__							=    MegatronBertForQuestionAnswering(config=__UpperCamelCase      )
						model.to(__UpperCamelCase      )
						model.eval()
						__magic_name__							=    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 	_SCREAMING_SNAKE_CASE							(   self: List[str]  , __UpperCamelCase: List[Any]  , __UpperCamelCase: Dict  , __UpperCamelCase: List[Any]  , __UpperCamelCase: str  , __UpperCamelCase: Tuple  , __UpperCamelCase: List[Any]  , __UpperCamelCase: List[str]      ):

						'''simple docstring'''





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







	def 	_SCREAMING_SNAKE_CASE							(   self: Dict  , __UpperCamelCase: Any  , __UpperCamelCase: Optional[int]  , __UpperCamelCase: Optional[int]  , __UpperCamelCase: Union[str, Any]  , __UpperCamelCase: str  , __UpperCamelCase: str  , __UpperCamelCase: Optional[int]      ):

						'''simple docstring'''





						__magic_name__							=    self.num_labels
						__magic_name__							=    MegatronBertForTokenClassification(config=__UpperCamelCase      )
						model.to(__UpperCamelCase      )
						model.eval()
						__magic_name__							=    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 	_SCREAMING_SNAKE_CASE							(   self: Any  , __UpperCamelCase: Tuple  , __UpperCamelCase: List[str]  , __UpperCamelCase: Dict  , __UpperCamelCase: Union[str, Any]  , __UpperCamelCase: Optional[int]  , __UpperCamelCase: int  , __UpperCamelCase: Optional[int]      ):

						'''simple docstring'''





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







	def 	_SCREAMING_SNAKE_CASE							(   self: str      ):

						'''simple docstring'''





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




@require_torch
class 				__UpperCamelCase				(					SCREAMING_SNAKE_CASE  ,						SCREAMING_SNAKE_CASE  ,						unittest.TestCase			):
	_lowercase  :   str       				=  (
	    (
	        MegatronBertModel,
	        MegatronBertForMaskedLM,
	        MegatronBertForCausalLM,
	        MegatronBertForMultipleChoice,
	        MegatronBertForNextSentencePrediction,
	        MegatronBertForPreTraining,
	        MegatronBertForQuestionAnswering,
	        MegatronBertForSequenceClassification,
	        MegatronBertForTokenClassification,
	    )
	    if is_torch_available()
	    else ()
	)
	_lowercase  :   Optional[int]       				=  (
	    {
	        "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 {}
	)
	_lowercase  :   Union[str, Any]       				=  True
	# test_resize_embeddings = False
	_lowercase  :   Dict       				=  False







	def 	_SCREAMING_SNAKE_CASE							(   self: Any  , __UpperCamelCase: Tuple  , __UpperCamelCase: List[str]  , __UpperCamelCase: Dict=False      ):

						'''simple docstring'''





						__magic_name__							=    super()._prepare_for_class(__UpperCamelCase  , __UpperCamelCase  , return_labels=__UpperCamelCase      )

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







	def 	_SCREAMING_SNAKE_CASE							(   self: Tuple      ):

						'''simple docstring'''





						__magic_name__							=    MegatronBertModelTester(self      )
						__magic_name__							=    ConfigTester(self  , config_class=__UpperCamelCase  , hidden_size=37      )







	def 	_SCREAMING_SNAKE_CASE							(   self: int      ):

						'''simple docstring'''





						self.config_tester.run_common_tests()







	def 	_SCREAMING_SNAKE_CASE							(   self: Tuple      ):

						'''simple docstring'''





						__magic_name__							=    self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_model(*__UpperCamelCase      )







	def 	_SCREAMING_SNAKE_CASE							(   self: Any      ):

						'''simple docstring'''





						__magic_name__							=    self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_for_masked_lm(*__UpperCamelCase      )







	def 	_SCREAMING_SNAKE_CASE							(   self: List[Any]      ):

						'''simple docstring'''





						__magic_name__							=    self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*__UpperCamelCase      )







	def 	_SCREAMING_SNAKE_CASE							(   self: List[Any]      ):

						'''simple docstring'''





						__magic_name__							=    self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*__UpperCamelCase      )







	def 	_SCREAMING_SNAKE_CASE							(   self: List[str]      ):

						'''simple docstring'''





						__magic_name__							=    self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_for_pretraining(*__UpperCamelCase      )







	def 	_SCREAMING_SNAKE_CASE							(   self: List[Any]      ):

						'''simple docstring'''





						__magic_name__							=    self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_for_question_answering(*__UpperCamelCase      )







	def 	_SCREAMING_SNAKE_CASE							(   self: Dict      ):

						'''simple docstring'''





						__magic_name__							=    self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*__UpperCamelCase      )







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

						'''simple docstring'''





						__magic_name__							=    self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_for_token_classification(*__UpperCamelCase      )




def 					_lowercase     (				a_				:					str				)				->	Any:




					'''simple docstring'''







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


A__						=							1e-4


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







	@slow
	@unittest.skip('Model is not available.'      )
	def 	_SCREAMING_SNAKE_CASE							(   self: List[Any]      ):

						'''simple docstring'''





						__magic_name__							=    'nvidia/megatron-bert-uncased-345m'
						if "MYDIR" in os.environ:
											__magic_name__							=    os.path.join(os.environ['MYDIR']  , __UpperCamelCase      )
						__magic_name__							=    MegatronBertModel.from_pretrained(__UpperCamelCase      )
						model.to(__UpperCamelCase      )
						model.half()
						__magic_name__							=    _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():
											__magic_name__							=    model(__UpperCamelCase      )[0]
						__magic_name__							=    torch.Size((1, 9, 10_24)      )
						self.assertEqual(output.shape  , __UpperCamelCase      )

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







from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


A__						=							logging.get_logger(__name__)

A__						=							{
    "camembert-base": "https://huggingface.co/camembert-base/resolve/main/config.json",
    "umberto-commoncrawl-cased-v1": (
        "https://huggingface.co/Musixmatch/umberto-commoncrawl-cased-v1/resolve/main/config.json"
    ),
    "umberto-wikipedia-uncased-v1": (
        "https://huggingface.co/Musixmatch/umberto-wikipedia-uncased-v1/resolve/main/config.json"
    ),
}

class 				__UpperCamelCase				(					SCREAMING_SNAKE_CASE			):
	_lowercase  :   Optional[int]       				=  "camembert"







	def __init__(   self: Optional[Any]  , __UpperCamelCase: Any=3_05_22  , __UpperCamelCase: Tuple=7_68  , __UpperCamelCase: str=12  , __UpperCamelCase: Optional[int]=12  , __UpperCamelCase: List[str]=30_72  , __UpperCamelCase: Any="gelu"  , __UpperCamelCase: Optional[int]=0.1  , __UpperCamelCase: Any=0.1  , __UpperCamelCase: str=5_12  , __UpperCamelCase: Dict=2  , __UpperCamelCase: str=0.02  , __UpperCamelCase: List[Any]=1E-12  , __UpperCamelCase: List[Any]=1  , __UpperCamelCase: Any=0  , __UpperCamelCase: Union[str, Any]=2  , __UpperCamelCase: Dict="absolute"  , __UpperCamelCase: Any=True  , __UpperCamelCase: Any=None  , **__UpperCamelCase: Union[str, Any]  , ):

						'''simple docstring'''





						super().__init__(pad_token_id=__UpperCamelCase  , bos_token_id=__UpperCamelCase  , eos_token_id=__UpperCamelCase  , **__UpperCamelCase      )

						__magic_name__							=    vocab_size
						__magic_name__							=    hidden_size
						__magic_name__							=    num_hidden_layers
						__magic_name__							=    num_attention_heads
						__magic_name__							=    hidden_act
						__magic_name__							=    intermediate_size
						__magic_name__							=    hidden_dropout_prob
						__magic_name__							=    attention_probs_dropout_prob
						__magic_name__							=    max_position_embeddings
						__magic_name__							=    type_vocab_size
						__magic_name__							=    initializer_range
						__magic_name__							=    layer_norm_eps
						__magic_name__							=    position_embedding_type
						__magic_name__							=    use_cache
						__magic_name__							=    classifier_dropout




class 				__UpperCamelCase				(					SCREAMING_SNAKE_CASE			):







	@property
	def 	_SCREAMING_SNAKE_CASE							(   self: Dict      ):

						'''simple docstring'''





						if self.task == "multiple-choice":
											__magic_name__							=    {0: 'batch', 1: 'choice', 2: 'sequence'}
						else:
											__magic_name__							=    {0: 'batch', 1: 'sequence'}
						return OrderedDict(
						    [
						        ('input_ids', dynamic_axis),
						        ('attention_mask', dynamic_axis),
						    ]      )




