colab_quickstart.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/ThilinaRajapakse/pytorch-transformers-classification/blob/master/colab_quickstart.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "MrIwJ5QDz-e9"
   },
   "outputs": [],
   "source": [
    "!pip install pytorch-transformers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "0YLoS0hWz-ch"
   },
   "outputs": [],
   "source": [
    "%%writefile utils.py \n",
    "\n",
    "# coding=utf-8\n",
    "# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.\n",
    "# Copyright (c) 2018, NVIDIA CORPORATION.  All rights reserved.\n",
    "#\n",
    "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "# you may not use this file except in compliance with the License.\n",
    "# You may obtain a copy of the License at\n",
    "#\n",
    "#     http://www.apache.org/licenses/LICENSE-2.0\n",
    "#\n",
    "# Unless required by applicable law or agreed to in writing, software\n",
    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "# See the License for the specific language governing permissions and\n",
    "# limitations under the License.\n",
    "\"\"\" BERT classification fine-tuning: utilities to work with GLUE tasks \"\"\"\n",
    "\n",
    "from __future__ import absolute_import, division, print_function\n",
    "\n",
    "import csv\n",
    "import logging\n",
    "import os\n",
    "import sys\n",
    "from io import open\n",
    "\n",
    "from scipy.stats import pearsonr, spearmanr\n",
    "from sklearn.metrics import matthews_corrcoef, f1_score\n",
    "\n",
    "from multiprocessing import Pool, cpu_count\n",
    "from tqdm import tqdm\n",
    "\n",
    "logger = logging.getLogger(__name__)\n",
    "csv.field_size_limit(2147483647)\n",
    "\n",
    "class InputExample(object):\n",
    "    \"\"\"A single training/test example for simple sequence classification.\"\"\"\n",
    "\n",
    "    def __init__(self, guid, text_a, text_b=None, label=None):\n",
    "        \"\"\"Constructs a InputExample.\n",
    "\n",
    "        Args:\n",
    "            guid: Unique id for the example.\n",
    "            text_a: string. The untokenized text of the first sequence. For single\n",
    "            sequence tasks, only this sequence must be specified.\n",
    "            text_b: (Optional) string. The untokenized text of the second sequence.\n",
    "            Only must be specified for sequence pair tasks.\n",
    "            label: (Optional) string. The label of the example. This should be\n",
    "            specified for train and dev examples, but not for test examples.\n",
    "        \"\"\"\n",
    "        self.guid = guid\n",
    "        self.text_a = text_a\n",
    "        self.text_b = text_b\n",
    "        self.label = label\n",
    "\n",
    "\n",
    "class InputFeatures(object):\n",
    "    \"\"\"A single set of features of data.\"\"\"\n",
    "\n",
    "    def __init__(self, input_ids, input_mask, segment_ids, label_id):\n",
    "        self.input_ids = input_ids\n",
    "        self.input_mask = input_mask\n",
    "        self.segment_ids = segment_ids\n",
    "        self.label_id = label_id\n",
    "\n",
    "\n",
    "class DataProcessor(object):\n",
    "    \"\"\"Base class for data converters for sequence classification data sets.\"\"\"\n",
    "\n",
    "    def get_train_examples(self, data_dir):\n",
    "        \"\"\"Gets a collection of `InputExample`s for the train set.\"\"\"\n",
    "        raise NotImplementedError()\n",
    "\n",
    "    def get_dev_examples(self, data_dir):\n",
    "        \"\"\"Gets a collection of `InputExample`s for the dev set.\"\"\"\n",
    "        raise NotImplementedError()\n",
    "\n",
    "    def get_labels(self):\n",
    "        \"\"\"Gets the list of labels for this data set.\"\"\"\n",
    "        raise NotImplementedError()\n",
    "\n",
    "    @classmethod\n",
    "    def _read_tsv(cls, input_file, quotechar=None):\n",
    "        \"\"\"Reads a tab separated value file.\"\"\"\n",
    "        with open(input_file, \"r\", encoding=\"utf-8-sig\") as f:\n",
    "            reader = csv.reader(f, delimiter=\"\\t\", quotechar=quotechar)\n",
    "            lines = []\n",
    "            for line in reader:\n",
    "                if sys.version_info[0] == 2:\n",
    "                    line = list(unicode(cell, 'utf-8') for cell in line)\n",
    "                lines.append(line)\n",
    "            return lines\n",
    "\n",
    "\n",
    "class BinaryProcessor(DataProcessor):\n",
    "    \"\"\"Processor for the binary data sets\"\"\"\n",
    "\n",
    "    def get_train_examples(self, data_dir):\n",
    "        \"\"\"See base class.\"\"\"\n",
    "        return self._create_examples(\n",
    "            self._read_tsv(os.path.join(data_dir, \"train.tsv\")), \"train\")\n",
    "\n",
    "    def get_dev_examples(self, data_dir):\n",
    "        \"\"\"See base class.\"\"\"\n",
    "        return self._create_examples(\n",
    "            self._read_tsv(os.path.join(data_dir, \"dev.tsv\")), \"dev\")\n",
    "\n",
    "    def get_labels(self):\n",
    "        \"\"\"See base class.\"\"\"\n",
    "        return [\"0\", \"1\"]\n",
    "\n",
    "    def _create_examples(self, lines, set_type):\n",
    "        \"\"\"Creates examples for the training and dev sets.\"\"\"\n",
    "        examples = []\n",
    "        for (i, line) in enumerate(lines):\n",
    "            guid = \"%s-%s\" % (set_type, i)\n",
    "            text_a = line[3]\n",
    "            label = line[1]\n",
    "            examples.append(\n",
    "                InputExample(guid=guid, text_a=text_a, text_b=None, label=label))\n",
    "        return examples\n",
    "\n",
    "\n",
    "def convert_example_to_feature(example_row, pad_token=0,\n",
    "sequence_a_segment_id=0, sequence_b_segment_id=1,\n",
    "cls_token_segment_id=1, pad_token_segment_id=0,\n",
    "mask_padding_with_zero=True):\n",
    "    example, label_map, max_seq_length, tokenizer, output_mode, cls_token_at_end, cls_token, sep_token, cls_token_segment_id, pad_on_left, pad_token_segment_id = example_row\n",
    "\n",
    "    tokens_a = tokenizer.tokenize(example.text_a)\n",
    "\n",
    "    tokens_b = None\n",
    "    if example.text_b:\n",
    "        tokens_b = tokenizer.tokenize(example.text_b)\n",
    "        # Modifies `tokens_a` and `tokens_b` in place so that the total\n",
    "        # length is less than the specified length.\n",
    "        # Account for [CLS], [SEP], [SEP] with \"- 3\"\n",
    "        _truncate_seq_pair(tokens_a, tokens_b, max_seq_length - 3)\n",
    "    else:\n",
    "        # Account for [CLS] and [SEP] with \"- 2\"\n",
    "        if len(tokens_a) > max_seq_length - 2:\n",
    "            tokens_a = tokens_a[:(max_seq_length - 2)]\n",
    "\n",
    "    # The convention in BERT is:\n",
    "    # (a) For sequence pairs:\n",
    "    #  tokens:   [CLS] is this jack ##son ##ville ? [SEP] no it is not . [SEP]\n",
    "    #  type_ids:   0   0  0    0    0     0       0   0   1  1  1  1   1   1\n",
    "    # (b) For single sequences:\n",
    "    #  tokens:   [CLS] the dog is hairy . [SEP]\n",
    "    #  type_ids:   0   0   0   0  0     0   0\n",
    "    #\n",
    "    # Where \"type_ids\" are used to indicate whether this is the first\n",
    "    # sequence or the second sequence. The embedding vectors for `type=0` and\n",
    "    # `type=1` were learned during pre-training and are added to the wordpiece\n",
    "    # embedding vector (and position vector). This is not *strictly* necessary\n",
    "    # since the [SEP] token unambiguously separates the sequences, but it makes\n",
    "    # it easier for the model to learn the concept of sequences.\n",
    "    #\n",
    "    # For classification tasks, the first vector (corresponding to [CLS]) is\n",
    "    # used as as the \"sentence vector\". Note that this only makes sense because\n",
    "    # the entire model is fine-tuned.\n",
    "    tokens = tokens_a + [sep_token]\n",
    "    segment_ids = [sequence_a_segment_id] * len(tokens)\n",
    "\n",
    "    if tokens_b:\n",
    "        tokens += tokens_b + [sep_token]\n",
    "        segment_ids += [sequence_b_segment_id] * (len(tokens_b) + 1)\n",
    "\n",
    "    if cls_token_at_end:\n",
    "        tokens = tokens + [cls_token]\n",
    "        segment_ids = segment_ids + [cls_token_segment_id]\n",
    "    else:\n",
    "        tokens = [cls_token] + tokens\n",
    "        segment_ids = [cls_token_segment_id] + segment_ids\n",
    "\n",
    "    input_ids = tokenizer.convert_tokens_to_ids(tokens)\n",
    "\n",
    "    # The mask has 1 for real tokens and 0 for padding tokens. Only real\n",
    "    # tokens are attended to.\n",
    "    input_mask = [1 if mask_padding_with_zero else 0] * len(input_ids)\n",
    "\n",
    "    # Zero-pad up to the sequence length.\n",
    "    padding_length = max_seq_length - len(input_ids)\n",
    "    if pad_on_left:\n",
    "        input_ids = ([pad_token] * padding_length) + input_ids\n",
    "        input_mask = ([0 if mask_padding_with_zero else 1] * padding_length) + input_mask\n",
    "        segment_ids = ([pad_token_segment_id] * padding_length) + segment_ids\n",
    "    else:\n",
    "        input_ids = input_ids + ([pad_token] * padding_length)\n",
    "        input_mask = input_mask + ([0 if mask_padding_with_zero else 1] * padding_length)\n",
    "        segment_ids = segment_ids + ([pad_token_segment_id] * padding_length)\n",
    "\n",
    "    assert len(input_ids) == max_seq_length\n",
    "    assert len(input_mask) == max_seq_length\n",
    "    assert len(segment_ids) == max_seq_length\n",
    "\n",
    "    if output_mode == \"classification\":\n",
    "        label_id = label_map[example.label]\n",
    "    elif output_mode == \"regression\":\n",
    "        label_id = float(example.label)\n",
    "    else:\n",
    "        raise KeyError(output_mode)\n",
    "\n",
    "    return InputFeatures(input_ids=input_ids,\n",
    "                        input_mask=input_mask,\n",
    "                        segment_ids=segment_ids,\n",
    "                        label_id=label_id)\n",
    "    \n",
    "\n",
    "def convert_examples_to_features(examples, label_list, max_seq_length,\n",
    "                                 tokenizer, output_mode,\n",
    "                                 cls_token_at_end=False, pad_on_left=False,\n",
    "                                 cls_token='[CLS]', sep_token='[SEP]', pad_token=0,\n",
    "                                 sequence_a_segment_id=0, sequence_b_segment_id=1,\n",
    "                                 cls_token_segment_id=1, pad_token_segment_id=0,\n",
    "                                 mask_padding_with_zero=True,\n",
    "                                 process_count=cpu_count() - 2):\n",
    "    \"\"\" Loads a data file into a list of `InputBatch`s\n",
    "        `cls_token_at_end` define the location of the CLS token:\n",
    "            - False (Default, BERT/XLM pattern): [CLS] + A + [SEP] + B + [SEP]\n",
    "            - True (XLNet/GPT pattern): A + [SEP] + B + [SEP] + [CLS]\n",
    "        `cls_token_segment_id` define the segment id associated to the CLS token (0 for BERT, 2 for XLNet)\n",
    "    \"\"\"\n",
    "\n",
    "    label_map = {label : i for i, label in enumerate(label_list)}\n",
    "\n",
    "    examples = [(example, label_map, max_seq_length, tokenizer, output_mode, cls_token_at_end, cls_token, sep_token, cls_token_segment_id, pad_on_left, pad_token_segment_id) for example in examples]\n",
    "\n",
    "    with Pool(process_count) as p:\n",
    "        features = list(tqdm(p.imap(convert_example_to_feature, examples, chunksize=100), total=len(examples)))\n",
    "\n",
    "    return features\n",
    "\n",
    "\n",
    "def _truncate_seq_pair(tokens_a, tokens_b, max_length):\n",
    "    \"\"\"Truncates a sequence pair in place to the maximum length.\"\"\"\n",
    "\n",
    "    # This is a simple heuristic which will always truncate the longer sequence\n",
    "    # one token at a time. This makes more sense than truncating an equal percent\n",
    "    # of tokens from each, since if one sequence is very short then each token\n",
    "    # that's truncated likely contains more information than a longer sequence.\n",
    "    while True:\n",
    "        total_length = len(tokens_a) + len(tokens_b)\n",
    "        if total_length <= max_length:\n",
    "            break\n",
    "        if len(tokens_a) > len(tokens_b):\n",
    "            tokens_a.pop()\n",
    "        else:\n",
    "            tokens_b.pop()\n",
    "\n",
    "\n",
    "processors = {\n",
    "    \"binary\": BinaryProcessor\n",
    "}\n",
    "\n",
    "output_modes = {\n",
    "    \"binary\": \"classification\"\n",
    "}\n",
    "\n",
    "GLUE_TASKS_NUM_LABELS = {\n",
    "    \"binary\": 2\n",
    "}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "iDJUEfruz-aF"
   },
   "outputs": [],
   "source": [
    "!mkdir data\n",
    "!wget https://s3.amazonaws.com/fast-ai-nlp/yelp_review_polarity_csv.tgz -O data/data.tgz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "aepMB1Z_2u0V"
   },
   "outputs": [],
   "source": [
    "!tar -xvzf data/data.tgz -C data/\n",
    "!mv data/yelp_review_polarity_csv/* data/\n",
    "!rm -r data/yelp_review_polarity_csv/\n",
    "!rm data/data.tgz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "qYv4tGNs4Tg8"
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "from tqdm import tqdm_notebook\n",
    "\n",
    "prefix = 'data/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "BiC1S1xq4XnN"
   },
   "outputs": [],
   "source": [
    "train_df = pd.read_csv(prefix + 'train.csv', header=None)\n",
    "train_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "um-kSl434Xk6"
   },
   "outputs": [],
   "source": [
    "test_df = pd.read_csv(prefix + 'test.csv', header=None)\n",
    "test_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "_Xe5tm-2441r"
   },
   "outputs": [],
   "source": [
    "train_df[0] = (train_df[0] == 2).astype(int)\n",
    "test_df[0] = (test_df[0] == 2).astype(int)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "hVsnPNzr4XjB"
   },
   "outputs": [],
   "source": [
    "train_df = pd.DataFrame({\n",
    "    'id':range(len(train_df)),\n",
    "    'label':train_df[0],\n",
    "    'alpha':['a']*train_df.shape[0],\n",
    "    'text': train_df[1].replace(r'\\n', ' ', regex=True)\n",
    "})\n",
    "\n",
    "train_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "IkUFgZeh4Xf2"
   },
   "outputs": [],
   "source": [
    "dev_df = pd.DataFrame({\n",
    "    'id':range(len(test_df)),\n",
    "    'label':test_df[0],\n",
    "    'alpha':['a']*test_df.shape[0],\n",
    "    'text': test_df[1].replace(r'\\n', ' ', regex=True)\n",
    "})\n",
    "\n",
    "dev_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "pFlvoH234XdO"
   },
   "outputs": [],
   "source": [
    "train_df.to_csv('data/train.tsv', sep='\\t', index=False, header=False, columns=train_df.columns)\n",
    "dev_df.to_csv('data/dev.tsv', sep='\\t', index=False, header=False, columns=dev_df.columns)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "VeXuXWylz7BD"
   },
   "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",
    "\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",
    "\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": {
    "colab": {},
    "colab_type": "code",
    "id": "F93_pIopz7BG"
   },
   "outputs": [],
   "source": [
    "args = {\n",
    "    'data_dir': 'data/',\n",
    "    'model_type':  'roberta',\n",
    "    'model_name': 'roberta-base',\n",
    "    'task_name': 'binary',\n",
    "    'output_dir': 'outputs/',\n",
    "    'cache_dir': 'cache/',\n",
    "    'do_train': True,\n",
    "    'do_eval': True,\n",
    "    'fp16': False,\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_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': False,\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": {
    "colab": {},
    "colab_type": "code",
    "id": "atGwIw3iz7BJ"
   },
   "outputs": [],
   "source": [
    "args"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Uzr2RwGLz7BL"
   },
   "outputs": [],
   "source": [
    "with open('args.json', 'w') as f:\n",
    "    json.dump(args, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ymjmIyOhz7BN"
   },
   "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": {
    "colab": {},
    "colab_type": "code",
    "id": "LAHYiiLMz7BP"
   },
   "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": {
    "colab": {},
    "colab_type": "code",
    "id": "qm5AguwFz7BR"
   },
   "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": {
    "colab": {},
    "colab_type": "code",
    "id": "IGZHNvKAz7BU"
   },
   "outputs": [],
   "source": [
    "model = model_class.from_pretrained(args['model_name'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "xyxKpk_6z7BW"
   },
   "outputs": [],
   "source": [
    "model.to(device);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "bsPmRyGE8GnR"
   },
   "outputs": [],
   "source": [
    "device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Xe4P94Bfz7Ba"
   },
   "outputs": [],
   "source": [
    "task = args['task_name']\n",
    "\n",
    "processor = processors[task]()\n",
    "label_list = processor.get_labels()\n",
    "num_labels = len(label_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "xqr_fwM3z7Bd"
   },
   "outputs": [],
   "source": [
    "def load_and_cache_examples(task, tokenizer, evaluate=False, undersample_scale_factor=0.01):\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",
    "        print(len(examples))\n",
    "        examples  = [example for example in examples if np.random.rand() < undersample_scale_factor]\n",
    "        print(len(examples))\n",
    "        \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",
    "            sep_token=tokenizer.sep_token,\n",
    "            cls_token_segment_id=2 if args['model_type'] in ['xlnet'] else 0,\n",
    "            pad_on_left=bool(args['model_type'] in ['xlnet']),                 # pad on the left for xlnet\n",
    "            pad_token_segment_id=4 if args['model_type'] in ['xlnet'] else 0,\n",
    "            process_count=2)\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": {
    "colab": {},
    "colab_type": "code",
    "id": "oCul6vvCz7Bg"
   },
   "outputs": [],
   "source": [
    "def train(train_dataset, model, tokenizer):\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",
    "    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",
    "                scheduler.step()  # Update learning rate schedule\n",
    "                optimizer.step()\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",
    "\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": {
    "colab": {},
    "colab_type": "code",
    "id": "tUvkEBZUz7Bk"
   },
   "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": {
    "colab": {},
    "colab_type": "code",
    "id": "MlaeXY9sz7Bm"
   },
   "outputs": [],
   "source": [
    "# IMPORTANT #\n",
    "# Due to the 12 hour limit on Google Colab and the time it would take to convert the dataset into features, the load_and_cache_examples() function has been modified\n",
    "# to randomly undersample the dataset by a scale of 0.1\n",
    "\n",
    "if args['do_train']:\n",
    "    train_dataset = load_and_cache_examples(task, tokenizer, undersample_scale_factor=0.1)\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": {
    "colab": {},
    "colab_type": "code",
    "id": "1On6YjIULf7v"
   },
   "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'))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "tqiWWPA0z7Bo"
   },
   "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": {
    "colab": {},
    "colab_type": "code",
    "id": "AMb25x63z7Bq"
   },
   "outputs": [],
   "source": [
    "results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "eyvWYNjRLHrI"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "include_colab_link": true,
   "name": "Copy of train.ipynb",
   "provenance": [],
   "version": "0.3.2"
  },
  "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
}