File size: 5,993 Bytes
20b15f3 | 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 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 | """
reranking.py
Reranks papers by semantic similarity to the Research Intent using SPECTER
(allenai-specter) — an embedding model trained specifically on academic
paper title/abstract pairs, rather than a generic sentence embedding model.
Install once:
pip install sentence-transformers
"""
import logging
from typing import Optional
from sentence_transformers import SentenceTransformer
import numpy as np
# Loaded once at module import time — NOT inside the function.
# Loading a transformer model from disk/HF hub takes a few seconds; if this
# were inside the function, every call would reload it, which is wasteful
# if this function gets called more than once in a session.
#
# device is pinned explicitly instead of left to SentenceTransformer's
# auto-detect. On HF ZeroGPU, torch reports a GPU at import time but only
# actually grants one inside an @spaces.GPU window — and this model is called
# from a LangGraph node, which is outside any such window. Auto-detect would
# therefore load it onto "cuda" and fail on first encode. Reranking ~30 abstracts
# is a couple of seconds on CPU.
_MODEL = SentenceTransformer("sentence-transformers/allenai-specter", device="cpu")
_fallback_logger = logging.getLogger(__name__)
def rerank_by_relevance(
research_intent: str,
papers: dict[str, str],
top_n: int = 10,
logger: Optional[logging.Logger] = None,
) -> dict[str, str]:
"""
Rerank papers by semantic similarity to the Research Intent.
Parameters
----------
research_intent : str
The full Research Intent text (Problem + Objective + Additional
Context, or however you've combined it) — used as the query vector.
papers : dict[str, str]
{normalized_title: abstract} — the aggregator's title/abstract pairs.
top_n : int
How many top-ranked papers to keep. Default 15.
logger : logging.Logger, optional
Node-scoped logger from the caller. Falls back to a module logger
when this function is used standalone.
Returns
-------
dict[str, str]
A NEW dict, same {normalized_title: abstract} shape, containing only
the top_n most relevant entries, ordered from most to least relevant.
Insertion order is preserved (Python 3.7+ dicts are ordered), so
iterating this dict gives you the ranking directly.
"""
log = logger or _fallback_logger
if not papers:
return {}
# --- Guard against empty/whitespace-only abstracts ---
# These can't be meaningfully embedded for relevance comparison. Rather
# than crash or silently mis-rank them, exclude them from ranking and
# log which ones were skipped so nothing disappears without a trace.
valid_titles = []
valid_abstracts = []
skipped_no_abstract = []
for normalized_title, abstract in papers.items():
if abstract and abstract.strip():
valid_titles.append(normalized_title)
valid_abstracts.append(abstract)
else:
skipped_no_abstract.append(normalized_title)
if skipped_no_abstract:
log.info("Skipped %d paper(s) with no abstract: %s", len(skipped_no_abstract), skipped_no_abstract)
if not valid_abstracts:
log.info("No papers had usable abstracts — returning empty result.")
return {}
# --- Embed the query (Research Intent) and all candidate abstracts ---
# normalize_embeddings=True means each vector has unit length, so a
# simple dot product between two vectors IS the cosine similarity —
# no separate cosine-similarity library call needed.
query_embedding = _MODEL.encode(
research_intent,
normalize_embeddings=True,
convert_to_numpy=True,
)
paper_embeddings = _MODEL.encode(
valid_abstracts,
normalize_embeddings=True,
convert_to_numpy=True,
batch_size=32,
show_progress_bar=False,
)
# --- Cosine similarity of every paper against the query, in one shot ---
similarities = paper_embeddings @ query_embedding # shape: (num_papers,)
# --- Sort by similarity, descending ---
ranked_indices = np.argsort(-similarities)
# --- Log the full ranking for visibility/debugging before truncating ---
ranking_lines = "\n".join(
f" {similarities[idx]:.4f} {valid_titles[idx]}" for idx in ranked_indices
)
log.info("Full relevance ranking (%d papers, title : similarity score):\n%s", len(valid_titles), ranking_lines)
top_indices = ranked_indices[:top_n]
dropped_count = len(valid_titles) - len(top_indices)
if dropped_count > 0:
log.info("Kept top %d, dropped %d lower-relevance paper(s).", len(top_indices), dropped_count)
# --- Build the result dict in ranked order ---
result = {
valid_titles[idx]: valid_abstracts[idx]
for idx in top_indices
}
return result
if __name__ == "__main__":
# Quick standalone test
test_intent = (
"Identify a robust methodology for comparing the fuel efficiency of "
"human-driven and reinforcement-learning-controlled vehicles in "
"car-following maneuvers, accounting for speed, acceleration, and headway."
)
test_papers = {
"ecofollower an environmentfriendly car following model": (
"This study introduces EcoFollower, a novel eco-car-following "
"model developed using reinforcement learning to optimize fuel "
"consumption in car-following scenarios."
),
"predicting fuel research octane number using spectra": (
"We show that an accurate statistical model for the Research "
"Octane Number of gasoline can be constructed using infrared "
"absorbance spectroscopy data."
),
}
reranked = rerank_by_relevance(test_intent, test_papers, top_n=15)
print("\nFinal reranked result:")
for title, abstract in reranked.items():
print(f"- {title}")
|