


'''simple docstring'''



import json
import pathlib
import unittest

import numpy as np

from transformers.testing_utils import require_torch, require_vision, slow
from transformers.utils import is_torch_available, is_vision_available

from ...test_image_processing_common import ImageProcessingSavingTestMixin, prepare_image_inputs


if is_torch_available():
      import torch

if is_vision_available():
      from PIL import Image

      from transformers import DetaImageProcessor





class   UpperCAmelCase_    (unittest.TestCase    ):





    """simple docstring"""






    def __init__(				self     ,	SCREAMING_SNAKE_CASE_     ,	SCREAMING_SNAKE_CASE_=7     ,	SCREAMING_SNAKE_CASE_=3     ,	SCREAMING_SNAKE_CASE_=30     ,	SCREAMING_SNAKE_CASE_=4_00     ,	SCREAMING_SNAKE_CASE_=True     ,	SCREAMING_SNAKE_CASE_=None     ,	SCREAMING_SNAKE_CASE_=True     ,	SCREAMING_SNAKE_CASE_=[0.5, 0.5, 0.5]     ,	SCREAMING_SNAKE_CASE_=[0.5, 0.5, 0.5]     ,	SCREAMING_SNAKE_CASE_=True     ,	SCREAMING_SNAKE_CASE_=1 / 2_55     ,	SCREAMING_SNAKE_CASE_=True     ,	)							->	int:
         # by setting size["longest_edge"] > max_resolution we're effectively not testing this :p
         __lowerCamelCase  :      int														=						size if size is not None else {'shortest_edge': 18, 'longest_edge': 13_33}
         __lowerCamelCase  :      Optional[int]														=						parent
         __lowerCamelCase  :      Union[str, Any]														=						batch_size
         __lowerCamelCase  :      str														=						num_channels
         __lowerCamelCase  :      Union[str, Any]														=						min_resolution
         __lowerCamelCase  :      int														=						max_resolution
         __lowerCamelCase  :      str														=						do_resize
         __lowerCamelCase  :      Tuple														=						size
         __lowerCamelCase  :      str														=						do_normalize
         __lowerCamelCase  :      Optional[int]														=						image_mean
         __lowerCamelCase  :      int														=						image_std
         __lowerCamelCase  :      Union[str, Any]														=						do_rescale
         __lowerCamelCase  :      Optional[Any]														=						rescale_factor
         __lowerCamelCase  :      Optional[int]														=						do_pad







    def  lowercase_							(				self      )							->	str:
         return {
             "do_resize": self.do_resize,
             "size": self.size,
             "do_normalize": self.do_normalize,
             "image_mean": self.image_mean,
             "image_std": self.image_std,
             "do_rescale": self.do_rescale,
             "rescale_factor": self.rescale_factor,
             "do_pad": self.do_pad,
         }







    def  lowercase_							(				self     ,	SCREAMING_SNAKE_CASE_     ,	SCREAMING_SNAKE_CASE_=False      )							->	Any:
         if not batched:
              __lowerCamelCase  :      str														=						image_inputs[0]
              if isinstance(SCREAMING_SNAKE_CASE_     ,	Image.Image      ):
                   __lowerCamelCase							,			__lowerCamelCase  :      Union[str, Any]														=						image.size
              else:
                   __lowerCamelCase							,			__lowerCamelCase  :      Optional[int]														=						image.shape[1], image.shape[2]
              if w < h:
                   __lowerCamelCase  :      List[Any]														=						int(self.size['shortest_edge'] * h / w      )
                   __lowerCamelCase  :      str														=						self.size['shortest_edge']
              elif w > h:
                   __lowerCamelCase  :      Optional[int]														=						self.size['shortest_edge']
                   __lowerCamelCase  :      List[str]														=						int(self.size['shortest_edge'] * w / h      )
              else:
                   __lowerCamelCase  :      Any														=						self.size['shortest_edge']
                   __lowerCamelCase  :      List[str]														=						self.size['shortest_edge']

         else:
              __lowerCamelCase  :      Optional[Any]														=						[]
              for image in image_inputs:
                   __lowerCamelCase							,			__lowerCamelCase  :      Dict														=						self.get_expected_values([image]      )
                   expected_values.append((expected_height, expected_width)      )
              __lowerCamelCase  :      Any														=						max(SCREAMING_SNAKE_CASE_     ,	key=lambda SCREAMING_SNAKE_CASE_      : item[0]      )[0]
              __lowerCamelCase  :      Any														=						max(SCREAMING_SNAKE_CASE_     ,	key=lambda SCREAMING_SNAKE_CASE_      : item[1]      )[1]

         return expected_height, expected_width







