





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




 def __init__(						self      ,  _lowerCAmelCase      ,  _lowerCAmelCase=13      ,  _lowerCAmelCase=7      ,  _lowerCAmelCase=True      ,  _lowerCAmelCase=True      ,  _lowerCAmelCase=True      ,  _lowerCAmelCase=True      ,  _lowerCAmelCase=99      ,  _lowerCAmelCase=64      ,  _lowerCAmelCase=32      ,  _lowerCAmelCase=5      ,  _lowerCAmelCase=4      ,  _lowerCAmelCase=37      ,  _lowerCAmelCase="gelu"      ,  _lowerCAmelCase=0.1      ,  _lowerCAmelCase=0.1      ,  _lowerCAmelCase=512      ,  _lowerCAmelCase=16      ,  _lowerCAmelCase=2      ,  _lowerCAmelCase=0.02      ,  _lowerCAmelCase=3      ,  _lowerCAmelCase=4      ,  _lowerCAmelCase=None      ,  )     ->				int:
     _lowerCAmelCase						      =   parent
     _lowerCAmelCase						      =   batch_size
     _lowerCAmelCase						      =   seq_length
     _lowerCAmelCase						      =   is_training
     _lowerCAmelCase						      =   use_input_mask
     _lowerCAmelCase						      =   use_token_type_ids
     _lowerCAmelCase						      =   use_labels
     _lowerCAmelCase						      =   vocab_size
     _lowerCAmelCase						      =   hidden_size
     _lowerCAmelCase						      =   embedding_size
     _lowerCAmelCase						      =   num_hidden_layers
     _lowerCAmelCase						      =   num_attention_heads
     _lowerCAmelCase						      =   intermediate_size
     _lowerCAmelCase						      =   hidden_act
     _lowerCAmelCase						      =   hidden_dropout_prob
     _lowerCAmelCase						      =   attention_probs_dropout_prob
     _lowerCAmelCase						      =   max_position_embeddings
     _lowerCAmelCase						      =   type_vocab_size
     _lowerCAmelCase						      =   type_sequence_label_size
     _lowerCAmelCase						      =   initializer_range
     _lowerCAmelCase						      =   num_labels
     _lowerCAmelCase						      =   num_choices
     _lowerCAmelCase						      =   scope




 def 				_snake_case (						self		)     ->				Any:
     _lowerCAmelCase						      =   ids_tensor([self.batch_size, self.seq_length]      ,  self.vocab_size		)

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

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

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

     _lowerCAmelCase						      =   self.get_config()

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




 def 				_snake_case (						self		)     ->				str:
     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=_lowerCAmelCase      ,  initializer_range=self.initializer_range      ,  )




 def 				_snake_case (						self      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase		)     ->				List[Any]:
     _lowerCAmelCase						      =   MegatronBertModel(config=_lowerCAmelCase		)
     model.to(_lowerCAmelCase		)
     model.eval()
     _lowerCAmelCase						      =   model(_lowerCAmelCase      ,  attention_mask=_lowerCAmelCase      ,  token_type_ids=_lowerCAmelCase		)
     _lowerCAmelCase						      =   model(_lowerCAmelCase      ,  token_type_ids=_lowerCAmelCase		)
     _lowerCAmelCase						      =   model(_lowerCAmelCase		)

     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 				_snake_case (						self      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase		)     ->				Any:
     _lowerCAmelCase						      =   MegatronBertForMaskedLM(config=_lowerCAmelCase		)
     model.to(_lowerCAmelCase		)
     model.eval()
     _lowerCAmelCase						      =   model(_lowerCAmelCase      ,  attention_mask=_lowerCAmelCase      ,  token_type_ids=_lowerCAmelCase      ,  labels=_lowerCAmelCase		)
     self.parent.assertEqual(result.logits.shape      ,  (self.batch_size, self.seq_length, self.vocab_size)		)




 def 				_snake_case (						self      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase		)     ->				Optional[Any]:
     _lowerCAmelCase						      =   MegatronBertForCausalLM(config=_lowerCAmelCase		)
     model.to(_lowerCAmelCase		)
     model.eval()
     _lowerCAmelCase						      =   model(_lowerCAmelCase      ,  attention_mask=_lowerCAmelCase      ,  token_type_ids=_lowerCAmelCase      ,  labels=_lowerCAmelCase		)
     self.parent.assertEqual(result.logits.shape      ,  (self.batch_size, self.seq_length, self.vocab_size)		)




 def 				_snake_case (						self      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase		)     ->				str:
     _lowerCAmelCase						      =   MegatronBertForNextSentencePrediction(config=_lowerCAmelCase		)
     model.to(_lowerCAmelCase		)
     model.eval()
     _lowerCAmelCase						      =   model(
         _lowerCAmelCase      ,  attention_mask=_lowerCAmelCase      ,  token_type_ids=_lowerCAmelCase      ,  labels=_lowerCAmelCase      ,  )
     self.parent.assertEqual(result.logits.shape      ,  (self.batch_size, 2)		)




 def 				_snake_case (						self      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase		)     ->				str:
     _lowerCAmelCase						      =   MegatronBertForPreTraining(config=_lowerCAmelCase		)
     model.to(_lowerCAmelCase		)
     model.eval()
     _lowerCAmelCase						      =   model(
         _lowerCAmelCase      ,  attention_mask=_lowerCAmelCase      ,  token_type_ids=_lowerCAmelCase      ,  labels=_lowerCAmelCase      ,  next_sentence_label=_lowerCAmelCase      ,  )
     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 				_snake_case (						self      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase		)     ->				Any:
     _lowerCAmelCase						      =   MegatronBertForQuestionAnswering(config=_lowerCAmelCase		)
     model.to(_lowerCAmelCase		)
     model.eval()
     _lowerCAmelCase						      =   model(
         _lowerCAmelCase      ,  attention_mask=_lowerCAmelCase      ,  token_type_ids=_lowerCAmelCase      ,  start_positions=_lowerCAmelCase      ,  end_positions=_lowerCAmelCase      ,  )
     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 				_snake_case (						self      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase		)     ->				Optional[int]:
     _lowerCAmelCase						      =   self.num_labels
     _lowerCAmelCase						      =   MegatronBertForSequenceClassification(_lowerCAmelCase		)
     model.to(_lowerCAmelCase		)
     model.eval()
     _lowerCAmelCase						      =   model(_lowerCAmelCase      ,  attention_mask=_lowerCAmelCase      ,  token_type_ids=_lowerCAmelCase      ,  labels=_lowerCAmelCase		)
     self.parent.assertEqual(result.logits.shape      ,  (self.batch_size, self.num_labels)		)




 def 				_snake_case (						self      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase		)     ->				Dict:
     _lowerCAmelCase						      =   self.num_labels
     _lowerCAmelCase						      =   MegatronBertForTokenClassification(config=_lowerCAmelCase		)
     model.to(_lowerCAmelCase		)
     model.eval()
     _lowerCAmelCase						      =   model(_lowerCAmelCase      ,  attention_mask=_lowerCAmelCase      ,  token_type_ids=_lowerCAmelCase      ,  labels=_lowerCAmelCase		)
     self.parent.assertEqual(result.logits.shape      ,  (self.batch_size, self.seq_length, self.num_labels)		)




 def 				_snake_case (						self      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase		)     ->				int:
     _lowerCAmelCase						      =   self.num_choices
     _lowerCAmelCase						      =   MegatronBertForMultipleChoice(config=_lowerCAmelCase		)
     model.to(_lowerCAmelCase		)
     model.eval()
     _lowerCAmelCase						      =   input_ids.unsqueeze(1		).expand(-1      ,  self.num_choices      ,  -1		).contiguous()
     _lowerCAmelCase						      =   token_type_ids.unsqueeze(1		).expand(-1      ,  self.num_choices      ,  -1		).contiguous()
     _lowerCAmelCase						      =   input_mask.unsqueeze(1		).expand(-1      ,  self.num_choices      ,  -1		).contiguous()
     _lowerCAmelCase						      =   model(
         _lowerCAmelCase      ,  attention_mask=_lowerCAmelCase      ,  token_type_ids=_lowerCAmelCase      ,  labels=_lowerCAmelCase      ,  )
     self.parent.assertEqual(result.logits.shape      ,  (self.batch_size, self.num_choices)		)





 def 				_snake_case (						self		)     ->				Union[str, Any]:
     _lowerCAmelCase						      =   self.prepare_config_and_inputs()
     (
         (
         _lowerCAmelCase
     )      ,      (
         _lowerCAmelCase
     )      ,      (
         _lowerCAmelCase
     )      ,      (
         _lowerCAmelCase
     )      ,      (
         _lowerCAmelCase
     )      ,      (
         _lowerCAmelCase
     )      ,      (
         _lowerCAmelCase
     )      ,      
     )						      =   config_and_inputs
     _lowerCAmelCase						      =   {"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
     return config, inputs_dict







@require_torch
class  lowerCAmelCase_		(		__magic_name__    ,__magic_name__    ,unittest.TestCase		):
 __lowerCamelCase    :				Optional[int]								=       (
     (
         MegatronBertModel,
         MegatronBertForMaskedLM,
         MegatronBertForCausalLM,
         MegatronBertForMultipleChoice,
         MegatronBertForNextSentencePrediction,
         MegatronBertForPreTraining,
         MegatronBertForQuestionAnswering,
         MegatronBertForSequenceClassification,
         MegatronBertForTokenClassification,
     )
     if is_torch_available()
     else ()
 )
 __lowerCamelCase    :				Tuple								=       (
     {
         "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    :				Any								=       True
 # test_resize_embeddings = False
 __lowerCamelCase    :				str								=       False




 def 				_snake_case (						self      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase=False		)     ->				Tuple:
     _lowerCAmelCase						      =   super()._prepare_for_class(_lowerCAmelCase      ,  _lowerCAmelCase      ,  return_labels=_lowerCAmelCase		)

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




 def 				_snake_case (						self		)     ->				str:
     _lowerCAmelCase						      =   MegatronBertModelTester(self		)
     _lowerCAmelCase						      =   ConfigTester(self      ,  config_class=_lowerCAmelCase      ,  hidden_size=37		)




 def 				_snake_case (						self		)     ->				List[str]:
     self.config_tester.run_common_tests()




 def 				_snake_case (						self		)     ->				Optional[int]:
     _lowerCAmelCase						      =   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_model(*_lowerCAmelCase		)




 def 				_snake_case (						self		)     ->				Tuple:
     _lowerCAmelCase						      =   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_for_masked_lm(*_lowerCAmelCase		)




 def 				_snake_case (						self		)     ->				Tuple:
     _lowerCAmelCase						      =   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*_lowerCAmelCase		)




 def 				_snake_case (						self		)     ->				Optional[Any]:
     _lowerCAmelCase						      =   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*_lowerCAmelCase		)




 def 				_snake_case (						self		)     ->				int:
     _lowerCAmelCase						      =   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_for_pretraining(*_lowerCAmelCase		)




 def 				_snake_case (						self		)     ->				List[str]:
     _lowerCAmelCase						      =   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_for_question_answering(*_lowerCAmelCase		)




 def 				_snake_case (						self		)     ->				Any:
     _lowerCAmelCase						      =   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*_lowerCAmelCase		)





 def 				_snake_case (						self		)     ->				str:
     _lowerCAmelCase						      =   self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_megatron_bert_for_token_classification(*_lowerCAmelCase		)







