

"""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		:			List[Any]					,A		:			List[str]					,A		:			Union[str, Any]=13					,A		:			List[str]=7					,A		:			str=True					,A		:			Optional[int]=True					,A		:			List[str]=True					,A		:			Tuple=True					,A		:			str=99					,A		:			str=64					,A		:			Dict=32					,A		:			str=5					,A		:			Dict=4					,A		:			int=37					,A		:			Union[str, Any]="gelu"					,A		:			Any=0.1					,A		:			List[str]=0.1					,A		:			str=512					,A		:			Optional[Any]=16					,A		:			Optional[Any]=2					,A		:			Union[str, Any]=0.0_2					,A		:			Optional[Any]=3					,A		:			Tuple=4					,A		:			int=None					,):






   '''simple docstring'''
   UpperCAmelCase__					:  List[str]					   =					parent
   UpperCAmelCase__					:  Dict					   =					batch_size
   UpperCAmelCase__					:  Optional[Any]					   =					seq_length
   UpperCAmelCase__					:  Tuple					   =					is_training
   UpperCAmelCase__					:  Tuple					   =					use_input_mask
   UpperCAmelCase__					:  str					   =					use_token_type_ids
   UpperCAmelCase__					:  int					   =					use_labels
   UpperCAmelCase__					:  Tuple					   =					vocab_size
   UpperCAmelCase__					:  Optional[Any]					   =					hidden_size
   UpperCAmelCase__					:  Tuple					   =					embedding_size
   UpperCAmelCase__					:  str					   =					num_hidden_layers
   UpperCAmelCase__					:  Optional[int]					   =					num_attention_heads
   UpperCAmelCase__					:  Union[str, Any]					   =					intermediate_size
   UpperCAmelCase__					:  Union[str, Any]					   =					hidden_act
   UpperCAmelCase__					:  Tuple					   =					hidden_dropout_prob
   UpperCAmelCase__					:  Union[str, Any]					   =					attention_probs_dropout_prob
   UpperCAmelCase__					:  List[Any]					   =					max_position_embeddings
   UpperCAmelCase__					:  Tuple					   =					type_vocab_size
   UpperCAmelCase__					:  str					   =					type_sequence_label_size
   UpperCAmelCase__					:  Optional[int]					   =					initializer_range
   UpperCAmelCase__					:  Union[str, Any]					   =					num_labels
   UpperCAmelCase__					:  str					   =					num_choices
   UpperCAmelCase__					:  List[Any]					   =					scope

  def    __lowercase       (  self		:			Any    ):






   '''simple docstring'''
   UpperCAmelCase__					:  List[Any]					   =					ids_tensor([self.batch_size, self.seq_length]					,self.vocab_size    )

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

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

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

   UpperCAmelCase__					:  List[str]					   =					self.get_config()

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

  def    __lowercase       (  self		:			int    ):






   '''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=A					,initializer_range=self.initializer_range					,)

  def    __lowercase       (  self		:			Union[str, Any]					,A		:			List[str]					,A		:			Optional[Any]					,A		:			List[str]					,A		:			List[str]					,A		:			Tuple					,A		:			int					,A		:			int    ):






   '''simple docstring'''
   UpperCAmelCase__					:  Optional[int]					   =					MegatronBertModel(config=A    )
   model.to(A    )
   model.eval()
   UpperCAmelCase__					:  Union[str, Any]					   =					model(A					,attention_mask=A					,token_type_ids=A    )
   UpperCAmelCase__					:  str					   =					model(A					,token_type_ids=A    )
   UpperCAmelCase__					:  Optional[int]					   =					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    __lowercase       (  self		:			str					,A		:			Union[str, Any]					,A		:			int					,A		:			Dict					,A		:			str					,A		:			Dict					,A		:			Optional[int]					,A		:			Optional[Any]    ):






   '''simple docstring'''
   UpperCAmelCase__					:  Optional[int]					   =					MegatronBertForMaskedLM(config=A    )
   model.to(A    )
   model.eval()
   UpperCAmelCase__					:  Dict					   =					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    __lowercase       (  self		:			str					,A		:			Optional[int]					,A		:			Any					,A		:			Dict					,A		:			Union[str, Any]					,A		:			int					,A		:			str					,A		:			Optional[Any]    ):






   '''simple docstring'''
   UpperCAmelCase__					:  List[Any]					   =					MegatronBertForCausalLM(config=A    )
   model.to(A    )
   model.eval()
   UpperCAmelCase__					:  List[str]					   =					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    __lowercase       (  self		:			List[Any]					,A		:			Any					,A		:			Optional[Any]					,A		:			Union[str, Any]					,A		:			Tuple					,A		:			List[Any]					,A		:			Union[str, Any]					,A		:			Dict    ):






   '''simple docstring'''
   UpperCAmelCase__					:  str					   =					MegatronBertForNextSentencePrediction(config=A    )
   model.to(A    )
   model.eval()
   UpperCAmelCase__					:  Optional[Any]					   =					model(
       A					,attention_mask=A					,token_type_ids=A					,labels=A					,)
   self.parent.assertEqual(result.logits.shape					,(self.batch_size, 2)    )

  def    __lowercase       (  self		:			Optional[int]					,A		:			Optional[int]					,A		:			Dict					,A		:			Any					,A		:			List[Any]					,A		:			Any					,A		:			List[Any]					,A		:			Any    ):






   '''simple docstring'''
   UpperCAmelCase__					:  List[Any]					   =					MegatronBertForPreTraining(config=A    )
   model.to(A    )
   model.eval()
   UpperCAmelCase__					:  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    __lowercase       (  self		:			List[Any]					,A		:			Dict					,A		:			List[str]					,A		:			int					,A		:			Union[str, Any]					,A		:			Tuple					,A		:			str					,A		:			Tuple    ):






   '''simple docstring'''
   UpperCAmelCase__					:  Tuple					   =					MegatronBertForQuestionAnswering(config=A    )
   model.to(A    )
   model.eval()
   UpperCAmelCase__					:  Tuple					   =					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    __lowercase       (  self		:			Optional[Any]					,A		:			List[str]					,A		:			Optional[int]					,A		:			Optional[int]					,A		:			Optional[Any]					,A		:			Tuple					,A		:			Union[str, Any]					,A		:			Optional[int]    ):






   '''simple docstring'''
   UpperCAmelCase__					:  Optional[int]					   =					self.num_labels
   UpperCAmelCase__					:  Optional[Any]					   =					MegatronBertForSequenceClassification(A    )
   model.to(A    )
   model.eval()
   UpperCAmelCase__					:  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    __lowercase       (  self		:			Optional[int]					,A		:			Optional[int]					,A		:			Union[str, Any]					,A		:			Union[str, Any]					,A		:			Tuple					,A		:			Any					,A		:			Any					,A		:			List[str]    ):






   '''simple docstring'''
   UpperCAmelCase__					:  List[Any]					   =					self.num_labels
   UpperCAmelCase__					:  Tuple					   =					MegatronBertForTokenClassification(config=A    )
   model.to(A    )
   model.eval()
   UpperCAmelCase__					:  Optional[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.num_labels)    )

  def    __lowercase       (  self		:			Tuple					,A		:			Tuple					,A		:			List[Any]					,A		:			List[str]					,A		:			Union[str, Any]					,A		:			List[str]					,A		:			Optional[Any]					,A		:			int    ):






   '''simple docstring'''
   UpperCAmelCase__					:  List[Any]					   =					self.num_choices
   UpperCAmelCase__					:  int					   =					MegatronBertForMultipleChoice(config=A    )
   model.to(A    )
   model.eval()
   UpperCAmelCase__					:  List[str]					   =					input_ids.unsqueeze(1    ).expand(-1					,self.num_choices					,-1    ).contiguous()
   UpperCAmelCase__					:  Tuple					   =					token_type_ids.unsqueeze(1    ).expand(-1					,self.num_choices					,-1    ).contiguous()
   UpperCAmelCase__					:  Union[str, Any]					   =					input_mask.unsqueeze(1    ).expand(-1					,self.num_choices					,-1    ).contiguous()
   UpperCAmelCase__					:  Dict					   =					model(
       A					,attention_mask=A					,token_type_ids=A					,labels=A					,)
   self.parent.assertEqual(result.logits.shape					,(self.batch_size, self.num_choices)    )




  def    __lowercase       (  self		:			Tuple    ):






   '''simple docstring'''
   UpperCAmelCase__					:  Union[str, Any]					   =					self.prepare_config_and_inputs()
   (
       (
       UpperCAmelCase__
   )   ,      (
       UpperCAmelCase__
   )   ,      (
       UpperCAmelCase__
   )   ,      (
       UpperCAmelCase__
   )   ,      (
       UpperCAmelCase__
   )   ,      (
       UpperCAmelCase__
   )   ,      (
       UpperCAmelCase__
   )   ,      
   )					:  List[str]					   =					config_and_inputs
   UpperCAmelCase__					:  Any					   =					{"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
   return config, inputs_dict







@require_torch
class      __lowercase			(    __lowerCamelCase	,			__lowerCamelCase	,			unittest.TestCase					):
  snake_case_				       =       (
      (
          MegatronBertModel,
          MegatronBertForMaskedLM,
          MegatronBertForCausalLM,
          MegatronBertForMultipleChoice,
          MegatronBertForNextSentencePrediction,
          MegatronBertForPreTraining,
          MegatronBertForQuestionAnswering,
          MegatronBertForSequenceClassification,
          MegatronBertForTokenClassification,
      )
      if is_torch_available()
      else ()
  )
  snake_case_				       =       (
      {
          """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 {}
  )
  snake_case_				       =       True
  # test_resize_embeddings = False
  snake_case_				       =       False

  def    __lowercase       (  self		:			List[Any]					,A		:			Optional[Any]					,A		:			Any					,A		:			Dict=False    ):






   '''simple docstring'''
   UpperCAmelCase__					:  Optional[Any]					   =					super()._prepare_for_class(A					,A					,return_labels=A    )

   if return_labels:
    if model_class in get_values(A    ):
     UpperCAmelCase__					:  Any					   =					torch.zeros(
         (self.model_tester.batch_size, self.model_tester.seq_length)					,dtype=torch.long					,device=A    )
     UpperCAmelCase__					:  int					   =					torch.zeros(
         self.model_tester.batch_size					,dtype=torch.long					,device=A    )
   return inputs_dict

  def    __lowercase       (  self		:			Union[str, Any]    ):






   '''simple docstring'''
   UpperCAmelCase__					:  List[str]					   =					MegatronBertModelTester(self    )
   UpperCAmelCase__					:  Tuple					   =					ConfigTester(self					,config_class=A					,hidden_size=37    )

  def    __lowercase       (  self		:			List[str]    ):






   '''simple docstring'''
   self.config_tester.run_common_tests()

  def    __lowercase       (  self		:			List[Any]    ):






   '''simple docstring'''
   UpperCAmelCase__					:  Any					   =					self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_megatron_bert_model(*A    )

  def    __lowercase       (  self		:			int    ):






   '''simple docstring'''
   UpperCAmelCase__					:  Any					   =					self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_megatron_bert_for_masked_lm(*A    )

  def    __lowercase       (  self		:			Union[str, Any]    ):






   '''simple docstring'''
   UpperCAmelCase__					:  Optional[int]					   =					self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*A    )

  def    __lowercase       (  self		:			Union[str, Any]    ):






   '''simple docstring'''
   UpperCAmelCase__					:  Optional[Any]					   =					self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*A    )

  def    __lowercase       (  self		:			List[Any]    ):






   '''simple docstring'''
   UpperCAmelCase__					:  List[Any]					   =					self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_megatron_bert_for_pretraining(*A    )

  def    __lowercase       (  self		:			Tuple    ):






   '''simple docstring'''
   UpperCAmelCase__					:  Union[str, Any]					   =					self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_megatron_bert_for_question_answering(*A    )

  def    __lowercase       (  self		:			str    ):






   '''simple docstring'''
   UpperCAmelCase__					:  Union[str, Any]					   =					self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*A    )




  def    __lowercase       (  self		:			Union[str, Any]    ):






   '''simple docstring'''
   UpperCAmelCase__					:  int					   =					self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_megatron_bert_for_token_classification(*A    )







