

'''simple docstring'''

import unittest

from transformers import ElectraConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
    import torch

    from transformers import (
        MODEL_FOR_PRETRAINING_MAPPING,
        ElectraForCausalLM,
        ElectraForMaskedLM,
        ElectraForMultipleChoice,
        ElectraForPreTraining,
        ElectraForQuestionAnswering,
        ElectraForSequenceClassification,
        ElectraForTokenClassification,
        ElectraModel,
    )
    from transformers.models.electra.modeling_electra import ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST







class   UpperCamelCase__   :


  """simple docstring"""

  def __init__(			self   ,      snake_case   ,      snake_case=1_3   ,      snake_case=7   ,      snake_case=True   ,      snake_case=True   ,      snake_case=True   ,      snake_case=True   ,      snake_case=9_9   ,      snake_case=3_2   ,      snake_case=5   ,      snake_case=4   ,      snake_case=3_7   ,      snake_case="gelu"   ,      snake_case=0.1   ,      snake_case=0.1   ,      snake_case=5_1_2   ,      snake_case=1_6   ,      snake_case=2   ,      snake_case=0.02   ,      snake_case=3   ,      snake_case=4   ,      snake_case=None   ,      ):


       '''simple docstring'''

       UpperCAmelCase		:							Dict        =       parent
       UpperCAmelCase		:							List[str]        =       batch_size
       UpperCAmelCase		:							Dict        =       seq_length
       UpperCAmelCase		:							List[str]        =       is_training
       UpperCAmelCase		:							List[Any]        =       use_input_mask
       UpperCAmelCase		:							Optional[Any]        =       use_token_type_ids
       UpperCAmelCase		:							Optional[int]        =       use_labels
       UpperCAmelCase		:							List[str]        =       vocab_size
       UpperCAmelCase		:							Any        =       hidden_size
       UpperCAmelCase		:							List[Any]        =       num_hidden_layers
       UpperCAmelCase		:							Any        =       num_attention_heads
       UpperCAmelCase		:							Dict        =       intermediate_size
       UpperCAmelCase		:							Optional[Any]        =       hidden_act
       UpperCAmelCase		:							Optional[int]        =       hidden_dropout_prob
       UpperCAmelCase		:							List[Any]        =       attention_probs_dropout_prob
       UpperCAmelCase		:							Optional[Any]        =       max_position_embeddings
       UpperCAmelCase		:							Optional[Any]        =       type_vocab_size
       UpperCAmelCase		:							List[Any]        =       type_sequence_label_size
       UpperCAmelCase		:							Tuple        =       initializer_range
       UpperCAmelCase		:							List[str]        =       num_labels
       UpperCAmelCase		:							Optional[Any]        =       num_choices
       UpperCAmelCase		:							List[str]        =       scope







  def 				A_	(			self      ):


       '''simple docstring'''

       UpperCAmelCase		:							List[str]        =       ids_tensor([self.batch_size, self.seq_length]   ,      self.vocab_size      )

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

       UpperCAmelCase		:							List[Any]        =       None
       if self.use_token_type_ids:
            UpperCAmelCase		:							Optional[Any]        =       ids_tensor([self.batch_size, self.seq_length]   ,      self.type_vocab_size      )

       UpperCAmelCase		:							str        =       None
       UpperCAmelCase		:							Optional[int]        =       None
       UpperCAmelCase		:							Tuple        =       None
       if self.use_labels:
            UpperCAmelCase		:							Optional[int]        =       ids_tensor([self.batch_size]   ,      self.type_sequence_label_size      )
            UpperCAmelCase		:							List[Any]        =       ids_tensor([self.batch_size, self.seq_length]   ,      self.num_labels      )
            UpperCAmelCase		:							str        =       ids_tensor([self.batch_size]   ,      self.num_choices      )
            UpperCAmelCase		:							Tuple        =       ids_tensor([self.batch_size, self.seq_length]   ,      1      )

       UpperCAmelCase		:							Any        =       self.get_config()

       return (
           config,
           input_ids,
           token_type_ids,
           input_mask,
           sequence_labels,
           token_labels,
           choice_labels,
           fake_token_labels,
       )







  def 				A_	(			self      ):


       '''simple docstring'''

       return ElectraConfig(
           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   ,      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=snake_case_   ,      initializer_range=self.initializer_range   ,      )







  def 				A_	(			self      ):


       '''simple docstring'''

       (
           UpperCAmelCase   
       )		:							int        =       self.prepare_config_and_inputs()

       UpperCAmelCase		:							str        =       True
       UpperCAmelCase		:							List[str]        =       floats_tensor([self.batch_size, self.seq_length, self.hidden_size]      )
       UpperCAmelCase		:							Optional[Any]        =       ids_tensor([self.batch_size, self.seq_length]   ,      vocab_size=2      )

       return (
           config,
           input_ids,
           token_type_ids,
           input_mask,
           sequence_labels,
           token_labels,
           choice_labels,
           encoder_hidden_states,
           encoder_attention_mask,
       )







  def 				A_	(			self   ,      snake_case   ,      snake_case   ,      snake_case   ,      snake_case   ,      snake_case   ,      snake_case   ,      snake_case   ,      snake_case   ,      ):


       '''simple docstring'''

       UpperCAmelCase		:							List[Any]        =       ElectraModel(config=snake_case_      )
       model.to(snake_case_      )
       model.eval()
       UpperCAmelCase		:							List[Any]        =       model(snake_case_   ,      attention_mask=snake_case_   ,      token_type_ids=snake_case_      )
       UpperCAmelCase		:							Optional[int]        =       model(snake_case_   ,      token_type_ids=snake_case_      )
       UpperCAmelCase		:							List[str]        =       model(snake_case_      )
       self.parent.assertEqual(result.last_hidden_state.shape   ,      (self.batch_size, self.seq_length, self.hidden_size)      )







  def 				A_	(			self   ,      snake_case   ,      snake_case   ,      snake_case   ,      snake_case   ,      snake_case   ,      snake_case   ,      snake_case   ,      snake_case   ,      snake_case   ,      ):


       '''simple docstring'''

       UpperCAmelCase		:							Optional[int]        =       True
       UpperCAmelCase		:							List[str]        =       ElectraModel(snake_case_      )
       model.to(snake_case_      )
       model.eval()
       UpperCAmelCase		:							Tuple        =       model(
           snake_case_   ,      attention_mask=snake_case_   ,      token_type_ids=snake_case_   ,      encoder_hidden_states=snake_case_   ,      encoder_attention_mask=snake_case_   ,      )
       UpperCAmelCase		:							Dict        =       model(
           snake_case_   ,      attention_mask=snake_case_   ,      token_type_ids=snake_case_   ,      encoder_hidden_states=snake_case_   ,      )
       UpperCAmelCase		:							Any        =       model(snake_case_   ,      attention_mask=snake_case_   ,      token_type_ids=snake_case_      )
       self.parent.assertEqual(result.last_hidden_state.shape   ,      (self.batch_size, self.seq_length, self.hidden_size)      )







  def 				A_	(			self   ,      snake_case   ,      snake_case   ,      snake_case   ,      snake_case   ,      snake_case   ,      snake_case   ,      snake_case   ,      snake_case   ,      ):


       '''simple docstring'''

       UpperCAmelCase		:							Tuple        =       ElectraForMaskedLM(config=snake_case_      )
       model.to(snake_case_      )
       model.eval()
       UpperCAmelCase		:							Optional[int]        =       model(snake_case_   ,      attention_mask=snake_case_   ,      token_type_ids=snake_case_   ,      labels=snake_case_      )
       self.parent.assertEqual(result.logits.shape   ,      (self.batch_size, self.seq_length, self.vocab_size)      )







  def 				A_	(			self   ,      snake_case   ,      snake_case   ,      snake_case   ,      snake_case   ,      snake_case   ,      snake_case   ,      snake_case   ,      snake_case   ,      snake_case   ,      ):


       '''simple docstring'''

       UpperCAmelCase		:							Union[str, Any]        =       ElectraForCausalLM(config=snake_case_      )
       model.to(snake_case_      )
       model.eval()
       UpperCAmelCase		:							Any        =       model(snake_case_   ,      attention_mask=snake_case_   ,      token_type_ids=snake_case_   ,      labels=snake_case_      )
       self.parent.assertEqual(result.logits.shape   ,      (self.batch_size, self.seq_length, self.vocab_size)      )







  def 				A_	(			self   ,      snake_case   ,      snake_case   ,      snake_case   ,      snake_case   ,      snake_case   ,      snake_case   ,      snake_case   ,      snake_case   ,      ):


       '''simple docstring'''

       UpperCAmelCase		:							Dict        =       self.num_labels
       UpperCAmelCase		:							int        =       ElectraForTokenClassification(config=snake_case_      )
       model.to(snake_case_      )
       model.eval()
       UpperCAmelCase		:							Optional[Any]        =       model(snake_case_   ,      attention_mask=snake_case_   ,      token_type_ids=snake_case_   ,      labels=snake_case_      )
       self.parent.assertEqual(result.logits.shape   ,      (self.batch_size, self.seq_length, self.num_labels)      )







  def 				A_	(			self   ,      snake_case   ,      snake_case   ,      snake_case   ,      snake_case   ,      snake_case   ,      snake_case   ,      snake_case   ,      snake_case   ,      ):


       '''simple docstring'''

       UpperCAmelCase		:							Tuple        =       self.num_labels
       UpperCAmelCase		:							Optional[int]        =       ElectraForPreTraining(config=snake_case_      )
       model.to(snake_case_      )
       model.eval()
       UpperCAmelCase		:							str        =       model(snake_case_   ,      attention_mask=snake_case_   ,      token_type_ids=snake_case_   ,      labels=snake_case_      )
       self.parent.assertEqual(result.logits.shape   ,      (self.batch_size, self.seq_length)      )







  def 				A_	(			self   ,      snake_case   ,      snake_case   ,      snake_case   ,      snake_case   ,      snake_case   ,      snake_case   ,      snake_case   ,      snake_case   ,      ):


       '''simple docstring'''

       UpperCAmelCase		:							List[str]        =       self.num_labels
       UpperCAmelCase		:							Union[str, Any]        =       ElectraForSequenceClassification(snake_case_      )
       model.to(snake_case_      )
       model.eval()
       UpperCAmelCase		:							Dict        =       model(snake_case_   ,      attention_mask=snake_case_   ,      token_type_ids=snake_case_   ,      labels=snake_case_      )
       self.parent.assertEqual(result.logits.shape   ,      (self.batch_size, self.num_labels)      )







  def 				A_	(			self   ,      snake_case   ,      snake_case   ,      snake_case   ,      snake_case   ,      snake_case   ,      snake_case   ,      snake_case   ,      snake_case   ,      ):


       '''simple docstring'''

       UpperCAmelCase		:							List[str]        =       ElectraForQuestionAnswering(config=snake_case_      )
       model.to(snake_case_      )
       model.eval()
       UpperCAmelCase		:							Any        =       model(
           snake_case_   ,      attention_mask=snake_case_   ,      token_type_ids=snake_case_   ,      start_positions=snake_case_   ,      end_positions=snake_case_   ,      )
       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   ,      snake_case   ,      snake_case   ,      snake_case   ,      snake_case   ,      snake_case   ,      snake_case   ,      snake_case   ,      snake_case   ,      ):


       '''simple docstring'''

       UpperCAmelCase		:							Optional[Any]        =       self.num_choices
       UpperCAmelCase		:							List[Any]        =       ElectraForMultipleChoice(config=snake_case_      )
       model.to(snake_case_      )
       model.eval()
       UpperCAmelCase		:							Optional[int]        =       input_ids.unsqueeze(1      ).expand(-1   ,      self.num_choices   ,      -1      ).contiguous()
       UpperCAmelCase		:							Any        =       token_type_ids.unsqueeze(1      ).expand(-1   ,      self.num_choices   ,      -1      ).contiguous()
       UpperCAmelCase		:							Optional[Any]        =       input_mask.unsqueeze(1      ).expand(-1   ,      self.num_choices   ,      -1      ).contiguous()
       UpperCAmelCase		:							Dict        =       model(
           snake_case_   ,      attention_mask=snake_case_   ,      token_type_ids=snake_case_   ,      labels=snake_case_   ,      )
       self.parent.assertEqual(result.logits.shape   ,      (self.batch_size, self.num_choices)      )







  def 				A_	(			self      ):


       '''simple docstring'''

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







