

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     lowercase						(       pl.LightningModule					):



 '''simple docstring'''


 def __init__(self       ,       __a       ,       __a=None       ,       __a="base"       ,       __a=None       ,       __a=None       ,       __a=None       ,       **__a       ,       ) ->					Tuple:






  """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   UpperCamelCase__							(self       ,       *__a       ,       **__a ) ->					List[Any]:






  """simple docstring"""





  UpperCAmelCase__							=				self.model_type.from_pretrained(*__a       ,       **__a )




 def   UpperCamelCase__							(self ) ->					List[str]:






  """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   UpperCamelCase__							(self ) ->					Optional[int]:






  """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   UpperCamelCase__							(self       ,       __a       ,       __a ) ->					List[str]:






  """simple docstring"""





  return self.validation_step(__a       ,       __a )




 def   UpperCamelCase__							(self       ,       __a ) ->					Tuple:






  """simple docstring"""





  return self.validation_end(__a )




 def   UpperCamelCase__							(self ) ->					int:






  """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   UpperCamelCase__							(self       ,       __a ) ->					Any:






  """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   UpperCamelCase__							(self       ,       __a       ,       __a       ,       __a = False ) ->					Optional[Any]:






  """simple docstring"""





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




 def   UpperCamelCase__							(self ) ->					Dict:






  """simple docstring"""





  return self.train_loader




 def   UpperCamelCase__							(self ) ->					Any:






  """simple docstring"""





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




 def   UpperCamelCase__							(self ) ->					int:






  """simple docstring"""





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




 def   UpperCamelCase__							(self       ,       __a ) ->					int:






  """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   UpperCamelCase__							(self       ,       __a ) ->					None:






  """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   UpperCamelCase__							(__a       ,       __a ) ->					Dict:






  """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     lowercase						(       pl.Callback					):



 '''simple docstring'''


 def   UpperCamelCase__							(self       ,       __a       ,       __a ) ->					Union[str, Any]:






  """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   UpperCamelCase__							(self       ,       __a       ,       __a ) ->					int:






  """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   UpperCamelCase__							(self       ,       __a       ,       __a ) ->					Union[str, Any]:






  """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 						UpperCamelCase_(			snake_case__:  Dict  ,    snake_case__:  int   )       -> None:
 #  To allow all pl args uncomment the following line
 #  parser = pl.Trainer.add_argparse_args(parser)
 parser.add_argument(
     '--output_dir'  ,    default=snake_case__  ,    type=snake_case__  ,    required=snake_case__  ,    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=snake_case__  ,    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=snake_case__   )
 parser.add_argument('--max_grad_norm'  ,    dest='gradient_clip_val'  ,    default=1.0  ,    type=snake_case__  ,    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=snake_case__  ,    default=1  ,    help='Number of updates steps to accumulate before performing a backward/update pass.'  ,    )
 parser.add_argument('--seed'  ,    type=snake_case__  ,    default=42  ,    help='random seed for initialization'   )
 parser.add_argument(
     '--data_dir'  ,    default=snake_case__  ,    type=snake_case__  ,    required=snake_case__  ,    help='The input data dir. Should contain the training files for the CoNLL-2003 NER task.'  ,    )







def 						UpperCamelCase_(			snake_case__:  BaseTransformer  ,    snake_case__:  argparse.Namespace  ,    snake_case__:  int=None  ,    snake_case__:  Tuple=True  ,    snake_case__:  int=[]  ,    snake_case__:  Dict=None  ,    snake_case__:  int=None  ,    **snake_case__:  List[str]  ,    )       -> Optional[Any]:
 pl.seed_everything(args.seed   )

 # init model
 UpperCAmelCase__							=				Path(model.hparams.output_dir   )
 odir.mkdir(exist_ok=snake_case__   )

 # 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(snake_case__   )
 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'  ,    snake_case__   )
 UpperCAmelCase__							=				extra_train_kwargs.get('profiler'  ,    snake_case__   )

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

 if args.do_train:
  trainer.fit(snake_case__   )

 return trainer






from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


_UpperCamelCase			       =			logging.get_logger(__name__)

_UpperCamelCase			       =			{
    '''google/rembert''': '''https://huggingface.co/google/rembert/resolve/main/config.json''',
    # See all RemBERT models at https://huggingface.co/models?filter=rembert
}



class     lowercase						(       _UpperCamelCase					):



 '''simple docstring'''


 __SCREAMING_SNAKE_CASE												=				"""rembert"""




 def __init__(self       ,       __a=250300       ,       __a=1152       ,       __a=32       ,       __a=18       ,       __a=256       ,       __a=1664       ,       __a=4608       ,       __a="gelu"       ,       __a=0.0       ,       __a=0.0       ,       __a=0.1       ,       __a=512       ,       __a=2       ,       __a=0.02       ,       __a=1E-1_2       ,       __a=True       ,       __a=0       ,       __a=312       ,       __a=313       ,       **__a       ,       ) ->					str:






  """simple docstring"""





  super().__init__(pad_token_id=__a       ,       bos_token_id=__a       ,       eos_token_id=__a       ,       **__a )

  UpperCAmelCase__							=				vocab_size
  UpperCAmelCase__							=				input_embedding_size
  UpperCAmelCase__							=				output_embedding_size
  UpperCAmelCase__							=				max_position_embeddings
  UpperCAmelCase__							=				hidden_size
  UpperCAmelCase__							=				num_hidden_layers
  UpperCAmelCase__							=				num_attention_heads
  UpperCAmelCase__							=				intermediate_size
  UpperCAmelCase__							=				hidden_act
  UpperCAmelCase__							=				hidden_dropout_prob
  UpperCAmelCase__							=				attention_probs_dropout_prob
  UpperCAmelCase__							=				classifier_dropout_prob
  UpperCAmelCase__							=				initializer_range
  UpperCAmelCase__							=				type_vocab_size
  UpperCAmelCase__							=				layer_norm_eps
  UpperCAmelCase__							=				use_cache
  UpperCAmelCase__							=				False







class     lowercase						(       _UpperCamelCase					):



 '''simple docstring'''


 @property
 def   UpperCamelCase__							(self ) ->					Mapping[str, Mapping[int, str]]:






  """simple docstring"""





  if self.task == "multiple-choice":
   UpperCAmelCase__							=				{0: 'batch', 1: 'choice', 2: 'sequence'}
  else:
   UpperCAmelCase__							=				{0: 'batch', 1: 'sequence'}
  return OrderedDict(
      [
          ('input_ids', dynamic_axis),
          ('attention_mask', dynamic_axis),
          ('token_type_ids', dynamic_axis),
      ] )




 @property
 def   UpperCamelCase__							(self ) ->					float:






  """simple docstring"""





  return 1E-4



