
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


__lowerCAmelCase      :					Union[str, Any]   	=       logging.getLogger(__name__)

require_version('pytorch_lightning>=1.0.4')

__lowerCAmelCase      :					str   	=       {
    '''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
__lowerCAmelCase      :					str   	=       {
    '''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
}
__lowerCAmelCase      :					Dict   	=       sorted(arg_to_scheduler.keys())
__lowerCAmelCase      :					Optional[int]   	=       '''{''' + ''', '''.join(arg_to_scheduler_choices) + '''}'''




class 	UpperCAmelCase_      (			pl.LightningModule					):


  '''simple docstring'''







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







    """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(_A   )
    __magic_name__    					=	0
    __magic_name__    					=	Path(self.hparams.output_dir   )
    __magic_name__    					=	self.hparams.cache_dir if self.hparams.cache_dir else None
    if config is None:
      __magic_name__    					=	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=_A      ,							**_A      ,							)
    else:
      __magic_name__    					=	config

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

    if tokenizer is None:
      __magic_name__    					=	AutoTokenizer.from_pretrained(
          self.hparams.tokenizer_name if self.hparams.tokenizer_name else self.hparams.model_name_or_path      ,							cache_dir=_A      ,							)
    else:
      __magic_name__    					=	tokenizer
    __magic_name__    					=	MODEL_MODES[mode]
    if model is None:
      __magic_name__    					=	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=_A      ,							)
    else:
      __magic_name__    					=	model
  def 							_lowercase							(			self      :			int      ,							*UpperCamelCase__      :			List[Any]      ,							**UpperCamelCase__      :			int   )      ->     List[str]:







    """simple docstring"""






    __magic_name__    					=	self.model_type.from_pretrained(*_A      ,							**_A   )
  def 							_lowercase							(			self      :			List[Any]   )      ->     Tuple:







    """simple docstring"""






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







    """simple docstring"""






    __magic_name__    					=	self.model
    __magic_name__    					=	["bias", "LayerNorm.weight"]
    __magic_name__    					=	[
        {
            "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:
      __magic_name__    					=	Adafactor(
          _A      ,							lr=self.hparams.learning_rate      ,							scale_parameter=_A      ,							relative_step=_A   )

    else:
      __magic_name__    					=	AdamW(
          _A      ,							lr=self.hparams.learning_rate      ,							eps=self.hparams.adam_epsilon   )
    __magic_name__    					=	optimizer

    __magic_name__    					=	self.get_lr_scheduler()

    return [optimizer], [scheduler]
  def 							_lowercase							(			self      :			Any      ,							UpperCamelCase__      :			Union[str, Any]      ,							UpperCamelCase__      :			Any   )      ->     Union[str, Any]:







    """simple docstring"""






    return self.validation_step(_A      ,							_A   )
  def 							_lowercase							(			self      :			Any      ,							UpperCamelCase__      :			Union[str, Any]   )      ->     str:







    """simple docstring"""






    return self.validation_end(_A   )
  def 							_lowercase							(			self      :			Dict   )      ->     Any:







    """simple docstring"""






    __magic_name__    					=	max(1      ,							self.hparams.gpus   )  # TODO: consider num_tpu_cores
    __magic_name__    					=	self.hparams.train_batch_size * self.hparams.accumulate_grad_batches * num_devices
    return (self.dataset_size / effective_batch_size) * self.hparams.max_epochs
  def 							_lowercase							(			self      :			Optional[int]      ,							UpperCamelCase__      :			Optional[int]   )      ->     Tuple:







    """simple docstring"""






    if mode == "test":
      __magic_name__    					=	len(self.test_dataloader().dataset   )
    else:
      __magic_name__    					=	self.get_dataloader("""train"""      ,							self.hparams.train_batch_size      ,							shuffle=_A   )
      __magic_name__    					=	len(self.train_dataloader().dataset   )
  def 							_lowercase							(			self      :			Optional[int]      ,							UpperCamelCase__      :			str      ,							UpperCamelCase__      :			int      ,							UpperCamelCase__      :			bool = False   )      ->     str:







    """simple docstring"""






    raise NotImplementedError("""You must implement this for your task"""   )
  def 							_lowercase							(			self      :			List[str]   )      ->     Optional[int]:







    """simple docstring"""






    return self.train_loader
  def 							_lowercase							(			self      :			List[Any]   )      ->     str:







    """simple docstring"""






    return self.get_dataloader("""dev"""      ,							self.hparams.eval_batch_size      ,							shuffle=_A   )
  def 							_lowercase							(			self      :			int   )      ->     Union[str, Any]:







    """simple docstring"""






    return self.get_dataloader("""test"""      ,							self.hparams.eval_batch_size      ,							shuffle=_A   )
  def 							_lowercase							(			self      :			Union[str, Any]      ,							UpperCamelCase__      :			str   )      ->     str:







    """simple docstring"""






    return os.path.join(
        self.hparams.data_dir      ,							"""cached_{}_{}_{}""".format(
            _A      ,							list(filter(_A      ,							self.hparams.model_name_or_path.split("""/"""   )   )   ).pop()      ,							str(self.hparams.max_seq_length   )      ,							)      ,							)
  @pl.utilities.rank_zero_only
  def 							_lowercase							(			self      :			List[Any]      ,							UpperCamelCase__      :			Dict[str, Any]   )      ->     Union[str, Any]:







    """simple docstring"""






    __magic_name__    					=	self.output_dir.joinpath("""best_tfmr"""   )
    __magic_name__    					=	self.step_count
    self.model.save_pretrained(_A   )
    self.tokenizer.save_pretrained(_A   )





  @staticmethod
  def 							_lowercase							(			UpperCamelCase__      :			Any      ,							UpperCamelCase__      :			List[Any]   )      ->     Union[str, Any]:







    """simple docstring"""






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







class 	UpperCAmelCase_      (			pl.Callback					):


  '''simple docstring'''







  def 							_lowercase							(			self      :			int      ,							UpperCamelCase__      :			List[str]      ,							UpperCamelCase__      :			List[Any]   )      ->     str:







    """simple docstring"""






    __magic_name__    					=	trainer.lr_schedulers[0]["scheduler"]
    __magic_name__    					=	{F'''lr_group_{i}''': lr for i, lr in enumerate(lr_scheduler.get_lr()   )}
    pl_module.logger.log_metrics(_A   )
  def 							_lowercase							(			self      :			Optional[int]      ,							UpperCamelCase__      :			pl.Trainer      ,							UpperCamelCase__      :			pl.LightningModule   )      ->     Union[str, Any]:







    """simple docstring"""






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





  def 							_lowercase							(			self      :			Tuple      ,							UpperCamelCase__      :			pl.Trainer      ,							UpperCamelCase__      :			pl.LightningModule   )      ->     str:







    """simple docstring"""






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







def a__   (	A_,					A_ ):






  '''simple docstring'''





  parser.add_argument(
      """--output_dir""",					default=A_,					type=A_,					required=A_,					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=A_,					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=A_ )
  parser.add_argument("""--max_grad_norm""",					dest="""gradient_clip_val""",					default=1.0,					type=A_,					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=A_,					default=1,					help="""Number of updates steps to accumulate before performing a backward/update pass.""",					)
  parser.add_argument("""--seed""",					type=A_,					default=42,					help="""random seed for initialization""" )
  parser.add_argument(
      """--data_dir""",					default=A_,					type=A_,					required=A_,					help="""The input data dir. Should contain the training files for the CoNLL-2003 NER task.""",					)





def a__   (	A_,					A_,					A_=None,					A_=True,					A_=[],					A_=None,					A_=None,					**A_,					):






  '''simple docstring'''





  pl.seed_everything(args.seed )

  # init model
  __magic_name__    					=	Path(model.hparams.output_dir )
  odir.mkdir(exist_ok=A_ )

  # add custom checkpoints
  if checkpoint_callback is None:
    __magic_name__    					=	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(A_ )
  if logging_callback is None:
    __magic_name__    					=	LoggingCallback()

  __magic_name__    					=	{}

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

  if args.gpus > 1:
    __magic_name__    					=	"ddp"

  __magic_name__    					=	args.accumulate_grad_batches
  __magic_name__    					=	extra_train_kwargs.get("""accelerator""",					A_ )
  __magic_name__    					=	extra_train_kwargs.get("""profiler""",					A_ )

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

  if args.do_train:
    trainer.fit(A_ )

  return trainer



import torch.nn as nn

from .imports import is_fpa_available


if is_fpa_available():
    import transformer_engine.pytorch as te

def a__   (	A_,					A_=True,					A_=True,					A_=True ):






  '''simple docstring'''





  if not is_fpa_available():
    raise ImportError("""Using `convert_model` requires transformer_engine to be installed.""" )
  for name, module in model.named_children():
    if isinstance(A_,					nn.Linear ) and to_transformer_engine and _convert_linear:
      # Return early if the linear layer weights are not multiples of 16
      if any(p % 16 != 0 for p in module.weight.shape ):
        return
      __magic_name__    					=	module.bias is not None
      __magic_name__    					=	te.Linear(
          module.in_features,					module.out_features,					bias=A_,					params_dtype=module.weight.dtype )
      __magic_name__    					=	module.weight.data.clone()
      if has_bias:
        __magic_name__    					=	module.bias.data.clone()

      setattr(A_,					A_,					A_ )
    elif isinstance(A_,					nn.LayerNorm ) and to_transformer_engine and _convert_ln:
      __magic_name__    					=	te.LayerNorm(module.normalized_shape[0],					eps=module.eps,					params_dtype=module.weight.dtype )
      __magic_name__    					=	module.weight.data.clone()
      __magic_name__    					=	module.bias.data.clone()

      setattr(A_,					A_,					A_ )
    elif isinstance(A_,					te.Linear ) and not to_transformer_engine and _convert_linear:
      __magic_name__    					=	module.bias is not None
      __magic_name__    					=	nn.Linear(
          module.in_features,					module.out_features,					bias=A_,					params_dtype=module.weight.dtype )
      __magic_name__    					=	module.weight.data.clone()
      if has_bias:
        __magic_name__    					=	module.bias.data.clone()

      setattr(A_,					A_,					A_ )
    elif isinstance(A_,					te.LayerNorm ) and not to_transformer_engine and _convert_ln:
      __magic_name__    					=	nn.LayerNorm(module.normalized_shape[0],					eps=module.eps,					params_dtype=module.weight.dtype )
      __magic_name__    					=	module.weight.data.clone()
      __magic_name__    					=	module.bias.data.clone()

      setattr(A_,					A_,					A_ )
    else:
      convert_model(
          A_,					to_transformer_engine=A_,					_convert_linear=_convert_linear,					_convert_ln=_convert_ln,					)





def a__   (	A_ ):






  '''simple docstring'''





  if not is_fpa_available():
    raise ImportError("""Using `has_transformer_engine_layers` requires transformer_engine to be installed.""" )
  for m in model.modules():
    if isinstance(A_,					(te.LayerNorm, te.Linear) ):
      return True
  return False

