01-idp-document-classification.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0d4f7af5",
   "metadata": {},
   "source": [
    "# Intelligent Document Processing\n",
    "\n",
    "Documents contain valuable information and come in various shapes and forms. In most cases, you are manually processing these documents which is time consuming, prone to error, and costly. Not only do you want this information extracted quickly but can also automate business processes that presently relies on manual inputs and intervention across various file types and formats.\n",
    "\n",
    "To help you overcome these challenges, AWS Machine Learning (ML) now provides you choices when it comes to extracting information from complex content in any document format such as insurance claims, mortgages, healthcare claims, contracts, and legal contracts.\n",
    "\n",
    "The diagram below shows an architecture for an Intelligent document processing workflow. It starts with data capture stage to securely store and aggregate different types (PDF, PNG, JPEG, and TIFF), formats, and layouts of documents. Followed by accurate classification of documents and extracting text and key insights from documents and perform further enrichments of the documents (such as identity entities, redaction etc.). Finally, the verification and review stage involves manual review of the documents for quality and accuracy, followed by consumption of the documents and extracted information into downstream databases/applications.\n",
    "\n",
    "In this workshop, we will explore the various aspects of this workflow such as the document classification, text and insights extraction, enrichments, and human review.\n",
    "\n",
    "![Arch](./images/idp.png)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f328a8b6",
   "metadata": {},
   "source": [
    "# Document Classification\n",
    "In this lab we will walk you through an hands-on lab on document classification using Amazon Comprehend\n",
    "Custom Classifier. We will use Amazon Textract to first extract the text out of our documents and then label them and then use the data for training our Amazon comprehend custom classifier. We will create an Amazon Comprehend real time endpoint with the custom classifier to classify our documents.\n",
    "\n",
    "![IDP Classify](./images/idp-classify.png)\n",
    "\n",
    "- [Step 1: Setup notebook and upload sample documents to Amazon S3](#step1)\n",
    "- [Step 2: Extract text from sample documents using Amazon Textract](#step2)\n",
    "- [Step 3: Prepare a CSV training dataset for Amazon Comprehend custom classifier training](#step3)\n",
    "- [Step 4: Create Amazon Comprehend Classification training job](#step4)\n",
    "- [Step 5: Create Amazon Comprehend real-time endpoint](#step5)\n",
    "- [Step 6: Classify Documents using the real-time endpoint](#step6)\n",
    "\n",
    "\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e3c0d52",
   "metadata": {},
   "source": [
    "# Step 1: Setup notebook and upload  sample documents to Amazon S3 <a id=\"step1\"></a>\n",
    "\n",
    "In this step, we will import some necessary libraries that will be used throughout this notebook. We will then upload all the documents from the `/classification-training` folder to SageMaker's default bucket."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a67449e",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "!python -m pip install -q amazon-textract-response-parser --upgrade\n",
    "!python -m pip install -q amazon-textract-caller --upgrade\n",
    "!python -m pip install -q amazon-textract-prettyprinter --upgrade"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7f8a235",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "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"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5bf31866",
   "metadata": {},
   "source": [
    "If the import statements above fails then please restart the notebook kernel by clicking the circular arrow button at the top of the notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66660511",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import boto3\n",
    "import botocore\n",
    "import sagemaker\n",
    "import os\n",
    "import io\n",
    "import datetime\n",
    "import pandas as pd\n",
    "from PIL import Image\n",
    "from pathlib import Path\n",
    "import multiprocessing as mp\n",
    "from IPython.display import Image, display, HTML, JSON\n",
    "\n",
    "# variables\n",
    "data_bucket = sagemaker.Session().default_bucket()\n",
    "region = boto3.session.Session().region_name\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",
   "id": "a8a843f1",
   "metadata": {},
   "source": [
    "### Download and Unzip the sample data `classification-training.zip`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc74349d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!curl https://idp-assets-wwso.s3.us-east-2.amazonaws.com/workshop-data/classification-training.zip --output classification-training.zip"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "582ed56f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import shutil\n",
    "\n",
    "try:\n",
    "    shutil.unpack_archive(\"./classification-training.zip\", extract_dir=\"classification-training\")\n",
    "    print(\"Document archive extracted successfully...\")\n",
    "    for path, subdirs, files in os.walk('./classification-training'):\n",
    "        for name in files:\n",
    "            if name.startswith('.'):\n",
    "                hidden = os.path.join(path, name)\n",
    "                print(f'Removing hidden files/directories: {hidden}')\n",
    "                os.system(f\"rm -rf {hidden}\")\n",
    "        for dirs in subdirs:\n",
    "            if dirs.startswith('.'):\n",
    "                if dirs.startswith('.'):\n",
    "                    hidden = os.path.join(path, dirs)\n",
    "                    print(f'Removing hidden files/directories: {hidden}')\n",
    "                    os.system(f\"rm -rf {hidden}\")\n",
    "except Exception as e:\n",
    "    print(\"Please upload the document zip file classification-training.zip\")\n",
    "    raise e"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0edfeef",
   "metadata": {},
   "source": [
    "### Upload sample data to S3 bucket\n",
    "\n",
    "The sample documents are in `/classification-training` directory. For this workshop, we will be using sample bank statements, invoices, and receipts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33b5e387",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Upload images to S3 bucket:\n",
    "!aws s3 cp classification-training s3://{data_bucket}/idp/textract --recursive --only-show-errors"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14da0c56",
   "metadata": {},
   "source": [
    "### Validate the documents in S3\n",
    "\n",
    "We will create a small utility function to verify that our documents have been uploaded to the S3 bucket. This function will also be used to collect the document paths (S3 keys) into an array that we will user later to extract text using Amazon Textract."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba331a70",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def get_s3_bucket_items(bucket, prefix, start_after):\n",
    "    list_items=[]\n",
    "    \n",
    "    paginator = s3.get_paginator('list_objects_v2')\n",
    "    operation_parameters = {'Bucket': bucket,\n",
    "                            'Prefix': prefix,\n",
    "                            'StartAfter':start_after}\n",
    "    page_iterator = paginator.paginate(**operation_parameters)\n",
    "    for page in page_iterator:\n",
    "        for item in page['Contents']:\n",
    "            list_items.append(item['Key'])\n",
    "    names=list(set([os.path.dirname(x)+'/' for x in list_items]))\n",
    "    images=[x for x in list_items if x not in names and '.ipynb_checkpoints' not in x ]\n",
    "    names=[x.replace(prefix,'').strip('/') for x in names if  '.ipynb_checkpoints' not in x]\n",
    "    return list_items, names, images"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab0ba5b8",
   "metadata": {},
   "source": [
    "The code cell below will list a few documents that have been uploaded to our S3 bucket."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b58e56f6",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "docs=[]\n",
    "\n",
    "train_objects, names, train_images=get_s3_bucket_items(data_bucket, 'idp/textract', 'idp/textract/') \n",
    "docs.append(train_images)\n",
    "\n",
    "if type(docs[0]) is list:\n",
    "    docs=[item for sublist in docs for item in sublist]\n",
    "    \n",
    "names, docs[-10:], docs[:10]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "df6a6649",
   "metadata": {},
   "source": [
    "---\n",
    "# Step 2: Extract text from sample documents using Amazon Textract and label<a id=\"step2\"></a>\n",
    "\n",
    "In this section we  use Amazon Textract's `detect_document_text` API to extract the raw text information for all the documents in S3. We will also label the data according to the document type. This labeled data will be used to train a custom Amazon Comprehend classifier. We define a utility function that uses the `textract_extract_text` API to extract text from a document and find which category (or directory in S3) it belongs to and then label the data and return an array `[<label>, <document_text>]`. \n",
    "\n",
    "In order to extract text from a document using textract we use the `DetectDocumentText` API. You can use the Boto3 version of the API as `textract.detect_document_text`, however in this notebook we will use the `call_textract` tool that we installed earlier in the Notebook ([refer to `amazon-textract-caller`](https://pypi.org/project/amazon-textract-caller/) for more info)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f93877e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def textract_extract_text(document, bucket=data_bucket):        \n",
    "    try:\n",
    "        print(f'Processing document: {document}')\n",
    "        lines = \"\"\n",
    "        row = []\n",
    "        \n",
    "        # using amazon-textract-caller\n",
    "        response = call_textract(input_document=f's3://{bucket}/{document}') \n",
    "        # using pretty printer to get all the lines\n",
    "        lines = get_string(textract_json=response, output_type=[Textract_Pretty_Print.LINES])\n",
    "        \n",
    "        label = [name for name in names if(name in document)]  \n",
    "        row.append(label[0])\n",
    "        row.append(lines)        \n",
    "        return row\n",
    "    except Exception as e:\n",
    "        print (e)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "93b012b8",
   "metadata": {},
   "source": [
    "Call the Textract function defined above for all the documents to extract text --\n",
    "\n",
    "<div class=\"alert alert-block alert-warning\">\n",
    "    <b>⚠️ Note:</b> The code below uses <a href=\"https://docs.python.org/3/library/multiprocessing.html\" target=\"_blank\">multiprocessing</a> and can cause Amazon Textract API throttling due to <a href=\"https://docs.aws.amazon.com/general/latest/gr/textract.html\" target=\"_blank\">Amazon Textract soft limits</a>. The code below should work fine for this workshop in an <b>ml.t3.medium</b> SageMaker instance with 2vCPU. Please excercise caution in using similar patterns in production and make sure to implement proper <a href=\"https://docs.aws.amazon.com/general/latest/gr/api-retries.html\" target=\"_blank\">exponential back off</a> in your code. This code below is for demonstration purposes only.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd4608d2",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "pool = mp.Pool(mp.cpu_count())\n",
    "pool_results = [pool.apply_async(textract_extract_text, (document,data_bucket)) for document in docs]\n",
    "labeled_collection = [res.get() for res in pool_results]\n",
    "pool.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "461b57a3",
   "metadata": {},
   "source": [
    "---\n",
    "# Step 3: Prepare a CSV training dataset for Amazon Comprehend custom classifier training<a id=\"step3\"></a>\n",
    "\n",
    "Now that we have text extracted from our documents and have also labeled them, we will create the training data in order to train an [Amazon Comprehend custom classification model](https://docs.aws.amazon.com/comprehend/latest/dg/how-document-classification.html). Let's take a look at the labeled data. We have 100 sample of each document, so we should have about 300 rows of labeled data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31957382",
   "metadata": {},
   "outputs": [],
   "source": [
    "comprehend_df = pd.DataFrame(labeled_collection, columns=['label','document'])\n",
    "comprehend_df"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f0fdafc",
   "metadata": {},
   "source": [
    "We will create a training dataset from extracted text and upload it to Amazon S3. The training data file will be written in `CSV` format and will be named `comprehend_train_data.csv`. Note that you can have more than one `CSV` file in an S3 bucket for training a Comprehend custom classifier. If you have more than one file, you can specify only the bucket/prefix in call to train the custom classifier. Amazon Comprehend will automatically use all the files under the bucket/prefix for training purposes.\n",
    "\n",
    "The following code cells will upload the training data to the S3 bucket, and create a Custom Comprehend Classifier. You can also create a custom classifier manually, please see the subsequent sections for instructions on how to do that."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac163d35",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Upload Comprehend training data to S3\n",
    "key='idp/comprehend/comprehend_train_data.csv'\n",
    "\n",
    "comprehend_df.to_csv(\"comprehend_train_data.csv\", index=False, header=False)\n",
    "s3.upload_file(Filename='comprehend_train_data.csv', \n",
    "               Bucket=data_bucket, \n",
    "               Key=key)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "addd0e8a",
   "metadata": {},
   "source": [
    "---\n",
    "# Step 4: Create Amazon Comprehend Classification training job <a id=\"step4\"></a>\n",
    "\n",
    "Once we have a labeled dataset ready we are going to create and train a [Amazon Comprehend custom classification model](https://docs.aws.amazon.com/comprehend/latest/dg/how-document-classification.html) with the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9bf534d-5031-4e0a-9514-c08312482414",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# If you have imported a custom classifier 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 5.\n",
    "\n",
    "# document_classifier_arn = \"<model_arn>\"\n",
    "# %store document_classifier_arn"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9447b020",
   "metadata": {},
   "source": [
    "### Create Amazon Comprehend custom classification Training Job\n",
    "\n",
    "We will use Amazon Comprehend's Custom Classification to train our own model for classifying the documents. We will use Amazon Comprehend `CreateDocumentClassifier` API to create a classifier which will train a custom model using the labeled data CSV file we created above. The training data contains extracted text, that was extracted using Amazon Textract, and then labeled."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a1236a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "f's3://{data_bucket}/{key}'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8ca3cd6",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Create a document classifier\n",
    "account_id = boto3.client('sts').get_caller_identity().get('Account')\n",
    "id = str(datetime.datetime.now().strftime(\"%s\"))\n",
    "\n",
    "document_classifier_name = 'Sample-Doc-Classifier-IDP'\n",
    "document_classifier_version = 'Sample-Doc-Classifier-IDP-v1'\n",
    "document_classifier_arn = ''\n",
    "response = None\n",
    "\n",
    "try:\n",
    "    create_response = comprehend.create_document_classifier(\n",
    "        InputDataConfig={\n",
    "            'DataFormat': 'COMPREHEND_CSV',\n",
    "            'S3Uri': f's3://{data_bucket}/{key}'\n",
    "        },\n",
    "        DataAccessRoleArn=role,\n",
    "        DocumentClassifierName=document_classifier_name,\n",
    "        VersionName=document_classifier_version,\n",
    "        LanguageCode='en',\n",
    "        Mode='MULTI_CLASS'\n",
    "    )\n",
    "    \n",
    "    document_classifier_arn = create_response['DocumentClassifierArn']\n",
    "    \n",
    "    print(f\"Comprehend Custom Classifier created with ARN: {document_classifier_arn}\")\n",
    "except Exception as error:\n",
    "    if error.response['Error']['Code'] == 'ResourceInUseException':\n",
    "        print(f'A classifier with the name \"{document_classifier_name}\" already exists.')\n",
    "        document_classifier_arn = f'arn:aws:comprehend:{region}:{account_id}:document-classifier/{document_classifier_name}/version/{document_classifier_version}'\n",
    "        print(f'The classifier ARN is: \"{document_classifier_arn}\"')\n",
    "    else:\n",
    "        print(error)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32253449",
   "metadata": {},
   "outputs": [],
   "source": [
    "%store document_classifier_arn\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9963109f",
   "metadata": {},
   "source": [
    "Check status of the Comprehend Custom Classification Job"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61cf9177",
   "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['DocumentClassifierArn']\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",
    "    describe_custom_classifier = comprehend.describe_document_classifier(\n",
    "        DocumentClassifierArn = jobArn\n",
    "    )\n",
    "    status = describe_custom_classifier[\"DocumentClassifierProperties\"][\"Status\"]\n",
    "    clear_output(wait=True)\n",
    "    print(f\"{current_time} : Custom document classifier: {status}\")\n",
    "    \n",
    "    if status == \"TRAINED\" or status == \"IN_ERROR\":\n",
    "        break\n",
    "        \n",
    "    time.sleep(60)\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fa1f4860",
   "metadata": {},
   "source": [
    "\n",
    "Alternatively, to create a Comprehend Custom Classifier Job manually using the console go to [Amazon Comprehend Console](https://console.aws.amazon.com/comprehend/v2/home?region=us-east-1#classification)\n",
    "  \n",
    "- On the left menu click \"Custom Classification\"\n",
    "- In the \"Classifier models\" section, click on \"Create new model\"\n",
    "- In Model Setting for Model name, enter a name \n",
    "- In Data Specification; select \"Using Single-label\" mode and for Data format select CSV file\n",
    "- For Training dataset browse to your data-bucket created above and select the file `comprehend_train_data.csv`\n",
    "- For IAM role select \"Create an IAM role\" and specify a prefix (this will create a new IAM Role for Comprehend)\n",
    "- Click create"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a4daf51",
   "metadata": {},
   "source": [
    "This job can take ~30 minutes to complete. Once the training job is completed move on to next step."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c889c270-a37f-4f06-bcc6-0da188508b90",
   "metadata": {},
   "source": [
    "---\n",
    "# Step 5: Classify documents with Amazon Comprehend custom classifier\n",
    "\n",
    "In this step we will use Amazon Comprehend custom classification model to classify sample documents. We will use `start_document_classification_job` API to launch an asynchronous job to classify the documents. This API supports documents in their native format (PDF/PNG/JPG/TIF) and can use Amazon Textract behind the scenes to read the text from the documents and subsequently determine the document class. Let's start by uploading our sample documents to the S3 bucket"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eaddbf3d-0e60-42bb-bd95-4f14a3a84645",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!aws s3 cp ./samples s3://{data_bucket}/idp/comprehend --recursive --only-show-errors"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "83d448a4-e17a-46d2-aad7-a6af1ed1132e",
   "metadata": {},
   "source": [
    "Amazon Comprehend Async classification works with PDF, PNG, JPEG, as well as UTF-8 encoded plaintext files. Since our sample documents under the `samples` directory are of PNG format, we will specify a `DocumentReadAction` and use Amazon Textract with the `TEXTRACT_DETECT_DOCUMENT_TEXT` option. This will tell Amazon Comprehend to use Amazon Textract `DetectDocumentText` API behind the scenes to extract the text and then perform classification. For `InputFormat`, we will use `ONE_DOC_PER_FILE` mode which signifies that each file is a single document (the other mode is `ONE_DOC_PER_LINE` which means every line in the plaintext file is a document, this is best suited for small documents such as product reviews or customer service chat transcripts etc.). More on this, see [documentation](https://docs.aws.amazon.com/comprehend/latest/dg/what-is.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57db2b90-5832-425f-b8d8-6b58027d06ad",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import uuid\n",
    "\n",
    "jobname = f'doc-classification-job-{uuid.uuid1()}'\n",
    "print(f'Starting Comprehend Classification job {jobname} with model {document_classifier_arn}')\n",
    "\n",
    "response = comprehend.start_document_classification_job(\n",
    "    JobName=jobname,\n",
    "    DocumentClassifierArn=document_classifier_arn,\n",
    "    InputDataConfig={\n",
    "        'S3Uri': f's3://{data_bucket}/idp/comprehend/mixedbag/',\n",
    "        'InputFormat': 'ONE_DOC_PER_FILE',\n",
    "        'DocumentReaderConfig': {\n",
    "            'DocumentReadAction': 'TEXTRACT_DETECT_DOCUMENT_TEXT',\n",
    "            'DocumentReadMode': 'FORCE_DOCUMENT_READ_ACTION'\n",
    "        }\n",
    "    },\n",
    "    OutputDataConfig={\n",
    "        'S3Uri': f's3://{data_bucket}/idp/comprehend/doc-class-output/'\n",
    "    },\n",
    "    DataAccessRoleArn=role\n",
    ")\n",
    "\n",
    "response"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d49c8019-a033-405e-b7fa-20448959f458",
   "metadata": {},
   "source": [
    "## Check status of the classification job\n",
    "\n",
    "The code block below will check the status of the classification job. If the job completes then it will download the output predictions. The output is a zip file which will contain the inference result for each of the documents being classified. The zip will also contain the output of the Textract operation performed by Amazon Comprehend."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc504e97-d914-4265-be72-9f705b025cc7",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "# Loop through and wait for the training to complete . Takes up to 10 mins \n",
    "import time\n",
    "import json\n",
    "from datetime import datetime\n",
    "import tarfile\n",
    "import os\n",
    "\n",
    "classify_response=response\n",
    "max_time = time.time() + 3*60*60 # 3 hours\n",
    "documents=[]\n",
    "\n",
    "while time.time() < max_time:\n",
    "    now = datetime.now()\n",
    "    current_time = now.strftime(\"%H:%M:%S\")\n",
    "    describe_job = comprehend.describe_document_classification_job(\n",
    "        JobId=classify_response['JobId']\n",
    "    )\n",
    "    status = describe_job[\"DocumentClassificationJobProperties\"][\"JobStatus\"]\n",
    "\n",
    "    print(f\"{current_time} : Custom document classifier Job: {status}\")\n",
    "    \n",
    "    if status == \"COMPLETED\" or status == \"FAILED\":\n",
    "        if status == \"COMPLETED\":\n",
    "            classify_output_file = describe_job[\"DocumentClassificationJobProperties\"][\"OutputDataConfig\"][\"S3Uri\"]\n",
    "            print(f'Output generated - {classify_output_file}')\n",
    "            !mkdir -p classification-output\n",
    "            !aws s3 cp {classify_output_file} ./classification-output\n",
    "            \n",
    "            opfile = os.path.basename(classify_output_file)\n",
    "            # open file\n",
    "            file = tarfile.open(f'./classification-output/{opfile}')\n",
    "            # extracting file\n",
    "            file.extractall('./classification-output')\n",
    "            file.close()\n",
    "            \n",
    "            for file in os.listdir('./classification-output'):\n",
    "                if file.endswith('.out'):\n",
    "                    with open(f'./classification-output/{file}', 'r') as f:\n",
    "                        documents.append(dict(file=file, classification_output=json.load(f)['Classes']))        \n",
    "        else:\n",
    "            print(\"Classification job failed\")\n",
    "            print(describe_job)\n",
    "        break\n",
    "        \n",
    "    time.sleep(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3eb6d7c-5f7b-475a-af20-c8b3737434e6",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "documents"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bba396ca-25f0-472f-8ee2-3418a1e7aa34",
   "metadata": {},
   "source": [
    "Let's take a look at the Amazon Comprehend classification output. We have collected the output for all the files in a documents variable. The script above will download and un-zip the zip file locally, so you can navigate into the classification-output directory from the file browser panel on the left and inspect the files manually."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b6d4fa5-06d5-4600-a215-72e88ce18716",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "classification = []\n",
    "for doc in documents:\n",
    "    document = []    \n",
    "    classes_df = pd.DataFrame(doc['classification_output'])\n",
    "    result = classes_df.iloc[classes_df['Score'].idxmax()]\n",
    "    document.extend([doc['file'].replace(\".out\",\"\"), result.Name, result.Score])    \n",
    "    classification.append(document)\n",
    "    \n",
    "doc_class_df = pd.DataFrame(classification, columns = ['Document', 'DocType', 'Confidence'])\n",
    "doc_class_df                                                      "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8ff7cd5-537d-4a8a-8683-e14c503926b9",
   "metadata": {},
   "source": [
    "Our documents under the `samples/mixedbag` folder has now been classified. We will upload them into S3 with proper prefix label."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a621ec0-ffe7-40d7-8928-bc184da8475e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(f'Default bucket is : s3://{data_bucket}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a8e4dd3-99a1-4e12-96c0-8a0c0da1c0ce",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "root='idp/comprehend/classified-docs'\n",
    "\n",
    "def upload_classified_docs(filename,prefix):\n",
    "    document = os.path.basename(filename)\n",
    "    key = f'{root}/{prefix}/{document}'\n",
    "    print(f'Uploading: {filename}...')\n",
    "    res = s3.upload_file(Filename=f\"./samples/mixedbag/{filename}\", \n",
    "                   Bucket=data_bucket, \n",
    "                   Key=key)\n",
    "    return f'{root}/{prefix}/{document}'\n",
    "    \n",
    "doc_class_df['s3path'] = doc_class_df.apply(lambda row : upload_classified_docs(row['Document'],row['DocType']), axis = 1)\n",
    "\n",
    "#verify uploads\n",
    "[objects['Key'] for objects in s3.list_objects(Bucket=data_bucket, Prefix=f\"{root}/\")['Contents']]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "334c7e55-b316-463e-97b3-1908a71d7371",
   "metadata": {},
   "source": [
    "The Comprehend classification process has also generated the Textract output from the documents (present under the `classification-output/amazon-textract-output`). This directory contains a folder for each document with the Amazon Textract JSON response. Let's load the plain text of each of these documents into the data frame. We use the pretty printer tool to get the LINES out of the documents."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37558117-6121-4ed0-90a3-765d82de0a03",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from textractprettyprinter.t_pretty_print import Textract_Pretty_Print, get_string\n",
    "import json\n",
    "\n",
    "def get_text(doc):\n",
    "    with open(f'classification-output/amazon-textract-output/{doc}/1', 'r') as myfile:\n",
    "        data=myfile.read()\n",
    "    obj = json.loads(data)\n",
    "    text = get_string(textract_json=obj, output_type=[Textract_Pretty_Print.LINES])\n",
    "    return text\n",
    "\n",
    "doc_class_df['DocText'] = doc_class_df.apply(lambda row : get_text(row['Document']), axis = 1)\n",
    "doc_class_df"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9f8f94ee-6a2b-4a6e-8e44-9644982186dd",
   "metadata": {},
   "source": [
    " Finally, we keep a copy of the extracted text of the documents in the mixedbag directory for use later in Notebook 3 (Document Enrichment)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6687cbb-977f-4436-8138-ff55854e9b94",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "doc_class_df.to_csv('extracted_doc.csv')\n",
    "#Upload dataframe as csv to S3\n",
    "s3.upload_file(Filename='extracted_doc.csv', \n",
    "               Bucket=data_bucket, \n",
    "               Key=f'idp/comprehend/extracted/extracted_doc.csv')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6b22ce7",
   "metadata": {},
   "source": [
    "---\n",
    "# Step 6: Create Amazon Comprehend real time endpoint _(optional)_ <a id=\"step5\"></a>\n",
    "\n",
    "Once our Comprehend custom classifier is fully trained (i.e. status = `TRAINED`). We can create a real-time endpoint. We will use this endpoint to classify documents in real time. The following code cells use the `comprehend` Boto3 client to create an endpoint, but you can also create one manually via the console. Instructions on how to do that can be found in the subsequent section."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "154c092c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "#create comprehend endpoint\n",
    "model_arn = document_classifier_arn\n",
    "ep_name = 'idp-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",
    "    ENDPOINT_ARN=endpoint_response['EndpointArn']\n",
    "    print(f'Endpoint created with ARN: {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",
    "        ENDPOINT_ARN = f'arn:aws:comprehend:{region}:{account_id}:document-classifier-endpoint/{ep_name}'\n",
    "        print(f'The classifier endpoint ARN is: \"{ENDPOINT_ARN}\"')\n",
    "        %store ENDPOINT_ARN\n",
    "    else:\n",
    "        print(error)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10a2b1c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "%store ENDPOINT_ARN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b455293",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "display(endpoint_response)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "affa81cc",
   "metadata": {},
   "source": [
    "Alternatively, use the steps below to create a Comprehend endpoint using the AWS console.\n",
    "\n",
    "- Go to [Comprehend on AWS Console](https://console.aws.amazon.com/comprehend/v2/home?region=us-east-1#endpoints) and click on Endpoints in the left menu.\n",
    "- Click on \"Create endpoint\"\n",
    "- Give an Endpoint name; for Custom model type select Custom classification; for version select no version or the latest version of the model.\n",
    "- For Classifier model select from the drop down menu\n",
    "- For Inference Unit select 1\n",
    "- Check \"Acknowledge\"\n",
    "- Click \"Create endpoint\"\n",
    "\n",
    "[It may take ~15 minutes](https://console.aws.amazon.com/comprehend/v2/home?region=us-east-1#endpoints) for the endpoint to get created. The code cell below checks the creation status.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b78e4e0e",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "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",
    "    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 document classifier: {status}\")\n",
    "    \n",
    "    if status == \"IN_SERVICE\" or status == \"FAILED\":\n",
    "        break\n",
    "        \n",
    "    time.sleep(10)\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "130bb4fe",
   "metadata": {},
   "source": [
    "---\n",
    "# Step 7: Classify a document with the real-time endpoint _(optional)_ <a id=\"step6\"></a>\n",
    "\n",
    "Once the endpoint has been created, we will use a mix of documents under the `/samples/mixedbag/` directory and try to classify them to bank statement, invoice, and receipt documents respectively."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b81e0a61",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import random\n",
    "\n",
    "root = \"./samples/mixedbag\"\n",
    "files = []\n",
    "\n",
    "for file in os.listdir(root):\n",
    "    if not file.startswith('.'):\n",
    "        files.append(f'./samples/mixedbag/{file}')\n",
    "\n",
    "files_df = pd.DataFrame(files, columns=[\"Document\"])\n",
    "files_df\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba88021e",
   "metadata": {},
   "source": [
    "Let's view one of the documents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "088e1148",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "file = files_df.sample().iloc[0]['Document']\n",
    "display(Image(filename=file, width=400, height=500))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c002971b",
   "metadata": {},
   "source": [
    "To classify this sample document, we will first convert the documents to ByteArray and then use Textract `classify_document` API to classify it. Since `classify_document` is a real time (synchronous) API we will call it with the document bytes of the above sample document. Again, as before we will let Amazon Comprehend utilize Amazon Textract behind the scenes to read the document and then classify it. Since we are allowing Comprehend to use Amazon Textract behind the scenes to extract the text, we will be limited to use single page documents."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac3b2ee2",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import time \n",
    "from datetime import datetime\n",
    "\n",
    "def classify_doc(document):\n",
    "    imageBytes = None\n",
    "    with open(document, 'rb') as document:\n",
    "        imageBytes = bytearray(document.read())\n",
    "    try:\n",
    "        response = comprehend.classify_document(\n",
    "            Bytes = imageBytes,\n",
    "            EndpointArn=ENDPOINT_ARN,\n",
    "            DocumentReaderConfig={\n",
    "                'DocumentReadAction': 'TEXTRACT_DETECT_DOCUMENT_TEXT',\n",
    "                'DocumentReadMode': 'FORCE_DOCUMENT_READ_ACTION'\n",
    "            }\n",
    "        )\n",
    "               \n",
    "        return response      # return the corresponding response\n",
    "    except Exception as e:\n",
    "        print(e)\n",
    "        return 'error'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3f7ad0f4",
   "metadata": {},
   "source": [
    "Lets now run the inference on our sample document"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b974e35",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "classification_op = classify_doc(file)\n",
    "JSON(classification_op)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4857433a",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "# Cleanup\n",
    "\n",
    "Cleanup is optional if you want to execute subsequent notebooks. \n",
    "\n",
    "Refer to the `05-idp-cleanup.ipynb` for cleanup and deletion of resources."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7635771a",
   "metadata": {},
   "source": [
    "---\n",
    "# Conclusion\n",
    "\n",
    "In this notebook we have trained an Amazon Comprehend custom classifier using our sample documents by extracting the text from the documents using Amazon Textract and labeling the data into a CSV file format training dataset. We then trained an Amazon Comprehend custom classifier with the extracted text and created an Amazon Comprehend Classifier real time endpoint to performe classification of documents.\n",
    "\n",
    "In the next notebook we will look at a few methods to perfrom extraction of key insights from our documents using Amazon Textract."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "470616f6",
   "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": 5
}