103 lines
3.8 KiB
Python
Executable file
103 lines
3.8 KiB
Python
Executable file
#!/usr/bin/env python
|
|
"""Build a T²-RAGBench leaderboard submission file (JSONL) for one subset.
|
|
|
|
Joins QA predictions with retrieval rankings by question text and emits one
|
|
object per question: ``{id, subset, context_id, prediction}`` — the format the
|
|
leaderboard expects (NM is scored on ``prediction``, MRR@3 on ``context_id``).
|
|
See https://t2ragbench.demo.hcds.uni-hamburg.de/submission.html.
|
|
|
|
Predictions come from a QA-run CSV (the merged master export: needs columns
|
|
``id``, ``question``, ``output``). Retrieval rankings come from a Logfire
|
|
retrieval trace (case spans store the ranked retrieved context ids in
|
|
``output`` and the question in ``inputs``) or, with --retrieval-csv, a CSV with
|
|
``question`` and a JSON-list ``output``.
|
|
|
|
Logfire access needs LOGFIRE_READ_TOKEN in the environment (EU project).
|
|
|
|
Example:
|
|
LOGFIRE_READ_TOKEN=... uv run python scripts/build_t2_submission.py \
|
|
--predictions ../t2_finqa_qwen3.6_019e982d.csv \
|
|
--retrieval-trace 019e9296e64d38b33f3592beff6660a9 \
|
|
--subset FinQA --topk 3 --out finqa_submission.jsonl
|
|
"""
|
|
|
|
import argparse
|
|
import csv
|
|
import datetime
|
|
import json
|
|
import os
|
|
import sys
|
|
|
|
from evaluations.submission import build_submission_rows
|
|
|
|
|
|
def _load_predictions(path: str) -> list[dict[str, str]]:
|
|
with open(path) as f:
|
|
return list(csv.DictReader(f))
|
|
|
|
|
|
def _retrieval_from_csv(path: str) -> dict[str, list[str]]:
|
|
out: dict[str, list[str]] = {}
|
|
with open(path) as f:
|
|
for row in csv.DictReader(f):
|
|
out[row["question"]] = json.loads(row["output"])
|
|
return out
|
|
|
|
|
|
def _retrieval_from_trace(trace_id: str) -> dict[str, list[str]]:
|
|
from logfire.experimental.query_client import LogfireQueryClient
|
|
|
|
token = os.environ.get("LOGFIRE_READ_TOKEN")
|
|
if not token:
|
|
sys.exit("LOGFIRE_READ_TOKEN not set (needed to read the retrieval trace).")
|
|
client = LogfireQueryClient(
|
|
read_token=token, base_url="https://logfire-eu.pydantic.dev"
|
|
)
|
|
rows = client.query_json_rows(
|
|
sql=(
|
|
"SELECT attributes->>'inputs' AS question, attributes->>'output' AS ranked "
|
|
f"FROM records WHERE trace_id = '{trace_id}' AND span_name LIKE 'case:%'"
|
|
),
|
|
min_timestamp=datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc),
|
|
limit=10000,
|
|
)["rows"]
|
|
return {r["question"]: json.loads(r["ranked"]) for r in rows if r.get("ranked")}
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--predictions", required=True, help="QA-run CSV.")
|
|
parser.add_argument("--retrieval-trace", help="Logfire retrieval trace id.")
|
|
parser.add_argument("--retrieval-csv", help="Retrieval CSV (question, output).")
|
|
parser.add_argument("--subset", required=True, help="Subset name, e.g. FinQA.")
|
|
parser.add_argument("--topk", type=int, default=3, help="Ranked context ids.")
|
|
parser.add_argument("--out", required=True, help="Output .jsonl path.")
|
|
args = parser.parse_args()
|
|
|
|
if bool(args.retrieval_trace) == bool(args.retrieval_csv):
|
|
sys.exit("Pass exactly one of --retrieval-trace or --retrieval-csv.")
|
|
|
|
predictions = _load_predictions(args.predictions)
|
|
retrieval = (
|
|
_retrieval_from_csv(args.retrieval_csv)
|
|
if args.retrieval_csv
|
|
else _retrieval_from_trace(args.retrieval_trace)
|
|
)
|
|
|
|
rows = build_submission_rows(
|
|
predictions, retrieval, subset=args.subset, topk=args.topk
|
|
)
|
|
with open(args.out, "w") as f:
|
|
for row in rows:
|
|
f.write(json.dumps(row) + "\n")
|
|
|
|
answered = sum(1 for r in rows if r["prediction"])
|
|
matched = sum(1 for r in rows if r["context_id"])
|
|
print(
|
|
f"wrote {args.out}: {len(rows)} rows "
|
|
f"({answered} with a prediction, {matched} with a retrieved context_id)"
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|