
"""simple docstring"""


from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


A_    							=					{
    """configuration_time_series_transformer""": [
        """TIME_SERIES_TRANSFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP""",
        """TimeSeriesTransformerConfig""",
    ],
}

try:
						if not is_torch_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						A_    							=					[
						    """TIME_SERIES_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST""",
						    """TimeSeriesTransformerForPrediction""",
						    """TimeSeriesTransformerModel""",
						    """TimeSeriesTransformerPreTrainedModel""",
						]


if TYPE_CHECKING:
						from .configuration_time_series_transformer import (
						    TIME_SERIES_TRANSFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP,
						    TimeSeriesTransformerConfig,
						)

						try:
												if not is_torch_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_time_series_transformer import (
												    TIME_SERIES_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
												    TimeSeriesTransformerForPrediction,
												    TimeSeriesTransformerModel,
												    TimeSeriesTransformerPreTrainedModel,
												)

else:
						import sys

						A_    							=					_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)






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        _a		:


					'''simple docstring'''






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






					def     __A     (					self     ):
						A__				:       Tuple	  =						ids_tensor([self.batch_size, self.seq_length]       ,			self.vocab_size     )

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

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

						A__				:       Any	  =						None
						A__				:       str	  =						None
						A__				:       Dict	  =						None
						if self.use_labels:
							A__				:       str	  =						ids_tensor([self.batch_size]       ,			self.type_sequence_label_size     )
							A__				:       Optional[int]	  =						ids_tensor([self.batch_size, self.seq_length]       ,			self.num_labels     )
							A__				:       Union[str, Any]	  =						ids_tensor([self.batch_size]       ,			self.num_choices     )

						A__				:       Optional[int]	  =						self.get_config()

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






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






					def     __A     (					self       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__     ):
						A__				:       Dict	  =						MegatronBertModel(config=A__     )
						model.to(A__     )
						model.eval()
						A__				:       str	  =						model(A__       ,			attention_mask=A__       ,			token_type_ids=A__     )
						A__				:       Optional[Any]	  =						model(A__       ,			token_type_ids=A__     )
						A__				:       Dict	  =						model(A__     )

						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       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__     ):
						A__				:       str	  =						MegatronBertForMaskedLM(config=A__     )
						model.to(A__     )
						model.eval()
						A__				:       List[Any]	  =						model(A__       ,			attention_mask=A__       ,			token_type_ids=A__       ,			labels=A__     )
						self.parent.assertEqual(result.logits.shape       ,			(self.batch_size, self.seq_length, self.vocab_size)     )






					def     __A     (					self       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__     ):
						A__				:       Tuple	  =						MegatronBertForCausalLM(config=A__     )
						model.to(A__     )
						model.eval()
						A__				:       int	  =						model(A__       ,			attention_mask=A__       ,			token_type_ids=A__       ,			labels=A__     )
						self.parent.assertEqual(result.logits.shape       ,			(self.batch_size, self.seq_length, self.vocab_size)     )






					def     __A     (					self       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__     ):
						A__				:       Dict	  =						MegatronBertForNextSentencePrediction(config=A__     )
						model.to(A__     )
						model.eval()
						A__				:       Optional[int]	  =						model(
						    A__       ,			attention_mask=A__       ,			token_type_ids=A__       ,			labels=A__       ,			)
						self.parent.assertEqual(result.logits.shape       ,			(self.batch_size, 2)     )






					def     __A     (					self       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__     ):
						A__				:       Optional[Any]	  =						MegatronBertForPreTraining(config=A__     )
						model.to(A__     )
						model.eval()
						A__				:       List[str]	  =						model(
						    A__       ,			attention_mask=A__       ,			token_type_ids=A__       ,			labels=A__       ,			next_sentence_label=A__       ,			)
						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       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__     ):
						A__				:       Optional[Any]	  =						MegatronBertForQuestionAnswering(config=A__     )
						model.to(A__     )
						model.eval()
						A__				:       Optional[Any]	  =						model(
						    A__       ,			attention_mask=A__       ,			token_type_ids=A__       ,			start_positions=A__       ,			end_positions=A__       ,			)
						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       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__     ):
						A__				:       Optional[int]	  =						self.num_labels
						A__				:       Union[str, Any]	  =						MegatronBertForSequenceClassification(A__     )
						model.to(A__     )
						model.eval()
						A__				:       List[str]	  =						model(A__       ,			attention_mask=A__       ,			token_type_ids=A__       ,			labels=A__     )
						self.parent.assertEqual(result.logits.shape       ,			(self.batch_size, self.num_labels)     )






					def     __A     (					self       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__       ,			A__     ):
						A__				:       Union[str, Any]	  =						self.num_labels
						A__				:       int	  =						MegatronBertForTokenClassification(config=A__     )
						model.to(A__     )
						model.eval()
						A__				:       Optional[int]	  =						model(A__       ,			attention_mask=A__       ,			token_type_ids=A__       ,			labels=A__     )
						self.parent.assertEqual(result.logits.shape       ,			(self.batch_size, self.seq_length, self.num_labels)     )






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






					def     __A     (					self     ):
						A__				:       Optional[Any]	  =						self.prepare_config_and_inputs()
						(
						    (
						    A__
						)  ,  (
						    A__
						)  ,  (
						    A__
						)  ,  (
						    A__
						)  ,  (
						    A__
						)  ,  (
						    A__
						)  ,  (
						    A__
						)  ,  
						)				:       Any	  =						config_and_inputs
						A__				:       Any	  =						{"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
						return config, inputs_dict



@require_torch
class        _a		(__magic_name__						,  __magic_name__						,  unittest.TestCase							):


					'''simple docstring'''






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






					def     __A     (					self       ,			A__       ,			A__       ,			A__=False     ):
						A__				:       Union[str, Any]	  =						super()._prepare_for_class(A__       ,			A__       ,			return_labels=A__     )

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






					def     __A     (					self     ):
						A__				:       Union[str, Any]	  =						MegatronBertModelTester(self     )
						A__				:       Union[str, Any]	  =						ConfigTester(self       ,			config_class=A__       ,			hidden_size=37     )






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






					def     __A     (					self     ):
						A__				:       List[str]	  =						self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_model(*A__     )






					def     __A     (					self     ):
						A__				:       Dict	  =						self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_for_masked_lm(*A__     )






					def     __A     (					self     ):
						A__				:       Tuple	  =						self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*A__     )






					def     __A     (					self     ):
						A__				:       Dict	  =						self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*A__     )






					def     __A     (					self     ):
						A__				:       Union[str, Any]	  =						self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_for_pretraining(*A__     )






					def     __A     (					self     ):
						A__				:       Dict	  =						self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_for_question_answering(*A__     )






					def     __A     (					self     ):
						A__				:       int	  =						self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*A__     )






					def     __A     (					self     ):
						A__				:       Tuple	  =						self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_for_token_classification(*A__     )



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


A_			:					int    			=     1E-4



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


					'''simple docstring'''






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

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



