scripts/build_e2e_structures.py

#!/usr/bin/env python3
"""Build per-question E2E v3 structure shards for a static viewer."""

from __future__ import annotations

import argparse
import base64
import gzip
import json
import shutil
from collections import defaultdict
from pathlib import Path
from typing import Any, Iterator


ROOT = Path(__file__).resolve().parent.parent
DEFAULT_SHAPES = [
    "tabular_records",
    "chronology_and_timeline_indexes",
    "claim_and_theme_summaries",
    "qa_shortcuts_and_templates",
    "relation_graphs_and_mappings",
]


def iter_jsonl(path: Path) -> Iterator[dict[str, Any]]:
    with path.open(encoding="utf-8") as handle:
        for line_number, line in enumerate(handle, 1):
            if not line.strip():
                continue
            value = json.loads(line)
            if not isinstance(value, dict):
                raise ValueError(f"{path}:{line_number}: expected an object")
            yield value


def load_shape_index(
    scaffolds_dir: Path, shape: str
) -> tuple[str, dict[str, list[dict[str, str]]]]:
    index_path = scaffolds_dir / shape / "_index.json"
    if not index_path.exists():
        raise FileNotFoundError(f"missing shape index: {index_path}")
    payload = json.loads(index_path.read_text(encoding="utf-8"))
    files: dict[str, list[dict[str, str]]] = defaultdict(list)
    for entry in payload.get("entries", []):
        if entry.get("doc_id") is None or not entry.get("file"):
            continue
        files[str(entry["doc_id"])].append(
            {
                "file": str(entry["file"]),
                "unit_name": str(entry.get("unit_name") or ""),
                "unit_description": str(entry.get("unit_description") or ""),
            }
        )
    return str(payload.get("description") or ""), dict(files)


def load_bundle_map(path: Path | None) -> dict[str, str]:
    if path is None:
        return {}
    mapping: dict[str, str] = {}
    for row in iter_jsonl(path):
        physical_id = str(row["physical_id"])
        for source_id in row.get("source_doc_ids", []):
            source_id = str(source_id)
            if source_id in mapping:
                raise ValueError(f"duplicate bundle mapping for {source_id}")
            mapping[source_id] = physical_id
    return mapping


