03-idp-document-enrichment.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Document Enrichment\n",
    "\n",
    "<div class=\"alert alert-block alert-warning\">\n",
    "    <b>⚠️ PRE-REQUISITE:</b> In order to execute this notebook, make sure you have completed the first notebook 01-idp-document-classification.ipynb\n",
    "</div>\n",
    "\n",
    "So far, in the previous two notebooks, we have categorized documents and identified the bank statements and receipt documents. We have also extracted the text from these documents and have viewed the entities extraced by Amazon Comprehend. However, we want specific entities to be recognized for our use case so that we can perform certain enrichments on our documents. \n",
    "\n",
    "In this notebook we will train an Amazon Comprehend custom entity recognizer and deploy an endpoint with it. We will then identify the specific entities and generate custom metadata about our document in CSV format to be later analyzed by the business use case, and we will also identify any ADDRESS entity in bank statements and perform redaction on it, since that is a customer private information.\n",
    "\n",
    "![IDP Entity](./images/idp-entity.png)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "In this notebook we will - \n",
    "\n",
    "- [Step 1: Setup notebook](#step1)\n",
    "- [Step 2: Perform entity recocnition with Amazon Comprehend](#step2)\n",
    "- [Step 3: Train a custom Amazon Comprehend entity recognizer](#step3)\n",
    "- [Step 4: Create custom entity recognizer real-time endpoint](#step4)\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Step 1: Setup notebook <a id=\"step1\"></a>\n",
    "\n",
    "In this step, we will import some necessary libraries that will be used throughout this notebook. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "import botocore\n",
    "import sagemaker\n",
    "import time\n",
    "import os\n",
    "import json\n",
    "import datetime\n",
    "import io\n",
    "import uuid\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "from pytz import timezone\n",
    "from PIL import Image, ImageDraw, ImageFont\n",
    "import multiprocessing as mp\n",
    "from pathlib import Path\n",
    "from textractcaller.t_call import call_textract, Textract_Features\n",
    "from textractprettyprinter.t_pretty_print import Textract_Pretty_Print, get_string\n",
    "from trp import Document\n",
    "\n",
    "# Document\n",
    "from IPython.display import Image, display, HTML, JSON\n",
    "from PIL import Image as PImage, ImageDraw\n",
    "\n",
    "\n",
    "# variables\n",
    "data_bucket = sagemaker.Session().default_bucket()\n",
    "region = boto3.session.Session().region_name\n",
    "account_id = boto3.client('sts').get_caller_identity().get('Account')\n",
    "\n",
    "os.environ[\"BUCKET\"] = data_bucket\n",
    "os.environ[\"REGION\"] = region\n",
    "role = sagemaker.get_execution_role()\n",
    "\n",
    "print(f\"SageMaker role is: {role}\\nDefault SageMaker Bucket: s3://{data_bucket}\")\n",
    "\n",
    "s3=boto3.client('s3')\n",
    "textract = boto3.client('textract', region_name=region)\n",
    "comprehend=boto3.client('comprehend', region_name=region)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "# Step 2: Perform Name Entity Recognition using Amazon Comprehend <a id=\"step1\"></a>\n",
    "\n",
    "We have categorized our documents according to their respective document types and stored them in S3. Next, we will perform name entity recognition for 1 bank statement and 1 receipt using [Amazon Comprehend NER](https://docs.aws.amazon.com/comprehend/latest/dg/how-entities.html), in this case Comprehend will extract the prebuilt generic entity types from the documents. \n",
    "\n",
    "We will start the process by loading the extracted document text from S3 into a dataframe and subsequently using Amazon Comprehend [DetectEntities](https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/comprehend.html#Comprehend.Client.detect_entities) API. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "extracted_df = pd.read_csv('extracted_doc.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from trp import Document\n",
    "\n",
    "def get_entities(text):\n",
    "    try:\n",
    "        #detect entities\n",
    "        entities = comprehend.detect_entities(LanguageCode=\"en\", Text=text)  \n",
    "        df = pd.DataFrame(entities[\"Entities\"], columns = ['Text', 'Type'])\n",
    "        display(HTML(df.to_html(index=False)))\n",
    "    except Exception as e:\n",
    "        print(e)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will grab one of bank statement and receipt each from our S3 location."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "bank_df = extracted_df[extracted_df['DocType'] == 'bank-statements']\n",
    "\n",
    "bank_document = bank_df.iloc[0]['DocText']\n",
    "print(bank_document)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Execute entity extraction on Bank Statement"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "get_entities(bank_document)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Although entity extraction worked fairly well in identifying the _generic_ entity types for everything in the documents, we want _specific_ entities to be recognized for our use case. More specifically, we need to identify the customer's Savings and Checking bank account numbers, for example we want the entity types to be \"CHECKING_AC\" and \"SAVINGS_AC\".\n",
    "\n",
    "Amazon Comprehend's default prebuilt entity recognizer isn't aware of these entity types, so we will need to train and use a custom entity recognizer in this notebook. We will also perform some document enrichments for example, in the bank statement we want to redact the customer's account numbers. We will discuss more and do all of this in the next notebook."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "---\n",
    "# Step 3: Train a custom Amazon Comprehend entity recognizer <a id=\"step3\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# If you have imported a custom entity recognizer model, uncomment the follwing code lines and replace the <model_arn> with the model's ARN and execute this code cell. \n",
    "# If you execute this code cell then you can skip the model training and move to Step 4.\n",
    "\n",
    "# entity_recognizer_arn = \"<model_arn>\"\n",
    "# %store entity_recognizer_arn"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will be training a custom [Amazon Comprehend entity recognizer](https://docs.aws.amazon.com/comprehend/latest/dg/custom-entity-recognition.html). There are two ways a custom recognizer can be trained -\n",
    "\n",
    "- [Using Annotations](https://docs.aws.amazon.com/comprehend/latest/dg/cer-annotation.html)\n",
    "- [Using Entity Lists](https://docs.aws.amazon.com/comprehend/latest/dg/cer-entity-list.html)\n",
    "\n",
    "Annotations uses a large set of PDF files that have been annotated. These annotations can be created with service such as Amazon Ground Truth where real human workers can review your files and annotate them. This method is quite involved and if you are interested to learn more refer to [this blog](https://aws.amazon.com/blogs/machine-learning/custom-document-annotation-for-extracting-named-entities-in-documents-using-amazon-comprehend/) and [this blog](https://aws.amazon.com/blogs/machine-learning/extract-custom-entities-from-documents-in-their-native-format-with-amazon-comprehend/). \n",
    "\n",
    "In our case, we will use Entity Lists, which is a CSV file that should contain the texts and it's corresponding entity type. The entities in this file is going to be specific to our business needs. For the purposes of this exercise, we have provided an entity list in CSV format in the `/entity-training/` directory called `entitylist.csv`. This file contains a custom entity _Type_ for customer account numbers. We have used _CHECKING_AC_ and _SAVINGS_AC_ as the custom entity types. With this, we ultimately need the custom entity recognizer to recognize the savings and checking bank account numbers.\n",
    "\n",
    "Let's take a look at our entity list."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "entities_df = pd.read_csv('./entity-training/entitylist.csv', dtype={'Text': object})\n",
    "entities = entities_df[\"Type\"].unique().tolist()\n",
    "print(f'Custom entities : {entities}')\n",
    "print(f'\\nTotal Custom entities: {entities_df[\"Type\"].nunique()}')\n",
    "display(HTML(entities_df.to_html(index=False)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's train a custom entity recognizer with Amazon Comprehend. In order to train a custom entity recognizer we will need the entity list and the set of documents to train the model. We will use the same set of documents that we used earlier to train the custom classifer for this purpose.\n",
    "\n",
    "Each custom entity needs atleast 100 samples in the data corpus (documents) for training purposes, meaning you should have atleast a 100 documents containing examples of each of the custom entities in your training dataset. Also, a minimum of 250 entity matches are needed per entity in the entity list to train a model for custom entity recognition. We have provided a training corpus named `entity_training_corpus.csv` which can be used to train the entity recognizer along with the entity list. Note that this corpus was generated the same way we generated training data for training a custom classifier in the first notebook. With these two data sets we will use Amazon Comprehend's [`CreateEntityRecognizer` API](https://docs.aws.amazon.com/comprehend/latest/dg/API_CreateEntityRecognizer.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Upload entity list CSV to S3\n",
    "entities_key='idp/comprehend/entities/entitylist.csv'\n",
    "training_data_key='idp/comprehend/entities/entity_training_corpus.csv'\n",
    "\n",
    "s3.upload_file(Filename='./entity-training/entitylist.csv', \n",
    "               Bucket=data_bucket, \n",
    "               Key=entities_key)\n",
    "\n",
    "s3.upload_file(Filename='./entity-training/entity_training_corpus.csv', \n",
    "               Bucket=data_bucket, \n",
    "               Key=training_data_key)\n",
    "\n",
    "entities_uri = f's3://{data_bucket}/{entities_key}'\n",
    "training_data_uri = f's3://{data_bucket}/{training_data_key}'\n",
    "\n",
    "print(f'Entity List CSV File: {entities_uri}')\n",
    "print(f'Training Data File: {training_data_uri}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Let's now train a custom entity recognizer with this data and the entity list of savings and checking account numbers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a custom entity recognizer\n",
    "account_id = boto3.client('sts').get_caller_identity().get('Account')\n",
    "id = str(datetime.datetime.now().strftime(\"%s\"))\n",
    "\n",
    "entity_recognizer_name = 'Sample-Entity-Recognizer-IDP'\n",
    "entity_recognizer_version = 'Sample-Entity-Recognizer-IDP-v1'\n",
    "entity_recognizer_arn = ''\n",
    "create_response = None\n",
    "EntityTypes = [ {'Type': entity} for entity in entities]\n",
    "\n",
    "try:\n",
    "    create_response = comprehend.create_entity_recognizer(\n",
    "        InputDataConfig={\n",
    "            'DataFormat': 'COMPREHEND_CSV',\n",
    "            'EntityTypes': EntityTypes,\n",
    "            'Documents': {\n",
    "                'S3Uri': training_data_uri\n",
    "            },\n",
    "            'EntityList': {\n",
    "                'S3Uri': entities_uri\n",
    "            }\n",
    "        },\n",
    "        DataAccessRoleArn=role,\n",
    "        RecognizerName=entity_recognizer_name,\n",
    "        VersionName=entity_recognizer_version,\n",
    "        LanguageCode='en'\n",
    "    )\n",
    "    \n",
    "    entity_recognizer_arn = create_response['EntityRecognizerArn']\n",
    "    \n",
    "    print(f\"Comprehend Custom entity recognizer created with ARN: {entity_recognizer_arn}\")\n",
    "except Exception as error:\n",
    "    if error.response['Error']['Code'] == 'ResourceInUseException':\n",
    "        print(f'An entity recognizer with the name \"{entity_recognizer_name}\" already exists.')\n",
    "        entity_recognizer_arn = f'arn:aws:comprehend:{region}:{account_id}:entity-recognizer/{entity_recognizer_name}/version/{entity_recognizer_version}'\n",
    "        print(f'The entity recognizer ARN is: \"{entity_recognizer_arn}\"')\n",
    "    else:\n",
    "        print(error)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%store entity_recognizer_arn"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check status of the Comprehend custom entity recognizer job"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "# Loop through and wait for the training to complete . Takes up to 10 mins \n",
    "from IPython.display import clear_output\n",
    "import time\n",
    "from datetime import datetime\n",
    "\n",
    "jobArn = create_response['EntityRecognizerArn']\n",
    "\n",
    "max_time = time.time() + 3*60*60 # 3 hours\n",
    "while time.time() < max_time:\n",
    "    now = datetime.now()\n",
    "    current_time = now.strftime(\"%H:%M:%S\")\n",
    "    \n",
    "    describe_custom_recognizer = comprehend.describe_entity_recognizer(\n",
    "        EntityRecognizerArn = jobArn\n",
    "    )\n",
    "    status = describe_custom_recognizer[\"EntityRecognizerProperties\"][\"Status\"]\n",
    "    clear_output(wait=True)\n",
    "    print(f\"{current_time} : Custom document entity recognizer: {status}\")\n",
    "    \n",
    "    if status == \"TRAINED\" or status == \"IN_ERROR\":\n",
    "        break\n",
    "    time.sleep(10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "# Step 4: Create custom entity recognizer real-time endpoint <a id=\"step4\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will create a real time entity recognizer endpoint with the trained entity recognizer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#create comprehend endpoint\n",
    "model_arn = entity_recognizer_arn\n",
    "ep_name = 'idp-er-endpoint'\n",
    "\n",
    "try:\n",
    "    endpoint_response = comprehend.create_endpoint(\n",
    "        EndpointName=ep_name,\n",
    "        ModelArn=model_arn,\n",
    "        DesiredInferenceUnits=1,    \n",
    "        DataAccessRoleArn=role\n",
    "    )\n",
    "    ER_ENDPOINT_ARN=endpoint_response['EndpointArn']\n",
    "    print(f'Endpoint created with ARN: {ER_ENDPOINT_ARN}')\n",
    "    %store ER_ENDPOINT_ARN\n",
    "except Exception as error:\n",
    "    if error.response['Error']['Code'] == 'ResourceInUseException':\n",
    "        print(f'An endpoint with the name \"{ep_name}\" already exists.')\n",
    "        ER_ENDPOINT_ARN = f'arn:aws:comprehend:{region}:{account_id}:entity-recognizer-endpoint/{ep_name}'\n",
    "        print(f'The classifier endpoint ARN is: \"{ER_ENDPOINT_ARN}\"')\n",
    "        %store ER_ENDPOINT_ARN\n",
    "    else:\n",
    "        print(error)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check creation status of the entity recognizer endpoint."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "# Loop through and wait for the training to complete . Takes up to 10 mins \n",
    "from IPython.display import clear_output\n",
    "import time\n",
    "from datetime import datetime\n",
    "\n",
    "ep_arn = endpoint_response[\"EndpointArn\"]\n",
    "\n",
    "max_time = time.time() + 3*60*60 # 3 hours\n",
    "while time.time() < max_time:\n",
    "    now = datetime.now()\n",
    "    current_time = now.strftime(\"%H:%M:%S\")\n",
    "    \n",
    "    describe_endpoint_resp = comprehend.describe_endpoint(\n",
    "        EndpointArn=ep_arn\n",
    "    )\n",
    "    status = describe_endpoint_resp[\"EndpointProperties\"][\"Status\"]\n",
    "    clear_output(wait=True)\n",
    "    print(f\"{current_time} : Custom entity recognizer classifier: {status}\")\n",
    "    \n",
    "    if status == \"IN_SERVICE\" or status == \"FAILED\":\n",
    "        break\n",
    "        \n",
    "    time.sleep(10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "Once we train a custom entity recognizer, we will use to extract some enriched information out of the document and then perform document redaction with the help of the custom entities recognized by Amazon Comprehend and bounding box information from Amazon Textract\n",
    "\n",
    "- [Enrichment 1: Extract custom entity out of the documents](#er1)\n",
    "- [Enrichment 2: Perform redaction document enrichment ](#er2)\n",
    "\n",
    "# Enrichment 1: Detect Custom Entities  <a id=\"er1\"></a>\n",
    "\n",
    "<p align=\"center\">\n",
    "  <img src=\"./images/idp-custom-entity.png\" alt=\"cfn1\" style=\"width:500px;\"/>\n",
    "</p>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "extracted_df = pd.read_csv('extracted_doc.csv')\n",
    "extracted_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from trp import Document\n",
    "\n",
    "def get_entities(text):\n",
    "    try:\n",
    "        #detect entities\n",
    "        entities_custom = comprehend.detect_entities(LanguageCode=\"en\", Text=text, EndpointArn=ER_ENDPOINT_ARN)    \n",
    "        df_custom = pd.DataFrame(entities_custom[\"Entities\"], columns = ['Text', 'Type'])\n",
    "        df_custom = df_custom.drop_duplicates(subset=['Text']).reset_index()\n",
    "        return df_custom\n",
    "    except Exception as e:\n",
    "        print(e)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will grab a bank statement from our S3 location."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bank_df = extracted_df[extracted_df['DocType'] == 'bank-statements']\n",
    "\n",
    "bank_document = bank_df.iloc[3]['DocText']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Execute entity extraction on Bank Statement"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "entities_df = get_entities(bank_document)\n",
    "entities_df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Much better! Now instead of returning us generic entities, Amazon Comprehend is returning us the entities from our scanned document that we are interested, i.e. the checking and savings account number. We will now save this as a CSV file, which can be further consumed into a database or datawarehouse for additional analysis. For our excercise, we will save it to a csv file and upload to S3."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "entities_df.to_csv('final_output.csv', index=False)\n",
    "\n",
    "#Upload dataframe as csv to S3\n",
    "s3.upload_file(Filename='final_output.csv', \n",
    "               Bucket=data_bucket, \n",
    "               Key=f'idp/comprehend/output-entities/final_output.csv')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "# Enrichment 2: Perform redaction document enrichment  <a id=\"er2\"></a>\n",
    "We still need to perform some enrichments on the document. Since the document contains the customers savings and checking account numbers, we would like to redact those. Since we already know, by means of our custom entity, which is the name and which is the address, we can easily use Amazon Textract's geometry data to redact that information in the document.\n",
    "\n",
    "<p align=\"center\">\n",
    "  <img src=\"./images/idp-redaction.png\" alt=\"cfn1\" style=\"width:500px;\"/>\n",
    "</p>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's grab a bank statement from our classified list of documents, we will pick the S3 location of the document and then perform the actions below-\n",
    "\n",
    "- Use Amazon Textract to get the geometry information i.e. the bounding boxes, of all the lines in the document \n",
    "- Use the extracted text above to identify the entities CHECKING_AC and SAVINGS_AC, using Comprehend custom entity recognizer\n",
    "- Find the bounding box for the CHECKING_AC and SAVINGS_AC words from the Textract response\n",
    "- Use the bounding box geometry to annotate the document and redact the customer name and address.\n",
    "\n",
    "In order to obtain the bounding box geometry of all words from the document we will use a tool called `amazon-textract-overlayer`. See [documentation](https://github.com/aws-samples/amazon-textract-textractor/tree/master/overlayer) for learn more about `amazon-textract-overlayer`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python -m pip install -q amazon-textract-overlayer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from textractoverlayer.t_overlay import DocumentDimensions, get_bounding_boxes\n",
    "from textractcaller.t_call import Textract_Types\n",
    "\n",
    "def redact_doc(s3document, localpath):\n",
    "    try:\n",
    "        img = PImage.open(localpath)\n",
    "        draw = ImageDraw.Draw(img)\n",
    "\n",
    "        # Use call_textract to get bounding boxes\n",
    "        resp = call_textract(input_document = f's3://{data_bucket}/{s3document}', features = [Textract_Features.TABLES, Textract_Features.FORMS])\n",
    "        document_dimension:DocumentDimensions = DocumentDimensions(doc_width=img.size[0], doc_height=img.size[1])\n",
    "        overlay=[Textract_Types.WORD, Textract_Types.FORM, Textract_Types.CELL]\n",
    "        bounding_box_list = get_bounding_boxes(textract_json=resp, document_dimensions=[document_dimension], overlay_features=overlay)\n",
    "\n",
    "        print('Detecting entities...')\n",
    "                \n",
    "        text = get_string(textract_json=resp, output_type=[Textract_Pretty_Print.LINES])\n",
    "                \n",
    "        #entity recognizer\n",
    "        entity_resp = comprehend.detect_entities(LanguageCode=\"en\", Text=text, EndpointArn=ER_ENDPOINT_ARN)         \n",
    "        entities = [{'Type': entity['Type'], 'Text': entity['Text']} for entity in entity_resp['Entities']]\n",
    "        redactions = []\n",
    "\n",
    "        #collect the bounding boxes for the custom entities\n",
    "        for entity in entities:\n",
    "            entity_text = entity['Text']\n",
    "            print(f'Found Entity: {entity_text}')\n",
    "            for bbox in bounding_box_list:                \n",
    "                if bbox.text == entity_text:\n",
    "                    redactions.append(bbox)\n",
    "\n",
    "        for box in redactions:\n",
    "            draw.rectangle(xy=[box.xmin, box.ymin, box.xmax, box.ymax], fill=\"Black\")\n",
    "\n",
    "        opfile = Path(localpath).stem\n",
    "        opfile = f'{opfile}_redacted.png'        \n",
    "        img.save(opfile)           \n",
    "        print(f'Done.... Redacted file saved: {opfile}')\n",
    "        return opfile\n",
    "    except Exception as e:\n",
    "        print(e)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The function above finds the custom entities in the document, finds the corresponding geometry information of the custom entity text and perform redaction on the document. Let's call it for a sample bank statement."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Select a bank document\n",
    "\n",
    "bank_df = extracted_df[extracted_df['DocType'] == 'bank-statements']\n",
    "bank_document_s3 = bank_df.iloc[3]['s3path']\n",
    "bank_document_local = bank_df.iloc[3]['Document']\n",
    "bank_document_local, bank_document_s3\n",
    "\n",
    "#perform redaction enrichment\n",
    "redacted_file = redact_doc(bank_document_s3, f'./samples/mixedbag/{bank_document_local}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once our redacted file has been generated, lets take a look at it..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f'\\nUnredacted Document\\t\\t\\t\\t\\t\\t\\tRedacted Document \\n')\n",
    "\n",
    "HTML(f\"\"\"\n",
    "    <div class=\"row\">\n",
    "            <img src=./samples/mixedbag/{bank_document_local} style=\"width:45%\"></img>\n",
    "            <img src={redacted_file} style=\"width:45%\">  </img>\n",
    "    </div>\n",
    "    \"\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "# Conclusion\n",
    "\n",
    "In this notebook we trained an Amazon Comprehend custom entity recognizer using our own entity list so that we can extract those entities from our documents. We used 2 entities CHECKING_AC, and SAVINGS_AC. We then created an endpoint with the custom entity recognizer and performed a detect_entities with the endpoint with one of the bank statements. Finally, we saved the extracted entities into a CSV file and uploaded it to S3 for further analysis. \n",
    "\n",
    "We still needed to perform some enrichments on the document. Since the document contains the customers checking and savings account numbers, we would like to redact those. Since we already know, by means of our custom entity, the customer's checking and savings bank account numbers, we used Amazon Textract's geometry data to redact that information in the document."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "# Cleanup\n",
    "\n",
    "Cleanup is optional at this point if you want to execute the next notebook (Module-4). \n",
    "\n",
    "Refer to the `05-idp-cleanup.ipynb` for cleanup and deletion of resources."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "instance_type": "ml.t3.medium",
  "kernelspec": {
   "display_name": "Python 3 (Data Science)",
   "language": "python",
   "name": "python3__SAGEMAKER_INTERNAL__arn:aws:sagemaker:us-east-2:429704687514:image/datascience-1.0"
  },
  "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}