Add cross-encoder reranking (two-stage retrieval pipeline)

Stage 1 retrieves top_k*4 candidates via bi-encoder, Stage 2
reranks with ms-marco-MiniLM-L-6-v2 cross-encoder. Toggle in UI
lets users compare rankings with dual-score display.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Dockerfile

@@ -2,6 +2,7 @@ FROM python:3.11-slim
 WORKDIR /app
 COPY requirements.txt .
 RUN pip install --no-cache-dir -r requirements.txt
+RUN python -c "from sentence_transformers import CrossEncoder; CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')"
 COPY . .
 EXPOSE 7860
 CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

app.py

@@ -9,24 +9,28 @@ import os
 from fastapi import FastAPI, Query
 from fastapi.responses import FileResponse
 from fastapi.staticfiles import StaticFiles
-from sentence_transformers import SentenceTransformer
+from sentence_transformers import SentenceTransformer, CrossEncoder
 import chromadb
 
 MODEL_NAME = "intfloat/multilingual-e5-small"
+RERANKER_MODEL = "cross-encoder/ms-marco-MiniLM-L-6-v2"
 CHROMA_PATH = os.path.join(os.path.dirname(__file__), "data", "chroma_db")
 COLLECTION_NAME = "scifact"
 
 app = FastAPI(title="SciFact Multilingual Semantic Search")
 
 model: SentenceTransformer = None
+reranker: CrossEncoder = None
 collection: chromadb.Collection = None
 
 
 @app.on_event("startup")
 def startup():
-    global model, collection
+    global model, reranker, collection
     print(f"Loading model: {MODEL_NAME}")
     model = SentenceTransformer(MODEL_NAME)
+    print(f"Loading reranker: {RERANKER_MODEL}")
+    reranker = CrossEncoder(RERANKER_MODEL)
     print(f"Loading ChromaDB from: {CHROMA_PATH}")
     client = chromadb.PersistentClient(path=CHROMA_PATH)
     collection = client.get_collection(COLLECTION_NAME)
@@ -41,36 +45,62 @@ def index():
 
 
 @app.get("/search")
-def search(q: str = Query(..., min_length=1), top_k: int = Query(5, ge=1, le=20)):
+def search(
+    q: str = Query(..., min_length=1),
+    top_k: int = Query(5, ge=1, le=20),
+    rerank: bool = Query(True),
+):
+    raw_query = q.strip()
+
+    # Stage 1: bi-encoder retrieval
+    n_candidates = min(top_k * 4, 40) if rerank else top_k
     query_embedding = model.encode(
-        [f"query: {q.strip()}"],
+        [f"query: {raw_query}"],
         normalize_embeddings=True,
     ).tolist()
 
     results = collection.query(
         query_embeddings=query_embedding,
-        n_results=top_k,
+        n_results=n_candidates,
         include=["metadatas", "distances", "documents"],
     )
 
