





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


UpperCamelCase__          =logging.getLogger(__name__)

require_version('pytorch_lightning>=1.0.4')

UpperCamelCase__          ={
    '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
UpperCamelCase__          ={
    '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
}
UpperCamelCase__          =sorted(arg_to_scheduler.keys())
UpperCamelCase__          ='{' + ', '.join(arg_to_scheduler_choices) + '}'





class 	lowerCAmelCase__(       pl.LightningModule			):






       '''simple docstring'''


       def __init__(					self			,  __lowerCamelCase			,  __lowerCamelCase=None			,  __lowerCamelCase="base"			,  __lowerCamelCase=None			,  __lowerCamelCase=None			,  __lowerCamelCase=None			,  **__lowerCamelCase			,  )	->   List[str]:
            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				:		int  			=							0
            _SCREAMING_SNAKE_CASE				:		List[str]  			=							Path(self.hparams.output_dir							)
            _SCREAMING_SNAKE_CASE				:		int  			=							self.hparams.cache_dir if self.hparams.cache_dir else None
            if config is None:
                 _SCREAMING_SNAKE_CASE				:		Union[str, Any]  			=							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				:		PretrainedConfig  			=							config

            _SCREAMING_SNAKE_CASE				:		List[str]  			=							("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				:		Tuple  			=							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				:		PreTrainedTokenizer  			=							tokenizer
            _SCREAMING_SNAKE_CASE				:		int  			=							MODEL_MODES[mode]
            if model is None:
                 _SCREAMING_SNAKE_CASE				:		Any  			=							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				:		Union[str, Any]  			=							model




       def 			UpperCamelCase_	(					self			,  *__lowerCamelCase			,  **__lowerCamelCase							)	->   List[str]:
            _SCREAMING_SNAKE_CASE				:		str  			=							self.model_type.from_pretrained(*__lowerCamelCase			,  **__lowerCamelCase							)




       def 			UpperCamelCase_	(					self							)	->   str:
            _SCREAMING_SNAKE_CASE				:		Union[str, Any]  			=							arg_to_scheduler[self.hparams.lr_scheduler]
            _SCREAMING_SNAKE_CASE				:		Dict  			=							get_schedule_func(
                self.opt			,  num_warmup_steps=self.hparams.warmup_steps			,  num_training_steps=self.total_steps()							)
            _SCREAMING_SNAKE_CASE				:		Optional[int]  			=							{"scheduler": scheduler, "interval": "step", "frequency": 1}
            return scheduler




       def 			UpperCamelCase_	(					self							)	->   Union[str, Any]:
            _SCREAMING_SNAKE_CASE				:		str  			=							self.model
            _SCREAMING_SNAKE_CASE				:		List[str]  			=							["bias", "LayerNorm.weight"]
            _SCREAMING_SNAKE_CASE				:		List[Any]  			=							[
                {
                    "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				:		str  			=							Adafactor(
                     __lowerCamelCase			,  lr=self.hparams.learning_rate			,  scale_parameter=__lowerCamelCase			,  relative_step=__lowerCamelCase							)

            else:
                 _SCREAMING_SNAKE_CASE				:		int  			=							AdamW(
                     __lowerCamelCase			,  lr=self.hparams.learning_rate			,  eps=self.hparams.adam_epsilon							)
            _SCREAMING_SNAKE_CASE				:		int  			=							optimizer

            _SCREAMING_SNAKE_CASE				:		Optional[int]  			=							self.get_lr_scheduler()

            return [optimizer], [scheduler]




       def 			UpperCamelCase_	(					self			,  __lowerCamelCase			,  __lowerCamelCase							)	->   Tuple:
            return self.validation_step(__lowerCamelCase			,  __lowerCamelCase							)




       def 			UpperCamelCase_	(					self			,  __lowerCamelCase							)	->   Dict:
            return self.validation_end(__lowerCamelCase							)




       def 			UpperCamelCase_	(					self							)	->   int:
            _SCREAMING_SNAKE_CASE				:		Optional[int]  			=							max(1			,  self.hparams.gpus							)  # TODO: consider num_tpu_cores
            _SCREAMING_SNAKE_CASE				:		Union[str, Any]  			=							self.hparams.train_batch_size * self.hparams.accumulate_grad_batches * num_devices
            return (self.dataset_size / effective_batch_size) * self.hparams.max_epochs




       def 			UpperCamelCase_	(					self			,  __lowerCamelCase							)	->   Dict:
            if mode == "test":
                 _SCREAMING_SNAKE_CASE				:		str  			=							len(self.test_dataloader().dataset							)
            else:
                 _SCREAMING_SNAKE_CASE				:		Optional[Any]  			=							self.get_dataloader("train"			,  self.hparams.train_batch_size			,  shuffle=__lowerCamelCase							)
                 _SCREAMING_SNAKE_CASE				:		Optional[int]  			=							len(self.train_dataloader().dataset							)




       def 			UpperCamelCase_	(					self			,  __lowerCamelCase			,  __lowerCamelCase			,  __lowerCamelCase = False							)	->   Dict:
            raise NotImplementedError("You must implement this for your task"							)




       def 			UpperCamelCase_	(					self							)	->   int:
            return self.train_loader




       def 			UpperCamelCase_	(					self							)	->   Tuple:
            return self.get_dataloader("dev"			,  self.hparams.eval_batch_size			,  shuffle=__lowerCamelCase							)




       def 			UpperCamelCase_	(					self							)	->   List[Any]:
            return self.get_dataloader("test"			,  self.hparams.eval_batch_size			,  shuffle=__lowerCamelCase							)




       def 			UpperCamelCase_	(					self			,  __lowerCamelCase							)	->   int:
            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 			UpperCamelCase_	(					self			,  __lowerCamelCase							)	->   None:
            _SCREAMING_SNAKE_CASE				:		int  			=							self.output_dir.joinpath("best_tfmr"							)
            _SCREAMING_SNAKE_CASE				:		str  			=							self.step_count
            self.model.save_pretrained(__lowerCamelCase							)
            self.tokenizer.save_pretrained(__lowerCamelCase							)







       @staticmethod
       def 			UpperCamelCase_	(					__lowerCamelCase			,  __lowerCamelCase							)	->   str:
            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=3_2			,  type=__lowerCamelCase							)
            parser.add_argument("--eval_batch_size"			,  default=3_2			,  type=__lowerCamelCase							)
            parser.add_argument("--adafactor"			,  action="store_true"							)





class 	lowerCAmelCase__(       pl.Callback			):






       '''simple docstring'''


       def 			UpperCamelCase_	(					self			,  __lowerCamelCase			,  __lowerCamelCase							)	->   Tuple:
            _SCREAMING_SNAKE_CASE				:		Union[str, Any]  			=							trainer.lr_schedulers[0]["scheduler"]
            _SCREAMING_SNAKE_CASE				:		Union[str, Any]  			=							{F"""lr_group_{i}""": lr for i, lr in enumerate(lr_scheduler.get_lr()							)}
            pl_module.logger.log_metrics(__lowerCamelCase							)




       def 			UpperCamelCase_	(					self			,  __lowerCamelCase			,  __lowerCamelCase							)	->   Union[str, Any]:
            rank_zero_info("***** Validation results *****"							)
            _SCREAMING_SNAKE_CASE				:		Any  			=							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 			UpperCamelCase_	(					self			,  __lowerCamelCase			,  __lowerCamelCase							)	->   str:
            rank_zero_info("***** Test results *****"							)
            _SCREAMING_SNAKE_CASE				:		int  			=							trainer.callback_metrics
            # Log and save results to file
            _SCREAMING_SNAKE_CASE				:		str  			=							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 			lowerCamelCase__	(__lowerCamelCase,				__lowerCamelCase		):
     #  To allow all pl args uncomment the following line
     #  parser = pl.Trainer.add_argparse_args(parser)
     parser.add_argument(
         "--output_dir",				default=__lowerCamelCase,				type=__lowerCamelCase,				required=__lowerCamelCase,				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=__lowerCamelCase,				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=__lowerCamelCase		)
     parser.add_argument("--max_grad_norm",				dest="gradient_clip_val",				default=1.0,				type=__lowerCamelCase,				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=__lowerCamelCase,				default=1,				help="Number of updates steps to accumulate before performing a backward/update pass.",				)
     parser.add_argument("--seed",				type=__lowerCamelCase,				default=42,				help="random seed for initialization"		)
     parser.add_argument(
         "--data_dir",				default=__lowerCamelCase,				type=__lowerCamelCase,				required=__lowerCamelCase,				help="The input data dir. Should contain the training files for the CoNLL-2003 NER task.",				)






def 			lowerCamelCase__	(__lowerCamelCase,				__lowerCamelCase,				__lowerCamelCase=None,				__lowerCamelCase=True,				__lowerCamelCase=[],				__lowerCamelCase=None,				__lowerCamelCase=None,				**__lowerCamelCase,				):
     pl.seed_everything(args.seed		)

     # init model
     _SCREAMING_SNAKE_CASE				:		List[str]  			=							Path(model.hparams.output_dir		)
     odir.mkdir(exist_ok=__lowerCamelCase		)

     # add custom checkpoints
     if checkpoint_callback is None:
          _SCREAMING_SNAKE_CASE				:		List[str]  			=							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(__lowerCamelCase		)
     if logging_callback is None:
          _SCREAMING_SNAKE_CASE				:		Any  			=							LoggingCallback()

     _SCREAMING_SNAKE_CASE				:		List[Any]  			=							{}

     # TODO: remove with PyTorch 1.6 since pl uses native amp
     if args.fpaa:
          _SCREAMING_SNAKE_CASE				:		Union[str, Any]  			=							16
          _SCREAMING_SNAKE_CASE				:		int  			=							args.fpaa_opt_level

     if args.gpus > 1:
          _SCREAMING_SNAKE_CASE				:		str  			=							"ddp"

     _SCREAMING_SNAKE_CASE				:		int  			=							args.accumulate_grad_batches
     _SCREAMING_SNAKE_CASE				:		Optional[int]  			=							extra_train_kwargs.get("accelerator",				__lowerCamelCase		)
     _SCREAMING_SNAKE_CASE				:		Optional[int]  			=							extra_train_kwargs.get("profiler",				__lowerCamelCase		)

     _SCREAMING_SNAKE_CASE				:		List[Any]  			=							pl.Trainer.from_argparse_args(
         __lowerCamelCase,				weights_summary=__lowerCamelCase,				callbacks=[logging_callback] + extra_callbacks,				logger=__lowerCamelCase,				checkpoint_callback=__lowerCamelCase,				**__lowerCamelCase,				)

     if args.do_train:
          trainer.fit(__lowerCamelCase		)

     return trainer






import warnings
from dataclasses import dataclass, field
from typing import Optional, Tuple

from .training_args import TrainingArguments
from .utils import cached_property, is_tf_available, logging, requires_backends


UpperCamelCase__          =logging.get_logger(__name__)

if is_tf_available():
       import tensorflow as tf





@dataclass
class 	lowerCAmelCase__(       __lowercase			):






       '''simple docstring'''


       __snake_case       	=			'tf'
       __snake_case       	=			field(
           default=__lowercase     ,   metadata={'help': 'Name of TPU'}     ,   )

       __snake_case       	=			field(
           default=__lowercase     ,   metadata={'help': 'Zone of TPU'}     ,   )

       __snake_case       	=			field(
           default=__lowercase     ,   metadata={'help': 'Name of Cloud TPU-enabled project'}     ,   )

       __snake_case       	=			field(
           default=1.0     ,   metadata={'help': 'Power for the Polynomial decay LR scheduler.'}     ,   )

       __snake_case       	=			field(default=__lowercase     ,   metadata={'help': 'Whether to activate the XLA compilation or not'}			)




       @cached_property
       def 			UpperCamelCase_	(					self							)	->   Tuple["tf.distribute.Strategy", int]:
            requires_backends(self			,  ["tf"]							)
            logger.info("Tensorflow: setting up strategy"							)

            _SCREAMING_SNAKE_CASE				:		List[Any]  			=							tf.config.list_physical_devices("GPU"							)

            # Set to float16 at first
            if self.fpaa:
                 tf.keras.mixed_precision.set_global_policy("mixed_float16"							)

            if self.no_cuda:
                 _SCREAMING_SNAKE_CASE				:		List[Any]  			=							tf.distribute.OneDeviceStrategy(device="/cpu:0"							)
            else:
                 try:
                      if self.tpu_name:
                           _SCREAMING_SNAKE_CASE				:		Any  			=							tf.distribute.cluster_resolver.TPUClusterResolver(
                               self.tpu_name			,  zone=self.tpu_zone			,  project=self.gcp_project							)
                      else:
                           _SCREAMING_SNAKE_CASE				:		str  			=							tf.distribute.cluster_resolver.TPUClusterResolver()
                 except ValueError:
                      if self.tpu_name:
                           raise RuntimeError(F"""Couldn't connect to TPU {self.tpu_name}!"""							)
                      else:
                           _SCREAMING_SNAKE_CASE				:		Any  			=							None

                 if tpu:
                      # Set to bfloat16 in case of TPU
                      if self.fpaa:
                           tf.keras.mixed_precision.set_global_policy("mixed_bfloat16"							)

                      tf.config.experimental_connect_to_cluster(__lowerCamelCase							)
                      tf.tpu.experimental.initialize_tpu_system(__lowerCamelCase							)

                      _SCREAMING_SNAKE_CASE				:		List[str]  			=							tf.distribute.TPUStrategy(__lowerCamelCase							)

                 elif len(__lowerCamelCase							) == 0:
                      _SCREAMING_SNAKE_CASE				:		Dict  			=							tf.distribute.OneDeviceStrategy(device="/cpu:0"							)
                 elif len(__lowerCamelCase							) == 1:
                      _SCREAMING_SNAKE_CASE				:		Dict  			=							tf.distribute.OneDeviceStrategy(device="/gpu:0"							)
                 elif len(__lowerCamelCase							) > 1:
                      # If you only want to use a specific subset of GPUs use `CUDA_VISIBLE_DEVICES=0`
                      _SCREAMING_SNAKE_CASE				:		Any  			=							tf.distribute.MirroredStrategy()
                 else:
                      raise ValueError("Cannot find the proper strategy, please check your environment properties."							)

            return strategy




       @property
       def 			UpperCamelCase_	(					self							)	->   "tf.distribute.Strategy":
            requires_backends(self			,  ["tf"]							)
            return self._setup_strategy




       @property
       def 			UpperCamelCase_	(					self							)	->   int:
            requires_backends(self			,  ["tf"]							)
            return self._setup_strategy.num_replicas_in_sync




       @property
       def 			UpperCamelCase_	(					self							)	->   List[Any]:
            return False  # TF Logging is handled by Keras not the Trainer




       @property
       def 			UpperCamelCase_	(					self							)	->   int:
            if self.per_gpu_train_batch_size:
                 logger.warning(
                     "Using deprecated `--per_gpu_train_batch_size` argument which will be removed in a future "
                     "version. Using `--per_device_train_batch_size` is preferred."							)
            _SCREAMING_SNAKE_CASE				:		Tuple  			=							self.per_gpu_train_batch_size or self.per_device_train_batch_size
            return per_device_batch_size * self.n_replicas




       @property
       def 			UpperCamelCase_	(					self							)	->   int:
            if self.per_gpu_eval_batch_size:
                 logger.warning(
                     "Using deprecated `--per_gpu_eval_batch_size` argument which will be removed in a future "
                     "version. Using `--per_device_eval_batch_size` is preferred."							)
            _SCREAMING_SNAKE_CASE				:		int  			=							self.per_gpu_eval_batch_size or self.per_device_eval_batch_size
            return per_device_batch_size * self.n_replicas







       @property
       def 			UpperCamelCase_	(					self							)	->   int:
            requires_backends(self			,  ["tf"]							)
            warnings.warn(
                "The n_gpu argument is deprecated and will be removed in a future version, use n_replicas instead."			,  __lowerCamelCase			,  )
            return self._setup_strategy.num_replicas_in_sync