Skip to content

Commit ece96bb

Browse files
ikatsovikatsov
committed
Extending the RAG referece
1 parent 36dea80 commit ece96bb

File tree

1 file changed

+196
-40
lines changed

1 file changed

+196
-40
lines changed

‎search/retrieval-augmented-generation-llm.ipynb‎

Lines changed: 196 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -13,9 +13,10 @@
1313
"source": [
1414
"# Retrieval-augmented Generation (RAG) Using LLMs \n",
1515
"\n",
16-
"In this notebook, we create a basic prototype of a retrieval-augmented generation system. This prototype demonstates how traditional information retrieval methods can be combined with LLM-based text summarization to search through large documents or collection of documents.\n",
16+
"In this notebook, we create several prototypes of a Retrieval-Augmented Generation (RAG) system. 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.\n",
1717
"\n",
18-
"We use a Google Pixel 7 release notes as the input document. It is available in the `tensor-house-data` repository."
18+
"### Data\n",
19+
"We use small input documents that are availbale in the `tensor-house-data` repository."
1920
]
2021
},
2122
{
@@ -34,7 +35,7 @@
3435
},
3536
{
3637
"cell_type": "code",
37-
"execution_count": 2,
38+
"execution_count": 6,
3839
"id": "2a97eb6a-ccec-4527-b9fb-36b69c6ff3de",
3940
"metadata": {
4041
"editable": true,
@@ -56,19 +57,165 @@
5657
")"
5758
]
5859
},
60+
{
61+
"cell_type": "code",
62+
"execution_count": 41,
63+
"id": "a239de43-7783-46f3-b919-5ad39b0a2daf",
64+
"metadata": {},
65+
"outputs": [],
66+
"source": [
67+
"#\n",
68+
"# Imports\n",
69+
"#\n",
70+
"from langchain.llms import VertexAI\n",
71+
"\n",
72+
"from langchain.vectorstores import Chroma\n",
73+
"from langchain.chains import RetrievalQA, ConversationalRetrievalChain\n",
74+
"from langchain.memory import ConversationBufferMemory\n",
75+
"\n",
76+
"from langchain import PromptTemplate\n",
77+
"from langchain.chains.question_answering import load_qa_chain\n",
78+
"\n",
79+
"from langchain.document_loaders import UnstructuredHTMLLoader, TextLoader\n",
80+
"from langchain.embeddings.vertexai import VertexAIEmbeddings\n",
81+
"from langchain.text_splitter import CharacterTextSplitter"
82+
]
83+
},
84+
{
85+
"cell_type": "markdown",
86+
"id": "1c28f863-c2e2-4c56-b459-df05ff12ad8f",
87+
"metadata": {},
88+
"source": [
89+
"## Question Answering Using a Single Prompt\n",
90+
"\n",
91+
"The most basic scenario is querying small documents that fit the LLM prompt."
92+
]
93+
},
94+
{
95+
"cell_type": "code",
96+
"execution_count": 68,
97+
"id": "bf134242-352d-4281-8f57-099ca387143b",
98+
"metadata": {},
99+
"outputs": [
100+
{
101+
"name": "stdout",
102+
"output_type": "stream",
103+
"text": [
104+
"The text states that \"Lemons and limes have particularly 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)\". So the correct answer is lemons and limes.\n"
105+
]
106+
}
107+
],
108+
"source": [
109+
"query_prompt = \"\"\"\n",
110+
"Please read the following text and answer which fruits have the highest concentration of citric acid.\n",
111+
"\n",
112+
"TEXT:\n",
113+
"Citric acid occurs in a variety of fruits and vegetables, most notably citrus fruits. Lemons and limes have particularly \n",
114+
"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",
115+
"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",
116+
"and limes; these values vary within species depending upon the cultivar and the circumstances under which the fruit was grown.\n",
117+
"\"\"\"\n",
118+
"\n",
119+
"llm = VertexAI(temperature=0.7)\n",
120+
"response = llm(query_prompt)\n",
121+
"print(response)"
122+
]
123+
},
124+
{
125+
"cell_type": "markdown",
126+
"id": "6869bf50-e4f2-4ba6-b7b9-2278c1a072a0",
127+
"metadata": {},
128+
"source": [
129+
"## Question Answering Using MapReduce\n",
130+
"\n",
131+
"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."
132+
]
133+
},
134+
{
135+
"cell_type": "code",
136+
"execution_count": 84,
137+
"id": "c5a43885-5f2e-49d3-b5cf-bc6204b10d3a",
138+
"metadata": {},
139+
"outputs": [
140+
{
141+
"name": "stdout",
142+
"output_type": "stream",
143+
"text": [
144+
"The input document has been split into 2 chunks\n",
145+
"\n",
146+
"The three most important applications of citric acid are:\n",
147+
"1. As a flavoring and preservative in food and beverages, especially soft drinks.\n",
148+
"2. In soaps and laundry detergents to remove hard water stains.\n",
149+
"3. In the biotechnology and pharmaceutical industry to passivate high purity process piping.\n"
150+
]
151+
}
152+
],
153+
"source": [
154+
"#\n",
155+
"# Load the input document\n",
156+
"#\n",
157+
"loader = TextLoader(\"../../tensor-house-data/search/food-additives/citric-acid-applications.txt\")\n",
158+
"documents = loader.load()\n",
159+
"\n",
160+
"#\n",
161+
"# Splitting\n",
162+
"#\n",
163+
"text_splitter = CharacterTextSplitter(chunk_size=2000, chunk_overlap=0)\n",
164+
"texts = text_splitter.split_documents(documents)\n",
165+
"print(f'The input document has been split into {len(texts)} chunks\\n')\n",
166+
"\n",
167+
"#\n",
168+
"# Querying\n",
169+
"#\n",
170+
"question_prompt = \"\"\"\n",
171+
"Use the following portion of a long document to see if any of the text is relevant to answer the question. \n",
172+
"Return bullet points that help to answer the question.\n",
173+
"\n",
174+
"{context}\n",
175+
"\n",
176+
"Question: {question}\n",
177+
"Bullet points:\n",
178+
"\"\"\"\n",
179+
"question_prompt_template = PromptTemplate(template=question_prompt, input_variables=[\"context\", \"question\"])\n",
180+
"\n",
181+
"combine_prompt = \"\"\"\n",
182+
"Given the following bullet points extracted from a long document and a question, create a final answer.\n",
183+
"Question: {question}\n",
184+
"\n",
185+
"=========\n",
186+
"{summaries}\n",
187+
"=========\n",
188+
"\n",
189+
"Final answer:\n",
190+
"\"\"\"\n",
191+
"combine_prompt_template = PromptTemplate(template=combine_prompt, input_variables=[\"summaries\", \"question\"])\n",
192+
"\n",
193+
"llm = VertexAI(temperature=0.7)\n",
194+
"qa_chain = load_qa_chain(llm=llm,\n",
195+
" chain_type='map_reduce',\n",
196+
" question_prompt=question_prompt_template,\n",
197+
" combine_prompt=combine_prompt_template,\n",
198+
" verbose=False)\n",
199+
"\n",
200+
"question = \"What are the three most important applications of citric acid? Provide a short justification for each application.\"\n",
201+
"response = qa_chain({\"input_documents\": texts, \"question\": question})\n",
202+
"\n",
203+
"print(response['output_text'])"
204+
]
205+
},
59206
{
60207
"cell_type": "markdown",
61208
"id": "b3eb6e94-a2f2-4548-8d0c-f4934f0224f6",
62209
"metadata": {},
63210
"source": [
64-
"## Question Answering Over a Large Document\n",
211+
"## Question Answering Using Vector Search\n",
65212
"\n",
66-
"In this section, we demonstrate how standalone questions can be answered. The input document(s) is split inot 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."
213+
"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."
67214
]
68215
},
69216
{
70217
"cell_type": "code",
71-
"execution_count": 27,
218+
"execution_count": 17,
72219
"id": "3857ccc4-2559-493d-9e44-babb744576db",
73220
"metadata": {
74221
"editable": true,
@@ -82,37 +229,23 @@
82229
"name": "stdout",
83230
"output_type": "stream",
84231
"text": [
85-
"The input document has been split into 2 chunks\n",
86-
"\n",
87-
"The three coolest Pixel 7 features are:\n",
232+
"The input document has been split into 5 chunks\n",
88233
"\n",
89-
"* Next-generation Super Res Zoom up to 8x on Pixel 7 and up to 30x on Pixel 7 Pro.\n",
90-
"* Macro Focus, which delivers Pixel HDR+ photo quality from as close as three centimeters away.\n",
91-
"* Photo Unblur, a Google Photos feature only on Pixel 7 and Pixel 7 Pro. Photo Unblur uses machine learning to improve your blurry pictures – even old ones.\n"
234+
"The melting point of citric acid is approximately 153 °C (307 °F).\n"
92235
]
93236
}
94237
],
95238
"source": [
96-
"from langchain.llms import VertexAI\n",
97-
"\n",
98-
"from langchain.chains import RetrievalQA, ConversationalRetrievalChain\n",
99-
"from langchain.memory import ConversationBufferMemory\n",
100-
"\n",
101-
"from langchain.document_loaders import UnstructuredHTMLLoader\n",
102-
"from langchain.embeddings.vertexai import VertexAIEmbeddings\n",
103-
"from langchain.text_splitter import CharacterTextSplitter\n",
104-
"from langchain.vectorstores import Chroma\n",
105-
"\n",
106239
"#\n",
107240
"# Load the input document\n",
108241
"#\n",
109-
"loader = UnstructuredHTMLLoader(\"../../tensor-house-data/search/news-feeds/pixel-7-release.html\")\n",
242+
"loader = TextLoader(\"../../tensor-house-data/search/food-additives/food-additives.txt\")\n",
110243
"documents = loader.load()\n",
111244
"\n",
112245
"#\n",
113246
"# Splitting\n",
114247
"#\n",
115-
"text_splitter = CharacterTextSplitter(chunk_size=4000, chunk_overlap=0)\n",
248+
"text_splitter = CharacterTextSplitter(chunk_size=3000, chunk_overlap=0)\n",
116249
"texts = text_splitter.split_documents(documents)\n",
117250
"print(f'The input document has been split into {len(texts)} chunks\\n')\n",
118251
"\n",
@@ -126,9 +259,12 @@
126259
"# Querying\n",
127260
"#\n",
128261
"llm = VertexAI(temperature=0.7, verbose=True)\n",
129-
"qa_chain = RetrievalQA.from_chain_type(llm=llm, chain_type=\"stuff\", retriever=docsearch.as_retriever(search_kwargs={\"k\": 2}), return_source_documents=True)\n",
262+
"qa_chain = RetrievalQA.from_chain_type(llm=llm, \n",
263+
" chain_type=\"stuff\", \n",
264+
" retriever=docsearch.as_retriever(search_kwargs={\"k\": 1}), \n",
265+
" return_source_documents=True)\n",
130266
"\n",
131-
"question = \"What are the three coolest Pixel 7 features? Please provide a list with short summaries.\"\n",
267+
"question = \"What is the melting point of citric acid?\"\n",
132268
"response = qa_chain({\"query\": question})\n",
133269
"\n",
134270
"print(response['result'])"
@@ -146,8 +282,8 @@
146282
},
147283
{
148284
"cell_type": "code",
149-
"execution_count": 26,
150-
"id": "96697fd3-4efc-4e98-a61d-ae722d142289",
285+
"execution_count": 65,
286+
"id": "141daeb7-a343-4ac8-af55-fda046326a60",
151287
"metadata": {
152288
"editable": true,
153289
"slideshow": {
@@ -160,41 +296,61 @@
160296
"name": "stdout",
161297
"output_type": "stream",
162298
"text": [
163-
"What are the three coolest Pixel features? \n",
299+
"The input document has been split into 5 chunks\n",
164300
"\n",
165-
"The three coolest Pixel features are:\n",
301+
"How much times aspartam is sweeter than table sugar? \n",
166302
"\n",
167-
"1. Super Res Zoom up to 8x on Pixel 7 and up to 30x on Pixel 7 Pro.\n",
168-
"2. Macro Focus, which delivers Pixel HDR+ photo quality from as close as three centimeters away.\n",
169-
"3. Photo Unblur, a Google Photos feature only on Pixel 7 and Pixel 7 Pro. \n",
303+
"Aspartame is approximately 150 to 200 times sweeter than sucrose (table sugar), making it an intense sweetener. \n",
170304
"\n",
171-
"Can you explain the first feature in more detail? \n",
305+
"What is its caloric value? \n",
172306
"\n",
173-
"Super Res Zoom is a camera feature that uses machine learning to improve the quality of zoomed-in images. It works by taking multiple photos at different focal lengths and then stitching them together to create a single, high-resolution image. This results in images that are sharper and have less noise than images taken with a traditional zoom lens. \n",
307+
"Aspartame is virtually calorie-free, as the human body does not metabolize it into energy like sugars or carbohydrates. \n",
174308
"\n"
175309
]
176310
}
177311
],
178312
"source": [
313+
"#\n",
314+
"# Load the input document\n",
315+
"#\n",
316+
"loader = TextLoader(\"../../tensor-house-data/search/food-additives/food-additives.txt\")\n",
317+
"documents = loader.load()\n",
318+
"\n",
319+
"#\n",
320+
"# Splitting\n",
321+
"#\n",
322+
"text_splitter = CharacterTextSplitter(chunk_size=3000, chunk_overlap=0)\n",
323+
"texts = text_splitter.split_documents(documents)\n",
324+
"print(f'The input document has been split into {len(texts)} chunks\\n')\n",
325+
"\n",
326+
"#\n",
327+
"# Indexing and storing\n",
328+
"#\n",
329+
"embeddings = VertexAIEmbeddings()\n",
330+
"docsearch = Chroma.from_documents(texts, embeddings)\n",
331+
"\n",
179332
"#\n",
180333
"# Initialize new chat\n",
181334
"#\n",
182335
"llm = VertexAI(verbose=True)\n",
183336
"memory = ConversationBufferMemory(memory_key=\"chat_history\", return_messages=True)\n",
184-
"chat = ConversationalRetrievalChain.from_llm(llm, retriever=docsearch.as_retriever(search_kwargs={\"k\": 2}), memory=memory, verbose=False, return_generated_question=False)\n",
337+
"chat = ConversationalRetrievalChain.from_llm(llm, \n",
338+
" retriever=docsearch.as_retriever(search_kwargs={\"k\": 1}), \n",
339+
" memory=memory, \n",
340+
" verbose=False, \n",
341+
" return_generated_question=False)\n",
185342
"\n",
186343
"#\n",
187-
"# Ask questions with continuous context\n",
344+
"# Ask questions with a continuous context\n",
188345
"#\n",
189-
"\n",
190346
"def print_last_chat_turn(chat_history):\n",
191347
" print(chat_history[-2].content, '\\n')\n",
192348
" print(chat_history[-1].content, '\\n') \n",
193349
"\n",
194-
"result = chat({\"question\": \"What are the three coolest Pixel features?\"})\n",
350+
"result = chat({\"question\": \"How much times aspartam is sweeter than table sugar?\"})\n",
195351
"print_last_chat_turn(result['chat_history'])\n",
196352
"\n",
197-
"result = chat({\"question\": \"Can you explain the first feature in more detail?\"})\n",
353+
"result = chat({\"question\": \"What is its caloric value?\"})\n",
198354
"print_last_chat_turn(result['chat_history'])"
199355
]
200356
}

0 commit comments

Comments
 (0)