





import argparse
import logging
import os
from pathlib import Path
from typing import Any, Dict

import pytorch_lightning as pl
from pytorch_lightning.utilities import rank_zero_info

from transformers import (
    AdamW,
    AutoConfig,
    AutoModel,
    AutoModelForPreTraining,
    AutoModelForQuestionAnswering,
    AutoModelForSeqaSeqLM,
    AutoModelForSequenceClassification,
    AutoModelForTokenClassification,
    AutoModelWithLMHead,
    AutoTokenizer,
    PretrainedConfig,
    PreTrainedTokenizer,
)
from transformers.optimization import (
    Adafactor,
    get_cosine_schedule_with_warmup,
    get_cosine_with_hard_restarts_schedule_with_warmup,
    get_linear_schedule_with_warmup,
    get_polynomial_decay_schedule_with_warmup,
)
from transformers.utils.versions import require_version


SCREAMING_SNAKE_CASE_		  =			logging.getLogger(__name__)

require_version("""pytorch_lightning>=1.0.4""")

SCREAMING_SNAKE_CASE_		  =			{
    'base': AutoModel,
    'sequence-classification': AutoModelForSequenceClassification,
    'question-answering': AutoModelForQuestionAnswering,
    'pretraining': AutoModelForPreTraining,
    'token-classification': AutoModelForTokenClassification,
    'language-modeling': AutoModelWithLMHead,
    'summarization': AutoModelForSeqaSeqLM,
    'translation': AutoModelForSeqaSeqLM,
}


# update this and the import above to support new schedulers from transformers.optimization
SCREAMING_SNAKE_CASE_		  =			{
    'linear': get_linear_schedule_with_warmup,
    'cosine': get_cosine_schedule_with_warmup,
    'cosine_w_restarts': get_cosine_with_hard_restarts_schedule_with_warmup,
    'polynomial': get_polynomial_decay_schedule_with_warmup,
    # '': get_constant_schedule,             # not supported for now
    # '': get_constant_schedule_with_warmup, # not supported for now
}
SCREAMING_SNAKE_CASE_		  =			sorted(arg_to_scheduler.keys())
SCREAMING_SNAKE_CASE_		  =			'{' + ', '.join(arg_to_scheduler_choices) + '}'

