J2_Light_example_model_use.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Using Jurassic-2 Light on SageMaker through Model Packages\n",
    "\n",
    "This sample notebook shows you how to deploy **Jurassic-2 Light** 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-roz6zicyvi666\">https://aws.amazon.com/marketplace/pp/prodview-roz6zicyvi666</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",
    "Jurassic-2 Light is the quickest large language model (LLM) by AI21 Labs. Small but mighty, Jurassic-2 Light is ideal for simple language tasks that require maximum affordability and minimal latency in your private environment. Common use cases include keyword extraction, sentence classification, named entity recognition (NER), short-form copy generation, sentiment analysis, and more. It follows natural language instructions and supports non-English languages including Spanish, French, German, Portuguese, Italian and Dutch.\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. [Prompt with instructions](#C.-Prompt-with-instructions)\n",
    "   1. [Prompt with examples](#D.-Prompt-with-examples)\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=\"j2-light\", 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:#00178E'> 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:#00178E'>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#jurassic-2-light)."
   ]
  },
  {
   "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 = \"j2-light-internal\"\n",
    "\n",
    "content_type = \"application/json\"\n",
    "\n",
    "real_time_inference_instance_type = (\n",
    "    \"ml.g5.12xlarge\"    # Recommended instance\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": [
    "You can think of Jurassic-2 light as a smart auto-completion algorithm: give it some text as input and it will generate relevant text to naturally complete your input.\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": [
    "First, create an AI21 Sagemaker client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ai21 import AI21SageMakerClient\n",
    "\n",
    "client = AI21SageMakerClient(endpoint_name=endpoint_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Enter a simple prompt: \"To be, or\", and let the model complete it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = client.completion.create(\n",
    "    prompt=\"To be, or\",\n",
    "    max_tokens=4,\n",
    "    temperature=0,\n",
    "    num_results=1,\n",
    ")\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. Prompt with instructions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Why**? This model was specifically trained to follow natural language instructions. It is the most natural way to interact with large language models: simply tell the model what you want it to do, and it will follow.\n",
    "\n",
    "**When?** Drafting, seeking for inspiration, or when the format and guidelines are \"work in progress\".\n",
    "\n",
    "**How?** Just provide an instruction.\n",
    "\n",
    "For this notebook, we will apply the model to extract entities from news headlines. We will start with providing the model a simple instruction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "headline = \"Homer Simpson to sign executive order on January 2024 to increase the number of gun background checks in Springfield\"\n",
    "\n",
    "instruction = f\"\"\"Extract the entities from the following headline.\n",
    "Headline: {headline}\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = client.completion.create(\n",
    "    prompt=instruction,\n",
    "    max_tokens=10,\n",
    "    temperature=0,\n",
    "    num_results=1,\n",
    ")\n",
    "\n",
    "print(response.completions[0].data.text)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Adjust the parameters\n",
    "A useful parameter is the temperature. **You can increase creativity by tweaking the temperature.** With temperature 0, the model will always choose the most probable completion, so it will always be the same. Increasing the temperature will provide varying completions, where the completion may be different with every generation.\n",
    "*Note: in tasks such as NER, you should use low temperature, somewhere between 0-0.2*."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = client.completion.create(\n",
    "    prompt=instruction,\n",
    "    max_tokens=100,\n",
    "    temperature=0.2,\n",
    "    num_results=2, # this will make the model generate 2 optional completions\n",
    ")\n",
    "\n",
    "for comp in response.completions:\n",
    "    print(comp.data.text.strip())\n",
    "    print(\"=============\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Be specific in your prompt\n",
    "You may want to extract specific entities. You can ask is specifically from the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "specific_instruction = f\"\"\"Extract the entities from the following headline.\n",
    "Headline: {headline}\n",
    "Name, Date, Location:\"\"\"\n",
    "\n",
    "response = client.completion.create(\n",
    "    prompt=specific_instruction,\n",
    "    max_tokens=100,\n",
    "    temperature=0,\n",
    ")\n",
    "\n",
    "print(response.completions[0].data.text)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### D. Prompt with examples"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Why?** Examples are helpful in assisting the model to comprehend and generate responses that adhere to the intended format.\n",
    "\n",
    "**When?** Examples are particularly useful when there are stringent format constraints, a well-defined objective, and an overall structure to be maintained.\n",
    "\n",
    "**How?** To establish a pattern for the model to follow, present a few instances (“shots”) of input-output pairs in the prompt. This enables the model to mimic the pattern. Then, provide the input for a query example and allow the model to generate a suitable completion. This approach is commonly referred to as a \"*few-shot prompt*\"."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Create a few-shot prompt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "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."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, we collect some example data for the prompt prefix:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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": null,
   "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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_ner_prompt(headline):\n",
    "    \"\"\"\n",
    "    Create a few-shot prompt to extract named entities with Jurassic-2 Light 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 the t-shirt with the specific features from before:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "few_shot_prompt = create_ner_prompt(headline=headline)\n",
    "\n",
    "print(few_shot_prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = client.completion.create(\n",
    "    prompt=few_shot_prompt,\n",
    "    max_tokens=30,\n",
    "    temperature=0,\n",
    "    stop_sequences=['##'],\n",
    "    num_results=1,\n",
    ")\n",
    "\n",
    "print(response.completions[0].data.text)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you can see, the completions follow a similar pattern to the examples in the few-shot prompt."
   ]
  },
  {
   "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": 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
}