


import inspect
import unittest
import warnings
from math import ceil, floor

from transformers import LevitConfig
from transformers.file_utils import cached_property, is_torch_available, is_vision_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, require_vision, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
						import torch

						from transformers import (
						    MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING,
						    MODEL_MAPPING,
						    LevitForImageClassification,
						    LevitForImageClassificationWithTeacher,
						    LevitModel,
						)
						from transformers.models.levit.modeling_levit import LEVIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
						from PIL import Image

						from transformers import LevitImageProcessor







class        lowerCamelCase__	(							UpperCAmelCase      ):






			"""simple docstring"""






			def 				snake_case__				(							self     ):

						'''simple docstring'''


						UpperCamelCase__					=       self.config_class(**self.inputs_dict     )
						self.parent.assertTrue(hasattr(snake_case      ,						"hidden_sizes"     )     )
						self.parent.assertTrue(hasattr(snake_case      ,						"num_attention_heads"     )     )







class        lowerCamelCase__	:






			"""simple docstring"""






			def __init__(							self      ,						snake_case      ,						snake_case=13      ,						snake_case=64      ,						snake_case=3      ,						snake_case=3      ,						snake_case=2      ,						snake_case=1      ,						snake_case=16      ,						snake_case=[128, 256, 384]      ,						snake_case=[4, 6, 8]      ,						snake_case=[2, 3, 4]      ,						snake_case=[16, 16, 16]      ,						snake_case=0      ,						snake_case=[2, 2, 2]      ,						snake_case=[2, 2, 2]      ,						snake_case=0.02      ,						snake_case=True      ,						snake_case=True      ,						snake_case=2      ,						):

						'''simple docstring'''


						UpperCamelCase__					=       parent
						UpperCamelCase__					=       batch_size
						UpperCamelCase__					=       image_size
						UpperCamelCase__					=       num_channels
						UpperCamelCase__					=       kernel_size
						UpperCamelCase__					=       stride
						UpperCamelCase__					=       padding
						UpperCamelCase__					=       hidden_sizes
						UpperCamelCase__					=       num_attention_heads
						UpperCamelCase__					=       depths
						UpperCamelCase__					=       key_dim
						UpperCamelCase__					=       drop_path_rate
						UpperCamelCase__					=       patch_size
						UpperCamelCase__					=       attention_ratio
						UpperCamelCase__					=       mlp_ratio
						UpperCamelCase__					=       initializer_range
						UpperCamelCase__					=       [
						    ["Subsample", key_dim[0], hidden_sizes[0] // key_dim[0], 4, 2, 2],
						    ["Subsample", key_dim[0], hidden_sizes[1] // key_dim[0], 4, 2, 2],
						]
						UpperCamelCase__					=       is_training
						UpperCamelCase__					=       use_labels
						UpperCamelCase__					=       num_labels
						UpperCamelCase__					=       initializer_range




			def 				snake_case__				(							self     ):

						'''simple docstring'''


						UpperCamelCase__					=       floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]     )

						UpperCamelCase__					=       None
						if self.use_labels:
									UpperCamelCase__					=       ids_tensor([self.batch_size]      ,						self.num_labels     )

						UpperCamelCase__					=       self.get_config()
						return config, pixel_values, labels




			def 				snake_case__				(							self     ):

						'''simple docstring'''


						return LevitConfig(
						    image_size=self.image_size      ,						num_channels=self.num_channels      ,						kernel_size=self.kernel_size      ,						stride=self.stride      ,						padding=self.padding      ,						patch_size=self.patch_size      ,						hidden_sizes=self.hidden_sizes      ,						num_attention_heads=self.num_attention_heads      ,						depths=self.depths      ,						key_dim=self.key_dim      ,						drop_path_rate=self.drop_path_rate      ,						mlp_ratio=self.mlp_ratio      ,						attention_ratio=self.attention_ratio      ,						initializer_range=self.initializer_range      ,						down_ops=self.down_ops      ,						)




			def 				snake_case__				(							self      ,						snake_case      ,						snake_case      ,						snake_case     ):

						'''simple docstring'''


						UpperCamelCase__					=       LevitModel(config=snake_case     )
						model.to(snake_case     )
						model.eval()
						UpperCamelCase__					=       model(snake_case     )
						UpperCamelCase__					=       (self.image_size, self.image_size)
						UpperCamelCase__,      UpperCamelCase__					=       image_size[0], image_size[1]
						for _ in range(4     ):
									UpperCamelCase__					=       floor(((height + 2 * self.padding - self.kernel_size) / self.stride) + 1     )
									UpperCamelCase__					=       floor(((width + 2 * self.padding - self.kernel_size) / self.stride) + 1     )
						self.parent.assertEqual(
						    result.last_hidden_state.shape      ,						(self.batch_size, ceil(height / 4     ) * ceil(width / 4     ), self.hidden_sizes[-1])      ,						)




			def 				snake_case__				(							self      ,						snake_case      ,						snake_case      ,						snake_case     ):

						'''simple docstring'''


						UpperCamelCase__					=       self.num_labels
						UpperCamelCase__					=       LevitForImageClassification(snake_case     )
						model.to(snake_case     )
						model.eval()
						UpperCamelCase__					=       model(snake_case      ,						labels=snake_case     )
						self.parent.assertEqual(result.logits.shape      ,						(self.batch_size, self.num_labels)     )







			def 				snake_case__				(							self     ):

						'''simple docstring'''


						UpperCamelCase__					=       self.prepare_config_and_inputs()
						UpperCamelCase__,      UpperCamelCase__,      UpperCamelCase__					=       config_and_inputs
						UpperCamelCase__					=       {"pixel_values": pixel_values}
						return config, inputs_dict







