





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





   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=512					,			snake_case=16					,			snake_case=2					,			snake_case=0.02					,			snake_case=3					,			snake_case=4					,			snake_case=None					,			)				->  List[str]:
      """simple docstring"""





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





   def 	_snake_case ( self      )				->  Optional[int]:
      """simple docstring"""





      a__	:						Optional[Any]      		=  ids_tensor([self.batch_size, self.seq_length]					,			self.vocab_size      )

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

      a__	:						str      		=  None
      if self.use_token_type_ids:
         a__	:						Optional[int]      		=  ids_tensor([self.batch_size, self.seq_length]					,			self.type_vocab_size      )

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

      a__	:						int      		=  self.get_config()

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





   def 	_snake_case ( self      )				->  Dict:
      """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=lowerCamelCase__					,			initializer_range=self.initializer_range					,			)





   def 	_snake_case ( self					,			snake_case					,			snake_case					,			snake_case					,			snake_case					,			snake_case					,			snake_case					,			snake_case      )				->  Dict:
      """simple docstring"""





      a__	:						int      		=  MegatronBertModel(config=lowerCamelCase__      )
      model.to(lowerCamelCase__      )
      model.eval()
      a__	:						List[str]      		=  model(lowerCamelCase__					,			attention_mask=lowerCamelCase__					,			token_type_ids=lowerCamelCase__      )
      a__	:						Optional[int]      		=  model(lowerCamelCase__					,			token_type_ids=lowerCamelCase__      )
      a__	:						int      		=  model(lowerCamelCase__      )

      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 	_snake_case ( self					,			snake_case					,			snake_case					,			snake_case					,			snake_case					,			snake_case					,			snake_case					,			snake_case      )				->  Union[str, Any]:
      """simple docstring"""





      a__	:						List[str]      		=  MegatronBertForMaskedLM(config=lowerCamelCase__      )
      model.to(lowerCamelCase__      )
      model.eval()
      a__	:						int      		=  model(lowerCamelCase__					,			attention_mask=lowerCamelCase__					,			token_type_ids=lowerCamelCase__					,			labels=lowerCamelCase__      )
      self.parent.assertEqual(result.logits.shape					,			(self.batch_size, self.seq_length, self.vocab_size)      )





   def 	_snake_case ( self					,			snake_case					,			snake_case					,			snake_case					,			snake_case					,			snake_case					,			snake_case					,			snake_case      )				->  str:
      """simple docstring"""





      a__	:						Any      		=  MegatronBertForCausalLM(config=lowerCamelCase__      )
      model.to(lowerCamelCase__      )
      model.eval()
      a__	:						Optional[int]      		=  model(lowerCamelCase__					,			attention_mask=lowerCamelCase__					,			token_type_ids=lowerCamelCase__					,			labels=lowerCamelCase__      )
      self.parent.assertEqual(result.logits.shape					,			(self.batch_size, self.seq_length, self.vocab_size)      )





   def 	_snake_case ( self					,			snake_case					,			snake_case					,			snake_case					,			snake_case					,			snake_case					,			snake_case					,			snake_case      )				->  Tuple:
      """simple docstring"""





      a__	:						str      		=  MegatronBertForNextSentencePrediction(config=lowerCamelCase__      )
      model.to(lowerCamelCase__      )
      model.eval()
      a__	:						Tuple      		=  model(
          lowerCamelCase__					,			attention_mask=lowerCamelCase__					,			token_type_ids=lowerCamelCase__					,			labels=lowerCamelCase__					,			)
      self.parent.assertEqual(result.logits.shape					,			(self.batch_size, 2)      )





   def 	_snake_case ( self					,			snake_case					,			snake_case					,			snake_case					,			snake_case					,			snake_case					,			snake_case					,			snake_case      )				->  Tuple:
      """simple docstring"""





      a__	:						Union[str, Any]      		=  MegatronBertForPreTraining(config=lowerCamelCase__      )
      model.to(lowerCamelCase__      )
      model.eval()
      a__	:						Optional[int]      		=  model(
          lowerCamelCase__					,			attention_mask=lowerCamelCase__					,			token_type_ids=lowerCamelCase__					,			labels=lowerCamelCase__					,			next_sentence_label=lowerCamelCase__					,			)
      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 	_snake_case ( self					,			snake_case					,			snake_case					,			snake_case					,			snake_case					,			snake_case					,			snake_case					,			snake_case      )				->  Optional[int]:
      """simple docstring"""





      a__	:						Dict      		=  MegatronBertForQuestionAnswering(config=lowerCamelCase__      )
      model.to(lowerCamelCase__      )
      model.eval()
      a__	:						int      		=  model(
          lowerCamelCase__					,			attention_mask=lowerCamelCase__					,			token_type_ids=lowerCamelCase__					,			start_positions=lowerCamelCase__					,			end_positions=lowerCamelCase__					,			)
      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 	_snake_case ( self					,			snake_case					,			snake_case					,			snake_case					,			snake_case					,			snake_case					,			snake_case					,			snake_case      )				->  Tuple:
      """simple docstring"""





      a__	:						Dict      		=  self.num_labels
      a__	:						Any      		=  MegatronBertForSequenceClassification(lowerCamelCase__      )
      model.to(lowerCamelCase__      )
      model.eval()
      a__	:						Any      		=  model(lowerCamelCase__					,			attention_mask=lowerCamelCase__					,			token_type_ids=lowerCamelCase__					,			labels=lowerCamelCase__      )
      self.parent.assertEqual(result.logits.shape					,			(self.batch_size, self.num_labels)      )





   def 	_snake_case ( self					,			snake_case					,			snake_case					,			snake_case					,			snake_case					,			snake_case					,			snake_case					,			snake_case      )				->  int:
      """simple docstring"""





      a__	:						List[str]      		=  self.num_labels
      a__	:						str      		=  MegatronBertForTokenClassification(config=lowerCamelCase__      )
      model.to(lowerCamelCase__      )
      model.eval()
      a__	:						int      		=  model(lowerCamelCase__					,			attention_mask=lowerCamelCase__					,			token_type_ids=lowerCamelCase__					,			labels=lowerCamelCase__      )
      self.parent.assertEqual(result.logits.shape					,			(self.batch_size, self.seq_length, self.num_labels)      )





   def 	_snake_case ( self					,			snake_case					,			snake_case					,			snake_case					,			snake_case					,			snake_case					,			snake_case					,			snake_case      )				->  Optional[Any]:
      """simple docstring"""





      a__	:						Optional[Any]      		=  self.num_choices
      a__	:						Tuple      		=  MegatronBertForMultipleChoice(config=lowerCamelCase__      )
      model.to(lowerCamelCase__      )
      model.eval()
      a__	:						Union[str, Any]      		=  input_ids.unsqueeze(1      ).expand(-1					,			self.num_choices					,			-1      ).contiguous()
      a__	:						List[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__	:						Any      		=  model(
          lowerCamelCase__					,			attention_mask=lowerCamelCase__					,			token_type_ids=lowerCamelCase__					,			labels=lowerCamelCase__					,			)
      self.parent.assertEqual(result.logits.shape					,			(self.batch_size, self.num_choices)      )





   def 	_snake_case ( self      )				->  Optional[Any]:
      """simple docstring"""





      a__	:						str      		=  self.prepare_config_and_inputs()
      (
          a__				
      )	:						Any      		=  config_and_inputs
      a__	:						Tuple      		=  {'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
      return config, inputs_dict







@require_torch
class   __lowerCAmelCase			(					__lowerCAmelCase		,__lowerCAmelCase		,unittest.TestCase  ):
   _UpperCamelCase :						List[Any]     		=				(
       (
           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 :						Tuple     		=				True
   # test_resize_embeddings = False
   _UpperCamelCase :						List[str]     		=				False





   def 	_snake_case ( self					,			snake_case					,			snake_case					,			snake_case=False      )				->  List[Any]:
      """simple docstring"""





      a__	:						List[Any]      		=  super()._prepare_for_class(lowerCamelCase__					,			lowerCamelCase__					,			return_labels=lowerCamelCase__      )

      if return_labels:
         if model_class in get_values(lowerCamelCase__      ):
            a__	:						Dict      		=  torch.zeros(
                (self.model_tester.batch_size, self.model_tester.seq_length)					,			dtype=torch.long					,			device=lowerCamelCase__      )
            a__	:						Optional[Any]      		=  torch.zeros(
                self.model_tester.batch_size					,			dtype=torch.long					,			device=lowerCamelCase__      )
      return inputs_dict





   def 	_snake_case ( self      )				->  Dict:
      """simple docstring"""





      a__	:						int      		=  MegatronBertModelTester(self      )
      a__	:						Union[str, Any]      		=  ConfigTester(self					,			config_class=lowerCamelCase__					,			hidden_size=37      )





   def 	_snake_case ( self      )				->  Union[str, Any]:
      """simple docstring"""





      self.config_tester.run_common_tests()





   def 	_snake_case ( self      )				->  Tuple:
      """simple docstring"""





      a__	:						Union[str, Any]      		=  self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_megatron_bert_model(*lowerCamelCase__      )





   def 	_snake_case ( self      )				->  Optional[Any]:
      """simple docstring"""





      a__	:						List[Any]      		=  self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_megatron_bert_for_masked_lm(*lowerCamelCase__      )





   def 	_snake_case ( self      )				->  List[Any]:
      """simple docstring"""





      a__	:						str      		=  self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*lowerCamelCase__      )





   def 	_snake_case ( self      )				->  Union[str, Any]:
      """simple docstring"""





      a__	:						Union[str, Any]      		=  self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*lowerCamelCase__      )





   def 	_snake_case ( self      )				->  Tuple:
      """simple docstring"""





      a__	:						Optional[Any]      		=  self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_megatron_bert_for_pretraining(*lowerCamelCase__      )





   def 	_snake_case ( self      )				->  Optional[Any]:
      """simple docstring"""





      a__	:						List[str]      		=  self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_megatron_bert_for_question_answering(*lowerCamelCase__      )





   def 	_snake_case ( self      )				->  Union[str, Any]:
      """simple docstring"""





      a__	:						List[str]      		=  self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*lowerCamelCase__      )





   def 	_snake_case ( self      )				->  Optional[int]:
      """simple docstring"""





      a__	:						Any      		=  self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_megatron_bert_for_token_classification(*lowerCamelCase__      )







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


SCREAMING_SNAKE_CASE__       :			Dict			      =						1e-4







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





   @slow
   @unittest.skip("Model is not available."      )
   def 	_snake_case ( self      )				->  Dict:
      """simple docstring"""





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

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










import json
from typing import List, Optional, Tuple

from tokenizers import normalizers

from ...tokenization_utils_fast import PreTrainedTokenizerFast
from .tokenization_lxmert import LxmertTokenizer


SCREAMING_SNAKE_CASE__       :			Optional[Any]			      =						{"""vocab_file""": """vocab.txt""", """tokenizer_file""": """tokenizer.json"""}

SCREAMING_SNAKE_CASE__       :			List[Any]			      =						{
    """vocab_file""": {
        """unc-nlp/lxmert-base-uncased""": """https://huggingface.co/unc-nlp/lxmert-base-uncased/resolve/main/vocab.txt""",
    },
    """tokenizer_file""": {
        """unc-nlp/lxmert-base-uncased""": (
            """https://huggingface.co/unc-nlp/lxmert-base-uncased/resolve/main/tokenizer.json"""
        ),
    },
}