class 						UpperCamelCase__					( pl.LightningModule  ):



     '''simple docstring'''

     def __init__(				self		:       Any		,lowerCamelCase__		:       argparse.Namespace		,lowerCamelCase__		:       List[str]=None		,lowerCamelCase__		:       int="base"		,lowerCamelCase__		:       int=None		,lowerCamelCase__		:       List[str]=None		,lowerCamelCase__		:       Union[str, Any]=None		,**lowerCamelCase__		:       List[str]		,)	->						str:




           '''simple docstring'''


           super().__init__()
           # TODO: move to self.save_hyperparameters()
           # self.save_hyperparameters()
           # can also expand arguments into trainer signature for easier reading

           self.save_hyperparameters(lowerCamelCase__			)
           SCREAMING_SNAKE_CASE						=		0
           SCREAMING_SNAKE_CASE						=		Path(self.hparams.output_dir			)
           SCREAMING_SNAKE_CASE						=		self.hparams.cache_dir if self.hparams.cache_dir else None
           if config is None:
                 SCREAMING_SNAKE_CASE						=		AutoConfig.from_pretrained(
                     self.hparams.config_name if self.hparams.config_name else self.hparams.model_name_or_path		,**({"""num_labels""": num_labels} if num_labels is not None else {})		,cache_dir=lowerCamelCase__		,**lowerCamelCase__		,)
           else:
                 SCREAMING_SNAKE_CASE						=		config

           SCREAMING_SNAKE_CASE						=		('encoder_layerdrop', 'decoder_layerdrop', 'dropout', 'attention_dropout')
           for p in extra_model_params:
                 if getattr(self.hparams		,lowerCamelCase__		,lowerCamelCase__			):
                       assert hasattr(self.config		,lowerCamelCase__			), F"""model config doesn\'t have a `{p}` attribute"""
                       setattr(self.config		,lowerCamelCase__		,getattr(self.hparams		,lowerCamelCase__			)			)

           if tokenizer is None:
                 SCREAMING_SNAKE_CASE						=		AutoTokenizer.from_pretrained(
                     self.hparams.tokenizer_name if self.hparams.tokenizer_name else self.hparams.model_name_or_path		,cache_dir=lowerCamelCase__		,)
           else:
                 SCREAMING_SNAKE_CASE						=		tokenizer
           SCREAMING_SNAKE_CASE						=		MODEL_MODES[mode]
           if model is None:
                 SCREAMING_SNAKE_CASE						=		self.model_type.from_pretrained(
                     self.hparams.model_name_or_path		,from_tf=bool(""".ckpt""" in self.hparams.model_name_or_path			)		,config=self.config		,cache_dir=lowerCamelCase__		,)
           else:
                 SCREAMING_SNAKE_CASE						=		model






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




           '''simple docstring'''


           SCREAMING_SNAKE_CASE						=		self.model_type.from_pretrained(*lowerCamelCase__		,**lowerCamelCase__			)






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




           '''simple docstring'''


           SCREAMING_SNAKE_CASE						=		arg_to_scheduler[self.hparams.lr_scheduler]
           SCREAMING_SNAKE_CASE						=		get_schedule_func(
               self.opt		,num_warmup_steps=self.hparams.warmup_steps		,num_training_steps=self.total_steps()			)
           SCREAMING_SNAKE_CASE						=		{'scheduler': scheduler, 'interval': 'step', 'frequency': 1}
           return scheduler






     def 							SCREAMING_SNAKE_CASE__    (				self		:       Tuple			)	->						Any:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE						=		self.model
           SCREAMING_SNAKE_CASE						=		['bias', 'LayerNorm.weight']
           SCREAMING_SNAKE_CASE						=		[
               {
                   'params': [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay			)],
                   'weight_decay': self.hparams.weight_decay,
               },
               {
                   'params': [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay			)],
                   'weight_decay': 0.0,
               },
           ]
           if self.hparams.adafactor:
                 SCREAMING_SNAKE_CASE						=		Adafactor(
                     lowerCamelCase__		,lr=self.hparams.learning_rate		,scale_parameter=lowerCamelCase__		,relative_step=lowerCamelCase__			)

           else:
                 SCREAMING_SNAKE_CASE						=		AdamW(
                     lowerCamelCase__		,lr=self.hparams.learning_rate		,eps=self.hparams.adam_epsilon			)
           SCREAMING_SNAKE_CASE						=		optimizer

           SCREAMING_SNAKE_CASE						=		self.get_lr_scheduler()

           return [optimizer], [scheduler]






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




           '''simple docstring'''


           return self.validation_step(lowerCamelCase__		,lowerCamelCase__			)






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




           '''simple docstring'''


           return self.validation_end(lowerCamelCase__			)






     def 							SCREAMING_SNAKE_CASE__    (				self		:       Any			)	->						int:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE						=		max(1		,self.hparams.gpus			)  # TODO: consider num_tpu_cores
           SCREAMING_SNAKE_CASE						=		self.hparams.train_batch_size * self.hparams.accumulate_grad_batches * num_devices
           return (self.dataset_size / effective_batch_size) * self.hparams.max_epochs






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




           '''simple docstring'''


           if mode == "test":
                 SCREAMING_SNAKE_CASE						=		len(self.test_dataloader().dataset			)
           else:
                 SCREAMING_SNAKE_CASE						=		self.get_dataloader("""train"""		,self.hparams.train_batch_size		,shuffle=lowerCamelCase__			)
                 SCREAMING_SNAKE_CASE						=		len(self.train_dataloader().dataset			)






     def 							SCREAMING_SNAKE_CASE__    (				self		:       Optional[int]		,lowerCamelCase__		:       str		,lowerCamelCase__		:       int		,lowerCamelCase__		:       bool = False			)	->						List[str]:




           '''simple docstring'''


           raise NotImplementedError("""You must implement this for your task"""			)






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




           '''simple docstring'''


           return self.train_loader






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




           '''simple docstring'''


           return self.get_dataloader("""dev"""		,self.hparams.eval_batch_size		,shuffle=lowerCamelCase__			)






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




           '''simple docstring'''


           return self.get_dataloader("""test"""		,self.hparams.eval_batch_size		,shuffle=lowerCamelCase__			)






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




           '''simple docstring'''


           return os.path.join(
               self.hparams.data_dir		,"""cached_{}_{}_{}""".format(
                   lowerCamelCase__		,list(filter(lowerCamelCase__		,self.hparams.model_name_or_path.split("""/"""			)			)			).pop()		,str(self.hparams.max_seq_length			)		,)		,)






     @pl.utilities.rank_zero_only
     def 							SCREAMING_SNAKE_CASE__    (				self		:       str		,lowerCamelCase__		:       Dict[str, Any]			)	->						None:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE						=		self.output_dir.joinpath("""best_tfmr"""			)
           SCREAMING_SNAKE_CASE						=		self.step_count
           self.model.save_pretrained(lowerCamelCase__			)
           self.tokenizer.save_pretrained(lowerCamelCase__			)






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




           '''simple docstring'''


           parser.add_argument(
               """--model_name_or_path"""		,default=lowerCamelCase__		,type=lowerCamelCase__		,required=lowerCamelCase__		,help="""Path to pretrained model or model identifier from huggingface.co/models"""		,)
           parser.add_argument(
               """--config_name"""		,default=""""""		,type=lowerCamelCase__		,help="""Pretrained config name or path if not the same as model_name"""			)
           parser.add_argument(
               """--tokenizer_name"""		,default=lowerCamelCase__		,type=lowerCamelCase__		,help="""Pretrained tokenizer name or path if not the same as model_name"""		,)
           parser.add_argument(
               """--cache_dir"""		,default=""""""		,type=lowerCamelCase__		,help="""Where do you want to store the pre-trained models downloaded from huggingface.co"""		,)
           parser.add_argument(
               """--encoder_layerdrop"""		,type=lowerCamelCase__		,help="""Encoder layer dropout probability (Optional). Goes into model.config"""		,)
           parser.add_argument(
               """--decoder_layerdrop"""		,type=lowerCamelCase__		,help="""Decoder layer dropout probability (Optional). Goes into model.config"""		,)
           parser.add_argument(
               """--dropout"""		,type=lowerCamelCase__		,help="""Dropout probability (Optional). Goes into model.config"""		,)
           parser.add_argument(
               """--attention_dropout"""		,type=lowerCamelCase__		,help="""Attention dropout probability (Optional). Goes into model.config"""		,)
           parser.add_argument("""--learning_rate"""		,default=5e-5		,type=lowerCamelCase__		,help="""The initial learning rate for Adam."""			)
           parser.add_argument(
               """--lr_scheduler"""		,default="""linear"""		,choices=lowerCamelCase__		,metavar=lowerCamelCase__		,type=lowerCamelCase__		,help="""Learning rate scheduler"""		,)
           parser.add_argument("""--weight_decay"""		,default=0.0		,type=lowerCamelCase__		,help="""Weight decay if we apply some."""			)
           parser.add_argument("""--adam_epsilon"""		,default=1e-8		,type=lowerCamelCase__		,help="""Epsilon for Adam optimizer."""			)
           parser.add_argument("""--warmup_steps"""		,default=0		,type=lowerCamelCase__		,help="""Linear warmup over warmup_steps."""			)
           parser.add_argument("""--num_workers"""		,default=4		,type=lowerCamelCase__		,help="""kwarg passed to DataLoader"""			)
           parser.add_argument("""--num_train_epochs"""		,dest="""max_epochs"""		,default=3		,type=lowerCamelCase__			)
           parser.add_argument("""--train_batch_size"""		,default=32		,type=lowerCamelCase__			)
           parser.add_argument("""--eval_batch_size"""		,default=32		,type=lowerCamelCase__			)
           parser.add_argument("""--adafactor"""		,action="""store_true"""			)

