search/retrieval-augmented-generation-llm.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "67cf6cc2-b3cf-4843-b502-71872d3dde35",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "# Retrieval-augmented Generation (RAG)\n",
    "\n",
    "This notebook provides several prototypes of a Retrieval-Augmented Generation (RAG) system. We use product data analytics use cases as a running example.\n",
    "\n",
    "### Use Case\n",
    "We have a large collection of documents and want to use LLM to summarize these documents, answer standalone questions based on the document content, or answer questions in a conversational mode. Examples include a sales assistant that answers customers' questions about company's products, coding assistant that answers developers' questions about the codebase, and legal assistant that answers questions about regulations.   \n",
    "\n",
    "### Prototype: Approach and Data\n",
    "We start with a basic case where the input documents are small enough to fit the LLM context, and then develop more advanced solutions that can handle large document collections. We use small input documents that are available in the `tensor-house-data` repository.\n",
    "\n",
    "### Usage and Productization\n",
    "The implementation uses a production-grade framework, but external embedding storage (vector store) and additional components such as caching are typically needed to create production grade applications."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "49c0804d-029e-416d-8014-c7828d96e4d9",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "## Environment Setup and Initialization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "2a97eb6a-ccec-4527-b9fb-36b69c6ff3de",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": [],
    "ExecuteTime": {
     "end_time": "2024-01-14T16:38:51.407953Z",
     "start_time": "2024-01-14T16:38:51.282527Z"
    }
   },
   "outputs": [],
   "source": [
    "#\n",
    "# Imports\n",
    "#\n",
    "from langchain_community.llms.vertexai import VertexAI\n",
    "\n",
    "from langchain_core.prompts import PromptTemplate\n",
    "\n",
    "from langchain.text_splitter import CharacterTextSplitter\n",
    "from langchain_community.document_loaders import TextLoader\n",
    "\n",
    "from langchain_community.embeddings import VertexAIEmbeddings\n",
    "from langchain_community.vectorstores.chroma import Chroma\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "from langchain_core.runnables import RunnablePassthrough\n",
    "\n",
    "#\n",
    "# Initialize LLM provider\n",
    "# (google-cloud-aiplatform must be installed)\n",
    "#\n",
    "from google.cloud import aiplatform\n",
    "aiplatform.init(\n",
    "    project='<< specify your project name here >>',\n",
    "    location='us-central1'\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c28f863-c2e2-4c56-b459-df05ff12ad8f",
   "metadata": {},
   "source": [
    "## Question Answering Using In-Prompt Documents\n",
    "\n",
    "The most basic scenario is querying small documents that fit the LLM prompt."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "bf134242-352d-4281-8f57-099ca387143b",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-14T15:35:31.201941Z",
     "start_time": "2024-01-14T15:35:29.085146Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Lemons and limes have the highest concentration of citric acid among the mentioned fruits.\n"
     ]
    }
   ],
   "source": [
    "query_prompt = \"\"\"\n",
    "Please read the following text and answer which fruits have the highest concentration of citric acid.\n",
    "\n",
    "TEXT:\n",
    "Citric acid occurs in a variety of fruits and vegetables, most notably citrus fruits. Lemons and limes have particularly \n",
    "high concentrations of the acid; it can constitute as much as 8% of the dry weight of these fruits (about 47 g/L in the juices).\n",
    "The concentrations of citric acid in citrus fruits range from 0.005 mol/L for oranges and grapefruits to 0.30 mol/L in lemons \n",
    "and limes; these values vary within species depending upon the cultivar and the circumstances under which the fruit was grown.\n",
    "\"\"\"\n",
    "\n",
    "llm = VertexAI(temperature=0.7)\n",
    "response = llm(query_prompt)\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6869bf50-e4f2-4ba6-b7b9-2278c1a072a0",
   "metadata": {},
   "source": [
    "## Question Answering Using MapReduce\n",
    "\n",
    "For large documents and collections of documents that do not fit the LLM context, we can apply the MapReduce pattern to independently extract relevant summaries from document parts, and then merge these summaries into the final answer. This approach is appropriate for summarization and summarization-like queries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "c5a43885-5f2e-49d3-b5cf-bc6204b10d3a",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-14T16:23:44.559333Z",
     "start_time": "2024-01-14T16:23:37.703018Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The input document has been split into 2 chunks\n",
      " The three most important applications of citric acid are:\n",
      "\n",
      "1. **Food and beverage industry:** Citric acid enhances flavors and acts as a preservative, making it a popular choice in various food and beverage products, especially soft drinks.\n",
      "\n",
      "2. **Cleaning products:** Citric acid's ability to chelate metals and remove hard water stains makes it an effective component in soaps, detergents, and household cleaners, improving their performance in hard water conditions.\n",
      "\n",
      "3. **Pharmaceutical and biotechnology industry:** Citric acid is used in various pharmaceutical and biotechnological applications, including hair care products to open hair cuticles and enhance treatment penetration, photography\n"
     ]
    }
   ],
   "source": [
    "from langchain.chains import LLMChain, StuffDocumentsChain, ReduceDocumentsChain, MapReduceDocumentsChain\n",
    "\n",
    "#\n",
    "# Load the input document\n",
    "#\n",
    "loader = TextLoader(\"../../tensor-house-data/search/food-additives/citric-acid-applications.txt\")\n",
    "documents = loader.load()\n",
    "\n",
    "#\n",
    "# Splitting\n",
    "#\n",
    "text_splitter = CharacterTextSplitter(chunk_size=2000, chunk_overlap=0)\n",
    "texts = text_splitter.split_documents(documents)\n",
    "print(f'The input document has been split into {len(texts)} chunks\\n')\n",
    "\n",
    "#\n",
    "# Querying\n",
    "#\n",
    "map_prompt = \"\"\"\n",
    "Use the following portion of a long document to see if any of the text is relevant to answer the question. \n",
    "Return bullet points that help to answer the question.\n",
    "\n",
    "{context}\n",
    "\n",
    "Question: {question}\n",
    "Bullet points:\n",
    "\"\"\"\n",
    "map_prompt_template = PromptTemplate(template=map_prompt, input_variables=[\"context\", \"question\"])\n",
    "\n",
    "combine_prompt = \"\"\"\n",
    "Given the following bullet points extracted from a long document and a question, create a final answer.\n",
    "Question: {question}\n",
    "\n",
    "=========\n",
    "{context}\n",
    "=========\n",
    "\n",
    "Final answer:\n",
    "\"\"\"\n",
    "combine_prompt_template = PromptTemplate(template=combine_prompt, input_variables=[\"summaries\", \"question\"])\n",
    "\n",
    "llm = VertexAI(temperature=0.7)\n",
    "\n",
    "document_prompt = PromptTemplate(\n",
    "    input_variables=[\"page_content\"],\n",
    "    template=\"{page_content}\"\n",
    ")\n",
    "\n",
    "document_variable_name = \"context\"\n",
    "\n",
    "map_llm_chain = LLMChain(llm=llm, prompt=map_prompt_template)\n",
    "\n",
    "reduce_llm_chain = LLMChain(llm=llm, prompt=combine_prompt_template)\n",
    "\n",
    "combine_documents_chain = StuffDocumentsChain(\n",
    "    llm_chain=reduce_llm_chain,\n",
    "    document_prompt=document_prompt,\n",
    "    document_variable_name=document_variable_name\n",
    ")\n",
    "reduce_documents_chain = ReduceDocumentsChain(\n",
    "    combine_documents_chain=combine_documents_chain,\n",
    ")\n",
    "chain = MapReduceDocumentsChain(\n",
    "    llm_chain=map_llm_chain,\n",
    "    reduce_documents_chain=reduce_documents_chain,\n",
    "    document_variable_name=document_variable_name\n",
    ")\n",
    "\n",
    "question = \"What are the three most important applications of citric acid? Provide a short justification for each application.\"\n",
    "\n",
    "print(chain.invoke({'input_documents': texts, 'question': question})['output_text'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b3eb6e94-a2f2-4548-8d0c-f4934f0224f6",
   "metadata": {},
   "source": [
    "## Question Answering Using Vector Search\n",
    "\n",
    "For large documents and point questions that require only specific document parts to be answered, LLMs can be combined with traditional information retrieval techniques. The input document(s) is split into chunks which are then indexed in a vector store. To answer the user question, the most relevant chunks are retrieved and passed to the LLM."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "3857ccc4-2559-493d-9e44-babb744576db",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": [],
    "ExecuteTime": {
     "end_time": "2024-01-14T16:48:35.054246Z",
     "start_time": "2024-01-14T16:48:32.371222Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The input document has been split into 5 chunks\n"
     ]
    },
    {
     "data": {
      "text/plain": "' The melting point of citric acid is approximately 153 °C (307 °F).'"
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#\n",
    "# Load the input document\n",
    "#\n",
    "loader = TextLoader(\"../../tensor-house-data/search/food-additives/food-additives.txt\")\n",
    "documents = loader.load()\n",
    "\n",
    "#\n",
    "# Splitting\n",
    "#\n",
    "text_splitter = CharacterTextSplitter(chunk_size=3000, chunk_overlap=0)\n",
    "texts = text_splitter.split_documents(documents)\n",
    "print(f'The input document has been split into {len(texts)} chunks\\n')\n",
    "\n",
    "#\n",
    "# Indexing and storing\n",
    "#\n",
    "embeddings = VertexAIEmbeddings(model_name=\"textembedding-gecko\")\n",
    "vectorstore = Chroma.from_documents(texts, embeddings)\n",
    "retriever = vectorstore.as_retriever(search_type=\"similarity\", search_kwargs={\"k\": 1})\n",
    "\n",
    "#\n",
    "# Querying\n",
    "#\n",
    "llm = VertexAI(temperature=0.7, verbose=True)\n",
    "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",
    "\n",
    "{context}\n",
    "\n",
    "Question: {question}\n",
    "\n",
    "Helpful Answer:\"\"\"\n",
    "custom_rag_prompt = PromptTemplate.from_template(template)\n",
    "\n",
    "rag_chain = (\n",
    "    {\"context\": retriever, \"question\": RunnablePassthrough()}\n",
    "    | custom_rag_prompt\n",
    "    | llm\n",
    "    | StrOutputParser()\n",
    ")\n",
    "\n",
    "query = \"What is the melting point of citric acid?\"\n",
    "rag_chain.invoke(query)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae1eb2aa-9c11-429d-94eb-a718d44a3979",
   "metadata": {},
   "source": [
    "## Conversational Retrieval\n",
    "\n",
    "In this section, we prototype a conversational retrieval system. It combines the chat history with the retrieved documents to answer the question."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "141daeb7-a343-4ac8-af55-fda046326a60",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": [],
    "ExecuteTime": {
     "end_time": "2024-01-14T16:59:45.463477Z",
     "start_time": "2024-01-14T16:59:41.424605Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The input document has been split into 5 chunks\n",
      "human: How much times aspartam is sweeter than table sugar?\n",
      "ai:  Aspartame is approximately 150 to 200 times sweeter than sucrose (table sugar), making it an intense sweetener.\n",
      "human: What is its caloric value?\n",
      "ai:  AI: Aspartame is virtually calorie-free, as the human body does not metabolize it into energy like sugars or carbohydrates.\n"
     ]
    }
   ],
   "source": [
    "from langchain_core.messages import HumanMessage, AIMessage\n",
    "from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder\n",
    "\n",
    "#\n",
    "# Load the input document\n",
    "#\n",
    "loader = TextLoader(\"../../tensor-house-data/search/food-additives/food-additives.txt\")\n",
    "documents = loader.load()\n",
    "\n",
    "#\n",
    "# Splitting\n",
    "#\n",
    "text_splitter = CharacterTextSplitter(chunk_size=3000, chunk_overlap=0)\n",
    "texts = text_splitter.split_documents(documents)\n",
    "print(f'The input document has been split into {len(texts)} chunks\\n')\n",
    "\n",
    "#\n",
    "# Indexing and storing\n",
    "#\n",
    "embeddings = VertexAIEmbeddings(model_name=\"textembedding-gecko\")\n",
    "vectorstore = Chroma.from_documents(texts, embeddings)\n",
    "retriever = vectorstore.as_retriever(search_type=\"similarity\", search_kwargs={\"k\": 1})\n",
    "\n",
    "#\n",
    "# Initialize new chat\n",
    "#\n",
    "\n",
    "contextualize_q_system_prompt = \"\"\"Given a chat history and the latest user question \\\n",
    "which might reference context in the chat history, formulate a standalone question \\\n",
    "which can be understood without the chat history. Do NOT answer the question, \\\n",
    "just reformulate it if needed and otherwise return it as is.\"\"\"\n",
    "contextualize_q_prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", contextualize_q_system_prompt),\n",
    "        MessagesPlaceholder(variable_name=\"chat_history\"),\n",
    "        (\"human\", \"{question}\"),\n",
    "    ]\n",
    ")\n",
    "contextualize_q_chain = contextualize_q_prompt | llm | StrOutputParser()\n",
    "\n",
    "qa_system_prompt = \"\"\"You are an assistant for question-answering tasks. \\\n",
    "Use the following pieces of retrieved context to answer the question. \\\n",
    "If you don't know the answer, just say that you don't know. \\\n",
    "Use three sentences maximum and keep the answer concise.\\\n",
    "\n",
    "{context}\"\"\"\n",
    "qa_prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", qa_system_prompt),\n",
    "        MessagesPlaceholder(variable_name=\"chat_history\"),\n",
    "        (\"human\", \"{question}\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "def contextualized_question(input: dict):\n",
    "    if input.get(\"chat_history\"):\n",
    "        return contextualize_q_chain\n",
    "    else:\n",
    "        return input[\"question\"]\n",
    "\n",
    "\n",
    "rag_chain = (\n",
    "    RunnablePassthrough.assign(\n",
    "        context=contextualized_question | retriever\n",
    "    )\n",
    "    | qa_prompt\n",
    "    | llm\n",
    ")\n",
    "\n",
    "chat_history = []\n",
    "\n",
    "question1 = \"How much times aspartam is sweeter than table sugar?\"\n",
    "ai_msg1 = rag_chain.invoke({\"question\": question, \"chat_history\": chat_history})\n",
    "chat_history.extend([HumanMessage(content=question1), AIMessage(content=ai_msg1)])\n",
    "\n",
    "question2 = \"What is its caloric value?\"\n",
    "ai_msg2 = rag_chain.invoke({\"question\": question2, \"chat_history\": chat_history})\n",
    "chat_history.extend([HumanMessage(content=question2), AIMessage(content=ai_msg2)])\n",
    "\n",
    "for m in chat_history:\n",
    "    print(f'{m.type}: {m.content}')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py310",
   "language": "python",
   "name": "py310"
  },
  "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
}