Amazon Augmented AI (A2I) and Amazon Translate.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Amazon Augmented AI (Amazon A2I) integration with Amazon Translate [Example]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Introduction\n",
    "\n",
    "Amazon Translate is constantly learning and evolving to provide the “perfect” output. In domain sensitive applications such as legal, medical, construction, engineering, etc., customers can always improve the translation quality by using custom terminology (https://aws.amazon.com/blogs/machine-learning/introducing-amazon-translate-custom-terminology/). This is a great approach for most of the cases but there are some outliers which might require light post-editing by human teams. The post-editing process helps businesses to understand the needs of their customers better by capturing the nuances of local language that can be lost in translation.\n",
    "\n",
    "For such businesses and organizations who want to augment the output of Amazon Translate (and other Amazon AI services), Amazon Augmented AI (https://aws.amazon.com/augmented-ai/) (A2I) provides a managed approach to build human driven post-editing workflows. Amazon A2I brings human review to all developers, removing the undifferentiated heavy lifting associated with building human review systems or managing large numbers of human reviewers.\n",
    "\n",
    "In this tutorial, we will show how you can use **Amazon Augmented A2I and Amazon Translate to create a human review workflow which allows your private workforce to effectively review, correct and tag the documents translated by Amazon Translate, at scale**.\n",
    "\n",
    "To incorporate A2I in your Amazon Translate Workflows, you will the following resources:\n",
    "\n",
    "1. An **S3 Bucket** to store the files that you need to translate and process the output generated from the Human Review Workflow after the Human Loop has completed. \n",
    "\n",
    "2. A **Worker Team** to review and improve the translations done using Amazon Translate. To learn more about Private Worker Teams, see https://docs.aws.amazon.com/sagemaker/latest/dg/sms-workforce-private.html\n",
    "\n",
    "3. A **Worker Task Template** to create a worker UI. The worker UI displays your input data, such as documents or images, and instructions to workers. It also provides interactive tools that the worker uses to complete your tasks. For more information, see https://docs.aws.amazon.com/sagemaker/latest/dg/a2i-instructions-overview.html\n",
    "\n",
    "4. A **Human Review Workflow**, also referred to as a flow definition. You use the flow definition to configure your human workforce and provide information about how to accomplish the human review task. You can create a flow definition in the Amazon Augmented AI console or with Amazon A2I APIs. To learn more about both of these options, see https://docs.aws.amazon.com/sagemaker/latest/dg/a2i-create-flow-definition.html\n",
    "    \n",
    "When using a custom task type, as this tutorial will show, you start a human loop using the Amazon Augmented AI Runtime API. When you call `start_human_loop()` in your custom application, a task is sent to human reviewers."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prerequisite Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# First, let's get the latest installations of our dependencies\n",
    "!pip install --upgrade pip\n",
    "!pip install boto3 --upgrade\n",
    "!pip install -U botocore"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Environment Setup\n",
    "\n",
    "We need to set up the following data:\n",
    "* `REGION` - Region to call A2I.\n",
    "* `BUCKET_NAME` - A S3 bucket accessible by the given role\n",
    "    * Used to store the input files and output results\n",
    "    * Must be within the same region A2I is called from\n",
    "* `WORKTEAM_ARN` - To create your **Private Workteam**, visit the instructions here: https://docs.aws.amazon.com/sagemaker/latest/dg/sms-workforce-private.html After you have created your workteam, replace *\\<YOUR-WORKTEAM-ARN\\>* below\n",
    "* `ROLE` - The IAM role used as part of StartHumanLoop. By default, this notebook will use the execution role. You can learn more about IAM Policies here https://docs.aws.amazon.com/IAM/latest/UserGuide/access_policies.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "REGION = '<REGION-ID>'\n",
    "BUCKET_NAME = '<BUCKET-NAME>'\n",
    "WORKTEAM_ARN= \"<YOUR-WORKTEAM-ARN>\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Role and Permissions\n",
    "\n",
    "The AWS IAM Role used to execute the notebook needs to have the following policies attached:\n",
    "\n",
    "* SagemakerFullAccess\n",
    "* TranslateFullAccess  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker import get_execution_role\n",
    "import sagemaker\n",
    "\n",
    "# Setting Role to the default SageMaker Execution Role\n",
    "ROLE = get_execution_role()\n",
    "display(ROLE)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setup Bucket and Paths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import boto3\n",
    "import botocore\n",
    "\n",
    "sess = sagemaker.Session()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Client Setup"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's setup the clients for Amazon S3, Amazon SageMaker A2I Runtime and Amazon Translate."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "import io\n",
    "import json\n",
    "import uuid\n",
    "import botocore\n",
    "import time\n",
    "import botocore\n",
    "\n",
    "# Amazon SageMaker client\n",
    "sagemaker = boto3.client('sagemaker', REGION)\n",
    "\n",
    "# Amazon Translate client\n",
    "translate = boto3.client('translate', REGION)\n",
    "\n",
    "# S3 client\n",
    "s3 = boto3.client('s3', REGION)\n",
    "\n",
    "# A2I Runtime client\n",
    "a2i_runtime_client = boto3.client('sagemaker-a2i-runtime', REGION)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set up a pretty printer for the AWS SDK responses"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pprint\n",
    "\n",
    "# Pretty print setup\n",
    "pp = pprint.PrettyPrinter(indent=2)\n",
    "\n",
    "# Function to pretty-print AWS SDK responses\n",
    "def print_response(response):\n",
    "    if 'ResponseMetadata' in response:\n",
    "        del response['ResponseMetadata']\n",
    "    pp.pprint(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sample Data\n",
    "\n",
    "Let's create some sample text that we would test our translation with and store it in S3."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "translation_text = \"\"\"\n",
    "Just then another visitor entered the drawing room: Prince Andrew Bolkónski, the little princess’ husband. He was a very handsome young man, of medium height, with firm, clearcut features. Everything about him, from his weary, bored expression to his quiet, measured step, offered a most striking contrast to his quiet, little wife. It was evident that he not only knew everyone in the drawing room, but had found them to be so tiresome that it wearied him to look at or listen to them. And among all these faces that he found so tedious, none seemed to bore him so much as that of his pretty wife. He turned away from her with a grimace that distorted his handsome face, kissed Anna Pávlovna’s hand, and screwing up his eyes scanned the whole company.\n",
    "\"\"\"\n",
    "\n",
    "key = \"input/test.txt\"\n",
    "\n",
    "s3.put_object(Bucket=BUCKET_NAME, Key=key, Body=translation_text)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create Control Plane Resources"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create a Worker Task Tempalte\n",
    "\n",
    "Create a human task UI resource, giving a UI template in liquid html. This template will be rendered to the human workers whenever human loop is required.\n",
    "\n",
    "For over 70 pre built UIs, check: https://github.com/aws-samples/amazon-a2i-sample-task-uis.\n",
    "\n",
    "We will be taking [translation review and correction UI](https://github.com/aws-samples/amazon-a2i-sample-task-uis/blob/master/text/translation-review-and-correction.liquid.html) and filling in the object categories in the labels variable in the template."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "template = \"\"\"\n",
    "\n",
    "<script src=\"https://assets.crowd.aws/crowd-html-elements.js\"></script>\n",
    "\n",
    "<style>\n",
    "  table, tr, th, td {\n",
    "    border: 1px solid black;\n",
    "    border-collapse: collapse;\n",
    "    padding: 5px;\n",
    "  }\n",
    "</style>\n",
    "\n",
    "<crowd-form>\n",
    "    <div>\n",
    "        <h1>Instructions</h1>\n",
    "        <p>Please review the below translations and make corrections and improvements.</p>\n",
    "        <p>Your corrections should:\n",
    "          <ol>\n",
    "           <li>Make the translated text more accurately express the meaning of the original text</li>\n",
    "           <li>Make the translated text read more like something a person would write rather than an automated translation</li>\n",
    "          </ol>\n",
    "        </p>\n",
    "    </div>\n",
    "\n",
    "    <table>\n",
    "      <tr>\n",
    "        <th>Original</th>\n",
    "        <th>Translation</th>\n",
    "        <th style=\"width: 70px\">Rating</th>\n",
    "      </tr>\n",
    "\n",
    "      {% for pair in task.input.translationPairs %}\n",
    "\n",
    "        <tr>\n",
    "          <td>{{ pair.originalText }}</td>\n",
    "          <td><crowd-text-area name=\"translation{{ forloop.index }}\" value=\"{{ pair.translation }}\"></crowd-text-area></td>\n",
    "          <td>\n",
    "            <p>\n",
    "              <input type=\"radio\" id=\"good{{ forloop.index }}\" name=\"rating{{ forloop.index }}\" value=\"good\" required>\n",
    "              <label for=\"good{{ forloop.index }}\">Good</label>\n",
    "            </p>\n",
    "            <p>\n",
    "              <input type=\"radio\" id=\"bad{{ forloop.index }}\" name=\"rating{{ forloop.index }}\" value=\"bad\" required>\n",
    "              <label for=\"bad{{ forloop.index }}\">Bad</label>       \n",
    "            </p>\n",
    "          </td>\n",
    "        </tr>\n",
    "\n",
    "      {% endfor %}\n",
    "\n",
    "    </table>\n",
    "</crowd-form>\n",
    "\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create a Worker Task Template Creator Function\n",
    "\n",
    "This function would be a higher level abstration, on the SageMaker package's method to create the Worker Task Template which we will use in the next step to create a human review workflow."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_task_ui(task_ui_name, template):\n",
    "    '''\n",
    "    Creates a Human Task UI resource.\n",
    "\n",
    "    Returns:\n",
    "    struct: HumanTaskUiArn\n",
    "    '''\n",
    "    response = sagemaker.create_human_task_ui(\n",
    "        HumanTaskUiName=task_ui_name,\n",
    "        UiTemplate={'Content': template})\n",
    "    return response"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Task UI name - this value is unique per account and region. You can also provide your own value here.\n",
    "taskUIName = 'a2i-translate-test-01-ue-1'\n",
    "\n",
    "# Create task UI\n",
    "humanTaskUiResponse = create_task_ui(taskUIName, template)\n",
    "humanTaskUiArn = humanTaskUiResponse['HumanTaskUiArn']\n",
    "print(humanTaskUiArn)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Creating the Flow Definition"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this section, we're going to create a flow definition definition. Flow Definitions allow us to specify:\n",
    "\n",
    "* The workforce that your tasks will be sent to.\n",
    "* The instructions that your workforce will receive. This is called a worker task template.\n",
    "* Where your output data will be stored.\n",
    "\n",
    "This demo is going to use the API, but you can optionally create this workflow definition in the console as well. \n",
    "\n",
    "For more details and instructions, see: https://docs.aws.amazon.com/sagemaker/latest/dg/a2i-create-flow-definition.html."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_flow_definition(flow_definition_name):\n",
    "    '''\n",
    "    Creates a Flow Definition resource\n",
    "\n",
    "    Returns:\n",
    "    struct: FlowDefinitionArn\n",
    "    '''\n",
    "    response = sagemaker.create_flow_definition(\n",
    "            FlowDefinitionName= flow_definition_name,\n",
    "            RoleArn= ROLE,\n",
    "            HumanLoopConfig= {\n",
    "                \"WorkteamArn\": WORKTEAM_ARN,\n",
    "                \"HumanTaskUiArn\": humanTaskUiArn,\n",
    "                \"TaskCount\": 1,\n",
    "                \"TaskDescription\": \"Please review the translations done using Amazon Translate and make corrections and improvements.\",\n",
    "                \"TaskTitle\": \"Review and Improve translations.\"\n",
    "            },\n",
    "            OutputConfig={\n",
    "                \"S3OutputPath\" : \"s3://\"+BUCKET_NAME+\"/\"\n",
    "            }\n",
    "        )\n",
    "    \n",
    "    return response['FlowDefinitionArn']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we are ready to create our flow definition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Flow definition name - this value is unique per account and region. You can also provide your own value here.\n",
    "uniqueId = str(uuid.uuid4())\n",
    "flowDefinitionName = f'translate-a2i-{uniqueId}' \n",
    "\n",
    "flowDefinitionArn = create_flow_definition(flowDefinitionName)\n",
    "print(flowDefinitionArn)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Translate Documents\n",
    "\n",
    "Now that we have the Human Review Workflow set up, we can translate our documents and pass them over to a Human Loop for review."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get file from S3 and load it into a variable\n",
    "file_contents = s3.get_object(Bucket=BUCKET_NAME, Key=key)['Body'].read().decode(\"utf-8\", 'ignore')\n",
    "\n",
    "# Get just the filename without prefix or suffix\n",
    "fileName = key[key.rindex('/')+1:key.rindex('.')]\n",
    "\n",
    "# Create the human loop input JSON object\n",
    "humanLoopInput = {\n",
    "    'SourceLanguage' : 'English',\n",
    "    'TargetLanguage' : 'Spanish',\n",
    "    'sourceLanguageCode':'en',\n",
    "    'targetLanguageCode' : 'es',\n",
    "    'translationPairs' : [],\n",
    "    'rowCount': 0,\n",
    "    'bucketName': BUCKET_NAME,\n",
    "    'keyName': key\n",
    "}\n",
    "\n",
    "translatedText = ''\n",
    "rowCount = 0\n",
    "\n",
    "print('Splitting file and performing translation')    \n",
    "\n",
    "# split the body by period to get individual sentences\n",
    "for sentence in file_contents.split('.'):\n",
    "    if len(sentence.lstrip()) > 0:\n",
    "        # call translation\n",
    "        translate_response = translate.translate_text(\n",
    "                                Text=sentence + '.',\n",
    "                                SourceLanguageCode='en',\n",
    "                                TargetLanguageCode='es'\n",
    "                            )\n",
    "\n",
    "        translatedSentence = translate_response['TranslatedText']\n",
    "\n",
    "        translationPair = {\n",
    "                            'originalText': sentence + '.',\n",
    "                            'translation': translatedSentence\n",
    "                            }\n",
    "        humanLoopInput['translationPairs'].append(translationPair)\n",
    "        rowCount+=1\n",
    "        translatedText = translatedText + translatedSentence + ' '\n",
    "\n",
    "humanLoopInput['rowCount'] = rowCount\n",
    "\n",
    "humanLoopName = 'Translate-A2I-Text' + str(int(round(time.time() * 1000)))\n",
    "print('Starting human loop - ' + humanLoopName)\n",
    "response = a2i_runtime_client.start_human_loop(\n",
    "                            HumanLoopName=humanLoopName,\n",
    "                            FlowDefinitionArn= flowDefinitionArn,\n",
    "                            HumanLoopInput={\n",
    "                                'InputContent': json.dumps(humanLoopInput)\n",
    "                                }\n",
    "                            )\n",
    "\n",
    "# write the machine translated file to S3 bucket.\n",
    "targetKey = ('machine_output/MO-{0}.txt').format(fileName)\n",
    "print ('Writing translated text to '+ BUCKET_NAME + '/' + targetKey)\n",
    "s3.put_object(Bucket=BUCKET_NAME, Key=targetKey, Body=translatedText.encode('utf-8'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Check Status of Human Loop\n",
    "\n",
    "Let's define a function that allows us to check the status of Human Loop progress."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "resp = a2i_runtime_client.describe_human_loop(HumanLoopName=humanLoopName)\n",
    "print(f'HumanLoop Name: {humanLoopName}')\n",
    "print(f'HumanLoop Status: {resp[\"HumanLoopStatus\"]}')\n",
    "print(f'HumanLoop Output Destination: {resp[\"HumanLoopOutput\"]}')\n",
    "print('\\n')\n",
    "\n",
    "humanLoopStatus = resp[\"HumanLoopStatus\"]\n",
    "outputFilePath = resp[\"HumanLoopOutput\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Wait For Work Team to Complete Task"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "workteamName = WORKTEAM_ARN[WORKTEAM_ARN.rfind('/') + 1:]\n",
    "print(\"Navigate to the private worker portal and do the tasks. Make sure you've invited yourself to your workteam!\")\n",
    "print('https://' + sagemaker.describe_workteam(WorkteamName=workteamName)['Workteam']['SubDomain'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Check Status of Human Loop Again and process Task Results\n",
    "\n",
    "Once the Human Loop Status has changed to completed, you can post process the results to build the final file, with Human Reviewed corrections, for future use. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "resp = a2i_runtime_client.describe_human_loop(HumanLoopName=humanLoopName)\n",
    "humanLoopStatus = resp[\"HumanLoopStatus\"]\n",
    "outputFilePath = resp[\"HumanLoopOutput\"]['OutputS3Uri']\n",
    "\n",
    "if humanLoopStatus == \"Completed\":\n",
    "    # Remove s3:// from S3 File Path\n",
    "    outputFilePath = outputFilePath.replace(\"s3://\", \"\")\n",
    "\n",
    "    # recreate the output text document, including post edits.\n",
    "    tmsFile = s3.get_object(Bucket=outputFilePath.split('/')[0],\n",
    "                                Key=\"/\".join(outputFilePath.split('/')[1:]))['Body'].read()\n",
    "\n",
    "    tmsFile = json.loads(tmsFile.decode('utf-8'))\n",
    "    inputContent = tmsFile['inputContent']\n",
    "    rowcount = inputContent['rowCount']\n",
    "    answerContent = tmsFile['humanAnswers'][0]['answerContent']\n",
    "    editedContent = ''\n",
    "    for index in range(1, rowcount):\n",
    "        editedContent += (answerContent['translation'+str(index)] + \" \")\n",
    "\n",
    "    # extract the file name\n",
    "    targetKeyName = inputContent['keyName']\n",
    "    targetKeyName = targetKeyName[targetKeyName.index('/') + 1: len(targetKeyName)]\n",
    "\n",
    "    # save the file.\n",
    "    s3.put_object(Bucket=BUCKET_NAME,\n",
    "                      Key='post_edits/PO-{0}'.format(targetKeyName),\n",
    "                    Body=editedContent.encode('utf-8'))\n",
    "\n",
    "    print(\"Output File successfully stored in s3://{0}/post_edits/PO-{1}\".format(BUCKET_NAME,targetKeyName))\n",
    "elif humanLoopStatus == \"InProgress\":\n",
    "    print(\"Navigate to the private worker portal and do the tasks. Make sure you've invited yourself to your workteam!\")\n",
    "    print('https://' + sagemaker.describe_workteam(WorkteamName=workteamName)['Workteam']['SubDomain'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Your translated and human reviewed files, are now available in your S3 Bucket "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The End"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_python3",
   "language": "python",
   "name": "conda_python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}