J2_Jumbo_example_model_use.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Using Jurassic-2 Jumbo 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 Ultra](https://github.com/AI21Labs/SageMaker/blob/main/J2_Ultra_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 Jumbo** 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-jumbo-v1-0-43-bb4674f324d7340e9fb7d8d4abb54dd9\",\n",
    "    \"us-east-2\": \"arn:aws:sagemaker:us-east-2:057799348421:model-package/j2-jumbo-v1-0-43-bb4674f324d7340e9fb7d8d4abb54dd9\",\n",
    "    \"us-west-1\": \"arn:aws:sagemaker:us-west-1:382657785993:model-package/j2-jumbo-v1-0-43-bb4674f324d7340e9fb7d8d4abb54dd9\",\n",
    "    \"us-west-2\": \"arn:aws:sagemaker:us-west-2:594846645681:model-package/j2-jumbo-v1-0-43-bb4674f324d7340e9fb7d8d4abb54dd9\",\n",
    "    \"ca-central-1\": \"arn:aws:sagemaker:ca-central-1:470592106596:model-package/j2-jumbo-v1-0-43-bb4674f324d7340e9fb7d8d4abb54dd9\",\n",
    "    \"eu-central-1\": \"arn:aws:sagemaker:eu-central-1:446921602837:model-package/j2-jumbo-v1-0-43-bb4674f324d7340e9fb7d8d4abb54dd9\",\n",
    "    \"eu-west-1\": \"arn:aws:sagemaker:eu-west-1:985815980388:model-package/j2-jumbo-v1-0-43-bb4674f324d7340e9fb7d8d4abb54dd9\",\n",
    "    \"eu-west-2\": \"arn:aws:sagemaker:eu-west-2:856760150666:model-package/j2-jumbo-v1-0-43-bb4674f324d7340e9fb7d8d4abb54dd9\",\n",
    "    \"eu-west-3\": \"arn:aws:sagemaker:eu-west-3:843114510376:model-package/j2-jumbo-v1-0-43-bb4674f324d7340e9fb7d8d4abb54dd9\",\n",
    "    \"eu-north-1\": \"arn:aws:sagemaker:eu-north-1:136758871317:model-package/j2-jumbo-v1-0-43-bb4674f324d7340e9fb7d8d4abb54dd9\",\n",
    "    \"ap-southeast-1\": \"arn:aws:sagemaker:ap-southeast-1:192199979996:model-package/j2-jumbo-v1-0-43-bb4674f324d7340e9fb7d8d4abb54dd9\",\n",
    "    \"ap-southeast-2\": \"arn:aws:sagemaker:ap-southeast-2:666831318237:model-package/j2-jumbo-v1-0-43-bb4674f324d7340e9fb7d8d4abb54dd9\",\n",
    "    \"ap-northeast-2\": \"arn:aws:sagemaker:ap-northeast-2:745090734665:model-package/j2-jumbo-v1-0-43-bb4674f324d7340e9fb7d8d4abb54dd9\",\n",
    "    \"ap-northeast-1\": \"arn:aws:sagemaker:ap-northeast-1:977537786026:model-package/j2-jumbo-v1-0-43-bb4674f324d7340e9fb7d8d4abb54dd9\",\n",
    "    \"ap-south-1\": \"arn:aws:sagemaker:ap-south-1:077584701553:model-package/j2-jumbo-v1-0-43-bb4674f324d7340e9fb7d8d4abb54dd9\",\n",
    "    \"sa-east-1\": \"arn:aws:sagemaker:sa-east-1:270155090741:model-package/j2-jumbo-v1-0-43-bb4674f324d7340e9fb7d8d4abb54dd9\"\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: 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: 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: 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: 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: 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-jumbo\"\n",
    "\n",
    "content_type = \"application/json\"\n",
    "\n",
    "real_time_inference_instance_type = (\n",
    "    \"ml.p4d.24xlarge\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### A. Create an endpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "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 Jumbo 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": 9,
   "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 that generates product descriptions for a fashion eCommerce website."
   ]
  },
  {
   "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 inputs (a product name and some features to incorporate) and the output (an engaging product description). They are separated by \"##\".\n",
    "\n",
    "2. The query inputs. An unseen product name and set of features for which we would like the model to output a new product description. These 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": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "EXAMPLES_DATA = [\n",
    "    {\"product\": \"Humor Men's Graphic T-Shirt\", \n",
    "     \"features\": [\"Funny print\", \"Soft cotton\", \"Short sleeve\"], \n",
    "     \"description\": \"Funny is cool with this Graphic Print T-Shirt. Crafted in soft cotton fabric and featuring a classic short sleeve silhouette for a laid-back look and style, this men’s tee pairs with your favorite jeans or shorts and easily earns a spot on the graphic tee rotation—perfect for fun weekends or your everyday look!\"},\n",
    "    {\"product\": \"Women's Boho Beach Dress\", \n",
    "     \"features\": [\"Midi dress\", \"Swing hem\", \"Slimming\", \"Special for summer\"], \n",
    "     \"description\": \"Women's Boho Beach Dress The style of midi dress with swing hem is slimming, It will make your tummy look flatter. This dress for women could not only make you feel cool, but also perfectly present the curve of female waist. Even in hot summer, it still keeps your elegant.\"},\n",
    "    {\"product\":\"Men's Soho Trousers\", \n",
    "     \"features\": [\"Designed by Retrolicious\", \"Stretch cotton fabric\", \"Side pockets\", \"Stripes print\"],\n",
    "     \"description\": \"Designed with a side stripe and a modern slim fit, these Retrolicious Soho Trousers are the perfect pair for everyday wear. Its stretch cotton fabric makes it comfortable to wear all day long, while side pockets offer convenient storage.\"}\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then, we use the following helper functions to construct the prefix:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def list_features(features):\n",
    "    return \"\\n\".join(\n",
    "        f\"- {feat}\" for feat in features\n",
    "    )\n",
    "\n",
    "def make_single_example(product, features, description):\n",
    "    example = \"Write an engaging product description for clothing eCommerce site based on a list of features.\\n\"\n",
    "    example += f\"Product: {product}\\n\"\n",
    "    example += \"Features:\\n\"\n",
    "    example += list_features(features)\n",
    "    example += \"\\nDescription:\\n\"\n",
    "    example += description\n",
    "    \n",
    "    return example\n",
    "\n",
    "SEPARATOR = \"\\n\\n##\\n\\n\"\n",
    "\n",
    "FEW_SHOT_PREFIX = SEPARATOR.join(\n",
    "    make_single_example(x[\"product\"], x[\"features\"], x[\"description\"]) 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": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_product_description_prompt(product, features):\n",
    "    \"\"\"\n",
    "    Create a few-shot prompt to generate a product description with Jurassic-2 Jumbo given a name and features\n",
    "    The prompt contains a preset sequence of examples followed by the query product name and features\n",
    "    \"\"\"\n",
    "    return FEW_SHOT_PREFIX + SEPARATOR + make_single_example(product, features, '')  # keep the description blank and let the model generate"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's see how this looks for a director bag role with specific features:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Write an engaging product description for clothing eCommerce site based on a list of features.\n",
      "Product: Humor Men's Graphic T-Shirt\n",
      "Features:\n",
      "- Funny print\n",
      "- Soft cotton\n",
      "- Short sleeve\n",
      "Description:\n",
      "Funny is cool with this Graphic Print T-Shirt. Crafted in soft cotton fabric and featuring a classic short sleeve silhouette for a laid-back look and style, this men’s tee pairs with your favorite jeans or shorts and easily earns a spot on the graphic tee rotation—perfect for fun weekends or your everyday look!\n",
      "\n",
      "##\n",
      "\n",
      "Write an engaging product description for clothing eCommerce site based on a list of features.\n",
      "Product: Women's Boho Beach Dress\n",
      "Features:\n",
      "- Midi dress\n",
      "- Swing hem\n",
      "- Slimming\n",
      "- Special for summer\n",
      "Description:\n",
      "Women's Boho Beach Dress The style of midi dress with swing hem is slimming, It will make your tummy look flatter. This dress for women could not only make you feel cool, but also perfectly present the curve of female waist. Even in hot summer, it still keeps your elegant.\n",
      "\n",
      "##\n",
      "\n",
      "Write an engaging product description for clothing eCommerce site based on a list of features.\n",
      "Product: Men's Soho Trousers\n",
      "Features:\n",
      "- Designed by Retrolicious\n",
      "- Stretch cotton fabric\n",
      "- Side pockets\n",
      "- Stripes print\n",
      "Description:\n",
      "Designed with a side stripe and a modern slim fit, these Retrolicious Soho Trousers are the perfect pair for everyday wear. Its stretch cotton fabric makes it comfortable to wear all day long, while side pockets offer convenient storage.\n",
      "\n",
      "##\n",
      "\n",
      "Write an engaging product description for clothing eCommerce site based on a list of features.\n",
      "Product: Camp Director Crossbody Bag\n",
      "Features:\n",
      "- Black canvas purse\n",
      "- Rainbow space print\n",
      "- Leather trim\n",
      "- Two securely-zipped compartments\n",
      "Description:\n",
      "\n"
     ]
    }
   ],
   "source": [
    "prompt = create_product_description_prompt(\n",
    "    product=\"Camp Director Crossbody Bag\", \n",
    "    features=[\"Black canvas purse\", \"Rainbow space print\", \"Leather trim\", \"Two securely-zipped compartments\"]\n",
    ")\n",
    "\n",
    "print(prompt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### D. Perform real-time inference"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Time to put Jurassic-2 Jumbo to work!\n",
    "\n",
    "Let it generate a product description for the given input. Since the generation is not deterministic, you can ask for several completions per request and pick the one you prefer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = ai21.Completion.execute(sm_endpoint=endpoint_name,\n",
    "                                   prompt=prompt,\n",
    "                                   maxTokens=100,\n",
    "                                   temperature=0.7,\n",
    "                                   stopSequences=['##'],\n",
    "                                   numResults=3)\n",
    "\n",
    "for comp in response['completions']:\n",
    "    print(comp['data']['text'].strip())\n",
    "    print(\"=============\")"
   ]
  },
  {
   "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": 37,
   "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": 38,
   "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
}