






import inspect
import unittest

from transformers import MobileNetVaConfig
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

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 MobileNetVaForImageClassification, MobileNetVaForSemanticSegmentation, MobileNetVaModel
   from transformers.models.mobilenet_va.modeling_mobilenet_va import MOBILENET_V2_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
   from PIL import Image

   from transformers import MobileNetVaImageProcessor




class 							_lowerCAmelCase(			UpperCAmelCase_   ):





       """simple docstring"""







       def      _a       (   self				):
         UpperCamelCase_:   Dict        =							self.config_class(**self.inputs_dict				)
         self.parent.assertTrue(hasattr(_lowerCamelCase    ,     'tf_padding'				)				)
         self.parent.assertTrue(hasattr(_lowerCamelCase    ,     'depth_multiplier'				)				)







class 							_lowerCAmelCase:





       """simple docstring"""







       def __init__(   self    ,     _lowerCamelCase    ,     _lowerCamelCase=1_3    ,     _lowerCamelCase=3    ,     _lowerCamelCase=3_2    ,     _lowerCamelCase=0.2_5    ,     _lowerCamelCase=8    ,     _lowerCamelCase=8    ,     _lowerCamelCase=6    ,     _lowerCamelCase=3_2    ,     _lowerCamelCase=True    ,     _lowerCamelCase=True    ,     _lowerCamelCase=True    ,     _lowerCamelCase="relu6"    ,     _lowerCamelCase=1_2_8_0    ,     _lowerCamelCase=0.1    ,     _lowerCamelCase=0.0_2    ,     _lowerCamelCase=True    ,     _lowerCamelCase=True    ,     _lowerCamelCase=1_0    ,     _lowerCamelCase=None    ,     ):
         UpperCamelCase_:   List[Any]        =							parent
         UpperCamelCase_:   List[str]        =							batch_size
         UpperCamelCase_:   int        =							num_channels
         UpperCamelCase_:   Union[str, Any]        =							image_size
         UpperCamelCase_:   int        =							depth_multiplier
         UpperCamelCase_:   Optional[int]        =							depth_divisible_by
         UpperCamelCase_:   Optional[int]        =							min_depth
         UpperCamelCase_:   List[Any]        =							expand_ratio
         UpperCamelCase_:   List[Any]        =							tf_padding
         UpperCamelCase_:   str        =							output_stride
         UpperCamelCase_:   Any        =							first_layer_is_expansion
         UpperCamelCase_:   Optional[int]        =							finegrained_output
         UpperCamelCase_:   Optional[int]        =							hidden_act
         UpperCamelCase_:   Any        =							last_hidden_size if finegrained_output else int(last_hidden_size * depth_multiplier				)
         UpperCamelCase_:   Optional[int]        =							classifier_dropout_prob
         UpperCamelCase_:   Optional[Any]        =							use_labels
         UpperCamelCase_:   Optional[Any]        =							is_training
         UpperCamelCase_:   Optional[Any]        =							num_labels
         UpperCamelCase_:   List[str]        =							initializer_range
         UpperCamelCase_:   Union[str, Any]        =							scope





       def      _a       (   self				):
         UpperCamelCase_:   Any        =							floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]				)

         UpperCamelCase_:   Optional[Any]        =							None
         UpperCamelCase_:   str        =							None
         if self.use_labels:
           UpperCamelCase_:   List[str]        =							ids_tensor([self.batch_size]    ,     self.num_labels				)
           UpperCamelCase_:   Union[str, Any]        =							ids_tensor([self.batch_size, self.image_size, self.image_size]    ,     self.num_labels				)

         UpperCamelCase_:   List[Any]        =							self.get_config()

         return config, pixel_values, labels, pixel_labels





       def      _a       (   self				):
         return MobileNetVaConfig(
             num_channels=self.num_channels    ,     image_size=self.image_size    ,     depth_multiplier=self.depth_multiplier    ,     depth_divisible_by=self.depth_divisible_by    ,     min_depth=self.min_depth    ,     expand_ratio=self.expand_ratio    ,     output_stride=self.output_stride    ,     first_layer_is_expansion=self.first_layer_is_expansion    ,     finegrained_output=self.finegrained_output    ,     hidden_act=self.hidden_act    ,     tf_padding=self.tf_padding    ,     classifier_dropout_prob=self.classifier_dropout_prob    ,     initializer_range=self.initializer_range    ,     )





       def      _a       (   self    ,     _lowerCamelCase    ,     _lowerCamelCase    ,     _lowerCamelCase    ,     _lowerCamelCase				):
         UpperCamelCase_:   Any        =							MobileNetVaModel(config=_lowerCamelCase				)
         model.to(_lowerCamelCase				)
         model.eval()
         UpperCamelCase_:   List[Any]        =							model(_lowerCamelCase				)
         self.parent.assertEqual(
             result.last_hidden_state.shape    ,     (
                 self.batch_size,
                 self.last_hidden_size,
                 self.image_size // self.output_stride,
                 self.image_size // self.output_stride,
             )    ,     )
         self.parent.assertEqual(
             result.pooler_output.shape    ,     (self.batch_size, self.last_hidden_size)    ,     )





       def      _a       (   self    ,     _lowerCamelCase    ,     _lowerCamelCase    ,     _lowerCamelCase    ,     _lowerCamelCase				):
         UpperCamelCase_:   Optional[int]        =							self.num_labels
         UpperCamelCase_:   Dict        =							MobileNetVaForImageClassification(_lowerCamelCase				)
         model.to(_lowerCamelCase				)
         model.eval()
         UpperCamelCase_:   Optional[int]        =							model(_lowerCamelCase    ,     labels=_lowerCamelCase				)
         self.parent.assertEqual(result.logits.shape    ,     (self.batch_size, self.num_labels)				)





       def      _a       (   self    ,     _lowerCamelCase    ,     _lowerCamelCase    ,     _lowerCamelCase    ,     _lowerCamelCase				):
         UpperCamelCase_:   List[str]        =							self.num_labels
         UpperCamelCase_:   Union[str, Any]        =							MobileNetVaForSemanticSegmentation(_lowerCamelCase				)
         model.to(_lowerCamelCase				)
         model.eval()
         UpperCamelCase_:   Optional[Any]        =							model(_lowerCamelCase				)
         self.parent.assertEqual(
             result.logits.shape    ,     (
                 self.batch_size,
                 self.num_labels,
                 self.image_size // self.output_stride,
                 self.image_size // self.output_stride,
             )    ,     )
         UpperCamelCase_:   Dict        =							model(_lowerCamelCase    ,     labels=_lowerCamelCase				)
         self.parent.assertEqual(
             result.logits.shape    ,     (
                 self.batch_size,
                 self.num_labels,
                 self.image_size // self.output_stride,
                 self.image_size // self.output_stride,
             )    ,     )





       def      _a       (   self				):
         UpperCamelCase_:   Dict        =							self.prepare_config_and_inputs()
         UpperCamelCase_ ,UpperCamelCase_ ,UpperCamelCase_ ,UpperCamelCase_:   Any        =							config_and_inputs
         UpperCamelCase_:   Dict        =							{'pixel_values': pixel_values}
         return config, inputs_dict







