





import argparse
import glob
import logging
import os
import time
from argparse import Namespace

import numpy as np
import torch
from lightning_base import BaseTransformer, add_generic_args, generic_train
from torch.utils.data import DataLoader, TensorDataset

from transformers import glue_compute_metrics as compute_metrics
from transformers import glue_convert_examples_to_features as convert_examples_to_features
from transformers import glue_output_modes, glue_tasks_num_labels
from transformers import glue_processors as processors


SCREAMING_SNAKE_CASE_		     = logging.getLogger(__name__)

class 	UpperCamelCase__				(   lowerCAmelCase_     ):



     '''simple docstring'''

     __snake_case					:			Any				=    "sequence-classification"






     def __init__(							self		:					Dict		,lowerCamelCase__		:					Any )  ->	List[Any]:




           '''simple docstring'''


           if type(lowerCamelCase__ ) == dict:
                 SCREAMING_SNAKE_CASE	=		Namespace(**lowerCamelCase__ )
           SCREAMING_SNAKE_CASE	=		glue_output_modes[hparams.task]
           SCREAMING_SNAKE_CASE	=		glue_tasks_num_labels[hparams.task]

           super().__init__(lowerCamelCase__		,lowerCamelCase__		,self.mode )






     def        SCREAMING_SNAKE_CASE__						(							self		:					Optional[int]		,**lowerCamelCase__		:					Any )  ->	Optional[int]:




           '''simple docstring'''


           return self.model(**lowerCamelCase__ )






     def        SCREAMING_SNAKE_CASE__						(							self		:					Any		,lowerCamelCase__		:					Optional[Any]		,lowerCamelCase__		:					List[Any] )  ->	Optional[int]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE	=		{"""input_ids""": batch[0], """attention_mask""": batch[1], """labels""": batch[3]}

           if self.config.model_type not in ["distilbert", "bart"]:
                 SCREAMING_SNAKE_CASE	=		batch[2] if self.config.model_type in ["""bert""", """xlnet""", """albert"""] else None

           SCREAMING_SNAKE_CASE	=		self(**lowerCamelCase__ )
           SCREAMING_SNAKE_CASE	=		outputs[0]

           SCREAMING_SNAKE_CASE	=		self.trainer.lr_schedulers[0]["""scheduler"""]
           SCREAMING_SNAKE_CASE	=		{"""loss""": loss, """rate""": lr_scheduler.get_last_lr()[-1]}
           return {"loss": loss, "log": tensorboard_logs}






     def        SCREAMING_SNAKE_CASE__						(							self		:					str )  ->	Union[str, Any]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE	=		self.hparams
           SCREAMING_SNAKE_CASE	=		processors[args.task]()
           SCREAMING_SNAKE_CASE	=		processor.get_labels()

           for mode in ["train", "dev"]:
                 SCREAMING_SNAKE_CASE	=		self._feature_file(lowerCamelCase__ )
                 if os.path.exists(lowerCamelCase__ ) and not args.overwrite_cache:
                       logger.info("""Loading features from cached file %s"""		,lowerCamelCase__ )
                 else:
                       logger.info("""Creating features from dataset file at %s"""		,args.data_dir )
                       SCREAMING_SNAKE_CASE	=		(
                           processor.get_dev_examples(args.data_dir )
                           if mode == """dev"""
                           else processor.get_train_examples(args.data_dir )
                       )
                       SCREAMING_SNAKE_CASE	=		convert_examples_to_features(
                           lowerCamelCase__		,self.tokenizer		,max_length=args.max_seq_length		,label_list=self.labels		,output_mode=args.glue_output_mode		,)
                       logger.info("""Saving features into cached file %s"""		,lowerCamelCase__ )
                       torch.save(lowerCamelCase__		,lowerCamelCase__ )






     def        SCREAMING_SNAKE_CASE__						(							self		:					Any		,lowerCamelCase__		:					str		,lowerCamelCase__		:					int		,lowerCamelCase__		:					bool = False )  ->	DataLoader:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE	=		"""dev""" if mode == """test""" else mode

           SCREAMING_SNAKE_CASE	=		self._feature_file(lowerCamelCase__ )
           logger.info("""Loading features from cached file %s"""		,lowerCamelCase__ )
           SCREAMING_SNAKE_CASE	=		torch.load(lowerCamelCase__ )
           SCREAMING_SNAKE_CASE	=		torch.tensor([f.input_ids for f in features]		,dtype=torch.long )
           SCREAMING_SNAKE_CASE	=		torch.tensor([f.attention_mask for f in features]		,dtype=torch.long )
           SCREAMING_SNAKE_CASE	=		torch.tensor([f.token_type_ids for f in features]		,dtype=torch.long )
           if self.hparams.glue_output_mode == "classification":
                 SCREAMING_SNAKE_CASE	=		torch.tensor([f.label for f in features]		,dtype=torch.long )
           elif self.hparams.glue_output_mode == "regression":
                 SCREAMING_SNAKE_CASE	=		torch.tensor([f.label for f in features]		,dtype=torch.float )

           return DataLoader(
               TensorDataset(lowerCamelCase__		,lowerCamelCase__		,lowerCamelCase__		,lowerCamelCase__ )		,batch_size=lowerCamelCase__		,shuffle=lowerCamelCase__		,)






     def        SCREAMING_SNAKE_CASE__						(							self		:					List[str]		,lowerCamelCase__		:					Union[str, Any]		,lowerCamelCase__		:					Union[str, Any] )  ->	str:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE	=		{"""input_ids""": batch[0], """attention_mask""": batch[1], """labels""": batch[3]}

           if self.config.model_type not in ["distilbert", "bart"]:
                 SCREAMING_SNAKE_CASE	=		batch[2] if self.config.model_type in ["""bert""", """xlnet""", """albert"""] else None

           SCREAMING_SNAKE_CASE	=		self(**lowerCamelCase__ )
           SCREAMING_SNAKE_CASE,      SCREAMING_SNAKE_CASE	=		outputs[:2]
           SCREAMING_SNAKE_CASE	=		logits.detach().cpu().numpy()
           SCREAMING_SNAKE_CASE	=		inputs["""labels"""].detach().cpu().numpy()

           return {"val_loss": tmp_eval_loss.detach().cpu(), "pred": preds, "target": out_label_ids}






     def        SCREAMING_SNAKE_CASE__						(							self		:					Dict		,lowerCamelCase__		:					Any )  ->	tuple:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE	=		torch.stack([x["""val_loss"""] for x in outputs] ).mean().detach().cpu().item()
           SCREAMING_SNAKE_CASE	=		np.concatenate([x["""pred"""] for x in outputs]		,axis=0 )

           if self.hparams.glue_output_mode == "classification":
                 SCREAMING_SNAKE_CASE	=		np.argmax(lowerCamelCase__		,axis=1 )
           elif self.hparams.glue_output_mode == "regression":
                 SCREAMING_SNAKE_CASE	=		np.squeeze(lowerCamelCase__ )

           SCREAMING_SNAKE_CASE	=		np.concatenate([x["""target"""] for x in outputs]		,axis=0 )
           SCREAMING_SNAKE_CASE	=		[[] for _ in range(out_label_ids.shape[0] )]
           SCREAMING_SNAKE_CASE	=		[[] for _ in range(out_label_ids.shape[0] )]

           SCREAMING_SNAKE_CASE	=		{**{"""val_loss""": val_loss_mean}, **compute_metrics(self.hparams.task		,lowerCamelCase__		,lowerCamelCase__ )}

           SCREAMING_SNAKE_CASE	=		dict(results.items() )
           SCREAMING_SNAKE_CASE	=		results
           return ret, preds_list, out_label_list






     def        SCREAMING_SNAKE_CASE__						(							self		:					Optional[Any]		,lowerCamelCase__		:					list )  ->	dict:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE,      SCREAMING_SNAKE_CASE,      SCREAMING_SNAKE_CASE	=		self._eval_end(lowerCamelCase__ )
           SCREAMING_SNAKE_CASE	=		ret["""log"""]
           return {"val_loss": logs["val_loss"], "log": logs, "progress_bar": logs}






     def        SCREAMING_SNAKE_CASE__						(							self		:					Tuple		,lowerCamelCase__		:					Tuple )  ->	dict:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE,      SCREAMING_SNAKE_CASE,      SCREAMING_SNAKE_CASE	=		self._eval_end(lowerCamelCase__ )
           SCREAMING_SNAKE_CASE	=		ret["""log"""]
           # `val_loss` is the key returned by `self._eval_end()` but actually refers to `test_loss`
           return {"avg_test_loss": logs["val_loss"], "log": logs, "progress_bar": logs}






     @staticmethod
     def        SCREAMING_SNAKE_CASE__						(							lowerCamelCase__		:					Tuple		,lowerCamelCase__		:					Optional[Any] )  ->	Any:




           '''simple docstring'''


           BaseTransformer.add_model_specific_args(lowerCamelCase__		,lowerCamelCase__ )
           parser.add_argument(
               """--max_seq_length"""		,default=128		,type=lowerCamelCase__		,help=(
                   """The maximum total input sequence length after tokenization. Sequences longer """
                   """than this will be truncated, sequences shorter will be padded."""
               )		,)

           parser.add_argument(
               """--task"""		,default=""""""		,type=lowerCamelCase__		,required=lowerCamelCase__		,help="""The GLUE task to run"""		,)
           parser.add_argument(
               """--gpus"""		,default=0		,type=lowerCamelCase__		,help="""The number of GPUs allocated for this, it is by default 0 meaning none"""		,)

           parser.add_argument(
               """--overwrite_cache"""		,action="""store_true"""		,help="""Overwrite the cached training and evaluation sets""" )

           return parser



