Add cross-encoder rerank button + dual embedding models (general/biomedical) with live switch

README.md

@@ -7,37 +7,56 @@ sdk: gradio
 sdk_version: 6.17.3
 app_file: app.py
 pinned: false
-short_description: BM25 vs Dense vs Hybrid retrieval, side by side, on ZeroGPU
+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 **three methods** and shows them
-side by side over a corpus of ~4,000 Europe PMC abstracts (microRNA & disease):
+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 `all-MiniLM-L6-v2` embeddings
+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, fits the PCA, 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*.
+  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 functions. The 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 before presenting.**
-- **Swappable model:** `EMBEDDING_MODEL` in `config.py` is the single source of truth for
-  both the offline build and the live query. Swap the one line (a biomedical
-  PubMedBERT sentence model is included as a comment) and re-run `build_index.py`.
+  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
 

app.py

@@ -1,17 +1,23 @@
-"""Retrieval-methods comparison demo (BM25 vs Dense vs Hybrid/RRF).
+"""Retrieval-methods comparison demo (BM25 vs Dense vs Hybrid/RRF, + optional reranking).
 
 A teaching app for EMBL-EBI researchers. Reads pre-built artifacts from ./data/ (made by
 build_index.py) and only embeds the user's LIVE query at runtime. Read top-to-bottom:
 
     artifacts  ->  embed query (GPU)  ->  filter  ->  3 rankings  ->  RRF  ->  UMAP plot
+                                                           |
+                                              optional: cross-encoder rerank (GPU)
 
