A straightforward implementation of Retrieval-Augmented Generation (RAG) using ChromaDB vector database and Google's Gemini models for intelligent document querying.
This project demonstrates a basic RAG pipeline that processes PDF documents, creates vector embeddings, and enables intelligent querying using semantic search. The system extracts text from PDFs, chunks it into meaningful paragraphs, stores them in a vector database, and retrieves relevant context to answer user queries.
Naive RAG Pipeline showing the flow from query to answer generation
- PDF Text Extraction: Automated extraction and preprocessing of text from PDF documents
- Smart Text Chunking: Intelligent paragraph segmentation for optimal retrieval
- Vector Storage: Efficient storage using ChromaDB with Google's embedding models
- Semantic Search: Context-aware retrieval based on query similarity
- Contextual Responses: AI-generated answers using retrieved relevant passages
- Python 3.7+
- Google API Key for Gemini models
- PDF documents to process
- Install required packages:
pip install -U chromadb google-genai PyPDF2- Set up Google API Key:
import os
os.environ["GOOGLE_API_KEY"] = "your_api_key_here"- Initialize ChromaDB Client:
import chromadb
from chromadb.utils.embedding_functions.google_embedding_function import GoogleGenerativeAiEmbeddingFunction
chroma_client = chromadb.PersistentClient("/db/")
google_collection = chroma_client.create_collection(
name="persistent_google_collection",
embedding_function=GoogleGenerativeAiEmbeddingFunction(
api_key=os.getenv("GOOGLE_API_KEY"),
model_name="gemini-embedding-001"
)
)- Process and Store Documents:
# Extract and preprocess text from PDF
reader = PdfReader("/path/to/your/document.pdf")
text = ""
for page in reader.pages:
text += page.extract_text() + "\n"
# Clean and chunk text
# Remove hyperlinks
text = re.sub(r'http\S+|www\S+', '', text)
# Remove unwanted punctuation (keep .,!? for readability)
text = re.sub(r'[^\w\s.,!?]', ' ', text)
# Replace multiple spaces/newlines with a single space
text = re.sub(r'\s+', ' ', text)
# Strip leading/trailing spaces
text = text.strip()
# Split into paragraphs (heuristic: double newlines OR section keywords)
paragraphs = re.split(r'(?<=\.)(?=\s+[A-Z])', text)
# Final paragraphs
clean_paragraphs = [p.strip() for p in paragraphs if p.strip()]
# Store in vector database
for i, para in enumerate(clean_paragraphs):
google_collection.upsert(documents=[para], ids=[f"id{i}"])- Query the System:
query = "What is taxonomy?"
answer = generate(query)
print(answer)simple-rag-implementation/
│
├── Simple_RAG_implementation.ipynb # Main notebook with complete implementation
├── README.md # This file
├── requirements.txt # Python dependencies
└── bio1.pdf # Class 11 Biology Textbook Chapter 1
- Hyperlink Removal: Automatically removes URLs and web links
- Punctuation Handling: Preserves essential punctuation (.,!?) for readability
- Whitespace Normalization: Consolidates multiple spaces/newlines
- Paragraph Segmentation: Smart splitting based on sentence endings and capitalization
- Embedding Model:
gemini-embedding-001 - Similarity Search: Top 5 most relevant chunks
- Storage: Persistent ChromaDB instance
- Model:
gemma-3-1b-it - Temperature: 0.2 (focused responses)
- Max Tokens: 500
- Context Window: Top 5 retrieved passages
# Basic factual questions
generate("What is biology?")
generate("Define taxonomy")
generate("What are the taxonomic categories?")
# Specific information retrieval
generate("Who is Ernst Mayr?")
generate("What is binomial nomenclature?")
generate("Explain the classification hierarchy")- Document Processing: PDF text is extracted, cleaned, and segmented into meaningful paragraphs
- Vector Embedding: Each paragraph is converted to a vector using Google's embedding model
- Storage: Vectors are stored in ChromaDB with persistent storage
- Query Processing: User queries are embedded and matched against stored vectors
- Context Retrieval: Most similar passages are retrieved based on semantic similarity
- Response Generation: Retrieved context is fed to the language model for answer generation
Built with ❤️ for educational purposes and learning RAG implementations
