


"""simple docstring"""






from typing import Dict, List, Optional, Tuple, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import (
    center_crop,
    get_resize_output_image_size,
    normalize,
    rescale,
    resize,
    to_channel_dimension_format,
)
from ...image_utils import (
    IMAGENET_STANDARD_MEAN,
    IMAGENET_STANDARD_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_torch_available, is_torch_tensor, logging


if is_torch_available():
					import torch


UpperCamelCase__				:Any			   =	logging.get_logger(__name__)






class      A(  lowerCAmelCase__		):

		"""simple docstring"""





		A    =							["pixel_values"]
		def __init__( self					,				SCREAMING_SNAKE_CASE__ = True					,				SCREAMING_SNAKE_CASE__ = None					,				SCREAMING_SNAKE_CASE__ = PILImageResampling.BILINEAR					,				SCREAMING_SNAKE_CASE__ = True					,				SCREAMING_SNAKE_CASE__ = None					,				SCREAMING_SNAKE_CASE__ = True					,				SCREAMING_SNAKE_CASE__ = 1 / 2_55					,				SCREAMING_SNAKE_CASE__ = True					,				SCREAMING_SNAKE_CASE__ = None					,				SCREAMING_SNAKE_CASE__ = None					,				**SCREAMING_SNAKE_CASE__					,				)      ->	None:

					"""simple docstring"""



					super().__init__(**snake_case_   )
					_UpperCamelCase					:Union[str, Any]					     =	size if size is not None else {'''shortest_edge''': 2_56}
					_UpperCamelCase					:Union[str, Any]					     =	get_size_dict(snake_case_					,				default_to_square=snake_case_   )
					_UpperCamelCase					:Optional[Any]					     =	crop_size if crop_size is not None else {'''height''': 2_24, '''width''': 2_24}
					_UpperCamelCase					:Any					     =	get_size_dict(snake_case_					,				param_name='''crop_size'''   )
					_UpperCamelCase					:Optional[Any]					     =	do_resize
					_UpperCamelCase					:List[Any]					     =	size
					_UpperCamelCase					:str					     =	resample
					_UpperCamelCase					:str					     =	do_center_crop
					_UpperCamelCase					:Union[str, Any]					     =	crop_size
					_UpperCamelCase					:Any					     =	do_rescale
					_UpperCamelCase					:List[Any]					     =	rescale_factor
					_UpperCamelCase					:str					     =	do_normalize
					_UpperCamelCase					:Optional[Any]					     =	image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
					_UpperCamelCase					:Tuple					     =	image_std if image_std is not None else IMAGENET_STANDARD_STD
		def   _UpperCamelCase( self					,				SCREAMING_SNAKE_CASE__					,				SCREAMING_SNAKE_CASE__					,				SCREAMING_SNAKE_CASE__ = PILImageResampling.BICUBIC					,				SCREAMING_SNAKE_CASE__ = None					,				**SCREAMING_SNAKE_CASE__					,				)      ->	np.ndarray:

					"""simple docstring"""



					_UpperCamelCase					:Union[str, Any]					     =	get_size_dict(snake_case_					,				default_to_square=snake_case_   )
					if "shortest_edge" not in size:
								raise ValueError(f"The `size` parameter must contain the key `shortest_edge`. Got {size.keys()}"   )
					_UpperCamelCase					:List[str]					     =	get_resize_output_image_size(snake_case_					,				size=size['''shortest_edge''']					,				default_to_square=snake_case_   )
					return resize(snake_case_					,				size=snake_case_					,				resample=snake_case_					,				data_format=snake_case_					,				**snake_case_   )
		def   _UpperCamelCase( self					,				SCREAMING_SNAKE_CASE__					,				SCREAMING_SNAKE_CASE__					,				SCREAMING_SNAKE_CASE__ = None					,				**SCREAMING_SNAKE_CASE__					,				)      ->	np.ndarray:

					"""simple docstring"""



					_UpperCamelCase					:int					     =	get_size_dict(snake_case_   )
					if "height" not in size or "width" not in size:
								raise ValueError(f"The `size` parameter must contain the keys `height` and `width`. Got {size.keys()}"   )
					return center_crop(snake_case_					,				size=(size['''height'''], size['''width'''])					,				data_format=snake_case_					,				**snake_case_   )
		def   _UpperCamelCase( self					,				SCREAMING_SNAKE_CASE__					,				SCREAMING_SNAKE_CASE__					,				SCREAMING_SNAKE_CASE__ = None					,				**SCREAMING_SNAKE_CASE__   )      ->	np.ndarray:

					"""simple docstring"""



					return rescale(snake_case_					,				scale=snake_case_					,				data_format=snake_case_					,				**snake_case_   )
		def   _UpperCamelCase( self					,				SCREAMING_SNAKE_CASE__					,				SCREAMING_SNAKE_CASE__					,				SCREAMING_SNAKE_CASE__					,				SCREAMING_SNAKE_CASE__ = None					,				**SCREAMING_SNAKE_CASE__					,				)      ->	np.ndarray:

					"""simple docstring"""



					return normalize(snake_case_					,				mean=snake_case_					,				std=snake_case_					,				data_format=snake_case_					,				**snake_case_   )
		def   _UpperCamelCase( self					,				SCREAMING_SNAKE_CASE__					,				SCREAMING_SNAKE_CASE__ = None					,				SCREAMING_SNAKE_CASE__ = None					,				SCREAMING_SNAKE_CASE__ = None					,				SCREAMING_SNAKE_CASE__ = None					,				SCREAMING_SNAKE_CASE__ = None					,				SCREAMING_SNAKE_CASE__ = None					,				SCREAMING_SNAKE_CASE__ = None					,				SCREAMING_SNAKE_CASE__ = None					,				SCREAMING_SNAKE_CASE__ = None					,				SCREAMING_SNAKE_CASE__ = None					,				SCREAMING_SNAKE_CASE__ = None					,				SCREAMING_SNAKE_CASE__ = ChannelDimension.FIRST					,				**SCREAMING_SNAKE_CASE__					,				)      ->	int:

					"""simple docstring"""



					_UpperCamelCase					:List[str]					     =	do_resize if do_resize is not None else self.do_resize
					_UpperCamelCase					:int					     =	size if size is not None else self.size
					_UpperCamelCase					:str					     =	get_size_dict(snake_case_					,				default_to_square=snake_case_   )
					_UpperCamelCase					:Tuple					     =	resample if resample is not None else self.resample
					_UpperCamelCase					:List[Any]					     =	do_center_crop if do_center_crop is not None else self.do_center_crop
					_UpperCamelCase					:Optional[int]					     =	crop_size if crop_size is not None else self.crop_size
					_UpperCamelCase					:Any					     =	get_size_dict(snake_case_					,				param_name='''crop_size'''   )
					_UpperCamelCase					:int					     =	do_rescale if do_rescale is not None else self.do_rescale
					_UpperCamelCase					:str					     =	rescale_factor if rescale_factor is not None else self.rescale_factor
					_UpperCamelCase					:int					     =	do_normalize if do_normalize is not None else self.do_normalize
					_UpperCamelCase					:Any					     =	image_mean if image_mean is not None else self.image_mean
					_UpperCamelCase					:Tuple					     =	image_std if image_std is not None else self.image_std

					_UpperCamelCase					:Optional[int]					     =	make_list_of_images(snake_case_   )

					if not valid_images(snake_case_   ):
								raise ValueError(
								    '''Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, '''
								    '''torch.Tensor, tf.Tensor or jax.ndarray.'''   )

					if do_resize and size is None:
								raise ValueError('''Size must be specified if do_resize is True.'''   )

					if do_center_crop and crop_size is None:
								raise ValueError('''Crop size must be specified if do_center_crop is True.'''   )

					if do_rescale and rescale_factor is None:
								raise ValueError('''Rescale factor must be specified if do_rescale is True.'''   )

					if do_normalize and (image_mean is None or image_std is None):
								raise ValueError('''Image mean and std must be specified if do_normalize is True.'''   )

					# All transformations expect numpy arrays.
					_UpperCamelCase					:int					     =	[to_numpy_array(snake_case_   ) for image in images]

					if do_resize:
								_UpperCamelCase					:Optional[Any]					     =	[self.resize(image=snake_case_					,				size=snake_case_					,				resample=snake_case_   ) for image in images]

					if do_center_crop:
								_UpperCamelCase					:List[Any]					     =	[self.center_crop(image=snake_case_					,				size=snake_case_   ) for image in images]

					if do_rescale:
								_UpperCamelCase					:Any					     =	[self.rescale(image=snake_case_					,				scale=snake_case_   ) for image in images]

					if do_normalize:
								_UpperCamelCase					:List[Any]					     =	[self.normalize(image=snake_case_					,				mean=snake_case_					,				std=snake_case_   ) for image in images]

					_UpperCamelCase					:Optional[Any]					     =	[to_channel_dimension_format(snake_case_					,				snake_case_   ) for image in images]

					_UpperCamelCase					:int					     =	{'''pixel_values''': images}
					return BatchFeature(data=snake_case_					,				tensor_type=snake_case_   )





		def   _UpperCamelCase( self					,				SCREAMING_SNAKE_CASE__					,				SCREAMING_SNAKE_CASE__ = None   )      ->	List[Any]:

					"""simple docstring"""



					_UpperCamelCase					:Dict					     =	outputs.logits

					# Resize logits and compute semantic segmentation maps
					if target_sizes is not None:
								if len(snake_case_   ) != len(snake_case_   ):
											raise ValueError(
											    '''Make sure that you pass in as many target sizes as the batch dimension of the logits'''   )

								if is_torch_tensor(snake_case_   ):
											_UpperCamelCase					:Tuple					     =	target_sizes.numpy()

								_UpperCamelCase					:int					     =	[]

								for idx in range(len(snake_case_   )   ):
											_UpperCamelCase					:str					     =	torch.nn.functional.interpolate(
											    logits[idx].unsqueeze(dim=0   )					,				size=target_sizes[idx]					,				mode='''bilinear'''					,				align_corners=snake_case_   )
											_UpperCamelCase					:Optional[int]					     =	resized_logits[0].argmax(dim=0   )
											semantic_segmentation.append(snake_case_   )
					else:
								_UpperCamelCase					:Union[str, Any]					     =	logits.argmax(dim=1   )
								_UpperCamelCase					:Union[str, Any]					     =	[semantic_segmentation[i] for i in range(semantic_segmentation.shape[0]   )]

					return semantic_segmentation





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_    )