@require_torch
class 							_lowerCAmelCase(			UpperCAmelCase_				,      UpperCAmelCase_				,      unittest.TestCase   ):





       """simple docstring"""







       a							:						List[str]										=(
           (MobileNetVaModel, MobileNetVaForImageClassification, MobileNetVaForSemanticSegmentation)
           if is_torch_available()
           else ()
       )
       a							:						Dict										=(
           {
               '''feature-extraction''': MobileNetVaModel,
               '''image-classification''': MobileNetVaForImageClassification,
               '''image-segmentation''': MobileNetVaForSemanticSegmentation,
           }
           if is_torch_available()
           else {}
       )

       a							:						List[str]										=False
       a							:						Tuple										=False
       a							:						List[Any]										=False
       a							:						List[str]										=False





       def      _a       (   self				):
         UpperCamelCase_:   Any        =							MobileNetVaModelTester(self				)
         UpperCamelCase_:   Tuple        =							MobileNetVaConfigTester(self    ,     config_class=_lowerCamelCase    ,     has_text_modality=_lowerCamelCase				)





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





       @unittest.skip(reason='MobileNetV2 does not use inputs_embeds'				)
       def      _a       (   self				):
         pass





       @unittest.skip(reason='MobileNetV2 does not support input and output embeddings'				)
       def      _a       (   self				):
         pass





       @unittest.skip(reason='MobileNetV2 does not output attentions'				)
       def      _a       (   self				):
         pass





       def      _a       (   self				):
         UpperCamelCase_ ,UpperCamelCase_:   List[Any]        =							self.model_tester.prepare_config_and_inputs_for_common()

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

           UpperCamelCase_:   Union[str, Any]        =							['pixel_values']
           self.assertListEqual(arg_names[:1]    ,     _lowerCamelCase				)





       def      _a       (   self				):
         UpperCamelCase_:   List[Any]        =							self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*_lowerCamelCase				)





       def      _a       (   self				):
         def check_hidden_states_output(_lowerCamelCase    ,     _lowerCamelCase    ,     _lowerCamelCase				):
           UpperCamelCase_:   Dict        =							model_class(_lowerCamelCase				)
           model.to(_lowerCamelCase				)
           model.eval()

           with torch.no_grad():
             UpperCamelCase_:   List[Any]        =							model(**self._prepare_for_class(_lowerCamelCase    ,     _lowerCamelCase				)				)

           UpperCamelCase_:   Tuple        =							outputs.hidden_states

           UpperCamelCase_:   Union[str, Any]        =							1_6
           self.assertEqual(len(_lowerCamelCase				)    ,     _lowerCamelCase				)

         UpperCamelCase_ ,UpperCamelCase_:   Union[str, Any]        =							self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
           UpperCamelCase_:   List[str]        =							True
           check_hidden_states_output(_lowerCamelCase    ,     _lowerCamelCase    ,     _lowerCamelCase				)

           # check that output_hidden_states also work using config
           del inputs_dict["output_hidden_states"]
           UpperCamelCase_:   List[str]        =							True

           check_hidden_states_output(_lowerCamelCase    ,     _lowerCamelCase    ,     _lowerCamelCase				)





       def      _a       (   self				):
         UpperCamelCase_:   List[str]        =							self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_image_classification(*_lowerCamelCase				)





       def      _a       (   self				):
         UpperCamelCase_:   List[Any]        =							self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_semantic_segmentation(*_lowerCamelCase				)





       @slow
       def      _a       (   self				):
         for model_name in MOBILENET_V2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
           UpperCamelCase_:   int        =							MobileNetVaModel.from_pretrained(_lowerCamelCase				)
           self.assertIsNotNone(_lowerCamelCase				)







