
"""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 			_SCREAMING_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.0_2       ,   __A=3       ,   __A=4       ,   __A=None       ,   )			->     Optional[Any]:
						lowerCAmelCase_    :Optional[int]			     =  parent
						lowerCAmelCase_    :Any			     =  batch_size
						lowerCAmelCase_    :Optional[Any]			     =  seq_length
						lowerCAmelCase_    :List[str]			     =  is_training
						lowerCAmelCase_    :List[Any]			     =  use_input_mask
						lowerCAmelCase_    :List[Any]			     =  use_token_type_ids
						lowerCAmelCase_    :Any			     =  use_labels
						lowerCAmelCase_    :str			     =  vocab_size
						lowerCAmelCase_    :str			     =  hidden_size
						lowerCAmelCase_    :List[str]			     =  embedding_size
						lowerCAmelCase_    :Union[str, Any]			     =  num_hidden_layers
						lowerCAmelCase_    :Dict			     =  num_attention_heads
						lowerCAmelCase_    :Dict			     =  intermediate_size
						lowerCAmelCase_    :int			     =  hidden_act
						lowerCAmelCase_    :Optional[int]			     =  hidden_dropout_prob
						lowerCAmelCase_    :Tuple			     =  attention_probs_dropout_prob
						lowerCAmelCase_    :List[str]			     =  max_position_embeddings
						lowerCAmelCase_    :List[str]			     =  type_vocab_size
						lowerCAmelCase_    :List[Any]			     =  type_sequence_label_size
						lowerCAmelCase_    :str			     =  initializer_range
						lowerCAmelCase_    :List[Any]			     =  num_labels
						lowerCAmelCase_    :Tuple			     =  num_choices
						lowerCAmelCase_    :Any			     =  scope






			def  __lowerCAmelCase (					self       )			->     List[str]:
						lowerCAmelCase_    :Any			     =  ids_tensor([self.batch_size, self.seq_length]       ,   self.vocab_size       )

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

						lowerCAmelCase_    :Optional[int]			     =  None
						if self.use_token_type_ids:
									lowerCAmelCase_    :Dict			     =  ids_tensor([self.batch_size, self.seq_length]       ,   self.type_vocab_size       )

						lowerCAmelCase_    :Union[str, Any]			     =  None
						lowerCAmelCase_    :Tuple			     =  None
						lowerCAmelCase_    :str			     =  None
						if self.use_labels:
									lowerCAmelCase_    :Optional[int]			     =  ids_tensor([self.batch_size]       ,   self.type_sequence_label_size       )
									lowerCAmelCase_    :List[Any]			     =  ids_tensor([self.batch_size, self.seq_length]       ,   self.num_labels       )
									lowerCAmelCase_    :Optional[int]			     =  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  __lowerCAmelCase (					self       )			->     Optional[Any]:
						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=SCREAMING_SNAKE_CASE_       ,   initializer_range=self.initializer_range       ,   )






			def  __lowerCAmelCase (					self       ,   __A       ,   __A       ,   __A       ,   __A       ,   __A       ,   __A       ,   __A       )			->     List[str]:
						lowerCAmelCase_    :Union[str, Any]			     =  MegatronBertModel(config=SCREAMING_SNAKE_CASE_       )
						model.to(SCREAMING_SNAKE_CASE_       )
						model.eval()
						lowerCAmelCase_    :Union[str, Any]			     =  model(SCREAMING_SNAKE_CASE_       ,   attention_mask=SCREAMING_SNAKE_CASE_       ,   token_type_ids=SCREAMING_SNAKE_CASE_       )
						lowerCAmelCase_    :Dict			     =  model(SCREAMING_SNAKE_CASE_       ,   token_type_ids=SCREAMING_SNAKE_CASE_       )
						lowerCAmelCase_    :List[str]			     =  model(SCREAMING_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  __lowerCAmelCase (					self       ,   __A       ,   __A       ,   __A       ,   __A       ,   __A       ,   __A       ,   __A       )			->     Dict:
						lowerCAmelCase_    :List[str]			     =  MegatronBertForMaskedLM(config=SCREAMING_SNAKE_CASE_       )
						model.to(SCREAMING_SNAKE_CASE_       )
						model.eval()
						lowerCAmelCase_    :Tuple			     =  model(SCREAMING_SNAKE_CASE_       ,   attention_mask=SCREAMING_SNAKE_CASE_       ,   token_type_ids=SCREAMING_SNAKE_CASE_       ,   labels=SCREAMING_SNAKE_CASE_       )
						self.parent.assertEqual(result.logits.shape       ,   (self.batch_size, self.seq_length, self.vocab_size)       )






			def  __lowerCAmelCase (					self       ,   __A       ,   __A       ,   __A       ,   __A       ,   __A       ,   __A       ,   __A       )			->     Any:
						lowerCAmelCase_    :List[Any]			     =  MegatronBertForCausalLM(config=SCREAMING_SNAKE_CASE_       )
						model.to(SCREAMING_SNAKE_CASE_       )
						model.eval()
						lowerCAmelCase_    :Optional[int]			     =  model(SCREAMING_SNAKE_CASE_       ,   attention_mask=SCREAMING_SNAKE_CASE_       ,   token_type_ids=SCREAMING_SNAKE_CASE_       ,   labels=SCREAMING_SNAKE_CASE_       )
						self.parent.assertEqual(result.logits.shape       ,   (self.batch_size, self.seq_length, self.vocab_size)       )






			def  __lowerCAmelCase (					self       ,   __A       ,   __A       ,   __A       ,   __A       ,   __A       ,   __A       ,   __A       )			->     str:
						lowerCAmelCase_    :Dict			     =  MegatronBertForNextSentencePrediction(config=SCREAMING_SNAKE_CASE_       )
						model.to(SCREAMING_SNAKE_CASE_       )
						model.eval()
						lowerCAmelCase_    :Optional[Any]			     =  model(
						    SCREAMING_SNAKE_CASE_       ,   attention_mask=SCREAMING_SNAKE_CASE_       ,   token_type_ids=SCREAMING_SNAKE_CASE_       ,   labels=SCREAMING_SNAKE_CASE_       ,   )
						self.parent.assertEqual(result.logits.shape       ,   (self.batch_size, 2)       )






			def  __lowerCAmelCase (					self       ,   __A       ,   __A       ,   __A       ,   __A       ,   __A       ,   __A       ,   __A       )			->     List[Any]:
						lowerCAmelCase_    :int			     =  MegatronBertForPreTraining(config=SCREAMING_SNAKE_CASE_       )
						model.to(SCREAMING_SNAKE_CASE_       )
						model.eval()
						lowerCAmelCase_    :int			     =  model(
						    SCREAMING_SNAKE_CASE_       ,   attention_mask=SCREAMING_SNAKE_CASE_       ,   token_type_ids=SCREAMING_SNAKE_CASE_       ,   labels=SCREAMING_SNAKE_CASE_       ,   next_sentence_label=SCREAMING_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  __lowerCAmelCase (					self       ,   __A       ,   __A       ,   __A       ,   __A       ,   __A       ,   __A       ,   __A       )			->     Optional[int]:
						lowerCAmelCase_    :List[str]			     =  MegatronBertForQuestionAnswering(config=SCREAMING_SNAKE_CASE_       )
						model.to(SCREAMING_SNAKE_CASE_       )
						model.eval()
						lowerCAmelCase_    :int			     =  model(
						    SCREAMING_SNAKE_CASE_       ,   attention_mask=SCREAMING_SNAKE_CASE_       ,   token_type_ids=SCREAMING_SNAKE_CASE_       ,   start_positions=SCREAMING_SNAKE_CASE_       ,   end_positions=SCREAMING_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  __lowerCAmelCase (					self       ,   __A       ,   __A       ,   __A       ,   __A       ,   __A       ,   __A       ,   __A       )			->     List[Any]:
						lowerCAmelCase_    :List[str]			     =  self.num_labels
						lowerCAmelCase_    :Any			     =  MegatronBertForSequenceClassification(SCREAMING_SNAKE_CASE_       )
						model.to(SCREAMING_SNAKE_CASE_       )
						model.eval()
						lowerCAmelCase_    :Optional[int]			     =  model(SCREAMING_SNAKE_CASE_       ,   attention_mask=SCREAMING_SNAKE_CASE_       ,   token_type_ids=SCREAMING_SNAKE_CASE_       ,   labels=SCREAMING_SNAKE_CASE_       )
						self.parent.assertEqual(result.logits.shape       ,   (self.batch_size, self.num_labels)       )






			def  __lowerCAmelCase (					self       ,   __A       ,   __A       ,   __A       ,   __A       ,   __A       ,   __A       ,   __A       )			->     int:
						lowerCAmelCase_    :Union[str, Any]			     =  self.num_labels
						lowerCAmelCase_    :Optional[int]			     =  MegatronBertForTokenClassification(config=SCREAMING_SNAKE_CASE_       )
						model.to(SCREAMING_SNAKE_CASE_       )
						model.eval()
						lowerCAmelCase_    :Dict			     =  model(SCREAMING_SNAKE_CASE_       ,   attention_mask=SCREAMING_SNAKE_CASE_       ,   token_type_ids=SCREAMING_SNAKE_CASE_       ,   labels=SCREAMING_SNAKE_CASE_       )
						self.parent.assertEqual(result.logits.shape       ,   (self.batch_size, self.seq_length, self.num_labels)       )






			def  __lowerCAmelCase (					self       ,   __A       ,   __A       ,   __A       ,   __A       ,   __A       ,   __A       ,   __A       )			->     Optional[int]:
						lowerCAmelCase_    :Optional[Any]			     =  self.num_choices
						lowerCAmelCase_    :str			     =  MegatronBertForMultipleChoice(config=SCREAMING_SNAKE_CASE_       )
						model.to(SCREAMING_SNAKE_CASE_       )
						model.eval()
						lowerCAmelCase_    :Dict			     =  input_ids.unsqueeze(1       ).expand(-1       ,   self.num_choices       ,   -1       ).contiguous()
						lowerCAmelCase_    :Tuple			     =  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[Any]			     =  model(
						    SCREAMING_SNAKE_CASE_       ,   attention_mask=SCREAMING_SNAKE_CASE_       ,   token_type_ids=SCREAMING_SNAKE_CASE_       ,   labels=SCREAMING_SNAKE_CASE_       ,   )
						self.parent.assertEqual(result.logits.shape       ,   (self.batch_size, self.num_choices)       )






			def  __lowerCAmelCase (					self       )			->     Tuple:
						lowerCAmelCase_    :int			     =  self.prepare_config_and_inputs()
						(
						    lowerCAmelCase_       
						)    :Tuple			     =  config_and_inputs
						lowerCAmelCase_    :Dict			     =  {'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
						return config, inputs_dict




@require_torch
class 			_SCREAMING_SNAKE_CASE					(			_UpperCAmelCase	, _UpperCAmelCase	, unittest.TestCase							):
			UpperCAmelCase_  :Dict				 =						(
			    (
			        MegatronBertModel,
			        MegatronBertForMaskedLM,
			        MegatronBertForCausalLM,
			        MegatronBertForMultipleChoice,
			        MegatronBertForNextSentencePrediction,
			        MegatronBertForPreTraining,
			        MegatronBertForQuestionAnswering,
			        MegatronBertForSequenceClassification,
			        MegatronBertForTokenClassification,
			    )
			    if is_torch_available()
			    else ()
			)
			UpperCAmelCase_  :Union[str, 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_  :List[str]				 =						False






			def  __lowerCAmelCase (					self       ,   __A       ,   __A       ,   __A=False       )			->     Union[str, Any]:
						lowerCAmelCase_    :Union[str, Any]			     =  super()._prepare_for_class(SCREAMING_SNAKE_CASE_       ,   SCREAMING_SNAKE_CASE_       ,   return_labels=SCREAMING_SNAKE_CASE_       )

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






			def  __lowerCAmelCase (					self       )			->     Union[str, Any]:
						lowerCAmelCase_    :Union[str, Any]			     =  MegatronBertModelTester(self       )
						lowerCAmelCase_    :Union[str, Any]			     =  ConfigTester(self       ,   config_class=SCREAMING_SNAKE_CASE_       ,   hidden_size=37       )






			def  __lowerCAmelCase (					self       )			->     Any:
						self.config_tester.run_common_tests()






			def  __lowerCAmelCase (					self       )			->     Dict:
						lowerCAmelCase_    :Optional[Any]			     =  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_model(*SCREAMING_SNAKE_CASE_       )






			def  __lowerCAmelCase (					self       )			->     Optional[Any]:
						lowerCAmelCase_    :Tuple			     =  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_for_masked_lm(*SCREAMING_SNAKE_CASE_       )






			def  __lowerCAmelCase (					self       )			->     List[str]:
						lowerCAmelCase_    :Optional[int]			     =  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*SCREAMING_SNAKE_CASE_       )






			def  __lowerCAmelCase (					self       )			->     int:
						lowerCAmelCase_    :Dict			     =  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*SCREAMING_SNAKE_CASE_       )






			def  __lowerCAmelCase (					self       )			->     Optional[Any]:
						lowerCAmelCase_    :Any			     =  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_for_pretraining(*SCREAMING_SNAKE_CASE_       )






			def  __lowerCAmelCase (					self       )			->     Union[str, Any]:
						lowerCAmelCase_    :Dict			     =  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_for_question_answering(*SCREAMING_SNAKE_CASE_       )






			def  __lowerCAmelCase (					self       )			->     str:
						lowerCAmelCase_    :Dict			     =  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*SCREAMING_SNAKE_CASE_       )






			def  __lowerCAmelCase (					self       )			->     List[str]:
						lowerCAmelCase_    :Any			     =  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_for_token_classification(*SCREAMING_SNAKE_CASE_       )






def 						_snake_case       (	lowercase__		:  Optional[int]		)  ->				Any:


			'''simple docstring'''





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


__UpperCAmelCase       =  1e-4




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






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

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










'''simple docstring'''



import numpy as np
from cva import destroyAllWindows, imread, imshow, waitKey





class    UpperCAmelCase_  :





    """simple docstring"""






    def __init__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     int:
         if dst_width < 0 or dst_height < 0:
              raise ValueError('Destination width/height should be > 0'	)

         __lowerCamelCase   :	Dict								=		img
         __lowerCamelCase   :	Any								=		img.shape[1]
         __lowerCamelCase   :	Optional[int]								=		img.shape[0]
         __lowerCamelCase   :	Dict								=		dst_width
         __lowerCamelCase   :	str								=		dst_height

         __lowerCamelCase   :	Dict								=		self.src_w / self.dst_w
         __lowerCamelCase   :	List[Any]								=		self.src_h / self.dst_h

         __lowerCamelCase   :	Optional[int]								=		(
             np.ones((self.dst_h, self.dst_w, 3)			,     np.uinta	) * 2_55
         )







    def       lowercase_				(							self	)				->     List[Any]:
         for i in range(self.dst_h	):
              for j in range(self.dst_w	):
                   __lowerCamelCase   :	Union[str, Any]								=		self.img[self.get_y(SCREAMING_SNAKE_CASE_	)][self.get_x(SCREAMING_SNAKE_CASE_	)]







    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_	)				->     int:
         return int(self.ratio_x * x	)







    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_	)				->     int:
         return int(self.ratio_y * y	)


if __name__ == "__main__":
      A__						,					A__		:					Optional[Any]												=  800, 600
      A__		:					List[str]												=  imread("""image_data/lena.jpg""", 1)
      A__		:					List[Any]												=  NearestNeighbour(im, dst_w, dst_h)
      n.process()

      imshow(
          f'''Image resized from: {im.shape[1]}x{im.shape[0]} to {dst_w}x{dst_h}''', n.output
      )
      waitKey(0)
      destroyAllWindows()