def      __lowercase     (   )				->     Tuple:


      '''simple docstring'''


      SCREAMING_SNAKE_CASE	=		argparse.ArgumentParser()
      add_generic_args(_SCREAMING_SNAKE_CASE				,		os.getcwd()			)
      SCREAMING_SNAKE_CASE	=		GLUETransformer.add_model_specific_args(_SCREAMING_SNAKE_CASE				,		os.getcwd()			)
      SCREAMING_SNAKE_CASE	=		parser.parse_args()

      # If output_dir not provided, a folder will be generated in pwd
      if args.output_dir is None:
            SCREAMING_SNAKE_CASE	=		os.path.join(
                """./results"""				,		F"""{args.task}_{time.strftime('%Y%m%d_%H%M%S'			)}"""				,		)
            os.makedirs(args.output_dir			)

      SCREAMING_SNAKE_CASE	=		GLUETransformer(_SCREAMING_SNAKE_CASE			)
      SCREAMING_SNAKE_CASE	=		generic_train(_SCREAMING_SNAKE_CASE				,		_SCREAMING_SNAKE_CASE			)

      # Optionally, predict on dev set and write to output_dir
      if args.do_predict:
            SCREAMING_SNAKE_CASE	=		sorted(glob.glob(os.path.join(args.output_dir				,		"""checkpoint-epoch=*.ckpt"""			)				,		recursive=_SCREAMING_SNAKE_CASE			)			)
            SCREAMING_SNAKE_CASE	=		model.load_from_checkpoint(checkpoints[-1]			)
            return trainer.test(_SCREAMING_SNAKE_CASE			)


