File size: 3,130 Bytes
bba5457 | 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 | #!/usr/bin/env python3
"""Build open-wikitable `unified/` shards from the unified eval bundle.
Mirrors monaco-benchmark-viewer's `unified/` layout: one JSON shard per qid plus
an `index.json` array used by the sidebar. The source is the canonical
unified-format file ``test_with_chunks.unified.jsonl`` (one row per question,
schema = {qid, question, answers, docs:[{id, contents}], ...]}).
Usage:
python scripts/build_unified.py
python scripts/build_unified.py --src /path/to/test_with_chunks.unified.jsonl --out unified
"""
from __future__ import annotations
import argparse
import json
import os
import sys
DATASET = "wiki_opentable"
HERE = os.path.dirname(os.path.abspath(__file__))
DEFAULT_OUT = os.path.join(os.path.dirname(HERE), "unified")
DEFAULT_DATA_ROOT = os.environ.get("DATA_ROOT", "/mnt/ramdisk/blobstore/timchen0618/data")
DEFAULT_SRC = os.path.join(DEFAULT_DATA_ROOT, "eval", DATASET, "unified", "test_with_chunks.unified.jsonl")
def main() -> int:
ap = argparse.ArgumentParser()
ap.add_argument("--src", default=DEFAULT_SRC, help=f"Source unified jsonl (default: {DEFAULT_SRC})")
ap.add_argument("--out", default=DEFAULT_OUT, help="Output unified/ directory (default: ./unified)")
args = ap.parse_args()
if not os.path.exists(args.src):
print(f"error: source file not found: {args.src}", file=sys.stderr)
return 2
records = []
with open(args.src) as f:
for line in f:
line = line.strip()
if not line:
continue
obj = json.loads(line)
qid = str(obj["qid"])
docs = obj.get("docs") or []
records.append({
"qid": qid,
"question": obj.get("question", ""),
"answers": obj.get("answers") or [],
"sql": obj.get("sql"),
"original_table_id": obj.get("original_table_id"),
"dataset_origin": obj.get("dataset_origin"),
"docs": [{"id": d.get("id"), "contents": d.get("contents", "")} for d in docs],
})
records.sort(key=lambda r: r["qid"])
rec_dir = os.path.join(args.out, "records")
os.makedirs(rec_dir, exist_ok=True)
for stale in os.listdir(rec_dir):
if stale.endswith(".json"):
os.remove(os.path.join(rec_dir, stale))
with open(os.path.join(args.out, "index.json"), "w") as f:
json.dump(
[
{
"qid": r["qid"],
"question": r["question"],
"dataset_origin": r["dataset_origin"],
"n_docs": len(r["docs"]),
"n_answers": len(r["answers"]),
}
for r in records
],
f,
ensure_ascii=False,
)
for r in records:
with open(os.path.join(rec_dir, f"{r['qid']}.json"), "w") as f:
json.dump(r, f, ensure_ascii=False)
print(f"wrote {len(records)} {DATASET} unified record(s) to {args.out}", file=sys.stderr)
return 0
if __name__ == "__main__":
raise SystemExit(main())
|