{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Building a Llama 3 chatbot with Retrieval Augmented Generation (RAG)\n",
    "\n",
    "This notebook shows a complete example of how to build a Llama 2 chatbot hosted on your browser that can answer questions based on your own data. We'll cover:\n",
    "* How to run Llama 3 in the cloud hosted on OctoAI\n",
    "* A chatbot example built with [Gradio](https://github.com/gradio-app/gradio) and wired to the server\n",
    "* Adding RAG capability with Llama 3 specific knowledge based on our Getting Started [guide](https://ai.meta.com/llama/get-started/)\n",
    "\n",
    "\n",
    "**Note** We will be using OctoAI to run the examples here. You will need to first sign into [OctoAI](https://octoai.cloud/) with your Github or Google account, then create a free API token [here](https://octo.ai/docs/getting-started/how-to-create-an-octoai-access-token) that you can use for a while (a month or $10 in OctoAI credits, whichever one runs out first).\n",
    "After the free trial ends, you will need to enter billing info to continue to use Llama 3 hosted on OctoAI."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## RAG Architecture\n",
    "\n",
    "LLMs have unprecedented capabilities in NLU (Natural Language Understanding) & NLG (Natural Language Generation), but they have a knowledge cutoff date, and are only trained on publicly available data before that date.\n",
    "\n",
    "RAG, invented by [Meta](https://ai.meta.com/blog/retrieval-augmented-generation-streamlining-the-creation-of-intelligent-natural-language-processing-models/) in 2020, is one of the most popular methods to augment LLMs. RAG allows enterprises to keep sensitive data on-prem and get more relevant answers from generic models without fine-tuning models for specific roles.\n",
    "\n",
    "RAG is a method that:\n",
    "* Retrieves data from outside a foundation model\n",
    "* Augments your questions or prompts to LLMs by adding the retrieved relevant data as context\n",
    "* Allows LLMs to answer questions about your own data, or data not publicly available when LLMs were trained\n",
    "* Greatly reduces the hallucination in  model's response generation\n",
    "\n",
    "The following diagram shows the general RAG components and process:"
   ]
  },
  {
   "attachments": {
    "image.png": {
     "image/png": "[encoded data, 99588 characters]"
    }
   },
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![image.png](attachment:image.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## How to Develop a RAG Powered Llama 3 Chatbot\n",
    "\n",
    "The easiest way to develop RAG-powered Llama 3 chatbots is to use frameworks such as [**LangChain**](https://www.langchain.com/) and [**LlamaIndex**](https://www.llamaindex.ai/), two leading open-source frameworks for building LLM apps. Both offer convenient APIs for implementing RAG with Llama 3 including:\n",
    "\n",
    "* Load and split documents\n",
    "* Embed and store document splits\n",
    "* Retrieve the relevant context based on the user query\n",
    "* Call Llama 3 with query and context to generate the answer\n",
    "\n",
    "LangChain is a more general purpose and flexible framework for developing LLM apps with RAG capabilities, while LlamaIndex as a data framework focuses on connecting custom data sources to LLMs. The integration of the two may provide the best performant and effective solution to building real world RAG apps.\n",
    "In our example, for simplicifty, we will use LangChain alone with locally stored PDF data."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Install Dependencies\n",
    "\n",
    "For this demo, we will be using the Gradio for chatbot UI, Text-generation-inference framework for model serving.\n",
    "For vector storage and similarity search, we will be using [FAISS](https://github.com/facebookresearch/faiss).\n",
    "In this example, we will be running everything in a AWS EC2 instance (i.e. [g5.2xlarge]( https://aws.amazon.com/ec2/instance-types/g5/)). g5.2xlarge features one A10G GPU. We recommend running this notebook with at least one GPU equivalent to A10G with at least 16GB video memory.\n",
    "There are certain techniques to downsize the Llama 3 7B model, so it can fit into smaller GPUs. But it is out of scope here.\n",
    "\n",
    "First, let's install all dependencies with PIP. We also recommend you start a dedicated Conda environment for better package management.\n",
    "\n",
    "And let's set up the OctoAI token."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -r requirements.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from getpass import getpass\n",
    "import os\n",
    "\n",
    "OCTOAI_API_TOKEN = getpass()\n",
    "os.environ[\"OCTOAI_API_TOKEN\"] = OCTOAI_API_TOKEN"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Data Processing\n",
    "\n",
    "First run all the imports and define the path of the data and vector storage after processing.\n",
    "For the data, we will be using a raw pdf crawled from \"Llama 2 Getting Started\" guide on [Meta AI website](https://ai.meta.com/llama/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.embeddings import OctoAIEmbeddings\n",
    "from langchain.vectorstores import FAISS\n",
    "from langchain.document_loaders import PyPDFDirectoryLoader\n",
    "from langchain.text_splitter import RecursiveCharacterTextSplitter\n",
    "\n",
    "DATA_PATH = 'data' #Your root data folder path\n",
    "DB_FAISS_PATH = 'vectorstore/db_faiss'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then we use the `PyPDFDirectoryLoader` to load the entire directory. You can also use `PyPDFLoader` for loading one single file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loader = PyPDFDirectoryLoader(DATA_PATH)\n",
    "documents = loader.load()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check the length and content of the doc to ensure we have loaded the right document with number of pages as 37."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(documents), documents[0].page_content[0:100])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Split the loaded documents into smaller chunks.\n",
    "[`RecursiveCharacterTextSplitter`](https://api.python.langchain.com/en/latest/text_splitter/langchain.text_splitter.RecursiveCharacterTextSplitter.html) is one common splitter that splits long pieces of text into smaller, semantically meaningful chunks.\n",
    "Other splitters include:\n",
    "* SpacyTextSplitter\n",
    "* NLTKTextSplitter\n",
    "* SentenceTransformersTokenTextSplitter\n",
    "* CharacterTextSplitter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=10)\n",
    "splits = text_splitter.split_documents(documents)\n",
    "print(len(splits), splits[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that we have set `chunk_size` to 500 and `chunk_overlap` to 10. In the spliting, these two parameters can directly affects the quality of the LLM's answers.\n",
    "Here is a good [guide](https://dev.to/peterabel/what-chunk-size-and-chunk-overlap-should-you-use-4338) on how you should carefully set these two parameters."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we will need to choose an embedding model for our splited documents.\n",
    "**Embeddings are numerial representations of text**. The default embedding model in OctoAI Embeddings is GTE-Large with a 1024 vector length."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "embeddings = OctoAIEmbeddings(endpoint_url=\"https://text.octoai.run/v1/embeddings\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Lastly, with splits and choice of the embedding model ready, we want to index them and store all the split chunks as embeddings into the vector storage.\n",
    "\n",
    "Vector stores are databases storing embeddings. There're at least 60 [vector stores](https://python.langchain.com/docs/integrations/vectorstores) supported by LangChain, and two of the most popular open source ones are:\n",
    "* [Chroma](https://www.trychroma.com/): a light-weight and in memory so it's easy to get started with and use for **local development**.\n",
    "* [FAISS](https://python.langchain.com/docs/integrations/vectorstores/faiss) (Facebook AI Similarity Search): a vector store that supports search in vectors that may not fit in RAM and is appropriate for **production use**.\n",
    "\n",
    "Since we are running on a EC2 instance with abundant CPU resources and RAM, we will use FAISS in this example. Note that FAISS can also run on GPUs, where some of the most useful algorithms are implemented there. In that case, install `faiss-gpu` package with PIP instead."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "db = FAISS.from_documents(splits, embeddings)\n",
    "db.save_local(DB_FAISS_PATH)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once you saved database into local path. You can find them as `index.faiss` and `index.pkl`. In the chatbot example, you can then load this database from local and plug it into our retrival process."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Building the Chatbot UI\n",
    "\n",
    "Now we are ready to build the chatbot UI to wire up RAG data and API server. In our example we will be using Gradio to build the Chatbot UI.\n",
    "Gradio is an open-source Python library that is used to build machine learning and data science demos and web applications. It has been widely used by the community. Other alternatives are:\n",
    "* [Streamlit](https://streamlit.io/)\n",
    "* [Dash](https://plotly.com/dash/)\n",
    "* [Flask](https://flask.palletsprojects.com/en/3.0.x/)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Again, we start by adding all the imports, paths, constants and set LangChain in debug mode, so it shows clear actions within the chain process."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import langchain\n",
    "from queue import Queue\n",
    "from typing import Any\n",
    "from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler\n",
    "from langchain.schema import LLMResult\n",
    "from langchain.embeddings import OctoAIEmbeddings\n",
    "from langchain.vectorstores import FAISS\n",
    "from langchain.chains import RetrievalQA\n",
    "from langchain.prompts.prompt import PromptTemplate\n",
    "from anyio.from_thread import start_blocking_portal #For model callback streaming\n",
    "\n",
    "# Vector db path\n",
    "DB_FAISS_PATH = 'vectorstore/db_faiss'\n",
    "\n",
    "model_dict = {\n",
    "    \"8b-instruct\" : \"meta-llama-3-8b-instruct\",\n",
    "    \"70b-instruct\" : \"meta-llama-3-70b-instruct\",\n",
    "}\n",
    "\n",
    "system_message = {\"role\": \"system\", \"content\": \"You are a helpful assistant.\"}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then we load the FAISS vector store"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "embeddings = OctoAIEmbeddings(endpoint_url=\"https://text.octoai.run/v1/embeddings\")\n",
    "db = FAISS.load_local(DB_FAISS_PATH, embeddings, allow_dangerous_deserialization=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we call the Llama 3 model from OctoAI. In this example we will use the Llama 3 8b instruct model. You can find more on Llama models on the [OctoAI text generation solution page](https://octoai.cloud/text).\n",
    "\n",
    "At the time of writing this notebook the following Llama models are available on OctoAI:\n",
    "* meta-llama-3-8b-instruct\n",
    "* meta-llama-3-70b-instruct\n",
    "* codellama-7b-instruct\n",
    "* codellama-13b-instruct\n",
    "* codellama-34b-instruct\n",
    "* llama-2-13b-chat\n",
    "* llama-2-70b-chat\n",
    "* llamaguard-7b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.llms.octoai_endpoint import OctoAIEndpoint\n",
    "\n",
    "llm = OctoAIEndpoint(\n",
    "    model=model_dict[\"8b-instruct\"],\n",
    "    max_tokens=500,\n",
    "    temperature=0.01\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we define the retriever and template for our RetrivalQA chain. For each call of the RetrievalQA, LangChain performs a semantic similarity search of the query in the vector database, then passes the search results as the context to Llama to answer the query about the data stored in the verctor database.\n",
    "Whereas for the template, this defines the format of the question along with context that we will be sent into Llama for generation. In general, Llama 3 has special prompt format to handle special tokens. In some cases, the serving framework might already have taken care of it. Otherwise, you will need to write customized template to properly handle that."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "template = \"\"\"\n",
    "[INST]Use the following pieces of context to answer the question. If no context provided, answer like a AI assistant.\n",
    "{context}\n",
    "Question: {question} [/INST]\n",
    "\"\"\"\n",
    "\n",
    "retriever = db.as_retriever(\n",
    "        search_kwargs={\"k\": 6}\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Lastly, we can define the retrieval chain for QA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "qa_chain = RetrievalQA.from_chain_type(\n",
    "    llm=llm,\n",
    "    retriever=retriever,\n",
    "    chain_type_kwargs={\n",
    "        \"prompt\": PromptTemplate(\n",
    "            template=template,\n",
    "            input_variables=[\"context\", \"question\"],\n",
    "        ),\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we should have a working chain for QA. Let's test it out before wire it up with UI blocks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result = qa_chain.invoke({\"query\": \"Why choose Llama?\"})\n",
    "print(result[\"result\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After confirming the validity, we can start building the UI. We'll use a simple interface built out of Gradio's ChatInterface."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gradio as gr\n",
    "\n",
    "def predict(message, history):\n",
    "    llm_response = qa_chain.invoke(message)[\"result\"]\n",
    "    return llm_response\n",
    "\n",
    "gr.ChatInterface(predict).launch()"
   ]
  }
 ],
 "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.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