def structure_format(filename: str) -> str:
    suffixes = Path(filename).suffixes
    if filename.endswith(".timeline.json"):
        return "json"
    if filename.endswith(".edges.jsonl"):
        return "jsonl"
    return suffixes[-1].lstrip(".").lower() if suffixes else "txt"


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--dataset", required=True)
    parser.add_argument("--unified", type=Path, required=True)
    parser.add_argument("--scaffolds-dir", type=Path, required=True)
    parser.add_argument("--bundle-manifest", type=Path)
    parser.add_argument("--out", type=Path, default=ROOT / "e2e_structures_v3")
    parser.add_argument("--label", default="E2E Structures v3")
    parser.add_argument("--shapes", default=",".join(DEFAULT_SHAPES))
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    shapes = [value.strip() for value in args.shapes.split(",") if value.strip()]
    if not shapes:
        raise ValueError("--shapes must contain at least one shape")

    shape_descriptions: dict[str, str] = {}
    shape_files: dict[str, dict[str, list[dict[str, str]]]] = {}
    for shape in shapes:
        description, files = load_shape_index(args.scaffolds_dir, shape)
        shape_descriptions[shape] = description
        shape_files[shape] = files
        print(f"{shape:38s} {len(files):>6} extraction docs")

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

    index_rows: list[dict[str, Any]] = []
    extraction_docs_seen: set[str] = set()
    source_docs_seen: set[str] = set()
    source_docs_unmapped: set[str] = set()
    extraction_docs_without_structures: set[str] = set()
    total_structures = 0

    for row in iter_jsonl(args.unified):
        grouped: dict[str, list[dict[str, Any]]] = {}
        group_order: list[str] = []
        for doc in row.get("docs", []) or []:
            source_id = str(doc.get("id"))
            source_docs_seen.add(source_id)
            extraction_id = bundle_map.get(source_id, source_id)
            if bundle_map and source_id not in bundle_map:
                source_docs_unmapped.add(source_id)
            if extraction_id not in grouped:
                grouped[extraction_id] = []
                group_order.append(extraction_id)
            grouped[extraction_id].append(
                {"doc_id": source_id, "contents": str(doc.get("contents") or "")}
            )

        per_doc: list[dict[str, Any]] = []
        question_structures = 0
        for extraction_id in group_order:
            extraction_docs_seen.add(extraction_id)
            structures: list[dict[str, Any]] = []
            for shape in shapes:
                for entry in shape_files[shape].get(extraction_id, []):
                    filename = entry["file"]
                    path = args.scaffolds_dir / shape / filename
                    structures.append(
                        {
                            "shape_id": shape,
                            "description": shape_descriptions[shape],
                            "file": filename,
                            "unit_name": entry["unit_name"],
                            "unit_description": entry["unit_description"],
                            "format": structure_format(filename),
                            "content": path.read_text(encoding="utf-8")
                            if path.exists()
                            else "",
                        }
                    )
            if not structures:
                extraction_docs_without_structures.add(extraction_id)
            question_structures += len(structures)
            source_docs = grouped[extraction_id]
            per_doc.append(
                {
                    "doc_id": extraction_id,
                    "source_doc_ids": [doc["doc_id"] for doc in source_docs],
                    "source_docs": source_docs,
                    "is_supporting": True,
                    "n_structures": len(structures),
                    "contents": "\n\n".join(
                        f"===== {doc['doc_id']} =====\n{doc['contents']}"
                        for doc in source_docs
                    ),
                    "structures": structures,
                }
            )

        qid = str(row["qid"])
        record = {
            "qid": qid,
            "dataset": args.dataset,
            "question": row.get("question"),
            "gold_answers": row.get("answers", []),
            "n_docs": len(per_doc),
            "n_source_docs": sum(len(doc["source_doc_ids"]) for doc in per_doc),
            "n_structures": question_structures,
            "docs": per_doc,
        }
        raw = json.dumps(record, ensure_ascii=False, separators=(",", ":")).encode()
        encoded = base64.b64encode(gzip.compress(raw, compresslevel=9)).decode("ascii")
        relative_path = f"records/{qid}.json.gzip.b64"
        (args.out / relative_path).write_text(encoded, encoding="ascii")
        index_rows.append(
            {
                "qid": qid,
                "question": row.get("question"),
                "n_docs": len(per_doc),
                "n_source_docs": record["n_source_docs"],
                "n_structures": question_structures,
                "doc_ids": [doc["doc_id"] for doc in per_doc],
                "source_doc_ids": [
                    source_id
                    for doc in per_doc
                    for source_id in doc["source_doc_ids"]
                ],
                "path": relative_path,
            }
        )
        total_structures += question_structures

    meta = {
        "label": args.label,
        "dataset": args.dataset,
        "n_qids": len(index_rows),
        "n_extraction_docs": len(extraction_docs_seen),
        "n_source_docs": len(source_docs_seen),
        "n_source_docs_unmapped": len(source_docs_unmapped),
        "n_extraction_docs_without_structures": len(
            extraction_docs_without_structures
        ),
        "n_structures_total": total_structures,
        "shapes": shapes,
        "shape_descriptions": shape_descriptions,
        "scaffolds_dir": str(args.scaffolds_dir),
        "unified": str(args.unified),
        "bundle_manifest": str(args.bundle_manifest) if args.bundle_manifest else None,
    }
    args.out.mkdir(parents=True, exist_ok=True)
    (args.out / "index.json").write_text(
        json.dumps({"meta": meta, "rows": index_rows}, ensure_ascii=False),
        encoding="utf-8",
    )
    print(json.dumps(meta, indent=2))


if __name__ == "__main__":
    main()