code: 
  
  
  
  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  				=					logging.getLogger(__name__)
  
  require_version("""pytorch_lightning>=1.0.4""")
  
  __lowerCamelCase  				=					{
      """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  				=					{
      """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  				=					sorted(arg_to_scheduler.keys())
  __lowerCamelCase  				=					"""{""" + """, """.join(arg_to_scheduler_choices) + """}"""
  
  
  
  class 		_snake_case      (			pl.LightningModule				):
  
  
  
  
  
   """simple docstring"""
  
  
  
   def __init__(						self	,    a	,    a=None	,    a="base"	,    a=None	,    a=None	,    a=None	,    **a	,    )	->							List[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(lowerCAmelCase__					)
     _A           =    0
     _A           =    Path(self.hparams.output_dir					)
     _A           =    self.hparams.cache_dir if self.hparams.cache_dir else None
     if config is None:
       _A           =    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=lowerCAmelCase__	,    **lowerCAmelCase__	,    )
     else:
       _A           =    config
  
     _A           =    ("encoder_layerdrop", "decoder_layerdrop", "dropout", "attention_dropout")
     for p in extra_model_params:
       if getattr(self.hparams	,    lowerCAmelCase__	,    lowerCAmelCase__					):
         assert hasattr(self.config	,    lowerCAmelCase__					), f'''model config doesn\'t have a `{p}` attribute'''
         setattr(self.config	,    lowerCAmelCase__	,    getattr(self.hparams	,    lowerCAmelCase__					)					)
  
     if tokenizer is None:
       _A           =    AutoTokenizer.from_pretrained(
           self.hparams.tokenizer_name if self.hparams.tokenizer_name else self.hparams.model_name_or_path	,    cache_dir=lowerCAmelCase__	,    )
     else:
       _A           =    tokenizer
     _A           =    MODEL_MODES[mode]
     if model is None:
       _A           =    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=lowerCAmelCase__	,    )
     else:
       _A           =    model
   def lowercase_		(						self	,    *a	,    **a					)	->							List[str]:
  
  
  
  
     """simple docstring"""
  
  
  
  
  
  
  
     _A           =    self.model_type.from_pretrained(*lowerCAmelCase__	,    **lowerCAmelCase__					)
   def lowercase_		(						self					)	->							List[Any]:
  
  
  
  
     """simple docstring"""
  
  
  
  
  
  
  
     _A           =    arg_to_scheduler[self.hparams.lr_scheduler]
     _A           =    get_schedule_func(
         self.opt	,    num_warmup_steps=self.hparams.warmup_steps	,    num_training_steps=self.total_steps()					)
     _A           =    {"scheduler": scheduler, "interval": "step", "frequency": 1}
     return scheduler
   def lowercase_		(						self					)	->							Any:
  
  
  
  
     """simple docstring"""
  
  
  
  
  
  
  
     _A           =    self.model
     _A           =    ["bias", "LayerNorm.weight"]
     _A           =    [
         {
             "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:
       _A           =    Adafactor(
           lowerCAmelCase__	,    lr=self.hparams.learning_rate	,    scale_parameter=lowerCAmelCase__	,    relative_step=lowerCAmelCase__					)
  
     else:
       _A           =    AdamW(
           lowerCAmelCase__	,    lr=self.hparams.learning_rate	,    eps=self.hparams.adam_epsilon					)
     _A           =    optimizer
  
     _A           =    self.get_lr_scheduler()
  
     return [optimizer], [scheduler]
   def lowercase_		(						self	,    a	,    a					)	->							Dict:
  
  
  
  
     """simple docstring"""
  
  
  
  
  
  
  
     return self.validation_step(lowerCAmelCase__	,    lowerCAmelCase__					)
   def lowercase_		(						self	,    a					)	->							str:
  
  
  
  
     """simple docstring"""
  
  
  
  
  
  
  
     return self.validation_end(lowerCAmelCase__					)
   def lowercase_		(						self					)	->							int:
  
  
  
  
     """simple docstring"""
  
  
  
  
  
  
  
     _A           =    max(1	,    self.hparams.gpus					)  # TODO: consider num_tpu_cores
     _A           =    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	,    a					)	->							Optional[int]:
  
  
  
  
     """simple docstring"""
  
  
  
  
  
  
  
     if mode == "test":
       _A           =    len(self.test_dataloader().dataset					)
     else:
       _A           =    self.get_dataloader('''train'''	,    self.hparams.train_batch_size	,    shuffle=lowerCAmelCase__					)
       _A           =    len(self.train_dataloader().dataset					)
   def lowercase_		(						self	,    a	,    a	,    a = False					)	->							List[Any]:
  
  
  
  
     """simple docstring"""
  
  
  
  
  
  
  
     raise NotImplementedError('''You must implement this for your task'''					)
   def lowercase_		(						self					)	->							Dict:
  
  
  
  
     """simple docstring"""
  
  
  
  
  
  
  
     return self.train_loader
   def lowercase_		(						self					)	->							str:
  
  
  
  
     """simple docstring"""
  
  
  
  
  
  
  
     return self.get_dataloader('''dev'''	,    self.hparams.eval_batch_size	,    shuffle=lowerCAmelCase__					)
   def lowercase_		(						self					)	->							str:
  
  
  
  
     """simple docstring"""
  
  
  
  
  
  
  
     return self.get_dataloader('''test'''	,    self.hparams.eval_batch_size	,    shuffle=lowerCAmelCase__					)
   def lowercase_		(						self	,    a					)	->							List[Any]:
  
  
  
  
     """simple docstring"""
  
  
  
  
  
  
  
     return os.path.join(
         self.hparams.data_dir	,    '''cached_{}_{}_{}'''.format(
             lowerCAmelCase__	,    list(filter(lowerCAmelCase__	,    self.hparams.model_name_or_path.split('''/'''					)					)					).pop()	,    str(self.hparams.max_seq_length					)	,    )	,    )
   @pl.utilities.rank_zero_only
   def lowercase_		(						self	,    a					)	->							None:
  
  
  
  
     """simple docstring"""
  
  
  
  
  
  
  
     _A           =    self.output_dir.joinpath('''best_tfmr'''					)
     _A           =    self.step_count
     self.model.save_pretrained(lowerCAmelCase__					)
     self.tokenizer.save_pretrained(lowerCAmelCase__					)
  
  
  
   @staticmethod
   def lowercase_		(						a	,    a					)	->							Tuple:
  
  
  
  
     """simple docstring"""
  
  
  
  
  
  
  
     parser.add_argument(
         '''--model_name_or_path'''	,    default=lowerCAmelCase__	,    type=lowerCAmelCase__	,    required=lowerCAmelCase__	,    help='''Path to pretrained model or model identifier from huggingface.co/models'''	,    )
     parser.add_argument(
         '''--config_name'''	,    default=''''''	,    type=lowerCAmelCase__	,    help='''Pretrained config name or path if not the same as model_name'''					)
     parser.add_argument(
         '''--tokenizer_name'''	,    default=lowerCAmelCase__	,    type=lowerCAmelCase__	,    help='''Pretrained tokenizer name or path if not the same as model_name'''	,    )
     parser.add_argument(
         '''--cache_dir'''	,    default=''''''	,    type=lowerCAmelCase__	,    help='''Where do you want to store the pre-trained models downloaded from huggingface.co'''	,    )
     parser.add_argument(
         '''--encoder_layerdrop'''	,    type=lowerCAmelCase__	,    help='''Encoder layer dropout probability (Optional). Goes into model.config'''	,    )
     parser.add_argument(
         '''--decoder_layerdrop'''	,    type=lowerCAmelCase__	,    help='''Decoder layer dropout probability (Optional). Goes into model.config'''	,    )
     parser.add_argument(
         '''--dropout'''	,    type=lowerCAmelCase__	,    help='''Dropout probability (Optional). Goes into model.config'''	,    )
     parser.add_argument(
         '''--attention_dropout'''	,    type=lowerCAmelCase__	,    help='''Attention dropout probability (Optional). Goes into model.config'''	,    )
     parser.add_argument('''--learning_rate'''	,    default=5e-5	,    type=lowerCAmelCase__	,    help='''The initial learning rate for Adam.'''					)
     parser.add_argument(
         '''--lr_scheduler'''	,    default='''linear'''	,    choices=lowerCAmelCase__	,    metavar=lowerCAmelCase__	,    type=lowerCAmelCase__	,    help='''Learning rate scheduler'''	,    )
     parser.add_argument('''--weight_decay'''	,    default=0.0	,    type=lowerCAmelCase__	,    help='''Weight decay if we apply some.'''					)
     parser.add_argument('''--adam_epsilon'''	,    default=1e-8	,    type=lowerCAmelCase__	,    help='''Epsilon for Adam optimizer.'''					)
     parser.add_argument('''--warmup_steps'''	,    default=0	,    type=lowerCAmelCase__	,    help='''Linear warmup over warmup_steps.'''					)
     parser.add_argument('''--num_workers'''	,    default=4	,    type=lowerCAmelCase__	,    help='''kwarg passed to DataLoader'''					)
     parser.add_argument('''--num_train_epochs'''	,    dest='''max_epochs'''	,    default=3	,    type=lowerCAmelCase__					)
     parser.add_argument('''--train_batch_size'''	,    default=3_2	,    type=lowerCAmelCase__					)
     parser.add_argument('''--eval_batch_size'''	,    default=3_2	,    type=lowerCAmelCase__					)
     parser.add_argument('''--adafactor'''	,    action='''store_true'''					)
  
  
  
  
  
  
  
  class 		_snake_case      (			pl.Callback				):
  
  
  
  
  
   """simple docstring"""
  
  
  
   def lowercase_		(						self	,    a	,    a					)	->							str:
  
  
  
  
     """simple docstring"""
  
  
  
  
  
  
  
     _A           =    trainer.lr_schedulers[0]["scheduler"]
     _A           =    {f'''lr_group_{i}''': lr for i, lr in enumerate(lr_scheduler.get_lr()					)}
     pl_module.logger.log_metrics(lowerCAmelCase__					)
   def lowercase_		(						self	,    a	,    a					)	->							Any:
  
  
  
  
     """simple docstring"""
  
  
  
  
  
  
  
     rank_zero_info('''***** Validation results *****'''					)
     _A           =    trainer.callback_metrics
     # Log results
     for key in sorted(lowerCAmelCase__					):
       if key not in ["log", "progress_bar"]:
         rank_zero_info('''{} = {}\n'''.format(lowerCAmelCase__	,    str(metrics[key]					)					)					)
  
  
  
   def lowercase_		(						self	,    a	,    a					)	->							Optional[int]:
  
  
  
  
     """simple docstring"""
  
  
  
  
  
  
  
     rank_zero_info('''***** Test results *****'''					)
     _A           =    trainer.callback_metrics
     # Log and save results to file
     _A           =    os.path.join(pl_module.hparams.output_dir	,    '''test_results.txt'''					)
     with open(lowerCAmelCase__	,    '''w'''					) as writer:
       for key in sorted(lowerCAmelCase__					):
         if key not in ["log", "progress_bar"]:
           rank_zero_info('''{} = {}\n'''.format(lowerCAmelCase__	,    str(metrics[key]					)					)					)
           writer.write('''{} = {}\n'''.format(lowerCAmelCase__	,    str(metrics[key]					)					)					)
  
  
  
  
  
  
  
  def     UpperCAmelCase__ (							__snake_case						,				__snake_case							)   -> Tuple:
    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						,				__snake_case						,				__snake_case=None						,				__snake_case=True						,				__snake_case=[]						,				__snake_case=None						,				__snake_case=None						,				**__snake_case						,				)   -> List[Any]:
    pl.seed_everything(args.seed							)
  
    # init model
    _A           =    Path(model.hparams.output_dir							)
    odir.mkdir(exist_ok=__snake_case							)
  
    # add custom checkpoints
    if checkpoint_callback is None:
      _A           =    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:
      _A           =    LoggingCallback()
  
    _A           =    {}
  
    # TODO: remove with PyTorch 1.6 since pl uses native amp
    if args.fpaa:
      _A           =    16
      _A           =    args.fpaa_opt_level
  
    if args.gpus > 1:
      _A           =    "ddp"
  
    _A           =    args.accumulate_grad_batches
    _A           =    extra_train_kwargs.get('''accelerator'''						,				__snake_case							)
    _A           =    extra_train_kwargs.get('''profiler'''						,				__snake_case							)
  
    _A           =    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
