






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






   def __init__(       self		,					a		,					a=13		,					a=7		,					a=True		,					a=True		,					a=True		,					a=True		,					a=99		,					a=64		,					a=32		,					a=5		,					a=4		,					a=37		,					a="gelu"		,					a=0.1		,					a=0.1		,					a=512		,					a=16		,					a=2		,					a=0.02		,					a=3		,					a=4		,					a=None		,					)      ->       str:
      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    SCREAMING_SNAKE_CASE__    (       self)      ->       Union[str, Any]:
      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    SCREAMING_SNAKE_CASE__    (       self)      ->       List[str]:
      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=_lowercase		,					initializer_range=self.initializer_range		,					)






   def    SCREAMING_SNAKE_CASE__    (       self		,					a		,					a		,					a		,					a		,					a		,					a		,					a)      ->       Optional[Any]:
      SCREAMING_SNAKE_CASE						= MegatronBertModel(config=_lowercase)
      model.to(_lowercase)
      model.eval()
      SCREAMING_SNAKE_CASE						= model(_lowercase		,					attention_mask=_lowercase		,					token_type_ids=_lowercase)
      SCREAMING_SNAKE_CASE						= model(_lowercase		,					token_type_ids=_lowercase)
      SCREAMING_SNAKE_CASE						= model(_lowercase)

      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    SCREAMING_SNAKE_CASE__    (       self		,					a		,					a		,					a		,					a		,					a		,					a		,					a)      ->       Optional[int]:
      SCREAMING_SNAKE_CASE						= MegatronBertForMaskedLM(config=_lowercase)
      model.to(_lowercase)
      model.eval()
      SCREAMING_SNAKE_CASE						= model(_lowercase		,					attention_mask=_lowercase		,					token_type_ids=_lowercase		,					labels=_lowercase)
      self.parent.assertEqual(result.logits.shape		,					(self.batch_size, self.seq_length, self.vocab_size))






   def    SCREAMING_SNAKE_CASE__    (       self		,					a		,					a		,					a		,					a		,					a		,					a		,					a)      ->       int:
      SCREAMING_SNAKE_CASE						= MegatronBertForCausalLM(config=_lowercase)
      model.to(_lowercase)
      model.eval()
      SCREAMING_SNAKE_CASE						= model(_lowercase		,					attention_mask=_lowercase		,					token_type_ids=_lowercase		,					labels=_lowercase)
      self.parent.assertEqual(result.logits.shape		,					(self.batch_size, self.seq_length, self.vocab_size))






   def    SCREAMING_SNAKE_CASE__    (       self		,					a		,					a		,					a		,					a		,					a		,					a		,					a)      ->       Union[str, Any]:
      SCREAMING_SNAKE_CASE						= MegatronBertForNextSentencePrediction(config=_lowercase)
      model.to(_lowercase)
      model.eval()
      SCREAMING_SNAKE_CASE						= model(
          _lowercase		,					attention_mask=_lowercase		,					token_type_ids=_lowercase		,					labels=_lowercase		,					)
      self.parent.assertEqual(result.logits.shape		,					(self.batch_size, 2))






   def    SCREAMING_SNAKE_CASE__    (       self		,					a		,					a		,					a		,					a		,					a		,					a		,					a)      ->       Any:
      SCREAMING_SNAKE_CASE						= MegatronBertForPreTraining(config=_lowercase)
      model.to(_lowercase)
      model.eval()
      SCREAMING_SNAKE_CASE						= model(
          _lowercase		,					attention_mask=_lowercase		,					token_type_ids=_lowercase		,					labels=_lowercase		,					next_sentence_label=_lowercase		,					)
      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    SCREAMING_SNAKE_CASE__    (       self		,					a		,					a		,					a		,					a		,					a		,					a		,					a)      ->       Optional[int]:
      SCREAMING_SNAKE_CASE						= MegatronBertForQuestionAnswering(config=_lowercase)
      model.to(_lowercase)
      model.eval()
      SCREAMING_SNAKE_CASE						= model(
          _lowercase		,					attention_mask=_lowercase		,					token_type_ids=_lowercase		,					start_positions=_lowercase		,					end_positions=_lowercase		,					)
      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    SCREAMING_SNAKE_CASE__    (       self		,					a		,					a		,					a		,					a		,					a		,					a		,					a)      ->       List[str]:
      SCREAMING_SNAKE_CASE						= self.num_labels
      SCREAMING_SNAKE_CASE						= MegatronBertForSequenceClassification(_lowercase)
      model.to(_lowercase)
      model.eval()
      SCREAMING_SNAKE_CASE						= model(_lowercase		,					attention_mask=_lowercase		,					token_type_ids=_lowercase		,					labels=_lowercase)
      self.parent.assertEqual(result.logits.shape		,					(self.batch_size, self.num_labels))






   def    SCREAMING_SNAKE_CASE__    (       self		,					a		,					a		,					a		,					a		,					a		,					a		,					a)      ->       Union[str, Any]:
      SCREAMING_SNAKE_CASE						= self.num_labels
      SCREAMING_SNAKE_CASE						= MegatronBertForTokenClassification(config=_lowercase)
      model.to(_lowercase)
      model.eval()
      SCREAMING_SNAKE_CASE						= model(_lowercase		,					attention_mask=_lowercase		,					token_type_ids=_lowercase		,					labels=_lowercase)
      self.parent.assertEqual(result.logits.shape		,					(self.batch_size, self.seq_length, self.num_labels))






   def    SCREAMING_SNAKE_CASE__    (       self		,					a		,					a		,					a		,					a		,					a		,					a		,					a)      ->       Union[str, Any]:
      SCREAMING_SNAKE_CASE						= self.num_choices
      SCREAMING_SNAKE_CASE						= MegatronBertForMultipleChoice(config=_lowercase)
      model.to(_lowercase)
      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(
          _lowercase		,					attention_mask=_lowercase		,					token_type_ids=_lowercase		,					labels=_lowercase		,					)
      self.parent.assertEqual(result.logits.shape		,					(self.batch_size, self.num_choices))







   def    SCREAMING_SNAKE_CASE__    (       self)      ->       List[Any]:
      SCREAMING_SNAKE_CASE						= self.prepare_config_and_inputs()
      (
          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 		_snake_case  (  UpperCAmelCase_   ,    UpperCAmelCase_   ,    unittest.TestCase					):
   _lowercase		:    Dict			     =							(
       (
           MegatronBertModel,
           MegatronBertForMaskedLM,
           MegatronBertForCausalLM,
           MegatronBertForMultipleChoice,
           MegatronBertForNextSentencePrediction,
           MegatronBertForPreTraining,
           MegatronBertForQuestionAnswering,
           MegatronBertForSequenceClassification,
           MegatronBertForTokenClassification,
       )
       if is_torch_available()
       else ()
   )
   _lowercase		:    List[str]			     =							(
       {
           '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 {}
   )
   _lowercase		:    Any			     =							True
   # test_resize_embeddings = False
   _lowercase		:    List[str]			     =							False






   def    SCREAMING_SNAKE_CASE__    (       self		,					a		,					a		,					a=False)      ->       Dict:
      SCREAMING_SNAKE_CASE						= super()._prepare_for_class(_lowercase		,					_lowercase		,					return_labels=_lowercase)

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






   def    SCREAMING_SNAKE_CASE__    (       self)      ->       Optional[int]:
      SCREAMING_SNAKE_CASE						= MegatronBertModelTester(self)
      SCREAMING_SNAKE_CASE						= ConfigTester(self		,					config_class=_lowercase		,					hidden_size=37)






   def    SCREAMING_SNAKE_CASE__    (       self)      ->       List[Any]:
      self.config_tester.run_common_tests()






   def    SCREAMING_SNAKE_CASE__    (       self)      ->       Dict:
      SCREAMING_SNAKE_CASE						= self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_megatron_bert_model(*_lowercase)






   def    SCREAMING_SNAKE_CASE__    (       self)      ->       List[str]:
      SCREAMING_SNAKE_CASE						= self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_megatron_bert_for_masked_lm(*_lowercase)






   def    SCREAMING_SNAKE_CASE__    (       self)      ->       Tuple:
      SCREAMING_SNAKE_CASE						= self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*_lowercase)






   def    SCREAMING_SNAKE_CASE__    (       self)      ->       Any:
      SCREAMING_SNAKE_CASE						= self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*_lowercase)






   def    SCREAMING_SNAKE_CASE__    (       self)      ->       Tuple:
      SCREAMING_SNAKE_CASE						= self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_megatron_bert_for_pretraining(*_lowercase)






   def    SCREAMING_SNAKE_CASE__    (       self)      ->       Optional[int]:
      SCREAMING_SNAKE_CASE						= self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_megatron_bert_for_question_answering(*_lowercase)






   def    SCREAMING_SNAKE_CASE__    (       self)      ->       Dict:
      SCREAMING_SNAKE_CASE						= self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*_lowercase)







   def    SCREAMING_SNAKE_CASE__    (       self)      ->       List[Any]:
      SCREAMING_SNAKE_CASE						= self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_megatron_bert_for_token_classification(*_lowercase)






