Unlocking the Power of Retrieval Augmented Generation (RAG) with Llama Index

In the evolving landscape of artificial intelligence, the ability to efficiently retrieve and generate information is paramount. Retrieval Augmented Generation (RAG) is a cutting-edge technique that combines the strengths of retrieval-based models and generation-based models to provide accurate and contextually relevant responses. Today, we'll explore how to implement RAG using the Llama Index library, a versatile tool that simplifies the process of creating and querying vector indices.

What is Retrieval Augmented Generation (RAG)?

Working with Large Language Models (LLMs) presents several challenges, including gaps in domain knowledge, factual inaccuracies, and hallucinations. Retrieval Augmented Generation (RAG) helps mitigate these issues by enhancing LLMs with external knowledge sources, such as databases. This makes RAG especially valuable in knowledge-intensive scenarios or domain-specific applications that require constantly updated information. One significant advantage of RAG is that it doesn't require retraining the LLM for specific tasks. Recently, RAG has gained popularity for its application in conversational agents

At its core, Retrieval Augmented Generation (RAG) is an innovative method that integrates two powerful approaches in AI: retrieval and generation. Here’s an intuitive explanation:

• Retrieval: Think of it as searching for information. Imagine you have a massive library of documents, and you need to find relevant pieces of information based on a query. Retrieval models excel at this, quickly identifying documents that contain relevant information.
• Generation: This involves creating new text based on input data. Imagine you have a skilled writer who can take the retrieved documents and craft a coherent and contextually appropriate response. Generation models excel at understanding the context and producing human-like text.

RAG combines these two approaches by first using a retrieval model to gather relevant information and then passing this information to a generation model to produce a well-formed response. This synergy ensures that the responses are both accurate (thanks to the retrieval model) and contextually rich (thanks to the generation model).

The inspiration for this project came from this link: https://punkx.org/jackdoe/30.html. As shared by John Carmack, Ilya Sutskever of OpenAI provided him with an essential reading list of around 30 research papers, remarking, "If you really learn all of these, you’ll know 90% of what matters today in AI." This project uses these research papers as the knowledge base to create an AI query machine!

Note: The resources listed on this link were downloaded as PDFs and stored in a '/data' folder. The GitHub link to the code and vectors: https://github.com/raktimparashar-upenn/LLM_RAG

Setting Up the Environment

First, we need to ensure our environment is ready. We'll import necessary libraries, including os, llama_index, dotenv, and openai. Loading environment variables from a .env file ensures that sensitive information, like API keys, is securely managed.

# Import the os module to interact with the operating system
import os

# Import the llama_index library 
import llama_index

# Import the load_dotenv function from the dotenv library to load environment variables from a .env file
from dotenv import load_dotenv

# Import the openai library to interact with OpenAI's API
import openai

# Load environment variables from a .env file into the environment
load_dotenv()        

# Set the 'OPENAI_API_KEY' environment variable in the current environment to the value retrieved from the environment variables
os.environ['OPENAI_API_KEY'] = os.getenv('OPENAI_API_KEY')        

Loading Data and Creating the Index

Using SimpleDirectoryReader from the llama_index.core module, we can load documents from a specified directory. This example assumes PDF files are located in the "data" directory.

# Import the VectorStoreIndex and SimpleDirectoryReader classes from the llama_index.core module
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader

# Create an instance of SimpleDirectoryReader to read files from the "data" directory
pdfs = SimpleDirectoryReader("data").load_data()

# The load_data() method reads and loads the data from the specified directory

# Create an instance of VectorStoreIndex from the documents loaded into pdfs
# The from_documents method is used to build the index from the provided documents
# The show_progress parameter, when set to True, displays the progress of the indexing process

index = VectorStoreIndex.from_documents(pdfs, show_progress=True)        

With the data loaded, we create a VectorStoreIndex, which transforms our documents into word embeddings.

Query Engine Setup

To handle queries, we convert the VectorStoreIndex instance into a query engine. This enables us to perform natural language queries on our indexed data.

query_engine = index.as_query_engine()        

Customizing the Query Engine

from llama_index.core.retrievers import VectorIndexRetriever
from llama_index.core.query_engine import RetrieverQueryEngine
from llama_index.core.indices.postprocessor import SimilarityPostprocessor

retriever = VectorIndexRetriever(index=index, similarity_top_k=4)
postprocessor = SimilarityPostprocessor(similarity_cutoff=0.70)

query_engine = RetrieverQueryEngine(
    retriever=retriever,
    node_postprocessors=[postprocessor]
)        

For more refined query results, we employ VectorIndexRetriever, SimilarityPostprocessor, and RetrieverQueryEngine:

• VectorIndexRetriever: Retrieves the top similar documents based on the query.
• SimilarityPostprocessor: Filters out documents that do not meet a similarity threshold.
• RetrieverQueryEngine: Integrates the retriever and postprocessor for advanced querying.

