Understanding RAGs - Give LLMs superpowers

An introduction to Retrieval Augmented Generation with real life examples


RAG

Introduction

As an engineer who has been working with Retrieval Augmented Generation (RAG) pipelines for around a year, the speed at which innovations happening in this field has caught me off guard. The future is AI, and RAG will have a big impact on specialized fields. So, I am putting this article out in an effort to spread the word on what RAG is and how to quickly get started in this area.

RAG is a powerful way to make Large Language Models (LLMs) even better. It combines the knowledge of LLMs with the ability to fetch and use external data. This makes responses more accurate, up-to-date, and relevant. In this article, I'll show you how to use RAG with LlamaIndex, a great data framework for LLM applications. We'll also look at some examples to see how it works in real life.

Understanding RAG

RAG is a mix of two things:

Retrieval: Finding useful information from a knowledge base.
Generation: Using an LLM to create responses based on the retrieved information and the input query.

This approach fixes some issues with traditional LLMs:

  • Knowledge cutoff: RAG can use up-to-date information, not just what the LLM was trained on.

  • Hallucination reduction: By using retrieved facts, RAG can reduce the chance of generating false or inconsistent information.

  • Customization: RAG lets you add specific knowledge without retraining the whole model.

  • Efficiency: RAG can handle large datasets more efficiently by retrieving only the relevant information.

  • Scalability: It allows for easy updates and additions to the knowledge base without needing to retrain the model.

Implementing RAG with LlamaIndex

LlamaIndex gives you the tools to build powerful RAG systems. Please make sure that you have correctly installed LlamaIndex by following the instructions from here. Now, we will build our RAG application. Let's break down the process into four steps:

1. Data Ingestion: Feeding Your RAG System

Data ingestion is all about getting information into your system. It's the foundation of your RAG setup.

Here's how to load data from different sources:

from llama_index.core import SimpleDirectoryReader
# Grab all the documents
all_docs = SimpleDirectoryReader('data/pdfs/').load_data()

2. Indexing: Organizing Your Data with Vector Indexes

Indexing is like giving your data a smart filing system. We'll focus on vector indexes, which are particularly powerful for RAG systems. LlamaIndex uses an OpenAI model by default for creating these vector indexes. So, you will need an OpenAI account for this step to work. You will also need at least $5 in your OpenAI account. After making sure that you have all this, you need to add your OpenAI API key as an environment variable. Click here to know how to set it up.

What are vector indexes? Vector indexes represent your documents as high-dimensional vectors. These vectors capture the semantic meaning of the text, not just keywords. To simplify, they are the mathematical representation of your data. They capture the true meaning behind the text which will help you filter data in a MUCH better way.

Here's how to create a vector index:

from llama_index.core import VectorStoreIndex
# Create a vector index from your documents
vector_index = VectorStoreIndex.from_documents(all_docs)
# Save your work
vector_index.storage_context.persist()

3. Querying: Digging Up Relevant Info

Querying is where you actually find the info you need. It's the "Retrieval" in RAG. With vector indexes, this becomes a similarity search in the vector space.

Let's look at how to query a vector index:

from llama_index.core.retrievers import VectorIndexRetriever
# Set up your vector retriever
vector_retriever = VectorIndexRetriever(index=vector_index)
# Create a query engine
query_engine = vector_index.as_query_engine()
# Fire away!
response = query_engine.query("What's new in AI?")
print(response)

4. Augmented Generation: Crafting Smart Responses

This is where the magic happens. We combine retrieved info with the LLM's smarts to generate responses.

Why it's a game-changer:

  • Your LLM now has access to fresh, relevant info from your vector index.

  • It blends broad LLM knowledge with specific retrieved data.

  • You can tweak how much the LLM relies on retrieved info vs. its own knowledge.

We are going to see an example where we use an OpenAI model. But you can use other LLMs such as Claude, llama, Mistral and many more. For more details, click here.

Let's see an example now. LlamaIndex handles a lot behind the scenes, but you can customize:

# Initialize the LLM
llm = OpenAI(model="gpt-3.5", temperature=0.7)
# Create a vector index with your custom setup
custom_index = VectorStoreIndex.from_documents(
    all_docs, settings={"llm": llm, "chunk_size": 512}
)
# Create a custom query engine
custom_query_engine = custom_index.as_query_engine()
# Query and print the response
response = custom_query_engine.query("Break down transfer learning in AI")
print(response)

