scripts/build_e2e_structures.py

#!/usr/bin/env python3
"""
Build `e2e_structures_v*/` shards for the monaco-dev-viewer.

For each of the 100 monaco_dev test qids, collect:
  - the question + gold answers (from the local unified eval shards)
  - every supporting doc (id + raw markdown from unified corpus)
  - 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:
  e2e_structures_v0/index.json         # or v2, v3, ...
  e2e_structures_v0/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).

Usage:
    # e2e v0
    python scripts/build_e2e_structures.py \
        --scaffolds-dir /home/azureuser/run_logs/e2e-monaco_dev-comparison/outputs/v0/named-outputs/scaffolds_dir \
        --out           e2e_structures_v0 \
        --label         "e2e v0 (monaco_dev)"

    # e2e v2 / v3 — same shape, just point at the matching scaffolds_dir
    python scripts/build_e2e_structures.py --scaffolds-dir .../v2/named-outputs/scaffolds_dir --out e2e_structures_v2 --label "e2e v2 (monaco_dev)"
    python scripts/build_e2e_structures.py --scaffolds-dir .../v3/named-outputs/scaffolds_dir --out e2e_structures_v3 --label "e2e v3 (monaco_dev)"

Self-contained — only stdlib + filesystem reads.
"""

from __future__ import annotations

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


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

HERE = Path(__file__).resolve().parent
REPO = HERE.parent

DEFAULT_UNIFIED_DIR = REPO / "unified" / "records"
DEFAULT_SCAFFOLDS_DIR = (
    "/home/azureuser/run_logs/e2e-monaco_dev-comparison/outputs/v0/named-outputs/scaffolds_dir"
)
DEFAULT_OUT = REPO / "e2e_structures_v0"

SHAPES = [
    "tabular_records",
    "chronology_and_timeline_indexes",
    "claim_and_theme_summaries",
    "qa_shortcuts_and_templates",
    "relation_graphs_and_mappings",
]


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

def load_unified_dir(path: Path) -> list[dict[str, Any]]:
    """Smoke layout: one JSON file per qid in `unified/records/`."""
    rows = []
    for p in sorted(path.glob("*.json")):
        rows.append(json.loads(p.read_text()))
    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 from final extension token.

    Filenames may use multi-part suffixes such as `.edges.jsonl` or
    `.timeline.json`; the renderer keys off the trailing token so this
    yields the right rendering (jsonl / json / csv).
    """
    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=str(DEFAULT_UNIFIED_DIR),
                    help="Path to unified shards dir (one *.json per qid).")
    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("--label", default=None,
                    help="Human label for the run (stored in meta.label; default: derive from scaffolds-dir).")
    args = ap.parse_args()

    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_dir(Path(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 []

        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": "monaco_dev",
            "question": row.get("question"),
            "gold_answers": row.get("answers", []),
            "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,
            "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": str(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())