RAG from Scratch — Retrieval-Augmented Generation for Knowledge-Intensive NLP

A from-scratch implementation of the RAG paper:

Lewis et al. (2020). Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks. arXiv:2005.11401.

This project implements every core concept from the paper:

Paper Concept	Implementation
DPR bi-encoder retriever (non-parametric memory)	src/retriever.py
BART seq2seq generator (parametric memory)	src/generator.py
RAG-Sequence (same doc for whole sequence)	src/rag_model.py
RAG-Token (different doc per token)	src/rag_model.py
Latent variable marginalization over top-K docs	src/rag_model.py
End-to-end joint training (query encoder + generator)	scripts/train.py
Document encoder frozen, query encoder fine-tuned	scripts/train.py
Thorough Decoding & Fast Decoding	src/decoding.py
FAISS MIPS index over 100-word Wikipedia chunks	src/document_store.py
Knowledge hot-swapping (swap index, no retraining)	scripts/hotswap_demo.py
Open-domain QA, Jeopardy generation, FEVER fact verification	src/data.py
Exact Match, BLEU, Rouge-L, label accuracy metrics	src/metrics.py

Architecture

                    ┌──────────────────────────────────────────────┐
                    │                 RAG Model                     │
                    │                                              │
  Query x ─────────┼──► Query Encoder (BERT) ──► q(x)             │
                    │         │                                    │
                    │         ▼                                    │
                    │   FAISS MIPS Search (top-K docs)             │
                    │         │                                    │
                    │         ▼                                    │
                    │   z₁, z₂, ..., z_K  (retrieved passages)     │
                    │         │                                    │
                    │    ┌────┴────┐                               │
                    │    ▼         ▼                               │
                    │  RAG-Seq   RAG-Token                         │
                    │  (1 doc    (K docs per                       │
                    │   per seq)  token)                           │
                    │    │         │                               │
                    │    ▼         ▼                               │
                    │  BART Generator (parametric memory)          │
                    │    │         │                               │
                    │    ▼         ▼                               │
                    │  Marginalize p(z|x) × p(y|x,z)              │
                    │    │                                         │
                    │    ▼                                         │
                    │  Output y                                    │
                    └──────────────────────────────────────────────┘

  Non-parametric memory: FAISS index of 21M Wikipedia 100-word chunks
  Parametric memory: BART-large (400M params) pre-trained seq2seq

Tested Environment

Component	Version	Notes
Python	3.11.9	Installed from python.org (not Windows Store)
PyTorch	2.6.0+cu124	GPU build with CUDA 12.4
GPU	NVIDIA RTX 3050 6GB	CUDA detected, torch.cuda.is_available() = True
Transformers	5.12.1	HuggingFace transformers
FAISS	faiss-cpu 1.14.3	MIPS vector search
NumPy	2.4.4
Tests	22/22 passed, 0 skipped	python -m unittest tests.test_rag -v

PyTorch GPU Install

# For NVIDIA RTX 3050 (CUDA 12.4)
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124

# CPU-only fallback
pip install torch torchvision torchaudio

Quickstart

# Install dependencies
pip install -r requirements.txt

# 1. Build a document index from Wikipedia (or custom corpus)
python scripts/build_index.py --source wikipedia --chunk-size 100

# 2. Download datasets (NQ, TriviaQA, FEVER, Jeopardy)
python scripts/download_data.py

# 3. Train RAG end-to-end (joint retriever + generator)
python scripts/train.py --model rag-sequence --task qa --epochs 10

# 4. Evaluate
python scripts/evaluate.py --checkpoint checkpoints/rag-sequence-best.pt

# 5. Run inference
python scripts/inference.py --checkpoint checkpoints/rag-sequence-best.pt \
    --query "Who wrote The Divine Comedy?"

# 6. Demo knowledge hot-swapping
python scripts/hotswap_demo.py

# 7. Launch web interface
python app.py

# 8. Run tests
python -m unittest tests.test_rag -v

Project Structure

