'''simple docstring'''







import warnings
from typing import List, Optional, Union

from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding, PaddingStrategy, PreTokenizedInput, TextInput, TruncationStrategy
from ...utils import TensorType




class _SCREAMING_SNAKE_CASE(							_SCREAMING_SNAKE_CASE					):
						A_			:							Tuple     		= ['image_processor', 'tokenizer']
						A_			:							Tuple     		= 'LayoutLMv2ImageProcessor'
						A_			:							Optional[int]     		= ('LayoutXLMTokenizer', 'LayoutXLMTokenizerFast')




						def __init__(    self :       Optional[int]							,							UpperCamelCase_ :       Any=None							,							UpperCamelCase_ :       List[str]=None							,							**UpperCamelCase_ :       int      )     ->				int:
													if "feature_extractor" in kwargs:
																				warnings.warn(
																				    'The `feature_extractor` argument is deprecated and will be removed in v5, use `image_processor`'
																				    ' instead.'							,							UpperCamelCase_							,							)
																				SCREAMING_SNAKE_CASE__     :Tuple						    = kwargs.pop('feature_extractor'      )

													SCREAMING_SNAKE_CASE__     :Optional[Any]						    = image_processor if image_processor is not None else feature_extractor
													if image_processor is None:
																				raise ValueError('You need to specify an `image_processor`.'      )
													if tokenizer is None:
																				raise ValueError('You need to specify a `tokenizer`.'      )

													super().__init__(UpperCamelCase_							,							UpperCamelCase_      )




						def __call__(    self :       Optional[int]							,							UpperCamelCase_ :       Any							,							UpperCamelCase_ :       Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]] = None							,							UpperCamelCase_ :       Optional[Union[PreTokenizedInput, List[PreTokenizedInput]]] = None							,							UpperCamelCase_ :       Union[List[List[int]], List[List[List[int]]]] = None							,							UpperCamelCase_ :       Optional[Union[List[int], List[List[int]]]] = None							,							UpperCamelCase_ :       bool = True							,							UpperCamelCase_ :       Union[bool, str, PaddingStrategy] = False							,							UpperCamelCase_ :       Union[bool, str, TruncationStrategy] = None							,							UpperCamelCase_ :       Optional[int] = None							,							UpperCamelCase_ :       int = 0							,							UpperCamelCase_ :       Optional[int] = None							,							UpperCamelCase_ :       Optional[bool] = None							,							UpperCamelCase_ :       Optional[bool] = None							,							UpperCamelCase_ :       bool = False							,							UpperCamelCase_ :       bool = False							,							UpperCamelCase_ :       bool = False							,							UpperCamelCase_ :       bool = False							,							UpperCamelCase_ :       bool = True							,							UpperCamelCase_ :       Optional[Union[str, TensorType]] = None							,							**UpperCamelCase_ :       List[str]							,							)     ->				BatchEncoding:
													# verify input
													if self.image_processor.apply_ocr and (boxes is not None):
																				raise ValueError(
																				    'You cannot provide bounding boxes '
																				    'if you initialized the image processor with apply_ocr set to True.'      )

													if self.image_processor.apply_ocr and (word_labels is not None):
																				raise ValueError(
																				    'You cannot provide word labels if you initialized the image processor with apply_ocr set to True.'      )

													if return_overflowing_tokens is True and return_offsets_mapping is False:
																				raise ValueError('You cannot return overflowing tokens without returning the offsets mapping.'      )

													# first, apply the image processor
													SCREAMING_SNAKE_CASE__     :Union[str, Any]						    = self.image_processor(images=UpperCamelCase_							,							return_tensors=UpperCamelCase_      )

													# second, apply the tokenizer
													if text is not None and self.image_processor.apply_ocr and text_pair is None:
																				if isinstance(UpperCamelCase_							,							UpperCamelCase_      ):
																											SCREAMING_SNAKE_CASE__     :Optional[Any]						    = [text]  # add batch dimension (as the image processor always adds a batch dimension)
																				SCREAMING_SNAKE_CASE__     :Union[str, Any]						    = features['words']

													SCREAMING_SNAKE_CASE__     :Optional[int]						    = self.tokenizer(
													    text=text if text is not None else features['words']							,							text_pair=text_pair if text_pair is not None else None							,							boxes=boxes if boxes is not None else features['boxes']							,							word_labels=UpperCamelCase_							,							add_special_tokens=UpperCamelCase_							,							padding=UpperCamelCase_							,							truncation=UpperCamelCase_							,							max_length=UpperCamelCase_							,							stride=UpperCamelCase_							,							pad_to_multiple_of=UpperCamelCase_							,							return_token_type_ids=UpperCamelCase_							,							return_attention_mask=UpperCamelCase_							,							return_overflowing_tokens=UpperCamelCase_							,							return_special_tokens_mask=UpperCamelCase_							,							return_offsets_mapping=UpperCamelCase_							,							return_length=UpperCamelCase_							,							verbose=UpperCamelCase_							,							return_tensors=UpperCamelCase_							,							**UpperCamelCase_							,							)

													# add pixel values
													SCREAMING_SNAKE_CASE__     :Optional[int]						    = features.pop('pixel_values'      )
													if return_overflowing_tokens is True:
																				SCREAMING_SNAKE_CASE__     :Optional[int]						    = self.get_overflowing_images(UpperCamelCase_							,							encoded_inputs['overflow_to_sample_mapping']      )
													SCREAMING_SNAKE_CASE__     :Dict						    = images

													return encoded_inputs




						def      __lowerCamelCase       (    self :       Any							,							UpperCamelCase_ :       str							,							UpperCamelCase_ :       List[Any]      )     ->				Dict:
													# in case there's an overflow, ensure each `input_ids` sample is mapped to its corresponding image
													SCREAMING_SNAKE_CASE__     :str						    = []
													for sample_idx in overflow_to_sample_mapping:
																				images_with_overflow.append(images[sample_idx]      )

													if len(UpperCamelCase_      ) != len(UpperCamelCase_      ):
																				raise ValueError(
																				    'Expected length of images to be the same as the length of `overflow_to_sample_mapping`, but got'
																				    f''' {len(UpperCamelCase_      )} and {len(UpperCamelCase_      )}'''      )

													return images_with_overflow




						def      __lowerCamelCase       (    self :       Optional[Any]							,							*UpperCamelCase_ :       int							,							**UpperCamelCase_ :       Dict      )     ->				int:
													return self.tokenizer.batch_decode(*UpperCamelCase_							,							**UpperCamelCase_      )




						def      __lowerCamelCase       (    self :       Any							,							*UpperCamelCase_ :       List[str]							,							**UpperCamelCase_ :       Dict      )     ->				Optional[Any]:
													return self.tokenizer.decode(*UpperCamelCase_							,							**UpperCamelCase_      )




						@property
						def      __lowerCamelCase       (    self :       Dict      )     ->				Tuple:
													return ["input_ids", "bbox", "attention_mask", "image"]




						@property
						def      __lowerCamelCase       (    self :       List[str]      )     ->				Any:
													warnings.warn(
													    '`feature_extractor_class` is deprecated and will be removed in v5. Use `image_processor_class` instead.'							,							UpperCamelCase_							,							)
													return self.image_processor_class




						@property
						def      __lowerCamelCase       (    self :       Any      )     ->				List[Any]:
													warnings.warn(
													    '`feature_extractor` is deprecated and will be removed in v5. Use `image_processor` instead.'							,							UpperCamelCase_							,							)
													return self.image_processor