class 						UpperCamelCase__					( pl.Callback  ):



     '''simple docstring'''

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




           '''simple docstring'''


           SCREAMING_SNAKE_CASE						=		trainer.lr_schedulers[0]['scheduler']
           SCREAMING_SNAKE_CASE						=		{F"""lr_group_{i}""": lr for i, lr in enumerate(lr_scheduler.get_lr()			)}
           pl_module.logger.log_metrics(lowerCamelCase__			)






     def 							SCREAMING_SNAKE_CASE__    (				self		:       List[Any]		,lowerCamelCase__		:       pl.Trainer		,lowerCamelCase__		:       pl.LightningModule			)	->						Optional[Any]:




           '''simple docstring'''


           rank_zero_info("""***** Validation results *****"""			)
           SCREAMING_SNAKE_CASE						=		trainer.callback_metrics
           # Log results
           for key in sorted(lowerCamelCase__			):
                 if key not in ["log", "progress_bar"]:
                       rank_zero_info("""{} = {}\n""".format(lowerCamelCase__		,str(metrics[key]			)			)			)






     def 							SCREAMING_SNAKE_CASE__    (				self		:       str		,lowerCamelCase__		:       pl.Trainer		,lowerCamelCase__		:       pl.LightningModule			)	->						List[str]:




           '''simple docstring'''


           rank_zero_info("""***** Test results *****"""			)
           SCREAMING_SNAKE_CASE						=		trainer.callback_metrics
           # Log and save results to file
           SCREAMING_SNAKE_CASE						=		os.path.join(pl_module.hparams.output_dir		,"""test_results.txt"""			)
           with open(lowerCamelCase__		,"""w"""			) as writer:
                 for key in sorted(lowerCamelCase__			):
                       if key not in ["log", "progress_bar"]:
                             rank_zero_info("""{} = {}\n""".format(lowerCamelCase__		,str(metrics[key]			)			)			)
                             writer.write("""{} = {}\n""".format(lowerCamelCase__		,str(metrics[key]			)			)			)



