{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "79ebfce1-f5ac-4e39-83db-11416e310e8e",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Jupyter AI with the SageMaker endpoint\n",
    "\n",
    "This demo showcases the IPython magics Jupyter AI provides out-of-the-box for Amazon SageMaker.\n",
    "\n",
    "First, make sure that you've set your `AWS_ACCESS_KEY_ID` and `AWS_SECRET_ACCESS_KEY` environment variables either before starting JupyterLab or using the `%env` magic command within JupyterLab.\n",
    "\n",
    "Then, load the IPython extension:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "24f3f446-2b1d-4802-a47c-d298c06fc86e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%load_ext jupyter_ai"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9f2b0270-1c33-4918-b534-4ec104f90141",
   "metadata": {},
   "source": [
    "Jupyter AI supports language models hosted on SageMaker endpoints that use JSON APIs. Authenticate with AWS via the `boto3` SDK and have the credentials stored in the `default` profile.  Guidance on how to do this can be found in the [`boto3` documentation](https://boto3.amazonaws.com/v1/documentation/api/latest/guide/credentials.html).\n",
    "\n",
    "You will need to deploy a model in SageMaker, then provide it as your model name (as `sagemaker-endpoint:my-model-name`). See the [documentation on how to deploy a JumpStart model](https://docs.aws.amazon.com/sagemaker/latest/dg/jumpstart-deploy.html).\n",
    "\n",
    "All SageMaker endpoint requests require you to specify the `--region-name`, `--request-schema`, and `--response-path` options.\n",
    "\n",
    "The `--region-name` parameter is set to the [AWS region code](https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/using-regions-availability-zones.html), such as `us-east-1` or `eu-west-1`.\n",
    "\n",
    "The `--request-schema` parameter is the JSON object the endpoint expects as input, with the prompt being substituted into any value that matches the string literal `\"<prompt>\"`. For example, the request schema `{\"text_inputs\":\"<prompt>\"}` will submit a JSON object with the prompt stored under the `text_inputs` key.\n",
    "\n",
    "The `--response-path` option is a [JSONPath](https://goessner.net/articles/JsonPath/index.html) string that retrieves the language model's output from the endpoint's JSON response. For example, if your endpoint returns an object with the schema `{\"generated_texts\":[\"<output>\"]}`, its response path is `generated_texts.[0]`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "31f3e6e3-48cf-4e60-96d3-8b8e1dd34bec",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "What is the weather like in the US today?"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "execution_count": 2,
     "metadata": {
      "text/markdown": {
       "jupyter_ai": {
        "model_id": "jumpstart-dft-hf-text2text-flan-t5-xl",
        "provider_id": "sagemaker-endpoint"
       }
      }
     },
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%ai sagemaker-endpoint:jumpstart-dft-hf-text2text-flan-t5-xl --region-name=us-east-1 --request-schema={\"text_inputs\":\"<prompt>\"} --response-path=generated_texts.[0]\n",
    "Generate a question about weather today."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ad8c62e-b0a5-4091-94e3-4067ed8d6c4a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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.10.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}