scripts/build_unified.py

#!/usr/bin/env python3
"""Build open-wikitable-smoke `unified/` shards from the unified eval bundle.

Mirrors the open-wikitable-viewer `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}], ...]}).

Defaults to the wiki_opentable_smoke 100-qid bundle.

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_smoke"
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())