def   __lowercase  ( _SCREAMING_SNAKE_CASE					,		_SCREAMING_SNAKE_CASE      )					->  None:


      '''simple docstring'''


      parser.add_argument(
          """--output_dir"""					,		default=_UpperCAmelCase					,		type=_UpperCAmelCase					,		required=_UpperCAmelCase					,		help="""The output directory where the model predictions and checkpoints will be written."""					,		)
      parser.add_argument(
          """--fp16"""					,		action="""store_true"""					,		help="""Whether to use 16-bit (mixed) precision (through NVIDIA apex) instead of 32-bit"""					,		)

      parser.add_argument(
          """--fp16_opt_level"""					,		type=_UpperCAmelCase					,		default="""O2"""					,		help=(
              """For fp16: Apex AMP optimization level selected in [\'O0\', \'O1\', \'O2\', and \'O3\']."""
              """See details at https://nvidia.github.io/apex/amp.html"""
          )					,		)
      parser.add_argument("""--n_tpu_cores"""					,		dest="""tpu_cores"""					,		type=_UpperCAmelCase      )
      parser.add_argument("""--max_grad_norm"""					,		dest="""gradient_clip_val"""					,		default=1.0					,		type=_UpperCAmelCase					,		help="""Max gradient norm"""      )
      parser.add_argument("""--do_train"""					,		action="""store_true"""					,		help="""Whether to run training."""      )
      parser.add_argument("""--do_predict"""					,		action="""store_true"""					,		help="""Whether to run predictions on the test set."""      )
      parser.add_argument(
          """--gradient_accumulation_steps"""					,		dest="""accumulate_grad_batches"""					,		type=_UpperCAmelCase					,		default=1					,		help="""Number of updates steps to accumulate before performing a backward/update pass."""					,		)
      parser.add_argument("""--seed"""					,		type=_UpperCAmelCase					,		default=42					,		help="""random seed for initialization"""      )
      parser.add_argument(
          """--data_dir"""					,		default=_UpperCAmelCase					,		type=_UpperCAmelCase					,		required=_UpperCAmelCase					,		help="""The input data dir. Should contain the training files for the CoNLL-2003 NER task."""					,		)



