scripts/build_e2e_structures.py

#!/usr/bin/env python3
"""
Build E2E structure shards for the open-wikitable-viewer.

For each of the 500 wiki_opentable test qids, collect:
  - the question + gold answers (from the unified eval bundle)
  - every supporting doc (id + raw markdown)
  - for each supporting doc, the per-shape structure files produced by the
    information-scaffolds E2E pipeline (`scaffolds_dir/<shape_id>/<doc_id>.<ext>`)

Output layout:
  <out>/index.json
  <out>/records/<qid>.json

The shape-id → human description mapping comes from each shape's
`_index.json` (key "description"). File contents are embedded verbatim
(small: ~200-2000 chars each).

Compare-tab style: this script is self-contained — no dependency on the
`information-scaffolds` package. Only stdlib + filesystem reads.
"""

from __future__ import annotations

import argparse
import json
import os
import shutil
import sys
from pathlib import Path
from typing import Any


# ---------------------------------------------------------------------------
# defaults

DEFAULT_UNIFIED = (
    "/mnt/ramdisk/blobstore/timchen0618/data/eval/wiki_opentable/unified/"
    "test_with_chunks.unified.jsonl"
)
DEFAULT_SCAFFOLDS_DIR = (
    "/home/azureuser/projects/information-scaffolds/outputs/e2e_runs/"
    "wiki-opentable-fullcorpus-fulleval-allshapes16k-rpm300-conc500-symlink-20260622/"
    "named-outputs/scaffolds_dir"
)
DEFAULT_OUT = Path(__file__).resolve().parent.parent / "e2e_structures"

DEFAULT_SHAPES = [
    "entity_fact_records",
    "chronology_and_timeline_indexes",
    "claim_and_theme_summaries",
    "qa_shortcuts_and_templates",
    "relation_graphs_and_mappings",
]


# ---------------------------------------------------------------------------
# loaders

def load_unified(path: str) -> list[dict[str, Any]]:
    rows = []
    with open(path) as f:
        for line in f:
            line = line.strip()
            if line:
                rows.append(json.loads(line))
    return rows


def load_shape_index(scaffolds_dir: str, shape: str) -> dict[str, Any]:
    """Return {description, doc_id → file_basename}."""
    ix_path = Path(scaffolds_dir) / shape / "_index.json"
    if not ix_path.exists():
        print(f"  warn: missing {ix_path} — skipping shape", file=sys.stderr)
        return {"description": "", "files": {}}
    ix = json.loads(ix_path.read_text())
    files: dict[str, str] = {}
    for e in ix.get("entries", []):
        doc_id = str(e.get("doc_id"))
        fname = e.get("file")
        if doc_id and fname:
            files[doc_id] = fname
    return {"description": ix.get("description", ""), "files": files}


def read_structure_file(scaffolds_dir: str, shape: str, fname: str) -> tuple[str, str]:
    """Return (format, content). format ∈ {csv, json, jsonl}, falls back to ext."""
    p = Path(scaffolds_dir) / shape / fname
    if not p.exists():
        return ("missing", "")
    ext = fname.rsplit(".", 1)[-1].lower()
    fmt = ext if ext in {"csv", "json", "jsonl", "md"} else "txt"
    try:
        return (fmt, p.read_text())
    except Exception as e:
        return ("error", f"<read error: {e}>")


# ---------------------------------------------------------------------------
# main