@require_torch
class        lowerCamelCase__	(							UpperCAmelCase					,	UpperCAmelCase					,	unittest.TestCase      ):






			"""simple docstring"""






			_UpperCamelCase		:					Dict    =		(
			    (LevitModel, LevitForImageClassification, LevitForImageClassificationWithTeacher)
			    if is_torch_available()
			    else ()
			)
			_UpperCamelCase		:					Tuple    =		(
			    {
			        'feature-extraction': LevitModel,
			        'image-classification': (LevitForImageClassification, LevitForImageClassificationWithTeacher),
			    }
			    if is_torch_available()
			    else {}
			)

			_UpperCamelCase		:					Tuple    =		False
			_UpperCamelCase		:					Dict    =		False
			_UpperCamelCase		:					Union[str, Any]    =		False
			_UpperCamelCase		:					Tuple    =		False
			_UpperCamelCase		:					Dict    =		False




			def 				snake_case__				(							self     ):

						'''simple docstring'''


						UpperCamelCase__					=       LevitModelTester(self     )
						UpperCamelCase__					=       ConfigTester(self      ,						config_class=snake_case      ,						has_text_modality=snake_case      ,						hidden_size=37     )




			def 				snake_case__				(							self     ):

						'''simple docstring'''


						self.create_and_test_config_common_properties()
						self.config_tester.create_and_test_config_to_json_string()
						self.config_tester.create_and_test_config_to_json_file()
						self.config_tester.create_and_test_config_from_and_save_pretrained()
						self.config_tester.create_and_test_config_with_num_labels()
						self.config_tester.check_config_can_be_init_without_params()
						self.config_tester.check_config_arguments_init()




			def 				snake_case__				(							self     ):

						'''simple docstring'''


						return




			@unittest.skip(reason="Levit does not use inputs_embeds"     )
			def 				snake_case__				(							self     ):

						'''simple docstring'''


						pass




			@unittest.skip(reason="Levit does not support input and output embeddings"     )
			def 				snake_case__				(							self     ):

						'''simple docstring'''


						pass




			@unittest.skip(reason="Levit does not output attentions"     )
			def 				snake_case__				(							self     ):

						'''simple docstring'''


						pass




			def 				snake_case__				(							self     ):

						'''simple docstring'''


						UpperCamelCase__,      UpperCamelCase__					=       self.model_tester.prepare_config_and_inputs_for_common()

						for model_class in self.all_model_classes:
									UpperCamelCase__					=       model_class(snake_case     )
									UpperCamelCase__					=       inspect.signature(model.forward     )
									# signature.parameters is an OrderedDict => so arg_names order is deterministic
									UpperCamelCase__					=       [*signature.parameters.keys()]

									UpperCamelCase__					=       ["pixel_values"]
									self.assertListEqual(arg_names[:1]      ,						snake_case     )




			def 				snake_case__				(							self     ):

						'''simple docstring'''


						def check_hidden_states_output(snake_case      ,						snake_case      ,						snake_case     ):
									UpperCamelCase__					=       model_class(snake_case     )
									model.to(snake_case     )
									model.eval()

									with torch.no_grad():
												UpperCamelCase__					=       model(**self._prepare_for_class(snake_case      ,						snake_case     )     )

									UpperCamelCase__					=       outputs.hidden_states

									UpperCamelCase__					=       len(self.model_tester.depths     ) + 1
									self.assertEqual(len(snake_case     )      ,						snake_case     )

									UpperCamelCase__					=       (self.model_tester.image_size, self.model_tester.image_size)
									UpperCamelCase__,      UpperCamelCase__					=       image_size[0], image_size[1]
									for _ in range(4     ):
												UpperCamelCase__					=       floor(
												    (
												        (height + 2 * self.model_tester.padding - self.model_tester.kernel_size)
												        / self.model_tester.stride
												    )
												    + 1     )
												UpperCamelCase__					=       floor(
												    (
												        (width + 2 * self.model_tester.padding - self.model_tester.kernel_size)
												        / self.model_tester.stride
												    )
												    + 1     )
									# verify the first hidden states (first block)
									self.assertListEqual(
									    list(hidden_states[0].shape[-2:]     )      ,						[
									        height * width,
									        self.model_tester.hidden_sizes[0],
									    ]      ,						)

						UpperCamelCase__,      UpperCamelCase__					=       self.model_tester.prepare_config_and_inputs_for_common()

						for model_class in self.all_model_classes:
									UpperCamelCase__					=       True
									check_hidden_states_output(snake_case      ,						snake_case      ,						snake_case     )

									# check that output_hidden_states also work using config
									del inputs_dict["output_hidden_states"]
									UpperCamelCase__					=       True

									check_hidden_states_output(snake_case      ,						snake_case      ,						snake_case     )




			@unittest.skip("Will be fixed soon by reducing the size of the model used for common tests."     )
			def 				snake_case__				(							self     ):

						'''simple docstring'''


						pass




			def 				snake_case__				(							self      ,						snake_case      ,						snake_case      ,						snake_case=False     ):

						'''simple docstring'''


						UpperCamelCase__					=       super()._prepare_for_class(snake_case      ,						snake_case      ,						return_labels=snake_case     )

						if return_labels:
									if model_class.__name__ == "LevitForImageClassificationWithTeacher":
												del inputs_dict["labels"]

						return inputs_dict




			def 				snake_case__				(							self     ):

						'''simple docstring'''


						UpperCamelCase__					=       self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*snake_case     )




			def 				snake_case__				(							self     ):

						'''simple docstring'''


						UpperCamelCase__					=       self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_image_classification(*snake_case     )




			def 				snake_case__				(							self     ):

						'''simple docstring'''


						if not self.model_tester.is_training:
									return

						UpperCamelCase__,      UpperCamelCase__					=       self.model_tester.prepare_config_and_inputs_for_common()
						UpperCamelCase__					=       True

						for model_class in self.all_model_classes:
									# LevitForImageClassificationWithTeacher supports inference-only
									if (
									    model_class in get_values(snake_case     )
									    or model_class.__name__ == "LevitForImageClassificationWithTeacher"
									):
												continue
									UpperCamelCase__					=       model_class(snake_case     )
									model.to(snake_case     )
									model.train()
									UpperCamelCase__					=       self._prepare_for_class(snake_case      ,						snake_case      ,						return_labels=snake_case     )
									UpperCamelCase__					=       model(**snake_case     ).loss
									loss.backward()




			def 				snake_case__				(							self     ):

						'''simple docstring'''


						UpperCamelCase__,      UpperCamelCase__					=       self.model_tester.prepare_config_and_inputs_for_common()
						if not self.model_tester.is_training:
									return

						UpperCamelCase__					=       False
						UpperCamelCase__					=       True

						for model_class in self.all_model_classes:
									if model_class in get_values(snake_case     ) or not model_class.supports_gradient_checkpointing:
												continue
									# LevitForImageClassificationWithTeacher supports inference-only
									if model_class.__name__ == "LevitForImageClassificationWithTeacher":
												continue
									UpperCamelCase__					=       model_class(snake_case     )
									model.gradient_checkpointing_enable()
									model.to(snake_case     )
									model.train()
									UpperCamelCase__					=       self._prepare_for_class(snake_case      ,						snake_case      ,						return_labels=snake_case     )
									UpperCamelCase__					=       model(**snake_case     ).loss
									loss.backward()




			def 				snake_case__				(							self     ):

						'''simple docstring'''


						UpperCamelCase__,      UpperCamelCase__					=       self.model_tester.prepare_config_and_inputs_for_common()

						UpperCamelCase__					=       [
						    {"title": "multi_label_classification", "num_labels": 2, "dtype": torch.float},
						    {"title": "single_label_classification", "num_labels": 1, "dtype": torch.long},
						    {"title": "regression", "num_labels": 1, "dtype": torch.float},
						]

						for model_class in self.all_model_classes:
									if (
									    model_class
									    not in [
									        *get_values(snake_case     ),
									    ]
									    or model_class.__name__ == "LevitForImageClassificationWithTeacher"
									):
												continue

									for problem_type in problem_types:
												with self.subTest(msg=F'''Testing {model_class} with {problem_type['title']}'''     ):
															UpperCamelCase__					=       problem_type["title"]
															UpperCamelCase__					=       problem_type["num_labels"]

															UpperCamelCase__					=       model_class(snake_case     )
															model.to(snake_case     )
															model.train()

															UpperCamelCase__					=       self._prepare_for_class(snake_case      ,						snake_case      ,						return_labels=snake_case     )

															if problem_type["num_labels"] > 1:
																		UpperCamelCase__					=       inputs["labels"].unsqueeze(1     ).repeat(1      ,						problem_type["num_labels"]     )

															UpperCamelCase__					=       inputs["labels"].to(problem_type["dtype"]     )

															# This tests that we do not trigger the warning form PyTorch "Using a target size that is different
															# to the input size. This will likely lead to incorrect results due to broadcasting. Please ensure
															# they have the same size." which is a symptom something in wrong for the regression problem.
															# See https://github.com/huggingface/transformers/issues/11780
															with warnings.catch_warnings(record=snake_case     ) as warning_list:
																		UpperCamelCase__					=       model(**snake_case     ).loss
															for w in warning_list:
																		if "Using a target size that is different to the input size" in str(w.message     ):
																					raise ValueError(
																					    F'''Something is going wrong in the regression problem: intercepted {w.message}'''     )

															loss.backward()







			@slow
			def 				snake_case__				(							self     ):

						'''simple docstring'''


						for model_name in LEVIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
									UpperCamelCase__					=       LevitModel.from_pretrained(snake_case     )
									self.assertIsNotNone(snake_case     )



