





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     :					Dict,			_UpperCAmelCase     :					List[Any],			_UpperCAmelCase     :					Dict=1_3,			_UpperCAmelCase     :					List[str]=7,			_UpperCAmelCase     :					List[Any]=True,			_UpperCAmelCase     :					List[Any]=True,			_UpperCAmelCase     :					Optional[Any]=True,			_UpperCAmelCase     :					Union[str, Any]=True,			_UpperCAmelCase     :					List[Any]=9_9,			_UpperCAmelCase     :					Any=6_4,			_UpperCAmelCase     :					Optional[int]=3_2,			_UpperCAmelCase     :					int=5,			_UpperCAmelCase     :					Optional[int]=4,			_UpperCAmelCase     :					Tuple=3_7,			_UpperCAmelCase     :					Optional[int]="gelu",			_UpperCAmelCase     :					List[str]=0.1,			_UpperCAmelCase     :					Dict=0.1,			_UpperCAmelCase     :					Optional[Any]=5_1_2,			_UpperCAmelCase     :					Union[str, Any]=1_6,			_UpperCAmelCase     :					Optional[int]=2,			_UpperCAmelCase     :					str=0.02,			_UpperCAmelCase     :					Union[str, Any]=3,			_UpperCAmelCase     :					Optional[int]=4,			_UpperCAmelCase     :					Union[str, Any]=None,			)				->			Any:





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      Tuple     =						parent
														SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						batch_size
														SCREAMING_SNAKE_CASE__       :      str     =						seq_length
														SCREAMING_SNAKE_CASE__       :      str     =						is_training
														SCREAMING_SNAKE_CASE__       :      List[str]     =						use_input_mask
														SCREAMING_SNAKE_CASE__       :      Any     =						use_token_type_ids
														SCREAMING_SNAKE_CASE__       :      Optional[int]     =						use_labels
														SCREAMING_SNAKE_CASE__       :      List[Any]     =						vocab_size
														SCREAMING_SNAKE_CASE__       :      Any     =						hidden_size
														SCREAMING_SNAKE_CASE__       :      Optional[int]     =						embedding_size
														SCREAMING_SNAKE_CASE__       :      int     =						num_hidden_layers
														SCREAMING_SNAKE_CASE__       :      int     =						num_attention_heads
														SCREAMING_SNAKE_CASE__       :      List[Any]     =						intermediate_size
														SCREAMING_SNAKE_CASE__       :      List[Any]     =						hidden_act
														SCREAMING_SNAKE_CASE__       :      List[Any]     =						hidden_dropout_prob
														SCREAMING_SNAKE_CASE__       :      Dict     =						attention_probs_dropout_prob
														SCREAMING_SNAKE_CASE__       :      Dict     =						max_position_embeddings
														SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						type_vocab_size
														SCREAMING_SNAKE_CASE__       :      List[str]     =						type_sequence_label_size
														SCREAMING_SNAKE_CASE__       :      Optional[Any]     =						initializer_range
														SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						num_labels
														SCREAMING_SNAKE_CASE__       :      int     =						num_choices
														SCREAMING_SNAKE_CASE__       :      Optional[int]     =						scope






							def  A_     ( self     :					Dict						)				->			Optional[Any]:





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      Optional[int]     =						ids_tensor([self.batch_size, self.seq_length],			self.vocab_size						)

														SCREAMING_SNAKE_CASE__       :      Optional[int]     =						None
														if self.use_input_mask:
																					SCREAMING_SNAKE_CASE__       :      List[Any]     =						random_attention_mask([self.batch_size, self.seq_length]						)

														SCREAMING_SNAKE_CASE__       :      List[str]     =						None
														if self.use_token_type_ids:
																					SCREAMING_SNAKE_CASE__       :      Dict     =						ids_tensor([self.batch_size, self.seq_length],			self.type_vocab_size						)

														SCREAMING_SNAKE_CASE__       :      Any     =						None
														SCREAMING_SNAKE_CASE__       :      List[str]     =						None
														SCREAMING_SNAKE_CASE__       :      str     =						None
														if self.use_labels:
																					SCREAMING_SNAKE_CASE__       :      List[str]     =						ids_tensor([self.batch_size],			self.type_sequence_label_size						)
																					SCREAMING_SNAKE_CASE__       :      List[str]     =						ids_tensor([self.batch_size, self.seq_length],			self.num_labels						)
																					SCREAMING_SNAKE_CASE__       :      Optional[int]     =						ids_tensor([self.batch_size],			self.num_choices						)

														SCREAMING_SNAKE_CASE__       :      Optional[Any]     =						self.get_config()

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






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





														"""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=__SCREAMING_SNAKE_CASE,			initializer_range=self.initializer_range,			)






							def  A_     ( self     :					List[Any],			_UpperCAmelCase     :					List[str],			_UpperCAmelCase     :					Dict,			_UpperCAmelCase     :					int,			_UpperCAmelCase     :					Tuple,			_UpperCAmelCase     :					Optional[Any],			_UpperCAmelCase     :					List[Any],			_UpperCAmelCase     :					Tuple						)				->			Tuple:





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      Any     =						MegatronBertModel(config=__SCREAMING_SNAKE_CASE						)
														model.to(__SCREAMING_SNAKE_CASE						)
														model.eval()
														SCREAMING_SNAKE_CASE__       :      Any     =						model(__SCREAMING_SNAKE_CASE,			attention_mask=__SCREAMING_SNAKE_CASE,			token_type_ids=__SCREAMING_SNAKE_CASE						)
														SCREAMING_SNAKE_CASE__       :      int     =						model(__SCREAMING_SNAKE_CASE,			token_type_ids=__SCREAMING_SNAKE_CASE						)
														SCREAMING_SNAKE_CASE__       :      Optional[Any]     =						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  A_     ( self     :					List[str],			_UpperCAmelCase     :					List[Any],			_UpperCAmelCase     :					int,			_UpperCAmelCase     :					Optional[Any],			_UpperCAmelCase     :					List[Any],			_UpperCAmelCase     :					str,			_UpperCAmelCase     :					Any,			_UpperCAmelCase     :					Optional[Any]						)				->			Union[str, Any]:





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      Tuple     =						MegatronBertForMaskedLM(config=__SCREAMING_SNAKE_CASE						)
														model.to(__SCREAMING_SNAKE_CASE						)
														model.eval()
														SCREAMING_SNAKE_CASE__       :      int     =						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  A_     ( self     :					List[str],			_UpperCAmelCase     :					str,			_UpperCAmelCase     :					List[str],			_UpperCAmelCase     :					Optional[int],			_UpperCAmelCase     :					List[str],			_UpperCAmelCase     :					Union[str, Any],			_UpperCAmelCase     :					Dict,			_UpperCAmelCase     :					str						)				->			int:





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      Any     =						MegatronBertForCausalLM(config=__SCREAMING_SNAKE_CASE						)
														model.to(__SCREAMING_SNAKE_CASE						)
														model.eval()
														SCREAMING_SNAKE_CASE__       :      Optional[Any]     =						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  A_     ( self     :					Dict,			_UpperCAmelCase     :					str,			_UpperCAmelCase     :					List[Any],			_UpperCAmelCase     :					List[str],			_UpperCAmelCase     :					List[Any],			_UpperCAmelCase     :					int,			_UpperCAmelCase     :					List[str],			_UpperCAmelCase     :					Optional[Any]						)				->			Dict:





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      Dict     =						MegatronBertForNextSentencePrediction(config=__SCREAMING_SNAKE_CASE						)
														model.to(__SCREAMING_SNAKE_CASE						)
														model.eval()
														SCREAMING_SNAKE_CASE__       :      Tuple     =						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  A_     ( self     :					Optional[int],			_UpperCAmelCase     :					List[str],			_UpperCAmelCase     :					str,			_UpperCAmelCase     :					int,			_UpperCAmelCase     :					int,			_UpperCAmelCase     :					Any,			_UpperCAmelCase     :					Optional[int],			_UpperCAmelCase     :					str						)				->			str:





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      Tuple     =						MegatronBertForPreTraining(config=__SCREAMING_SNAKE_CASE						)
														model.to(__SCREAMING_SNAKE_CASE						)
														model.eval()
														SCREAMING_SNAKE_CASE__       :      Any     =						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  A_     ( self     :					Union[str, Any],			_UpperCAmelCase     :					List[Any],			_UpperCAmelCase     :					Union[str, Any],			_UpperCAmelCase     :					List[str],			_UpperCAmelCase     :					List[str],			_UpperCAmelCase     :					Union[str, Any],			_UpperCAmelCase     :					Tuple,			_UpperCAmelCase     :					Union[str, Any]						)				->			str:





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      Optional[int]     =						MegatronBertForQuestionAnswering(config=__SCREAMING_SNAKE_CASE						)
														model.to(__SCREAMING_SNAKE_CASE						)
														model.eval()
														SCREAMING_SNAKE_CASE__       :      List[str]     =						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  A_     ( self     :					Tuple,			_UpperCAmelCase     :					str,			_UpperCAmelCase     :					Optional[Any],			_UpperCAmelCase     :					Optional[Any],			_UpperCAmelCase     :					Optional[int],			_UpperCAmelCase     :					Tuple,			_UpperCAmelCase     :					Optional[int],			_UpperCAmelCase     :					Optional[Any]						)				->			Optional[int]:





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      Optional[int]     =						self.num_labels
														SCREAMING_SNAKE_CASE__       :      Any     =						MegatronBertForSequenceClassification(__SCREAMING_SNAKE_CASE						)
														model.to(__SCREAMING_SNAKE_CASE						)
														model.eval()
														SCREAMING_SNAKE_CASE__       :      Optional[int]     =						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  A_     ( self     :					Dict,			_UpperCAmelCase     :					Any,			_UpperCAmelCase     :					Optional[Any],			_UpperCAmelCase     :					Optional[Any],			_UpperCAmelCase     :					int,			_UpperCAmelCase     :					Optional[Any],			_UpperCAmelCase     :					Any,			_UpperCAmelCase     :					int						)				->			int:





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      Any     =						self.num_labels
														SCREAMING_SNAKE_CASE__       :      Tuple     =						MegatronBertForTokenClassification(config=__SCREAMING_SNAKE_CASE						)
														model.to(__SCREAMING_SNAKE_CASE						)
														model.eval()
														SCREAMING_SNAKE_CASE__       :      List[str]     =						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  A_     ( self     :					Tuple,			_UpperCAmelCase     :					Any,			_UpperCAmelCase     :					Optional[Any],			_UpperCAmelCase     :					Union[str, Any],			_UpperCAmelCase     :					str,			_UpperCAmelCase     :					Optional[Any],			_UpperCAmelCase     :					str,			_UpperCAmelCase     :					Dict						)				->			Any:





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      Any     =						self.num_choices
														SCREAMING_SNAKE_CASE__       :      Dict     =						MegatronBertForMultipleChoice(config=__SCREAMING_SNAKE_CASE						)
														model.to(__SCREAMING_SNAKE_CASE						)
														model.eval()
														SCREAMING_SNAKE_CASE__       :      Optional[int]     =						input_ids.unsqueeze(1						).expand(-1,			self.num_choices,			-1						).contiguous()
														SCREAMING_SNAKE_CASE__       :      Dict     =						token_type_ids.unsqueeze(1						).expand(-1,			self.num_choices,			-1						).contiguous()
														SCREAMING_SNAKE_CASE__       :      str     =						input_mask.unsqueeze(1						).expand(-1,			self.num_choices,			-1						).contiguous()
														SCREAMING_SNAKE_CASE__       :      Optional[Any]     =						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  A_     ( self     :					Union[str, Any]						)				->			Dict:





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      int     =						self.prepare_config_and_inputs()
														(
														    (
														    SCREAMING_SNAKE_CASE__	
														)   ,(
														    SCREAMING_SNAKE_CASE__	
														)   ,(
														    SCREAMING_SNAKE_CASE__	
														)   ,(
														    SCREAMING_SNAKE_CASE__	
														)   ,(
														    SCREAMING_SNAKE_CASE__	
														)   ,(
														    SCREAMING_SNAKE_CASE__	
														)   ,(
														    SCREAMING_SNAKE_CASE__	
														)   ,	
														)       :      List[str]     =						config_and_inputs
														SCREAMING_SNAKE_CASE__       :      Tuple     =						{"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     :					int,			_UpperCAmelCase     :					Any,			_UpperCAmelCase     :					int,			_UpperCAmelCase     :					Tuple=False						)				->			List[str]:





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      List[str]     =						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						):
																												SCREAMING_SNAKE_CASE__       :      Any     =						torch.zeros(
																												    (self.model_tester.batch_size, self.model_tester.seq_length),			dtype=torch.long,			device=__SCREAMING_SNAKE_CASE						)
																												SCREAMING_SNAKE_CASE__       :      Optional[int]     =						torch.zeros(
																												    self.model_tester.batch_size,			dtype=torch.long,			device=__SCREAMING_SNAKE_CASE						)
														return inputs_dict






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





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      Optional[Any]     =						MegatronBertModelTester(self						)
														SCREAMING_SNAKE_CASE__       :      Optional[Any]     =						ConfigTester(self,			config_class=__SCREAMING_SNAKE_CASE,			hidden_size=3_7						)






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





														"""simple docstring"""
														self.config_tester.run_common_tests()






							def  A_     ( self     :					int						)				->			int:





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      List[Any]     =						self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_megatron_bert_model(*__SCREAMING_SNAKE_CASE						)






							def  A_     ( self     :					Dict						)				->			Optional[int]:





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      str     =						self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_megatron_bert_for_masked_lm(*__SCREAMING_SNAKE_CASE						)






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





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      int     =						self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*__SCREAMING_SNAKE_CASE						)






							def  A_     ( self     :					Dict						)				->			Tuple:





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      Tuple     =						self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*__SCREAMING_SNAKE_CASE						)






							def  A_     ( self     :					int						)				->			Optional[Any]:





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_megatron_bert_for_pretraining(*__SCREAMING_SNAKE_CASE						)






							def  A_     ( self     :					List[str]						)				->			Any:





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      Optional[Any]     =						self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_megatron_bert_for_question_answering(*__SCREAMING_SNAKE_CASE						)






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





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      List[Any]     =						self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*__SCREAMING_SNAKE_CASE						)






							def  A_     ( self     :					Dict						)				->			Dict:





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      Any     =						self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_megatron_bert_for_token_classification(*__SCREAMING_SNAKE_CASE						)




def _a (  SCREAMING_SNAKE_CASE__      :  Union[str, Any]       )	->    Optional[Any]:

							'''simple docstring'''






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


_lowerCamelCase :    Tuple    							= 1e-4







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







							"""simple docstring"""




							@slow
							@unittest.skip("Model is not available."						)
							def  A_     ( self     :					int						)				->			Tuple:





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      List[Any]     =						"nvidia/megatron-bert-uncased-345m"
														if "MYDIR" in os.environ:
																					SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						os.path.join(os.environ["MYDIR"],			__SCREAMING_SNAKE_CASE						)
														SCREAMING_SNAKE_CASE__       :      Any     =						MegatronBertModel.from_pretrained(__SCREAMING_SNAKE_CASE						)
														model.to(__SCREAMING_SNAKE_CASE						)
														model.half()
														SCREAMING_SNAKE_CASE__       :      Dict     =						_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():
																					SCREAMING_SNAKE_CASE__       :      int     =						model(__SCREAMING_SNAKE_CASE						)[0]
														SCREAMING_SNAKE_CASE__       :      List[str]     =						torch.Size((1, 9, 1_0_2_4)						)
														self.assertEqual(output.shape,			__SCREAMING_SNAKE_CASE						)

														SCREAMING_SNAKE_CASE__       :      List[str]     =						[-0.6040, -0.2517, -0.1025, 0.3420, -0.6758, -0.0017, -0.1089, -0.1990, 0.5728]
														for ii in range(3						):
																					for jj in range(3						):
																												SCREAMING_SNAKE_CASE__       :      List[str]     =						output[0, ii, jj]
																												SCREAMING_SNAKE_CASE__       :      List[str]     =						expected[3 * ii + jj]
																												SCREAMING_SNAKE_CASE__       :      Optional[int]     =						"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 pytest

from datasets.splits import SplitDict, SplitInfo
from datasets.utils.py_utils import asdict




@pytest.mark.parametrize(
    "split_dict"			,				[
        SplitDict(),
        SplitDict({"train": SplitInfo(name="train"			,				num_bytes=13_37			,				num_examples=42			,				dataset_name="my_dataset"       )}       ),
        SplitDict({"train": SplitInfo(name="train"			,				num_bytes=13_37			,				num_examples=42       )}       ),
        SplitDict({"train": SplitInfo()}       ),
    ]			,				)
def _a (  SCREAMING_SNAKE_CASE__      :  SplitDict       )	->    str:

							'''simple docstring'''






							SCREAMING_SNAKE_CASE__       :      Dict     =						split_dict._to_yaml_list()
							assert len(SCREAMING_SNAKE_CASE__       ) == len(SCREAMING_SNAKE_CASE__       )
							SCREAMING_SNAKE_CASE__       :      List[Any]     =						SplitDict._from_yaml_list(SCREAMING_SNAKE_CASE__       )
							for split_name, split_info in split_dict.items():
														# dataset_name field is deprecated, and is therefore not part of the YAML dump
														SCREAMING_SNAKE_CASE__       :      str     =						None
														# the split name of split_dict takes over the name of the split info object
														SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						split_name
							assert split_dict == reloaded







@pytest.mark.parametrize(
    "split_info"			,				[SplitInfo(), SplitInfo(dataset_name=SCREAMING_SNAKE_CASE__       ), SplitInfo(dataset_name="my_dataset"       )]       )
def _a (  SCREAMING_SNAKE_CASE__      :  Any       )	->    Dict:

							'''simple docstring'''






							SCREAMING_SNAKE_CASE__       :      str     =						asdict(SplitDict({"train": split_info}       )       )
							assert "dataset_name" in split_dict_asdict["train"]
							assert split_dict_asdict["train"]["dataset_name"] == split_info.dataset_name