def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--unified", default=DEFAULT_UNIFIED,
                    help="Path to test_with_chunks.unified.jsonl (qid → docs[id, contents]).")
    ap.add_argument("--scaffolds-dir", default=DEFAULT_SCAFFOLDS_DIR,
                    help="Path to e2e named-outputs/scaffolds_dir (has 5 shape subdirs).")
    ap.add_argument("--out", default=str(DEFAULT_OUT),
                    help="Output dir (will hold index.json + records/<qid>.json).")
    ap.add_argument(
        "--shapes",
        default=",".join(DEFAULT_SHAPES),
        help="Comma-separated scaffold shape directories to include.",
    )
    ap.add_argument("--label", default=None,
                    help="Human label for the run (stored in meta.label; default: derive from scaffolds-dir).")
    args = ap.parse_args()
    shapes = [shape.strip() for shape in args.shapes.split(",") if shape.strip()]
    if not shapes:
        ap.error("--shapes must contain at least one shape directory")

    out_dir = Path(args.out)
    records_dir = out_dir / "records"
    if records_dir.exists():
        shutil.rmtree(records_dir)
    records_dir.mkdir(parents=True, exist_ok=True)

    # Pre-load all shape indexes once.
    print("Loading shape indexes …")
    shape_data: dict[str, dict[str, Any]] = {}
    for shape in shapes:
        sd = load_shape_index(args.scaffolds_dir, shape)
        shape_data[shape] = sd
        print(f"  {shape:38s} {len(sd['files']):>6} docs indexed")

    # Walk unified eval, build per-qid shards.
    rows = load_unified(args.unified)
    print(f"\nUnified rows: {len(rows)}")

    index_rows: list[dict[str, Any]] = []
    total_structures = 0
    docs_seen: set[str] = set()
    docs_with_no_structures: set[str] = set()

    for row in rows:
        qid = row["qid"]
        docs_in = row.get("docs", []) or []
        dataset_origin = row.get("dataset_origin")
        if dataset_origin is None:
            dataset_origin = "wikisql" if qid.startswith("wikisql") else "wikitq"

        per_doc: list[dict[str, Any]] = []
        n_structures_qid = 0
        for d in docs_in:
            doc_id = str(d.get("id"))
            contents = d.get("contents", "")
            docs_seen.add(doc_id)

            structs: list[dict[str, Any]] = []
            for shape in shapes:
                fname = shape_data[shape]["files"].get(doc_id)
                if not fname:
                    continue
                fmt, content = read_structure_file(args.scaffolds_dir, shape, fname)
                structs.append({
                    "shape_id": shape,
                    "description": shape_data[shape]["description"],
                    "file": fname,
                    "format": fmt,
                    "content": content,
                })
            if not structs:
                docs_with_no_structures.add(doc_id)
            n_structures_qid += len(structs)
            per_doc.append({
                "doc_id": doc_id,
                "is_supporting": True,
                "n_structures": len(structs),
                "contents": contents,
                "structures": structs,
            })

        rec = {
            "qid": qid,
            "dataset_origin": dataset_origin,
            "original_table_id": row.get("original_table_id"),
            "question": row.get("question"),
            "gold_answers": row.get("answers", []),
            "sql": row.get("sql"),
            "n_docs": len(per_doc),
            "n_structures": n_structures_qid,
            "docs": per_doc,
        }
        (records_dir / f"{qid}.json").write_text(json.dumps(rec, ensure_ascii=False))

        index_rows.append({
            "qid": qid,
            "dataset_origin": dataset_origin,
            "original_table_id": row.get("original_table_id"),
            "question": row.get("question"),
            "n_docs": len(per_doc),
            "n_structures": n_structures_qid,
            "doc_ids": [d["doc_id"] for d in per_doc],
        })
        total_structures += n_structures_qid

    # Run label
    if args.label is None:
        run_dirname = Path(args.scaffolds_dir).resolve().parent.parent.name
        label = f"e2e-pipeline · {run_dirname} · scaffolds_dir"
    else:
        label = args.label

    # Meta
    n_with = sum(1 for r in index_rows if r["n_structures"] > 0)
    avg = (total_structures / n_with) if n_with else 0
    meta = {
        "label": label,
        "scaffolds_dir": args.scaffolds_dir,
        "unified": args.unified,
        "n_qids": len(index_rows),
        "n_docs_unique": len(docs_seen),
        "n_docs_with_no_structures": len(docs_with_no_structures),
        "n_structures_total": total_structures,
        "n_structures_avg_per_qid": round(avg, 2),
        "shapes": shapes,
        "shape_descriptions": {s: shape_data[s]["description"] for s in shapes},
    }
    index = {"meta": meta, "rows": index_rows}
    (out_dir / "index.json").write_text(json.dumps(index, ensure_ascii=False))

    print(f"\n✓ Wrote {out_dir}/index.json + {len(index_rows)} record shards")
    print(f"   n_qids                  = {len(index_rows)}")
    print(f"   n_docs_unique           = {len(docs_seen)}")
    print(f"   docs w/ no structures   = {len(docs_with_no_structures)}")
    print(f"   total structures        = {total_structures}")
    print(f"   avg structures / qid    = {round(avg, 2)} (over {n_with} qids with ≥1)")

    return 0


if __name__ == "__main__":
    sys.exit(main())