J2_Large_example_model_use.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Using Jurassic-2 Large on SageMaker through Model Packages\n",
    "\n",
    "## <span style='color:Red'>THIS MODEL HAS BEEN DEPRECATED</span>\n",
    "### In its place, we recommend [Jurassic-2 Light](https://github.com/AI21Labs/SageMaker/blob/main/J2_Light_example_model_use.ipynb), an improved model that provides quality results both with natural language instructions and with few shot prompts.\n",
    "\n",
    "===============================================\n",
    "\n",
    "This sample notebook shows you how to deploy **Jurassic-2 Large** using Amazon SageMaker.\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. [Create a few-shot prompt](#C.-Create-a-few-shot-prompt)\n",
    "   1. [Perform real-time inference!](#D.-Perform-real-time-inference)\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": [
    "## 1. Select model package\n",
    "Confirm that you recieved this notebook from model catalog on SageMaker AWS Management Console."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_package_map = {\n",
    "    \"us-east-1\": \"arn:aws:sagemaker:us-east-1:865070037744:model-package/j2-large-v1-0-43-03eb8175aa02346c9bd81decfd2b1e56\",\n",
    "    \"us-east-2\": \"arn:aws:sagemaker:us-east-2:057799348421:model-package/j2-large-v1-0-43-03eb8175aa02346c9bd81decfd2b1e56\",\n",
    "    \"us-west-1\": \"arn:aws:sagemaker:us-west-1:382657785993:model-package/j2-large-v1-0-43-03eb8175aa02346c9bd81decfd2b1e56\",\n",
    "    \"us-west-2\": \"arn:aws:sagemaker:us-west-2:594846645681:model-package/j2-large-v1-0-43-03eb8175aa02346c9bd81decfd2b1e56\",\n",
    "    \"ca-central-1\": \"arn:aws:sagemaker:ca-central-1:470592106596:model-package/j2-large-v1-0-43-03eb8175aa02346c9bd81decfd2b1e56\",\n",
    "    \"eu-central-1\": \"arn:aws:sagemaker:eu-central-1:446921602837:model-package/j2-large-v1-0-43-03eb8175aa02346c9bd81decfd2b1e56\",\n",
    "    \"eu-west-1\": \"arn:aws:sagemaker:eu-west-1:985815980388:model-package/j2-large-v1-0-43-03eb8175aa02346c9bd81decfd2b1e56\",\n",
    "    \"eu-west-2\": \"arn:aws:sagemaker:eu-west-2:856760150666:model-package/j2-large-v1-0-43-03eb8175aa02346c9bd81decfd2b1e56\",\n",
    "    \"eu-west-3\": \"arn:aws:sagemaker:eu-west-3:843114510376:model-package/j2-large-v1-0-43-03eb8175aa02346c9bd81decfd2b1e56\",\n",
    "    \"eu-north-1\": \"arn:aws:sagemaker:eu-north-1:136758871317:model-package/j2-large-v1-0-43-03eb8175aa02346c9bd81decfd2b1e56\",\n",
    "    \"ap-southeast-1\": \"arn:aws:sagemaker:ap-southeast-1:192199979996:model-package/j2-large-v1-0-43-03eb8175aa02346c9bd81decfd2b1e56\",\n",
    "    \"ap-southeast-2\": \"arn:aws:sagemaker:ap-southeast-2:666831318237:model-package/j2-large-v1-0-43-03eb8175aa02346c9bd81decfd2b1e56\",\n",
    "    \"ap-northeast-2\": \"arn:aws:sagemaker:ap-northeast-2:745090734665:model-package/j2-large-v1-0-43-03eb8175aa02346c9bd81decfd2b1e56\",\n",
    "    \"ap-northeast-1\": \"arn:aws:sagemaker:ap-northeast-1:977537786026:model-package/j2-large-v1-0-43-03eb8175aa02346c9bd81decfd2b1e56\",\n",
    "    \"ap-south-1\": \"arn:aws:sagemaker:ap-south-1:077584701553:model-package/j2-large-v1-0-43-03eb8175aa02346c9bd81decfd2b1e56\",\n",
    "    \"sa-east-1\": \"arn:aws:sagemaker:sa-east-1:270155090741:model-package/j2-large-v1-0-43-03eb8175aa02346c9bd81decfd2b1e56\"\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "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": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'1.26.74'"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "boto3.__version__"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Install ai21 python SDK"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Looking in indexes: https://pypi.org/simple, https://pip.repos.neuron.amazonaws.com\n",
      "Requirement already satisfied: ai21[SM] in /home/ec2-user/anaconda3/envs/pytorch_p39/lib/python3.9/site-packages (1.0.5)\n",
      "Requirement already satisfied: requests in /home/ec2-user/anaconda3/envs/pytorch_p39/lib/python3.9/site-packages (from ai21[SM]) (2.28.1)\n",
      "Requirement already satisfied: boto3 in /home/ec2-user/anaconda3/envs/pytorch_p39/lib/python3.9/site-packages (from ai21[SM]) (1.26.74)\n",
      "Requirement already satisfied: s3transfer<0.7.0,>=0.6.0 in /home/ec2-user/anaconda3/envs/pytorch_p39/lib/python3.9/site-packages (from boto3->ai21[SM]) (0.6.0)\n",
      "Requirement already satisfied: botocore<1.30.0,>=1.29.74 in /home/ec2-user/anaconda3/envs/pytorch_p39/lib/python3.9/site-packages (from boto3->ai21[SM]) (1.29.74)\n",
      "Requirement already satisfied: jmespath<2.0.0,>=0.7.1 in /home/ec2-user/anaconda3/envs/pytorch_p39/lib/python3.9/site-packages (from boto3->ai21[SM]) (1.0.1)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /home/ec2-user/anaconda3/envs/pytorch_p39/lib/python3.9/site-packages (from requests->ai21[SM]) (3.4)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /home/ec2-user/anaconda3/envs/pytorch_p39/lib/python3.9/site-packages (from requests->ai21[SM]) (2022.12.7)\n",
      "Requirement already satisfied: charset-normalizer<3,>=2 in /home/ec2-user/anaconda3/envs/pytorch_p39/lib/python3.9/site-packages (from requests->ai21[SM]) (2.1.1)\n",
      "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /home/ec2-user/anaconda3/envs/pytorch_p39/lib/python3.9/site-packages (from requests->ai21[SM]) (1.26.8)\n",
      "Requirement already satisfied: python-dateutil<3.0.0,>=2.1 in /home/ec2-user/anaconda3/envs/pytorch_p39/lib/python3.9/site-packages (from botocore<1.30.0,>=1.29.74->boto3->ai21[SM]) (2.8.2)\n",
      "Requirement already satisfied: six>=1.5 in /home/ec2-user/anaconda3/envs/pytorch_p39/lib/python3.9/site-packages (from python-dateutil<3.0.0,>=2.1->botocore<1.30.0,>=1.29.74->boto3->ai21[SM]) (1.16.0)\n"
     ]
    }
   ],
   "source": [
    "! pip install -U \"ai21[SM]\"\n",
    "import ai21"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "region = boto3.Session().region_name\n",
    "if region not in model_package_map.keys():\n",
    "    raise (\"UNSUPPORTED REGION\")\n",
    "\n",
    "model_package_arn = model_package_map[region]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "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": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "endpoint_name = \"j2-large\"\n",
    "\n",
    "content_type = \"application/json\"\n",
    "\n",
    "real_time_inference_instance_type = (\n",
    "    \"ml.g5.48xlarge\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### A. Create an endpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "----------------"
     ]
    }
   ],
   "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": [
    "You can think of Jurassic-2 Large as a smart auto-completion algorithm: it is very good at latching onto hints and patterns expressed in plain English, and generating text that follows the same patterns.\n",
    "\n",
    "These two helpful concepts are worth being familiar with:\n",
    "- **Prompt** - the input you provide to the model.\n",
    "- **Completion** - the output text the model generates."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Enter a simple prompt: \"To be, or\", and let the model complete it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " not to be: that is the question\n"
     ]
    }
   ],
   "source": [
    "response = ai21.Completion.execute(sm_endpoint=endpoint_name,\n",
    "                                   prompt=\"To be, or\",\n",
    "                                   maxTokens=4,\n",
    "                                   temperature=0,\n",
    "                                   numResults=1)\n",
    "\n",
    "print(response['completions'][0]['data']['text'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you can see, the model identifies the beginning of a famous quote, and completes it correctly."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### C. Create a few-shot prompt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The best way to guide the model is to provide several examples of input-output pairs in the prompt. This establishes a pattern for the model to mimic. Then add the input for a query example and let the model complete it with an appropriate generation. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this example, we will create a prompt for performing named entity recognition (NER) on news headlines."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To do this, we will build a few-shot prompt comprised of the following:\n",
    "\n",
    "1. Prefix with 3 examples. Each example contains the relevant input (a news headline) and the output (extraction of entity, location and time). They are separated by \"##\".\n",
    "\n",
    "2. The query input. An unseen headline for which we would like the model to output the entity, location and time. This should follow the same format of the inputs in the prefix.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, we collect some example data for the prompt prefix:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "EXAMPLES_DATA = [\n",
    "    {\"headline\": \"Inflation cooled to 6% in February 2023 as the Federal Reserve weighs next steps on interest rates\", \n",
    "     \"entity\": \"the Federal Reserve\", \n",
    "     \"time\": \"February 2023\",\n",
    "     \"location\": \"NA\"},\n",
    "    {\"headline\": \"Novo Nordisk to lower list price of some of its insulin by up to 75% in the U.S.\", \n",
    "     \"entity\": \"Novo Nordisk\", \n",
    "     \"time\": \"NA\",\n",
    "     \"location\": \"the U.S.\"},\n",
    "   {\"headline\": \"John Snow says protecting Winterfell is not a 'vital' national interest\", \n",
    "     \"entity\": \"John Snow\", \n",
    "     \"time\": \"NA\",\n",
    "     \"location\": \"Winterfell\"}\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then, we use the following helper functions to construct the prefix:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_single_example(headline, entity, time, location):\n",
    "    example = \"Extract from the following headline these properties: Entity, Time, Location. In the case where it doesn't appear in the sentence, write NA.\\n\"\n",
    "    example += f\"Headline: {headline}\\n\"\n",
    "    if entity:\n",
    "        example += f\"Entity: {entity}\\n\"\n",
    "        example += f\"Time: {time}\\n\"\n",
    "        example += f\"Location: {location}\"\n",
    "    \n",
    "    return example\n",
    "\n",
    "SEPARATOR = \"\\n##\\n\"\n",
    "\n",
    "FEW_SHOT_PREFIX = SEPARATOR.join(\n",
    "    make_single_example(x[\"headline\"], x[\"entity\"], x[\"time\"], x[\"location\"]) for x in EXAMPLES_DATA\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And finally, we create a function to handle query inputs and create the full prompt:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_ner_prompt(headline):\n",
    "    \"\"\"\n",
    "    Create a few-shot prompt to extract named entities with Jurassic-2 Large given a headline\n",
    "    The prompt contains a preset sequence of examples followed by the query headline\n",
    "    \"\"\"\n",
    "    return FEW_SHOT_PREFIX + SEPARATOR + make_single_example(headline, '', '', '')  # keep the entities blank and let the model generate"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's see how this looks for a a new headline about what's happening in Springfield:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Extract from the following headline these properties: Entity, Time, Location. In the case where it doesn't appear in the sentence, write NA.\n",
      "Headline: Inflation cooled to 6% in February 2023 as the Federal Reserve weighs next steps on interest rates\n",
      "Entity: the Federal Reserve\n",
      "Time: February 2023\n",
      "Location: NA\n",
      "##\n",
      "Extract from the following headline these properties: Entity, Time, Location. In the case where it doesn't appear in the sentence, write NA.\n",
      "Headline: Novo Nordisk to lower list price of some of its insulin by up to 75% in the U.S.\n",
      "Entity: Novo Nordisk\n",
      "Time: NA\n",
      "Location: the U.S.\n",
      "##\n",
      "Extract from the following headline these properties: Entity, Time, Location. In the case where it doesn't appear in the sentence, write NA.\n",
      "Headline: John Snow says protecting Winterfell is not a 'vital' national interest\n",
      "Entity: John Snow\n",
      "Time: NA\n",
      "Location: Winterfell\n",
      "##\n",
      "Extract from the following headline these properties: Entity, Time, Location. In the case where it doesn't appear in the sentence, write NA.\n",
      "Headline: Homer Simpson to sign executive order on January 2024 to increase the number of gun background checks in Springfield\n",
      "\n"
     ]
    }
   ],
   "source": [
    "prompt = create_ner_prompt(headline=\"Homer Simpson to sign executive order on January 2024 to increase the number of gun background checks in Springfield\")\n",
    "\n",
    "print(prompt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### D. Perform real-time inference"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Time to put Jurassic-2 Large to work!\n",
    "\n",
    "Let it extract the entity, time and location from our provided headline."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Entity: Homer Simpson\n",
      "Time: January 2024\n",
      "Location: Springfield\n",
      "\n"
     ]
    }
   ],
   "source": [
    "response = ai21.Completion.execute(sm_endpoint=endpoint_name,\n",
    "                                   prompt=prompt,\n",
    "                                   maxTokens=30,\n",
    "                                   temperature=0,\n",
    "                                   stopSequences=['##'],\n",
    "                                   numResults=1)\n",
    "\n",
    "print(response['completions'][0]['data']['text'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Interested in learning more?\n",
    "Take a look at our [blog post](https://www.ai21.com/blog/building-cv-profile-generator-using-ai21-studio) to understand the process of building a good prompt."
   ]
  },
  {
   "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": 38,
   "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": 39,
   "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
}