Querying the Index

With our query engine in place, we can now perform queries to retrieve information. For example:

response = query_engine.query("What is a CNN?")
print(response)
response = query_engine.query("What is a RNN?")
print(response)
response = query_engine.query("Explain the transformer architecture.")
print(response)
response = query_engine.query("What is Attention is All You Need?")
print(response)

from llama_index.core.response.pprint_utils import pprint_response
pprint_response(response, show_source=True)        

Here are the query answers:

A CNN, or Convolutional Neural Network, is a type of neural network that is specifically designed for processing and analyzing visual data, such as images. It consists of neurons with learnable weights and biases, where each neuron receives inputs, performs a dot product operation, and may apply a non-linearity. CNNs are structured to make assumptions about the input data being images, allowing for efficient implementation and a reduction in the number of parameters in the network compared to traditional neural networks.

A Recurrent Neural Network (RNN) is a type of neural network that is designed to operate over sequences of vectors. Unlike Vanilla Neural Networks, which accept fixed-sized inputs and produce fixed-sized outputs using a fixed number of computational steps, RNNs can process sequences in the input, output, or both. This capability allows RNNs to learn patterns and dependencies in sequential data, making them particularly effective for tasks involving sequences like text generation, speech recognition, and time series prediction.

The Transformer model architecture consists of an encoder and a decoder. The encoder is made up of a stack of identical layers, each containing two sub-layers: a multi-head self-attention mechanism and a position-wise fully connected feed-forward network. Residual connections and layer normalization are applied around each sub-layer. The decoder also consists of a stack of identical layers, with an additional third sub-layer that performs multi-head attention over the output of the encoder stack. The self-attention mechanism in the decoder is modified to prevent positions from attending to subsequent positions. This architecture allows for parallelization and draws global dependencies between input and output using self-attention without relying on recurrent neural networks or convolution.

Storing and Reloading the Index

Persisting the index allows for efficient storage and retrieval of data without rebuilding the index each time. The following code checks for an existing storage directory and either creates a new index or loads an existing one:

# Import necessary modules and classes
import os.path
from llama_index.core import (
    VectorStoreIndex,  # For creating and handling vector store indices
    SimpleDirectoryReader,  # For reading documents from a directory
    StorageContext,  # For managing storage contexts
    load_index_from_storage,  # For loading an index from storage
)

# Define the directory where the storage will be persisted
PERSIST_DIR = "./storage"

# Check if the storage directory already exists
if not os.path.exists(PERSIST_DIR):
    # If the storage directory does not exist, load the documents and create the index
    documents = SimpleDirectoryReader("data").load_data()  # Read documents from the "data" directory
    index = VectorStoreIndex.from_documents(documents)  # Create an index from the loaded documents
    
    # Store the created index for later use
    index.storage_context.persist(persist_dir=PERSIST_DIR)
else:
    # If the storage directory exists, load the existing index
    storage_context = StorageContext.from_defaults(persist_dir=PERSIST_DIR)  # Create a storage context from the existing directory
    index = load_index_from_storage(storage_context)  # Load the index from the storage context

# Create a query engine from the index, regardless of whether it was newly created or loaded from storage
query_engine = index.as_query_engine()

# Query the index with a specific question
response = query_engine.query("Summarize Attention is all you need in 250 words.")

# Print the response from the query engine
print(response)        

Business Use Case: Enhancing Customer Support

Imagine a large enterprise with a vast repository of customer support documents, including FAQs, troubleshooting guides, and user manuals. Traditional search systems may struggle to deliver precise answers quickly. Here’s how RAG can revolutionize this scenario:

1. Efficiency in Information Retrieval: Using RAG, the system can rapidly retrieve the most relevant documents based on a customer query. For example, if a customer asks, "How do I reset my password?", the retrieval model can quickly identify relevant sections from various documents.
2. Contextual Response Generation: Once the relevant documents are retrieved, the generation model can synthesize the information and provide a coherent, contextually appropriate response. This means the customer receives a well-formatted answer that directly addresses their issue, rather than sifting through multiple documents.
3. Continuous Improvement: The RAG system can continuously learn from new data, improving its retrieval and generation capabilities over time. This ensures that the system remains up-to-date with the latest information and can handle an evolving set of customer queries.

Conclusion

RAG is a powerful technique that leverages the capabilities of both retrieval and generation models to provide comprehensive and context-aware responses. Using the llama_index library, we can efficiently create, query, and manage vector indices, unlocking the potential of our data. Whether it's for academic research, industry applications, or enhancing user interactions, RAG stands out as a vital tool in the AI toolkit. By integrating RAG into business processes, companies can significantly improve efficiency and customer satisfaction, ultimately driving better business outcomes.