
'''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		:
					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,    ):

									"""simple docstring"""


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

									"""simple docstring"""


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

									lowerCamelCase					:			str									= None
									if self.use_input_mask:
													lowerCamelCase					:			Optional[Any]									= random_attention_mask([self.batch_size, self.seq_length]      )

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

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

									lowerCamelCase					:			Dict									= self.get_config()

									return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
					def     UpperCAmelCase_							(     self      ):

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

									"""simple docstring"""


									lowerCamelCase					:			List[Any]									= MegatronBertModel(config=A      )
									model.to(A      )
									model.eval()
									lowerCamelCase					:			Optional[int]									= model(A,    attention_mask=A,    token_type_ids=A      )
									lowerCamelCase					:			int									= model(A,    token_type_ids=A      )
									lowerCamelCase					:			Optional[Any]									= model(A      )

									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,    A,    A,    A,    A,    A,    A,    A      ):

									"""simple docstring"""


									lowerCamelCase					:			str									= MegatronBertForMaskedLM(config=A      )
									model.to(A      )
									model.eval()
									lowerCamelCase					:			Tuple									= model(A,    attention_mask=A,    token_type_ids=A,    labels=A      )
									self.parent.assertEqual(result.logits.shape,    (self.batch_size, self.seq_length, self.vocab_size)      )
					def     UpperCAmelCase_							(     self,    A,    A,    A,    A,    A,    A,    A      ):

									"""simple docstring"""


									lowerCamelCase					:			str									= MegatronBertForCausalLM(config=A      )
									model.to(A      )
									model.eval()
									lowerCamelCase					:			Optional[Any]									= model(A,    attention_mask=A,    token_type_ids=A,    labels=A      )
									self.parent.assertEqual(result.logits.shape,    (self.batch_size, self.seq_length, self.vocab_size)      )
					def     UpperCAmelCase_							(     self,    A,    A,    A,    A,    A,    A,    A      ):

									"""simple docstring"""


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

									"""simple docstring"""


									lowerCamelCase					:			str									= MegatronBertForPreTraining(config=A      )
									model.to(A      )
									model.eval()
									lowerCamelCase					:			List[str]									= model(
									    A,    attention_mask=A,    token_type_ids=A,    labels=A,    next_sentence_label=A,    )
									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,    A,    A,    A,    A,    A,    A,    A      ):

									"""simple docstring"""


									lowerCamelCase					:			Optional[int]									= MegatronBertForQuestionAnswering(config=A      )
									model.to(A      )
									model.eval()
									lowerCamelCase					:			str									= model(
									    A,    attention_mask=A,    token_type_ids=A,    start_positions=A,    end_positions=A,    )
									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,    A,    A,    A,    A,    A,    A,    A      ):

									"""simple docstring"""


									lowerCamelCase					:			Dict									= self.num_labels
									lowerCamelCase					:			List[str]									= MegatronBertForSequenceClassification(A      )
									model.to(A      )
									model.eval()
									lowerCamelCase					:			Tuple									= model(A,    attention_mask=A,    token_type_ids=A,    labels=A      )
									self.parent.assertEqual(result.logits.shape,    (self.batch_size, self.num_labels)      )
					def     UpperCAmelCase_							(     self,    A,    A,    A,    A,    A,    A,    A      ):

									"""simple docstring"""


									lowerCamelCase					:			Union[str, Any]									= self.num_labels
									lowerCamelCase					:			str									= MegatronBertForTokenClassification(config=A      )
									model.to(A      )
									model.eval()
									lowerCamelCase					:			int									= model(A,    attention_mask=A,    token_type_ids=A,    labels=A      )
									self.parent.assertEqual(result.logits.shape,    (self.batch_size, self.seq_length, self.num_labels)      )
					def     UpperCAmelCase_							(     self,    A,    A,    A,    A,    A,    A,    A      ):

									"""simple docstring"""


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




					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			Tuple									= self.prepare_config_and_inputs()
									(
									    lowerCamelCase 
									)					:			List[str]									= config_and_inputs
									lowerCamelCase					:			Tuple									= {'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
									return config, inputs_dict





@require_torch
class     __snake_case		(   a__    ,   a__    ,   unittest.TestCase):
					_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,    A,    A,    A=False      ):

									"""simple docstring"""


									lowerCamelCase					:			Dict									= super()._prepare_for_class(A,    A,    return_labels=A      )

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

									"""simple docstring"""


									lowerCamelCase					:			Dict									= MegatronBertModelTester(self      )
									lowerCamelCase					:			Tuple									= ConfigTester(self,    config_class=A,    hidden_size=37      )
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


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

									"""simple docstring"""


									lowerCamelCase					:			Dict									= self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_megatron_bert_model(*A      )
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			Optional[Any]									= self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_megatron_bert_for_masked_lm(*A      )
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			Union[str, Any]									= self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*A      )
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			Optional[int]									= self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*A      )
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			Dict									= self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_megatron_bert_for_pretraining(*A      )
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			Any									= self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_megatron_bert_for_question_answering(*A      )
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			str									= self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*A      )




					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			List[str]									= self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_megatron_bert_for_token_classification(*A      )





def    UpperCAmelCase							(  UpperCAmelCase__							:							str):
				return torch.tensor(
				    UpperCAmelCase__       ,							dtype=torch.long       ,							device=UpperCAmelCase__       ,							)


A			      =      1E-4




@require_torch
@require_sentencepiece
@require_tokenizers
class     __snake_case		(   unittest.TestCase):
					@slow
					@unittest.skip('Model is not available.'      )
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			Any									= 'nvidia/megatron-bert-uncased-345m'
									if "MYDIR" in os.environ:
													lowerCamelCase					:			Union[str, Any]									= os.path.join(os.environ['MYDIR'],    A      )
									lowerCamelCase					:			Any									= MegatronBertModel.from_pretrained(A      )
									model.to(A      )
									model.half()
									lowerCamelCase					:			Optional[int]									= _long_tensor([[101, 7110, 1005, 1056, 2023, 1_1333, 1_7413, 1029, 102]]      )
									with torch.no_grad():
													lowerCamelCase					:			Optional[Any]									= model(A      )[0]
									lowerCamelCase					:			Optional[Any]									= torch.Size((1, 9, 1024)      )
									self.assertEqual(output.shape,    A      )

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


'''simple docstring'''






import warnings
from collections import OrderedDict
from typing import Mapping

from packaging import version

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


A			      =      logging.get_logger(__name__)

A			      =      {
    'nvidia/segformer-b0-finetuned-ade-512-512': (
        'https://huggingface.co/nvidia/segformer-b0-finetuned-ade-512-512/resolve/main/config.json'
    ),
    # See all SegFormer models at https://huggingface.co/models?filter=segformer
}




class     __snake_case		(   a__):
					_lowerCAmelCase			     =	'''segformer'''
					def __init__(     self,    A=3,    A=4,    A=[2, 2, 2, 2],    A=[8, 4, 2, 1],    A=[32, 64, 160, 256],    A=[7, 3, 3, 3],    A=[4, 2, 2, 2],    A=[1, 2, 5, 8],    A=[4, 4, 4, 4],    A="gelu",    A=0.0,    A=0.0,    A=0.1,    A=0.02,    A=0.1,    A=1e-6,    A=256,    A=255,    **A,    ):

									"""simple docstring"""


									super().__init__(**A      )

									if "reshape_last_stage" in kwargs and kwargs["reshape_last_stage"] is False:
													warnings.warn(
													    'Reshape_last_stage is set to False in this config. This argument is deprecated and will soon be'
													    ' removed, as the behaviour will default to that of reshape_last_stage = True.',    A,    )

									lowerCamelCase					:			Tuple									= num_channels
									lowerCamelCase					:			Dict									= num_encoder_blocks
									lowerCamelCase					:			Dict									= depths
									lowerCamelCase					:			Any									= sr_ratios
									lowerCamelCase					:			List[Any]									= hidden_sizes
									lowerCamelCase					:			List[Any]									= patch_sizes
									lowerCamelCase					:			Tuple									= strides
									lowerCamelCase					:			int									= mlp_ratios
									lowerCamelCase					:			Dict									= num_attention_heads
									lowerCamelCase					:			Union[str, Any]									= hidden_act
									lowerCamelCase					:			Tuple									= hidden_dropout_prob
									lowerCamelCase					:			Optional[int]									= attention_probs_dropout_prob
									lowerCamelCase					:			List[str]									= classifier_dropout_prob
									lowerCamelCase					:			Union[str, Any]									= initializer_range
									lowerCamelCase					:			Union[str, Any]									= drop_path_rate
									lowerCamelCase					:			int									= layer_norm_eps
									lowerCamelCase					:			Any									= decoder_hidden_size
									lowerCamelCase					:			Optional[int]									= kwargs.get('reshape_last_stage',    A      )
									lowerCamelCase					:			Any									= semantic_loss_ignore_index





class     __snake_case		(   a__):
					_lowerCAmelCase			     =	version.parse('''1.11''')
					@property
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									return OrderedDict(
									    [
									        ('pixel_values', {0: 'batch', 1: 'num_channels', 2: 'height', 3: 'width'}),
									    ]      )
					@property
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									return 1e-4




					@property
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									return 12
