named_entity_recognition_using_spacy_and_tensorflow.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "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.8"
    },
    "colab": {
      "name": "named_entity_recognition_using_spacy_and_tensorflow.ipynb",
      "provenance": []
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pe3A1dyBKpBB",
        "colab_type": "text"
      },
      "source": [
        "# Overview\n",
        "\n",
        "This Jupyter notebook demonstrates training and serving custom named entity recognition (NER) models, which are used to identify named entities such as locations, times, and people in text documents. NER is used in a number of business applications such as powering recommender systems, simplifying customer support, and optimizing internal search engines. \n",
        "\n",
        "The notebook is broken down into the following three sections:\n",
        "   * NER packages: An overview of the language, license, and methodology of commercially available NER packages\n",
        "   * NER with SpaCy : Code examples for training and serving a custom NER model in SpaCy\n",
        "   * NER with Tensorflow : Code examples for creating, training, and serving a custom deep-learning NER model with Tensorflow\n",
        "\n",
        "# 1. Named Entity Recognition Research & Examples\n",
        "\n",
        "\n",
        "\n",
        "## Named entity recognition packages\n",
        "\n",
        "NER can be implemented with either statistical or rule-based methods, both of which require a large amount of labeled training data and are typically trained in a fully or semi-supervised manner. Statistical approaches to NER include Hidden Markov Models, Maximum Entropy, and Conditional Random Fields, as well as deep learning approaches with Recurrent Neural Networks, such as Seq2Seq. All of these processes involve sentence inputs and annotated sentence outputs. Many of these processes also involve additional feature engineering, providing as input summary statistics of the sentences.\n",
        "\n",
        "Many of the production-ready NER packages are written in Java and served as Docker containers, such as GATE, OpenNLP, and DBPedia spotlight. SpaCy is perhaps the most frequently used NER package in Python. \n",
        "\n",
        "\n",
        "| Name | Language   | License | Method        |\n",
        "|------|------------|--------|----------------|\n",
        "| [GATE](https://github.com/GateNLP/gateplugin-Java) | Java       |   GPLv3     |  Hidden Markov |\n",
        "|[OpenNLP](https://opennlp.apache.org)| Java      |     Apache 2.0   | Maximum Entropy / Rule-based |\n",
        "|[DBPedia](https://opennlp.apache.org)| Java      |    Apache 2.0    | Rule-based |\n",
        "|[SpaCy](https://spacy.io)| Python/Cython      |    MIT  | Convolutional NN |\n",
        "\n",
        "## NER Methods\n",
        "\n",
        "NER may be implemented with a variety of statistical and rule-based methods with varying amounts of feature engineering. All production-ready NER methods are at least semi-supervised, though unsupervised approaches are an emerging research topic.\n",
        "\n",
        "#### Supervised statistical\n",
        "\n",
        "Supervised statistical approaches to NER typically use either Hidden Markov Models (HMM), Maximum Entropy (ME), or Conditional Random Fields (CRF). OpenNLP's statistical NER relies on ME. GATE relies on HMM.\n",
        "\n",
        "Typical feature engineering approaches for NER include such approaches as orthography, n-grams, lexicons, suffixes and prefixes, unsupervised cluster features, and trigger words for named entities (such as river or lake). These features are generated algorithmically in a rule-based manner.\n",
        "\n",
        "    \n",
        "#### Supervised rule-based\n",
        "\n",
        "OpenNLP contains rule based (as well as statistical) NER. The rule-based approach relies on a series of regular expression matches. The feature generation seems to be done with a beam search to determine the word context.\n",
        "\n",
        "DBPedia spotlight performs NER with substring matching using the Aho-Corasick algorithm. The approach only uses tokenization with no other feature engineering. The two-step approach first involves generating all possible candidate annotations that form known labels. This is rule-based in that it involves identifying nouns, prepositions, capitalized words, and known entities. This is based on OpenNLP under the hood. The second step selects the best candidates from the proposed candidates. Each candidate is scored based on annotation probability using a version of tf-idf with article links and anchor texts instead of documents and terms.\n",
        "\n",
        "\n",
        "#### Supervised deep learning\n",
        "\n",
        "**[SpaCy](https://spacy.io),** which is one of the most popular productionized NER environments, **uses residual convolutional neural networks (CNN) and incremental parsing with Bloom embeddings for NER.** See [this](https://www.youtube.com/watch?v=sqDHBH9IjRU) Youtube explanation from the developers for more detail. To summarize the algorithm, 1D convolutional filters are applied over the input text to predict how the upcoming words may change the current entity tags. Upcoming words may either shift (change the entity), reduce (make the entity more granular), or output the entity. The input sequence is embedded with bloom embeddings, which model the characters, prefix, suffix, and part of speech of each word. Residual blocks are used for the CNNs, andn the filter sizes are chosen with beam search.\n",
        "\n",
        "Recurrent neural network (RNN) approaches to NER also exist, typically comprising long short term memory networks (LSTM) at either the word- or character-level, relying on word or character embeddings, respectively (e.g. word2vec, gloVe, FASTtext).\n",
        "    \n",
        "## NER Datasets\n",
        "\n",
        "Although there are a number of datasets for NER in other languages, here we will focus on English datasets. Many of the NER datasets are domain-specific (i.e. Twitter, biomedical, advertising, news). A few standard NER datasets are described below to show the range of domain applications of NER.\n",
        "   * [i2b2](https://www.i2b2.org/NLP/DataSets/) - Medication, treatments, diseases, risk factors, and medications\n",
        "   * [CoNLL 2003](https://www.clips.uantwerpen.be/conll2003/ner/) - English and german news articles annotated with location, organization, person, and miscellaneous\n",
        "    \n",
        "    \n",
        "## NER Evaluation metrics\n",
        "\n",
        "NER is most commonly evaluated with precision, recall, and F1-score. F1-score can either be relaxed or strict, with the latter requiring the character offsets to match exactly. \n",
        "    \n",
        "# 2. Named entity recognition with SpaCy\n",
        "\n",
        "This section will focus on the Python package SpaCy for demonstrating named entity recognition. SpaCy is an open-source python library for NLP written in Python and Cython. It offers pre-trained models for multi-language NER, as well as allowing developers to train and deploy custom NER models on domain specific corpuses. SpaCy models are designed to be production-ready. \n",
        "\n",
        "SpaCy's pretrained models are trained on the [OntoNotes 5](https://catalog.ldc.upenn.edu/LDC2013T19) corpus, and support the identification of the [following entities](https://spacy.io/models/en): \n",
        "\n",
        "\n",
        "|TYPE\t| DESCRIPTION\n",
        "------------ | ------------\n",
        "|PERSON\t| People, including fictional\n",
        "|NORP\t|Nationalities or religious or political groups\n",
        "|FAC\t|Buildings, airports, highways, bridges, etc\n",
        "|ORG\t|Companies, agencies, institutions, etc\n",
        "|GPE\t|Countries, cities, states\n",
        "|LOC\t|Non-GPE locations, mountain ranges, bodies of water\n",
        "|PRODUCT|\tObjects, vehicles, foods, etc. (Not services.)\n",
        "|EVENT\t|Named hurricanes, battles, wars, sports events, etc\n",
        "|WORK_OF_ART|\tTitles of books, songs, etc\n",
        "|LAW\t|Named documents made into laws\n",
        "|LANGUAGE|\tAny named language\n",
        "|DATE\t|Absolute or relative dates or periods\n",
        "|TIME\t|Times smaller than a day\n",
        "|PERCENT|\tPercentage, including ”%“\n",
        "|MONEY\t|Monetary values, including unit\n",
        "|QUANTITY|\tMeasurements, as of weight or distance\n",
        "|ORDINAL|\t“first”, “second”, etc\n",
        "|CARDINAL|\tNumerals that do not fall under another type\n",
        "\n",
        "While these pretrained models are often sufficient for general applications, we will consider a domain-specific application of NER on the [MIT Movies corpus](https://groups.csail.mit.edu/sls/downloads/movie/), which contains 10,000 queries about various aspects of movies, with the following entity labels:\n",
        "\n",
        "| Type | Example |\n",
        "------- | ------- |\n",
        "| ACTOR | Matt Damon |\n",
        "| YEAR | 1980s |\n",
        "| TITLE | Pulp Fiction\n",
        "| GENRE | science fiction\n",
        "| DIRECTOR | George Lucas |\n",
        "| SONG | Aerosmith |\n",
        "| PLOT | Flying cars |\n",
        "| REVIEW | must see |\n",
        "| CHARACTER | Queen Elizabeth |\n",
        "|RATING | PG-13 |\n",
        "|RATINGS_AVERAGE | best rated |\n",
        "| TRAILER | preview\n",
        "\n",
        "As these tables show, the pretrained SpaCy models would not be sufficient to identify entities to help answer a question such as \"did george clooney make a science fiction movie in the 1980s?\" While the pre-trained entities may identify the presence of `PERSON`, `DATE`, and `PRODUCT`, a custom model should be able to detect `ACTOR`, `GENRE`, and `DATE`. In the following sections, we will compare the results of applying a pre-trained and a custom-trained model to the MIT movies corpus."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9kG7xCTRKpBG",
        "colab_type": "text"
      },
      "source": [
        "### Install and import required packages"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "RCWDkqVyKpBI",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 384
        },
        "outputId": "3422d85a-c472-4bc8-d69c-eda8c44bf176"
      },
      "source": [
        "import sys\n",
        "!{sys.executable} -m pip install spacy # !{sys.executable} ensures package installation in conda env"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Requirement already satisfied: spacy in /usr/local/lib/python3.6/dist-packages (2.2.4)\n",
            "Requirement already satisfied: srsly<1.1.0,>=1.0.2 in /usr/local/lib/python3.6/dist-packages (from spacy) (1.0.2)\n",
            "Requirement already satisfied: setuptools in /usr/local/lib/python3.6/dist-packages (from spacy) (47.3.1)\n",
            "Requirement already satisfied: thinc==7.4.0 in /usr/local/lib/python3.6/dist-packages (from spacy) (7.4.0)\n",
            "Requirement already satisfied: preshed<3.1.0,>=3.0.2 in /usr/local/lib/python3.6/dist-packages (from spacy) (3.0.2)\n",
            "Requirement already satisfied: cymem<2.1.0,>=2.0.2 in /usr/local/lib/python3.6/dist-packages (from spacy) (2.0.3)\n",
            "Requirement already satisfied: catalogue<1.1.0,>=0.0.7 in /usr/local/lib/python3.6/dist-packages (from spacy) (1.0.0)\n",
            "Requirement already satisfied: murmurhash<1.1.0,>=0.28.0 in /usr/local/lib/python3.6/dist-packages (from spacy) (1.0.2)\n",
            "Requirement already satisfied: wasabi<1.1.0,>=0.4.0 in /usr/local/lib/python3.6/dist-packages (from spacy) (0.6.0)\n",
            "Requirement already satisfied: blis<0.5.0,>=0.4.0 in /usr/local/lib/python3.6/dist-packages (from spacy) (0.4.1)\n",
            "Requirement already satisfied: requests<3.0.0,>=2.13.0 in /usr/local/lib/python3.6/dist-packages (from spacy) (2.23.0)\n",
            "Requirement already satisfied: numpy>=1.15.0 in /usr/local/lib/python3.6/dist-packages (from spacy) (1.18.5)\n",
            "Requirement already satisfied: tqdm<5.0.0,>=4.38.0 in /usr/local/lib/python3.6/dist-packages (from spacy) (4.41.1)\n",
            "Requirement already satisfied: plac<1.2.0,>=0.9.6 in /usr/local/lib/python3.6/dist-packages (from spacy) (1.1.3)\n",
            "Requirement already satisfied: importlib-metadata>=0.20; python_version < \"3.8\" in /usr/local/lib/python3.6/dist-packages (from catalogue<1.1.0,>=0.0.7->spacy) (1.6.1)\n",
            "Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.6/dist-packages (from requests<3.0.0,>=2.13.0->spacy) (3.0.4)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.6/dist-packages (from requests<3.0.0,>=2.13.0->spacy) (2020.4.5.2)\n",
            "Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.6/dist-packages (from requests<3.0.0,>=2.13.0->spacy) (1.24.3)\n",
            "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.6/dist-packages (from requests<3.0.0,>=2.13.0->spacy) (2.9)\n",
            "Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.6/dist-packages (from importlib-metadata>=0.20; python_version < \"3.8\"->catalogue<1.1.0,>=0.0.7->spacy) (3.1.0)\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PB3BrOMTKpBV",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import spacy\n",
        "import random\n",
        "import time\n",
        "import numpy as np\n",
        "from spacy.util import minibatch, compounding"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "g1_hXJraKpBe",
        "colab_type": "text"
      },
      "source": [
        "### Load and transform data"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ve7aFtu3KpBf",
        "colab_type": "text"
      },
      "source": [
        "Create the data directory if it doesn't exist"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KI40XqPAKpBg",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from os import path, mkdir\n",
        "if not path.isdir(\"data/\"):\n",
        "    mkdir(\"data/\")\n",
        "if not path.isdir(\"models/\"):\n",
        "    mkdir(\"models/\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AEiQhBkZKpBo",
        "colab_type": "text"
      },
      "source": [
        "Download the test and training dataset from MIT's Computer Science and Aritficial Intelligence Laboratory (CSAIL)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6qrEKhhhKpBp",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 122
        },
        "outputId": "c1383bc0-53c7-448f-984e-a7900d6c6a64"
      },
      "source": [
        "!curl https://groups.csail.mit.edu/sls/downloads/movie/engtest.bio -o data/test.txt\n",
        "!curl https://groups.csail.mit.edu/sls/downloads/movie/engtrain.bio -o data/train.txt"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "  % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current\n",
            "                                 Dload  Upload   Total   Spent    Left  Speed\n",
            "100  246k  100  246k    0     0   236k      0  0:00:01  0:00:01 --:--:--  236k\n",
            "  % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current\n",
            "                                 Dload  Upload   Total   Spent    Left  Speed\n",
            "100  989k  100  989k    0     0   812k      0  0:00:01  0:00:01 --:--:--  811k\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IEwVe7NVKpBw",
        "colab_type": "text"
      },
      "source": [
        "<img src=\"https://storage.googleapis.com/kf-pipeline-contrib-public/release-0.1.3/kfp-components/notebooks/entity_extraction/assets/fig1.png\" width=\"700\" align = \"left\"/>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3Fcss52PKpBx",
        "colab_type": "text"
      },
      "source": [
        "SpaCy requires training data to be in the format of `TRAIN_DATA = [(Sentence, {entities: [(start, end, label)]}, ...]`. The `load_data` function parses and transforms the input data into the required format for spaCy."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Kmjumt2BKpBy",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def load_data_spacy(file_path):\n",
        "    ''' Converts data from:\n",
        "    label \\t word \\n label \\t word \\n \\n label \\t word\n",
        "    to: sentence, {entities : [(start, end, label), (stard, end, label)]}\n",
        "    '''\n",
        "    file = open(file_path, 'r')\n",
        "    training_data, entities, sentence, unique_labels = [], [], [], []\n",
        "    current_annotation = None\n",
        "    end = 0 # initialize counter to keep track of start and end characters\n",
        "    for line in file:\n",
        "        line = line.strip(\"\\n\").split(\"\\t\")\n",
        "        # lines with len > 1 are words\n",
        "        if len(line) > 1:\n",
        "            label = line[0][2:]     # the .txt is formatted: label \\t word, label[0:2] = label_type\n",
        "            label_type = line[0][0] # beginning of annotations - \"B\", intermediate - \"I\"\n",
        "            word = line[1]\n",
        "            sentence.append(word)\n",
        "            end += (len(word) + 1)  # length of the word + trailing space\n",
        "            \n",
        "            if label_type != 'I' and current_annotation:  # if at the end of an annotation\n",
        "                entities.append((start, end - 2 - len(word), current_annotation))  # append the annotation\n",
        "                current_annotation = None                 # reset the annotation\n",
        "            if label_type == 'B':                         # if beginning new annotation\n",
        "                start = end - len(word) - 1  # start annotation at beginning of word\n",
        "                current_annotation = label   # append the word to the current annotation\n",
        "            if label_type == 'I':            # if the annotation is multi-word\n",
        "                current_annotation = label   # append the word\n",
        "            \n",
        "            if label != 'O' and label not in unique_labels:\n",
        "                unique_labels.append(label)\n",
        " \n",
        "        # lines with len == 1 are breaks between sentences\n",
        "        if len(line) == 1: \n",
        "            if current_annotation:\n",
        "                entities.append((start, end - 1, current_annotation))\n",
        "            sentence = \" \".join(sentence)\n",
        "            training_data.append([sentence, {'entities' : entities}])\n",
        "            # reset the counters and temporary lists\n",
        "            end = 0            \n",
        "            entities, sentence = [], []\n",
        "            current_annotation = None\n",
        "    file.close()\n",
        "    return training_data, unique_labels            \n",
        "            \n",
        "TRAIN_DATA, LABELS = load_data_spacy(\"data/train.txt\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dvipD72OLhIg",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 70
        },
        "outputId": "cb88a10b-c96b-4bc5-e525-1c75b0308d54"
      },
      "source": [
        "TRAIN_DATA[:2]"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[['what movies star bruce willis', {'entities': [(17, 29, 'ACTOR')]}],\n",
              " ['show me films with drew barrymore from the 1980s',\n",
              "  {'entities': [(19, 33, 'ACTOR'), (43, 48, 'YEAR')]}]]"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 6
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BybUHTMFKpB3",
        "colab_type": "text"
      },
      "source": [
        "### Data overview"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5CqA5U3dKpB3",
        "colab_type": "text"
      },
      "source": [
        "Sample sentences from the training data, which contains queries about movie information"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-IH0M9d3KpB4",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 174
        },
        "outputId": "c019220e-1803-4eae-e79c-504472b00402"
      },
      "source": [
        "[x[0] for x in TRAIN_DATA[1:10]]"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "['show me films with drew barrymore from the 1980s',\n",
              " 'what movies starred both al pacino and robert deniro',\n",
              " 'find me all of the movies that starred harold ramis and bill murray',\n",
              " 'find me a movie with a quote about baseball in it',\n",
              " 'what movies have mississippi in the title',\n",
              " 'show me science fiction films directed by steven spielberg',\n",
              " 'do you have any thrillers directed by sofia coppola',\n",
              " 'what leonard cohen songs have been used in a movie',\n",
              " 'show me films elvis films set in hawaii']"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 6
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hJwBsK-JKpB7",
        "colab_type": "text"
      },
      "source": [
        "<img src=\"https://storage.googleapis.com/kf-pipeline-contrib-public/release-0.1.3/kfp-components/notebooks/entity_extraction/assets/fig2.png\" align = \"left\" width=\"600\"/>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "g-ONh1rZKpB8",
        "colab_type": "text"
      },
      "source": [
        "Sample labeled annotations for the training data"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NgazyoVeKpB9",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 174
        },
        "outputId": "0f4312a4-3f21-4aff-c69d-d1f43d526231"
      },
      "source": [
        "[x[1] for x in TRAIN_DATA[1:10]]"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[{'entities': [(19, 33, 'ACTOR'), (43, 48, 'YEAR')]},\n",
              " {'entities': [(25, 34, 'ACTOR'), (39, 52, 'ACTOR')]},\n",
              " {'entities': [(39, 51, 'ACTOR'), (56, 67, 'ACTOR')]},\n",
              " {'entities': []},\n",
              " {'entities': [(17, 28, 'TITLE')]},\n",
              " {'entities': [(8, 29, 'GENRE'), (42, 58, 'DIRECTOR')]},\n",
              " {'entities': [(16, 25, 'GENRE'), (38, 51, 'DIRECTOR')]},\n",
              " {'entities': [(5, 24, 'SONG')]},\n",
              " {'entities': [(14, 19, 'ACTOR'), (26, 39, 'PLOT')]}]"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 7
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mhZujz4xKpCA",
        "colab_type": "text"
      },
      "source": [
        "<img src=\"https://storage.googleapis.com/kf-pipeline-contrib-public/release-0.1.3/kfp-components/notebooks/entity_extraction/assets/fig3.png\" align = \"left\" width=\"500\"/>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VMbkV-UuKpCA",
        "colab_type": "text"
      },
      "source": [
        "### Test pre-trained NER Model\n",
        "\n",
        "First, download the pre-trained model with a subprocess call."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IstyQlEBKpCB",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 506
        },
        "outputId": "54674c05-6d25-47f1-e75a-f484a36463b0"
      },
      "source": [
        "!{sys.executable} -m spacy download en"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Requirement already satisfied: en_core_web_sm==2.2.5 from https://github.com/explosion/spacy-models/releases/download/en_core_web_sm-2.2.5/en_core_web_sm-2.2.5.tar.gz#egg=en_core_web_sm==2.2.5 in /usr/local/lib/python3.6/dist-packages (2.2.5)\n",
            "Requirement already satisfied: spacy>=2.2.2 in /usr/local/lib/python3.6/dist-packages (from en_core_web_sm==2.2.5) (2.2.4)\n",
            "Requirement already satisfied: murmurhash<1.1.0,>=0.28.0 in /usr/local/lib/python3.6/dist-packages (from spacy>=2.2.2->en_core_web_sm==2.2.5) (1.0.2)\n",
            "Requirement already satisfied: requests<3.0.0,>=2.13.0 in /usr/local/lib/python3.6/dist-packages (from spacy>=2.2.2->en_core_web_sm==2.2.5) (2.23.0)\n",
            "Requirement already satisfied: srsly<1.1.0,>=1.0.2 in /usr/local/lib/python3.6/dist-packages (from spacy>=2.2.2->en_core_web_sm==2.2.5) (1.0.2)\n",
            "Requirement already satisfied: blis<0.5.0,>=0.4.0 in /usr/local/lib/python3.6/dist-packages (from spacy>=2.2.2->en_core_web_sm==2.2.5) (0.4.1)\n",
            "Requirement already satisfied: numpy>=1.15.0 in /usr/local/lib/python3.6/dist-packages (from spacy>=2.2.2->en_core_web_sm==2.2.5) (1.18.5)\n",
            "Requirement already satisfied: tqdm<5.0.0,>=4.38.0 in /usr/local/lib/python3.6/dist-packages (from spacy>=2.2.2->en_core_web_sm==2.2.5) (4.41.1)\n",
            "Requirement already satisfied: setuptools in /usr/local/lib/python3.6/dist-packages (from spacy>=2.2.2->en_core_web_sm==2.2.5) (47.3.1)\n",
            "Requirement already satisfied: plac<1.2.0,>=0.9.6 in /usr/local/lib/python3.6/dist-packages (from spacy>=2.2.2->en_core_web_sm==2.2.5) (1.1.3)\n",
            "Requirement already satisfied: cymem<2.1.0,>=2.0.2 in /usr/local/lib/python3.6/dist-packages (from spacy>=2.2.2->en_core_web_sm==2.2.5) (2.0.3)\n",
            "Requirement already satisfied: wasabi<1.1.0,>=0.4.0 in /usr/local/lib/python3.6/dist-packages (from spacy>=2.2.2->en_core_web_sm==2.2.5) (0.6.0)\n",
            "Requirement already satisfied: catalogue<1.1.0,>=0.0.7 in /usr/local/lib/python3.6/dist-packages (from spacy>=2.2.2->en_core_web_sm==2.2.5) (1.0.0)\n",
            "Requirement already satisfied: preshed<3.1.0,>=3.0.2 in /usr/local/lib/python3.6/dist-packages (from spacy>=2.2.2->en_core_web_sm==2.2.5) (3.0.2)\n",
            "Requirement already satisfied: thinc==7.4.0 in /usr/local/lib/python3.6/dist-packages (from spacy>=2.2.2->en_core_web_sm==2.2.5) (7.4.0)\n",
            "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.6/dist-packages (from requests<3.0.0,>=2.13.0->spacy>=2.2.2->en_core_web_sm==2.2.5) (2.9)\n",
            "Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.6/dist-packages (from requests<3.0.0,>=2.13.0->spacy>=2.2.2->en_core_web_sm==2.2.5) (1.24.3)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.6/dist-packages (from requests<3.0.0,>=2.13.0->spacy>=2.2.2->en_core_web_sm==2.2.5) (2020.4.5.2)\n",
            "Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.6/dist-packages (from requests<3.0.0,>=2.13.0->spacy>=2.2.2->en_core_web_sm==2.2.5) (3.0.4)\n",
            "Requirement already satisfied: importlib-metadata>=0.20; python_version < \"3.8\" in /usr/local/lib/python3.6/dist-packages (from catalogue<1.1.0,>=0.0.7->spacy>=2.2.2->en_core_web_sm==2.2.5) (1.6.1)\n",
            "Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.6/dist-packages (from importlib-metadata>=0.20; python_version < \"3.8\"->catalogue<1.1.0,>=0.0.7->spacy>=2.2.2->en_core_web_sm==2.2.5) (3.1.0)\n",
            "\u001b[38;5;2m✔ Download and installation successful\u001b[0m\n",
            "You can now load the model via spacy.load('en_core_web_sm')\n",
            "\u001b[38;5;2m✔ Linking successful\u001b[0m\n",
            "/usr/local/lib/python3.6/dist-packages/en_core_web_sm -->\n",
            "/usr/local/lib/python3.6/dist-packages/spacy/data/en\n",
            "You can now load the model via spacy.load('en')\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BTa5mT2PKpCF",
        "colab_type": "text"
      },
      "source": [
        "The pretrained model fails to identify any genres, plots, actors, directors, characters, movie titles, or ratings present in the movie queries. Interestingly, it also fails to identify persons, works of art, and products! Clearly, the pretrained model does not fit this domain application, so we will train our own model from scratch.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0W3mPOhsKpCF",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 538
        },
        "outputId": "b5be12e1-cc95-4949-9c2c-1c2f21229604"
      },
      "source": [
        "from spacy import displacy\n",
        "import warnings\n",
        "warnings.filterwarnings(\"ignore\")\n",
        "nlp = spacy.load('en')\n",
        "TEST_DATA, _ = load_data_spacy(\"data/test.txt\")\n",
        "\n",
        "test_sentences = [x[0] for x in TEST_DATA[0:15]] # extract the sentences from [sentence, entity]\n",
        "for x in test_sentences:\n",
        "    doc = nlp(x)\n",
        "    displacy.render(doc, jupyter = True, style = \"ent\")\n",
        "warnings.filterwarnings(\"default\")"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "<span class=\"tex2jax_ignore\"><div class=\"entities\" style=\"line-height: 2.5; direction: ltr\">are there any good romantic comedies out right now</div></span>"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "<span class=\"tex2jax_ignore\"><div class=\"entities\" style=\"line-height: 2.5; direction: ltr\">show me a movie about cars that talk</div></span>"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "<span class=\"tex2jax_ignore\"><div class=\"entities\" style=\"line-height: 2.5; direction: ltr\">list the \n",
              "<mark class=\"entity\" style=\"background: #e4e7d2; padding: 0.45em 0.6em; margin: 0 0.25em; line-height: 1; border-radius: 0.35em;\">\n",
              "    five\n",
              "    <span style=\"font-size: 0.8em; font-weight: bold; line-height: 1; border-radius: 0.35em; text-transform: uppercase; vertical-align: middle; margin-left: 0.5rem\">CARDINAL</span>\n",
              "</mark>\n",
              " star rated movies starring \n",
              "<mark class=\"entity\" style=\"background: #aa9cfc; padding: 0.45em 0.6em; margin: 0 0.25em; line-height: 1; border-radius: 0.35em;\">\n",
              "    mel gibson\n",
              "    <span style=\"font-size: 0.8em; font-weight: bold; line-height: 1; border-radius: 0.35em; text-transform: uppercase; vertical-align: middle; margin-left: 0.5rem\">PERSON</span>\n",
              "</mark>\n",
              "</div></span>"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "<span class=\"tex2jax_ignore\"><div class=\"entities\" style=\"line-height: 2.5; direction: ltr\">what science fiction films have come out recently</div></span>"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "<span class=\"tex2jax_ignore\"><div class=\"entities\" style=\"line-height: 2.5; direction: ltr\">did the same director make all of the harry potter movies</div></span>"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "<span class=\"tex2jax_ignore\"><div class=\"entities\" style=\"line-height: 2.5; direction: ltr\">show me \n",
              "<mark class=\"entity\" style=\"background: #bfe1d9; padding: 0.45em 0.6em; margin: 0 0.25em; line-height: 1; border-radius: 0.35em;\">\n",
              "    1980s\n",
              "    <span style=\"font-size: 0.8em; font-weight: bold; line-height: 1; border-radius: 0.35em; text-transform: uppercase; vertical-align: middle; margin-left: 0.5rem\">DATE</span>\n",
              "</mark>\n",
              " action movies</div></span>"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "<span class=\"tex2jax_ignore\"><div class=\"entities\" style=\"line-height: 2.5; direction: ltr\">what is the name of the \n",
              "<mark class=\"entity\" style=\"background: #e4e7d2; padding: 0.45em 0.6em; margin: 0 0.25em; line-height: 1; border-radius: 0.35em;\">\n",
              "    third\n",
              "    <span style=\"font-size: 0.8em; font-weight: bold; line-height: 1; border-radius: 0.35em; text-transform: uppercase; vertical-align: middle; margin-left: 0.5rem\">ORDINAL</span>\n",
              "</mark>\n",
              " movie in the star trek series</div></span>"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "<span class=\"tex2jax_ignore\"><div class=\"entities\" style=\"line-height: 2.5; direction: ltr\">can you get a soundtrac for the harry potter films</div></span>"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "<span class=\"tex2jax_ignore\"><div class=\"entities\" style=\"line-height: 2.5; direction: ltr\">find me science fiction movies since \n",
              "<mark class=\"entity\" style=\"background: #bfe1d9; padding: 0.45em 0.6em; margin: 0 0.25em; line-height: 1; border-radius: 0.35em;\">\n",
              "    2005\n",
              "    <span style=\"font-size: 0.8em; font-weight: bold; line-height: 1; border-radius: 0.35em; text-transform: uppercase; vertical-align: middle; margin-left: 0.5rem\">DATE</span>\n",
              "</mark>\n",
              "</div></span>"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "<span class=\"tex2jax_ignore\"><div class=\"entities\" style=\"line-height: 2.5; direction: ltr\">what is the most current movie featuring mat damon</div></span>"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "<span class=\"tex2jax_ignore\"><div class=\"entities\" style=\"line-height: 2.5; direction: ltr\">show me films where \n",
              "<mark class=\"entity\" style=\"background: #aa9cfc; padding: 0.45em 0.6em; margin: 0 0.25em; line-height: 1; border-radius: 0.35em;\">\n",
              "    jim carrey\n",
              "    <span style=\"font-size: 0.8em; font-weight: bold; line-height: 1; border-radius: 0.35em; text-transform: uppercase; vertical-align: middle; margin-left: 0.5rem\">PERSON</span>\n",
              "</mark>\n",
              " is a detective</div></span>"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "<span class=\"tex2jax_ignore\"><div class=\"entities\" style=\"line-height: 2.5; direction: ltr\">did \n",
              "<mark class=\"entity\" style=\"background: #aa9cfc; padding: 0.45em 0.6em; margin: 0 0.25em; line-height: 1; border-radius: 0.35em;\">\n",
              "    george clooney\n",
              "    <span style=\"font-size: 0.8em; font-weight: bold; line-height: 1; border-radius: 0.35em; text-transform: uppercase; vertical-align: middle; margin-left: 0.5rem\">PERSON</span>\n",
              "</mark>\n",
              " make a musical in \n",
              "<mark class=\"entity\" style=\"background: #bfe1d9; padding: 0.45em 0.6em; margin: 0 0.25em; line-height: 1; border-radius: 0.35em;\">\n",
              "    the 1980s\n",
              "    <span style=\"font-size: 0.8em; font-weight: bold; line-height: 1; border-radius: 0.35em; text-transform: uppercase; vertical-align: middle; margin-left: 0.5rem\">DATE</span>\n",
              "</mark>\n",
              "</div></span>"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "<span class=\"tex2jax_ignore\"><div class=\"entities\" style=\"line-height: 2.5; direction: ltr\">show me films with both \n",
              "<mark class=\"entity\" style=\"background: #aa9cfc; padding: 0.45em 0.6em; margin: 0 0.25em; line-height: 1; border-radius: 0.35em;\">\n",
              "    matt damon\n",
              "    <span style=\"font-size: 0.8em; font-weight: bold; line-height: 1; border-radius: 0.35em; text-transform: uppercase; vertical-align: middle; margin-left: 0.5rem\">PERSON</span>\n",
              "</mark>\n",
              " ad \n",
              "<mark class=\"entity\" style=\"background: #aa9cfc; padding: 0.45em 0.6em; margin: 0 0.25em; line-height: 1; border-radius: 0.35em;\">\n",
              "    ben affleck\n",
              "    <span style=\"font-size: 0.8em; font-weight: bold; line-height: 1; border-radius: 0.35em; text-transform: uppercase; vertical-align: middle; margin-left: 0.5rem\">PERSON</span>\n",
              "</mark>\n",
              "</div></span>"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "<span class=\"tex2jax_ignore\"><div class=\"entities\" style=\"line-height: 2.5; direction: ltr\">what is the borrowers movie</div></span>"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "<span class=\"tex2jax_ignore\"><div class=\"entities\" style=\"line-height: 2.5; direction: ltr\">have u movie hm about to pg \n",
              "<mark class=\"entity\" style=\"background: #e4e7d2; padding: 0.45em 0.6em; margin: 0 0.25em; line-height: 1; border-radius: 0.35em;\">\n",
              "    18\n",
              "    <span style=\"font-size: 0.8em; font-weight: bold; line-height: 1; border-radius: 0.35em; text-transform: uppercase; vertical-align: middle; margin-left: 0.5rem\">CARDINAL</span>\n",
              "</mark>\n",
              "</div></span>"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xDO2w6OOKpCI",
        "colab_type": "text"
      },
      "source": [
        "<img src=\"https://storage.googleapis.com/kf-pipeline-contrib-public/release-0.1.3/kfp-components/notebooks/entity_extraction/assets/fig4.png\" align = \"left\" width=\"500\"/>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YcXsLg_gKpCI",
        "colab_type": "text"
      },
      "source": [
        "### Train and save custom NER model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GyCH7eDuKpCJ",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 194
        },
        "outputId": "31b0c0c3-6e2d-4e83-b421-87c31f738c31"
      },
      "source": [
        "# A simple decorator to log function processing time\n",
        "def timer(method):\n",
        "    def timed(*args, **kw):\n",
        "        ts = time.time()\n",
        "        result = method(*args, **kw)\n",
        "        te = time.time()\n",
        "        print(\"Completed in {} seconds\".format(int(te - ts)))\n",
        "        return result\n",
        "    return timed\n",
        "\n",
        "# Data must be of the form (sentence, {entities: [start, end, label]})\n",
        "@timer\n",
        "def train_spacy(train_data, labels, iterations, dropout = 0.2, display_freq = 1):\n",
        "    ''' Train a spacy NER model, which can be queried against with test data\n",
        "    \n",
        "    train_data : training data in the format of (sentence, {entities: [(start, end, label)]})\n",
        "    labels : a list of unique annotations\n",
        "    iterations : number of training iterations\n",
        "    dropout : dropout proportion for training\n",
        "    display_freq : number of epochs between logging losses to console\n",
        "    '''\n",
        "    nlp = spacy.blank('en') \n",
        "    if 'ner' not in nlp.pipe_names:\n",
        "        ner = nlp.create_pipe('ner')\n",
        "        nlp.add_pipe(ner)\n",
        "    \n",
        "    # Add entity labels to the NER pipeline\n",
        "    for i in labels:\n",
        "        ner.add_label(i)\n",
        "\n",
        "    # Disable other pipelines in SpaCy to only train NER\n",
        "    other_pipes = [pipe for pipe in nlp.pipe_names if pipe != 'ner']\n",
        "    with nlp.disable_pipes(*other_pipes):\n",
        "        nlp.vocab.vectors.name = 'spacy_model' # without this, spaCy throws an \"unnamed\" error\n",
        "        optimizer = nlp.begin_training()\n",
        "        for itr in range(iterations):\n",
        "            random.shuffle(train_data) # shuffle the training data before each iteration\n",
        "            losses = {}\n",
        "            batches = minibatch(train_data, size = compounding(4., 32., 1.001))\n",
        "            for batch in batches:\n",
        "                texts, annotations = zip(*batch)\n",
        "                nlp.update(           \n",
        "                    texts,\n",
        "                    annotations, \n",
        "                    drop = dropout,   \n",
        "                    sgd = optimizer,\n",
        "                    losses = losses)\n",
        "            if itr % display_freq == 0:\n",
        "                print(\"Iteration {} Loss: {}\".format(itr + 1, losses))\n",
        "    return nlp\n",
        "\n",
        "# Train (and save) the NER model\n",
        "ner = train_spacy(TRAIN_DATA, LABELS,2)\n",
        "ner.to_disk(\"models/spacy_example\")"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "/usr/lib/python3.6/importlib/_bootstrap.py:219: RuntimeWarning: numpy.ufunc size changed, may indicate binary incompatibility. Expected 176 from C header, got 216 from PyObject\n",
            "  return f(*args, **kwds)\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "Iteration 1 Loss: {'ner': 19436.620473670202}\n",
            "Iteration 2 Loss: {'ner': 12923.654009589925}\n",
            "Completed in 72 seconds\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "/usr/local/lib/python3.6/dist-packages/spacy/language.py:900: ResourceWarning: unclosed file <_io.TextIOWrapper name='models/spacy_example/meta.json' mode='w' encoding='UTF-8'>\n",
            "  srsly.json_dumps(self.meta)\n",
            "/usr/local/lib/python3.6/dist-packages/spacy/util.py:645: ResourceWarning: unclosed file <_io.BufferedWriter name='models/spacy_example/vocab/vectors'>\n",
            "  writer(path / key)\n"
          ],
          "name": "stderr"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CLJSi1wXKpCL",
        "colab_type": "text"
      },
      "source": [
        "<img src=\"https://storage.googleapis.com/kf-pipeline-contrib-public/release-0.1.3/kfp-components/notebooks/entity_extraction/assets/fig5.png\" align = \"left\" width=\"400\"/>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Dzu3CrkYKpCL",
        "colab_type": "text"
      },
      "source": [
        "### Test model on new sentences"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5R6fcVyjKpCM",
        "colab_type": "text"
      },
      "source": [
        "To test the model on new sentences, the `load_model` function is used to reload the trained model weights, and `load_data` is called to load and transform the test data. The spacy function `displacy` is used to visualize the predictions of the first 15 test sentences. As the results show, the architecture has learned good representations of the entities. However, there still exist a few errors. While some of these may be mitigated with increased training time (the loss was still decreasing rapidly after 5 iterations), others may require additional pre-processing, such as fixing spelling mistakes."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3fSFQD_ZKpCM",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 592
        },
        "outputId": "446997aa-98aa-4bbe-e7f1-813951dd2784"
      },
      "source": [
        "from spacy import displacy\n",
        "\n",
        "def load_model(model_path):\n",
        "    ''' Loads a pre-trained model for prediction on new test sentences\n",
        "    \n",
        "    model_path : directory of model saved by spacy.to_disk\n",
        "    '''\n",
        "    nlp = spacy.blank('en') \n",
        "    if 'ner' not in nlp.pipe_names:\n",
        "        ner = nlp.create_pipe('ner')\n",
        "        nlp.add_pipe(ner)\n",
        "    ner = nlp.from_disk(model_path)\n",
        "    return ner\n",
        "\n",
        "ner = load_model(\"models/spacy_example\")\n",
        "\n",
        "TEST_DATA, _ = load_data_spacy(\"data/test.txt\")\n",
        "\n",
        "test_sentences = [x[0] for x in TEST_DATA[0:15]] # extract the sentences from [sentence, entity]\n",
        "for x in test_sentences:\n",
        "    doc = ner(x)\n",
        "    displacy.render(doc, jupyter = True, style = \"ent\")"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "<span class=\"tex2jax_ignore\"><div class=\"entities\" style=\"line-height: 2.5; direction: ltr\">are there any good \n",
              "<mark class=\"entity\" style=\"background: #ddd; padding: 0.45em 0.6em; margin: 0 0.25em; line-height: 1; border-radius: 0.35em;\">\n",
              "    romantic comedies\n",
              "    <span style=\"font-size: 0.8em; font-weight: bold; line-height: 1; border-radius: 0.35em; text-transform: uppercase; vertical-align: middle; margin-left: 0.5rem\">GENRE</span>\n",
              "</mark>\n",
              " out right now</div></span>"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "<span class=\"tex2jax_ignore\"><div class=\"entities\" style=\"line-height: 2.5; direction: ltr\">show me a movie about \n",
              "<mark class=\"entity\" style=\"background: #ddd; padding: 0.45em 0.6em; margin: 0 0.25em; line-height: 1; border-radius: 0.35em;\">\n",
              "    cars\n",
              "    <span style=\"font-size: 0.8em; font-weight: bold; line-height: 1; border-radius: 0.35em; text-transform: uppercase; vertical-align: middle; margin-left: 0.5rem\">PLOT</span>\n",
              "</mark>\n",
              " that talk</div></span>"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "<span class=\"tex2jax_ignore\"><div class=\"entities\" style=\"line-height: 2.5; direction: ltr\">list the \n",
              "<mark class=\"entity\" style=\"background: #ddd; padding: 0.45em 0.6em; margin: 0 0.25em; line-height: 1; border-radius: 0.35em;\">\n",
              "    five star\n",
              "    <span style=\"font-size: 0.8em; font-weight: bold; line-height: 1; border-radius: 0.35em; text-transform: uppercase; vertical-align: middle; margin-left: 0.5rem\">RATINGS_AVERAGE</span>\n",
              "</mark>\n",
              " rated movies starring \n",
              "<mark class=\"entity\" style=\"background: #ddd; padding: 0.45em 0.6em; margin: 0 0.25em; line-height: 1; border-radius: 0.35em;\">\n",
              "    mel gibson\n",
              "    <span style=\"font-size: 0.8em; font-weight: bold; line-height: 1; border-radius: 0.35em; text-transform: uppercase; vertical-align: middle; margin-left: 0.5rem\">ACTOR</span>\n",
              "</mark>\n",
              "</div></span>"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "<span class=\"tex2jax_ignore\"><div class=\"entities\" style=\"line-height: 2.5; direction: ltr\">what \n",
              "<mark class=\"entity\" style=\"background: #ddd; padding: 0.45em 0.6em; margin: 0 0.25em; line-height: 1; border-radius: 0.35em;\">\n",
              "    science fiction\n",
              "    <span style=\"font-size: 0.8em; font-weight: bold; line-height: 1; border-radius: 0.35em; text-transform: uppercase; vertical-align: middle; margin-left: 0.5rem\">GENRE</span>\n",
              "</mark>\n",
              " films have come out recently</div></span>"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "<span class=\"tex2jax_ignore\"><div class=\"entities\" style=\"line-height: 2.5; direction: ltr\">did the same director make all of the \n",
              "<mark class=\"entity\" style=\"background: #ddd; padding: 0.45em 0.6em; margin: 0 0.25em; line-height: 1; border-radius: 0.35em;\">\n",
              "    harry potter\n",
              "    <span style=\"font-size: 0.8em; font-weight: bold; line-height: 1; border-radius: 0.35em; text-transform: uppercase; vertical-align: middle; margin-left: 0.5rem\">TITLE</span>\n",
              "</mark>\n",
              " movies</div></span>"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "<span class=\"tex2jax_ignore\"><div class=\"entities\" style=\"line-height: 2.5; direction: ltr\">show me \n",
              "<mark class=\"entity\" style=\"background: #ddd; padding: 0.45em 0.6em; margin: 0 0.25em; line-height: 1; border-radius: 0.35em;\">\n",
              "    1980s\n",
              "    <span style=\"font-size: 0.8em; font-weight: bold; line-height: 1; border-radius: 0.35em; text-transform: uppercase; vertical-align: middle; margin-left: 0.5rem\">YEAR</span>\n",
              "</mark>\n",
              " \n",
              "<mark class=\"entity\" style=\"background: #ddd; padding: 0.45em 0.6em; margin: 0 0.25em; line-height: 1; border-radius: 0.35em;\">\n",
              "    action\n",
              "    <span style=\"font-size: 0.8em; font-weight: bold; line-height: 1; border-radius: 0.35em; text-transform: uppercase; vertical-align: middle; margin-left: 0.5rem\">GENRE</span>\n",
              "</mark>\n",
              " movies</div></span>"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "<span class=\"tex2jax_ignore\"><div class=\"entities\" style=\"line-height: 2.5; direction: ltr\">what is the name of the third movie in the \n",
              "<mark class=\"entity\" style=\"background: #ddd; padding: 0.45em 0.6em; margin: 0 0.25em; line-height: 1; border-radius: 0.35em;\">\n",
              "    star trek series\n",
              "    <span style=\"font-size: 0.8em; font-weight: bold; line-height: 1; border-radius: 0.35em; text-transform: uppercase; vertical-align: middle; margin-left: 0.5rem\">TITLE</span>\n",
              "</mark>\n",
              "</div></span>"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "<span class=\"tex2jax_ignore\"><div class=\"entities\" style=\"line-height: 2.5; direction: ltr\">can you get a soundtrac for the \n",
              "<mark class=\"entity\" style=\"background: #ddd; padding: 0.45em 0.6em; margin: 0 0.25em; line-height: 1; border-radius: 0.35em;\">\n",
              "    harry potter\n",
              "    <span style=\"font-size: 0.8em; font-weight: bold; line-height: 1; border-radius: 0.35em; text-transform: uppercase; vertical-align: middle; margin-left: 0.5rem\">TITLE</span>\n",
              "</mark>\n",
              " films</div></span>"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "<span class=\"tex2jax_ignore\"><div class=\"entities\" style=\"line-height: 2.5; direction: ltr\">find me \n",
              "<mark class=\"entity\" style=\"background: #ddd; padding: 0.45em 0.6em; margin: 0 0.25em; line-height: 1; border-radius: 0.35em;\">\n",
              "    science fiction\n",
              "    <span style=\"font-size: 0.8em; font-weight: bold; line-height: 1; border-radius: 0.35em; text-transform: uppercase; vertical-align: middle; margin-left: 0.5rem\">GENRE</span>\n",
              "</mark>\n",
              " movies \n",
              "<mark class=\"entity\" style=\"background: #ddd; padding: 0.45em 0.6em; margin: 0 0.25em; line-height: 1; border-radius: 0.35em;\">\n",
              "    since 2005\n",
              "    <span style=\"font-size: 0.8em; font-weight: bold; line-height: 1; border-radius: 0.35em; text-transform: uppercase; vertical-align: middle; margin-left: 0.5rem\">YEAR</span>\n",
              "</mark>\n",
              "</div></span>"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "<span class=\"tex2jax_ignore\"><div class=\"entities\" style=\"line-height: 2.5; direction: ltr\">what is the most current movie featuring \n",
              "<mark class=\"entity\" style=\"background: #ddd; padding: 0.45em 0.6em; margin: 0 0.25em; line-height: 1; border-radius: 0.35em;\">\n",
              "    mat damon\n",
              "    <span style=\"font-size: 0.8em; font-weight: bold; line-height: 1; border-radius: 0.35em; text-transform: uppercase; vertical-align: middle; margin-left: 0.5rem\">ACTOR</span>\n",
              "</mark>\n",
              "</div></span>"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "<span class=\"tex2jax_ignore\"><div class=\"entities\" style=\"line-height: 2.5; direction: ltr\">show me films where \n",
              "<mark class=\"entity\" style=\"background: #ddd; padding: 0.45em 0.6em; margin: 0 0.25em; line-height: 1; border-radius: 0.35em;\">\n",
              "    jim carrey\n",
              "    <span style=\"font-size: 0.8em; font-weight: bold; line-height: 1; border-radius: 0.35em; text-transform: uppercase; vertical-align: middle; margin-left: 0.5rem\">ACTOR</span>\n",
              "</mark>\n",
              " is a detective</div></span>"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "<span class=\"tex2jax_ignore\"><div class=\"entities\" style=\"line-height: 2.5; direction: ltr\">did \n",
              "<mark class=\"entity\" style=\"background: #ddd; padding: 0.45em 0.6em; margin: 0 0.25em; line-height: 1; border-radius: 0.35em;\">\n",
              "    george clooney\n",
              "    <span style=\"font-size: 0.8em; font-weight: bold; line-height: 1; border-radius: 0.35em; text-transform: uppercase; vertical-align: middle; margin-left: 0.5rem\">ACTOR</span>\n",
              "</mark>\n",
              " make a \n",
              "<mark class=\"entity\" style=\"background: #ddd; padding: 0.45em 0.6em; margin: 0 0.25em; line-height: 1; border-radius: 0.35em;\">\n",
              "    musical\n",
              "    <span style=\"font-size: 0.8em; font-weight: bold; line-height: 1; border-radius: 0.35em; text-transform: uppercase; vertical-align: middle; margin-left: 0.5rem\">GENRE</span>\n",
              "</mark>\n",
              " in the \n",
              "<mark class=\"entity\" style=\"background: #ddd; padding: 0.45em 0.6em; margin: 0 0.25em; line-height: 1; border-radius: 0.35em;\">\n",
              "    1980s\n",
              "    <span style=\"font-size: 0.8em; font-weight: bold; line-height: 1; border-radius: 0.35em; text-transform: uppercase; vertical-align: middle; margin-left: 0.5rem\">YEAR</span>\n",
              "</mark>\n",
              "</div></span>"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "<span class=\"tex2jax_ignore\"><div class=\"entities\" style=\"line-height: 2.5; direction: ltr\">show me films with both \n",
              "<mark class=\"entity\" style=\"background: #ddd; padding: 0.45em 0.6em; margin: 0 0.25em; line-height: 1; border-radius: 0.35em;\">\n",
              "    matt damon\n",
              "    <span style=\"font-size: 0.8em; font-weight: bold; line-height: 1; border-radius: 0.35em; text-transform: uppercase; vertical-align: middle; margin-left: 0.5rem\">ACTOR</span>\n",
              "</mark>\n",
              " ad \n",
              "<mark class=\"entity\" style=\"background: #ddd; padding: 0.45em 0.6em; margin: 0 0.25em; line-height: 1; border-radius: 0.35em;\">\n",
              "    ben affleck\n",
              "    <span style=\"font-size: 0.8em; font-weight: bold; line-height: 1; border-radius: 0.35em; text-transform: uppercase; vertical-align: middle; margin-left: 0.5rem\">ACTOR</span>\n",
              "</mark>\n",
              "</div></span>"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "/usr/lib/python3.6/runpy.py:193: UserWarning: [W006] No entities to visualize found in Doc object. If this is surprising to you, make sure the Doc was processed using a model that supports named entity recognition, and check the `doc.ents` property manually if necessary.\n",
            "  \"__main__\", mod_spec)\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "<span class=\"tex2jax_ignore\"><div class=\"entities\" style=\"line-height: 2.5; direction: ltr\">what is the borrowers movie</div></span>"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "<span class=\"tex2jax_ignore\"><div class=\"entities\" style=\"line-height: 2.5; direction: ltr\">have u movie hm about to \n",
              "<mark class=\"entity\" style=\"background: #ddd; padding: 0.45em 0.6em; margin: 0 0.25em; line-height: 1; border-radius: 0.35em;\">\n",
              "    pg 18\n",
              "    <span style=\"font-size: 0.8em; font-weight: bold; line-height: 1; border-radius: 0.35em; text-transform: uppercase; vertical-align: middle; margin-left: 0.5rem\">RATING</span>\n",
              "</mark>\n",
              "</div></span>"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lcGgoTYHKpCP",
        "colab_type": "text"
      },
      "source": [
        "<img src=\"https://storage.googleapis.com/kf-pipeline-contrib-public/release-0.1.3/kfp-components/notebooks/entity_extraction/assets/fig6.png\" align = \"left\" width=\"500\"/>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PUm2JuW3KpCP",
        "colab_type": "text"
      },
      "source": [
        "### Evaluation Metrics\n",
        "\n",
        "Model performance is assessed on the entirety of the test dataset (2,443 sentences) based on the following metrics and their definitions.\n",
        "\n",
        "   * Precision: true positives / (true positives + false positives)\n",
        "   * Recall: true positives / (true positives + false negatives)\n",
        "   * F1-score: harmonic average of precision and recall"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jOzt9dG3KpCP",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def calc_precision(pred, true):        \n",
        "    precision = len([x for x in pred if x in true]) / (len(pred) + 1e-20) # true positives / total pred\n",
        "    return precision\n",
        "\n",
        "def calc_recall(pred, true):\n",
        "    recall = len([x for x in true if x in pred]) / (len(true) + 1e-20)    # true positives / total test\n",
        "    return recall\n",
        "\n",
        "def calc_f1(precision, recall):\n",
        "    f1 = 2 * ((precision * recall) / (precision + recall + 1e-20))\n",
        "    return f1"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cT1SIRHUKpCR",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 70
        },
        "outputId": "c02061da-ad05-4dd2-cbc1-b441495b36f6"
      },
      "source": [
        "from itertools import chain\n",
        "\n",
        "# run the predictions on each sentence in the test dataset, and return the spacy object\n",
        "preds = [ner(x[0]) for x in TEST_DATA]\n",
        "\n",
        "precisions, recalls, f1s = [], [], [] \n",
        "\n",
        "# iterate over predictions and test data and calculate precision, recall, and F1-score\n",
        "for pred, true in zip(preds, TEST_DATA): \n",
        "    true = [x[2] for x in list(chain.from_iterable(true[1].values()))] # x[2] = annotation, true[1] = (start, end, annot)\n",
        "    pred = [i.label_ for i in pred.ents] # i.label_ = annotation label, pred.ents = list of annotations\n",
        "    precision = calc_precision(true, pred)\n",
        "    precisions.append(precision)\n",
        "    recall = calc_recall(true, pred)\n",
        "    recalls.append(recall)\n",
        "    f1s.append(calc_f1(precision, recall))\n",
        "    \n",
        "print(\"Precision: {} \\nRecall: {} \\nF1-score: {}\".format(np.around(np.mean(precisions), 3),\n",
        "                                                         np.around(np.mean(recalls), 3),\n",
        "                                                         np.around(np.mean(f1s), 3)))"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Precision: 0.855 \n",
            "Recall: 0.871 \n",
            "F1-score: 0.855\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "scrolled": true,
        "id": "8Puzx8QjKpCV",
        "colab_type": "text"
      },
      "source": [
        "<img src=\"https://storage.googleapis.com/kf-pipeline-contrib-public/release-0.1.3/kfp-components/notebooks/entity_extraction/assets/fig7.png\" align = \"left\" width=\"200\"/>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Mmh4vRhNKpCW",
        "colab_type": "text"
      },
      "source": [
        "# 3. Named entity recognition with Tensorflow\n",
        "\n",
        "This section focuses on developing, training, and serving a custom NER architecture with Tensorflow 1.9.0. We will implement an LSTM-CRF model as described in [Huang, Xu, and Yu, 2015](https://arxiv.org/pdf/1508.01991.pdf).\n",
        "\n",
        "This approach can be broken down into its constituent parts as follows:\n",
        "   * Embedding: Generating a dense vector representation of words\n",
        "   * LSTM: Incorporating past and future features to generate a representation of each time step\n",
        "   * CRF: Make use of neighboring information to predict current tags. The CRF approach has been shown to provide higher accuracy than maximum entropy models because CRF considers the entire sentence rather than relying on beam search to find optimal context sizes. \n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9WbJrVW_KpCW",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        },
        "outputId": "40b362f8-9eef-49cf-a6b8-07b1458667a7"
      },
      "source": [
        "import tensorflow as tf\n",
        "from functools import partial\n",
        "from keras.preprocessing.text import Tokenizer\n",
        "from keras.preprocessing.sequence import pad_sequences\n",
        "from keras.utils import to_categorical\n",
        "import pickle\n",
        "\n",
        "params = {\n",
        "    'dim' : 300,            # dimension of embeddings\n",
        "    'maximum_steps' : 1000, # number of training steps        \n",
        "    'lstm_size' : 150,      # dimension of LSTM\n",
        "    'batch_size' : 25,      # batch size\n",
        "    'max_words' : 10000,    # maximum number of words to embed\n",
        "    'padding_size' : 20,    # maximum sentence size\n",
        "    'num_classes' : 14,     # number of unique classes\n",
        "    'save_dir' : 'models/' # directory to save hash tables, model weights, etc.\n",
        "}"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Using TensorFlow backend.\n"
          ],
          "name": "stderr"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MWu89J-LKpCa",
        "colab_type": "text"
      },
      "source": [
        "The first step in implementing a tensorflow named entity recognition architecture is to specify the data loading and transformation process. The words and labels need to be transformed to an integer vector format that tensorflow can process. Tokenization is used to do this, where unique words and labels are mapped to integers and the mapping is stored in a hashtable for back-conversion. \n",
        "\n",
        "For this process to work, however, we have to see all of the training data all at once to prevent overlapping hashes. This means that this tokenization process needs to happen separately from the training process. The `make_tokenizer` function takes in the training data and labels and returns two dictionaries, `word_index`, and `labels_index`. The former specifies integer mappings for the words, and the latter specifies integer mappings for the labels. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KpVnYzzJzo9J",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 55
        },
        "outputId": "26e47e6a-2d43-4657-fb17-cc319d081e4a"
      },
      "source": [
        "#visualizing raw data\n",
        "f = open('data/train.txt')\n",
        "lines = list()\n",
        "for line in f.readlines():\n",
        "  lines.append(line)\n",
        "print(lines[:10])\n"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "['O\\twhat\\n', 'O\\tmovies\\n', 'O\\tstar\\n', 'B-ACTOR\\tbruce\\n', 'I-ACTOR\\twillis\\n', '\\n', 'O\\tshow\\n', 'O\\tme\\n', 'O\\tfilms\\n', 'O\\twith\\n']\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZO5nlFeZKpCa",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def save_obj(directory, obj, name):\n",
        "    '''Helper function using pickle to save and load objects'''\n",
        "    with open(directory + name + '.pkl', 'wb+') as f:\n",
        "        pickle.dump(obj, f, pickle.HIGHEST_PROTOCOL)\n",
        "\n",
        "def load_obj(directory, name):\n",
        "    '''Helper function using pickle to save and load objects'''\n",
        "    with open(directory + name + \".pkl\", \"rb\") as f:\n",
        "        return pickle.load(f)\n",
        "    \n",
        "def load_data(file = \"data/train.txt\"):\n",
        "    '''Helper function to load and transform inputs and labels\n",
        "    included as a separate function due to NER-specific evaluation needs:\n",
        "        tensorflow does not have multi-class precision/accuracy as a metric\n",
        "        so data_y is needed to manually calculate evaluations'''\n",
        "    file = open(file, 'r')\n",
        "    sentence, labels = [], []\n",
        "    data_x, data_y = [], []\n",
        "    for line in file:\n",
        "        line = line.strip(\"\\n\").split(\"\\t\")\n",
        "        \n",
        "        # lines with len > 1 are words\n",
        "        if len(line) > 1:\n",
        "            sentence.append(line[1])\n",
        "            labels.append(line[0][2:]) if len(line[0]) > 1 else labels.append(line[0])\n",
        "        \n",
        "        # lins with len == 1 are sentence breaks\n",
        "        if len(line) == 1: \n",
        "            data_x.append(' '.join(sentence))\n",
        "            data_y.append(labels)\n",
        "            sentence, labels = [], []\n",
        "    return data_x, data_y\n",
        "\n",
        "def make_tokenizer(file = \"data/train.txt\", params = params):\n",
        "    ''' In order for one hot encoding of words and labels to work, \n",
        "    every word and label has to be seen at least once to make a hashing table.\n",
        "    This function outputs hash tables for the words and the labels\n",
        "    that can be used to one-hot-encode them in the generator\n",
        "    '''\n",
        "    # Load parameters and data\n",
        "    max_words = params['max_words']\n",
        "    padding_size = params['padding_size']\n",
        "    save_dir = params['save_dir']\n",
        "    data_x, data_y = load_data(file)\n",
        "            \n",
        "    # Use the Keras tokenizer API to generate hashing table for data_x\n",
        "    tokenizer = Tokenizer(num_words = max_words)\n",
        "    \n",
        "    tokenizer.fit_on_texts(data_x)\n",
        "    word_index = tokenizer.word_index\n",
        "    \n",
        "    # Flatten data_y and create hashing table using set logic\n",
        "    data_y_flattened = [item for sublist in data_y for item in sublist]\n",
        "    data_x_flattened = [item for sublist in data_x for item in sublist]\n",
        "    \n",
        "    labels_index = dict([(y, x + 1) for x, y in enumerate(sorted(set(data_y_flattened)))])\n",
        "    labels = []\n",
        "    for item in data_y:\n",
        "        labels.append([labels_index.get(i) for i in item])\n",
        "    labels_lookup = {v : k for k, v in labels_index.items()} # reverse dictionary for lookup\n",
        "    # save hash tables to disk for model serving\n",
        "    for item, name in zip([word_index, labels_index, labels_lookup],\n",
        "                          [\"word_index\", \"labels_index\", \"labels_lookup\"]):\n",
        "        save_obj(save_dir, item, name)\n",
        "    return word_index, labels_index, labels_lookup\n",
        "\n",
        "word_index, labels_index, labels_lookup = make_tokenizer()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "281Nma0ONCaC",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "data_x, data_y= load_data(file = \"data/train.txt\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ynMtFytTNIxp",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 399
        },
        "outputId": "f15b500f-32a4-4882-b7cc-32b37c67eb78"
      },
      "source": [
        "data_y[:10]"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[['O', 'O', 'O', 'ACTOR', 'ACTOR'],\n",
              " ['O', 'O', 'O', 'O', 'ACTOR', 'ACTOR', 'O', 'O', 'YEAR'],\n",
              " ['O', 'O', 'O', 'O', 'ACTOR', 'ACTOR', 'O', 'ACTOR', 'ACTOR'],\n",
              " ['O',\n",
              "  'O',\n",
              "  'O',\n",
              "  'O',\n",
              "  'O',\n",
              "  'O',\n",
              "  'O',\n",
              "  'O',\n",
              "  'ACTOR',\n",
              "  'ACTOR',\n",
              "  'O',\n",
              "  'ACTOR',\n",
              "  'ACTOR'],\n",
              " ['O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O'],\n",
              " ['O', 'O', 'O', 'TITLE', 'O', 'O', 'O'],\n",
              " ['O', 'O', 'GENRE', 'GENRE', 'GENRE', 'O', 'O', 'DIRECTOR', 'DIRECTOR'],\n",
              " ['O', 'O', 'O', 'O', 'GENRE', 'O', 'O', 'DIRECTOR', 'DIRECTOR'],\n",
              " ['O', 'SONG', 'SONG', 'SONG', 'O', 'O', 'O', 'O', 'O', 'O'],\n",
              " ['O', 'O', 'O', 'ACTOR', 'O', 'PLOT', 'PLOT', 'PLOT']]"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 16
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "F0BuF5Aw1Xrp",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 191
        },
        "outputId": "f8eee8d7-08ee-48eb-d776-c9b73e84deff"
      },
      "source": [
        "data_x[:10]"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "['what movies star bruce willis',\n",
              " 'show me films with drew barrymore from the 1980s',\n",
              " 'what movies starred both al pacino and robert deniro',\n",
              " 'find me all of the movies that starred harold ramis and bill murray',\n",
              " 'find me a movie with a quote about baseball in it',\n",
              " 'what movies have mississippi in the title',\n",
              " 'show me science fiction films directed by steven spielberg',\n",
              " 'do you have any thrillers directed by sofia coppola',\n",
              " 'what leonard cohen songs have been used in a movie',\n",
              " 'show me films elvis films set in hawaii']"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 13
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xeQBajm7KpCe",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 243
        },
        "outputId": "5326af21-611e-41f8-d4f1-02bee055e053"
      },
      "source": [
        "labels_index"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'ACTOR': 1,\n",
              " 'CHARACTER': 2,\n",
              " 'DIRECTOR': 3,\n",
              " 'GENRE': 4,\n",
              " 'O': 5,\n",
              " 'PLOT': 6,\n",
              " 'RATING': 7,\n",
              " 'RATINGS_AVERAGE': 8,\n",
              " 'REVIEW': 9,\n",
              " 'SONG': 10,\n",
              " 'TITLE': 11,\n",
              " 'TRAILER': 12,\n",
              " 'YEAR': 13}"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 16
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Lg9kO4a0KpCg",
        "colab_type": "text"
      },
      "source": [
        "<img src=\"https://storage.googleapis.com/kf-pipeline-contrib-public/release-0.1.3/kfp-components/notebooks/entity_extraction/assets/fig8.png\" align = \"left\" width=\"200\"/>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "q8E4F0ylKpCg",
        "colab_type": "text"
      },
      "source": [
        "Next, we create a generator to serve as input to the tensorflow DataSet API. The `generate_batches` function takes training data in BIO format and yields batches as input to the model function. The DataSet API requires two inputs - features and labels. For a recurrent neural network, we also need to specify sequence lengths to mask variable length sequences. This length is returned as a tuple in the features, as `(batch_x, lengths)`."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "j_BgbdM3KpCh",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def generate_batches(file = \"data/train.txt\", params = params, train = True):\n",
        "    ''' Generate minibatch with dimensions:\n",
        "    batch_x : (batch_size, max_len)\n",
        "    lengths : (batch_size,)\n",
        "    batch_y : (batch_size, num_classes)\n",
        "    \n",
        "    file : path to .txt containing training data in BIO format\n",
        "    '''\n",
        "    \n",
        "    batch_size = params['batch_size']\n",
        "    max_len = params['padding_size']\n",
        "    save_dir = params['save_dir']\n",
        "    \n",
        "    # load hash tables for tokenization\n",
        "    for item, name in zip([word_index, labels_index, labels_lookup],\n",
        "                          [\"word_index\", \"labels_index\", \"labels_lookup\"]):\n",
        "        item = load_obj(save_dir, name)\n",
        "    \n",
        "    while True:\n",
        "        with open(file, 'r') as f:\n",
        "            batch_x, lengths, batch_y = [], [], []\n",
        "            words, labels = [], []\n",
        "            for line in f:\n",
        "                line = line.strip(\"\\n\").split(\"\\t\")\n",
        "                # lines with len > 1 are words\n",
        "                if len(line) > 1:\n",
        "                    labels.append(line[0][2:]) if len(line[0]) > 1 else labels.append(line[0])\n",
        "                    words.append(line[1])\n",
        "\n",
        "                # lines with len == 1 are breaks between sentences\n",
        "                if len(line) == 1: \n",
        "                    words = [word_index.get(x) if x in word_index.keys() else 0 for x in words]\n",
        "                    labels = [labels_index.get(y) for y in labels]\n",
        "                    batch_x.append(words)\n",
        "                    batch_y.append(labels)\n",
        "                    lengths.append(min(len(words), max_len))\n",
        "                    words, labels = [], []\n",
        "\n",
        "                if len(batch_x) == batch_size:\n",
        "                    batch_x = pad_sequences(batch_x, maxlen = max_len, value = 0, padding = \"post\")\n",
        "                    batch_y = pad_sequences(batch_y, maxlen = max_len, value = 0, padding = \"post\")\n",
        "                    yield (batch_x, lengths), batch_y \n",
        "                    batch_x, lengths, batch_y = [], [], []\n",
        "            if train == False:\n",
        "                break\n",
        "    print('batch_x: ', batch_x)\n",
        "    print('batch_y: ', batch_y)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JFwssv9xL0FE",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        },
        "outputId": "4f6e4f9d-e37f-4dc3-a87b-9a3030bbef76"
      },
      "source": [
        "generate_batches(file = \"data/train.txt\", params = params, train = False)\n"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<generator object generate_batches at 0x7f47c7779410>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 25
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "X2lA4sOSKpCl",
        "colab_type": "text"
      },
      "source": [
        "The estimator API requires an input function and a model function. The `input_fn` maps the `generate_batches` generator to a tensorflow Dataset."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vkIodZPbKpCl",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# For model training, we need an input function that will feed a tf.Dataset\n",
        "def input_fn(file, params = None, train = True):\n",
        "    params = params if params is not None else {}\n",
        "    shapes = (([None, None], [None]), [None, None]) # batch_x, lengths, batch_y shapes\n",
        "    types = ((tf.int32, tf.int32), tf.int32)        # batch_x, lengths, batch_y data types\n",
        "    \n",
        "    generator = partial(generate_batches, file, train = train)\n",
        "    dataset = tf.data.Dataset.from_generator(generator, types, shapes)\n",
        "    return dataset\n",
        "\n",
        "# For model serving, we need a serving function that will feed tf.placeholders\n",
        "def serving_input_fn():\n",
        "    words = tf.placeholder(dtype=tf.int32, shape=[None, None], name='words')\n",
        "    length = tf.placeholder(dtype=tf.int32, shape=[None], name='length')\n",
        "    receiver_tensors = {'words': words, 'length': length}\n",
        "    features = {'words': words, 'length': length}\n",
        "    return tf.estimator.export.ServingInputReceiver(features, receiver_tensors)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nfaf9fVUKpCo",
        "colab_type": "text"
      },
      "source": [
        "The `model_fn` unpacks features and labels to create the specified model architecture, which is an LSTM-CRF."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4Gr9X-dvKpCq",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def model_fn(features, labels, mode, params = params):\n",
        "    # import the data and unpack the features\n",
        "    # serving input_fn returns a dict, convert to multivalue obj\n",
        "    if isinstance(features, dict):\n",
        "        features = features['words'], features['length']\n",
        "    \n",
        "    words, length = features\n",
        "    \n",
        "    # Embedding\n",
        "    embedding = tf.Variable(tf.random_normal([params['max_words'], params['dim']]))\n",
        "    embedding_lookup_for_x = tf.nn.embedding_lookup(embedding, words)\n",
        "    \n",
        "    # LSTM\n",
        "    lstm_cell_fw = tf.nn.rnn_cell.BasicLSTMCell(params['lstm_size'], state_is_tuple = True)\n",
        "    lstm_cell_bw = tf.nn.rnn_cell.BasicLSTMCell(params['lstm_size'], state_is_tuple = True)\n",
        "    states, final_state = tf.nn.bidirectional_dynamic_rnn(\n",
        "                                        cell_fw = lstm_cell_fw, \n",
        "                                        cell_bw = lstm_cell_bw,\n",
        "                                        inputs = embedding_lookup_for_x, \n",
        "                                        dtype = tf.float32,\n",
        "                                        time_major = False,\n",
        "                                        sequence_length = length)\n",
        "    lstm_out = tf.concat([states[0], states[1]], axis = 2)\n",
        "        \n",
        "    # Conditional random fields\n",
        "    logits = tf.layers.dense(lstm_out, params['num_classes'])\n",
        "    crf_params = tf.get_variable(\"crf\", [params['num_classes'], params['num_classes']],\n",
        "                                 dtype=tf.float32)\n",
        "    pred_ids, _ = tf.contrib.crf.crf_decode(logits, crf_params, length)\n",
        "    training = (mode == tf.estimator.ModeKeys.TRAIN)\n",
        "    \n",
        "    # Prediction\n",
        "    if mode == tf.estimator.ModeKeys.PREDICT:\n",
        "        predictions = { \n",
        "            'pred_ids': pred_ids,\n",
        "            'tags': words,\n",
        "            'length' : length,\n",
        "        }\n",
        "        export_outputs = {\n",
        "          'prediction': tf.estimator.export.PredictOutput(predictions)\n",
        "      }\n",
        "        \n",
        "        return tf.estimator.EstimatorSpec(mode, predictions=predictions,\n",
        "                                          export_outputs=export_outputs)\n",
        "    \n",
        "    # Loss functions and optimizers\n",
        "    log_likelihood, _ = tf.contrib.crf.crf_log_likelihood(\n",
        "        logits, labels, length, crf_params)\n",
        "    \n",
        "    loss = tf.reduce_mean(-log_likelihood)\n",
        "    train_op = tf.train.AdamOptimizer().minimize(\n",
        "        loss, global_step = tf.train.get_or_create_global_step())\n",
        "        \n",
        "    # Training\n",
        "    if mode == tf.estimator.ModeKeys.TRAIN:\n",
        "        return tf.estimator.EstimatorSpec(mode = mode,\n",
        "                                           loss = loss,\n",
        "                                           train_op = train_op)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GbBnxNKAKpCs",
        "colab_type": "text"
      },
      "source": [
        "The Estimator API requires separate \"Spec\" objects, through `tf.estimator.TrainSpec` and `EvalSpec` for training and evaluation configuration. We use `functools.partial` to modify the input to the `input_fn` to create separate training and evaluation inputs, and then create separate `Spec` objects for training and evaluation."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4nmuDDVQKpCs",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 159
        },
        "outputId": "bcb6371d-be8c-4d93-e6f0-efb59e57e17b"
      },
      "source": [
        "# Spin up the estimator\n",
        "config = tf.estimator.RunConfig()\n",
        "estimator = tf.estimator.Estimator(model_fn, 'models/model', config, params)\n",
        "\n",
        "# Create train spec\n",
        "train_input_fn = partial(input_fn, \"data/train.txt\", params = params)\n",
        "train_spec = tf.estimator.TrainSpec(train_input_fn)\n",
        "\n",
        "# Create evaluation spec\n",
        "eval_input_fn = partial(input_fn, \"data/test.txt\", params = params, train = False)\n",
        "eval_spec = tf.estimator.EvalSpec(eval_input_fn)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:Using config: {'_model_dir': 'models/model', '_tf_random_seed': None, '_save_summary_steps': 100, '_save_checkpoints_steps': None, '_save_checkpoints_secs': 600, '_session_config': allow_soft_placement: true\n",
            "graph_options {\n",
            "  rewrite_options {\n",
            "    meta_optimizer_iterations: ONE\n",
            "  }\n",
            "}\n",
            ", '_keep_checkpoint_max': 5, '_keep_checkpoint_every_n_hours': 10000, '_log_step_count_steps': 100, '_train_distribute': None, '_device_fn': None, '_protocol': None, '_eval_distribute': None, '_experimental_distribute': None, '_experimental_max_worker_delay_secs': None, '_session_creation_timeout_secs': 7200, '_service': None, '_cluster_spec': ClusterSpec({}), '_task_type': 'worker', '_task_id': 0, '_global_id_in_cluster': 0, '_master': '', '_evaluation_master': '', '_is_chief': True, '_num_ps_replicas': 0, '_num_worker_replicas': 1}\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "txJB3Sf-KpCv",
        "colab_type": "text"
      },
      "source": [
        "### Training"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7J1h1ATDKpCw",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 399
        },
        "outputId": "78888b23-9249-4802-868c-681b7b699793"
      },
      "source": [
        "#ts = time.time()\n",
        "estimator.train(input_fn = train_input_fn, max_steps = 1000)\n",
        "#te = time.time()\n",
        "#print(\"Completed in {} seconds\".format(int(te - ts)))\n",
        "estimator.export_savedmodel('models/saved_model/', serving_input_fn)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:Calling model_fn.\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "error",
          "ename": "AttributeError",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-22-023a88ec4c0a>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0;31m#ts = time.time()\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m \u001b[0mestimator\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrain\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput_fn\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtrain_input_fn\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmax_steps\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;36m1000\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      3\u001b[0m \u001b[0;31m#te = time.time()\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m \u001b[0;31m#print(\"Completed in {} seconds\".format(int(te - ts)))\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m \u001b[0mestimator\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mexport_savedmodel\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'models/saved_model/'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mserving_input_fn\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.6/dist-packages/tensorflow_estimator/python/estimator/estimator.py\u001b[0m in \u001b[0;36mtrain\u001b[0;34m(self, input_fn, hooks, steps, max_steps, saving_listeners)\u001b[0m\n\u001b[1;32m    347\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    348\u001b[0m       \u001b[0msaving_listeners\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_check_listeners_type\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0msaving_listeners\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 349\u001b[0;31m       \u001b[0mloss\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_train_model\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput_fn\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mhooks\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0msaving_listeners\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    350\u001b[0m       \u001b[0mlogging\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0minfo\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'Loss for final step: %s.'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mloss\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    351\u001b[0m       \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.6/dist-packages/tensorflow_estimator/python/estimator/estimator.py\u001b[0m in \u001b[0;36m_train_model\u001b[0;34m(self, input_fn, hooks, saving_listeners)\u001b[0m\n\u001b[1;32m   1180\u001b[0m       \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_train_model_distributed\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput_fn\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mhooks\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0msaving_listeners\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1181\u001b[0m     \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1182\u001b[0;31m       \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_train_model_default\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput_fn\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mhooks\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0msaving_listeners\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1183\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1184\u001b[0m   \u001b[0;32mdef\u001b[0m \u001b[0m_train_model_default\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput_fn\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mhooks\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0msaving_listeners\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.6/dist-packages/tensorflow_estimator/python/estimator/estimator.py\u001b[0m in \u001b[0;36m_train_model_default\u001b[0;34m(self, input_fn, hooks, saving_listeners)\u001b[0m\n\u001b[1;32m   1209\u001b[0m       \u001b[0mworker_hooks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mextend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput_hooks\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1210\u001b[0m       estimator_spec = self._call_model_fn(features, labels, ModeKeys.TRAIN,\n\u001b[0;32m-> 1211\u001b[0;31m                                            self.config)\n\u001b[0m\u001b[1;32m   1212\u001b[0m       \u001b[0mglobal_step_tensor\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtf\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcompat\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mv1\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrain\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mget_global_step\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mg\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1213\u001b[0m       return self._train_with_estimator_spec(estimator_spec, worker_hooks,\n",
            "\u001b[0;32m/usr/local/lib/python3.6/dist-packages/tensorflow_estimator/python/estimator/estimator.py\u001b[0m in \u001b[0;36m_call_model_fn\u001b[0;34m(self, features, labels, mode, config)\u001b[0m\n\u001b[1;32m   1168\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1169\u001b[0m     \u001b[0mlogging\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0minfo\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'Calling model_fn.'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1170\u001b[0;31m     \u001b[0mmodel_fn_results\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_model_fn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfeatures\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mfeatures\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1171\u001b[0m     \u001b[0mlogging\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0minfo\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'Done calling model_fn.'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1172\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m<ipython-input-19-72159b25290a>\u001b[0m in \u001b[0;36mmodel_fn\u001b[0;34m(features, labels, mode, params)\u001b[0m\n\u001b[1;32m      8\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      9\u001b[0m     \u001b[0;31m# Embedding\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 10\u001b[0;31m     \u001b[0membedding\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtf\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mVariable\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtf\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrandom_normal\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mparams\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'max_words'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mparams\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'dim'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     11\u001b[0m     \u001b[0membedding_lookup_for_x\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtf\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnn\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0membedding_lookup\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0membedding\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mwords\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     12\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;31mAttributeError\u001b[0m: module 'tensorflow' has no attribute 'random_normal'"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r3oaSvSIKpCy",
        "colab_type": "text"
      },
      "source": [
        "### Evaluation"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MyFkbYajKpCy",
        "colab_type": "text"
      },
      "source": [
        "Although one might normally use `if mode == tf.estimator.ModeKeys.EVAL` in the `model_fn` to specify evaluation metrics with `tf.metrics`, NER requires multi-class precision, recall, and F1-score which are not available in `tf.metrics`. Instead, we load the true test labels and calculate precision, recall, and F1-score based upon the model predictions for each sentence at the entity-level (discarding non-entity words)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "56p3pSMKKpCy",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Generate predictions\n",
        "predictions = estimator.predict(eval_input_fn)\n",
        "\n",
        "# Load hash tables and true labels\n",
        "labels_index = load_obj(params['save_dir'], \"labels_index\")\n",
        "_, true = load_data(\"data/test.txt\")\n",
        "\n",
        "# Specify which label_index is non-entity\n",
        "dummy_label = labels_index.get(\"O\") \n",
        "\n",
        "# Convert [[string, string], [string, string] ...] to [[int, int], [int, int]]\n",
        "# with hashing table for label indexes\n",
        "labels = []\n",
        "for row in true:\n",
        "    labels.append([labels_index.get(y) for y in row])\n",
        "    \n",
        "# Loop through preds, labels and calculate metrics\n",
        "precisions, recalls, f1s = [], [], []\n",
        "for pred, true in zip(predictions, labels):\n",
        "    pred = pred['pred_ids'][:pred['length']] # undo pad_sequences\n",
        "    pred = [x for x in pred if x != dummy_label] # remove preds that aren't entities\n",
        "    true = np.asarray([x for x in true if x != dummy_label])\n",
        "    recall = calc_recall(true, pred)\n",
        "    recalls.append(recall)\n",
        "    precision = calc_precision(true, pred)\n",
        "    precisions.append(precision)\n",
        "    f1s.append(calc_f1(precision, recall))\n",
        "    \n",
        "print(\"Precision: {} \\nRecall: {} \\nF1-score: {}\".format(np.around(np.mean(precisions), 3),\n",
        "                                                         np.around(np.mean(recalls), 3),\n",
        "                                                         np.around(np.mean(f1s), 3)))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xZn9DMYXKpC0",
        "colab_type": "text"
      },
      "source": [
        "<img src=\"https://storage.googleapis.com/kf-pipeline-contrib-public/release-0.1.3/kfp-components/notebooks/entity_extraction/assets/fig9.png\" align = \"left\" width=\"200\"/>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KCmZ90E0KpC0",
        "colab_type": "text"
      },
      "source": [
        "### Serving model for on-the-fly predictions\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "425NPzNQKpC0",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from pathlib import Path\n",
        "from tensorflow.contrib import predictor\n",
        "\n",
        "LINE = 'did george clooney make a science fiction movie in the 1980s'\n",
        "\n",
        "\n",
        "def predict(line, export_dir = 'models/saved_model/', params = params):\n",
        "    # Load hash tables\n",
        "    word_index = load_obj(params['save_dir'], \"word_index\")\n",
        "    labels_lookup = load_obj(params['save_dir'], \"labels_lookup\")\n",
        "    \n",
        "    # Identify and load model weights\n",
        "    subdirs = [x for x in Path(export_dir).iterdir()\n",
        "                   if x.is_dir() and 'temp' not in str(x)]\n",
        "    latest_model = str(sorted(subdirs)[-1])\n",
        "    predict_fn = predictor.from_saved_model(latest_model)\n",
        "                \n",
        "    # Preprocess sentence input\n",
        "    line = line.strip().split()\n",
        "    vector = [word_index.get(x) if x in word_index.keys() else 0 for x in line] # tokenize\n",
        "    vector[len(vector):20] = [0] * (20 - len(vector)) # pad prediction\n",
        "        \n",
        "    # Calculate precision and transform for display\n",
        "    predictions = predict_fn({'words': [vector], 'length': [len(line)]})\n",
        "    tags = predictions.get('tags')\n",
        "    preds = predictions.get('pred_ids')\n",
        "    for tag, pred in zip(tags, preds):\n",
        "        tag = [word for word in tag if word != 0] # unpad\n",
        "        pred = pred[:len(tag)]\n",
        "        pred = [labels_lookup.get(num) for num in pred] #untokenize\n",
        "        print(line, \"\\n\", pred)\n",
        "    \n",
        "predict(LINE)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wGRsGvMlKpC2",
        "colab_type": "text"
      },
      "source": [
        "<img src=\"https://storage.googleapis.com/kf-pipeline-contrib-public/release-0.1.3/kfp-components/notebooks/entity_extraction/assets/fig10.png\" align = \"left\" width=\"800\"/>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fRSfTw04KpC2",
        "colab_type": "text"
      },
      "source": [
        "## Conclusion\n",
        "\n",
        "This notebook shows two end-to-end approaches of training and serving custom NER models including loading and transforming data, creating an NER pipeline, and calculating performance metrics. The custom architecture developed in Tensorflow was competitive with out-of-the-box algorithms while being an order of magnitude faster to train and being able to use the powerful high level APIs in tensorflow like Dataset and Estimator for scalable serving.\n",
        "\n",
        "Both approaches perform well on diverse queries about movies with spelling mistakes and complicated query structures. NER pipelines like the ones presented in this notebook can be integrated into recommender systems, search engines, NLP feature engineering, and customer support / chatbots, among many other business applications."
      ]
    }
  ]
}