@require_torch
class   UpperCamelCase__   (       _a							,							_a							,							unittest.TestCase       ):


  """simple docstring"""

  SCREAMING_SNAKE_CASE__  :     Union[str, Any]       =							(
      (
          ElectraModel,
          ElectraForPreTraining,
          ElectraForMaskedLM,
          ElectraForCausalLM,
          ElectraForMultipleChoice,
          ElectraForTokenClassification,
          ElectraForSequenceClassification,
          ElectraForQuestionAnswering,
      )
      if is_torch_available()
      else ()
  )
  SCREAMING_SNAKE_CASE__  :     Dict       =							(
      {
          """feature-extraction""": ElectraModel,
          """fill-mask""": ElectraForMaskedLM,
          """question-answering""": ElectraForQuestionAnswering,
          """text-classification""": ElectraForSequenceClassification,
          """text-generation""": ElectraForCausalLM,
          """token-classification""": ElectraForTokenClassification,
          """zero-shot""": ElectraForSequenceClassification,
      }
      if is_torch_available()
      else {}
  )
  SCREAMING_SNAKE_CASE__  :     Union[str, Any]       =							True







  def 				A_	(			self   ,      snake_case   ,      snake_case   ,      snake_case=False      ):


       '''simple docstring'''

       UpperCAmelCase		:							Tuple        =       super()._prepare_for_class(snake_case_   ,      snake_case_   ,      return_labels=snake_case_      )

       if return_labels:
            if model_class in get_values(snake_case_      ):
                 UpperCAmelCase		:							List[Any]        =       torch.zeros(
                     (self.model_tester.batch_size, self.model_tester.seq_length)   ,      dtype=torch.long   ,      device=snake_case_      )
       return inputs_dict







  def 				A_	(			self      ):


       '''simple docstring'''

       UpperCAmelCase		:							Tuple        =       ElectraModelTester(self      )
       UpperCAmelCase		:							List[str]        =       ConfigTester(self   ,      config_class=snake_case_   ,      hidden_size=3_7      )







  def 				A_	(			self      ):


       '''simple docstring'''

       self.config_tester.run_common_tests()







  def 				A_	(			self      ):


       '''simple docstring'''

       UpperCAmelCase		:							Dict        =       self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_electra_model(*snake_case_      )







  def 				A_	(			self      ):


       '''simple docstring'''

       UpperCAmelCase		:							List[Any]        =       self.model_tester.prepare_config_and_inputs_for_decoder()
       self.model_tester.create_and_check_electra_model_as_decoder(*snake_case_      )







  def 				A_	(			self      ):


       '''simple docstring'''

       UpperCAmelCase		:							Dict        =       self.model_tester.prepare_config_and_inputs()
       for type in ["absolute", "relative_key", "relative_key_query"]:
            UpperCAmelCase		:							str        =       type
            self.model_tester.create_and_check_electra_model(*snake_case_      )







  def 				A_	(			self      ):


       '''simple docstring'''

       UpperCAmelCase		:							Union[str, Any]        =       self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_electra_for_masked_lm(*snake_case_      )







  def 				A_	(			self      ):


       '''simple docstring'''

       UpperCAmelCase		:							Dict        =       self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_electra_for_token_classification(*snake_case_      )







  def 				A_	(			self      ):


       '''simple docstring'''

       UpperCAmelCase		:							List[str]        =       self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_electra_for_pretraining(*snake_case_      )







  def 				A_	(			self      ):


       '''simple docstring'''

       UpperCAmelCase		:							int        =       self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_electra_for_sequence_classification(*snake_case_      )







  def 				A_	(			self      ):


       '''simple docstring'''

       UpperCAmelCase		:							Optional[Any]        =       self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_electra_for_question_answering(*snake_case_      )







  def 				A_	(			self      ):


       '''simple docstring'''

       UpperCAmelCase		:							List[Any]        =       self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_electra_for_multiple_choice(*snake_case_      )







  @slow
  def 				A_	(			self      ):


       '''simple docstring'''

       for model_name in ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
            UpperCAmelCase		:							str        =       ElectraModel.from_pretrained(snake_case_      )
            self.assertIsNotNone(snake_case_      )







  def 				A_	(			self      ):


       '''simple docstring'''

       UpperCAmelCase		:							Any        =       self.model_tester.prepare_config_and_inputs_for_decoder()
       self.model_tester.create_and_check_electra_for_causal_lm(*snake_case_      )







