#!/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()