demo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# demo \n",
    "\n",
    "We will go through 5 examples: \n",
    "\n",
    "* **[text classification ](#text_classification)** - the goal is to classify a single sentence or short text.\n",
    "\n",
    "\n",
    "* **[text pair classification ](#text_pair_classification)** - the goal is to to classify a pair of sentences or short texts.\n",
    "\n",
    "\n",
    "* **[text pair regression ](#text_pair_regression)** - the goal is to predict a numerical value for a pair of sentences or short texts.\n",
    "\n",
    "\n",
    "* **[Named Entity Recognition (NER)](#ner_conll_eng)** - the goal is to tag each token in a list of tokens as a person, location, organization,etc.\n",
    "\n",
    "\n",
    "\n",
    "* **[Biomedical NER](#ner_ncbi)** - another NER task but focused on the biomedical domain . We will finetune [**BERT**](#NCBI_BERT), [**SciBERT**](#NCBI_SciBERT), and  [**BioBERT**](#NCBI_BioBERT) based models.\n",
    "\n",
    "\n",
    "### A note on GPU cards and memory\n",
    "\n",
    "While its possible, it would be very slow to run the examples without a GPU card of some sort. In addition, the BERT models (especially the large model) are pretty big so it helps to have more GPU memory. All the examples in this demo were run on a laptop with a Nvidia GTX-1070 card that has 8Gb of memory.\n",
    "\n",
    "The three biggest parameters you can change which will reduce the GPU memory requirements significantly are:\n",
    "\n",
    "* **`bert_model`** - BERT models come in 2 sizes : `base` and `large`. As you would expect the large model demands more GPU memory and takes longer to train. If you have a small GPU, start with the any of the `base` models first. The default is set to `'bert-base-uncased'`\n",
    "\n",
    "> `base(110M parameter models)` : `'bert-base-uncased'`, `'bert-base-cased'`, `'bert-base-multilingual-uncased'`, `'bert-base-multilingual-cased'`, `'bert-base-chinese'`, and all the **`BioBERT`** and **`SciBERT`** models.\n",
    "\n",
    "> `large(340M parameter models)`: `'bert-large-uncased'` and `'bert-large-cased'`\n",
    "\n",
    "\n",
    "* **`max_seq_length`** - the defualt is 128 with a max value of 512. But seting it to a smaller value like 96 or even 64  saves a lot of GPU memory and still gets good results on a lot of tasks.\n",
    "\n",
    "\n",
    "* **`train_batch_size`** - the default is 32. Cutting it in half will save memory and should also still give good results.\n",
    "\n",
    "\n",
    "In addition to these two parameters,  [huggingface/pytorch-pretrained-BERT](https://github.com/huggingface/pytorch-pretrained-BERT#Training-large-models-introduction,-tools-and-examples) has several options to reduce the GPU memory requirements which are passed through in `bert-sklearn`:\n",
    "\n",
    "* **`gradient_accumulation_steps`** - this is the number of update steps to accumulate gradients before performing an update step with the optimizer. The default is 1. Setting it to a higher integer(i.e 2, 4, up to the **`train_batch_size`** ) will trade GPU memory for compute time. I use this a lot when I train BERT models on my laptop GPU.\n",
    "\n",
    "\n",
    "* **`fp16`** - this is whether to use 16-bit float precision instead of the 32-bit. The default is set to `False`. To enable half precision, you must install [Nvidia apex](https://github.com/NVIDIA/apex). Then setting this option to `True` will cut the model memory load in half. I use this when I train on my laptop GPU as well.\n",
    "\n",
    "\n",
    "Finally the two other system setups that will help reduce the memory requirement: \n",
    "\n",
    "* `multiple gpus` - for a single machine with multiple GPUs, following the huggingface port, the GPUs will be detected and will split the load onto the multiple cards. \n",
    "\n",
    "\n",
    "* `distributed training` - the huggingface port allows you to train across distributed GPUs. The parameter,  **`local_rank`**, is exposed in `bert-sklearn`. But this option has not been tested yet.\n",
    "\n",
    "First, let's import some libraries and define a few functions we will use in the rest of the demo:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import math\n",
    "import random\n",
    "import csv\n",
    "import sys\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from sklearn import metrics\n",
    "from sklearn.metrics import f1_score\n",
    "from sklearn.metrics import classification_report\n",
    "import statistics as stats\n",
    "\n",
    "from bert_sklearn import BertClassifier\n",
    "from bert_sklearn import BertRegressor\n",
    "from bert_sklearn import BertTokenClassifier\n",
    "from bert_sklearn import load_model\n",
    "\n",
    "def read_tsv(filename, quotechar=None):\n",
    "    with open(filename, \"r\", encoding='utf-8') as f:\n",
    "        return list(csv.reader(f, delimiter=\"\\t\", quotechar=quotechar))   \n",
    "\n",
    "def flatten(l):\n",
    "    return [item for sublist in l for item in sublist]\n",
    "\n",
    "def read_CoNLL2003_format(filename, idx=3):\n",
    "    \"\"\"Read file in CoNLL-2003 shared task format\"\"\"\n",
    "    \n",
    "    # read file\n",
    "    lines =  open(filename).read().strip()   \n",
    "    \n",
    "    # find sentence-like boundaries\n",
    "    lines = lines.split(\"\\n\\n\")  \n",
    "    \n",
    "     # split on newlines\n",
    "    lines = [line.split(\"\\n\") for line in lines]\n",
    "    \n",
    "    # get tokens\n",
    "    tokens = [[l.split()[0] for l in line] for line in lines]\n",
    "    \n",
    "    # get labels/tags\n",
    "    labels = [[l.split()[idx] for l in line] for line in lines]\n",
    "    \n",
    "    #convert to df\n",
    "    data= {'tokens': tokens, 'labels': labels}\n",
    "    df=pd.DataFrame(data=data)\n",
    "    \n",
    "    return df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "<a id='text_classification'></a>\n",
    "# text classification \n",
    "\n",
    "For single text classification, we have the input data `X`, and target data `y` where:\n",
    "\n",
    "* `X` is a list, pandas Series, or numpy array of text data.\n",
    "\n",
    "\n",
    "* `y` is a list, pandas Series, or numpy array of text labels.\n",
    "\n",
    "For this example, we will use the **`Stanford Sentiment Treebank (SST-2)`** data set from the [GLUE benchmarks](https://gluebenchmark.com/). The **`SST-2`** task consists of sentences drawn from movie reviews and annotated with a sentiment label. \n",
    "\n",
    "See [website](https://nlp.stanford.edu/sentiment/code.html) and [paper](https://nlp.stanford.edu/~socherr/EMNLP2013_RNTN.pdf) for more info.\n",
    "\n",
    "The input features are short sentences and the labels are the standard sentiment polarity of:\n",
    "\n",
    "*    0 for negative \n",
    "\n",
    "\n",
    "*    1 for positive.\n",
    "\n",
    "## get data\n",
    "\n",
    "First download the data using the GLUE downloder:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading and extracting SST...\n",
      "\tCompleted!\n"
     ]
    }
   ],
   "source": [
    "%%bash\n",
    "python3 ./glue_examples/download_glue_data.py --data_dir ./glue_examples//glue_data --tasks SST "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SST-2 train data size: 67349 \n",
      "SST-2 dev data size: 872 \n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>text</th>\n",
       "      <th>label</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>hide new secretions from the parental units</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>contains no wit , only labored gags</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>that loves its characters and communicates som...</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>remains utterly satisfied to remain the same t...</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>on the worst revenge-of-the-nerds clichés the ...</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                text  label\n",
       "0       hide new secretions from the parental units       0\n",
       "1               contains no wit , only labored gags       0\n",
       "2  that loves its characters and communicates som...      1\n",
       "3  remains utterly satisfied to remain the same t...      0\n",
       "4  on the worst revenge-of-the-nerds clichés the ...      0"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\"\"\"\n",
    "SST-2 train data size: 67349 \n",
    "SST-2 dev data size: 872 \n",
    "\"\"\"\n",
    "DATADIR = './glue_examples/glue_data'\n",
    "\n",
    "def get_sst_data(train_file=DATADIR + '/SST-2/train.tsv',\n",
    "                 dev_file=DATADIR + '/SST-2/dev.tsv'):\n",
    "\n",
    "    train = pd.read_csv(train_file, sep='\\t', encoding='utf8', keep_default_na=False)\n",
    "    train.columns=['text', 'label']\n",
    "    print(\"SST-2 train data size: %d \"%(len(train)))\n",
    "    \n",
    "    dev = pd.read_csv(dev_file, sep='\\t', encoding='utf8', keep_default_na=False)\n",
    "    dev.columns=['text', 'label']\n",
    "    print(\"SST-2 dev data size: %d \"%(len(dev)))\n",
    "    label_list = np.unique(train['label'])\n",
    "\n",
    "    return train, dev, label_list\n",
    "\n",
    "train, dev, label_list = get_sst_data()\n",
    "train.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## setup data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will subsample the data for the demo. See the [SST-2.ipynb notebook](https://github.com/charles9n/bert-sklearn/blob/master/glue_examples/SST-2.ipynb) for a finetune demo on the full data set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# subsample data \n",
    "n = 1000\n",
    "train = train.sample(n, random_state=42)\n",
    "\n",
    "X_train = train['text']\n",
    "y_train = train['label']\n",
    "\n",
    "# use the dev set for testing\n",
    "test = dev\n",
    "X_test = test['text']\n",
    "y_test = test['label']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## define model\n",
    "\n",
    "We will set up a classifier with the defualt settings, but let's reduce **`max_sequence_length`** , and **`train_batch_size`**, so it can run on a smaller GPU. This config uses ~5Gb of GPU memory om my laptop 8GB GTX-1070:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Building sklearn text classifier...\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "BertClassifier(bert_config_json=None, bert_model='bert-base-uncased',\n",
       "        bert_vocab=None, do_lower_case=None, epochs=3, eval_batch_size=8,\n",
       "        fp16=False, from_tf=False, gradient_accumulation_steps=1,\n",
       "        ignore_label=None, label_list=None, learning_rate=2e-05,\n",
       "        local_rank=-1, logfile='bert_sklearn.log', loss_scale=0,\n",
       "        max_seq_length=64, num_mlp_hiddens=500, num_mlp_layers=0,\n",
       "        random_state=42, restore_file=None, train_batch_size=16,\n",
       "        use_cuda=True, validation_fraction=0.1, warmup_proportion=0.1)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = BertClassifier(max_seq_length=64, train_batch_size=16)\n",
    "model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## finetune model\n",
    "\n",
    "finetune = fit model on train data\n",
    "\n",
    "The `model.fit()` routine:\n",
    "\n",
    "* Loads the pretrained BERT model defined in `model.bert_model`. The first time this runs will be slower as it downloads the BERT model from the internet. Subsequent calls will be faster as the model is saved in a file cache locally.\n",
    "\n",
    "\n",
    "* Uses `model.validation_fraction`(defualt=0.1) of the data for validation and finetunes BERT on the remainder for `model.epochs`(default=3) epochs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading bert-base-uncased model...\n",
      "Defaulting to linear classifier/regressor\n",
      "Loading Pytorch checkpoint\n",
      "train data size: 900, validation data size: 100\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Training  : 100%|██████████| 57/57 [00:15<00:00,  4.17it/s, loss=0.529]\n",
      "Validating: 100%|██████████| 13/13 [00:00<00:00, 22.70it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1, Train loss: 0.5295, Val loss: 0.4408, Val accy: 81.00%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Training  : 100%|██████████| 57/57 [00:15<00:00,  4.16it/s, loss=0.167]\n",
      "Validating: 100%|██████████| 13/13 [00:00<00:00, 22.89it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 2, Train loss: 0.1668, Val loss: 0.4380, Val accy: 87.00%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Training  : 100%|██████████| 57/57 [00:15<00:00,  4.16it/s, loss=0.0434]\n",
      "Validating: 100%|██████████| 13/13 [00:00<00:00, 22.74it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 3, Train loss: 0.0434, Val loss: 0.5512, Val accy: 86.00%\n",
      "CPU times: user 36.1 s, sys: 16.5 s, total: 52.6 s\n",
      "Wall time: 54.1 s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "model = model.fit(X_train, y_train)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## score and make predictions on test data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Testing: 100%|██████████| 109/109 [00:03<00:00, 27.99it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Loss: 0.3717, Accuracy: 88.07%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Predicting: 100%|██████████| 109/109 [00:03<00:00, 27.94it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "class prob estimates:\n",
      " [[0.00176739 0.9982326 ]\n",
      " [0.978774   0.02122599]\n",
      " [0.00462427 0.99537575]\n",
      " ...\n",
      " [0.96313787 0.03686218]\n",
      " [0.1856012  0.81439877]\n",
      " [0.00501524 0.99498475]]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Predicting: 100%|██████████| 109/109 [00:03<00:00, 27.85it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy: 88.07%\n",
      "              precision    recall  f1-score   support\n",
      "\n",
      "    negative       0.88      0.87      0.88       428\n",
      "    positive       0.88      0.89      0.88       444\n",
      "\n",
      "   micro avg       0.88      0.88      0.88       872\n",
      "   macro avg       0.88      0.88      0.88       872\n",
      "weighted avg       0.88      0.88      0.88       872\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "from tqdm import tqdm\n",
    "# score model\n",
    "accy = model.score(X_test, y_test)\n",
    "\n",
    "# make class probability predictions\n",
    "y_prob = model.predict_proba(X_test)\n",
    "print(\"class prob estimates:\\n\", y_prob)\n",
    "\n",
    "# make predictions\n",
    "y_pred = model.predict(X_test)\n",
    "print(\"Accuracy: %0.2f%%\"%(metrics.accuracy_score(y_pred, y_test) * 100))\n",
    "\n",
    "target_names = ['negative', 'positive']\n",
    "print(classification_report(y_test, y_pred, target_names=target_names))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## save/load model from disk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading model from /data/test.bin...\n",
      "Defaulting to linear classifier/regressor\n",
      "Building sklearn text classifier...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Testing: 100%|██████████| 109/109 [00:03<00:00, 27.95it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Loss: 0.3717, Accuracy: 88.07%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "#save model to disk\n",
    "savefile = '/data/test.bin'\n",
    "model.save(savefile)\n",
    "\n",
    "# load model from disk\n",
    "new_model = load_model(savefile)\n",
    "\n",
    "# predict with new model\n",
    "accy = new_model.score(X_test, y_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### random seed\n",
    "The finetuned model weights will change depending on the random seeds we use for the pytorch and numpy RNGs. The variance in test accuracy is higher when the training data is small. If you want to check out the variability with a few random seeds the following cell  takes ~3min to run and uses ~6.5GB on my laptop GPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading bert-base-uncased model...\n",
      "Defaulting to linear classifier/regressor\n",
      "Loading Pytorch checkpoint\n",
      "train data size: 900, validation data size: 100\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Training  : 100%|██████████| 57/57 [00:16<00:00,  4.04it/s, loss=0.502]\n",
      "Validating: 100%|██████████| 13/13 [00:00<00:00, 22.24it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1, Train loss: 0.5023, Val loss: 0.5222, Val accy: 81.00%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Training  : 100%|██████████| 57/57 [00:16<00:00,  4.10it/s, loss=0.184]\n",
      "Validating: 100%|██████████| 13/13 [00:00<00:00, 21.74it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 2, Train loss: 0.1842, Val loss: 0.4111, Val accy: 87.00%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Training  : 100%|██████████| 57/57 [00:16<00:00,  3.99it/s, loss=0.0377]\n",
      "Validating: 100%|██████████| 13/13 [00:00<00:00, 21.17it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 3, Train loss: 0.0377, Val loss: 0.5228, Val accy: 88.00%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Testing: 100%|██████████| 109/109 [00:04<00:00, 28.08it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Loss: 0.3578, Accuracy: 88.07%\n",
      "Loading bert-base-uncased model...\n",
      "Defaulting to linear classifier/regressor\n",
      "Loading Pytorch checkpoint\n",
      "train data size: 900, validation data size: 100\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Training  : 100%|██████████| 57/57 [00:16<00:00,  3.73it/s, loss=0.533]\n",
      "Validating: 100%|██████████| 13/13 [00:00<00:00, 20.84it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1, Train loss: 0.5332, Val loss: 0.2840, Val accy: 86.00%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Training  : 100%|██████████| 57/57 [00:16<00:00,  3.94it/s, loss=0.204]\n",
      "Validating: 100%|██████████| 13/13 [00:00<00:00, 21.30it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 2, Train loss: 0.2042, Val loss: 0.2330, Val accy: 91.00%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Training  : 100%|██████████| 57/57 [00:16<00:00,  3.98it/s, loss=0.0656]\n",
      "Validating: 100%|██████████| 13/13 [00:00<00:00, 21.87it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 3, Train loss: 0.0656, Val loss: 0.2525, Val accy: 91.00%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Testing: 100%|██████████| 109/109 [00:04<00:00, 27.26it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Loss: 0.4295, Accuracy: 85.89%\n",
      "Loading bert-base-uncased model...\n",
      "Defaulting to linear classifier/regressor\n",
      "Loading Pytorch checkpoint\n",
      "train data size: 900, validation data size: 100\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Training  : 100%|██████████| 57/57 [00:16<00:00,  4.10it/s, loss=0.585]\n",
      "Validating: 100%|██████████| 13/13 [00:00<00:00, 21.12it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1, Train loss: 0.5846, Val loss: 0.4408, Val accy: 76.00%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Training  : 100%|██████████| 57/57 [00:16<00:00,  3.94it/s, loss=0.291]\n",
      "Validating: 100%|██████████| 13/13 [00:00<00:00, 20.98it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 2, Train loss: 0.2907, Val loss: 0.2925, Val accy: 90.00%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Training  : 100%|██████████| 57/57 [00:17<00:00,  3.86it/s, loss=0.0896]\n",
      "Validating: 100%|██████████| 13/13 [00:00<00:00, 20.87it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 3, Train loss: 0.0896, Val loss: 0.2609, Val accy: 91.00%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Testing: 100%|██████████| 109/109 [00:04<00:00, 22.76it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Loss: 0.4628, Accuracy: 84.75%\n",
      "CPU times: user 1min 58s, sys: 55.7 s, total: 2min 54s\n",
      "Wall time: 2min 58s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "scores = []; \n",
    "for seed in [4, 27, 33]:\n",
    "    model.random_state = seed\n",
    "    model.fit(X_train, y_train)\n",
    "    scores.append(model.score(X_test, y_test))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[88.0733945  85.89449541 84.74770642 88.0733945 ]\n",
      "86.70% (+/-3.313)\n"
     ]
    }
   ],
   "source": [
    "# lets add the accy from our earlier run as well that uses the default seed=42\n",
    "scores = np.array(scores + [accy])\n",
    "print(scores)\n",
    "print(\"%0.2f%% (+/-%0.03f)\"% (stats.mean(scores), stats.stdev(scores) * 2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='text_pair_classification'></a>\n",
    "\n",
    "# text pair classification\n",
    "\n",
    "For text pair classification, we have input data `X`, and target data `y` where :\n",
    "\n",
    "* `X` is a list, pandas dataframe, or numpy array of text pairs (`text_a`, `text_b`) .\n",
    "\n",
    "\n",
    "* `y` is a list, pandas Series, or numpy array of text labels\n",
    "\n",
    "For this example, we will use the **`Quora Question Pair(QQP)`** data set from the [GLUE benchmarks](https://gluebenchmark.com/). This data consists of sentence pairs from the Quora website labeled as duplicate or not. See the [original release post](https://data.quora.com/First-Quora-Dataset-Release-Question-Pairs) for more info.\n",
    "\n",
    "The input features are pairs of questions (text_a,text_b) along with the labels :\n",
    "*    0 if `text_a` and `text_b` are not duplicates\n",
    "\n",
    "*    1 if `text_a` and `text_b` are duplicates\n",
    "\n",
    "\n",
    "## get data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading and extracting QQP...\n",
      "\tCompleted!\n"
     ]
    }
   ],
   "source": [
    "%%bash\n",
    "python3 ./glue_examples/download_glue_data.py --data_dir ./glue_examples//glue_data --tasks QQP "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "QQP train data size: 363849 \n",
      "QQP dev data size: 40430 \n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>text_a</th>\n",
       "      <th>text_b</th>\n",
       "      <th>label</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>How is the life of a math student? Could you d...</td>\n",
       "      <td>Which level of prepration is enough for the ex...</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>How do I control my horny emotions?</td>\n",
       "      <td>How do you control your horniness?</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>What causes stool color to change to yellow?</td>\n",
       "      <td>What can cause stool to come out as little balls?</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>What can one do after MBBS?</td>\n",
       "      <td>What do i do after my MBBS ?</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Where can I find a power outlet for my laptop ...</td>\n",
       "      <td>Would a second airport in Sydney, Australia be...</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                              text_a  \\\n",
       "0  How is the life of a math student? Could you d...   \n",
       "1                How do I control my horny emotions?   \n",
       "2       What causes stool color to change to yellow?   \n",
       "3                        What can one do after MBBS?   \n",
       "4  Where can I find a power outlet for my laptop ...   \n",
       "\n",
       "                                              text_b label  \n",
       "0  Which level of prepration is enough for the ex...     0  \n",
       "1                 How do you control your horniness?     1  \n",
       "2  What can cause stool to come out as little balls?     0  \n",
       "3                       What do i do after my MBBS ?     1  \n",
       "4  Would a second airport in Sydney, Australia be...     0  "
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\"\"\"\n",
    "QQP train data size: 363849 \n",
    "QQP dev data size: 40430 \n",
    "\"\"\"\n",
    "\n",
    "DATADIR = './glue_examples/glue_data'\n",
    "\n",
    "def get_quora_df(filename):\n",
    "    rows = read_tsv(filename)\n",
    "    df=pd.DataFrame(rows[1:], columns=rows[0])\n",
    "    df=df[['question1', 'question2', 'is_duplicate']]\n",
    "    df = df[pd.notnull(df['is_duplicate'])]\n",
    "    df.columns=['text_a', 'text_b', 'label']\n",
    "    return df\n",
    "\n",
    "def get_quora_data(train_file=DATADIR+'/QQP/train.tsv', \n",
    "                   dev_file=DATADIR+'/QQP/dev.tsv'):\n",
    "    train = get_quora_df(train_file)\n",
    "    print(\"QQP train data size: %d \"%(len(train)))\n",
    "    dev = get_quora_df(dev_file)\n",
    "    print(\"QQP dev data size: %d \"%(len(dev)))\n",
    "\n",
    "    label_list = np.unique(train['label'].values)\n",
    "    return train, dev, label_list\n",
    "\n",
    "train, dev, label_list = get_quora_data()\n",
    "train.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## setup data\n",
    "\n",
    "We will subsample the data for the demo. See the [QQP.ipynb](https://github.com/charles9n/bert-sklearn/blob/master/glue_examples/QQP.ipynb) notebook for a finetune demo on the full data set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# subsample data \n",
    "n = 1000\n",
    "train = train.sample(n, random_state=42)\n",
    "dev = dev.sample(n, random_state=42)\n",
    "\n",
    "X_train = train[['text_a', 'text_b']]\n",
    "y_train = train['label']\n",
    "\n",
    "# use the dev set for testing...\n",
    "test = dev\n",
    "X_test = test[['text_a', 'text_b']]\n",
    "y_test = test['label']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## finetune"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Building sklearn text classifier...\n",
      "Loading bert-base-uncased model...\n",
      "Defaulting to linear classifier/regressor\n",
      "Loading Pytorch checkpoint\n",
      "train data size: 900, validation data size: 100\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Training  : 100%|██████████| 57/57 [00:15<00:00,  4.18it/s, loss=0.643]\n",
      "Validating: 100%|██████████| 13/13 [00:00<00:00, 21.79it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1, Train loss: 0.6428, Val loss: 0.5923, Val accy: 65.00%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Training  : 100%|██████████| 57/57 [00:16<00:00,  4.10it/s, loss=0.424]\n",
      "Validating: 100%|██████████| 13/13 [00:00<00:00, 21.50it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 2, Train loss: 0.4240, Val loss: 0.6151, Val accy: 64.00%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Training  : 100%|██████████| 57/57 [00:17<00:00,  3.72it/s, loss=0.225]\n",
      "Validating: 100%|██████████| 13/13 [00:00<00:00, 14.36it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 3, Train loss: 0.2246, Val loss: 0.7069, Val accy: 64.00%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Testing: 100%|██████████| 125/125 [00:05<00:00, 22.48it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Loss: 0.5384, Accuracy: 74.90%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Predicting: 100%|██████████| 125/125 [00:05<00:00, 21.51it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy: 74.90%\n",
      "               precision    recall  f1-score   support\n",
      "\n",
      "not duplicate       0.85      0.72      0.78       617\n",
      " is duplicate       0.64      0.80      0.71       383\n",
      "\n",
      "    micro avg       0.75      0.75      0.75      1000\n",
      "    macro avg       0.74      0.76      0.74      1000\n",
      " weighted avg       0.77      0.75      0.75      1000\n",
      "\n",
      "CPU times: user 44 s, sys: 21.1 s, total: 1min 5s\n",
      "Wall time: 1min 6s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "# define model\n",
    "model = BertClassifier(max_seq_length=64, train_batch_size=16)\n",
    "\n",
    "# finetune model\n",
    "model.fit(X_train, y_train)\n",
    "\n",
    "# score model\n",
    "model.score(X_test, y_test)\n",
    "\n",
    "# make predictions\n",
    "y_pred = model.predict(X_test)\n",
    "print(\"Accuracy: %0.2f%%\"%(metrics.accuracy_score(y_pred, y_test) * 100))\n",
    "\n",
    "target_names = ['not duplicate', 'is duplicate']\n",
    "print(classification_report(y_test, y_pred, target_names=target_names))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='text_pair_regression'></a>\n",
    "\n",
    "# text pair regression  \n",
    "\n",
    "For text pair regression, we have input data `X`, and target data `y` where :\n",
    "\n",
    "* `X` is a list, pandas dataframe, or numpy array of text pairs (`text_a`, `text_b`) .\n",
    "\n",
    "\n",
    "* `y` is a list, pandas Series, or numpy array of floats.\n",
    "\n",
    "\n",
    "For this example, we will use the **`STS-B`** data set from [GLUE benchmarks](https://gluebenchmark.com/). The data consists of sentence pairs drawn from news headlines and image captions with annotated similarity scores ranging from 1 to 5.\n",
    "\n",
    "See [website](http://ixa2.si.ehu.es/stswiki/index.php/STSbenchmark) and [paper](http://www.aclweb.org/anthology/S/S17/S17-2001.pdf) for more info.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### STS-B"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading and extracting STS...\n",
      "\tCompleted!\n"
     ]
    }
   ],
   "source": [
    "%%bash\n",
    "python3 ./glue_examples/download_glue_data.py --data_dir ./glue_examples//glue_data --tasks STS "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "STS-B train data size: 5749 \n",
      "STS-B dev data size: 1500 \n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>text_a</th>\n",
       "      <th>text_b</th>\n",
       "      <th>label</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>A plane is taking off.</td>\n",
       "      <td>An air plane is taking off.</td>\n",
       "      <td>5.00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>A man is playing a large flute.</td>\n",
       "      <td>A man is playing a flute.</td>\n",
       "      <td>3.80</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>A man is spreading shreded cheese on a pizza.</td>\n",
       "      <td>A man is spreading shredded cheese on an uncoo...</td>\n",
       "      <td>3.80</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Three men are playing chess.</td>\n",
       "      <td>Two men are playing chess.</td>\n",
       "      <td>2.60</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>A man is playing the cello.</td>\n",
       "      <td>A man seated is playing the cello.</td>\n",
       "      <td>4.25</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                          text_a  \\\n",
       "0                         A plane is taking off.   \n",
       "1                A man is playing a large flute.   \n",
       "2  A man is spreading shreded cheese on a pizza.   \n",
       "3                   Three men are playing chess.   \n",
       "4                    A man is playing the cello.   \n",
       "\n",
       "                                              text_b  label  \n",
       "0                        An air plane is taking off.   5.00  \n",
       "1                          A man is playing a flute.   3.80  \n",
       "2  A man is spreading shredded cheese on an uncoo...   3.80  \n",
       "3                         Two men are playing chess.   2.60  \n",
       "4                 A man seated is playing the cello.   4.25  "
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\"\"\"\n",
    "STS-B train data size: 5749 \n",
    "STS-B dev data size: 1500 \n",
    "\"\"\"\n",
    "\n",
    "DATADIR = './glue_examples/glue_data'\n",
    "\n",
    "def get_sts_b_df(filename):\n",
    "    rows = read_tsv(filename)\n",
    "    df=pd.DataFrame(rows[1:], columns=rows[0])\n",
    "    df=df[['sentence1', 'sentence2', 'score']]    \n",
    "    df.columns=['text_a', 'text_b', 'label']\n",
    "    df.label = pd.to_numeric(df.label)\n",
    "    df = df[pd.notnull(df['label'])]                \n",
    "    return df\n",
    "\n",
    "def get_sts_b_data(train_file=DATADIR + '/STS-B/train.tsv',\n",
    "                   dev_file=DATADIR + '/STS-B/dev.tsv'):\n",
    "    train = get_sts_b_df(train_file)\n",
    "    print(\"STS-B train data size: %d \"%(len(train)))    \n",
    "    dev   = get_sts_b_df(dev_file)\n",
    "    print(\"STS-B dev data size: %d \"%(len(dev)))  \n",
    "    return train,dev\n",
    "\n",
    "train, dev = get_sts_b_data()\n",
    "train.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## setup data\n",
    "\n",
    "We will subsample the data for the demo. See the [STS-B.ipynb](https://github.com/charles9n/bert-sklearn/blob/master/glue_examples/STS-B.ipynb) notebook for a finetune demo on the full data set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "# subsample data\n",
    "n = 1000\n",
    "train = train.sample(n, random_state=42)\n",
    "dev = dev.sample(n, random_state=42)\n",
    "\n",
    "X_train = train[['text_a', 'text_b']]\n",
    "y_train = train['label']\n",
    "\n",
    "# use the dev set for testing...\n",
    "test = dev\n",
    "X_test = test[['text_a', 'text_b']]\n",
    "y_test = test['label']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## finetune\n",
    "\n",
    "* For regression, we will report validation accuracy as pearson correlation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Building sklearn text regressor...\n",
      "Loading bert-base-uncased model...\n",
      "Defaulting to linear classifier/regressor\n",
      "Loading Pytorch checkpoint\n",
      "train data size: 900, validation data size: 100\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Training  : 100%|██████████| 29/29 [00:12<00:00,  2.90it/s, loss=3.03]\n",
      "Validating: 100%|██████████| 13/13 [00:00<00:00, 21.43it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1, Train loss: 3.0313, Val loss: 1.5151, Val accy: 71.77%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Training  : 100%|██████████| 29/29 [00:12<00:00,  2.87it/s, loss=0.713]\n",
      "Validating: 100%|██████████| 13/13 [00:00<00:00, 21.33it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 2, Train loss: 0.7135, Val loss: 1.0438, Val accy: 81.48%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Training  : 100%|██████████| 29/29 [00:12<00:00,  2.86it/s, loss=0.333]\n",
      "Validating: 100%|██████████| 13/13 [00:00<00:00, 21.54it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 3, Train loss: 0.3328, Val loss: 0.7230, Val accy: 82.61%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Testing: 100%|██████████| 125/125 [00:04<00:00, 27.48it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Loss: 0.5683, Accuracy: 86.77%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Predicting: 100%|██████████| 125/125 [00:04<00:00, 28.98it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Pearson : 86.77\n",
      "CPU times: user 32.6 s, sys: 17.1 s, total: 49.7 s\n",
      "Wall time: 51.2 s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "from scipy.stats import pearsonr\n",
    "\n",
    "# define model\n",
    "model = BertRegressor()\n",
    "model.max_seq_length = 64\n",
    "\n",
    "# finetune model\n",
    "model.fit(X_train, y_train)\n",
    "\n",
    "# score model\n",
    "model.score(X_test, y_test)\n",
    "\n",
    "# make predictions\n",
    "y_pred = model.predict(X_test)\n",
    "pearson_accy = pearsonr(y_pred, y_test)[0] * 100\n",
    "print(\"Pearson : %0.2f\"%(pearson_accy))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='ner_conll_eng'></a>\n",
    "\n",
    "\n",
    "## Named Entity Recognition (NER) :  CoNLL 2003 task\n",
    "\n",
    "\n",
    "For this example, we will use the  **`CoNLL 2003`** shared task which consists of data from the Reuters 1996 news corpus with annotations for 4 types of `Named Entities` (persons, locations, organizations, and miscellaneous entities). The data is in a [IOB2](https://en.wikipedia.org/wiki/Inside%E2%80%93outside%E2%80%93beginning_(tagging)) format. Each token enitity has a `'B-'` or `'I-'` tag indicating if its the start of the entity or if the token is inside the annotation. The `'O'` tag means the token is not a named entity.\n",
    "\n",
    "* **`Person`**: `'B-PER'` and  `'I-PER'`\n",
    "\n",
    "\n",
    "* **`Organization`**: `'B-ORG'` and `'I-ORG'`\n",
    "\n",
    "\n",
    "* **`Location`**: `'B-LOC'`  and `'I-LOC'`\n",
    "\n",
    "\n",
    "* **`Miscellaneous`**: `'B-MISC'` and `'I-MISC'`\n",
    "\n",
    "\n",
    "* **`Other(non-named entity)`**: `'O'`\n",
    "\n",
    "See [website](https://www.clips.uantwerpen.be/conll2003/ner/) and [paper](https://www.clips.uantwerpen.be/conll2003/pdf/14247tjo.pdf) for more info.\n",
    "\n",
    "The data is already tokenized and tagged:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# tokens: EU     rejects  German  call  to  boycott  British  lamb . \n",
    "# tags  : B-ORG  O        B-MISC  O     O   O        B-MISC   O    O"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So for the named entity recognition (NER) task the data consists of features:`X`and labels:`y`\n",
    "\n",
    "* **`X`** :  a list of list of tokens \n",
    "\n",
    "\n",
    "* **`y`** :  a list of list of NER tags\n",
    "\n",
    "\n",
    "### get data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "cd other_examples\n",
    "DATADIR=\"ner_english\"\n",
    "if test ! -d \"$DATADIR\";then\n",
    "    echo \"Creating $DATADIR dir\"\n",
    "    mkdir \"$DATADIR\"\n",
    "    cd \"$DATADIR\"\n",
    "    wget https://raw.githubusercontent.com/mxhofer/Named-Entity-Recognition-BidirectionalLSTM-CNN-CoNLL/master/data/train.txt\n",
    "    wget https://raw.githubusercontent.com/mxhofer/Named-Entity-Recognition-BidirectionalLSTM-CNN-CoNLL/master/data/test.txt\n",
    "    wget https://raw.githubusercontent.com/mxhofer/Named-Entity-Recognition-BidirectionalLSTM-CNN-CoNLL/master/data/dev.txt\n",
    "fi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train data: 14987 sentences, 204567 tokens\n",
      "Dev data: 3466 sentences, 51578 tokens\n",
      "Test data: 3684 sentences, 46666 tokens\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>tokens</th>\n",
       "      <th>labels</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>[-DOCSTART-]</td>\n",
       "      <td>[O]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>[EU, rejects, German, call, to, boycott, Briti...</td>\n",
       "      <td>[B-ORG, O, B-MISC, O, O, O, B-MISC, O, O]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>[Peter, Blackburn]</td>\n",
       "      <td>[B-PER, I-PER]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>[BRUSSELS, 1996-08-22]</td>\n",
       "      <td>[B-LOC, O]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>[The, European, Commission, said, on, Thursday...</td>\n",
       "      <td>[O, B-ORG, I-ORG, O, O, O, O, O, O, B-MISC, O,...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                              tokens  \\\n",
       "0                                       [-DOCSTART-]   \n",
       "1  [EU, rejects, German, call, to, boycott, Briti...   \n",
       "2                                 [Peter, Blackburn]   \n",
       "3                             [BRUSSELS, 1996-08-22]   \n",
       "4  [The, European, Commission, said, on, Thursday...   \n",
       "\n",
       "                                              labels  \n",
       "0                                                [O]  \n",
       "1          [B-ORG, O, B-MISC, O, O, O, B-MISC, O, O]  \n",
       "2                                     [B-PER, I-PER]  \n",
       "3                                         [B-LOC, O]  \n",
       "4  [O, B-ORG, I-ORG, O, O, O, O, O, O, B-MISC, O,...  "
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\"\"\"\n",
    "Train data: 14987 sentences, 204567 tokens\n",
    "Dev data: 3466 sentences, 51578 tokens\n",
    "Test data: 3684 sentences, 46666 tokens\n",
    "\"\"\"\n",
    "\n",
    "DATADIR = \"./other_examples/ner_english/\"\n",
    "\n",
    "def get_conll2003_data(trainfile=DATADIR + \"train.txt\",\n",
    "                       devfile=DATADIR + \"dev.txt\",\n",
    "                       testfile=DATADIR + \"test.txt\"):\n",
    "\n",
    "    train = read_CoNLL2003_format(trainfile)\n",
    "    print(\"Train data: %d sentences, %d tokens\"%(len(train), len(flatten(train.tokens))))\n",
    "    dev = read_CoNLL2003_format(devfile)\n",
    "    print(\"Dev data: %d sentences, %d tokens\"%(len(dev), len(flatten(dev.tokens))))\n",
    "    test = read_CoNLL2003_format(testfile)\n",
    "    print(\"Test data: %d sentences, %d tokens\"%(len(test), len(flatten(test.tokens))))\n",
    "    \n",
    "    return train, dev, test\n",
    "\n",
    "\n",
    "train, dev, test = get_conll2003_data()\n",
    "train.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## setup data\n",
    "\n",
    "We will subsample the data for the demo. See the [ner_english.ipynb](https://github.com/charles9n/bert-sklearn/blob/master/other_examples/ner_english.ipynb) notebook for finetune demo on the full data set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "NER tags: ['B-LOC', 'B-MISC', 'B-ORG', 'B-PER', 'I-LOC', 'I-MISC', 'I-ORG', 'I-PER', 'O']\n"
     ]
    }
   ],
   "source": [
    "X_train, y_train = train.tokens, train.labels\n",
    "X_dev, y_dev = dev.tokens, dev.labels\n",
    "X_test, y_test = test.tokens, test.labels\n",
    "\n",
    "label_list = np.unique(flatten(y_train))\n",
    "label_list = list(label_list)\n",
    "print(\"\\nNER tags:\",label_list)\n",
    "\n",
    "# take a subset of the data for demo\n",
    "n = 1000\n",
    "X_train, y_train = X_train[:n], y_train[:n]\n",
    "X_test, y_test = X_test[:n], y_test[:n]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## finetune \n",
    "\n",
    "\n",
    "Let's define our model using the **`BertTokenClassifier`** class\n",
    "\n",
    "* We will include an **`ignore_label`** option to exclude the `'O'` labels, i.e non named entities label, to calculate  `f1`. The non named entities are a huge majority of the labels, and typically `f1` is reported with non named entities excluded.\n",
    "\n",
    "\n",
    "* We will also use the cased model,`'bert-base-cased'`, as casing provides an important signal for NER. The first time you run this it will take a little longer to download the model into the cache.\n",
    "\n",
    "\n",
    "* With the BertTokenClassifier we should also be mindful to set the **` max_seq_len`**  high enough to cover lengths of the token lists. See the [ner_english.ipynb](https://github.com/charles9n/bert-sklearn/blob/master/other_examples/ner_english.ipynb) for more detail.\n",
    "\n",
    "This uses around 7GB GPU memory on my laptop. If this gives you OOM, then set the  **` max_seq_len`** to a lower number, i.e 128 or 96. You can also increase **`gradient_accumulation_steps`** up to 8 to further reduce the GPU memory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Building sklearn token classifier...\n",
      "BertTokenClassifier(bert_config_json=None, bert_model='bert-base-cased',\n",
      "          bert_vocab=None, do_lower_case=None, epochs=3,\n",
      "          eval_batch_size=16, fp16=False, from_tf=False,\n",
      "          gradient_accumulation_steps=2, ignore_label=['O'],\n",
      "          label_list=None, learning_rate=2e-05, local_rank=-1,\n",
      "          logfile='bert_sklearn.log', loss_scale=0, max_seq_length=173,\n",
      "          num_mlp_hiddens=500, num_mlp_layers=0, random_state=42,\n",
      "          restore_file=None, train_batch_size=16, use_cuda=True,\n",
      "          validation_fraction=0.1, warmup_proportion=0.1)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 213450/213450 [00:00<00:00, 891015.14B/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading bert-base-cased model...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 435779157/435779157 [02:58<00:00, 2437434.04B/s] \n",
      "100%|██████████| 313/313 [00:00<00:00, 62271.95B/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Defaulting to linear classifier/regressor\n",
      "Loading Pytorch checkpoint\n",
      "train data size: 900, validation data size: 100\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Training  : 100%|██████████| 113/113 [00:33<00:00,  3.91it/s, loss=0.0501]\n",
      "Validating: 100%|██████████| 7/7 [00:01<00:00,  5.63it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1, Train loss: 0.0501, Val loss: 0.0102, Val accy: 95.16%, f1: 83.12\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Training  : 100%|██████████| 113/113 [00:33<00:00,  3.87it/s, loss=0.00848]\n",
      "Validating: 100%|██████████| 7/7 [00:01<00:00,  5.60it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 2, Train loss: 0.0085, Val loss: 0.0071, Val accy: 96.93%, f1: 90.42\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Training  : 100%|██████████| 113/113 [00:34<00:00,  3.86it/s, loss=0.00417]\n",
      "Validating: 100%|██████████| 7/7 [00:01<00:00,  5.56it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 3, Train loss: 0.0042, Val loss: 0.0055, Val accy: 97.30%, f1: 91.74\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Predicting: 100%|██████████| 63/63 [00:11<00:00,  5.42it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test f1: 87.65\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Predicting: 100%|██████████| 63/63 [00:12<00:00,  5.38it/s]\n",
      "Predicting: 100%|██████████| 63/63 [00:13<00:00,  5.25it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "              precision    recall  f1-score   support\n",
      "\n",
      "       B-LOC       0.84      0.86      0.85       418\n",
      "      B-MISC       0.67      0.69      0.68       189\n",
      "       B-ORG       0.84      0.82      0.83       489\n",
      "       B-PER       0.99      0.92      0.95       655\n",
      "       I-LOC       0.86      0.55      0.67        66\n",
      "      I-MISC       0.59      0.72      0.65        82\n",
      "       I-ORG       0.87      0.88      0.87       205\n",
      "       I-PER       0.99      1.00      0.99       483\n",
      "           O       0.99      0.99      0.99      8479\n",
      "\n",
      "   micro avg       0.96      0.96      0.96     11066\n",
      "   macro avg       0.85      0.83      0.83     11066\n",
      "weighted avg       0.96      0.96      0.96     11066\n",
      "\n",
      "CPU times: user 1min 53s, sys: 1min 1s, total: 2min 55s\n",
      "Wall time: 5min 27s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "# define model\n",
    "model = BertTokenClassifier(bert_model='bert-base-cased',\n",
    "                            epochs=3,\n",
    "                            max_seq_length=173,\n",
    "                            learning_rate=2e-5,\n",
    "                            train_batch_size=16,\n",
    "                            eval_batch_size=16,\n",
    "                            gradient_accumulation_steps=2,\n",
    "                            ignore_label=['O'])\n",
    "\n",
    "\n",
    "print(model)\n",
    "\n",
    "# finetune model\n",
    "model.fit(X_train, y_train)\n",
    "\n",
    "# score model\n",
    "f1_test = model.score(X_test, y_test)\n",
    "print(\"Test f1: %0.02f\"%(f1_test))\n",
    "\n",
    "# make predictions\n",
    "y_preds = model.predict(X_test)\n",
    "\n",
    "# calculate the probability of each class\n",
    "y_probs = model.predict_proba(X_test)\n",
    "\n",
    "print(classification_report(flatten(y_test), flatten(y_preds)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### span level stats on test data\n",
    "\n",
    "If we want span level stats, we can run the [perl script](https://www.clips.uantwerpen.be/conll2003/ner/bin/conlleval) from the original `CoNLL-2000/2003 shared task`, to evaluate the results:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "processed 11066 tokens with 1751 phrases; found: 1765 phrases; correct: 1467.\r\n",
      "accuracy:  96.34%; precision:  83.12%; recall:  83.78%; FB1:  83.45\r\n",
      "              LOC: precision:  79.27%; recall:  83.25%; FB1:  81.21  439\r\n",
      "             MISC: precision:  55.96%; recall:  64.55%; FB1:  59.95  218\r\n",
      "              ORG: precision:  79.31%; recall:  79.96%; FB1:  79.63  493\r\n",
      "              PER: precision:  98.54%; recall:  92.52%; FB1:  95.43  615\r\n"
     ]
    }
   ],
   "source": [
    "# write out predictions to file for conlleval.pl\n",
    "iter_zip = zip(flatten(X_test),flatten(y_test),flatten(y_preds))\n",
    "preds = [\" \".join([token, y, y_pred]) for token, y, y_pred in iter_zip]\n",
    "with open(\"preds.txt\",'w') as f:\n",
    "    for x in preds:\n",
    "        f.write(str(x)+'\\n') \n",
    "\n",
    "# run conlleval perl script \n",
    "!perl ./other_examples/conlleval.pl < preds.txt\n",
    "!rm preds.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### check results on test data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "         token   label predict\n",
      "0        Dutch  B-MISC  B-MISC\n",
      "1      forward       O       O\n",
      "2       Reggie   B-PER   B-PER\n",
      "3      Blinker   I-PER   I-PER\n",
      "4          had       O       O\n",
      "5          his       O       O\n",
      "6   indefinite       O       O\n",
      "7   suspension       O       O\n",
      "8       lifted       O       O\n",
      "9           by       O       O\n",
      "10        FIFA   B-ORG   B-ORG\n",
      "11          on       O       O\n",
      "12      Friday       O       O\n",
      "13         and       O       O\n",
      "14         was       O       O\n",
      "15         set       O       O\n",
      "16          to       O       O\n",
      "17        make       O       O\n",
      "18         his       O       O\n",
      "19   Sheffield   B-ORG   B-LOC\n",
      "20   Wednesday   I-ORG   I-ORG\n",
      "21    comeback       O       O\n",
      "22     against       O       O\n",
      "23   Liverpool   B-ORG   B-ORG\n",
      "24          on       O       O\n",
      "25    Saturday       O       O\n",
      "26           .       O       O\n"
     ]
    }
   ],
   "source": [
    "i = 152\n",
    "tokens = X_test[i]\n",
    "labels = y_test[i]\n",
    "preds = y_preds[i]\n",
    "data = {\"token\": tokens,\"label\": labels,\"predict\": preds}\n",
    "df=pd.DataFrame(data=data)\n",
    "print(df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "         token  B-LOC  B-MISC  B-ORG  B-PER  I-LOC  I-MISC  I-ORG  I-PER    O\n",
      "0        Dutch   0.11    0.81   0.06   0.01   0.00    0.00   0.00   0.00 0.00\n",
      "1      forward   0.00    0.00   0.00   0.00   0.00    0.00   0.00   0.00 1.00\n",
      "2       Reggie   0.00    0.00   0.00   1.00   0.00    0.00   0.00   0.00 0.00\n",
      "3      Blinker   0.00    0.00   0.00   0.00   0.00    0.00   0.00   0.99 0.00\n",
      "4          had   0.00    0.00   0.00   0.00   0.00    0.00   0.00   0.00 1.00\n",
      "5          his   0.00    0.00   0.00   0.00   0.00    0.00   0.00   0.00 1.00\n",
      "6   indefinite   0.00    0.00   0.00   0.00   0.00    0.00   0.00   0.00 1.00\n",
      "7   suspension   0.00    0.00   0.00   0.00   0.00    0.00   0.00   0.00 1.00\n",
      "8       lifted   0.00    0.00   0.00   0.00   0.00    0.00   0.00   0.00 1.00\n",
      "9           by   0.00    0.00   0.00   0.00   0.00    0.00   0.00   0.00 1.00\n",
      "10        FIFA   0.30    0.07   0.60   0.01   0.01    0.00   0.01   0.00 0.00\n",
      "11          on   0.00    0.00   0.00   0.00   0.00    0.00   0.00   0.00 1.00\n",
      "12      Friday   0.00    0.00   0.00   0.00   0.00    0.00   0.00   0.00 1.00\n",
      "13         and   0.00    0.00   0.00   0.00   0.00    0.00   0.00   0.00 1.00\n",
      "14         was   0.00    0.00   0.00   0.00   0.00    0.00   0.00   0.00 1.00\n",
      "15         set   0.00    0.00   0.00   0.00   0.00    0.00   0.00   0.00 1.00\n",
      "16          to   0.00    0.00   0.00   0.00   0.00    0.00   0.00   0.00 1.00\n",
      "17        make   0.00    0.00   0.00   0.00   0.00    0.00   0.00   0.00 1.00\n",
      "18         his   0.00    0.00   0.00   0.00   0.00    0.00   0.00   0.00 1.00\n",
      "19   Sheffield   0.43    0.13   0.41   0.01   0.01    0.01   0.00   0.00 0.00\n",
      "20   Wednesday   0.01    0.01   0.01   0.01   0.28    0.16   0.51   0.01 0.00\n",
      "21    comeback   0.00    0.00   0.00   0.00   0.00    0.00   0.00   0.00 1.00\n",
      "22     against   0.00    0.00   0.00   0.00   0.00    0.00   0.00   0.00 1.00\n",
      "23   Liverpool   0.29    0.01   0.69   0.00   0.00    0.00   0.00   0.00 0.00\n",
      "24          on   0.00    0.00   0.00   0.00   0.00    0.00   0.00   0.00 1.00\n",
      "25    Saturday   0.00    0.00   0.00   0.00   0.00    0.00   0.00   0.00 1.00\n",
      "26           .   0.00    0.00   0.00   0.00   0.00    0.00   0.00   0.00 1.00\n"
     ]
    }
   ],
   "source": [
    "# pprint out probs for this observation\n",
    "prob = y_probs[i]\n",
    "tokens_prob = model.tokens_proba(tokens, prob)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, lets predict the tags and tag probabilities on some new text:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Predicting: 100%|██████████| 1/1 [00:00<00:00,  6.48it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "       token predicted tags\n",
      "0  Jefferson          B-PER\n",
      "1      wants              O\n",
      "2         to              O\n",
      "3         go              O\n",
      "4         to              O\n",
      "5     France          B-LOC\n",
      "6          .              O\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Predicting: 100%|██████████| 1/1 [00:00<00:00,  6.75it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "       token  B-LOC  B-MISC  B-ORG  B-PER  I-LOC  I-MISC  I-ORG  I-PER    O\n",
      "0  Jefferson   0.00    0.01   0.03   0.95   0.00    0.00   0.00   0.00 0.00\n",
      "1      wants   0.00    0.00   0.00   0.00   0.00    0.00   0.00   0.00 1.00\n",
      "2         to   0.00    0.00   0.00   0.00   0.00    0.00   0.00   0.00 1.00\n",
      "3         go   0.00    0.00   0.00   0.00   0.00    0.00   0.00   0.00 1.00\n",
      "4         to   0.00    0.00   0.00   0.00   0.00    0.00   0.00   0.00 1.00\n",
      "5     France   0.99    0.00   0.00   0.00   0.00    0.00   0.00   0.00 0.00\n",
      "6          .   0.00    0.00   0.00   0.00   0.00    0.00   0.00   0.00 1.00\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "text = \"Jefferson wants to go to France.\"       \n",
    "\n",
    "tag_predicts  = model.tag_text(text)       \n",
    "prob_predicts = model.tag_text_proba(text)    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='ner_ncbi'></a>\n",
    "\n",
    "\n",
    "##  Biomedical NER : NCBI Disease Corpus\n",
    "\n",
    "\n",
    "## NCBI Disease Corpus\n",
    "\n",
    "The NCBI disease corpus  task is a Named Entity Recognition(NER) task in the biomedical domain. The data is from  a collection of 793 PubMed abstracts with annotations for disease entities. Each token enitity has a `'B-'` or `'I-'` tags indicating if it is the start of the entity or if the token is inside the annotation. The `'O'` tag means the token is not a named entity. See this [paper](https://www.ncbi.nlm.nih.gov/pmc/articles/PMC3951655/) for more information\n",
    "\n",
    "\n",
    "NER tasks are token classification tasks where the data consists of features,`X`, and labels,`y`, where:\n",
    "\n",
    "* **`X`** :  is a list of list of tokens \n",
    "\n",
    "\n",
    "* **`y`** :  is a list of list of NER tags\n",
    "\n",
    "\n",
    "We will finetune models from:\n",
    "\n",
    "\n",
    "* [**BERT**](#NCBI_BERT) - this is the standard `BERT` base case model from Google  pretrained on the Books Corpus and English Wikipedia data.\n",
    "\n",
    "\n",
    "\n",
    "* [**SciBERT**](#NCBI_SciBERT) - `SciBERT` is a model from [AllenAI](https://allenai.org/) based on `BERT` but pretrained on large scientific text corpus.  For more information on `SciBERT`, see the [github repo](https://github.com/allenai/scibert) and [paper](https://arxiv.org/pdf/1903.10676.pdf).\n",
    "\n",
    "\n",
    "\n",
    "* [**BioBERT**](#NCBI_BioBERT) -  `BioBERT` is a model also based on `BERT` but pretrained on a large biomedical text corpus.  For more information on `BioBERT`, see the [ github repo](https://github.com/dmis-lab/biobert) and [paper](https://arxiv.org/pdf/1901.08746.pdf).\n",
    "\n",
    "\n",
    "For this demo we will use a subsample of the the data. See [ner_NCBI_disease_BioBERT_SciBERT](https://github.com/charles9n/bert-sklearn/blob/master/other_examples/ner_NCBI_disease_BioBERT_SciBERT.ipynb) for a finetune demo using all the training data.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### get data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "DATADIR=\"other_examples/NCBI_disease\"\n",
    "if test ! -d \"$DATADIR\";then\n",
    "    echo \"Creating $DATADIR dir\"\n",
    "    mkdir \"$DATADIR\"\n",
    "    cd \"$DATADIR\"\n",
    "    wget https://raw.githubusercontent.com/allenai/scibert/master/data/ner/NCBI-disease/dev.txt\n",
    "    wget https://raw.githubusercontent.com/allenai/scibert/master/data/ner/NCBI-disease/test.txt\n",
    "    wget https://raw.githubusercontent.com/allenai/scibert/master/data/ner/NCBI-disease/train.txt\n",
    "fi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train and dev data: 6347 sentences, 159670 tokens\n",
      "Test data: 940 sentences, 24497 tokens\n",
      "\n",
      "NER tags: ['B-Disease', 'I-Disease', 'O']\n"
     ]
    }
   ],
   "source": [
    "\"\"\"\n",
    "Train and dev data: 6347 sentences, 159670 tokens\n",
    "Test data: 940 sentences, 24497 tokens\n",
    "\"\"\"\n",
    "DATADIR = \"other_examples/NCBI_disease/\"\n",
    "\n",
    "def get_data(trainfile=DATADIR + \"train.txt\",\n",
    "             devfile=DATADIR + \"dev.txt\",\n",
    "             testfile=DATADIR + \"test.txt\"):\n",
    "\n",
    "    train = read_CoNLL2003_format(trainfile, idx=3)    \n",
    "    dev = read_CoNLL2003_format(devfile, idx=3)\n",
    "    \n",
    "    # combine train and dev\n",
    "    train = pd.concat([train, dev])\n",
    "    print(\"Train and dev data: %d sentences, %d tokens\"%(len(train), len(flatten(train.tokens))))\n",
    "\n",
    "    test = read_CoNLL2003_format(testfile, idx=3)\n",
    "    print(\"Test data: %d sentences, %d tokens\"%(len(test), len(flatten(test.tokens))))\n",
    "    \n",
    "    return train, test\n",
    "\n",
    "train, test = get_data()\n",
    "\n",
    "\n",
    "X_train, y_train = train.tokens, train.labels\n",
    "\n",
    "# subsample the data for the demo\n",
    "n=600\n",
    "X_train = X_train[:n]\n",
    "y_train = y_train[:n]\n",
    "\n",
    "X_test, y_test = test.tokens, test.labels\n",
    "\n",
    "label_list = np.unique(flatten(y_train))\n",
    "label_list = list(label_list)\n",
    "print(\"\\nNER tags:\",label_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>tokens</th>\n",
       "      <th>labels</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>[Identification, of, APC2, ,, a, homologue, of...</td>\n",
       "      <td>[O, O, O, O, O, O, O, O, B-Disease, I-Disease,...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>[The, adenomatous, polyposis, coli, (, APC, ),...</td>\n",
       "      <td>[O, B-Disease, I-Disease, I-Disease, I-Disease...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>[Complex, formation, induces, the, rapid, degr...</td>\n",
       "      <td>[O, O, O, O, O, O, O, O, O]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>[In, colon, carcinoma, cells, ,, loss, of, APC...</td>\n",
       "      <td>[O, B-Disease, I-Disease, O, O, O, O, O, O, O,...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>[Here, ,, we, report, the, identification, and...</td>\n",
       "      <td>[O, O, O, O, O, O, O, O, O, O, O, O, O]</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                              tokens  \\\n",
       "0  [Identification, of, APC2, ,, a, homologue, of...   \n",
       "1  [The, adenomatous, polyposis, coli, (, APC, ),...   \n",
       "2  [Complex, formation, induces, the, rapid, degr...   \n",
       "3  [In, colon, carcinoma, cells, ,, loss, of, APC...   \n",
       "4  [Here, ,, we, report, the, identification, and...   \n",
       "\n",
       "                                              labels  \n",
       "0  [O, O, O, O, O, O, O, O, B-Disease, I-Disease,...  \n",
       "1  [O, B-Disease, I-Disease, I-Disease, I-Disease...  \n",
       "2                        [O, O, O, O, O, O, O, O, O]  \n",
       "3  [O, B-Disease, I-Disease, O, O, O, O, O, O, O,...  \n",
       "4            [O, O, O, O, O, O, O, O, O, O, O, O, O]  "
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "         token      label\n",
      "0   Occasional          O\n",
      "1     missense          O\n",
      "2    mutations          O\n",
      "3           in          O\n",
      "4          ATM          O\n",
      "5         were          O\n",
      "6         also          O\n",
      "7        found          O\n",
      "8           in          O\n",
      "9       tumour  B-Disease\n",
      "10         DNA          O\n",
      "11        from          O\n",
      "12    patients          O\n",
      "13        with          O\n",
      "14           B  B-Disease\n",
      "15           -  I-Disease\n",
      "16        cell  I-Disease\n",
      "17         non  I-Disease\n",
      "18           -  I-Disease\n",
      "19    Hodgkins  I-Disease\n",
      "20   lymphomas  I-Disease\n",
      "21           (          O\n",
      "22           B  B-Disease\n",
      "23           -  I-Disease\n",
      "24         NHL  I-Disease\n",
      "25           )          O\n",
      "26         and          O\n",
      "27           a          O\n",
      "28           B  B-Disease\n",
      "29           -  I-Disease\n",
      "30         NHL  I-Disease\n",
      "31        cell          O\n",
      "32        line          O\n",
      "33           .          O\n"
     ]
    }
   ],
   "source": [
    "i = 9\n",
    "tokens = X_test[i]\n",
    "labels = y_test[i]\n",
    "\n",
    "data = {\"token\": tokens,\"label\": labels}\n",
    "df=pd.DataFrame(data=data)\n",
    "print(df)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='NCBI_BERT'></a>\n",
    "### finetune  BERT\n",
    "\n",
    "As in the CoNLL task, we define our model using the **`BertTokenClassifier`** class\n",
    "\n",
    "* We will include an **`ignore_label`** option to exclude the `'O'`, non named entities label, to calculate  `f1`. The non named entities are a huge majority of the labels, and typically `f1` is reported with non named entities excluded.\n",
    "\n",
    "\n",
    "* With the BertTokenClassifier we should also be mindful to set the **` max_seq_len`**  high enough to cover lengths of the token lists. See the extended demo for more detail.\n",
    "\n",
    "This uses around 7GB GPU memory on my laptop. If this gives you OOM, then set the  **` max_seq_len`** to a lower number, i.e 128 or 96. You can also increase **`gradient_accumulation_steps`** up to 8 to further reduce the GPU memory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Building sklearn token classifier...\n",
      "BertTokenClassifier(bert_config_json=None, bert_model='bert-base-cased',\n",
      "          bert_vocab=None, do_lower_case=None, epochs=3,\n",
      "          eval_batch_size=16, fp16=False, from_tf=False,\n",
      "          gradient_accumulation_steps=2, ignore_label=['O'],\n",
      "          label_list=['B-Disease', 'I-Disease', 'O'], learning_rate=3e-05,\n",
      "          local_rank=-1, logfile='bert_sklearn.log', loss_scale=0,\n",
      "          max_seq_length=178, num_mlp_hiddens=500, num_mlp_layers=0,\n",
      "          random_state=42, restore_file=None, train_batch_size=16,\n",
      "          use_cuda=True, validation_fraction=0.0, warmup_proportion=0.1)\n",
      "Loading bert-base-cased model...\n",
      "Defaulting to linear classifier/regressor\n",
      "Loading Pytorch checkpoint\n",
      "train data size: 600, validation data size: 0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Training  : 100%|██████████| 75/75 [00:23<00:00,  3.20it/s, loss=0.0514]\n",
      "Training  : 100%|██████████| 75/75 [00:25<00:00,  2.95it/s, loss=0.0105]\n",
      "Training  : 100%|██████████| 75/75 [00:29<00:00,  2.57it/s, loss=0.00384]\n",
      "Predicting: 100%|██████████| 59/59 [00:15<00:00,  3.75it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test f1: 76.72\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Predicting: 100%|██████████| 59/59 [00:17<00:00,  3.49it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "              precision    recall  f1-score   support\n",
      "\n",
      "   B-Disease       0.81      0.71      0.75       960\n",
      "   I-Disease       0.90      0.69      0.78      1087\n",
      "           O       0.98      0.99      0.98     22450\n",
      "\n",
      "   micro avg       0.97      0.97      0.97     24497\n",
      "   macro avg       0.89      0.80      0.84     24497\n",
      "weighted avg       0.97      0.97      0.97     24497\n",
      "\n",
      "CPU times: user 1min 10s, sys: 44.3 s, total: 1min 54s\n",
      "Wall time: 1min 55s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "model = BertTokenClassifier('bert-base-cased',\n",
    "                            max_seq_length=178,\n",
    "                            epochs=3,\n",
    "                            gradient_accumulation_steps=2,\n",
    "                            learning_rate=3e-5,\n",
    "                            train_batch_size=16,\n",
    "                            eval_batch_size=16,\n",
    "                            validation_fraction=0.,                            \n",
    "                            label_list=label_list,\n",
    "                            ignore_label=['O'])\n",
    "\n",
    "print(model)\n",
    "\n",
    "# finetune model on train data\n",
    "model.fit(X_train, y_train)\n",
    "\n",
    "# score model on test data\n",
    "f1_test = model.score(X_test, y_test,'macro')\n",
    "print(\"Test f1: %0.02f\"%(f1_test))\n",
    "\n",
    "# get predictions on test data\n",
    "y_preds = model.predict(X_test)\n",
    "\n",
    "# print report on classifier stats\n",
    "print(classification_report(flatten(y_test), flatten(y_preds)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For span level stats let's use the conlleval.pl script:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "processed 24497 tokens with 960 phrases; found: 904 phrases; correct: 631.\r\n",
      "accuracy:  96.72%; precision:  69.80%; recall:  65.73%; FB1:  67.70\r\n",
      "          Disease: precision:  69.80%; recall:  65.73%; FB1:  67.70  904\r\n"
     ]
    }
   ],
   "source": [
    "# write out predictions to file for conlleval.pl\n",
    "iter_zip = zip(flatten(X_test),flatten(y_test),flatten(y_preds))\n",
    "preds = [\" \".join([token, y, y_pred]) for token, y, y_pred in iter_zip]\n",
    "with open(\"preds.txt\",'w') as f:\n",
    "    for x in preds:\n",
    "        f.write(str(x)+'\\n') \n",
    "\n",
    "# run conlleval perl script \n",
    "!perl ./other_examples/conlleval.pl < preds.txt\n",
    "!rm preds.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "get the token predicts for our sample to compare with SCiBERT and BioBERT models later"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get the token predicts \n",
    "i = 9\n",
    "tokens = X_test[i]\n",
    "labels = y_test[i]\n",
    "bert_preds  = y_preds[i]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='NCBI_SciBERT'></a>\n",
    "### finetune SciBERT\n",
    "\n",
    "There are 4 SciBERT models available: \n",
    "\n",
    "\n",
    "* `scibert-scivocab-uncased` \n",
    "\n",
    "\n",
    "* `scibert-scivocab-cased`\n",
    "\n",
    "\n",
    "* `scibert-basevocab-uncased`\n",
    "\n",
    "\n",
    "* `scibert-basevocab-cased`\n",
    "\n",
    "\n",
    "See the [`SciBERT` github](https://github.com/allenai/scibert) and [paper](https://arxiv.org/pdf/1903.10676.pdf) for more info.\n",
    "\n",
    "While the  `scibert-scivocab-uncased` is the recomended model on the website, we will use `'scibert-scivocab-cased'` for the demo to have comparable results with the BERT and BioBERT models. Note as before, the first time you run this will take some extra time to download the model into the file cache:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Building sklearn token classifier...\n",
      "BertTokenClassifier(bert_config_json=None,\n",
      "          bert_model='scibert-scivocab-cased', bert_vocab=None,\n",
      "          do_lower_case=None, epochs=3, eval_batch_size=16, fp16=False,\n",
      "          from_tf=False, gradient_accumulation_steps=2, ignore_label=['O'],\n",
      "          label_list=['B-Disease', 'I-Disease', 'O'], learning_rate=3e-05,\n",
      "          local_rank=-1, logfile='bert_sklearn.log', loss_scale=0,\n",
      "          max_seq_length=178, num_mlp_hiddens=500, num_mlp_layers=0,\n",
      "          random_state=42, restore_file=None, train_batch_size=16,\n",
      "          use_cuda=True, validation_fraction=0.0, warmup_proportion=0.1)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 410521600/410521600 [00:21<00:00, 18680427.63B/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading scibert-scivocab-cased model...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 410521600/410521600 [00:18<00:00, 21932301.40B/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Defaulting to linear classifier/regressor\n",
      "Loading Pytorch checkpoint\n",
      "train data size: 600, validation data size: 0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Training  : 100%|██████████| 75/75 [00:23<00:00,  3.36it/s, loss=0.0342]\n",
      "Training  : 100%|██████████| 75/75 [00:23<00:00,  3.10it/s, loss=0.00651]\n",
      "Training  : 100%|██████████| 75/75 [00:24<00:00,  3.31it/s, loss=0.0025] \n",
      "Predicting: 100%|██████████| 59/59 [00:10<00:00,  6.12it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test f1: 85.33\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Predicting: 100%|██████████| 59/59 [00:11<00:00,  5.23it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "              precision    recall  f1-score   support\n",
      "\n",
      "   B-Disease       0.85      0.85      0.85       960\n",
      "   I-Disease       0.89      0.83      0.86      1087\n",
      "           O       0.99      0.99      0.99     22450\n",
      "\n",
      "   micro avg       0.98      0.98      0.98     24497\n",
      "   macro avg       0.91      0.89      0.90     24497\n",
      "weighted avg       0.98      0.98      0.98     24497\n",
      "\n",
      "CPU times: user 1min 22s, sys: 41.4 s, total: 2min 4s\n",
      "Wall time: 2min 25s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "model = BertTokenClassifier('scibert-scivocab-cased',\n",
    "                            max_seq_length=178,\n",
    "                            epochs=3,\n",
    "                            gradient_accumulation_steps=2,\n",
    "                            learning_rate=3e-5,\n",
    "                            train_batch_size=16,\n",
    "                            eval_batch_size=16,\n",
    "                            validation_fraction=0.,                            \n",
    "                            label_list=label_list,\n",
    "                            ignore_label=['O'])\n",
    "\n",
    "print(model)\n",
    "\n",
    "# finetune model on train data\n",
    "model.fit(X_train, y_train)\n",
    "\n",
    "# score model on test data\n",
    "f1_test = model.score(X_test, y_test,'macro')\n",
    "print(\"Test f1: %0.02f\"%(f1_test))\n",
    "\n",
    "# get predictions on test data\n",
    "y_preds = model.predict(X_test)\n",
    "\n",
    "# print report on classifier stats\n",
    "print(classification_report(flatten(y_test), flatten(y_preds)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "processed 24497 tokens with 960 phrases; found: 1013 phrases; correct: 777.\r\n",
      "accuracy:  97.84%; precision:  76.70%; recall:  80.94%; FB1:  78.76\r\n",
      "          Disease: precision:  76.70%; recall:  80.94%; FB1:  78.76  1013\r\n"
     ]
    }
   ],
   "source": [
    "# write out predictions to file for conlleval.pl\n",
    "iter_zip = zip(flatten(X_test),flatten(y_test),flatten(y_preds))\n",
    "preds = [\" \".join([token, y, y_pred]) for token, y, y_pred in iter_zip]\n",
    "with open(\"preds.txt\",'w') as f:\n",
    "    for x in preds:\n",
    "        f.write(str(x)+'\\n') \n",
    "\n",
    "# run conlleval perl script \n",
    "!perl ./other_examples/conlleval.pl < preds.txt\n",
    "!rm preds.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get the token predicts \n",
    "i = 9\n",
    "tokens = X_test[i]\n",
    "labels = y_test[i]\n",
    "scibert_preds  = y_preds[i]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='NCBI_BioBERT'></a>\n",
    "### finetune BioBERT\n",
    "\n",
    "There are 4 **`BioBERT`** models available:\n",
    "\n",
    "* `'biobert-v1.0-pmc-base-cased'`\n",
    "\n",
    "\n",
    "* `'biobert-v1.0-pubmed-base-cased'`\n",
    "\n",
    "\n",
    "* `'biobert-v1.0-pubmed-pmc-base-cased'` \n",
    "\n",
    "\n",
    "* `'biobert-v1.1-pubmed-base-cased'` \n",
    "\n",
    "See [BioBERT github](https://github.com/dmis-lab/biobert) and [paper](https://arxiv.org/pdf/1901.08746.pdf)  for more info.\n",
    "\n",
    "The **`BioBERT`** models are archived as tensorflow checkpoints. We will use the hugginface pytorch BERT port to convert them to a pytorch model, however tensorflow will need to be installed on the system for the inital conversion. If it is not installed run : \n",
    "\n",
    "`pip install tensorflow-gpu`\n",
    "\n",
    "We will use the most recent release,`biobert-v1.1-pubmed-base-cased`, for the demo:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Building sklearn token classifier...\n",
      "BertTokenClassifier(bert_config_json=None,\n",
      "          bert_model='biobert-v1.1-pubmed-base-cased', bert_vocab=None,\n",
      "          do_lower_case=None, epochs=3, eval_batch_size=16, fp16=False,\n",
      "          from_tf=False, gradient_accumulation_steps=2, ignore_label=['O'],\n",
      "          label_list=['B-Disease', 'I-Disease', 'O'], learning_rate=3e-05,\n",
      "          local_rank=-1, logfile='bert_sklearn.log', loss_scale=0,\n",
      "          max_seq_length=178, num_mlp_hiddens=500, num_mlp_layers=0,\n",
      "          random_state=42, restore_file=None, train_batch_size=16,\n",
      "          use_cuda=True, validation_fraction=0, warmup_proportion=0.1)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 401403346/401403346 [00:27<00:00, 14423439.71B/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading biobert-v1.1-pubmed-base-cased model...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 401403346/401403346 [00:23<00:00, 16765954.84B/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Defaulting to linear classifier/regressor\n",
      "Loading Tensorflow checkpoint from  model.ckpt-1000000\n",
      "train data size: 600, validation data size: 0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Training: 100%|██████████| 75/75 [00:25<00:00,  3.13it/s, loss=0.0471]\n",
      "Training: 100%|██████████| 75/75 [00:24<00:00,  3.13it/s, loss=0.00758]\n",
      "Training: 100%|██████████| 75/75 [00:25<00:00,  2.97it/s, loss=0.00332]\n",
      "Predicting: 100%|██████████| 59/59 [00:12<00:00,  4.39it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test f1: 85.36\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Predicting: 100%|██████████| 59/59 [00:14<00:00,  4.36it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "              precision    recall  f1-score   support\n",
      "\n",
      "   B-Disease       0.84      0.85      0.84       960\n",
      "   I-Disease       0.88      0.85      0.86      1087\n",
      "           O       0.99      0.99      0.99     22450\n",
      "\n",
      "   micro avg       0.98      0.98      0.98     24497\n",
      "   macro avg       0.90      0.90      0.90     24497\n",
      "weighted avg       0.98      0.98      0.98     24497\n",
      "\n",
      "CPU times: user 1min 31s, sys: 44.1 s, total: 2min 15s\n",
      "Wall time: 2min 50s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "model = BertTokenClassifier('biobert-v1.1-pubmed-base-cased',\n",
    "                            max_seq_length=178,\n",
    "                            epochs=3,\n",
    "                            gradient_accumulation_steps=2,\n",
    "                            learning_rate=3e-5,\n",
    "                            train_batch_size=16,\n",
    "                            eval_batch_size=16,\n",
    "                            validation_fraction=0,                            \n",
    "                            label_list=label_list,\n",
    "                            ignore_label=['O'])\n",
    "\n",
    "print(model)\n",
    "\n",
    "# finetune model on train data\n",
    "model.fit(X_train, y_train)\n",
    "\n",
    "# score model on test data\n",
    "f1_test = model.score(X_test, y_test,'macro')\n",
    "print(\"Test f1: %0.02f\"%(f1_test))\n",
    "\n",
    "# get predictions on test data\n",
    "y_preds = model.predict(X_test)\n",
    "\n",
    "# print report on classifier stats\n",
    "print(classification_report(flatten(y_test), flatten(y_preds)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "processed 24497 tokens with 960 phrases; found: 1006 phrases; correct: 776.\r\n",
      "accuracy:  97.84%; precision:  77.14%; recall:  80.83%; FB1:  78.94\r\n",
      "          Disease: precision:  77.14%; recall:  80.83%; FB1:  78.94  1006\r\n"
     ]
    }
   ],
   "source": [
    "# get span level stats with conlleval.pl script\n",
    "iter_zip = zip(flatten(X_test),flatten(y_test),flatten(y_preds))\n",
    "preds = [\" \".join([token, y, y_pred]) for token, y, y_pred in iter_zip]\n",
    "with open(\"preds.txt\",'w') as f:\n",
    "    for x in preds:\n",
    "        f.write(str(x)+'\\n') \n",
    "\n",
    "# run conlleval perl script \n",
    "!perl ./other_examples/conlleval.pl < preds.txt\n",
    "!rm preds.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get the token predicts\n",
    "i = 9\n",
    "tokens = X_test[i]\n",
    "labels = y_test[i]\n",
    "biobert_preds  = y_preds[i]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "###  examine the token tags from the 3 models"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compare the token predictions between BERT, SciBERT, and BioBERT. Remember these are results based on training on just 600 training samples. See [ner_NCBI_disease_BioBERT_SciBERT](https://github.com/charles9n/bert-sklearn/blob/master/other_examples/ner_NCBI_disease_BioBERT_SciBERT.ipynb) for a finetune demo on all the data with the various models. On the full dataset both SciBERT and BiBERT are pretty close and noticably better than standard BERT."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "         token      label       bert    scibert    biobert\n",
      "0   Occasional          O          O          O          O\n",
      "1     missense          O          O          O          O\n",
      "2    mutations          O          O          O          O\n",
      "3           in          O          O          O          O\n",
      "4          ATM          O          O          O          O\n",
      "5         were          O          O          O          O\n",
      "6         also          O          O          O          O\n",
      "7        found          O          O          O          O\n",
      "8           in          O          O          O          O\n",
      "9       tumour  B-Disease  B-Disease  B-Disease          O\n",
      "10         DNA          O          O          O          O\n",
      "11        from          O          O          O          O\n",
      "12    patients          O          O          O          O\n",
      "13        with          O          O          O          O\n",
      "14           B  B-Disease          O  B-Disease  B-Disease\n",
      "15           -  I-Disease          O  I-Disease  I-Disease\n",
      "16        cell  I-Disease          O  I-Disease  I-Disease\n",
      "17         non  I-Disease  B-Disease  I-Disease  I-Disease\n",
      "18           -  I-Disease  I-Disease  I-Disease  I-Disease\n",
      "19    Hodgkins  I-Disease  I-Disease  I-Disease  I-Disease\n",
      "20   lymphomas  I-Disease  I-Disease  I-Disease  I-Disease\n",
      "21           (          O          O          O          O\n",
      "22           B  B-Disease  B-Disease  B-Disease  B-Disease\n",
      "23           -  I-Disease          O  I-Disease  B-Disease\n",
      "24         NHL  I-Disease  I-Disease  I-Disease  I-Disease\n",
      "25           )          O          O          O          O\n",
      "26         and          O          O          O          O\n",
      "27           a          O          O          O          O\n",
      "28           B  B-Disease          O  B-Disease  B-Disease\n",
      "29           -  I-Disease          O  I-Disease  I-Disease\n",
      "30         NHL  I-Disease          O  I-Disease  I-Disease\n",
      "31        cell          O          O          O          O\n",
      "32        line          O          O          O          O\n",
      "33           .          O          O          O          O\n"
     ]
    }
   ],
   "source": [
    "data = {\"token\": tokens,\"label\": labels, \"bert\": bert_preds,\"scibert\": scibert_preds, \"biobert\": biobert_preds}\n",
    "df=pd.DataFrame(data=data)\n",
    "print(df)"
   ]
  }
 ],
 "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.6.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}