run_fgc_baseline.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "SIq-UqB4Zhtp"
   },
   "source": [
    "# Requirements\n",
    "1. Before running this notebook, make sure you have \n",
    "```\n",
    "DRCD_plus_FGC_training.json\n",
    "FGC_all_mocks_DRCD_format.json\n",
    "FGC_final_DRCD_format.json\n",
    "```\n",
    "in the current directory.\n",
    "2. When **not running on Google Colab**, comment out `pip install transformers`\n",
    "3. If you have `tensorboard` jupyter notebook plugin, use `%load_ext tensorboard`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ixchtS-EFx6Y"
   },
   "outputs": [],
   "source": [
    "# !pip install transformers\n",
    "%load_ext tensorboard\n",
    "import glob\n",
    "import logging\n",
    "import os\n",
    "import random\n",
    "import argparse\n",
    "import json\n",
    "import timeit\n",
    "import numpy as np\n",
    "import torch\n",
    "from torch.utils.data import DataLoader, RandomSampler, SequentialSampler\n",
    "from torch.utils.data.distributed import DistributedSampler\n",
    "from tqdm.notebook import tqdm, trange\n",
    "\n",
    "from transformers import (\n",
    "  BertTokenizer,\n",
    "  AlbertForQuestionAnswering,\n",
    "  AutoConfig,\n",
    "  AdamW,\n",
    "  get_linear_schedule_with_warmup,\n",
    "  squad_convert_examples_to_features,\n",
    ")\n",
    "\n",
    "from transformers.data.metrics.squad_metrics import (\n",
    "  compute_predictions_log_probs,\n",
    "  compute_predictions_logits,\n",
    "  squad_evaluate,\n",
    ")\n",
    "\n",
    "from transformers.data.processors.squad import SquadResult, SquadProcessor\n",
    "try:\n",
    "  from torch.utils.tensorboard import SummaryWriter\n",
    "except ImportError:\n",
    "  from tensorboardX import SummaryWriter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "vobmUQtA9vkW",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "pretrained = 'voidful/albert_chinese_base'\n",
    "tokenizer = BertTokenizer.from_pretrained(pretrained)\n",
    "model = AlbertForQuestionAnswering.from_pretrained(pretrained)\n",
    "config = AutoConfig.from_pretrained(pretrained)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ve4YD7sLGJQc"
   },
   "outputs": [],
   "source": [
    "class Args(object):\n",
    "  model_type = \"ALBERT\"\n",
    "  model_name_or_path = \".\"\n",
    "  output_dir = \"output_training\"\n",
    "  data_dir = \".\"\n",
    "  null_score_diff_threshold = 0.0 # If null_score - best_non_null is greater than the threshold predict null\n",
    "  version_2_with_negative = False\n",
    "  max_seq_length = 512\n",
    "  doc_stride = 128\n",
    "  max_query_length = 128\n",
    "  do_train = True\n",
    "  do_eval = True\n",
    "  evaluate_during_training = True\n",
    "  eval_all_checkpoints = True\n",
    "  do_lower_case = True\n",
    "  per_gpu_train_batch_size = 8\n",
    "  per_gpu_eval_batch_size = 4\n",
    "  no_cuda = False\n",
    "\n",
    "  learning_rate = 3e-5\n",
    "  gradient_accumulation_steps = 8\n",
    "  weight_decay = 0.0\n",
    "  adam_epsilon = 1e-8\n",
    "  max_grad_norm = 1.\n",
    "  num_train_epochs = 3.\n",
    "  max_steps = -1 # If > 0: set total number of training steps to perform. Override num_train_epochs.\n",
    "  warmup_steps = 200\n",
    "  n_best_size = 5\n",
    "  max_answer_length = 200\n",
    "  logging_steps = 1000\n",
    "  save_steps = 1000\n",
    "\n",
    "  seed = 1337\n",
    "  local_rank = -1\n",
    "  threads = 4\n",
    "  n_gpu = 1\n",
    "  overwrite_cache = False\n",
    "  overwrite_output_dir = True\n",
    "  verbose_logging = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "5RURpZ9kmCT7"
   },
   "outputs": [],
   "source": [
    "class DRCDProcessor(SquadProcessor):\n",
    "  train_file = \"./DRCD_plus_FGC_training.json\"\n",
    "  dev_file = \"./FGC_all_mocks_DRCD_format.json\"\n",
    "\n",
    "def set_seed(args):\n",
    "  random.seed(args.seed)\n",
    "  np.random.seed(args.seed)\n",
    "  torch.manual_seed(args.seed)\n",
    "  if args.n_gpu > 0:\n",
    "    torch.cuda.manual_seed_all(args.seed)\n",
    "\n",
    "def to_list(tensor):\n",
    "  return tensor.detach().cpu().tolist()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "wwFUL-xCmPoS"
   },
   "outputs": [],
   "source": [
    "def train(args, train_dataset, model, tokenizer, config):\n",
    "  \"\"\" Train the model \"\"\"\n",
    "  tb_writer = SummaryWriter()\n",
    "  logger = logging.getLogger(__name__)\n",
    "\n",
    "  args.train_batch_size = args.per_gpu_train_batch_size * max(1, args.n_gpu)\n",
    "  train_sampler = RandomSampler(train_dataset) if args.local_rank == -1 else DistributedSampler(train_dataset)\n",
    "  train_dataloader = DataLoader(train_dataset, sampler=train_sampler, batch_size=args.train_batch_size)\n",
    "\n",
    "  if args.max_steps > 0:\n",
    "    t_total = args.max_steps\n",
    "    args.num_train_epochs = args.max_steps // (len(train_dataloader) // args.gradient_accumulation_steps) + 1\n",
    "  else:\n",
    "    t_total = len(train_dataloader) // args.gradient_accumulation_steps * args.num_train_epochs\n",
    "\n",
    "  # Prepare optimizer and schedule (linear warmup and decay)\n",
    "  no_decay = [\"bias\", \"LayerNorm.weight\"]\n",
    "  optimizer_grouped_parameters = [\n",
    "    {\n",
    "      \"params\": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],\n",
    "      \"weight_decay\": args.weight_decay,\n",
    "    },\n",
    "    {\"params\": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)], \"weight_decay\": 0.0},\n",
    "  ]\n",
    "  optimizer = AdamW(optimizer_grouped_parameters, lr=args.learning_rate, eps=args.adam_epsilon)\n",
    "  scheduler = get_linear_schedule_with_warmup(\n",
    "    optimizer, num_warmup_steps=args.warmup_steps, num_training_steps=t_total\n",
    "  )\n",
    "\n",
    "  # Train!\n",
    "  logger.info(\"***** Running training *****\")\n",
    "  logger.info(\"  Num examples = %d\", len(train_dataset))\n",
    "  logger.info(\"  Num Epochs = %d\", args.num_train_epochs)\n",
    "  logger.info(\"  Instantaneous batch size per GPU = %d\", args.per_gpu_train_batch_size)\n",
    "  logger.info(\n",
    "    \"  Total train batch size (w. parallel, distributed & accumulation) = %d\",\n",
    "    args.train_batch_size\n",
    "    * args.gradient_accumulation_steps\n",
    "    * (torch.distributed.get_world_size() if args.local_rank != -1 else 1),\n",
    "  )\n",
    "  logger.info(\"  Gradient Accumulation steps = %d\", args.gradient_accumulation_steps)\n",
    "  logger.info(\"  Total optimization steps = %d\", t_total)\n",
    "\n",
    "  global_step = 1\n",
    "  epochs_trained = 0\n",
    "  steps_trained_in_current_epoch = 0\n",
    "\n",
    "  logger.info(\"  Starting fine-tuning.\")\n",
    "\n",
    "  tr_loss, logging_loss = 0.0, 0.0\n",
    "\n",
    "  model.zero_grad()\n",
    "  train_iterator = trange(\n",
    "    epochs_trained, int(args.num_train_epochs), desc=\"Epoch\", disable=args.local_rank not in [-1, 0], position=0, leave=False\n",
    "  )\n",
    "  # Added here for reproductibility\n",
    "  set_seed(args)\n",
    "\n",
    "  for _ in train_iterator:\n",
    "    epoch_iterator = tqdm(train_dataloader, desc=\"Iteration\", disable=args.local_rank not in [-1, 0], position=1, leave=False)\n",
    "    for step, batch in enumerate(epoch_iterator):\n",
    "\n",
    "      # Skip past any already trained steps if resuming training\n",
    "      if steps_trained_in_current_epoch > 0:\n",
    "        steps_trained_in_current_epoch -= 1\n",
    "        continue\n",
    "\n",
    "      model.train()\n",
    "      batch = tuple(t.to(args.device) for t in batch)\n",
    "\n",
    "      inputs = {\n",
    "        \"input_ids\": batch[0],\n",
    "        \"attention_mask\": batch[1],\n",
    "        \"token_type_ids\": batch[2],\n",
    "        \"start_positions\": batch[3],\n",
    "        \"end_positions\": batch[4],\n",
    "      }\n",
    "\n",
    "      if args.model_type in [\"xlm\", \"roberta\", \"distilbert\", \"camembert\"]:\n",
    "        del inputs[\"token_type_ids\"]\n",
    "\n",
    "      outputs = model(**inputs)\n",
    "      # model outputs are always tuple in transformers (see doc)\n",
    "      loss = outputs[0]\n",
    "\n",
    "      if args.n_gpu > 1:\n",
    "        loss = loss.mean()  # mean() to average on multi-gpu parallel (not distributed) training\n",
    "      if args.gradient_accumulation_steps > 1:\n",
    "        loss = loss / args.gradient_accumulation_steps\n",
    "\n",
    "\n",
    "      loss.backward()\n",
    "\n",
    "      tr_loss += loss.item()\n",
    "      if (step + 1) % args.gradient_accumulation_steps == 0:\n",
    "        torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)\n",
    "\n",
    "        optimizer.step()\n",
    "        scheduler.step()  # Update learning rate schedule\n",
    "        model.zero_grad()\n",
    "        global_step += 1\n",
    "\n",
    "        # Log metrics\n",
    "        if args.local_rank in [-1, 0] and args.logging_steps > 0 and global_step % args.logging_steps == 0:\n",
    "          # Only evaluate when single GPU otherwise metrics may not average well\n",
    "          if args.local_rank == -1 and args.evaluate_during_training:\n",
    "            results = evaluate(args, model, tokenizer, prefix=\"step_{}\".format(global_step))\n",
    "            for key, value in results.items():\n",
    "              tb_writer.add_scalar(\"eval_{}\".format(key), value, global_step)\n",
    "          tb_writer.add_scalar(\"lr\", scheduler.get_lr()[0], global_step)\n",
    "          tb_writer.add_scalar(\"loss\", (tr_loss - logging_loss) / args.logging_steps, global_step)\n",
    "          logging_loss = tr_loss\n",
    "\n",
    "        # Save model checkpoint\n",
    "        if args.local_rank in [-1, 0] and args.save_steps > 0 and global_step % args.save_steps == 0:\n",
    "          output_dir = os.path.join(args.output_dir, \"checkpoint-{}\".format(global_step))\n",
    "          # Take care of distributed/parallel training\n",
    "          model_to_save = model.module if hasattr(model, \"module\") else model\n",
    "          model_to_save.save_pretrained(output_dir)\n",
    "          tokenizer.save_pretrained(output_dir)\n",
    "\n",
    "          torch.save(args, os.path.join(output_dir, \"training_args.bin\"))\n",
    "          logger.info(\"Saving model checkpoint to %s\", output_dir)\n",
    "\n",
    "          torch.save(optimizer.state_dict(), os.path.join(output_dir, \"optimizer.pt\"))\n",
    "          torch.save(scheduler.state_dict(), os.path.join(output_dir, \"scheduler.pt\"))\n",
    "          logger.info(\"Saving optimizer and scheduler states to %s\", output_dir)\n",
    "\n",
    "      if args.max_steps > 0 and global_step > args.max_steps:\n",
    "        epoch_iterator.close()\n",
    "        break\n",
    "    if args.max_steps > 0 and global_step > args.max_steps:\n",
    "      train_iterator.close()\n",
    "      break\n",
    "\n",
    "  if args.local_rank in [-1, 0]:\n",
    "    tb_writer.close()\n",
    "\n",
    "  return global_step, tr_loss / global_step\n",
    "\n",
    "\n",
    "def evaluate(args, model, tokenizer, prefix=\"\"):\n",
    "  logger = logging.getLogger(__name__)\n",
    "  dataset, examples, features = load_and_cache_examples(args, tokenizer, evaluate=True, output_examples=True)\n",
    "\n",
    "  if not os.path.exists(args.output_dir) and args.local_rank in [-1, 0]:\n",
    "    os.makedirs(args.output_dir)\n",
    "\n",
    "  args.eval_batch_size = args.per_gpu_eval_batch_size * max(1, args.n_gpu)\n",
    "\n",
    "  # Note that DistributedSampler samples randomly\n",
    "  eval_sampler = SequentialSampler(dataset)\n",
    "  eval_dataloader = DataLoader(dataset, sampler=eval_sampler, batch_size=args.eval_batch_size)\n",
    "\n",
    "  # multi-gpu evaluate\n",
    "  if args.n_gpu > 1 and not isinstance(model, torch.nn.DataParallel):\n",
    "    model = torch.nn.DataParallel(model)\n",
    "\n",
    "  # Eval!\n",
    "  logger.info(\"***** Running evaluation {} *****\".format(prefix))\n",
    "  logger.info(\"  Num examples = %d\", len(dataset))\n",
    "  logger.info(\"  Batch size = %d\", args.eval_batch_size)\n",
    "\n",
    "  all_results = []\n",
    "  start_time = timeit.default_timer()\n",
    "\n",
    "  for batch in tqdm(eval_dataloader, desc=\"Evaluating\", position=0, leave=False):\n",
    "    model.eval()\n",
    "    batch = tuple(t.to(args.device) for t in batch)\n",
    "\n",
    "    with torch.no_grad():\n",
    "      inputs = {\n",
    "        \"input_ids\": batch[0],\n",
    "        \"attention_mask\": batch[1],\n",
    "        \"token_type_ids\": batch[2],\n",
    "      }\n",
    "\n",
    "      if args.model_type in [\"xlm\", \"roberta\", \"distilbert\", \"camembert\"]:\n",
    "        del inputs[\"token_type_ids\"]\n",
    "\n",
    "      feature_indices = batch[3]\n",
    "\n",
    "      outputs = model(**inputs)\n",
    "\n",
    "    for i, feature_index in enumerate(feature_indices):\n",
    "      eval_feature = features[feature_index.item()]\n",
    "      unique_id = int(eval_feature.unique_id)\n",
    "\n",
    "      output = [to_list(output[i]) for output in outputs]\n",
    "\n",
    "      # Some models (XLNet, XLM) use 5 arguments for their predictions, while the other \"simpler\"\n",
    "      # models only use two.\n",
    "      if len(output) >= 5:\n",
    "        start_logits = output[0]\n",
    "        start_top_index = output[1]\n",
    "        end_logits = output[2]\n",
    "        end_top_index = output[3]\n",
    "        cls_logits = output[4]\n",
    "\n",
    "        result = SquadResult(\n",
    "          unique_id,\n",
    "          start_logits,\n",
    "          end_logits,\n",
    "          start_top_index=start_top_index,\n",
    "          end_top_index=end_top_index,\n",
    "          cls_logits=cls_logits,\n",
    "        )\n",
    "\n",
    "      else:\n",
    "        start_logits, end_logits = output\n",
    "        result = SquadResult(unique_id, start_logits, end_logits)\n",
    "\n",
    "      all_results.append(result)\n",
    "\n",
    "  evalTime = timeit.default_timer() - start_time\n",
    "  logger.info(\"  Evaluation done in total %f secs (%f sec per example)\", evalTime, evalTime / len(dataset))\n",
    "\n",
    "  # Compute predictions\n",
    "  output_prediction_file = os.path.join(args.output_dir, \"predictions_{}.json\".format(prefix))\n",
    "  output_nbest_file = os.path.join(args.output_dir, \"nbest_predictions_{}.json\".format(prefix))\n",
    "  output_null_log_odds_file = None\n",
    "\n",
    "\n",
    "  predictions = compute_predictions_logits(\n",
    "    examples,\n",
    "    features,\n",
    "    all_results,\n",
    "    args.n_best_size,\n",
    "    args.max_answer_length,\n",
    "    args.do_lower_case,\n",
    "    output_prediction_file,\n",
    "    output_nbest_file,\n",
    "    output_null_log_odds_file,\n",
    "    args.verbose_logging,\n",
    "    args.version_2_with_negative,\n",
    "    args.null_score_diff_threshold,\n",
    "    tokenizer,\n",
    "  )\n",
    "\n",
    "  # Compute the F1 and exact scores.\n",
    "  results = squad_evaluate(examples, predictions)\n",
    "  return results\n",
    "\n",
    "\n",
    "def load_and_cache_examples(args, tokenizer, evaluate=False, output_examples=False, processor=DRCDProcessor()):\n",
    "  logger = logging.getLogger(__name__)\n",
    "  # Load data features from cache or dataset file\n",
    "  input_dir = args.data_dir if args.data_dir else \".\"\n",
    "  cached_features_file = os.path.join(\n",
    "    input_dir,\n",
    "    \"cached_{}_{}_{}\".format(\n",
    "      \"dev\" if evaluate else \"train\",\n",
    "      list(filter(None, (processor.dev_file if evaluate else processor.train_file).split('/'))).pop(),\n",
    "      str(args.max_seq_length),\n",
    "    ),\n",
    "  )\n",
    "\n",
    "  # Init features and dataset from cache if it exists\n",
    "  if os.path.exists(cached_features_file) and not args.overwrite_cache:\n",
    "    logger.info(\"Loading features from cached file %s\", cached_features_file)\n",
    "    features_and_dataset = torch.load(cached_features_file)\n",
    "    features, dataset, examples = (\n",
    "      features_and_dataset[\"features\"],\n",
    "      features_and_dataset[\"dataset\"],\n",
    "      features_and_dataset[\"examples\"],\n",
    "    )\n",
    "  else:\n",
    "    logger.info(\"Creating features from dataset file at %s\", input_dir)\n",
    "\n",
    "    if evaluate:\n",
    "      examples = processor.get_dev_examples(args.data_dir)\n",
    "    else:\n",
    "      examples = processor.get_train_examples(args.data_dir)\n",
    "\n",
    "    features, dataset = squad_convert_examples_to_features(\n",
    "      examples=examples,\n",
    "      tokenizer=tokenizer,\n",
    "      max_seq_length=args.max_seq_length,\n",
    "      doc_stride=args.doc_stride,\n",
    "      max_query_length=args.max_query_length,\n",
    "      is_training=not evaluate,\n",
    "      return_dataset=\"pt\",\n",
    "      threads=args.threads,\n",
    "    )\n",
    "\n",
    "    if args.local_rank in [-1, 0]:\n",
    "      logger.info(\"Saving features into cached file %s\", cached_features_file)\n",
    "      torch.save({\"features\": features, \"dataset\": dataset, \"examples\": examples}, cached_features_file)\n",
    "\n",
    "  if args.local_rank == 0 and not evaluate:\n",
    "    # Make sure only the first process in distributed training process the dataset, and the others will use the cache\n",
    "    torch.distributed.barrier()\n",
    "\n",
    "  if output_examples:\n",
    "    return dataset, examples, features\n",
    "  return dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "9inQ1ZPQ9MNe"
   },
   "outputs": [],
   "source": [
    "def main(args, model, tokenizer, config):\n",
    "  logger = logging.getLogger(__name__)\n",
    "  if args.doc_stride >= args.max_seq_length - args.max_query_length:\n",
    "    logger.warning(\n",
    "      \"WARNING - You've set a doc stride which may be superior to the document length in some \"\n",
    "      \"examples. This could result in errors when building features from the examples. Please reduce the doc \"\n",
    "      \"stride or increase the maximum length to ensure the features are correctly built.\"\n",
    "    )\n",
    "\n",
    "  if (\n",
    "    os.path.exists(args.output_dir)\n",
    "    and os.listdir(args.output_dir)\n",
    "    and args.do_train\n",
    "    and not args.overwrite_output_dir\n",
    "  ):\n",
    "    raise ValueError(\n",
    "      \"Output directory ({}) already exists and is not empty. Use --overwrite_output_dir to overcome.\".format(\n",
    "        args.output_dir\n",
    "      )\n",
    "    )\n",
    "\n",
    "  # Setup CUDA, GPU & distributed training\n",
    "  if args.local_rank == -1 or args.no_cuda:\n",
    "    device = torch.device(\"cuda\" if torch.cuda.is_available() and not args.no_cuda else \"cpu\")\n",
    "    args.n_gpu = 0 if args.no_cuda else torch.cuda.device_count()\n",
    "  else:  # Initializes the distributed backend which will take care of sychronizing nodes/GPUs\n",
    "    torch.cuda.set_device(args.local_rank)\n",
    "    device = torch.device(\"cuda\", args.local_rank)\n",
    "    args.n_gpu = 1\n",
    "  args.device = device\n",
    "\n",
    "  # Setup logging\n",
    "  logging.basicConfig(\n",
    "    format=\"%(asctime)s - %(levelname)s - %(name)s -   %(message)s\",\n",
    "    datefmt=\"%m/%d/%Y %H:%M:%S\",\n",
    "    level=logging.INFO if args.local_rank in [-1, 0] else logging.WARN,\n",
    "  )\n",
    "  logger.warning(\n",
    "    \"Process rank: %s, device: %s, n_gpu: %s, distributed training: %s\",\n",
    "    args.local_rank,\n",
    "    device,\n",
    "    args.n_gpu,\n",
    "    bool(args.local_rank != -1),\n",
    "  )\n",
    "\n",
    "  # Set seed\n",
    "  set_seed(args)\n",
    "  args.model_type = args.model_type.lower()\n",
    "\n",
    "  model.to(args.device)\n",
    "\n",
    "  # Training\n",
    "  if args.do_train:\n",
    "    train_dataset = load_and_cache_examples(args, tokenizer, evaluate=False, output_examples=False)\n",
    "    global_step, tr_loss = train(args, train_dataset, model, tokenizer, config)\n",
    "    logger.info(\" global_step = %s, average loss = %s\", global_step, tr_loss)\n",
    "\n",
    "  # Save the trained model and the tokenizer\n",
    "  if args.do_train and (args.local_rank == -1 or torch.distributed.get_rank() == 0):\n",
    "    logger.info(\"Saving model checkpoint to %s\", args.output_dir)\n",
    "    # Save a trained model, configuration and tokenizer using `save_pretrained()`.\n",
    "    # They can then be reloaded using `from_pretrained()`\n",
    "    # Take care of distributed/parallel training\n",
    "    model_to_save = model.module if hasattr(model, \"module\") else model\n",
    "    model_to_save.save_pretrained(args.output_dir)\n",
    "    tokenizer.save_pretrained(args.output_dir)\n",
    "\n",
    "    # Good practice: save your training arguments together with the trained model\n",
    "    torch.save(args, os.path.join(args.output_dir, \"training_args.bin\"))\n",
    "\n",
    "  # Evaluation - we can ask to evaluate all the checkpoints (sub-directories) in a directory\n",
    "  results = {}\n",
    "  if args.do_eval and args.local_rank in [-1, 0]:\n",
    "    if args.do_train:\n",
    "      logger.info(\"Loading checkpoints saved during training for evaluation\")\n",
    "      checkpoints = [args.output_dir]\n",
    "      if args.eval_all_checkpoints:\n",
    "        checkpoints = list(\n",
    "          os.path.dirname(c)\n",
    "          for c in sorted(glob.glob(args.output_dir + \"/**/\" + \"pytorch_model.bin\", recursive=True))\n",
    "        )\n",
    "        logging.getLogger(\"transformers.modeling_utils\").setLevel(logging.WARN)  # Reduce model loading logs\n",
    "    else:\n",
    "      logger.info(\"Loading checkpoint %s for evaluation\", args.model_name_or_path)\n",
    "      checkpoints = [args.model_name_or_path]\n",
    "\n",
    "    logger.info(\"Evaluate the following checkpoints: %s\", checkpoints)\n",
    "\n",
    "    for checkpoint in checkpoints:\n",
    "      # Reload the model\n",
    "      global_step = checkpoint.split(\"-\")[-1] if len(checkpoints) > 1 else \"\"\n",
    "      model = AlbertForQuestionAnswering.from_pretrained(checkpoint)\n",
    "      model.to(args.device)\n",
    "\n",
    "      # Evaluate\n",
    "      result = evaluate(args, model, tokenizer, prefix=global_step)\n",
    "\n",
    "      result = dict((k + (\"_{}\".format(global_step) if global_step else \"\"), v) for k, v in result.items())\n",
    "      results.update(result)\n",
    "\n",
    "  logger.info(\"Results: {}\".format(results))\n",
    "\n",
    "  return results\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "KFgNfCsFIGOP",
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "main(Args(), model, tokenizer, config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "class TestSetProcessor(DRCDProcessor):\n",
    "  dev_file = \"./FGC_final_DRCD_format.json\"\n",
    "\n",
    "test_args = Args()\n",
    "test_args.do_train=False\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() and not test_args.no_cuda else \"cpu\")\n",
    "test_args.n_gpu = 0 if test_args.no_cuda else torch.cuda.device_count()\n",
    "test_args.device = device\n",
    "\n",
    "def evaluate_test_set(args, model, tokenizer, prefix=\"\"):\n",
    "  logger = logging.getLogger(__name__)\n",
    "  dataset, examples, features = load_and_cache_examples(args, tokenizer, evaluate=True, output_examples=True, \\\n",
    "                                                        processor=TestSetProcessor())\n",
    "\n",
    "  if not os.path.exists(args.output_dir) and args.local_rank in [-1, 0]:\n",
    "    os.makedirs(args.output_dir)\n",
    "\n",
    "  args.eval_batch_size = args.per_gpu_eval_batch_size * max(1, args.n_gpu)\n",
    "\n",
    "  # Note that DistributedSampler samples randomly\n",
    "  eval_sampler = SequentialSampler(dataset)\n",
    "  eval_dataloader = DataLoader(dataset, sampler=eval_sampler, batch_size=args.eval_batch_size)\n",
    "\n",
    "  # multi-gpu evaluate\n",
    "  if args.n_gpu > 1 and not isinstance(model, torch.nn.DataParallel):\n",
    "    model = torch.nn.DataParallel(model)\n",
    "\n",
    "  # Eval!\n",
    "  logger.info(\"***** Running evaluation {} *****\".format(prefix))\n",
    "  logger.info(\"  Num examples = %d\", len(dataset))\n",
    "  logger.info(\"  Batch size = %d\", args.eval_batch_size)\n",
    "\n",
    "  all_results = []\n",
    "  start_time = timeit.default_timer()\n",
    "\n",
    "  for batch in tqdm(eval_dataloader, desc=\"Evaluating\", position=0, leave=False):\n",
    "    model.eval()\n",
    "    batch = tuple(t.to(args.device) for t in batch)\n",
    "\n",
    "    with torch.no_grad():\n",
    "      inputs = {\n",
    "        \"input_ids\": batch[0],\n",
    "        \"attention_mask\": batch[1],\n",
    "        \"token_type_ids\": batch[2],\n",
    "      }\n",
    "\n",
    "      if args.model_type in [\"xlm\", \"roberta\", \"distilbert\", \"camembert\"]:\n",
    "        del inputs[\"token_type_ids\"]\n",
    "\n",
    "      feature_indices = batch[3]\n",
    "\n",
    "      outputs = model(**inputs)\n",
    "\n",
    "    for i, feature_index in enumerate(feature_indices):\n",
    "      eval_feature = features[feature_index.item()]\n",
    "      unique_id = int(eval_feature.unique_id)\n",
    "\n",
    "      output = [to_list(output[i]) for output in outputs]\n",
    "\n",
    "      # Some models (XLNet, XLM) use 5 arguments for their predictions, while the other \"simpler\"\n",
    "      # models only use two.\n",
    "      if len(output) >= 5:\n",
    "        start_logits = output[0]\n",
    "        start_top_index = output[1]\n",
    "        end_logits = output[2]\n",
    "        end_top_index = output[3]\n",
    "        cls_logits = output[4]\n",
    "\n",
    "        result = SquadResult(\n",
    "          unique_id,\n",
    "          start_logits,\n",
    "          end_logits,\n",
    "          start_top_index=start_top_index,\n",
    "          end_top_index=end_top_index,\n",
    "          cls_logits=cls_logits,\n",
    "        )\n",
    "\n",
    "      else:\n",
    "        start_logits, end_logits = output\n",
    "        result = SquadResult(unique_id, start_logits, end_logits)\n",
    "\n",
    "      all_results.append(result)\n",
    "\n",
    "  evalTime = timeit.default_timer() - start_time\n",
    "  logger.info(\"  Evaluation done in total %f secs (%f sec per example)\", evalTime, evalTime / len(dataset))\n",
    "\n",
    "  # Compute predictions\n",
    "  output_prediction_file = os.path.join(args.output_dir, \"predictions_{}.json\".format(prefix))\n",
    "  output_nbest_file = os.path.join(args.output_dir, \"nbest_predictions_{}.json\".format(prefix))\n",
    "  output_null_log_odds_file = None\n",
    "\n",
    "\n",
    "  predictions = compute_predictions_logits(\n",
    "    examples,\n",
    "    features,\n",
    "    all_results,\n",
    "    args.n_best_size,\n",
    "    args.max_answer_length,\n",
    "    args.do_lower_case,\n",
    "    output_prediction_file,\n",
    "    output_nbest_file,\n",
    "    output_null_log_odds_file,\n",
    "    args.verbose_logging,\n",
    "    args.version_2_with_negative,\n",
    "    args.null_score_diff_threshold,\n",
    "    tokenizer,\n",
    "  )\n",
    "\n",
    "  # Compute the F1 and exact scores.\n",
    "  results = squad_evaluate(examples, predictions)\n",
    "  return results\n",
    "\n",
    "test_results = {}\n",
    "checkpoints = list(\n",
    "          os.path.dirname(c)\n",
    "          for c in sorted(glob.glob(test_args.output_dir + \"/**/\" + \"pytorch_model.bin\", recursive=True))\n",
    "        )\n",
    "logging.getLogger(\"transformers.modeling_utils\").setLevel(logging.WARN)  # Reduce model loading logs\n",
    "logging.getLogger(\"transformers.configuration_utils\").setLevel(logging.WARN)  # Reduce model loading logs\n",
    "logger = logging.getLogger(__name__)\n",
    "logger.info(\"Evaluate the following checkpoints: %s\", checkpoints)\n",
    "\n",
    "for checkpoint in checkpoints:\n",
    "  # Reload the model\n",
    "  global_step = checkpoint.split(\"-\")[-1] if len(checkpoints) > 1 else \"\"\n",
    "  model = AlbertForQuestionAnswering.from_pretrained(checkpoint)\n",
    "  model.to(test_args.device)\n",
    "\n",
    "  # Evaluate\n",
    "  result = evaluate_test_set(test_args, model, tokenizer, prefix=\"test_set_\" + str(global_step))\n",
    "\n",
    "  result = dict((k + (\"_{}\".format(global_step) if global_step else \"\"), v) for k, v in result.items())\n",
    "  test_results.update(result)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "step 1000 EM 20.0\n",
      "step training EM 20.0\n"
     ]
    }
   ],
   "source": [
    "for k in test_results.keys():\n",
    "    if k.startswith(\"exact_\"):\n",
    "        print(\"step {} EM {}\".format(k.split('_')[-1], test_results[k]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "run_fgc_baseline.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}