text: # coding=utf-8
  # Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
  # Copyright (c) 2018, NVIDIA CORPORATION.  All rights reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  """
  Fine-tuning the library models for language modeling on a text file (GPT, GPT-2, BERT, RoBERTa).
  GPT and GPT-2 are fine-tuned using a causal language modeling (CLM) loss while BERT and RoBERTa are fine-tuned
  using a masked language modeling (MLM) loss.
  """
  
  import re
  import os
  import glob
  import torch
  import random
  import shutil
  import logging
  import argparse
  import numpy as np
  
  from tqdm import tqdm, trange
  from typing import Dict, List, Tuple
  from torch.nn import CrossEntropyLoss
  from torch.nn.utils.rnn import pad_sequence
  from torch.utils.data import DataLoader, RandomSampler, SequentialSampler
  
  from colors.src.common import JsonTextDataset
  
  from transformers import (
      WEIGHTS_NAME,
      AdamW,
      AutoConfig,
      AutoModelWithLMHead,
      AutoTokenizer,
      get_linear_schedule_with_warmup
  )
  
  try:
      from torch.utils.tensorboard import SummaryWriter
  except ImportError:
      from tensorboardX import SummaryWriter
  
  
  logger = logging.getLogger(__name__)
  
  
  def load_and_cache_examples(args, tokenizer, evaluate=False):
      file_path = args.eval_data_file if evaluate else args.train_data_file
      return JsonTextDataset(tokenizer, args, file_path=file_path, block_size=args.block_size)
  
  
  def set_seed(args):
      random.seed(args.seed)
      np.random.seed(args.seed)
      torch.manual_seed(args.seed)
      if args.n_gpu > 0:
          torch.cuda.manual_seed_all(args.seed)
  
  
  def _sorted_checkpoints(args, checkpoint_prefix="checkpoint", use_mtime=False) -> List[str]:
      ordering_and_checkpoint_path = []
  
      glob_checkpoints = glob.glob(os.path.join(args.output_dir, "{}-*".format(checkpoint_prefix)))
  
      for path in glob_checkpoints:
          if use_mtime:
              ordering_and_checkpoint_path.append((os.path.getmtime(path), path))
          else:
              regex_match = re.match(".*{}-([0-9]+)".format(checkpoint_prefix), path)
              if regex_match and regex_match.groups():
                  ordering_and_checkpoint_path.append((int(regex_match.groups()[0]), path))
  
      checkpoints_sorted = sorted(ordering_and_checkpoint_path)
      checkpoints_sorted = [checkpoint[1] for checkpoint in checkpoints_sorted]
      return checkpoints_sorted
  
  
  def _rotate_checkpoints(args, checkpoint_prefix="checkpoint", use_mtime=False) -> None:
      if not args.save_total_limit:
          return
      if args.save_total_limit <= 0:
          return
  
      # Check if we should delete older checkpoint(s)
      checkpoints_sorted = _sorted_checkpoints(args, checkpoint_prefix, use_mtime)
      if len(checkpoints_sorted) <= args.save_total_limit:
          return
  
      number_of_checkpoints_to_delete = max(0, len(checkpoints_sorted) - args.save_total_limit)
      checkpoints_to_be_deleted = checkpoints_sorted[:number_of_checkpoints_to_delete]
      for checkpoint in checkpoints_to_be_deleted:
          logger.info("Deleting older checkpoint [{}] due to args.save_total_limit".format(checkpoint))
          shutil.rmtree(checkpoint)
  
  
  def mask_tokens(inputs: torch.Tensor,
                  tokenizer: AutoTokenizer,
                  label: torch.Tensor,
                  masked_indices: torch.Tensor,
                  args) -> Tuple[torch.Tensor, torch.Tensor]:
      """ Prepare masked tokens inputs/labels for masked language modeling: 80% MASK, 10% random, 10% original. """
      # We only compute loss on masked tokens
      labels = torch.ones_like(inputs) * -100.0
  
      if tokenizer.mask_token is None:
          raise ValueError(
              "This tokenizer does not have a mask token which is necessary for masked language modeling. Remove the --mlm flag if you want to use this tokenizer."
          )
  
      probability_matrix = torch.full(labels.shape, args.mlm_probability)
      special_tokens_mask = [
          tokenizer.get_special_tokens_mask(val, already_has_special_tokens=True) for val in labels.tolist()
      ]
      probability_matrix.masked_fill_(torch.tensor(special_tokens_mask, dtype=torch.bool), value=0.0)
      if tokenizer._pad_token is not None:
          padding_mask = labels.eq(tokenizer.pad_token_id)
          probability_matrix.masked_fill_(padding_mask, value=0.0)
  
      # target_labels = torch.cat(list(label), dim=0).float()
      labels.scatter_(1, masked_indices.unsqueeze(1), label.unsqueeze(1).float())
  
      # The rest of the time (10% of the time) we keep the masked input tokens unchanged
      return inputs, labels
  
  
  def train(args, train_dataset, model: AutoModelWithLMHead, tokenizer: AutoTokenizer) -> Tuple[int, float]:
      """ Train the model """
      tb_writer = SummaryWriter()
      args.train_batch_size = args.per_gpu_train_batch_size * max(1, args.n_gpu)
  
      def collate(examples: List[torch.Tensor]):
          instances, labels, masked_indices = zip(*examples)
          if tokenizer._pad_token is None:
              return pad_sequence(instances, batch_first=True), pad_sequence(labels, batch_first=True), masked_indices
  
          return pad_sequence(instances, batch_first=True, padding_value=tokenizer.pad_token_id), \
                 pad_sequence(labels, batch_first=True, padding_value=tokenizer.pad_token_id), masked_indices
  
      train_sampler = RandomSampler(train_dataset)
      train_dataloader = DataLoader(
          train_dataset, sampler=train_sampler, batch_size=args.train_batch_size, collate_fn=collate
      )
  
      if args.max_steps > 0:
          t_total = args.max_steps
          args.num_train_epochs = args.max_steps // (len(train_dataloader) // args.gradient_accumulation_steps) + 1
      else:
          t_total = len(train_dataloader) // args.gradient_accumulation_steps * args.num_train_epochs
  
      # Prepare optimizer and schedule (linear warmup and decay)
      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": args.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},
      ]
      optimizer = AdamW(optimizer_grouped_parameters, lr=args.learning_rate, eps=args.adam_epsilon)
      scheduler = get_linear_schedule_with_warmup(
          optimizer, num_warmup_steps=args.warmup_steps, num_training_steps=t_total
      )
  
      # Check if saved optimizer or scheduler states exist
      if (
          args.model_name_or_path
          and os.path.isfile(os.path.join(args.model_name_or_path, "optimizer.pt"))
          and os.path.isfile(os.path.join(args.model_name_or_path, "scheduler.pt"))
      ):
          # Load in optimizer and scheduler states
          optimizer.load_state_dict(torch.load(os.path.join(args.model_name_or_path, "optimizer.pt")))
          scheduler.load_state_dict(torch.load(os.path.join(args.model_name_or_path, "scheduler.pt")))
  
      if args.fp16:
          try:
              from apex import amp
          except ImportError:
              raise ImportError("Please install apex from https://www.github.com/nvidia/apex to use fp16 training.")
          model, optimizer = amp.initialize(model, optimizer, opt_level=args.fp16_opt_level)
  
      # multi-gpu training (should be after apex fp16 initialization)
      if args.n_gpu > 1:
          model = torch.nn.DataParallel(model)
  
      # Train!
      logger.info("***** Running training *****")
      logger.info("  Num examples = %d", len(train_dataset))
      logger.info("  Num Epochs = %d", args.num_train_epochs)
      logger.info("  Instantaneous batch size per GPU = %d", args.per_gpu_train_batch_size)
      logger.info(
          "  Total train batch size (w. accumulation) = %d",
          args.train_batch_size * args.gradient_accumulation_steps
      )
      logger.info("  Gradient Accumulation steps = %d", args.gradient_accumulation_steps)
      logger.info("  Total optimization steps = %d", t_total)
  
      global_step = 0
      epochs_trained = 0
      steps_trained_in_current_epoch = 0
      # Check if continuing training from a checkpoint
      if args.model_name_or_path and os.path.exists(args.model_name_or_path):
          try:
              # set global_step to gobal_step of last saved checkpoint from model path
              checkpoint_suffix = args.model_name_or_path.split("-")[-1].split("/")[0]
              global_step = int(checkpoint_suffix)
              epochs_trained = global_step // (len(train_dataloader) // args.gradient_accumulation_steps)
              steps_trained_in_current_epoch = global_step % (len(train_dataloader) // args.gradient_accumulation_steps)
  
              logger.info("  Continuing training from checkpoint, will skip to saved global_step")
              logger.info("  Continuing training from epoch %d", epochs_trained)
              logger.info("  Continuing training from global step %d", global_step)
              logger.info("  Will skip the first %d steps in the first epoch", steps_trained_in_current_epoch)
          except ValueError:
              logger.info("  Starting fine-tuning.")
  
      tr_loss, logging_loss = 0.0, 0.0
  
      model_to_resize = model.module if hasattr(model, "module") else model
      model_to_resize.resize_token_embeddings(len(tokenizer))
  
      model.zero_grad()
      train_iterator = trange(
          epochs_trained, int(args.num_train_epochs), desc="Epoch"
      )
      set_seed(args)  # Added here for reproducibility
      for _ in train_iterator:
          epoch_iterator = tqdm(train_dataloader, desc="Iteration")
          for step, (batch, batch_labels, masked_indices) in enumerate(epoch_iterator):
              batch_labels, masked_indices = batch_labels.squeeze(), torch.stack(masked_indices).squeeze()
              # Skip past any already trained steps if resuming training
              if steps_trained_in_current_epoch > 0:
                  steps_trained_in_current_epoch -= 1
                  continue
  
              inputs, labels = mask_tokens(batch, tokenizer, batch_labels, masked_indices, args)
              inputs = inputs.to(args.device)
              labels = labels.to(args.device).long()
              model.train()
              prediction_scores = model(inputs)[0]
              masked_lm_loss = CrossEntropyLoss()(
                  prediction_scores.view(-1, model.config.vocab_size), labels.view(-1))
              loss = masked_lm_loss
  
              if args.n_gpu > 1:
                  loss = loss.mean()  # mean() to average on multi-gpu parallel training
              if args.gradient_accumulation_steps > 1:
                  loss = loss / args.gradient_accumulation_steps
  
              loss.backward()
              tr_loss += loss.item()
              if (step + 1) % args.gradient_accumulation_steps == 0:
                  if args.fp16:
                      torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer), args.max_grad_norm)
                  else:
                      torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)
                  optimizer.step()
                  scheduler.step()  # Update learning rate schedule
                  model.zero_grad()
                  global_step += 1
  
                  if args.logging_steps > 0 and global_step % args.logging_steps == 0:
                      # Log metrics
                      if args.evaluate_during_training:
                          results = evaluate(args, model, tokenizer)
                          for key, value in results.items():
                              tb_writer.add_scalar("eval_{}".format(key), value, global_step)
                      tb_writer.add_scalar("lr", scheduler.get_lr()[0], global_step)
                      tb_writer.add_scalar("loss", (tr_loss - logging_loss) / args.logging_steps, global_step)
                      logging_loss = tr_loss
  
                  if args.save_steps > 0 and global_step % args.save_steps == 0:
                      checkpoint_prefix = "checkpoint"
                      # Save model checkpoint
                      output_dir = os.path.join(args.output_dir, "{}-{}".format(checkpoint_prefix, global_step))
                      os.makedirs(output_dir, exist_ok=True)
                      model_to_save = (
                          model.module if hasattr(model, "module") else model
                      )
                      model_to_save.save_pretrained(output_dir)
                      tokenizer.save_pretrained(output_dir)
  
                      torch.save(args, os.path.join(output_dir, "training_args.bin"))
                      logger.info("Saving model checkpoint to %s", output_dir)
  
                      _rotate_checkpoints(args, checkpoint_prefix)
  
                      torch.save(optimizer.state_dict(), os.path.join(output_dir, "optimizer.pt"))
                      torch.save(scheduler.state_dict(), os.path.join(output_dir, "scheduler.pt"))
                      logger.info("Saving optimizer and scheduler states to %s", output_dir)
  
              if args.max_steps > 0 and global_step > args.max_steps:
                  epoch_iterator.close()
                  break
  
          if args.max_steps > 0 and global_step > args.max_steps:
              train_iterator.close()
              break
  
      tb_writer.close()
      return global_step, tr_loss / global_step
  
  
  def evaluate(args, model: AutoModelWithLMHead, tokenizer: AutoTokenizer, prefix="") -> Dict:
      # Loop to handle MNLI double evaluation (matched, mis-matched)
      eval_output_dir = args.output_dir
      eval_dataset = load_and_cache_examples(args, tokenizer, evaluate=True)
      os.makedirs(eval_output_dir, exist_ok=True)
      args.eval_batch_size = args.per_gpu_eval_batch_size * max(1, args.n_gpu)
  
      def collate(examples: List[torch.Tensor]):
          instances, labels, masked_indices = zip(*examples)
          if tokenizer._pad_token is None:
              return pad_sequence(instances, batch_first=True), pad_sequence(labels, batch_first=True), masked_indices
  
          return pad_sequence(instances, batch_first=True, padding_value=tokenizer.pad_token_id), \
                 pad_sequence(labels, batch_first=True, padding_value=tokenizer.pad_token_id), masked_indices
  
      eval_sampler = SequentialSampler(eval_dataset)
      eval_dataloader = DataLoader(
          eval_dataset, sampler=eval_sampler, batch_size=args.eval_batch_size, collate_fn=collate
      )
  
      # multi-gpu evaluate
      if args.n_gpu > 1:
          model = torch.nn.DataParallel(model)
  
      # Eval!
      logger.info("***** Running evaluation {} *****".format(prefix))
      logger.info("  Num examples = %d", len(eval_dataset))
      logger.info("  Batch size = %d", args.eval_batch_size)
      eval_loss = 0.0
      nb_eval_steps = 0
      model.eval()
  
      for batch, batch_labels, masked_indices in tqdm(eval_dataloader, desc="Evaluating"):
          batch_labels, masked_indices = batch_labels.squeeze(), torch.stack(masked_indices).squeeze()
          inputs, labels = mask_tokens(batch, tokenizer, batch_labels, masked_indices, args)
          inputs = inputs.to(args.device)
          labels = labels.to(args.device)
  
          with torch.no_grad():
              prediction_scores = model(inputs)[0]
              masked_lm_loss = CrossEntropyLoss()(
                  prediction_scores.view(-1, model.config.vocab_size), labels.view(-1).long())
              loss = masked_lm_loss
  
              eval_loss += loss.mean().item()
          nb_eval_steps += 1
  
      eval_loss = eval_loss / nb_eval_steps
      perplexity = torch.exp(torch.tensor(eval_loss))
  
      result = {"perplexity": perplexity}
  
      output_eval_file = os.path.join(eval_output_dir, prefix, "eval_results.txt")
      with open(output_eval_file, "w") as writer:
          logger.info("***** Eval results {} *****".format(prefix))
          for key in sorted(result.keys()):
              logger.info("  %s = %s", key, str(result[key]))
              writer.write("%s = %s\n" % (key, str(result[key])))
  
      return result
  
  
  def main():
      parser = argparse.ArgumentParser()
  
      # Required parameters
      parser.add_argument(
          "--output_dir",
          type=str,
          required=True,
          help="The output directory where the model predictions and checkpoints will be written.",
      )
  
      # Other parameters
      parser.add_argument(
          "--train_data_file",
          default=None,
          type=str,
          required=False,
          help="The input training data file (a text file)."
      )
      parser.add_argument(
          "--eval_data_file",
          default=None,
          type=str,
          help="An optional input evaluation data file to evaluate the perplexity on (a text file).",
      )
      parser.add_argument(
          "--line_by_line",
          action="store_true",
          help="Whether distinct lines of text in the dataset are to be handled as distinct sequences.",
      )
      parser.add_argument(
          "--should_continue", action="store_true", help="Whether to continue from latest checkpoint in output_dir"
      )
      parser.add_argument(
          "--model_name_or_path",
          default=None,
          type=str,
          help="The model checkpoint for weights initialization. Leave None if you want to train a model from scratch.",
      )
  
      parser.add_argument(
          "--mlm_probability", type=float, default=0.15, help="Ratio of tokens to mask for masked language modeling loss"
      )
  
      parser.add_argument(
          "--config_name",
          default=None,
          type=str,
          help="Optional pretrained config name or path if not the same as model_name_or_path. If both are None, initialize a new config.",
      )
      parser.add_argument(
          "--tokenizer_name",
          default=None,
          type=str,
          help="Optional pretrained tokenizer name or path if not the same as model_name_or_path. If both are None, initialize a new tokenizer.",
      )
      parser.add_argument(
          "--cache_dir",
          default=None,
          type=str,
          help="Optional directory to store the pre-trained models downloaded from s3 (instead of the default one)",
      )
      parser.add_argument(
          "--block_size",
          default=-1,
          type=int,
          help="Optional input sequence length after tokenization."
          "The training dataset will be truncated in block of this size for training."
          "Default to the model max input length for single sentence inputs (take into account special tokens).",
      )
      parser.add_argument("--do_train", action="store_true", help="Whether to run training.")
      parser.add_argument("--do_eval", action="store_true", help="Whether to run eval on the dev set.")
      parser.add_argument(
          "--evaluate_during_training", action="store_true", help="Run evaluation during training at each logging step."
      )
  
      parser.add_argument("--per_gpu_train_batch_size", default=4, type=int, help="Batch size per GPU/CPU for training.")
      parser.add_argument(
          "--per_gpu_eval_batch_size", default=4, type=int, help="Batch size per GPU/CPU for evaluation."
      )
      parser.add_argument(
          "--gradient_accumulation_steps",
          type=int,
          default=1,
          help="Number of updates steps to accumulate before performing a backward/update pass.",
      )
      parser.add_argument("--learning_rate", default=5e-5, type=float, help="The initial learning rate for Adam.")
      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("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.")
      parser.add_argument(
          "--num_train_epochs", default=1.0, type=float, help="Total number of training epochs to perform."
      )
      parser.add_argument(
          "--max_steps",
          default=-1,
          type=int,
          help="If > 0: set total number of training steps to perform. Override num_train_epochs.",
      )
      parser.add_argument("--warmup_steps", default=0, type=int, help="Linear warmup over warmup_steps.")
  
      parser.add_argument("--logging_steps", type=int, default=500, help="Log every X updates steps.")
      parser.add_argument("--save_steps", type=int, default=500, help="Save checkpoint every X updates steps.")
      parser.add_argument(
          "--save_total_limit",
          type=int,
          default=None,
          help="Limit the total amount of checkpoints, delete the older checkpoints in the output_dir, does not delete by default",
      )
      parser.add_argument(
          "--eval_all_checkpoints",
          action="store_true",
          help="Evaluate all checkpoints starting with the same prefix as model_name_or_path ending and ending with step number",
      )
      parser.add_argument("--no_cuda", action="store_true", help="Avoid using CUDA when available")
      parser.add_argument(
          "--overwrite_output_dir", action="store_true", help="Overwrite the content of the output directory"
      )
      parser.add_argument(
          "--overwrite_cache", action="store_true", help="Overwrite the cached training and evaluation sets"
      )
      parser.add_argument("--seed", type=int, default=42, help="random seed for initialization")
  
      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="O1",
          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("--server_ip", type=str, default="", help="For distant debugging.")
      parser.add_argument("--server_port", type=str, default="", help="For distant debugging.")
      parser.add_argument("--max_length", type=int, default=50, help="")
      parser.add_argument("--device", default="cpu", required=False)
      args = parser.parse_args()
      args.mlm = True
  
      if args.eval_data_file is None and args.do_eval:
          raise ValueError(
              "Cannot do evaluation without an evaluation data file. Either supply a file to --eval_data_file "
              "or remove the --do_eval argument."
          )
      if args.should_continue:
          sorted_checkpoints = _sorted_checkpoints(args)
          if len(sorted_checkpoints) == 0:
              raise ValueError("Used --should_continue but no checkpoint was found in --output_dir.")
          else:
              args.model_name_or_path = sorted_checkpoints[-1]
  
      if (
          os.path.exists(args.output_dir)
          and os.listdir(args.output_dir)
          and args.do_train
          and not args.overwrite_output_dir
      ):
          raise ValueError(
              "Output directory ({}) already exists and is not empty. Use --overwrite_output_dir to overcome.".format(
                  args.output_dir
              )
          )
  
      # Setup distant debugging if needed
      if args.server_ip and args.server_port:
          # Distant debugging - see https://code.visualstudio.com/docs/python/debugging#_attach-to-a-local-script
          import ptvsd
  
          print("Waiting for debugger attach")
          ptvsd.enable_attach(address=(args.server_ip, args.server_port), redirect_output=True)
          ptvsd.wait_for_attach()
  
      if args.device != "cpu":
          args.device = f"cuda:{args.device}"
  
      device = torch.device(args.device)
      args.n_gpu = 1
      args.device = device
  
      # Setup logging
      logging.basicConfig(
          format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
          datefmt="%m/%d/%Y %H:%M:%S",
          level=logging.INFO
      )
  
      # Set seed
      set_seed(args)
  
      if args.config_name:
          config = AutoConfig.from_pretrained(args.config_name, cache_dir=args.cache_dir)
      else:
          config = AutoConfig.from_pretrained(args.model_name_or_path, cache_dir=args.cache_dir)
  
      if args.tokenizer_name:
          tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_name, cache_dir=args.cache_dir)
      elif args.model_name_or_path:
          tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path, cache_dir=args.cache_dir)
      else:
          raise ValueError(
              "You are instantiating a new {} tokenizer. This is not supported, but you can do it from another script, save it,"
              "and load it from here, using --tokenizer_name".format(AutoTokenizer.__name__)
          )
  
      if args.block_size <= 0:
          args.block_size = tokenizer.max_len
          # Our input block size will be the max possible for the model
      else:
          args.block_size = min(args.block_size, tokenizer.max_len)
  
      if args.model_name_or_path:
          model = AutoModelWithLMHead.from_pretrained(
              args.model_name_or_path,
              from_tf=bool(".ckpt" in args.model_name_or_path),
              config=config,
              cache_dir=args.cache_dir,
          )
      else:
          logger.info("Training new model from scratch")
          model = AutoModelWithLMHead(config=config)
  
      model.to(args.device)
      logger.info("Training/evaluation parameters %s", args)
  
      # Training
      if args.do_train:
          train_dataset = load_and_cache_examples(args, tokenizer, evaluate=False)
          global_step, tr_loss = train(args, train_dataset, model, tokenizer)
          logger.info(" global_step = %s, average loss = %s", global_step, tr_loss)
  
      # Saving best-practices: if you use save_pretrained for the model and tokenizer,
      # you can reload them using from_pretrained()
      if args.do_train:
  
          # Create output directory if needed
          os.makedirs(args.output_dir, exist_ok=True)
  
          logger.info("Saving model checkpoint to %s", args.output_dir)
          # Save a trained model, configuration and tokenizer using `save_pretrained()`.
          # They can then be reloaded using `from_pretrained()`
          model_to_save = (
              model.module if hasattr(model, "module") else model
          )
          model_to_save.save_pretrained(args.output_dir)
          tokenizer.save_pretrained(args.output_dir)
  
          # Good practice: save your training arguments together with the trained model
          torch.save(args, os.path.join(args.output_dir, "training_args.bin"))
  
          # Load a trained model and vocabulary that you have fine-tuned
          model = AutoModelWithLMHead.from_pretrained(args.output_dir)
          tokenizer = AutoTokenizer.from_pretrained(args.output_dir)
          model.to(args.device)
  
      # Evaluation
      results = {}
      if args.do_eval:
          checkpoints = [args.output_dir]
          if args.eval_all_checkpoints:
              checkpoints = list(
                  os.path.dirname(c) for c in sorted(glob.glob(args.output_dir + "/**/" + WEIGHTS_NAME, recursive=True))
              )
              logging.getLogger("transformers.modeling_utils").setLevel(logging.WARN)  # Reduce logging
          logger.info("Evaluate the following checkpoints: %s", checkpoints)
          for checkpoint in checkpoints:
              global_step = checkpoint.split("-")[-1] if len(checkpoints) > 1 else ""
              prefix = checkpoint.split("/")[-1] if checkpoint.find("checkpoint") != -1 else ""
  
              model = AutoModelWithLMHead.from_pretrained(checkpoint)
              model.to(args.device)
              result = evaluate(args, model, tokenizer, prefix=prefix)
              result = dict((k + "_{}".format(global_step), v) for k, v in result.items())
              results.update(result)
  
      return results
  
  
  if __name__ == "__main__":
      main()