INTRODUCTION
Welcome to this tutorial on building a powerful Large Language Model chatbot with Retrieval-Augmented Generation capabilities. This guide will take you from zero to hero, starting with a simple chatbot and progressively adding sophisticated features that make modern AI assistants so powerful.
Large Language Models have revolutionized how we interact with computers, but they have limitations. They can hallucinate facts, lack access to recent information, and cannot reference your private documents. Retrieval-Augmented Generation solves these problems by combining the reasoning capabilities of LLMs with the precision of information retrieval systems.
By the end of this tutorial, you will have built a production-ready chatbot system that can run on various hardware configurations including NVIDIA CUDA GPUs, AMD ROCm GPUs, Intel GPUs, and Apple Silicon MPS. You will understand every component deeply and be able to customize the system for your specific needs.
UNDERSTANDING THE FUNDAMENTALS
Before we write any code, let us understand what we are building and why each component matters.
What is a Large Language Model?
A Large Language Model is a neural network trained on vast amounts of text data to predict the next token in a sequence. The word "token" here refers to a piece of text, which could be a word, part of a word, or even a single character. When you give an LLM a prompt like "The capital of France is", it predicts that "Paris" is the most likely continuation based on patterns it learned during training.
Modern LLMs like GPT, LLaMA, Mistral, and others contain billions of parameters. These parameters are numerical weights that the model adjusts during training to capture language patterns, facts, reasoning abilities, and even some world knowledge. The larger the model, generally the more capable it is, but also the more computational resources it requires.
What is Retrieval-Augmented Generation?
Retrieval-Augmented Generation, commonly abbreviated as RAG, is a technique that enhances LLM responses by providing relevant context from external knowledge sources. Instead of relying solely on the model's training data, RAG systems first retrieve relevant documents or passages from a knowledge base, then feed these to the LLM along with the user's question.
Think of it like an open-book exam versus a closed-book exam. Without RAG, the LLM must answer from memory alone. With RAG, it can consult reference materials before answering, leading to more accurate and up-to-date responses.
The RAG process follows these steps. First, when a user asks a question, the system converts that question into a numerical representation called an embedding. Second, it searches a database of pre-computed embeddings to find the most similar documents. Third, it retrieves those documents and combines them with the user's question into a prompt. Fourth, it sends this enhanced prompt to the LLM. Finally, the LLM generates a response based on both its training and the retrieved context.
Why Use HuggingFace Transformers?
HuggingFace Transformers is the de facto standard library for working with transformer-based models in Python. It provides a unified interface to thousands of pre-trained models, handles the complexity of different model architectures, and includes optimizations for various hardware platforms.
The library abstracts away many low-level details while still giving you control when needed. You can load a model with just a few lines of code, but you can also customize tokenization, generation parameters, and inference optimizations. This balance makes it perfect for both beginners and advanced users.
SETTING UP YOUR DEVELOPMENT ENVIRONMENT
A proper development environment is crucial for success. We need to install the right packages and configure hardware acceleration correctly.
Hardware Acceleration Setup
Modern LLMs require significant computational power. While you can run small models on CPUs, GPU acceleration makes the experience much better. The challenge is that different GPU manufacturers require different software stacks.
For NVIDIA GPUs, you need CUDA and cuDNN. PyTorch automatically detects CUDA if installed correctly. For AMD GPUs, you need ROCm, which is AMD's open-source GPU computing platform. Intel GPUs use Intel Extension for PyTorch. Apple Silicon Macs use Metal Performance Shaders through PyTorch's MPS backend.
The good news is that PyTorch and HuggingFace Transformers handle most of this complexity automatically. We just need to detect which device is available and use it.
Installing Required Packages
We will need several Python packages. The core package is transformers from HuggingFace, which provides model loading and inference capabilities. We need torch, the PyTorch deep learning framework that powers transformers. For RAG functionality, we need sentence-transformers for creating embeddings, faiss-cpu or faiss-gpu for fast similarity search, and langchain for document processing utilities.
Here is how to install these packages:
# Install core dependencies
pip install torch torchvision torchaudio
pip install transformers accelerate
pip install sentence-transformers
pip install faiss-cpu # Use faiss-gpu if you have CUDA
pip install langchain langchain-community
pip install pypdf # For PDF processing
pip install chromadb # Alternative vector store
The accelerate library is particularly important. It automatically handles device placement, mixed precision training, and distributed computing. It makes your code work across different hardware without manual device management.
Device Detection and Configuration
Let us write code to detect and configure the best available hardware. This code will be the foundation for all our subsequent work.
import torch
import platform
def get_optimal_device():
"""
Detect and return the best available compute device.
This function checks for GPU availability in order of preference:
CUDA (NVIDIA), ROCm (AMD), MPS (Apple Silicon), then falls back to CPU.
Returns:
torch.device: The optimal device for computation
str: A human-readable description of the device
"""
if torch.cuda.is_available():
device = torch.device("cuda")
gpu_name = torch.cuda.get_device_name(0)
device_info = f"CUDA GPU: {gpu_name}"
print(f"Using {device_info}")
print(f"GPU Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.2f} GB")
return device, device_info
elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
device = torch.device("mps")
device_info = "Apple Silicon MPS"
print(f"Using {device_info}")
return device, device_info
elif torch.version.hip is not None:
# ROCm is available
device = torch.device("cuda") # ROCm uses cuda device name
device_info = "AMD ROCm GPU"
print(f"Using {device_info}")
return device, device_info
else:
device = torch.device("cpu")
device_info = f"CPU: {platform.processor()}"
print(f"Using {device_info}")
print("Warning: No GPU detected. Inference will be slower.")
return device, device_info
This function is more sophisticated than it might appear. It checks for CUDA first because NVIDIA GPUs are most common in deep learning. Then it checks for Apple's MPS backend, which is available on Apple Silicon Macs. The ROCm check looks at torch.version.hip because ROCm-enabled PyTorch builds include HIP support. Finally, it falls back to CPU if no GPU is available.
The function also prints useful information about the detected device. For CUDA, it shows the GPU name and memory, which helps you understand if your model will fit. This information is invaluable when debugging out-of-memory errors.
BUILDING THE BASIC LLM CHATBOT
Now we will build our first chatbot. We will start simple and add complexity gradually. This approach helps you understand each component before moving to the next.
Choosing a Model
Model selection is critical. You need to balance capability, size, and hardware requirements. For this tutorial, we will use models from the LLaMA family because they offer excellent performance at various sizes and are widely available.
Smaller models like LLaMA 2 7B or Mistral 7B can run on consumer GPUs with 8-16GB of VRAM. Larger models like LLaMA 2 13B or 70B require more powerful hardware or quantization techniques. Quantization reduces model precision from 16-bit or 32-bit floating point to 8-bit or even 4-bit integers, dramatically reducing memory usage with minimal quality loss.
For our examples, we will use a 7B parameter model, which strikes a good balance. If you have limited hardware, you can use smaller models or quantized versions.
Understanding the Model Loading Process
Loading a model involves several steps. First, we load the tokenizer, which converts text to numerical tokens. Second, we load the model weights, which can be several gigabytes. Third, we move the model to the appropriate device. Fourth, we configure generation parameters.
Let us examine each step in detail with code:
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
import torch
class LLMChatbot:
"""
A simple but powerful LLM chatbot implementation.
This class encapsulates model loading, conversation management,
and text generation with proper error handling and device management.
"""
def __init__(self, model_name="meta-llama/Llama-2-7b-chat-hf",
device=None, use_quantization=False):
"""
Initialize the chatbot with a specified model.
Args:
model_name: HuggingFace model identifier
device: Torch device to use (auto-detected if None)
use_quantization: Whether to use 4-bit quantization to save memory
"""
self.model_name = model_name
# Detect device if not provided
if device is None:
self.device, self.device_info = get_optimal_device()
else:
self.device = device
self.device_info = str(device)
print(f"Loading model: {model_name}")
print(f"Target device: {self.device_info}")
# Configure quantization if requested
quantization_config = None
if use_quantization:
quantization_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_compute_dtype=torch.float16,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_type="nf4"
)
print("Using 4-bit quantization")
# Load tokenizer
self.tokenizer = AutoTokenizer.from_pretrained(
model_name,
trust_remote_code=True
)
# Ensure tokenizer has a pad token
if self.tokenizer.pad_token is None:
self.tokenizer.pad_token = self.tokenizer.eos_token
# Load model
self.model = AutoModelForCausalLM.from_pretrained(
model_name,
quantization_config=quantization_config,
device_map="auto",
trust_remote_code=True,
torch_dtype=torch.float16 if not use_quantization else None
)
# Set model to evaluation mode
self.model.eval()
# Initialize conversation history
self.conversation_history = []
print("Model loaded successfully!")
Let us break down what happens in this initialization code. The AutoTokenizer.from_pretrained method downloads the tokenizer configuration from HuggingFace's model hub if not already cached locally. The tokenizer is responsible for converting text into token IDs that the model can process. Different models use different tokenization schemes, so we must use the tokenizer that matches our model.
The pad token check is important. Some models do not define a padding token by default, which causes errors during batch processing. We set it to the end-of-sequence token as a safe default.
The AutoModelForCausalLM.from_pretrained method loads the actual model weights. The device_map="auto" argument is powerful. It automatically distributes model layers across available devices, enabling you to run large models that do not fit on a single GPU. The torch_dtype=torch.float16 argument uses half-precision floating point, which halves memory usage with negligible quality impact for inference.
The quantization configuration uses the BitsAndBytes library to load the model in 4-bit precision. The nf4 quantization type is specifically designed for neural networks and provides better quality than standard 4-bit quantization. Double quantization further compresses the quantization constants themselves.
Implementing Text Generation
Now we implement the core functionality: generating responses. This involves careful prompt construction, token generation, and post-processing.
def generate_response(self, user_input, max_new_tokens=512,
temperature=0.7, top_p=0.9, top_k=50):
"""
Generate a response to user input.
Args:
user_input: The user's message
max_new_tokens: Maximum number of tokens to generate
temperature: Sampling temperature (higher = more random)
top_p: Nucleus sampling parameter
top_k: Top-k sampling parameter
Returns:
str: The generated response
"""
# Add user input to conversation history
self.conversation_history.append({
"role": "user",
"content": user_input
})
# Format the conversation for the model
prompt = self._format_conversation()
# Tokenize the prompt
inputs = self.tokenizer(
prompt,
return_tensors="pt",
padding=True,
truncation=True,
max_length=2048
).to(self.device)
# Generate response
with torch.no_grad():
outputs = self.model.generate(
**inputs,
max_new_tokens=max_new_tokens,
temperature=temperature,
top_p=top_p,
top_k=top_k,
do_sample=True,
pad_token_id=self.tokenizer.pad_token_id,
eos_token_id=self.tokenizer.eos_token_id
)
# Decode the generated tokens
generated_text = self.tokenizer.decode(
outputs[0],
skip_special_tokens=True
)
# Extract only the new response
response = self._extract_response(generated_text, prompt)
# Add assistant response to history
self.conversation_history.append({
"role": "assistant",
"content": response
})
return response
The generation parameters deserve detailed explanation. Temperature controls randomness in token selection. A temperature of 1.0 uses the model's raw probability distribution. Lower temperatures make the model more deterministic and focused, while higher temperatures increase creativity and randomness. A temperature of 0.7 is a good default for conversational responses.
Top-p sampling, also called nucleus sampling, selects from the smallest set of tokens whose cumulative probability exceeds p. This prevents the model from selecting very unlikely tokens while maintaining diversity. A top_p of 0.9 means we only consider tokens that make up the top 90 percent of the probability mass.
Top-k sampling limits selection to the k most likely tokens. This provides a hard cutoff on randomness. Combining top-k and top-p gives fine-grained control over generation quality.
The torch.no_grad() context manager is crucial for inference. It disables gradient computation, which we do not need during generation and which consumes significant memory. This allows us to run larger models or longer sequences.
Formatting Conversations
Different models expect different conversation formats. LLaMA 2 uses a specific format with special tokens. We need to format our conversation history correctly.
def _format_conversation(self):
"""
Format the conversation history into a prompt string.
Different models use different conversation formats. This method
handles the LLaMA 2 chat format with proper special tokens.
Returns:
str: Formatted prompt string
"""
if "llama-2" in self.model_name.lower() and "chat" in self.model_name.lower():
# LLaMA 2 chat format
formatted_parts = []
for message in self.conversation_history:
if message["role"] == "user":
formatted_parts.append(f"[INST] {message['content']} [/INST]")
elif message["role"] == "assistant":
formatted_parts.append(message["content"])
elif message["role"] == "system":
formatted_parts.append(f"<<SYS>>\n{message['content']}\n<</SYS>>")
return " ".join(formatted_parts)
else:
# Generic format for other models
formatted_parts = []
for message in self.conversation_history:
role = message["role"].capitalize()
content = message["content"]
formatted_parts.append(f"{role}: {content}")
return "\n".join(formatted_parts) + "\nAssistant:"
The LLaMA 2 chat format uses special tokens like [INST] and [/INST] to mark user instructions. The <
For other models, we use a simpler format with role labels. This works reasonably well for most instruction-tuned models. In production, you would want to check the model's documentation for its preferred format.
Extracting the Response
After generation, we need to extract just the new response, not the entire conversation history that was in the prompt.
def _extract_response(self, generated_text, prompt):
"""
Extract the new response from the generated text.
The model generates a continuation of the prompt, so we need to
remove the prompt portion to get just the assistant's response.
Args:
generated_text: The full generated text
prompt: The original prompt
Returns:
str: Just the new response
"""
# Remove the prompt from the generated text
if generated_text.startswith(prompt):
response = generated_text[len(prompt):].strip()
else:
# Fallback: try to find where the response starts
response = generated_text.strip()
# Clean up any trailing special tokens
for token in ["</s>", "<|endoftext|>", "[/INST]"]:
response = response.replace(token, "")
return response.strip()
This extraction is necessary because the model generates a continuation of the input sequence. If we prompted it with "User: Hello\nAssistant:", it generates "User: Hello\nAssistant: Hi there! How can I help you?" We only want the "Hi there! How can I help you?" part.
We also clean up special tokens that might appear at the end of generation. Different models use different end-of-sequence tokens, so we handle the most common ones.
Adding Conversation Management
A good chatbot needs to manage conversation history properly. We should allow clearing history, setting system prompts, and limiting history length to prevent context overflow.
def set_system_prompt(self, system_prompt):
"""
Set a system prompt that defines the assistant's behavior.
Args:
system_prompt: Instructions for the assistant's behavior
"""
# Remove any existing system prompt
self.conversation_history = [
msg for msg in self.conversation_history
if msg["role"] != "system"
]
# Add new system prompt at the beginning
self.conversation_history.insert(0, {
"role": "system",
"content": system_prompt
})
def clear_history(self):
"""
Clear the conversation history.
Preserves the system prompt if one exists.
"""
system_prompts = [
msg for msg in self.conversation_history
if msg["role"] == "system"
]
self.conversation_history = system_prompts
def get_history(self):
"""
Get the current conversation history.
Returns:
list: List of conversation messages
"""
return self.conversation_history.copy()
def trim_history(self, max_messages=10):
"""
Trim conversation history to the most recent messages.
This prevents the context from growing too long and exceeding
the model's maximum context length.
Args:
max_messages: Maximum number of message pairs to keep
"""
# Always keep system prompts
system_prompts = [
msg for msg in self.conversation_history
if msg["role"] == "system"
]
# Get non-system messages
other_messages = [
msg for msg in self.conversation_history
if msg["role"] != "system"
]
# Keep only the most recent messages
if len(other_messages) > max_messages * 2:
other_messages = other_messages[-(max_messages * 2):]
# Reconstruct history
self.conversation_history = system_prompts + other_messages
The trim_history method is particularly important for long conversations. LLMs have a maximum context length, typically 2048 or 4096 tokens. If the conversation history exceeds this, the model will either truncate it or fail. By keeping only recent messages, we ensure the context stays within limits while maintaining conversational coherence.
We multiply max_messages by two because each exchange consists of a user message and an assistant response. Keeping ten exchanges means twenty messages total.
UNDERSTANDING EMBEDDINGS AND VECTOR SEARCH
Before we add RAG capabilities, we need to understand embeddings and vector search. These are the foundation of retrieval systems.
What Are Embeddings?
An embedding is a numerical representation of text that captures its semantic meaning. Instead of treating text as a sequence of discrete tokens, we represent it as a point in a high-dimensional space. Texts with similar meanings are close together in this space, while unrelated texts are far apart.
For example, the sentences "The cat sat on the mat" and "A feline rested on the rug" would have similar embeddings despite using different words, because they convey similar meanings. Meanwhile, "Quantum physics is complex" would have a very different embedding.
Modern embedding models are neural networks trained specifically to produce these representations. They learn to map text to vectors such that semantic similarity corresponds to vector similarity. The most common similarity metric is cosine similarity, which measures the angle between vectors.
Sentence Transformers
The sentence-transformers library provides easy access to state-of-the-art embedding models. These models are based on transformer architectures like BERT but are fine-tuned specifically for generating high-quality sentence embeddings.
Let us create an embedding generator:
from sentence_transformers import SentenceTransformer
import numpy as np
class EmbeddingGenerator:
"""
Generate embeddings for text using sentence transformers.
This class handles loading an embedding model and generating
vector representations of text that can be used for similarity search.
"""
def __init__(self, model_name="all-MiniLM-L6-v2", device=None):
"""
Initialize the embedding generator.
Args:
model_name: Name of the sentence transformer model
device: Device to run the model on
"""
if device is None:
self.device, _ = get_optimal_device()
else:
self.device = device
print(f"Loading embedding model: {model_name}")
# Load the sentence transformer model
self.model = SentenceTransformer(model_name)
self.model.to(self.device)
# Get embedding dimension
self.embedding_dim = self.model.get_sentence_embedding_dimension()
print(f"Embedding dimension: {self.embedding_dim}")
print("Embedding model loaded successfully!")
def encode(self, texts, batch_size=32, show_progress=False):
"""
Generate embeddings for a list of texts.
Args:
texts: Single text string or list of text strings
batch_size: Number of texts to process at once
show_progress: Whether to show a progress bar
Returns:
numpy.ndarray: Array of embeddings
"""
# Handle single text input
if isinstance(texts, str):
texts = [texts]
# Generate embeddings
embeddings = self.model.encode(
texts,
batch_size=batch_size,
show_progress_bar=show_progress,
convert_to_numpy=True,
normalize_embeddings=True # Normalize for cosine similarity
)
return embeddings
The all-MiniLM-L6-v2 model is an excellent default choice. It is small and fast while producing high-quality embeddings. The model outputs 384-dimensional vectors, which is a good balance between expressiveness and computational efficiency.
We normalize embeddings by default. Normalization scales vectors to unit length, which makes cosine similarity equivalent to dot product. This simplifies and speeds up similarity calculations.
Vector Databases and FAISS
Once we have embeddings, we need a way to search them efficiently. A naive approach would compare the query embedding to every document embedding, but this becomes slow with thousands or millions of documents. Vector databases solve this with specialized indexing structures.
FAISS, developed by Facebook AI Research, is the gold standard for vector similarity search. It provides various index types optimized for different scenarios. For our purposes, we will use a flat index for small datasets and an IVF index for larger ones.
import faiss
import pickle
import os
class VectorStore:
"""
A vector store for efficient similarity search using FAISS.
This class manages a collection of document embeddings and provides
fast similarity search capabilities.
"""
def __init__(self, embedding_dim, use_gpu=False):
"""
Initialize the vector store.
Args:
embedding_dim: Dimension of the embeddings
use_gpu: Whether to use GPU acceleration for search
"""
self.embedding_dim = embedding_dim
self.use_gpu = use_gpu
# Initialize FAISS index
# We use a flat L2 index for exact search
self.index = faiss.IndexFlatL2(embedding_dim)
# Move to GPU if requested and available
if use_gpu and faiss.get_num_gpus() > 0:
self.index = faiss.index_cpu_to_gpu(
faiss.StandardGpuResources(),
0,
self.index
)
print("Using GPU-accelerated FAISS index")
# Store the actual documents
self.documents = []
self.metadata = []
def add_documents(self, documents, embeddings, metadata=None):
"""
Add documents and their embeddings to the store.
Args:
documents: List of document texts
embeddings: Numpy array of embeddings
metadata: Optional list of metadata dictionaries
"""
# Ensure embeddings are float32 (FAISS requirement)
embeddings = embeddings.astype(np.float32)
# Add to FAISS index
self.index.add(embeddings)
# Store documents
self.documents.extend(documents)
# Store metadata
if metadata is None:
metadata = [{} for _ in documents]
self.metadata.extend(metadata)
print(f"Added {len(documents)} documents. Total: {len(self.documents)}")
def search(self, query_embedding, top_k=5):
"""
Search for the most similar documents to a query.
Args:
query_embedding: The query embedding vector
top_k: Number of results to return
Returns:
list: List of tuples (document, score, metadata)
"""
# Ensure query is the right shape and type
query_embedding = query_embedding.astype(np.float32)
if len(query_embedding.shape) == 1:
query_embedding = query_embedding.reshape(1, -1)
# Search the index
distances, indices = self.index.search(query_embedding, top_k)
# Prepare results
results = []
for dist, idx in zip(distances[0], indices[0]):
if idx < len(self.documents): # Valid index
results.append({
"document": self.documents[idx],
"score": float(dist),
"metadata": self.metadata[idx]
})
return results
def save(self, directory):
"""
Save the vector store to disk.
Args:
directory: Directory to save the store
"""
os.makedirs(directory, exist_ok=True)
# Save FAISS index
if self.use_gpu:
# Move back to CPU for saving
cpu_index = faiss.index_gpu_to_cpu(self.index)
faiss.write_index(cpu_index, os.path.join(directory, "index.faiss"))
else:
faiss.write_index(self.index, os.path.join(directory, "index.faiss"))
# Save documents and metadata
with open(os.path.join(directory, "documents.pkl"), "wb") as f:
pickle.dump({
"documents": self.documents,
"metadata": self.metadata
}, f)
print(f"Vector store saved to {directory}")
def load(self, directory):
"""
Load the vector store from disk.
Args:
directory: Directory containing the saved store
"""
# Load FAISS index
self.index = faiss.read_index(os.path.join(directory, "index.faiss"))
# Move to GPU if requested
if self.use_gpu and faiss.get_num_gpus() > 0:
self.index = faiss.index_cpu_to_gpu(
faiss.StandardGpuResources(),
0,
self.index
)
# Load documents and metadata
with open(os.path.join(directory, "documents.pkl"), "rb") as f:
data = pickle.load(f)
self.documents = data["documents"]
self.metadata = data["metadata"]
print(f"Loaded vector store with {len(self.documents)} documents")
The IndexFlatL2 performs exact nearest neighbor search using L2 distance. This is perfect for small to medium datasets up to a few hundred thousand vectors. For larger datasets, you would use an approximate index like IndexIVFFlat, which trades some accuracy for much faster search.
The GPU acceleration in FAISS can provide significant speedups, especially for large-scale searches. However, it requires additional GPU memory, so you need to balance this with your model's memory requirements.
DOCUMENT PROCESSING FOR RAG
To build a RAG system, we need to process documents into chunks that can be embedded and retrieved. This involves loading documents, splitting them intelligently, and creating embeddings.
Document Loading
Different document types require different loading strategies. We will support plain text, PDF, and markdown files.
from typing import List, Dict
import os
from pathlib import Path
class DocumentLoader:
"""
Load documents from various file formats.
Supports text files, PDFs, and markdown files with appropriate
preprocessing for each format.
"""
def __init__(self):
"""Initialize the document loader."""
self.supported_extensions = {'.txt', '.md', '.pdf'}
def load_file(self, file_path):
"""
Load a single file.
Args:
file_path: Path to the file
Returns:
dict: Document dictionary with text and metadata
"""
file_path = Path(file_path)
if not file_path.exists():
raise FileNotFoundError(f"File not found: {file_path}")
extension = file_path.suffix.lower()
if extension not in self.supported_extensions:
raise ValueError(f"Unsupported file type: {extension}")
# Load based on file type
if extension == '.pdf':
text = self._load_pdf(file_path)
else:
# Text and markdown files
with open(file_path, 'r', encoding='utf-8') as f:
text = f.read()
# Create document dictionary
document = {
"text": text,
"metadata": {
"source": str(file_path),
"filename": file_path.name,
"extension": extension
}
}
return document
def _load_pdf(self, file_path):
"""
Load text from a PDF file.
Args:
file_path: Path to the PDF file
Returns:
str: Extracted text
"""
try:
from pypdf import PdfReader
except ImportError:
raise ImportError("pypdf is required for PDF support. Install with: pip install pypdf")
reader = PdfReader(file_path)
text_parts = []
for page_num, page in enumerate(reader.pages):
text = page.extract_text()
if text.strip():
text_parts.append(text)
return "\n\n".join(text_parts)
def load_directory(self, directory_path, recursive=True):
"""
Load all supported documents from a directory.
Args:
directory_path: Path to the directory
recursive: Whether to search subdirectories
Returns:
list: List of document dictionaries
"""
directory_path = Path(directory_path)
if not directory_path.is_dir():
raise NotADirectoryError(f"Not a directory: {directory_path}")
documents = []
# Get all files
if recursive:
files = directory_path.rglob("*")
else:
files = directory_path.glob("*")
# Load each supported file
for file_path in files:
if file_path.is_file() and file_path.suffix.lower() in self.supported_extensions:
try:
doc = self.load_file(file_path)
documents.append(doc)
except Exception as e:
print(f"Error loading {file_path}: {e}")
print(f"Loaded {len(documents)} documents from {directory_path}")
return documents
The PDF loading uses pypdf, which is a pure Python library for reading PDFs. It extracts text page by page and combines them. For production use, you might want to preserve page numbers in the metadata for better citation.
The directory loading function uses Path.rglob for recursive search, which finds all files in subdirectories. This is useful when you have a large document collection organized in folders.
Text Chunking
Raw documents are often too long to use as-is. We need to split them into smaller chunks that fit within the context window and provide focused, relevant information. However, we cannot just split on arbitrary boundaries because that would break semantic coherence.
Intelligent chunking considers sentence boundaries, paragraph breaks, and semantic coherence. We want each chunk to be self-contained and meaningful.
import re
class TextChunker:
"""
Split text into semantically coherent chunks.
This class implements intelligent text splitting that respects
sentence boundaries and maintains semantic coherence.
"""
def __init__(self, chunk_size=500, chunk_overlap=50):
"""
Initialize the text chunker.
Args:
chunk_size: Target size of each chunk in characters
chunk_overlap: Number of characters to overlap between chunks
"""
self.chunk_size = chunk_size
self.chunk_overlap = chunk_overlap
def chunk_text(self, text, metadata=None):
"""
Split text into chunks.
Args:
text: The text to chunk
metadata: Optional metadata to attach to each chunk
Returns:
list: List of chunk dictionaries
"""
if metadata is None:
metadata = {}
# Split into sentences first
sentences = self._split_into_sentences(text)
chunks = []
current_chunk = []
current_length = 0
for sentence in sentences:
sentence_length = len(sentence)
# If adding this sentence would exceed chunk size
if current_length + sentence_length > self.chunk_size and current_chunk:
# Save current chunk
chunk_text = " ".join(current_chunk)
chunks.append({
"text": chunk_text,
"metadata": {**metadata, "chunk_index": len(chunks)}
})
# Start new chunk with overlap
# Keep last few sentences for context
overlap_text = " ".join(current_chunk)
if len(overlap_text) > self.chunk_overlap:
# Find where to cut for overlap
overlap_sentences = []
overlap_length = 0
for s in reversed(current_chunk):
if overlap_length + len(s) <= self.chunk_overlap:
overlap_sentences.insert(0, s)
overlap_length += len(s)
else:
break
current_chunk = overlap_sentences
current_length = overlap_length
else:
current_chunk = []
current_length = 0
# Add sentence to current chunk
current_chunk.append(sentence)
current_length += sentence_length
# Add final chunk if not empty
if current_chunk:
chunk_text = " ".join(current_chunk)
chunks.append({
"text": chunk_text,
"metadata": {**metadata, "chunk_index": len(chunks)}
})
return chunks
def _split_into_sentences(self, text):
"""
Split text into sentences.
Uses a simple regex-based approach that handles common cases.
For production, consider using a proper sentence tokenizer like spaCy.
Args:
text: The text to split
Returns:
list: List of sentences
"""
# Simple sentence splitting pattern
# Matches periods, exclamation marks, question marks followed by space
sentence_pattern = r'(?<=[.!?])\s+'
sentences = re.split(sentence_pattern, text)
# Clean up sentences
sentences = [s.strip() for s in sentences if s.strip()]
return sentences
def chunk_documents(self, documents):
"""
Chunk multiple documents.
Args:
documents: List of document dictionaries
Returns:
list: List of chunk dictionaries
"""
all_chunks = []
for doc in documents:
chunks = self.chunk_text(doc["text"], doc.get("metadata", {}))
all_chunks.extend(chunks)
print(f"Created {len(all_chunks)} chunks from {len(documents)} documents")
return all_chunks
The chunking strategy uses a sliding window approach with overlap. Overlap is crucial because it ensures that information near chunk boundaries appears in multiple chunks, reducing the risk of missing relevant context.
The sentence-based splitting is better than character-based splitting because it preserves semantic units. Breaking a sentence in the middle can make the chunk harder to understand and reduce retrieval quality.
For production systems, you might want to use more sophisticated sentence splitting with libraries like spaCy or NLTK, which handle edge cases like abbreviations and decimal numbers better.
BUILDING THE RAG SYSTEM
Now we combine all the pieces to create a complete RAG system. This system will load documents, create embeddings, store them in a vector database, and use them to enhance LLM responses.
The RAG Pipeline
The RAG pipeline consists of two main phases: indexing and retrieval. During indexing, we process documents, create embeddings, and store them. During retrieval, we find relevant documents for a query and use them to augment the LLM's context.
class RAGSystem:
"""
A complete Retrieval-Augmented Generation system.
This class combines document loading, embedding generation, vector search,
and LLM generation into a cohesive system that can answer questions based
on a knowledge base.
"""
def __init__(self, llm_model_name="meta-llama/Llama-2-7b-chat-hf",
embedding_model_name="all-MiniLM-L6-v2",
chunk_size=500, chunk_overlap=50,
device=None, use_quantization=False):
"""
Initialize the RAG system.
Args:
llm_model_name: Name of the LLM to use
embedding_model_name: Name of the embedding model
chunk_size: Size of text chunks
chunk_overlap: Overlap between chunks
device: Device to use for models
use_quantization: Whether to quantize the LLM
"""
print("Initializing RAG system...")
# Initialize components
self.document_loader = DocumentLoader()
self.text_chunker = TextChunker(chunk_size, chunk_overlap)
self.embedding_generator = EmbeddingGenerator(embedding_model_name, device)
# Initialize vector store
self.vector_store = VectorStore(
self.embedding_generator.embedding_dim,
use_gpu=False # Keep embeddings on CPU to save GPU memory for LLM
)
# Initialize LLM chatbot
self.chatbot = LLMChatbot(llm_model_name, device, use_quantization)
# Set a default system prompt for RAG
self.chatbot.set_system_prompt(
"You are a helpful assistant. When answering questions, use the provided context "
"to give accurate and detailed responses. If the context does not contain enough "
"information to answer the question, say so clearly. Always cite the source of "
"your information when possible."
)
print("RAG system initialized successfully!")
def index_documents(self, source, is_directory=False, recursive=True):
"""
Index documents from a file or directory.
This processes documents, creates embeddings, and stores them
in the vector database for later retrieval.
Args:
source: Path to a file or directory
is_directory: Whether the source is a directory
recursive: If directory, whether to search recursively
"""
print(f"Indexing documents from: {source}")
# Load documents
if is_directory:
documents = self.document_loader.load_directory(source, recursive)
else:
documents = [self.document_loader.load_file(source)]
if not documents:
print("No documents found to index")
return
# Chunk documents
chunks = self.text_chunker.chunk_documents(documents)
# Generate embeddings
print("Generating embeddings...")
chunk_texts = [chunk["text"] for chunk in chunks]
embeddings = self.embedding_generator.encode(
chunk_texts,
batch_size=32,
show_progress=True
)
# Store in vector database
chunk_metadata = [chunk["metadata"] for chunk in chunks]
self.vector_store.add_documents(chunk_texts, embeddings, chunk_metadata)
print(f"Indexing complete! Indexed {len(chunks)} chunks.")
def retrieve_context(self, query, top_k=3):
"""
Retrieve relevant context for a query.
Args:
query: The user's question
top_k: Number of chunks to retrieve
Returns:
list: List of relevant document chunks
"""
# Generate query embedding
query_embedding = self.embedding_generator.encode(query)
# Search vector store
results = self.vector_store.search(query_embedding, top_k)
return results
def generate_response(self, query, top_k=3, max_new_tokens=512,
temperature=0.7, include_sources=True):
"""
Generate a response using RAG.
This retrieves relevant context and uses it to augment the LLM's
response to the user's query.
Args:
query: The user's question
top_k: Number of context chunks to retrieve
max_new_tokens: Maximum tokens to generate
temperature: Sampling temperature
include_sources: Whether to include source citations
Returns:
dict: Response with text and sources
"""
# Retrieve relevant context
context_chunks = self.retrieve_context(query, top_k)
# Format context for the prompt
context_text = self._format_context(context_chunks)
# Create augmented prompt
augmented_query = self._create_rag_prompt(query, context_text)
# Generate response
response = self.chatbot.generate_response(
augmented_query,
max_new_tokens=max_new_tokens,
temperature=temperature
)
# Prepare result
result = {
"response": response,
"sources": []
}
if include_sources:
result["sources"] = [
{
"text": chunk["document"],
"score": chunk["score"],
"metadata": chunk["metadata"]
}
for chunk in context_chunks
]
return result
def _format_context(self, context_chunks):
"""
Format retrieved context chunks into a readable string.
Args:
context_chunks: List of retrieved chunks
Returns:
str: Formatted context
"""
if not context_chunks:
return "No relevant context found."
formatted_parts = []
for i, chunk in enumerate(context_chunks, 1):
source = chunk["metadata"].get("source", "Unknown")
text = chunk["document"]
formatted_parts.append(f"[Source {i}: {source}]\n{text}")
return "\n\n".join(formatted_parts)
def _create_rag_prompt(self, query, context):
"""
Create a prompt that includes retrieved context.
Args:
query: The user's question
context: The retrieved context
Returns:
str: The augmented prompt
"""
prompt = f"""Based on the following context, please answer the question.
Context:
{context}
Question: {query}
Please provide a detailed answer based on the context above. If the context does not contain enough information to fully answer the question, please state that clearly."""
return prompt
def save_index(self, directory):
"""
Save the vector store index to disk.
Args:
directory: Directory to save the index
"""
self.vector_store.save(directory)
def load_index(self, directory):
"""
Load a previously saved vector store index.
Args:
directory: Directory containing the saved index
"""
self.vector_store.load(directory)
The RAG prompt engineering is critical. We clearly separate the context from the question and instruct the model to base its answer on the provided context. We also tell it to acknowledge when it does not have enough information, which reduces hallucinations.
The source tracking allows users to verify the information and explore the original documents. This transparency is essential for building trust in AI systems, especially in professional or academic contexts.
Advanced RAG Techniques
The basic RAG implementation works well, but there are several advanced techniques that can improve quality further.
One important technique is query expansion. Sometimes the user's question does not match the terminology used in the documents. We can generate alternative phrasings of the question and retrieve documents for each, then combine the results.
Another technique is re-ranking. After initial retrieval, we can use a more sophisticated model to re-score the results based on relevance to the specific question. This is more expensive than vector search but more accurate.
We can also implement hybrid search, combining vector similarity with traditional keyword search. This catches cases where exact keyword matches are important, such as searching for specific names or technical terms.
Here is an implementation of query expansion:
def expand_query(self, query, num_expansions=2):
"""
Generate alternative phrasings of a query.
This helps retrieve relevant documents even when the query
uses different terminology than the documents.
Args:
query: The original query
num_expansions: Number of alternative phrasings to generate
Returns:
list: List of query variations including the original
"""
expansion_prompt = f"""Generate {num_expansions} alternative ways to phrase this question,
keeping the same meaning but using different words:
Original question: {query}
Alternative phrasings:"""
# Generate expansions using the LLM
response = self.chatbot.generate_response(
expansion_prompt,
max_new_tokens=200,
temperature=0.8
)
# Parse the response to extract alternatives
# This is a simple parsing; production code would be more robust
alternatives = [line.strip() for line in response.split('\n')
if line.strip() and not line.strip().startswith('-')]
# Combine with original query
all_queries = [query] + alternatives[:num_expansions]
return all_queries
def retrieve_context_expanded(self, query, top_k=3, num_expansions=2):
"""
Retrieve context using query expansion.
Args:
query: The user's question
top_k: Number of chunks to retrieve per query
num_expansions: Number of query variations to generate
Returns:
list: Deduplicated list of relevant chunks
"""
# Generate query variations
queries = self.expand_query(query, num_expansions)
# Retrieve for each query
all_results = []
seen_texts = set()
for q in queries:
results = self.retrieve_context(q, top_k)
# Deduplicate based on text content
for result in results:
text = result["document"]
if text not in seen_texts:
seen_texts.add(text)
all_results.append(result)
# Sort by score and take top results
all_results.sort(key=lambda x: x["score"])
return all_results[:top_k * 2] # Return more results due to expansion
This query expansion technique uses the LLM itself to generate alternative phrasings. This is a form of self-improvement where the system uses its own capabilities to enhance its performance.
PRODUCTION BEST PRACTICES
Building a working prototype is one thing, but creating a production-ready system requires attention to many additional concerns. Let us discuss the most important best practices for deploying LLM applications.
Error Handling and Robustness
Production systems must handle errors gracefully. Network failures, out-of-memory errors, malformed inputs, and many other issues can occur. We need comprehensive error handling.
import logging
from functools import wraps
import traceback
# Configure logging
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)
def handle_errors(func):
"""
Decorator for comprehensive error handling.
This catches exceptions, logs them, and returns a safe error response
instead of crashing the application.
"""
@wraps(func)
def wrapper(*args, **kwargs):
try:
return func(*args, **kwargs)
except torch.cuda.OutOfMemoryError:
logger.error("GPU out of memory error")
logger.error(traceback.format_exc())
# Clear GPU cache and retry with smaller batch
if torch.cuda.is_available():
torch.cuda.empty_cache()
return {"error": "Out of memory. Try reducing batch size or model size."}
except Exception as e:
logger.error(f"Error in {func.__name__}: {str(e)}")
logger.error(traceback.format_exc())
return {"error": f"An error occurred: {str(e)}"}
return wrapper
This decorator can be applied to any function that might fail. It catches exceptions, logs detailed information for debugging, and returns a safe error response. For GPU memory errors, it even attempts recovery by clearing the cache.
Performance Monitoring
Understanding your system's performance is crucial for optimization and capacity planning. We should track metrics like response time, token generation speed, and memory usage.
import time
from contextlib import contextmanager
class PerformanceMonitor:
"""
Monitor and log performance metrics.
This class tracks timing, memory usage, and other metrics
to help optimize the system.
"""
def __init__(self):
"""Initialize the performance monitor."""
self.metrics = {
"generation_times": [],
"retrieval_times": [],
"embedding_times": [],
"tokens_generated": []
}
@contextmanager
def measure_time(self, operation):
"""
Context manager for timing operations.
Usage:
with monitor.measure_time("generation"):
result = model.generate(...)
"""
start_time = time.time()
try:
yield
finally:
elapsed = time.time() - start_time
self.metrics[f"{operation}_times"].append(elapsed)
logger.info(f"{operation} took {elapsed:.2f} seconds")
def log_memory_usage(self):
"""Log current GPU memory usage."""
if torch.cuda.is_available():
allocated = torch.cuda.memory_allocated() / 1e9
reserved = torch.cuda.memory_reserved() / 1e9
logger.info(f"GPU Memory - Allocated: {allocated:.2f}GB, Reserved: {reserved:.2f}GB")
def get_statistics(self):
"""
Get performance statistics.
Returns:
dict: Statistics for each tracked metric
"""
stats = {}
for metric, values in self.metrics.items():
if values:
stats[metric] = {
"mean": sum(values) / len(values),
"min": min(values),
"max": max(values),
"count": len(values)
}
return stats
This monitor can be integrated throughout the system to track performance. The statistics help identify bottlenecks and guide optimization efforts.
Caching and Optimization
Caching can dramatically improve performance for repeated queries. We can cache embeddings, retrieved contexts, and even generated responses.
from functools import lru_cache
import hashlib
class CachedRAGSystem(RAGSystem):
"""
RAG system with caching for improved performance.
This extends the base RAG system with caching capabilities
for embeddings and retrieved contexts.
"""
def __init__(self, *args, cache_size=1000, **kwargs):
"""
Initialize the cached RAG system.
Args:
cache_size: Maximum number of items to cache
"""
super().__init__(*args, **kwargs)
self.cache_size = cache_size
self.embedding_cache = {}
self.retrieval_cache = {}
def _hash_text(self, text):
"""
Create a hash of text for cache keys.
Args:
text: The text to hash
Returns:
str: Hash of the text
"""
return hashlib.md5(text.encode()).hexdigest()
def retrieve_context(self, query, top_k=3):
"""
Retrieve context with caching.
Args:
query: The user's question
top_k: Number of chunks to retrieve
Returns:
list: List of relevant chunks
"""
# Create cache key
cache_key = f"{self._hash_text(query)}_{top_k}"
# Check cache
if cache_key in self.retrieval_cache:
logger.info("Cache hit for retrieval")
return self.retrieval_cache[cache_key]
# Compute and cache
results = super().retrieve_context(query, top_k)
# Manage cache size
if len(self.retrieval_cache) >= self.cache_size:
# Remove oldest entry
self.retrieval_cache.pop(next(iter(self.retrieval_cache)))
self.retrieval_cache[cache_key] = results
return results
Caching is especially effective for RAG systems because the same questions often get asked multiple times. The cache key includes both the query and the top_k parameter to ensure we return the correct number of results.
Security Considerations
LLM applications have unique security concerns. We need to protect against prompt injection, ensure safe content generation, and handle sensitive data properly.
Prompt injection is when a user crafts input that manipulates the model into ignoring its instructions or revealing system prompts. We can mitigate this by clearly separating user input from system instructions and validating inputs.
class SecureRAGSystem(RAGSystem):
"""
RAG system with security enhancements.
This adds input validation, content filtering, and other
security measures to protect against misuse.
"""
def __init__(self, *args, **kwargs):
"""Initialize the secure RAG system."""
super().__init__(*args, **kwargs)
# Define forbidden patterns
self.forbidden_patterns = [
r"ignore previous instructions",
r"disregard all previous",
r"you are now",
r"new instructions",
]
def validate_input(self, text):
"""
Validate user input for security issues.
Args:
text: The input text to validate
Returns:
bool: True if input is safe, False otherwise
"""
import re
# Check for forbidden patterns
for pattern in self.forbidden_patterns:
if re.search(pattern, text, re.IGNORECASE):
logger.warning(f"Blocked potentially malicious input: {pattern}")
return False
# Check length
if len(text) > 10000:
logger.warning("Input too long")
return False
return True
def generate_response(self, query, **kwargs):
"""
Generate response with input validation.
Args:
query: The user's question
**kwargs: Additional arguments
Returns:
dict: Response or error
"""
# Validate input
if not self.validate_input(query):
return {
"error": "Input validation failed. Please rephrase your question.",
"response": None,
"sources": []
}
# Generate response
return super().generate_response(query, **kwargs)
This validation is not foolproof, as determined attackers can often find ways around filters. However, it provides a reasonable defense against common attack patterns.
Scalability and Deployment
For production deployment, we need to consider how the system will scale to handle multiple concurrent users. This typically involves deploying the system as a web service.
A common architecture uses FastAPI for the web server, with the RAG system running in the background. We can use async processing to handle multiple requests concurrently.
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import asyncio
from concurrent.futures import ThreadPoolExecutor
app = FastAPI(title="RAG API")
# Initialize RAG system (do this once at startup)
rag_system = None
executor = ThreadPoolExecutor(max_workers=4)
class QueryRequest(BaseModel):
"""Request model for queries."""
query: str
top_k: int = 3
max_tokens: int = 512
temperature: float = 0.7
class QueryResponse(BaseModel):
"""Response model for queries."""
response: str
sources: list
@app.on_event("startup")
async def startup_event():
"""Initialize the RAG system on startup."""
global rag_system
rag_system = RAGSystem(
llm_model_name="meta-llama/Llama-2-7b-chat-hf",
use_quantization=True
)
# Load pre-built index
rag_system.load_index("./vector_store")
@app.post("/query", response_model=QueryResponse)
async def query_endpoint(request: QueryRequest):
"""
Query the RAG system.
Args:
request: The query request
Returns:
QueryResponse: The generated response with sources
"""
if rag_system is None:
raise HTTPException(status_code=503, detail="System not initialized")
# Run generation in thread pool to avoid blocking
loop = asyncio.get_event_loop()
result = await loop.run_in_executor(
executor,
lambda: rag_system.generate_response(
request.query,
top_k=request.top_k,
max_new_tokens=request.max_tokens,
temperature=request.temperature
)
)
return QueryResponse(
response=result["response"],
sources=result["sources"]
)
@app.get("/health")
async def health_check():
"""Health check endpoint."""
return {"status": "healthy", "model_loaded": rag_system is not None}
This API design allows multiple clients to query the system concurrently. The thread pool executor prevents the LLM generation from blocking the async event loop. For even better performance, you could use a queue-based system with dedicated worker processes.
COMPLETE PRODUCTION-READY IMPLEMENTATION
Now let us put everything together into a complete, production-ready implementation. This code includes all the components we have discussed, with proper error handling, logging, and optimization.
#!/usr/bin/env python3
"""
Complete RAG System Implementation
A production-ready Retrieval-Augmented Generation system that combines
document processing, vector search, and large language model generation
to answer questions based on a knowledge base.
This implementation supports multiple GPU architectures (NVIDIA CUDA, AMD ROCm,
Intel, Apple MPS) and includes comprehensive error handling, caching, and
performance monitoring.
Author: Michael Stal
License: MIT
"""
import torch
import numpy as np
import faiss
import pickle
import os
import re
import time
import hashlib
import logging
import traceback
from pathlib import Path
from typing import List, Dict, Optional, Tuple, Union
from functools import wraps, lru_cache
from contextlib import contextmanager
import platform
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
from sentence_transformers import SentenceTransformer
# Configure logging
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
handlers=[
logging.FileHandler('rag_system.log'),
logging.StreamHandler()
]
)
logger = logging.getLogger(__name__)
# ============================================================================
# UTILITY FUNCTIONS AND DECORATORS
# ============================================================================
def handle_errors(func):
"""
Decorator for comprehensive error handling.
This catches exceptions, logs them with full tracebacks, and returns
a safe error response instead of crashing the application. Special
handling is provided for GPU out-of-memory errors.
Args:
func: The function to wrap
Returns:
Wrapped function with error handling
"""
@wraps(func)
def wrapper(*args, **kwargs):
try:
return func(*args, **kwargs)
except torch.cuda.OutOfMemoryError:
logger.error("GPU out of memory error in %s", func.__name__)
logger.error(traceback.format_exc())
if torch.cuda.is_available():
torch.cuda.empty_cache()
return {
"error": "GPU out of memory. Try reducing batch size, model size, or enable quantization.",
"success": False
}
except FileNotFoundError as e:
logger.error("File not found in %s: %s", func.__name__, str(e))
return {"error": f"File not found: {str(e)}", "success": False}
except ValueError as e:
logger.error("Value error in %s: %s", func.__name__, str(e))
return {"error": f"Invalid value: {str(e)}", "success": False}
except Exception as e:
logger.error("Unexpected error in %s: %s", func.__name__, str(e))
logger.error(traceback.format_exc())
return {"error": f"An unexpected error occurred: {str(e)}", "success": False}
return wrapper
def get_optimal_device() -> Tuple[torch.device, str]:
"""
Detect and return the best available compute device.
This function checks for GPU availability in order of preference:
CUDA (NVIDIA), ROCm (AMD), MPS (Apple Silicon), then falls back to CPU.
It also logs detailed information about the detected device.
Returns:
tuple: (torch.device, str) - The optimal device and its description
"""
if torch.cuda.is_available():
device = torch.device("cuda")
gpu_name = torch.cuda.get_device_name(0)
gpu_memory = torch.cuda.get_device_properties(0).total_memory / 1e9
device_info = f"CUDA GPU: {gpu_name} ({gpu_memory:.2f} GB)"
logger.info("Using %s", device_info)
return device, device_info
elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
device = torch.device("mps")
device_info = "Apple Silicon MPS"
logger.info("Using %s", device_info)
return device, device_info
elif torch.version.hip is not None:
device = torch.device("cuda")
device_info = "AMD ROCm GPU"
logger.info("Using %s", device_info)
return device, device_info
else:
device = torch.device("cpu")
device_info = f"CPU: {platform.processor()}"
logger.info("Using %s", device_info)
logger.warning("No GPU detected. Inference will be significantly slower.")
return device, device_info
# ============================================================================
# PERFORMANCE MONITORING
# ============================================================================
class PerformanceMonitor:
"""
Monitor and log performance metrics for the RAG system.
This class tracks timing information, memory usage, and other metrics
to help identify bottlenecks and optimize system performance.
Attributes:
metrics (dict): Dictionary storing lists of metric values
"""
def __init__(self):
"""Initialize the performance monitor with empty metric storage."""
self.metrics = {
"generation_times": [],
"retrieval_times": [],
"embedding_times": [],
"indexing_times": [],
"tokens_generated": [],
"documents_processed": []
}
@contextmanager
def measure_time(self, operation: str):
"""
Context manager for timing operations.
This measures the elapsed time for any operation and stores it
in the metrics dictionary. It also logs the timing information.
Args:
operation: Name of the operation being timed
Yields:
None
Example:
with monitor.measure_time("generation"):
result = model.generate(...)
"""
start_time = time.time()
try:
yield
finally:
elapsed = time.time() - start_time
metric_key = f"{operation}_times"
if metric_key in self.metrics:
self.metrics[metric_key].append(elapsed)
logger.info("%s took %.2f seconds", operation, elapsed)
def log_memory_usage(self):
"""
Log current GPU memory usage.
This provides visibility into GPU memory consumption, which is
critical for debugging out-of-memory errors and optimizing
batch sizes.
"""
if torch.cuda.is_available():
allocated = torch.cuda.memory_allocated() / 1e9
reserved = torch.cuda.memory_reserved() / 1e9
logger.info("GPU Memory - Allocated: %.2fGB, Reserved: %.2fGB",
allocated, reserved)
def record_metric(self, metric_name: str, value: float):
"""
Record a custom metric value.
Args:
metric_name: Name of the metric
value: Value to record
"""
if metric_name not in self.metrics:
self.metrics[metric_name] = []
self.metrics[metric_name].append(value)
def get_statistics(self) -> Dict[str, Dict[str, float]]:
"""
Calculate and return statistics for all tracked metrics.
Returns:
dict: Statistics (mean, min, max, count) for each metric
"""
stats = {}
for metric, values in self.metrics.items():
if values:
stats[metric] = {
"mean": sum(values) / len(values),
"min": min(values),
"max": max(values),
"count": len(values),
"total": sum(values)
}
return stats
def print_statistics(self):
"""Print formatted statistics for all metrics."""
stats = self.get_statistics()
logger.info("=== Performance Statistics ===")
for metric, values in stats.items():
logger.info("%s:", metric)
logger.info(" Mean: %.2f, Min: %.2f, Max: %.2f, Count: %d",
values["mean"], values["min"], values["max"], values["count"])
# ============================================================================
# DOCUMENT LOADING
# ============================================================================
class DocumentLoader:
"""
Load documents from various file formats.
This class supports loading text files, PDFs, and markdown files with
appropriate preprocessing for each format. It can load individual files
or entire directories recursively.
Attributes:
supported_extensions (set): Set of supported file extensions
"""
def __init__(self):
"""Initialize the document loader with supported file types."""
self.supported_extensions = {'.txt', '.md', '.pdf', '.markdown'}
@handle_errors
def load_file(self, file_path: Union[str, Path]) -> Dict:
"""
Load a single file and return its content with metadata.
Args:
file_path: Path to the file to load
Returns:
dict: Document dictionary with 'text' and 'metadata' keys
Raises:
FileNotFoundError: If the file does not exist
ValueError: If the file type is not supported
"""
file_path = Path(file_path)
if not file_path.exists():
raise FileNotFoundError(f"File not found: {file_path}")
extension = file_path.suffix.lower()
if extension not in self.supported_extensions:
raise ValueError(f"Unsupported file type: {extension}. "
f"Supported types: {self.supported_extensions}")
logger.info("Loading file: %s", file_path)
# Load based on file type
if extension == '.pdf':
text = self._load_pdf(file_path)
else:
# Text and markdown files
with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
text = f.read()
# Create document dictionary
document = {
"text": text,
"metadata": {
"source": str(file_path),
"filename": file_path.name,
"extension": extension,
"size_bytes": file_path.stat().st_size
}
}
logger.info("Loaded %d characters from %s", len(text), file_path.name)
return document
def _load_pdf(self, file_path: Path) -> str:
"""
Load text from a PDF file.
This uses the pypdf library to extract text from all pages of a PDF.
Page breaks are preserved with double newlines.
Args:
file_path: Path to the PDF file
Returns:
str: Extracted text from the PDF
Raises:
ImportError: If pypdf is not installed
"""
try:
from pypdf import PdfReader
except ImportError:
raise ImportError(
"pypdf is required for PDF support. "
"Install with: pip install pypdf"
)
reader = PdfReader(file_path)
text_parts = []
for page_num, page in enumerate(reader.pages):
text = page.extract_text()
if text.strip():
text_parts.append(text)
logger.debug("Extracted %d characters from page %d",
len(text), page_num + 1)
full_text = "\n\n".join(text_parts)
logger.info("Extracted %d pages, %d total characters",
len(reader.pages), len(full_text))
return full_text
@handle_errors
def load_directory(self, directory_path: Union[str, Path],
recursive: bool = True) -> List[Dict]:
"""
Load all supported documents from a directory.
Args:
directory_path: Path to the directory
recursive: Whether to search subdirectories
Returns:
list: List of document dictionaries
Raises:
NotADirectoryError: If the path is not a directory
"""
directory_path = Path(directory_path)
if not directory_path.is_dir():
raise NotADirectoryError(f"Not a directory: {directory_path}")
logger.info("Loading documents from: %s (recursive=%s)",
directory_path, recursive)
documents = []
# Get all files
if recursive:
files = directory_path.rglob("*")
else:
files = directory_path.glob("*")
# Load each supported file
for file_path in files:
if file_path.is_file() and file_path.suffix.lower() in self.supported_extensions:
try:
doc = self.load_file(file_path)
if doc.get("success", True):
documents.append(doc)
except Exception as e:
logger.error("Error loading %s: %s", file_path, str(e))
logger.info("Loaded %d documents from %s", len(documents), directory_path)
return documents
# ============================================================================
# TEXT CHUNKING
# ============================================================================
class TextChunker:
"""
Split text into semantically coherent chunks.
This class implements intelligent text splitting that respects sentence
boundaries and maintains semantic coherence. It uses a sliding window
approach with overlap to ensure context is preserved across chunks.
Attributes:
chunk_size (int): Target size of each chunk in characters
chunk_overlap (int): Number of characters to overlap between chunks
"""
def __init__(self, chunk_size: int = 500, chunk_overlap: int = 50):
"""
Initialize the text chunker.
Args:
chunk_size: Target size of each chunk in characters
chunk_overlap: Number of characters to overlap between chunks
"""
self.chunk_size = chunk_size
self.chunk_overlap = chunk_overlap
if chunk_overlap >= chunk_size:
logger.warning("Chunk overlap (%d) is >= chunk size (%d). "
"Setting overlap to chunk_size // 4",
chunk_overlap, chunk_size)
self.chunk_overlap = chunk_size // 4
def chunk_text(self, text: str, metadata: Optional[Dict] = None) -> List[Dict]:
"""
Split text into chunks with overlap.
This method splits text at sentence boundaries to maintain semantic
coherence. Each chunk overlaps with the previous one to preserve
context across boundaries.
Args:
text: The text to chunk
metadata: Optional metadata to attach to each chunk
Returns:
list: List of chunk dictionaries with 'text' and 'metadata' keys
"""
if metadata is None:
metadata = {}
# Split into sentences first
sentences = self._split_into_sentences(text)
if not sentences:
logger.warning("No sentences found in text")
return []
chunks = []
current_chunk = []
current_length = 0
for sentence in sentences:
sentence_length = len(sentence)
# If adding this sentence would exceed chunk size
if current_length + sentence_length > self.chunk_size and current_chunk:
# Save current chunk
chunk_text = " ".join(current_chunk)
chunks.append({
"text": chunk_text,
"metadata": {
**metadata,
"chunk_index": len(chunks),
"chunk_size": len(chunk_text)
}
})
# Start new chunk with overlap
overlap_sentences = []
overlap_length = 0
for s in reversed(current_chunk):
if overlap_length + len(s) <= self.chunk_overlap:
overlap_sentences.insert(0, s)
overlap_length += len(s)
else:
break
current_chunk = overlap_sentences
current_length = overlap_length
# Add sentence to current chunk
current_chunk.append(sentence)
current_length += sentence_length
# Add final chunk if not empty
if current_chunk:
chunk_text = " ".join(current_chunk)
chunks.append({
"text": chunk_text,
"metadata": {
**metadata,
"chunk_index": len(chunks),
"chunk_size": len(chunk_text)
}
})
logger.debug("Created %d chunks from text of length %d",
len(chunks), len(text))
return chunks
def _split_into_sentences(self, text: str) -> List[str]:
"""
Split text into sentences using regex patterns.
This is a simple but effective sentence splitter that handles
common cases. For production use with complex documents, consider
using a more sophisticated tokenizer like spaCy or NLTK.
Args:
text: The text to split
Returns:
list: List of sentences
"""
# Pattern matches sentence-ending punctuation followed by whitespace
# Handles common abbreviations like "Dr.", "Mr.", etc.
sentence_pattern = r'(?<!\w\.\w.)(?<![A-Z][a-z]\.)(?<=\.|\?|\!)\s+'
sentences = re.split(sentence_pattern, text)
# Clean up sentences
sentences = [s.strip() for s in sentences if s.strip()]
return sentences
def chunk_documents(self, documents: List[Dict]) -> List[Dict]:
"""
Chunk multiple documents.
Args:
documents: List of document dictionaries
Returns:
list: List of chunk dictionaries
"""
all_chunks = []
for doc in documents:
chunks = self.chunk_text(doc["text"], doc.get("metadata", {}))
all_chunks.extend(chunks)
logger.info("Created %d chunks from %d documents",
len(all_chunks), len(documents))
return all_chunks
# ============================================================================
# EMBEDDING GENERATION
# ============================================================================
class EmbeddingGenerator:
"""
Generate embeddings for text using sentence transformers.
This class handles loading an embedding model and generating vector
representations of text that can be used for similarity search. It
supports batch processing and automatic device placement.
Attributes:
model: The sentence transformer model
device: The device the model is running on
embedding_dim (int): Dimension of the embedding vectors
"""
def __init__(self, model_name: str = "all-MiniLM-L6-v2",
device: Optional[torch.device] = None):
"""
Initialize the embedding generator.
Args:
model_name: Name of the sentence transformer model to use
device: Device to run the model on (auto-detected if None)
"""
if device is None:
self.device, _ = get_optimal_device()
else:
self.device = device
logger.info("Loading embedding model: %s", model_name)
# Load the sentence transformer model
self.model = SentenceTransformer(model_name)
self.model.to(self.device)
# Get embedding dimension
self.embedding_dim = self.model.get_sentence_embedding_dimension()
logger.info("Embedding model loaded. Dimension: %d", self.embedding_dim)
@handle_errors
def encode(self, texts: Union[str, List[str]],
batch_size: int = 32,
show_progress: bool = False) -> np.ndarray:
"""
Generate embeddings for text or list of texts.
This method handles both single texts and batches of texts. It
normalizes embeddings by default to enable efficient cosine
similarity computation via dot product.
Args:
texts: Single text string or list of text strings
batch_size: Number of texts to process at once
show_progress: Whether to show a progress bar
Returns:
numpy.ndarray: Array of embeddings with shape (n_texts, embedding_dim)
"""
# Handle single text input
if isinstance(texts, str):
texts = [texts]
logger.debug("Encoding %d texts with batch size %d",
len(texts), batch_size)
# Generate embeddings
embeddings = self.model.encode(
texts,
batch_size=batch_size,
show_progress_bar=show_progress,
convert_to_numpy=True,
normalize_embeddings=True # Normalize for cosine similarity
)
logger.debug("Generated embeddings with shape %s", embeddings.shape)
return embeddings
# ============================================================================
# VECTOR STORE
# ============================================================================
class VectorStore:
"""
A vector store for efficient similarity search using FAISS.
This class manages a collection of document embeddings and provides
fast similarity search capabilities. It supports saving and loading
the index to/from disk for persistence.
Attributes:
embedding_dim (int): Dimension of the embeddings
use_gpu (bool): Whether to use GPU acceleration
index: The FAISS index
documents (list): List of document texts
metadata (list): List of metadata dictionaries
"""
def __init__(self, embedding_dim: int, use_gpu: bool = False):
"""
Initialize the vector store.
Args:
embedding_dim: Dimension of the embeddings
use_gpu: Whether to use GPU acceleration for search
"""
self.embedding_dim = embedding_dim
self.use_gpu = use_gpu
# Initialize FAISS index
# We use a flat L2 index for exact search
self.index = faiss.IndexFlatL2(embedding_dim)
# Move to GPU if requested and available
if use_gpu and faiss.get_num_gpus() > 0:
res = faiss.StandardGpuResources()
self.index = faiss.index_cpu_to_gpu(res, 0, self.index)
logger.info("Using GPU-accelerated FAISS index")
# Store the actual documents and metadata
self.documents = []
self.metadata = []
@handle_errors
def add_documents(self, documents: List[str],
embeddings: np.ndarray,
metadata: Optional[List[Dict]] = None):
"""
Add documents and their embeddings to the store.
Args:
documents: List of document texts
embeddings: Numpy array of embeddings with shape (n_docs, embedding_dim)
metadata: Optional list of metadata dictionaries
"""
# Validate inputs
if len(documents) != embeddings.shape[0]:
raise ValueError(f"Number of documents ({len(documents)}) does not match "
f"number of embeddings ({embeddings.shape[0]})")
# Ensure embeddings are float32 (FAISS requirement)
embeddings = embeddings.astype(np.float32)
# Add to FAISS index
self.index.add(embeddings)
# Store documents
self.documents.extend(documents)
# Store metadata
if metadata is None:
metadata = [{} for _ in documents]
self.metadata.extend(metadata)
logger.info("Added %d documents. Total documents: %d",
len(documents), len(self.documents))
@handle_errors
def search(self, query_embedding: np.ndarray,
top_k: int = 5) -> List[Dict]:
"""
Search for the most similar documents to a query.
This performs k-nearest neighbor search in the embedding space
to find the most relevant documents.
Args:
query_embedding: The query embedding vector
top_k: Number of results to return
Returns:
list: List of dictionaries with 'document', 'score', and 'metadata' keys
"""
# Ensure query is the right shape and type
query_embedding = query_embedding.astype(np.float32)
if len(query_embedding.shape) == 1:
query_embedding = query_embedding.reshape(1, -1)
# Validate top_k
if top_k > len(self.documents):
logger.warning("top_k (%d) is greater than number of documents (%d). "
"Returning all documents.", top_k, len(self.documents))
top_k = len(self.documents)
# Search the index
distances, indices = self.index.search(query_embedding, top_k)
# Prepare results
results = []
for dist, idx in zip(distances[0], indices[0]):
if idx < len(self.documents): # Valid index
results.append({
"document": self.documents[idx],
"score": float(dist),
"metadata": self.metadata[idx]
})
logger.debug("Found %d results for query", len(results))
return results
@handle_errors
def save(self, directory: Union[str, Path]):
"""
Save the vector store to disk.
This saves both the FAISS index and the document/metadata store
to the specified directory.
Args:
directory: Directory to save the store
"""
directory = Path(directory)
directory.mkdir(parents=True, exist_ok=True)
# Save FAISS index
if self.use_gpu:
# Move back to CPU for saving
cpu_index = faiss.index_gpu_to_cpu(self.index)
faiss.write_index(cpu_index, str(directory / "index.faiss"))
else:
faiss.write_index(self.index, str(directory / "index.faiss"))
# Save documents and metadata
with open(directory / "documents.pkl", "wb") as f:
pickle.dump({
"documents": self.documents,
"metadata": self.metadata,
"embedding_dim": self.embedding_dim
}, f)
logger.info("Vector store saved to %s", directory)
@handle_errors
def load(self, directory: Union[str, Path]):
"""
Load the vector store from disk.
Args:
directory: Directory containing the saved store
"""
directory = Path(directory)
if not directory.exists():
raise FileNotFoundError(f"Directory not found: {directory}")
# Load FAISS index
self.index = faiss.read_index(str(directory / "index.faiss"))
# Move to GPU if requested
if self.use_gpu and faiss.get_num_gpus() > 0:
res = faiss.StandardGpuResources()
self.index = faiss.index_cpu_to_gpu(res, 0, self.index)
# Load documents and metadata
with open(directory / "documents.pkl", "rb") as f:
data = pickle.load(f)
self.documents = data["documents"]
self.metadata = data["metadata"]
# Verify embedding dimension matches
if data["embedding_dim"] != self.embedding_dim:
logger.warning("Loaded embedding dimension (%d) does not match "
"current dimension (%d)",
data["embedding_dim"], self.embedding_dim)
logger.info("Loaded vector store with %d documents from %s",
len(self.documents), directory)
# ============================================================================
# LLM CHATBOT
# ============================================================================
class LLMChatbot:
"""
A simple but powerful LLM chatbot implementation.
This class encapsulates model loading, conversation management, and text
generation with proper error handling and device management. It supports
various quantization options to reduce memory usage.
Attributes:
model_name (str): HuggingFace model identifier
device: Torch device the model is running on
device_info (str): Human-readable device description
tokenizer: The model's tokenizer
model: The language model
conversation_history (list): List of conversation messages
"""
def __init__(self, model_name: str = "meta-llama/Llama-2-7b-chat-hf",
device: Optional[torch.device] = None,
use_quantization: bool = False,
quantization_bits: int = 4):
"""
Initialize the chatbot with a specified model.
Args:
model_name: HuggingFace model identifier
device: Torch device to use (auto-detected if None)
use_quantization: Whether to use quantization to save memory
quantization_bits: Number of bits for quantization (4 or 8)
"""
self.model_name = model_name
# Detect device if not provided
if device is None:
self.device, self.device_info = get_optimal_device()
else:
self.device = device
self.device_info = str(device)
logger.info("Initializing LLM chatbot")
logger.info("Model: %s", model_name)
logger.info("Device: %s", self.device_info)
logger.info("Quantization: %s", "enabled" if use_quantization else "disabled")
# Configure quantization if requested
quantization_config = None
if use_quantization:
if quantization_bits == 4:
quantization_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_compute_dtype=torch.float16,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_type="nf4"
)
logger.info("Using 4-bit NF4 quantization")
elif quantization_bits == 8:
quantization_config = BitsAndBytesConfig(
load_in_8bit=True
)
logger.info("Using 8-bit quantization")
else:
raise ValueError(f"Unsupported quantization bits: {quantization_bits}. "
"Use 4 or 8.")
# Load tokenizer
logger.info("Loading tokenizer...")
self.tokenizer = AutoTokenizer.from_pretrained(
model_name,
trust_remote_code=True
)
# Ensure tokenizer has a pad token
if self.tokenizer.pad_token is None:
self.tokenizer.pad_token = self.tokenizer.eos_token
logger.debug("Set pad_token to eos_token")
# Load model
logger.info("Loading model... (this may take a few minutes)")
self.model = AutoModelForCausalLM.from_pretrained(
model_name,
quantization_config=quantization_config,
device_map="auto",
trust_remote_code=True,
torch_dtype=torch.float16 if not use_quantization else None,
low_cpu_mem_usage=True
)
# Set model to evaluation mode
self.model.eval()
# Initialize conversation history
self.conversation_history = []
logger.info("Model loaded successfully!")
@handle_errors
def generate_response(self, user_input: str,
max_new_tokens: int = 512,
temperature: float = 0.7,
top_p: float = 0.9,
top_k: int = 50,
repetition_penalty: float = 1.1) -> str:
"""
Generate a response to user input.
This method handles the complete generation pipeline: formatting the
conversation, tokenizing, generating tokens, and decoding the result.
Args:
user_input: The user's message
max_new_tokens: Maximum number of tokens to generate
temperature: Sampling temperature (higher = more random)
top_p: Nucleus sampling parameter
top_k: Top-k sampling parameter
repetition_penalty: Penalty for repeating tokens
Returns:
str: The generated response
"""
# Add user input to conversation history
self.conversation_history.append({
"role": "user",
"content": user_input
})
# Format the conversation for the model
prompt = self._format_conversation()
# Tokenize the prompt
inputs = self.tokenizer(
prompt,
return_tensors="pt",
padding=True,
truncation=True,
max_length=2048
).to(self.device)
input_length = inputs["input_ids"].shape[1]
logger.debug("Input length: %d tokens", input_length)
# Generate response
with torch.no_grad():
outputs = self.model.generate(
**inputs,
max_new_tokens=max_new_tokens,
temperature=temperature,
top_p=top_p,
top_k=top_k,
repetition_penalty=repetition_penalty,
do_sample=True,
pad_token_id=self.tokenizer.pad_token_id,
eos_token_id=self.tokenizer.eos_token_id
)
# Decode the generated tokens
generated_text = self.tokenizer.decode(
outputs[0],
skip_special_tokens=True
)
# Extract only the new response
response = self._extract_response(generated_text, prompt)
# Add assistant response to history
self.conversation_history.append({
"role": "assistant",
"content": response
})
output_length = outputs.shape[1] - input_length
logger.debug("Generated %d new tokens", output_length)
return response
def _format_conversation(self) -> str:
"""
Format the conversation history into a prompt string.
Different models use different conversation formats. This method
handles the LLaMA 2 chat format and provides a generic fallback.
Returns:
str: Formatted prompt string
"""
if "llama-2" in self.model_name.lower() and "chat" in self.model_name.lower():
# LLaMA 2 chat format
formatted_parts = []
for message in self.conversation_history:
if message["role"] == "system":
formatted_parts.append(f"<<SYS>>\n{message['content']}\n<</SYS>>")
elif message["role"] == "user":
formatted_parts.append(f"[INST] {message['content']} [/INST]")
elif message["role"] == "assistant":
formatted_parts.append(message["content"])
return " ".join(formatted_parts)
else:
# Generic format for other models
formatted_parts = []
for message in self.conversation_history:
role = message["role"].capitalize()
content = message["content"]
formatted_parts.append(f"{role}: {content}")
return "\n".join(formatted_parts) + "\nAssistant:"
def _extract_response(self, generated_text: str, prompt: str) -> str:
"""
Extract the new response from the generated text.
The model generates a continuation of the prompt, so we need to
remove the prompt portion to get just the assistant's response.
Args:
generated_text: The full generated text
prompt: The original prompt
Returns:
str: Just the new response
"""
# Remove the prompt from the generated text
if generated_text.startswith(prompt):
response = generated_text[len(prompt):].strip()
else:
# Fallback: try to find where the response starts
response = generated_text.strip()
# Clean up any trailing special tokens
for token in ["</s>", "<|endoftext|>", "[/INST]", "<|im_end|>"]:
response = response.replace(token, "")
return response.strip()
def set_system_prompt(self, system_prompt: str):
"""
Set a system prompt that defines the assistant's behavior.
Args:
system_prompt: Instructions for the assistant's behavior
"""
# Remove any existing system prompt
self.conversation_history = [
msg for msg in self.conversation_history
if msg["role"] != "system"
]
# Add new system prompt at the beginning
self.conversation_history.insert(0, {
"role": "system",
"content": system_prompt
})
logger.info("System prompt set")
def clear_history(self):
"""
Clear the conversation history.
This preserves any system prompt that has been set.
"""
system_prompts = [
msg for msg in self.conversation_history
if msg["role"] == "system"
]
self.conversation_history = system_prompts
logger.info("Conversation history cleared")
def get_history(self) -> List[Dict]:
"""
Get the current conversation history.
Returns:
list: Copy of the conversation history
"""
return self.conversation_history.copy()
def trim_history(self, max_messages: int = 10):
"""
Trim conversation history to the most recent messages.
This prevents the context from growing too long and exceeding
the model's maximum context length. System prompts are always preserved.
Args:
max_messages: Maximum number of message pairs to keep
"""
# Always keep system prompts
system_prompts = [
msg for msg in self.conversation_history
if msg["role"] == "system"
]
# Get non-system messages
other_messages = [
msg for msg in self.conversation_history
if msg["role"] != "system"
]
# Keep only the most recent messages
if len(other_messages) > max_messages * 2:
other_messages = other_messages[-(max_messages * 2):]
logger.info("Trimmed conversation history to %d messages",
len(other_messages))
# Reconstruct history
self.conversation_history = system_prompts + other_messages
# ============================================================================
# COMPLETE RAG SYSTEM
# ============================================================================
class RAGSystem:
"""
A complete Retrieval-Augmented Generation system.
This class combines document loading, embedding generation, vector search,
and LLM generation into a cohesive system that can answer questions based
on a knowledge base. It includes performance monitoring, caching, and
comprehensive error handling.
Attributes:
document_loader: Document loading component
text_chunker: Text chunking component
embedding_generator: Embedding generation component
vector_store: Vector storage and search component
chatbot: LLM chatbot component
monitor: Performance monitoring component
"""
def __init__(self,
llm_model_name: str = "meta-llama/Llama-2-7b-chat-hf",
embedding_model_name: str = "all-MiniLM-L6-v2",
chunk_size: int = 500,
chunk_overlap: int = 50,
device: Optional[torch.device] = None,
use_quantization: bool = False):
"""
Initialize the RAG system.
Args:
llm_model_name: Name of the LLM to use
embedding_model_name: Name of the embedding model
chunk_size: Size of text chunks in characters
chunk_overlap: Overlap between chunks in characters
device: Device to use for models
use_quantization: Whether to quantize the LLM
"""
logger.info("=" * 70)
logger.info("Initializing RAG System")
logger.info("=" * 70)
# Initialize performance monitor
self.monitor = PerformanceMonitor()
# Initialize components
logger.info("Initializing document loader...")
self.document_loader = DocumentLoader()
logger.info("Initializing text chunker...")
self.text_chunker = TextChunker(chunk_size, chunk_overlap)
logger.info("Initializing embedding generator...")
self.embedding_generator = EmbeddingGenerator(embedding_model_name, device)
logger.info("Initializing vector store...")
self.vector_store = VectorStore(
self.embedding_generator.embedding_dim,
use_gpu=False # Keep embeddings on CPU to save GPU memory for LLM
)
logger.info("Initializing LLM chatbot...")
self.chatbot = LLMChatbot(llm_model_name, device, use_quantization)
# Set a default system prompt for RAG
self.chatbot.set_system_prompt(
"You are a helpful and knowledgeable assistant. When answering questions, "
"use the provided context to give accurate and detailed responses. "
"If the context does not contain enough information to fully answer the "
"question, clearly state what information is missing or uncertain. "
"Always cite the source of your information when possible. "
"Be concise but thorough in your explanations."
)
logger.info("=" * 70)
logger.info("RAG System initialized successfully!")
logger.info("=" * 70)
@handle_errors
def index_documents(self, source: Union[str, Path],
is_directory: bool = False,
recursive: bool = True):
"""
Index documents from a file or directory.
This processes documents, creates embeddings, and stores them in the
vector database for later retrieval. The process includes loading,
chunking, embedding generation, and storage.
Args:
source: Path to a file or directory
is_directory: Whether the source is a directory
recursive: If directory, whether to search recursively
"""
logger.info("=" * 70)
logger.info("Starting document indexing")
logger.info("Source: %s", source)
logger.info("=" * 70)
with self.monitor.measure_time("indexing"):
# Load documents
if is_directory:
documents = self.document_loader.load_directory(source, recursive)
else:
doc = self.document_loader.load_file(source)
documents = [doc] if doc.get("success", True) else []
if not documents:
logger.warning("No documents found to index")
return {"success": False, "error": "No documents found"}
self.monitor.record_metric("documents_processed", len(documents))
# Chunk documents
logger.info("Chunking documents...")
chunks = self.text_chunker.chunk_documents(documents)
if not chunks:
logger.warning("No chunks created from documents")
return {"success": False, "error": "No chunks created"}
# Generate embeddings
logger.info("Generating embeddings for %d chunks...", len(chunks))
chunk_texts = [chunk["text"] for chunk in chunks]
with self.monitor.measure_time("embedding"):
embeddings = self.embedding_generator.encode(
chunk_texts,
batch_size=32,
show_progress=True
)
# Store in vector database
logger.info("Storing embeddings in vector database...")
chunk_metadata = [chunk["metadata"] for chunk in chunks]
self.vector_store.add_documents(chunk_texts, embeddings, chunk_metadata)
logger.info("=" * 70)
logger.info("Indexing complete!")
logger.info("Documents: %d", len(documents))
logger.info("Chunks: %d", len(chunks))
logger.info("=" * 70)
return {
"success": True,
"documents": len(documents),
"chunks": len(chunks)
}
@handle_errors
def retrieve_context(self, query: str, top_k: int = 3) -> List[Dict]:
"""
Retrieve relevant context for a query.
This generates an embedding for the query and searches the vector
store for the most similar document chunks.
Args:
query: The user's question
top_k: Number of chunks to retrieve
Returns:
list: List of relevant document chunks with scores
"""
with self.monitor.measure_time("retrieval"):
# Generate query embedding
query_embedding = self.embedding_generator.encode(query)
# Search vector store
results = self.vector_store.search(query_embedding, top_k)
logger.info("Retrieved %d context chunks for query", len(results))
return results
@handle_errors
def generate_response(self, query: str,
top_k: int = 3,
max_new_tokens: int = 512,
temperature: float = 0.7,
include_sources: bool = True) -> Dict:
"""
Generate a response using RAG.
This is the main entry point for the RAG system. It retrieves relevant
context and uses it to augment the LLM's response to the user's query.
Args:
query: The user's question
top_k: Number of context chunks to retrieve
max_new_tokens: Maximum tokens to generate
temperature: Sampling temperature
include_sources: Whether to include source citations
Returns:
dict: Response with 'response', 'sources', and metadata
"""
logger.info("Processing query: %s", query[:100])
with self.monitor.measure_time("generation"):
# Retrieve relevant context
context_chunks = self.retrieve_context(query, top_k)
# Format context for the prompt
context_text = self._format_context(context_chunks)
# Create augmented prompt
augmented_query = self._create_rag_prompt(query, context_text)
# Generate response
response = self.chatbot.generate_response(
augmented_query,
max_new_tokens=max_new_tokens,
temperature=temperature
)
# Prepare result
result = {
"success": True,
"response": response,
"sources": [],
"query": query,
"num_sources": len(context_chunks)
}
if include_sources:
result["sources"] = [
{
"text": chunk["document"][:200] + "...", # Truncate for display
"full_text": chunk["document"],
"score": chunk["score"],
"metadata": chunk["metadata"]
}
for chunk in context_chunks
]
logger.info("Response generated successfully")
return result
def _format_context(self, context_chunks: List[Dict]) -> str:
"""
Format retrieved context chunks into a readable string.
Args:
context_chunks: List of retrieved chunks
Returns:
str: Formatted context
"""
if not context_chunks:
return "No relevant context found in the knowledge base."
formatted_parts = []
for i, chunk in enumerate(context_chunks, 1):
source = chunk["metadata"].get("source", "Unknown")
filename = chunk["metadata"].get("filename", "Unknown")
text = chunk["document"]
formatted_parts.append(
f"[Source {i}: {filename}]\n{text}\n"
)
return "\n".join(formatted_parts)
def _create_rag_prompt(self, query: str, context: str) -> str:
"""
Create a prompt that includes retrieved context.
This carefully structures the prompt to make it clear to the model
what information it should use and how to cite sources.
Args:
query: The user's question
context: The retrieved context
Returns:
str: The augmented prompt
"""
prompt = f"""Based on the following context from the knowledge base, please answer the question.
Context:
{context}
Question: {query}
Instructions:
- Provide a detailed and accurate answer based primarily on the context above
- If the context contains relevant information, use it to support your answer
- If the context does not contain enough information, clearly state what is missing
- Cite specific sources when making claims (e.g., "According to Source 1...")
- Be concise but thorough
Answer:"""
return prompt
@handle_errors
def save_index(self, directory: Union[str, Path]):
"""
Save the vector store index to disk.
Args:
directory: Directory to save the index
"""
self.vector_store.save(directory)
logger.info("Index saved successfully")
@handle_errors
def load_index(self, directory: Union[str, Path]):
"""
Load a previously saved vector store index.
Args:
directory: Directory containing the saved index
"""
self.vector_store.load(directory)
logger.info("Index loaded successfully")
def get_statistics(self) -> Dict:
"""
Get performance statistics for the system.
Returns:
dict: Performance statistics
"""
return self.monitor.get_statistics()
def print_statistics(self):
"""Print formatted performance statistics."""
self.monitor.print_statistics()
# ============================================================================
# EXAMPLE USAGE AND DEMONSTRATION
# ============================================================================
def main():
"""
Demonstrate the RAG system with example usage.
This function shows how to initialize the system, index documents,
and generate responses to queries.
"""
logger.info("=" * 70)
logger.info("RAG SYSTEM DEMONSTRATION")
logger.info("=" * 70)
# Initialize the RAG system
# For this demo, we'll use a smaller model if available
# In production, use the full model name
rag = RAGSystem(
llm_model_name="meta-llama/Llama-2-7b-chat-hf",
embedding_model_name="all-MiniLM-L6-v2",
chunk_size=500,
chunk_overlap=50,
use_quantization=True # Enable quantization to reduce memory usage
)
# Example 1: Index a single document
logger.info("\n" + "=" * 70)
logger.info("EXAMPLE 1: Indexing a single document")
logger.info("=" * 70)
# Create a sample document
sample_doc_path = Path("sample_document.txt")
sample_content = """
Artificial Intelligence and Machine Learning
Artificial Intelligence (AI) is the simulation of human intelligence processes
by machines, especially computer systems. These processes include learning,
reasoning, and self-correction.
Machine Learning (ML) is a subset of AI that focuses on the development of
algorithms that can learn from and make predictions or decisions based on data.
ML algorithms build a model based on sample data, known as training data, in
order to make predictions or decisions without being explicitly programmed.
Deep Learning is a subset of machine learning that uses neural networks with
multiple layers. These deep neural networks can learn hierarchical representations
of data, making them particularly effective for tasks like image recognition,
natural language processing, and speech recognition.
Natural Language Processing (NLP) is a branch of AI that helps computers
understand, interpret, and manipulate human language. NLP draws from many
disciplines, including computer science and computational linguistics.
"""
with open(sample_doc_path, "w") as f:
f.write(sample_content)
# Index the document
result = rag.index_documents(sample_doc_path, is_directory=False)
logger.info("Indexing result: %s", result)
# Example 2: Query the system
logger.info("\n" + "=" * 70)
logger.info("EXAMPLE 2: Querying the RAG system")
logger.info("=" * 70)
queries = [
"What is machine learning?",
"How does deep learning differ from traditional machine learning?",
"What is the relationship between AI and NLP?"
]
for query in queries:
logger.info("\nQuery: %s", query)
response = rag.generate_response(
query,
top_k=2,
max_new_tokens=200,
temperature=0.7
)
if response.get("success"):
logger.info("\nResponse: %s", response["response"])
logger.info("\nNumber of sources used: %d", response["num_sources"])
if response.get("sources"):
logger.info("\nSources:")
for i, source in enumerate(response["sources"], 1):
logger.info(" %d. %s (score: %.4f)",
i, source["text"], source["score"])
else:
logger.error("Error: %s", response.get("error"))
# Example 3: Save and load the index
logger.info("\n" + "=" * 70)
logger.info("EXAMPLE 3: Saving and loading the index")
logger.info("=" * 70)
index_dir = Path("vector_store_index")
rag.save_index(index_dir)
logger.info("Index saved to %s", index_dir)
# Create a new RAG system and load the index
rag2 = RAGSystem(
llm_model_name="meta-llama/Llama-2-7b-chat-hf",
use_quantization=True
)
rag2.load_index(index_dir)
logger.info("Index loaded successfully")
# Test with a query
test_query = "What is artificial intelligence?"
response = rag2.generate_response(test_query, top_k=2)
if response.get("success"):
logger.info("\nTest query: %s", test_query)
logger.info("Response: %s", response["response"])
# Print performance statistics
logger.info("\n" + "=" * 70)
logger.info("PERFORMANCE STATISTICS")
logger.info("=" * 70)
rag.print_statistics()
# Cleanup
sample_doc_path.unlink()
logger.info("\n" + "=" * 70)
logger.info("DEMONSTRATION COMPLETE")
logger.info("=" * 70)
if __name__ == "__main__":
main()
CONCLUSION AND NEXT STEPS
Congratulations! You have now built a complete, production-ready RAG system from scratch. This system combines the power of large language models with the precision of information retrieval to create an AI assistant that can answer questions based on your own documents.
Let us recap what we have accomplished. We started by understanding the fundamentals of LLMs and RAG. We learned how embeddings represent text as vectors and how vector databases enable fast similarity search. We built a simple chatbot and progressively added RAG capabilities. We implemented best practices including error handling, performance monitoring, caching, and security measures. Finally, we created a complete system that supports multiple GPU architectures and can be deployed in production.
The system we built includes several key components. The document loader handles various file formats and can process entire directories. The text chunker splits documents intelligently while preserving semantic coherence. The embedding generator creates vector representations of text using state-of-the-art models. The vector store provides fast similarity search using FAISS. The LLM chatbot generates responses with proper conversation management. The RAG system ties everything together into a cohesive whole.
For next steps, consider these enhancements. You could implement more sophisticated chunking strategies that understand document structure like headings and sections. You could add support for more file formats like Word documents, HTML, or structured data. You could implement hybrid search combining vector similarity with keyword matching. You could add re-ranking to improve retrieval quality. You could implement query expansion to handle terminology mismatches. You could add conversation memory that persists across sessions. You could build a web interface using FastAPI or Streamlit. You could implement user authentication and access control. You could add logging and monitoring for production deployment. You could optimize for specific domains by fine-tuning the embedding model.
The field of LLM applications is evolving rapidly. New models, techniques, and best practices emerge constantly. The foundation you have built here will serve you well as you explore more advanced topics. You now understand the core concepts and have working code that you can adapt and extend for your specific needs.
Remember that building AI systems is as much art as science. Experiment with different models, parameters, and techniques. Measure performance carefully. Listen to user feedback. Iterate and improve continuously. The most successful AI applications are those that solve real problems for real users.
Thank you for following this tutorial. I hope it has given you both the knowledge and the confidence to build your own LLM applications. The code provided is production-ready and follows industry best practices. You can use it as-is or as a starting point for your own projects. Good luck with your AI journey!
No comments:
Post a Comment