-    items = []
-    for i, (meta, dist, doc) in enumerate(
-        zip(
-            results["metadatas"][0],
-            results["distances"][0],
-            results["documents"][0],
-        )
+    candidates = []
+    for meta, dist, doc in zip(
+        results["metadatas"][0],
+        results["distances"][0],
+        results["documents"][0],
     ):
-        items.append(
+        candidates.append(
             {
-                "rank": i + 1,
-                "score": round(1 - dist, 4),
+                "bi_score": round(1 - dist, 4),
                 "title": meta.get("title", ""),
-                "text": doc[:300] if doc else meta.get("text", "")[:300],
+                "full_text": doc if doc else meta.get("text", ""),
             }
         )
 
-    return {"query": q, "results": items}
+    # Stage 2: cross-encoder reranking
+    if rerank and candidates:
+        pairs = [(raw_query, c["full_text"]) for c in candidates]
+        ce_scores = reranker.predict(pairs).tolist()
+        for c, score in zip(candidates, ce_scores):
+            c["ce_score"] = round(score, 4)
+        candidates.sort(key=lambda c: c["ce_score"], reverse=True)
+
+    items = []
+    for i, c in enumerate(candidates[:top_k]):
+        item = {
+            "rank": i + 1,
+            "bi_score": c["bi_score"],
+            "title": c["title"],
+            "text": c["full_text"][:300],
+            "reranked": rerank,
+        }
+        if rerank:
+            item["ce_score"] = c.get("ce_score")
+        items.append(item)
+
+    return {"query": raw_query, "results": items}
 
 
 app.mount("/static", StaticFiles(directory=os.path.join(os.path.dirname(__file__), "static")), name="static")

static/index.html

@@ -162,6 +162,77 @@
   .search-row button:hover { background: #C42A20; }
   .search-row button:disabled { background: var(--gray-2); color: var(--gray-3); cursor: not-allowed; }
 
+  /* ── Rerank Toggle ── */
+  .toggle-row {
+    display: flex;
+    align-items: center;
+    gap: 12px;
+    margin-top: 14px;
+  }
+  .toggle-label {
+    font-family: 'JetBrains Mono', monospace;
+    font-size: .7rem;
+    color: var(--gray-4);
+    letter-spacing: .5px;
+    text-transform: uppercase;
+  }
+  .toggle-track {
+    position: relative;
+    width: 42px;
+    height: 22px;
+    border: 2px solid var(--white);
+    background: transparent;
+    cursor: pointer;
+    transition: all .2s;
+    flex-shrink: 0;
+  }
+  .toggle-track.active {
+    border-color: var(--red);
+    background: var(--red);
+  }
+  .toggle-thumb {
+    position: absolute;
+    top: 2px;
+    left: 2px;
+    width: 14px;
+    height: 14px;
+    background: var(--white);
+    transition: transform .2s;
+  }
+  .toggle-track.active .toggle-thumb {
+    transform: translateX(20px);
+  }
+  .toggle-status {
+    font-family: 'JetBrains Mono', monospace;
+    font-size: .62rem;
+    color: var(--gray-3);
+    letter-spacing: .5px;
+  }
+
+  /* ── CE Score / Badge ── */
+  .ce-row {
+    display: flex;
+    align-items: center;
+    gap: 10px;
+    margin-bottom: 4px;
+  }
+  .ce-badge {
+    font-family: 'JetBrains Mono', monospace;
+    font-size: .58rem;
+    letter-spacing: 1px;
+    color: var(--black);
+    background: var(--red);
+    padding: 2px 7px;
+    font-weight: 600;
+  }
+  .ce-val {
+    font-family: 'JetBrains Mono', monospace;
+    font-size: .72rem;
+    color: var(--white);
+    font-weight: 600;
+    letter-spacing: .5px;
+  }
+
   /* ── Chips ── */
   .chips {
     display: flex;
@@ -446,6 +517,13 @@
       <input id="q" type="text" placeholder="Enter a query in any language..." autocomplete="off" spellcheck="false">
       <button id="btn" onclick="doSearch()">SEARCH</button>
     </div>
+    <div class="toggle-row">
+      <div class="toggle-track active" id="rerankToggle" onclick="toggleRerank()">
+        <div class="toggle-thumb"></div>
+      </div>
+      <span class="toggle-label">Cross-encoder reranking</span>
+      <span class="toggle-status" id="rerankStatus">ON — ms-marco-MiniLM-L-6-v2</span>
+    </div>
     <div class="chips">
       <div class="chip" onclick="chipSearch(this)"><span class="lang">EN</span>effects of vaccination</div>
       <div class="chip" onclick="chipSearch(this)"><span class="lang">FR</span>effets de la vaccination</div>
@@ -469,7 +547,7 @@
 
   <!-- Footer -->
   <div class="footer">
-    <span>intfloat/multilingual-e5-small</span>
+    <span>bi-encoder: multilingual-e5-small<span class="dot">/</span>cross-encoder: ms-marco-MiniLM-L-6-v2</span>
     <span>chromadb<span class="dot">/</span>fastapi<span class="dot">/</span>cpu</span>
   </div>
 
@@ -479,9 +557,20 @@
 const input = document.getElementById('q');
 const btn = document.getElementById('btn');
 const results = document.getElementById('results');
+let rerankEnabled = true;
 
 input.addEventListener('keydown', e => { if (e.key === 'Enter') doSearch(); });
 
+function toggleRerank() {
+  rerankEnabled = !rerankEnabled;
+  const toggle = document.getElementById('rerankToggle');
+  const status = document.getElementById('rerankStatus');
+  toggle.classList.toggle('active', rerankEnabled);
+  status.textContent = rerankEnabled ? 'ON — ms-marco-MiniLM-L-6-v2' : 'OFF — bi-encoder only';
+  // Re-run search if there are existing results
+  if (input.value.trim()) doSearch();
+}
+
 function chipSearch(el) {
   const text = el.textContent.replace(/^[A-Z]{2}/, '').trim();
   input.value = text;
@@ -497,7 +586,7 @@ async function doSearch() {
   results.innerHTML = skeletonHTML();
 
   try {
-    const res = await fetch('/search?q=' + encodeURIComponent(q) + '&top_k=5');
+    const res = await fetch('/search?q=' + encodeURIComponent(q) + '&top_k=5&rerank=' + rerankEnabled);
     if (!res.ok) throw new Error(res.statusText);
     const data = await res.json();
     renderResults(data.results);
@@ -526,11 +615,20 @@ function renderResults(items) {
     return;
   }
   results.innerHTML = items.map((r, i) => {
-    const pct = Math.max(0, Math.min(100, r.score * 100));
+    const pct = Math.max(0, Math.min(100, r.bi_score * 100));
+    const rankStr = r.rank < 10 ? '0' + r.rank : '' + r.rank;
+    let ceHTML = '';
+    if (r.reranked && r.ce_score != null) {
+      ceHTML = '<div class="ce-row">' +
+        '<span class="ce-badge">RERANKED</span>' +
+        '<span class="ce-val">CE ' + r.ce_score.toFixed(2) + '</span>' +
+      '</div>';
+    }
     return '<div class="result-card" style="animation-delay:' + (i * .08) + 's">' +
-      '<div class="rank-col"><div class="rank-num">0' + r.rank + '</div><div class="rank-line"></div></div>' +
+      '<div class="rank-col"><div class="rank-num">' + rankStr + '</div><div class="rank-line"></div></div>' +
       '<div class="result-body">' +
         '<div class="result-title">' + esc(r.title) + '</div>' +
+        ceHTML +
         '<div class="score-row">' +
           '<div class="score-track"><div class="score-fill" style="width:' + pct + '%"></div></div>' +
           '<div class="score-val">' + pct.toFixed(1) + '%</div>' +