SCREAMING_SNAKE_CASE__       :			Any			      =						{
    """unc-nlp/lxmert-base-uncased""": 5_1_2,
}

SCREAMING_SNAKE_CASE__       :			Optional[int]			      =						{
    """unc-nlp/lxmert-base-uncased""": {"""do_lower_case""": True},
}







class   __lowerCAmelCase			(					_UpperCamelCase  ):
   _UpperCamelCase :						List[str]     		=				VOCAB_FILES_NAMES
   _UpperCamelCase :						str     		=				PRETRAINED_VOCAB_FILES_MAP
   _UpperCamelCase :						List[str]     		=				PRETRAINED_INIT_CONFIGURATION
   _UpperCamelCase :						Optional[int]     		=				PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
   _UpperCamelCase :						Optional[Any]     		=				LxmertTokenizer





   def __init__( self					,			snake_case=None					,			snake_case=None					,			snake_case=True					,			snake_case="[UNK]"					,			snake_case="[SEP]"					,			snake_case="[PAD]"					,			snake_case="[CLS]"					,			snake_case="[MASK]"					,			snake_case=True					,			snake_case=None					,			**snake_case					,			)				->  Any:
      """simple docstring"""





      super().__init__(
          snake_case					,			tokenizer_file=snake_case					,			do_lower_case=snake_case					,			unk_token=snake_case					,			sep_token=snake_case					,			pad_token=snake_case					,			cls_token=snake_case					,			mask_token=snake_case					,			tokenize_chinese_chars=snake_case					,			strip_accents=snake_case					,			**snake_case					,			)

      a__	:						Optional[Any]      		=  json.loads(self.backend_tokenizer.normalizer.__getstate__()      )
      if (
          normalizer_state.get("lowercase"					,			snake_case      ) != do_lower_case
          or normalizer_state.get("strip_accents"					,			snake_case      ) != strip_accents
          or normalizer_state.get("handle_chinese_chars"					,			snake_case      ) != tokenize_chinese_chars
      ):
         a__	:						str      		=  getattr(snake_case					,			normalizer_state.pop("type"      )      )
         a__	:						Tuple      		=  do_lower_case
         a__	:						Union[str, Any]      		=  strip_accents
         a__	:						str      		=  tokenize_chinese_chars
         a__	:						List[str]      		=  normalizer_class(**snake_case      )

      a__	:						str      		=  do_lower_case





   def 	_snake_case ( self					,			snake_case					,			snake_case=None      )				->  List[str]:
      """simple docstring"""





      a__	:						Optional[int]      		=  [self.cls_token_id] + token_ids_a + [self.sep_token_id]

      if token_ids_a:
         output += token_ids_a + [self.sep_token_id]

      return output





   def 	_snake_case ( self					,			snake_case					,			snake_case = None      )				->  List[int]:
      """simple docstring"""





      a__	:						Tuple      		=  [self.sep_token_id]
      a__	:						str      		=  [self.cls_token_id]
      if token_ids_a is None:
         return len(cls + token_ids_a + sep      ) * [0]
      return len(cls + token_ids_a + sep      ) * [0] + len(token_ids_a + sep      ) * [1]





   def 	_snake_case ( self					,			snake_case					,			snake_case = None      )				->  Tuple[str]:
      """simple docstring"""





      a__	:						Optional[int]      		=  self._tokenizer.model.save(snake_case					,			name=snake_case      )
      return tuple(snake_case      )

