Retrieval-Augmented Generation (RAG) bridges the best of two worlds: leveraging the accurate, contextually relevant answers from Large Language Models (LLMs) while sourcing custom and real-time information from your own datasets. In this tutorial, you will learn how to build your own RAG system, understand its core components, implement query expansion techniques, and address issues found in naive implementations.
Prerequisites for Building an RAG System
Before diving into building a Retrieval-Augmented Generation (RAG) system, ensure you have the necessary tools and resources prepared.
prerequisites
- Python Installed: Ensure you have Python 3.8+ installed on your system (Windows, macOS, or Linux).
- Virtual Environment Setup: Create and activate a Python virtual environment for the project.
- Required Libraries: Install dependencies such as
openai,langchain,chromadb,sentence_transformers,matplotlib,pypdf, andumap-learn. - OpenAI API Key: Obtain an OpenAI API key for LLM integration.
- Text Dataset: Prepare a collection of documents for retrieval. These can be JSON files, text files, or even PDFs.
Setting Up and Installing Necessary Tools
Follow these step-by-step instructions to set up a Python environment for building a RAG-based system.
steps
Create and Activate Virtual Environment:
bashpython -m venv rag_env source rag_env/bin/activate # For macOS/Linux .\rag_env\Scripts\activate # For WindowsInstall Required Libraries: Install all the necessary Python packages:
bashpip install openai langchain chromadb sentence-transformers matplotlib pypdf umap-learnSet up OpenAI API Key: Create a
.envfile in your project directory with your OpenAI API key:OPENAI_API_KEY=your-api-key
Understanding the Basics of RAG
Retrieval-Augmented Generation (RAG) combines two key components:
- Retriever: Responsible for fetching relevant documents or pieces of information from a preprocessed vectorized dataset.
- Generator: Utilizes the retrieved documents as additional context to generate more accurate responses using a Large Language Model (LLM).
Typical RAG Workflow:
- Documents are split into smaller "chunks."
- Each chunk is transformed into numerical embeddings using an embedding model (e.g., OpenAI embeddings or Sentence Transformers).
- Embeddings are indexed and stored in a vector database, such as ChromaDB.
- Queries are converted into embeddings and used to identify relevant documents stored in the database.
- Retrieved chunks are used to provide custom contextual data to the LLM for generating responses.
Building the Retrieval Pipeline
In this section, we create functions to preprocess documents, generate embeddings, and store them in a vector database.
steps
1. Load and Split Text
We begin by splitting input documents into manageable chunks for embedding.
from langchain.text_splitter import RecursiveCharacterTextSplitter
def load_and_preprocess_documents(file_paths):
"""
Load and split documents into chunks.
"""
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
documents = []
for file_path in file_paths:
with open(file_path, "r", encoding="utf-8") as file:
full_text = file.read()
documents += text_splitter.split_text(full_text)
return documents
# Example: Load and split text files
file_paths = ["./data/document1.txt", "./data/document2.txt"]
documents = load_and_preprocess_documents(file_paths)
print(f"Total chunks: {len(documents)}")2. Embed and Store in a Vector Database
from chromadb import Client
from chromadb.config import Settings
from chromadb.utils import embedding_functions
import os
# It is recommended to load your OpenAI API key from environment variables
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
# Initialize the ChromaDB client.
chroma_client = Client(Settings(persist_directory="./vector_db"))
# Load OpenAI embeddings function
embedding_function = embedding_functions.OpenAIEmbeddingFunction(api_key=OPENAI_API_KEY)
# Create a collection for storing document embeddings
collection = chroma_client.get_or_create_collection(
name="my_documents",
embedding_function=embedding_function,
)
# Add chunks into the vector database
for idx, doc in enumerate(documents):
embedding = embedding_function.embed_text(doc)
collection.add(
embeddings=[embedding],
documents=[doc],
metadatas=[{"id": idx}],
ids=[str(idx)],
)
print(f"Stored {collection.count()} document chunks.")Designing the Augmentation Workflow
To improve the accuracy of retrieved data, it's important to enhance both the initial query and the data search mechanisms (pre-retrieval) and refine output post-retrieval.
steps
1. Query Expansion for RAG
You can augment queries either by:
- Expanding the query with relevant terms generated by the LLM.
- Generating related additional queries.
Here’s how to implement basic query expansion with one answer:
def augment_query(query, model="gpt-3.5-turbo"):
"""
Use LLM to augment the query.
"""
import openai
prompt = (
"You are an assistant for information retrieval. "
"Suggest additional related terms or angles to improve the following query: "
f"{query}"
)
response = openai.ChatCompletion.create(
model=model,
messages=[{"role": "system", "content": "Assistant"}, {"role": "user", "content": prompt}],
)
return response['choices'][0]['message']['content']
# Example use case for augmenting:
original_query = "What were the causes of revenue growth in 2022?"
augmented_query = augment_query(original_query)
print("Augmented Query:", augmented_query)2. Query Results Reranking and Visualization
To evaluate similarity and ensure relevancy, visualization tools map embeddings to a graph.
import umap
import matplotlib.pyplot as plt
def visualize_embeddings(embeddings, labels):
"""
Project embeddings to 2D space and visualize them.
"""
projections = umap.UMAP(n_neighbors=10, min_dist=0.1, n_components=2).fit_transform(embeddings)
plt.scatter(projections[:, 0], projections[:, 1], c='blue', label='Documents')
for i, label in enumerate(labels):
plt.annotate(str(label), (projections[i, 0], projections[i, 1]))
plt.legend()
plt.title("Document Embeddings")
plt.show()
# Extract document embeddings and map them.
embeddings = collection.get_embeddings()
labels = [meta.get("id") for meta in collection.get_metadatas()]
visualize_embeddings(embeddings, labels)Common Pitfalls with Naïve RAG Systems
While RAG systems are powerful, naive implementations come with limitations. IT_GUIDES_COMPONENT_4
Summary and Next Steps
FAQ
What is a retrieval-augmented generation system?
A RAG system combines information retrieval techniques with generative Large Language Models to produce accurate, contextual, and customized answers by integrating and augmenting retrieved data.
How does a vector database differ from a traditional database?
A vector database stores embeddings (numerical representations of data) instead of raw data, allowing for similarity search and fast, relevant queries based on the vector proximity, ideal for complex search scenarios.
Why is query expansion important in RAG systems?
Query expansion enhances retrieval accuracy by including related terms and contexts, increasing the chances of retrieving semantically relevant documents.
Official reference: OpenAI API documentation.