

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 :      List[Any]							,    snake_case__ :      Dict							,    snake_case__ :      List[Any]=1_3							,    snake_case__ :      Optional[int]=7							,    snake_case__ :      Dict=True							,    snake_case__ :      List[str]=True							,    snake_case__ :      Union[str, Any]=True							,    snake_case__ :      Optional[Any]=True							,    snake_case__ :      Union[str, Any]=9_9							,    snake_case__ :      int=6_4							,    snake_case__ :      Optional[int]=3_2							,    snake_case__ :      Union[str, Any]=5							,    snake_case__ :      str=4							,    snake_case__ :      Dict=3_7							,    snake_case__ :      List[str]="gelu"							,    snake_case__ :      int=0.1							,    snake_case__ :      Optional[int]=0.1							,    snake_case__ :      Dict=5_1_2							,    snake_case__ :      Optional[int]=1_6							,    snake_case__ :      List[Any]=2							,    snake_case__ :      int=0.02							,    snake_case__ :      Any=3							,    snake_case__ :      int=4							,    snake_case__ :      str=None							,    ):

     """simple docstring"""






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

    def UpperCamelCase			(			self :      Optional[Any]    ):

     """simple docstring"""






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

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

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

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

     SCREAMING_SNAKE_CASE								=			self.get_config()

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

    def UpperCamelCase			(			self :      List[str]    ):

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

    def UpperCamelCase			(			self :      int							,    snake_case__ :      List[str]							,    snake_case__ :      Optional[Any]							,    snake_case__ :      Tuple							,    snake_case__ :      Optional[Any]							,    snake_case__ :      Dict							,    snake_case__ :      List[Any]							,    snake_case__ :      Any    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			MegatronBertModel(config=snake_case__    )
     model.to(snake_case__    )
     model.eval()
     SCREAMING_SNAKE_CASE								=			model(snake_case__							,    attention_mask=snake_case__							,    token_type_ids=snake_case__    )
     SCREAMING_SNAKE_CASE								=			model(snake_case__							,    token_type_ids=snake_case__    )
     SCREAMING_SNAKE_CASE								=			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[int]							,    snake_case__ :      Optional[Any]							,    snake_case__ :      List[Any]							,    snake_case__ :      List[str]							,    snake_case__ :      List[Any]							,    snake_case__ :      Union[str, Any]							,    snake_case__ :      int    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			MegatronBertForMaskedLM(config=snake_case__    )
     model.to(snake_case__    )
     model.eval()
     SCREAMING_SNAKE_CASE								=			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__ :      List[Any]							,    snake_case__ :      List[Any]							,    snake_case__ :      Union[str, Any]							,    snake_case__ :      List[Any]							,    snake_case__ :      int							,    snake_case__ :      Optional[int]							,    snake_case__ :      Any    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			MegatronBertForCausalLM(config=snake_case__    )
     model.to(snake_case__    )
     model.eval()
     SCREAMING_SNAKE_CASE								=			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 :      Any							,    snake_case__ :      Optional[int]							,    snake_case__ :      Tuple							,    snake_case__ :      Tuple							,    snake_case__ :      int							,    snake_case__ :      Optional[Any]							,    snake_case__ :      List[Any]							,    snake_case__ :      Tuple    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			MegatronBertForNextSentencePrediction(config=snake_case__    )
     model.to(snake_case__    )
     model.eval()
     SCREAMING_SNAKE_CASE								=			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 :      List[str]							,    snake_case__ :      Optional[Any]							,    snake_case__ :      Any							,    snake_case__ :      Union[str, Any]							,    snake_case__ :      Optional[int]							,    snake_case__ :      str							,    snake_case__ :      int							,    snake_case__ :      Optional[Any]    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			MegatronBertForPreTraining(config=snake_case__    )
     model.to(snake_case__    )
     model.eval()
     SCREAMING_SNAKE_CASE								=			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 :      Any							,    snake_case__ :      Optional[Any]							,    snake_case__ :      List[Any]							,    snake_case__ :      Union[str, Any]							,    snake_case__ :      Union[str, Any]							,    snake_case__ :      List[Any]							,    snake_case__ :      Optional[int]							,    snake_case__ :      Optional[int]    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			MegatronBertForQuestionAnswering(config=snake_case__    )
     model.to(snake_case__    )
     model.eval()
     SCREAMING_SNAKE_CASE								=			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 :      Union[str, Any]							,    snake_case__ :      Union[str, Any]							,    snake_case__ :      List[Any]							,    snake_case__ :      Optional[Any]							,    snake_case__ :      Optional[Any]							,    snake_case__ :      int							,    snake_case__ :      int							,    snake_case__ :      Optional[int]    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			self.num_labels
     SCREAMING_SNAKE_CASE								=			MegatronBertForSequenceClassification(snake_case__    )
     model.to(snake_case__    )
     model.eval()
     SCREAMING_SNAKE_CASE								=			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 :      Any							,    snake_case__ :      List[str]							,    snake_case__ :      Optional[int]							,    snake_case__ :      str							,    snake_case__ :      List[str]							,    snake_case__ :      List[Any]							,    snake_case__ :      List[Any]							,    snake_case__ :      List[str]    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			self.num_labels
     SCREAMING_SNAKE_CASE								=			MegatronBertForTokenClassification(config=snake_case__    )
     model.to(snake_case__    )
     model.eval()
     SCREAMING_SNAKE_CASE								=			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 :      Dict							,    snake_case__ :      Optional[int]							,    snake_case__ :      int							,    snake_case__ :      List[Any]							,    snake_case__ :      Any							,    snake_case__ :      Optional[Any]							,    snake_case__ :      List[Any]							,    snake_case__ :      Dict    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			self.num_choices
     SCREAMING_SNAKE_CASE								=			MegatronBertForMultipleChoice(config=snake_case__    )
     model.to(snake_case__    )
     model.eval()
     SCREAMING_SNAKE_CASE								=			input_ids.unsqueeze(1    ).expand(-1							,    self.num_choices							,    -1    ).contiguous()
     SCREAMING_SNAKE_CASE								=			token_type_ids.unsqueeze(1    ).expand(-1							,    self.num_choices							,    -1    ).contiguous()
     SCREAMING_SNAKE_CASE								=			input_mask.unsqueeze(1    ).expand(-1							,    self.num_choices							,    -1    ).contiguous()
     SCREAMING_SNAKE_CASE								=			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 :      str    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			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
     )						,	
     )								=			config_and_inputs
     SCREAMING_SNAKE_CASE								=			{'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
     return config, inputs_dict




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

    def UpperCamelCase			(			self :      int							,    snake_case__ :      Optional[Any]							,    snake_case__ :      Any							,    snake_case__ :      str=False    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			super()._prepare_for_class(snake_case__							,    snake_case__							,    return_labels=snake_case__    )

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

    def UpperCamelCase			(			self :      List[str]    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			MegatronBertModelTester(self    )
     SCREAMING_SNAKE_CASE								=			ConfigTester(self							,    config_class=snake_case__							,    hidden_size=3_7    )

    def UpperCamelCase			(			self :      Any    ):

     """simple docstring"""






     self.config_tester.run_common_tests()

    def UpperCamelCase			(			self :      str    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_model(*snake_case__    )

    def UpperCamelCase			(			self :      int    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_for_masked_lm(*snake_case__    )

    def UpperCamelCase			(			self :      Any    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*snake_case__    )

    def UpperCamelCase			(			self :      List[str]    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*snake_case__    )

    def UpperCamelCase			(			self :      List[Any]    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_for_pretraining(*snake_case__    )

    def UpperCamelCase			(			self :      int    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_for_question_answering(*snake_case__    )

    def UpperCamelCase			(			self :      List[str]    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*snake_case__    )




    def UpperCamelCase			(			self :      List[Any]    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_for_token_classification(*snake_case__    )


def        __lowerCAmelCase    (       _UpperCamelCase		:  str	)		->							Dict:
 '''simple docstring'''


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


a_    :						Tuple    		=      1e-4




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

    @slow
    @unittest.skip('Model is not available.'    )
    def UpperCamelCase			(			self :      Optional[Any]    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			'nvidia/megatron-bert-uncased-345m'
     if "MYDIR" in os.environ:
      SCREAMING_SNAKE_CASE								=			os.path.join(os.environ['MYDIR']							,    snake_case__    )
     SCREAMING_SNAKE_CASE								=			MegatronBertModel.from_pretrained(snake_case__    )
     model.to(snake_case__    )
     model.half()
     SCREAMING_SNAKE_CASE								=			_long_tensor([[1_0_1, 7_1_1_0, 1_0_0_5, 1_0_5_6, 2_0_2_3, 1_1_3_3_3, 1_7_4_1_3, 1_0_2_9, 1_0_2]]    )
     with torch.no_grad():
      SCREAMING_SNAKE_CASE								=			model(snake_case__    )[0]
     SCREAMING_SNAKE_CASE								=			torch.Size((1, 9, 1_0_2_4)    )
     self.assertEqual(output.shape							,    snake_case__    )

     SCREAMING_SNAKE_CASE								=			[-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    ):
       SCREAMING_SNAKE_CASE								=			output[0, ii, jj]
       SCREAMING_SNAKE_CASE								=			expected[3 * ii + jj]
       SCREAMING_SNAKE_CASE								=			'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__    )




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


a_    :						Any    		=      logging.get_logger(__name__)

a_    :						Dict    		=      {
    "Visual-Attention-Network/van-base": (
        "https://huggingface.co/Visual-Attention-Network/van-base/blob/main/config.json"
    ),
}




class      UpperCamelCase (     SCREAMING_SNAKE_CASE   ):
    __UpperCamelCase         ="van"

    def __init__(			self :      Optional[Any]							,    snake_case__ :      Tuple=2_2_4							,    snake_case__ :      Dict=3							,    snake_case__ :      Union[str, Any]=[7, 3, 3, 3]							,    snake_case__ :      str=[4, 2, 2, 2]							,    snake_case__ :      Optional[Any]=[6_4, 1_2_8, 3_2_0, 5_1_2]							,    snake_case__ :      Optional[Any]=[3, 3, 1_2, 3]							,    snake_case__ :      Tuple=[8, 8, 4, 4]							,    snake_case__ :      Any="gelu"							,    snake_case__ :      Dict=0.02							,    snake_case__ :      List[str]=1E-6							,    snake_case__ :      int=1E-2							,    snake_case__ :      Any=0.0							,    snake_case__ :      Tuple=0.0							,    **snake_case__ :      Any							,    ):

     """simple docstring"""






     super().__init__(**snake_case__    )
     SCREAMING_SNAKE_CASE								=			image_size
     SCREAMING_SNAKE_CASE								=			num_channels
     SCREAMING_SNAKE_CASE								=			patch_sizes
     SCREAMING_SNAKE_CASE								=			strides
     SCREAMING_SNAKE_CASE								=			hidden_sizes
     SCREAMING_SNAKE_CASE								=			depths
     SCREAMING_SNAKE_CASE								=			mlp_ratios
     SCREAMING_SNAKE_CASE								=			hidden_act
     SCREAMING_SNAKE_CASE								=			initializer_range
     SCREAMING_SNAKE_CASE								=			layer_norm_eps
     SCREAMING_SNAKE_CASE								=			layer_scale_init_value
     SCREAMING_SNAKE_CASE								=			drop_path_rate
     SCREAMING_SNAKE_CASE								=			dropout_rate

