vector_demo / README.md
afg1's picture
Add cross-encoder rerank button + dual embedding models (general/biomedical) with live switch
a2c6853 verified
|
Raw
History Blame Contribute Delete
4.16 kB

A newer version of the Gradio SDK is available: 6.20.0

Upgrade
metadata
title: RAG Retrieval Compare
emoji: πŸ”
colorFrom: indigo
colorTo: green
sdk: gradio
sdk_version: 6.17.3
app_file: app.py
pinned: false
short_description: BM25 vs Dense vs Hybrid vs Rerank, side by side, on ZeroGPU

πŸ” Retrieval methods, side by side

A teaching demo for technical researchers learning vector-database fundamentals for RAG. Given a query, it retrieves the top-k results from several methods and shows them side by side over a corpus of ~4,000 Europe PMC abstracts (microRNA & disease):

  1. BM25 (lexical) β€” rank-bm25
  2. Dense (vector) β€” exact cosine similarity over sentence-transformer embeddings
  3. Hybrid β€” Reciprocal Rank Fusion (RRF, k=60) over the BM25 and dense rankings
  4. Reranked (optional, on-demand button) β€” a cross-encoder that scores (query, document) pairs jointly and reorders the hybrid shortlist. It can't be precomputed, so it runs live and the UI times it to make the cost visible.

Two embedding models are shipped β€” general (all-MiniLM-L6-v2, 384-dim) and biomedical (pritamdeka/S-PubMedBert-MS-MARCO, 768-dim) β€” and a radio switches which one Dense, Hybrid and the plot use. BM25 is lexical, so it's unaffected by the switch.

Plus a metadata filter (year / journal) that restricts the candidate pool before retrieval β€” the distinctly vector-DB feature β€” and a UMAP scatter plot for spatial intuition, with connector lines from the query to the documents each method retrieved. Click any result to expand its full abstract.

Retrieve-then-rerank, and how it differs from RRF

  • RRF (Hybrid) is a fusion of rankings β€” it only uses BM25's and Dense's rank positions, no model, essentially free. It improves recall by combining lexical + semantic signal.
  • The reranker is a second-stage precision step. A bi-encoder (Dense) embeds query and document separately; a cross-encoder runs them through a transformer together, so it judges relevance far more accurately β€” but at one forward pass per candidate, so it only runs on a shortlist (here, the hybrid top-30). It can only reorder what stage 1 surfaced, so recall still matters.

How it's built (read this before the talk)

  • Everything is precomputed offline by build_index.py, which fetches the abstracts, embeds the corpus with each model, fits a UMAP per model, and tokenises for BM25. The artifacts in ./data/ are committed. The app never embeds the corpus or calls an external API at startup β€” it only embeds your live query (and, on demand, runs the reranker).
  • ZeroGPU: GPU is allocated on demand only inside the @spaces.GPU-decorated query embedding and reranking functions. Every model is loaded on CPU at import; nothing touches CUDA at startup. ⚠️ The first GPU call after idle has a cold start of a few seconds β€” pre-warm with a dummy query (and a dummy rerank) before presenting.
  • Swappable models: config.py holds the EMBEDDING_MODELS registry and RERANKER_MODEL β€” the single source of truth for both the offline build and the live app. Add/swap a model there and re-run build_index.py.

Exact search, honestly

The corpus is tiny, so dense retrieval uses exact cosine similarity (a plain NumPy dot product) β€” instant and correct. At production scale you'd reach for an ANN index (e.g. HNSW) to keep search fast over millions of vectors. We deliberately don't show a fake HNSW timing comparison here: at this scale the difference is invisible, so it would mislead rather than teach.

Teaching caveat on the plot

The plot uses UMAP for the 2D projection (it shows local cluster structure better than PCA). Two honest caveats: UMAP warps global distances, and the live query is placed by an approximate out-of-sample transform, so its position is only indicative. Retrieved points may not be the visually-closest dots. The ranked lists are authoritative; the plot is only for intuition.

Rebuild the index locally

uv venv
uv pip install -r requirements.txt
uv run python build_index.py   # writes ./data/