Spaces:
Sleeping
Sleeping
| title: RAG Document Q&A | |
| emoji: 📄 | |
| colorFrom: indigo | |
| colorTo: blue | |
| sdk: docker | |
| app_port: 7860 | |
| pinned: false | |
| # RAG Document Q&A System | |
| Upload PDF documents and get cited, grounded answers — with hybrid retrieval, cross-encoder re-ranking, and confidence-aware generation. | |
| ## Overview | |
| This project implements a production-grade Retrieval-Augmented Generation pipeline that goes beyond the typical "embed + nearest-neighbor + prompt" tutorial. It combines semantic and keyword search with cross-encoder re-ranking, applies Anthropic's contextual retrieval pattern to enrich embeddings with document metadata, and gates LLM calls on retrieval confidence to avoid hallucinated answers. The system ships with both a streaming Gradio UI and a FastAPI backend with Server-Sent Events, plus an evaluation harness that benchmarks chunking strategies and scores end-to-end answer quality using an LLM-as-judge. | |
| ## Architecture | |
| ``` | |
| ┌─────────────────────────────────────────┐ | |
| │ INGESTION │ | |
| │ │ | |
| PDF ──► pymupdf4llm ──► │ Chunking (3 strategies) │ | |
| (pdfplumber │ ├─ fixed_size sliding window │ | |
| fallback) │ ├─ recursive_char LangChain splits │ | |
| │ └─ semantic embedding sim │ | |
| │ │ │ | |
| │ ▼ │ | |
| │ Contextual Enrichment │ | |
| │ "filename | Page N | Section\ntext" │ | |
| │ │ │ | |
| │ ▼ │ | |
| │ all-MiniLM-L6-v2 (384-dim, L2-norm) │ | |
| │ │ │ | |
| │ ┌─────┴──────┐ │ | |
| │ ▼ ▼ │ | |
| │ FAISS BM25Okapi │ | |
| │ IndexFlatIP rank-bm25 │ | |
| └─────────────────────────────────────────┘ | |
| ┌─────────────────────────────────────────┐ | |
| │ QUERY │ | |
| │ │ | |
| Question ──────────────►│ Adaptive Expansion (Groq) │ | |
| │ only if top_score < 0.45 │ | |
| │ │ │ | |
| │ ▼ │ | |
| │ Hybrid Search │ | |
| │ score = 0.7·semantic + 0.3·bm25 │ | |
| │ (top 20 candidates) │ | |
| │ │ │ | |
| │ ▼ │ | |
| │ Cross-Encoder Re-ranking │ | |
| │ ms-marco-MiniLM-L-6-v2 │ | |
| │ (20 → top k) │ | |
| │ │ │ | |
| │ ▼ │ | |
| │ Confidence Gating │ | |
| │ cosine sim < 0.3 → skip LLM │ | |
| │ │ │ | |
| │ ▼ │ | |
| │ Groq (Llama 3.3 70B) streaming │ | |
| │ │ │ | |
| │ ▼ │ | |
| │ Cited answer [Source N] notation │ | |
| └─────────────────────────────────────────┘ | |
| ┌─────────────────────────────────────────┐ | |
| │ EVALUATION │ | |
| │ │ | |
| │ Retrieval: Acc@1/3/5, MRR │ | |
| │ End-to-end: faithfulness + relevance │ | |
| │ (LLM-as-judge, temp=0.0) │ | |
| │ Config sweep: chunking strategy grid │ | |
| └─────────────────────────────────────────┘ | |
| ``` | |
| ## Key Features | |
| - **Cross-encoder re-ranking** — bi-encoder retrieval (FAISS) produces 20 candidates; a cross-encoder (`ms-marco-MiniLM-L-6-v2`) scores each `(query, chunk)` pair jointly and re-orders them. Cross-encoders are too slow for full-corpus search but dramatically improve precision over the shortlist. | |
| - **Hybrid search** — FAISS inner-product search and BM25 run in parallel. Both score sets are normalized to `[0, 1]` then fused (`0.7·semantic + 0.3·keyword`), catching exact-match terms that dense embeddings tend to smooth over. | |
| - **Adaptive query expansion** — the pipeline runs an initial search first. Only if the top cosine similarity falls below 0.45 does it call the LLM to generate 3 rephrased variants, search each, and merge results. High-confidence queries pay zero expansion cost. | |
| - **Contextual chunk enrichment** — each chunk is embedded as `"<filename> | Page N | <section>\n<text>"` rather than raw text, following Anthropic's contextual retrieval pattern. The embedding captures document location and section context, not just lexical content. | |
| - **Confidence-based response gating** — `max_score` is taken from cosine similarity *before* re-ranking (cross-encoder scores on a −10 to +10 scale would break the threshold). If `max_score < 0.3`, the LLM is skipped entirely and the user receives an honest "insufficient context" message. Scores between 0.3–0.5 append a low-confidence warning. | |
| - **Streaming responses** — Gradio UI streams tokens via a generator; the FastAPI backend streams as Server-Sent Events (`text/event-stream`) with `{"text": chunk}` events and a terminal `{"done": true}`. | |
| - **Layout-aware PDF parsing** — `pymupdf4llm` extracts structured markdown with headings and tables preserved. `pdfplumber` serves as a fallback for PDFs that `pymupdf4llm` cannot parse cleanly. | |
| - **End-to-end evaluation** — the eval harness measures retrieval hit rate (Acc@K, MRR), then pipes results through the LLM-as-judge to score faithfulness (are claims grounded in context?) and relevance (does the answer address the question?) at temperature 0 for determinism. | |
| ## Tech Stack | |
| | Component | Technology | | |
| |---|---| | |
| | Embedding model | `all-MiniLM-L6-v2` (sentence-transformers, 384-dim) | | |
| | Vector index | FAISS `IndexFlatIP` (exact inner-product search) | | |
| | Keyword index | BM25 (`rank-bm25`, `BM25Okapi`) | | |
| | Re-ranker | `cross-encoder/ms-marco-MiniLM-L-6-v2` | | |
| | LLM | Llama 3.3 70B via Groq API | | |
| | PDF parsing | `pymupdf4llm` + `pdfplumber` fallback | | |
| | Text splitting | LangChain `RecursiveCharacterTextSplitter` | | |
| | Backend | FastAPI + Uvicorn | | |
| | Frontend | Gradio (Blocks, streaming) | | |
| | Containerization | Docker | | |
| | Cloud deployment | HuggingFace Spaces | | |
| ## Quick Start | |
| ```bash | |
| git clone <repo-url> | |
| cd rag-qa | |
| python -m venv venv && source venv/bin/activate | |
| pip install -r requirements.txt | |
| ``` | |
| Create a `.env` file: | |
| ``` | |
| GROQ_API_KEY=your_key_here | |
| ``` | |
| Get a free key at [console.groq.com](https://console.groq.com) — the free tier allows 14,400 requests/day. | |
| Launch the Gradio UI: | |
| ```bash | |
| python app_gradio.py | |
| # Opens at http://127.0.0.1:7860 | |
| ``` | |
| Or run the FastAPI server: | |
| ```bash | |
| uvicorn main:app --reload | |
| # API docs at http://127.0.0.1:8000/docs | |
| ``` | |
| ## API Endpoints | |
| | Method | Endpoint | Description | | |
| |---|---|---| | |
| | `GET` | `/health` | Liveness check | | |
| | `GET` | `/stats` | Index size, chunks ingested, model info | | |
| | `POST` | `/ingest` | Upload a PDF; returns chunk count and IDs | | |
| | `POST` | `/query` | Ask a question; returns answer + sources + confidence | | |
| | `POST` | `/query/stream` | Same as `/query` but streams as SSE | | |
| | `POST` | `/debug/chunks` | Return raw retrieved chunks without calling the LLM | | |
| All `POST` query endpoints accept `{"question": "...", "top_k": 5}`. | |
| The `/debug/chunks` endpoint is useful for diagnosing retrieval issues — it shows each chunk's rank, cosine similarity score, source, page, section header, and a 300-character text preview. | |
| ## Evaluation | |
| The evaluation framework has three layers: | |
| **Retrieval metrics** (`evaluation/eval.py`) — given a test set of `(question, expected_source, expected_page)` tuples, measures whether the correct chunk appears in the top-k results: | |
| | Metric | Definition | | |
| |---|---| | |
| | Accuracy@K | Fraction of queries where correct chunk is in top K | | |
| | MRR | Mean Reciprocal Rank of the first correct result | | |
| **End-to-end metrics** (`evaluation/eval_e2e.py`) — runs the full pipeline and scores each answer: | |
| | Metric | Scorer | | |
| |---|---| | |
| | Faithfulness | LLM judge: are all claims grounded in retrieved context? | | |
| | Relevance | LLM judge: does the answer address the question? | | |
| | Keyword coverage | Fraction of expected answer keywords found in output | | |
| **Chunking strategy comparison** (`evaluation/compare.py`) — grid search over chunk sizes, overlap values, and re-ranking on/off. Results table (fill in with your benchmark numbers): | |
| | Strategy | Chunk size | Overlap | Reranking | Acc@5 | MRR | Faithfulness | Relevance | | |
| |---|---|---|---|---|---|---|---| | |
| | recursive_char | 300 | 0 | No | — | — | — | — | | |
| | recursive_char | 500 | 50 | No | — | — | — | — | | |
| | recursive_char | 500 | 50 | Yes | — | — | — | — | | |
| | recursive_char | 1000 | 100 | Yes | — | — | — | — | | |
| | semantic | 500 | — | Yes | — | — | — | — | | |
| ## Project Structure | |
| ``` | |
| rag-qa/ | |
| ├── main.py # FastAPI server | |
| ├── app_gradio.py # Gradio web UI | |
| ├── requirements.txt | |
| ├── Dockerfile | |
| ├── .env.example | |
| │ | |
| ├── ingestion/ | |
| │ ├── pdf_reader.py # pymupdf4llm + pdfplumber extraction | |
| │ ├── chunker.py # fixed_size, recursive_char, semantic strategies | |
| │ ├── embedder.py # SentenceTransformer + contextual enrichment | |
| │ └── pipeline.py # Orchestrates extract → chunk → embed → index | |
| │ | |
| ├── retrieval/ | |
| │ ├── index.py # FAISS IndexFlatIP wrapper | |
| │ ├── bm25_index.py # BM25Okapi wrapper | |
| │ ├── reranker.py # Cross-encoder re-ranking | |
| │ └── searcher.py # Hybrid search + expansion + confidence score | |
| │ | |
| ├── generation/ | |
| │ └── generator.py # Groq client, confidence gating, streaming | |
| │ | |
| ├── evaluation/ | |
| │ ├── judge.py # LLM-as-judge (faithfulness + relevance) | |
| │ ├── eval.py # Retrieval metrics: Acc@K, MRR | |
| │ ├── eval_e2e.py # End-to-end pipeline evaluation | |
| │ └── compare.py # Configuration grid comparison | |
| │ | |
| └── data/ # Sample PDFs for testing | |
| ``` | |
| ## How It Works | |
| **PDF Extraction** — `pymupdf4llm` converts each page to structured markdown, preserving headings and table layout. If it fails, `pdfplumber` extracts flat text. Pages are processed independently so that `page_num` metadata stays accurate for citations. | |
| **Chunking** — three strategies are available. `fixed_size` is a naive sliding window (fast baseline). `recursive_character` uses LangChain's splitter with markdown-aware separators (`\n## `, `\n\n`, `. `) and extracts the nearest heading above each chunk as `section_header` metadata. `semantic` embeds individual sentences and splits at cosine similarity drops below a threshold, grouping topically coherent content together. | |
| **Embedding** — before embedding, each chunk's text is prefixed with `"<filename> | Page N [| Section]\n"`. This means the stored vector encodes not just the chunk's content but where in the document it came from, which improves retrieval precision for location-specific queries. | |
| **Retrieval** — a query vector is searched against FAISS (semantic) and BM25 (keyword) simultaneously. Both result sets are min-max normalized to `[0, 1]` and fused by weighted sum. The top 20 fused candidates are then re-scored by the cross-encoder, which reads the full `(query, chunk)` string pair and produces a relevance score independent of embedding geometry. | |
| **Generation** — `max_score` (cosine similarity of the top-ranked chunk before re-ranking) determines whether to call the LLM. Below 0.3, the call is skipped. Between 0.3 and 0.5, the answer is appended with a confidence warning. Above 0.5, the model receives a numbered context block and is prompted to cite every claim with `[Source N]` notation and synthesize across sources. | |
| ## Deployment | |
| **Docker:** | |
| ```bash | |
| docker build -t rag-qa . | |
| docker run -p 7860:7860 -e GROQ_API_KEY=your_key rag-qa | |
| ``` | |
| **HuggingFace Spaces:** | |
| Set `GROQ_API_KEY` in your Space's Settings → Repository secrets, then push the repository. The `Dockerfile` is picked up automatically by Spaces. | |
| ## Future Improvements | |
| - **OCR support** — integrate `pytesseract` or `surya-ocr` for scanned PDFs that contain no extractable text layer | |
| - **Persistent vector store** — replace the in-memory FAISS index with a disk-backed store (FAISS `write_index` with reload on startup, or a dedicated vector DB) so the index survives server restarts | |
| - **Multi-document cross-referencing** — surface when the same fact is corroborated across multiple uploaded documents rather than citing only one source | |
| - **GPU-accelerated embedding** — the sentence-transformers model runs on CPU by default; passing `device="cuda"` cuts batch embedding time significantly for large document sets | |