---
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