UMAP projection + filtered plot + expandable abstracts + cardiovascular example queries

README.md

@@ -21,8 +21,9 @@ side by side over a corpus of ~4,000 Europe PMC abstracts (microRNA & disease):
 3. **Hybrid** — Reciprocal Rank Fusion (RRF, k=60) over the BM25 and dense rankings
 
 Plus a **metadata filter** (year / journal) that restricts the candidate pool *before*
-retrieval — the distinctly vector-DB feature — and a **PCA scatter plot** for spatial
+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.
 
 ## How it's built (read this before the talk)
 
@@ -48,9 +49,11 @@ mislead rather than teach.
 
 ## Teaching caveat on the plot
 
-The 2D PCA projection distorts true high-dimensional distances, so the retrieved points
-may not be the visually-closest dots. **The ranked lists are authoritative; the plot is
-only for intuition.**
+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
 

app.py

@@ -3,12 +3,13 @@
 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  ->  plot
+    artifacts  ->  embed query (GPU)  ->  filter  ->  3 rankings  ->  RRF  ->  UMAP plot
 
 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.
 """
 
+import html
 import json
 import os
 
@@ -30,8 +31,8 @@ from text_utils import tokenize
 # ====================================================================================
 D = config.DATA_DIR
 EMBEDDINGS = np.load(os.path.join(D, "embeddings.npy"))      # (N, dim), L2-normalised
-PCA_COORDS = np.load(os.path.join(D, "pca_coords.npy"))      # (N, 2)
-PCA = joblib.load(os.path.join(D, "pca.joblib"))
+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
 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))
@@ -54,7 +55,9 @@ METHOD_COLORS = {"BM25": "#ff7f0e", "Dense": "#1f77b4", "Hybrid": "#2ca02c"}
 # Embedding model — loaded on CPU. GPU is used ONLY inside @spaces.GPU below.
 # ====================================================================================
 MODEL = SentenceTransformer(config.EMBEDDING_MODEL, device="cpu")
-assert MODEL.get_sentence_embedding_dimension() == EMBEDDINGS.shape[1], (
+# 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?"
 )
@@ -124,29 +127,42 @@ 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 Markdown. This list is the SOURCE OF TRUTH for retrieval."""
+    """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.
+    """
     if not results:
-        return f"### {method}\n\n_No results._"
-    lines = [f"### {method}"]
+        return f"<h3>{method}</h3><p><em>No results.</em></p>"
+    parts = [f"<h3>{method}</h3>"]
     for rank, (doc, score) in enumerate(results, start=1):
