'''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							,					_a							,					_a=13							,					_a=7							,					_a=True							,					_a=True							,					_a=True							,					_a=True							,					_a=99							,					_a=32							,					_a=5							,					_a=4							,					_a=37							,					_a="gelu"							,					_a=0.1							,					_a=0.1							,					_a=512							,					_a=16							,					_a=2							,					_a=0.02							,					_a=3							,					_a=4							,					_a=None							,					)     ->  Optional[int]:
        lowercase_       :							str  					=       parent
        lowercase_       :							Tuple  					=       batch_size
        lowercase_       :							Any  					=       seq_length
        lowercase_       :							Tuple  					=       is_training
        lowercase_       :							Tuple  					=       use_input_mask
        lowercase_       :							str  					=       use_token_type_ids
        lowercase_       :							List[str]  					=       use_labels
        lowercase_       :							Tuple  					=       vocab_size
        lowercase_       :							List[Any]  					=       hidden_size
        lowercase_       :							List[Any]  					=       num_hidden_layers
        lowercase_       :							Dict  					=       num_attention_heads
        lowercase_       :							str  					=       intermediate_size
        lowercase_       :							Optional[int]  					=       hidden_act
        lowercase_       :							int  					=       hidden_dropout_prob
        lowercase_       :							Dict  					=       attention_probs_dropout_prob
        lowercase_       :							int  					=       max_position_embeddings
        lowercase_       :							Union[str, Any]  					=       type_vocab_size
        lowercase_       :							Union[str, Any]  					=       type_sequence_label_size
        lowercase_       :							str  					=       initializer_range
        lowercase_       :							Tuple  					=       num_labels
        lowercase_       :							int  					=       num_choices
        lowercase_       :							Union[str, Any]  					=       scope
   def 		_lowerCamelCase			(self					)     ->  Any:
        lowercase_       :							Union[str, Any]  					=       ids_tensor([self.batch_size, self.seq_length]							,					self.vocab_size					)

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

        lowercase_       :							Tuple  					=       None
        if self.use_token_type_ids:
             lowercase_       :							List[str]  					=       ids_tensor([self.batch_size, self.seq_length]							,					self.type_vocab_size					)

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

        lowercase_       :							Optional[int]  					=       self.get_config()

        return (
            config,
            input_ids,
            token_type_ids,
            input_mask,
            sequence_labels,
            token_labels,
            choice_labels,
            fake_token_labels,
        )
   def 		_lowerCamelCase			(self					)     ->  List[str]:
        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=_a							,					initializer_range=self.initializer_range							,					)
   def 		_lowerCamelCase			(self					)     ->  int:
        (
            (
            lowercase_
        )					,(
            lowercase_
        )					,(
            lowercase_
        )					,(
            lowercase_
        )					,(
            lowercase_
        )					,(
            lowercase_
        )					,(
            lowercase_
        )					,(
            lowercase_
        )					,
        )       :							List[str]  					=       self.prepare_config_and_inputs()

        lowercase_       :							str  					=       True
        lowercase_       :							Optional[Any]  					=       floats_tensor([self.batch_size, self.seq_length, self.hidden_size]					)
        lowercase_       :							List[str]  					=       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 		_lowerCamelCase			(self							,					_a							,					_a							,					_a							,					_a							,					_a							,					_a							,					_a							,					_a							,					)     ->  Union[str, Any]:
        lowercase_       :							int  					=       ElectraModel(config=_a					)
        model.to(_a					)
        model.eval()
        lowercase_       :							List[str]  					=       model(_a							,					attention_mask=_a							,					token_type_ids=_a					)
        lowercase_       :							Optional[Any]  					=       model(_a							,					token_type_ids=_a					)
        lowercase_       :							Union[str, Any]  					=       model(_a					)
        self.parent.assertEqual(result.last_hidden_state.shape							,					(self.batch_size, self.seq_length, self.hidden_size)					)
   def 		_lowerCamelCase			(self							,					_a							,					_a							,					_a							,					_a							,					_a							,					_a							,					_a							,					_a							,					_a							,					)     ->  str:
        lowercase_       :							List[Any]  					=       True
        lowercase_       :							Any  					=       ElectraModel(_a					)
        model.to(_a					)
        model.eval()
        lowercase_       :							Union[str, Any]  					=       model(
            _a							,					attention_mask=_a							,					token_type_ids=_a							,					encoder_hidden_states=_a							,					encoder_attention_mask=_a							,					)
        lowercase_       :							List[str]  					=       model(
            _a							,					attention_mask=_a							,					token_type_ids=_a							,					encoder_hidden_states=_a							,					)
        lowercase_       :							Optional[Any]  					=       model(_a							,					attention_mask=_a							,					token_type_ids=_a					)
        self.parent.assertEqual(result.last_hidden_state.shape							,					(self.batch_size, self.seq_length, self.hidden_size)					)
   def 		_lowerCamelCase			(self							,					_a							,					_a							,					_a							,					_a							,					_a							,					_a							,					_a							,					_a							,					)     ->  List[Any]:
        lowercase_       :							List[Any]  					=       ElectraForMaskedLM(config=_a					)
        model.to(_a					)
        model.eval()
        lowercase_       :							int  					=       model(_a							,					attention_mask=_a							,					token_type_ids=_a							,					labels=_a					)
        self.parent.assertEqual(result.logits.shape							,					(self.batch_size, self.seq_length, self.vocab_size)					)
   def 		_lowerCamelCase			(self							,					_a							,					_a							,					_a							,					_a							,					_a							,					_a							,					_a							,					_a							,					_a							,					)     ->  List[str]:
        lowercase_       :							Union[str, Any]  					=       ElectraForCausalLM(config=_a					)
        model.to(_a					)
        model.eval()
        lowercase_       :							Optional[int]  					=       model(_a							,					attention_mask=_a							,					token_type_ids=_a							,					labels=_a					)
        self.parent.assertEqual(result.logits.shape							,					(self.batch_size, self.seq_length, self.vocab_size)					)
   def 		_lowerCamelCase			(self							,					_a							,					_a							,					_a							,					_a							,					_a							,					_a							,					_a							,					_a							,					)     ->  Any:
        lowercase_       :							Optional[int]  					=       self.num_labels
        lowercase_       :							Optional[int]  					=       ElectraForTokenClassification(config=_a					)
        model.to(_a					)
        model.eval()
        lowercase_       :							Union[str, Any]  					=       model(_a							,					attention_mask=_a							,					token_type_ids=_a							,					labels=_a					)
        self.parent.assertEqual(result.logits.shape							,					(self.batch_size, self.seq_length, self.num_labels)					)
   def 		_lowerCamelCase			(self							,					_a							,					_a							,					_a							,					_a							,					_a							,					_a							,					_a							,					_a							,					)     ->  Any:
        lowercase_       :							Dict  					=       self.num_labels
        lowercase_       :							Dict  					=       ElectraForPreTraining(config=_a					)
        model.to(_a					)
        model.eval()
        lowercase_       :							Tuple  					=       model(_a							,					attention_mask=_a							,					token_type_ids=_a							,					labels=_a					)
        self.parent.assertEqual(result.logits.shape							,					(self.batch_size, self.seq_length)					)
   def 		_lowerCamelCase			(self							,					_a							,					_a							,					_a							,					_a							,					_a							,					_a							,					_a							,					_a							,					)     ->  List[str]:
        lowercase_       :							Optional[int]  					=       self.num_labels
        lowercase_       :							Optional[Any]  					=       ElectraForSequenceClassification(_a					)
        model.to(_a					)
        model.eval()
        lowercase_       :							str  					=       model(_a							,					attention_mask=_a							,					token_type_ids=_a							,					labels=_a					)
        self.parent.assertEqual(result.logits.shape							,					(self.batch_size, self.num_labels)					)
   def 		_lowerCamelCase			(self							,					_a							,					_a							,					_a							,					_a							,					_a							,					_a							,					_a							,					_a							,					)     ->  Optional[Any]:
        lowercase_       :							Dict  					=       ElectraForQuestionAnswering(config=_a					)
        model.to(_a					)
        model.eval()
        lowercase_       :							Optional[Any]  					=       model(
            _a							,					attention_mask=_a							,					token_type_ids=_a							,					start_positions=_a							,					end_positions=_a							,					)
        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 		_lowerCamelCase			(self							,					_a							,					_a							,					_a							,					_a							,					_a							,					_a							,					_a							,					_a							,					)     ->  Union[str, Any]:
        lowercase_       :							str  					=       self.num_choices
        lowercase_       :							List[Any]  					=       ElectraForMultipleChoice(config=_a					)
        model.to(_a					)
        model.eval()
        lowercase_       :							int  					=       input_ids.unsqueeze(1					).expand(-1							,					self.num_choices							,					-1					).contiguous()
        lowercase_       :							List[Any]  					=       token_type_ids.unsqueeze(1					).expand(-1							,					self.num_choices							,					-1					).contiguous()
        lowercase_       :							Any  					=       input_mask.unsqueeze(1					).expand(-1							,					self.num_choices							,					-1					).contiguous()
        lowercase_       :							List[Any]  					=       model(
            _a							,					attention_mask=_a							,					token_type_ids=_a							,					labels=_a							,					)
        self.parent.assertEqual(result.logits.shape							,					(self.batch_size, self.num_choices)					)





   def 		_lowerCamelCase			(self					)     ->  Dict:
        lowercase_       :							List[str]  					=       self.prepare_config_and_inputs()
        (
            (
            lowercase_
        )					,(
            lowercase_
        )					,(
            lowercase_
        )					,(
            lowercase_
        )					,(
            lowercase_
        )					,(
            lowercase_
        )					,(
            lowercase_
        )					,(
            lowercase_
        )					,
        )       :							Optional[Any]  					=       config_and_inputs
        lowercase_       :							Optional[Any]  					=       {'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
        return config, inputs_dict



@require_torch
class    UpperCAmelCase__	(							_snake_case       ,  _snake_case       ,  unittest.TestCase      ):






   """simple docstring"""
   A				:							Tuple			     =     (
       (
           ElectraModel,
           ElectraForPreTraining,
           ElectraForMaskedLM,
           ElectraForCausalLM,
           ElectraForMultipleChoice,
           ElectraForTokenClassification,
           ElectraForSequenceClassification,
           ElectraForQuestionAnswering,
       )
       if is_torch_available()
       else ()
   )
   A				:							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 {}
   )
   A				:							str			     =     True
   def 		_lowerCamelCase			(self							,					_a							,					_a							,					_a=False					)     ->  Optional[Any]:
        lowercase_       :							List[Any]  					=       super()._prepare_for_class(_a							,					_a							,					return_labels=_a					)

        if return_labels:
             if model_class in get_values(_a					):
                  lowercase_       :							Optional[int]  					=       torch.zeros(
                      (self.model_tester.batch_size, self.model_tester.seq_length)							,					dtype=torch.long							,					device=_a					)
        return inputs_dict
   def 		_lowerCamelCase			(self					)     ->  List[Any]:
        lowercase_       :							int  					=       ElectraModelTester(self					)
        lowercase_       :							Optional[int]  					=       ConfigTester(self							,					config_class=_a							,					hidden_size=37					)
   def 		_lowerCamelCase			(self					)     ->  Dict:
        self.config_tester.run_common_tests()
   def 		_lowerCamelCase			(self					)     ->  Dict:
        lowercase_       :							Optional[Any]  					=       self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_electra_model(*_a					)
   def 		_lowerCamelCase			(self					)     ->  List[Any]:
        lowercase_       :							Optional[int]  					=       self.model_tester.prepare_config_and_inputs_for_decoder()
        self.model_tester.create_and_check_electra_model_as_decoder(*_a					)
   def 		_lowerCamelCase			(self					)     ->  str:
        lowercase_       :							int  					=       self.model_tester.prepare_config_and_inputs()
        for type in ["absolute", "relative_key", "relative_key_query"]:
             lowercase_       :							Any  					=       type
             self.model_tester.create_and_check_electra_model(*_a					)
   def 		_lowerCamelCase			(self					)     ->  Optional[int]:
        lowercase_       :							Tuple  					=       self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_electra_for_masked_lm(*_a					)
   def 		_lowerCamelCase			(self					)     ->  int:
        lowercase_       :							Tuple  					=       self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_electra_for_token_classification(*_a					)
   def 		_lowerCamelCase			(self					)     ->  List[Any]:
        lowercase_       :							int  					=       self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_electra_for_pretraining(*_a					)
   def 		_lowerCamelCase			(self					)     ->  Union[str, Any]:
        lowercase_       :							Any  					=       self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_electra_for_sequence_classification(*_a					)
   def 		_lowerCamelCase			(self					)     ->  Any:
        lowercase_       :							int  					=       self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_electra_for_question_answering(*_a					)
   def 		_lowerCamelCase			(self					)     ->  Tuple:
        lowercase_       :							Union[str, Any]  					=       self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_electra_for_multiple_choice(*_a					)
   @slow
   def 		_lowerCamelCase			(self					)     ->  Any:
        for model_name in ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             lowercase_       :							List[Any]  					=       ElectraModel.from_pretrained(_a					)
             self.assertIsNotNone(_a					)





   def 		_lowerCamelCase			(self					)     ->  Tuple:
        lowercase_       :							Optional[int]  					=       self.model_tester.prepare_config_and_inputs_for_decoder()
        self.model_tester.create_and_check_electra_for_causal_lm(*_a					)



@require_torch
class    UpperCAmelCase__	(							unittest.TestCase      ):






   """simple docstring"""
   @slow
   def 		_lowerCamelCase			(self					)     ->  int:
        lowercase_       :							Tuple  					=       ElectraModel.from_pretrained('google/electra-small-discriminator'					)
        lowercase_       :							int  					=       torch.tensor([[0, 345, 232, 328, 740, 140, 1_695, 69, 6_078, 1_588, 2]]					)
        lowercase_       :							List[str]  					=       torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]					)
        lowercase_       :							Optional[int]  					=       model(_a							,					attention_mask=_a					)[0]
        lowercase_       :							str  					=       torch.Size((1, 11, 256)					)
        self.assertEqual(output.shape							,					_a					)
        lowercase_       :							Dict  					=       torch.tensor(
            [[[0.44_71, 0.68_21, -0.32_65], [0.46_27, 0.52_55, -0.36_68], [0.45_32, 0.33_13, -0.43_44]]]					)

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





'''simple docstring'''




import logging
from pathlib import Path

import numpy as np
import pytorch_lightning as pl
import torch
from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
from pytorch_lightning.utilities import rank_zero_only

from utils import save_json

def _UpperCamelCase (							SCREAMING_SNAKE_CASE_				):
     lowercase_       :							Optional[Any]  					=       filter(lambda SCREAMING_SNAKE_CASE_				: p.requires_grad	,							model.parameters()				)
     lowercase_       :							Tuple  					=       sum([np.prod(p.size()				) for p in model_parameters]				)
     return params


_A										= logging.getLogger(__name__)


class    UpperCAmelCase__	(							pl.Callback      ):






   """simple docstring"""
   def 		_lowerCamelCase			(self							,					_a							,					_a					)     ->  str:
        lowercase_       :							List[str]  					=       {f'''lr_group_{i}''': param['lr'] for i, param in enumerate(pl_module.trainer.optimizers[0].param_groups					)}
        pl_module.logger.log_metrics(_a					)
   @rank_zero_only
   def 		_lowerCamelCase			(self							,					_a							,					_a							,					_a							,					_a=True					)     ->  None:
        logger.info(f'''***** {type_path} results at step {trainer.global_step:05d} *****'''					)
        lowercase_       :							Union[str, Any]  					=       trainer.callback_metrics
        trainer.logger.log_metrics({k: v for k, v in metrics.items() if k not in ['log', 'progress_bar', 'preds']}					)
        # Log results
        lowercase_       :							Optional[int]  					=       Path(pl_module.hparams.output_dir					)
        if type_path == "test":
             lowercase_       :							Union[str, Any]  					=       od / 'test_results.txt'
             lowercase_       :							Tuple  					=       od / 'test_generations.txt'
        else:
             # this never gets hit. I prefer not to save intermediate generations, and results are in metrics.json
             # If people want this it will be easy enough to add back.
             lowercase_       :							List[str]  					=       od / f'''{type_path}_results/{trainer.global_step:05d}.txt'''
             lowercase_       :							int  					=       od / f'''{type_path}_generations/{trainer.global_step:05d}.txt'''
             results_file.parent.mkdir(exist_ok=_a					)
             generations_file.parent.mkdir(exist_ok=_a					)
        with open(_a							,					'a+'					) as writer:
             for key in sorted(_a					):
                  if key in ["log", "progress_bar", "preds"]:
                       continue
                  lowercase_       :							Any  					=       metrics[key]
                  if isinstance(_a							,					torch.Tensor					):
                       lowercase_       :							Any  					=       val.item()
                  lowercase_       :							List[str]  					=       f'''{key}: {val:.6f}\n'''
                  writer.write(_a					)

        if not save_generations:
             return

        if "preds" in metrics:
             lowercase_       :							Optional[Any]  					=       '\n'.join(metrics['preds']					)
             generations_file.open('w+'					).write(_a					)
   @rank_zero_only
   def 		_lowerCamelCase			(self							,					_a							,					_a					)     ->  Tuple:
        try:
             lowercase_       :							List[Any]  					=       pl_module.model.model.num_parameters()
        except AttributeError:
             lowercase_       :							Dict  					=       pl_module.model.num_parameters()

        lowercase_       :							Optional[Any]  					=       count_trainable_parameters(_a					)
        # mp stands for million parameters
        trainer.logger.log_metrics({'n_params': npars, 'mp': npars / 1e6, 'grad_mp': n_trainable_pars / 1e6}					)
   @rank_zero_only
   def 		_lowerCamelCase			(self							,					_a							,					_a					)     ->  int:
        save_json(pl_module.metrics							,					pl_module.metrics_save_path					)
        return self._write_logs(_a							,					_a							,					'test'					)





   @rank_zero_only
   def 		_lowerCamelCase			(self							,					_a							,					_a					)     ->  List[str]:
        save_json(pl_module.metrics							,					pl_module.metrics_save_path					)
        # Uncommenting this will save val generations
        # return self._write_logs(trainer, pl_module, "valid")



def _UpperCamelCase (							SCREAMING_SNAKE_CASE_	,							SCREAMING_SNAKE_CASE_	,							SCREAMING_SNAKE_CASE_=1	,							SCREAMING_SNAKE_CASE_=False				):
     if metric == "rouge2":
          lowercase_       :							Optional[int]  					=       '{val_avg_rouge2:.4f}-{step_count}'
     elif metric == "bleu":
          lowercase_       :							int  					=       '{val_avg_bleu:.4f}-{step_count}'
     elif metric == "loss":
          lowercase_       :							List[Any]  					=       '{val_avg_loss:.4f}-{step_count}'
     else:
          raise NotImplementedError(
              f'''seq2seq callbacks only support rouge2, bleu and loss, got {metric}, You can make your own by adding to'''
              ' this function.'				)

     lowercase_       :							Tuple  					=       ModelCheckpoint(
         dirpath=SCREAMING_SNAKE_CASE_	,							filename=SCREAMING_SNAKE_CASE_	,							monitor=f'''val_{metric}'''	,							mode='min' if 'loss' in metric else 'max'	,							save_top_k=SCREAMING_SNAKE_CASE_	,							)
     return checkpoint_callback






def _UpperCamelCase (							SCREAMING_SNAKE_CASE_	,							SCREAMING_SNAKE_CASE_				):
     return EarlyStopping(
         monitor=f'''val_{metric}'''	,							mode='min' if 'loss' in metric else 'max'	,							patience=SCREAMING_SNAKE_CASE_	,							verbose=SCREAMING_SNAKE_CASE_	,							)