rag-from-scratch/
├── README.md
├── requirements.txt
├── configs/
│   └── default.yaml          # All hyperparameters (matches paper)
├── src/
│   ├── __init__.py
│   ├── config.py             # Config loader
│   ├── document_store.py     # FAISS index + document storage (non-parametric memory)
│   ├── retriever.py          # DPR bi-encoder: BERT query + BERT document encoders
│   ├── generator.py          # BART-large seq2seq generator (parametric memory)
│   ├── rag_model.py          # RAG-Sequence + RAG-Token with latent marginalization
│   ├── decoding.py           # Thorough Decoding & Fast Decoding
│   ├── data.py               # Dataset classes: QA, FEVER, Jeopardy generation
│   ├── metrics.py            # Exact Match, BLEU-1, Rouge-L, label accuracy
│   └── utils.py              # Logging, seeding, checkpoint utils
├── scripts/
│   ├── build_index.py        # Build FAISS index from Wikipedia / custom corpus
│   ├── download_data.py      # Download and prepare datasets
│   ├── train.py              # End-to-end training loop
│   ├── evaluate.py           # Full evaluation on any task
│   ├── inference.py          # CLI inference
│   └── hotswap_demo.py       # Knowledge hot-swapping demonstration
├── tests/
│   └── test_rag.py           # Self-tests for all components
└── app.py                    # Gradio web interface

Key Design Decisions (faithful to the paper)

1. Document encoder frozen during training — updating it requires re-indexing 21M docs. Only the query encoder and generator are fine-tuned.
2. Retrieved documents treated as latent variables — marginalized with top-K approximation, not hard-selected.
3. No retrieval supervision — the model learns what to retrieve purely from the downstream task loss.
4. Wikipedia split into 100-word disjoint chunks — exactly as in the paper (21M passages from Dec 2018 dump).
5. FAISS with HNSW approximation for sub-linear MIPS, matching the paper's FAISS configuration.
6. RAG-Sequence uses Thorough Decoding for short outputs (QA), Fast Decoding for long outputs (generation).

Test Results

test_dataset_items              ok
test_dummy_data_creation        ok
test_01_create_and_add          ok
test_02_retrieve_document       ok
test_03_build_and_search_index  ok
test_04_hotswap                 ok
test_full_pipeline              ok
test_accuracy                   ok
test_bleu_1                     ok
test_diversity_ratio            ok
test_evaluate_fever             ok
test_evaluate_jeopardy          ok
test_evaluate_qa                ok
test_exact_match                ok
test_exact_match_batch          ok
test_f1_score                   ok
test_q_bleu_1                   ok
test_rouge_l                    ok
test_model_initialization       ok
test_chunk_text                 ok
test_chunk_text_short           ok
test_set_seed                   ok

Ran 22 tests in 17.012s
OK

Tests cover:

• Metrics (10 tests): Exact Match, F1, BLEU-1, Q-BLEU-1, Rouge-L, diversity ratio, accuracy
• Document Store (4 tests): SQLite storage, FAISS index build/search, hot-swapping
• Datasets (2 tests): Dummy data creation, dataset item structure
• RAG Model (1 test): Full model initialization (downloads BERT + BART from HuggingFace)
• End-to-end pipeline (1 test): Build index → search → retrieve documents
• Utils (3 tests): Text chunking, reproducibility seeding
• Config (1 test): YAML loading

Resume Description

RAG from Scratch — Implemented RAG (Lewis et al., 2020) in PyTorch: DPR bi-encoder retriever + FAISS MIPS over 21M Wikipedia docs, BART-large generator, RAG-Sequence & RAG-Token with latent marginalization. Trained on QA, FEVER & Jeopardy. 22/22 tests passing on CUDA (RTX 3050).

Citation

@article{lewis2020rag,
  title={Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks},
  author={Lewis, Patrick and Perez, Ethan and Piktus, Aleksandra and Petroni, Fabio and
          Karpukhin, Vladimir and Goyal, Naman and K{\"u}ttler, Heinrich and Lewis, Mike and
          Yih, Wen-tau and Rockt{\"a}schel, Tim and Riedel, Sebastian and Kiela, Douwe},
  journal={arXiv preprint arXiv:2005.11401},
  year={2020}
}

License

MIT