def        UpperCamelCase_(				)->		Tuple:
			UpperCamelCase__					=       Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png")
			return image







@require_torch
@require_vision
class        lowerCamelCase__	(							unittest.TestCase      ):






			"""simple docstring"""






			@cached_property
			def 				snake_case__				(							self     ):

						'''simple docstring'''


						return LevitImageProcessor.from_pretrained(LEVIT_PRETRAINED_MODEL_ARCHIVE_LIST[0]     )







			@slow
			def 				snake_case__				(							self     ):

						'''simple docstring'''


						UpperCamelCase__					=       LevitForImageClassificationWithTeacher.from_pretrained(LEVIT_PRETRAINED_MODEL_ARCHIVE_LIST[0]     ).to(
						    snake_case     )

						UpperCamelCase__					=       self.default_image_processor
						UpperCamelCase__					=       prepare_img()
						UpperCamelCase__					=       image_processor(images=snake_case      ,						return_tensors="pt"     ).to(snake_case     )

						# forward pass
						with torch.no_grad():
									UpperCamelCase__					=       model(**snake_case     )

						# verify the logits
						UpperCamelCase__					=       torch.Size((1, 1000)     )
						self.assertEqual(outputs.logits.shape      ,						snake_case     )

						UpperCamelCase__					=       torch.tensor([1.0448, -0.3745, -1.8317]     ).to(snake_case     )

						self.assertTrue(torch.allclose(outputs.logits[0, :3]      ,						snake_case      ,						atol=1E-4     )     )






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      ,						snake_case      ,						snake_case=13      ,						snake_case=7      ,						snake_case=True      ,						snake_case=True      ,						snake_case=True      ,						snake_case=True      ,						snake_case=99      ,						snake_case=64      ,						snake_case=32      ,						snake_case=5      ,						snake_case=4      ,						snake_case=37      ,						snake_case="gelu"      ,						snake_case=0.1      ,						snake_case=0.1      ,						snake_case=512      ,						snake_case=16      ,						snake_case=2      ,						snake_case=0.02      ,						snake_case=3      ,						snake_case=4      ,						snake_case=None      ,						):

						'''simple docstring'''


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




			def 				snake_case__				(							self     ):

						'''simple docstring'''


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

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

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

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

						UpperCamelCase__					=       self.get_config()

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




			def 				snake_case__				(							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=snake_case      ,						initializer_range=self.initializer_range      ,						)




			def 				snake_case__				(							self      ,						snake_case      ,						snake_case      ,						snake_case      ,						snake_case      ,						snake_case      ,						snake_case      ,						snake_case     ):

						'''simple docstring'''


						UpperCamelCase__					=       MegatronBertModel(config=snake_case     )
						model.to(snake_case     )
						model.eval()
						UpperCamelCase__					=       model(snake_case      ,						attention_mask=snake_case      ,						token_type_ids=snake_case     )
						UpperCamelCase__					=       model(snake_case      ,						token_type_ids=snake_case     )
						UpperCamelCase__					=       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 				snake_case__				(							self      ,						snake_case      ,						snake_case      ,						snake_case      ,						snake_case      ,						snake_case      ,						snake_case      ,						snake_case     ):

						'''simple docstring'''


						UpperCamelCase__					=       MegatronBertForMaskedLM(config=snake_case     )
						model.to(snake_case     )
						model.eval()
						UpperCamelCase__					=       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 				snake_case__				(							self      ,						snake_case      ,						snake_case      ,						snake_case      ,						snake_case      ,						snake_case      ,						snake_case      ,						snake_case     ):

						'''simple docstring'''


						UpperCamelCase__					=       MegatronBertForCausalLM(config=snake_case     )
						model.to(snake_case     )
						model.eval()
						UpperCamelCase__					=       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 				snake_case__				(							self      ,						snake_case      ,						snake_case      ,						snake_case      ,						snake_case      ,						snake_case      ,						snake_case      ,						snake_case     ):

						'''simple docstring'''


						UpperCamelCase__					=       MegatronBertForNextSentencePrediction(config=snake_case     )
						model.to(snake_case     )
						model.eval()
						UpperCamelCase__					=       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 				snake_case__				(							self      ,						snake_case      ,						snake_case      ,						snake_case      ,						snake_case      ,						snake_case      ,						snake_case      ,						snake_case     ):

						'''simple docstring'''


						UpperCamelCase__					=       MegatronBertForPreTraining(config=snake_case     )
						model.to(snake_case     )
						model.eval()
						UpperCamelCase__					=       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 				snake_case__				(							self      ,						snake_case      ,						snake_case      ,						snake_case      ,						snake_case      ,						snake_case      ,						snake_case      ,						snake_case     ):

						'''simple docstring'''


						UpperCamelCase__					=       MegatronBertForQuestionAnswering(config=snake_case     )
						model.to(snake_case     )
						model.eval()
						UpperCamelCase__					=       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 				snake_case__				(							self      ,						snake_case      ,						snake_case      ,						snake_case      ,						snake_case      ,						snake_case      ,						snake_case      ,						snake_case     ):

						'''simple docstring'''


						UpperCamelCase__					=       self.num_labels
						UpperCamelCase__					=       MegatronBertForSequenceClassification(snake_case     )
						model.to(snake_case     )
						model.eval()
						UpperCamelCase__					=       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 				snake_case__				(							self      ,						snake_case      ,						snake_case      ,						snake_case      ,						snake_case      ,						snake_case      ,						snake_case      ,						snake_case     ):

						'''simple docstring'''


						UpperCamelCase__					=       self.num_labels
						UpperCamelCase__					=       MegatronBertForTokenClassification(config=snake_case     )
						model.to(snake_case     )
						model.eval()
						UpperCamelCase__					=       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 				snake_case__				(							self      ,						snake_case      ,						snake_case      ,						snake_case      ,						snake_case      ,						snake_case      ,						snake_case      ,						snake_case     ):

						'''simple docstring'''


						UpperCamelCase__					=       self.num_choices
						UpperCamelCase__					=       MegatronBertForMultipleChoice(config=snake_case     )
						model.to(snake_case     )
						model.eval()
						UpperCamelCase__					=       input_ids.unsqueeze(1     ).expand(-1      ,						self.num_choices      ,						-1     ).contiguous()
						UpperCamelCase__					=       token_type_ids.unsqueeze(1     ).expand(-1      ,						self.num_choices      ,						-1     ).contiguous()
						UpperCamelCase__					=       input_mask.unsqueeze(1     ).expand(-1      ,						self.num_choices      ,						-1     ).contiguous()
						UpperCamelCase__					=       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 				snake_case__				(							self     ):

						'''simple docstring'''


						UpperCamelCase__					=       self.prepare_config_and_inputs()
						(
						    (
						    UpperCamelCase__
						),      (
						    UpperCamelCase__
						),      (
						    UpperCamelCase__
						),      (
						    UpperCamelCase__
						),      (
						    UpperCamelCase__
						),      (
						    UpperCamelCase__
						),      (
						    UpperCamelCase__
						),      
						)					=       config_and_inputs
						UpperCamelCase__					=       {"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"""






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




			def 				snake_case__				(							self      ,						snake_case      ,						snake_case      ,						snake_case=False     ):

						'''simple docstring'''


						UpperCamelCase__					=       super()._prepare_for_class(snake_case      ,						snake_case      ,						return_labels=snake_case     )

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




			def 				snake_case__				(							self     ):

						'''simple docstring'''


						UpperCamelCase__					=       MegatronBertModelTester(self     )
						UpperCamelCase__					=       ConfigTester(self      ,						config_class=snake_case      ,						hidden_size=37     )




			def 				snake_case__				(							self     ):

						'''simple docstring'''


						self.config_tester.run_common_tests()




			def 				snake_case__				(							self     ):

						'''simple docstring'''


						UpperCamelCase__					=       self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_model(*snake_case     )




			def 				snake_case__				(							self     ):

						'''simple docstring'''


						UpperCamelCase__					=       self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_for_masked_lm(*snake_case     )




			def 				snake_case__				(							self     ):

						'''simple docstring'''


						UpperCamelCase__					=       self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*snake_case     )




			def 				snake_case__				(							self     ):

						'''simple docstring'''


						UpperCamelCase__					=       self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*snake_case     )




			def 				snake_case__				(							self     ):

						'''simple docstring'''


						UpperCamelCase__					=       self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_for_pretraining(*snake_case     )




			def 				snake_case__				(							self     ):

						'''simple docstring'''


						UpperCamelCase__					=       self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_for_question_answering(*snake_case     )




			def 				snake_case__				(							self     ):

						'''simple docstring'''


						UpperCamelCase__					=       self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*snake_case     )







			def 				snake_case__				(							self     ):

						'''simple docstring'''


						UpperCamelCase__					=       self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_for_token_classification(*snake_case     )