def 				lowerCamelCase__				(_UpperCAmelCase):
   return torch.tensor(
       _UpperCAmelCase							,       dtype=torch.long							,       device=_UpperCAmelCase							,       )


a_		:						Any             =    1E-4







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






   @slow
   @unittest.skip('Model is not available.')
   def    SCREAMING_SNAKE_CASE__    (       self)      ->       Tuple:
      SCREAMING_SNAKE_CASE						= 'nvidia/megatron-bert-uncased-345m'
      if "MYDIR" in os.environ:
         SCREAMING_SNAKE_CASE						= os.path.join(os.environ['MYDIR']		,					_lowercase)
      SCREAMING_SNAKE_CASE						= MegatronBertModel.from_pretrained(_lowercase)
      model.to(_lowercase)
      model.half()
      SCREAMING_SNAKE_CASE						= _long_tensor([[101, 7110, 1005, 1056, 2023, 1_1333, 1_7413, 1029, 102]])
      with torch.no_grad():
         SCREAMING_SNAKE_CASE						= model(_lowercase)[0]
      SCREAMING_SNAKE_CASE						= torch.Size((1, 9, 1024))
      self.assertEqual(output.shape		,					_lowercase)

      SCREAMING_SNAKE_CASE						= [-0.60_40, -0.25_17, -0.10_25, 0.34_20, -0.67_58, -0.00_17, -0.10_89, -0.19_90, 0.57_28]
      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(_lowercase		,					_lowercase		,					_lowercase		,					_lowercase)
            self.assertTrue(math.isclose(_lowercase		,					_lowercase		,					rel_tol=_lowercase		,					abs_tol=_lowercase)		,					msg=_lowercase)












