



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			:

 """simple docstring"""



 def __init__(	self		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE=1_3		,				SCREAMING_SNAKE_CASE=7		,				SCREAMING_SNAKE_CASE=True		,				SCREAMING_SNAKE_CASE=True		,				SCREAMING_SNAKE_CASE=True		,				SCREAMING_SNAKE_CASE=True		,				SCREAMING_SNAKE_CASE=9_9		,				SCREAMING_SNAKE_CASE=6_4		,				SCREAMING_SNAKE_CASE=3_2		,				SCREAMING_SNAKE_CASE=5		,				SCREAMING_SNAKE_CASE=4		,				SCREAMING_SNAKE_CASE=3_7		,				SCREAMING_SNAKE_CASE="gelu"		,				SCREAMING_SNAKE_CASE=0.1		,				SCREAMING_SNAKE_CASE=0.1		,				SCREAMING_SNAKE_CASE=5_1_2		,				SCREAMING_SNAKE_CASE=1_6		,				SCREAMING_SNAKE_CASE=2		,				SCREAMING_SNAKE_CASE=0.02		,				SCREAMING_SNAKE_CASE=3		,				SCREAMING_SNAKE_CASE=4		,				SCREAMING_SNAKE_CASE=None		,				) ->			Optional[int]:
  a__												= parent
  a__												= batch_size
  a__												= seq_length
  a__												= is_training
  a__												= use_input_mask
  a__												= use_token_type_ids
  a__												= use_labels
  a__												= vocab_size
  a__												= hidden_size
  a__												= embedding_size
  a__												= num_hidden_layers
  a__												= num_attention_heads
  a__												= intermediate_size
  a__												= hidden_act
  a__												= hidden_dropout_prob
  a__												= attention_probs_dropout_prob
  a__												= max_position_embeddings
  a__												= type_vocab_size
  a__												= type_sequence_label_size
  a__												= initializer_range
  a__												= num_labels
  a__												= num_choices
  a__												= scope

 def 						_UpperCAmelCase       (	self      ) ->			Any:
  a__												= ids_tensor([self.batch_size, self.seq_length]		,				self.vocab_size      )

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

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

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

  a__												= self.get_config()

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

 def 						_UpperCAmelCase       (	self      ) ->			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=SCREAMING_SNAKE_CASE		,				initializer_range=self.initializer_range		,				)

 def 						_UpperCAmelCase       (	self		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE      ) ->			str:
  a__												= MegatronBertModel(config=SCREAMING_SNAKE_CASE      )
  model.to(SCREAMING_SNAKE_CASE      )
  model.eval()
  a__												= model(SCREAMING_SNAKE_CASE		,				attention_mask=SCREAMING_SNAKE_CASE		,				token_type_ids=SCREAMING_SNAKE_CASE      )
  a__												= model(SCREAMING_SNAKE_CASE		,				token_type_ids=SCREAMING_SNAKE_CASE      )
  a__												= model(SCREAMING_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		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE      ) ->			str:
  a__												= MegatronBertForMaskedLM(config=SCREAMING_SNAKE_CASE      )
  model.to(SCREAMING_SNAKE_CASE      )
  model.eval()
  a__												= model(SCREAMING_SNAKE_CASE		,				attention_mask=SCREAMING_SNAKE_CASE		,				token_type_ids=SCREAMING_SNAKE_CASE		,				labels=SCREAMING_SNAKE_CASE      )
  self.parent.assertEqual(result.logits.shape		,				(self.batch_size, self.seq_length, self.vocab_size)      )

 def 						_UpperCAmelCase       (	self		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE      ) ->			Dict:
  a__												= MegatronBertForCausalLM(config=SCREAMING_SNAKE_CASE      )
  model.to(SCREAMING_SNAKE_CASE      )
  model.eval()
  a__												= model(SCREAMING_SNAKE_CASE		,				attention_mask=SCREAMING_SNAKE_CASE		,				token_type_ids=SCREAMING_SNAKE_CASE		,				labels=SCREAMING_SNAKE_CASE      )
  self.parent.assertEqual(result.logits.shape		,				(self.batch_size, self.seq_length, self.vocab_size)      )

 def 						_UpperCAmelCase       (	self		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE      ) ->			List[str]:
  a__												= MegatronBertForNextSentencePrediction(config=SCREAMING_SNAKE_CASE      )
  model.to(SCREAMING_SNAKE_CASE      )
  model.eval()
  a__												= model(
      SCREAMING_SNAKE_CASE		,				attention_mask=SCREAMING_SNAKE_CASE		,				token_type_ids=SCREAMING_SNAKE_CASE		,				labels=SCREAMING_SNAKE_CASE		,				)
  self.parent.assertEqual(result.logits.shape		,				(self.batch_size, 2)      )

 def 						_UpperCAmelCase       (	self		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE      ) ->			Any:
  a__												= MegatronBertForPreTraining(config=SCREAMING_SNAKE_CASE      )
  model.to(SCREAMING_SNAKE_CASE      )
  model.eval()
  a__												= model(
      SCREAMING_SNAKE_CASE		,				attention_mask=SCREAMING_SNAKE_CASE		,				token_type_ids=SCREAMING_SNAKE_CASE		,				labels=SCREAMING_SNAKE_CASE		,				next_sentence_label=SCREAMING_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		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE      ) ->			Tuple:
  a__												= MegatronBertForQuestionAnswering(config=SCREAMING_SNAKE_CASE      )
  model.to(SCREAMING_SNAKE_CASE      )
  model.eval()
  a__												= model(
      SCREAMING_SNAKE_CASE		,				attention_mask=SCREAMING_SNAKE_CASE		,				token_type_ids=SCREAMING_SNAKE_CASE		,				start_positions=SCREAMING_SNAKE_CASE		,				end_positions=SCREAMING_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		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE      ) ->			Dict:
  a__												= self.num_labels
  a__												= MegatronBertForSequenceClassification(SCREAMING_SNAKE_CASE      )
  model.to(SCREAMING_SNAKE_CASE      )
  model.eval()
  a__												= model(SCREAMING_SNAKE_CASE		,				attention_mask=SCREAMING_SNAKE_CASE		,				token_type_ids=SCREAMING_SNAKE_CASE		,				labels=SCREAMING_SNAKE_CASE      )
  self.parent.assertEqual(result.logits.shape		,				(self.batch_size, self.num_labels)      )

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

 def 						_UpperCAmelCase       (	self		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE      ) ->			List[str]:
  a__												= self.num_choices
  a__												= MegatronBertForMultipleChoice(config=SCREAMING_SNAKE_CASE      )
  model.to(SCREAMING_SNAKE_CASE      )
  model.eval()
  a__												= input_ids.unsqueeze(1      ).expand(-1		,				self.num_choices		,				-1      ).contiguous()
  a__												= token_type_ids.unsqueeze(1      ).expand(-1		,				self.num_choices		,				-1      ).contiguous()
  a__												= input_mask.unsqueeze(1      ).expand(-1		,				self.num_choices		,				-1      ).contiguous()
  a__												= model(
      SCREAMING_SNAKE_CASE		,				attention_mask=SCREAMING_SNAKE_CASE		,				token_type_ids=SCREAMING_SNAKE_CASE		,				labels=SCREAMING_SNAKE_CASE		,				)
  self.parent.assertEqual(result.logits.shape		,				(self.batch_size, self.num_choices)      )



 def 						_UpperCAmelCase       (	self      ) ->			int:
  a__												= self.prepare_config_and_inputs()
  (
      (
      a__
  )							,			(
      a__
  )							,			(
      a__
  )							,			(
      a__
  )							,			(
      a__
  )							,			(
      a__
  )							,			(
      a__
  )							,			
  )												= config_and_inputs
  a__												= {'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
  return config, inputs_dict





@require_torch
class    __UpperCamelCase			(       _lowercase	,	_lowercase	,	unittest.TestCase				):

 """simple docstring"""



 _lowercase   :    Any      						=					(
     (
         MegatronBertModel,
         MegatronBertForMaskedLM,
         MegatronBertForCausalLM,
         MegatronBertForMultipleChoice,
         MegatronBertForNextSentencePrediction,
         MegatronBertForPreTraining,
         MegatronBertForQuestionAnswering,
         MegatronBertForSequenceClassification,
         MegatronBertForTokenClassification,
     )
     if is_torch_available()
     else ()
 )
 _lowercase   :    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 {}
 )
 _lowercase   :    Any      						=					True
 # test_resize_embeddings = False
 _lowercase   :    Tuple      						=					False

 def 						_UpperCAmelCase       (	self		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE=False      ) ->			List[Any]:
  a__												= super()._prepare_for_class(SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				return_labels=SCREAMING_SNAKE_CASE      )

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

 def 						_UpperCAmelCase       (	self      ) ->			Union[str, Any]:
  a__												= MegatronBertModelTester(self      )
  a__												= ConfigTester(self		,				config_class=SCREAMING_SNAKE_CASE		,				hidden_size=3_7      )

 def 						_UpperCAmelCase       (	self      ) ->			Tuple:
  self.config_tester.run_common_tests()

 def 						_UpperCAmelCase       (	self      ) ->			Optional[int]:
  a__												= self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_megatron_bert_model(*SCREAMING_SNAKE_CASE      )

 def 						_UpperCAmelCase       (	self      ) ->			Dict:
  a__												= self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_megatron_bert_for_masked_lm(*SCREAMING_SNAKE_CASE      )

 def 						_UpperCAmelCase       (	self      ) ->			Dict:
  a__												= self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*SCREAMING_SNAKE_CASE      )

 def 						_UpperCAmelCase       (	self      ) ->			int:
  a__												= self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*SCREAMING_SNAKE_CASE      )

 def 						_UpperCAmelCase       (	self      ) ->			Any:
  a__												= self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_megatron_bert_for_pretraining(*SCREAMING_SNAKE_CASE      )

 def 						_UpperCAmelCase       (	self      ) ->			List[Any]:
  a__												= self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_megatron_bert_for_question_answering(*SCREAMING_SNAKE_CASE      )

 def 						_UpperCAmelCase       (	self      ) ->			Tuple:
  a__												= self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*SCREAMING_SNAKE_CASE      )



 def 						_UpperCAmelCase       (	self      ) ->			Optional[Any]:
  a__												= self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_megatron_bert_for_token_classification(*SCREAMING_SNAKE_CASE      )






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


