LSTM-WiNER-EN.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LSTM for Entity Extraction on WikiNER (English)\n",
    "\n",
    "WikiNER is a dataset of annotated sentences for Entity Extraction taken from Wikipedia. In this notebook, we train and evaluate a Bidirectional LSTM neural network model on the english WikiNER dataset to recognize Person, Locations and Organizations.\n",
    "\n",
    "We use `tf.keras.preprocessing.text.Tokenizer` for text preprocessing, we pad all the santences to the same length and load Glove embeddings for token encoding, then we use `tensorflow.keras` to build the model. Evaluation is made with the `seqeval` package.\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "from pprint import pprint\n",
    "from utils import dataio, kerasutils, modelutils\n",
    "from sklearn.model_selection import train_test_split\n",
    "from tensorflow.keras.preprocessing.text import Tokenizer\n",
    "from tensorflow.keras.preprocessing.sequence import pad_sequences\n",
    "from tensorflow.keras.utils import to_categorical, plot_model\n",
    "from tensorflow.keras.callbacks import ModelCheckpoint, EarlyStopping\n",
    "from seqeval.metrics import classification_report"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load Dataset\n",
    "\n",
    "Thanks to the author of [this repo](https://github.com/dice-group/FOX/blob/master/input/Wikiner/aij-wikiner-en-wp3.bz2) that makes WikiNER data easily available."
   ]
  },
  {
   "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, token_only=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Labels: {'I-MISC', 'B-ORG', 'B-PER', 'O', 'I-ORG', 'B-MISC', 'I-LOC', 'B-LOC', 'I-PER'}\n"
     ]
    }
   ],
   "source": [
    "print(\"Labels:\", output_labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Sentence Example:\n",
      "In               |  O\n",
      "the              |  O\n",
      "end              |  O\n",
      ",                |  O\n",
      "for              |  O\n",
      "anarchist        |  O\n",
      "historian        |  O\n",
      "Daniel           |  I-PER\n",
      "Guerin           |  I-PER\n",
      "\"                |  O\n",
      "Some             |  O\n",
      "anarchists       |  O\n",
      "are              |  O\n",
      "more             |  O\n",
      "individualistic  |  O\n",
      "than             |  O\n",
      "social           |  O\n",
      ",                |  O\n",
      "some             |  O\n",
      "more             |  O\n",
      "social           |  O\n",
      "than             |  O\n",
      "individualistic  |  O\n",
      ".                |  O\n"
     ]
    }
   ],
   "source": [
    "print(\"Sentence Example:\")\n",
    "for i in range(len(sentences[1])):\n",
    "    print(f'{sentences[1][i]:15}  |  {tags[1][i]}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Text preprocessing and token encoding"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Token Ordinal Encoding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'I-MISC': 0, 'B-ORG': 1, 'B-PER': 2, 'O': 3, 'I-ORG': 4, 'B-MISC': 5, 'I-LOC': 6, 'B-LOC': 7, 'I-PER': 8}\n"
     ]
    }
   ],
   "source": [
    "tokenizer = Tokenizer()\n",
    "tokenizer.fit_on_texts(sentences)\n",
    "X = tokenizer.texts_to_sequences(sentences)\n",
    "\n",
    "tag2idx = { tag: idx for idx, tag in enumerate(output_labels) }\n",
    "print(tag2idx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "idx2tag = { idx: tag for tag, idx in tag2idx.items() }\n",
    "tags = [[tag2idx[tag] for tag in sentence] for sentence in tags]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['In', 'the', 'end', ',', 'for', 'anarchist', 'historian', 'Daniel', 'Guerin', '\"', 'Some', 'anarchists', 'are', 'more', 'individualistic', 'than', 'social', ',', 'some', 'more', 'social', 'than', 'individualistic', '.']\n",
      "[5, 1, 160, 2, 14, 4838, 2337, 2371, 55954, 10, 56, 7102, 31, 54, 21657, 71, 414, 2, 56, 54, 414, 71, 21657, 3]\n",
      "     5 | in\n",
      "     1 | the\n",
      "   160 | end\n",
      "     2 | ,\n",
      "    14 | for\n",
      "  4838 | anarchist\n",
      "  2337 | historian\n",
      "  2371 | daniel\n",
      " 55954 | guerin\n",
      "    10 | \"\n",
      "    56 | some\n",
      "  7102 | anarchists\n",
      "    31 | are\n",
      "    54 | more\n",
      " 21657 | individualistic\n",
      "    71 | than\n",
      "   414 | social\n",
      "     2 | ,\n",
      "    56 | some\n",
      "    54 | more\n",
      "   414 | social\n",
      "    71 | than\n",
      " 21657 | individualistic\n",
      "     3 | .\n"
     ]
    }
   ],
   "source": [
    "print(sentences[1])\n",
    "print(X[1])\n",
    "for i in X[1]:\n",
    "    print(f'{i:6} | {tokenizer.index_word[i]}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "108276\n"
     ]
    }
   ],
   "source": [
    "vocabulary_size = len(tokenizer.word_counts)\n",
    "print(vocabulary_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['In', 'the', 'end', ',', 'for', 'anarchist', 'historian', 'Daniel', 'Guerin', '\"', 'Some', 'anarchists', 'are', 'more', 'individualistic', 'than', 'social', ',', 'some', 'more', 'social', 'than', 'individualistic', '.']\n",
      "[3, 3, 3, 3, 3, 3, 3, 8, 8, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3]\n",
      "3 : O\n",
      "3 : O\n",
      "3 : O\n",
      "3 : O\n",
      "3 : O\n",
      "3 : O\n",
      "3 : O\n",
      "8 : I-PER\n",
      "8 : I-PER\n",
      "3 : O\n",
      "3 : O\n",
      "3 : O\n",
      "3 : O\n",
      "3 : O\n",
      "3 : O\n",
      "3 : O\n",
      "3 : O\n",
      "3 : O\n",
      "3 : O\n",
      "3 : O\n",
      "3 : O\n",
      "3 : O\n",
      "3 : O\n",
      "3 : O\n"
     ]
    }
   ],
   "source": [
    "print(sentences[1])\n",
    "print(tags[1])\n",
    "for i in tags[1]:\n",
    "    print(f'{i} : {idx2tag[i]}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Sequence Padding\n",
    "\n",
    "The input sequence of an LSTM model must have a fixed length. We choose the most appropriate seqence length given the length of the sentences of the dataset, than we pad shorter sentences and truncate the longer ones."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Longest sequence: 173\n",
      "[('75%', 31.0), ('80%', 33.0), ('85%', 36.0), ('90%', 40.0), ('95%', 46.0), ('100%', 173.0)]\n"
     ]
    }
   ],
   "source": [
    "sequence_len = np.array([len(s) for s in sentences])\n",
    "longest_sequence = sequence_len.max()\n",
    "print(f'Longest sequence: {longest_sequence}')\n",
    "\n",
    "print([(str(p) + '%', np.percentile(sequence_len, p)) for p in range(75,101, 5)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "max_sequence_len = 50\n",
    "X = pad_sequences(X, maxlen=max_sequence_len, padding='post', truncating='post')\n",
    "\n",
    "y = pad_sequences(tags, maxlen=max_sequence_len, value=tag2idx['O'], padding='post', truncating='post')\n",
    "y = to_categorical(y, num_classes=len(output_labels), dtype='int32')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer.index_word[0] = '_PAD_'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = np.array(X)\n",
    "y = np.array(y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(142153, 50)\n",
      "(142153, 50, 9)\n"
     ]
    }
   ],
   "source": [
    "print(X.shape)\n",
    "print(y.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Build, train and evaluate an LSTM model\n",
    "\n",
    "The function that loads Glove embedding and the function which creates the LSTM model can be found in the `utils/kerasutils.py` module. Training stopping criterion is Early Stopping with patience on the loss value on validation set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "USE_GLOVE=True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found 400001 word vectors.\n"
     ]
    }
   ],
   "source": [
    "glove_matrix=None\n",
    "if USE_GLOVE:\n",
    "    glove_embedding_path = os.path.join('embeddings', 'glove.6B.100d.txt')\n",
    "    embedding_dim = 100\n",
    "    glove_matrix = kerasutils.load_glove_embedding_matrix(glove_embedding_path, tokenizer.word_index, embedding_dim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"sequential\"\n",
      "_________________________________________________________________\n",
      "Layer (type)                 Output Shape              Param #   \n",
      "=================================================================\n",
      "embedding (Embedding)        (None, 50, 100)           10827700  \n",
      "_________________________________________________________________\n",
      "dropout (Dropout)            (None, 50, 100)           0         \n",
      "_________________________________________________________________\n",
      "bidirectional (Bidirectional (None, 50, 400)           481600    \n",
      "_________________________________________________________________\n",
      "dense (Dense)                (None, 50, 9)             3609      \n",
      "=================================================================\n",
      "Total params: 11,312,909\n",
      "Trainable params: 11,312,909\n",
      "Non-trainable params: 0\n",
      "_________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "model = kerasutils.create_paper_BiLSTM(vocabulary_size+1, max_sequence_len, len(output_labels), \n",
    "                                 use_glove=USE_GLOVE, glove_matrix=glove_matrix)\n",
    "\n",
    "# Early Stopping on validation loss\n",
    "early_stopping_callback = EarlyStopping(\n",
    "    monitor=\"val_loss\", \n",
    "    min_delta=0.01, \n",
    "    patience=3, \n",
    "    verbose=1, \n",
    "    mode=\"auto\", \n",
    "    restore_best_weights=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create train-test-dev split\n",
    "X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=3791)\n",
    "X_train, X_valid, y_train, y_valid = train_test_split(X_train, y_train, test_size=0.2, random_state=3791)\n",
    "batch_size = 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model size: 44.317 MB\n"
     ]
    }
   ],
   "source": [
    "kerasutils.print_model_memory_usage(batch_size, model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/50\n",
      "9098/9098 [==============================] - 1361s 150ms/step - loss: 0.1021 - accuracy: 0.9334 - precision: 0.9472 - recall: 0.9215 - val_loss: 0.0628 - val_accuracy: 0.9593 - val_precision: 0.9648 - val_recall: 0.9545\n",
      "Epoch 2/50\n",
      "9098/9098 [==============================] - 1334s 147ms/step - loss: 0.0712 - accuracy: 0.9534 - precision: 0.9607 - recall: 0.9471 - val_loss: 0.0565 - val_accuracy: 0.9637 - val_precision: 0.9673 - val_recall: 0.9607\n",
      "Epoch 3/50\n",
      "9098/9098 [==============================] - 1319s 145ms/step - loss: 0.0615 - accuracy: 0.9599 - precision: 0.9655 - recall: 0.9552 - val_loss: 0.0533 - val_accuracy: 0.9659 - val_precision: 0.9691 - val_recall: 0.9634\n",
      "Epoch 4/50\n",
      "9098/9098 [==============================] - 1319s 145ms/step - loss: 0.0557 - accuracy: 0.9637 - precision: 0.9683 - recall: 0.9596 - val_loss: 0.0526 - val_accuracy: 0.9666 - val_precision: 0.9698 - val_recall: 0.9643\n",
      "Epoch 5/50\n",
      "9098/9098 [==============================] - 1318s 145ms/step - loss: 0.0516 - accuracy: 0.9663 - precision: 0.9704 - recall: 0.9628 - val_loss: 0.0514 - val_accuracy: 0.9676 - val_precision: 0.9702 - val_recall: 0.9658\n",
      "Epoch 6/50\n",
      "9098/9098 [==============================] - 1317s 145ms/step - loss: 0.0489 - accuracy: 0.9680 - precision: 0.9717 - recall: 0.9648 - val_loss: 0.0510 - val_accuracy: 0.9683 - val_precision: 0.9708 - val_recall: 0.9662\n",
      "Epoch 7/50\n",
      "9098/9098 [==============================] - ETA: 0s - loss: 0.0466 - accuracy: 0.9694 - precision: 0.9729 - recall: 0.9664Restoring model weights from the end of the best epoch.\n",
      "9098/9098 [==============================] - 1317s 145ms/step - loss: 0.0466 - accuracy: 0.9694 - precision: 0.9729 - recall: 0.9664 - val_loss: 0.0510 - val_accuracy: 0.9682 - val_precision: 0.9707 - val_recall: 0.9663\n",
      "Epoch 00007: early stopping\n",
      "Wall time: 2h 34min 54s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "history = model.fit(X_train, \n",
    "          y_train, \n",
    "          batch_size=batch_size, \n",
    "          epochs=50,\n",
    "          verbose=1,\n",
    "          callbacks=[early_stopping_callback],\n",
    "          validation_data=(X_valid, y_valid)\n",
    "         )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## Evaluation\n",
    "We evaluate three aspects of the model:\n",
    "* **Memory consumption** using the `kerasutils.print_model_memory_usage()` function (found [here](https://stackoverflow.com/questions/43137288/how-to-determine-needed-memory-of-keras-model));\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": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model latency in prediction: 0.00389 s\n"
     ]
    }
   ],
   "source": [
    "print(f'Model latency in prediction: {modelutils.compute_prediction_latency(X_test, model):.3} s')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "datasets = [('Training Set', X_train, y_train), ('Test Set', X_test, y_test)]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training Set\n",
      "           precision    recall  f1-score   support\n",
      "\n",
      "      LOC      0.816     0.878     0.846     54367\n",
      "      ORG      0.752     0.733     0.743     31449\n",
      "      PER      0.943     0.951     0.947     61086\n",
      "     MISC      0.747     0.695     0.720     46733\n",
      "\n",
      "micro avg      0.831     0.833     0.832    193635\n",
      "macro avg      0.829     0.833     0.831    193635\n",
      "\n",
      "\n",
      "\n",
      "Test Set\n",
      "           precision    recall  f1-score   support\n",
      "\n",
      "     MISC      0.685     0.626     0.654     14427\n",
      "      PER      0.913     0.931     0.922     19192\n",
      "      LOC      0.776     0.841     0.807     17119\n",
      "      ORG      0.695     0.676     0.685      9760\n",
      "\n",
      "micro avg      0.788     0.792     0.790     60498\n",
      "macro avg      0.785     0.792     0.787     60498\n",
      "\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "for title, X, Y in datasets:\n",
    "    Y_pred = model.predict(X, batch_size=128)\n",
    "    Y_pred = np.array(np.argmax(Y_pred, axis=-1))\n",
    "    Y = np.array(np.argmax(Y, axis=-1))\n",
    "    Y, Y_pred = kerasutils.remove_seq_padding(X, Y, Y_pred)\n",
    "    Y, Y_pred = modelutils.from_encode_to_literal_labels(Y, Y_pred, idx2tag)\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
}