| """Offline evaluation: do multi-query search and reranking actually help? |
| |
| Two steps, so OpenAlex/LLM costs are paid once and runs are reproducible: |
| |
| python evaluate.py build # fetch candidates + judge relevance -> eval/dataset.json |
| python evaluate.py run # rank the saved candidates every way, report metrics |
| |
| Relevance labels come from an LLM judge (0 = not, 1 = partly, 2 = highly |
| relevant) over a pool of every system's top results (TREC-style pooling). |
| They are a starting point: correct them by hand in eval/dataset.json. |
| """ |
| import json |
| import math |
| from pathlib import Path |
|
|
| from analyzer import MAX_ABSTRACT_CHARS, LLMError, chat_json, decompose |
| from reranker import DEFAULT_MODEL, load_model, rerank |
| from retriever import fetch_papers, multi_search |
|
|
|
|
| DATASET = Path("eval/dataset.json") |
| QUERIES = Path("eval/queries.txt") |
| JUDGE_BATCH = 5 |
|
|
|
|
| |
|
|
| def dcg_at_k(relevances: list[float], k: int) -> float: |
| """Discounted Cumulative Gain at rank k.""" |
| return sum(rel / math.log2(i + 1) for i, rel in enumerate(relevances[:k], start=1)) |
|
|
|
|
| def ndcg_at_k(ranked_relevances: list[float], all_relevances: list[float], k: int) -> float: |
| """NDCG at rank k, normalized by the best possible top-k over ALL judged papers. |
| |
| Normalizing by every judged paper (not just the ones this ranking returned) |
| means a system that fails to retrieve relevant papers is penalized for it. |
| """ |
| idcg = dcg_at_k(sorted(all_relevances, reverse=True), k) |
| return dcg_at_k(ranked_relevances, k) / idcg if idcg else 0.0 |
|
|
|
|
| def mrr(ranked_relevances: list[float], threshold: float = 2) -> float: |
| """Reciprocal rank of the first paper with relevance >= threshold (default: highly relevant).""" |
| for rank, rel in enumerate(ranked_relevances, start=1): |
| if rel >= threshold: |
| return 1.0 / rank |
| return 0.0 |
|
|
|
|
| |
|
|
| def judge(question: str, papers: list[dict]) -> dict[str, int]: |
| """Grade each paper 0/1/2 for the question. Returns {paperId: grade}.""" |
| labels = {} |
| for i in range(0, len(papers), JUDGE_BATCH): |
| batch = papers[i:i + JUDGE_BATCH] |
| papers_text = "\n\n".join( |
| f"[{j}] Title: {p['title']}\nAbstract: {p['abstract'][:MAX_ABSTRACT_CHARS]}" |
| for j, p in enumerate(batch) |
| ) |
| prompt = f"""You are an expert assessor judging search results for a researcher. |
| |
| Research question: {question} |
| |
| Grade each paper's relevance using its title and abstract: |
| 2 = highly relevant: directly addresses the question; the researcher would definitely want it |
| 1 = partly relevant: related, addresses some parts of the question |
| 0 = not relevant |
| |
| {papers_text} |
| |
| Respond with JSON only: {{"grades": [{{"index": 0, "grade": 2}}]}}""" |
| data = chat_json(prompt, max_tokens=60 * len(batch)) |
| for g in data.get("grades", []): |
| j, grade = g.get("index"), g.get("grade") |
| if isinstance(j, int) and 0 <= j < len(batch) and grade in (0, 1, 2): |
| labels[batch[j]["paperId"]] = grade |
| return labels |
|
|
|
|
| def _system_rankings(question: str, entry: dict, papers: dict, models: dict) -> dict[str, list[str]]: |
| """Every system's ranking of the saved candidates, as lists of paperIds.""" |
| def docs(ids: list[str]) -> list[dict]: |
| return [{"paperId": pid, **papers[pid]} for pid in ids] |
|
|
| rankings = { |
| "OpenAlex (1 search)": entry["single"], |
| "Multi-search (merged)": entry["multi"], |
| } |
| for name, model in models.items(): |
| short = name.split("/")[-1] |
| rankings[f"1 search + {short}"] = [p["paperId"] for p in rerank(question, docs(entry["single"]), model)] |
| rankings[f"Multi + {short}"] = [p["paperId"] for p in rerank(question, docs(entry["multi"]), model)] |
| return rankings |
|
|
|
|
| def build(model_names: list[str], pool_depth: int, refresh: bool) -> None: |
| data = json.loads(DATASET.read_text("utf-8")) if DATASET.exists() else {"queries": [], "papers": {}} |
| done = {q["question"]: q for q in data["queries"]} |
| questions = [q.strip() for q in QUERIES.read_text("utf-8").splitlines() if q.strip() and not q.startswith("#")] |
| models = {name: load_model(name) for name in model_names} |
|
|
| for n, question in enumerate(questions, 1): |
| entry = done.get(question) |
| if entry is None or refresh: |
| print(f"[{n}/{len(questions)}] fetching: {question}") |
| plan = decompose(question) |
| single = fetch_papers(question) |
| multi = multi_search(plan["sub_queries"]) |
| for p in single + multi: |
| data["papers"][p["paperId"]] = {"title": p["title"], "abstract": p["abstract"]} |
| entry = { |
| "question": question, |
| "sub_queries": plan["sub_queries"], |
| "single": [p["paperId"] for p in single], |
| "multi": [p["paperId"] for p in multi], |
| "labels": (entry or {}).get("labels", {}), |
| } |
| if question not in done: |
| data["queries"].append(entry) |
| else: |
| data["queries"][data["queries"].index(done[question])] = entry |
| done[question] = entry |
|
|
| |
| rankings = _system_rankings(question, entry, data["papers"], models) |
| pool = {pid for ranking in rankings.values() for pid in ranking[:pool_depth]} |
| todo = [{"paperId": pid, **data["papers"][pid]} for pid in sorted(pool) if pid not in entry["labels"]] |
| if todo: |
| print(f" judging {len(todo)} papers ...") |
| try: |
| entry["labels"].update(judge(question, todo)) |
| except LLMError as e: |
| print(f" judge failed ({e}); rerun build to resume") |
| DATASET.parent.mkdir(exist_ok=True) |
| DATASET.write_text(json.dumps(data, indent=1, ensure_ascii=False), "utf-8") |
|
|
| print(f"\nSaved {len(data['queries'])} queries, {len(data['papers'])} papers -> {DATASET}") |
|
|
|
|
| |
|
|
| def run(model_names: list[str], k: int, verbose: bool) -> None: |
| data = json.loads(DATASET.read_text("utf-8")) |
| models = {name: load_model(name) for name in model_names} |
| totals: dict[str, dict[str, float]] = {} |
|
|
| for entry in data["queries"]: |
| labels = entry["labels"] |
| if not labels: |
| continue |
| all_rels = list(labels.values()) |
| rankings = _system_rankings(entry["question"], entry, data["papers"], models) |
| if verbose: |
| print(f"\n{entry['question']}") |
| for system, ranking in rankings.items(): |
| rels = [labels.get(pid, 0) for pid in ranking] |
| scores = { |
| f"NDCG@{k}": ndcg_at_k(rels, all_rels, k), |
| "MRR": mrr(rels), |
| f"Judged@{k}": sum(pid in labels for pid in ranking[:k]) / k, |
| } |
| for metric, value in scores.items(): |
| totals.setdefault(system, {}).setdefault(metric, 0.0) |
| totals[system][metric] += value |
| if verbose: |
| print(f" {system:<34}" + "".join(f"{v:>10.3f}" for v in scores.values())) |
|
|
| n = sum(1 for q in data["queries"] if q["labels"]) |
| metrics = list(next(iter(totals.values()))) |
| print(f"\nMean over {n} queries (labels: LLM judge, 0/1/2; MRR = first highly relevant paper)\n") |
| print(f" {'System':<34}" + "".join(f"{m:>10}" for m in metrics)) |
| print(" " + "-" * (34 + 10 * len(metrics))) |
| for system, scores in totals.items(): |
| print(f" {system:<34}" + "".join(f"{scores[m] / n:>10.3f}" for m in metrics)) |
| print(f"\n Judged@{k} < 1.000 means some top results were never labeled (unlabeled count as 0);") |
| print(" rerun build with the same --models to judge them.") |
|
|
|
|
| if __name__ == "__main__": |
| import argparse |
| import sys |
|
|
| from dotenv import load_dotenv |
|
|
| load_dotenv() |
| sys.stdout.reconfigure(encoding="utf-8") |
|
|
| parser = argparse.ArgumentParser(description="Evaluate retrieval + reranking against judged queries") |
| parser.add_argument("command", choices=["build", "run"]) |
| parser.add_argument("--models", nargs="+", default=[DEFAULT_MODEL], help="Reranker models to compare") |
| parser.add_argument("--k", type=int, default=10, help="Metric cutoff") |
| parser.add_argument("--pool-depth", type=int, default=20, help="build: judge each system's top N") |
| parser.add_argument("--refresh", action="store_true", help="build: re-fetch candidates for every query") |
| parser.add_argument("--verbose", action="store_true", help="run: show per-query scores") |
| args = parser.parse_args() |
|
|
| if args.command == "build": |
| build(args.models, args.pool_depth, args.refresh) |
| else: |
| run(args.models, args.k, args.verbose) |
|
|