-        snippet = ABSTRACTS[doc][:200].rsplit(" ", 1)[0] + "…"
-        lines.append(
-            f"**{rank}. {TITLES[doc]}**  \n"
-            f"`{score_label}={score:.3f}` · {int(YEARS[doc])} · *{JOURNALS[doc]}*  \n"
-            f"{snippet}\n"
+        title = html.escape(TITLES[doc])          # abstracts contain <, >, & — must escape
+        abstract = html.escape(ABSTRACTS[doc])
+        journal = html.escape(str(JOURNALS[doc]))
+        parts.append(
+            "<details style='margin-bottom:10px;border-bottom:1px solid #ddd;padding-bottom:6px;'>"
+            f"<summary style='cursor:pointer;'><b>{rank}. {title}</b><br>"
+            f"<code>{score_label}={score:.3f}</code> · {int(YEARS[doc])} · <em>{journal}</em>"
+            "</summary>"
+            f"<p style='margin-top:6px;font-size:0.9em;line-height:1.4;'>{abstract}</p>"
+            "</details>"
         )
-    return "\n".join(lines)
+    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(query_coord=None, retrieved=None, extra_point=None) -> go.Figure:
-    """PCA scatter of the whole corpus, 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.
+    """
     fig = go.Figure()
-    # Whole corpus as grey background — "the space the database searches through".
+    bg = np.arange(len(COORDS)) if cand is None else cand
+    # Grey background — "the space the database searches through" (after filtering).
     fig.add_trace(
         go.Scatter(
-            x=PCA_COORDS[:, 0], y=PCA_COORDS[:, 1], mode="markers",
-            marker=dict(size=4, color="lightgrey"), text=TITLES, hoverinfo="text",
-            name="corpus",
+            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:
@@ -154,16 +170,16 @@ def make_plot(query_coord=None, retrieved=None, extra_point=None) -> go.Figure:
             color = METHOD_COLORS[method]
             xs, ys = [], []
             for doc, _ in results:  # connector lines query -> each hit
-                xs += [query_coord[0], PCA_COORDS[doc, 0], None]
-                ys += [query_coord[1], PCA_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=1),
                            opacity=0.5, name=f"{method} links", hoverinfo="skip")
             )
             fig.add_trace(
                 go.Scatter(
-                    x=[PCA_COORDS[d, 0] for d, _ in results],
-                    y=[PCA_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=11, color=color, symbol="circle-open", line=dict(width=2)),
                     text=[TITLES[d] for d, _ in results], hoverinfo="text", name=f"{method} hits",
@@ -212,8 +228,10 @@ def run_search(query, k, year_lo, year_hi, journal):
     dense_top = top_k(dense_scores, cand, k)
     hybrid_top = rrf_fuse(bm25_scores, dense_scores, cand, k)
 
-    query_coord = PCA.transform(qvec.reshape(1, -1))[0]
-    fig = make_plot(query_coord, {"BM25": bm25_top, "Dense": dense_top, "Hybrid": hybrid_top})
+    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
+    )
 
     return (
         format_results("BM25", bm25_top, "score"),
@@ -229,39 +247,45 @@ def run_embed_own(text):
     if not text:
         return make_plot()
     vec = embed_text(text)                           # <-- GPU work
-    coord = PCA.transform(vec.reshape(1, -1))[0]
+    coord = REDUCER.transform(vec.reshape(1, -1))[0]
     return make_plot(extra_point=(coord, text[:80]))
 
 
 # ====================================================================================
-# Example queries — TODO: the presenter fills these in.
-# Each placeholder is meant to make ONE method visibly win; pick real strings against the
-# current microRNA/disease corpus before the talk.
+# 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 should win: a precise token/ID that appears VERBATIM in abstracts.
-EXACT_ID_QUERY = "TODO_EXACT_ID: e.g. a specific miRNA like 'miR-21' or a gene symbol"
-# Dense should win: a conceptual paraphrase using NONE of the corpus's exact words.
-PARAPHRASE_QUERY = "TODO_PARAPHRASE: e.g. 'small RNAs that switch genes off in tumours'"
-# Lexical vs semantic gap: an acronym whose expansion is what's written in the text.
-ACRONYM_QUERY = "TODO_ACRONYM: e.g. an acronym vs its full form"
-# BM25 strong on rare exact tokens.
-RARE_TERM_QUERY = "TODO_RARE_TERM: e.g. a rare, very specific technical term"
-# Hybrid should win: a broad topic where fusing lexical + semantic beats either alone.
-BROAD_CONCEPT_QUERY = "TODO_BROAD: e.g. 'microRNA biomarkers for early cancer detection'"
+# 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.
+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.
+BROAD_CONCEPT_QUERY = "circulating microRNA biomarkers for cardiovascular disease"
 
 EXAMPLES = [
     ("Exact ID (BM25)", EXACT_ID_QUERY),
     ("Paraphrase (Dense)", PARAPHRASE_QUERY),
-    ("Acronym", ACRONYM_QUERY),
-    ("Rare term (BM25)", RARE_TERM_QUERY),
+    ("Acronym (BM25)", ACRONYM_QUERY),
+    ("Specific molecule (BM25)", RARE_TERM_QUERY),
     ("Broad concept (Hybrid)", BROAD_CONCEPT_QUERY),
 ]
 
 PLOT_CAVEAT = (
-    "⚠️ **The 2D projection distorts true distances.** This PCA view captures only a "
-    "sliver of the 384-dimensional space, so retrieved points may *not* be the "
-    "visually-closest dots. **The ranked lists above are authoritative** — the plot is "
-    "only for spatial intuition."
+    "⚠️ **The 2D UMAP projection distorts true distances.** UMAP preserves rough local "
+    "neighbourhoods but warps global distances, and the live query is placed by an "
+    "*approximate* out-of-sample fit — so its position (and the connector lines) are only "
+    "indicative. Retrieved points may *not* be the visually-closest dots. "
+    "**The ranked lists above are authoritative**; the plot is only for intuition."
 )
 
 # ====================================================================================
@@ -294,12 +318,13 @@ with gr.Blocks(title="RAG retrieval: BM25 vs Dense vs Hybrid") as demo:
     search_btn = gr.Button("Search", variant="primary")
     filter_info = gr.Markdown("—")
 
+    gr.Markdown("_Click any result to expand its full abstract._")
     with gr.Row():
-        bm25_out = gr.Markdown(label="BM25")
-        dense_out = gr.Markdown(label="Dense")
-        hybrid_out = gr.Markdown(label="Hybrid")
+        bm25_out = gr.HTML(label="BM25")
+        dense_out = gr.HTML(label="Dense")
+        hybrid_out = gr.HTML(label="Hybrid")
 
-    gr.Markdown("## Vector space (PCA projection)")
+    gr.Markdown("## Vector space (UMAP projection)")
     plot = gr.Plot(value=make_plot())
     gr.Markdown(PLOT_CAVEAT)
 

build_index.py

@@ -15,8 +15,8 @@ import joblib
 import numpy as np
 import pandas as pd
 import requests
+import umap
 from sentence_transformers import SentenceTransformer
-from sklearn.decomposition import PCA
 
 import config
 from text_utils import tokenize
@@ -106,13 +106,18 @@ def main():
     np.save(os.path.join(config.DATA_DIR, "embeddings.npy"), embeddings)
     print(f"Saved embeddings: {embeddings.shape}")
 
-    # 3. PCA ------------------------------------------------------------------------
-    print("Fitting 2D PCA...")
-    pca = PCA(n_components=2, random_state=0)
-    coords = pca.fit_transform(embeddings).astype(np.float32)
-    joblib.dump(pca, os.path.join(config.DATA_DIR, "pca.joblib"))
-    np.save(os.path.join(config.DATA_DIR, "pca_coords.npy"), coords)
-    print(f"Saved PCA + corpus coords: {coords.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 ----------------------------------------------------------------
     print("Tokenising for BM25...")

data/{pca.joblib → umap.joblib}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:f13f7a85bd6fe24bf75709961bcdefb4e2ca4587e3e377cf67f29a7b765baefd
-size 5583
+oid sha256:2690f51123d5179bf011fad8d62e5ee5856f31fb0ed78f2a1038b2d4d356b24a
+size 6914713

data/{pca_coords.npy → umap_coords.npy}

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

requirements.txt

@@ -3,6 +3,7 @@ spaces
 sentence-transformers
 rank-bm25
 scikit-learn
+umap-learn
 numpy
 pandas
 pyarrow