


"""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      __A     :





							'''simple docstring'''

							def __init__(				self							:      Dict		,       UpperCAmelCase_							:      Any		,       UpperCAmelCase_							:      Dict=13		,       UpperCAmelCase_							:      Union[str, Any]=7		,       UpperCAmelCase_							:      Tuple=True		,       UpperCAmelCase_							:      List[str]=True		,       UpperCAmelCase_							:      str=True		,       UpperCAmelCase_							:      str=True		,       UpperCAmelCase_							:      Tuple=99		,       UpperCAmelCase_							:      List[str]=64		,       UpperCAmelCase_							:      Optional[int]=32		,       UpperCAmelCase_							:      Dict=5		,       UpperCAmelCase_							:      Optional[Any]=4		,       UpperCAmelCase_							:      Dict=37		,       UpperCAmelCase_							:      List[Any]="gelu"		,       UpperCAmelCase_							:      List[Any]=0.1		,       UpperCAmelCase_							:      int=0.1		,       UpperCAmelCase_							:      str=512		,       UpperCAmelCase_							:      Tuple=16		,       UpperCAmelCase_							:      Optional[int]=2		,       UpperCAmelCase_							:      List[str]=0.02		,       UpperCAmelCase_							:      Optional[Any]=3		,       UpperCAmelCase_							:      Tuple=4		,       UpperCAmelCase_							:      Tuple=None		,       )					->Dict:



														"""simple docstring"""



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



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



														"""simple docstring"""



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

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

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

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

														snake_case_           =		self.get_config()

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



							def 		lowerCAmelCase       (				self							:      Optional[int]       )					->Tuple:



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



							def 		lowerCAmelCase       (				self							:      Optional[Any]		,       UpperCAmelCase_							:      Dict		,       UpperCAmelCase_							:      Dict		,       UpperCAmelCase_							:      List[str]		,       UpperCAmelCase_							:      Optional[Any]		,       UpperCAmelCase_							:      Dict		,       UpperCAmelCase_							:      List[Any]		,       UpperCAmelCase_							:      Optional[int]       )					->Optional[Any]:



														"""simple docstring"""



														snake_case_           =		MegatronBertModel(config=UpperCAmelCase_       )
														model.to(UpperCAmelCase_       )
														model.eval()
														snake_case_           =		model(UpperCAmelCase_		,       attention_mask=UpperCAmelCase_		,       token_type_ids=UpperCAmelCase_       )
														snake_case_           =		model(UpperCAmelCase_		,       token_type_ids=UpperCAmelCase_       )
														snake_case_           =		model(UpperCAmelCase_       )

														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 		lowerCAmelCase       (				self							:      Dict		,       UpperCAmelCase_							:      Optional[int]		,       UpperCAmelCase_							:      List[Any]		,       UpperCAmelCase_							:      Optional[int]		,       UpperCAmelCase_							:      List[str]		,       UpperCAmelCase_							:      Any		,       UpperCAmelCase_							:      Tuple		,       UpperCAmelCase_							:      int       )					->List[str]:



														"""simple docstring"""



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



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



														"""simple docstring"""



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



							def 		lowerCAmelCase       (				self							:      List[Any]		,       UpperCAmelCase_							:      str		,       UpperCAmelCase_							:      Optional[int]		,       UpperCAmelCase_							:      Any		,       UpperCAmelCase_							:      List[Any]		,       UpperCAmelCase_							:      Optional[Any]		,       UpperCAmelCase_							:      str		,       UpperCAmelCase_							:      str       )					->str:



														"""simple docstring"""



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



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



														"""simple docstring"""



														snake_case_           =		MegatronBertForPreTraining(config=UpperCAmelCase_       )
														model.to(UpperCAmelCase_       )
														model.eval()
														snake_case_           =		model(
														    UpperCAmelCase_		,       attention_mask=UpperCAmelCase_		,       token_type_ids=UpperCAmelCase_		,       labels=UpperCAmelCase_		,       next_sentence_label=UpperCAmelCase_		,       )
														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 		lowerCAmelCase       (				self							:      Dict		,       UpperCAmelCase_							:      List[Any]		,       UpperCAmelCase_							:      List[Any]		,       UpperCAmelCase_							:      List[Any]		,       UpperCAmelCase_							:      Dict		,       UpperCAmelCase_							:      List[Any]		,       UpperCAmelCase_							:      Union[str, Any]		,       UpperCAmelCase_							:      int       )					->Tuple:



														"""simple docstring"""



														snake_case_           =		MegatronBertForQuestionAnswering(config=UpperCAmelCase_       )
														model.to(UpperCAmelCase_       )
														model.eval()
														snake_case_           =		model(
														    UpperCAmelCase_		,       attention_mask=UpperCAmelCase_		,       token_type_ids=UpperCAmelCase_		,       start_positions=UpperCAmelCase_		,       end_positions=UpperCAmelCase_		,       )
														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 		lowerCAmelCase       (				self							:      Tuple		,       UpperCAmelCase_							:      Optional[int]		,       UpperCAmelCase_							:      Any		,       UpperCAmelCase_							:      Any		,       UpperCAmelCase_							:      Tuple		,       UpperCAmelCase_							:      List[Any]		,       UpperCAmelCase_							:      Optional[Any]		,       UpperCAmelCase_							:      List[str]       )					->Tuple:



														"""simple docstring"""



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



							def 		lowerCAmelCase       (				self							:      str		,       UpperCAmelCase_							:      List[Any]		,       UpperCAmelCase_							:      Any		,       UpperCAmelCase_							:      str		,       UpperCAmelCase_							:      Tuple		,       UpperCAmelCase_							:      Tuple		,       UpperCAmelCase_							:      int		,       UpperCAmelCase_							:      Optional[int]       )					->List[str]:



														"""simple docstring"""



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



							def 		lowerCAmelCase       (				self							:      List[str]		,       UpperCAmelCase_							:      Union[str, Any]		,       UpperCAmelCase_							:      Optional[Any]		,       UpperCAmelCase_							:      Union[str, Any]		,       UpperCAmelCase_							:      Dict		,       UpperCAmelCase_							:      Tuple		,       UpperCAmelCase_							:      Tuple		,       UpperCAmelCase_							:      List[Any]       )					->int:



														"""simple docstring"""



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



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



														"""simple docstring"""



														snake_case_           =		self.prepare_config_and_inputs()
														(
														    (
														    snake_case_      
														)       ,						(
														    snake_case_      
														)       ,						(
														    snake_case_      
														)       ,						(
														    snake_case_      
														)       ,						(
														    snake_case_      
														)       ,						(
														    snake_case_      
														)       ,						(
														    snake_case_      
														)       ,						      
														)           =		config_and_inputs
														snake_case_           =		{"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
														return config, inputs_dict





@require_torch
class      __A     (snake_case_						,		snake_case_						,		unittest.TestCase):





							'''simple docstring'''

							__lowercase:		int							      =    (
							    (
							        MegatronBertModel,
							        MegatronBertForMaskedLM,
							        MegatronBertForCausalLM,
							        MegatronBertForMultipleChoice,
							        MegatronBertForNextSentencePrediction,
							        MegatronBertForPreTraining,
							        MegatronBertForQuestionAnswering,
							        MegatronBertForSequenceClassification,
							        MegatronBertForTokenClassification,
							    )
							    if is_torch_available()
							    else ()
							)
							__lowercase:		Any							      =    (
							    {
							        """feature-extraction""": MegatronBertModel,
							        """fill-mask""": MegatronBertForMaskedLM,
							        """question-answering""": MegatronBertForQuestionAnswering,
							        """text-classification""": MegatronBertForSequenceClassification,
							        """text-generation""": MegatronBertForCausalLM,
							        """token-classification""": MegatronBertForTokenClassification,
							        """zero-shot""": MegatronBertForSequenceClassification,
							    }
							    if is_torch_available()
							    else {}
							)
							__lowercase:		Dict							      =    True
							# test_resize_embeddings = False
							__lowercase:		str							      =    False



							def 		lowerCAmelCase       (				self							:      Dict		,       UpperCAmelCase_							:      List[Any]		,       UpperCAmelCase_							:      int		,       UpperCAmelCase_							:      Optional[int]=False       )					->str:



														"""simple docstring"""



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

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



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



														"""simple docstring"""



														snake_case_           =		MegatronBertModelTester(self       )
														snake_case_           =		ConfigTester(self		,       config_class=UpperCAmelCase_		,       hidden_size=37       )



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



														"""simple docstring"""



														self.config_tester.run_common_tests()



							def 		lowerCAmelCase       (				self							:      Tuple       )					->Any:



														"""simple docstring"""



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



							def 		lowerCAmelCase       (				self							:      Any       )					->int:



														"""simple docstring"""



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



							def 		lowerCAmelCase       (				self							:      str       )					->int:



														"""simple docstring"""



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



							def 		lowerCAmelCase       (				self							:      List[Any]       )					->Tuple:



														"""simple docstring"""



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



							def 		lowerCAmelCase       (				self							:      Optional[Any]       )					->str:



														"""simple docstring"""



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



							def 		lowerCAmelCase       (				self							:      List[str]       )					->int:



														"""simple docstring"""



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



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



														"""simple docstring"""



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



							def 		lowerCAmelCase       (				self							:      Dict       )					->int:



														"""simple docstring"""



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







def 						_a				(      _SCREAMING_SNAKE_CASE				)    ->				int:
							return torch.tensor(
							    _SCREAMING_SNAKE_CASE				, dtype=torch.long				, device=_SCREAMING_SNAKE_CASE				, )


__SCREAMING_SNAKE_CASE							:    Optional[int]		=    1E-4





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





							'''simple docstring'''

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



														"""simple docstring"""



														snake_case_           =		"""nvidia/megatron-bert-uncased-345m"""
														if "MYDIR" in os.environ:
																					snake_case_           =		os.path.join(os.environ["""MYDIR"""]		,       UpperCAmelCase_       )
														snake_case_           =		MegatronBertModel.from_pretrained(UpperCAmelCase_       )
														model.to(UpperCAmelCase_       )
														model.half()
														snake_case_           =		_long_tensor([[101, 7_110, 1_005, 1_056, 2_023, 11_333, 17_413, 1_029, 102]]       )
														with torch.no_grad():
																					snake_case_           =		model(UpperCAmelCase_       )[0]
														snake_case_           =		torch.Size((1, 9, 1_024)       )
														self.assertEqual(output.shape		,       UpperCAmelCase_       )

														snake_case_           =		[-0.6_040, -0.2_517, -0.1_025, 0.3_420, -0.6_758, -0.0_017, -0.1_089, -0.1_990, 0.5_728]
														for ii in range(3       ):
																					for jj in range(3       ):
																												snake_case_           =		output[0, ii, jj]
																												snake_case_           =		expected[3 * ii + jj]
																												snake_case_           =		"""ii={} jj={} a={} b={}""".format(UpperCAmelCase_		,       UpperCAmelCase_		,       UpperCAmelCase_		,       UpperCAmelCase_       )
																												self.assertTrue(math.isclose(UpperCAmelCase_		,       UpperCAmelCase_		,       rel_tol=UpperCAmelCase_		,       abs_tol=UpperCAmelCase_       )		,       msg=UpperCAmelCase_       )







"""simple docstring"""
import copy
from dataclasses import dataclass, field
from typing import ClassVar, Dict

from ..features import Audio, Features, Value
from .base import TaskTemplate





@dataclass(frozen=snake_case__)
class      __A     (snake_case__):





							'''simple docstring'''

							__lowercase:		str							      =    field(default="""automatic-speech-recognition"""						,		metadata={"""include_in_asdict_even_if_is_default""": True})
							__lowercase:		ClassVar[Features]							      =    Features({"""audio""": Audio()})
							__lowercase:		ClassVar[Features]							      =    Features({"""transcription""": Value("""string""")})
							__lowercase:		str							      =    "audio"
							__lowercase:		str							      =    "transcription"



							def 		lowerCAmelCase       (				self							:      List[Any]		,       UpperCAmelCase_							:      Any       )					->int:



														"""simple docstring"""



														if self.audio_column not in features:
																					raise ValueError(F"""Column {self.audio_column} is not present in features."""       )
														if not isinstance(features[self.audio_column]		,       UpperCAmelCase_       ):
																					raise ValueError(F"""Column {self.audio_column} is not an Audio type."""       )
														snake_case_           =		copy.deepcopy(self       )
														snake_case_           =		self.input_schema.copy()
														snake_case_           =		features[self.audio_column]
														snake_case_           =		input_schema
														return task_template



							@property
							def 		lowerCAmelCase       (				self							:      List[str]       )					->Dict[str, str]:



														"""simple docstring"""



														return {self.audio_column: "audio", self.transcription_column: "transcription"}



