





'''simple docstring'''

from typing import TYPE_CHECKING

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


UpperCAmelCase					:		List[Any]		=							{'processing_layoutxlm': ['LayoutXLMProcessor']}

try:
   if not is_sentencepiece_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   UpperCAmelCase					:		Tuple		=							['LayoutXLMTokenizer']

try:
   if not is_tokenizers_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   UpperCAmelCase					:		Optional[Any]		=							['LayoutXLMTokenizerFast']

if TYPE_CHECKING:
   from .processing_layoutxlm import LayoutXLMProcessor

   try:
      if not is_sentencepiece_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .tokenization_layoutxlm import LayoutXLMTokenizer

   try:
      if not is_tokenizers_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .tokenization_layoutxlm_fast import LayoutXLMTokenizerFast

else:
   import sys

   UpperCAmelCase					:		List[Any]		=							_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)









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


UpperCAmelCase					:		Tuple		=							logging.getLogger(__name__)

require_version('pytorch_lightning>=1.0.4')

UpperCAmelCase					:		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
UpperCAmelCase					:		Optional[Any]		=							{
    '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					:		Tuple		=							sorted(arg_to_scheduler.keys())
UpperCAmelCase					:		Optional[Any]		=							'{' + ', '.join(arg_to_scheduler_choices) + '}'





class lowerCamelCase						(pl.LightningModule					):







  def __init__(		self     ,     lowercase__     ,     lowercase__=None     ,     lowercase__="base"     ,     lowercase__=None     ,     lowercase__=None     ,     lowercase__=None     ,     **lowercase__     ,     )		->						Optional[int]:






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

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







  def 				UpperCAmelCase_				(		self     ,     *lowercase__     ,     **lowercase__		)		->						List[str]:






        """simple docstring"""





        _snake_case		:	Dict					     =     self.model_type.from_pretrained(*lowercase__     ,     **lowercase__		)







  def 				UpperCAmelCase_				(		self		)		->						List[Any]:






        """simple docstring"""





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







  def 				UpperCAmelCase_				(		self		)		->						int:






        """simple docstring"""





        _snake_case		:	Any					     =     self.model
        _snake_case		:	List[Any]					     =     ['''bias''', '''LayerNorm.weight''']
        _snake_case		:	List[str]					     =     [
            {
                '''params''': [
                    p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay		)
                ],  # check this named paramters
                '''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:
              _snake_case		:	Any					     =     Adafactor(
                  lowercase__     ,     lr=self.hparams.learning_rate     ,     scale_parameter=lowercase__     ,     relative_step=lowercase__		)

        else:
              _snake_case		:	List[str]					     =     AdamW(
                  lowercase__     ,     lr=self.hparams.learning_rate     ,     eps=self.hparams.adam_epsilon		)
        _snake_case		:	List[str]					     =     optimizer

        _snake_case		:	Any					     =     self.get_lr_scheduler()

        return [optimizer], [scheduler]







  def 				UpperCAmelCase_				(		self     ,     lowercase__     ,     lowercase__		)		->						Any:






        """simple docstring"""





        return self.validation_step(lowercase__     ,     lowercase__		)







  def 				UpperCAmelCase_				(		self     ,     lowercase__		)		->						Tuple:






        """simple docstring"""





        return self.validation_end(lowercase__		)







  def 				UpperCAmelCase_				(		self		)		->						int:






        """simple docstring"""





        _snake_case		:	Any					     =     max(1     ,     self.hparams.gpus		)  # TODO: consider num_tpu_cores
        _snake_case		:	Optional[int]					     =     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__		)		->						Any:






        """simple docstring"""





        if stage == "test":
              _snake_case		:	Any					     =     len(self.test_dataloader().dataset		)
        else:
              _snake_case		:	Dict					     =     self.get_dataloader('''train'''     ,     self.hparams.train_batch_size     ,     shuffle=lowercase__		)
              _snake_case		:	Optional[int]					     =     len(self.train_dataloader().dataset		)







  def 				UpperCAmelCase_				(		self     ,     lowercase__     ,     lowercase__     ,     lowercase__ = False		)		->						str:






        """simple docstring"""





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







  def 				UpperCAmelCase_				(		self		)		->						Optional[int]:






        """simple docstring"""





        return self.train_loader







  def 				UpperCAmelCase_				(		self		)		->						Dict:






        """simple docstring"""





        return self.get_dataloader('''dev'''     ,     self.hparams.eval_batch_size     ,     shuffle=lowercase__		)







  def 				UpperCAmelCase_				(		self		)		->						Optional[Any]:






        """simple docstring"""





        return self.get_dataloader('''test'''     ,     self.hparams.eval_batch_size     ,     shuffle=lowercase__		)







  def 				UpperCAmelCase_				(		self     ,     lowercase__		)		->						Optional[int]:






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






        """simple docstring"""





        _snake_case		:	Dict					     =     self.output_dir.joinpath('''best_tfmr'''		)
        _snake_case		:	Tuple					     =     self.step_count
        self.model.save_pretrained(lowercase__		)
        self.tokenizer.save_pretrained(lowercase__		)







  @staticmethod
  def 				UpperCAmelCase_				(		lowercase__     ,     lowercase__		)		->						Tuple:






        """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=str(Path(lowercase__		).parent / '''test_run''' / '''cache'''		)     ,     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 lowerCamelCase						(pl.Callback					):







  def 				UpperCAmelCase_				(		self     ,     lowercase__     ,     lowercase__		)		->						str:






        """simple docstring"""





        if (
            trainer.is_global_zero and trainer.global_rank == 0
        ):  # we initialize the retriever only on master worker with RAY. In new pytorch-lightning accelorators are removed.
              pl_module.model.rag.retriever.init_retrieval()  # better to use hook functions.





class lowerCamelCase						(pl.Callback					):







  def 				UpperCAmelCase_				(		self     ,     lowercase__     ,     lowercase__		)		->						List[str]:






        """simple docstring"""





        for name, param in pl_module.model.rag.named_parameters():
              if param.grad is None:
                    print(lowercase__		)





class lowerCamelCase						(pl.Callback					):







  def 				UpperCAmelCase_				(		self     ,     lowercase__     ,     lowercase__		)		->						Any:






        """simple docstring"""





        _snake_case		:	Any					     =     trainer.lr_schedulers[0]['''scheduler''']
        _snake_case		:	Optional[int]					     =     {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__		)		->						List[str]:






        """simple docstring"""





        rank_zero_info('''***** Validation results *****'''		)
        _snake_case		:	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__		)		->						Dict:






        """simple docstring"""





        rank_zero_info('''***** Test results *****'''		)
        _snake_case		:	Dict					     =     trainer.callback_metrics
        # Log and save results to file
        _snake_case		:	str					     =     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_	):





      """simple docstring"""



      parser.add_argument(
          '''--output_dir'''			,      default=str(Path(lowerCAmelCase_	).parent / '''test_run''' / '''model_checkpoints'''	)			,      type=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=str(Path(lowerCAmelCase_	).parent / '''test_run''' / '''dummy-train-data'''	)			,      type=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_			,      ):





      """simple docstring"""



      pl.seed_everything(args.seed	)

      # init model
      _snake_case		:	Union[str, Any]					     =     Path(model.hparams.output_dir	)
      odir.mkdir(exist_ok=lowerCAmelCase_	)

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

      _snake_case		:	Tuple					     =     {}

      if args.fpaa:
            _snake_case		:	Union[str, Any]					     =     16

      if args.gpus > 1:
            _snake_case		:	Optional[Any]					     =     '''auto'''
            _snake_case		:	Tuple					     =     '''ddp'''

      _snake_case		:	Optional[Any]					     =     args.accumulate_grad_batches
      _snake_case		:	Tuple					     =     None
      _snake_case		:	str					     =     '''auto'''

      _snake_case		:	int					     =     pl.Trainer.from_argparse_args(
          lowerCAmelCase_			,      weights_summary=lowerCAmelCase_			,      callbacks=[logging_callback] + extra_callbacks + [InitCallback()] + [checkpoint_callback]			,      logger=lowerCAmelCase_			,      val_check_interval=1			,      num_sanity_val_steps=2			,      **lowerCAmelCase_			,      )

      if args.do_train:
            trainer.fit(lowerCAmelCase_	)

      else:
            print('''RAG modeling tests with new set functions successfuly executed!'''	)
      return trainer


