Contextual Retrieval with Any LLM: A Step-by-Step Guide

In a previous video we looked at contextual retrieval, which is a new chunking strategy from Anthropic which has shown to improve the failure rates by 35% in contextual embeddings. And when combined with PM25, the failure rate can be reduced by 49%. These are great numbers for any retrieval system. In this video I want to show you how you can use the same strategy for any other LLM that you want. I covered the theory in one of my previous videos. I'll put a link to that video, but here's a quick recap. In standard RAC system, we take each document independently, we convert those into chunks, and usually those chunks do not have any contextual information. Each of the independent chunks are embedded using an embedding model and stored in a vector store. Similarly, we can use a search index such as BM25 to create a keyword based search index for each of the chunks. And at runtime we compare the user query with the chunks in the database. As you can see, we lose a lot of information because each chunk is independent and it doesn't really contain any contextual information. The contextual retrieval system adds a pre processing step in which we take each chunk, feed that into a specific prompt which also looks at the whole document, and then that prompt identify or locate that chunk in the whole document, and add contextual information related to that specific chunk. So in general, this will add 50 to 100 tokens depending on the context in which a certain chunk belongs in the the document, and then we take these updated contextual chunks, compute embeddings, put them in a vector store, and at the same point we do the index creation based on these updated chunks. Now you can use a re ranking step at runtime or inference time to further increase the accuracy of your retrieval. So here's the whole implementation that we are going to explore in this video. In order to implement the contextual retrieval system, Anthropic suggests to use this prompt. In their case they use the Cloud3 Haiku, which is a long context model with close to 200,000 tokens context window. And for this approach to work, you actually need a long context window models. So the OpenAI models has a context window of 128,000 tokens, Gemini has 2 million tokens context window. Even the llama 3.2 has a context window of 128,001 tokens. So any of those models will work. And the way we do it is we pick the whole document. So let's say you have a PDF file you create chunks out of that PDF file, then you take the whole PDF file, provide that to the model as a part of the prompt, and then you say, here's a chunk we want to situate within the whole document. Then you provide one chunk at a time and simply tell the model to situate that in context of the whole document. So ideally you should be adding just 50 to 100 different tokens to each individual chunk. And visually again, the process is going to look something like this. So you first take single document, create chunks out of those, feed those chunks individually to the prompt that we saw along with the original document, and, and it will add this context to each of the chunks and the rest of the embedding process is going to remain the same. So let me show you how you can implement this whole thing in a Google Collab notebook. For this experiment, I'm using the OpenAI models, but Gemini models or even the Llama models will work in exactly the same fashion. So first we need to install all the different packages that we will need. This include LangChain. We're specifically using this for splitting our documents into chunks and computing the embeddings. But if you have your own splitter, you don't really need LangChain. OpenAI is going to be used for the LLM and the embedding model, but you can replace it with any other embedding models in LLM. And we're going to use FiOS for storing our embeddings in a Vector store. For BM25, which is the a traditional keyword based search Mechanism, we're using rank BM25 package. So first we're just importing all the things that we will need. I'm setting the environment variable because as I said, we'll be using the OpenAI models right now. Here we created a contextual retrieval class for chunk creation. I'm using a chunk size of 800 characters with an overlap of 100 characters. But you can convert this to tokens as well, depending on the size of your document. And we're using the recursive character text splitter. So basically it will take paragraphs. Try to chunk your document based on paragraphs. If it's not able to do that, it's going to do it based on sentences and so on and so forth. For embeddings, we're using the OpenAI embedding. The model is going to be GPT4. Okay. You can also use the GPT4 ohm mini. And we're setting the temperature to zero and at maximum, it can make two retries if it run into any issues when you it's making an API call. So first here is a dummy data that's basically Tesla's financial statements. And these are made up numbers. So this is going to be our base document. In your case, if you have multiple files, you are going to provide each file independently because we want to take a single file and then contextualize each chunk in that file. Coming back to our code, there is a function called ProcessDocument which receives the document that we just created. And based on that document, we are going to convert it into chunks. Each chunk is going to be size 800 characters with an overlap of 100 characters. So you will get chunks out of it. And then there is another function called Generate contextualized Chunks which receives the original document that we provided along with the chunks that we just created. And this will return the contextualized chunks. So let's look at how this contextualized chunks works. So we take each chunk individually and then there's another function called Generate context. So for each chunk, we provide the original document plus that given chunk and pass it on to generate context. Here's the actual function that we're using to generate the context for each chunk. Now, the prompt that you see here is a little different, the one that we saw from Anthropic. And the reason is I'm trying to make it customized to our specific application. So if you are working on a very specific application or specific types of documents, you want to customize this contextualizing prompt to that specific set of documents. So here is my attempt with the help of an LLM. You are an AI assistant specializing in financial analysis, particularly for Tesla. Your task is to provide brief relevant context for a chunk of text from the Tesla's Q3 2023 financial report. So depending on your own needs, you can adjust this. So we provide here is the financial report, then here's the chunk we want to situate within the whole document. So this is going to get individual chunks. And then we say provide a concise two to three sentences for this chunk, considering the following guidelines. Identify the main financial topics or metrics discussed, then mention any relevant time periods or comparisons. Then, if applicable, note how this information relates to Tesla's overall financial health strategy or market position. Include any key figures or percentages that provides important context and do not use phrases like this chunks, discusses or this section provides. So this is going to generate the contextualized prompts along with the original prompts, and then we Simply take those, append it to this contextualized chunks list. Then we have other functions for creating the vector stores for these chunks. And the way we do the vector stores is we create vector stores for both the original and the contextualized chunks. I just want to show you an example of how they are going to look like. We are also creating the BM25 index which is keyword based search index. You can create a cache key for the documents that you are processing if you want. And then there is another function that was going to use these chunks plus the original query to generate answers. Now in this implementation we don't have re ranking or query expansion. These are the topics that I cover in my Rack Beyond Basics course. So if you're looking for advanced rag techniques, I highly recommend to check that course out. Link and details are in the video description. So anyways, this is the whole class that you just need. And in, in my case I'm using OpenAI. But you can replace the LLM and the embedding models with any one of your choice, right? The only change that you're going to make is going to be here. So this is basically the embedding model and your LLM. Now let's look at a couple of examples in action. So we create an object for that class, right? So CR is the contextual retrieval. Then we are providing documents to the process document function. This is basically the financial information that we created. And you're going to get both the original chunks plus the contextualized chunks. Now for this document we have 15 chunks that are created based on our criteria. Okay, so here is the first chunk. It's basically the executive summary. So the original chunk is Tesla Financial Analysis and market overview for Q3 2023. Now it talks about the executive summary or what type of results we got in Q3 2023. But as you can see, there are no actual revenue numbers or growth numbers in this specific chunk. However, if you look at the contextualized chunk, it adds this at the top which says the executive Summary and financial performance overview highlights Tesla's leadership in the EV market, reporting 9% year over year revenue increase to 23.36 billion in Q3 2023. So even with this chunk, the model was able to add a lot of contextual information and it also added some of the key Metrics, which is 4.1 billion and net income of 1.85 billion. Right. So you can already see that this basically adds a lot more information to each given chunk. Now for the, for this Specific chunk. There is already a lot of information in the original chunk, but the contextual chunk added some further information. Specifically, it summarizes everything that is happening around that chunk that we are looking at. Now, I think we can still customize our prompt to add some further information if we need to, but based on the original chunks plus the contextualized chunks, we are just creating another indices so we will have two separate vector stores. We are also creating the search indices for both the original chunks and contextualized chunks. And this is based on the BM25. You can just provide your original queries or the queries that you want. So here are just some example queries that I was looking at. And, and these are again generated with the help of an LLM. And the way we do it is that we are doing a retrieval on the original vector store, contextualized vector store, and then the BM25 index created on the original chunks plus the contextualized chunks. For this example, you don't really see differences in terms of the answers, but you can see that the chunks that are being used are a little different. So, for example, here's a query. What was the Tesla's total revenue in Q3 2023? What was the gross profit and cash position? These are the original chunks. Now, based on this, these original chunks, I think it has enough information that it can make up the answer that the Total revenue was 23.35 billion and the profit or the gross profit for that quarter was 4.14 billion. Right. You can look at the contextual chunks, which I think provide a lot more information in the big beginning of the chunks. This is the information that is added by this contextualized summary. And with that you can get pretty much the same answers Again. This was a quick overview of how you can potentially implement this contextual retrieval for any LLM and any embedded model of your choice. Now, the results that you are going to get are going to be highly dependent on the prompt that you're using. So you can potentially use a simple prompt like this, which is basically a copy of the prompt provided by the Anthropic team, or you can tailor or customize it for your own needs. Here is another version that simply talks about focus on the financial metrics, company names and time periods, etc. And then I expanded this prompt for my final notebook that I was showing you. It's a very neat idea, very similar to some of the other approaches that we have seen. But rather than adding summaries of the whole document, this basically tries to situate a given chunk in reference to the surrounding chunks. Now apart from this there is also the concept of late chunking and long context embedding models. This is an approach proposed by Jenna AI who are creating some really interesting embedding models and re ranking models. So I'll put a link to this blog post in the video description. If there is interest I can create a more detailed video on this. But there are some really interesting things happening in the rag ecosystem and if you don't want to deal with any chunking strategy you can just go end to end vision language model based solutions such as Local GPT Vision which uses a vision model to identify the pages where the information is contained and then uses a multimodal model to generate answers. I'll put a link to my repo that you can look at how this could be implemented. Anyways, a link to the notebook is going to be in the video description. Give it a try and see if you can replicate this for your own applications. I hope you found this video useful. Thanks for watching and as always see you in the next one.