


'''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


__lowercase:		Any							=     logging.getLogger(__name__)

require_version("pytorch_lightning>=1.0.4")

__lowercase:		int							=     {
    "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
__lowercase:		Dict							=     {
    "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
}
__lowercase:		List[Any]							=     sorted(arg_to_scheduler.keys())
__lowercase:		Tuple							=     "{" + ", ".join(arg_to_scheduler_choices) + "}"







class  UpperCAmelCase (       pl.LightningModule):

 def __init__( self :	Dict,					a_ :	argparse.Namespace,					a_ :	Dict=None,					a_ :	Tuple="base",					a_ :	int=None,					a_ :	int=None,					a_ :	Any=None,					**a_ :	Any,					):





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

    UpperCamelCase__ 						=						("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:
       UpperCamelCase__ 						=						AutoTokenizer.from_pretrained(
           self.hparams.tokenizer_name if self.hparams.tokenizer_name else self.hparams.model_name_or_path,					cache_dir=a_,					)
    else:
       UpperCamelCase__ 						=						tokenizer
    UpperCamelCase__ 						=						MODEL_MODES[mode]
    if model is None:
       UpperCamelCase__ 						=						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:
       UpperCamelCase__ 						=						model

 def 		lowercase_						( self :	List[str],					*a_ :	Union[str, Any],					**a_ :	Optional[int]					):





    """simple docstring"""



    UpperCamelCase__ 						=						self.model_type.from_pretrained(*a_,					**a_					)

 def 		lowercase_						( self :	Optional[int]					):





    """simple docstring"""



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

 def 		lowercase_						( self :	Optional[Any]					):





    """simple docstring"""



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

    else:
       UpperCamelCase__ 						=						AdamW(
           a_,					lr=self.hparams.learning_rate,					eps=self.hparams.adam_epsilon					)
    UpperCamelCase__ 						=						optimizer

    UpperCamelCase__ 						=						self.get_lr_scheduler()

    return [optimizer], [scheduler]

 def 		lowercase_						( self :	str,					a_ :	List[str],					a_ :	Union[str, Any]					):





    """simple docstring"""



    return self.validation_step(a_,					a_					)

 def 		lowercase_						( self :	int,					a_ :	Dict					):





    """simple docstring"""



    return self.validation_end(a_					)

 def 		lowercase_						( self :	Any					):





    """simple docstring"""



    UpperCamelCase__ 						=						max(1,					self.hparams.gpus					)  # TODO: consider num_tpu_cores
    UpperCamelCase__ 						=						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 :	str,					a_ :	List[str]					):





    """simple docstring"""



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

 def 		lowercase_						( self :	str,					a_ :	str,					a_ :	int,					a_ :	bool = False					):





    """simple docstring"""



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

 def 		lowercase_						( self :	Dict					):





    """simple docstring"""



    return self.train_loader

 def 		lowercase_						( self :	int					):





    """simple docstring"""



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

 def 		lowercase_						( self :	List[Any]					):





    """simple docstring"""



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

 def 		lowercase_						( self :	int,					a_ :	Union[str, Any]					):





    """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 :	Dict,					a_ :	Dict[str, Any]					):





    """simple docstring"""



    UpperCamelCase__ 						=						self.output_dir.joinpath("best_tfmr"					)
    UpperCamelCase__ 						=						self.step_count
    self.model.save_pretrained(a_					)
    self.tokenizer.save_pretrained(a_					)

 @staticmethod
 def 		lowercase_						( a_ :	str,					a_ :	List[str]					):





    """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):

 def 		lowercase_						( self :	str,					a_ :	str,					a_ :	str					):





    """simple docstring"""



    UpperCamelCase__ 						=						trainer.lr_schedulers[0]["scheduler"]
    UpperCamelCase__ 						=						{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],					a_ :	pl.Trainer,					a_ :	pl.LightningModule					):





    """simple docstring"""



    rank_zero_info("***** Validation results *****"					)
    UpperCamelCase__ 						=						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 :	List[Any],					a_ :	pl.Trainer,					a_ :	pl.LightningModule					):





    """simple docstring"""



    rank_zero_info("***** Test results *****"					)
    UpperCamelCase__ 						=						trainer.callback_metrics
    # Log and save results to file
    UpperCamelCase__ 						=						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        SCREAMING_SNAKE_CASE__(		_UpperCamelCase      :				str    ,		_UpperCamelCase      :				Any   )		->					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        SCREAMING_SNAKE_CASE__(		_UpperCamelCase      :				BaseTransformer    ,		_UpperCamelCase      :				argparse.Namespace    ,		_UpperCamelCase      :				str=None    ,		_UpperCamelCase      :				int=True    ,		_UpperCamelCase      :				Dict=[]    ,		_UpperCamelCase      :				Tuple=None    ,		_UpperCamelCase      :				Dict=None    ,		**_UpperCamelCase      :				str    ,		)		->					Optional[Any]:





   '''simple docstring'''







   pl.seed_everything(args.seed   )

   # init model
   UpperCamelCase__ 						=						Path(model.hparams.output_dir   )
   odir.mkdir(exist_ok=_UpperCamelCase   )

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

   UpperCamelCase__ 						=						{}

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

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

   UpperCamelCase__ 						=						args.accumulate_grad_batches
   UpperCamelCase__ 						=						extra_train_kwargs.get("accelerator"    ,		_UpperCamelCase   )
   UpperCamelCase__ 						=						extra_train_kwargs.get("profiler"    ,		_UpperCamelCase   )

   UpperCamelCase__ 						=						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'''

from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_sentencepiece_available,
    is_tokenizers_available,
    is_torch_available,
)


__lowercase:		int							=     {
    "configuration_open_llama": ["OPEN_LLAMA_PRETRAINED_CONFIG_ARCHIVE_MAP", "OpenLlamaConfig"],
}

try:
       if not is_sentencepiece_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       __lowercase:		str							=     ["LlamaTokenizer"]

try:
       if not is_tokenizers_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       __lowercase:		List[Any]							=     ["LlamaTokenizerFast"]

try:
       if not is_torch_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       __lowercase:		Union[str, Any]							=     [
           "OpenLlamaForCausalLM",
           "OpenLlamaModel",
           "OpenLlamaPreTrainedModel",
           "OpenLlamaForSequenceClassification",
       ]


if TYPE_CHECKING:
       from .configuration_open_llama import OPEN_LLAMA_PRETRAINED_CONFIG_ARCHIVE_MAP, OpenLlamaConfig

       try:
              if not is_sentencepiece_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from transformers import LlamaTokenizer

       try:
              if not is_tokenizers_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from transformers import LlamaTokenizerFast

       try:
              if not is_torch_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_open_llama import (
                  OpenLlamaForCausalLM,
                  OpenLlamaForSequenceClassification,
                  OpenLlamaModel,
                  OpenLlamaPreTrainedModel,
              )


else:
       import sys

       __lowercase:		Tuple							=     _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)