Powerful Use Cases for RAG

Let's take a look at how we can use the things that we learned today.

1. Intelligent Code Assistant

Think of it as a super-powered engineer. It can learn about your whole codebase and suggest improvements. It can also catch bugs and tell you how to solve them. Isn't that awesome?

Let's see an example:

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.core.node_parser import SimpleNodeParser
 
# Load your codebase
codebase_documents = SimpleDirectoryReader('./code').load_data()
 
# Parse the documents into nodes
parser = SimpleNodeParser.from_defaults()
nodes = parser.get_nodes_from_documents(codebase_documents)
 
# Create an index of your codebase
code_index = VectorStoreIndex(nodes)
code_query_engine = code_index.as_query_engine()
 
def code_assistant(current_file, current_function, task):
    context = f"Current file: {current_file}\nCurrent function: {
        current_function}\nTask: {task}"
    response = code_query_engine.query(context)
    return response
 
# Example usage
current_file = "data_processing.py"
current_function = """
def process_data(data):
    # TODO: Implement data normalization
    pass
"""
task = "Implement data normalization in the process_data function. Consider any existing utility functions in our codebase for normalization. provide me the code."
 
print(code_assistant(current_file, current_function, task))

In this example, we're:

  1. Loading the entire codebase into our RAG system.

  2. Creating an index that allows us to search and understand the codebase.

  3. Implementing a code assistant that can provide suggestions based on the current file, function, and task.

  4. Asking for help with implementing a specific function, while considering the context of the existing codebase.

The RAG system would search the codebase for relevant information (like existing normalization functions), understand the context of the current file and function, and provide a suggestion that's tailored to your specific codebase and needs. It's like having a senior developer who knows your entire codebase inside out, always ready to help. You can even customize it to make your own Github Co-Pilot!

2. Personalized Customer Support

No matter what kind of product you have, you can always customize your customer support chatbot according to your needs. You can inform it about all your products, your policies and whatnot. Why would you need to hire a customer support person then?

Here's a quick example:

from llama_index.core import SimpleDirectoryReader, VectorStoreIndex
# Load the support documents
support_documents = SimpleDirectoryReader(
    './data/pdfs').load_data()
# Create an index from the support documents
support_index = VectorStoreIndex.from_documents(support_documents)
# Create a query engine from the index
support_query_engine = support_index.as_query_engine()
 
def handle_customer_query(customer_query, customer_history):
    # Combine the customer history and query into a single context
    context = f"Customer history: {
        customer_history}\n\nCustomer query: {customer_query}"
    # Use the query engine to generate a response based on the context
    response = support_query_engine.query(context)
    return response
 
# Test the function with a sample query and history
customer_query = "How do I reset my password?"
customer_history = "Previous interactions: Billing inquiry, Product return"
print(handle_customer_query(customer_query, customer_history))

3. Intelligent Research Assistant

Imagine that you are doing your thesis and you have collected a bunch of research papers. You can now build a research assistant that knows everything about your papers and can answer any of your questions regarding this.

Let's see it in action:

from llama_index.core import SimpleDirectoryReader, VectorStoreIndex
# Load the research papers
research_papers = SimpleDirectoryReader(
    './data/pdfs').load_data()
# Create an index from the research papers
research_index = VectorStoreIndex.from_documents(research_papers)
# Create a query engine from the index
research_query_engine = research_index.as_query_engine()
 
def research_assistant(topic):
    # Construct a query that summarizes the latest findings on the topic, highlights key controversies, and suggests potential future research directions
    query = f"Summarize the latest findings on {
        topic}, highlight key controversies, and suggest potential future research directions."
    # Use the query engine to generate a response
    response = research_query_engine.query(query)
    return response
 
# Test the function with a sample topic
print(research_assistant("CRISPR gene editing in cancer therapy"))

Conclusion

As I've worked with RAG, I've seen how powerful it can be. I think it's going to change the way we use AI in many fields. I'm excited to see where RAG will go from here and how it will impact the work we do. I hope this article has been helpful in giving you a sense of what RAG is and how it works, and that you'll start exploring its possibilities for yourself.