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 (NER) from a text with OpenVINO. It uses a [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 a larger BERT-large model. The model comes from [Open Model Zoo](https://github.com/openvinotoolkit/open_model_zoo/). Final part of this notebook includes live inference results from the inputs and templates. The notebook shows 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",
    "Use `omz_downloader`, which is a command-line tool from the `openvino-dev` package. It 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 in the code below. Any of these models are already converted to OpenVINO Intermediate Representation (OpenVINO IR), so there is no need to use `omz_converter`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# A directory where the model will be downloaded.\n",
    "base_model_dir = \"model\"\n",
    "\n",
    "# The desired precision.\n",
    "precision = \"FP16-INT8\"\n",
    "\n",
    "# A 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 such as dynamic input shapes. Hence:\n",
    "\n",
    "1. Input dimension with dynamic input shapes needs to be specified before loading entity extraction model.\n",
    "2. A 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 this notebook, the upper bound of the dynamic input and longest input text allowed is 384, that is 380 tokens for content + 1 for entity + 3 special (separation) tokens. It is highly recommended to assign dynamic shape, using `Dimension(, upper bound)` (in this case, `Dimension(1, 384)` ) so it will use memory more efficiently with OpenVINO 2022.1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Initialize OpenVINO Runtime.\n",
    "ie_core = Core()\n",
    "# Read the network and corresponding weights from the 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 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, there are 4 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 a standard input. A token is a single word converted to some integer. To provide the proper input, you need the vocabulary for such mapping. You also define some special tokens like separators and a function to load the content. The content is loaded from a simple text."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# A path to a vocabulary file.\n",
    "vocab_file_path = \"data/vocab.txt\"\n",
    "\n",
    "# Create a 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`) of used BERT model consists of two parts: entity tokens and context tokens, separated by some special tokens. You also need to provide: \n",
    "- `attention_mask` - a sequence of integer values representing the mask of valid values in the input, \n",
    "- `token_type_ids`- a sequence of integer values representing the segmentation of `input_ids` into entity and context, \n",
    "- `position_ids`- a sequence of integer values from 0 to length of input, extended by separation tokens, representing the position index for each input token. \n",
    "\n",
    "For more information, refer to the **Input** section of [BERT model documentation](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": [
    "# A 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 an 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). Use the softmax function to get the probability distribution. Then, find the best entity extraction in the current part of the context (the highest score) and 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 the context.\n",
    "    score_start = get_score(output_start)\n",
    "    score_end = get_score(output_end)\n",
    "\n",
    "    # Index of the first context token in a tensor.\n",
    "    context_start_idx = len(entity_tokens) + 2\n",
    "    # Index of last+1 context token in a tensor.\n",
    "    context_end_idx = input_size - 1\n",
    "\n",
    "    # Find the 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": [
    "First, create a list of tokens from the context and the entity. Then, find 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 the context string to tokens.\n",
    "    context_tokens, context_tokens_end = tokens.text_to_tokens(\n",
    "        text=context.lower(), vocab=vocab)\n",
    "    # Convert the 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 changed by setting the `confidence_threshold` variable 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 the 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 the longer the context, the longer 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 they are not supported by the template. In such cases, random results will appear.\n",
    "\n",
    "Sample source: [Intel on Wikipedia](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 on Wikipedia](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: a 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
}