text: """Train and test methods for MLM-MLM pretraining.
  
  Author:
      Jeffrey Shen
  """
  
  import random
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  import torch.optim as optim
  import torch.utils.data as data
  import torch.cuda.amp as amp
  
  
  from collections import OrderedDict
  from tqdm import tqdm
  from ujson import load as json_load
  
  from models import RoBERTa
  from datasets.bpe_squad import collate_fn, MLM
  from preprocess.bpe import BPE
  import trainer.trainer as base_trainer
  import trainer.util as util
  import trainer.stats as stats
  import models.transformer as T
  import trainer.scheduler as sched
  
  
  def add_special_tokens(args):
      args.ignore_idx = -1
      args.padding_idx = 0
      args.cls_idx = 1
      args.sep_idx = 2
      args.mask_idx = 3
  
  
  def get_args(args):
      # Compute derived args values
      device, args.gpu_ids = util.get_available_devices()
  
      args.batch_size_per_gpu = args.batch_size
      args.batch_size *= max(1, len(args.gpu_ids))
      args.num_steps = args.epoch_size // args.batch_size // args.gradient_accumulation
      if args.num_epochs >= 0:
          args.num_steps *= args.num_epochs
  
      if args.decay_forever:
          args.num_steps = float("inf")
  
      return args, device
  
  
  def get_bpe(args):
      bpe = BPE()
      with open(args.bpe_file, "r") as file:
          bpe.load_state_dict(json_load(file))
      add_special_tokens(args)
      return bpe
  
  
  def get_dataset(args, epoch_size, file, bpe):
      # Don't need to supply special idxs, since they're the same.
      dataset = MLM(
          file,
          max_tokens=len(bpe),
          epoch_size=epoch_size,
          mask_prob=args.mask_prob,
          unmask_prob=args.unmask_prob,
          randomize_prob=args.randomize_prob,
          block_size=args.max_positions,
      )
      loader = data.DataLoader(
          dataset,
          batch_size=args.batch_size,
          num_workers=args.num_workers,
          collate_fn=collate_fn,
      )
      return dataset, loader
  
  
  def get_model(args, bpe):
      model = RoBERTa(
          dim=args.dim,
          n_heads=args.n_heads,
          ff_dim=args.ff_dim,
          activation=args.activation,
          dropout=args.dropout,
          attn_dropout=args.attn_dropout,
          act_dropout=args.act_dropout,
          n_layers=args.n_layers,
          max_positions=args.max_positions,
          max_tokens=len(bpe),
          padding_idx=args.padding_idx,
          ignore_idx=args.ignore_idx,
          prenorm=args.prenorm,
          qa=False,
      )
      return model
  
  
  def train(args):
      trainer = base_trainer.Trainer()
      args, device = get_args(args)
      args, log, tbx = trainer.setup(args)
  
      # Get BPE
      log.info("Loading BPE...")
      bpe = get_bpe(args)
      log.info("Loaded {} BPE tokens".format(len(bpe)))
  
      # Get data loader
      log.info("Building dataset...")
      train_dataset, train_loader = get_dataset(
          args, args.epoch_size, args.train_record_file, bpe
      )
      dev_dataset, dev_loader = get_dataset(
          args, args.dev_epoch_size, args.dev_record_file, bpe
      )
  
      # Get model
      log.info("Building model...")
      model = get_model(args, bpe)
      model = trainer.setup_model(model, device)
  
      # Get optimizer, scheduler, and scaler
      optimizer = optim.AdamW(
          model.parameters(),
          args.lr,
          betas=(args.beta_1, args.beta_2),
          eps=args.eps,
          weight_decay=args.l2_wd,
      )
  
      scheduler = sched.get_linear_warmup_power_decay_scheduler(
          optimizer, args.warmup_steps, args.num_steps, power=args.power_decay
      )
  
      scaler = amp.GradScaler()
      optimizer, scheduler, scaler = trainer.setup_optimizer(optimizer, scheduler, scaler)
  
      # Train
      log.info("Training...")
      model.train()
      sample_num = 0
      samples_till_eval = args.eval_per_n_samples
      epoch = 0
      step = 0
      trainer.setup_saver()
      trainer.setup_random()
      sample_num, samples_till_eval, epoch, step = trainer.setup_step(
          step_vars=(sample_num, samples_till_eval, epoch, step)
      )
      trainer.setup_close()
  
      while epoch != args.num_epochs:
          trainer.save_checkpoint(step_vars=(sample_num, samples_till_eval, epoch, step))
          epoch += 1
          log.info(f"Starting epoch {epoch}...")
          # Print histogram of weights every epoch
          for tags, params in model.named_parameters():
              tbx.add_histogram(tags, params.data, epoch)
          with torch.enable_grad(), tqdm(total=len(train_loader.dataset)) as progress_bar:
              for x, y in train_loader:
                  loss_weight = args.gradient_accumulation * (args.mlm_samples * args.lambda_weight + 1)
                  batch_size = x.size(0)
                  loss, loss_val, scores = forward(x, y, args, device, model)
                  loss = loss / loss_weight
                  scaler.scale(loss).backward()
  
                  # ELECTRA MLM
                  x, y = x.to(device), y.to(device)
                  x, y = sample_mlm_pred(
                      model.module, x, y, scores, args.mlm_samples, args
                  )
                  # Try to free up scores
                  del scores
  
                  loss_val_e = 0
                  for i in range(args.mlm_samples):
                      loss, loss_val_item, _ = forward(x[i], y, args, device, model)
                      loss = loss / loss_weight * args.lambda_weight
                      loss_val_e += loss_val_item / x.size(0)
                      scaler.scale(loss).backward()
  
                  # Optimizer step
                  if (step + 1) % args.gradient_accumulation == 0:
                      scaler.unscale_(optimizer)
                      nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)
                      scaler.step(optimizer)
                      scaler.update()
                      scheduler.step()
                      optimizer.zero_grad()
  
                  # Log info
                  step += 1
                  sample_num += batch_size
                  progress_bar.update(batch_size)
                  progress_bar.set_postfix(epoch=epoch, NLL=loss_val, NLL_E=loss_val_e)
                  tbx.add_scalar("train/NLL", loss_val, sample_num)
                  tbx.add_scalar("train/NLL_E", loss_val_e, sample_num)
                  tbx.add_scalar("train/LR", optimizer.param_groups[0]["lr"], sample_num)
                  tbx.add_scalar(
                      "train/steps", step // args.gradient_accumulation, sample_num
                  )
  
                  samples_till_eval -= batch_size
                  if samples_till_eval <= 0:
                      samples_till_eval = args.eval_per_n_samples
  
                      # Evaluate and save checkpoint
                      log.info(f"Evaluating at sample step {sample_num}...")
                      results, preds, preds_e = evaluate(model, dev_loader, device, args)
                      trainer.save_best(sample_num, results[args.metric_name])
  
                      # Log to console
                      results_str = ", ".join(
                          f"{k}: {v:05.2f}" for k, v in results.items()
                      )
                      log.info(f"Dev {results_str}")
  
                      # Log to TensorBoard
                      log.info("Visualizing in TensorBoard...")
                      for k, v in results.items():
                          tbx.add_scalar(f"dev/{k}", v, sample_num)
                      visualize(
                          tbx,
                          preds=preds,
                          bpe=bpe,
                          sample_num=sample_num,
                          split="dev",
                          num_visuals=args.num_visuals,
                      )
                      visualize(
                          tbx,
                          preds=preds_e,
                          bpe=bpe,
                          sample_num=sample_num,
                          split="dev_e",
                          num_visuals=args.num_visuals,
                      )
  
  
  def forward(x, y, args, device, model, autocast=True):
      # Setup for forward
      x = x.to(device)
      padding_mask = T.get_padding_mask(x, args.padding_idx)
  
      # Forward
      with amp.autocast(enabled=autocast):
          scores = model(x, padding_mask=padding_mask)
          scores = model.module.mask_scores(scores, padding_mask)
          y = y.to(device)
          loss = model.module.get_loss(scores, y)
          loss_val = loss.item()
  
      return loss, loss_val, scores
  
  
  def visualize(tbx, preds, bpe, sample_num, split, num_visuals):
      """Visualize text examples to TensorBoard.
  
      Args:
          tbx (SummaryWriter): Summary writer.
          preds (list(tuple)): list of (pred, ques, ans)
          bpe (BPE): bpe encoder/decoder
          sample_num (int): Number of examples seen so far during training.
          split (str): Name of data split being visualized.
          num_visuals (int): Number of visuals to select at random from preds.
      """
      if num_visuals <= 0:
          return
      if num_visuals > len(preds):
          num_visuals = len(preds)
  
      visuals = random.sample(preds, k=num_visuals)
  
      for i, (pred, ques, ans) in enumerate(visuals):
          pred = bpe.decode(pred)
          ques = bpe.decode(ques)
          ans = bpe.decode(ans)
          tbl_fmt = (
              f"- **Answer:** {ans}\n"
              + f"- **MLM Question:** {ques}\n"
              + f"- **Prediction:** {pred}"
          )
          tbx.add_text(
              tag=f"{split}/{i+1}_of_{num_visuals}",
              text_string=tbl_fmt,
              global_step=sample_num,
          )
  
  
  def get_mlm_pred(model, x, y, scores, args):
      pred = model.get_top(scores)
      mask = y != args.ignore_idx
      acc = (pred[mask] == y[mask]).float().mean().item()
      ans = y.clone().detach()
      ans[~mask] = x[~mask]
      pred[~mask] = x[~mask]
      pred = pred.tolist()
      pred = [[token for token in sample if token != args.padding_idx] for sample in pred]
      ques = x.tolist()
      ques = [[token for token in sample if token != args.padding_idx] for sample in ques]
      ans = ans.tolist()
      ans = [[token for token in sample if token != args.padding_idx] for sample in ans]
      return pred, ques, ans, acc
  
  
  def sample_mlm_pred(model, x, y, scores, k, args):
      y = y.clone().detach()
      x = x.clone().detach()
      scores = scores.clone().detach()
      mask = y != args.ignore_idx
      y[~mask] = x[~mask]
  
      # Just in case, ignore the padding idx
      padding_mask = T.get_padding_mask(x, args.padding_idx)
      y[padding_mask] = args.ignore_idx
  
      x = x.unsqueeze(-1)
      x = x.repeat(1, 1, k)
      # Don't choose the padding idx!
      scores[:, :, args.padding_idx] = float("-inf")
      x[mask, :] = model.sample(scores[mask, :], k, alpha=args.sample_temperature)
      return x.permute(-1, 0, 1), y
  
  
  def evaluate(model, data_loader, device, args):
      nll_meter = stats.AverageMeter()
      acc_meter = stats.AverageMeter()
      nll_e_meter = stats.AverageMeter()
      acc_e_meter = stats.AverageMeter()
  
      model.eval()
      preds = []
      preds_e = []
      with torch.no_grad():
          for x, y in data_loader:
              batch_size = x.size(0)
              _, loss_val, scores = forward(x, y, args, device, model)
              nll_meter.update(loss_val, batch_size)
  
              x = x.to(device)
              y = y.to(device)
              pred, ques, ans, acc = get_mlm_pred(model.module, x, y, scores, args)
  
              acc_meter.update(acc, batch_size)
              preds += zip(pred, ques, ans)
  
              x, y = sample_mlm_pred(model.module, x, y, scores, 1, args)
              x = x.squeeze(0)
              _, loss_val, scores = forward(x, y, args, device, model)
              nll_e_meter.update(loss_val, batch_size)
  
              pred, ques, ans, acc = get_mlm_pred(model.module, x, y, scores, args)
              acc_e_meter.update(acc, batch_size)
              preds_e += zip(pred, ques, ans)
  
      model.train()
  
      results_list = [("NLL", nll_meter.avg), ("acc", acc_meter.avg)]
      results_list += [("NLL_E", nll_e_meter.avg), ("acc_e", acc_e_meter.avg)]
      results = OrderedDict(results_list)
  
      return results, preds, preds_e
  
  
  def add_mlm_args(parser):
      parser.add_argument(
          "--mask_prob", type=float, default=0.25, help="Mask probability."
      )
      parser.add_argument(
          "--unmask_prob",
          type=float,
          default=0.1,
          help="Probability to leave mask unchanged.",
      )
      parser.add_argument(
          "--randomize_prob",
          type=float,
          default=0.1,
          help="Probability to use a random token instead of mask.",
      )
  
  
  def add_electra_mlm_args(parser):
      parser.add_argument(
          "--sample_temperature", type=float, default=2.0, help="Sample temperature."
      )
      parser.add_argument(
          "--mlm_samples", type=int, default=3, help="How many to sample for replace MLM"
      )
      parser.add_argument(
          "--lambda_weight", type=float, default=16, help="Multiplier for the replace MLM"
      )
  
  
  def add_train_args(parser):
      """Add arguments needed in train.py."""
      add_train_test_args(parser)
      base_trainer.add_train_args(parser)
      add_mlm_args(parser)
      add_electra_mlm_args(parser)
  
      parser.add_argument(
          "--epoch_size", type=int, default=25000, help="Number of samples per epoch."
      )
      parser.add_argument(
          "--dev_epoch_size",
          type=int,
          default=1000,
          help="Number of samples for a dev evaluation.",
      )
      parser.add_argument(
          "--eval_per_n_samples",
          type=int,
          default=12500,
          help="Number of samples between successive evaluations.",
      )
      parser.add_argument(
          "--gradient_accumulation",
          type=int,
          default=4,
      )
      parser.add_argument("--lr", type=float, default=0.025, help="Learning rate.")
      parser.add_argument(
          "--warmup_steps", type=float, default=7500, help="Warmup optimizer steps."
      )
      parser.add_argument(
          "--power_decay", type=float, default=-0.5, help="Power of the decay."
      )
      parser.add_argument(
          "--decay_forever",
          type=lambda s: s.lower().startswith("t"),
          default=True,
          help="Whether the decay should reach end_lr at the end of training, or in the limit to infinity",
      )
  
      parser.add_argument("--l2_wd", type=float, default=0.01, help="AdamW weight decay.")
      parser.add_argument("--eps", type=float, default=1e-6, help="Adam epsilon.")
      parser.add_argument("--beta_1", type=float, default=0.9, help="Adam beta_1.")
      parser.add_argument("--beta_2", type=float, default=0.98, help="Adam beta_2.")
      parser.add_argument(
          "--num_epochs",
          type=int,
          default=30,
          help="Number of epochs for which to train. Negative means forever.",
      )
      parser.add_argument(
          "--metric_name",
          type=str,
          default="NLL_E",
          choices=("NLL", "NLL_E"),
          help="Name of dev metric to determine best checkpoint.",
      )
      parser.add_argument(
          "--max_grad_norm",
          type=float,
          default=5.0,
          help="Maximum gradient norm for gradient clipping.",
      )
  
  
  def add_train_test_args(parser):
      parser.add_argument(
          "--num_workers",
          type=int,
          default=4,
          help="Number of sub-processes to use per data loader.",
      )
      parser.add_argument(
          "--batch_size",
          type=int,
          default=16,
          help="Batch size per GPU. Scales automatically when \
                                multiple GPUs are available.",
      )
      parser.add_argument(
          "--num_visuals",
          type=int,
          default=10,
          help="Number of examples to visualize in TensorBoard.",
      )
  
      # Model params
      parser.add_argument(
          "--dim",
          type=int,
          default=768,
          help="Embedding dimension.",
      )
      parser.add_argument(
          "--n_heads",
          type=int,
          default=12,
          help="Attention heads.",
      )
      parser.add_argument(
          "--ff_dim",
          type=int,
          default=3072,
          help="Feedforward dimension.",
      )
      parser.add_argument(
          "--activation",
          choices=["relu", "gelu"],
          default="gelu",
          help="Feedforward activation function.",
      )
      parser.add_argument(
          "--dropout",
          type=float,
          default=0.1,
          help="Dropout probability.",
      )
      parser.add_argument(
          "--attn_dropout",
          type=float,
          default=0.1,
          help="Dropout probability for attention weights within self attn.",
      )
      parser.add_argument(
          "--act_dropout",
          type=float,
          default=0.0,
          help="Dropout probability after activation within FF.",
      )
      parser.add_argument(
          "--n_layers",
          type=int,
          default=12,
          help="Number of layers.",
      )
      parser.add_argument(
          "--max_positions",
          type=int,
          default=512,
          help="Maximum number of tokens.",
      )
      parser.add_argument(
          "--prenorm",
          type=lambda s: s.lower().startswith("t"),
          default=False,
          help="Whether to put LayerNorm after the residual or before.",
      )