Building a RAG agent
Document ingestion, late chunking, RAPTOR retrieval, BM25, RRF, and reranking for a RAG agent.
What is RAG and why we even need it?
Retrievel Augmented Generation(RAG) is an approach that combines the language capabilites of LLMs with the ability to access dynamic and specific data. It allows us to combine world knowledge with specialized knowledge to deliver contextful and precise answers. LLMs retrieves their data from a training dataset which is fitted to them a long while back suppose if you asked chatgpt in 2023 when he didn't had access to internet anything of your time he couldn't answer right? he can right now coz he got the ability to access the internet that is also a type of RAG you can say dynamic rag but we can not gonna talk about that one, in this blog we are gonna discuss about providing specific data source like pdf, image, url, database,etc.
Workflow
- ▸Query - User asks / submits their queries
- ▸Retrieval - Retrieval system searches for match in the defined sourches such as pdf and databases.
- ▸Augmentation - This Retrieved information is passed to the model and used as reference / contextual foundation for generation.
- ▸Generation - LLMs uses these additinal data to provide the very contextualized information / response.
So basically you can understand it as a extension to the LLM which we plugs in to enhance it's ability and get very specialized and present / timeless data.
Ingestion
We understood how the full workflow works but now let's understand how we can provide high quality data. Document ingestion is very important if we don't do it right then LLM won't understand a thing so we need to make sure we provides our data in best form.
What is Data Ingestion?
- ▸Data Ingestion is parsing of information from external sources like pdf, webpages,etc. so that it can be embedded into a search space for later retrieval.
- ▸But what is the problem even we can use any parser and parse it right? yeah, we can definitely do it for simple texts but when comes to word files or pdfs things gets messy as it contains various headings, tables, etc.
We have two options we can either implement it byself or use some library.
if you wanna implement it you can read this Blog: https://aiengineering.academy/RAG/01_Data_Ingestion/data_ingestion/#key-steps-in-data-ingestion. I had implemented last time so gonna use library this time.
- ▸There are various options available for us some are paid and some are free. like Docling by IBM, Langchain Data Loaders, LlamaIndex, Vectorize.io, Unstructured
Docling parser is an open source tool / model by IBM it parses the documents with structured content gracefully. So we are gonna use docling for document processing in structured format(ex- JSON) and applying boundaries.
Chunking
Early chunking vs Late chunking
In early chunking, we splits / segments the large document into pre-defined text units(chunks) and then we encodes every chunk in their own vector embedding(every text has meaning and embedding is way to represent that meaning in form of number in vector).It's very fast, but, while doing so we often times lose the context.
Suppose, we have these sentences :- Kavya loves that flower. She lives near film city. We will create seperate embeddings for both the sentences then tell me when we search then how do we know who is She?
For solving this we use Late chunking: https://arxiv.org/pdf/2409.04701
In Late chunking, whole document is processed and embed using a long-context model to capture all cross-chunk relationships and meaning, and then applies the chunking boundries to those rich, full-context embeddings, which provides higher retrieval precision and semantic accuracy at the cost of higher initial computation.
You can also read this for reference Blog: https://medium.com/@visrow/what-is-late-chunking-in-rag-how-can-you-improve-your-rag-with-late-chunking-f981a0cb39bb
Jina v3
I am gonna use jina v3 model for this purpose: https://jina.ai/news/late-chunking-in-long-context-embedding-models/
So it's gonna take the whole document and process it and build embedding then instead of text we chunks tokens which ofcourse knows the whole context then create chunk embeddings. The Jina Segmenter is applied to the raw text to establish structural boundaries (e.g., paragraphs, sentences, or sections
So we have processed the data and stored it in database, here i am using postgres(Prisma) for this case. you can use any vector databases.
Now after processing and storage there comes the searching / matching right. we wants that whenever user enters some queries we can quickly search through our database and able to provide the relevant data to the LLM right.
Suppose we have a very huge research paper then we chunk it and embed it and user queries like what is the methodology of this paper? then our previous approach would return like best 5-10 most similar chunks but chances are high it would miss the real theme as it would spread across the paper. So for this problem we can try a different approach, RAPTOR
RAPTOR
Recursive Abstractive Processing for Tree Organized Retrieval or RAPTOR recursively embeds, clusters and summarizes the text chunks to build a tree structure with a different level of summarization from the bottom up. https://arxiv.org/html/2401.18059v1
Workflow
Leaf nodes are the original document chunks(fine-grained). Intermediate notes are cluster summaries(medium abstraction). The root node is the global summary(high abstraction). Every node(leaf and summary both) is embeded and stored and when you query / search, you search all of them at once.
How the tree is built.
- ▸Embed all leaf chunks: Each chunk from docling -> Jina v3 -> 1024-dimensional vector
- ▸Reduce dimension with UMAP(Uniform Manifold Approximation and Projection): UMAP is a dimensionality reduction algorithm while preserving the relationships and structure of the original data. Why reduce? because everything is so far in higher dimension. It is hard to find meaningful clusters.
- ▸Cluster with HDBSCAN: Hierarchical Density-Based Spatial Clustering of Application with Noise or HDBSCAN is an unsupervised algorithm to find clusters or groups of chunks are semantically close. https://hdbscan.readthedocs.io/en/latest/how_hdbscan_works.html
- ▸Summarize each cluster with Gemini: For each cluster, concentage all chunk texts and send to the gemini and let it generate the summary and use that summary to create a new node one level up.
- ▸Repeat Recursively: Embed summaries and cluster them. Summarize again. and keep doing this until you are left with just one cluster(root).
Retrieval Flow
At query time we don't traverse the tree top-down. We are gonna use collapsed retrieval technique, flattening all nodes from all levels into one pool ans search everything at once. By, doing this we are gonna get big picture as well as specific details simultaneously.
We are gonna store all the nodes in our vector database
BM25
Suppose we wanna some term like what is cosine similarity? then searching from previous strategy is inefficient so we can use another algorithm which is designed for this specific purpose BM25(Best Matching 25).
RRF(Retrieval Rank Fusion)
It is a powerful algorithm used to merge and rerank search results from different retrieval methods so we can gonna give it our vector and keyword search. Hence it would help us in getting better results.
Jina Reranker
After we got our result we can pass it through this which will again find the most relavant and near to the context data among them which will again help in sending the best of best related data to the LLM.
LLM
Finally we will Pass the final result to the LLM and get our answer.
I know i am not good at writing blogs. If you still make it till here, Thanks for reading. :)