


'''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       ,			__lowercase       ,			__lowercase=13       ,			__lowercase=7       ,			__lowercase=True       ,			__lowercase=True       ,			__lowercase=True       ,			__lowercase=True       ,			__lowercase=99       ,			__lowercase=64       ,			__lowercase=32       ,			__lowercase=5       ,			__lowercase=4       ,			__lowercase=37       ,			__lowercase="gelu"       ,			__lowercase=0.1       ,			__lowercase=0.1       ,			__lowercase=5_12       ,			__lowercase=16       ,			__lowercase=2       ,			__lowercase=0.0_2       ,			__lowercase=3       ,			__lowercase=4       ,			__lowercase=None       ,			):
          __lowerCAmelCase							       =				parent
          __lowerCAmelCase							       =				batch_size
          __lowerCAmelCase							       =				seq_length
          __lowerCAmelCase							       =				is_training
          __lowerCAmelCase							       =				use_input_mask
          __lowerCAmelCase							       =				use_token_type_ids
          __lowerCAmelCase							       =				use_labels
          __lowerCAmelCase							       =				vocab_size
          __lowerCAmelCase							       =				hidden_size
          __lowerCAmelCase							       =				embedding_size
          __lowerCAmelCase							       =				num_hidden_layers
          __lowerCAmelCase							       =				num_attention_heads
          __lowerCAmelCase							       =				intermediate_size
          __lowerCAmelCase							       =				hidden_act
          __lowerCAmelCase							       =				hidden_dropout_prob
          __lowerCAmelCase							       =				attention_probs_dropout_prob
          __lowerCAmelCase							       =				max_position_embeddings
          __lowerCAmelCase							       =				type_vocab_size
          __lowerCAmelCase							       =				type_sequence_label_size
          __lowerCAmelCase							       =				initializer_range
          __lowerCAmelCase							       =				num_labels
          __lowerCAmelCase							       =				num_choices
          __lowerCAmelCase							       =				scope




    def 					_snake_case  (self		):
          __lowerCAmelCase							       =				ids_tensor([self.batch_size, self.seq_length]       ,			self.vocab_size		)

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

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

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

          __lowerCAmelCase							       =				self.get_config()

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




    def 					_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=__UpperCamelCase       ,			initializer_range=self.initializer_range       ,			)




    def 					_snake_case  (self       ,			__lowercase       ,			__lowercase       ,			__lowercase       ,			__lowercase       ,			__lowercase       ,			__lowercase       ,			__lowercase		):
          __lowerCAmelCase							       =				MegatronBertModel(config=__UpperCamelCase		)
          model.to(__UpperCamelCase		)
          model.eval()
          __lowerCAmelCase							       =				model(__UpperCamelCase       ,			attention_mask=__UpperCamelCase       ,			token_type_ids=__UpperCamelCase		)
          __lowerCAmelCase							       =				model(__UpperCamelCase       ,			token_type_ids=__UpperCamelCase		)
          __lowerCAmelCase							       =				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 					_snake_case  (self       ,			__lowercase       ,			__lowercase       ,			__lowercase       ,			__lowercase       ,			__lowercase       ,			__lowercase       ,			__lowercase		):
          __lowerCAmelCase							       =				MegatronBertForMaskedLM(config=__UpperCamelCase		)
          model.to(__UpperCamelCase		)
          model.eval()
          __lowerCAmelCase							       =				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 					_snake_case  (self       ,			__lowercase       ,			__lowercase       ,			__lowercase       ,			__lowercase       ,			__lowercase       ,			__lowercase       ,			__lowercase		):
          __lowerCAmelCase							       =				MegatronBertForCausalLM(config=__UpperCamelCase		)
          model.to(__UpperCamelCase		)
          model.eval()
          __lowerCAmelCase							       =				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 					_snake_case  (self       ,			__lowercase       ,			__lowercase       ,			__lowercase       ,			__lowercase       ,			__lowercase       ,			__lowercase       ,			__lowercase		):
          __lowerCAmelCase							       =				MegatronBertForNextSentencePrediction(config=__UpperCamelCase		)
          model.to(__UpperCamelCase		)
          model.eval()
          __lowerCAmelCase							       =				model(
              __UpperCamelCase       ,			attention_mask=__UpperCamelCase       ,			token_type_ids=__UpperCamelCase       ,			labels=__UpperCamelCase       ,			)
          self.parent.assertEqual(result.logits.shape       ,			(self.batch_size, 2)		)




    def 					_snake_case  (self       ,			__lowercase       ,			__lowercase       ,			__lowercase       ,			__lowercase       ,			__lowercase       ,			__lowercase       ,			__lowercase		):
          __lowerCAmelCase							       =				MegatronBertForPreTraining(config=__UpperCamelCase		)
          model.to(__UpperCamelCase		)
          model.eval()
          __lowerCAmelCase							       =				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 					_snake_case  (self       ,			__lowercase       ,			__lowercase       ,			__lowercase       ,			__lowercase       ,			__lowercase       ,			__lowercase       ,			__lowercase		):
          __lowerCAmelCase							       =				MegatronBertForQuestionAnswering(config=__UpperCamelCase		)
          model.to(__UpperCamelCase		)
          model.eval()
          __lowerCAmelCase							       =				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 					_snake_case  (self       ,			__lowercase       ,			__lowercase       ,			__lowercase       ,			__lowercase       ,			__lowercase       ,			__lowercase       ,			__lowercase		):
          __lowerCAmelCase							       =				self.num_labels
          __lowerCAmelCase							       =				MegatronBertForSequenceClassification(__UpperCamelCase		)
          model.to(__UpperCamelCase		)
          model.eval()
          __lowerCAmelCase							       =				model(__UpperCamelCase       ,			attention_mask=__UpperCamelCase       ,			token_type_ids=__UpperCamelCase       ,			labels=__UpperCamelCase		)
          self.parent.assertEqual(result.logits.shape       ,			(self.batch_size, self.num_labels)		)




    def 					_snake_case  (self       ,			__lowercase       ,			__lowercase       ,			__lowercase       ,			__lowercase       ,			__lowercase       ,			__lowercase       ,			__lowercase		):
          __lowerCAmelCase							       =				self.num_labels
          __lowerCAmelCase							       =				MegatronBertForTokenClassification(config=__UpperCamelCase		)
          model.to(__UpperCamelCase		)
          model.eval()
          __lowerCAmelCase							       =				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 					_snake_case  (self       ,			__lowercase       ,			__lowercase       ,			__lowercase       ,			__lowercase       ,			__lowercase       ,			__lowercase       ,			__lowercase		):
          __lowerCAmelCase							       =				self.num_choices
          __lowerCAmelCase							       =				MegatronBertForMultipleChoice(config=__UpperCamelCase		)
          model.to(__UpperCamelCase		)
          model.eval()
          __lowerCAmelCase							       =				input_ids.unsqueeze(1		).expand(-1       ,			self.num_choices       ,			-1		).contiguous()
          __lowerCAmelCase							       =				token_type_ids.unsqueeze(1		).expand(-1       ,			self.num_choices       ,			-1		).contiguous()
          __lowerCAmelCase							       =				input_mask.unsqueeze(1		).expand(-1       ,			self.num_choices       ,			-1		).contiguous()
          __lowerCAmelCase							       =				model(
              __UpperCamelCase       ,			attention_mask=__UpperCamelCase       ,			token_type_ids=__UpperCamelCase       ,			labels=__UpperCamelCase       ,			)
          self.parent.assertEqual(result.logits.shape       ,			(self.batch_size, self.num_choices)		)




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


@require_torch
class       a__				( __snake_case							,						__snake_case							,						unittest.TestCase      ):







    """simple docstring"""



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




    def 					_snake_case  (self       ,			__lowercase       ,			__lowercase       ,			__lowercase=False		):
          __lowerCAmelCase							       =				super()._prepare_for_class(__UpperCamelCase       ,			__UpperCamelCase       ,			return_labels=__UpperCamelCase		)

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




    def 					_snake_case  (self		):
          __lowerCAmelCase							       =				MegatronBertModelTester(self		)
          __lowerCAmelCase							       =				ConfigTester(self       ,			config_class=__UpperCamelCase       ,			hidden_size=37		)




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




    def 					_snake_case  (self		):
          __lowerCAmelCase							       =				self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_megatron_bert_model(*__UpperCamelCase		)




    def 					_snake_case  (self		):
          __lowerCAmelCase							       =				self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_megatron_bert_for_masked_lm(*__UpperCamelCase		)




    def 					_snake_case  (self		):
          __lowerCAmelCase							       =				self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*__UpperCamelCase		)




    def 					_snake_case  (self		):
          __lowerCAmelCase							       =				self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*__UpperCamelCase		)




    def 					_snake_case  (self		):
          __lowerCAmelCase							       =				self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_megatron_bert_for_pretraining(*__UpperCamelCase		)




    def 					_snake_case  (self		):
          __lowerCAmelCase							       =				self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_megatron_bert_for_question_answering(*__UpperCamelCase		)




    def 					_snake_case  (self		):
          __lowerCAmelCase							       =				self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*__UpperCamelCase		)




    def 					_snake_case  (self		):
          __lowerCAmelCase							       =				self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_megatron_bert_for_token_classification(*__UpperCamelCase		)

def 							__magic_name__(						lowerCamelCase):
      return torch.tensor(
          lowercase__,				dtype=torch.long,				device=lowercase__,				)


_UpperCAmelCase				: str							    =				1E-4




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







    """simple docstring"""



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

          __lowerCAmelCase							       =				[-0.6_0_4_0, -0.2_5_1_7, -0.1_0_2_5, 0.3_4_2_0, -0.6_7_5_8, -0.0_0_1_7, -0.1_0_8_9, -0.1_9_9_0, 0.5_7_2_8]
          for ii in range(3		):
                for jj in range(3		):
                      __lowerCAmelCase							       =				output[0, ii, jj]
                      __lowerCAmelCase							       =				expected[3 * ii + jj]
                      __lowerCAmelCase							       =				'''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'''

from PIL import Image

def 							__magic_name__(						lowerCamelCase,				lowerCamelCase):

						def brightness(lowerCamelCase) -> float:
												return 1_2_8 + level + (c - 1_2_8)

						if not -2_55.0 <= level <= 2_55.0:
												raise ValueError('''level must be between -255.0 (black) and 255.0 (white)''')
						return img.point(lowerCamelCase)


if __name__ == "__main__":
	# Load image
	with Image.open("""image_data/lena.jpg""") as img:
		# Change brightness to 100
		_UpperCAmelCase				: str							    =				change_brightness(img, 1_0_0)
		brigt_img.save("""image_data/lena_brightness.png""", format="""png""")

