






'''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    _SCREAMING_SNAKE_CASE							:


			'''simple docstring'''





			def __init__(   self							:						Tuple ,      lowercase							:						str ,      lowercase							:						Optional[Any]=1_3 ,      lowercase							:						Tuple=7 ,      lowercase							:						str=True ,      lowercase							:						List[Any]=True ,      lowercase							:						Tuple=True ,      lowercase							:						Optional[Any]=True ,      lowercase							:						Tuple=9_9 ,      lowercase							:						Optional[int]=6_4 ,      lowercase							:						int=3_2 ,      lowercase							:						Dict=5 ,      lowercase							:						Dict=4 ,      lowercase							:						Dict=3_7 ,      lowercase							:						Any="gelu" ,      lowercase							:						Dict=0.1 ,      lowercase							:						List[str]=0.1 ,      lowercase							:						int=5_1_2 ,      lowercase							:						Any=1_6 ,      lowercase							:						Any=2 ,      lowercase							:						Optional[int]=0.0_2 ,      lowercase							:						Dict=3 ,      lowercase							:						Any=4 ,      lowercase							:						Tuple=None ,      )							->    Tuple:

									'''simple docstring'''







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






			def 				A      (   self							:						int			)							->    Optional[int]:

									'''simple docstring'''







									UpperCamelCase__         = ids_tensor([self.batch_size, self.seq_length] ,      self.vocab_size			)

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

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

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

									UpperCamelCase__         = self.get_config()

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






			def 				A      (   self							:						Union[str, Any]			)							->    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=lowercase ,      initializer_range=self.initializer_range ,      )






			def 				A      (   self							:						Union[str, Any] ,      lowercase							:						int ,      lowercase							:						List[Any] ,      lowercase							:						Tuple ,      lowercase							:						List[str] ,      lowercase							:						Tuple ,      lowercase							:						Any ,      lowercase							:						Union[str, Any]			)							->    Optional[int]:

									'''simple docstring'''







									UpperCamelCase__         = MegatronBertModel(config=lowercase			)
									model.to(lowercase			)
									model.eval()
									UpperCamelCase__         = model(lowercase ,      attention_mask=lowercase ,      token_type_ids=lowercase			)
									UpperCamelCase__         = model(lowercase ,      token_type_ids=lowercase			)
									UpperCamelCase__         = 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 				A      (   self							:						Union[str, Any] ,      lowercase							:						int ,      lowercase							:						List[Any] ,      lowercase							:						Any ,      lowercase							:						Optional[int] ,      lowercase							:						str ,      lowercase							:						Union[str, Any] ,      lowercase							:						List[str]			)							->    Any:

									'''simple docstring'''







									UpperCamelCase__         = MegatronBertForMaskedLM(config=lowercase			)
									model.to(lowercase			)
									model.eval()
									UpperCamelCase__         = 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 				A      (   self							:						Tuple ,      lowercase							:						Any ,      lowercase							:						List[str] ,      lowercase							:						str ,      lowercase							:						Tuple ,      lowercase							:						List[Any] ,      lowercase							:						Optional[int] ,      lowercase							:						int			)							->    int:

									'''simple docstring'''







									UpperCamelCase__         = MegatronBertForCausalLM(config=lowercase			)
									model.to(lowercase			)
									model.eval()
									UpperCamelCase__         = 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 				A      (   self							:						Optional[Any] ,      lowercase							:						List[str] ,      lowercase							:						List[str] ,      lowercase							:						int ,      lowercase							:						Optional[int] ,      lowercase							:						int ,      lowercase							:						Any ,      lowercase							:						Tuple			)							->    Tuple:

									'''simple docstring'''







									UpperCamelCase__         = MegatronBertForNextSentencePrediction(config=lowercase			)
									model.to(lowercase			)
									model.eval()
									UpperCamelCase__         = model(
									    lowercase ,      attention_mask=lowercase ,      token_type_ids=lowercase ,      labels=lowercase ,      )
									self.parent.assertEqual(result.logits.shape ,      (self.batch_size, 2)			)






			def 				A      (   self							:						int ,      lowercase							:						int ,      lowercase							:						List[str] ,      lowercase							:						Any ,      lowercase							:						List[Any] ,      lowercase							:						Optional[int] ,      lowercase							:						Union[str, Any] ,      lowercase							:						int			)							->    str:

									'''simple docstring'''







									UpperCamelCase__         = MegatronBertForPreTraining(config=lowercase			)
									model.to(lowercase			)
									model.eval()
									UpperCamelCase__         = 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 				A      (   self							:						Union[str, Any] ,      lowercase							:						Tuple ,      lowercase							:						List[Any] ,      lowercase							:						Optional[Any] ,      lowercase							:						List[str] ,      lowercase							:						List[Any] ,      lowercase							:						Any ,      lowercase							:						int			)							->    Dict:

									'''simple docstring'''







									UpperCamelCase__         = MegatronBertForQuestionAnswering(config=lowercase			)
									model.to(lowercase			)
									model.eval()
									UpperCamelCase__         = 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 				A      (   self							:						str ,      lowercase							:						Any ,      lowercase							:						Any ,      lowercase							:						str ,      lowercase							:						Dict ,      lowercase							:						Tuple ,      lowercase							:						Union[str, Any] ,      lowercase							:						Tuple			)							->    Dict:

									'''simple docstring'''







									UpperCamelCase__         = self.num_labels
									UpperCamelCase__         = MegatronBertForSequenceClassification(lowercase			)
									model.to(lowercase			)
									model.eval()
									UpperCamelCase__         = model(lowercase ,      attention_mask=lowercase ,      token_type_ids=lowercase ,      labels=lowercase			)
									self.parent.assertEqual(result.logits.shape ,      (self.batch_size, self.num_labels)			)






			def 				A      (   self							:						str ,      lowercase							:						List[str] ,      lowercase							:						Union[str, Any] ,      lowercase							:						List[str] ,      lowercase							:						Optional[int] ,      lowercase							:						Any ,      lowercase							:						str ,      lowercase							:						Any			)							->    Union[str, Any]:

									'''simple docstring'''







									UpperCamelCase__         = self.num_labels
									UpperCamelCase__         = MegatronBertForTokenClassification(config=lowercase			)
									model.to(lowercase			)
									model.eval()
									UpperCamelCase__         = 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 				A      (   self							:						Optional[Any] ,      lowercase							:						Any ,      lowercase							:						Any ,      lowercase							:						Dict ,      lowercase							:						List[str] ,      lowercase							:						Any ,      lowercase							:						Any ,      lowercase							:						Dict			)							->    Optional[int]:

									'''simple docstring'''







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







			def 				A      (   self							:						Union[str, Any]			)							->    int:

									'''simple docstring'''







									UpperCamelCase__         = self.prepare_config_and_inputs()
									(
									    (
									    UpperCamelCase__
									)    ,   (
									    UpperCamelCase__
									)    ,   (
									    UpperCamelCase__
									)    ,   (
									    UpperCamelCase__
									)    ,   (
									    UpperCamelCase__
									)    ,   (
									    UpperCamelCase__
									)    ,   (
									    UpperCamelCase__
									)    ,   
									)         = config_and_inputs
									UpperCamelCase__         = {"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
									return config, inputs_dict
@require_torch
class    _SCREAMING_SNAKE_CASE							(		SCREAMING_SNAKE_CASE       ,SCREAMING_SNAKE_CASE       ,unittest.TestCase	):


			'''simple docstring'''





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






			def 				A      (   self							:						Dict ,      lowercase							:						Optional[Any] ,      lowercase							:						str ,      lowercase							:						Dict=False			)							->    str:

									'''simple docstring'''







									UpperCamelCase__         = super()._prepare_for_class(lowercase ,      lowercase ,      return_labels=lowercase			)

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






			def 				A      (   self							:						List[str]			)							->    List[str]:

									'''simple docstring'''







									UpperCamelCase__         = MegatronBertModelTester(self			)
									UpperCamelCase__         = ConfigTester(self ,      config_class=lowercase ,      hidden_size=3_7			)






			def 				A      (   self							:						Optional[int]			)							->    Optional[int]:

									'''simple docstring'''







									self.config_tester.run_common_tests()






			def 				A      (   self							:						Tuple			)							->    Tuple:

									'''simple docstring'''







									UpperCamelCase__         = self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_megatron_bert_model(*lowercase			)






			def 				A      (   self							:						Union[str, Any]			)							->    Tuple:

									'''simple docstring'''







									UpperCamelCase__         = self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_megatron_bert_for_masked_lm(*lowercase			)






			def 				A      (   self							:						Dict			)							->    List[Any]:

									'''simple docstring'''







									UpperCamelCase__         = self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*lowercase			)






			def 				A      (   self							:						Tuple			)							->    int:

									'''simple docstring'''







									UpperCamelCase__         = self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*lowercase			)






			def 				A      (   self							:						int			)							->    Optional[int]:

									'''simple docstring'''







									UpperCamelCase__         = self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_megatron_bert_for_pretraining(*lowercase			)






			def 				A      (   self							:						List[str]			)							->    Tuple:

									'''simple docstring'''







									UpperCamelCase__         = self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_megatron_bert_for_question_answering(*lowercase			)






			def 				A      (   self							:						Dict			)							->    str:

									'''simple docstring'''







									UpperCamelCase__         = self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*lowercase			)







			def 				A      (   self							:						int			)							->    Optional[int]:

									'''simple docstring'''







									UpperCamelCase__         = self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_megatron_bert_for_token_classification(*lowercase			)







def 	__magic_name__( _A					):
						'''simple docstring'''





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


lowerCamelCase_				: Optional[Any]    	=							1e-4







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


			'''simple docstring'''





			@slow
			@unittest.skip("""Model is not available."""			)
			def 				A      (   self							:						Optional[int]			)							->    Optional[Any]:

									'''simple docstring'''







									UpperCamelCase__         = """nvidia/megatron-bert-uncased-345m"""
									if "MYDIR" in os.environ:
															UpperCamelCase__         = os.path.join(os.environ["""MYDIR"""] ,      lowercase			)
									UpperCamelCase__         = MegatronBertModel.from_pretrained(lowercase			)
									model.to(lowercase			)
									model.half()
									UpperCamelCase__         = _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():
															UpperCamelCase__         = model(lowercase			)[0]
									UpperCamelCase__         = torch.Size((1, 9, 1_0_2_4)			)
									self.assertEqual(output.shape ,      lowercase			)

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








'''simple docstring'''


from typing import List, Optional, Union

import numpy as np
import PIL.Image

from ...image_processing_utils import BaseImageProcessor, BatchFeature
from ...image_transforms import rescale, resize, to_channel_dimension_format
from ...image_utils import (
    ChannelDimension,
    PILImageResampling,
    get_image_size,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, logging


lowerCamelCase_				: Optional[int]    	=							logging.get_logger(__name__)
class    _SCREAMING_SNAKE_CASE							(		SCREAMING_SNAKE_CASE	):


			'''simple docstring'''





			__a							:							Optional[Any]					     =					["pixel_values"]






			def __init__(   self							:						str ,      lowercase							:						bool = True ,      lowercase							:						int = 3_2 ,      lowercase							:						List[Any]=PILImageResampling.BILINEAR ,      lowercase							:						bool = True ,      **lowercase							:						str ,      )							->    None:

									'''simple docstring'''







									UpperCamelCase__         = do_resize
									UpperCamelCase__         = do_rescale
									UpperCamelCase__         = size_divisor
									UpperCamelCase__         = resample
									super().__init__(**lowercase			)






			def 				A      (   self							:						Optional[int] ,      lowercase							:						np.ndarray ,      lowercase							:						int ,      lowercase							:						List[Any] ,      lowercase							:						Optional[ChannelDimension] = None ,      **lowercase							:						Any			)							->    np.ndarray:

									'''simple docstring'''







									UpperCamelCase__    ,   UpperCamelCase__         = get_image_size(lowercase			)
									# Rounds the height and width down to the closest multiple of size_divisor
									UpperCamelCase__         = height // size_divisor * size_divisor
									UpperCamelCase__         = width // size_divisor * size_divisor
									UpperCamelCase__         = resize(lowercase ,      (new_h, new_w) ,      resample=lowercase ,      data_format=lowercase ,      **lowercase			)
									return image






			def 				A      (   self							:						int ,      lowercase							:						np.ndarray ,      lowercase							:						float ,      lowercase							:						Optional[ChannelDimension] = None ,      **lowercase							:						List[str]			)							->    np.ndarray:

									'''simple docstring'''







									return rescale(image=lowercase ,      scale=lowercase ,      data_format=lowercase ,      **lowercase			)







			def 				A      (   self							:						int ,      lowercase							:						Union["PIL.Image.Image", TensorType, List["PIL.Image.Image"], List[TensorType]] ,      lowercase							:						Optional[bool] = None ,      lowercase							:						Optional[int] = None ,      lowercase							:						Optional[Any]=None ,      lowercase							:						Optional[bool] = None ,      lowercase							:						Optional[Union[TensorType, str]] = None ,      lowercase							:						ChannelDimension = ChannelDimension.FIRST ,      **lowercase							:						str ,      )							->    BatchFeature:

									'''simple docstring'''







									UpperCamelCase__         = do_resize if do_resize is not None else self.do_resize
									UpperCamelCase__         = do_rescale if do_rescale is not None else self.do_rescale
									UpperCamelCase__         = size_divisor if size_divisor is not None else self.size_divisor
									UpperCamelCase__         = resample if resample is not None else self.resample

									if do_resize and size_divisor is None:
															raise ValueError("""size_divisor is required for resizing"""			)

									UpperCamelCase__         = make_list_of_images(lowercase			)

									if not valid_images(lowercase			):
															raise ValueError("""Invalid image(s)"""			)

									# All transformations expect numpy arrays.
									UpperCamelCase__         = [to_numpy_array(lowercase			) for img in images]

									if do_resize:
															UpperCamelCase__         = [self.resize(lowercase ,      size_divisor=lowercase ,      resample=lowercase			) for image in images]

									if do_rescale:
															UpperCamelCase__         = [self.rescale(lowercase ,      scale=1 / 2_5_5			) for image in images]

									UpperCamelCase__         = [to_channel_dimension_format(lowercase ,      lowercase			) for image in images]

									UpperCamelCase__         = {"""pixel_values""": images}
									return BatchFeature(data=lowercase ,      tensor_type=lowercase			)
