


"""simple docstring"""


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       ,  __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=5_12       ,  __a=16       ,  __a=2       ,  __a=0.0_2       ,  __a=3       ,  __a=4       ,  __a=None       ,  ):
										__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      snake_case							(	self      ):
										__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      snake_case							(	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=__a       ,  initializer_range=self.initializer_range       ,  )

							def      snake_case							(	self       ,  __a       ,  __a       ,  __a       ,  __a       ,  __a       ,  __a       ,  __a      ):
										__lowerCAmelCase							       =     MegatronBertModel(config=__a      )
										model.to(__a      )
										model.eval()
										__lowerCAmelCase							       =     model(__a       ,  attention_mask=__a       ,  token_type_ids=__a      )
										__lowerCAmelCase							       =     model(__a       ,  token_type_ids=__a      )
										__lowerCAmelCase							       =     model(__a      )

										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       ,  __a       ,  __a       ,  __a       ,  __a       ,  __a       ,  __a       ,  __a      ):
										__lowerCAmelCase							       =     MegatronBertForMaskedLM(config=__a      )
										model.to(__a      )
										model.eval()
										__lowerCAmelCase							       =     model(__a       ,  attention_mask=__a       ,  token_type_ids=__a       ,  labels=__a      )
										self.parent.assertEqual(result.logits.shape       ,  (self.batch_size, self.seq_length, self.vocab_size)      )

							def      snake_case							(	self       ,  __a       ,  __a       ,  __a       ,  __a       ,  __a       ,  __a       ,  __a      ):
										__lowerCAmelCase							       =     MegatronBertForCausalLM(config=__a      )
										model.to(__a      )
										model.eval()
										__lowerCAmelCase							       =     model(__a       ,  attention_mask=__a       ,  token_type_ids=__a       ,  labels=__a      )
										self.parent.assertEqual(result.logits.shape       ,  (self.batch_size, self.seq_length, self.vocab_size)      )

							def      snake_case							(	self       ,  __a       ,  __a       ,  __a       ,  __a       ,  __a       ,  __a       ,  __a      ):
										__lowerCAmelCase							       =     MegatronBertForNextSentencePrediction(config=__a      )
										model.to(__a      )
										model.eval()
										__lowerCAmelCase							       =     model(
										    __a       ,  attention_mask=__a       ,  token_type_ids=__a       ,  labels=__a       ,  )
										self.parent.assertEqual(result.logits.shape       ,  (self.batch_size, 2)      )

							def      snake_case							(	self       ,  __a       ,  __a       ,  __a       ,  __a       ,  __a       ,  __a       ,  __a      ):
										__lowerCAmelCase							       =     MegatronBertForPreTraining(config=__a      )
										model.to(__a      )
										model.eval()
										__lowerCAmelCase							       =     model(
										    __a       ,  attention_mask=__a       ,  token_type_ids=__a       ,  labels=__a       ,  next_sentence_label=__a       ,  )
										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       ,  __a       ,  __a       ,  __a       ,  __a       ,  __a       ,  __a       ,  __a      ):
										__lowerCAmelCase							       =     MegatronBertForQuestionAnswering(config=__a      )
										model.to(__a      )
										model.eval()
										__lowerCAmelCase							       =     model(
										    __a       ,  attention_mask=__a       ,  token_type_ids=__a       ,  start_positions=__a       ,  end_positions=__a       ,  )
										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       ,  __a       ,  __a       ,  __a       ,  __a       ,  __a       ,  __a       ,  __a      ):
										__lowerCAmelCase							       =     self.num_labels
										__lowerCAmelCase							       =     MegatronBertForSequenceClassification(__a      )
										model.to(__a      )
										model.eval()
										__lowerCAmelCase							       =     model(__a       ,  attention_mask=__a       ,  token_type_ids=__a       ,  labels=__a      )
										self.parent.assertEqual(result.logits.shape       ,  (self.batch_size, self.num_labels)      )

							def      snake_case							(	self       ,  __a       ,  __a       ,  __a       ,  __a       ,  __a       ,  __a       ,  __a      ):
										__lowerCAmelCase							       =     self.num_labels
										__lowerCAmelCase							       =     MegatronBertForTokenClassification(config=__a      )
										model.to(__a      )
										model.eval()
										__lowerCAmelCase							       =     model(__a       ,  attention_mask=__a       ,  token_type_ids=__a       ,  labels=__a      )
										self.parent.assertEqual(result.logits.shape       ,  (self.batch_size, self.seq_length, self.num_labels)      )

							def      snake_case							(	self       ,  __a       ,  __a       ,  __a       ,  __a       ,  __a       ,  __a       ,  __a      ):
										__lowerCAmelCase							       =     self.num_choices
										__lowerCAmelCase							       =     MegatronBertForMultipleChoice(config=__a      )
										model.to(__a      )
										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(
										    __a       ,  attention_mask=__a       ,  token_type_ids=__a       ,  labels=__a       ,  )
										self.parent.assertEqual(result.logits.shape       ,  (self.batch_size, self.num_choices)      )



							def      snake_case							(	self      ):
										__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     _UpperCamelCase    (  lowerCAmelCase__			,lowerCAmelCase__			,unittest.TestCase						):
							'''simple docstring'''







							__UpperCAmelCase		:       str							 =(
							    (
							        MegatronBertModel,
							        MegatronBertForMaskedLM,
							        MegatronBertForCausalLM,
							        MegatronBertForMultipleChoice,
							        MegatronBertForNextSentencePrediction,
							        MegatronBertForPreTraining,
							        MegatronBertForQuestionAnswering,
							        MegatronBertForSequenceClassification,
							        MegatronBertForTokenClassification,
							    )
							    if is_torch_available()
							    else ()
							)
							__UpperCAmelCase		:       Optional[int]							 =(
							    {
							        """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		:       Union[str, Any]							 =True
							# test_resize_embeddings = False
							__UpperCAmelCase		:       str							 =False

							def      snake_case							(	self       ,  __a       ,  __a       ,  __a=False      ):
										__lowerCAmelCase							       =     super()._prepare_for_class(__a       ,  __a       ,  return_labels=__a      )

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

							def      snake_case							(	self      ):
										__lowerCAmelCase							       =     MegatronBertModelTester(self      )
										__lowerCAmelCase							       =     ConfigTester(self       ,  config_class=__a       ,  hidden_size=37      )

							def      snake_case							(	self      ):
										self.config_tester.run_common_tests()

							def      snake_case							(	self      ):
										__lowerCAmelCase							       =     self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_megatron_bert_model(*__a      )

							def      snake_case							(	self      ):
										__lowerCAmelCase							       =     self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_megatron_bert_for_masked_lm(*__a      )

							def      snake_case							(	self      ):
										__lowerCAmelCase							       =     self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*__a      )

							def      snake_case							(	self      ):
										__lowerCAmelCase							       =     self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*__a      )

							def      snake_case							(	self      ):
										__lowerCAmelCase							       =     self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_megatron_bert_for_pretraining(*__a      )

							def      snake_case							(	self      ):
										__lowerCAmelCase							       =     self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_megatron_bert_for_question_answering(*__a      )

							def      snake_case							(	self      ):
										__lowerCAmelCase							       =     self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*__a      )



							def      snake_case							(	self      ):
										__lowerCAmelCase							       =     self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_megatron_bert_for_token_classification(*__a      )


def _lowerCamelCase      (						_UpperCamelCase ):

			'''simple docstring'''




			return torch.tensor(
			    _UpperCamelCase       ,					dtype=torch.long       ,					device=_UpperCamelCase       ,					)


A	:  int            =     1e-4


@require_torch
@require_sentencepiece
@require_tokenizers
class     _UpperCamelCase    (  unittest.TestCase						):
							'''simple docstring'''







							@slow
							@unittest.skip("Model is not available."      )
							def      snake_case							(	self      ):
										__lowerCAmelCase							       =     "nvidia/megatron-bert-uncased-345m"
										if "MYDIR" in os.environ:
													__lowerCAmelCase							       =     os.path.join(os.environ["MYDIR"]       ,  __a      )
										__lowerCAmelCase							       =     MegatronBertModel.from_pretrained(__a      )
										model.to(__a      )
										model.half()
										__lowerCAmelCase							       =     _long_tensor([[1_01, 71_10, 10_05, 10_56, 20_23, 1_13_33, 1_74_13, 10_29, 1_02]]      )
										with torch.no_grad():
													__lowerCAmelCase							       =     model(__a      )[0]
										__lowerCAmelCase							       =     torch.Size((1, 9, 10_24)      )
										self.assertEqual(output.shape       ,  __a      )

										__lowerCAmelCase							       =     [-0.6_0_4_0, -0.2_5_1_7, -0.1_0_2_5, 0.3_4_2_0, -0.6_7_5_8, -0.0_0_1_7, -0.1_0_8_9, -0.1_9_9_0, 0.5_7_2_8]
										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(__a       ,  __a       ,  __a       ,  __a      )
																self.assertTrue(math.isclose(__a       ,  __a       ,  rel_tol=__a       ,  abs_tol=__a      )       ,  msg=__a      )








"""simple docstring"""


import warnings

from ...utils import logging
from .image_processing_yolos import YolosImageProcessor


A	:  Optional[Any]            =     logging.get_logger(__name__)
class     _UpperCamelCase    (  lowerCAmelCase__						):
							'''simple docstring'''







							def __init__(	self       ,  *__a       ,  **__a      ):
										warnings.warn(
										    "The class YolosFeatureExtractor is deprecated and will be removed in version 5 of Transformers. Please"
										    " use YolosImageProcessor instead."       ,  __a       ,  )
										super().__init__(*__a       ,  **__a      )




