



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	,				__UpperCAmelCase	,				__UpperCAmelCase=13	,				__UpperCAmelCase=7	,				__UpperCAmelCase=True	,				__UpperCAmelCase=True	,				__UpperCAmelCase=True	,				__UpperCAmelCase=True	,				__UpperCAmelCase=99	,				__UpperCAmelCase=64	,				__UpperCAmelCase=32	,				__UpperCAmelCase=5	,				__UpperCAmelCase=4	,				__UpperCAmelCase=37	,				__UpperCAmelCase="gelu"	,				__UpperCAmelCase=0.1	,				__UpperCAmelCase=0.1	,				__UpperCAmelCase=512	,				__UpperCAmelCase=16	,				__UpperCAmelCase=2	,				__UpperCAmelCase=0.02	,				__UpperCAmelCase=3	,				__UpperCAmelCase=4	,				__UpperCAmelCase=None	,				):



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




 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    __lowerCamelCase													=    ids_tensor([self.batch_size, self.seq_length]	,				self.vocab_size				)

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

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

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

    __lowerCamelCase													=    self.get_config()

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




 def        lowerCamelCase      (						self				):



    '''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        lowerCamelCase      (						self	,				__UpperCAmelCase	,				__UpperCAmelCase	,				__UpperCAmelCase	,				__UpperCAmelCase	,				__UpperCAmelCase	,				__UpperCAmelCase	,				__UpperCAmelCase				):



    '''simple docstring'''
    __lowerCamelCase													=    MegatronBertModel(config=__UpperCAmelCase				)
    model.to(__UpperCAmelCase				)
    model.eval()
    __lowerCamelCase													=    model(__UpperCAmelCase	,				attention_mask=__UpperCAmelCase	,				token_type_ids=__UpperCAmelCase				)
    __lowerCamelCase													=    model(__UpperCAmelCase	,				token_type_ids=__UpperCAmelCase				)
    __lowerCamelCase													=    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        lowerCamelCase      (						self	,				__UpperCAmelCase	,				__UpperCAmelCase	,				__UpperCAmelCase	,				__UpperCAmelCase	,				__UpperCAmelCase	,				__UpperCAmelCase	,				__UpperCAmelCase				):



    '''simple docstring'''
    __lowerCamelCase													=    MegatronBertForMaskedLM(config=__UpperCAmelCase				)
    model.to(__UpperCAmelCase				)
    model.eval()
    __lowerCamelCase													=    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        lowerCamelCase      (						self	,				__UpperCAmelCase	,				__UpperCAmelCase	,				__UpperCAmelCase	,				__UpperCAmelCase	,				__UpperCAmelCase	,				__UpperCAmelCase	,				__UpperCAmelCase				):



    '''simple docstring'''
    __lowerCamelCase													=    MegatronBertForCausalLM(config=__UpperCAmelCase				)
    model.to(__UpperCAmelCase				)
    model.eval()
    __lowerCamelCase													=    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        lowerCamelCase      (						self	,				__UpperCAmelCase	,				__UpperCAmelCase	,				__UpperCAmelCase	,				__UpperCAmelCase	,				__UpperCAmelCase	,				__UpperCAmelCase	,				__UpperCAmelCase				):



    '''simple docstring'''
    __lowerCamelCase													=    MegatronBertForNextSentencePrediction(config=__UpperCAmelCase				)
    model.to(__UpperCAmelCase				)
    model.eval()
    __lowerCamelCase													=    model(
        __UpperCAmelCase	,				attention_mask=__UpperCAmelCase	,				token_type_ids=__UpperCAmelCase	,				labels=__UpperCAmelCase	,				)
    self.parent.assertEqual(result.logits.shape	,				(self.batch_size, 2)				)




 def        lowerCamelCase      (						self	,				__UpperCAmelCase	,				__UpperCAmelCase	,				__UpperCAmelCase	,				__UpperCAmelCase	,				__UpperCAmelCase	,				__UpperCAmelCase	,				__UpperCAmelCase				):



    '''simple docstring'''
    __lowerCamelCase													=    MegatronBertForPreTraining(config=__UpperCAmelCase				)
    model.to(__UpperCAmelCase				)
    model.eval()
    __lowerCamelCase													=    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        lowerCamelCase      (						self	,				__UpperCAmelCase	,				__UpperCAmelCase	,				__UpperCAmelCase	,				__UpperCAmelCase	,				__UpperCAmelCase	,				__UpperCAmelCase	,				__UpperCAmelCase				):



    '''simple docstring'''
    __lowerCamelCase													=    MegatronBertForQuestionAnswering(config=__UpperCAmelCase				)
    model.to(__UpperCAmelCase				)
    model.eval()
    __lowerCamelCase													=    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        lowerCamelCase      (						self	,				__UpperCAmelCase	,				__UpperCAmelCase	,				__UpperCAmelCase	,				__UpperCAmelCase	,				__UpperCAmelCase	,				__UpperCAmelCase	,				__UpperCAmelCase				):



    '''simple docstring'''
    __lowerCamelCase													=    self.num_labels
    __lowerCamelCase													=    MegatronBertForSequenceClassification(__UpperCAmelCase				)
    model.to(__UpperCAmelCase				)
    model.eval()
    __lowerCamelCase													=    model(__UpperCAmelCase	,				attention_mask=__UpperCAmelCase	,				token_type_ids=__UpperCAmelCase	,				labels=__UpperCAmelCase				)
    self.parent.assertEqual(result.logits.shape	,				(self.batch_size, self.num_labels)				)




 def        lowerCamelCase      (						self	,				__UpperCAmelCase	,				__UpperCAmelCase	,				__UpperCAmelCase	,				__UpperCAmelCase	,				__UpperCAmelCase	,				__UpperCAmelCase	,				__UpperCAmelCase				):



    '''simple docstring'''
    __lowerCamelCase													=    self.num_labels
    __lowerCamelCase													=    MegatronBertForTokenClassification(config=__UpperCAmelCase				)
    model.to(__UpperCAmelCase				)
    model.eval()
    __lowerCamelCase													=    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        lowerCamelCase      (						self	,				__UpperCAmelCase	,				__UpperCAmelCase	,				__UpperCAmelCase	,				__UpperCAmelCase	,				__UpperCAmelCase	,				__UpperCAmelCase	,				__UpperCAmelCase				):



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




 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    __lowerCamelCase													=    self.prepare_config_and_inputs()
    (
        (
        __lowerCamelCase
    )							,(
        __lowerCamelCase
    )							,(
        __lowerCamelCase
    )							,(
        __lowerCamelCase
    )							,(
        __lowerCamelCase
    )							,(
        __lowerCamelCase
    )							,(
        __lowerCamelCase
    )							,
    )													=    config_and_inputs
    __lowerCamelCase													=    {'''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   ):
 lowerCAmelCase__       					=  (
     (
         MegatronBertModel,
         MegatronBertForMaskedLM,
         MegatronBertForCausalLM,
         MegatronBertForMultipleChoice,
         MegatronBertForNextSentencePrediction,
         MegatronBertForPreTraining,
         MegatronBertForQuestionAnswering,
         MegatronBertForSequenceClassification,
         MegatronBertForTokenClassification,
     )
     if is_torch_available()
     else ()
 )
 lowerCAmelCase__       					=  (
     {
         """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 {}
 )
 lowerCAmelCase__       					=  True
 # test_resize_embeddings = False
 lowerCAmelCase__       					=  False




 def        lowerCamelCase      (						self	,				__UpperCAmelCase	,				__UpperCAmelCase	,				__UpperCAmelCase=False				):



    '''simple docstring'''
    __lowerCamelCase													=    super()._prepare_for_class(__UpperCAmelCase	,				__UpperCAmelCase	,				return_labels=__UpperCAmelCase				)

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




 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    __lowerCamelCase													=    MegatronBertModelTester(self				)
    __lowerCamelCase													=    ConfigTester(self	,				config_class=__UpperCAmelCase	,				hidden_size=37				)




 def        lowerCamelCase      (						self				):



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




 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    __lowerCamelCase													=    self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_megatron_bert_model(*__UpperCAmelCase				)




 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    __lowerCamelCase													=    self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_megatron_bert_for_masked_lm(*__UpperCAmelCase				)




 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    __lowerCamelCase													=    self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*__UpperCAmelCase				)




 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    __lowerCamelCase													=    self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*__UpperCAmelCase				)




 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    __lowerCamelCase													=    self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_megatron_bert_for_pretraining(*__UpperCAmelCase				)




 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    __lowerCamelCase													=    self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_megatron_bert_for_question_answering(*__UpperCAmelCase				)




 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    __lowerCamelCase													=    self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*__UpperCAmelCase				)




 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    __lowerCamelCase													=    self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_megatron_bert_for_token_classification(*__UpperCAmelCase				)





def 		a__			(		_UpperCamelCase			: str							):
   return torch.tensor(
       _UpperCamelCase ,dtype=torch.long ,device=_UpperCamelCase ,)


a_			=					1E-4





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




 @slow
 @unittest.skip('''Model is not available.'''				)
 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    __lowerCamelCase													=    '''nvidia/megatron-bert-uncased-345m'''
    if "MYDIR" in os.environ:
       __lowerCamelCase													=    os.path.join(os.environ['''MYDIR''']	,				__UpperCAmelCase				)
    __lowerCamelCase													=    MegatronBertModel.from_pretrained(__UpperCAmelCase				)
    model.to(__UpperCAmelCase				)
    model.half()
    __lowerCamelCase													=    _long_tensor([[101, 7110, 1005, 1056, 2023, 11333, 17413, 1029, 102]]				)
    with torch.no_grad():
       __lowerCamelCase													=    model(__UpperCAmelCase				)[0]
    __lowerCamelCase													=    torch.Size((1, 9, 1024)				)
    self.assertEqual(output.shape	,				__UpperCAmelCase				)

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











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


a_			=					logging.get_logger(__name__)

a_			=					{
    """unc-nlp/lxmert-base-uncased""": """https://huggingface.co/unc-nlp/lxmert-base-uncased/resolve/main/config.json""",
}



class  __lowerCAmelCase     (	lowerCAmelCase__   ):
 lowerCAmelCase__       					=  """lxmert"""
 lowerCAmelCase__       					=  {}




 def __init__(						self	,				__UpperCAmelCase=30522	,				__UpperCAmelCase=768	,				__UpperCAmelCase=12	,				__UpperCAmelCase=9500	,				__UpperCAmelCase=1600	,				__UpperCAmelCase=400	,				__UpperCAmelCase=3072	,				__UpperCAmelCase="gelu"	,				__UpperCAmelCase=0.1	,				__UpperCAmelCase=0.1	,				__UpperCAmelCase=512	,				__UpperCAmelCase=2	,				__UpperCAmelCase=0.02	,				__UpperCAmelCase=1E-1_2	,				__UpperCAmelCase=9	,				__UpperCAmelCase=5	,				__UpperCAmelCase=5	,				__UpperCAmelCase=2048	,				__UpperCAmelCase=4	,				__UpperCAmelCase=6.67	,				__UpperCAmelCase=True	,				__UpperCAmelCase=True	,				__UpperCAmelCase=True	,				__UpperCAmelCase=True	,				__UpperCAmelCase=True	,				__UpperCAmelCase=True	,				__UpperCAmelCase=True	,				**__UpperCAmelCase	,				):



    '''simple docstring'''
    __lowerCamelCase													=    vocab_size
    __lowerCamelCase													=    hidden_size
    __lowerCamelCase													=    num_attention_heads
    __lowerCamelCase													=    hidden_act
    __lowerCamelCase													=    intermediate_size
    __lowerCamelCase													=    hidden_dropout_prob
    __lowerCamelCase													=    attention_probs_dropout_prob
    __lowerCamelCase													=    max_position_embeddings
    __lowerCamelCase													=    type_vocab_size
    __lowerCamelCase													=    initializer_range
    __lowerCamelCase													=    layer_norm_eps
    __lowerCamelCase													=    num_qa_labels
    __lowerCamelCase													=    num_object_labels
    __lowerCamelCase													=    num_attr_labels
    __lowerCamelCase													=    l_layers
    __lowerCamelCase													=    x_layers
    __lowerCamelCase													=    r_layers
    __lowerCamelCase													=    visual_feat_dim
    __lowerCamelCase													=    visual_pos_dim
    __lowerCamelCase													=    visual_loss_normalizer
    __lowerCamelCase													=    task_matched
    __lowerCamelCase													=    task_mask_lm
    __lowerCamelCase													=    task_obj_predict
    __lowerCamelCase													=    task_qa
    __lowerCamelCase													=    visual_obj_loss
    __lowerCamelCase													=    visual_attr_loss
    __lowerCamelCase													=    visual_feat_loss
    __lowerCamelCase													=    {'''vision''': r_layers, '''cross_encoder''': x_layers, '''language''': l_layers}
    super().__init__(**__UpperCAmelCase				)