a_      :      str    =     1E-4

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

 """simple docstring"""



 @slow
 @unittest.skip('''Model is not available.'''      )
 def 						_UpperCAmelCase       (	self      ) ->			int:
  a__												= '''nvidia/megatron-bert-uncased-345m'''
  if "MYDIR" in os.environ:
   a__												= os.path.join(os.environ['''MYDIR''']		,				SCREAMING_SNAKE_CASE      )
  a__												= MegatronBertModel.from_pretrained(SCREAMING_SNAKE_CASE      )
  model.to(SCREAMING_SNAKE_CASE      )
  model.half()
  a__												= _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():
   a__												= model(SCREAMING_SNAKE_CASE      )[0]
  a__												= torch.Size((1, 9, 1_0_2_4)      )
  self.assertEqual(output.shape		,				SCREAMING_SNAKE_CASE      )

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










import warnings

from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding
class    __UpperCamelCase			(       _lowercase				):

 """simple docstring"""



 _lowercase   :    Optional[Any]      						=					['''image_processor''', '''tokenizer''']
 _lowercase   :    Optional[int]      						=					'''ChineseCLIPImageProcessor'''
 _lowercase   :    List[Any]      						=					('''BertTokenizer''', '''BertTokenizerFast''')

 def __init__(	self		,				SCREAMING_SNAKE_CASE=None		,				SCREAMING_SNAKE_CASE=None		,				**SCREAMING_SNAKE_CASE      ) ->			Union[str, Any]:
  a__												= None
  if "feature_extractor" in kwargs:
   warnings.warn(
       '''The `feature_extractor` argument is deprecated and will be removed in v5, use `image_processor`'''
       ''' instead.'''		,				SCREAMING_SNAKE_CASE		,				)
   a__												= kwargs.pop('''feature_extractor'''      )

  a__												= image_processor if image_processor is not None else feature_extractor
  if image_processor is None:
   raise ValueError('''You need to specify an `image_processor`.'''      )
  if tokenizer is None:
   raise ValueError('''You need to specify a `tokenizer`.'''      )

  super().__init__(SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE      )
  a__												= self.image_processor

 def __call__(	self		,				SCREAMING_SNAKE_CASE=None		,				SCREAMING_SNAKE_CASE=None		,				SCREAMING_SNAKE_CASE=None		,				**SCREAMING_SNAKE_CASE      ) ->			List[str]:

  if text is None and images is None:
   raise ValueError('''You have to specify either text or images. Both cannot be none.'''      )

  if text is not None:
   a__												= self.tokenizer(SCREAMING_SNAKE_CASE		,				return_tensors=SCREAMING_SNAKE_CASE		,				**SCREAMING_SNAKE_CASE      )

  if images is not None:
   a__												= self.image_processor(SCREAMING_SNAKE_CASE		,				return_tensors=SCREAMING_SNAKE_CASE		,				**SCREAMING_SNAKE_CASE      )

  if text is not None and images is not None:
   a__												= image_features.pixel_values
   return encoding
  elif text is not None:
   return encoding
  else:
   return BatchEncoding(data=dict(**SCREAMING_SNAKE_CASE      )		,				tensor_type=SCREAMING_SNAKE_CASE      )

 def 						_UpperCAmelCase       (	self		,				*SCREAMING_SNAKE_CASE		,				**SCREAMING_SNAKE_CASE      ) ->			List[str]:
  return self.tokenizer.batch_decode(*SCREAMING_SNAKE_CASE		,				**SCREAMING_SNAKE_CASE      )

 def 						_UpperCAmelCase       (	self		,				*SCREAMING_SNAKE_CASE		,				**SCREAMING_SNAKE_CASE      ) ->			int:
  return self.tokenizer.decode(*SCREAMING_SNAKE_CASE		,				**SCREAMING_SNAKE_CASE      )

 @property
 def 						_UpperCAmelCase       (	self      ) ->			List[str]:
  a__												= self.tokenizer.model_input_names
  a__												= self.image_processor.model_input_names
  return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names      )      )



 @property
 def 						_UpperCAmelCase       (	self      ) ->			Optional[int]:
  warnings.warn(
      '''`feature_extractor_class` is deprecated and will be removed in v5. Use `image_processor_class` instead.'''		,				SCREAMING_SNAKE_CASE		,				)
  return self.image_processor_class





