File size: 4,173 Bytes
4d9eec9 0ccd8ab 483f7b5 0ccd8ab 2e38dc6 4d9eec9 df1df12 0ccd8ab 4d9eec9 0ccd8ab 4d9eec9 0ccd8ab 4d9eec9 0ccd8ab 4d9eec9 0ccd8ab 4d9eec9 0ccd8ab 4d9eec9 0ccd8ab 4d9eec9 df1df12 0ccd8ab 4d9eec9 df1df12 0ccd8ab 4d9eec9 df1df12 0ccd8ab df1df12 4d9eec9 0ccd8ab 4d9eec9 0ccd8ab | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 | import torch
from sentence_transformers import CrossEncoder
DEFAULT_MODEL = "cross-encoder/ms-marco-MiniLM-L-6-v2"
AVAILABLE_MODELS = [
"cross-encoder/ms-marco-MiniLM-L-6-v2",
"BAAI/bge-reranker-v2-m3",
"cross-encoder/ms-marco-electra-base",
]
# (label, minimum score). Checked top-down; first match wins.
TIERS = [
("Perfectly Relevant", 0.90),
("Relevant", 0.70),
("Somewhat Relevant", 0.50),
("Not Relevant", float("-inf")),
]
def load_model(model_name: str = DEFAULT_MODEL) -> CrossEncoder:
"""Load a CrossEncoder model from HuggingFace (cached on disk after first download).
Sigmoid is forced so every model returns scores in [0, 1]. Without it,
the ms-marco models return raw logits (their saved config uses Identity)
while BGE returns probabilities, and the tier thresholds stop meaning anything.
"""
if model_name not in AVAILABLE_MODELS:
raise ValueError(f"Unknown model {model_name!r}. Choose from: {AVAILABLE_MODELS}")
return CrossEncoder(model_name, max_length=512, activation_fn=torch.nn.Sigmoid())
def relevance_tier(score: float) -> str:
"""Map a [0, 1] CrossEncoder score to a relevance tier label."""
for label, threshold in TIERS:
if score >= threshold:
return label
return TIERS[-1][0]
def build_pairs(query: str, papers: list[dict]) -> list[tuple[str, str]]:
"""Build (query, document) pairs for CrossEncoder scoring.
The document is title + abstract concatenated — both fields carry
relevance signal and the CrossEncoder attends over the full pair.
"""
pairs = []
for paper in papers:
title = paper.get("title") or ""
abstract = paper.get("abstract") or ""
document = f"{title}. {abstract}".strip()
pairs.append((query, document))
return pairs
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.
Returns copies of the paper dicts, each annotated with `ce_score`,
`ce_rank` and `tier`, sorted by CrossEncoder score descending.
"""
if not papers:
return []
pairs = build_pairs(query, papers)
scores = model.predict(pairs, batch_size=batch_size, show_progress_bar=False)
ranked = sorted(
zip(papers, scores),
key=lambda x: x[1],
reverse=True,
)
results = []
for rank, (paper, score) in enumerate(ranked, start=1):
paper = dict(paper)
paper["ce_score"] = round(float(score), 4)
paper["ce_rank"] = rank
paper["tier"] = relevance_tier(paper["ce_score"])
results.append(paper)
return results
if __name__ == "__main__":
import argparse
import sys
from retriever import fetch_papers
sys.stdout.reconfigure(encoding="utf-8") # Windows consoles default to cp1252
parser = argparse.ArgumentParser(description="Rerank OpenAlex candidates with a CrossEncoder")
parser.add_argument("query", type=str, help="Search query")
parser.add_argument("--model", type=str, default=DEFAULT_MODEL, choices=AVAILABLE_MODELS)
parser.add_argument("--limit", type=int, default=50, help="Number of OpenAlex candidates to fetch")
parser.add_argument("--top", type=int, default=10, help="Number of reranked results to display")
args = parser.parse_args()
papers = fetch_papers(args.query, limit=args.limit)
print(f"Loading {args.model} ...")
model = load_model(args.model)
reranked = rerank(args.query, papers, model)
print(f"\nQuery: {args.query!r}")
print(f"Showing top {args.top} after CrossEncoder reranking:\n")
for paper in reranked[:args.top]:
authors = ", ".join(a["name"] for a in paper.get("authors", [])[:2])
print(f"[CE #{paper['ce_rank']}] {paper['ce_score']:.3f} {paper['tier']} (was OpenAlex #{paper['openalex_rank']})")
print(f" Title : {paper['title']}")
print(f" Year : {paper.get('year', 'N/A')} | Citations: {paper.get('citationCount', 'N/A')}")
print(f" Authors: {authors}")
print()
|