AI21_ContextualAnswers_example_model_use.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Using AI21 Contextual Answers on SageMaker through Model Packages\n",
    "\n",
    "This sample notebook shows you how to deploy **AI21 Contextual Answers** 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-gwbjdp3tmh3bw\">https://aws.amazon.com/marketplace/pp/prodview-gwbjdp3tmh3bw</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",
    "You can provide your users with relevant, immediate answers in natural language to questions they ask about documents from your knowledge base or website. The AI21 Contextual Answers model ensures accurate and reliable question answering capabilities based entirely on the context of a specific document or article. This model is fully grounded in your provided context, therefore avoiding all factual inaccuracies including hallucinations and distortions. This means that if the answer to a question is not specified in the document, the model will indicate this, rather than returning a potentially inaccurate answer. Easily integrate this model into existing systems **without requiring any prompt engineering.**\n",
    "\n",
    "\n",
    "**AI21 Contextual Answers** allows you to provide your users with relevant, immediate answers in natural language to questions they ask about documents from your knowledge base or website. It ensures accurate and reliable question answering capabilities based entirely on the context of a specific document or article. This means that if the answer to a question is not specified in the document, the model will indicate this, rather than returning a potentially inaccurate answer.\n",
    "As a task-specific model, it can be easily integrated into existing systems without requiring prompt engineering: simply provide the context and the question, and the model will provide the answer.\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. [Interact with the model](#B.-Interact-with-the-model)\n",
    "   1. [Ask about financial reports](#C.-Ask-about-financial-reports)\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=\"contextual-answers\", 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": [
    "### <span style='color:Blue'> How to choose the best instance for my use case?</span>\n",
    "<span style='color:#0057FF'> When you create your endpoint, you need to choose the instance type to run the model on. Choosing the right instance is mainly a matter of economics. Depending on your use case, you probably want the most cost-effective instance possible. In this notebook we use one of the supported instances.</span>\n",
    "\n",
    "<span style='color:#0057FF'>Looking for the list of all supported instances? See</span> [here](https://docs.ai21.com/docs/choosing-the-right-instance-type-for-amazon-sagemaker-models#ai21-contextual-answers)."
   ]
  },
  {
   "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 = \"contextual-answers\"\n",
    "\n",
    "content_type = \"application/json\"\n",
    "\n",
    "real_time_inference_instance_type = (\n",
    "    \"ml.p4d.24xlarge\"    # Recommended instance\n",
    "#     \"ml.g5.48xlarge\"   # Cheaper and faster - recommended for relatively short inputs/outputs\n",
    "#     \"ml.g5.12xlarge\"   # Even more cheaper and faster - up to 10K characters\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. Interact with the model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**AI21 Studio Contextual Answers model** allows you to access our high-quality question answering technology. It was designed to answer questions based on a specific document context provided by the customer. This avoids any factual issues that language models may have and makes sure the answers it provides are grounded in that context document.\n",
    "\n",
    "This model receives document text, serving as a context, and a question and returns an answer based entirely on this context. This means that if the answer to your question is not in the document, the model will indicate it (instead of providing a false answer)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To get a sense of the model's behavior, let's use this toy example of asking what is the Eiffel tower height. Most language models will simply answer according to their training data.\n",
    "\n",
    "This model, however, bases its answer solely on the context you provide. Let's use the following [Wikipedia paragraph](https://en.wikipedia.org/wiki/Eiffel_Tower#:~:text=The%20Eiffel%20Tower%20(%2F%CB%88a%C9%AA,from%20the%20Champ%20de%20Mars) as context, with small modifications:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Actual paragraph\n",
    "context = \"The tower is 330 metres (1,083 ft) tall,[6] about the same height as an 81-storey building, and the tallest structure in Paris. Its base is square, measuring 125 metres (410 ft) on each side. During its construction, the Eiffel Tower surpassed the Washington Monument to become the tallest human-made structure in the world, a title it held for 41 years until the Chrysler Building in New York City was finished in 1930. It was the first structure in the world to surpass both the 200-metre and 300-metre mark in height. Due to the addition of a broadcasting aerial at the top of the tower in 1957, it is now taller than the Chrysler Building by 5.2 metres (17 ft). Excluding transmitters, the Eiffel Tower is the second tallest free-standing structure in France after the Millau Viaduct.\"\n",
    "\n",
    "# The paragraph with manual changes of the height\n",
    "false_context = \"The tower is 3 metres (10 ft) tall,[6] about the same height as an 81-storey building, and the tallest structure in Paris. Its base is square, measuring 125 metres (410 ft) on each side. During its construction, the Eiffel Tower surpassed the Washington Monument to become the tallest human-made structure in the world, a title it held for 41 years until the Chrysler Building in New York City was finished in 1930. It was the first structure in the world to surpass both the 200-metre and 300-metre mark in height. Due to the addition of a broadcasting aerial at the top of the tower in 1957, it is now taller than the Chrysler Building by 5.2 metres (17 ft). Excluding transmitters, the Eiffel Tower is the second tallest free-standing structure in France after the Millau Viaduct.\"\n",
    "\n",
    "# The paragraph with the height omitted\n",
    "partial_context = \"Its base is square, measuring 125 metres (410 ft) on each side. During its construction, the Eiffel Tower surpassed the Washington Monument to become the tallest human-made structure in the world, a title it held for 41 years until the Chrysler Building in New York City was finished in 1930. It was the first structure in the world to surpass both the 200-metre and 300-metre mark in height. Due to the addition of a broadcasting aerial at the top of the tower in 1957, it is now taller than the Chrysler Building by 5.2 metres (17 ft). Excluding transmitters, the Eiffel Tower is the second tallest free-standing structure in France after the Millau Viaduct.\""
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "First, create an AI21 Sagemaker client"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "from ai21 import AI21SageMakerClient\n",
    "client = AI21SageMakerClient(endpoint_name=endpoint_name)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is what the model will say when asked the same question in each context.\n",
    "\n",
    "With the true context, the model will answer the correct response (which is indeed indicated in the provided context):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# True context\n",
    "response = client.answer.create(\n",
    "    context=context,\n",
    "    question=\"What is the height of the Eiffel tower?\",\n",
    ")\n",
    "\n",
    "print(response.answer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we feed the model with context which is false in real life, it will still answer according to provided context:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# False context\n",
    "response = client.answer.create(\n",
    "    context=false_context,\n",
    "    question=\"What is the height of the Eiffel tower?\",\n",
    ")\n",
    "\n",
    "print(response.answer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we omitted the first line from the context, hence providing the model a context without any information on the height. Instead of making up an answer, the model will simply say that the answer is not in the provided document:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Irrelevant context\n",
    "response = client.answer.create(\n",
    "    context=partial_context,\n",
    "    question=\"What is the height of the Eiffel tower?\",\n",
    ")\n",
    "\n",
    "print(response.answer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Important - Model Usage\n",
    "The input to the model should be simple and straight-forward, no prompting needed.\n",
    "\n",
    "**Good question:** *What is the height of the Eiffel tower?*\n",
    "\n",
    "**Bad question:** *Please answer the following question: What is the height of the Eiffel tower?*\n",
    "\n",
    "**Another bad question:** *You are an expert model, answer the question according to the context. What is the height of the Eiffel tower?*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### C. Ask about financial reports"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The document context should be **no more than 10,000 characters**, and the question can be up to 160 characters."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Imagine you are performing research and rely on financial reports to base your findings. Let's take the following part from [JPMorgan Chase & Co. 2021 annual report](https://www.jpmorganchase.com/content/dam/jpmc/jpmorgan-chase-and-co/investor-relations/documents/annualreport-2021.pdf):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "financial_context = \"\"\"In 2020 and 2021, enormous QE — approximately $4.4 trillion, or 18%, of 2021 gross domestic product (GDP) — and enormous fiscal stimulus (which has been and always will be inflationary) — approximately $5 trillion, or 21%, of 2021 GDP — stabilized markets and allowed companies to raise enormous amounts of capital. In addition, this infusion of capital saved many small businesses and put more than $2.5 trillion in the hands of consumers and almost $1 trillion into state and local coffers. These actions led to a rapid decline in unemployment, dropping from 15% to under 4% in 20 months — the magnitude and speed of which were both unprecedented. Additionally, the economy grew 7% in 2021 despite the arrival of the Delta and Omicron variants and the global supply chain shortages, which were largely fueled by the dramatic upswing in consumer spending and the shift in that spend from services to goods. Fortunately, during these two years, vaccines for COVID-19 were also rapidly developed and distributed.\n",
    "In today's economy, the consumer is in excellent financial shape (on average), with leverage among the lowest on record, excellent mortgage underwriting (even though we've had home price appreciation), plentiful jobs with wage increases and more than $2 trillion in excess savings, mostly due to government stimulus. Most consumers and companies (and states) are still flush with the money generated in 2020 and 2021, with consumer spending over the last several months 12% above pre-COVID-19 levels. (But we must recognize that the account balances in lower-income households, smaller to begin with, are going down faster and that income for those households is not keeping pace with rising inflation.)\n",
    "Today's economic landscape is completely different from the 2008 financial crisis when the consumer was extraordinarily overleveraged, as was the financial system as a whole — from banks and investment banks to shadow banks, hedge funds, private equity, Fannie Mae and many other entities. In addition, home price appreciation, fed by bad underwriting and leverage in the mortgage system, led to excessive speculation, which was missed by virtually everyone — eventually leading to nearly $1 trillion in actual losses.\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Rather than reading the entire report, just ask what you want to know:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "question = \"Did the economy shrink after the Omicron variant arrived?\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The model will answer based on the provided report:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = client.answer.create(\n",
    "    context=financial_context,\n",
    "    question=question,\n",
    ")\n",
    "\n",
    "print(response.answer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In addition, you can ask more complex questions, where the answer requires deductions rather than just extracting the correct sentence from the document context. This will result in abstractive, rather than extractive, answers that draw on several different parts of the document. For example, look at the following question:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "harder_question = \"Did COVID-19 eventually help the economy?\"\n",
    "\n",
    "response = client.answer.create(\n",
    "    context=financial_context,\n",
    "    question=harder_question,\n",
    ")\n",
    "\n",
    "print(response.answer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now present the model with the following question. You may be confused to answer something based on the last paragraph without delving into the text. However, if you read the provided document context properly, you will discover that the answer does not appear there. The model handles this as expected:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "irrelevant_question = \"How did COVID-19 affect the financial crisis of 2008?\"\n",
    "\n",
    "response = client.answer.create(\n",
    "    context=financial_context,\n",
    "    question=irrelevant_question,\n",
    ")\n",
    "\n",
    "print(response.answer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Interested in learning more?\n",
    "Take a look at our [guide](https://docs.ai21.com/docs/contextual-answers-api) to understand all the capabilities of AI21 Contextual Answers 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
}