#!/usr/bin/env python3 """ Build the `e2e_structures/` shards for the open-wikitable-smoke-viewer. For each of the 100 wiki_opentable_smoke test qids, collect: - the question + gold answers (from the local unified eval shards) - every supporting doc (id + raw markdown) - for each supporting doc, the per-shape structure files produced by the information-scaffolds E2E pipeline (`scaffolds_dir//__*.`) Output layout: e2e_structures/index.json e2e_structures/records/.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). Smoke-specific differences from the parent open-wikitable-viewer builder: - Reads unified gold from `unified/records/.json` (per-qid shards committed in this repo) instead of a single .jsonl file — the parent's blobfuse path isn't mounted on this VM. - Shape set includes `tabular_records` (the e2e v1 smoke run) instead of `entity_fact_records` (the parent's full-corpus run). 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-slug-desc-dev-v0/" "e2e-wiki_opentable_dev-slug-desc-dev-v0/named-outputs/scaffolds_dir" ) DEFAULT_OUT = REPO / "e2e_structures" 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"") # --------------------------------------------------------------------------- # 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/.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 [] 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": 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())