def     snake_case   ()			->							List[str]:
  UpperCamelCase_:   int        =							Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'    )
  return image




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





       """simple docstring"""







       @cached_property
       def      _a       (   self				):
         return (
             MobileNetVaImageProcessor.from_pretrained('google/mobilenet_v2_1.0_224'				) if is_vision_available() else None
         )





       @slow
       def      _a       (   self				):
         UpperCamelCase_:   Optional[Any]        =							MobileNetVaForImageClassification.from_pretrained('google/mobilenet_v2_1.0_224'				).to(_lowerCamelCase				)

         UpperCamelCase_:   Any        =							self.default_image_processor
         UpperCamelCase_:   Optional[int]        =							prepare_img()
         UpperCamelCase_:   int        =							image_processor(images=_lowerCamelCase    ,     return_tensors='pt'				).to(_lowerCamelCase				)

         # forward pass
         with torch.no_grad():
           UpperCamelCase_:   Tuple        =							model(**_lowerCamelCase				)

         # verify the logits
         UpperCamelCase_:   Any        =							torch.Size((1, 1_0_0_1)				)
         self.assertEqual(outputs.logits.shape    ,     _lowerCamelCase				)

         UpperCamelCase_:   Dict        =							torch.tensor([0.2_4_4_5, -1.1_9_9_3, 0.1_9_0_5]				).to(_lowerCamelCase				)

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





       @slow
       def      _a       (   self				):
         UpperCamelCase_:   Optional[Any]        =							MobileNetVaForSemanticSegmentation.from_pretrained('google/deeplabv3_mobilenet_v2_1.0_513'				)
         UpperCamelCase_:   List[Any]        =							model.to(_lowerCamelCase				)

         UpperCamelCase_:   int        =							MobileNetVaImageProcessor.from_pretrained('google/deeplabv3_mobilenet_v2_1.0_513'				)

         UpperCamelCase_:   List[Any]        =							prepare_img()
         UpperCamelCase_:   Dict        =							image_processor(images=_lowerCamelCase    ,     return_tensors='pt'				).to(_lowerCamelCase				)

         # forward pass
         with torch.no_grad():
           UpperCamelCase_:   Dict        =							model(**_lowerCamelCase				)
         UpperCamelCase_:   Optional[Any]        =							outputs.logits

         # verify the logits
         UpperCamelCase_:   Optional[int]        =							torch.Size((1, 2_1, 6_5, 6_5)				)
         self.assertEqual(logits.shape    ,     _lowerCamelCase				)

         UpperCamelCase_:   int        =							torch.tensor(
             [
                 [[1_7.5_7_9_0, 1_7.7_5_8_1, 1_8.3_3_5_5], [1_8.3_2_5_7, 1_8.4_2_3_0, 1_8.8_9_7_3], [1_8.6_1_6_9, 1_8.8_6_5_0, 1_9.2_1_8_7]],
                 [[-2.1_5_9_5, -2.0_9_7_7, -2.3_7_4_1], [-2.4_2_2_6, -2.3_0_2_8, -2.6_8_3_5], [-2.7_8_1_9, -2.5_9_9_1, -2.7_7_0_6]],
                 [[4.2_0_5_8, 4.8_3_1_7, 4.7_6_3_8], [4.4_1_3_6, 5.0_3_6_1, 4.9_3_8_3], [4.5_0_2_8, 4.9_6_4_4, 4.8_7_3_4]],
             ]    ,     device=_lowerCamelCase    ,     )

         self.assertTrue(torch.allclose(logits[0, :3, :3, :3]    ,     _lowerCamelCase    ,     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_=1_3			,							snake_case_=7			,							snake_case_=True			,							snake_case_=True			,							snake_case_=True			,							snake_case_=True			,							snake_case_=9_9			,							snake_case_=6_4			,							snake_case_=3_2			,							snake_case_=5			,							snake_case_=4			,							snake_case_=3_7			,							snake_case_="gelu"			,							snake_case_=0.1			,							snake_case_=0.1			,							snake_case_=5_1_2			,							snake_case_=1_6			,							snake_case_=2			,							snake_case_=0.0_2			,							snake_case_=3			,							snake_case_=4			,							snake_case_=None			,							) ->      Tuple:







								'''simple docstring'''

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






					def A   (					self    ) ->      List[str]:







								'''simple docstring'''

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

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

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

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

								__lowercase			       = self.get_config()

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






					def A   (					self    ) ->      List[Any]:







								'''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 A   (					self			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_    ) ->      List[str]:







								'''simple docstring'''

								__lowercase			       = MegatronBertModel(config=snake_case_    )
								model.to(snake_case_    )
								model.eval()
								__lowercase			       = model(snake_case_			,							attention_mask=snake_case_			,							token_type_ids=snake_case_    )
								__lowercase			       = model(snake_case_			,							token_type_ids=snake_case_    )
								__lowercase			       = 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 A   (					self			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_    ) ->      List[str]:







								'''simple docstring'''

								__lowercase			       = MegatronBertForMaskedLM(config=snake_case_    )
								model.to(snake_case_    )
								model.eval()
								__lowercase			       = 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 A   (					self			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_    ) ->      int:







								'''simple docstring'''

								__lowercase			       = MegatronBertForCausalLM(config=snake_case_    )
								model.to(snake_case_    )
								model.eval()
								__lowercase			       = 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 A   (					self			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_    ) ->      Tuple:







								'''simple docstring'''

								__lowercase			       = MegatronBertForNextSentencePrediction(config=snake_case_    )
								model.to(snake_case_    )
								model.eval()
								__lowercase			       = 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 A   (					self			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_    ) ->      Optional[int]:







								'''simple docstring'''

								__lowercase			       = MegatronBertForPreTraining(config=snake_case_    )
								model.to(snake_case_    )
								model.eval()
								__lowercase			       = 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 A   (					self			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_    ) ->      Any:







								'''simple docstring'''

								__lowercase			       = MegatronBertForQuestionAnswering(config=snake_case_    )
								model.to(snake_case_    )
								model.eval()
								__lowercase			       = 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 A   (					self			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_    ) ->      Optional[int]:







								'''simple docstring'''

								__lowercase			       = self.num_labels
								__lowercase			       = MegatronBertForSequenceClassification(snake_case_    )
								model.to(snake_case_    )
								model.eval()
								__lowercase			       = 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 A   (					self			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_    ) ->      Dict:







								'''simple docstring'''

								__lowercase			       = self.num_labels
								__lowercase			       = MegatronBertForTokenClassification(config=snake_case_    )
								model.to(snake_case_    )
								model.eval()
								__lowercase			       = 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 A   (					self			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_    ) ->      Dict:







								'''simple docstring'''

								__lowercase			       = self.num_choices
								__lowercase			       = MegatronBertForMultipleChoice(config=snake_case_    )
								model.to(snake_case_    )
								model.eval()
								__lowercase			       = input_ids.unsqueeze(1    ).expand(-1			,							self.num_choices			,							-1    ).contiguous()
								__lowercase			       = token_type_ids.unsqueeze(1    ).expand(-1			,							self.num_choices			,							-1    ).contiguous()
								__lowercase			       = input_mask.unsqueeze(1    ).expand(-1			,							self.num_choices			,							-1    ).contiguous()
								__lowercase			       = 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 A   (					self    ) ->      Union[str, Any]:







								'''simple docstring'''

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




@require_torch
class 						lowerCamelCase_							(					lowerCAmelCase__  ,  lowerCAmelCase__  ,  unittest.TestCase					):



					'''simple docstring'''







					__UpperCAmelCase							  =					(
					    (
					        MegatronBertModel,
					        MegatronBertForMaskedLM,
					        MegatronBertForCausalLM,
					        MegatronBertForMultipleChoice,
					        MegatronBertForNextSentencePrediction,
					        MegatronBertForPreTraining,
					        MegatronBertForQuestionAnswering,
					        MegatronBertForSequenceClassification,
					        MegatronBertForTokenClassification,
					    )
					    if is_torch_available()
					    else ()
					)
					__UpperCAmelCase							  =					(
					    {
					        "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							  =					True
					# test_resize_embeddings = False
					__UpperCAmelCase							  =					False






					def A   (					self			,							snake_case_			,							snake_case_			,							snake_case_=False    ) ->      Tuple:







								'''simple docstring'''

								__lowercase			       = super()._prepare_for_class(snake_case_			,							snake_case_			,							return_labels=snake_case_    )

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






					def A   (					self    ) ->      List[Any]:







								'''simple docstring'''

								__lowercase			       = MegatronBertModelTester(self    )
								__lowercase			       = ConfigTester(self			,							config_class=snake_case_			,							hidden_size=3_7    )






					def A   (					self    ) ->      int:







								'''simple docstring'''

								self.config_tester.run_common_tests()






					def A   (					self    ) ->      List[str]:







								'''simple docstring'''

								__lowercase			       = self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_megatron_bert_model(*snake_case_    )






					def A   (					self    ) ->      Optional[int]:







								'''simple docstring'''

								__lowercase			       = self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_megatron_bert_for_masked_lm(*snake_case_    )






					def A   (					self    ) ->      str:







								'''simple docstring'''

								__lowercase			       = self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*snake_case_    )






					def A   (					self    ) ->      List[Any]:







								'''simple docstring'''

								__lowercase			       = self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*snake_case_    )






					def A   (					self    ) ->      Dict:







								'''simple docstring'''

								__lowercase			       = self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_megatron_bert_for_pretraining(*snake_case_    )






					def A   (					self    ) ->      Optional[Any]:







								'''simple docstring'''

								__lowercase			       = self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_megatron_bert_for_question_answering(*snake_case_    )






					def A   (					self    ) ->      List[Any]:







								'''simple docstring'''

								__lowercase			       = self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*snake_case_    )






					def A   (					self    ) ->      Dict:







								'''simple docstring'''

								__lowercase			       = self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_megatron_bert_for_token_classification(*snake_case_    )




def 		lowercase_	(  _UpperCamelCase	):


			'''simple docstring'''







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


a			:			Optional[Any]												=					1E-4





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



					'''simple docstring'''







					@slow
					@unittest.skip('''Model is not available.'''    )
					def A   (					self    ) ->      Tuple:







								'''simple docstring'''

								__lowercase			       = '''nvidia/megatron-bert-uncased-345m'''
								if "MYDIR" in os.environ:
											__lowercase			       = os.path.join(os.environ['''MYDIR''']			,							snake_case_    )
								__lowercase			       = MegatronBertModel.from_pretrained(snake_case_    )
								model.to(snake_case_    )
								model.half()
								__lowercase			       = _long_tensor([[1_0_1, 7_1_1_0, 1_0_0_5, 1_0_5_6, 2_0_2_3, 1_1_3_3_3, 1_7_4_1_3, 1_0_2_9, 1_0_2]]    )
								with torch.no_grad():
											__lowercase			       = model(snake_case_    )[0]
								__lowercase			       = torch.Size((1, 9, 1_0_2_4)    )
								self.assertEqual(output.shape			,							snake_case_    )

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


