Jamba_1_5_Large_example_model_use.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Using Jamba 1.5 Large on SageMaker through Model Packages\n",
    "\n",
    "This sample notebook shows you how to deploy **Jamba 1.5 Large** 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-bf26px7gdisek\">https://aws.amazon.com/marketplace/pp/prodview-bf26px7gdisek</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",
    "Jamba 1.5 Large is the first of its kind hybrid Mamba-Transformer architecture at a production grade level offering unmatched efficiency. With an unprecedented context window length (256K) in a smaller model, it offers high quality output for tasks needing large input context & low latency, at a competitive price point.\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.35.68** or higher.\n",
    "\n",
    "## Contents:\n",
    "1. [Select a model package](#1.-Select-a-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](#Create-an-endpoint)\n",
    "   1. [Interact with the model](#Interact-with-the-model)\n",
    "1. [Clean-up](#3.-Clean-up)\n",
    "   1. [Delete the endpoint](#Delete-the-endpoint)\n",
    "   1. [Delete the endpoint configuration](#Delete-the-endpoint-configuration)\n",
    "   1. [Delete the model](#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": [
    "## Installations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install boto3 requests --upgrade"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import boto3\n",
    "import requests"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Check the version of boto3 - must be v1.35.68 or higher"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "boto3.__version__"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Select a model package\n",
    "Confirm that you received this notebook from the model catalog in SageMaker AWS Management Console."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Get latest model package ARN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "region = boto3.Session().region_name\n",
    "model_name = \"jamba-1-5-large\"\n",
    "version = \"latest\"\n",
    "\n",
    "# Get the updated ARN\n",
    "model_package_arn_payload = {\"modelName\": model_name, \"version\": version, \"region\": region}\n",
    "get_model_package_arn_url = \"https://api.ai21.com/studio/v1/jumpstart/get_model_version_arn\"\n",
    "model_package_arn_response = requests.post(get_model_package_arn_url, json=model_package_arn_payload)\n",
    "model_package_arn = model_package_arn_response.json()[\"arn\"]"
   ]
  },
  {
   "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": [
    "# Set SageMaker execution role ARN\n",
    "sagemaker_execution_role_arn = \"<<INSERT-YOUR-SAGEMAKER-EXECUTION-ROLE>>\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sagemaker_base_name = \"jamba-1-5-large\"\n",
    "sagemaker_model_name = f\"{sagemaker_base_name}-model\"\n",
    "sagemaker_endpoint_config_name = f\"{sagemaker_base_name}-endpoint-config\"\n",
    "sagemaker_endpoint_name = f\"{sagemaker_base_name}-endpoint\"\n",
    "\n",
    "content_type = \"application/json\"\n",
    "\n",
    "real_time_inference_instance_type = (\n",
    "    \"ml.p4de.24xlarge\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create an endpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sm_client = boto3.client('sagemaker', region_name=region)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create model\n",
    "sm_client.create_model(\n",
    "    ModelName=sagemaker_model_name,\n",
    "    ExecutionRoleArn=sagemaker_execution_role_arn,\n",
    "    PrimaryContainer={\n",
    "        'ModelPackageName': model_package_arn,\n",
    "    },\n",
    "    EnableNetworkIsolation=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create endpoint config\n",
    "endpoint_config = sm_client.create_endpoint_config(\n",
    "    EndpointConfigName=sagemaker_endpoint_config_name,\n",
    "    ProductionVariants=[{\n",
    "        'InstanceType': real_time_inference_instance_type,\n",
    "        'InitialInstanceCount': 1,\n",
    "        'ModelName': sagemaker_model_name,\n",
    "        'VariantName': 'variant1',\n",
    "        'ModelDataDownloadTimeoutInSeconds': 3600,\n",
    "        'InferenceAmiVersion': 'al2-ami-sagemaker-inference-gpu-2',\n",
    "    }]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create endpoint\n",
    "endpoint = sm_client.create_endpoint(\n",
    "    EndpointName=sagemaker_endpoint_name,\n",
    "    EndpointConfigName=sagemaker_endpoint_config_name,\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": [
    "### Interact with the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sm_runtime = boto3.client(\"sagemaker-runtime\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For advanced usage, you can check out the [AI21 Jamba 1.5 documentation](https://docs.ai21.com/reference/jamba-15-api-ref)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Non-streaming"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "messages = [{\"role\": \"user\", \"content\": \"Tell me a joke about pokemons\"}]\n",
    "input_json = json.dumps({\"messages\": messages})\n",
    "\n",
    "response = sm_runtime.invoke_endpoint(\n",
    "    EndpointName=sagemaker_endpoint_name,\n",
    "    Accept=\"application/json\",\n",
    "    Body=input_json,\n",
    ")\n",
    "\n",
    "print(json.load(response[\"Body\"]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Streaming"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "messages = [{\"role\": \"user\", \"content\": \"How are you?\"}]\n",
    "input_json = json.dumps({\"messages\": messages, \"max_tokens\": 10})\n",
    "\n",
    "response = sm_runtime.invoke_endpoint_with_response_stream(\n",
    "    EndpointName=sagemaker_endpoint_name,\n",
    "    Accept=\"application/json\",\n",
    "    Body=input_json,\n",
    ")\n",
    "\n",
    "for event in response['Body']:\n",
    "    print(event['PayloadPart']['Bytes'].decode('utf-8'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Clean-up"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Delete the endpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sm_client.delete_endpoint(EndpointName=sagemaker_endpoint_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Delete the endpoint configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sm_client.delete_endpoint_config(EndpointConfigName=sagemaker_endpoint_config_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Delete the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sm_client.delete_model(ModelName=sagemaker_model_name)"
   ]
  }
 ],
 "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
}