
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__		:
      def __init__(      self    ,  __magic_name__    ,  __magic_name__=1_3    ,  __magic_name__=7    ,  __magic_name__=True    ,  __magic_name__=True    ,  __magic_name__=True    ,  __magic_name__=True    ,  __magic_name__=9_9    ,  __magic_name__=6_4    ,  __magic_name__=3_2    ,  __magic_name__=5    ,  __magic_name__=4    ,  __magic_name__=3_7    ,  __magic_name__="gelu"    ,  __magic_name__=0.1    ,  __magic_name__=0.1    ,  __magic_name__=5_1_2    ,  __magic_name__=1_6    ,  __magic_name__=2    ,  __magic_name__=0.02    ,  __magic_name__=3    ,  __magic_name__=4    ,  __magic_name__=None    ,  ):
       lowerCamelCase		:	Dict							=      parent
       lowerCamelCase		:	Any							=      batch_size
       lowerCamelCase		:	Dict							=      seq_length
       lowerCamelCase		:	int							=      is_training
       lowerCamelCase		:	int							=      use_input_mask
       lowerCamelCase		:	str							=      use_token_type_ids
       lowerCamelCase		:	Optional[Any]							=      use_labels
       lowerCamelCase		:	Optional[Any]							=      vocab_size
       lowerCamelCase		:	str							=      hidden_size
       lowerCamelCase		:	Tuple							=      embedding_size
       lowerCamelCase		:	Optional[int]							=      num_hidden_layers
       lowerCamelCase		:	Optional[int]							=      num_attention_heads
       lowerCamelCase		:	Dict							=      intermediate_size
       lowerCamelCase		:	Union[str, Any]							=      hidden_act
       lowerCamelCase		:	int							=      hidden_dropout_prob
       lowerCamelCase		:	Dict							=      attention_probs_dropout_prob
       lowerCamelCase		:	Tuple							=      max_position_embeddings
       lowerCamelCase		:	int							=      type_vocab_size
       lowerCamelCase		:	Optional[int]							=      type_sequence_label_size
       lowerCamelCase		:	Optional[int]							=      initializer_range
       lowerCamelCase		:	Any							=      num_labels
       lowerCamelCase		:	Any							=      num_choices
       lowerCamelCase		:	Optional[Any]							=      scope
      def 						UpperCamelCase__							(      self      ):
       lowerCamelCase		:	str							=      ids_tensor([self.batch_size, self.seq_length]    ,  self.vocab_size      )

       lowerCamelCase		:	List[Any]							=      None
       if self.use_input_mask:
        lowerCamelCase		:	int							=      random_attention_mask([self.batch_size, self.seq_length]      )

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

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

       lowerCamelCase		:	str							=      self.get_config()

       return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
      def 						UpperCamelCase__							(      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=__magic_name__    ,  initializer_range=self.initializer_range    ,  )
      def 						UpperCamelCase__							(      self    ,  __magic_name__    ,  __magic_name__    ,  __magic_name__    ,  __magic_name__    ,  __magic_name__    ,  __magic_name__    ,  __magic_name__      ):
       lowerCamelCase		:	Optional[int]							=      MegatronBertModel(config=__magic_name__      )
       model.to(__magic_name__      )
       model.eval()
       lowerCamelCase		:	Optional[int]							=      model(__magic_name__    ,  attention_mask=__magic_name__    ,  token_type_ids=__magic_name__      )
       lowerCamelCase		:	Any							=      model(__magic_name__    ,  token_type_ids=__magic_name__      )
       lowerCamelCase		:	Optional[int]							=      model(__magic_name__      )

       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    ,  __magic_name__    ,  __magic_name__    ,  __magic_name__    ,  __magic_name__    ,  __magic_name__    ,  __magic_name__    ,  __magic_name__      ):
       lowerCamelCase		:	List[str]							=      MegatronBertForMaskedLM(config=__magic_name__      )
       model.to(__magic_name__      )
       model.eval()
       lowerCamelCase		:	List[Any]							=      model(__magic_name__    ,  attention_mask=__magic_name__    ,  token_type_ids=__magic_name__    ,  labels=__magic_name__      )
       self.parent.assertEqual(result.logits.shape    ,  (self.batch_size, self.seq_length, self.vocab_size)      )
      def 						UpperCamelCase__							(      self    ,  __magic_name__    ,  __magic_name__    ,  __magic_name__    ,  __magic_name__    ,  __magic_name__    ,  __magic_name__    ,  __magic_name__      ):
       lowerCamelCase		:	Dict							=      MegatronBertForCausalLM(config=__magic_name__      )
       model.to(__magic_name__      )
       model.eval()
       lowerCamelCase		:	List[str]							=      model(__magic_name__    ,  attention_mask=__magic_name__    ,  token_type_ids=__magic_name__    ,  labels=__magic_name__      )
       self.parent.assertEqual(result.logits.shape    ,  (self.batch_size, self.seq_length, self.vocab_size)      )
      def 						UpperCamelCase__							(      self    ,  __magic_name__    ,  __magic_name__    ,  __magic_name__    ,  __magic_name__    ,  __magic_name__    ,  __magic_name__    ,  __magic_name__      ):
       lowerCamelCase		:	str							=      MegatronBertForNextSentencePrediction(config=__magic_name__      )
       model.to(__magic_name__      )
       model.eval()
       lowerCamelCase		:	Dict							=      model(
           __magic_name__    ,  attention_mask=__magic_name__    ,  token_type_ids=__magic_name__    ,  labels=__magic_name__    ,  )
       self.parent.assertEqual(result.logits.shape    ,  (self.batch_size, 2)      )
      def 						UpperCamelCase__							(      self    ,  __magic_name__    ,  __magic_name__    ,  __magic_name__    ,  __magic_name__    ,  __magic_name__    ,  __magic_name__    ,  __magic_name__      ):
       lowerCamelCase		:	Union[str, Any]							=      MegatronBertForPreTraining(config=__magic_name__      )
       model.to(__magic_name__      )
       model.eval()
       lowerCamelCase		:	Optional[int]							=      model(
           __magic_name__    ,  attention_mask=__magic_name__    ,  token_type_ids=__magic_name__    ,  labels=__magic_name__    ,  next_sentence_label=__magic_name__    ,  )
       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    ,  __magic_name__    ,  __magic_name__    ,  __magic_name__    ,  __magic_name__    ,  __magic_name__    ,  __magic_name__    ,  __magic_name__      ):
       lowerCamelCase		:	List[str]							=      MegatronBertForQuestionAnswering(config=__magic_name__      )
       model.to(__magic_name__      )
       model.eval()
       lowerCamelCase		:	Optional[int]							=      model(
           __magic_name__    ,  attention_mask=__magic_name__    ,  token_type_ids=__magic_name__    ,  start_positions=__magic_name__    ,  end_positions=__magic_name__    ,  )
       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    ,  __magic_name__    ,  __magic_name__    ,  __magic_name__    ,  __magic_name__    ,  __magic_name__    ,  __magic_name__    ,  __magic_name__      ):
       lowerCamelCase		:	List[Any]							=      self.num_labels
       lowerCamelCase		:	Union[str, Any]							=      MegatronBertForSequenceClassification(__magic_name__      )
       model.to(__magic_name__      )
       model.eval()
       lowerCamelCase		:	List[Any]							=      model(__magic_name__    ,  attention_mask=__magic_name__    ,  token_type_ids=__magic_name__    ,  labels=__magic_name__      )
       self.parent.assertEqual(result.logits.shape    ,  (self.batch_size, self.num_labels)      )
      def 						UpperCamelCase__							(      self    ,  __magic_name__    ,  __magic_name__    ,  __magic_name__    ,  __magic_name__    ,  __magic_name__    ,  __magic_name__    ,  __magic_name__      ):
       lowerCamelCase		:	Any							=      self.num_labels
       lowerCamelCase		:	int							=      MegatronBertForTokenClassification(config=__magic_name__      )
       model.to(__magic_name__      )
       model.eval()
       lowerCamelCase		:	Any							=      model(__magic_name__    ,  attention_mask=__magic_name__    ,  token_type_ids=__magic_name__    ,  labels=__magic_name__      )
       self.parent.assertEqual(result.logits.shape    ,  (self.batch_size, self.seq_length, self.num_labels)      )
      def 						UpperCamelCase__							(      self    ,  __magic_name__    ,  __magic_name__    ,  __magic_name__    ,  __magic_name__    ,  __magic_name__    ,  __magic_name__    ,  __magic_name__      ):
       lowerCamelCase		:	Union[str, Any]							=      self.num_choices
       lowerCamelCase		:	Optional[int]							=      MegatronBertForMultipleChoice(config=__magic_name__      )
       model.to(__magic_name__      )
       model.eval()
       lowerCamelCase		:	Tuple							=      input_ids.unsqueeze(1      ).expand(-1    ,  self.num_choices    ,  -1      ).contiguous()
       lowerCamelCase		:	List[str]							=      token_type_ids.unsqueeze(1      ).expand(-1    ,  self.num_choices    ,  -1      ).contiguous()
       lowerCamelCase		:	Any							=      input_mask.unsqueeze(1      ).expand(-1    ,  self.num_choices    ,  -1      ).contiguous()
       lowerCamelCase		:	Dict							=      model(
           __magic_name__    ,  attention_mask=__magic_name__    ,  token_type_ids=__magic_name__    ,  labels=__magic_name__    ,  )
       self.parent.assertEqual(result.logits.shape    ,  (self.batch_size, self.num_choices)      )






      def 						UpperCamelCase__							(      self      ):
       lowerCamelCase		:	Tuple							=      self.prepare_config_and_inputs()
       (
           lowerCamelCase
       )		:	Dict							=      config_and_inputs
       lowerCamelCase		:	Union[str, Any]							=      {'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
       return config, inputs_dict





@require_torch
class 							A__		(	__lowerCamelCase  ,   __lowerCamelCase  ,   unittest.TestCase):
      _UpperCAmelCase				:    Optional[int]           =			(
          (
              MegatronBertModel,
              MegatronBertForMaskedLM,
              MegatronBertForCausalLM,
              MegatronBertForMultipleChoice,
              MegatronBertForNextSentencePrediction,
              MegatronBertForPreTraining,
              MegatronBertForQuestionAnswering,
              MegatronBertForSequenceClassification,
              MegatronBertForTokenClassification,
          )
          if is_torch_available()
          else ()
      )
      _UpperCAmelCase				:    Optional[Any]           =			(
          {
              "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				:    Any           =			True
      # test_resize_embeddings = False
      _UpperCAmelCase				:    Union[str, Any]           =			False
      def 						UpperCamelCase__							(      self    ,  __magic_name__    ,  __magic_name__    ,  __magic_name__=False      ):
       lowerCamelCase		:	List[Any]							=      super()._prepare_for_class(__magic_name__    ,  __magic_name__    ,  return_labels=__magic_name__      )

       if return_labels:
        if model_class in get_values(__magic_name__      ):
         lowerCamelCase		:	List[str]							=      torch.zeros(
             (self.model_tester.batch_size, self.model_tester.seq_length)    ,  dtype=torch.long    ,  device=__magic_name__      )
         lowerCamelCase		:	List[str]							=      torch.zeros(
             self.model_tester.batch_size    ,  dtype=torch.long    ,  device=__magic_name__      )
       return inputs_dict
      def 						UpperCamelCase__							(      self      ):
       lowerCamelCase		:	List[str]							=      MegatronBertModelTester(self      )
       lowerCamelCase		:	List[Any]							=      ConfigTester(self    ,  config_class=__magic_name__    ,  hidden_size=3_7      )
      def 						UpperCamelCase__							(      self      ):
       self.config_tester.run_common_tests()
      def 						UpperCamelCase__							(      self      ):
       lowerCamelCase		:	Union[str, Any]							=      self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_model(*__magic_name__      )
      def 						UpperCamelCase__							(      self      ):
       lowerCamelCase		:	List[Any]							=      self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_masked_lm(*__magic_name__      )
      def 						UpperCamelCase__							(      self      ):
       lowerCamelCase		:	Tuple							=      self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*__magic_name__      )
      def 						UpperCamelCase__							(      self      ):
       lowerCamelCase		:	Any							=      self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*__magic_name__      )
      def 						UpperCamelCase__							(      self      ):
       lowerCamelCase		:	List[str]							=      self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_pretraining(*__magic_name__      )
      def 						UpperCamelCase__							(      self      ):
       lowerCamelCase		:	Dict							=      self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_question_answering(*__magic_name__      )
      def 						UpperCamelCase__							(      self      ):
       lowerCamelCase		:	int							=      self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*__magic_name__      )






      def 						UpperCamelCase__							(      self      ):
       lowerCamelCase		:	Optional[Any]							=      self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_token_classification(*__magic_name__      )





