Enhancing RAG with Retrieval Augmented Fine-tuning

Introduction

While Large Language Models (LLMs) can answer questions on many topics, they may not correctly answer about the topics that are not included in the training data such as recent events, or deep web (i.e. data that is not indexed by search engines). Another missing piece is not getting exact source of the answer, making verification challenging. This is where Retrieval Augmented Generation (RAG) can be useful. RAG combines the generative capabilities of LLMs with information retrieval from external sources of data and also can cite the exact sources of its answers, greatly improving verifiability and reliability. In this article we will enhance RAG with Retrieval Augmented Fine-tuning.

Learning Objectives

• Learners will identify the limitations of Large Language Models (LLMs) and understand how Retrieval Augmented Generation (RAG) enhances answer accuracy and reliability by incorporating external data.
• Participants will learn how to prepare data for LLM fine-tuning, including chunking, question generation, and the selection of oracle and distractor contexts.
• Learners will become familiar with configuring the RAFTDatasetPack parameters for optimal question and answer generation from LLMs.
• Learn to comprehend how the Semantic Splitter Node Parser processes data into chunks and the significance of cosine dissimilarity in improving model understanding.

Adapting LLM to RAG

In RAG, we split the data into chunks, find the top-K most similar chunks to the query, and present those chunks of content to the LLM to generate an answer. However, those top-K chunks can contain a mix of relevant and irrelevant content for the given query. The LLM should be able to find the relevant content for the query among those chunks given to it to generate the answer. So, if we can finetune the LLM for this specific task of generating an answer given both relevant and irrelevant content in the prompt, it can improve the accuracy of RAG.

What is Retrieval Augmented Fine-tuning?

As shown in the above picture, generating an answer for a query based on only the training data is like a “closed book” exam. An “Open book” exam is where the answer is generated using external data, which RAG.