def 						__a(SCREAMING_SNAKE_CASE_	:       Optional[Any]     ):



    '''simple docstring'''





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


_SCREAMING_SNAKE_CASE									=				1e-4





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




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

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









'''simple docstring'''


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


_SCREAMING_SNAKE_CASE									=				logging.get_logger(__name__)

_SCREAMING_SNAKE_CASE									=				{
    "tiiuae/falcon-40b": "https://huggingface.co/tiiuae/falcon-40b/resolve/main/config.json",
    "tiiuae/falcon-7b": "https://huggingface.co/tiiuae/falcon-7b/resolve/main/config.json",
}

class  lowerCAmelCase_		(		__magic_name__		):
 __lowerCamelCase    :				Any								=       "falcon"
 __lowerCamelCase    :				List[str]								=       ["past_key_values"]




 def __init__(						self      ,  _lowerCAmelCase=65024      ,  _lowerCAmelCase=4544      ,  _lowerCAmelCase=32      ,  _lowerCAmelCase=71      ,  _lowerCAmelCase=1E-5      ,  _lowerCAmelCase=0.02      ,  _lowerCAmelCase=True      ,  _lowerCAmelCase=0.0      ,  _lowerCAmelCase=0.0      ,  _lowerCAmelCase=None      ,  _lowerCAmelCase=False      ,  _lowerCAmelCase=False      ,  _lowerCAmelCase=True      ,  _lowerCAmelCase=True      ,  _lowerCAmelCase=False      ,  _lowerCAmelCase=11      ,  _lowerCAmelCase=11      ,  **_lowerCAmelCase      ,  )     ->				Union[str, Any]:
     _lowerCAmelCase						      =   vocab_size
     # Backward compatibility with n_embed kwarg
     _lowerCAmelCase						      =   kwargs.pop("n_embed"      ,  _lowerCAmelCase		)
     _lowerCAmelCase						      =   hidden_size if n_embed is None else n_embed
     _lowerCAmelCase						      =   num_hidden_layers
     _lowerCAmelCase						      =   num_attention_heads
     _lowerCAmelCase						      =   layer_norm_epsilon
     _lowerCAmelCase						      =   initializer_range
     _lowerCAmelCase						      =   use_cache
     _lowerCAmelCase						      =   hidden_dropout
     _lowerCAmelCase						      =   attention_dropout

     _lowerCAmelCase						      =   bos_token_id
     _lowerCAmelCase						      =   eos_token_id
     _lowerCAmelCase						      =   num_attention_heads if num_kv_heads is None else num_kv_heads
     _lowerCAmelCase						      =   alibi
     _lowerCAmelCase						      =   new_decoder_architecture
     _lowerCAmelCase						      =   multi_query  # Ignored when new_decoder_architecture is True
     _lowerCAmelCase						      =   parallel_attn
     _lowerCAmelCase						      =   bias

     super().__init__(bos_token_id=_lowerCAmelCase      ,  eos_token_id=_lowerCAmelCase      ,  **_lowerCAmelCase		)




 @property
 def 				_snake_case (						self		)     ->				Optional[Any]:
     return self.hidden_size // self.num_attention_heads





 @property
 def 				_snake_case (						self		)     ->				Optional[Any]:
     return not self.alibi