if __name__ == "__main__":
      main()









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


SCREAMING_SNAKE_CASE_		     = logging.get_logger(__name__)

SCREAMING_SNAKE_CASE_		     = {
    """google/pegasus-large""": """https://huggingface.co/google/pegasus-large/resolve/main/config.json""",
    # See all PEGASUS models at https://huggingface.co/models?filter=pegasus
}

class 	UpperCamelCase__				(   lowerCAmelCase_     ):



     '''simple docstring'''

     __snake_case					:			str				=    "pegasus"
     __snake_case					:			List[Any]				=    ["past_key_values"]
     __snake_case					:			Any				=    {"num_attention_heads": "encoder_attention_heads", "hidden_size": "d_model"}






     def __init__(							self		:					Dict		,lowerCamelCase__		:					List[str]=50265		,lowerCamelCase__		:					int=1024		,lowerCamelCase__		:					Optional[Any]=12		,lowerCamelCase__		:					Tuple=4096		,lowerCamelCase__		:					Dict=16		,lowerCamelCase__		:					Optional[Any]=12		,lowerCamelCase__		:					str=4096		,lowerCamelCase__		:					Dict=16		,lowerCamelCase__		:					Tuple=0.0		,lowerCamelCase__		:					Dict=0.0		,lowerCamelCase__		:					str=True		,lowerCamelCase__		:					Any=True		,lowerCamelCase__		:					List[Any]="gelu"		,lowerCamelCase__		:					Tuple=1024		,lowerCamelCase__		:					Optional[int]=0.1		,lowerCamelCase__		:					Tuple=0.0		,lowerCamelCase__		:					Dict=0.0		,lowerCamelCase__		:					Optional[int]=0.02		,lowerCamelCase__		:					Optional[int]=0		,lowerCamelCase__		:					List[Any]=False		,lowerCamelCase__		:					Union[str, Any]=0		,lowerCamelCase__		:					Optional[int]=1		,lowerCamelCase__		:					Any=1		,**lowerCamelCase__		:					Dict		,)  ->	str:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE	=		vocab_size
           SCREAMING_SNAKE_CASE	=		max_position_embeddings
           SCREAMING_SNAKE_CASE	=		d_model
           SCREAMING_SNAKE_CASE	=		encoder_ffn_dim
           SCREAMING_SNAKE_CASE	=		encoder_layers
           SCREAMING_SNAKE_CASE	=		encoder_attention_heads
           SCREAMING_SNAKE_CASE	=		decoder_ffn_dim
           SCREAMING_SNAKE_CASE	=		decoder_layers
           SCREAMING_SNAKE_CASE	=		decoder_attention_heads
           SCREAMING_SNAKE_CASE	=		dropout
           SCREAMING_SNAKE_CASE	=		attention_dropout
           SCREAMING_SNAKE_CASE	=		activation_dropout
           SCREAMING_SNAKE_CASE	=		activation_function
           SCREAMING_SNAKE_CASE	=		init_std
           SCREAMING_SNAKE_CASE	=		encoder_layerdrop
           SCREAMING_SNAKE_CASE	=		decoder_layerdrop
           SCREAMING_SNAKE_CASE	=		use_cache
           SCREAMING_SNAKE_CASE	=		encoder_layers
           SCREAMING_SNAKE_CASE	=		scale_embedding  # scale factor will be sqrt(d_model) if True
           super().__init__(
               pad_token_id=lowerCamelCase__		,eos_token_id=lowerCamelCase__		,is_encoder_decoder=lowerCamelCase__		,decoder_start_token_id=lowerCamelCase__		,forced_eos_token_id=lowerCamelCase__		,**lowerCamelCase__		,)






     @property
     def        SCREAMING_SNAKE_CASE__						(							self		:					Tuple )  ->	int:




           '''simple docstring'''


           return self.encoder_attention_heads






     @property
     def        SCREAMING_SNAKE_CASE__						(							self		:					List[str] )  ->	int:




           '''simple docstring'''


           return self.d_model