In a new method, we train the LLM about how to effectively use the external data. This method significantly improved the RAG performance (https://arxiv.org/pdf/2403.10131.pdf)

How to Prepare the Data for Fine-tuning the LLM?

Let us now dive deeper on how to prepare the data for Fine-tuning the LLM:

• Dividing Sample Data: Segments the sample data into chunks. Each chunk represents a potential source of information or context for generating questions.
• Generating Questions: Creates corresponding questions for every chunk of data. These questions are designed to be answerable using the information within the chunk.
• Generate answer with Oracle Context: The ‘oracle context’ refers to the chunk of data that contains the precise information needed to answer a given question. Uses this context alongside the question to generate the answer using Chain of Thought prompting.
• Selecting Distractor Contexts: In addition to the oracle context, a few random chunks of data are chosen as ‘distractor contexts’.These simulate noise and irrelevant information, challenging the model to focus on the relevant context.
• Compiling Training Data: Compilers include the question, oracle context, distractor contexts, and the generated answer, alongside explicit instructions on how the model should discern and utilize the relevant context to answer questions, into a comprehensive training dataset.
• Fine-Tuning the Model: Utilizing this dataset, the model undergoes fine-tuning, learning to accurately distinguish relevant from irrelevant information and to generate precise answers based on the context provided.

Implementation of Retrieval Augmented Fine-tuning

Let us now learn the implementation of Retrieval Augmented Fine-tuning. Initially we start with installing the required libraries using the following commands:

• pip install llama-index
• pip install llama-index-packs-raft-dataset

Then you can import the RAFTDataset:

from llama_index.packs.raft_dataset import RAFTDatasetPack

For the data preparation process for Q/A generation, the RAFTDatasetPack is configured with the following parameters:

• filepath: Specifies the path of the file used to generate questions and answers. This file acts as the primary source of content for the dataset
• llm: Defines the Large Language Model (LLM) employed for generating questions and answers. GPT-4 is used by default if no model is specified. Choose the model carefully by considering the costs.
• embed-model: The embedding model used to calculate the similarity between a query and its context, essential for selecting relevant context chunks.
• num_questions_per_chunk: It determines the number of questions to be created for each data chunk, directly affecting the comprehensiveness of the training dataset.
• num_distract_docs: Sets the number of random context chunks used as distractors for each question, challenging the model to identify relevant information
• chunk_size: Llama-index uses SemanticSplitterNodeParser to split the dataset into chunks. So, this parameter is not useful.
• default_breakpoint_percentile_threshold: Controls the threshold for combining chunks based on their dissimilarity. A higher value results in fewer, larger chunks, affecting the granularity of the data used for training.

Semantic Splitter Node Parser

The SemanticNodeParser operates by dissecting the data at the sentence level, initially dividing the text into smaller segments or ‘chunks’. Here’s how the process unfolds:

• Initial Chunk Formation: The system splits each sentence of the data into initial chunks.
• Cosine Dissimilarity Calculation: For each pair of adjacent chunks, the parser calculates the cosine dissimilarity, which is (1- cosine similarity). This metric quantifies how different the chunks are from each other, based on their vector representations in a multi-dimensional semantic space.
• Threshold for Concatenation: The system sets a pre-defined threshold for dissimilarity. When the dissimilarity between adjacent chunks exceeds this threshold, it indicates that the chunks are significantly different from each other.
• Chunk Concatenation: When the dissimilarity does not exceed the threshold, indicating similarity or relevance between chunks, the system concatenates those chunks to form a larger, unified chunk. This process aims to ensure that each chunk represents a cohesive piece of information, enhancing the model’s ability to understand and process the data effectively.

Step1: Import Necessary Libraries

from llama_index.llms.openai import OpenAI
from llama_index.embeddings.huggingface import HuggingFaceEmbedding

Step2: Load OpenAI API Key Using Environment Variable File

import os
OPENAI_API_KEY = os.environ['OPENAI_API_KEY']

Step3: Define llm and Embedding Models

llm = OpenAI(model="gpt-3.5-turbo")
embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-small-en-v1.5")

Step4: Download the paul graham Dataset

!wget --user-agent "Mozilla" "https://raw.githubusercontent.com/run-llama/llama_index/main
/docs/docs/examples/data/paul_graham/paul_graham_essay.txt" -O './paul_graham_essay.txt'

# create RAFT Dataset object
raft_dataset = RAFTDatasetPack(file_path="./paul_graham_essay.txt", 
                               llm = llm, embed_model=embed_model, 
                               num_questions_per_chunk=1, num_distract_docs=2, chunk_size=1024, 
                               default_breakpoint_percentile_threshold=99)
                              
# create the dataset
dataset = raft_dataset.run()

# create the dataset
dataset = raft_dataset.run()

# save the dataset in jsonl format
output_path = './raft_dataset'
dataset.to_json(output_path + ".jsonl")

Step5: Loading Dataset 

with open('./raft_dataset.jsonl', 'r') as json_file:
    dataset = list(json_file)

# We can access the dataset with the following
json.loads(dataset[0]).keys()
# output
# dict_keys(['id', 'type', 'question', 'context', 'oracle_context', 'cot_answer', 'instruction'])

json.loads(dataset[0])['question']
# output
# 'What were the two main things the author worked on before college?'

Key Takeaways

• Learned that RAG significantly improves the performance of LLMs by integrating external data sources, enabling the models to generate more accurate, verifiable, and reliable answers, particularly for queries on topics not covered in their initial training data.
• Fine-tuning LLMs with a carefully prepared dataset, including oracle and distractor contexts, enhances the model’s ability to discern relevant from irrelevant information, leading to more precise responses to complex queries.
• Examined that Semantic Splitter Node Parser plays a crucial role in data preprocessing for LLM fine-tuning, by dividing data into semantically coherent chunks based on cosine dissimilarity, thereby optimizing the model’s training and subsequent performance.

Conclusion

The adoption of RAG alongside Large Language Models significantly mitigates their limitations by enabling accurate, verifiable responses to queries beyond their initial training scope. Fine-tuning LLMs with specifically prepared datasets and leveraging preprocessing techniques like the Semantic Splitter Node Parser enhances model performance. This approach marks a significant step forward in the evolution of AI applications, highlighting the importance of innovation in artificial intelligence for more reliable and sophisticated solutions.

Frequently Asked Questions

Santhosh Reddy Dandavolu

I am working as an Associate Data Scientist at Analytics Vidhya, a platform dedicated to building the Data Science ecosystem. My interests lie in the fields of Natural Language Processing (NLP), Deep Learning, and AI Agents.