@require_torch
@require_vision
class   UpperCAmelCase_    (_UpperCAmelCase   ,						unittest.TestCase    ):





    """simple docstring"""






    lowerCamelCase    :   Union[str, Any]					=     DetaImageProcessor if is_vision_available() else None







    def  lowercase_							(				self      )							->	str:
         __lowerCamelCase  :      Dict														=						DetaImageProcessingTester(self      )







    @property
    def  lowercase_							(				self      )							->	Tuple:
         return self.image_processor_tester.prepare_image_processor_dict()







    def  lowercase_							(				self      )							->	str:
         __lowerCamelCase  :      Dict														=						self.image_processing_class(**self.image_processor_dict      )
         self.assertTrue(hasattr(SCREAMING_SNAKE_CASE_     ,	'image_mean'      )      )
         self.assertTrue(hasattr(SCREAMING_SNAKE_CASE_     ,	'image_std'      )      )
         self.assertTrue(hasattr(SCREAMING_SNAKE_CASE_     ,	'do_normalize'      )      )
         self.assertTrue(hasattr(SCREAMING_SNAKE_CASE_     ,	'do_resize'      )      )
         self.assertTrue(hasattr(SCREAMING_SNAKE_CASE_     ,	'do_rescale'      )      )
         self.assertTrue(hasattr(SCREAMING_SNAKE_CASE_     ,	'do_pad'      )      )
         self.assertTrue(hasattr(SCREAMING_SNAKE_CASE_     ,	'size'      )      )







    def  lowercase_							(				self      )							->	int:
         __lowerCamelCase  :      Dict														=						self.image_processing_class.from_dict(self.image_processor_dict      )
         self.assertEqual(image_processor.size     ,	{'shortest_edge': 18, 'longest_edge': 13_33}      )
         self.assertEqual(image_processor.do_pad     ,	SCREAMING_SNAKE_CASE_      )







    def  lowercase_							(				self      )							->	List[str]:
         pass







    def  lowercase_							(				self      )							->	Union[str, Any]:
         # Initialize image_processing
         __lowerCamelCase  :      List[Any]														=						self.image_processing_class(**self.image_processor_dict      )
         # create random PIL images
         __lowerCamelCase  :      Dict														=						prepare_image_inputs(self.image_processor_tester     ,	equal_resolution=SCREAMING_SNAKE_CASE_      )
         for image in image_inputs:
              self.assertIsInstance(SCREAMING_SNAKE_CASE_     ,	Image.Image      )

         # Test not batched input
         __lowerCamelCase  :      Dict														=						image_processing(image_inputs[0]     ,	return_tensors='pt'      ).pixel_values

         __lowerCamelCase							,			__lowerCamelCase  :      List[str]														=						self.image_processor_tester.get_expected_values(SCREAMING_SNAKE_CASE_      )

         self.assertEqual(
             encoded_images.shape     ,	(1, self.image_processor_tester.num_channels, expected_height, expected_width)     ,	)

         # Test batched
         __lowerCamelCase							,			__lowerCamelCase  :      Optional[int]														=						self.image_processor_tester.get_expected_values(SCREAMING_SNAKE_CASE_     ,	batched=SCREAMING_SNAKE_CASE_      )

         __lowerCamelCase  :      str														=						image_processing(SCREAMING_SNAKE_CASE_     ,	return_tensors='pt'      ).pixel_values
         self.assertEqual(
             encoded_images.shape     ,	(
                 self.image_processor_tester.batch_size,
                 self.image_processor_tester.num_channels,
                 expected_height,
                 expected_width,
             )     ,	)







    def  lowercase_							(				self      )							->	Tuple:
         # Initialize image_processing
         __lowerCamelCase  :      List[Any]														=						self.image_processing_class(**self.image_processor_dict      )
         # create random numpy tensors
         __lowerCamelCase  :      Union[str, Any]														=						prepare_image_inputs(self.image_processor_tester     ,	equal_resolution=SCREAMING_SNAKE_CASE_     ,	numpify=SCREAMING_SNAKE_CASE_      )
         for image in image_inputs:
              self.assertIsInstance(SCREAMING_SNAKE_CASE_     ,	np.ndarray      )

         # Test not batched input
         __lowerCamelCase  :      List[Any]														=						image_processing(image_inputs[0]     ,	return_tensors='pt'      ).pixel_values

         __lowerCamelCase							,			__lowerCamelCase  :      int														=						self.image_processor_tester.get_expected_values(SCREAMING_SNAKE_CASE_      )

         self.assertEqual(
             encoded_images.shape     ,	(1, self.image_processor_tester.num_channels, expected_height, expected_width)     ,	)

         # Test batched
         __lowerCamelCase  :      Tuple														=						image_processing(SCREAMING_SNAKE_CASE_     ,	return_tensors='pt'      ).pixel_values

         __lowerCamelCase							,			__lowerCamelCase  :      List[str]														=						self.image_processor_tester.get_expected_values(SCREAMING_SNAKE_CASE_     ,	batched=SCREAMING_SNAKE_CASE_      )

         self.assertEqual(
             encoded_images.shape     ,	(
                 self.image_processor_tester.batch_size,
                 self.image_processor_tester.num_channels,
                 expected_height,
                 expected_width,
             )     ,	)







    def  lowercase_							(				self      )							->	Optional[Any]:
         # Initialize image_processing
         __lowerCamelCase  :      List[Any]														=						self.image_processing_class(**self.image_processor_dict      )
         # create random PyTorch tensors
         __lowerCamelCase  :      Optional[Any]														=						prepare_image_inputs(self.image_processor_tester     ,	equal_resolution=SCREAMING_SNAKE_CASE_     ,	torchify=SCREAMING_SNAKE_CASE_      )
         for image in image_inputs:
              self.assertIsInstance(SCREAMING_SNAKE_CASE_     ,	torch.Tensor      )

         # Test not batched input
         __lowerCamelCase  :      Optional[Any]														=						image_processing(image_inputs[0]     ,	return_tensors='pt'      ).pixel_values

         __lowerCamelCase							,			__lowerCamelCase  :      Optional[Any]														=						self.image_processor_tester.get_expected_values(SCREAMING_SNAKE_CASE_      )

         self.assertEqual(
             encoded_images.shape     ,	(1, self.image_processor_tester.num_channels, expected_height, expected_width)     ,	)

         # Test batched
         __lowerCamelCase  :      Optional[Any]														=						image_processing(SCREAMING_SNAKE_CASE_     ,	return_tensors='pt'      ).pixel_values

         __lowerCamelCase							,			__lowerCamelCase  :      str														=						self.image_processor_tester.get_expected_values(SCREAMING_SNAKE_CASE_     ,	batched=SCREAMING_SNAKE_CASE_      )

         self.assertEqual(
             encoded_images.shape     ,	(
                 self.image_processor_tester.batch_size,
                 self.image_processor_tester.num_channels,
                 expected_height,
                 expected_width,
             )     ,	)







    @slow
    def  lowercase_							(				self      )							->	List[Any]:
         # prepare image and target
         __lowerCamelCase  :      Dict														=						Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'      )
         with open('./tests/fixtures/tests_samples/COCO/coco_annotations.txt'     ,	'r'      ) as f:
              __lowerCamelCase  :      List[Any]														=						json.loads(f.read()      )

         __lowerCamelCase  :      List[str]														=						{'image_id': 3_97_69, 'annotations': target}

         # encode them
         __lowerCamelCase  :      Union[str, Any]														=						DetaImageProcessor()
         __lowerCamelCase  :      int														=						image_processing(images=SCREAMING_SNAKE_CASE_     ,	annotations=SCREAMING_SNAKE_CASE_     ,	return_tensors='pt'      )

         # verify pixel values
         __lowerCamelCase  :      Optional[int]														=						torch.Size([1, 3, 8_00, 10_66]      )
         self.assertEqual(encoding['pixel_values'].shape     ,	SCREAMING_SNAKE_CASE_      )

         __lowerCamelCase  :      int														=						torch.tensor([0.2_7_9_6, 0.3_1_3_8, 0.3_4_8_1]      )
         self.assertTrue(torch.allclose(encoding['pixel_values'][0, 0, 0, :3]     ,	SCREAMING_SNAKE_CASE_     ,	atol=1E-4      )      )

         # verify area
         __lowerCamelCase  :      Optional[Any]														=						torch.tensor([5_8_8_7.9_6_0_0, 1_1_2_5_0.2_0_6_1, 4_8_9_3_5_3.8_4_3_8, 8_3_7_1_2_2.7_5_0_0, 1_4_7_9_6_7.5_1_5_6, 1_6_5_7_3_2.3_4_3_8]      )
         self.assertTrue(torch.allclose(encoding['labels'][0]['area']     ,	SCREAMING_SNAKE_CASE_      )      )
         # verify boxes
         __lowerCamelCase  :      Optional[Any]														=						torch.Size([6, 4]      )
         self.assertEqual(encoding['labels'][0]['boxes'].shape     ,	SCREAMING_SNAKE_CASE_      )
         __lowerCamelCase  :      int														=						torch.tensor([0.5_5_0_3, 0.2_7_6_5, 0.0_6_0_4, 0.2_2_1_5]      )
         self.assertTrue(torch.allclose(encoding['labels'][0]['boxes'][0]     ,	SCREAMING_SNAKE_CASE_     ,	atol=1E-3      )      )
         # verify image_id
         __lowerCamelCase  :      Union[str, Any]														=						torch.tensor([3_97_69]      )
         self.assertTrue(torch.allclose(encoding['labels'][0]['image_id']     ,	SCREAMING_SNAKE_CASE_      )      )
         # verify is_crowd
         __lowerCamelCase  :      List[Any]														=						torch.tensor([0, 0, 0, 0, 0, 0]      )
         self.assertTrue(torch.allclose(encoding['labels'][0]['iscrowd']     ,	SCREAMING_SNAKE_CASE_      )      )
         # verify class_labels
         __lowerCamelCase  :      Optional[int]														=						torch.tensor([75, 75, 63, 65, 17, 17]      )
         self.assertTrue(torch.allclose(encoding['labels'][0]['class_labels']     ,	SCREAMING_SNAKE_CASE_      )      )
         # verify orig_size
         __lowerCamelCase  :      List[Any]														=						torch.tensor([4_80, 6_40]      )
         self.assertTrue(torch.allclose(encoding['labels'][0]['orig_size']     ,	SCREAMING_SNAKE_CASE_      )      )
         # verify size
         __lowerCamelCase  :      Optional[Any]														=						torch.tensor([8_00, 10_66]      )
         self.assertTrue(torch.allclose(encoding['labels'][0]['size']     ,	SCREAMING_SNAKE_CASE_      )      )







    @slow
    def  lowercase_							(				self      )							->	Optional[Any]:
         # prepare image, target and masks_path
         __lowerCamelCase  :      List[str]														=						Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'      )
         with open('./tests/fixtures/tests_samples/COCO/coco_panoptic_annotations.txt'     ,	'r'      ) as f:
              __lowerCamelCase  :      Any														=						json.loads(f.read()      )

         __lowerCamelCase  :      int														=						{'file_name': '000000039769.png', 'image_id': 3_97_69, 'segments_info': target}

         __lowerCamelCase  :      List[Any]														=						pathlib.Path('./tests/fixtures/tests_samples/COCO/coco_panoptic'      )

         # encode them
         __lowerCamelCase  :      Union[str, Any]														=						DetaImageProcessor(format='coco_panoptic'      )
         __lowerCamelCase  :      Optional[Any]														=						image_processing(images=SCREAMING_SNAKE_CASE_     ,	annotations=SCREAMING_SNAKE_CASE_     ,	masks_path=SCREAMING_SNAKE_CASE_     ,	return_tensors='pt'      )

         # verify pixel values
         __lowerCamelCase  :      List[str]														=						torch.Size([1, 3, 8_00, 10_66]      )
         self.assertEqual(encoding['pixel_values'].shape     ,	SCREAMING_SNAKE_CASE_      )

         __lowerCamelCase  :      int														=						torch.tensor([0.2_7_9_6, 0.3_1_3_8, 0.3_4_8_1]      )
         self.assertTrue(torch.allclose(encoding['pixel_values'][0, 0, 0, :3]     ,	SCREAMING_SNAKE_CASE_     ,	atol=1E-4      )      )

         # verify area
         __lowerCamelCase  :      Optional[Any]														=						torch.tensor([1_4_7_9_7_9.6_8_7_5, 1_6_5_5_2_7.0_4_6_9, 4_8_4_6_3_8.5_9_3_8, 1_1_2_9_2.9_3_7_5, 5_8_7_9.6_5_6_2, 7_6_3_4.1_1_4_7]      )
         self.assertTrue(torch.allclose(encoding['labels'][0]['area']     ,	SCREAMING_SNAKE_CASE_      )      )
         # verify boxes
         __lowerCamelCase  :      int														=						torch.Size([6, 4]      )
         self.assertEqual(encoding['labels'][0]['boxes'].shape     ,	SCREAMING_SNAKE_CASE_      )
         __lowerCamelCase  :      Optional[Any]														=						torch.tensor([0.2_6_2_5, 0.5_4_3_7, 0.4_6_8_8, 0.8_6_2_5]      )
         self.assertTrue(torch.allclose(encoding['labels'][0]['boxes'][0]     ,	SCREAMING_SNAKE_CASE_     ,	atol=1E-3      )      )
         # verify image_id
         __lowerCamelCase  :      Optional[int]														=						torch.tensor([3_97_69]      )
         self.assertTrue(torch.allclose(encoding['labels'][0]['image_id']     ,	SCREAMING_SNAKE_CASE_      )      )
         # verify is_crowd
         __lowerCamelCase  :      Union[str, Any]														=						torch.tensor([0, 0, 0, 0, 0, 0]      )
         self.assertTrue(torch.allclose(encoding['labels'][0]['iscrowd']     ,	SCREAMING_SNAKE_CASE_      )      )
         # verify class_labels
         __lowerCamelCase  :      Optional[Any]														=						torch.tensor([17, 17, 63, 75, 75, 93]      )
         self.assertTrue(torch.allclose(encoding['labels'][0]['class_labels']     ,	SCREAMING_SNAKE_CASE_      )      )
         # verify masks
         __lowerCamelCase  :      Any														=						82_28_73
         self.assertEqual(encoding['labels'][0]['masks'].sum().item()     ,	SCREAMING_SNAKE_CASE_      )
         # verify orig_size
         __lowerCamelCase  :      Tuple														=						torch.tensor([4_80, 6_40]      )
         self.assertTrue(torch.allclose(encoding['labels'][0]['orig_size']     ,	SCREAMING_SNAKE_CASE_      )      )
         # verify size
         __lowerCamelCase  :      Tuple														=						torch.tensor([8_00, 10_66]      )
         self.assertTrue(torch.allclose(encoding['labels'][0]['size']     ,	SCREAMING_SNAKE_CASE_      )      )










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



					'''simple docstring'''







					def __init__(					self			,							snake_case_			,							snake_case_=1_3			,							snake_case_=7			,							snake_case_=True			,							snake_case_=True			,							snake_case_=True			,							snake_case_=True			,							snake_case_=9_9			,							snake_case_=6_4			,							snake_case_=3_2			,							snake_case_=5			,							snake_case_=4			,							snake_case_=3_7			,							snake_case_="gelu"			,							snake_case_=0.1			,							snake_case_=0.1			,							snake_case_=5_1_2			,							snake_case_=1_6			,							snake_case_=2			,							snake_case_=0.0_2			,							snake_case_=3			,							snake_case_=4			,							snake_case_=None			,							) ->      Tuple:







								'''simple docstring'''

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






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







								'''simple docstring'''

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

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

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

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

								__lowercase			       = self.get_config()

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






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







								'''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=snake_case_			,							initializer_range=self.initializer_range			,							)






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







								'''simple docstring'''

								__lowercase			       = MegatronBertModel(config=snake_case_    )
								model.to(snake_case_    )
								model.eval()
								__lowercase			       = model(snake_case_			,							attention_mask=snake_case_			,							token_type_ids=snake_case_    )
								__lowercase			       = model(snake_case_			,							token_type_ids=snake_case_    )
								__lowercase			       = model(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 A   (					self			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_    ) ->      List[str]:







								'''simple docstring'''

								__lowercase			       = MegatronBertForMaskedLM(config=snake_case_    )
								model.to(snake_case_    )
								model.eval()
								__lowercase			       = model(snake_case_			,							attention_mask=snake_case_			,							token_type_ids=snake_case_			,							labels=snake_case_    )
								self.parent.assertEqual(result.logits.shape			,							(self.batch_size, self.seq_length, self.vocab_size)    )






					def A   (					self			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_    ) ->      int:







								'''simple docstring'''

								__lowercase			       = MegatronBertForCausalLM(config=snake_case_    )
								model.to(snake_case_    )
								model.eval()
								__lowercase			       = model(snake_case_			,							attention_mask=snake_case_			,							token_type_ids=snake_case_			,							labels=snake_case_    )
								self.parent.assertEqual(result.logits.shape			,							(self.batch_size, self.seq_length, self.vocab_size)    )






					def A   (					self			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_    ) ->      Tuple:







								'''simple docstring'''

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






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







								'''simple docstring'''

								__lowercase			       = MegatronBertForPreTraining(config=snake_case_    )
								model.to(snake_case_    )
								model.eval()
								__lowercase			       = model(
								    snake_case_			,							attention_mask=snake_case_			,							token_type_ids=snake_case_			,							labels=snake_case_			,							next_sentence_label=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 A   (					self			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_    ) ->      Any:







								'''simple docstring'''

								__lowercase			       = MegatronBertForQuestionAnswering(config=snake_case_    )
								model.to(snake_case_    )
								model.eval()
								__lowercase			       = model(
								    snake_case_			,							attention_mask=snake_case_			,							token_type_ids=snake_case_			,							start_positions=snake_case_			,							end_positions=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 A   (					self			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_    ) ->      Optional[int]:







								'''simple docstring'''

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






					def A   (					self			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_    ) ->      Dict:







								'''simple docstring'''

								__lowercase			       = self.num_labels
								__lowercase			       = MegatronBertForTokenClassification(config=snake_case_    )
								model.to(snake_case_    )
								model.eval()
								__lowercase			       = model(snake_case_			,							attention_mask=snake_case_			,							token_type_ids=snake_case_			,							labels=snake_case_    )
								self.parent.assertEqual(result.logits.shape			,							(self.batch_size, self.seq_length, self.num_labels)    )






					def A   (					self			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_    ) ->      Dict:







								'''simple docstring'''

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






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







								'''simple docstring'''

								__lowercase			       = self.prepare_config_and_inputs()
								(
								    (
								    __lowercase
								)       ,			(
								    __lowercase
								)       ,			(
								    __lowercase
								)       ,			(
								    __lowercase
								)       ,			(
								    __lowercase
								)       ,			(
								    __lowercase
								)       ,			(
								    __lowercase
								)       ,			
								)			       = config_and_inputs
								__lowercase			       = {'''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					):



					'''simple docstring'''







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






					def A   (					self			,							snake_case_			,							snake_case_			,							snake_case_=False    ) ->      Tuple:







								'''simple docstring'''

								__lowercase			       = super()._prepare_for_class(snake_case_			,							snake_case_			,							return_labels=snake_case_    )

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






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







								'''simple docstring'''

								__lowercase			       = MegatronBertModelTester(self    )
								__lowercase			       = ConfigTester(self			,							config_class=snake_case_			,							hidden_size=3_7    )






					def A   (					self    ) ->      int:







								'''simple docstring'''

								self.config_tester.run_common_tests()






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







								'''simple docstring'''

								__lowercase			       = self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_megatron_bert_model(*snake_case_    )






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







								'''simple docstring'''

								__lowercase			       = self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_megatron_bert_for_masked_lm(*snake_case_    )






					def A   (					self    ) ->      str:







								'''simple docstring'''

								__lowercase			       = self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*snake_case_    )






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







								'''simple docstring'''

								__lowercase			       = self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*snake_case_    )






					def A   (					self    ) ->      Dict:







								'''simple docstring'''

								__lowercase			       = self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_megatron_bert_for_pretraining(*snake_case_    )






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







								'''simple docstring'''

								__lowercase			       = self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_megatron_bert_for_question_answering(*snake_case_    )






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







								'''simple docstring'''

								__lowercase			       = self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*snake_case_    )






					def A   (					self    ) ->      Dict:







								'''simple docstring'''

								__lowercase			       = self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_megatron_bert_for_token_classification(*snake_case_    )




def 		lowercase_	(  _UpperCamelCase	):


			'''simple docstring'''







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


a			:			Optional[Any]												=					1E-4





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



					'''simple docstring'''







					@slow
					@unittest.skip('''Model is not available.'''    )
					def A   (					self    ) ->      Tuple:







								'''simple docstring'''

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

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


