Spaces:
Sleeping
Sleeping
File size: 4,193 Bytes
005e9fd | 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 | """Measure the retriever against the generated dataset.
Scores every retrieval mode on hit rate and MRR.
Free to run, because embedding is local and no API calls are made, so we can
re-run it after every retriever change.
uv run python -m eval.run_eval
uv run python -m eval.run_eval --retriever hybrid --retriever bm25
"""
import argparse
import json
import statistics
import sys
from llama_index.core.evaluation import HitRate, MRR
from rag.config import CHUNKS_PATH, EVAL_DATASET_PATH, RETRIEVAL_MODE, RETRIEVAL_TOP_K
from rag.index import RETRIEVAL_MODES, load_retriever, make_embed_model
from rag.types import EvalRow
# What counts as correct: only the chunk the question came from, or any chunk of
# the same section. We score both, because a neighbouring chunk often answers the
# question just as well and a citation points at the section anyway.
LEVELS = ("strict", "source")
METRICS = tuple(f"{level}_{metric}" for level in LEVELS for metric in ("found", "rank"))
# One score per question, per metric, averaged once every question is scored.
Scores = dict[str, list[float]]
HIT_RATE, RECIPROCAL_RANK = HitRate(), MRR()
def rank_all(retriever, rows: list[EvalRow]) -> list[list[str]]:
return [[node.node_id for node in retriever.retrieve(row["question"])] for row in rows]
def same_source_ids() -> dict[str, list[str]]:
by_url: dict[str, list[str]] = {}
ids: dict[str, str] = {}
with CHUNKS_PATH.open(encoding="utf-8") as handle:
for line in handle:
chunk = json.loads(line)
url = chunk["metadata"]["url"]
ids[chunk["id"]] = url
by_url.setdefault(url, []).append(chunk["id"])
return {chunk_id: by_url[url] for chunk_id, url in ids.items()}
def score_questions(
rows: list[EvalRow], rankings: list[list[str]], expand: dict[str, list[str]]
) -> Scores:
scores: Scores = {metric: [] for metric in METRICS}
for row, ranked in zip(rows, rankings):
retrieved = ranked[:RETRIEVAL_TOP_K]
expected = {"strict": [row["chunk_id"]], "source": expand[row["chunk_id"]]}
for level in LEVELS:
ids = expected[level]
scores[f"{level}_found"].append(
HIT_RATE.compute(expected_ids=ids, retrieved_ids=retrieved).score
)
scores[f"{level}_rank"].append(
RECIPROCAL_RANK.compute(expected_ids=ids, retrieved_ids=retrieved).score
)
return scores
def report(scores: dict[str, Scores]) -> None:
print(f"\n{'':12}{'exact':>18}{'section':>20}")
print(f"{'':12}{'hit_rate':>9}{'mrr':>9}{'hit_rate':>11}{'mrr':>9}")
for mode, metrics in scores.items():
marker = "*" if mode == RETRIEVAL_MODE else " "
cells = "".join(
f"{statistics.fmean(metrics[m]):>9.2f}" + (" " if m == "strict_rank" else "")
for m in METRICS
)
print(f"{marker} {mode:10}{cells}")
print(f"\n* scored on the top {RETRIEVAL_TOP_K} results")
def main(argv: list[str]) -> int:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument(
"--retriever", action="append", choices=RETRIEVAL_MODES,
help="modes to measure; repeatable, defaults to all",
)
parser.add_argument("--rerank-model", help="override the cross-encoder for rerank mode")
args = parser.parse_args(argv[1:])
modes = args.retriever or list(RETRIEVAL_MODES)
if not EVAL_DATASET_PATH.exists():
print("No dataset. Run: uv run python -m eval.build_dataset", file=sys.stderr)
return 2
rows = json.loads(EVAL_DATASET_PATH.read_text(encoding="utf-8"))
# Same embedding model and query instruction the app uses
print(f"{len(rows)} questions | query instruction: {make_embed_model().query_instruction!r}")
expand = same_source_ids()
scores: dict[str, Scores] = {}
for mode in modes:
retriever = load_retriever(mode, RETRIEVAL_TOP_K, None, args.rerank_model)
scores[mode] = score_questions(rows, rank_all(retriever, rows), expand)
print(f" scored {mode}", flush=True)
report(scores)
return 0
if __name__ == "__main__":
sys.exit(main(sys.argv))
|