CRF-CoNLL.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# CRF for Entity Extraction on CoNLL2003\n",
    "\n",
    "In this notebook we build a CRF model for Name Entity Recognition over the CONLL2003 english dataset. \n",
    "We will use the `sklearn-crfsuite` package for implementing our model and `seqeval` for f1-score evaluation.\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from utils import dataio, modelutils\n",
    "from pprint import pprint\n",
    "from seqeval.metrics import classification_report\n",
    "from sklearn_crfsuite import CRF"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load Dataset\n",
    "We load CONLL2003 dataset from [this GitHub repo](https://github.com/davidsbatista/NER-datasets/tree/master/CONLL2003). \n",
    "For each token it reports Part-of-Speech tag, Dependency tag and Entity (with BIO notation). One token per line, features separated with a whitespace, sentences are separated with an empty line."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Reading file data\\conll03\\train.txt\n",
      "Read 14027 sentences\n",
      "Reading file data\\conll03\\valid.txt\n",
      "Read 3249 sentences\n",
      "Reading file data\\conll03\\test.txt\n",
      "Read 3452 sentences\n"
     ]
    }
   ],
   "source": [
    "data_dir = os.path.join('data', 'conll03')\n",
    "raw_train, Y_train, output_labels = dataio.load_conll_data('train.txt', dir_path=data_dir)\n",
    "raw_valid, Y_valid, _ = dataio.load_conll_data('valid.txt', dir_path=data_dir)\n",
    "raw_test, Y_test, _ = dataio.load_conll_data('test.txt', dir_path=data_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Labels: {'I-ORG', 'I-PER', 'B-MISC', 'B-PER', 'I-MISC', 'B-ORG', 'B-LOC', 'O', 'I-LOC'}\n"
     ]
    }
   ],
   "source": [
    "print(\"Labels:\", output_labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Sentence Example:\n",
      "[('German', 'JJ', 'B-NP'),\n",
      " ('call', 'NN', 'I-NP'),\n",
      " ('to', 'TO', 'B-VP'),\n",
      " ('boycott', 'VB', 'I-VP'),\n",
      " ('British', 'JJ', 'B-NP'),\n",
      " ('lamb', 'NN', 'I-NP'),\n",
      " ('.', '.', 'O')]\n",
      "==============================\n",
      "['B-MISC', 'O', 'O', 'O', 'B-MISC', 'O', 'O']\n"
     ]
    }
   ],
   "source": [
    "print(\"Sentence Example:\")\n",
    "pprint(raw_train[0])\n",
    "print(\"=\"*30)\n",
    "print(Y_train[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## Feature Functions\n",
    "\n",
    "In this section we define the features to extract from each token. Each token will be represent with a vector that contain:\n",
    "* The lowercase token string*;\n",
    "* The token suffix;\n",
    "* If the token is capitalized*;\n",
    "* If the token is uppercase*;\n",
    "* If the token is a number;\n",
    "* Complete Part-of-Speech tag of the token*;\n",
    "* More general Part-of-Speech tag of the token*;\n",
    "* Complete Dependency tag of the token*;\n",
    "* More general Dependency tag of the token*;\n",
    "* If the token is the first of the sentence;\n",
    "* If the token is the last of the sentence.\n",
    "\n",
    "\\* also for previous and next tokens, if there are.  \n",
    "\n",
    "> Note: categorical features are one-hot encoded."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def word_features(sentence, idx):\n",
    "    \"\"\"Extract features related to a word and its neighbours\"\"\"\n",
    "    word, pos, dep = sentence[idx]\n",
    "    \n",
    "    features = {\n",
    "        'bias': 1.0,\n",
    "        'word.lower()': word.lower(),\n",
    "        'word[-3:]': word[-3:],\n",
    "        'word[-2:]': word[-2:],\n",
    "        'word.isupper()': word.isupper(),\n",
    "        'word.istitle()': word.istitle(),\n",
    "        'word.isdigit()': word.isdigit(),\n",
    "        'postag': pos,\n",
    "        'postag[:2]': pos[:2],\n",
    "        'deptag': dep,\n",
    "        'deptag[-2:]': dep[-2:]\n",
    "    }\n",
    "    if idx > 0:\n",
    "        word1, pos1, dep1 = sentence[idx-1]\n",
    "        features.update({\n",
    "            '-1:word.lower()': word1.lower(),\n",
    "            '-1:word.istitle()': word1.istitle(),\n",
    "            '-1:word.isupper()': word1.isupper(),\n",
    "            '-1:postag': pos1,\n",
    "            '-1:postag[:2]': pos1[:2],\n",
    "            '-1:deptag': dep1,\n",
    "            '-1:deptag[-2:]': dep1[-2:],\n",
    "        })\n",
    "    else:\n",
    "        features['BOS'] = True\n",
    "        \n",
    "    if idx < len(sentence)-1:\n",
    "        word1, pos1, dep1 = sentence[idx+1]\n",
    "        features.update({\n",
    "            '+1:word.lower()': word1.lower(),\n",
    "            '+1:word.istitle()': word1.istitle(),\n",
    "            '+1:word.isupper()': word1.isupper(),\n",
    "            '+1:postag': pos1,\n",
    "            '+1:postag[:2]': pos1[:2],\n",
    "            '+1:deptag': dep1,\n",
    "            '+1:deptag[-2:]': dep1[-2:],\n",
    "        })\n",
    "    else:\n",
    "        features['EOS'] = True\n",
    "                \n",
    "    return features\n",
    "\n",
    "\n",
    "def sentence_features(sentence):\n",
    "    return tuple(word_features(sentence, index) for index in range(len(sentence)))\n",
    "\n",
    "X_train = [sentence_features(sentence) for sentence in raw_train]\n",
    "X_valid = [sentence_features(sentence) for sentence in raw_valid]\n",
    "X_test = [sentence_features(sentence) for sentence in raw_test]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Token features example:\n",
      "{'-1:deptag': 'B-NP',\n",
      " '-1:deptag[-2:]': 'NP',\n",
      " '-1:postag': 'NNP',\n",
      " '-1:postag[:2]': 'NN',\n",
      " '-1:word.istitle()': True,\n",
      " '-1:word.isupper()': False,\n",
      " '-1:word.lower()': 'peter',\n",
      " 'EOS': True,\n",
      " 'bias': 1.0,\n",
      " 'deptag': 'I-NP',\n",
      " 'deptag[-2:]': 'NP',\n",
      " 'postag': 'NNP',\n",
      " 'postag[:2]': 'NN',\n",
      " 'word.isdigit()': False,\n",
      " 'word.istitle()': True,\n",
      " 'word.isupper()': False,\n",
      " 'word.lower()': 'blackburn',\n",
      " 'word[-2:]': 'rn',\n",
      " 'word[-3:]': 'urn'}\n",
      "==============================\n",
      "I-PER\n"
     ]
    }
   ],
   "source": [
    "print(\"Token features example:\")\n",
    "pprint(X_train[1][1])\n",
    "print(\"=\"*30)\n",
    "print(Y_train[1][1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Wall time: 4min 23s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\srichetti\\AppData\\Local\\Continuum\\anaconda3\\envs\\ner-suite\\lib\\site-packages\\sklearn\\base.py:213: FutureWarning: From version 0.24, get_params will raise an AttributeError if a parameter cannot be retrieved as an instance attribute. Previously it would return None.\n",
      "  FutureWarning)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "CRF(algorithm='lbfgs', all_possible_transitions=True, c1=0.1, c2=0.5,\n",
       "    keep_tempfiles=None, max_iterations=800)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "crf = CRF(\n",
    "    algorithm = 'lbfgs',\n",
    "    c1 = 0.1,\n",
    "    c2 = 0.5,\n",
    "    max_iterations = 800,\n",
    "    all_possible_transitions = True,\n",
    "    verbose = False\n",
    ")\n",
    "\n",
    "crf.fit(X_train, Y_train, X_dev=X_valid, y_dev=Y_valid)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluation\n",
    "\n",
    "We evaluate:\n",
    "* **Memory consumption** using the attribute `crf.size_`;\n",
    "* **Latency in prediction** using the function `time.process_time()`;\n",
    "* **F1-score** _on entities_ on the test set using `seqeval`;"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model size: 1.91M\n"
     ]
    }
   ],
   "source": [
    "print(f'Model size: {crf.size_ / 1000000:0.2f}M')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model latency in prediction: 0.00019 s\n"
     ]
    }
   ],
   "source": [
    "print(f'Model latency in prediction: {modelutils.compute_prediction_latency(X_test, crf):.3} s')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training Set\n",
      "           precision    recall  f1-score   support\n",
      "\n",
      "      ORG      0.961     0.943     0.952      6318\n",
      "      PER      0.977     0.974     0.975      6600\n",
      "      LOC      0.971     0.971     0.971      7140\n",
      "     MISC      0.964     0.916     0.940      3438\n",
      "\n",
      "micro avg      0.969     0.956     0.963     23496\n",
      "macro avg      0.969     0.956     0.963     23496\n",
      "\n",
      "\n",
      "\n",
      "Test Set\n",
      "           precision    recall  f1-score   support\n",
      "\n",
      "      PER      0.809     0.850     0.829      1616\n",
      "      ORG      0.737     0.697     0.717      1660\n",
      "      LOC      0.848     0.809     0.828      1667\n",
      "     MISC      0.801     0.728     0.762       701\n",
      "\n",
      "micro avg      0.799     0.778     0.788      5644\n",
      "macro avg      0.798     0.778     0.787      5644\n",
      "\n",
      "\n",
      "\n",
      "Validation Set\n",
      "           precision    recall  f1-score   support\n",
      "\n",
      "      ORG      0.818     0.772     0.794      1340\n",
      "      LOC      0.903     0.868     0.885      1837\n",
      "     MISC      0.907     0.808     0.855       922\n",
      "      PER      0.888     0.894     0.891      1842\n",
      "\n",
      "micro avg      0.880     0.845     0.862      5941\n",
      "macro avg      0.880     0.845     0.862      5941\n",
      "\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "datasets = [('Training Set', X_train, Y_train), \n",
    "            ('Test Set', X_test, Y_test), \n",
    "            ('Validation Set', X_valid, Y_valid)]\n",
    "\n",
    "for title, X, Y in datasets:\n",
    "    Y_pred = crf.predict(X)\n",
    "    print(title)\n",
    "    print(classification_report(Y, Y_pred, digits=3))\n",
    "    print('\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  }
 ],
 "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.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}