samwjya
Claude Opus 5.5
Reranker: consistent 0-1 scores via sigmoid, restore all three models
4d9eec9 | 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() | |