CRF-WiNER-EN.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# CRF for Entity Extraction on WikiNER (English)\n",
    "WiNER is a dataset of annotated sentences for Entity Extraction taken from Wikipedia. In this notebook we train and evaluate a CRF model on the english data to recognize entities such as Persons, Locations and Orgnizations from text.\n",
    "\n",
    "We 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 sklearn_crfsuite import CRF\n",
    "from sklearn.model_selection import train_test_split\n",
    "from seqeval.metrics import classification_report"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data Preparation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We load the dataset from the `data/` directory. For each token, the datatset reports word, Part of Speech tag and entity tag."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Read 142153 sentences.\n"
     ]
    }
   ],
   "source": [
    "file_path = os.path.join('data', 'wikiner-en-wp3-raw.txt')\n",
    "sentences, tags, output_labels = dataio.load_wikiner(file_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Labels: {'B-PER', 'B-MISC', 'B-LOC', 'I-PER', 'I-ORG', 'I-MISC', 'I-LOC', 'O', 'B-ORG'}\n"
     ]
    }
   ],
   "source": [
    "print(\"Labels:\", output_labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Sentence Example:\n",
      "[('In', 'IN'),\n",
      " ('the', 'DT'),\n",
      " ('end', 'NN'),\n",
      " (',', ','),\n",
      " ('for', 'IN'),\n",
      " ('anarchist', 'JJ'),\n",
      " ('historian', 'JJ'),\n",
      " ('Daniel', 'NNP'),\n",
      " ('Guerin', 'NNP'),\n",
      " ('\"', 'LQU'),\n",
      " ('Some', 'DT'),\n",
      " ('anarchists', 'NNS'),\n",
      " ('are', 'VBP'),\n",
      " ('more', 'RBR'),\n",
      " ('individualistic', 'JJ'),\n",
      " ('than', 'IN'),\n",
      " ('social', 'JJ'),\n",
      " (',', ','),\n",
      " ('some', 'DT'),\n",
      " ('more', 'JJR'),\n",
      " ('social', 'JJ'),\n",
      " ('than', 'IN'),\n",
      " ('individualistic', 'JJ'),\n",
      " ('.', '.')]\n",
      "==============================\n",
      "['O', 'O', 'O', 'O', 'O', 'O', 'O', 'I-PER', 'I-PER', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O']\n"
     ]
    }
   ],
   "source": [
    "print(\"Sentence Example:\")\n",
    "pprint(sentences[1])\n",
    "print(\"=\"*30)\n",
    "print(tags[1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Features Engineering\n",
    "\n",
    "In this section, we build our feature vector for each token. It is composed by:\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",
    "* 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 = 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",
    "    }\n",
    "    if idx > 0:\n",
    "        word1, pos1 = 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",
    "        })\n",
    "    else:\n",
    "        features['BOS'] = True\n",
    "        \n",
    "    if idx < len(sentence)-1:\n",
    "        word1, pos1 = 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",
    "        })\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)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = [sentence_features(sentence) for sentence in sentences]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Token features example:\n",
      "{'+1:postag': 'NN',\n",
      " '+1:postag[:2]': 'NN',\n",
      " '+1:word.istitle()': False,\n",
      " '+1:word.isupper()': False,\n",
      " '+1:word.lower()': 'end',\n",
      " '-1:postag': 'IN',\n",
      " '-1:postag[:2]': 'IN',\n",
      " '-1:word.istitle()': True,\n",
      " '-1:word.isupper()': False,\n",
      " '-1:word.lower()': 'in',\n",
      " 'bias': 1.0,\n",
      " 'postag': 'DT',\n",
      " 'postag[:2]': 'DT',\n",
      " 'word.isdigit()': False,\n",
      " 'word.istitle()': False,\n",
      " 'word.isupper()': False,\n",
      " 'word.lower()': 'the',\n",
      " 'word[-2:]': 'he',\n",
      " 'word[-3:]': 'the'}\n",
      "==============================\n",
      "O\n"
     ]
    }
   ],
   "source": [
    "print(\"Token features example:\")\n",
    "pprint(X[1][1])\n",
    "print(\"=\"*30)\n",
    "print(tags[1][1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train, X_test, y_train, y_test = train_test_split(X, tags, test_size=0.2, \n",
    "                                                    random_state=3791)\n",
    "X_train, X_valid, y_train, y_valid = train_test_split(X_train, y_train, \n",
    "                                                      test_size=0.2, \n",
    "                                                      random_state=3791)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Wall time: 22min 2s\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": 9,
     "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": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model size: 9.71M\n"
     ]
    }
   ],
   "source": [
    "print('Model size: {:0.2f}M'.format(crf.size_ / 1000000))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model latency in prediction: 0.000278 s\n"
     ]
    }
   ],
   "source": [
    "print(f'Model latency in prediction: {modelutils.compute_prediction_latency(X_test, crf):.3} s')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training Set\n",
      "           precision    recall  f1-score   support\n",
      "\n",
      "      LOC      0.861     0.898     0.879     54998\n",
      "     MISC      0.850     0.783     0.815     47222\n",
      "      PER      0.932     0.950     0.941     61603\n",
      "      ORG      0.879     0.784     0.829     31794\n",
      "\n",
      "micro avg      0.885     0.868     0.876    195617\n",
      "macro avg      0.884     0.868     0.875    195617\n",
      "\n",
      "\n",
      "\n",
      "Test Set\n",
      "           precision    recall  f1-score   support\n",
      "\n",
      "      ORG      0.809     0.704     0.753      9885\n",
      "      LOC      0.813     0.846     0.829     17267\n",
      "     MISC      0.767     0.696     0.730     14562\n",
      "      PER      0.889     0.921     0.905     19345\n",
      "\n",
      "micro avg      0.828     0.811     0.819     61059\n",
      "macro avg      0.825     0.811     0.817     61059\n",
      "\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "datasets = [('Training Set', X_train, y_train), ('Test Set', X_test, y_test)]\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
}