LSTM-CoNLL.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LSTM for Entity Extraction on CoNLL2003\n",
    "\n",
    "In this notebook, we perform Entity Extraction on the CoNLL03 dataset using a LSTM-based neural network. 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 urllib\n",
    "import sklearn\n",
    "import logging\n",
    "import os\n",
    "import numpy as np\n",
    "from pprint import pprint\n",
    "from utils import dataio, kerasutils, modelutils\n",
    "from seqeval.metrics import classification_report\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"
   ]
  },
  {
   "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 we keep only the string of the word and and the Entity tag (in BIO notation), we discard PoS and Dependency tags. 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, ner_train, output_labels = dataio.load_conll_data('train.txt', dir_path=data_dir, only_tokens=True)\n",
    "raw_valid, ner_valid, _ = dataio.load_conll_data('valid.txt', dir_path=data_dir, only_tokens=True)\n",
    "raw_test, ner_test, _ = dataio.load_conll_data('test.txt', dir_path=data_dir, only_tokens=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Labels: {'I-PER', 'B-MISC', 'O', 'B-ORG', 'I-LOC', 'B-LOC', 'I-ORG', 'I-MISC', 'B-PER'}\n"
     ]
    }
   ],
   "source": [
    "print(\"Labels:\", output_labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Sentence Example:\n",
      "German   |  B-MISC\n",
      "call     |  O\n",
      "to       |  O\n",
      "boycott  |  O\n",
      "British  |  B-MISC\n",
      "lamb     |  O\n",
      ".        |  O\n"
     ]
    }
   ],
   "source": [
    "print(\"Sentence Example:\")\n",
    "for i in range(len(raw_train[0])):\n",
    "    print(f'{raw_train[0][i]:7}  |  {ner_train[0][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": [],
   "source": [
    "# integer encode sequences of words\n",
    "token_tokenizer = Tokenizer()    # Automatically lowers tokens\n",
    "token_tokenizer.fit_on_texts(raw_train + raw_valid + raw_test)\n",
    "train_sequences = token_tokenizer.texts_to_sequences(raw_train)\n",
    "test_sequences = token_tokenizer.texts_to_sequences(raw_test)\n",
    "valid_sequences = token_tokenizer.texts_to_sequences(raw_valid)\n",
    "\n",
    "# Dictionaries for id <-> string conversation of labels\n",
    "tag2idx = { tag: idx for idx, tag in enumerate(output_labels) }\n",
    "idx2tag = { idx: tag for tag, idx in tag2idx.items() }\n",
    "\n",
    "# Label encoding\n",
    "ner_train_sequences = [[tag2idx[tag] for tag in sentence] for sentence in ner_train]\n",
    "ner_test_sequences  = [[tag2idx[tag] for tag in sentence] for sentence in ner_test ]\n",
    "ner_valid_sequences = [[tag2idx[tag] for tag in sentence] for sentence in ner_valid]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['JAPAN', 'GET', 'LUCKY', 'WIN', ',', 'CHINA', 'IN', 'SURPRISE', 'DEFEAT', '.']\n",
      "[210, 481, 4284, 161, 2, 175, 5, 2047, 946, 3]\n",
      "   210 | japan\n",
      "   481 | get\n",
      "  4284 | lucky\n",
      "   161 | win\n",
      "     2 | ,\n",
      "   175 | china\n",
      "     5 | in\n",
      "  2047 | surprise\n",
      "   946 | defeat\n",
      "     3 | .\n"
     ]
    }
   ],
   "source": [
    "print(raw_test[0])\n",
    "print(test_sequences[0])\n",
    "for i in test_sequences[0]:\n",
    "    print(f'{i:6} | {token_tokenizer.index_word[i]}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Vocabulary dimension: 26861\n"
     ]
    }
   ],
   "source": [
    "vocabulary_size = len(token_tokenizer.word_counts)\n",
    "print('Vocabulary dimension:', vocabulary_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['German', 'call', 'to', 'boycott', 'British', 'lamb', '.']\n",
      "[1, 2, 2, 2, 1, 2, 2]\n",
      "1 : B-MISC\n",
      "2 : O\n",
      "2 : O\n",
      "2 : O\n",
      "1 : B-MISC\n",
      "2 : O\n",
      "2 : O\n"
     ]
    }
   ],
   "source": [
    "print(raw_train[0])\n",
    "print(ner_train_sequences[0])\n",
    "for i in ner_train_sequences[0]:\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": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Longest sequence: 113\n",
      "[('75%', 22.0), ('80%', 26.0), ('85%', 29.0), ('90%', 32.0), ('95%', 37.69999999999891), ('100%', 113.0)]\n"
     ]
    }
   ],
   "source": [
    "sequence_len = np.array([len(s) for s in train_sequences])\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": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "max_sequence_len = 50\n",
    "X_train = pad_sequences(train_sequences, maxlen=max_sequence_len, padding='post', truncating='post')\n",
    "X_test = pad_sequences(test_sequences, maxlen=max_sequence_len, padding='post', truncating='post')\n",
    "X_valid = pad_sequences(valid_sequences, maxlen=max_sequence_len, padding='post', truncating='post')\n",
    "\n",
    "Y_train = pad_sequences(ner_train_sequences, maxlen=max_sequence_len, value=tag2idx['O'], padding='post', truncating='post')\n",
    "Y_test = pad_sequences(ner_test_sequences, maxlen=max_sequence_len, value=tag2idx['O'], padding='post', truncating='post')\n",
    "Y_valid = pad_sequences(ner_valid_sequences, maxlen=max_sequence_len, value=tag2idx['O'], padding='post', truncating='post')\n",
    "\n",
    "# Convert labels from ids to one-hot vectors\n",
    "Y_train = to_categorical(Y_train, num_classes=len(output_labels), dtype='int32')\n",
    "Y_test = to_categorical(Y_test, num_classes=len(output_labels), dtype='int32')\n",
    "Y_valid = to_categorical(Y_valid, num_classes=len(output_labels), dtype='int32')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "   207 | [0 1 0 0 0 0 0 0 0]\n",
      "   709 | [0 0 1 0 0 0 0 0 0]\n",
      "     6 | [0 0 1 0 0 0 0 0 0]\n",
      "  3628 | [0 0 1 0 0 0 0 0 0]\n",
      "   228 | [0 1 0 0 0 0 0 0 0]\n",
      "  7656 | [0 0 1 0 0 0 0 0 0]\n",
      "     3 | [0 0 1 0 0 0 0 0 0]\n",
      "     0 | [0 0 1 0 0 0 0 0 0]\n",
      "     0 | [0 0 1 0 0 0 0 0 0]\n",
      "     0 | [0 0 1 0 0 0 0 0 0]\n",
      "     0 | [0 0 1 0 0 0 0 0 0]\n",
      "     0 | [0 0 1 0 0 0 0 0 0]\n",
      "     0 | [0 0 1 0 0 0 0 0 0]\n",
      "     0 | [0 0 1 0 0 0 0 0 0]\n",
      "     0 | [0 0 1 0 0 0 0 0 0]\n",
      "     0 | [0 0 1 0 0 0 0 0 0]\n",
      "     0 | [0 0 1 0 0 0 0 0 0]\n",
      "     0 | [0 0 1 0 0 0 0 0 0]\n",
      "     0 | [0 0 1 0 0 0 0 0 0]\n",
      "     0 | [0 0 1 0 0 0 0 0 0]\n",
      "     0 | [0 0 1 0 0 0 0 0 0]\n",
      "     0 | [0 0 1 0 0 0 0 0 0]\n",
      "     0 | [0 0 1 0 0 0 0 0 0]\n",
      "     0 | [0 0 1 0 0 0 0 0 0]\n",
      "     0 | [0 0 1 0 0 0 0 0 0]\n",
      "     0 | [0 0 1 0 0 0 0 0 0]\n",
      "     0 | [0 0 1 0 0 0 0 0 0]\n",
      "     0 | [0 0 1 0 0 0 0 0 0]\n",
      "     0 | [0 0 1 0 0 0 0 0 0]\n",
      "     0 | [0 0 1 0 0 0 0 0 0]\n",
      "     0 | [0 0 1 0 0 0 0 0 0]\n",
      "     0 | [0 0 1 0 0 0 0 0 0]\n",
      "     0 | [0 0 1 0 0 0 0 0 0]\n",
      "     0 | [0 0 1 0 0 0 0 0 0]\n",
      "     0 | [0 0 1 0 0 0 0 0 0]\n",
      "     0 | [0 0 1 0 0 0 0 0 0]\n",
      "     0 | [0 0 1 0 0 0 0 0 0]\n",
      "     0 | [0 0 1 0 0 0 0 0 0]\n",
      "     0 | [0 0 1 0 0 0 0 0 0]\n",
      "     0 | [0 0 1 0 0 0 0 0 0]\n",
      "     0 | [0 0 1 0 0 0 0 0 0]\n",
      "     0 | [0 0 1 0 0 0 0 0 0]\n",
      "     0 | [0 0 1 0 0 0 0 0 0]\n",
      "     0 | [0 0 1 0 0 0 0 0 0]\n",
      "     0 | [0 0 1 0 0 0 0 0 0]\n",
      "     0 | [0 0 1 0 0 0 0 0 0]\n",
      "     0 | [0 0 1 0 0 0 0 0 0]\n",
      "     0 | [0 0 1 0 0 0 0 0 0]\n",
      "     0 | [0 0 1 0 0 0 0 0 0]\n",
      "     0 | [0 0 1 0 0 0 0 0 0]\n"
     ]
    }
   ],
   "source": [
    "for i in range(len(X_train[0])):\n",
    "    print(f'{X_train[0][i]:6} | {Y_train[0][i]}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "token_tokenizer.index_word[0] = '_PAD_'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train = np.array(X_train)\n",
    "Y_train = np.array(Y_train)\n",
    "X_test = np.array(X_test)\n",
    "Y_test = np.array(Y_test)\n",
    "X_valid = np.array(X_valid)\n",
    "Y_valid = np.array(Y_valid)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(14027, 50)\n",
      "(14027, 50, 9)\n"
     ]
    }
   ],
   "source": [
    "# Final training set dimensionalities\n",
    "print(X_train.shape)\n",
    "print(Y_train.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, token_tokenizer.word_index, embedding_dim)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The code that build the model is the following:\n",
    "```python\n",
    "def create_paper_BiLSTM(vocabulary_size, seq_len, n_classes, hidden_cells=200, \n",
    "                  embed_dim=100, drop=0.4, use_glove=False, glove_matrix=None):\n",
    "    \"\"\"Create a BiLSTM model using keras, given its parameters\"\"\"\n",
    "    \n",
    "    model = Sequential()\n",
    "    if use_glove:\n",
    "        model.add(Embedding(vocabulary_size, embed_dim, \n",
    "                            weights=[glove_matrix], input_length=seq_len,\n",
    "                            mask_zero=True, trainable=True))\n",
    "    else:\n",
    "        model.add(Embedding(vocabulary_size, embed_dim, input_length=seq_len, \n",
    "                            mask_zero=True))\n",
    "    model.add(Dropout(drop))\n",
    "    \n",
    "    model.add(Bidirectional(LSTM(hidden_cells, return_sequences=True, \n",
    "                                 dropout=drop)))\n",
    "\n",
    "    model.add(Dense(n_classes, activation='softmax'))\n",
    "    \n",
    "    model.compile(loss='categorical_crossentropy', \n",
    "                  optimizer='adam',\n",
    "                  metrics=['accuracy',\n",
    "                           Precision(),\n",
    "                           Recall()])\n",
    "    model.summary()\n",
    "    return model\n",
    "```"
   ]
  },
  {
   "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)           2686200   \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: 3,171,409\n",
      "Trainable params: 3,171,409\n",
      "Non-trainable params: 0\n",
      "_________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "model = kerasutils.create_paper_BiLSTM(vocabulary_size+1, max_sequence_len, \n",
    "                                       len(output_labels),\n",
    "                                       use_glove=USE_GLOVE, \n",
    "                                       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": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/50\n",
      "1403/1403 [==============================] - 104s 74ms/step - loss: 0.0865 - accuracy: 0.9154 - precision: 0.9413 - recall: 0.8884 - val_loss: 0.0441 - val_accuracy: 0.9612 - val_precision: 0.9680 - val_recall: 0.9550\n",
      "Epoch 2/50\n",
      "1403/1403 [==============================] - 102s 73ms/step - loss: 0.0448 - accuracy: 0.9535 - precision: 0.9638 - recall: 0.9445 - val_loss: 0.0351 - val_accuracy: 0.9696 - val_precision: 0.9740 - val_recall: 0.9664\n",
      "Epoch 3/50\n",
      "1403/1403 [==============================] - 102s 73ms/step - loss: 0.0339 - accuracy: 0.9645 - precision: 0.9714 - recall: 0.9587 - val_loss: 0.0306 - val_accuracy: 0.9729 - val_precision: 0.9762 - val_recall: 0.9706\n",
      "Epoch 4/50\n",
      "1403/1403 [==============================] - 102s 73ms/step - loss: 0.0269 - accuracy: 0.9716 - precision: 0.9766 - recall: 0.9675 - val_loss: 0.0279 - val_accuracy: 0.9752 - val_precision: 0.9782 - val_recall: 0.9728\n",
      "Epoch 5/50\n",
      "1403/1403 [==============================] - 102s 73ms/step - loss: 0.0226 - accuracy: 0.9755 - precision: 0.9795 - recall: 0.9723 - val_loss: 0.0273 - val_accuracy: 0.9751 - val_precision: 0.9780 - val_recall: 0.9734\n",
      "Epoch 6/50\n",
      "1403/1403 [==============================] - ETA: 0s - loss: 0.0194 - accuracy: 0.9790 - precision: 0.9821 - recall: 0.9765Restoring model weights from the end of the best epoch.\n",
      "1403/1403 [==============================] - 102s 73ms/step - loss: 0.0194 - accuracy: 0.9790 - precision: 0.9821 - recall: 0.9765 - val_loss: 0.0273 - val_accuracy: 0.9765 - val_precision: 0.9783 - val_recall: 0.9752\n",
      "Epoch 00006: early stopping\n",
      "Wall time: 10min 23s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "batch_size = 10\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": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model size: 13.26 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": [
      "Model latency in predictions: 0.00489 s\n"
     ]
    }
   ],
   "source": [
    "print(f'Model latency in predictions: {modelutils.compute_prediction_latency(X_test, model):.3} s')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training Set\n",
      "           precision    recall  f1-score   support\n",
      "\n",
      "      PER      0.971     0.973     0.972      6589\n",
      "      ORG      0.890     0.834     0.861      6312\n",
      "     MISC      0.870     0.813     0.841      3435\n",
      "      LOC      0.935     0.949     0.942      7134\n",
      "\n",
      "micro avg      0.925     0.905     0.915     23470\n",
      "macro avg      0.924     0.905     0.914     23470\n",
      "\n",
      "\n",
      "\n",
      "Test Set\n",
      "           precision    recall  f1-score   support\n",
      "\n",
      "      ORG      0.826     0.715     0.766      1657\n",
      "      PER      0.932     0.904     0.918      1579\n",
      "      LOC      0.831     0.886     0.858      1655\n",
      "     MISC      0.734     0.693     0.713       700\n",
      "\n",
      "micro avg      0.846     0.816     0.831      5591\n",
      "macro avg      0.846     0.816     0.829      5591\n",
      "\n",
      "\n",
      "\n",
      "Validation Set\n",
      "           precision    recall  f1-score   support\n",
      "\n",
      "      LOC      0.904     0.929     0.916      1834\n",
      "      ORG      0.842     0.743     0.790      1338\n",
      "      PER      0.940     0.937     0.938      1796\n",
      "     MISC      0.827     0.754     0.789       919\n",
      "\n",
      "micro avg      0.891     0.862     0.876      5887\n",
      "macro avg      0.889     0.862     0.874      5887\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",
    "    # Get predictions: for each token we have as prediction a vector \n",
    "    # of probabilites\n",
    "    Y_pred = model.predict(X, batch_size=batch_size)\n",
    "    # We choose as category the one with the highest probability\n",
    "    Y_pred = np.array(np.argmax(Y_pred, axis=-1))\n",
    "    # Also flatten true labels\n",
    "    Y = np.array(np.argmax(Y, axis=-1))\n",
    "    # Remove padding from predictions and labels\n",
    "    Y, Y_pred = kerasutils.remove_seq_padding(X, Y, Y_pred)\n",
    "    # Restore strings instead that entity idss\n",
    "    let_y_true, let_y_pred = modelutils.from_encode_to_literal_labels(Y, Y_pred, idx2tag)\n",
    "    \n",
    "    print(title)\n",
    "    print(classification_report(let_y_true, let_y_pred, digits=3))\n",
    "    print('\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "### Bonus: visualize results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "      TOKEN      TRUE Y | PRED Y\n",
      "==================================\n",
      "china            B-LOC  | B-LOC\n",
      "controlled       O      | O\n",
      "most             O      | O\n",
      "of               O      | O\n",
      "the              O      | O\n",
      "match            O      | O\n",
      "and              O      | O\n",
      "saw              O      | O\n",
      "several          O      | O\n",
      "chances          O      | O\n",
      "missed           O      | O\n",
      "until            O      | O\n",
      "the              O      | O\n",
      "78th             O      | O\n",
      "minute           O      | O\n",
      "when             O      | O\n",
      "uzbek            B-MISC | B-MISC\n",
      "striker          O      | O\n",
      "igor             B-PER  | B-PER\n",
      "shkvyrin         I-PER  | I-PER\n",
      "took             O      | O\n",
      "advantage        O      | O\n",
      "of               O      | O\n",
      "a                O      | O\n",
      "misdirected      O      | O\n",
      "defensive        O      | O\n",
      "header           O      | O\n",
      "to               O      | O\n",
      "lob              O      | O\n",
      "the              O      | O\n",
      "ball             O      | O\n",
      "over             O      | O\n",
      "the              O      | O\n",
      "advancing        O      | O\n",
      "chinese          B-MISC | B-MISC\n",
      "keeper           O      | O\n",
      "and              O      | O\n",
      "into             O      | O\n",
      "an               O      | O\n",
      "empty            O      | O\n",
      "net              O      | O\n",
      ".                O      | O\n",
      "_PAD_            O      | O\n",
      "_PAD_            O      | O\n",
      "_PAD_            O      | O\n",
      "_PAD_            O      | O\n",
      "_PAD_            O      | O\n",
      "_PAD_            O      | O\n",
      "_PAD_            O      | O\n",
      "_PAD_            O      | O\n"
     ]
    }
   ],
   "source": [
    "i = 5\n",
    "sentence = X_test[i]\n",
    "y_pred = model.predict(X_test)\n",
    "y_pred = np.argmax(y_pred, axis=-1)\n",
    "y_pred = y_pred[i]\n",
    "y_true = np.argmax(Y_test, axis=-1)[i]\n",
    "\n",
    "print('      TOKEN      TRUE Y | PRED Y')\n",
    "print('='*34)\n",
    "for idx in range(len(sentence)):\n",
    "    print(f'{token_tokenizer.index_word[sentence[idx]]:15}  {idx2tag[y_true[idx]]:6} | {idx2tag[y_pred[idx]]}')\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
}