update ui

requirements.txt

@@ -3,3 +3,4 @@ rank_bm25
 sentence-transformers
 pandas
 torch
+einops

reranker.py

@@ -1,20 +1,19 @@
-from concurrent.futures import ThreadPoolExecutor
 from sentence_transformers import CrossEncoder
 
 
-DEFAULT_MODEL = "jinaai/jina-reranker-v2-base-multilingual"
+DEFAULT_MODEL = "BAAI/bge-reranker-v2-m3"
 
-ENSEMBLE_MODELS = [
-    "jinaai/jina-reranker-v2-base-multilingual",
-    "Alibaba-NLP/gte-multilingual-reranker-base",
+AVAILABLE_MODELS = [
     "BAAI/bge-reranker-v2-m3",
+    "cross-encoder/ms-marco-MiniLM-L-6-v2",
+    "cross-encoder/ms-marco-electra-base",
 ]
 
 
 def load_model(model_name: str = DEFAULT_MODEL) -> CrossEncoder:
     """Load a CrossEncoder model from HuggingFace (cached after first download)."""
     print(f"Loading reranker model: {model_name} ...")
-    model = CrossEncoder(model_name, max_length=512)
+    model = CrossEncoder(model_name, max_length=512, trust_remote_code=True)
     print("Model loaded.\n")
     return model
 
@@ -37,15 +36,8 @@ def build_pairs(query: str, papers: list[dict]) -> list[tuple[str, str]]:
 def rerank(query: str, papers: list[dict], model: CrossEncoder, batch_size: int = 16) -> list[dict]:
     """Score every (query, paper) pair with the CrossEncoder and return papers sorted by score.
 
-    Args:
-        query:      The original search query string.
-        papers:     List of paper dicts (already BM25-ranked; order does not matter here).
-        model:      A loaded CrossEncoder instance.
-        batch_size: Number of pairs to score in one forward pass.
-
-    Returns:
-        The same paper dicts, each annotated with `ce_score` and `ce_rank`,
-        sorted by CrossEncoder score descending.
+    Returns the same paper dicts, each annotated with `ce_score` and `ce_rank`,
+    sorted by CrossEncoder score descending.
     """
     pairs = build_pairs(query, papers)
     scores = model.predict(pairs, batch_size=batch_size, show_progress_bar=True)
@@ -64,64 +56,6 @@ def rerank(query: str, papers: list[dict], model: CrossEncoder, batch_size: int
     return results
 
 
-def load_models(model_names: list[str] = ENSEMBLE_MODELS) -> list[CrossEncoder]:
-    """Load all CrossEncoder models in parallel using ThreadPoolExecutor."""
-    with ThreadPoolExecutor(max_workers=len(model_names)) as executor:
-        futures = [executor.submit(load_model, name) for name in model_names]
-        return [f.result() for f in futures]
-
-
-def _normalize(scores: list[float]) -> list[float]:
-    """Min-max normalize scores to [0, 1]. Returns all 1.0 when all scores are equal."""
-    lo, hi = min(scores), max(scores)
-    if hi == lo:
-        return [1.0] * len(scores)
-    return [(s - lo) / (hi - lo) for s in scores]
-
-
-def ensemble_rerank(
-    query: str,
-    papers: list[dict],
-    models: list[CrossEncoder],
-    batch_size: int = 16,
-) -> list[dict]:
-    """Score papers with multiple CrossEncoders, normalize each model's scores to
-    [0, 1] via min-max, then average across models (equal weights) to produce an
-    ensemble score. Returns papers sorted by ensemble score descending.
-
-    Each returned paper dict gains:
-      ce_score_model1 / 2 / 3  — raw score from each model
-      ce_score_ensemble         — mean of the three normalized scores
-      ce_score                  — alias for ce_score_ensemble (pipeline compat)
-      ce_rank                   — 1-based position in the ensemble ranking
-    """
-    pairs = build_pairs(query, papers)
-
-    with ThreadPoolExecutor(max_workers=len(models)) as executor:
-        futures = [executor.submit(m.predict, pairs) for m in models]
-        raw_scores = [list(f.result()) for f in futures]
-    norm_scores = [_normalize(s) for s in raw_scores]
-
-    n_models = len(models)
-    ensemble = [
-        sum(norm_scores[m][i] for m in range(n_models)) / n_models
-        for i in range(len(papers))
-    ]
-
-    order = sorted(range(len(papers)), key=lambda i: ensemble[i], reverse=True)
-
-    results = []
-    for rank, i in enumerate(order, start=1):
-        paper = dict(papers[i])
-        for m_idx, raw in enumerate(raw_scores, start=1):
-            paper[f"ce_score_model{m_idx}"] = round(float(raw[i]), 4)
-        paper["ce_score_ensemble"] = round(float(ensemble[i]), 4)
-        paper["ce_score"] = paper["ce_score_ensemble"]
-        paper["ce_rank"] = rank
-        results.append(paper)
-    return results
-
-
 if __name__ == "__main__":
     import argparse
     from retriever import fetch_papers

ui/app.py

@@ -39,8 +39,10 @@ def tier_badge(score: float) -> str:
 
 # ── Sidebar ────────────────────────────────────────────────────────────────────
 with st.sidebar:
-    st.title("PaperRank")
-    st.caption("CrossEncoder reranking on live academic papers.")
+    # st.title("PaperRank")
+    # st.caption("CrossEncoder reranking on live academic papers.")
+    st.title("Task1")
+    st.caption("Retrieval and Reranking")
     st.divider()
 
     query = st.text_input("Search query", placeholder="e.g. transformer models for information retrieval")
@@ -115,9 +117,10 @@ if (
 
 
 # ── Main area ──────────────────────────────────────────────────────────────────
-st.title("PaperRank")
-st.caption("CrossEncoder reranking on live academic papers.")
-
+# st.title("PaperRank")
+# st.caption("CrossEncoder reranking on live academic papers.")
+st.title("Task1")
+st.caption("Retrieval and Reranking")
 if "reranked_papers" not in st.session_state:
     st.info("Enter a query in the sidebar and click **Run pipeline** to get started.")
     st.stop()
@@ -177,7 +180,7 @@ def render_ce_card(paper: dict) -> None:
     score = paper.get("ce_score", 0.0)
 
     oa_r = oa_rank_by_id.get(pid)
-    note = (f" {'↑' if rank < oa_r else '↓'} from OA #{oa_r}") if oa_r and oa_r != rank else ""
+    note = title
 
     st.markdown(
         f"**#{rank}**{note} {tier_badge(score)}  \n{title_md}  \n"