RAG-tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0a211289",
   "metadata": {},
   "source": [
    "# RAG Tutorial\n",
    "\n",
    "Guest lecture for DS 3891 Generative AI\n",
    "\n",
    "*By Myranda Uselton Shirk, Senior Data Scientist, Data Science Institute*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1544305c",
   "metadata": {},
   "source": [
    "## What is RAG?\n",
    "\n",
    "RAG (Retrieval Augmented Generation) is a method of interacting with a large language model (LLM) by providing the model with a corpus of documents. In code, this process has two main steps (each listed here with their own sub-steps): \n",
    "\n",
    "**1. Retrieve information relevant to query**\n",
    "\n",
    "    a) Load the documents\n",
    "\n",
    "    b) Split/Chunk text\n",
    "\n",
    "    c) Embed each split\n",
    "\n",
    "    d) Semantic Search over text embeddings\n",
    "\n",
    "**2. Generate a response based on the retrieved information**\n",
    "\n",
    "    a) Give prompt + retrieved information to model\n",
    "\n",
    "    b) Generate response"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a5fe25f",
   "metadata": {},
   "source": [
    "This tutorial will walk through how to implement RAG using the Python library [Langchain](https://python.langchain.com/docs/get_started/introduction)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b44b0ca",
   "metadata": {},
   "source": [
    "## Libraries"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2364674",
   "metadata": {},
   "source": [
    "### Google Colab \n",
    "\n",
    "If you are using this notebook in Google Colab, uncomment and run this code cell:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eae31c24",
   "metadata": {},
   "outputs": [],
   "source": [
    "#!pip install -q langchain langchain-community langchain-core langchain-openai getpass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "1b85c17e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# import libraries\n",
    "\n",
    "from langchain_community.vectorstores import FAISS\n",
    "from langchain_community.document_loaders import TextLoader, PyPDFLoader\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "from langchain_core.prompts import PromptTemplate, ChatPromptTemplate\n",
    "from langchain_core.runnables import RunnableLambda, RunnablePassthrough\n",
    "from langchain_openai import ChatOpenAI, OpenAIEmbeddings\n",
    "from langchain.text_splitter import CharacterTextSplitter"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "485bf656",
   "metadata": {},
   "source": [
    "### Set up OpenAI API Key\n",
    "\n",
    "You will need an OpenAI API key to run this notebook, unless you change the model to inference. If you have an OpenAI API key, run this cell and enter it when prompted."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "8dd46452",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "········\n"
     ]
    }
   ],
   "source": [
    "import getpass\n",
    "import os\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = getpass.getpass()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b1af5cfb",
   "metadata": {},
   "source": [
    "## Walkthrough: Text Files\n",
    "\n",
    "First we will walk through the simplest case for information storage: a text file. The text file \"state_of_the_union.txt\" should be in your working directory (if it is not, find it in the GitHub repo and move it to where you can access it)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73c30216",
   "metadata": {},
   "source": [
    "### Load Files\n",
    "\n",
    "The first step in RAG is to load our information corpus - in this case, \"state_of_the_union.txt\", which contains the state of the union address given by President Biden.\n",
    "\n",
    "Langchain has many different types of document loaders. For text files we can use `TextLoader`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "a97b87bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load a text file\n",
    "\n",
    "text_loader = TextLoader(\"./state_of_the_union.txt\")\n",
    "docs = text_loader.load()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c8517b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# view the docs\n",
    "docs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7af1778d",
   "metadata": {},
   "source": [
    "### Split Text\n",
    "\n",
    "We need to split our text into individual chunks of information. There are several ways to do this. We will start with the simplest one, `CharacterTextSplitter`, which splits our text by number of characters. Then, we will embed our data using the OpenAI Embeddings, and store all of that into a database (we're using [FAISS database](https://faiss.ai/))."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "59bfdc00",
   "metadata": {},
   "outputs": [],
   "source": [
    "#chunk/split text\n",
    "text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)\n",
    "documents = text_splitter.split_documents(docs)\n",
    "\n",
    "# save vector embeddings in a database\n",
    "db = FAISS.from_documents(documents, OpenAIEmbeddings())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "85b41cd7",
   "metadata": {},
   "source": [
    "### Retreive information\n",
    "\n",
    "Next, we need to set up a retriever to retrieve relevant information for us. Again, there are several different methods (see the Resources tab in the README for those), but we will be using a vector similarity search. This will compare our query to each chunk of text and pull out the most relevant one(s). \n",
    "\n",
    "The below cell shows how you can search across a database. Then, we set up our retriever in the next cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "18c7be78",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tonight. I call on the Senate to: Pass the Freedom to Vote Act. Pass the John Lewis Voting Rights Act. And while you’re at it, pass the Disclose Act so Americans can know who is funding our elections. \n",
      "\n",
      "Tonight, I’d like to honor someone who has dedicated his life to serve this country: Justice Stephen Breyer—an Army veteran, Constitutional scholar, and retiring Justice of the United States Supreme Court. Justice Breyer, thank you for your service. \n",
      "\n",
      "One of the most serious constitutional responsibilities a President has is nominating someone to serve on the United States Supreme Court. \n",
      "\n",
      "And I did that 4 days ago, when I nominated Circuit Court of Appeals Judge Ketanji Brown Jackson. One of our nation’s top legal minds, who will continue Justice Breyer’s legacy of excellence.\n"
     ]
    }
   ],
   "source": [
    "# find relevant documents\n",
    "query = \"What did the president say about Ketanji Brown Jackson\"\n",
    "relevant_docs = db.similarity_search(query)\n",
    "print(relevant_docs[0].page_content)\n",
    "\n",
    "# Note this is NOT querying a chat model - only finding the relevant info"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "667a46a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# set up retriever\n",
    "\n",
    "retriever = db.as_retriever(search_type=\"similarity\", search_kwargs={\"k\": 6})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c380b258",
   "metadata": {},
   "source": [
    "### Prompting the LLM\n",
    "\n",
    "Now that we have set up our information corpus and a retriever, we can set up a prompt template for our model to use. This template allows us to organize our instructions to the model and insert the retrieved information and the query in the appropriate place."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "1f2dee63",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "PromptTemplate(input_variables=['context', 'question'], template='Use the following pieces of context to answer the question at the end.\\nIf you don\\'t know the answer, just say that you don\\'t know, don\\'t try to make up an answer.\\nUse three sentences maximum and keep the answer as concise as possible.\\nAlways say \"thanks for asking!\" at the end of the answer.\\n\\n{context}\\n\\nQuestion: {question}\\n\\nHelpful Answer:')"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "template = \"\"\"Use the following pieces of context to answer the question at the end.\n",
    "If you don't know the answer, just say that you don't know, don't try to make up an answer.\n",
    "Use three sentences maximum and keep the answer as concise as possible.\n",
    "Always say \"thanks for asking!\" at the end of the answer.\n",
    "\n",
    "{context}\n",
    "\n",
    "Question: {question}\n",
    "\n",
    "Helpful Answer:\"\"\"\n",
    "custom_rag_prompt = PromptTemplate.from_template(template)\n",
    "custom_rag_prompt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c112cf1",
   "metadata": {},
   "source": [
    "Finally, we set up our RAG chain and pass in our model, retriever, prompt template, and query."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd11c0a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# helper function to format our documents\n",
    "def format_docs(docs):\n",
    "    return \"\\n\\n\".join(doc.page_content for doc in docs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ab1d367",
   "metadata": {},
   "outputs": [],
   "source": [
    "# load model\n",
    "llm = ChatOpenAI(model_name=\"gpt-3.5-turbo\", temperature=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88ac002b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# set up RAG chain\n",
    "rag_chain = (\n",
    "    {\"context\": retriever | format_docs, \"question\": RunnablePassthrough()}\n",
    "    | custom_rag_prompt\n",
    "    | llm\n",
    "    | StrOutputParser()\n",
    ")\n",
    "\n",
    "rag_chain.invoke(query)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9e1bbb4",
   "metadata": {},
   "source": [
    "## Extension: QA over a PDF\n",
    "\n",
    "Most of the time, our text corpus may not be in .txt files. Here is a tutorial on creating a corpus from a PDF. In this case, we use the `PyPDFLoader`. You might recognize the paper I am using for our data!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "8445d06e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load a PDF\n",
    "\n",
    "pdf_loader = PyPDFLoader(\"./prompt_engineering.pdf\")\n",
    "pdf_pages = pdf_loader.load_and_split()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "d31df303",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3: TABLE I\n",
      "CLASSIFYING PROMPT PATTERNS\n",
      "Pattern Category Prompt Pattern\n",
      "Input Semantics Meta Language Creation\n",
      "Output Output Automater\n",
      "Customization Persona\n",
      "Visualization Generator\n",
      "Recipe\n",
      "Template\n",
      "Error Identiﬁcation Fact Check List\n",
      "Reﬂection\n",
      "Prompt Question Reﬁnement\n",
      "Improvement Alternative Approaches\n",
      "\n",
      "2: textual statements approach is that it is intentionally int uitive\n",
      "to users. In particular, we expect users will understand how to\n",
      "express and adapt the statements in a contextually appropri ate\n",
      "way for their domain. Moreover, since the underlying ideas o f\n",
      "the prompt are captured, these same ideas \n"
     ]
    }
   ],
   "source": [
    "# retrieve info from documents\n",
    "\n",
    "faiss_index = FAISS.from_documents(pdf_pages, OpenAIEmbeddings())\n",
    "docs = faiss_index.similarity_search(\"What are the different categories of prompts?\", k=2)\n",
    "for doc in docs:\n",
    "    print(str(doc.metadata[\"page\"]) + \":\", doc.page_content[:300])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "591247e8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'The different categories of prompts are Input Semantics, Output Customization, Error Identification, Prompt Improvement, Interaction, and Context Control. Each category focuses on different aspects of prompt patterns in the context of conversational LLMs. Thanks for asking!'"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "query = \"What are the different categories of prompts?\"\n",
    "\n",
    "retriever = faiss_index.as_retriever(search_type=\"similarity\", search_kwargs={\"k\": 6})\n",
    "\n",
    "rag_chain = (\n",
    "    {\"context\": retriever | format_docs, \"question\": RunnablePassthrough()}\n",
    "    | custom_rag_prompt\n",
    "    | llm\n",
    "    | StrOutputParser()\n",
    ")\n",
    "\n",
    "rag_chain.invoke(query)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8204ec1a",
   "metadata": {},
   "source": [
    "## Playground\n",
    "\n",
    "Now that you know how to use RAG, try it out on your own documents below. Experiment with different document loaders, load multiple documents, or change the vector store/prompt/model you use. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aedc283d",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}