style_context: 
  
  
  
  import argparse
  
  import torch
  
  from transformers import ElectraConfig, ElectraForMaskedLM, ElectraForPreTraining, load_tf_weights_in_electra
  from transformers.utils import logging
  
  
  logging.set_verbosity_info()
  
  
  
  
  def     UpperCAmelCase__ (							__snake_case						,				__snake_case						,				__snake_case						,				__snake_case							)   -> Tuple:
    # Initialise PyTorch model
    _A           =    ElectraConfig.from_json_file(__snake_case							)
    print(F'''Building PyTorch model from configuration: {config}'''							)
  
    if discriminator_or_generator == "discriminator":
      _A           =    ElectraForPreTraining(__snake_case							)
    elif discriminator_or_generator == "generator":
      _A           =    ElectraForMaskedLM(__snake_case							)
    else:
      raise ValueError('''The discriminator_or_generator argument should be either \'discriminator\' or \'generator\''''							)
  
    # Load weights from tf checkpoint
    load_tf_weights_in_electra(
        __snake_case						,				__snake_case						,				__snake_case						,				discriminator_or_generator=__snake_case							)
  
    # Save pytorch-model
    print(F'''Save PyTorch model to {pytorch_dump_path}'''							)
    torch.save(model.state_dict()						,				__snake_case							)
  
  
  if __name__ == "__main__":
      __lowerCamelCase  				=					argparse.ArgumentParser()
      # Required parameters
      parser.add_argument(
          """--tf_checkpoint_path""", default=None, type=str, required=True, help="""Path to the TensorFlow checkpoint path."""
      )
      parser.add_argument(
          """--config_file""",
          default=None,
          type=str,
          required=True,
          help="""The config json file corresponding to the pre-trained model. \nThis specifies the model architecture.""",
      )
      parser.add_argument(
          """--pytorch_dump_path""", default=None, type=str, required=True, help="""Path to the output PyTorch model."""
      )
      parser.add_argument(
          """--discriminator_or_generator""",
          default=None,
          type=str,
          required=True,
          help=(
              """Whether to export the generator or the discriminator. Should be a string, either 'discriminator' or """
              """'generator'."""
          ),
      )
      __lowerCamelCase  				=					parser.parse_args()
      convert_tf_checkpoint_to_pytorch(
          args.tf_checkpoint_path, args.config_file, args.pytorch_dump_path, args.discriminator_or_generator
      )