def   lowerCAmelCase						(					__UpperCamelCase      ):
 '''simple docstring'''





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


__UpperCAmelCase =				1E-4







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

  @slow
  @unittest.skip("""Model is not available."""    )
  def    __lowercase       (  self		:			List[str]    ):






   '''simple docstring'''
   UpperCAmelCase__					:  List[Any]					   =					"""nvidia/megatron-bert-uncased-345m"""
   if "MYDIR" in os.environ:
    UpperCAmelCase__					:  int					   =					os.path.join(os.environ["""MYDIR"""]					,A    )
   UpperCAmelCase__					:  Dict					   =					MegatronBertModel.from_pretrained(A    )
   model.to(A    )
   model.half()
   UpperCAmelCase__					:  Optional[int]					   =					_long_tensor([[101, 7_110, 1_005, 1_056, 2_023, 11_333, 17_413, 1_029, 102]]    )
   with torch.no_grad():
    UpperCAmelCase__					:  List[str]					   =					model(A    )[0]
   UpperCAmelCase__					:  str					   =					torch.Size((1, 9, 1_024)    )
   self.assertEqual(output.shape					,A    )

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








"""simple docstring"""




from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


__UpperCAmelCase =				{
    'configuration_table_transformer': [
        'TABLE_TRANSFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP',
        'TableTransformerConfig',
        'TableTransformerOnnxConfig',
    ]
}

try:
      if not is_torch_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      __UpperCAmelCase =				[
          'TABLE_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST',
          'TableTransformerForObjectDetection',
          'TableTransformerModel',
          'TableTransformerPreTrainedModel',
      ]


if TYPE_CHECKING:
      from .configuration_table_transformer import (
          TABLE_TRANSFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP,
          TableTransformerConfig,
          TableTransformerOnnxConfig,
      )

      try:
            if not is_torch_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_table_transformer import (
                TABLE_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
                TableTransformerForObjectDetection,
                TableTransformerModel,
                TableTransformerPreTrainedModel,
            )

else:
      import sys

      __UpperCAmelCase =				_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)