'''simple docstring'''







import math
import os
import unittest

from transformers import MegatronBertConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device

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


if is_torch_available():
		import torch

		from transformers import (
		    MODEL_FOR_PRETRAINING_MAPPING,
		    MegatronBertForCausalLM,
		    MegatronBertForMaskedLM,
		    MegatronBertForMultipleChoice,
		    MegatronBertForNextSentencePrediction,
		    MegatronBertForPreTraining,
		    MegatronBertForQuestionAnswering,
		    MegatronBertForSequenceClassification,
		    MegatronBertForTokenClassification,
		    MegatronBertModel,
		)




class _SCREAMING_SNAKE_CASE:




						def __init__(    self :       Optional[int]							,							UpperCamelCase_ :       Any							,							UpperCamelCase_ :       Tuple=13							,							UpperCamelCase_ :       Union[str, Any]=7							,							UpperCamelCase_ :       Union[str, Any]=True							,							UpperCamelCase_ :       Union[str, Any]=True							,							UpperCamelCase_ :       Any=True							,							UpperCamelCase_ :       int=True							,							UpperCamelCase_ :       List[Any]=99							,							UpperCamelCase_ :       List[str]=64							,							UpperCamelCase_ :       List[str]=32							,							UpperCamelCase_ :       Union[str, Any]=5							,							UpperCamelCase_ :       Dict=4							,							UpperCamelCase_ :       List[str]=37							,							UpperCamelCase_ :       str="gelu"							,							UpperCamelCase_ :       List[Any]=0.1							,							UpperCamelCase_ :       Union[str, Any]=0.1							,							UpperCamelCase_ :       Optional[Any]=5_12							,							UpperCamelCase_ :       Union[str, Any]=16							,							UpperCamelCase_ :       Union[str, Any]=2							,							UpperCamelCase_ :       Optional[int]=0.02							,							UpperCamelCase_ :       Optional[Any]=3							,							UpperCamelCase_ :       Any=4							,							UpperCamelCase_ :       int=None							,							)     ->				Optional[int]:
													SCREAMING_SNAKE_CASE__     :Optional[int]						    = parent
													SCREAMING_SNAKE_CASE__     :Any						    = batch_size
													SCREAMING_SNAKE_CASE__     :Tuple						    = seq_length
													SCREAMING_SNAKE_CASE__     :Any						    = is_training
													SCREAMING_SNAKE_CASE__     :int						    = use_input_mask
													SCREAMING_SNAKE_CASE__     :Dict						    = use_token_type_ids
													SCREAMING_SNAKE_CASE__     :str						    = use_labels
													SCREAMING_SNAKE_CASE__     :str						    = vocab_size
													SCREAMING_SNAKE_CASE__     :Optional[Any]						    = hidden_size
													SCREAMING_SNAKE_CASE__     :List[Any]						    = embedding_size
													SCREAMING_SNAKE_CASE__     :int						    = num_hidden_layers
													SCREAMING_SNAKE_CASE__     :int						    = num_attention_heads
													SCREAMING_SNAKE_CASE__     :Union[str, Any]						    = intermediate_size
													SCREAMING_SNAKE_CASE__     :str						    = hidden_act
													SCREAMING_SNAKE_CASE__     :Any						    = hidden_dropout_prob
													SCREAMING_SNAKE_CASE__     :Any						    = attention_probs_dropout_prob
													SCREAMING_SNAKE_CASE__     :Optional[int]						    = max_position_embeddings
													SCREAMING_SNAKE_CASE__     :Optional[int]						    = type_vocab_size
													SCREAMING_SNAKE_CASE__     :Tuple						    = type_sequence_label_size
													SCREAMING_SNAKE_CASE__     :Optional[Any]						    = initializer_range
													SCREAMING_SNAKE_CASE__     :int						    = num_labels
													SCREAMING_SNAKE_CASE__     :Optional[Any]						    = num_choices
													SCREAMING_SNAKE_CASE__     :List[Any]						    = scope




						def      __lowerCamelCase       (    self :       Any      )     ->				Optional[Any]:
													SCREAMING_SNAKE_CASE__     :Optional[int]						    = ids_tensor([self.batch_size, self.seq_length]							,							self.vocab_size      )

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

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

													SCREAMING_SNAKE_CASE__     :List[str]						    = None
													SCREAMING_SNAKE_CASE__     :int						    = None
													SCREAMING_SNAKE_CASE__     :Tuple						    = None
													if self.use_labels:
																				SCREAMING_SNAKE_CASE__     :Any						    = ids_tensor([self.batch_size]							,							self.type_sequence_label_size      )
																				SCREAMING_SNAKE_CASE__     :Optional[Any]						    = ids_tensor([self.batch_size, self.seq_length]							,							self.num_labels      )
																				SCREAMING_SNAKE_CASE__     :Any						    = ids_tensor([self.batch_size]							,							self.num_choices      )

													SCREAMING_SNAKE_CASE__     :str						    = self.get_config()

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




						def      __lowerCamelCase       (    self :       List[str]      )     ->				Any:
													return MegatronBertConfig(
													    vocab_size=self.vocab_size							,							hidden_size=self.hidden_size							,							num_hidden_layers=self.num_hidden_layers							,							num_attention_heads=self.num_attention_heads							,							intermediate_size=self.intermediate_size							,							embedding_size=self.embedding_size							,							hidden_act=self.hidden_act							,							hidden_dropout_prob=self.hidden_dropout_prob							,							attention_probs_dropout_prob=self.attention_probs_dropout_prob							,							max_position_embeddings=self.max_position_embeddings							,							type_vocab_size=self.type_vocab_size							,							is_decoder=UpperCamelCase_							,							initializer_range=self.initializer_range							,							)




						def      __lowerCamelCase       (    self :       Optional[int]							,							UpperCamelCase_ :       List[str]							,							UpperCamelCase_ :       List[str]							,							UpperCamelCase_ :       str							,							UpperCamelCase_ :       Optional[int]							,							UpperCamelCase_ :       List[str]							,							UpperCamelCase_ :       Any							,							UpperCamelCase_ :       int      )     ->				int:
													SCREAMING_SNAKE_CASE__     :Optional[int]						    = MegatronBertModel(config=UpperCamelCase_      )
													model.to(UpperCamelCase_      )
													model.eval()
													SCREAMING_SNAKE_CASE__     :int						    = model(UpperCamelCase_							,							attention_mask=UpperCamelCase_							,							token_type_ids=UpperCamelCase_      )
													SCREAMING_SNAKE_CASE__     :str						    = model(UpperCamelCase_							,							token_type_ids=UpperCamelCase_      )
													SCREAMING_SNAKE_CASE__     :List[str]						    = model(UpperCamelCase_      )

													self.parent.assertEqual(result.last_hidden_state.shape							,							(self.batch_size, self.seq_length, self.hidden_size)      )
													self.parent.assertEqual(result.pooler_output.shape							,							(self.batch_size, self.hidden_size)      )




						def      __lowerCamelCase       (    self :       str							,							UpperCamelCase_ :       Dict							,							UpperCamelCase_ :       Optional[int]							,							UpperCamelCase_ :       Optional[Any]							,							UpperCamelCase_ :       str							,							UpperCamelCase_ :       List[str]							,							UpperCamelCase_ :       int							,							UpperCamelCase_ :       Union[str, Any]      )     ->				List[str]:
													SCREAMING_SNAKE_CASE__     :str						    = MegatronBertForMaskedLM(config=UpperCamelCase_      )
													model.to(UpperCamelCase_      )
													model.eval()
													SCREAMING_SNAKE_CASE__     :Optional[int]						    = model(UpperCamelCase_							,							attention_mask=UpperCamelCase_							,							token_type_ids=UpperCamelCase_							,							labels=UpperCamelCase_      )
													self.parent.assertEqual(result.logits.shape							,							(self.batch_size, self.seq_length, self.vocab_size)      )




						def      __lowerCamelCase       (    self :       Any							,							UpperCamelCase_ :       List[Any]							,							UpperCamelCase_ :       Optional[int]							,							UpperCamelCase_ :       Any							,							UpperCamelCase_ :       Optional[int]							,							UpperCamelCase_ :       int							,							UpperCamelCase_ :       int							,							UpperCamelCase_ :       int      )     ->				Dict:
													SCREAMING_SNAKE_CASE__     :Optional[int]						    = MegatronBertForCausalLM(config=UpperCamelCase_      )
													model.to(UpperCamelCase_      )
													model.eval()
													SCREAMING_SNAKE_CASE__     :List[str]						    = model(UpperCamelCase_							,							attention_mask=UpperCamelCase_							,							token_type_ids=UpperCamelCase_							,							labels=UpperCamelCase_      )
													self.parent.assertEqual(result.logits.shape							,							(self.batch_size, self.seq_length, self.vocab_size)      )




						def      __lowerCamelCase       (    self :       List[str]							,							UpperCamelCase_ :       Dict							,							UpperCamelCase_ :       Optional[int]							,							UpperCamelCase_ :       Dict							,							UpperCamelCase_ :       Optional[Any]							,							UpperCamelCase_ :       List[str]							,							UpperCamelCase_ :       Union[str, Any]							,							UpperCamelCase_ :       Optional[int]      )     ->				List[Any]:
													SCREAMING_SNAKE_CASE__     :Optional[int]						    = MegatronBertForNextSentencePrediction(config=UpperCamelCase_      )
													model.to(UpperCamelCase_      )
													model.eval()
													SCREAMING_SNAKE_CASE__     :List[str]						    = model(
													    UpperCamelCase_							,							attention_mask=UpperCamelCase_							,							token_type_ids=UpperCamelCase_							,							labels=UpperCamelCase_							,							)
													self.parent.assertEqual(result.logits.shape							,							(self.batch_size, 2)      )




						def      __lowerCamelCase       (    self :       Dict							,							UpperCamelCase_ :       List[Any]							,							UpperCamelCase_ :       Dict							,							UpperCamelCase_ :       str							,							UpperCamelCase_ :       List[str]							,							UpperCamelCase_ :       Tuple							,							UpperCamelCase_ :       int							,							UpperCamelCase_ :       Optional[int]      )     ->				Any:
													SCREAMING_SNAKE_CASE__     :Tuple						    = MegatronBertForPreTraining(config=UpperCamelCase_      )
													model.to(UpperCamelCase_      )
													model.eval()
													SCREAMING_SNAKE_CASE__     :Union[str, Any]						    = model(
													    UpperCamelCase_							,							attention_mask=UpperCamelCase_							,							token_type_ids=UpperCamelCase_							,							labels=UpperCamelCase_							,							next_sentence_label=UpperCamelCase_							,							)
													self.parent.assertEqual(result.prediction_logits.shape							,							(self.batch_size, self.seq_length, self.vocab_size)      )
													self.parent.assertEqual(result.seq_relationship_logits.shape							,							(self.batch_size, 2)      )




						def      __lowerCamelCase       (    self :       Union[str, Any]							,							UpperCamelCase_ :       int							,							UpperCamelCase_ :       Any							,							UpperCamelCase_ :       Dict							,							UpperCamelCase_ :       Dict							,							UpperCamelCase_ :       Tuple							,							UpperCamelCase_ :       List[Any]							,							UpperCamelCase_ :       List[Any]      )     ->				Tuple:
													SCREAMING_SNAKE_CASE__     :int						    = MegatronBertForQuestionAnswering(config=UpperCamelCase_      )
													model.to(UpperCamelCase_      )
													model.eval()
													SCREAMING_SNAKE_CASE__     :Tuple						    = model(
													    UpperCamelCase_							,							attention_mask=UpperCamelCase_							,							token_type_ids=UpperCamelCase_							,							start_positions=UpperCamelCase_							,							end_positions=UpperCamelCase_							,							)
													self.parent.assertEqual(result.start_logits.shape							,							(self.batch_size, self.seq_length)      )
													self.parent.assertEqual(result.end_logits.shape							,							(self.batch_size, self.seq_length)      )




						def      __lowerCamelCase       (    self :       List[Any]							,							UpperCamelCase_ :       Tuple							,							UpperCamelCase_ :       int							,							UpperCamelCase_ :       List[str]							,							UpperCamelCase_ :       str							,							UpperCamelCase_ :       Any							,							UpperCamelCase_ :       Tuple							,							UpperCamelCase_ :       Any      )     ->				Dict:
													SCREAMING_SNAKE_CASE__     :List[Any]						    = self.num_labels
													SCREAMING_SNAKE_CASE__     :Optional[int]						    = MegatronBertForSequenceClassification(UpperCamelCase_      )
													model.to(UpperCamelCase_      )
													model.eval()
													SCREAMING_SNAKE_CASE__     :int						    = model(UpperCamelCase_							,							attention_mask=UpperCamelCase_							,							token_type_ids=UpperCamelCase_							,							labels=UpperCamelCase_      )
													self.parent.assertEqual(result.logits.shape							,							(self.batch_size, self.num_labels)      )




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




						def      __lowerCamelCase       (    self :       str							,							UpperCamelCase_ :       Any							,							UpperCamelCase_ :       List[Any]							,							UpperCamelCase_ :       Union[str, Any]							,							UpperCamelCase_ :       int							,							UpperCamelCase_ :       str							,							UpperCamelCase_ :       Any							,							UpperCamelCase_ :       List[Any]      )     ->				Dict:
													SCREAMING_SNAKE_CASE__     :Optional[Any]						    = self.num_choices
													SCREAMING_SNAKE_CASE__     :Tuple						    = MegatronBertForMultipleChoice(config=UpperCamelCase_      )
													model.to(UpperCamelCase_      )
													model.eval()
													SCREAMING_SNAKE_CASE__     :Any						    = input_ids.unsqueeze(1      ).expand(-1							,							self.num_choices							,							-1      ).contiguous()
													SCREAMING_SNAKE_CASE__     :List[Any]						    = token_type_ids.unsqueeze(1      ).expand(-1							,							self.num_choices							,							-1      ).contiguous()
													SCREAMING_SNAKE_CASE__     :Dict						    = input_mask.unsqueeze(1      ).expand(-1							,							self.num_choices							,							-1      ).contiguous()
													SCREAMING_SNAKE_CASE__     :Union[str, Any]						    = model(
													    UpperCamelCase_							,							attention_mask=UpperCamelCase_							,							token_type_ids=UpperCamelCase_							,							labels=UpperCamelCase_							,							)
													self.parent.assertEqual(result.logits.shape							,							(self.batch_size, self.num_choices)      )




						def      __lowerCamelCase       (    self :       Any      )     ->				int:
													SCREAMING_SNAKE_CASE__     :Dict						    = self.prepare_config_and_inputs()
													(
													    (
													    SCREAMING_SNAKE_CASE__
													)    ,      (
													    SCREAMING_SNAKE_CASE__
													)    ,      (
													    SCREAMING_SNAKE_CASE__
													)    ,      (
													    SCREAMING_SNAKE_CASE__
													)    ,      (
													    SCREAMING_SNAKE_CASE__
													)    ,      (
													    SCREAMING_SNAKE_CASE__
													)    ,      (
													    SCREAMING_SNAKE_CASE__
													)    ,      
													)     :Union[str, Any]						    = config_and_inputs
													SCREAMING_SNAKE_CASE__     :Tuple						    = {'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
													return config, inputs_dict




@require_torch
class _SCREAMING_SNAKE_CASE(							_SCREAMING_SNAKE_CASE						,						_SCREAMING_SNAKE_CASE						,						unittest.TestCase					):
						A_			:							List[Any]     		= (
						    (
						        MegatronBertModel,
						        MegatronBertForMaskedLM,
						        MegatronBertForCausalLM,
						        MegatronBertForMultipleChoice,
						        MegatronBertForNextSentencePrediction,
						        MegatronBertForPreTraining,
						        MegatronBertForQuestionAnswering,
						        MegatronBertForSequenceClassification,
						        MegatronBertForTokenClassification,
						    )
						    if is_torch_available()
						    else ()
						)
						A_			:							List[Any]     		= (
						    {
						        'feature-extraction': MegatronBertModel,
						        'fill-mask': MegatronBertForMaskedLM,
						        'question-answering': MegatronBertForQuestionAnswering,
						        'text-classification': MegatronBertForSequenceClassification,
						        'text-generation': MegatronBertForCausalLM,
						        'token-classification': MegatronBertForTokenClassification,
						        'zero-shot': MegatronBertForSequenceClassification,
						    }
						    if is_torch_available()
						    else {}
						)
						A_			:							int     		= True
						# test_resize_embeddings = False
						A_			:							Dict     		= False




						def      __lowerCamelCase       (    self :       Union[str, Any]							,							UpperCamelCase_ :       int							,							UpperCamelCase_ :       int							,							UpperCamelCase_ :       str=False      )     ->				List[Any]:
													SCREAMING_SNAKE_CASE__     :Optional[Any]						    = super()._prepare_for_class(UpperCamelCase_							,							UpperCamelCase_							,							return_labels=UpperCamelCase_      )

													if return_labels:
																				if model_class in get_values(UpperCamelCase_      ):
																											SCREAMING_SNAKE_CASE__     :str						    = torch.zeros(
																											    (self.model_tester.batch_size, self.model_tester.seq_length)							,							dtype=torch.long							,							device=UpperCamelCase_      )
																											SCREAMING_SNAKE_CASE__     :List[str]						    = torch.zeros(
																											    self.model_tester.batch_size							,							dtype=torch.long							,							device=UpperCamelCase_      )
													return inputs_dict




						def      __lowerCamelCase       (    self :       Union[str, Any]      )     ->				Any:
													SCREAMING_SNAKE_CASE__     :Dict						    = MegatronBertModelTester(self      )
													SCREAMING_SNAKE_CASE__     :Union[str, Any]						    = ConfigTester(self							,							config_class=UpperCamelCase_							,							hidden_size=37      )




						def      __lowerCamelCase       (    self :       List[str]      )     ->				Optional[int]:
													self.config_tester.run_common_tests()




						def      __lowerCamelCase       (    self :       List[str]      )     ->				List[str]:
													SCREAMING_SNAKE_CASE__     :Dict						    = self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_megatron_bert_model(*UpperCamelCase_      )




						def      __lowerCamelCase       (    self :       Any      )     ->				Tuple:
													SCREAMING_SNAKE_CASE__     :List[Any]						    = self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_megatron_bert_for_masked_lm(*UpperCamelCase_      )




						def      __lowerCamelCase       (    self :       Union[str, Any]      )     ->				str:
													SCREAMING_SNAKE_CASE__     :Any						    = self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*UpperCamelCase_      )




						def      __lowerCamelCase       (    self :       str      )     ->				Dict:
													SCREAMING_SNAKE_CASE__     :Dict						    = self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*UpperCamelCase_      )




						def      __lowerCamelCase       (    self :       int      )     ->				Any:
													SCREAMING_SNAKE_CASE__     :Dict						    = self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_megatron_bert_for_pretraining(*UpperCamelCase_      )




						def      __lowerCamelCase       (    self :       Dict      )     ->				Union[str, Any]:
													SCREAMING_SNAKE_CASE__     :Union[str, Any]						    = self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_megatron_bert_for_question_answering(*UpperCamelCase_      )




						def      __lowerCamelCase       (    self :       Any      )     ->				Union[str, Any]:
													SCREAMING_SNAKE_CASE__     :Union[str, Any]						    = self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*UpperCamelCase_      )




						def      __lowerCamelCase       (    self :       List[Any]      )     ->				Dict:
													SCREAMING_SNAKE_CASE__     :Dict						    = self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_megatron_bert_for_token_classification(*UpperCamelCase_      )


