






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



  def __init__(      self ,		snake_case__ ,		snake_case__=13 ,		snake_case__=7 ,		snake_case__=True ,		snake_case__=True ,		snake_case__=True ,		snake_case__=True ,		snake_case__=99 ,		snake_case__=64 ,		snake_case__=32 ,		snake_case__=5 ,		snake_case__=4 ,		snake_case__=37 ,		snake_case__="gelu" ,		snake_case__=0.1 ,		snake_case__=0.1 ,		snake_case__=5_12 ,		snake_case__=16 ,		snake_case__=2 ,		snake_case__=0.02 ,		snake_case__=3 ,		snake_case__=4 ,		snake_case__=None ,		)			->			Union[str, Any]:



         """simple docstring"""





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



  def 		UpperCamelCase_				(      self				)			->			str:



         """simple docstring"""





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

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

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

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

         UpperCAmelCase	    =			self.get_config()

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



  def 		UpperCamelCase_				(      self				)			->			Tuple:



         """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=UpperCamelCase__ ,		initializer_range=self.initializer_range ,		)



  def 		UpperCamelCase_				(      self ,		snake_case__ ,		snake_case__ ,		snake_case__ ,		snake_case__ ,		snake_case__ ,		snake_case__ ,		snake_case__				)			->			Optional[Any]:



         """simple docstring"""





         UpperCAmelCase	    =			MegatronBertModel(config=UpperCamelCase__				)
         model.to(UpperCamelCase__				)
         model.eval()
         UpperCAmelCase	    =			model(UpperCamelCase__ ,		attention_mask=UpperCamelCase__ ,		token_type_ids=UpperCamelCase__				)
         UpperCAmelCase	    =			model(UpperCamelCase__ ,		token_type_ids=UpperCamelCase__				)
         UpperCAmelCase	    =			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 		UpperCamelCase_				(      self ,		snake_case__ ,		snake_case__ ,		snake_case__ ,		snake_case__ ,		snake_case__ ,		snake_case__ ,		snake_case__				)			->			Optional[int]:



         """simple docstring"""





         UpperCAmelCase	    =			MegatronBertForMaskedLM(config=UpperCamelCase__				)
         model.to(UpperCamelCase__				)
         model.eval()
         UpperCAmelCase	    =			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 		UpperCamelCase_				(      self ,		snake_case__ ,		snake_case__ ,		snake_case__ ,		snake_case__ ,		snake_case__ ,		snake_case__ ,		snake_case__				)			->			Tuple:



         """simple docstring"""





         UpperCAmelCase	    =			MegatronBertForCausalLM(config=UpperCamelCase__				)
         model.to(UpperCamelCase__				)
         model.eval()
         UpperCAmelCase	    =			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 		UpperCamelCase_				(      self ,		snake_case__ ,		snake_case__ ,		snake_case__ ,		snake_case__ ,		snake_case__ ,		snake_case__ ,		snake_case__				)			->			str:



         """simple docstring"""





         UpperCAmelCase	    =			MegatronBertForNextSentencePrediction(config=UpperCamelCase__				)
         model.to(UpperCamelCase__				)
         model.eval()
         UpperCAmelCase	    =			model(
             UpperCamelCase__ ,		attention_mask=UpperCamelCase__ ,		token_type_ids=UpperCamelCase__ ,		labels=UpperCamelCase__ ,		)
         self.parent.assertEqual(result.logits.shape ,		(self.batch_size, 2)				)



  def 		UpperCamelCase_				(      self ,		snake_case__ ,		snake_case__ ,		snake_case__ ,		snake_case__ ,		snake_case__ ,		snake_case__ ,		snake_case__				)			->			Optional[int]:



         """simple docstring"""





         UpperCAmelCase	    =			MegatronBertForPreTraining(config=UpperCamelCase__				)
         model.to(UpperCamelCase__				)
         model.eval()
         UpperCAmelCase	    =			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 		UpperCamelCase_				(      self ,		snake_case__ ,		snake_case__ ,		snake_case__ ,		snake_case__ ,		snake_case__ ,		snake_case__ ,		snake_case__				)			->			List[str]:



         """simple docstring"""





         UpperCAmelCase	    =			MegatronBertForQuestionAnswering(config=UpperCamelCase__				)
         model.to(UpperCamelCase__				)
         model.eval()
         UpperCAmelCase	    =			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 		UpperCamelCase_				(      self ,		snake_case__ ,		snake_case__ ,		snake_case__ ,		snake_case__ ,		snake_case__ ,		snake_case__ ,		snake_case__				)			->			Optional[Any]:



         """simple docstring"""





         UpperCAmelCase	    =			self.num_labels
         UpperCAmelCase	    =			MegatronBertForSequenceClassification(UpperCamelCase__				)
         model.to(UpperCamelCase__				)
         model.eval()
         UpperCAmelCase	    =			model(UpperCamelCase__ ,		attention_mask=UpperCamelCase__ ,		token_type_ids=UpperCamelCase__ ,		labels=UpperCamelCase__				)
         self.parent.assertEqual(result.logits.shape ,		(self.batch_size, self.num_labels)				)



  def 		UpperCamelCase_				(      self ,		snake_case__ ,		snake_case__ ,		snake_case__ ,		snake_case__ ,		snake_case__ ,		snake_case__ ,		snake_case__				)			->			Any:



         """simple docstring"""





         UpperCAmelCase	    =			self.num_labels
         UpperCAmelCase	    =			MegatronBertForTokenClassification(config=UpperCamelCase__				)
         model.to(UpperCamelCase__				)
         model.eval()
         UpperCAmelCase	    =			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 		UpperCamelCase_				(      self ,		snake_case__ ,		snake_case__ ,		snake_case__ ,		snake_case__ ,		snake_case__ ,		snake_case__ ,		snake_case__				)			->			Tuple:



         """simple docstring"""





         UpperCAmelCase	    =			self.num_choices
         UpperCAmelCase	    =			MegatronBertForMultipleChoice(config=UpperCamelCase__				)
         model.to(UpperCamelCase__				)
         model.eval()
         UpperCAmelCase	    =			input_ids.unsqueeze(1				).expand(-1 ,		self.num_choices ,		-1				).contiguous()
         UpperCAmelCase	    =			token_type_ids.unsqueeze(1				).expand(-1 ,		self.num_choices ,		-1				).contiguous()
         UpperCAmelCase	    =			input_mask.unsqueeze(1				).expand(-1 ,		self.num_choices ,		-1				).contiguous()
         UpperCAmelCase	    =			model(
             UpperCamelCase__ ,		attention_mask=UpperCamelCase__ ,		token_type_ids=UpperCamelCase__ ,		labels=UpperCamelCase__ ,		)
         self.parent.assertEqual(result.logits.shape ,		(self.batch_size, self.num_choices)				)



  def 		UpperCamelCase_				(      self				)			->			str:



         """simple docstring"""





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





@require_torch
class   UpperCamelCase_ (       __A   ,					__A   ,					unittest.TestCase  ):
  _A       :					Any				=	(
      (
          MegatronBertModel,
          MegatronBertForMaskedLM,
          MegatronBertForCausalLM,
          MegatronBertForMultipleChoice,
          MegatronBertForNextSentencePrediction,
          MegatronBertForPreTraining,
          MegatronBertForQuestionAnswering,
          MegatronBertForSequenceClassification,
          MegatronBertForTokenClassification,
      )
      if is_torch_available()
      else ()
  )
  _A       :					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 {}
  )
  _A       :					str				=	True
  # test_resize_embeddings = False
  _A       :					int				=	False



  def 		UpperCamelCase_				(      self ,		snake_case__ ,		snake_case__ ,		snake_case__=False				)			->			str:



         """simple docstring"""





         UpperCAmelCase	    =			super()._prepare_for_class(UpperCamelCase__ ,		UpperCamelCase__ ,		return_labels=UpperCamelCase__				)

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



  def 		UpperCamelCase_				(      self				)			->			str:



         """simple docstring"""





         UpperCAmelCase	    =			MegatronBertModelTester(self				)
         UpperCAmelCase	    =			ConfigTester(self ,		config_class=UpperCamelCase__ ,		hidden_size=37				)



  def 		UpperCamelCase_				(      self				)			->			str:



         """simple docstring"""





         self.config_tester.run_common_tests()



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



         """simple docstring"""





         UpperCAmelCase	    =			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_megatron_bert_model(*UpperCamelCase__				)



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



         """simple docstring"""





         UpperCAmelCase	    =			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_megatron_bert_for_masked_lm(*UpperCamelCase__				)



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



         """simple docstring"""





         UpperCAmelCase	    =			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*UpperCamelCase__				)



  def 		UpperCamelCase_				(      self				)			->			Dict:



         """simple docstring"""





         UpperCAmelCase	    =			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*UpperCamelCase__				)



  def 		UpperCamelCase_				(      self				)			->			int:



         """simple docstring"""





         UpperCAmelCase	    =			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_megatron_bert_for_pretraining(*UpperCamelCase__				)



  def 		UpperCamelCase_				(      self				)			->			Any:



         """simple docstring"""





         UpperCAmelCase	    =			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_megatron_bert_for_question_answering(*UpperCamelCase__				)



  def 		UpperCamelCase_				(      self				)			->			Tuple:



         """simple docstring"""





         UpperCAmelCase	    =			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*UpperCamelCase__				)



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



         """simple docstring"""





         UpperCAmelCase	    =			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_megatron_bert_for_token_classification(*UpperCamelCase__				)






def       _lowerCAmelCase       (      lowerCAmelCase		):







       '''simple docstring'''
       return torch.tensor(
           lowerCAmelCase__	,      dtype=torch.long	,      device=lowerCAmelCase__	,      )


lowerCAmelCase_	:    Optional[Any]            =      1E-4





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



  @slow
  @unittest.skip("""Model is not available."""				)
  def 		UpperCamelCase_				(      self				)			->			Optional[int]:



         """simple docstring"""





         UpperCAmelCase	    =			'nvidia/megatron-bert-uncased-345m'
         if "MYDIR" in os.environ:
                UpperCAmelCase	    =			os.path.join(os.environ["""MYDIR"""] ,		UpperCamelCase__				)
         UpperCAmelCase	    =			MegatronBertModel.from_pretrained(UpperCamelCase__				)
         model.to(UpperCamelCase__				)
         model.half()
         UpperCAmelCase	    =			_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():
                UpperCAmelCase	    =			model(UpperCamelCase__				)[0]
         UpperCAmelCase	    =			torch.Size((1, 9, 10_24)				)
         self.assertEqual(output.shape ,		UpperCamelCase__				)

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








"""simple docstring"""







from typing import List, Optional, Union

from ...configuration_utils import PretrainedConfig
from ...utils import logging


lowerCAmelCase_	:    str            =      logging.get_logger(__name__)

lowerCAmelCase_	:    Optional[int]            =      {
    '''huggingface/time-series-transformer-tourism-monthly''': (
        '''https://huggingface.co/huggingface/time-series-transformer-tourism-monthly/resolve/main/config.json'''
    ),
    # See all TimeSeriesTransformer models at https://huggingface.co/models?filter=time_series_transformer
}





class   UpperCamelCase_ (       a_  ):
  _A       :					int				=	'time_series_transformer'
  _A       :					List[Any]				=	{
      'hidden_size': 'd_model',
      'num_attention_heads': 'encoder_attention_heads',
      'num_hidden_layers': 'encoder_layers',
  }



  def __init__(      self ,		snake_case__ = None ,		snake_case__ = None ,		snake_case__ = "student_t" ,		snake_case__ = "nll" ,		snake_case__ = 1 ,		snake_case__ = [1, 2, 3, 4, 5, 6, 7] ,		snake_case__ = "mean" ,		snake_case__ = 0 ,		snake_case__ = 0 ,		snake_case__ = 0 ,		snake_case__ = 0 ,		snake_case__ = None ,		snake_case__ = None ,		snake_case__ = 32 ,		snake_case__ = 32 ,		snake_case__ = 2 ,		snake_case__ = 2 ,		snake_case__ = 2 ,		snake_case__ = 2 ,		snake_case__ = True ,		snake_case__ = "gelu" ,		snake_case__ = 64 ,		snake_case__ = 0.1 ,		snake_case__ = 0.1 ,		snake_case__ = 0.1 ,		snake_case__ = 0.1 ,		snake_case__ = 0.1 ,		snake_case__ = 1_00 ,		snake_case__ = 0.02 ,		snake_case__=True ,		**snake_case__ ,		)			->			Union[str, Any]:



         """simple docstring"""





         UpperCAmelCase	    =			prediction_length
         UpperCAmelCase	    =			context_length or prediction_length
         UpperCAmelCase	    =			distribution_output
         UpperCAmelCase	    =			loss
         UpperCAmelCase	    =			input_size
         UpperCAmelCase	    =			num_time_features
         UpperCAmelCase	    =			lags_sequence
         UpperCAmelCase	    =			scaling
         UpperCAmelCase	    =			num_dynamic_real_features
         UpperCAmelCase	    =			num_static_real_features
         UpperCAmelCase	    =			num_static_categorical_features
         if cardinality and num_static_categorical_features > 0:
                if len(snake_case__				) != num_static_categorical_features:
                       raise ValueError(
                           """The cardinality should be a list of the same length as `num_static_categorical_features`"""				)
                UpperCAmelCase	    =			cardinality
         else:
                UpperCAmelCase	    =			[0]
         if embedding_dimension and num_static_categorical_features > 0:
                if len(snake_case__				) != num_static_categorical_features:
                       raise ValueError(
                           """The embedding dimension should be a list of the same length as `num_static_categorical_features`"""				)
                UpperCAmelCase	    =			embedding_dimension
         else:
                UpperCAmelCase	    =			[min(50 ,		(cat + 1) // 2				) for cat in self.cardinality]
         UpperCAmelCase	    =			num_parallel_samples

         # Transformer architecture configuration
         UpperCAmelCase	    =			input_size * len(snake_case__				) + self._number_of_features
         UpperCAmelCase	    =			d_model
         UpperCAmelCase	    =			encoder_attention_heads
         UpperCAmelCase	    =			decoder_attention_heads
         UpperCAmelCase	    =			encoder_ffn_dim
         UpperCAmelCase	    =			decoder_ffn_dim
         UpperCAmelCase	    =			encoder_layers
         UpperCAmelCase	    =			decoder_layers

         UpperCAmelCase	    =			dropout
         UpperCAmelCase	    =			attention_dropout
         UpperCAmelCase	    =			activation_dropout
         UpperCAmelCase	    =			encoder_layerdrop
         UpperCAmelCase	    =			decoder_layerdrop

         UpperCAmelCase	    =			activation_function
         UpperCAmelCase	    =			init_std

         UpperCAmelCase	    =			use_cache

         super().__init__(is_encoder_decoder=snake_case__ ,		**snake_case__				)



  @property
  def 		UpperCamelCase_				(      self				)			->			int:



         """simple docstring"""





         return (
             sum(self.embedding_dimension				)
             + self.num_dynamic_real_features
             + self.num_time_features
             + self.num_static_real_features
             + self.input_size * 2  # the log1p(abs(loc)) and log(scale) features
         )
