notebooks/GenAI/Pubmed_RAG_chatbot.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "2edc6187-82ae-44e2-852f-2ad2712c93aa",
   "metadata": {},
   "source": [
    "# Creating a PubMed Chatbot with Llama2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6824ed09",
   "metadata": {},
   "source": [
    "## Overview"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ecea2ad-7c65-4367-87e1-b021167c3a1d",
   "metadata": {},
   "source": [
    "For this tutorial we are creating a PubMed chatbot that will answer questions by gathering information from documents we have provided via an index. The model we will be using today is a pretrained Llama2 model from Jumpstart."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "df0a5b23",
   "metadata": {},
   "source": [
    "## Learning Objectives\n",
    "- Introduce langchain\n",
    "- Explain the differences between zero-shot, one-shot, and few-shot prompting\n",
    "- Practice using different document retrievers"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41f7a40b",
   "metadata": {},
   "source": [
    "## Prerequisites\n",
    "You need access to SageMaker, Model Jumpstart, S3, and Kendra"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0da27877",
   "metadata": {},
   "source": [
    "## Get Started"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d01e74b-b5b4-4be9-b16e-ec55419318ef",
   "metadata": {},
   "source": [
    "### Deploy the Model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9dbd13e7-afc9-416b-94dc-418a93e14587",
   "metadata": {},
   "source": [
    "Identify which model we want to deploy from Jumpstart, in this case we are using Llama2 with 7 billion parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b51bf71-d2e5-4afc-8569-338767b43b9c",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_id, model_version = \"meta-textgeneration-llama-2-13b-f\", \"2.*\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "624c9bb8-3ce2-4240-b2b8-6b1bb93bb9f2",
   "metadata": {},
   "source": [
    "Create an endpoint to deploy our model locally, so that we can communicate with our model, send inputs, and retrieve outputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bf27747d-443f-47e7-9d2c-a8f5c5c6f3b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.jumpstart.model import JumpStartModel\n",
    "\n",
    "model = JumpStartModel(model_id=model_id, model_version=model_version)\n",
    "predictor = model.deploy()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "163fabcb-a267-4279-a5ae-c93d0143c139",
   "metadata": {},
   "source": [
    "Next we will print the endpoint name which we will need later to run our chatbot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ad71f0d-3be5-4b03-9c1c-eb4585721fc8",
   "metadata": {},
   "outputs": [],
   "source": [
    "endpoint_id=predictor.endpoint_name"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f3e3ab1-5f7e-4028-a66f-9619926a2afd",
   "metadata": {},
   "source": [
    "### PubMed API vs Kendra Index"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a820eea-1538-4f40-86c4-eb14fe09e127",
   "metadata": {},
   "source": [
    "Our chatbot respond to prompts based on the documents we supplied. This occurs via a **vector index**. A vector index is a data structure composed of vectorized embeddings (generated from our inputs) that enables fast and accurate search and retrieval from a large dataset of objects. We will explore using two methods to generate our index: PubMed API vs Kendra Index."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7314b115-9433-460d-b275-78aa50f0a858",
   "metadata": {},
   "source": [
    "**What is the difference?**\n",
    "\n",
    "The **PubMed API** is provided free by langchain to connect your model to more than **35 million citations** for biomedical literature from MEDLINE, life science journals, and online books. **Kendra index** is an AWS product that allows the user more **security and control** on which documents you wish to supply to your model. \n",
    "\n",
    "We will explore both methods to see which produces the best results!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bcf1690d-e93d-4cd3-89c6-8d06b5a071a8",
   "metadata": {},
   "source": [
    "#### Setting up a Kendra Index"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02053f4d-fad7-44ab-a7c3-cfa1c218240f",
   "metadata": {},
   "source": [
    "If you choose to use a Kendra index to supply documents to your model follow the instructions below:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d1c9de7-4a06-4f85-b9ff-c8c9e51f8c70",
   "metadata": {},
   "source": [
    "AWS marketplace provides a PubMed database named **PubMed Central® (PMC)** that contains free full-text archive of biomedical and life sciences journal article at the U.S. National Institutes of Health's National Library of Medicine (NIH/NLM). We will subset this database to add documents to our Kendra index. Ensure that you have the correct roles and policies to allow your environment to connect to S3 buckets, SageMaker, and Kendra."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "78418da3-f806-4b19-9980-1fa41083a75a",
   "metadata": {},
   "source": [
    "The first step will be to create a bucket that we will later use as our data source for our index."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99d49432-cf03-4f19-aa82-ef7f8bad5bde",
   "metadata": {},
   "outputs": [],
   "source": [
    "#make bucket\n",
    "bucket = 'pubmed-chat-docs'\n",
    "! aws s3 mb s3://{bucket}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6ad30ba-cee8-47f9-bc1e-ece8961ac66a",
   "metadata": {},
   "source": [
    "Next we will download the metadata file from the PMC bucket. This metadata file will list all of the articles within the PMC bucket and their paths. We will use these data to subset the database into our own bucket."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b395e34-062d-4f77-afee-3601d471954a",
   "metadata": {},
   "outputs": [],
   "source": [
    "#download the metadata file\n",
    "! aws s3 cp  s3://pmc-oa-opendata/oa_comm/txt/metadata/txt/oa_comm.filelist.txt . --sse"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "93a8595a-767f-4cad-9273-62d8e2cf60d1",
   "metadata": {},
   "source": [
    "We only want the metadata of the first 50 files to keep this tutorial short."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c26b0f29-2b07-43a6-800d-4aa5e957fe52",
   "metadata": {},
   "outputs": [],
   "source": [
    "#import the file as a dataframe\n",
    "import pandas as pd\n",
    "import os\n",
    "\n",
    "df = pd.read_csv('oa_comm.filelist.csv')\n",
    "\n",
    "#first 50 files\n",
    "first_50=df[0:50]\n",
    "#save new metadata\n",
    "first_50.to_csv('oa_comm.filelist_50.csv', index=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "abd1ae93-450e-4c79-83cc-ea46a1b507c1",
   "metadata": {},
   "source": [
    "Lets look at our metadata! We can see that the bucket path to the files are under the **Key** column. This column is what we will use to loop through the PMC bucket and copy the first 50 files to our bucket."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff77b2aa-ed1b-4d27-8163-fdaa7a304582",
   "metadata": {},
   "outputs": [],
   "source": [
    "first_50"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d63a7e2-dbf1-49ec-bc84-b8c2c8bde62d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "#gather path to files in bucket\n",
    "for i in first_50['Key']:\n",
    "    os.system(f'aws s3 cp s3://pmc-oa-opendata/{i} s3://{bucket}/docs/ --sse')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1b396c8-baa9-44d6-948c-2326dc514839",
   "metadata": {},
   "source": [
    "We will also save our new metadata file to our bucket to help Kendra index our files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "79e7076a-65c6-4e06-b84b-795ee7d4de00",
   "metadata": {},
   "outputs": [],
   "source": [
    "! aws s3 cp oa_comm.filelist_50.csv s3://{bucket}/docs/"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "79fb7fe8-4904-4b19-a411-95bca27ea87d",
   "metadata": {},
   "source": [
    "Now we can create our Kendra index, use the following instructions to create a index via the console or command line [Creating a Kendra index](https://docs.aws.amazon.com/kendra/latest/dg/create-index.html). To connect our bucket as a data source follow the instructions provided [here](https://docs.aws.amazon.com/kendra/latest/dg/data-source-s3.html) to do so via the console or AWS Python SDK."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07b3bc6b-8c43-476f-a662-abda830dc2da",
   "metadata": {},
   "source": [
    "### Creating a Inference Script "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ba2291e-109e-4120-ad10-5dbfd341a07b",
   "metadata": {},
   "source": [
    "For us to fluidly send input and receive outputs from our chatbot we need to create an [**inference script**](https://docs.aws.amazon.com/sagemaker/latest/dg/deploy-model.html#deploy-model-options) that will format inputs in a way that the chatbot can understand and format outputs in a way we can understand. We will also be supplying instructions to the chatbot through the script.\n",
    "\n",
    "Our script will utilize **LangChain** tools and packages to enable our model to:\n",
    "- **Connect to sources of context** (e.g. providing our model with tasks and examples)\n",
    "- **Rely on reason** (e.g. instruct our model how to answer based on provided context)\n",
    "\n",
    "**Warning**: The inference script must be run on the terminal via the command `python YOUR_SCRIPT.py`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "538f42ee-a502-4e56-9e85-4f6e3726ef7a",
   "metadata": {},
   "source": [
    "**Warning:** The following tools must be installed via your terminal"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "58582571-eae0-4440-afda-be1af6750e11",
   "metadata": {},
   "source": [
    "`pip install \"langchain\" \"xmltodict\"`"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad374085-c4b1-4083-85a5-90cba35846d6",
   "metadata": {},
   "source": [
    "The first part of our script will be to list all the tools that are required. \n",
    "-  **PubMedRetriever:** Utilizes the langchain retriever tool to specifically retrieve PubMed documents from the PubMed API.\n",
    "- **AmazonKendraRetriever:** Utilizes the langchain retriever tool to specifically retrieve documents stored in your Kendra index.\n",
    "- **ConversationalRetrievalChain:** Allows the user to construct a conversation with the model and retrieves the outputs while sending inputs to the model.\n",
    "- **PromptTemplate:** Allows the user to prompt the model to provide instructions, best method for zero and few shot prompting\n",
    "- **LLMContentHandler:** Handles the content to and from the model by transforming the input to a format that model can accept and transforms the output from the model to string that the LLM class expects.\n",
    "- **SagemakerEndpoint**: Connects to our endpoint in SageMaker and allows all of the tools listed above to connect to our model.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f0ad48d-c6c8-421a-a48b-88e979d15b57",
   "metadata": {},
   "source": [
    "```python\n",
    "from langchain_community.retrievers import PubMedRetriever\n",
    "from langchain.retrievers import AmazonKendraRetriever\n",
    "from langchain_community.llms import SagemakerEndpoint\n",
    "from langchain_community.llms.sagemaker_endpoint import LLMContentHandler\n",
    "from langchain.chains import ConversationalRetrievalChain\n",
    "from langchain.prompts import PromptTemplate\n",
    "import sys\n",
    "import json\n",
    "import os\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "900f4c31-71cd-4f39-8bfc-de098bdbaafc",
   "metadata": {},
   "source": [
    "Second will build a class that will hold the functions we need to send inputs and retrieve outputs from our model. For the beginning of our class we will establish some colors to our text conversation with our chatbot which we will utilize later."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "decbb901-f811-4b8e-a956-4c8c7f914ae2",
   "metadata": {},
   "source": [
    "```python\n",
    "class bcolors:\n",
    "    HEADER = '\\033[95m'\n",
    "    OKBLUE = '\\033[94m'\n",
    "    OKCYAN = '\\033[96m'\n",
    "    OKGREEN = '\\033[92m'\n",
    "    WARNING = '\\033[93m'\n",
    "    FAIL = '\\033[91m'\n",
    "    ENDC = '\\033[0m'\n",
    "    BOLD = '\\033[1m'\n",
    "    UNDERLINE = '\\033[4m'\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba36d057-5189-4075-a243-18996c6fc932",
   "metadata": {},
   "source": [
    "Next is to create a function that will gather the necessary information to connect to our model, which will be the:\n",
    "- Location\n",
    "- Kendra Index ID **(only if you are using Kendra instead of the PubMed API)**\n",
    "- Endpoint_name or ID"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3f7a244a-7e71-40d3-ae78-8e166dd3c7ee",
   "metadata": {},
   "source": [
    "```python\n",
    "def build_chain():\n",
    "  region = os.environ[\"AWS_REGION\"]\n",
    "  kendra_index_id = os.environ[\"KENDRA_INDEX_ID\"] #only needed is using a Kendra index instead of Pubmed API\n",
    "  endpoint_name = os.environ[\"LLAMA_2_ENDPOINT\"]\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e681908-df7a-4cfd-924d-5350eb06e770",
   "metadata": {},
   "source": [
    "Next we will create a class named **'ContentHandeler'** that will transforms our inputs and outputs into a json format. For Llama2 to understand and accept our inputs we need to structure in a specific manner, this is done by the **transform_input** function:\n",
    "```\n",
    "{\n",
    "\"inputs\": \n",
    "    [[\n",
    "        {\"role\": \"user\", \"content\": prompt},\n",
    "    ]],\n",
    "    **model_kwargs\n",
    "                                  }\n",
    "```\n",
    "Where `prompt` will be our instructions to our model (what the model is expected to do with our input) and `**model_kwargs` is where we provide our parameters.\n",
    "\n",
    "Our input is then encoded in a **UTF-8** format to convert our string into 0s and 1s. \n",
    "\n",
    "The next function in our class is named **transform_output**, this function will take the outputs sent from our model and decode them from 0s and 1s to strings."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d63f9b4c-202a-4170-9988-369bbe52fe6b",
   "metadata": {},
   "source": [
    "```python\n",
    "class ContentHandler(LLMContentHandler):\n",
    "      content_type = \"application/json\"\n",
    "      accepts = \"application/json\"\n",
    "\n",
    "      def transform_input(self, prompt: str, model_kwargs: dict) -> bytes:\n",
    "          input_str = json.dumps({\"inputs\": \n",
    "                                  [[\n",
    "                                    {\"role\": \"user\", \"content\": prompt},\n",
    "                                  ]],\n",
    "                                  **model_kwargs\n",
    "                                  })\n",
    "          \n",
    "          return input_str.encode('utf-8')\n",
    "      \n",
    "      def transform_output(self, output: bytes) -> str:\n",
    "          response_json = json.loads(output.read().decode(\"utf-8\"))\n",
    "          \n",
    "          return response_json[0]['generation']['content']\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dab1012f-ed20-47b9-9162-924e03e836d5",
   "metadata": {},
   "source": [
    "Now that we have a class that handles our input and outputs in a format that our model can understand we use **SageMakerEndpoint** tool to connect to our endpoint that we made in SageMaker. \n",
    "\n",
    "Notice that we set our class as the **content_handler** and we specified other **model_kwargs** as our parameters to control the temperature, top p, and max number of new tokens the model should generate to process our output.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8cadb1af-2c46-4ab1-92f9-6e0861f83324",
   "metadata": {},
   "source": [
    "```python\n",
    "content_handler = ContentHandler()\n",
    "\n",
    "  llm=SagemakerEndpoint(\n",
    "          endpoint_name=endpoint_name, \n",
    "          region_name=region, \n",
    "          model_kwargs={\"parameters\": {\"max_new_tokens\": 1000, \"top_p\": 0.9,\"temperature\":0.6}},\n",
    "          endpoint_kwargs={\"CustomAttributes\":\"accept_eula=true\"},\n",
    "          content_handler=content_handler,\n",
    "      )\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c44b4f91-0c64-459b-a6e9-8a955c0797c7",
   "metadata": {},
   "source": [
    "We specify what our retriever both the PubMed and Kendra retriever are list, please only add one per script."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "21c61724-23d3-4b49-8c72-cbd208bdb5df",
   "metadata": {},
   "source": [
    "```python\n",
    "retriever= PubMedRetriever()\n",
    "retriever = AmazonKendraRetriever(index_id=kendra_index_id,region_name=region) #only use if using Kendra Index\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec8e464a-0931-444a-aa58-09ee0c4c9884",
   "metadata": {},
   "source": [
    "Here we are constructing our **prompt_template**, this is where we can try zero-shot or few-shot prompting. Only add one method per script."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4431051e-0e84-408e-9821-f50a9b88c9c1",
   "metadata": {},
   "source": [
    "#### Zero-shot prompting\n",
    "\n",
    "Zero-shot prompting does not require any additional training more so it gives a pre-trained language model a task or query to generate text (our output). The model relies on its general language understanding and the patterns it has learned during its training to produce relevant output. In our script we have connect our model to a **retriever** to make sure it gathers information from that retriever (this can be the PubMed API or Kendra). \n",
    "\n",
    "See below that the task is more like instructions notifying our model they will be asked questions which it will answer based on the info of the scientific documents provided from the index provided (this can be the PubMed API or Kendra index). All of this information is established as a **prompt template** for our model to receive."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0316dc5-6274-4a5e-92e4-3d266ed6a4df",
   "metadata": {},
   "source": [
    "```python\n",
    "prompt_template = \"\"\"\n",
    "  Ignore everything before.\n",
    "  \n",
    "  Instruction:\n",
    "  Instructions:\n",
    "  I will provide you with research papers on a specific topic in English, and you will create a cumulative summary. \n",
    "  The summary should be concise and should accurately and objectively communicate the takeaway of the papers related to the topic. \n",
    "  You should not include any personal opinions or interpretations in your summary, but rather focus on objectively presenting the information from the papers. \n",
    "  Your summary should be written in your own words and ensure that your summary is clear, concise, and accurately reflects the content of the original papers. First, provide a concise summary then citations at the end.\n",
    "  \n",
    "  {question} Answer \"don't know\" if not present in the document. \n",
    "  {context}\n",
    "  Solution:\"\"\"\n",
    "  PROMPT = PromptTemplate(\n",
    "      template=prompt_template, input_variables=[\"context\", \"question\"],\n",
    "  )\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "edbe7032-8507-4d07-baab-1b3bf0e92074",
   "metadata": {},
   "source": [
    "#### One-shot and Few-shot Prompting"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5614ea04-e1f8-4941-ae16-4359f718f98f",
   "metadata": {},
   "source": [
    "One and few shot prompting are similar to one-shot prompting, in addition to giving our model a task just like before we have also supplied an example of how the our model structure our output.\n",
    "\n",
    "See below that we have implemented one-shot prompting to our script.  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ffb9669-5b77-4d9b-9f4e-a0d3a18b0fae",
   "metadata": {},
   "source": [
    "```python\n",
    "prompt_template = \"\"\"\n",
    "  Instructions:\n",
    "  I will provide you with research papers on a specific topic in English, and you will create a cumulative summary. \n",
    "  The summary should be concise and should accurately and objectively communicate the takeaway of the papers related to the topic. \n",
    "  You should not include any personal opinions or interpretations in your summary, but rather focus on objectively presenting the information from the papers. \n",
    "  Your summary should be written in your own words and ensure that your summary is clear, concise, and accurately reflects the content of the original papers. First, provide a concise summary then citations at the end. \n",
    "  Examples:\n",
    "  Question: What is a cell?\n",
    "  Answer: '''\n",
    "  Cell, in biology, the basic membrane-bound unit that contains the fundamental molecules of life and of which all living things are composed. \n",
    "  Sources: \n",
    "  Chow, Christopher , Laskey, Ronald A. , Cooper, John A. , Alberts, Bruce M. , Staehelin, L. Andrew , \n",
    "  Stein, Wilfred D. , Bernfield, Merton R. , Lodish, Harvey F. , Cuffe, Michael and Slack, Jonathan M.W.. \n",
    "  \"cell\". Encyclopedia Britannica, 26 Sep. 2023, https://www.britannica.com/science/cell-biology. Accessed 9 November 2023.\n",
    "  '''\n",
    "  \n",
    "  {question} Answer \"don't know\" if not present in the document. \n",
    "  {context}\n",
    "  \n",
    "\n",
    "  \n",
    "  Solution:\"\"\"\n",
    "  PROMPT = PromptTemplate(\n",
    "      template=prompt_template, input_variables=[\"context\", \"question\"],\n",
    "  )\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "82c66d53-97b2-46dc-a466-70a3d3bee4a7",
   "metadata": {},
   "source": [
    "The following set of commands control the chat history essentially telling the model to expect another question after it finishes answering the previous one. Follow up questions can contain references to past chat history so the **ConversationalRetrievalChain** combines the chat history and the followup question into a standalone question, then looks up relevant documents from the retriever, and finally passes those documents and the question to a question-answering chain to return a response.\n",
    "\n",
    "All of these pieces such as our conversational chain, prompt, and chat history are passed through a function called **run_chain** so that our model can return is response. We have also set the length of our chat history to one meaning that our model can only refer to the pervious conversation as a reference."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fda4d33b-60f2-4462-a8e6-bbce7f8a7b07",
   "metadata": {},
   "source": [
    "```python\n",
    "condense_qa_template = \"\"\"\n",
    "  Chat History:\n",
    "  {chat_history}\n",
    "  Here is a new question for you: {question}\n",
    "  Standalone question:\"\"\"\n",
    "  standalone_question_prompt = PromptTemplate.from_template(condense_qa_template)\n",
    " \n",
    "    qa = ConversationalRetrievalChain.from_llm(\n",
    "        llm=llm, \n",
    "        retriever=retriever, \n",
    "        condense_question_prompt=standalone_question_prompt, \n",
    "        return_source_documents=True, \n",
    "        combine_docs_chain_kwargs={\"prompt\":PROMPT},\n",
    "        )\n",
    "      return qa\n",
    "\n",
    "def run_chain(chain, prompt: str, history=[]):\n",
    "    print(prompt)\n",
    "    return chain({\"question\": prompt, \"chat_history\": history})\n",
    "\n",
    "MAX_HISTORY_LENGTH = 1 #increase to refer to more pervious chats\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8f1ef8d-66fe-4f84-933b-af2d730bd114",
   "metadata": {},
   "source": [
    "The final part of our script utilizes our class and incorporates colors to add a bit of flare to our conversation with our model. The model when first initialized should greet the user asking **\"Hello! How can I help you?\"** then instructs the user to ask a question or exit the session **\"Ask a question, start a New search: or CTRL-D to exit.\"**. With every question submitted to the model it is labeled as a **new search** we then run the run_chain function to get the models response or answer and add the response to the **chat history**. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1aa6ef65-ced4-445e-875c-7fee3483b81d",
   "metadata": {},
   "source": [
    "```python\n",
    "if __name__ == \"__main__\":\n",
    "  chat_history = []\n",
    "  qa = build_chain()\n",
    "  print(bcolors.OKBLUE + \"Hello! How can I help you?\" + bcolors.ENDC)\n",
    "  print(bcolors.OKCYAN + \"Ask a question, start a New search: or CTRL-D to exit.\" + bcolors.ENDC)\n",
    "  print(\">\", end=\" \", flush=True)\n",
    "  for query in sys.stdin:\n",
    "    if (query.strip().lower().startswith(\"new search:\")):\n",
    "      query = query.strip().lower().replace(\"new search:\",\"\")\n",
    "      chat_history = []\n",
    "    elif (len(chat_history) == MAX_HISTORY_LENGTH):\n",
    "      chat_history.pop(0)\n",
    "    result = run_chain(qa, query, chat_history)\n",
    "    chat_history.append((query, result[\"answer\"]))\n",
    "    print(bcolors.OKGREEN + result['answer'] + bcolors.ENDC)\n",
    "    ###this if statment is not needed for PubMed Retreiver users\n",
    "    if 'source_documents' in result: \n",
    "      print(bcolors.OKGREEN + 'Sources:')\n",
    "      for d in result['source_documents']:\n",
    "        print(d.metadata['source'])\n",
    "    ###\n",
    "    print(bcolors.ENDC)\n",
    "    print(bcolors.OKCYAN + \"Ask a question, start a New search: or CTRL-D to exit.\" + bcolors.ENDC)\n",
    "    print(\">\", end=\" \", flush=True)\n",
    "  print(bcolors.OKBLUE + \"Bye\" + bcolors.ENDC)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1abcbd48-bb84-4310-b8eb-ad87850a8649",
   "metadata": {},
   "source": [
    "Running our script in the terminal will require us to export the following global variables before running the python script. Don't forget to run you python script on the terminal using the command `python NAME_OF_YOUR_SCRIPT.py`. For more guidence take a look at our **example inference scripts** for the [PubMed API](/example_scripts/langchain_chat_llama_2_zeroshot.py) and [Kendra](/example_scripts/kendra_chat_llama_2.py)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba97df23-6893-438d-8a67-cb7dbf83e407",
   "metadata": {},
   "outputs": [],
   "source": [
    "#retreive our endpoint id\n",
    "endpoint_id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7eab00a3-54ff-4873-8d25-eaf8bd18a2e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "#enter these global variables in your terminal\n",
    "export AWS_REGION='<Enter_location>'\n",
    "export LLAMA_2_ENDPOINT='<Enter_endpoint_id>'\n",
    "export KENDRA_INDEX_ID='<Enter_kendra_index>'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bbe127e6-c0b1-4e07-ad56-38c30a9bf858",
   "metadata": {},
   "source": [
    "You should see similar results on the terminal. In this example we ask the chatbot to explain brain cancer!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "80c8fb4b-e74f-4e8d-892b-0f913eff747d",
   "metadata": {},
   "source": [
    "![PubMed Chatbot Results](../../docs/images/PubMed_chatbot_results.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ddeee006",
   "metadata": {},
   "source": [
    "## Conclusions\n",
    "Here you learned how to deploy a model, create a vector database (index) from PubMed documents, and then interact with a model to product predictions using an inference script. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a178c1c6-368a-48c5-8beb-278443b685a2",
   "metadata": {},
   "source": [
    "### Clean Up"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ec06a34-dc47-453f-b519-424804fa2748",
   "metadata": {},
   "source": [
    "**Warning:** Dont forget to delete the resources we just made to avoid accruing additional costs!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c307bb17-757a-4579-a0d8-698eb1bb3f2e",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Delete model and endpoint\n",
    "model.delete()\n",
    "endpoint.delete()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "280cea0a-a8fc-494e-8ce4-afb65847a222",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Delete bucket\n",
    "! aws s3 rb s3://$bucket --force  "
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 5
}