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




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



     """simple docstring"""






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




   def 			__a		(					self    :       Union[str, Any]   ):



     """simple docstring"""






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

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

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

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

     SCREAMING_SNAKE_CASE__    							=   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=_lowercase ,			initializer_range=self.initializer_range ,			)




   def 			__a		(					self    :       List[str] ,			_lowercase    :       Union[str, Any] ,			_lowercase    :       Union[str, Any] ,			_lowercase    :       Union[str, Any] ,			_lowercase    :       Optional[int] ,			_lowercase    :       List[Any] ,			_lowercase    :       Any ,			_lowercase    :       str   ):



     """simple docstring"""






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



     """simple docstring"""






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




   def 			__a		(					self    :       Optional[Any] ,			_lowercase    :       Dict ,			_lowercase    :       Any ,			_lowercase    :       Optional[int] ,			_lowercase    :       str ,			_lowercase    :       Optional[int] ,			_lowercase    :       Optional[Any] ,			_lowercase    :       Any   ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__    							=   MegatronBertForCausalLM(config=_lowercase   )
     model.to(_lowercase   )
     model.eval()
     SCREAMING_SNAKE_CASE__    							=   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    :       Dict ,			_lowercase    :       Any ,			_lowercase    :       int ,			_lowercase    :       Optional[Any] ,			_lowercase    :       List[str] ,			_lowercase    :       Tuple ,			_lowercase    :       Tuple ,			_lowercase    :       List[str]   ):



     """simple docstring"""






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




   def 			__a		(					self    :       List[str] ,			_lowercase    :       int ,			_lowercase    :       str ,			_lowercase    :       Union[str, Any] ,			_lowercase    :       Optional[int] ,			_lowercase    :       str ,			_lowercase    :       Optional[Any] ,			_lowercase    :       Optional[int]   ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__    							=   MegatronBertForPreTraining(config=_lowercase   )
     model.to(_lowercase   )
     model.eval()
     SCREAMING_SNAKE_CASE__    							=   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    :       Dict ,			_lowercase    :       int ,			_lowercase    :       Dict ,			_lowercase    :       int ,			_lowercase    :       Dict ,			_lowercase    :       List[Any] ,			_lowercase    :       Any ,			_lowercase    :       Any   ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__    							=   MegatronBertForQuestionAnswering(config=_lowercase   )
     model.to(_lowercase   )
     model.eval()
     SCREAMING_SNAKE_CASE__    							=   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    :       Any ,			_lowercase    :       Optional[Any] ,			_lowercase    :       Union[str, Any] ,			_lowercase    :       Any ,			_lowercase    :       Tuple ,			_lowercase    :       Optional[int] ,			_lowercase    :       Optional[int] ,			_lowercase    :       Optional[Any]   ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__    							=   self.num_labels
     SCREAMING_SNAKE_CASE__    							=   MegatronBertForSequenceClassification(_lowercase   )
     model.to(_lowercase   )
     model.eval()
     SCREAMING_SNAKE_CASE__    							=   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    :       Optional[int] ,			_lowercase    :       Optional[Any] ,			_lowercase    :       Optional[int] ,			_lowercase    :       Union[str, Any] ,			_lowercase    :       Any ,			_lowercase    :       List[str] ,			_lowercase    :       List[Any] ,			_lowercase    :       List[Any]   ):



     """simple docstring"""






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



     """simple docstring"""






     SCREAMING_SNAKE_CASE__    							=   self.num_choices
     SCREAMING_SNAKE_CASE__    							=   MegatronBertForMultipleChoice(config=_lowercase   )
     model.to(_lowercase   )
     model.eval()
     SCREAMING_SNAKE_CASE__    							=   input_ids.unsqueeze(1   ).expand(-1 ,			self.num_choices ,			-1   ).contiguous()
     SCREAMING_SNAKE_CASE__    							=   token_type_ids.unsqueeze(1   ).expand(-1 ,			self.num_choices ,			-1   ).contiguous()
     SCREAMING_SNAKE_CASE__    							=   input_mask.unsqueeze(1   ).expand(-1 ,			self.num_choices ,			-1   ).contiguous()
     SCREAMING_SNAKE_CASE__    							=   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    :       Optional[Any]   ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__    							=   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__  
     )							,    
     )    							=   config_and_inputs
     SCREAMING_SNAKE_CASE__    							=   {"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
     return config, inputs_dict







@require_torch
class __snake_case       (		lowerCamelCase_     ,	lowerCamelCase_     ,	unittest.TestCase	):
   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    :       List[str] ,			_lowercase    :       List[str] ,			_lowercase    :       Optional[int] ,			_lowercase    :       Tuple=False   ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__    							=   super()._prepare_for_class(_lowercase ,			_lowercase ,			return_labels=_lowercase   )

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




   def 			__a		(					self    :       str   ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__    							=   MegatronBertModelTester(self   )
     SCREAMING_SNAKE_CASE__    							=   ConfigTester(self ,			config_class=_lowercase ,			hidden_size=37   )




   def 			__a		(					self    :       Union[str, Any]   ):



     """simple docstring"""






     self.config_tester.run_common_tests()




   def 			__a		(					self    :       Any   ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__    							=   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_model(*_lowercase   )




   def 			__a		(					self    :       Dict   ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__    							=   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_for_masked_lm(*_lowercase   )




   def 			__a		(					self    :       int   ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__    							=   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*_lowercase   )




   def 			__a		(					self    :       List[Any]   ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__    							=   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*_lowercase   )




   def 			__a		(					self    :       Optional[int]   ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__    							=   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_for_pretraining(*_lowercase   )




   def 			__a		(					self    :       int   ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__    							=   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_for_question_answering(*_lowercase   )




   def 			__a		(					self    :       Dict   ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__    							=   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*_lowercase   )




   def 			__a		(					self    :       Dict   ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__    							=   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_for_token_classification(*_lowercase   )





def 			__SCREAMING_SNAKE_CASE						(      __UpperCamelCase					: Dict			)							->      Any:







  """simple docstring"""





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


__lowerCamelCase    :				List[Any]							     =	1e-4







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




   @slow
   @unittest.skip("""Model is not available."""   )
   def 			__a		(					self    :       str   ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__    							=   """nvidia/megatron-bert-uncased-345m"""
     if "MYDIR" in os.environ:
       SCREAMING_SNAKE_CASE__    							=   os.path.join(os.environ["""MYDIR"""] ,			_lowercase   )
     SCREAMING_SNAKE_CASE__    							=   MegatronBertModel.from_pretrained(_lowercase   )
     model.to(_lowercase   )
     model.half()
     SCREAMING_SNAKE_CASE__    							=   _long_tensor([[1_01, 71_10, 10_05, 10_56, 20_23, 1_13_33, 1_74_13, 10_29, 1_02]]   )
     with torch.no_grad():
       SCREAMING_SNAKE_CASE__    							=   model(_lowercase   )[0]
     SCREAMING_SNAKE_CASE__    							=   torch.Size((1, 9, 10_24)   )
     self.assertEqual(output.shape ,			_lowercase   )

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

from math import log

from scipy.constants import Boltzmann, physical_constants

__lowerCamelCase    :				int							     =	300  # TEMPERATURE (unit = K)



def 			__SCREAMING_SNAKE_CASE						(      __UpperCamelCase					: float				,  __UpperCamelCase					: float				,  __UpperCamelCase					: float				,  )							->      float:







  """simple docstring"""





  if donor_conc <= 0:
    raise ValueError("""Donor concentration should be positive"""			)
  elif acceptor_conc <= 0:
    raise ValueError("""Acceptor concentration should be positive"""			)
  elif intrinsic_conc <= 0:
    raise ValueError("""Intrinsic concentration should be positive"""			)
  elif donor_conc <= intrinsic_conc:
    raise ValueError(
        """Donor concentration should be greater than intrinsic concentration"""			)
  elif acceptor_conc <= intrinsic_conc:
    raise ValueError(
        """Acceptor concentration should be greater than intrinsic concentration"""			)
  else:
    return (
        Boltzmann
        * T
        * log((donor_conc * acceptor_conc) / intrinsic_conc**2			)
        / physical_constants["electron volt"][0]
    )


if __name__ == "__main__":
   import doctest

   doctest.testmod()