def        UpperCamelCase_(				_A   :List[Any]  )->		Optional[Any]:
			return torch.tensor(
			    _A				,     dtype=torch.long				,     device=_A				,     )


__UpperCamelCase    		=       1E-4







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






			"""simple docstring"""






			@slow
			@unittest.skip("Model is not available."     )
			def 				snake_case__				(							self     ):

						'''simple docstring'''


						UpperCamelCase__					=       "nvidia/megatron-bert-uncased-345m"
						if "MYDIR" in os.environ:
									UpperCamelCase__					=       os.path.join(os.environ["MYDIR"]      ,						snake_case     )
						UpperCamelCase__					=       MegatronBertModel.from_pretrained(snake_case     )
						model.to(snake_case     )
						model.half()
						UpperCamelCase__					=       _long_tensor([[101, 7110, 1005, 1056, 2023, 11333, 17413, 1029, 102]]     )
						with torch.no_grad():
									UpperCamelCase__					=       model(snake_case     )[0]
						UpperCamelCase__					=       torch.Size((1, 9, 1024)     )
						self.assertEqual(output.shape      ,						snake_case     )

						UpperCamelCase__					=       [-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     ):
												UpperCamelCase__					=       output[0, ii, jj]
												UpperCamelCase__					=       expected[3 * ii + jj]
												UpperCamelCase__					=       "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     )


