run_model.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import absolute_import, division, print_function\n",
    "\n",
    "import glob\n",
    "import logging\n",
    "import os\n",
    "import random\n",
    "import json\n",
    "import math\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "from torch.utils.data import (DataLoader, RandomSampler, SequentialSampler,\n",
    "                              TensorDataset)\n",
    "import random\n",
    "from torch.utils.data.distributed import DistributedSampler\n",
    "from tqdm import tqdm_notebook, trange\n",
    "from tensorboardX import SummaryWriter\n",
    "\n",
    "\n",
    "from pytorch_transformers import (WEIGHTS_NAME, BertConfig, BertForSequenceClassification, BertTokenizer,\n",
    "                                  XLMConfig, XLMForSequenceClassification, XLMTokenizer, \n",
    "                                  XLNetConfig, XLNetForSequenceClassification, XLNetTokenizer,\n",
    "                                  RobertaConfig, RobertaForSequenceClassification, RobertaTokenizer)\n",
    "\n",
    "from pytorch_transformers import AdamW, WarmupLinearSchedule\n",
    "\n",
    "from utils import (convert_examples_to_features,\n",
    "                        output_modes, processors)\n",
    "\n",
    "logging.basicConfig(level=logging.INFO)\n",
    "logger = logging.getLogger(__name__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "args = {\n",
    "    'data_dir': 'data/',\n",
    "    'model_type':  'xlnet',\n",
    "    'model_name': 'xlnet-base-cased',\n",
    "    'task_name': 'binary',\n",
    "    'output_dir': 'outputs/',\n",
    "    'cache_dir': 'cache/',\n",
    "    'do_train': True,\n",
    "    'do_eval': True,\n",
    "    'fp16': True,\n",
    "    'fp16_opt_level': 'O1',\n",
    "    'max_seq_length': 128,\n",
    "    'output_mode': 'classification',\n",
    "    'train_batch_size': 8,\n",
    "    'eval_batch_size': 8,\n",
    "\n",
    "    'gradient_accumulation_steps': 1,\n",
    "    'num_train_epochs': 1,\n",
    "    'weight_decay': 0,\n",
    "    'learning_rate': 4e-5,\n",
    "    'adam_epsilon': 1e-8,\n",
    "    'warmup_ratio': 0.06,\n",
    "    'warmup_steps': 0,\n",
    "    'max_grad_norm': 1.0,\n",
    "\n",
    "    'logging_steps': 50,\n",
    "    'evaluate_during_training': False,\n",
    "    'save_steps': 2000,\n",
    "    'eval_all_checkpoints': True,\n",
    "\n",
    "    'overwrite_output_dir': False,\n",
    "    'reprocess_input_data': True,\n",
    "    'notes': 'Using Yelp Reviews dataset'\n",
    "}\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "args"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('args.json', 'w') as f:\n",
    "    json.dump(args, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if os.path.exists(args['output_dir']) and os.listdir(args['output_dir']) and args['do_train'] and not args['overwrite_output_dir']:\n",
    "    raise ValueError(\"Output directory ({}) already exists and is not empty. Use --overwrite_output_dir to overcome.\".format(args['output_dir']))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "MODEL_CLASSES = {\n",
    "    'bert': (BertConfig, BertForSequenceClassification, BertTokenizer),\n",
    "    'xlnet': (XLNetConfig, XLNetForSequenceClassification, XLNetTokenizer),\n",
    "    'xlm': (XLMConfig, XLMForSequenceClassification, XLMTokenizer),\n",
    "    'roberta': (RobertaConfig, RobertaForSequenceClassification, RobertaTokenizer)\n",
    "}\n",
    "\n",
    "config_class, model_class, tokenizer_class = MODEL_CLASSES[args['model_type']]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config = config_class.from_pretrained(args['model_name'], num_labels=2, finetuning_task=args['task_name'])\n",
    "tokenizer = tokenizer_class.from_pretrained(args['model_name'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = model_class.from_pretrained(args['model_name'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.to(device);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "task = args['task_name']\n",
    "\n",
    "if task in processors.keys() and task in output_modes.keys():\n",
    "    processor = processors[task]()\n",
    "    label_list = processor.get_labels()\n",
    "    num_labels = len(label_list)\n",
    "else:\n",
    "    raise KeyError(f'{task} not found in processors or in output_modes. Please check utils.py.')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_and_cache_examples(task, tokenizer, evaluate=False):\n",
    "    processor = processors[task]()\n",
    "    output_mode = args['output_mode']\n",
    "    \n",
    "    mode = 'dev' if evaluate else 'train'\n",
    "    cached_features_file = os.path.join(args['data_dir'], f\"cached_{mode}_{args['model_name']}_{args['max_seq_length']}_{task}\")\n",
    "    \n",
    "    if os.path.exists(cached_features_file) and not args['reprocess_input_data']:\n",
    "        logger.info(\"Loading features from cached file %s\", cached_features_file)\n",
    "        features = torch.load(cached_features_file)\n",
    "               \n",
    "    else:\n",
    "        logger.info(\"Creating features from dataset file at %s\", args['data_dir'])\n",
    "        label_list = processor.get_labels()\n",
    "        examples = processor.get_dev_examples(args['data_dir']) if evaluate else processor.get_train_examples(args['data_dir'])\n",
    "        \n",
    "        if __name__ == \"__main__\":\n",
    "            features = convert_examples_to_features(examples, label_list, args['max_seq_length'], tokenizer, output_mode,\n",
    "                cls_token_at_end=bool(args['model_type'] in ['xlnet']),            # xlnet has a cls token at the end\n",
    "                cls_token=tokenizer.cls_token,\n",
    "                cls_token_segment_id=2 if args['model_type'] in ['xlnet'] else 0,\n",
    "                sep_token=tokenizer.sep_token,\n",
    "                sep_token_extra=bool(args['model_type'] in ['roberta']),           # roberta uses an extra separator b/w pairs of sentences, cf. github.com/pytorch/fairseq/commit/1684e166e3da03f5b600dbb7855cb98ddfcd0805\n",
    "                pad_on_left=bool(args['model_type'] in ['xlnet']),                 # pad on the left for xlnet\n",
    "                pad_token=tokenizer.convert_tokens_to_ids([tokenizer.pad_token])[0],\n",
    "                pad_token_segment_id=4 if args['model_type'] in ['xlnet'] else 0)\n",
    "        \n",
    "        logger.info(\"Saving features into cached file %s\", cached_features_file)\n",
    "        torch.save(features, cached_features_file)\n",
    "        \n",
    "    all_input_ids = torch.tensor([f.input_ids for f in features], dtype=torch.long)\n",
    "    all_input_mask = torch.tensor([f.input_mask for f in features], dtype=torch.long)\n",
    "    all_segment_ids = torch.tensor([f.segment_ids for f in features], dtype=torch.long)\n",
    "    if output_mode == \"classification\":\n",
    "        all_label_ids = torch.tensor([f.label_id for f in features], dtype=torch.long)\n",
    "    elif output_mode == \"regression\":\n",
    "        all_label_ids = torch.tensor([f.label_id for f in features], dtype=torch.float)\n",
    "\n",
    "    dataset = TensorDataset(all_input_ids, all_input_mask, all_segment_ids, all_label_ids)\n",
    "    return dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(train_dataset, model, tokenizer):\n",
    "    tb_writer = SummaryWriter()\n",
    "    \n",
    "    train_sampler = RandomSampler(train_dataset)\n",
    "    train_dataloader = DataLoader(train_dataset, sampler=train_sampler, batch_size=args['train_batch_size'])\n",
    "    \n",
    "    t_total = len(train_dataloader) // args['gradient_accumulation_steps'] * args['num_train_epochs']\n",
    "    \n",
    "    no_decay = ['bias', 'LayerNorm.weight']\n",
    "    optimizer_grouped_parameters = [\n",
    "        {'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']},\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",
    "    \n",
    "    warmup_steps = math.ceil(t_total * args['warmup_ratio'])\n",
    "    args['warmup_steps'] = warmup_steps if args['warmup_steps'] == 0 else args['warmup_steps']\n",
    "    \n",
    "    optimizer = AdamW(optimizer_grouped_parameters, lr=args['learning_rate'], eps=args['adam_epsilon'])\n",
    "    scheduler = WarmupLinearSchedule(optimizer, warmup_steps=args['warmup_steps'], t_total=t_total)\n",
    "    \n",
    "    if args['fp16']:\n",
    "        try:\n",
    "            from apex import amp\n",
    "        except ImportError:\n",
    "            raise ImportError(\"Please install apex from https://www.github.com/nvidia/apex to use fp16 training.\")\n",
    "        model, optimizer = amp.initialize(model, optimizer, opt_level=args['fp16_opt_level'])\n",
    "        \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(\"  Total train batch size  = %d\", args['train_batch_size'])\n",
    "    logger.info(\"  Gradient Accumulation steps = %d\", args['gradient_accumulation_steps'])\n",
    "    logger.info(\"  Total optimization steps = %d\", t_total)\n",
    "\n",
    "    global_step = 0\n",
    "    tr_loss, logging_loss = 0.0, 0.0\n",
    "    model.zero_grad()\n",
    "    train_iterator = trange(int(args['num_train_epochs']), desc=\"Epoch\")\n",
    "    \n",
    "    for _ in train_iterator:\n",
    "        epoch_iterator = tqdm_notebook(train_dataloader, desc=\"Iteration\")\n",
    "        for step, batch in enumerate(epoch_iterator):\n",
    "            model.train()\n",
    "            batch = tuple(t.to(device) for t in batch)\n",
    "            inputs = {'input_ids':      batch[0],\n",
    "                      'attention_mask': batch[1],\n",
    "                      'token_type_ids': batch[2] if args['model_type'] in ['bert', 'xlnet'] else None,  # XLM don't use segment_ids\n",
    "                      'labels':         batch[3]}\n",
    "            outputs = model(**inputs)\n",
    "            loss = outputs[0]  # model outputs are always tuple in pytorch-transformers (see doc)\n",
    "            print(\"\\r%f\" % loss, end='')\n",
    "\n",
    "            if args['gradient_accumulation_steps'] > 1:\n",
    "                loss = loss / args['gradient_accumulation_steps']\n",
    "\n",
    "            if args['fp16']:\n",
    "                with amp.scale_loss(loss, optimizer) as scaled_loss:\n",
    "                    scaled_loss.backward()\n",
    "                torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer), args['max_grad_norm'])\n",
    "                \n",
    "            else:\n",
    "                loss.backward()\n",
    "                torch.nn.utils.clip_grad_norm_(model.parameters(), args['max_grad_norm'])\n",
    "\n",
    "            tr_loss += loss.item()\n",
    "            if (step + 1) % args['gradient_accumulation_steps'] == 0:\n",
    "                optimizer.step()\n",
    "                scheduler.step()  # Update learning rate schedule\n",
    "                model.zero_grad()\n",
    "                global_step += 1\n",
    "\n",
    "                if args['logging_steps'] > 0 and global_step % args['logging_steps'] == 0:\n",
    "                    # Log metrics\n",
    "                    if args['evaluate_during_training']:  # Only evaluate when single GPU otherwise metrics may not average well\n",
    "                        results, _ = evaluate(model, tokenizer)\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",
    "                if args['save_steps'] > 0 and global_step % args['save_steps'] == 0:\n",
    "                    # Save model checkpoint\n",
    "                    output_dir = os.path.join(args['output_dir'], 'checkpoint-{}'.format(global_step))\n",
    "                    if not os.path.exists(output_dir):\n",
    "                        os.makedirs(output_dir)\n",
    "                    model_to_save = model.module if hasattr(model, 'module') else model  # Take care of distributed/parallel training\n",
    "                    model_to_save.save_pretrained(output_dir)\n",
    "                    logger.info(\"Saving model checkpoint to %s\", output_dir)\n",
    "\n",
    "\n",
    "    return global_step, tr_loss / global_step"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.metrics import mean_squared_error, matthews_corrcoef, confusion_matrix\n",
    "from scipy.stats import pearsonr\n",
    "\n",
    "def get_mismatched(labels, preds):\n",
    "    mismatched = labels != preds\n",
    "    examples = processor.get_dev_examples(args['data_dir'])\n",
    "    wrong = [i for (i, v) in zip(examples, mismatched) if v]\n",
    "    \n",
    "    return wrong\n",
    "\n",
    "def get_eval_report(labels, preds):\n",
    "    mcc = matthews_corrcoef(labels, preds)\n",
    "    tn, fp, fn, tp = confusion_matrix(labels, preds).ravel()\n",
    "    return {\n",
    "        \"mcc\": mcc,\n",
    "        \"tp\": tp,\n",
    "        \"tn\": tn,\n",
    "        \"fp\": fp,\n",
    "        \"fn\": fn\n",
    "    }, get_mismatched(labels, preds)\n",
    "\n",
    "def compute_metrics(task_name, preds, labels):\n",
    "    assert len(preds) == len(labels)\n",
    "    return get_eval_report(labels, preds)\n",
    "\n",
    "def evaluate(model, tokenizer, prefix=\"\"):\n",
    "    # Loop to handle MNLI double evaluation (matched, mis-matched)\n",
    "    eval_output_dir = args['output_dir']\n",
    "\n",
    "    results = {}\n",
    "    EVAL_TASK = args['task_name']\n",
    "\n",
    "    eval_dataset = load_and_cache_examples(EVAL_TASK, tokenizer, evaluate=True)\n",
    "    if not os.path.exists(eval_output_dir):\n",
    "        os.makedirs(eval_output_dir)\n",
    "\n",
    "\n",
    "    eval_sampler = SequentialSampler(eval_dataset)\n",
    "    eval_dataloader = DataLoader(eval_dataset, sampler=eval_sampler, batch_size=args['eval_batch_size'])\n",
    "\n",
    "    # Eval!\n",
    "    logger.info(\"***** Running evaluation {} *****\".format(prefix))\n",
    "    logger.info(\"  Num examples = %d\", len(eval_dataset))\n",
    "    logger.info(\"  Batch size = %d\", args['eval_batch_size'])\n",
    "    eval_loss = 0.0\n",
    "    nb_eval_steps = 0\n",
    "    preds = None\n",
    "    out_label_ids = None\n",
    "    for batch in tqdm_notebook(eval_dataloader, desc=\"Evaluating\"):\n",
    "        model.eval()\n",
    "        batch = tuple(t.to(device) for t in batch)\n",
    "\n",
    "        with torch.no_grad():\n",
    "            inputs = {'input_ids':      batch[0],\n",
    "                      'attention_mask': batch[1],\n",
    "                      'token_type_ids': batch[2] if args['model_type'] in ['bert', 'xlnet'] else None,  # XLM don't use segment_ids\n",
    "                      'labels':         batch[3]}\n",
    "            outputs = model(**inputs)\n",
    "            tmp_eval_loss, logits = outputs[:2]\n",
    "\n",
    "            eval_loss += tmp_eval_loss.mean().item()\n",
    "        nb_eval_steps += 1\n",
    "        if preds is None:\n",
    "            preds = logits.detach().cpu().numpy()\n",
    "            out_label_ids = inputs['labels'].detach().cpu().numpy()\n",
    "        else:\n",
    "            preds = np.append(preds, logits.detach().cpu().numpy(), axis=0)\n",
    "            out_label_ids = np.append(out_label_ids, inputs['labels'].detach().cpu().numpy(), axis=0)\n",
    "\n",
    "    eval_loss = eval_loss / nb_eval_steps\n",
    "    if args['output_mode'] == \"classification\":\n",
    "        preds = np.argmax(preds, axis=1)\n",
    "    elif args['output_mode'] == \"regression\":\n",
    "        preds = np.squeeze(preds)\n",
    "    result, wrong = compute_metrics(EVAL_TASK, preds, out_label_ids)\n",
    "    results.update(result)\n",
    "\n",
    "    output_eval_file = os.path.join(eval_output_dir, \"eval_results.txt\")\n",
    "    with open(output_eval_file, \"w\") as writer:\n",
    "        logger.info(\"***** Eval results {} *****\".format(prefix))\n",
    "        for key in sorted(result.keys()):\n",
    "            logger.info(\"  %s = %s\", key, str(result[key]))\n",
    "            writer.write(\"%s = %s\\n\" % (key, str(result[key])))\n",
    "\n",
    "    return results, wrong"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if args['do_train']:\n",
    "    train_dataset = load_and_cache_examples(task, tokenizer)\n",
    "    global_step, tr_loss = train(train_dataset, model, tokenizer)\n",
    "    logger.info(\" global_step = %s, average loss = %s\", global_step, tr_loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if args['do_train']:\n",
    "    if not os.path.exists(args['output_dir']):\n",
    "            os.makedirs(args['output_dir'])\n",
    "    logger.info(\"Saving model checkpoint to %s\", args['output_dir'])\n",
    "    \n",
    "    model_to_save = model.module if hasattr(model, 'module') else model  # Take care of distributed/parallel training\n",
    "    model_to_save.save_pretrained(args['output_dir'])\n",
    "    tokenizer.save_pretrained(args['output_dir'])\n",
    "    torch.save(args, os.path.join(args['output_dir'], 'training_args.bin'))    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results = {}\n",
    "if args['do_eval']:\n",
    "    checkpoints = [args['output_dir']]\n",
    "    if args['eval_all_checkpoints']:\n",
    "        checkpoints = list(os.path.dirname(c) for c in sorted(glob.glob(args['output_dir'] + '/**/' + WEIGHTS_NAME, recursive=True)))\n",
    "        logging.getLogger(\"pytorch_transformers.modeling_utils\").setLevel(logging.WARN)  # Reduce logging\n",
    "    logger.info(\"Evaluate the following checkpoints: %s\", checkpoints)\n",
    "    for checkpoint in checkpoints:\n",
    "        global_step = checkpoint.split('-')[-1] if len(checkpoints) > 1 else \"\"\n",
    "        model = model_class.from_pretrained(checkpoint)\n",
    "        model.to(device)\n",
    "        result, wrong_preds = evaluate(model, tokenizer, prefix=global_step)\n",
    "        result = dict((k + '_{}'.format(global_step), v) for k, v in result.items())\n",
    "        results.update(result)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results"
   ]
  }
 ],
 "metadata": {
  "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}