-ZeroGPU note: the embedding model is loaded on CPU at import. GPU is touched ONLY inside
-the @spaces.GPU functions. Do not move the model to CUDA anywhere else.
+Two embedding models are shipped (general MiniLM + biomedical PubMedBERT); a radio switches
+which one Dense/Hybrid/plot use (BM25 is lexical, so it's unaffected).
+
+ZeroGPU note: every model is loaded on CPU at import. GPU is touched ONLY inside the
+@spaces.GPU functions. Do not move a model to CUDA anywhere else.
 """
 
 import html
 import json
 import os
+import time
 
 import gradio as gr
 import joblib
@@ -21,7 +27,7 @@ import plotly.graph_objects as go
 import spaces
 import torch
 from rank_bm25 import BM25Okapi
-from sentence_transformers import SentenceTransformer
+from sentence_transformers import CrossEncoder, SentenceTransformer
 
 import config
 from text_utils import tokenize
@@ -30,14 +36,11 @@ from text_utils import tokenize
 # Load pre-built artifacts (fast; no embedding, no network)
 # ====================================================================================
 D = config.DATA_DIR
-EMBEDDINGS = np.load(os.path.join(D, "embeddings.npy"))      # (N, dim), L2-normalised
-COORDS = np.load(os.path.join(D, "umap_coords.npy"))         # (N, 2) UMAP projection
-REDUCER = joblib.load(os.path.join(D, "umap.joblib"))        # projects new points via .transform
+
+# Shared, model-independent artifacts.
 META = pd.read_parquet(os.path.join(D, "metadata.parquet"))
 with open(os.path.join(D, "bm25_tokens.json")) as f:
     BM25 = BM25Okapi(json.load(f))
-with open(os.path.join(D, "config.json")) as f:
-    DATA_CFG = json.load(f)
 
 TITLES = META["title"].tolist()
 ABSTRACTS = META["abstract"].tolist()
@@ -49,50 +52,74 @@ YEAR_MAX = int(YEARS.max())
 # 1300+ distinct journals — show only the most common ones in the dropdown.
 TOP_JOURNALS = META["journal"].value_counts().head(30).index.tolist()
 
-METHOD_COLORS = {"BM25": "#ff7f0e", "Dense": "#1f77b4", "Hybrid": "#2ca02c"}
+# Per-model artifacts + models, keyed by model_key (see config.EMBEDDING_MODELS).
+EMB, COORDS, REDUCER, MODELS = {}, {}, {}, {}
+for _key, (_label, _name) in config.EMBEDDING_MODELS.items():
+    EMB[_key] = np.load(os.path.join(D, f"embeddings_{_key}.npy"))   # (N, dim), L2-normalised
+    COORDS[_key] = np.load(os.path.join(D, f"umap_coords_{_key}.npy"))  # (N, 2)
+    REDUCER[_key] = joblib.load(os.path.join(D, f"umap_{_key}.joblib"))  # .transform new points
+    MODELS[_key] = SentenceTransformer(_name, device="cpu")           # CPU only at import!
+    _gd = getattr(MODELS[_key], "get_embedding_dimension", None) or MODELS[_key].get_sentence_embedding_dimension
+    assert _gd() == EMB[_key].shape[1], f"Model/embedding dim mismatch for '{_key}' — rebuild?"
 
-# ====================================================================================
-# Embedding model — loaded on CPU. GPU is used ONLY inside @spaces.GPU below.
-# ====================================================================================
-MODEL = SentenceTransformer(config.EMBEDDING_MODEL, device="cpu")
-# Prefer the new method name, fall back to the deprecated one across versions.
-_get_dim = getattr(MODEL, "get_embedding_dimension", None) or MODEL.get_sentence_embedding_dimension
-assert _get_dim() == EMBEDDINGS.shape[1], (
-    "Model dim doesn't match committed embeddings — did you swap EMBEDDING_MODEL "
-    "without re-running build_index.py?"
-)
+# Cross-encoder reranker — also CPU at import; moved to GPU inside @spaces.GPU.
+RERANKER = CrossEncoder(config.RERANKER_MODEL, device="cpu")
 
+# UI label <-> internal key.
+MODEL_LABELS = {k: lbl for k, (lbl, _) in config.EMBEDDING_MODELS.items()}
+LABEL_TO_KEY = {lbl: k for k, lbl in MODEL_LABELS.items()}
+DEFAULT_LABEL = MODEL_LABELS[config.DEFAULT_MODEL_KEY]
 
-def _encode(text: str) -> np.ndarray:
-    """Embed one string -> (dim,) float32, L2-normalised. Uses CUDA when available."""
-    device = "cuda" if torch.cuda.is_available() else "cpu"  # safe: only called under @spaces.GPU
-    MODEL.to(device)
-    vec = MODEL.encode(
-        [text], normalize_embeddings=True, convert_to_numpy=True, device=device
-    )
+# Per-method plot styling. Sizes/widths are tiered so that a doc retrieved by several
+# methods nests into concentric rings instead of one trace hiding another.
+METHOD_COLORS = {"BM25": "#ff7f0e", "Dense": "#1f77b4", "Hybrid": "#2ca02c", "Reranked": "#111111"}
+METHOD_SIZES = {"BM25": 20, "Dense": 15, "Hybrid": 10, "Reranked": 13}
+METHOD_WIDTHS = {"BM25": 4, "Dense": 2.5, "Hybrid": 1, "Reranked": 2.5}
+
+
+# ====================================================================================
+# GPU functions — the ONLY place CUDA is touched
+# ====================================================================================
+def _encode(model_key: str, text: str) -> np.ndarray:
+    """Embed one string with the chosen model -> (dim,) float32, L2-normalised."""
+    device = "cuda" if torch.cuda.is_available() else "cpu"  # safe: only under @spaces.GPU
+    model = MODELS[model_key]
+    model.to(device)
+    vec = model.encode([text], normalize_embeddings=True, convert_to_numpy=True, device=device)
     return vec[0].astype(np.float32)
 
 
 @spaces.GPU
-def embed_query(text: str) -> np.ndarray:
+def embed_query(model_key: str, text: str) -> np.ndarray:
     """GPU-backed query embedding for retrieval."""
-    return _encode(text)
+    return _encode(model_key, text)
 
 
 @spaces.GPU
-def embed_text(text: str) -> np.ndarray:
+def embed_text(model_key: str, text: str) -> np.ndarray:
     """GPU-backed embedding for the 'embed your own text' feature."""
-    return _encode(text)
+    return _encode(model_key, text)
+
+
+@spaces.GPU
+def rerank_scores(query: str, texts: list[str]) -> np.ndarray:
+    """Cross-encoder relevance scores for (query, doc) pairs. One transformer forward
+    pass PER candidate — this is the expensive, can't-precompute step."""
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    RERANKER.model.to(device)  # CrossEncoder.device follows the underlying model
+    pairs = [(query, t) for t in texts]
+    scores = RERANKER.predict(pairs, batch_size=16, convert_to_numpy=True, show_progress_bar=False)
+    return np.asarray(scores, dtype=np.float32)
 
 
 # ====================================================================================
-# Retrieval
+# Retrieval (pure CPU / numpy)
 # ====================================================================================
 def candidate_indices(year_lo: int, year_hi: int, journal: str) -> np.ndarray:
     """Apply the metadata filter BEFORE retrieval. Returns indices of allowed docs.
 
-    This pre-filtering of the candidate pool is the distinctly vector-DB feature: all
-    three methods then search only within these documents.
+    This pre-filtering of the candidate pool is the distinctly vector-DB feature: every
+    method then searches only within these documents.
     """
     mask = (YEARS >= int(year_lo)) & (YEARS <= int(year_hi))
     if journal and journal != "All journals":
@@ -128,8 +155,7 @@ def rrf_fuse(bm25: np.ndarray, dense: np.ndarray, cand: np.ndarray, k: int):
 
 def format_results(method: str, results: list[tuple[int, float]], score_label: str) -> str:
     """Render a ranked list as HTML. Each result is a <details> — click the title to
-    expand the FULL abstract (helpful for seeing *why* something ranked highly).
-    This list is the SOURCE OF TRUTH for retrieval.
+    expand the FULL abstract. This list is the SOURCE OF TRUTH for retrieval.
     """
     if not results:
         return f"<h3>{method}</h3><p><em>No results.</em></p>"
@@ -149,47 +175,40 @@ def format_results(method: str, results: list[tuple[int, float]], score_label: s
     return "".join(parts)
 
 
-def make_plot(query_coord=None, retrieved=None, extra_point=None, cand=None) -> go.Figure:
-    """UMAP scatter of the corpus, plus query point + connector lines to hits.
+def make_plot(model_key, query_coord=None, retrieved=None, extra_point=None, cand=None) -> go.Figure:
+    """UMAP scatter for the chosen model, plus query point + connector lines to hits.
 
-    `cand` (optional) = the indices that passed the metadata filter; when given, only
-    those points are drawn as the grey background, so the plot reflects the filter.
+    `cand` (optional) = indices that passed the metadata filter; when given, only those
+    points form the grey background, so the plot reflects the filter.
     """
+    coords = COORDS[model_key]
     fig = go.Figure()
-    bg = np.arange(len(COORDS)) if cand is None else cand
-    # Grey background — "the space the database searches through" (after filtering).
+    bg = np.arange(len(coords)) if cand is None else np.asarray(cand)
     fig.add_trace(
         go.Scatter(
-            x=COORDS[bg, 0], y=COORDS[bg, 1], mode="markers",
+            x=coords[bg, 0], y=coords[bg, 1], mode="markers",
             marker=dict(size=4, color="lightgrey"),
             text=[TITLES[i] for i in bg], hoverinfo="text", name=f"corpus ({len(bg)})",
         )
     )
     if retrieved and query_coord is not None:
-        # Each method is its own trace. A doc retrieved by several methods has ONE
-        # coordinate, so markers/lines land on the same spot. We tier marker size and
-        # line width per method (BM25 biggest, Hybrid smallest) so overlaps nest into
-        # concentric rings instead of hiding each other — and "found by all three" is
-        # then visually obvious. No position jitter (that would misrepresent distance).
-        SIZES = {"BM25": 20, "Dense": 15, "Hybrid": 10}
-        WIDTHS = {"BM25": 4, "Dense": 2.5, "Hybrid": 1}
         for method, results in retrieved.items():
             color = METHOD_COLORS[method]
             xs, ys = [], []
             for doc, _ in results:  # connector lines query -> each hit
-                xs += [query_coord[0], COORDS[doc, 0], None]
-                ys += [query_coord[1], COORDS[doc, 1], None]
+                xs += [query_coord[0], coords[doc, 0], None]
+                ys += [query_coord[1], coords[doc, 1], None]
             fig.add_trace(
                 go.Scatter(x=xs, y=ys, mode="lines",
-                           line=dict(color=color, width=WIDTHS[method]),
+                           line=dict(color=color, width=METHOD_WIDTHS[method]),
                            opacity=0.5, name=f"{method} links", hoverinfo="skip")
             )
             fig.add_trace(
                 go.Scatter(
-                    x=[COORDS[d, 0] for d, _ in results],
-                    y=[COORDS[d, 1] for d, _ in results],
+                    x=[coords[d, 0] for d, _ in results],
+                    y=[coords[d, 1] for d, _ in results],
                     mode="markers",
-                    marker=dict(size=SIZES[method], color=color, symbol="circle-open",
+                    marker=dict(size=METHOD_SIZES[method], color=color, symbol="circle-open",
                                 line=dict(width=2)),
                     text=[TITLES[d] for d, _ in results], hoverinfo="text", name=f"{method} hits",
                 )
@@ -209,7 +228,7 @@ def make_plot(query_coord=None, retrieved=None, extra_point=None, cand=None) ->
         )
     fig.update_layout(
         margin=dict(l=10, r=10, t=30, b=10), height=520,
-        xaxis_title="PC 1", yaxis_title="PC 2",
+        xaxis_title="UMAP-1", yaxis_title="UMAP-2",
         legend=dict(orientation="h", yanchor="bottom", y=1.0),
     )
     return fig
@@ -218,67 +237,97 @@ def make_plot(query_coord=None, retrieved=None, extra_point=None, cand=None) ->
 # ====================================================================================
 # Gradio handlers
 # ====================================================================================
-def run_search(query, k, year_lo, year_hi, journal):
+RERANK_HINT = "<p><em>Press <b>Rerank</b> to reorder the hybrid shortlist with the cross-encoder.</em></p>"
+
+
+def run_search(query, k, year_lo, year_hi, journal, model_label):
+    """Stage 1: BM25 + Dense + Hybrid. Fast. Stashes a shortlist for optional reranking."""
+    key = LABEL_TO_KEY[model_label]
     query = (query or "").strip()
     if not query:
-        return "Enter a query.", "", "", make_plot(), "—"
+        return "Enter a query.", "", "", "", make_plot(key), "—", None, ""
 
     cand = candidate_indices(year_lo, year_hi, journal)
     info = f"**Candidate pool: {len(cand)} / {len(META)} abstracts** after filtering."
     if len(cand) == 0:
-        return "No documents match the filter.", "", "", make_plot(), info
+        return "No documents match the filter.", "", "", "", make_plot(key, cand=cand), info, None, ""
 
     k = int(k)
-    qvec = embed_query(query)                       # <-- the only GPU work in search
-    dense_scores = EMBEDDINGS @ qvec                # exact cosine (vectors are normalised)
+    qvec = embed_query(key, query)                  # <-- GPU work (live query only)
+    dense_scores = EMB[key] @ qvec                  # exact cosine (vectors are normalised)
     bm25_scores = np.asarray(BM25.get_scores(tokenize(query)))
 
     bm25_top = top_k(bm25_scores, cand, k)
     dense_top = top_k(dense_scores, cand, k)
     hybrid_top = rrf_fuse(bm25_scores, dense_scores, cand, k)
-
-    query_coord = REDUCER.transform(qvec.reshape(1, -1))[0]
-    fig = make_plot(
-        query_coord, {"BM25": bm25_top, "Dense": dense_top, "Hybrid": hybrid_top}, cand=cand
-    )
-
+    # Shortlist for the reranker = hybrid's top-N (first-stage retrieval feeds stage 2).
+    shortlist = [d for d, _ in rrf_fuse(bm25_scores, dense_scores, cand, config.RERANK_CANDIDATES)]
+
+    query_coord = REDUCER[key].transform(qvec.reshape(1, -1))[0]
+    fig = make_plot(key, query_coord, {"BM25": bm25_top, "Dense": dense_top, "Hybrid": hybrid_top}, cand=cand)
+
+    state = {
+        "key": key, "query": query,
+        "query_coord": [float(query_coord[0]), float(query_coord[1])],
+        "cand": cand.tolist(), "shortlist": shortlist,
+        "hybrid_top": [[int(d), float(s)] for d, s in hybrid_top],
+    }
     return (
         format_results("BM25", bm25_top, "score"),
         format_results("Dense", dense_top, "cosine"),
         format_results("Hybrid", hybrid_top, "RRF"),
-        fig,
-        info,
+        RERANK_HINT,
+        fig, info, state, "",
     )
 
 
-def run_embed_own(text):
+def run_rerank(state, k):
+    """Stage 2: cross-encoder reranks the hybrid shortlist. Slow — and we time it."""
+    if not state:
+        return "<p><em>Run a search first.</em></p>", gr.update(), ""
+
+    key, query, shortlist = state["key"], state["query"], state["shortlist"]
+    texts = [f"{TITLES[i]}. {ABSTRACTS[i]}" for i in shortlist]
+
+    t0 = time.time()
+    scores = rerank_scores(query, texts)            # <-- GPU work, one pass per candidate
+    dt = time.time() - t0
+
+    order = np.argsort(-scores)[: int(k)]
+    reranked = [(int(shortlist[i]), float(scores[i])) for i in order]
+
+    hybrid_top = [(int(d), float(s)) for d, s in state["hybrid_top"]]
+    fig = make_plot(key, state["query_coord"], {"Hybrid": hybrid_top, "Reranked": reranked},
+                    cand=np.asarray(state["cand"]))
+    timer = (
+        f"⏱️ Reranked **{len(shortlist)}** candidates with the cross-encoder in "
+        f"**{dt:.2f}s**. It ran a transformer on every (query, document) pair at query "
+        f"time — contrast the instant BM25/Dense/Hybrid. The plot shows Hybrid (green) "
+        f"vs Reranked (black)."
+    )
+    return format_results("Reranked", reranked, "cross-enc"), fig, timer
+
+
+def run_embed_own(text, model_label):
+    key = LABEL_TO_KEY[model_label]
     text = (text or "").strip()
     if not text:
-        return make_plot()
-    vec = embed_text(text)                           # <-- GPU work
-    coord = REDUCER.transform(vec.reshape(1, -1))[0]
-    return make_plot(extra_point=(coord, text[:80]))
+        return make_plot(key)
+    vec = embed_text(key, text)                      # <-- GPU work
+    coord = REDUCER[key].transform(vec.reshape(1, -1))[0]
+    return make_plot(key, extra_point=(coord, text[:80]))
 
 
 # ====================================================================================
 # Example queries — tuned to the cardiovascular slice of this microRNA/disease corpus.
 # Each is chosen (and verified against the corpus) to make ONE method visibly win.
 # ====================================================================================
-# BM25 wins: an exact miRNA identifier. The literal token "miR-208a" is matched precisely;
-# dense retrieval drifts to other, merely-similar cardiac miRNA papers.
-EXACT_ID_QUERY = "miR-208a"
-# Dense wins: a plain-English paraphrase using NONE of the corpus jargon (no "microRNA",
-# "fibrosis", "myocardial"). BM25 has almost no tokens to match; dense maps the meaning
-# to miRNA-in-cardiac-fibrosis papers.
+EXACT_ID_QUERY = "miR-208a"  # BM25 wins: exact identifier token, matched verbatim.
+# Dense wins: plain-English paraphrase using none of the corpus jargon.
 PARAPHRASE_QUERY = "small RNA molecules that worsen scarring after a heart attack"
-# BM25 wins: a bare acronym (acute myocardial infarction). "AMI" is a literal token BM25
-# nails, but as a 3-letter string it carries little semantic signal, so dense flounders.
-ACRONYM_QUERY = "AMI"
-# BM25 wins: a specific molecule + mechanism. BM25 pins the exact miR-499 papers; dense
-# returns topically-similar cardiomyocyte-apoptosis papers about *other* miRNAs.
-RARE_TERM_QUERY = "miR-499 cardiomyocyte apoptosis"
-# Hybrid wins: a broad real query where lexical and semantic each surface good-but-
-# different papers, and RRF fuses them into the strongest combined ranking.
+ACRONYM_QUERY = "AMI"  # BM25 wins: bare acronym is a lexical token; dense has little to grip.
+RARE_TERM_QUERY = "miR-499 cardiomyocyte apoptosis"  # BM25 pins the specific molecule.
+# Hybrid wins: broad query where lexical + semantic each surface good-but-different papers.
 BROAD_CONCEPT_QUERY = "circulating microRNA biomarkers for cardiovascular disease"
 
 EXAMPLES = [
@@ -300,12 +349,12 @@ PLOT_CAVEAT = (
 # ====================================================================================
 # UI
 # ====================================================================================
-with gr.Blocks(title="RAG retrieval: BM25 vs Dense vs Hybrid") as demo:
+with gr.Blocks(title="RAG retrieval: BM25 vs Dense vs Hybrid vs Rerank") as demo:
     gr.Markdown(
         "# 🔍 Retrieval methods, side by side\n"
-        "Compare **BM25** (lexical), **Dense** (vector cosine), and **Hybrid** "
-        "(Reciprocal Rank Fusion) over a corpus of biomedical abstracts. "
-        f"Corpus: {len(META):,} Europe PMC abstracts on microRNA & disease."
+        "Compare **BM25** (lexical), **Dense** (vector cosine), **Hybrid** (Reciprocal "
+        "Rank Fusion), and an optional **cross-encoder rerank** over a corpus of "
+        f"{len(META):,} Europe PMC abstracts on microRNA & disease."
     )
 
     with gr.Row():
@@ -317,8 +366,13 @@ with gr.Blocks(title="RAG retrieval: BM25 vs Dense vs Hybrid") as demo:
             btn = gr.Button(label, size="sm")
             btn.click(lambda t=text: t, outputs=query)
 
+    model_radio = gr.Radio(
+        choices=list(LABEL_TO_KEY.keys()), value=DEFAULT_LABEL,
+        label="Embedding model — used by Dense, Hybrid and the plot (BM25 is lexical, so it never changes)",
+    )
+
     with gr.Group():
-        gr.Markdown("**Metadata filter** — restricts the candidate pool *before* retrieval (all 3 methods).")
+        gr.Markdown("**Metadata filter** — restricts the candidate pool *before* retrieval (all methods).")
         with gr.Row():
             year_lo = gr.Slider(YEAR_MIN, YEAR_MAX, value=YEAR_MIN, step=1, label="Year from")
             year_hi = gr.Slider(YEAR_MIN, YEAR_MAX, value=YEAR_MAX, step=1, label="Year to")
@@ -332,24 +386,31 @@ with gr.Blocks(title="RAG retrieval: BM25 vs Dense vs Hybrid") as demo:
         bm25_out = gr.HTML(label="BM25")
         dense_out = gr.HTML(label="Dense")
         hybrid_out = gr.HTML(label="Hybrid")
+        reranked_out = gr.HTML(label="Reranked")
+
+    with gr.Row():
+        rerank_btn = gr.Button("Rerank hybrid shortlist (cross-encoder)", variant="secondary")
+        rerank_timer = gr.Markdown("")
 
     gr.Markdown("## Vector space (UMAP projection)")
-    plot = gr.Plot(value=make_plot())
+    plot = gr.Plot(value=make_plot(config.DEFAULT_MODEL_KEY))
     gr.Markdown(PLOT_CAVEAT)
 
     with gr.Accordion("Embed your own text", open=False):
         gr.Markdown(
-            "Type anything — it gets embedded with the same model and dropped onto the "
+            "Type anything — it gets embedded with the selected model and dropped onto the "
             "map. This *is* the space the database searches through."
         )
         own_text = gr.Textbox(label="Your text", placeholder="e.g. a sentence about gene regulation…")
         own_btn = gr.Button("Embed & plot")
-        own_btn.click(run_embed_own, inputs=own_text, outputs=plot)
-
-    inputs = [query, k, year_lo, year_hi, journal]
-    outputs = [bm25_out, dense_out, hybrid_out, plot, filter_info]
-    search_btn.click(run_search, inputs=inputs, outputs=outputs)
-    query.submit(run_search, inputs=inputs, outputs=outputs)
+        own_btn.click(run_embed_own, inputs=[own_text, model_radio], outputs=plot)
+
+    search_state = gr.State()
+    search_inputs = [query, k, year_lo, year_hi, journal, model_radio]
+    search_outputs = [bm25_out, dense_out, hybrid_out, reranked_out, plot, filter_info, search_state, rerank_timer]
+    search_btn.click(run_search, inputs=search_inputs, outputs=search_outputs)
+    query.submit(run_search, inputs=search_inputs, outputs=search_outputs)
+    rerank_btn.click(run_rerank, inputs=[search_state, k], outputs=[reranked_out, plot, rerank_timer])
 
 
 if __name__ == "__main__":

build_index.py

@@ -92,51 +92,46 @@ def main():
     # The text we both embed AND tokenise for BM25 — keep them identical for a fair fight.
     docs = (df["title"] + ". " + df["abstract"]).tolist()
 
-    # 2. Embed ----------------------------------------------------------------------
-    print(f"Loading embedding model: {config.EMBEDDING_MODEL}")
-    model = SentenceTransformer(config.EMBEDDING_MODEL)
-    print("Embedding corpus (this is the slow part)...")
-    embeddings = model.encode(
-        docs,
-        batch_size=64,
-        show_progress_bar=True,
-        normalize_embeddings=True,  # so cosine similarity == dot product
-        convert_to_numpy=True,
-    ).astype(np.float32)
-    np.save(os.path.join(config.DATA_DIR, "embeddings.npy"), embeddings)
-    print(f"Saved embeddings: {embeddings.shape}")
-
-    # 3. UMAP -----------------------------------------------------------------------
-    # UMAP shows local cluster structure better than PCA. It also supports projecting
-    # new/out-of-sample points via .transform() — we use that to place the live query.
-    # Caveat: single-point transform is an approximate fit, so query placement is rough.
-    print("Fitting 2D UMAP (cosine metric)...")
-    reducer = umap.UMAP(
-        n_components=2, metric="cosine", n_neighbors=15, min_dist=0.1, random_state=42
-    )
-    coords = reducer.fit_transform(embeddings).astype(np.float32)
-    joblib.dump(reducer, os.path.join(config.DATA_DIR, "umap.joblib"))
-    np.save(os.path.join(config.DATA_DIR, "umap_coords.npy"), coords)
-    print(f"Saved UMAP reducer + corpus coords: {coords.shape}")
-
-    # 4. BM25 tokens ----------------------------------------------------------------
+    # --- Shared (model-independent) artifacts --------------------------------------
+    # BM25 is lexical and the metadata is just text, so both are shared across models.
     print("Tokenising for BM25...")
     tokens = [tokenize(doc) for doc in docs]
     with open(os.path.join(config.DATA_DIR, "bm25_tokens.json"), "w") as f:
         json.dump(tokens, f)
-    print(f"Saved {len(tokens)} token lists.")
-
-    # 5. Metadata -------------------------------------------------------------------
     df.to_parquet(os.path.join(config.DATA_DIR, "metadata.parquet"), index=False)
-    print(f"Saved metadata.parquet: {df.shape}")
+    print(f"Saved bm25_tokens.json and metadata.parquet ({df.shape}).")
+
+    # --- Per-model artifacts: embeddings + UMAP ------------------------------------
+    cfg_record = {}
+    for key, (label, model_name) in config.EMBEDDING_MODELS.items():
+        print(f"\n[{key}] Loading embedding model: {model_name}")
+        model = SentenceTransformer(model_name)
+        print(f"[{key}] Embedding corpus (slow)...")
+        embeddings = model.encode(
+            docs,
+            batch_size=64,
+            show_progress_bar=True,
+            normalize_embeddings=True,  # so cosine similarity == dot product
+            convert_to_numpy=True,
+        ).astype(np.float32)
+        np.save(os.path.join(config.DATA_DIR, f"embeddings_{key}.npy"), embeddings)
+        print(f"[{key}] Saved embeddings: {embeddings.shape}")
+
+        # UMAP shows local cluster structure better than PCA, and projects new query
+        # points via .transform(). Caveat: single-point transform is an approximate fit.
+        print(f"[{key}] Fitting 2D UMAP (cosine metric)...")
+        reducer = umap.UMAP(
+            n_components=2, metric="cosine", n_neighbors=15, min_dist=0.1, random_state=42
+        )
+        coords = reducer.fit_transform(embeddings).astype(np.float32)
+        joblib.dump(reducer, os.path.join(config.DATA_DIR, f"umap_{key}.joblib"))
+        np.save(os.path.join(config.DATA_DIR, f"umap_coords_{key}.npy"), coords)
+        print(f"[{key}] Saved UMAP reducer + coords: {coords.shape}")
+        cfg_record[key] = {"model": model_name, "dim": int(embeddings.shape[1])}
 
-    # 6. Config record (startup sanity check in app.py) -----------------------------
+    # Config record (startup sanity check in app.py).
     with open(os.path.join(config.DATA_DIR, "config.json"), "w") as f:
-        json.dump(
-            {"embedding_model": config.EMBEDDING_MODEL, "dim": int(embeddings.shape[1])},
-            f,
-            indent=2,
-        )
+        json.dump(cfg_record, f, indent=2)
 
     print("\nDone. Artifacts written to", config.DATA_DIR)
 

config.py

@@ -5,12 +5,23 @@ are embedded with the *same* model. If they ever diverge, the query vector and t
 vectors live in different spaces and every retrieval result is silently wrong.
 """
 
-# --- Embedding model -----------------------------------------------------------------
-# Small, fast, 384-dim. The default for this demo.
-EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
-# Biomedical alternative (768-dim) — swap the line above for this one and rebuild to use a
-# domain model. ZeroGPU gives us the headroom for it.
-# EMBEDDING_MODEL = "pritamdeka/S-PubMedBert-MS-MARCO"
+# --- Embedding models ----------------------------------------------------------------
+# We build & ship artifacts for BOTH models so the Space can switch between them live.
+# model_key -> (display label, HuggingFace model name). The key is used in artifact
+# filenames (embeddings_<key>.npy, umap_<key>.joblib, umap_coords_<key>.npy), so keep it
+# filesystem-safe. Add a model here + rebuild to offer it in the UI.
+EMBEDDING_MODELS = {
+    "general": ("General — MiniLM-L6 (384-dim)", "sentence-transformers/all-MiniLM-L6-v2"),
+    "biomedical": ("Biomedical — S-PubMedBert (768-dim)", "pritamdeka/S-PubMedBert-MS-MARCO"),
+}
+DEFAULT_MODEL_KEY = "general"
+
+# --- Reranker (cross-encoder, runtime only) ------------------------------------------
+# A cross-encoder scores (query, document) PAIRS jointly — more accurate than the
+# bi-encoder dense retrieval, but it can't be precomputed, so it only runs on a shortlist.
+RERANKER_MODEL = "cross-encoder/ms-marco-MiniLM-L-6-v2"
+# Biomedical alternative: "ncbi/MedCPT-Cross-Encoder"
+RERANK_CANDIDATES = 30  # size of the hybrid shortlist fed to the reranker
 
 # --- Corpus (offline build only) -----------------------------------------------------
 # Europe PMC search query. See https://europepmc.org/Help for the query syntax.

data/config.json

@@ -1,4 +1,10 @@
 {
-  "embedding_model": "sentence-transformers/all-MiniLM-L6-v2",
-  "dim": 384
+  "general": {
+    "model": "sentence-transformers/all-MiniLM-L6-v2",
+    "dim": 384
+  },
+  "biomedical": {
+    "model": "pritamdeka/S-PubMedBert-MS-MARCO",
+    "dim": 768
+  }
 }

data/embeddings_biomedical.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:192c1aa5aea5d13d32d86d78fca9e51b1a6e186d4df9df374445ba6053cfbb44
+size 12266624

data/{embeddings.npy → embeddings_general.npy}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:2f81a7b345de3c3173cd6b8219789968f061980d7d5d6dffe8445e82d84298c6
+oid sha256:f7ae07dfbae0a82d14b930f6c25f39ec667fe47ca95d5915ce9de4fcb9e61a74
 size 6133376

data/metadata.parquet

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:d002fef7b9475d2794670200a5ec53092d7816d769583330c6ec7dedfabc24a1
-size 3825335
+oid sha256:9a2236a72b5f1c6af19dc333c36f02567529fe38766ecf1c03b7bb9bf4017aeb
+size 3826858

data/umap_biomedical.joblib

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b65f32aa01e697b7f1e1376a14598dd61256f0b5b75a7a9bbc0180bd6e6b5f5e
+size 13078057

data/{umap_coords.npy → umap_coords_biomedical.npy}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:e94741ae6cd213b43a26442c133d5ab9988024cd2e548bb2729e272dfa05e7af
+oid sha256:79f661e76f9248189db9209afb1a9b49043149414c89b7e45c27880711b599fe
 size 32072

data/{umap.joblib → umap_coords_general.npy}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:b6a25ecb9954e1500337230a69921f65caff1c4b2c615fce4e5deff58b86e4a8
-size 6913369
+oid sha256:0be6e2de236caa6973df9fc0c7d6f03acc005146e842c23e71b747a8b063b36c
+size 32072

data/umap_general.joblib

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:92ed6eaef29f645eb22a6895bd2e99bb6b57cf2c4b5a7cd22336b649d6ad4340
+size 6913641