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,
      AutoModelForSeq2SeqLM,
      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_examples
  
  
  logger = logging.getLogger(__name__)
  
  require_version_examples("pytorch_lightning>=1.0.4")
  
  MODEL_MODES = {
      "base": AutoModel,
      "sequence-classification": AutoModelForSequenceClassification,
      "question-answering": AutoModelForQuestionAnswering,
      "pretraining": AutoModelForPreTraining,
      "token-classification": AutoModelForTokenClassification,
      "language-modeling": AutoModelWithLMHead,
      "summarization": AutoModelForSeq2SeqLM,
      "translation": AutoModelForSeq2SeqLM,
  }
  
  
  # update this and the import above to support new schedulers from transformers.optimization
  arg_to_scheduler = {
      "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
  }
  arg_to_scheduler_choices = sorted(arg_to_scheduler.keys())
  arg_to_scheduler_metavar = "{" + ", ".join(arg_to_scheduler_choices) + "}"
  
  
  class BaseTransformer(pl.LightningModule):
      def __init__(
          self,
          hparams: argparse.Namespace,
          num_labels=None,
          mode="base",
          config=None,
          tokenizer=None,
          model=None,
          **config_kwargs
      ):
          """Initialize a model, tokenizer and config."""
          super().__init__()
          # TODO: move to self.save_hyperparameters()
          # self.save_hyperparameters()
          # can also expand arguments into trainer signature for easier reading
  
          self.save_hyperparameters(hparams)
          self.step_count = 0
          self.output_dir = Path(self.hparams.output_dir)
          cache_dir = self.hparams.cache_dir if self.hparams.cache_dir else None
          if config is None:
              self.config = 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=cache_dir,
                  **config_kwargs,
              )
          else:
              self.config: PretrainedConfig = config
  
          extra_model_params = ("encoder_layerdrop", "decoder_layerdrop", "dropout", "attention_dropout")
          for p in extra_model_params:
              if getattr(self.hparams, p, None):
                  assert hasattr(self.config, p), f"model config doesn't have a `{p}` attribute"
                  setattr(self.config, p, getattr(self.hparams, p))
  
          if tokenizer is None:
              self.tokenizer = AutoTokenizer.from_pretrained(
                  self.hparams.tokenizer_name if self.hparams.tokenizer_name else self.hparams.model_name_or_path,
                  cache_dir=cache_dir,
              )
          else:
              self.tokenizer: PreTrainedTokenizer = tokenizer
          self.model_type = MODEL_MODES[mode]
          if model is None:
              self.model = 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=cache_dir,
              )
          else:
              self.model = model
  
      def load_hf_checkpoint(self, *args, **kwargs):
          self.model = self.model_type.from_pretrained(*args, **kwargs)
  
      def get_lr_scheduler(self):
          get_schedule_func = arg_to_scheduler[self.hparams.lr_scheduler]
          scheduler = get_schedule_func(
              self.opt, num_warmup_steps=self.hparams.warmup_steps, num_training_steps=self.total_steps()
          )
          scheduler = {"scheduler": scheduler, "interval": "step", "frequency": 1}
          return scheduler
  
      def configure_optimizers(self):
          """Prepare optimizer and schedule (linear warmup and decay)"""
          model = self.model
          no_decay = ["bias", "LayerNorm.weight"]
          optimizer_grouped_parameters = [
              {
                  "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:
              optimizer = Adafactor(
                  optimizer_grouped_parameters, lr=self.hparams.learning_rate, scale_parameter=False, relative_step=False
              )
  
          else:
              optimizer = AdamW(
                  optimizer_grouped_parameters, lr=self.hparams.learning_rate, eps=self.hparams.adam_epsilon
              )
          self.opt = optimizer
  
          scheduler = self.get_lr_scheduler()
  
          return [optimizer], [scheduler]
  
      def test_step(self, batch, batch_nb):
          return self.validation_step(batch, batch_nb)
  
      def test_epoch_end(self, outputs):
          return self.validation_end(outputs)
  
      def total_steps(self) -> int:
          """The number of total training steps that will be run. Used for lr scheduler purposes."""
          num_devices = max(1, self.hparams.gpus)  # TODO: consider num_tpu_cores
          effective_batch_size = self.hparams.train_batch_size * self.hparams.accumulate_grad_batches * num_devices
          return (self.dataset_size / effective_batch_size) * self.hparams.max_epochs
  
      def setup(self, mode):
          if mode == "test":
              self.dataset_size = len(self.test_dataloader().dataset)
          else:
              self.train_loader = self.get_dataloader("train", self.hparams.train_batch_size, shuffle=True)
              self.dataset_size = len(self.train_dataloader().dataset)
  
      def get_dataloader(self, type_path: str, batch_size: int, shuffle: bool = False):
          raise NotImplementedError("You must implement this for your task")
  
      def train_dataloader(self):
          return self.train_loader
  
      def val_dataloader(self):
          return self.get_dataloader("dev", self.hparams.eval_batch_size, shuffle=False)
  
      def test_dataloader(self):
          return self.get_dataloader("test", self.hparams.eval_batch_size, shuffle=False)
  
      def _feature_file(self, mode):
          return os.path.join(
              self.hparams.data_dir,
              "cached_{}_{}_{}".format(
                  mode,
                  list(filter(None, self.hparams.model_name_or_path.split("/"))).pop(),
                  str(self.hparams.max_seq_length),
              ),
          )
  
      @pl.utilities.rank_zero_only
      def on_save_checkpoint(self, checkpoint: Dict[str, Any]) -> None:
          save_path = self.output_dir.joinpath("best_tfmr")
          self.model.config.save_step = self.step_count
          self.model.save_pretrained(save_path)
          self.tokenizer.save_pretrained(save_path)
  
      @staticmethod
      def add_model_specific_args(parser, root_dir):
          parser.add_argument(
              "--model_name_or_path",
              default=None,
              type=str,
              required=True,
              help="Path to pretrained model or model identifier from huggingface.co/models",
          )
          parser.add_argument(
              "--config_name", default="", type=str, help="Pretrained config name or path if not the same as model_name"
          )
          parser.add_argument(
              "--tokenizer_name",
              default=None,
              type=str,
              help="Pretrained tokenizer name or path if not the same as model_name",
          )
          parser.add_argument(
              "--cache_dir",
              default="",
              type=str,
              help="Where do you want to store the pre-trained models downloaded from huggingface.co",
          )
          parser.add_argument(
              "--encoder_layerdrop",
              type=float,
              help="Encoder layer dropout probability (Optional). Goes into model.config",
          )
          parser.add_argument(
              "--decoder_layerdrop",
              type=float,
              help="Decoder layer dropout probability (Optional). Goes into model.config",
          )
          parser.add_argument(
              "--dropout",
              type=float,
              help="Dropout probability (Optional). Goes into model.config",
          )
          parser.add_argument(
              "--attention_dropout",
              type=float,
              help="Attention dropout probability (Optional). Goes into model.config",
          )
          parser.add_argument("--learning_rate", default=5e-5, type=float, help="The initial learning rate for Adam.")
          parser.add_argument(
              "--lr_scheduler",
              default="linear",
              choices=arg_to_scheduler_choices,
              metavar=arg_to_scheduler_metavar,
              type=str,
              help="Learning rate scheduler",
          )
          parser.add_argument("--weight_decay", default=0.0, type=float, help="Weight decay if we apply some.")
          parser.add_argument("--adam_epsilon", default=1e-8, type=float, help="Epsilon for Adam optimizer.")
          parser.add_argument("--warmup_steps", default=0, type=int, help="Linear warmup over warmup_steps.")
          parser.add_argument("--num_workers", default=4, type=int, help="kwarg passed to DataLoader")
          parser.add_argument("--num_train_epochs", dest="max_epochs", default=3, type=int)
          parser.add_argument("--train_batch_size", default=32, type=int)
          parser.add_argument("--eval_batch_size", default=32, type=int)
          parser.add_argument("--adafactor", action="store_true")
  
  
  class LoggingCallback(pl.Callback):
      def on_batch_end(self, trainer, pl_module):
          lr_scheduler = trainer.lr_schedulers[0]["scheduler"]
          lrs = {f"lr_group_{i}": lr for i, lr in enumerate(lr_scheduler.get_lr())}
          pl_module.logger.log_metrics(lrs)
  
      def on_validation_end(self, trainer: pl.Trainer, pl_module: pl.LightningModule):
          rank_zero_info("***** Validation results *****")
          metrics = trainer.callback_metrics
          # Log results
          for key in sorted(metrics):
              if key not in ["log", "progress_bar"]:
                  rank_zero_info("{} = {}\n".format(key, str(metrics[key])))
  
      def on_test_end(self, trainer: pl.Trainer, pl_module: pl.LightningModule):
          rank_zero_info("***** Test results *****")
          metrics = trainer.callback_metrics
          # Log and save results to file
          output_test_results_file = os.path.join(pl_module.hparams.output_dir, "test_results.txt")
          with open(output_test_results_file, "w") as writer:
              for key in sorted(metrics):
                  if key not in ["log", "progress_bar"]:
                      rank_zero_info("{} = {}\n".format(key, str(metrics[key])))
                      writer.write("{} = {}\n".format(key, str(metrics[key])))
  
  
  def add_generic_args(parser, root_dir) -> None:
      #  To allow all pl args uncomment the following line
      #  parser = pl.Trainer.add_argparse_args(parser)
      parser.add_argument(
          "--output_dir",
          default=None,
          type=str,
          required=True,
          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=str,
          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=int)
      parser.add_argument("--max_grad_norm", dest="gradient_clip_val", default=1.0, type=float, 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=int,
          default=1,
          help="Number of updates steps to accumulate before performing a backward/update pass.",
      )
      parser.add_argument("--seed", type=int, default=42, help="random seed for initialization")
      parser.add_argument(
          "--data_dir",
          default=None,
          type=str,
          required=True,
          help="The input data dir. Should contain the training files for the CoNLL-2003 NER task.",
      )
  
  
  def generic_train(
      model: BaseTransformer,
      args: argparse.Namespace,
      early_stopping_callback=None,
      logger=True,  # can pass WandbLogger() here
      extra_callbacks=[],
      checkpoint_callback=None,
      logging_callback=None,
      **extra_train_kwargs
  ):
      pl.seed_everything(args.seed)
  
      # init model
      odir = Path(model.hparams.output_dir)
      odir.mkdir(exist_ok=True)
  
      # add custom checkpoints
      if checkpoint_callback is None:
          checkpoint_callback = 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(early_stopping_callback)
      if logging_callback is None:
          logging_callback = LoggingCallback()
  
      train_params = {}
  
      # TODO: remove with PyTorch 1.6 since pl uses native amp
      if args.fp16:
          train_params["precision"] = 16
          train_params["amp_level"] = args.fp16_opt_level
  
      if args.gpus > 1:
          train_params["distributed_backend"] = "ddp"
  
      train_params["accumulate_grad_batches"] = args.accumulate_grad_batches
      train_params["accelerator"] = extra_train_kwargs.get("accelerator", None)
      train_params["profiler"] = extra_train_kwargs.get("profiler", None)
  
      trainer = pl.Trainer.from_argparse_args(
          args,
          weights_summary=None,
          callbacks=[logging_callback] + extra_callbacks,
          logger=logger,
          checkpoint_callback=checkpoint_callback,
          **train_params,
      )
  
      if args.do_train:
          trainer.fit(model)
  
      return trainer
  