Skip to content

pieroevcc/ProbabilityRAG

Repository files navigation

🎲 ProbabilityRAG

A local-first RAG tutor that answers probability questions straight from the textbook - with citations you can open to the exact page.

Ask a probability question and it retrieves the relevant passages from Grinstead & Snell (hybrid dense + sparse search, then a cross-encoder rerank), grounds an LLM on them, and returns a cited answer you can trace back to the source PDF page. It refuses instead of guessing when the material doesn't cover the question - and it runs entirely offline on a single 8 GB GPU.

Free-tier hosting, so be patient: ~40s wake-up if it's been sleeping, ~30–90s per answer (the same question answers in ~8 s on local GPU). Limited to 5 questions/day per visitor.

ProbabilityRAG demo - ask, sources render, answer streams with LaTeX, citation opens the PDF at the exact page


🎨 High-level design

Simplified (left) and detailed (right) views of the system - start at the circles and follow the arrows.

ProbabilityRAG architecture - a simplified and a detailed high-level design of the indexing and query pipelines


✨ Features

  • 🟢 Grounded, cited answers - every claim cites its source chunk; click a citation to open the source PDF at that exact page.
  • 🟢 Hybrid retrieval - BGE-M3 dense and sparse vectors fused (RRF), then reranked by a cross-encoder.
  • 🟢 Query decomposition - comparison / multi-part questions are split into sub-queries, each reranked on its own, then round-robin merged.
  • 🟢 Three answer registers - the default terse answer, "Why? ↳ go deeper," and a plain-language "Simplify."
  • 🟢 Refuses when unsupported - replies "I don't know based on the provided material" instead of hallucinating.
  • 🟢 Streaming - sources render first, then tokens stream in (SSE), with a stop button.
  • 🟢 Fully local & offline - retrieval models and the LLM co-reside on one 8 GB GPU; no cloud API needed.
  • 🟢 Model switcher - a server-side allowlist maps names to local Ollama models or any OpenAI-compatible API (the demo runs GLM-4.7-Flash); the UI selector appears only when there's a real choice.
  • 🟢 Deployed public demo - one Docker image on HF Spaces: index and model weights baked at build, per-IP + global daily rate limits (hashed IPs, SQLite), CORS pinned, LLM failures surfaced as SSE error frames.
  • 🟢 Retrieval eval harness - a 45-question golden set scored by Hit@k / MRR drives every change, plus 57 unit tests (~1 s, GPU-free) run by CI on every push.

🧰 Tech stack

Area Tooling
PDF → math parsing Marker (Markdown + LaTeX)
Embeddings BGE-M3 (dense + sparse) via FlagEmbedding
Vector store Qdrant (native hybrid + RRF fusion)
Reranker bge-reranker-v2-m3
Generation Ollama - qwen2.5:3b-instruct locally · GLM-4.7-Flash (free API) in the demo
Backend FastAPI + Uvicorn
Deployment Multi-stage DockerHF Spaces free tier (index baked at build)
CI GitHub Actions - pytest on every push (torch stubbed, no GPU wheels)
Frontend React + Vite + Tailwind
Math rendering KaTeX + remark-math / rehype-katex
Env uv + Python 3.12

🛠️ Why/How I built it

I'm taking a hard probability course this fall, and I wanted a tutor that answered from the actual textbook instead of confidently making things up. I also just wanted to learn RAG properly so I wrote every stage in raw Python instead of reaching for LlamaIndex. Slower, but I actually understand what each piece is doing now.

The constraint that made it fun was keeping everything local on an 8 GB GPU - the retrieval models and the LLM have to share the card, no cloud API to lean on. I built the riskiest stuff first (PDF→math parsing, then retrieval) and checked every change against a golden set before moving on. The part that took longest was comparison questions like "compare the binomial and geometric variance." A single embedding can't sit near two distributions at once, so it'd only ever retrieve one side. Splitting those into sub-queries and merging the results took multi-hop retrieval from 42/45 to 45/45 - and taught me the annoying lesson that my reranker was fine the whole time; I was just handing it too small a candidate pool to ever see the right chunk.


