text: """Train and test methods.
  
  Author:
      Chris Chute (chute@stanford.edu)
  """
  
  import numpy as np
  import random
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  import torch.optim as optim
  import torch.optim.lr_scheduler as sched
  import torch.utils.data as data
  import torch.cuda.amp as amp
  
  
  from collections import OrderedDict
  from json import dumps
  from torch.utils.tensorboard import SummaryWriter
  from tqdm import tqdm
  from ujson import load as json_load
  from os.path import join
  
  from models import BiDAF
  from datasets.bidaf_squad import collate_fn, SQuAD
  import eval
  import trainer.util as util
  import trainer.stats as stats
  
  
  def train(args):
      # Set up logging and devices
      log = util.get_logger(args.save_dir, args.name)
      tbx = SummaryWriter(args.save_dir)
      device, args.gpu_ids = util.get_available_devices()
      log.info(f"Args: {dumps(vars(args), indent=4, sort_keys=True)}")
      args.batch_size *= max(1, len(args.gpu_ids))
  
      # Set random seed
      log.info(f"Using random seed {args.seed}...")
      random.seed(args.seed)
      np.random.seed(args.seed)
      torch.manual_seed(args.seed)
      torch.cuda.manual_seed_all(args.seed)
  
      # Get embeddings
      log.info("Loading embeddings...")
      word_vectors = util.torch_from_json(args.word_emb_file)
  
      # Get model
      log.info("Building model...")
      model = BiDAF(
          word_vectors=word_vectors,
          hidden_size=args.hidden_size,
          drop_prob=args.drop_prob,
          use_glove=args.use_glove,
      )
      model = nn.DataParallel(model, args.gpu_ids)
      if args.load_path:
          log.info(f"Loading checkpoint from {args.load_path}...")
          model, step = util.load_model(model, args.load_path, args.gpu_ids)
      else:
          step = 0
      model = model.to(device)
      model.train()
      ema = stats.EMA(model, args.ema_decay)
  
      # Get saver
      saver = util.CheckpointSaver(
          args.save_dir,
          max_checkpoints=args.max_checkpoints,
          metric_name=args.metric_name,
          maximize_metric=args.maximize_metric,
          log=log,
      )
  
      # Get optimizer and scheduler
      optimizer = optim.Adam(model.parameters(), args.lr, weight_decay=args.l2_wd)
      scheduler = sched.LambdaLR(optimizer, lambda s: 1.0)  # Constant LR
  
      # Get data loader
      log.info("Building dataset...")
      train_dataset = SQuAD(args.train_record_file, args.use_squad_v2)
      train_loader = data.DataLoader(
          train_dataset,
          batch_size=args.batch_size,
          shuffle=True,
          num_workers=args.num_workers,
          collate_fn=collate_fn,
      )
      dev_dataset = SQuAD(args.dev_record_file, args.use_squad_v2)
      dev_loader = data.DataLoader(
          dev_dataset,
          batch_size=args.batch_size,
          shuffle=False,
          num_workers=args.num_workers,
          collate_fn=collate_fn,
      )
  
      # Train
      log.info("Training...")
      steps_till_eval = args.eval_steps
      epoch = step // len(train_dataset)
      scaler = amp.GradScaler()
  
      while epoch != args.num_epochs:
          epoch += 1
          log.info(f"Starting epoch {epoch}...")
          with torch.enable_grad(), tqdm(total=len(train_loader.dataset)) as progress_bar:
              for cw_idxs, cc_idxs, qw_idxs, qc_idxs, y1, y2, ids in train_loader:
                  # Setup for forward
                  cw_idxs = cw_idxs.to(device)
                  qw_idxs = qw_idxs.to(device)
                  batch_size = cw_idxs.size(0)
                  optimizer.zero_grad()
  
                  # Forward
                  with amp.autocast():
                      log_p1, log_p2 = model(cw_idxs, qw_idxs)
                      y1, y2 = y1.to(device), y2.to(device)
                      loss = F.nll_loss(log_p1, y1) + F.nll_loss(log_p2, y2)
                      loss_val = loss.item()
  
                  # Backward
                  scaler.scale(loss).backward()
                  scaler.unscale_(optimizer)
                  nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)
                  scaler.step(optimizer)
                  scaler.update()
                  scheduler.step(step // batch_size)
                  ema(model, step // batch_size)
  
                  # Log info
                  step += batch_size
                  progress_bar.update(batch_size)
                  progress_bar.set_postfix(epoch=epoch, NLL=loss_val)
                  tbx.add_scalar("train/NLL", loss_val, step)
                  tbx.add_scalar("train/LR", optimizer.param_groups[0]["lr"], step)
  
                  steps_till_eval -= batch_size
                  if steps_till_eval <= 0:
                      steps_till_eval = args.eval_steps
  
                      # Evaluate and save checkpoint
                      log.info(f"Evaluating at step {step}...")
                      ema.assign(model)
                      results, pred_dict = evaluate(
                          model,
                          dev_loader,
                          device,
                          args.dev_eval_file,
                          args.max_ans_len,
                          args.use_squad_v2,
                      )
                      saver.save(step, model, results[args.metric_name], device)
                      ema.resume(model)
  
                      # 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, step)
                      util.visualize(
                          tbx,
                          pred_dict=pred_dict,
                          eval_path=args.dev_eval_file,
                          step=step,
                          split="dev",
                          num_visuals=args.num_visuals,
                      )
  
  
  def evaluate(model, data_loader, device, eval_file, max_len, use_squad_v2):
      nll_meter = stats.AverageMeter()
  
      model.eval()
      pred_dict = {}
      with open(eval_file, "r") as fh:
          gold_dict = json_load(fh)
      with torch.no_grad(), tqdm(total=len(data_loader.dataset)) as progress_bar:
          for cw_idxs, cc_idxs, qw_idxs, qc_idxs, y1, y2, ids in data_loader:
              # Setup for forward
              cw_idxs = cw_idxs.to(device)
              qw_idxs = qw_idxs.to(device)
              batch_size = cw_idxs.size(0)
  
              # Forward
              log_p1, log_p2 = model(cw_idxs, qw_idxs)
              y1, y2 = y1.to(device), y2.to(device)
              loss = F.nll_loss(log_p1, y1) + F.nll_loss(log_p2, y2)
              nll_meter.update(loss.item(), batch_size)
  
              # Get F1 and EM scores
              p1, p2 = log_p1.exp(), log_p2.exp()
              starts, ends = util.discretize(p1, p2, max_len, use_squad_v2)
  
              # Log info
              progress_bar.update(batch_size)
              progress_bar.set_postfix(NLL=nll_meter.avg)
  
              preds, _ = util.convert_tokens(
                  gold_dict, ids.tolist(), starts.tolist(), ends.tolist(), use_squad_v2
              )
              pred_dict.update(preds)
  
      model.train()
  
      results = {"NLL": nll_meter.avg}
      results.update(eval.eval_dicts(gold_dict, pred_dict, use_squad_v2))
      return results, pred_dict
  
  
  def add_train_args(parser):
      """Add arguments needed in train.py."""
      add_train_test_args(parser)
  
      parser.add_argument(
          "--eval_steps",
          type=int,
          default=50000,
          help="Number of steps between successive evaluations.",
      )
      parser.add_argument("--lr", type=float, default=0.001, help="Learning rate.")
      parser.add_argument("--l2_wd", type=float, default=0, help="L2 weight decay.")
      parser.add_argument(
          "--num_epochs",
          type=int,
          default=30,
          help="Number of epochs for which to train. Negative means forever.",
      )
      parser.add_argument(
          "--drop_prob",
          type=float,
          default=0.2,
          help="Probability of zeroing an activation in dropout layers.",
      )
      parser.add_argument(
          "--metric_name",
          type=str,
          default="F1",
          choices=("NLL", "EM", "F1"),
          help="Name of dev metric to determine best checkpoint.",
      )
      parser.add_argument(
          "--max_checkpoints",
          type=int,
          default=5,
          help="Maximum number of checkpoints to keep on disk.",
      )
      parser.add_argument(
          "--max_grad_norm",
          type=float,
          default=5.0,
          help="Maximum gradient norm for gradient clipping.",
      )
      parser.add_argument(
          "--seed", type=int, default=224, help="Random seed for reproducibility."
      )
      parser.add_argument(
          "--ema_decay",
          type=float,
          default=0.999,
          help="Decay rate for exponential moving average of parameters.",
      )
  
  
  def test(args):
      # Set up logging
      log = util.get_logger(args.save_dir, args.name)
      log.info(f"Args: {dumps(vars(args), indent=4, sort_keys=True)}")
      device, gpu_ids = util.get_available_devices()
      args.batch_size *= max(1, len(gpu_ids))
  
      # Get embeddings
      log.info("Loading embeddings...")
      word_vectors = util.torch_from_json(args.word_emb_file)
  
      # Get model
      log.info("Building model...")
      model = BiDAF(
          word_vectors=word_vectors,
          hidden_size=args.hidden_size,
          use_glove=args.use_glove,
      )
      model = nn.DataParallel(model, gpu_ids)
      log.info(f"Loading checkpoint from {args.load_path}...")
      model = util.load_model(model, args.load_path, gpu_ids, return_step=False)
      model = model.to(device)
      model.eval()
  
      # Get data loader
      log.info("Building dataset...")
      record_file = vars(args)[f"{args.split}_record_file"]
      dataset = SQuAD(record_file, args.use_squad_v2)
      data_loader = data.DataLoader(
          dataset,
          batch_size=args.batch_size,
          shuffle=False,
          num_workers=args.num_workers,
          collate_fn=collate_fn,
      )
  
      # Evaluate
      log.info(f"Evaluating on {args.split} split...")
      nll_meter = stats.AverageMeter()
      pred_dict = {}  # Predictions for TensorBoard
      sub_dict = {}  # Predictions for submission
      eval_file = vars(args)[f"{args.split}_eval_file"]
      with open(eval_file, "r") as fh:
          gold_dict = json_load(fh)
      with torch.no_grad(), tqdm(total=len(dataset)) as progress_bar:
          for cw_idxs, cc_idxs, qw_idxs, qc_idxs, y1, y2, ids in data_loader:
              # Setup for forward
              cw_idxs = cw_idxs.to(device)
              qw_idxs = qw_idxs.to(device)
              batch_size = cw_idxs.size(0)
  
              # Forward
              log_p1, log_p2 = model(cw_idxs, qw_idxs)
              y1, y2 = y1.to(device), y2.to(device)
              loss = F.nll_loss(log_p1, y1) + F.nll_loss(log_p2, y2)
              nll_meter.update(loss.item(), batch_size)
  
              # Get F1 and EM scores
              p1, p2 = log_p1.exp(), log_p2.exp()
              starts, ends = util.discretize(p1, p2, args.max_ans_len, args.use_squad_v2)
  
              # Log info
              progress_bar.update(batch_size)
              if args.split != "test":
                  # No labels for the test set, so NLL would be invalid
                  progress_bar.set_postfix(NLL=nll_meter.avg)
  
              idx2pred, uuid2pred = util.convert_tokens(
                  gold_dict,
                  ids.tolist(),
                  starts.tolist(),
                  ends.tolist(),
                  args.use_squad_v2,
              )
              pred_dict.update(idx2pred)
              sub_dict.update(uuid2pred)
  
      # Log results (except for test set, since it does not come with labels)
      if args.split != "test":
          results = {"NLL": nll_meter.avg}
          results.update(eval.eval_dicts(gold_dict, pred_dict, args.use_squad_v2))
  
          # Log to console
          results_str = ", ".join(f"{k}: {v:05.2f}" for k, v in results.items())
          log.info(f"{args.split.title()} {results_str}")
  
          # Log to TensorBoard
          tbx = SummaryWriter(args.save_dir)
          util.visualize(
              tbx,
              pred_dict=pred_dict,
              eval_path=eval_file,
              step=0,
              split=args.split,
              num_visuals=args.num_visuals,
          )
  
      # Write submission file
      if args.split == "dev":
          sub_path = join(args.save_dir, "val" + "_" + args.sub_file)
      else:
          sub_path = join(args.save_dir, args.split + "_" + args.sub_file)
      log.info(f"Writing submission file to {sub_path}...")
      eval.write_submission(sub_path, sub_dict)
  
  
  def add_test_args(parser):
      """Get arguments needed in test.py."""
      add_train_test_args(parser)
  
  
  def add_train_test_args(parser):
      """Add arguments common to train.py and test.py"""
      parser.add_argument(
          "--max_ans_len",
          type=int,
          default=15,
          help="Maximum length of a predicted answer.",
      )
      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=64,
          help="Batch size per GPU. Scales automatically when \
                                multiple GPUs are available.",
      )
      parser.add_argument(
          "--use_squad_v2",
          type=lambda s: s.lower().startswith("t"),
          default=True,
          help="Whether to use SQuAD 2.0 (unanswerable) questions.",
      )
      parser.add_argument(
          "--hidden_size",
          type=int,
          default=100,
          help="Number of features in encoder hidden layers.",
      )
      parser.add_argument(
          "--use_glove",
          type=lambda s: s.lower().startswith("t"),
          default=True,
          help="Whether to use pretrained glove vectors.",
      )
      parser.add_argument(
          "--num_visuals",
          type=int,
          default=10,
          help="Number of examples to visualize in TensorBoard.",
      )
  