




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






			'''simple docstring'''



			def __init__(						self						,   lowercase						,   lowercase=13						,   lowercase=7						,   lowercase=True						,   lowercase=True						,   lowercase=True						,   lowercase=True						,   lowercase=99						,   lowercase=64						,   lowercase=32						,   lowercase=5						,   lowercase=4						,   lowercase=37						,   lowercase="gelu"						,   lowercase=0.1						,   lowercase=0.1						,   lowercase=512						,   lowercase=16						,   lowercase=2						,   lowercase=0.02						,   lowercase=3						,   lowercase=4						,   lowercase=None						,   ):
							_lowerCamelCase  :    Dict					=					parent
							_lowerCamelCase  :    List[Any]					=					batch_size
							_lowerCamelCase  :    Optional[int]					=					seq_length
							_lowerCamelCase  :    Optional[int]					=					is_training
							_lowerCamelCase  :    Union[str, Any]					=					use_input_mask
							_lowerCamelCase  :    Tuple					=					use_token_type_ids
							_lowerCamelCase  :    Tuple					=					use_labels
							_lowerCamelCase  :    Optional[Any]					=					vocab_size
							_lowerCamelCase  :    Optional[int]					=					hidden_size
							_lowerCamelCase  :    Tuple					=					embedding_size
							_lowerCamelCase  :    Any					=					num_hidden_layers
							_lowerCamelCase  :    Optional[Any]					=					num_attention_heads
							_lowerCamelCase  :    Union[str, Any]					=					intermediate_size
							_lowerCamelCase  :    Tuple					=					hidden_act
							_lowerCamelCase  :    Dict					=					hidden_dropout_prob
							_lowerCamelCase  :    int					=					attention_probs_dropout_prob
							_lowerCamelCase  :    Dict					=					max_position_embeddings
							_lowerCamelCase  :    Any					=					type_vocab_size
							_lowerCamelCase  :    Union[str, Any]					=					type_sequence_label_size
							_lowerCamelCase  :    Union[str, Any]					=					initializer_range
							_lowerCamelCase  :    Any					=					num_labels
							_lowerCamelCase  :    int					=					num_choices
							_lowerCamelCase  :    int					=					scope
			def       A_					(						self						):
							_lowerCamelCase  :    List[str]					=					ids_tensor([self.batch_size, self.seq_length]						,   self.vocab_size						)

							_lowerCamelCase  :    List[str]					=					None
							if self.use_input_mask:
											_lowerCamelCase  :    Any					=					random_attention_mask([self.batch_size, self.seq_length]						)

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

							_lowerCamelCase  :    Dict					=					None
							_lowerCamelCase  :    str					=					None
							_lowerCamelCase  :    Union[str, Any]					=					None
							if self.use_labels:
											_lowerCamelCase  :    Any					=					ids_tensor([self.batch_size]						,   self.type_sequence_label_size						)
											_lowerCamelCase  :    str					=					ids_tensor([self.batch_size, self.seq_length]						,   self.num_labels						)
											_lowerCamelCase  :    int					=					ids_tensor([self.batch_size]						,   self.num_choices						)

							_lowerCamelCase  :    List[str]					=					self.get_config()

							return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
			def       A_					(						self						):
							return MegatronBertConfig(
							    vocab_size=self.vocab_size						,   hidden_size=self.hidden_size						,   num_hidden_layers=self.num_hidden_layers						,   num_attention_heads=self.num_attention_heads						,   intermediate_size=self.intermediate_size						,   embedding_size=self.embedding_size						,   hidden_act=self.hidden_act						,   hidden_dropout_prob=self.hidden_dropout_prob						,   attention_probs_dropout_prob=self.attention_probs_dropout_prob						,   max_position_embeddings=self.max_position_embeddings						,   type_vocab_size=self.type_vocab_size						,   is_decoder=lowercase						,   initializer_range=self.initializer_range						,   )
			def       A_					(						self						,   lowercase						,   lowercase						,   lowercase						,   lowercase						,   lowercase						,   lowercase						,   lowercase						):
							_lowerCamelCase  :    List[Any]					=					MegatronBertModel(config=lowercase						)
							model.to(lowercase						)
							model.eval()
							_lowerCamelCase  :    str					=					model(lowercase						,   attention_mask=lowercase						,   token_type_ids=lowercase						)
							_lowerCamelCase  :    List[Any]					=					model(lowercase						,   token_type_ids=lowercase						)
							_lowerCamelCase  :    Dict					=					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						,   lowercase						,   lowercase						,   lowercase						,   lowercase						,   lowercase						,   lowercase						,   lowercase						):
							_lowerCamelCase  :    Any					=					MegatronBertForMaskedLM(config=lowercase						)
							model.to(lowercase						)
							model.eval()
							_lowerCamelCase  :    List[str]					=					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						,   lowercase						,   lowercase						,   lowercase						,   lowercase						,   lowercase						,   lowercase						,   lowercase						):
							_lowerCamelCase  :    Any					=					MegatronBertForCausalLM(config=lowercase						)
							model.to(lowercase						)
							model.eval()
							_lowerCamelCase  :    Tuple					=					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						,   lowercase						,   lowercase						,   lowercase						,   lowercase						,   lowercase						,   lowercase						,   lowercase						):
							_lowerCamelCase  :    List[Any]					=					MegatronBertForNextSentencePrediction(config=lowercase						)
							model.to(lowercase						)
							model.eval()
							_lowerCamelCase  :    Optional[int]					=					model(
							    lowercase						,   attention_mask=lowercase						,   token_type_ids=lowercase						,   labels=lowercase						,   )
							self.parent.assertEqual(result.logits.shape						,   (self.batch_size, 2)						)
			def       A_					(						self						,   lowercase						,   lowercase						,   lowercase						,   lowercase						,   lowercase						,   lowercase						,   lowercase						):
							_lowerCamelCase  :    Any					=					MegatronBertForPreTraining(config=lowercase						)
							model.to(lowercase						)
							model.eval()
							_lowerCamelCase  :    int					=					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						,   lowercase						,   lowercase						,   lowercase						,   lowercase						,   lowercase						,   lowercase						,   lowercase						):
							_lowerCamelCase  :    str					=					MegatronBertForQuestionAnswering(config=lowercase						)
							model.to(lowercase						)
							model.eval()
							_lowerCamelCase  :    str					=					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						,   lowercase						,   lowercase						,   lowercase						,   lowercase						,   lowercase						,   lowercase						,   lowercase						):
							_lowerCamelCase  :    Dict					=					self.num_labels
							_lowerCamelCase  :    Optional[Any]					=					MegatronBertForSequenceClassification(lowercase						)
							model.to(lowercase						)
							model.eval()
							_lowerCamelCase  :    Union[str, Any]					=					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						,   lowercase						,   lowercase						,   lowercase						,   lowercase						,   lowercase						,   lowercase						,   lowercase						):
							_lowerCamelCase  :    str					=					self.num_labels
							_lowerCamelCase  :    List[str]					=					MegatronBertForTokenClassification(config=lowercase						)
							model.to(lowercase						)
							model.eval()
							_lowerCamelCase  :    List[Any]					=					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						,   lowercase						,   lowercase						,   lowercase						,   lowercase						,   lowercase						,   lowercase						,   lowercase						):
							_lowerCamelCase  :    Optional[int]					=					self.num_choices
							_lowerCamelCase  :    int					=					MegatronBertForMultipleChoice(config=lowercase						)
							model.to(lowercase						)
							model.eval()
							_lowerCamelCase  :    Tuple					=					input_ids.unsqueeze(1						).expand(-1						,   self.num_choices						,   -1						).contiguous()
							_lowerCamelCase  :    Any					=					token_type_ids.unsqueeze(1						).expand(-1						,   self.num_choices						,   -1						).contiguous()
							_lowerCamelCase  :    List[Any]					=					input_mask.unsqueeze(1						).expand(-1						,   self.num_choices						,   -1						).contiguous()
							_lowerCamelCase  :    Optional[Any]					=					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						):
							_lowerCamelCase  :    Optional[int]					=					self.prepare_config_and_inputs()
							(
							    _lowerCamelCase				
							)  :    Optional[int]					=					config_and_inputs
							_lowerCamelCase  :    Dict					=					{'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
							return config, inputs_dict






@require_torch
class lowerCAmelCase__			(  lowercase,			lowercase,			unittest.TestCase		):






			'''simple docstring'''



			lowerCamelCase__   =    (
			    (
			        MegatronBertModel,
			        MegatronBertForMaskedLM,
			        MegatronBertForCausalLM,
			        MegatronBertForMultipleChoice,
			        MegatronBertForNextSentencePrediction,
			        MegatronBertForPreTraining,
			        MegatronBertForQuestionAnswering,
			        MegatronBertForSequenceClassification,
			        MegatronBertForTokenClassification,
			    )
			    if is_torch_available()
			    else ()
			)
			lowerCamelCase__   =    (
			    {
			        """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 {}
			)
			lowerCamelCase__   =    True
			# test_resize_embeddings = False
			lowerCamelCase__   =    False
			def       A_					(						self						,   lowercase						,   lowercase						,   lowercase=False						):
							_lowerCamelCase  :    Any					=					super()._prepare_for_class(lowercase						,   lowercase						,   return_labels=lowercase						)

							if return_labels:
											if model_class in get_values(lowercase						):
															_lowerCamelCase  :    Any					=					torch.zeros(
															    (self.model_tester.batch_size, self.model_tester.seq_length)						,   dtype=torch.long						,   device=lowercase						)
															_lowerCamelCase  :    Any					=					torch.zeros(
															    self.model_tester.batch_size						,   dtype=torch.long						,   device=lowercase						)
							return inputs_dict
			def       A_					(						self						):
							_lowerCamelCase  :    Dict					=					MegatronBertModelTester(self						)
							_lowerCamelCase  :    Tuple					=					ConfigTester(self						,   config_class=lowercase						,   hidden_size=37						)
			def       A_					(						self						):
							self.config_tester.run_common_tests()
			def       A_					(						self						):
							_lowerCamelCase  :    List[Any]					=					self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_megatron_bert_model(*lowercase						)
			def       A_					(						self						):
							_lowerCamelCase  :    List[Any]					=					self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_megatron_bert_for_masked_lm(*lowercase						)
			def       A_					(						self						):
							_lowerCamelCase  :    Tuple					=					self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*lowercase						)
			def       A_					(						self						):
							_lowerCamelCase  :    Tuple					=					self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*lowercase						)
			def       A_					(						self						):
							_lowerCamelCase  :    str					=					self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_megatron_bert_for_pretraining(*lowercase						)
			def       A_					(						self						):
							_lowerCamelCase  :    Optional[int]					=					self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_megatron_bert_for_question_answering(*lowercase						)
			def       A_					(						self						):
							_lowerCamelCase  :    List[Any]					=					self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*lowercase						)






			def       A_					(						self						):
							_lowerCamelCase  :    List[str]					=					self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_megatron_bert_for_token_classification(*lowercase						)







def   _snake_case    (       lowercase__	):
				return torch.tensor(
				    lowercase__							,					dtype=torch.long							,					device=lowercase__							,					)


lowercase__  		=		1E-4






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






			'''simple docstring'''



			@slow
			@unittest.skip('Model is not available.'						)
			def       A_					(						self						):
							_lowerCamelCase  :    Dict					=					'nvidia/megatron-bert-uncased-345m'
							if "MYDIR" in os.environ:
											_lowerCamelCase  :    Optional[int]					=					os.path.join(os.environ['MYDIR']						,   lowercase						)
							_lowerCamelCase  :    Tuple					=					MegatronBertModel.from_pretrained(lowercase						)
							model.to(lowercase						)
							model.half()
							_lowerCamelCase  :    Tuple					=					_long_tensor([[101, 7110, 1005, 1056, 2023, 11333, 17413, 1029, 102]]						)
							with torch.no_grad():
											_lowerCamelCase  :    Optional[Any]					=					model(lowercase						)[0]
							_lowerCamelCase  :    int					=					torch.Size((1, 9, 1024)						)
							self.assertEqual(output.shape						,   lowercase						)

							_lowerCamelCase  :    Union[str, Any]					=					[-0.60_40, -0.25_17, -0.10_25, 0.34_20, -0.67_58, -0.00_17, -0.10_89, -0.19_90, 0.57_28]
							for ii in range(3						):
											for jj in range(3						):
															_lowerCamelCase  :    str					=					output[0, ii, jj]
															_lowerCamelCase  :    Optional[int]					=					expected[3 * ii + jj]
															_lowerCamelCase  :    Any					=					'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"""






import functools
import operator

from ...configuration_utils import PretrainedConfig
from ...utils import logging


lowercase__  		=		logging.get_logger(__name__)

lowercase__  		=		{
    """asapp/sew-d-tiny-100k""": """https://huggingface.co/asapp/sew-d-tiny-100k/resolve/main/config.json""",
    # See all SEW-D models at https://huggingface.co/models?filter=sew-d
}






class lowerCAmelCase__			(  lowercase		):






			'''simple docstring'''



			lowerCamelCase__   =    """sew-d"""
			def __init__(						self						,   lowercase=32						,   lowercase=768						,   lowercase=12						,   lowercase=12						,   lowercase=3072						,   lowercase=2						,   lowercase=512						,   lowercase=256						,   lowercase=True						,   lowercase=True						,   lowercase=("p2c", "c2p")						,   lowercase="layer_norm"						,   lowercase="gelu_python"						,   lowercase=0.1						,   lowercase=0.1						,   lowercase=0.1						,   lowercase=0.0						,   lowercase=0.1						,   lowercase=0.02						,   lowercase=1E-7						,   lowercase=1E-5						,   lowercase="group"						,   lowercase="gelu"						,   lowercase=(64, 128, 128, 128, 128, 256, 256, 256, 256, 512, 512, 512, 512)						,   lowercase=(5, 2, 1, 2, 1, 2, 1, 2, 1, 2, 1, 2, 1)						,   lowercase=(10, 3, 1, 3, 1, 3, 1, 3, 1, 2, 1, 2, 1)						,   lowercase=False						,   lowercase=128						,   lowercase=16						,   lowercase=True						,   lowercase=0.05						,   lowercase=10						,   lowercase=2						,   lowercase=0.0						,   lowercase=10						,   lowercase=0						,   lowercase="mean"						,   lowercase=False						,   lowercase=False						,   lowercase=256						,   lowercase=0						,   lowercase=1						,   lowercase=2						,   **lowercase						,   ):
							super().__init__(**lowercase						,   pad_token_id=lowercase						,   bos_token_id=lowercase						,   eos_token_id=lowercase						)
							_lowerCamelCase  :    List[str]					=					hidden_size
							_lowerCamelCase  :    int					=					feat_extract_norm
							_lowerCamelCase  :    Optional[Any]					=					feat_extract_activation
							_lowerCamelCase  :    Optional[Any]					=					list(lowercase						)
							_lowerCamelCase  :    Union[str, Any]					=					list(lowercase						)
							_lowerCamelCase  :    Optional[Any]					=					list(lowercase						)
							_lowerCamelCase  :    List[str]					=					conv_bias
							_lowerCamelCase  :    str					=					num_conv_pos_embeddings
							_lowerCamelCase  :    Optional[Any]					=					num_conv_pos_embedding_groups
							_lowerCamelCase  :    Tuple					=					len(self.conv_dim						)
							_lowerCamelCase  :    Optional[int]					=					num_hidden_layers
							_lowerCamelCase  :    List[str]					=					intermediate_size
							_lowerCamelCase  :    Dict					=					squeeze_factor
							_lowerCamelCase  :    Any					=					max_position_embeddings
							_lowerCamelCase  :    List[str]					=					position_buckets
							_lowerCamelCase  :    Dict					=					share_att_key
							_lowerCamelCase  :    List[str]					=					relative_attention
							_lowerCamelCase  :    Optional[Any]					=					norm_rel_ebd
							_lowerCamelCase  :    List[Any]					=					list(lowercase						)
							_lowerCamelCase  :    List[Any]					=					hidden_act
							_lowerCamelCase  :    Optional[Any]					=					num_attention_heads
							_lowerCamelCase  :    Optional[int]					=					hidden_dropout
							_lowerCamelCase  :    Optional[int]					=					attention_dropout
							_lowerCamelCase  :    Any					=					activation_dropout
							_lowerCamelCase  :    Tuple					=					feat_proj_dropout
							_lowerCamelCase  :    Any					=					final_dropout
							_lowerCamelCase  :    str					=					layer_norm_eps
							_lowerCamelCase  :    Tuple					=					feature_layer_norm_eps
							_lowerCamelCase  :    Any					=					initializer_range
							_lowerCamelCase  :    Dict					=					vocab_size

							if (
							    (len(self.conv_stride						) != self.num_feat_extract_layers)
							    or (len(self.conv_kernel						) != self.num_feat_extract_layers)
							    or (len(self.conv_dim						) != self.num_feat_extract_layers)
							):
											raise ValueError(
											    'Configuration for convolutional layers is incorrect.'
											    'It is required that `len(config.conv_dim)` == `len(config.conv_stride)` == `len(config.conv_kernel)`,'
											    F'''but is `len(config.conv_dim) = {len(self.conv_dim						)}`, `len(config.conv_stride)'''
											    F'''= {len(self.conv_stride						)}`, `len(config.conv_kernel) = {len(self.conv_kernel						)}`.'''						)

							# fine-tuning config parameters for SpecAugment: https://arxiv.org/abs/1904.08779
							_lowerCamelCase  :    Dict					=					apply_spec_augment
							_lowerCamelCase  :    int					=					mask_time_prob
							_lowerCamelCase  :    Union[str, Any]					=					mask_time_length
							_lowerCamelCase  :    Optional[Any]					=					mask_time_min_masks
							_lowerCamelCase  :    List[Any]					=					mask_feature_prob
							_lowerCamelCase  :    Union[str, Any]					=					mask_feature_length
							_lowerCamelCase  :    List[Any]					=					mask_feature_min_masks

							# ctc loss
							_lowerCamelCase  :    List[str]					=					ctc_loss_reduction
							_lowerCamelCase  :    List[str]					=					ctc_zero_infinity

							# sequence classification
							_lowerCamelCase  :    Dict					=					use_weighted_layer_sum
							_lowerCamelCase  :    List[Any]					=					classifier_proj_size






			@property
			def       A_					(						self						):
							return functools.reduce(operator.mul						,   self.conv_stride						,   1						)