📚 What I learned

  • Most of my "the model got it wrong" moments were actually retrieval - the passage just wasn't in the top-k. Fix retrieval before you touch the prompt.
  • A reranker can only reorder what you give it - I spent a while "fixing ranking" when the real bug was the pool size upstream, starving it of the right chunk.
  • Without the eval set I was just guessing - Hit@k / MRR turned "this feels better" into a number I could actually point at.
  • Pin your transitive dependencies before deploying - the demo's first build crashed in production because unpinned transformers resolved to a new major version that removed an API my reranker library called. Worked on my machine; died in the container.
  • Silent failures are the most expensive kind - the hosted LLM call was failing and the response stream just… stopped, with nothing to debug from. Making the server emit the provider's actual error as an SSE frame turned a mystery into a 5-minute fix (wrong model id), and made every future failure diagnosable.

🚀 How it could be improved

  • Comparison queries are ~12s vs ~8s single-hop - the decomposition lanes run sequentially; parallelizing the independent embed→retrieve→rerank lanes would cut that materially.
  • Eval covers retrieval, not generation quality - add RAGAS Faithfulness / Answer-Relevancy with a local judge for an end-to-end number.
  • The free-tier demo is slow (~30–90 s/answer) - retrieval runs fp32 on 2 shared vCPUs behind one worker. Paid CPU/GPU hardware or quantized ONNX embeddings would cut it dramatically.
  • Covariance has no corpus coverage - Grinstead & Snell never defines it, so that query (correctly) refuses; ingesting a supplementary source would close the gap.

▶️ How to run the project

Just want to try it? Use the live demo - no setup. To deploy your own copy, see DEPLOY.md.

Prerequisites: Ollama running with qwen2.5:3b-instruct pulled, and the Grinstead & Snell PDF at data/raw/grinstead_snell.pdf (the corpus is git-ignored - licensing).

1. Install deps

uv venv && uv pip install -r requirements.txt   # Python 3.12 env (torch, FlagEmbedding, qdrant-client, fastapi…)
cd web && npm install                            # frontend

2. Build the index (one time)

.venv/Scripts/python -u scripts/embed_store.py --build

3. Run it

# backend (loads BGE-M3 + reranker once, ~40s):
.venv/Scripts/python -m uvicorn app.server:app --port 8000
# frontend:
cd web && npm run dev         # http://localhost:5173
CLI (no web UI)
.venv/Scripts/python -u scripts/ask.py "What is the mgf of a Poisson?"
.venv/Scripts/python -u scripts/ask.py --think "Derive Var of a binomial via its mgf"
Operational gotchas (important)
  • Exactly ONE Ollama server, on the GPU - the desktop app OR ollama serve, never both (the loser serves on CPU and everything looks mysteriously slow). Verify: ollama psPROCESSOR = GPU.
  • Qdrant single-writer lock - stop the FastAPI server before running any script that opens qdrant_storage/ (ask.py, eval_retrieval.py, embed_store.py). The eval also needs Ollama up (query decomposition calls it).
  • Higher-quality model: PROBRAG_MODEL=deepseek-r1:7b (better answers, but spills to CPU on 8 GB → minutes/answer). Or enable several and pick in the UI: PROBRAG_MODELS=qwen,glm-local (full env reference in DEPLOY.md).
  • Full dev notes in eval/IMPROVEMENTS.md.
Tests / eval
# unit tests - pure pipeline logic (chunking, retrieval merge, refusal, stream protocol); ~1s, no GPU needed
.venv/Scripts/python -m pytest tests/ -q

# retrieval eval - Hit@k / MRR on the 45-question golden set (stop the backend first; Ollama up)
.venv/Scripts/python scripts/eval_retrieval.py

Machine: Dell XPS 16 · Win 11 · Intel Core Ultra 9 · 32 GB RAM · RTX 4060 (8 GB VRAM).

About

Local-first RAG tutor over Grinstead & Snell probability text (BGE-M3 hybrid + reranker, Qdrant, Ollama, FastAPI/React).

Topics

Resources

Stars

0 stars

Watchers

0 watching

Forks

Contributors