from typing import List, Optional, Union

from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding, PaddingStrategy, PreTokenizedInput, TextInput, TruncationStrategy
from ...utils import TensorType

class 		_snake_case  (  A__					):
			_lowercase		:    Union[str, Any]			     =							['''image_processor''', '''tokenizer''']
			_lowercase		:    Optional[int]			     =							'''Pix2StructImageProcessor'''
			_lowercase		:    str			     =							('''T5Tokenizer''', '''T5TokenizerFast''')






			def __init__(       self		,					a		,					a)      ->       Any:
						SCREAMING_SNAKE_CASE						= False
						super().__init__(a		,					a)






			def __call__(       self		,					a=None		,					a = None		,					a = True		,					a = False		,					a = None		,					a = None		,					a = 2048		,					a = 0		,					a = None		,					a = None		,					a = False		,					a = False		,					a = False		,					a = False		,					a = False		,					a = True		,					a = None		,					**a		,					)      ->       BatchEncoding:
						if images is None and text is None:
									raise ValueError('You have to specify either images or text.')

						# Get only text
						if images is None and not self.image_processor.is_vqa:
									SCREAMING_SNAKE_CASE						= self.tokenizer
									SCREAMING_SNAKE_CASE						= self.tokenizer(
									    text=a		,					add_special_tokens=a		,					padding=a		,					truncation=a		,					max_length=a		,					stride=a		,					pad_to_multiple_of=a		,					return_attention_mask=a		,					return_overflowing_tokens=a		,					return_special_tokens_mask=a		,					return_offsets_mapping=a		,					return_token_type_ids=a		,					return_length=a		,					verbose=a		,					return_tensors=a		,					**a		,					)
									return text_encoding

						if not self.image_processor.is_vqa:
									# add pixel_values
									SCREAMING_SNAKE_CASE						= self.image_processor(
									    a		,					return_tensors=a		,					max_patches=a		,					**a)
						else:
									# add pixel_values and bbox
									SCREAMING_SNAKE_CASE						= self.image_processor(
									    a		,					return_tensors=a		,					max_patches=a		,					header_text=a		,					**a)

						if text is not None and not self.image_processor.is_vqa:
									SCREAMING_SNAKE_CASE						= self.tokenizer(
									    text=a		,					add_special_tokens=a		,					padding=a		,					truncation=a		,					max_length=a		,					stride=a		,					pad_to_multiple_of=a		,					return_attention_mask=a		,					return_overflowing_tokens=a		,					return_special_tokens_mask=a		,					return_offsets_mapping=a		,					return_token_type_ids=a		,					return_length=a		,					verbose=a		,					return_tensors=a		,					**a		,					)

									if "attention_mask" in text_encoding:
												SCREAMING_SNAKE_CASE						= text_encoding.pop('attention_mask')
									if "input_ids" in text_encoding:
												SCREAMING_SNAKE_CASE						= text_encoding.pop('input_ids')
						else:
									SCREAMING_SNAKE_CASE						= None

						if text_encoding is not None:
									encoding_image_processor.update(a)

						return encoding_image_processor






			def    SCREAMING_SNAKE_CASE__    (       self		,					*a		,					**a)      ->       List[Any]:
						return self.tokenizer.batch_decode(*a		,					**a)






			def    SCREAMING_SNAKE_CASE__    (       self		,					*a		,					**a)      ->       Optional[int]:
						return self.tokenizer.decode(*a		,					**a)







			@property
			def    SCREAMING_SNAKE_CASE__    (       self)      ->       Tuple:
						SCREAMING_SNAKE_CASE						= self.tokenizer.model_input_names
						SCREAMING_SNAKE_CASE						= self.image_processor.model_input_names
						return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names))