def     _a					(						lowerCamelCase							):
 return torch.tensor(
     __UpperCAmelCase,	dtype=torch.long,	device=__UpperCAmelCase,	)


_lowerCamelCase					   =1E-4





@require_torch
@require_sentencepiece
@require_tokenizers
class 							A__		(	unittest.TestCase):
      @slow
      @unittest.skip("""Model is not available."""      )
      def 						UpperCamelCase__							(      self      ):
       lowerCamelCase		:	Any							=      '''nvidia/megatron-bert-uncased-345m'''
       if "MYDIR" in os.environ:
        lowerCamelCase		:	List[str]							=      os.path.join(os.environ["""MYDIR"""]    ,  __magic_name__      )
       lowerCamelCase		:	List[str]							=      MegatronBertModel.from_pretrained(__magic_name__      )
       model.to(__magic_name__      )
       model.half()
       lowerCamelCase		:	str							=      _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():
        lowerCamelCase		:	int							=      model(__magic_name__      )[0]
       lowerCamelCase		:	int							=      torch.Size((1, 9, 1_0_2_4)      )
       self.assertEqual(output.shape    ,  __magic_name__      )

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











from collections import OrderedDict
from typing import Mapping

from packaging import version

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


UpperCAmelCase_     							= logging.get_logger(__name__)

