Build a RAG Chatbot with FAISS
Build a local RAG pipeline using PyTorch, Hugging Face, and FAISS.
Introduction
30 Second Summary
Have you ever asked a chatbot a question and gotten a confident but completely wrong answer? The problem is that most language models can only draw on their training data, which might be outdated or missing the specific information you need.
In this project, you will build a Retrieval-Augmented Generation (RAG) chatbot that answers questions using your own documents as its knowledge base. Your chatbot will search for relevant information first, then generate answers grounded in what it actually found.
What You'll Build
A local RAG chatbot that ingests text documents, converts them into vector embeddings, indexes them with FAISS for fast similarity search, and uses a Hugging Face Transformers model to answer questions grounded in retrieved context.
By the end of this project, you'll have:
- A document chunking and embedding pipeline that turns plain text files into searchable vectors using sentence-transformers.
- A FAISS vector index that performs fast cosine similarity search across your knowledge base.
- An interactive Q&A loop powered by a local text-generation model that produces context-grounded answers with source attribution.
- Secret Mission: Add a relevance threshold so the chatbot says "I don't know" when the knowledge base has no good answer.
Are there any prerequisites?
You only need basic Python knowledge (variables, functions, and loops). No machine learning experience is required.
Everything runs locally on your CPU. No GPU, no cloud account, and no API keys needed.
Set Up Your Python Environment
Your RAG chatbot depends on several Python libraries working together. PyTorch provides the tensor operations, sentence-transformers handles embedding generation, FAISS powers vector search, and Hugging Face Transformers runs the text generation model.
To keep these dependencies isolated from other Python projects on your system, you'll create a virtual environment. Then you'll install everything your chatbot needs.
In this step, get ready to:
- Create and activate a Python virtual environment.
- Install PyTorch, Hugging Face Transformers, sentence-transformers, and FAISS.
- Verify all packages import successfully.
Create a project folder and virtual environment
A virtual environment is an isolated copy of Python where you can install packages without affecting the rest of your system. This keeps your RAG project's dependencies separate from anything else.
- Open a terminal window. Press Cmd+Space (macOS) or the Windows key (Windows) to open your search bar.
- Type Terminal (macOS) or cmd (Windows) and press Enter to open it.
- Move to your Desktop by running this command:
cd ~/Desktop
- Create your project folder and move into it by running these commands:
mkdir rag-chatbot
cd rag-chatbot
- Confirm the folder exists by running:
ls
You should see an empty output since the folder is brand new.
- Create the virtual environment by running this command:
python -m venv rag-env
What does this command do?
python -m venv rag-env tells Python to run its built-in venv module and create a new virtual environment in a folder called rag-env.
This creates an isolated Python installation inside rag-env/ with its own pip and package directory.
- Activate the virtual environment. The command differs by operating system:
macOS
- Activate by running:
source rag-env/bin/activate
Windows
- Activate by running:
rag-env\Scripts\activate
You should see (rag-env) appear at the beginning of your terminal prompt. This confirms the virtual environment is active.
Don't see (rag-env) in your prompt?
Make sure you ran the activate command from inside the rag-chatbot folder where you created the environment.
On Windows, if you get a permissions error in PowerShell, try using Command Prompt (cmd) instead.
help me activate my Python virtual environment
Install the required packages
Your RAG chatbot uses four main libraries. You'll install them in two commands because PyTorch needs a special index URL to pull the CPU-only build.
- Install PyTorch (CPU version) by running this command:
pip install torch --index-url https://download.pytorch.org/whl/cpu
Why the CPU-only build?
PyTorch ships with GPU support by default, which makes the download much larger (several GB). The CPU build is smaller and works on any machine.
The --index-url flag tells pip to fetch packages from PyTorch's CPU-specific package index instead of the default.
This download may take a minute or two depending on your internet speed.
- Install the remaining packages by running this command:
pip install transformers sentence-transformers faiss-cpu numpy
What are these packages?
- transformers is the Hugging Face library for loading pre-trained language models.
- sentence-transformers provides an easy API for generating text embeddings.
- faiss-cpu is Facebook's library for fast similarity search over vectors.
- numpy handles the numerical arrays that FAISS and PyTorch operate on.
Wait for the installation to finish. You should see a success message for each package.
Seeing installation errors?
Make sure (rag-env) is visible in your terminal prompt. If it is not, you are installing into your system Python instead of the virtual environment.
If you see a network timeout error, try running the command again. Large packages sometimes fail on slow connections.
help me fix my pip install errors
Verify the installation
Before moving on, confirm that all four libraries import without errors.
- Test that all packages are installed correctly by running this command:
python -c "import torch; import transformers; from sentence_transformers import SentenceTransformer; import faiss; print('PyTorch:', torch.__version__); print('Transformers:', transformers.__version__); print('All imports successful!')"
What should I see?
You should see two version numbers printed (one for PyTorch, one for Transformers) followed by the message All imports successful!.
The exact version numbers may vary, but as long as no errors appear, your environment is ready.
Seeing a ModuleNotFoundError?
Double-check that your virtual environment is still active. Look for (rag-env) at the start of your terminal prompt.
If you accidentally closed your terminal, navigate back to the rag-chatbot folder and reactivate the environment before retrying.
help me fix my ModuleNotFoundError
Your Python environment is locked and loaded. Next up, you'll create a knowledge base of text documents and build the chunking logic that turns them into bite-sized pieces for your chatbot to search through.
Create Your Knowledge Base and Chunk Documents
Your environment is ready with all the packages installed. Now it's time to give your chatbot something to work with.
A RAG chatbot retrieves relevant text from a knowledge base before generating answers. In this step, you'll create that knowledge base from plain text files and write a chunker that splits them into searchable pieces.
In this step, get ready to:
- Create a data/ directory and knowledge base files.
- Build a document loader and text chunker in rag.py.
- Test the chunking by running the script.
Create the project folder and knowledge base files
Your chatbot needs a collection of documents to search through. You'll create a small knowledge base of plain text files that the chatbot can reference when answering questions.
- In your terminal (with rag-env activated), move to your Desktop by running this command:
cd ~/Desktop
- Create your project folder and move into it by running these commands:
mkdir rag-chatbot
cd rag-chatbot
- Create a data directory inside your project folder by running:
mkdir data
- Confirm the folder structure exists by running:
ls
You should see data listed in the output.
Now you need some text files for your knowledge base. You'll create three small files about a fictional product called "NextBot."
- Create the first knowledge base file by running:
cat > data/overview.txt << 'EOF'
NextBot is an AI-powered assistant designed for small businesses. It can answer customer questions, schedule appointments, and process simple orders. NextBot runs entirely in the cloud and requires no hardware installation. The service costs $29 per month for up to 1000 conversations. NextBot supports English, Spanish, and French. It integrates with popular platforms including Shopify, WordPress, and Slack.
EOF
- Create the second file by running:
cat > data/setup.txt << 'EOF'
To set up NextBot, first create an account at nextbot.example.com. After signing in, click Create New Bot and choose your industry template. NextBot offers templates for retail, healthcare, restaurants, and professional services. Upload your FAQ document or product catalog so NextBot can learn your business details. The initial training takes about 10 minutes. Once training is complete, you can test your bot using the built-in chat simulator before deploying it to your website.
EOF
- Create the third file by running:
cat > data/troubleshooting.txt << 'EOF'
If NextBot is not responding, first check your internet connection and verify the service status at status.nextbot.example.com. Common issues include expired API keys, which can be regenerated from the dashboard under Settings then API Keys. If the bot gives incorrect answers, retrain it by uploading updated documents in the Training section. Response times over 5 seconds usually indicate high server load and resolve automatically. For billing issues, contact support at help@nextbot.example.com.
EOF
Use your own content
These example files work great for learning, but feel free to replace them with notes on any topic you care about. The RAG system works with any plain text content.
- Verify your knowledge base files exist by running:
ls data/
You should see overview.txt, setup.txt, and troubleshooting.txt listed.
- Now open the rag-chatbot folder in VS Code. Click File in the top menu bar.
- Click Open Folder.
- Navigate to your Desktop, select the rag-chatbot folder, and click Open.
Don't see rag-chatbot on your Desktop?
Make sure you ran cd ~/Desktop before mkdir rag-chatbot. If you skipped that step, the folder was created in a different location. Run pwd to check where you are.
Help me find my rag-chatbot folder
- Create a new file called rag.py in the root of your project. In VS Code, click the New File icon at the top of the Explorer sidebar.
- Type rag.py and press Enter.
Write the document loader and chunker
Your rag.py file will handle the entire RAG pipeline. You'll start by setting up the imports and configuration, then build two functions: one to load documents and one to split them into chunks.
- In rag.py, add the imports and configuration constants by pasting this code:
import os
import numpy as np
import faiss
from sentence_transformers import SentenceTransformer
from transformers import pipeline
# --- Configuration ---
DATA_DIR = "data"
CHUNK_SIZE = 300
CHUNK_OVERLAP = 50
TOP_K = 3
RELEVANCE_THRESHOLD = 0.3
What does this code do?
The imports bring in all the libraries you'll use throughout the project. You won't use all of them in this step, but setting them up now means the file is ready for later.
The configuration constants control how your RAG system behaves. CHUNK_SIZE = 300 means each text chunk is about 300 characters. CHUNK_OVERLAP = 50 means consecutive chunks share 50 characters, so important sentences at chunk boundaries don't get lost.
- Below the configuration section, add the document loading function:
# --- Document Loading and Chunking ---
def load_documents(data_dir):
"""Load all .txt files from the data directory."""
documents = []
for filename in os.listdir(data_dir):
if filename.endswith(".txt"):
filepath = os.path.join(data_dir, filename)
with open(filepath, "r", encoding="utf-8") as f:
text = f.read()
documents.append({"text": text, "source": filename})
return documents
What does this function do?
The load_documents() function scans the data/ directory for any file ending in .txt. For each file it finds, it reads the full contents and stores both the text and the source filename in a dictionary.
Tracking the source filename alongside the text gives your chatbot the ability to cite where it found its answers.
- Below load_documents, add the chunking function:
def chunk_documents(documents, chunk_size=CHUNK_SIZE, overlap=CHUNK_OVERLAP):
"""Split documents into overlapping chunks."""
chunks = []
for doc in documents:
text = doc["text"]
start = 0
while start < len(text):
end = start + chunk_size
chunk_text = text[start:end]
chunks.append({"text": chunk_text, "source": doc["source"]})
start += chunk_size - overlap
return chunks
Why split documents into chunks?
Embedding models work best on short passages of text. A full document is too long to capture in a single vector. Chunking breaks each document into bite-sized pieces that can be individually embedded and retrieved.
The overlap ensures that if an important sentence sits right at a chunk boundary, it still appears in at least one chunk fully intact. Each chunk also carries its source filename so you can trace answers back to their origin document.
- Save rag.py.
Test the chunking by running the script
You have the loader and chunker ready. Now you need a main() function that ties them together and prints results so you can verify everything works.
- At the bottom of rag.py, add the main function and entry point:
# --- Main ---
def main():
print("Loading documents...")
documents = load_documents(DATA_DIR)
print(f"Loaded {len(documents)} documents")
print("Chunking documents...")
chunks = chunk_documents(documents)
print(f"Created {len(chunks)} chunks")
if chunks:
print(f"Sample chunk: {chunks[0]}")
if __name__ == "__main__":
main()
What does this code do?
The main() function calls your loader and chunker in sequence, printing progress as it goes. The if __name__ == "__main__" guard ensures main() only runs when you execute the file directly, not when it's imported by another script.
Printing the document count, chunk count, and a sample chunk lets you confirm the pipeline is working before adding the more complex embedding and search steps.
- Save rag.py.
- Open a terminal in VS Code by pressing Ctrl+` (backtick key, same on macOS and Windows).
- Make sure rag-env is activated in this terminal. If you see (rag-env) at the start of your prompt, you're good.
- Run the script:
python rag.py
You should see output similar to this:
Loading documents...
Loaded 3 documents
Chunking documents...
Created 12 chunks
Sample chunk: {'text': 'NextBot is an AI-powered assistant designed...', 'source': 'overview.txt'}
What should I see?
The exact number of chunks depends on how much text is in your files. With the three example files above, you'll get around 8-12 chunks. Each chunk is approximately 300 characters with 50 characters of overlap between neighbors.
The sample chunk shows both the text content and the source filename, confirming that attribution tracking is working.
Seeing a FileNotFoundError?
This means Python can't find the data/ directory. Make sure you're running the script from inside the rag-chatbot folder, not from outside it. Run ls and check that you see both data and rag.py in the output.
Help me fix this FileNotFoundError
✔️ Awesome, I've got everything!
Great. Make sure you've saved rag.py before moving on to the next step.
ⓧ I'd like to double check the full code
Here's the complete rag.py file at this point:
import os
import numpy as np
import faiss
from sentence_transformers import SentenceTransformer
from transformers import pipeline
# --- Configuration ---
DATA_DIR = "data"
CHUNK_SIZE = 300
CHUNK_OVERLAP = 50
TOP_K = 3
RELEVANCE_THRESHOLD = 0.3
# --- Document Loading and Chunking ---
def load_documents(data_dir):
"""Load all .txt files from the data directory."""
documents = []
for filename in os.listdir(data_dir):
if filename.endswith(".txt"):
filepath = os.path.join(data_dir, filename)
with open(filepath, "r", encoding="utf-8") as f:
text = f.read()
documents.append({"text": text, "source": filename})
return documents
def chunk_documents(documents, chunk_size=CHUNK_SIZE, overlap=CHUNK_OVERLAP):
"""Split documents into overlapping chunks."""
chunks = []
for doc in documents:
text = doc["text"]
start = 0
while start < len(text):
end = start + chunk_size
chunk_text = text[start:end]
chunks.append({"text": chunk_text, "source": doc["source"]})
start += chunk_size - overlap
return chunks
# --- Main ---
def main():
print("Loading documents...")
documents = load_documents(DATA_DIR)
print(f"Loaded {len(documents)} documents")
print("Chunking documents...")
chunks = chunk_documents(documents)
print(f"Created {len(chunks)} chunks")
if chunks:
print(f"Sample chunk: {chunks[0]}")
if __name__ == "__main__":
main()
Your knowledge base is loaded and chunked. Next up, you'll convert those chunks into vector embeddings and build a FAISS index so your chatbot can search them by meaning instead of keywords.
Generate Embeddings and Build the FAISS Index
Your knowledge base is chunked and ready. Now you need a way to find which chunks are relevant to any given question.
The answer is embeddings. By converting text into numerical vectors, you can measure how similar two pieces of text are using math instead of keyword matching. In this step, you'll embed all your chunks into a FAISS index for lightning-fast similarity search.
In this step, get ready to:
- Build an embedding and indexing function using sentence-transformers and FAISS.
- Write a retrieval function that finds the most relevant chunks for any query.
- Test the full pipeline by retrieving chunks with similarity scores.
Create the build_index() function
Each text chunk is just a string right now. To search by meaning, you need to convert each chunk into a dense vector (an embedding) that captures its semantic content.
- In rag.py, find the line that says # --- Main ---. Copy the following code directly above that line:
# --- Embedding and Indexing ---
def build_index(chunks, model):
"""Generate embeddings and build a FAISS index."""
texts = [chunk["text"] for chunk in chunks]
embeddings = model.encode(texts, normalize_embeddings=True)
embeddings = np.array(embeddings, dtype=np.float32)
dimension = embeddings.shape[1]
index = faiss.IndexFlatIP(dimension)
index.add(embeddings)
print(f"Indexed {index.ntotal} vectors of dimension {dimension}")
return index
What does this code do?
This function takes your list of chunks and a sentence-transformers model, then converts every chunk into a 384-dimensional vector.
- model.encode(texts, normalize_embeddings=True) converts each chunk's text into a normalized vector. Normalizing means each vector has a length of 1, which makes inner product equivalent to cosine similarity.
- faiss.IndexFlatIP(dimension) creates a FAISS index that uses inner product (IP) for similarity search. Because the vectors are normalized, this gives you cosine similarity for free.
- index.add(embeddings) stores all chunk vectors in the index so they can be searched later.
Why IndexFlatIP instead of other index types?
FAISS offers many index types optimized for different scales. IndexFlatIP performs exact (brute-force) search, which is perfect for small knowledge bases like yours.
For collections under a million vectors, exact search is fast enough and gives you perfect accuracy. More complex index types add training steps and approximation that you don't need here.
Write the retrieve() function
With your chunks indexed, you need a function that takes a question, converts it into the same vector space, and finds the closest matches.
- In rag.py, find the line # --- Main ---. Copy the following code directly above it (below the build_index function you just added):
# --- Retrieval ---
def retrieve(query, model, index, chunks, top_k=TOP_K):
"""Find the most relevant chunks for a query."""
query_embedding = model.encode([query], normalize_embeddings=True)
query_embedding = np.array(query_embedding, dtype=np.float32)
scores, indices = index.search(query_embedding, top_k)
results = []
for i, idx in enumerate(indices[0]):
results.append({
"text": chunks[idx]["text"],
"source": chunks[idx]["source"],
"score": float(scores[0][i]),
})
return results
What does this code do?
- model.encode([query], normalize_embeddings=True) converts the user's question into the same 384-dimensional vector space as the chunks. This is key. The query and the documents must live in the same space for comparison to work.
- index.search(query_embedding, top_k) asks FAISS to find the top_k most similar vectors. It returns both the similarity scores and the indices of those vectors.
- The function packages each result into a dictionary with the chunk text, source filename, and similarity score for easy downstream use.
Update main() to test retrieval
Now wire everything together in your main() function to load the embedding model, build the index, and run a test query.
- In rag.py, find the existing main() function and replace it entirely with the following:
# --- Main ---
def main():
print("Loading documents...")
documents = load_documents(DATA_DIR)
print(f"Loaded {len(documents)} documents")
print("Chunking documents...")
chunks = chunk_documents(documents)
print(f"Created {len(chunks)} chunks")
print("Loading embedding model...")
embed_model = SentenceTransformer("all-MiniLM-L6-v2")
print("Building FAISS index...")
index = build_index(chunks, embed_model)
# Test retrieval
test_results = retrieve("your test question", embed_model, index, chunks, top_k=3)
for r in test_results:
print(r)
if __name__ == "__main__":
main()
What does this updated main() do?
- SentenceTransformer("all-MiniLM-L6-v2") downloads and loads the embedding model. This model is optimized for speed (5x faster than larger alternatives) while maintaining good quality for semantic search.
- build_index(chunks, embed_model) passes the model to your indexing function, which encodes all chunks and stores them in FAISS.
- retrieve("your test question", ...) runs a test search to confirm the full pipeline works end-to-end.
- Save rag.py.
- Run the script by entering this command in your terminal:
python rag.py
The first run will download the all-MiniLM-L6-v2 model (about 80MB). After that, you should see your documents loaded, chunks created, vectors indexed, and 3 retrieval results printed with similarity scores between 0 and 1.
Not seeing retrieval results?
- If you see an import error for sentence_transformers, make sure your virtual environment is activated. You should see (rag-env) in your terminal prompt.
- If you see a timeout or connection error, the model needs to download from the internet. Check your network connection and try again.
- If scores are all showing 0 or NaN, make sure your data/ folder has .txt files with actual text content (not empty files).
Why aren't my retrieval results showing similarity scores?
✔️ Awesome, I've got everything!
Great. Double check you've saved rag.py and your retrieval test is printing results with scores.
ⓧ I'd like to double check the full code
Here is the complete rag.py file at this point:
import os
import numpy as np
import faiss
from sentence_transformers import SentenceTransformer
from transformers import pipeline
# --- Configuration ---
DATA_DIR = "data"
CHUNK_SIZE = 300
CHUNK_OVERLAP = 50
TOP_K = 3
RELEVANCE_THRESHOLD = 0.3
# --- Document Loading and Chunking ---
def load_documents(data_dir):
"""Load all .txt files from the data directory."""
documents = []
for filename in os.listdir(data_dir):
if filename.endswith(".txt"):
filepath = os.path.join(data_dir, filename)
with open(filepath, "r", encoding="utf-8") as f:
text = f.read()
documents.append({"text": text, "source": filename})
return documents
def chunk_documents(documents, chunk_size=CHUNK_SIZE, overlap=CHUNK_OVERLAP):
"""Split documents into overlapping chunks."""
chunks = []
for doc in documents:
text = doc["text"]
start = 0
while start < len(text):
end = start + chunk_size
chunk_text = text[start:end]
chunks.append({"text": chunk_text, "source": doc["source"]})
start += chunk_size - overlap
return chunks
# --- Embedding and Indexing ---
def build_index(chunks, model):
"""Generate embeddings and build a FAISS index."""
texts = [chunk["text"] for chunk in chunks]
embeddings = model.encode(texts, normalize_embeddings=True)
embeddings = np.array(embeddings, dtype=np.float32)
dimension = embeddings.shape[1]
index = faiss.IndexFlatIP(dimension)
index.add(embeddings)
print(f"Indexed {index.ntotal} vectors of dimension {dimension}")
return index
# --- Retrieval ---
def retrieve(query, model, index, chunks, top_k=TOP_K):
"""Find the most relevant chunks for a query."""
query_embedding = model.encode([query], normalize_embeddings=True)
query_embedding = np.array(query_embedding, dtype=np.float32)
scores, indices = index.search(query_embedding, top_k)
results = []
for i, idx in enumerate(indices[0]):
results.append({
"text": chunks[idx]["text"],
"source": chunks[idx]["source"],
"score": float(scores[0][i]),
})
return results
# --- Main ---
def main():
print("Loading documents...")
documents = load_documents(DATA_DIR)
print(f"Loaded {len(documents)} documents")
print("Chunking documents...")
chunks = chunk_documents(documents)
print(f"Created {len(chunks)} chunks")
print("Loading embedding model...")
embed_model = SentenceTransformer("all-MiniLM-L6-v2")
print("Building FAISS index...")
index = build_index(chunks, embed_model)
# Test retrieval
test_results = retrieve("your test question", embed_model, index, chunks, top_k=3)
for r in test_results:
print(r)
if __name__ == "__main__":
main()
Your documents are now searchable by meaning, not just keywords. Next up, you'll connect a language model that reads the retrieved chunks and generates a natural language answer.
Wire Up the Generation Pipeline and Build the RAG Loop
Your retrieval system is working. Given any query, it finds the most relevant chunks from your knowledge base. But raw text chunks aren't useful on their own.
In this step, you'll connect a language model to your retrieval function. The model reads the retrieved chunks and generates a coherent answer grounded in your documents. This is the "generation" half of Retrieval-Augmented Generation.
In this step, get ready to:
- Load a Hugging Face text-generation pipeline for fast CPU inference.
- Write a rag_answer() function that combines retrieval and generation.
- Build an interactive terminal loop for chatting with your knowledge base.
Load the text-generation model
The Hugging Face Transformers library provides a pipeline API that wraps model loading, tokenization, and inference into a single function call. You already imported pipeline at the top of your file. Now you'll use it to load a lightweight generation model.
- In rag.py, scroll down to the main() function.
- Find the test retrieval code at the bottom of main():
# Test retrieval
test_results = retrieve("your test question", embed_model, index, chunks, top_k=3)
for r in test_results:
print(r)
- Delete those four lines and replace them with the following:
print("Loading generation model...")
generator = pipeline("text-generation", model="distilgpt2")
What does this code do?
pipeline("text-generation", model="distilgpt2") downloads and loads the distilgpt2 model. This is a small, distilled version of GPT-2 that runs quickly on CPU without needing a GPU.
The pipeline handles tokenization, inference, and decoding in one call. You'll pass it a prompt and get generated text back.
You won't see a change yet. The next chunk adds the function and loop that use this generator.
Build the rag_answer() function
This function ties retrieval and generation together. It takes the user's question, retrieves relevant chunks, builds a prompt with those chunks as context, and asks the model to answer based only on that context.
- In rag.py, find the # --- Main --- comment line. Add the following function above it:
# --- Generation ---
def rag_answer(question, model, index, chunks, generator):
"""Retrieve context and generate a grounded answer."""
results = retrieve(question, model, index, chunks, top_k=TOP_K)
# Build context from retrieved chunks
context = "\n".join(f"- {r['text']}" for r in results)
prompt = (
f"Answer the question based only on the provided context.\n\n"
f"Context:\n{context}\n\n"
f"Question: {question}\n"
f"Answer:"
)
output = generator(prompt, max_new_tokens=150, do_sample=False, return_full_text=False)
answer = output[0]["generated_text"].strip()
return answer, results
What does this code do?
- retrieve() fetches the top-3 most similar chunks from your FAISS index.
- The prompt instructs the model to answer using only the provided context. This grounds the output in your documents instead of letting the model hallucinate.
- max_new_tokens=150 caps the response length. do_sample=False makes the output deterministic (same input always produces same output).
- return_full_text=False tells the pipeline to return only the generated text, not the entire prompt repeated back.
- The function returns both the answer and the retrieved results, so you can show source attribution.
You won't see a change yet. The next chunk builds the interactive loop that calls this function.
Create the interactive Q&A loop
The final piece is an input loop that lets you type questions and see the chatbot respond in real time, complete with source citations.
- In rag.py, find the generator = pipeline("text-generation", model="distilgpt2") line you just added inside main(). Add the following code directly below it:
print("\nRAG Chatbot ready! Type 'quit' to exit.\n")
while True:
question = input("You: ").strip()
if question.lower() in ("quit", "exit", "q"):
print("Goodbye!")
break
if not question:
continue
answer, sources = rag_answer(question, embed_model, index, chunks, generator)
print(f"\nBot: {answer}")
if sources:
print("\nSources:")
for s in sources:
print(f" [{s['source']}] (score: {s['score']:.3f}) {s['text'][:80]}...")
print()
What does this code do?
- input("You: ") waits for the user to type a question.
- Typing quit, exit, or q breaks out of the loop and ends the program.
- rag_answer() retrieves relevant chunks and generates the grounded answer.
- The sources section shows which document each chunk came from, its similarity score, and a preview of the chunk text. This is source attribution in action.
- Save rag.py.
- Run the script by entering this command in your terminal:
python rag.py
The first run downloads the distilgpt2 model (about 300MB). After loading completes, you should see:
RAG Chatbot ready! Type 'quit' to exit.
You:
- Type a question related to your knowledge base documents and press Enter.
- You should see the bot's answer followed by a Sources section listing the document filename, similarity score, and a chunk preview.
- Type quit to exit when you're done testing.
Not seeing any output or getting an error?
- If you see an IndentationError, make sure the new code inside main() is indented with 4 spaces (one level inside the function).
- If you see NameError: name 'rag_answer' is not defined, make sure the rag_answer() function is placed above the # --- Main --- section in your file.
- If the model download stalls, check your internet connection. The distilgpt2 model needs to download once.
Still stuck? Help me debug my RAG chatbot - it's not producing output when I run python rag.py
Why is the answer quality not perfect?
distilgpt2 is a tiny model designed for speed, not quality. Its answers may be repetitive or incomplete. The point of this project is learning the RAG architecture, not producing production-grade text.
You can swap in a larger model later by changing the model parameter. The pipeline pattern stays the same regardless of model size.
✔️ Awesome, I've got everything!
Great. Double check you've saved rag.py and your chatbot responds to questions with answers and source citations.
ⓧ I'd like to double check the full code
import os
import numpy as np
import faiss
from sentence_transformers import SentenceTransformer
from transformers import pipeline
# --- Configuration ---
DATA_DIR = "data"
CHUNK_SIZE = 300
CHUNK_OVERLAP = 50
TOP_K = 3
RELEVANCE_THRESHOLD = 0.3
# --- Document Loading and Chunking ---
def load_documents(data_dir):
"""Load all .txt files from the data directory."""
documents = []
for filename in os.listdir(data_dir):
if filename.endswith(".txt"):
filepath = os.path.join(data_dir, filename)
with open(filepath, "r", encoding="utf-8") as f:
text = f.read()
documents.append({"text": text, "source": filename})
return documents
def chunk_documents(documents, chunk_size=CHUNK_SIZE, overlap=CHUNK_OVERLAP):
"""Split documents into overlapping chunks."""
chunks = []
for doc in documents:
text = doc["text"]
start = 0
while start < len(text):
end = start + chunk_size
chunk_text = text[start:end]
chunks.append({"text": chunk_text, "source": doc["source"]})
start += chunk_size - overlap
return chunks
# --- Embedding and Indexing ---
def build_index(chunks, model):
"""Generate embeddings and build a FAISS index."""
texts = [chunk["text"] for chunk in chunks]
embeddings = model.encode(texts, normalize_embeddings=True)
embeddings = np.array(embeddings, dtype=np.float32)
dimension = embeddings.shape[1]
index = faiss.IndexFlatIP(dimension)
index.add(embeddings)
print(f"Indexed {index.ntotal} vectors of dimension {dimension}")
return index
# --- Retrieval ---
def retrieve(query, model, index, chunks, top_k=TOP_K):
"""Find the most relevant chunks for a query."""
query_embedding = model.encode([query], normalize_embeddings=True)
query_embedding = np.array(query_embedding, dtype=np.float32)
scores, indices = index.search(query_embedding, top_k)
results = []
for i, idx in enumerate(indices[0]):
results.append({
"text": chunks[idx]["text"],
"source": chunks[idx]["source"],
"score": float(scores[0][i]),
})
return results
# --- Generation ---
def rag_answer(question, model, index, chunks, generator):
"""Retrieve context and generate a grounded answer."""
results = retrieve(question, model, index, chunks, top_k=TOP_K)
# Build context from retrieved chunks
context = "\n".join(f"- {r['text']}" for r in results)
prompt = (
f"Answer the question based only on the provided context.\n\n"
f"Context:\n{context}\n\n"
f"Question: {question}\n"
f"Answer:"
)
output = generator(prompt, max_new_tokens=150, do_sample=False, return_full_text=False)
answer = output[0]["generated_text"].strip()
return answer, results
# --- Main ---
def main():
print("Loading documents...")
documents = load_documents(DATA_DIR)
print(f"Loaded {len(documents)} documents")
print("Chunking documents...")
chunks = chunk_documents(documents)
print(f"Created {len(chunks)} chunks")
print("Loading embedding model...")
embed_model = SentenceTransformer("all-MiniLM-L6-v2")
print("Building FAISS index...")
index = build_index(chunks, embed_model)
print("Loading generation model...")
generator = pipeline("text-generation", model="distilgpt2")
print("\nRAG Chatbot ready! Type 'quit' to exit.\n")
while True:
question = input("You: ").strip()
if question.lower() in ("quit", "exit", "q"):
print("Goodbye!")
break
if not question:
continue
answer, sources = rag_answer(question, embed_model, index, chunks, generator)
print(f"\nBot: {answer}")
if sources:
print("\nSources:")
for s in sources:
print(f" [{s['source']}] (score: {s['score']:.3f}) {s['text'][:80]}...")
print()
if __name__ == "__main__":
main()
Your RAG chatbot is alive. It retrieves relevant context from your documents and generates grounded answers. Next up, the Secret Mission adds intelligence so the chatbot knows when to say "I don't know."
Secret mission
Add a Relevance Threshold
Your chatbot always generates an answer, even when the knowledge base has nothing relevant. In this secret mission, you'll add a relevance threshold so the bot gracefully says "I don't know" instead of hallucinating from irrelevant context.
Clean Up Your Resources
Clean Up Your Resources
This project runs entirely locally, so there are no ongoing costs. Decide whether to keep your resources on disk, pause them to come back later, or delete them entirely.
Resources you used:
- Project files (rag.py, data/ directory).
- Virtual environment (rag-env/ folder with all installed packages).
- Hugging Face model cache (downloaded all-MiniLM-L6-v2 and distilgpt2 models stored at ~/.cache/huggingface/hub).
Keep everything running
No action needed. Choose this if you plan to keep experimenting with your RAG chatbot or want to add more documents to your knowledge base.
- Your project folder stays on your Desktop, ready to use whenever you run rag.py again.
- The Hugging Face model cache stays at ~/.cache/huggingface/hub. This means future runs skip the model download step entirely.
Pause - I'll come back to this later
Shut down running processes to free up memory. Your files stay on disk so you can pick up where you left off.
- If your chatbot is still running in the terminal, type quit to exit the interactive loop.
- Deactivate your virtual environment by running:
deactivate
- Close VS Code. When you come back, reactivate the environment and run rag.py again.
Delete - I don't want to use this again
Remove all project resources and reclaim disk space. This deletes your code, knowledge base, virtual environment, and cached models.
- If your chatbot is still running, type quit to exit.
- Deactivate your virtual environment by running:
deactivate
- Delete your project folder. Navigate to your Desktop and remove the folder that contains rag.py.
On macOS/Linux:
cd ~/Desktop
ls
- Find the folder name that contains your RAG project, then remove it:
rm -rf your-project-folder-name
On Windows (PowerShell):
cd ~\Desktop
ls
- Find the folder name that contains your RAG project, then remove it:
Remove-Item -Recurse -Force your-project-folder-name
- Remove the Hugging Face model cache to free up space taken by the downloaded models.
On macOS/Linux:
rm -rf ~/.cache/huggingface/hub
On Windows (PowerShell):
Remove-Item -Recurse -Force "$env:USERPROFILE\.cache\huggingface\hub"
What about the rest of the Hugging Face cache?
The command above only removes the hub subdirectory, which stores downloaded model files. If you want to remove all Hugging Face cached data, delete the entire ~/.cache/huggingface directory instead.
That's a Wrap!
That's a Wrap!
Nice work! You just built a fully local Retrieval-Augmented Generation chatbot using PyTorch, Hugging Face Transformers, and FAISS.
You've learned how to:
- Chunk documents and generate vector embeddings using sentence-transformers, then index them in a FAISS similarity search index for fast retrieval.
- Build a retrieval function that encodes a query and finds the most relevant chunks from your knowledge base using cosine similarity.
- Wire a Hugging Face text-generation pipeline into a RAG loop that produces context-grounded answers with source attribution.
- Add a relevance threshold so the chatbot gracefully handles out-of-scope questions instead of generating hallucinated answers (Secret Mission).
Ready to quiz yourself?