
"""simple docstring"""





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


_a									=				logging.getLogger(__name__)

require_version('pytorch_lightning>=1.0.4')

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


class A_      (pl.LightningModule ):




  '''simple docstring'''







  def __init__(						self							,  lowercase_							,  lowercase_=None							,  lowercase_="base"							,  lowercase_=None							,  lowercase_=None							,  lowercase_=None							,  **lowercase_							,  ):

     """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(lowercase_    )
     UpperCAmelCase_   :   Tuple     	=       0
     UpperCAmelCase_   :   Optional[int]     	=       Path(self.hparams.output_dir    )
     UpperCAmelCase_   :   Optional[Any]     	=       self.hparams.cache_dir if self.hparams.cache_dir else None
     if config is None:
        UpperCAmelCase_   :   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=lowercase_							,  **lowercase_							,  )
     else:
        UpperCAmelCase_   :   PretrainedConfig     	=       config

     UpperCAmelCase_   :   List[str]     	=       ("encoder_layerdrop", "decoder_layerdrop", "dropout", "attention_dropout")
     for p in extra_model_params:
        if getattr(self.hparams							,  lowercase_							,  lowercase_    ):
           assert hasattr(self.config							,  lowercase_    ), F"""model config doesn't have a `{p}` attribute"""
           setattr(self.config							,  lowercase_							,  getattr(self.hparams							,  lowercase_    )    )

     if tokenizer is None:
        UpperCAmelCase_   :   Optional[Any]     	=       AutoTokenizer.from_pretrained(
            self.hparams.tokenizer_name if self.hparams.tokenizer_name else self.hparams.model_name_or_path							,  cache_dir=lowercase_							,  )
     else:
        UpperCAmelCase_   :   PreTrainedTokenizer     	=       tokenizer
     UpperCAmelCase_   :   str     	=       MODEL_MODES[mode]
     if model is None:
        UpperCAmelCase_   :   List[str]     	=       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=lowercase_							,  )
     else:
        UpperCAmelCase_   :   int     	=       model







  def 			UpperCamelCase__   (						self							,  *lowercase_							,  **lowercase_    ):

     """simple docstring"""
     UpperCAmelCase_   :   Optional[int]     	=       self.model_type.from_pretrained(*lowercase_							,  **lowercase_    )







  def 			UpperCamelCase__   (						self    ):

     """simple docstring"""
     UpperCAmelCase_   :   List[str]     	=       arg_to_scheduler[self.hparams.lr_scheduler]
     UpperCAmelCase_   :   Optional[Any]     	=       get_schedule_func(
         self.opt							,  num_warmup_steps=self.hparams.warmup_steps							,  num_training_steps=self.total_steps()    )
     UpperCAmelCase_   :   Optional[Any]     	=       {"scheduler": scheduler, "interval": "step", "frequency": 1}
     return scheduler







  def 			UpperCamelCase__   (						self    ):

     """simple docstring"""
     UpperCAmelCase_   :   Dict     	=       self.model
     UpperCAmelCase_   :   Any     	=       ["bias", "LayerNorm.weight"]
     UpperCAmelCase_   :   Dict     	=       [
         {
             "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:
        UpperCAmelCase_   :   Any     	=       Adafactor(
            lowercase_							,  lr=self.hparams.learning_rate							,  scale_parameter=lowercase_							,  relative_step=lowercase_    )

     else:
        UpperCAmelCase_   :   str     	=       AdamW(
            lowercase_							,  lr=self.hparams.learning_rate							,  eps=self.hparams.adam_epsilon    )
     UpperCAmelCase_   :   Union[str, Any]     	=       optimizer

     UpperCAmelCase_   :   Union[str, Any]     	=       self.get_lr_scheduler()

     return [optimizer], [scheduler]







  def 			UpperCamelCase__   (						self							,  lowercase_							,  lowercase_    ):

     """simple docstring"""
     return self.validation_step(lowercase_							,  lowercase_    )







  def 			UpperCamelCase__   (						self							,  lowercase_    ):

     """simple docstring"""
     return self.validation_end(lowercase_    )







  def 			UpperCamelCase__   (						self    ):

     """simple docstring"""
     UpperCAmelCase_   :   Optional[Any]     	=       max(1							,  self.hparams.gpus    )  # TODO: consider num_tpu_cores
     UpperCAmelCase_   :   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							,  lowercase_    ):

     """simple docstring"""
     if mode == "test":
        UpperCAmelCase_   :   Dict     	=       len(self.test_dataloader().dataset    )
     else:
        UpperCAmelCase_   :   Union[str, Any]     	=       self.get_dataloader("train"							,  self.hparams.train_batch_size							,  shuffle=lowercase_    )
        UpperCAmelCase_   :   Union[str, Any]     	=       len(self.train_dataloader().dataset    )







  def 			UpperCamelCase__   (						self							,  lowercase_							,  lowercase_							,  lowercase_ = False    ):

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







  def 			UpperCamelCase__   (						self    ):

     """simple docstring"""
     return self.train_loader







  def 			UpperCamelCase__   (						self    ):

     """simple docstring"""
     return self.get_dataloader("dev"							,  self.hparams.eval_batch_size							,  shuffle=lowercase_    )







  def 			UpperCamelCase__   (						self    ):

     """simple docstring"""
     return self.get_dataloader("test"							,  self.hparams.eval_batch_size							,  shuffle=lowercase_    )







  def 			UpperCamelCase__   (						self							,  lowercase_    ):

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







  @pl.utilities.rank_zero_only
  def 			UpperCamelCase__   (						self							,  lowercase_    ):

     """simple docstring"""
     UpperCAmelCase_   :   str     	=       self.output_dir.joinpath("best_tfmr"    )
     UpperCAmelCase_   :   Optional[int]     	=       self.step_count
     self.model.save_pretrained(lowercase_    )
     self.tokenizer.save_pretrained(lowercase_    )







  @staticmethod
  def 			UpperCamelCase__   (						lowercase_							,  lowercase_    ):

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




class A_      (pl.Callback ):




  '''simple docstring'''







  def 			UpperCamelCase__   (						self							,  lowercase_							,  lowercase_    ):

     """simple docstring"""
     UpperCAmelCase_   :   List[str]     	=       trainer.lr_schedulers[0]["scheduler"]
     UpperCAmelCase_   :   List[Any]     	=       {F"""lr_group_{i}""": lr for i, lr in enumerate(lr_scheduler.get_lr()    )}
     pl_module.logger.log_metrics(lowercase_    )







  def 			UpperCamelCase__   (						self							,  lowercase_							,  lowercase_    ):

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







  def 			UpperCamelCase__   (						self							,  lowercase_							,  lowercase_    ):

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







def       __a	(						__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       __a	(						__lowerCamelCase,						__lowerCamelCase,						__lowerCamelCase=None,						__lowerCamelCase=True,						__lowerCamelCase=[],						__lowerCamelCase=None,						__lowerCamelCase=None,						**__lowerCamelCase,						):
   pl.seed_everything(args.seed    )

   # init model
   UpperCAmelCase_   :   Any     	=       Path(model.hparams.output_dir    )
   odir.mkdir(exist_ok=__lowerCamelCase    )

   # add custom checkpoints
   if checkpoint_callback is None:
      UpperCAmelCase_   :   Optional[int]     	=       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:
      UpperCAmelCase_   :   Union[str, Any]     	=       LoggingCallback()

   UpperCAmelCase_   :   int     	=       {}

   # TODO: remove with PyTorch 1.6 since pl uses native amp
   if args.fpaa:
      UpperCAmelCase_   :   List[Any]     	=       16
      UpperCAmelCase_   :   str     	=       args.fpaa_opt_level

   if args.gpus > 1:
      UpperCAmelCase_   :   List[str]     	=       "ddp"

   UpperCAmelCase_   :   List[str]     	=       args.accumulate_grad_batches
   UpperCAmelCase_   :   Tuple     	=       extra_train_kwargs.get("accelerator",						__lowerCamelCase    )
   UpperCAmelCase_   :   List[str]     	=       extra_train_kwargs.get("profiler",						__lowerCamelCase    )

   UpperCAmelCase_   :   int     	=       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 os

from packaging import version

from .. import __version__
from .accelerate_utils import apply_forward_hook
from .constants import (
    CONFIG_NAME,
    DEPRECATED_REVISION_ARGS,
    DIFFUSERS_CACHE,
    DIFFUSERS_DYNAMIC_MODULE_NAME,
    FLAX_WEIGHTS_NAME,
    HF_MODULES_CACHE,
    HUGGINGFACE_CO_RESOLVE_ENDPOINT,
    ONNX_EXTERNAL_WEIGHTS_NAME,
    ONNX_WEIGHTS_NAME,
    SAFETENSORS_WEIGHTS_NAME,
    TEXT_ENCODER_ATTN_MODULE,
    WEIGHTS_NAME,
)
from .deprecation_utils import deprecate
from .doc_utils import replace_example_docstring
from .dynamic_modules_utils import get_class_from_dynamic_module
from .hub_utils import (
    HF_HUB_OFFLINE,
    _add_variant,
    _get_model_file,
    extract_commit_hash,
    http_user_agent,
)
from .import_utils import (
    BACKENDS_MAPPING,
    ENV_VARS_TRUE_AND_AUTO_VALUES,
    ENV_VARS_TRUE_VALUES,
    USE_JAX,
    USE_TF,
    USE_TORCH,
    DummyObject,
    OptionalDependencyNotAvailable,
    is_accelerate_available,
    is_accelerate_version,
    is_bsa_available,
    is_flax_available,
    is_ftfy_available,
    is_inflect_available,
    is_invisible_watermark_available,
    is_k_diffusion_available,
    is_k_diffusion_version,
    is_librosa_available,
    is_note_seq_available,
    is_omegaconf_available,
    is_onnx_available,
    is_safetensors_available,
    is_scipy_available,
    is_tensorboard_available,
    is_tf_available,
    is_torch_available,
    is_torch_version,
    is_torchsde_available,
    is_transformers_available,
    is_transformers_version,
    is_unidecode_available,
    is_wandb_available,
    is_xformers_available,
    requires_backends,
)
from .logging import get_logger
from .outputs import BaseOutput
from .pil_utils import PIL_INTERPOLATION, numpy_to_pil, pt_to_pil
from .torch_utils import is_compiled_module, randn_tensor


if is_torch_available():
		from .testing_utils import (
		    floats_tensor,
		    load_hf_numpy,
		    load_image,
		    load_numpy,
		    load_pt,
		    nightly,
		    parse_flag_from_env,
		    print_tensor_test,
		    require_torch_a,
		    require_torch_gpu,
		    skip_mps,
		    slow,
		    torch_all_close,
		    torch_device,
		)
		from .torch_utils import maybe_allow_in_graph

from .testing_utils import export_to_gif, export_to_video


__magic_name__       =     get_logger(__name__)





def 		__snake_case						(    _UpperCAmelCase				):
					"""simple docstring"""






					if version.parse(_UpperCAmelCase				) < version.parse(_UpperCAmelCase				):
										if "dev" in min_version:
															lowercase						   =       (
															    'This example requires a source install from HuggingFace diffusers (see '
															    '`https://huggingface.co/docs/diffusers/installation#install-from-source`),'
															)
										else:
															lowercase						   =       f"""This example requires a minimum version of {min_version},"""
										error_message += f""" but the version found is {__version__}.\n"""
										raise ImportError(_UpperCAmelCase				)