def    lowerCamelCase       (			UpperCAmelCase__				:  Dict		)			->	Any:


							'''simple docstring'''






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


UpperCamelCase_      	=	1e-4





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




						@slow
						@unittest.skip('Model is not available.'      )
						def      __lowerCamelCase       (    self :       str      )     ->				int:
													SCREAMING_SNAKE_CASE__     :Tuple						    = 'nvidia/megatron-bert-uncased-345m'
													if "MYDIR" in os.environ:
																				SCREAMING_SNAKE_CASE__     :Dict						    = os.path.join(os.environ['MYDIR']							,							UpperCamelCase_      )
													SCREAMING_SNAKE_CASE__     :Dict						    = MegatronBertModel.from_pretrained(UpperCamelCase_      )
													model.to(UpperCamelCase_      )
													model.half()
													SCREAMING_SNAKE_CASE__     :Optional[int]						    = _long_tensor([[1_01, 71_10, 10_05, 10_56, 20_23, 1_13_33, 1_74_13, 10_29, 1_02]]      )
													with torch.no_grad():
																				SCREAMING_SNAKE_CASE__     :str						    = model(UpperCamelCase_      )[0]
													SCREAMING_SNAKE_CASE__     :Dict						    = torch.Size((1, 9, 10_24)      )
													self.assertEqual(output.shape							,							UpperCamelCase_      )

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