UpperCAmelCase_     							= {
    'google/mobilenet_v2_1.4_224': 'https://huggingface.co/google/mobilenet_v2_1.4_224/resolve/main/config.json',
    'google/mobilenet_v2_1.0_224': 'https://huggingface.co/google/mobilenet_v2_1.0_224/resolve/main/config.json',
    'google/mobilenet_v2_0.75_160': 'https://huggingface.co/google/mobilenet_v2_0.75_160/resolve/main/config.json',
    'google/mobilenet_v2_0.35_96': 'https://huggingface.co/google/mobilenet_v2_0.35_96/resolve/main/config.json',
    # See all MobileNetV2 models at https://huggingface.co/models?filter=mobilenet_v2
}






class     lowercase__       (	__lowerCamelCase				):

   '''simple docstring'''


   a						:							List[str]   		=		"mobilenet_v2"




   def __init__(    self,		__magic_name__=3,		__magic_name__=224,		__magic_name__=1.0,		__magic_name__=8,		__magic_name__=8,		__magic_name__=6,		__magic_name__=32,		__magic_name__=True,		__magic_name__=True,		__magic_name__="relu6",		__magic_name__=True,		__magic_name__=0.8,		__magic_name__=0.02,		__magic_name__=0.001,		__magic_name__=255,		**__magic_name__,		) ->      List[Any]:






       """simple docstring"""
       super().__init__(**__magic_name__							)

       if depth_multiplier <= 0:
           raise ValueError('''depth_multiplier must be greater than zero.'''							)

       UpperCamelCase__							:				Union[str, Any]          =      num_channels
       UpperCamelCase__							:				int          =      image_size
       UpperCamelCase__							:				int          =      depth_multiplier
       UpperCamelCase__							:				Tuple          =      depth_divisible_by
       UpperCamelCase__							:				List[str]          =      min_depth
       UpperCamelCase__							:				Optional[int]          =      expand_ratio
       UpperCamelCase__							:				Optional[int]          =      output_stride
       UpperCamelCase__							:				Tuple          =      first_layer_is_expansion
       UpperCamelCase__							:				Union[str, Any]          =      finegrained_output
       UpperCamelCase__							:				str          =      hidden_act
       UpperCamelCase__							:				Optional[Any]          =      tf_padding
       UpperCamelCase__							:				Optional[int]          =      classifier_dropout_prob
       UpperCamelCase__							:				int          =      initializer_range
       UpperCamelCase__							:				Union[str, Any]          =      layer_norm_eps
       UpperCamelCase__							:				Tuple          =      semantic_loss_ignore_index






class     lowercase__       (	__lowerCamelCase				):

   '''simple docstring'''


   a						:							Union[str, Any]   		=		version.parse("1.11"				)




   @property
   def       UpperCamelCase__  (    self							) ->      Mapping[str, Mapping[int, str]]:






       """simple docstring"""
       return OrderedDict([('''pixel_values''', {0: '''batch'''})]							)




   @property
   def       UpperCamelCase__  (    self							) ->      Mapping[str, Mapping[int, str]]:






       """simple docstring"""
       if self.task == "image-classification":
           return OrderedDict([('''logits''', {0: '''batch'''})]							)
       else:
           return OrderedDict([('''last_hidden_state''', {0: '''batch'''}), ('''pooler_output''', {0: '''batch'''})]							)




   @property
   def       UpperCamelCase__  (    self							) ->      float:






       """simple docstring"""
       return 1E-4

