Blame for evaluate.py

"""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


# ── Metrics ───────────────────────────────────────────────────────────────────

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


# ── Build: candidates + LLM relevance judgments ───────────────────────────────

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

        # Pool: union of every system's top results; judge only papers not yet labeled.
        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")  # save after every query

    print(f"\nSaved {len(data['queries'])} queries, {len(data['papers'])} papers -> {DATASET}")


# ── Run: metrics per system ───────────────────────────────────────────────────

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")  # Windows consoles default to cp1252

    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)