 archived/question_answering_jumpstart_knn.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Retrieval-Augmented Generation: Question Answering based on Custom Dataset\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "This notebook's CI test result for us-west-2 is as follows. CI test results in other regions can be found at the end of the notebook. \n",
    "\n",
    "![This us-west-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/us-west-2/introduction_to_amazon_algorithms|jumpstart-foundation-models|question_answerIng_retrieval_augmented_generation_jumpstart|question_answering_jumpstart_knn.ipynb)\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "Many use cases such as building a chatbot require text (text2text) generation models like **[BloomZ 7B1](https://huggingface.co/bigscience/bloomz-7b1)**, **[Flan T5 XXL](https://huggingface.co/google/flan-t5-xxl)**, and **[Flan T5 UL2](https://huggingface.co/google/flan-ul2)** to respond to user questions with insightful answers. The **BloomZ 7B1**, **Flan T5 XXL**, and **Flan T5 UL2** models have picked up a lot of general knowledge in training, but we often need to ingest and use a large library of more specific information.\n",
    "\n",
    "In this notebook we will demonstrate how to use **BloomZ 7B1**, **Flan T5 XXL**, and **Flan T5 UL2** to answer questions using a library of documents as a reference, by using document embeddings and retrieval. The embeddings are generated from **GPT-J-6B** embedding model. \n",
    "\n",
    "**This notebook serves a template such that you can easily replace the example dataset by your own to build a custom question and asnwering application.**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 1. Deploy large language model (LLM) in SageMaker JumpStart\n",
    "\n",
    "To better illustrate the idea, let's first deploy all the models that are required to perform the demo. You can choose either deploying all three Flan T5 XL, BloomZ 7B1, and Flan UL2 models as the large language model (LLM) to compare their model performances, or select **subset** of the models based on your preference. To do that, you need modify the `_MODEL_CONFIG_` python dictionary defined as below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "!pip install --upgrade sagemaker --quiet\n",
    "!pip install ipywidgets==7.0.0 --quiet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "import sagemaker, boto3, json\n",
    "from sagemaker.session import Session\n",
    "from sagemaker.model import Model\n",
    "from sagemaker import image_uris, model_uris, script_uris, hyperparameters\n",
    "from sagemaker.predictor import Predictor\n",
    "from sagemaker.utils import name_from_base\n",
    "\n",
    "sagemaker_session = Session()\n",
    "aws_role = sagemaker_session.get_caller_identity_arn()\n",
    "aws_region = boto3.Session().region_name\n",
    "sess = sagemaker.Session()\n",
    "model_version = \"1.*\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def query_endpoint_with_json_payload(encoded_json, endpoint_name, content_type=\"application/json\"):\n",
    "    client = boto3.client(\"runtime.sagemaker\")\n",
    "    response = client.invoke_endpoint(\n",
    "        EndpointName=endpoint_name, ContentType=content_type, Body=encoded_json\n",
    "    )\n",
    "    return response\n",
    "\n",
    "\n",
    "def parse_response_model_flan_t5(query_response):\n",
    "    model_predictions = json.loads(query_response[\"Body\"].read())\n",
    "    generated_text = model_predictions[\"generated_texts\"]\n",
    "    return generated_text\n",
    "\n",
    "\n",
    "def parse_response_multiple_texts_bloomz(query_response):\n",
    "    generated_text = []\n",
    "    model_predictions = json.loads(query_response[\"Body\"].read())\n",
    "    for x in model_predictions[0]:\n",
    "        generated_text.append(x[\"generated_text\"])\n",
    "    return generated_text"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Please uncomment the entries as below if you want to deploy multiple LLM models to compare their performance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_MODEL_CONFIG_ = {\n",
    "    \"huggingface-text2text-flan-t5-xxl\": {\n",
    "        \"instance type\": \"ml.g5.12xlarge\",\n",
    "        \"env\": {\"SAGEMAKER_MODEL_SERVER_WORKERS\": \"1\", \"TS_DEFAULT_WORKERS_PER_MODEL\": \"1\"},\n",
    "        \"parse_function\": parse_response_model_flan_t5,\n",
    "        \"prompt\": \"\"\"Answer based on context:\\n\\n{context}\\n\\n{question}\"\"\",\n",
    "    },\n",
    "    # \"huggingface-textgeneration1-bloomz-7b1-fp16\": {\n",
    "    #     \"instance type\": \"ml.g5.12xlarge\",\n",
    "    #     \"env\": {},\n",
    "    #     \"parse_function\": parse_response_multiple_texts_bloomz,\n",
    "    #     \"prompt\": \"\"\"question: \\\"{question}\"\\\\n\\nContext: \\\"{context}\"\\\\n\\nAnswer:\"\"\",\n",
    "    # },\n",
    "    # \"huggingface-text2text-flan-ul2-bf16\": {\n",
    "    #     \"instance type\": \"ml.g5.24xlarge\",\n",
    "    #     \"env\": {\n",
    "    #         \"SAGEMAKER_MODEL_SERVER_WORKERS\": \"1\",\n",
    "    #         \"TS_DEFAULT_WORKERS_PER_MODEL\": \"1\"\n",
    "    #     },\n",
    "    #     \"parse_function\": parse_response_model_flan_t5,\n",
    "    #     \"prompt\": \"\"\"Answer based on context:\\n\\n{context}\\n\\n{question}\"\"\",\n",
    "    # }\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "newline, bold, unbold = \"\\n\", \"\\033[1m\", \"\\033[0m\"\n",
    "\n",
    "for model_id in _MODEL_CONFIG_:\n",
    "    endpoint_name = name_from_base(f\"jumpstart-example-ragknn-{model_id}\")\n",
    "    inference_instance_type = _MODEL_CONFIG_[model_id][\"instance type\"]\n",
    "\n",
    "    # Retrieve the inference container uri. This is the base HuggingFace container image for the default model above.\n",
    "    deploy_image_uri = image_uris.retrieve(\n",
    "        region=None,\n",
    "        framework=None,  # automatically inferred from model_id\n",
    "        image_scope=\"inference\",\n",
    "        model_id=model_id,\n",
    "        model_version=model_version,\n",
    "        instance_type=inference_instance_type,\n",
    "    )\n",
    "    # Retrieve the model uri.\n",
    "    model_uri = model_uris.retrieve(\n",
    "        model_id=model_id, model_version=model_version, model_scope=\"inference\"\n",
    "    )\n",
    "    model_inference = Model(\n",
    "        image_uri=deploy_image_uri,\n",
    "        model_data=model_uri,\n",
    "        role=aws_role,\n",
    "        predictor_cls=Predictor,\n",
    "        name=endpoint_name,\n",
    "        env=_MODEL_CONFIG_[model_id][\"env\"],\n",
    "    )\n",
    "    model_predictor_inference = model_inference.deploy(\n",
    "        initial_instance_count=1,\n",
    "        instance_type=inference_instance_type,\n",
    "        predictor_cls=Predictor,\n",
    "        endpoint_name=endpoint_name,\n",
    "    )\n",
    "    print(f\"{bold}Model {model_id} has been deployed successfully.{unbold}{newline}\")\n",
    "    _MODEL_CONFIG_[model_id][\"endpoint_name\"] = endpoint_name"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 2. Ask a question to LLM without providing the context\n",
    "\n",
    "To better illustrate why we need retrieval-augmented generation (RAG) based approach to solve the question and anwering problem. Let's directly ask the model a question and see how they respond."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "question = \"Which instances can I use with Managed Spot Training in SageMaker?\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "payload = {\n",
    "    \"text_inputs\": question,\n",
    "    \"max_length\": 100,\n",
    "    \"num_return_sequences\": 1,\n",
    "    \"top_k\": 50,\n",
    "    \"top_p\": 0.95,\n",
    "    \"do_sample\": True,\n",
    "}\n",
    "\n",
    "\n",
    "for model_id in _MODEL_CONFIG_:\n",
    "    endpoint_name = _MODEL_CONFIG_[model_id][\"endpoint_name\"]\n",
    "    query_response = query_endpoint_with_json_payload(\n",
    "        json.dumps(payload).encode(\"utf-8\"), endpoint_name=endpoint_name\n",
    "    )\n",
    "    generated_texts = _MODEL_CONFIG_[model_id][\"parse_function\"](query_response)\n",
    "    print(f\"For model: {model_id}, the generated output is: {generated_texts[0]}\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can see the generated answer is wrong or doesn't make much sense. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 3. Improve the answer to the same question using **prompt engineering** with insightful context\n",
    "\n",
    "\n",
    "To better answer the question well, we provide extra contextual information, combine it with a prompt, and send it to model together with the question. Below is an example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "context = \"\"\"Managed Spot Training can be used with all instances supported in Amazon SageMaker. Managed Spot Training is supported in all AWS Regions where Amazon SageMaker is currently available.\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "parameters = {\n",
    "    \"max_length\": 200,\n",
    "    \"num_return_sequences\": 1,\n",
    "    \"top_k\": 250,\n",
    "    \"top_p\": 0.95,\n",
    "    \"do_sample\": False,\n",
    "    \"temperature\": 1,\n",
    "}\n",
    "\n",
    "for model_id in _MODEL_CONFIG_:\n",
    "    endpoint_name = _MODEL_CONFIG_[model_id][\"endpoint_name\"]\n",
    "\n",
    "    prompt = _MODEL_CONFIG_[model_id][\"prompt\"]\n",
    "\n",
    "    text_input = prompt.replace(\"{context}\", context)\n",
    "    text_input = text_input.replace(\"{question}\", question)\n",
    "    payload = {\"text_inputs\": text_input, **parameters}\n",
    "\n",
    "    query_response = query_endpoint_with_json_payload(\n",
    "        json.dumps(payload).encode(\"utf-8\"), endpoint_name=endpoint_name\n",
    "    )\n",
    "    generated_texts = _MODEL_CONFIG_[model_id][\"parse_function\"](query_response)\n",
    "    print(\n",
    "        f\"{bold}For model: {model_id}, the generated output is: {generated_texts[0]}{unbold}{newline}\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 4. Use RAG based approach to identify the correct documents, and use them along with prompt and question to query LLM\n",
    "\n",
    "\n",
    "We plan to use document embeddings to fetch the most relevant documents in our document knowledge library and combine them with the prompt that we provide to LLM.\n",
    "\n",
    "To achieve that, we will do following.\n",
    "\n",
    "* **Generate embedings for each of document in the knowledge library with the GPT-J-6B embedding model.**\n",
    "* **Identify top K most relevant documents based on user query.**\n",
    "    * **For a query of your interest, generate the embedding of the query using the same embedding model.**\n",
    "    * **Search the indexes of top K most relevant documents in the embedding space using the SageMaker KNN algorithm.**\n",
    "    * **Use the indexes to retrieve the corresponded documents.**\n",
    "* **Combine the retrieved documents with prompt and question and send them into LLM.**\n",
    "\n",
    "\n",
    "\n",
    "Note: The retrieved document/text should be large enough to contain enough information to answer a question; but small enough to fit into the LLM prompt -- maximum sequence length of 1024 tokens. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.1 Deploying the model endpoint for GPT-J-6B embedding model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_id, model_version = \"huggingface-textembedding-gpt-j-6b\", \"1.*\"\n",
    "\n",
    "endpoint_name_embed = name_from_base(f\"jumpstart-example-ragknn-{model_id}\")\n",
    "\n",
    "inference_instance_type = \"ml.g5.24xlarge\"\n",
    "\n",
    "# Retrieve the inference container uri. This is the base HuggingFace container image for the default model above.\n",
    "deploy_image_uri = image_uris.retrieve(\n",
    "    region=None,\n",
    "    framework=None,  # automatically inferred from model_id\n",
    "    image_scope=\"inference\",\n",
    "    model_id=model_id,\n",
    "    model_version=model_version,\n",
    "    instance_type=inference_instance_type,\n",
    ")\n",
    "# Retrieve the model uri.\n",
    "model_uri = model_uris.retrieve(\n",
    "    model_id=model_id, model_version=model_version, model_scope=\"inference\"\n",
    ")\n",
    "model_inference = Model(\n",
    "    image_uri=deploy_image_uri,\n",
    "    model_data=model_uri,\n",
    "    role=aws_role,\n",
    "    name=endpoint_name_embed,\n",
    "    env={\"TS_DEFAULT_WORKERS_PER_MODEL\": \"1\"},\n",
    ")\n",
    "\n",
    "model_predictor_embed = model_inference.deploy(\n",
    "    initial_instance_count=1,\n",
    "    instance_type=inference_instance_type,\n",
    "    endpoint_name=endpoint_name_embed,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "\n",
    "\n",
    "def query_endpoint_with_json_payload(encoded_json, endpoint_name, content_type=\"application/json\"):\n",
    "    client = boto3.client(\"runtime.sagemaker\")\n",
    "    response = client.invoke_endpoint(\n",
    "        EndpointName=endpoint_name, ContentType=content_type, Body=encoded_json\n",
    "    )\n",
    "    return response\n",
    "\n",
    "\n",
    "def parse_response_multiple_texts(query_response):\n",
    "    model_predictions = json.loads(query_response[\"Body\"].read())\n",
    "    embeddings = model_predictions[\"embedding\"]\n",
    "    return embeddings\n",
    "\n",
    "\n",
    "def build_embed_table(df_knowledge, endpoint_name_embed, col_name_4_embed, batch_size=10):\n",
    "    res_embed = []\n",
    "    N = df_knowledge.shape[0]\n",
    "    for idx in tqdm(range(0, N, batch_size)):\n",
    "        content = df_knowledge.loc[idx : (idx + batch_size - 1)][\n",
    "            col_name_4_embed\n",
    "        ].tolist()  ## minus -1 as pandas loc slicing is end-inclusive\n",
    "        payload = {\"text_inputs\": content}\n",
    "        query_response = query_endpoint_with_json_payload(\n",
    "            json.dumps(payload).encode(\"utf-8\"), endpoint_name_embed\n",
    "        )\n",
    "        generated_embed = parse_response_multiple_texts(query_response)\n",
    "        res_embed.extend(generated_embed)\n",
    "    res_embed_df = pd.DataFrame(res_embed)\n",
    "    return res_embed_df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.2. Generate embedings for each of document in the knowledge library with the GPT-J-6B embedding model.\n",
    "\n",
    "For the purpose of the demo we will use [Amazon SageMaker FAQs](https://aws.amazon.com/sagemaker/faqs/) as knowledge library. The data are formatted in a CSV file with two columns Question and Answer. We use **only** the Answer column as the documents of knowledge library, from which relevant documents are retrieved based on a query. \n",
    "\n",
    "**Each row in the CSV format dataset corresponds to a textual document. \n",
    "We will iterate each document to get its embedding vector via the GPT-J-6B embedding models. \n",
    "For your purpose, you can replace the example dataset of your own to build a custom question and answering application.**\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, we download the dataset from our S3 bucket to the local."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s3_path = f\"s3://jumpstart-cache-prod-us-east-2/training-datasets/Amazon_SageMaker_FAQs/Amazon_SageMaker_FAQs.csv\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Downloading the Database\n",
    "!aws s3 cp $s3_path Amazon_SageMaker_FAQs.csv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "df_knowledge = pd.read_csv(\"Amazon_SageMaker_FAQs.csv\", header=None, names=[\"Question\", \"Answer\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_knowledge.head(5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Drop the `Question` column since it is not used in this notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_knowledge.drop([\"Question\"], axis=1, inplace=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_knowledge.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_knowledge_embed = build_embed_table(\n",
    "    df_knowledge, endpoint_name_embed, col_name_4_embed=\"Answer\", batch_size=20\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_knowledge_embed.head(5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert df_knowledge_embed.shape[0] == df_knowledge.shape[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Save the embedding data for further usage."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_knowledge_embed.to_csv(\"Amazon_SageMaker_FAQs_embedding.csv\", header=None, index=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.3. Index the embedding knowledge library using [SageMaker KNN algorithm](https://docs.aws.amazon.com/sagemaker/latest/dg/k-nearest-neighbors.html)\n",
    "\n",
    "The SageMaker KNN will conduct following.\n",
    "\n",
    "1. Start a training job to index the embedding knowledge data. The underlying algorithm used to index the data is [Faiss](https://github.com/facebookresearch/faiss).\n",
    "2. Start an endpoint to take the embedding of the query as input and return the top K nearest indexes of the documents.\n",
    "\n",
    "**Note.** For the KNN training job, the features are N by P matrix, where N is the number of documetns in the knowledge library, P is the embedding dimension, and each row corresponds to an embedding of a document. The labels are ordinal integers starting from 0. During inference, given an embedding of query, the labels of the top K nearest documents with respect to the query are used as indexes to retrieve the corresponded textual documents.\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We first upload the prepared dataset to the S3 bucket."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import os\n",
    "import io\n",
    "import sagemaker.amazon.common as smac\n",
    "\n",
    "\n",
    "train_features = np.array(df_knowledge_embed)\n",
    "\n",
    "# Providing each answer embedding label\n",
    "train_labels = np.array([i for i in range(len(train_features))])\n",
    "\n",
    "print(\"train_features shape = \", train_features.shape)\n",
    "print(\"train_labels shape = \", train_labels.shape)\n",
    "\n",
    "buf = io.BytesIO()\n",
    "smac.write_numpy_to_dense_tensor(buf, train_features, train_labels)\n",
    "buf.seek(0)\n",
    "\n",
    "\n",
    "bucket = sess.default_bucket()  # modify to your bucket name\n",
    "prefix = \"RAGDatabase\"\n",
    "key = \"Amazon-SageMaker-RAG\"\n",
    "\n",
    "boto3.resource(\"s3\").Bucket(bucket).Object(os.path.join(prefix, \"train\", key)).upload_fileobj(buf)\n",
    "s3_train_data = f\"s3://{bucket}/{prefix}/train/{key}\"\n",
    "print(f\"uploaded training data location: {s3_train_data}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We want to retrieve the top 5 most relevant documents. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "TOP_K = 5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.amazon.amazon_estimator import get_image_uri\n",
    "\n",
    "\n",
    "def trained_estimator_from_hyperparams(s3_train_data, hyperparams, output_path):\n",
    "    \"\"\"\n",
    "    Create an Estimator from the given hyperparams, fit to training data,\n",
    "    and return a deployed predictor\n",
    "\n",
    "    \"\"\"\n",
    "    # set up the estimator\n",
    "    knn = sagemaker.estimator.Estimator(\n",
    "        get_image_uri(boto3.Session().region_name, \"knn\"),\n",
    "        aws_role,\n",
    "        instance_count=1,\n",
    "        instance_type=\"ml.m5.2xlarge\",\n",
    "        output_path=output_path,\n",
    "        sagemaker_session=sess,\n",
    "    )\n",
    "    knn.set_hyperparameters(**hyperparams)\n",
    "\n",
    "    # train a model. fit_input contains the locations of the train data\n",
    "    fit_input = {\"train\": s3_train_data}\n",
    "    knn.fit(fit_input)\n",
    "    return knn\n",
    "\n",
    "\n",
    "hyperparams = {\n",
    "    \"feature_dim\": train_features.shape[1],\n",
    "    \"k\": TOP_K,\n",
    "    \"sample_size\": train_features.shape[0],\n",
    "    \"predictor_type\": \"classifier\",\n",
    "}\n",
    "output_path = f\"s3://{bucket}/{prefix}/default_example/output\"\n",
    "knn_estimator = trained_estimator_from_hyperparams(s3_train_data, hyperparams, output_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Deploy the KNN endpoint for retrieving indexes of top K most relevant docuemnts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.serializers import CSVSerializer\n",
    "from sagemaker.deserializers import JSONDeserializer\n",
    "\n",
    "\n",
    "def predictor_from_estimator(knn_estimator, instance_type, endpoint_name=None):\n",
    "    knn_predictor = knn_estimator.deploy(\n",
    "        initial_instance_count=1, instance_type=instance_type, endpoint_name=endpoint_name\n",
    "    )\n",
    "    knn_predictor.serializer = CSVSerializer()\n",
    "    knn_predictor.deserializer = JSONDeserializer()\n",
    "    return knn_predictor\n",
    "\n",
    "\n",
    "instance_type = \"ml.m4.xlarge\"\n",
    "endpoint_name = name_from_base(f\"jumpstart-example-ragknn\")\n",
    "\n",
    "knn_predictor = predictor_from_estimator(knn_estimator, instance_type, endpoint_name=endpoint_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.4 Retrieve the most relevant documents\n",
    "\n",
    "Given the embedding of a query, we will query the endpoint to get the indexes of top K most relevant documents and use the indexes to retrieve the corresponded textual documents.\n",
    "\n",
    "Next, the textual documents are concatenated with maximum length of `MAX_SECTION_LEN`. This is to make sure the context we send into the prompt contains a good enough amount of information all the while not exceeding model's capacity."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "MAX_SECTION_LEN = 500\n",
    "SEPARATOR = \"\\n* \"\n",
    "\n",
    "\n",
    "def construct_context(context_predictions_arr, df_knowledge) -> str:\n",
    "    chosen_sections = []\n",
    "    chosen_sections_len = 0\n",
    "\n",
    "    for index in context_predictions_arr:\n",
    "        # Add contexts until we run out of space.\n",
    "        document_section = df_knowledge.loc[index]\n",
    "        chosen_sections_len += len(document_section) + 2\n",
    "        if chosen_sections_len > MAX_SECTION_LEN:\n",
    "            break\n",
    "\n",
    "        chosen_sections.append(SEPARATOR + document_section.replace(\"\\n\", \" \"))\n",
    "    concatenated_doc = \"\".join(chosen_sections)\n",
    "    print(\n",
    "        f\"With maximum sequence length {MAX_SECTION_LEN}, selected top {len(chosen_sections)} document sections: {concatenated_doc}\"\n",
    "    )\n",
    "\n",
    "    return concatenated_doc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "start = time.time()\n",
    "\n",
    "query_response = query_endpoint_with_json_payload(\n",
    "    question, endpoint_name_embed, content_type=\"application/x-text\"\n",
    ")\n",
    "question_embedding = parse_response_multiple_texts(query_response)\n",
    "\n",
    "# Getting the most relevant context using KNN\n",
    "context_predictions_arr = knn_predictor.predict(\n",
    "    np.array(question_embedding),\n",
    "    initial_args={\"ContentType\": \"text/csv\", \"Accept\": \"application/json; verbose=true\"},\n",
    ")[\"predictions\"][0][\"labels\"]\n",
    "\n",
    "context_embed_retrieve = construct_context(context_predictions_arr, df_knowledge[\"Answer\"])\n",
    "\n",
    "print(\n",
    "    f\"{newline}{bold}Elastic time for computing the embedding of a query and retrieved the top K most relevant docuemnts: {time.time() - start} seconds.{unbold}{newline}\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.5 Combine the retrieved documents, prompt, and question to query the LLM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for model_id in _MODEL_CONFIG_:\n",
    "    endpoint_name = _MODEL_CONFIG_[model_id][\"endpoint_name\"]\n",
    "\n",
    "    prompt = _MODEL_CONFIG_[model_id][\"prompt\"]\n",
    "\n",
    "    text_input = prompt.replace(\"{context}\", context_embed_retrieve)\n",
    "    text_input = text_input.replace(\"{question}\", question)\n",
    "\n",
    "    payload = {\"text_inputs\": text_input, **parameters}\n",
    "\n",
    "    query_response = query_endpoint_with_json_payload(\n",
    "        json.dumps(payload).encode(\"utf-8\"), endpoint_name=endpoint_name\n",
    "    )\n",
    "    generated_texts = _MODEL_CONFIG_[model_id][\"parse_function\"](query_response)\n",
    "    print(f\"For model: {model_id}, the generated output is: {generated_texts[0]}\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Notebook CI Test Results\n",
    "\n",
    "This notebook was tested in multiple regions. The test results are as follows, except for us-west-2 which is shown at the top of the notebook.\n",
    "\n",
    "![This us-east-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/us-east-1/introduction_to_amazon_algorithms|jumpstart-foundation-models|question_answerIng_retrieval_augmented_generation_jumpstart|question_answering_jumpstart_knn.ipynb)\n",
    "\n",
    "![This us-east-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/us-east-2/introduction_to_amazon_algorithms|jumpstart-foundation-models|question_answerIng_retrieval_augmented_generation_jumpstart|question_answering_jumpstart_knn.ipynb)\n",
    "\n",
    "![This us-west-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/us-west-1/introduction_to_amazon_algorithms|jumpstart-foundation-models|question_answerIng_retrieval_augmented_generation_jumpstart|question_answering_jumpstart_knn.ipynb)\n",
    "\n",
    "![This ca-central-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ca-central-1/introduction_to_amazon_algorithms|jumpstart-foundation-models|question_answerIng_retrieval_augmented_generation_jumpstart|question_answering_jumpstart_knn.ipynb)\n",
    "\n",
    "![This sa-east-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/sa-east-1/introduction_to_amazon_algorithms|jumpstart-foundation-models|question_answerIng_retrieval_augmented_generation_jumpstart|question_answering_jumpstart_knn.ipynb)\n",
    "\n",
    "![This eu-west-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/eu-west-1/introduction_to_amazon_algorithms|jumpstart-foundation-models|question_answerIng_retrieval_augmented_generation_jumpstart|question_answering_jumpstart_knn.ipynb)\n",
    "\n",
    "![This eu-west-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/eu-west-2/introduction_to_amazon_algorithms|jumpstart-foundation-models|question_answerIng_retrieval_augmented_generation_jumpstart|question_answering_jumpstart_knn.ipynb)\n",
    "\n",
    "![This eu-west-3 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/eu-west-3/introduction_to_amazon_algorithms|jumpstart-foundation-models|question_answerIng_retrieval_augmented_generation_jumpstart|question_answering_jumpstart_knn.ipynb)\n",
    "\n",
    "![This eu-central-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/eu-central-1/introduction_to_amazon_algorithms|jumpstart-foundation-models|question_answerIng_retrieval_augmented_generation_jumpstart|question_answering_jumpstart_knn.ipynb)\n",
    "\n",
    "![This eu-north-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/eu-north-1/introduction_to_amazon_algorithms|jumpstart-foundation-models|question_answerIng_retrieval_augmented_generation_jumpstart|question_answering_jumpstart_knn.ipynb)\n",
    "\n",
    "![This ap-southeast-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ap-southeast-1/introduction_to_amazon_algorithms|jumpstart-foundation-models|question_answerIng_retrieval_augmented_generation_jumpstart|question_answering_jumpstart_knn.ipynb)\n",
    "\n",
    "![This ap-southeast-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ap-southeast-2/introduction_to_amazon_algorithms|jumpstart-foundation-models|question_answerIng_retrieval_augmented_generation_jumpstart|question_answering_jumpstart_knn.ipynb)\n",
    "\n",
    "![This ap-northeast-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ap-northeast-1/introduction_to_amazon_algorithms|jumpstart-foundation-models|question_answerIng_retrieval_augmented_generation_jumpstart|question_answering_jumpstart_knn.ipynb)\n",
    "\n",
    "![This ap-northeast-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ap-northeast-2/introduction_to_amazon_algorithms|jumpstart-foundation-models|question_answerIng_retrieval_augmented_generation_jumpstart|question_answering_jumpstart_knn.ipynb)\n",
    "\n",
    "![This ap-south-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ap-south-1/introduction_to_amazon_algorithms|jumpstart-foundation-models|question_answerIng_retrieval_augmented_generation_jumpstart|question_answering_jumpstart_knn.ipynb)\n"
   ]
  }
 ],
 "metadata": {
  "availableInstances": [
   {
    "_defaultOrder": 0,
    "_isFastLaunch": true,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 4,
    "name": "ml.t3.medium",
    "vcpuNum": 2
   },
   {
    "_defaultOrder": 1,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 8,
    "name": "ml.t3.large",
    "vcpuNum": 2
   },
   {
    "_defaultOrder": 2,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.t3.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 3,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.t3.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 4,
    "_isFastLaunch": true,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 8,
    "name": "ml.m5.large",
    "vcpuNum": 2
   },
   {
    "_defaultOrder": 5,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.m5.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 6,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.m5.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 7,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 64,
    "name": "ml.m5.4xlarge",
    "vcpuNum": 16
   },
   {
    "_defaultOrder": 8,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 128,
    "name": "ml.m5.8xlarge",
    "vcpuNum": 32
   },
   {
    "_defaultOrder": 9,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 192,
    "name": "ml.m5.12xlarge",
    "vcpuNum": 48
   },
   {
    "_defaultOrder": 10,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 256,
    "name": "ml.m5.16xlarge",
    "vcpuNum": 64
   },
   {
    "_defaultOrder": 11,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 384,
    "name": "ml.m5.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 12,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 8,
    "name": "ml.m5d.large",
    "vcpuNum": 2
   },
   {
    "_defaultOrder": 13,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.m5d.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 14,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.m5d.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 15,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 64,
    "name": "ml.m5d.4xlarge",
    "vcpuNum": 16
   },
   {
    "_defaultOrder": 16,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 128,
    "name": "ml.m5d.8xlarge",
    "vcpuNum": 32
   },
   {
    "_defaultOrder": 17,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 192,
    "name": "ml.m5d.12xlarge",
    "vcpuNum": 48
   },
   {
    "_defaultOrder": 18,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 256,
    "name": "ml.m5d.16xlarge",
    "vcpuNum": 64
   },
   {
    "_defaultOrder": 19,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 384,
    "name": "ml.m5d.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 20,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": true,
    "memoryGiB": 0,
    "name": "ml.geospatial.interactive",
    "supportedImageNames": [
     "sagemaker-geospatial-v1-0"
    ],
    "vcpuNum": 0
   },
   {
    "_defaultOrder": 21,
    "_isFastLaunch": true,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 4,
    "name": "ml.c5.large",
    "vcpuNum": 2
   },
   {
    "_defaultOrder": 22,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 8,
    "name": "ml.c5.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 23,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.c5.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 24,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.c5.4xlarge",
    "vcpuNum": 16
   },
   {
    "_defaultOrder": 25,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 72,
    "name": "ml.c5.9xlarge",
    "vcpuNum": 36
   },
   {
    "_defaultOrder": 26,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 96,
    "name": "ml.c5.12xlarge",
    "vcpuNum": 48
   },
   {
    "_defaultOrder": 27,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 144,
    "name": "ml.c5.18xlarge",
    "vcpuNum": 72
   },
   {
    "_defaultOrder": 28,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 192,
    "name": "ml.c5.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 29,
    "_isFastLaunch": true,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.g4dn.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 30,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.g4dn.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 31,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 64,
    "name": "ml.g4dn.4xlarge",
    "vcpuNum": 16
   },
   {
    "_defaultOrder": 32,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 128,
    "name": "ml.g4dn.8xlarge",
    "vcpuNum": 32
   },
   {
    "_defaultOrder": 33,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 4,
    "hideHardwareSpecs": false,
    "memoryGiB": 192,
    "name": "ml.g4dn.12xlarge",
    "vcpuNum": 48
   },
   {
    "_defaultOrder": 34,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 256,
    "name": "ml.g4dn.16xlarge",
    "vcpuNum": 64
   },
   {
    "_defaultOrder": 35,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 61,
    "name": "ml.p3.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 36,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 4,
    "hideHardwareSpecs": false,
    "memoryGiB": 244,
    "name": "ml.p3.8xlarge",
    "vcpuNum": 32
   },
   {
    "_defaultOrder": 37,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 8,
    "hideHardwareSpecs": false,
    "memoryGiB": 488,
    "name": "ml.p3.16xlarge",
    "vcpuNum": 64
   },
   {
    "_defaultOrder": 38,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 8,
    "hideHardwareSpecs": false,
    "memoryGiB": 768,
    "name": "ml.p3dn.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 39,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.r5.large",
    "vcpuNum": 2
   },
   {
    "_defaultOrder": 40,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.r5.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 41,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 64,
    "name": "ml.r5.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 42,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 128,
    "name": "ml.r5.4xlarge",
    "vcpuNum": 16
   },
   {
    "_defaultOrder": 43,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 256,
    "name": "ml.r5.8xlarge",
    "vcpuNum": 32
   },
   {
    "_defaultOrder": 44,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 384,
    "name": "ml.r5.12xlarge",
    "vcpuNum": 48
   },
   {
    "_defaultOrder": 45,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 512,
    "name": "ml.r5.16xlarge",
    "vcpuNum": 64
   },
   {
    "_defaultOrder": 46,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 768,
    "name": "ml.r5.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 47,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.g5.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 48,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.g5.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 49,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 64,
    "name": "ml.g5.4xlarge",
    "vcpuNum": 16
   },
   {
    "_defaultOrder": 50,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 128,
    "name": "ml.g5.8xlarge",
    "vcpuNum": 32
   },
   {
    "_defaultOrder": 51,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 256,
    "name": "ml.g5.16xlarge",
    "vcpuNum": 64
   },
   {
    "_defaultOrder": 52,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 4,
    "hideHardwareSpecs": false,
    "memoryGiB": 192,
    "name": "ml.g5.12xlarge",
    "vcpuNum": 48
   },
   {
    "_defaultOrder": 53,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 4,
    "hideHardwareSpecs": false,
    "memoryGiB": 384,
    "name": "ml.g5.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 54,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 8,
    "hideHardwareSpecs": false,
    "memoryGiB": 768,
    "name": "ml.g5.48xlarge",
    "vcpuNum": 192
   },
   {
    "_defaultOrder": 55,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 8,
    "hideHardwareSpecs": false,
    "memoryGiB": 1152,
    "name": "ml.p4d.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 56,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 8,
    "hideHardwareSpecs": false,
    "memoryGiB": 1152,
    "name": "ml.p4de.24xlarge",
    "vcpuNum": 96
   }
  ],
  "instance_type": "ml.t3.medium",
  "kernelspec": {
   "display_name": "Python 3 (Data Science)",
   "language": "python",
   "name": "python3__SAGEMAKER_INTERNAL__arn:aws:sagemaker:us-east-1:081325390199:image/datascience-1.0"
  },
  "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.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}