






"""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        lowercase__      :


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


			def      UpperCAmelCase__						(				self					:  Union[str, Any] ):
										lowerCamelCase_	:							Union[str, Any]									=ids_tensor([self.batch_size, self.seq_length]       ,	self.vocab_size )

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

										lowerCamelCase_	:							Optional[int]									=None
										if self.use_token_type_ids:
																	lowerCamelCase_	:							Union[str, Any]									=ids_tensor([self.batch_size, self.seq_length]       ,	self.type_vocab_size )

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

										lowerCamelCase_	:							Any									=self.get_config()

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


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


			def      UpperCAmelCase__						(				self					:  str       ,	snake_case__					:  List[Any]       ,	snake_case__					:  List[str]       ,	snake_case__					:  Tuple       ,	snake_case__					:  List[str]       ,	snake_case__					:  str       ,	snake_case__					:  str       ,	snake_case__					:  Optional[int] ):
										lowerCamelCase_	:							Optional[int]									=MegatronBertModel(config=snake_case__ )
										model.to(snake_case__ )
										model.eval()
										lowerCamelCase_	:							Any									=model(snake_case__       ,	attention_mask=snake_case__       ,	token_type_ids=snake_case__ )
										lowerCamelCase_	:							List[str]									=model(snake_case__       ,	token_type_ids=snake_case__ )
										lowerCamelCase_	:							str									=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      UpperCAmelCase__						(				self					:  Union[str, Any]       ,	snake_case__					:  Optional[Any]       ,	snake_case__					:  Optional[int]       ,	snake_case__					:  Optional[Any]       ,	snake_case__					:  str       ,	snake_case__					:  List[Any]       ,	snake_case__					:  List[str]       ,	snake_case__					:  Any ):
										lowerCamelCase_	:							List[str]									=MegatronBertForMaskedLM(config=snake_case__ )
										model.to(snake_case__ )
										model.eval()
										lowerCamelCase_	:							Union[str, Any]									=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      UpperCAmelCase__						(				self					:  Optional[Any]       ,	snake_case__					:  Dict       ,	snake_case__					:  Tuple       ,	snake_case__					:  List[str]       ,	snake_case__					:  Union[str, Any]       ,	snake_case__					:  Union[str, Any]       ,	snake_case__					:  List[str]       ,	snake_case__					:  str ):
										lowerCamelCase_	:							int									=MegatronBertForCausalLM(config=snake_case__ )
										model.to(snake_case__ )
										model.eval()
										lowerCamelCase_	:							int									=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      UpperCAmelCase__						(				self					:  Union[str, Any]       ,	snake_case__					:  Tuple       ,	snake_case__					:  List[str]       ,	snake_case__					:  int       ,	snake_case__					:  int       ,	snake_case__					:  int       ,	snake_case__					:  Any       ,	snake_case__					:  Dict ):
										lowerCamelCase_	:							str									=MegatronBertForNextSentencePrediction(config=snake_case__ )
										model.to(snake_case__ )
										model.eval()
										lowerCamelCase_	:							Union[str, Any]									=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      UpperCAmelCase__						(				self					:  Optional[int]       ,	snake_case__					:  Dict       ,	snake_case__					:  Optional[Any]       ,	snake_case__					:  Tuple       ,	snake_case__					:  List[Any]       ,	snake_case__					:  Optional[int]       ,	snake_case__					:  Any       ,	snake_case__					:  int ):
										lowerCamelCase_	:							Any									=MegatronBertForPreTraining(config=snake_case__ )
										model.to(snake_case__ )
										model.eval()
										lowerCamelCase_	:							Dict									=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      UpperCAmelCase__						(				self					:  List[str]       ,	snake_case__					:  str       ,	snake_case__					:  int       ,	snake_case__					:  Optional[Any]       ,	snake_case__					:  Tuple       ,	snake_case__					:  int       ,	snake_case__					:  List[Any]       ,	snake_case__					:  int ):
										lowerCamelCase_	:							Any									=MegatronBertForQuestionAnswering(config=snake_case__ )
										model.to(snake_case__ )
										model.eval()
										lowerCamelCase_	:							Dict									=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      UpperCAmelCase__						(				self					:  Any       ,	snake_case__					:  str       ,	snake_case__					:  List[str]       ,	snake_case__					:  int       ,	snake_case__					:  Optional[int]       ,	snake_case__					:  Optional[Any]       ,	snake_case__					:  Optional[int]       ,	snake_case__					:  str ):
										lowerCamelCase_	:							List[str]									=self.num_labels
										lowerCamelCase_	:							str									=MegatronBertForSequenceClassification(snake_case__ )
										model.to(snake_case__ )
										model.eval()
										lowerCamelCase_	:							Any									=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      UpperCAmelCase__						(				self					:  List[Any]       ,	snake_case__					:  Tuple       ,	snake_case__					:  Any       ,	snake_case__					:  Union[str, Any]       ,	snake_case__					:  int       ,	snake_case__					:  Any       ,	snake_case__					:  List[Any]       ,	snake_case__					:  Dict ):
										lowerCamelCase_	:							int									=self.num_labels
										lowerCamelCase_	:							Tuple									=MegatronBertForTokenClassification(config=snake_case__ )
										model.to(snake_case__ )
										model.eval()
										lowerCamelCase_	:							List[Any]									=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      UpperCAmelCase__						(				self					:  List[str]       ,	snake_case__					:  int       ,	snake_case__					:  List[str]       ,	snake_case__					:  Dict       ,	snake_case__					:  List[str]       ,	snake_case__					:  Optional[Any]       ,	snake_case__					:  Union[str, Any]       ,	snake_case__					:  str ):
										lowerCamelCase_	:							Union[str, Any]									=self.num_choices
										lowerCamelCase_	:							List[str]									=MegatronBertForMultipleChoice(config=snake_case__ )
										model.to(snake_case__ )
										model.eval()
										lowerCamelCase_	:							str									=input_ids.unsqueeze(1 ).expand(-1       ,	self.num_choices       ,	-1 ).contiguous()
										lowerCamelCase_	:							Optional[Any]									=token_type_ids.unsqueeze(1 ).expand(-1       ,	self.num_choices       ,	-1 ).contiguous()
										lowerCamelCase_	:							Dict									=input_mask.unsqueeze(1 ).expand(-1       ,	self.num_choices       ,	-1 ).contiguous()
										lowerCamelCase_	:							Optional[int]									=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      UpperCAmelCase__						(				self					:  Union[str, Any] ):
										lowerCamelCase_	:							Tuple									=self.prepare_config_and_inputs()
										(
										    (
										    lowerCamelCase_
										)							,						(
										    lowerCamelCase_
										)							,						(
										    lowerCamelCase_
										)							,						(
										    lowerCamelCase_
										)							,						(
										    lowerCamelCase_
										)							,						(
										    lowerCamelCase_
										)							,						(
										    lowerCamelCase_
										)							,						
										)	:							str									=config_and_inputs
										lowerCamelCase_	:							str									={"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
										return config, inputs_dict







@require_torch
class        lowercase__      (							snake_case__,       snake_case__,       unittest.TestCase       ):
			_UpperCAmelCase		:Tuple             =			(
			    (
			        MegatronBertModel,
			        MegatronBertForMaskedLM,
			        MegatronBertForCausalLM,
			        MegatronBertForMultipleChoice,
			        MegatronBertForNextSentencePrediction,
			        MegatronBertForPreTraining,
			        MegatronBertForQuestionAnswering,
			        MegatronBertForSequenceClassification,
			        MegatronBertForTokenClassification,
			    )
			    if is_torch_available()
			    else ()
			)
			_UpperCAmelCase		:Optional[int]             =			(
			    {
			        "feature-extraction": MegatronBertModel,
			        "fill-mask": MegatronBertForMaskedLM,
			        "question-answering": MegatronBertForQuestionAnswering,
			        "text-classification": MegatronBertForSequenceClassification,
			        "text-generation": MegatronBertForCausalLM,
			        "token-classification": MegatronBertForTokenClassification,
			        "zero-shot": MegatronBertForSequenceClassification,
			    }
			    if is_torch_available()
			    else {}
			)
			_UpperCAmelCase		:Union[str, Any]             =			True
			# test_resize_embeddings = False
			_UpperCAmelCase		:Tuple             =			False


			def      UpperCAmelCase__						(				self					:  Union[str, Any]       ,	snake_case__					:  Optional[Any]       ,	snake_case__					:  Optional[int]       ,	snake_case__					:  int=False ):
										lowerCamelCase_	:							List[Any]									=super()._prepare_for_class(snake_case__       ,	snake_case__       ,	return_labels=snake_case__ )

										if return_labels:
																	if model_class in get_values(snake_case__ ):
																								lowerCamelCase_	:							Tuple									=torch.zeros(
																								    (self.model_tester.batch_size, self.model_tester.seq_length)       ,	dtype=torch.long       ,	device=snake_case__ )
																								lowerCamelCase_	:							Any									=torch.zeros(
																								    self.model_tester.batch_size       ,	dtype=torch.long       ,	device=snake_case__ )
										return inputs_dict


			def      UpperCAmelCase__						(				self					:  Union[str, Any] ):
										lowerCamelCase_	:							int									=MegatronBertModelTester(self )
										lowerCamelCase_	:							str									=ConfigTester(self       ,	config_class=snake_case__       ,	hidden_size=37 )


			def      UpperCAmelCase__						(				self					:  Optional[int] ):
										self.config_tester.run_common_tests()


			def      UpperCAmelCase__						(				self					:  Any ):
										lowerCamelCase_	:							str									=self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_megatron_bert_model(*snake_case__ )


			def      UpperCAmelCase__						(				self					:  List[str] ):
										lowerCamelCase_	:							Optional[int]									=self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_megatron_bert_for_masked_lm(*snake_case__ )


			def      UpperCAmelCase__						(				self					:  List[str] ):
										lowerCamelCase_	:							Optional[Any]									=self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*snake_case__ )


			def      UpperCAmelCase__						(				self					:  List[Any] ):
										lowerCamelCase_	:							Tuple									=self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*snake_case__ )


			def      UpperCAmelCase__						(				self					:  Optional[int] ):
										lowerCamelCase_	:							int									=self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_megatron_bert_for_pretraining(*snake_case__ )


			def      UpperCAmelCase__						(				self					:  Tuple ):
										lowerCamelCase_	:							str									=self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_megatron_bert_for_question_answering(*snake_case__ )


			def      UpperCAmelCase__						(				self					:  Tuple ):
										lowerCamelCase_	:							List[Any]									=self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*snake_case__ )


			def      UpperCAmelCase__						(				self					:  int ):
										lowerCamelCase_	:							int									=self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_megatron_bert_for_token_classification(*snake_case__ )


def 			_snake_case    (					lowerCamelCase__     :   Dict  )    ->     Optional[int]:
							return torch.tensor(
							    lowerCamelCase__    ,						dtype=torch.long    ,						device=lowerCamelCase__    ,						)


A__	:    str			 = 1E-4







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


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

										lowerCamelCase_	:							Union[str, Any]									=[-0.6_040, -0.2_517, -0.1_025, 0.3_420, -0.6_758, -0.0_017, -0.1_089, -0.1_990, 0.5_728]
										for ii in range(3 ):
																	for jj in range(3 ):
																								lowerCamelCase_	:							Optional[int]									=output[0, ii, jj]
																								lowerCamelCase_	:							List[str]									=expected[3 * ii + jj]
																								lowerCamelCase_	:							Optional[int]									="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__ )











"""simple docstring"""



from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


A__	:    List[Any]			 = {
    'configuration_lilt': ['LILT_PRETRAINED_CONFIG_ARCHIVE_MAP', 'LiltConfig'],
}

try:
						if not is_torch_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						A__	:    List[str]			 = [
						    'LILT_PRETRAINED_MODEL_ARCHIVE_LIST',
						    'LiltForQuestionAnswering',
						    'LiltForSequenceClassification',
						    'LiltForTokenClassification',
						    'LiltModel',
						    'LiltPreTrainedModel',
						]

if TYPE_CHECKING:
						from .configuration_lilt import LILT_PRETRAINED_CONFIG_ARCHIVE_MAP, LiltConfig

						try:
												if not is_torch_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_lilt import (
												    LILT_PRETRAINED_MODEL_ARCHIVE_LIST,
												    LiltForQuestionAnswering,
												    LiltForSequenceClassification,
												    LiltForTokenClassification,
												    LiltModel,
												    LiltPreTrainedModel,
												)

else:
						import sys

						A__	:    Optional[Any]			 = _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)