@require_torch
class   UpperCamelCase__   (       unittest.TestCase       ):


  """simple docstring"""

  @slow
  def 				A_	(			self      ):


       '''simple docstring'''

       UpperCAmelCase		:							List[Any]        =       ElectraModel.from_pretrained("google/electra-small-discriminator"      )
       UpperCAmelCase		:							Union[str, Any]        =       torch.tensor([[0, 3_4_5, 2_3_2, 3_2_8, 7_4_0, 1_4_0, 1_6_9_5, 6_9, 6_0_7_8, 1_5_8_8, 2]]      )
       UpperCAmelCase		:							Dict        =       torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]      )
       UpperCAmelCase		:							Optional[Any]        =       model(snake_case_   ,      attention_mask=snake_case_      )[0]
       UpperCAmelCase		:							List[str]        =       torch.Size((1, 1_1, 2_5_6)      )
       self.assertEqual(output.shape   ,      snake_case_      )
       UpperCAmelCase		:							Optional[int]        =       torch.tensor(
           [[[0.4471, 0.6821, -0.3265], [0.4627, 0.5255, -0.3668], [0.4532, 0.3313, -0.4344]]]      )

       self.assertTrue(torch.allclose(output[:, 1:4, 1:4]   ,      snake_case_   ,      atol=1e-4      )      )



'''simple docstring'''

from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.neighbors import KNeighborsClassifier

# Load iris file
a       :      List[Any]  	=			load_iris()
iris.keys()


print(F'Target names: \n {iris.target_names} ')
print(F'\n Features: \n {iris.feature_names}')

# Train set e Test set
a	,						a	,						a	,						a       :      List[Any]  	=			train_test_split(
    iris["data"], iris["target"], random_state=4
)

# KNN

a       :      Tuple  	=			KNeighborsClassifier(n_neighbors=1)
knn.fit(X_train, y_train)

# new array to test
a       :      Optional[int]  	=			[[1, 2, 1, 4], [2, 3, 4, 5]]

a       :      Tuple  	=			knn.predict(X_new)

print(
    F'\nNew array: \n {X_new}\n\nTarget Names Prediction: \n'
    F' {iris["target_names"][prediction]}'
)
