{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Using AI21 Grammatical Error Correction (GEC) on SageMaker through Model Packages\n",
    "\n",
    "This sample notebook shows you how to deploy **AI21 Grammatical Error Correction (GEC)** using Amazon SageMaker.\n",
    "\n",
    "\n",
    "--------------------\n",
    "## <font color='orange'>Important:</font>\n",
    "Please visit model detail page in <a href=\"https://aws.amazon.com/marketplace/pp/prodview-dplsurctqbgf6 \">https://aws.amazon.com/marketplace/pp/prodview-dplsurctqbgf6</a> to learn more. <font color='orange'>If you do not have access to the link, please contact account admin for the help.</font>\n",
    "\n",
    "You will find details about the model including pricing, supported region, and end user license agreement. To use the model, please click “<font color='orange'>Continue to Subscribe</font>” from the detail page, come back here and learn how to deploy and inference.\n",
    "\n",
    "\n",
    "-------------------\n",
    "\n",
    "\n",
    "## Pre-requisites:\n",
    "1. Before running this notebook, please make sure you got this notebook from the model catalog on SageMaker AWS Management Console.\n",
    "1. **Note**: This notebook contains elements which render correctly in Jupyter interface. Open this notebook from an Amazon SageMaker Notebook Instance or Amazon SageMaker Studio.\n",
    "1. Ensure that IAM role used has **AmazonSageMakerFullAccess**.\n",
    "1. This notebook is intended to work with **boto3 v1.25.4** or higher.\n",
    "\n",
    "## Contents:\n",
    "1. [Select model package](#1.-Select-model-package)\n",
    "1. [Create an endpoint and perform real-time inference](#2.-Create-an-endpoint-and-perform-real-time-inference)\n",
    "   1. [Create an endpoint](#A.-Create-an-endpoint)\n",
    "   1. [Proofread and detect mistakes](#B.-Proofread-and-detect-mistakes)\n",
    "1. [Clean-up](#3.-Clean-up)\n",
    "   1. [Delete the endpoint](#A.-Delete-the-endpoint)\n",
    "   1. [Delete the model](#B.-Delete-the-model)\n",
    "\n",
    "\n",
    "## Usage instructions\n",
    "You can run this notebook one cell at a time (By using Shift+Enter for running a cell)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "from sagemaker import ModelPackage\n",
    "from sagemaker import get_execution_role\n",
    "import sagemaker as sage\n",
    "import boto3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Check the version of boto3 - must be v1.25.4 or higher\n",
    "If you see a lower version number, pick another kernel to run the notebook, with Python 3.8 or above"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "boto3.__version__"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Install ai21 python SDK"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install -U \"ai21[AWS]>=2.1.2\"\n",
    "import ai21"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Check the version of ai21 - must be 2.1.2 or higher"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ai21\n",
    "ai21.__version__"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Select model package\n",
    "Confirm that you received this notebook from the model catalog in SageMaker AWS Management Console."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "region = boto3.Session().region_name\n",
    "\n",
    "# Get the updated ARN\n",
    "model_package_arn = ai21.SageMaker.get_model_package_arn(model_name=\"gec\", region=region)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "role = get_execution_role()\n",
    "sagemaker_session = sage.Session()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Create an endpoint and perform real-time inference"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you want to understand how real-time inference with Amazon SageMaker works, see [Documentation](https://docs.aws.amazon.com/sagemaker/latest/dg/deploy-model.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "endpoint_name = \"gec\"\n",
    "\n",
    "content_type = \"application/json\"\n",
    "\n",
    "real_time_inference_instance_type = (\n",
    "    \"ml.g4dn.2xlarge\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### A. Create an endpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create a deployable model from the model package.\n",
    "model = ModelPackage(\n",
    "    role=role, model_package_arn=model_package_arn, sagemaker_session=sagemaker_session\n",
    ")\n",
    "\n",
    "# Deploy the model\n",
    "predictor = model.deploy(\n",
    "    initial_instance_count=1,\n",
    "    instance_type=real_time_inference_instance_type,\n",
    "    endpoint_name=endpoint_name, \n",
    "    model_data_download_timeout=3600,\n",
    "    container_startup_health_check_timeout=600,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once endpoint has been created, you would be able to perform real-time inference."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### B. Proofread and detect mistakes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**AI21 Grammatical Error Correction (GEC) model** offers access to our world-class proofreading engine. It has been specifically developed for providing suggestions on how to correct spelling mistakes, grammar, punctuation, word misuse, and more.\n",
    "\n",
    "This model takes a piece of text and returns a list of suggested corrections for detected mistakes, along with the type of mistake."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The input text should contain **no more than 500 characters**. The output will include all detected mistakes and suggested corrections."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following is a simple example of detecting different types of grammatical errors and suggesting corrections for them."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "First, create an AI21 Sagemaker client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from ai21 import AI21SageMakerClient\n",
    "client = AI21SageMakerClient(endpoint_name=endpoint_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "text = \"Your grammer is a reflection of you image. Good or bad: you have made an an impression. And like all impressions, you in total control.\"\n",
    "\n",
    "response = client.gec.create(\n",
    "    text=text,\n",
    ")\n",
    "\n",
    "print(response.corrections)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's print the suggested corrections in a more comfortable way:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Original text (with errors):\")\n",
    "print(text)\n",
    "print(\"============================\")\n",
    "print(\"Suggested corrections:\")\n",
    "print(\"\\n\".join([\n",
    "    \"- {correctionType} error detected at position ({startIndex},{endIndex}): {originalText} --> {suggestion}\".format(\n",
    "        correctionType=x.correction_type,\n",
    "        startIndex=x.start_index,\n",
    "        endIndex=x.end_index,\n",
    "        originalText=x.original_text,\n",
    "        suggestion=x.suggestion\n",
    "    ) for x in response.corrections]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Interested in learning more?\n",
    "Take a look at our [guide](https://docs.ai21.com/docs/gec-api) to understand all the capabilities of AI21 Grammatical Error model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Clean-up"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### A. Delete the endpoint"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that you have successfully performed a real-time inference, you do not need the endpoint any more. You can terminate the endpoint to avoid being charged."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.sagemaker_session.delete_endpoint(endpoint_name)\n",
    "model.sagemaker_session.delete_endpoint_config(endpoint_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### B. Delete the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.delete_model()"
   ]
  }
 ],
 "metadata": {
  "instance_type": "ml.t3.medium",
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}