Supercharge Your RAG with Contextualized Late Interactions

Dense embedding that you are using in your Rack pipelines have a critical flaw which probably is impacting the accuracy of your retrieval. In this video we will look at what this potential issue is and what is the potential solution. In order to explain this, let's look at a very simple example. So here I have three documents. Each one of them contains a single sentence. On the right hand side we have a query. We compute the embeddings FR documents or chunks and depending on the embedding model that we choose, we will get a vector for each of this chunk. Similarly for the query we will also get a vector and then we do similarity search. So basically we compare this embedding vector with each of the documents embedding vector and figure out which one is the closest and that is return as a result of our retrieval step. If you look here, there are different embedding models and each one of them have a different dimensions. One of the most widely used open source embedding model is this BGE Small, which has only 384 dimensions. And more recently we have started seeing much larger dimensional embeddings when it comes to open source embedding models. Now you might be thinking what is the issue with embeddings? Let's assume instead of this single sentence, your chunk contains multiple paragraphs because you have a much larger chunk size. In that case, if you're computing embeddings for your chunk, you're going to be compressing all that information into a single vector. That vector may not be able to represent everything that is present in a given chunk. In this case, the compression is basically loss of a lot of information. So what's the solution? A potential solution is in this paper Colbert v2 effective and efficient retrieval via lightweight late interaction. And we're going to look at a practical code example later in the video. Now this is a dense paper, but let me try to explain the concepts with a very simple example. So coming back to this, with dense embeddings that we usually use, we are representing a whole chunk with a single embedding vector. But Colbert or Cold Bear suggest a different techniques. In this case, the col is contextualized late interaction. And let me explain what it means. The first step in this case is that both for our documents and chunks as well as for a query, we tokenize it. That means depending on the tokenizer that we use, we are going to get individual tokens from our documents. So in this case, let's say our first document contains five different tokens. The second document Also has five different tokens. The third document has six tokens. And similarly we extract the individual tokens for our query as well. In the next step, we compute embeddings for each and every individual token. So every token is going to be represented by a different embedding vector. Just keep in mind this is a very simplified explanation of how this algorithm works. I'm skipping on a lot of technical details, but now, for each of the document as well as query, we are going to have a number of different embedding vectors depending on how many tokens are in each document. Next, for each token in the query, we will compute a similarity score with each token in your documents. So, for example, for the first token in the query, we compute similarity among all the tokens that are present in the first document. Then we repeat this whole process for the second documents as well for the third document. And we do this for each and every individual query token. And then based on these individual similarity scores, we compute an overall score for each document or each chunk. These are called late interactions which are shown here in this figure in the paper. Now, in this case, each token individually is contributing to the overall score. Rather than everything being compressed into a single vector, the embeddings that you get for each of these token are actually contextualized. So depending on the overall context in which this token is present, that is taken into account when the embeddings are computed. Okay, so I hope this makes sense. For more details, I'll highly recommend to check out the paper. I'll put a link to this paper in the video description. Now let's look at a practical example of how you can do semantic search using Colbert. Instead of the dense embedding models, there's a GitHub repo for Colbert v2. But instead of this, we are going to use another package. This is called regatui. So this is built on top of Colbert. Here they show you how you can train and fine tune Colbert models. Use it for embedding and indexing documents. And there are examples of retrieving documents as well. I'll highly recommend to check this out. So I put together this Google Colab for you. Link is going to be in the video description. First we need to install different packages that are going to be needed. For this example, we're going to be working with the Orca paper for retrieval here. You can download this paper if you want. But in my case I uploaded a copy of this paper to my Google Colab notebook. Now in this case, to load the PDF file I'm using the PDF reader from Llama Index. So basically here is our PDF file after reading it as a document. Next, instead of using the documents objects for Llama Index, we are extracting the original text from each of those document objects. But now this is basically a list of text. And simple text is a lot easier to work with in ragatui. All right, next we will import the pre trained model from ragatui. There are a couple of versions of Colbert. We're going to be using the second version. So this will basically load regatui RAG pipeline. Now you will see this warning message regarding hugging fist token, but you don't actually need it. Next we need to create our index. This is very similar to what you would do if you're building a vector db. So we need to provide a list of documents and this is basically the list of text. Then you need to give a name to your index. I'm just calling this orca. Then maximum documents length. Think about this as the chunk size. So you want to define the maximum document length here. I'm keeping it to 256 and whether we want to split individual documents based on this max document length or not. And I said yes. So by default, as a vector store it's using plate, but you can force it to use FIAS if you want. Okay, so what exactly is happening in here? So basically divided our text into 219 passages and for each of the tokens it will compute this embedding representation. Now, in order to do retrieval, all you need to do is. Now, in order to do retrieval, we're going to use the search function on top of our index that we just created. We'll pass on our query, then the number of documents that we want it to return and then the name of the index that we want to use. So you can have multiple indices. In this case, we have only one index and we are just using that. Now the output is actually a little different from what you would expect from a normal vector store. Now, the result that you get has multiple components. First you get the document or chunk that is retrieved, but along with that you also get the score as a rank and you get some document specific information. For example, passage ID document ID that is used for internal indexing. Now in this case, if you look at this rank, this is extremely important because it ranks documents based on the similarity to your query. Okay, so here I'm representing the same information but in a much better formatted way. So our query here was what is instruction tuning? And in this Case the first document that it return has a subtitle title of Instruction tuning. So this portion of the paper actually talks about instruction tuning. The second document that it returned also talks about instruction tuning for multimodal tasks. The third document that it returned also talks about instruction tuning specifically with GPT. 4 Responses so if you were to do rag on top of this, you're going to use these three documents as context to your LLM along with your query. Okay, now the question is going to be how does this retrieval compares to dense embedding models? For that we are going to be using LangChain. So I'm installing all the required packages. Now I want to do a comparison of Colbert with the OpenAI embedding model as well as one of the open source embedding model. Now in this case we're just setting up things so we import all the required packages. Then I'm using OpenAI so I need to provide my OpenAI API key since this is LangChain. So I need to just reload the paper. Again, embedding model that we're going to be using is the OpenAI embedding model. I put everything together in a single text. Now here I'm using the GPT4 tokenizer and again the chunk size is going to be 256, so I want to keep it similar to what we had for Colbert model. But keep in mind that the chunks that you get are not going to be exactly the same because the tokenizer is different in this case. Now we're using FIAS as a Vector DB or Vector star and we are going to run the same query what is instruction tuning? And here are the results. So the first chunk that it returned is exactly the same as what we saw before. So this is pretty good. The second chunk also talks about instruction tuning. So so far so good. But when we look at the third chunk, this is where the problem lies. There is instruction tuning in here, but this is basically the table of contents. Now in this case I returned four documents and the last one also talks about instruction tuning as well. So if you were to generate a response based on this context, all of this is going to go in to our LLM and that is going to generate response. But in here you can see that there's one chunk which has relatively irrelevant information and that could cause some issues. Now let's look at how the results looks for the open source embedding model. In this case I'm using the BGESMALL English. I wrote this function to load the embedding model. But again, I'm keeping the chunk size to 256 tokens. So we split everything based on the tokenizer that we have for BG small. Here we load the embedding model, then compute the embeddings, put them in a vector store, and we run it through the exact same query. Okay, so here's the response. The first chunk that it returns is again the table of content. There is instruction tuning at the top, but there are irrelevant things in here, which is not helpful at all. The second chunk does talk about instruction tuning. The third chunk also has some information regarding instruction tuning. And also the fourth chunk. Now, as you can see in this specific example, the cobalt is retrieving, I think much better context that we can pass on to our LLM. Now, if you recall from our initial discussion, depending on the size of the embedding model, it probably is going to start having more trouble if we increase the chunk size. So in this case I'm looking at a relatively complex query. We increase the max document length to 512 tokens. This is for the Colbert V2. The query in this case is what is a system message? And list all the system messages in the paper. So again, it actually needs to divide it into kind of two questions because we're looking at two different things. In this case it's able to retrieve the first chunk which actually talks about what a system message is. So this is pretty good. For example, here it says system instructions can specify the tone, task format and limitation of model responses. And system instructions are also a way of improving the safety of the model responses. Right, so the first chunk is relevant to our first part of the question. Now the second chunk actually talks about system messages that that are proposed within the paper. So there are a total of 16 different system messages. But since we define max document length of 512, so it was able to retrieve only 13 or 12 and a half of those system messages as a part of the second document. Now the third document is a table which actually talks about the performance based on different system messages. And the last document that it retrieved has the remaining system messages that are present in the paper. Now here, one thing probably that needs to be done is to improve the quality of chunking process. The great thing is this is integrated within LangChain. So let me know if you want me to create another video of how to use this as part of LangChain. So the responses are not bad at all because the first two chunks contains the exact information that we are looking for. Now let's look at the OpenAI embedding model. In this case, the first chunk that it retrieved actually lists some of the system messages that are proposed in the paper. The second passage or document, again for some reason retrieved a table of content. Now, the third document or chunk that it retrieves, it talks about the definition of system instruction or system message. So this is pretty good. And the last chunk talks about role of system instructions. All right, what about the open source embedding model? So this one has the most trouble. So the first chunk is actually partial list of the system messages that are proposed in the paper. The second document is basically a combination of a figure caption and a table in the paper. Now, the third one has a partial definition of system instruction, but the rest of the chunk talks about something else. The fourth one is not really relevant to the query whatsoever. It simply mentions system message, and it picked this whole paragraph or whole document based on that. Now, as you can see, Colbert can potentially give you better context when it comes to retrieval. And that context will be helpful to the LLM to generate a much better response. Okay, so I'll highly recommend to check it out. I don't think a lot of people are using this, but it has extremely high potential. Now, before we leave, a couple of things. So, first and foremost, the accuracy of retrieval also depends on how big your embedding model is. So for example, if you were to use something like this SFR embedding Mistral, this probably is going to give you a much better retrieval accuracy because it has a much better embedding size. And I think that's the reason since the OpenAI embedding models have much bigger size and they are much bigger models in general. That's why we were able to see like closer retrieved text both from the Colbert and from the OpenAI embedding model. I think by default we were using the small embedding model from OpenAI, which still has a dimension of 1536 vectors. Again, I think REGATUI has a lot of potential. Let me know if you want me to make subsequent videos in integrating this as a part of end to end RAC pipeline. If you want to learn more about advanced Rack concepts, make sure to sign up for my Advanced Rack course. Link is in the description. I hope you found this video useful. Thanks for watching and as always, see you in the next one.