110 lines
4.7 KiB
Python
110 lines
4.7 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
"""Score captured retrieval rows against independent relevance judgments."""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import argparse
|
||
|
|
import json
|
||
|
|
import statistics
|
||
|
|
from pathlib import Path
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
|
||
|
|
def score(results_path: Path, judgments_path: Path) -> dict[str, Any]:
|
||
|
|
lines = [line for line in results_path.read_text(encoding="utf-8").splitlines() if line]
|
||
|
|
if len(lines) < 2:
|
||
|
|
raise ValueError("Retrieval results must contain a manifest and at least one query")
|
||
|
|
manifest = json.loads(lines[0])
|
||
|
|
rows = [json.loads(line) for line in lines[1:]]
|
||
|
|
question_ids = [row.get("question", {}).get("id") for row in rows]
|
||
|
|
if any(not isinstance(question_id, str) for question_id in question_ids):
|
||
|
|
raise ValueError("Every retrieval row must identify its question")
|
||
|
|
if len(question_ids) != len(set(question_ids)):
|
||
|
|
raise ValueError("Retrieval results contain duplicate question IDs")
|
||
|
|
if len(rows) != manifest.get("question_count"):
|
||
|
|
raise ValueError("Captured query count does not match the run manifest")
|
||
|
|
|
||
|
|
judgments = json.loads(judgments_path.read_text(encoding="utf-8"))
|
||
|
|
relevant_by_question = judgments.get("fully_answerable_positive_questions")
|
||
|
|
if not isinstance(relevant_by_question, dict) or not relevant_by_question:
|
||
|
|
raise ValueError("Independent judgments contain no answerable positive questions")
|
||
|
|
result_by_question = {row["question"]["id"]: row for row in rows}
|
||
|
|
missing = sorted(set(relevant_by_question) - result_by_question.keys())
|
||
|
|
if missing:
|
||
|
|
raise ValueError(f"Retrieval run omitted judged questions: {', '.join(missing)}")
|
||
|
|
|
||
|
|
scored: list[dict[str, Any]] = []
|
||
|
|
for question_id, relevant_ids in relevant_by_question.items():
|
||
|
|
row = result_by_question[question_id]
|
||
|
|
if row.get("is_error"):
|
||
|
|
returned_ids: list[str] = []
|
||
|
|
else:
|
||
|
|
returned_ids = row.get("returned_node_ids", [])
|
||
|
|
relevant = set(relevant_ids)
|
||
|
|
hits = relevant & set(returned_ids[:5])
|
||
|
|
scored.append(
|
||
|
|
{
|
||
|
|
"question_id": question_id,
|
||
|
|
"relevant_node_ids": sorted(relevant),
|
||
|
|
"returned_node_ids": returned_ids[:5],
|
||
|
|
"hit_node_ids": sorted(hits),
|
||
|
|
"recall_at_5": len(hits) / len(relevant),
|
||
|
|
"is_error": bool(row.get("is_error")),
|
||
|
|
}
|
||
|
|
)
|
||
|
|
|
||
|
|
revisions = {
|
||
|
|
page.get("canonical_revision")
|
||
|
|
for row in rows
|
||
|
|
for page in row.get("response", [])
|
||
|
|
if page.get("canonical_revision") is not None
|
||
|
|
}
|
||
|
|
if len(revisions) > 1:
|
||
|
|
raise ValueError("Retrieval responses came from more than one canonical revision")
|
||
|
|
hit_count = sum(bool(row["hit_node_ids"]) for row in scored)
|
||
|
|
query_times = [float(row["seconds"]) for row in rows]
|
||
|
|
return {
|
||
|
|
"schema_version": 1,
|
||
|
|
"questions_sha256": manifest.get("questions_sha256"),
|
||
|
|
"question_count": len(rows),
|
||
|
|
"judged_answerable_questions": len(scored),
|
||
|
|
"mean_conditional_recall_at_5": statistics.mean(row["recall_at_5"] for row in scored),
|
||
|
|
"conditional_hit_rate": hit_count / len(scored),
|
||
|
|
"queries_with_any_rows": sum(bool(row.get("returned_node_ids")) for row in rows),
|
||
|
|
"mean_observed_seconds": statistics.mean(query_times),
|
||
|
|
"median_observed_seconds": statistics.median(query_times),
|
||
|
|
"canonical_revision": next(iter(revisions), None),
|
||
|
|
"partial_ingestion_questions": sorted(judgments.get("partial_ingestion", [])),
|
||
|
|
"missing_accepted_material_questions": sorted(
|
||
|
|
judgments.get("missing_accepted_material", [])
|
||
|
|
),
|
||
|
|
"pagination_blocked_queries": sum(bool(row.get("pagination_blocked")) for row in rows),
|
||
|
|
"judgments_method": judgments.get("method", "not supplied"),
|
||
|
|
"per_question": scored,
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def main() -> None:
|
||
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
||
|
|
parser.add_argument("--results", type=Path, required=True)
|
||
|
|
parser.add_argument("--judgments", type=Path, required=True)
|
||
|
|
parser.add_argument("--output", type=Path, required=True)
|
||
|
|
args = parser.parse_args()
|
||
|
|
result = score(args.results, args.judgments)
|
||
|
|
with args.output.open("x", encoding="utf-8") as output:
|
||
|
|
output.write(json.dumps(result, ensure_ascii=False, indent=2, sort_keys=True) + "\n")
|
||
|
|
print(
|
||
|
|
json.dumps(
|
||
|
|
{
|
||
|
|
"questions": result["question_count"],
|
||
|
|
"judged_answerable": result["judged_answerable_questions"],
|
||
|
|
"conditional_recall_at_5": result["mean_conditional_recall_at_5"],
|
||
|
|
"conditional_hit_rate": result["conditional_hit_rate"],
|
||
|
|
"output": str(args.output),
|
||
|
|
}
|
||
|
|
)
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|