def   __lowercase  ( _SCREAMING_SNAKE_CASE					,		_SCREAMING_SNAKE_CASE					,		_SCREAMING_SNAKE_CASE=None					,		_SCREAMING_SNAKE_CASE=True					,		_SCREAMING_SNAKE_CASE=[]					,		_SCREAMING_SNAKE_CASE=None					,		_SCREAMING_SNAKE_CASE=None					,		**_SCREAMING_SNAKE_CASE					,		)					->  str:


      '''simple docstring'''


      pl.seed_everything(args.seed      )

      # init model
      SCREAMING_SNAKE_CASE						=		Path(model.hparams.output_dir      )
      odir.mkdir(exist_ok=_UpperCAmelCase      )

      # add custom checkpoints
      if checkpoint_callback is None:
            SCREAMING_SNAKE_CASE						=		pl.callbacks.ModelCheckpoint(
                filepath=args.output_dir					,		prefix="""checkpoint"""					,		monitor="""val_loss"""					,		mode="""min"""					,		save_top_k=1      )
      if early_stopping_callback:
            extra_callbacks.append(_UpperCAmelCase      )
      if logging_callback is None:
            SCREAMING_SNAKE_CASE						=		LoggingCallback()

      SCREAMING_SNAKE_CASE						=		{}

      # TODO: remove with PyTorch 1.6 since pl uses native amp
      if args.fpaa:
            SCREAMING_SNAKE_CASE						=		16
            SCREAMING_SNAKE_CASE						=		args.fpaa_opt_level

      if args.gpus > 1:
            SCREAMING_SNAKE_CASE						=		'ddp'

      SCREAMING_SNAKE_CASE						=		args.accumulate_grad_batches
      SCREAMING_SNAKE_CASE						=		extra_train_kwargs.get("""accelerator"""					,		_UpperCAmelCase      )
      SCREAMING_SNAKE_CASE						=		extra_train_kwargs.get("""profiler"""					,		_UpperCAmelCase      )

      SCREAMING_SNAKE_CASE						=		pl.Trainer.from_argparse_args(
          _UpperCAmelCase					,		weights_summary=_UpperCAmelCase					,		callbacks=[logging_callback] + extra_callbacks					,		logger=_UpperCAmelCase					,		checkpoint_callback=_UpperCAmelCase					,		**_UpperCAmelCase					,		)

      if args.do_train:
            trainer.fit(_UpperCAmelCase      )

      return trainer







"""simple docstring"""




import unittest
from pathlib import Path

from transformers import is_vision_available, load_tool
from transformers.testing_utils import get_tests_dir

from .test_tools_common import ToolTesterMixin


if is_vision_available():
     from PIL import Image


class 	a__		(    unittest.TestCase,	SCREAMING_SNAKE_CASE__  ):


   def 							lowercase					(				self   :				Dict      )     ->							List[str]:
         lowercase	:   str											= load_tool('image-captioning'      )
         self.tool.setup()
         lowercase	:   List[str]											= load_tool('image-captioning',  remote=lowerCAmelCase      )


   def 							lowercase					(				self   :				str      )     ->							int:
         lowercase	:   Union[str, Any]											= Image.open(Path(get_tests_dir('fixtures/tests_samples/COCO'      )      ) / '000000039769.png'      )
         lowercase	:   List[str]											= self.tool(lowerCAmelCase      )
         self.assertEqual(lowerCAmelCase,  'two cats sleeping on a couch'      )


   def 							lowercase					(				self   :				str      )     ->							List[str]:
         lowercase	:   str											= Image.open(Path(get_tests_dir('fixtures/tests_samples/COCO'      )      ) / '000000039769.png'      )
         lowercase	:   Tuple											= self.remote_tool(lowerCAmelCase      )
         self.assertEqual(lowerCAmelCase,  'two cats sleeping on a couch'      )


   def 							lowercase					(				self   :				Tuple      )     ->							Union[str, Any]:
         lowercase	:   Tuple											= Image.open(Path(get_tests_dir('fixtures/tests_samples/COCO'      )      ) / '000000039769.png'      )
         lowercase	:   Union[str, Any]											= self.tool(image=lowerCAmelCase      )
         self.assertEqual(lowerCAmelCase,  'two cats sleeping on a couch'      )






   def 							lowercase					(				self   :				str      )     ->							Optional[int]:
         lowercase	:   List[str]											= Image.open(Path(get_tests_dir('fixtures/tests_samples/COCO'      )      ) / '000000039769.png'      )
         lowercase	:   str											= self.remote_tool(image=lowerCAmelCase      )
         self.assertEqual(lowerCAmelCase,  'two cats sleeping on a couch'      )
