


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





  """simple docstring"""




  def __init__(					self    ,						UpperCamelCase    ,						UpperCamelCase=13    ,						UpperCamelCase=7    ,						UpperCamelCase=True    ,						UpperCamelCase=True    ,						UpperCamelCase=True    ,						UpperCamelCase=True    ,						UpperCamelCase=99    ,						UpperCamelCase=64    ,						UpperCamelCase=32    ,						UpperCamelCase=5    ,						UpperCamelCase=4    ,						UpperCamelCase=37    ,						UpperCamelCase="gelu"    ,						UpperCamelCase=0.1    ,						UpperCamelCase=0.1    ,						UpperCamelCase=512    ,						UpperCamelCase=16    ,						UpperCamelCase=2    ,						UpperCamelCase=0.02    ,						UpperCamelCase=3    ,						UpperCamelCase=4    ,						UpperCamelCase=None    ,						):







   """simple docstring"""
   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  ):







   """simple docstring"""
   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  ):







   """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 						snake_case	(					self    ,						UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase  ):







   """simple docstring"""
   lowerCamelCase_  				=						MegatronBertModel(config=UpperCamelCase  )
   model.to(UpperCamelCase  )
   model.eval()
   lowerCamelCase_  				=						model(UpperCamelCase    ,						attention_mask=UpperCamelCase    ,						token_type_ids=UpperCamelCase  )
   lowerCamelCase_  				=						model(UpperCamelCase    ,						token_type_ids=UpperCamelCase  )
   lowerCamelCase_  				=						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 						snake_case	(					self    ,						UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase  ):







   """simple docstring"""
   lowerCamelCase_  				=						MegatronBertForMaskedLM(config=UpperCamelCase  )
   model.to(UpperCamelCase  )
   model.eval()
   lowerCamelCase_  				=						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 						snake_case	(					self    ,						UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase  ):







   """simple docstring"""
   lowerCamelCase_  				=						MegatronBertForCausalLM(config=UpperCamelCase  )
   model.to(UpperCamelCase  )
   model.eval()
   lowerCamelCase_  				=						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 						snake_case	(					self    ,						UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase  ):







   """simple docstring"""
   lowerCamelCase_  				=						MegatronBertForNextSentencePrediction(config=UpperCamelCase  )
   model.to(UpperCamelCase  )
   model.eval()
   lowerCamelCase_  				=						model(
       UpperCamelCase    ,						attention_mask=UpperCamelCase    ,						token_type_ids=UpperCamelCase    ,						labels=UpperCamelCase    ,						)
   self.parent.assertEqual(result.logits.shape    ,						(self.batch_size, 2)  )






  def 						snake_case	(					self    ,						UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase  ):







   """simple docstring"""
   lowerCamelCase_  				=						MegatronBertForPreTraining(config=UpperCamelCase  )
   model.to(UpperCamelCase  )
   model.eval()
   lowerCamelCase_  				=						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 						snake_case	(					self    ,						UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase  ):







   """simple docstring"""
   lowerCamelCase_  				=						MegatronBertForQuestionAnswering(config=UpperCamelCase  )
   model.to(UpperCamelCase  )
   model.eval()
   lowerCamelCase_  				=						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 						snake_case	(					self    ,						UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase  ):







   """simple docstring"""
   lowerCamelCase_  				=						self.num_labels
   lowerCamelCase_  				=						MegatronBertForSequenceClassification(UpperCamelCase  )
   model.to(UpperCamelCase  )
   model.eval()
   lowerCamelCase_  				=						model(UpperCamelCase    ,						attention_mask=UpperCamelCase    ,						token_type_ids=UpperCamelCase    ,						labels=UpperCamelCase  )
   self.parent.assertEqual(result.logits.shape    ,						(self.batch_size, self.num_labels)  )






  def 						snake_case	(					self    ,						UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase  ):







   """simple docstring"""
   lowerCamelCase_  				=						self.num_labels
   lowerCamelCase_  				=						MegatronBertForTokenClassification(config=UpperCamelCase  )
   model.to(UpperCamelCase  )
   model.eval()
   lowerCamelCase_  				=						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 						snake_case	(					self    ,						UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase  ):







   """simple docstring"""
   lowerCamelCase_  				=						self.num_choices
   lowerCamelCase_  				=						MegatronBertForMultipleChoice(config=UpperCamelCase  )
   model.to(UpperCamelCase  )
   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(
       UpperCamelCase    ,						attention_mask=UpperCamelCase    ,						token_type_ids=UpperCamelCase    ,						labels=UpperCamelCase    ,						)
   self.parent.assertEqual(result.logits.shape    ,						(self.batch_size, self.num_choices)  )







  def 						snake_case	(					self  ):







   """simple docstring"""
   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   snake_case					(					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 						snake_case	(					self    ,						UpperCamelCase    ,						UpperCamelCase    ,						UpperCamelCase=False  ):







   """simple docstring"""
   lowerCamelCase_  				=						super()._prepare_for_class(UpperCamelCase    ,						UpperCamelCase    ,						return_labels=UpperCamelCase  )

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






  def 						snake_case	(					self  ):







   """simple docstring"""
   lowerCamelCase_  				=						MegatronBertModelTester(self  )
   lowerCamelCase_  				=						ConfigTester(self    ,						config_class=UpperCamelCase    ,						hidden_size=37  )






  def 						snake_case	(					self  ):







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






  def 						snake_case	(					self  ):







   """simple docstring"""
   lowerCamelCase_  				=						self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_megatron_bert_model(*UpperCamelCase  )






  def 						snake_case	(					self  ):







   """simple docstring"""
   lowerCamelCase_  				=						self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_megatron_bert_for_masked_lm(*UpperCamelCase  )






  def 						snake_case	(					self  ):







   """simple docstring"""
   lowerCamelCase_  				=						self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*UpperCamelCase  )






  def 						snake_case	(					self  ):







   """simple docstring"""
   lowerCamelCase_  				=						self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*UpperCamelCase  )






  def 						snake_case	(					self  ):







   """simple docstring"""
   lowerCamelCase_  				=						self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_megatron_bert_for_pretraining(*UpperCamelCase  )






  def 						snake_case	(					self  ):







   """simple docstring"""
   lowerCamelCase_  				=						self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_megatron_bert_for_question_answering(*UpperCamelCase  )






  def 						snake_case	(					self  ):







   """simple docstring"""
   lowerCamelCase_  				=						self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*UpperCamelCase  )







  def 						snake_case	(					self  ):







   """simple docstring"""
   lowerCamelCase_  				=						self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_megatron_bert_for_token_classification(*UpperCamelCase  )





def __snake_case      (   UpperCAmelCase_		:		Any	):
 return torch.tensor(
     UpperCAmelCase_      ,       dtype=torch.long      ,       device=UpperCAmelCase_      ,       )


a_						:       Optional[int]      			=				1e-4







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





  """simple docstring"""




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







   """simple docstring"""
   lowerCamelCase_  				=						"nvidia/megatron-bert-uncased-345m"
   if "MYDIR" in os.environ:
    lowerCamelCase_  				=						os.path.join(os.environ["MYDIR"]    ,						UpperCamelCase  )
   lowerCamelCase_  				=						MegatronBertModel.from_pretrained(UpperCamelCase  )
   model.to(UpperCamelCase  )
   model.half()
   lowerCamelCase_  				=						_long_tensor([[101, 7110, 1005, 1056, 2023, 1_1333, 1_7413, 1029, 102]]  )
   with torch.no_grad():
    lowerCamelCase_  				=						model(UpperCamelCase  )[0]
   lowerCamelCase_  				=						torch.Size((1, 9, 1024)  )
   self.assertEqual(output.shape    ,						UpperCamelCase  )

   lowerCamelCase_  				=						[-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  ):
     lowerCamelCase_  				=						output[0, ii, jj]
     lowerCamelCase_  				=						expected[3 * ii + jj]
     lowerCamelCase_  				=						"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 numpy as np
from matplotlib import pyplot as plt
from sklearn.datasets import load_iris
from sklearn.metrics import ConfusionMatrixDisplay
from sklearn.model_selection import train_test_split
from xgboost import XGBClassifier





def __snake_case      (   UpperCAmelCase_		:		dict	):
 return (data["data"], data["target"])





def __snake_case      (   UpperCAmelCase_		:		np.ndarray      ,       UpperCAmelCase_		:		np.ndarray	):
 lowerCamelCase_  				=						XGBClassifier()
 classifier.fit(UpperCAmelCase_      ,       UpperCAmelCase_	)
 return classifier





def __snake_case      (   ):
 lowerCamelCase_  				=						load_iris()
 lowerCamelCase_				,lowerCamelCase_  				=						data_handling(UpperCAmelCase_	)
 lowerCamelCase_				,lowerCamelCase_				,lowerCamelCase_				,lowerCamelCase_  				=						train_test_split(
     UpperCAmelCase_      ,       UpperCAmelCase_      ,       test_size=0.25	)

 lowerCamelCase_  				=						iris["target_names"]

 # Create an XGBoost Classifier from the training data
 lowerCamelCase_  				=						xgboost(UpperCAmelCase_      ,       UpperCAmelCase_	)

 # Display the confusion matrix of the classifier with both training and test sets
 ConfusionMatrixDisplay.from_estimator(
     UpperCAmelCase_      ,       UpperCAmelCase_      ,       UpperCAmelCase_      ,       display_labels=UpperCAmelCase_      ,       cmap="Blues"      ,       normalize="true"      ,       )
 plt.title("Normalized Confusion Matrix - IRIS Dataset"	)
 plt.show()


if __name__ == "__main__":
     import doctest

     doctest.testmod(verbose=True)
     main()



