notebooks/204-named-entity-recognition/204-named-entity-recognition.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Document Entity Extraction with OpenVINO\n",
    "\n",
    "This demo shows named entity recognition from text with OpenVINO. We use [small BERT-large-like model](https://github.com/openvinotoolkit/open_model_zoo/tree/master/models/intel/bert-small-uncased-whole-word-masking-squad-int8-0002) distilled and quantized to INT8 on SQuAD v1.1 training set from larger BERT-large model. The model comes from [Open Model Zoo](https://github.com/openvinotoolkit/open_model_zoo/). At the bottom of this notebook, you will see live inference results from your inputs and templates. In the notebook we show how to create the following pipeline:\n",
    "\n",
    "<p align=\"center\" width=\"100%\">\n",
    "    <img width=\"80%\" src=\"https://user-images.githubusercontent.com/33627846/169465451-44bfecbb-b670-46ef-bb20-58d6da7d703b.png\"> \n",
    "</p>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import time\n",
    "import json\n",
    "\n",
    "import numpy as np\n",
    "import tokens_bert as tokens\n",
    "\n",
    "from openvino.runtime import Core\n",
    "from openvino.runtime import Dimension"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## The model\n",
    "\n",
    "### Download the model\n",
    "\n",
    "We use `omz_downloader`, which is a command-line tool from the `openvino-dev` package. `omz_downloader` automatically creates a directory structure and downloads the selected model. If the model is already downloaded, this step is skipped.\n",
    "\n",
    "You can download and use any of the following models: `bert-large-uncased-whole-word-masking-squad-0001`, `bert-large-uncased-whole-word-masking-squad-int8-0001`, `bert-small-uncased-whole-word-masking-squad-0001`, `bert-small-uncased-whole-word-masking-squad-0002`, `bert-small-uncased-whole-word-masking-squad-int8-0002`, just change the model name below. Any of these models are already converted to OpenVINO Intermediate Representation (IR), so there is no need to use `omz_converter`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# directory where model will be downloaded\n",
    "base_model_dir = \"model\"\n",
    "\n",
    "# desired precision\n",
    "precision = \"FP16-INT8\"\n",
    "\n",
    "# model name as named in Open Model Zoo\n",
    "model_name = \"bert-small-uncased-whole-word-masking-squad-int8-0002\"\n",
    "\n",
    "model_path = f\"model/intel/{model_name}/{precision}/{model_name}.xml\"\n",
    "model_weights_path = f\"model/intel/{model_name}/{precision}/{model_name}.bin\"\n",
    "\n",
    "download_command = f\"omz_downloader \" \\\n",
    "                   f\"--name {model_name} \" \\\n",
    "                   f\"--precision {precision} \" \\\n",
    "                   f\"--output_dir {base_model_dir} \" \\\n",
    "                   f\"--cache_dir {base_model_dir}\"\n",
    "! $download_command"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Load the model for Entity Extraction with Dynamic Shape\n",
    "\n",
    "Input to entity extraction model refers to text with different content sizes, i.e, dynamic into shapes. Hence:\n",
    "\n",
    "1. Input dimension with dynamic input shapes needs to be specified before loading entity extraction model\n",
    "2. Dynamic shape is specified by assigning -1 to the input dimension or by setting the upper bound of the input dimension using, for example, Dimension(1,384)\n",
    "\n",
    "In scope of this notebook, as we know the upper bound of the dynamic input and longest input text allowed is 384 i.e. 380 tokens for content + 1 for entity + 3 special (separation) tokens, it is more recommended to assign dynamic shape using Dimension(, upper bound) i.e. Dimension(1, 384) so it’ll use memory more efficiently with OpenVINO 2022.1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# initialize inference engine\n",
    "ie_core = Core()\n",
    "# read the network and corresponding weights from file\n",
    "model = ie_core.read_model(model=model_path, weights=model_weights_path)\n",
    "\n",
    "# assign dynamic shapes to every input layer on the last dimension\n",
    "for input_layer in model.inputs:\n",
    "    input_shape = input_layer.partial_shape\n",
    "    input_shape[1] = Dimension(1, 384)\n",
    "    model.reshape({input_layer: input_shape})\n",
    "\n",
    "# compile the model for the CPU\n",
    "compiled_model = ie_core.compile_model(model=model, device_name=\"CPU\")\n",
    "\n",
    "# get input names of nodes\n",
    "input_keys = list(compiled_model.inputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Input keys are the names of the network input nodes. In the case of the BERT-large-like model, we have four inputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "[i.any_name for i in input_keys]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Processing\n",
    "\n",
    "NLP models usually take a list of tokens as standard input. A token is a single word converted to some integer. To provide the proper input, we need the vocabulary for such mapping. We also define some special tokens like separators and a function to load the content. Content is loaded from simple text."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# path to vocabulary file\n",
    "vocab_file_path = \"data/vocab.txt\"\n",
    "\n",
    "# create dictionary with words and their indices\n",
    "vocab = tokens.load_vocab_file(vocab_file_path)\n",
    "\n",
    "# define special tokens\n",
    "cls_token = vocab[\"[CLS]\"]\n",
    "sep_token = vocab[\"[SEP]\"]\n",
    "\n",
    "# set a confidence score threshold\n",
    "confidence_threshold = 0.4"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "### Preprocessing\n",
    "\n",
    "The main input (`input_ids`) to used BERT model consist of two parts: entity tokens and context tokens separated by some special tokens. We also need to provide: `attention_mask`, which is a sequence of integer values representing the mask of valid values in the input; `token_type_ids`, which is a sequence of integer values representing the segmentation of the `input_ids` into entity and context; `position_ids`, which is a sequence of integer values from 0 to length of input, extended by separation tokens, representing the position index for each input token. To know more about input, please read [this](https://github.com/openvinotoolkit/open_model_zoo/tree/master/models/intel/bert-small-uncased-whole-word-masking-squad-int8-0002#input)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# generator of a sequence of inputs\n",
    "def prepare_input(entity_tokens, context_tokens):\n",
    "    input_ids = [cls_token] + entity_tokens + [sep_token] + \\\n",
    "        context_tokens + [sep_token]\n",
    "    # 1 for any index\n",
    "    attention_mask = [1] * len(input_ids)\n",
    "    # 0 for entity tokens, 1 for context part\n",
    "    token_type_ids = [0] * (len(entity_tokens) + 2) + \\\n",
    "        [1] * (len(context_tokens) + 1)\n",
    "\n",
    "    # create input to feed the model\n",
    "    input_dict = {\n",
    "        \"input_ids\": np.array([input_ids], dtype=np.int32),\n",
    "        \"attention_mask\": np.array([attention_mask], dtype=np.int32),\n",
    "        \"token_type_ids\": np.array([token_type_ids], dtype=np.int32),\n",
    "    }\n",
    "\n",
    "    # some models require additional position_ids\n",
    "    if \"position_ids\" in [i_key.any_name for i_key in input_keys]:\n",
    "        position_ids = np.arange(len(input_ids))\n",
    "        input_dict[\"position_ids\"] = np.array([position_ids], dtype=np.int32)\n",
    "\n",
    "    return input_dict"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Postprocessing\n",
    "\n",
    "The results from the network are raw (logits). We need to use the softmax function to get the probability distribution. Then, we are looking for the best entity extraction in the current part of the context (the highest score) and we return the score and the context range for the extracted entity."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def postprocess(output_start, output_end, entity_tokens,\n",
    "                context_tokens_start_end, input_size):\n",
    "\n",
    "    def get_score(logits):\n",
    "        out = np.exp(logits)\n",
    "        return out / out.sum(axis=-1)\n",
    "\n",
    "    # get start-end scores for context\n",
    "    score_start = get_score(output_start)\n",
    "    score_end = get_score(output_end)\n",
    "\n",
    "    # index of first context token in tensor\n",
    "    context_start_idx = len(entity_tokens) + 2\n",
    "    # index of last+1 context token in tensor\n",
    "    context_end_idx = input_size - 1\n",
    "\n",
    "    # find product of all start-end combinations to find the best one\n",
    "    max_score, max_start, max_end = find_best_entity_window(\n",
    "        start_score=score_start, end_score=score_end,\n",
    "        context_start_idx=context_start_idx, context_end_idx=context_end_idx\n",
    "    )\n",
    "\n",
    "    # convert to context text start-end index\n",
    "    max_start = context_tokens_start_end[max_start][0]\n",
    "    max_end = context_tokens_start_end[max_end][1]\n",
    "\n",
    "    return max_score, max_start, max_end\n",
    "\n",
    "\n",
    "def find_best_entity_window(start_score, end_score,\n",
    "                            context_start_idx, context_end_idx):\n",
    "    context_len = context_end_idx - context_start_idx\n",
    "    score_mat = np.matmul(\n",
    "        start_score[context_start_idx:context_end_idx].reshape(\n",
    "            (context_len, 1)),\n",
    "        end_score[context_start_idx:context_end_idx].reshape(\n",
    "            (1, context_len)),\n",
    "    )\n",
    "    # reset candidates with end before start\n",
    "    score_mat = np.triu(score_mat)\n",
    "    # reset long candidates (>16 words)\n",
    "    score_mat = np.tril(score_mat, 16)\n",
    "    # find the best start-end pair\n",
    "    max_s, max_e = divmod(score_mat.flatten().argmax(), score_mat.shape[1])\n",
    "    max_score = score_mat[max_s, max_e]\n",
    "\n",
    "    return max_score, max_s, max_e"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Firstly, we need to create a list of tokens from the context and the entity. Then, we are looking for the best extracted entity by trying different parts of the context. The best extracted entity should come with the highest score."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_best_entity(entity, context, vocab):\n",
    "    # convert context string to tokens\n",
    "    context_tokens, context_tokens_end = tokens.text_to_tokens(\n",
    "        text=context.lower(), vocab=vocab)\n",
    "    # convert entity string to tokens\n",
    "    entity_tokens, _ = tokens.text_to_tokens(text=entity.lower(), vocab=vocab)\n",
    "\n",
    "    network_input = prepare_input(entity_tokens, context_tokens)\n",
    "    input_size = len(context_tokens) + len(entity_tokens) + 3\n",
    "\n",
    "    # openvino inference\n",
    "    output_start_key = compiled_model.output(\"output_s\")\n",
    "    output_end_key = compiled_model.output(\"output_e\")\n",
    "    result = compiled_model(network_input)\n",
    "\n",
    "    # postprocess the result getting the score and context range for the answer\n",
    "    score_start_end = postprocess(output_start=result[output_start_key][0],\n",
    "                                  output_end=result[output_end_key][0],\n",
    "                                  entity_tokens=entity_tokens,\n",
    "                                  context_tokens_start_end=context_tokens_end,\n",
    "                                  input_size=input_size)\n",
    "\n",
    "    # return the part of the context, which is already an answer\n",
    "    return context[score_start_end[1]:score_start_end[2]], score_start_end[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Set the Entity Recognition Template\n",
    "\n",
    "Only the entities which have prediction confidence score more than 0.4 will be captured in the final output. This can be set by setting the variable confidence_threshold above. Sample entities supported by the application natural entity recognition template. \n",
    "- building, company, persons, city, state, height, floor and address"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "template = [\"building\", \"company\", \"persons\", \"city\",\n",
    "            \"state\", \"height\", \"floor\", \"address\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def run_analyze_entities(context):\n",
    "    print(f\"Context: {context}\\n\", flush=True)\n",
    "\n",
    "    if len(context) == 0:\n",
    "        print(\"Error: Empty context or outside paragraphs\")\n",
    "        return\n",
    "\n",
    "    if len(context) > 380:\n",
    "        print(\"Error: The context is too long for this particular model. \"\n",
    "              \"Try with context shorter than 380 words.\")\n",
    "        return\n",
    "\n",
    "    # measure processing time\n",
    "    start_time = time.perf_counter()\n",
    "    extract = []\n",
    "    for field in template:\n",
    "        entity_to_find = field + \"?\"\n",
    "        entity, score = get_best_entity(entity=entity_to_find,\n",
    "                                        context=context,\n",
    "                                        vocab=vocab)\n",
    "        if score >= confidence_threshold:\n",
    "            extract.append({\"Entity\": entity, \"Type\": field,\n",
    "                            \"Score\": f\"{score:.2f}\"})\n",
    "    end_time = time.perf_counter()\n",
    "    res = {\"Extraction\": extract, \"Time\": f\"{end_time - start_time:.2f}s\"}\n",
    "    print(\"\\nJSON Output:\")\n",
    "    print(json.dumps(res, sort_keys=False, indent=4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run on Simple Text\n",
    "\n",
    "### Sample 1\n",
    "\n",
    "Change sources to your own text, supported by the template, to perform entity extraction. It supports only one input text at a time. Usually, you need to wait a few seconds for the entities to be extracted, but longer the context, longer would be the waiting time. The model is very limited and sensitive for the input and predefined template. The answer can depend on whether it is supported by the template or not. The model will try to extract entities even if not supported by the template, so in that case, you can see random results.\n",
    "\n",
    "Sample source: Intel - Wikipedia (from [here](https://en.wikipedia.org/wiki/Intel))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "source_text = \"Intel Corporation is an American multinational and technology\" \\\n",
    "    \" company headquartered in Santa Clara, California.\"\n",
    "run_analyze_entities(source_text)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Sample 2\n",
    "\n",
    "Sample source: Intel - Wikipedia (from [here](https://en.wikipedia.org/wiki/Intel))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "source_text = \"Intel was founded in Mountain View, California, \" \\\n",
    "    \"in 1968 by Gordon E. Moore, a chemist, and Robert Noyce, \" \\\n",
    "    \"a physicist and co-inventor of the integrated circuit.\"\n",
    "run_analyze_entities(source_text)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Sample 3\n",
    "\n",
    "Sample source: Converted Paragraph (from [here](https://www.emporis.com/buildings/338440/robert-noyce-building-santa-clara-ca-usa))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "source_text = \"The Robert Noyce Building in Santa Clara, California, \" \\\n",
    "    \"is the headquarters for Intel Corporation. It was constructed in 1992 \" \\\n",
    "    \"and is located at 2200 Mission College Boulevard - 95054. It has an \" \\\n",
    "    \"estimated height of 22.20 meters and 6 floors above ground.\"\n",
    "run_analyze_entities(source_text)"
   ]
  }
 ],
 "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.7.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}