social-data-attribution-w2/scripts/slurm/dedup/finalize_unique_docs.py

#!/usr/bin/env python3
"""Build the SOC-127 manifest, stats, and dataset card."""

from __future__ import annotations

import argparse
import json
from collections import Counter
from pathlib import Path

from dolma.dedup.materialize import (
    REQUIRED_PREFLIGHT_FAMILIES,
    ensure_preflight_approved,
    iter_shard_records,
    load_preflight,
    resolve_record_doc_id,
    write_json,
)

DEFAULT_ROOT = Path("/storage/ice-shared/cs7634/staff/TDA/soc-90/unique_docs_work")


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Finalize SOC-127 outputs")
    parser.add_argument("--work-root", type=Path, default=DEFAULT_ROOT)
    parser.add_argument("--pool-dir", type=Path, default=None)
    parser.add_argument("--mix-raw-dir", type=Path, default=None)
    parser.add_argument("--mix-dir", type=Path, default=None)
    parser.add_argument("--preflight-file", type=Path, default=None)
    parser.add_argument("--manifest-output", type=Path, default=None)
    parser.add_argument("--stats-output", type=Path, default=None)
    parser.add_argument("--readme-output", type=Path, default=None)
    return parser.parse_args(argv)


def sum_field(paths: list[Path], field: str) -> int:
    return sum(
        int(json.loads(path.read_text(encoding="utf-8")).get(field, 0))
        for path in paths
    )


def write_manifest(
    data_files: list[Path], work_root: Path, manifest_output: Path
) -> tuple[int, Counter[str]]:
    import pyarrow as pa
    import pyarrow.parquet as pq

    rows: list[dict[str, object]] = []
    writer: pq.ParquetWriter | None = None
    family_counts: Counter[str] = Counter()
    row_count = 0
    for data_path in data_files:
        for _, record in iter_shard_records(data_path):
            if record is None or not isinstance(record, dict):
                raise ValueError(f"Invalid JSON in finalized output: {data_path}")
            metadata = record.get("_soc_127")
            if not isinstance(metadata, dict):
                raise ValueError(f"Missing manifest metadata in {data_path}")
            doc_id = metadata.get("doc_id")
            if not isinstance(doc_id, str) or not doc_id:
                doc_id, _ = resolve_record_doc_id(
                    record, str(metadata.get("input_shard", ""))
                )
            if not isinstance(doc_id, str) or not doc_id:
                raise ValueError(f"Missing manifest metadata in {data_path}")
            family = str(metadata["source_family"])
            rows.append(
                {
                    "doc_id": doc_id,
                    "source_family": family,
                    "source_folder": str(metadata["source_folder"]),
                    "input_shard": str(metadata["input_shard"]),
                    "phase": str(metadata["phase"]),
                    "output_path": str(data_path.relative_to(work_root)),
                    "text_length": len(str(record.get("text", ""))),
                }
            )
            family_counts[family] += 1
            row_count += 1
            if len(rows) >= 5_000:
                table = pa.Table.from_pylist(rows)
                if writer is None:
                    manifest_output.parent.mkdir(parents=True, exist_ok=True)
                    writer = pq.ParquetWriter(manifest_output, table.schema)
                writer.write_table(table)
                rows.clear()
    if rows:
        table = pa.Table.from_pylist(rows)
        if writer is None:
            manifest_output.parent.mkdir(parents=True, exist_ok=True)
            writer = pq.ParquetWriter(manifest_output, table.schema)
        writer.write_table(table)
    if writer is None:
        raise ValueError("No finalized records were found")
    writer.close()
    return row_count, family_counts


def build_readme(stats: dict[str, object]) -> str:
    return (
        "# dolma3-6t-unique\n\n"
        "This dataset contains SOC-127 materialized documents selected by the 6T Bloom filter.\n\n"
        "## Notes\n\n"
        "- Non-pool sources are gated by a fail-closed preflight on resolved document ID compatibility.\n"
        "- `stack_edu-Python` uses `blob_id` as the resolved document ID because those mix shards do not expose a top-level `id`.\n"
        "- Pool materialization is restricted to exact shard paths shared by the pool and the 6T mix.\n"
        "- Final counts may drift from issue estimates. Those estimates are advisory only.\n"
        f"- Total rows in this materialization: {stats['manifest_rows']:,}\n"
    )


def main(argv: list[str] | None = None) -> int:
    args = parse_args(argv)
    pool_dir = args.pool_dir or args.work_root / "pool"
    mix_raw_dir = args.mix_raw_dir or args.work_root / "mix_nonpool_raw"
    mix_dir = args.mix_dir or args.work_root / "mix_nonpool_final"
    preflight_file = args.preflight_file or args.work_root / "preflight.json"
    manifest_output = args.manifest_output or args.work_root / "manifest.parquet"
    stats_output = args.stats_output or args.work_root / "stats.json"
    readme_output = args.readme_output or args.work_root / "README.md"
    preflight = load_preflight(preflight_file)
    ensure_preflight_approved(preflight, REQUIRED_PREFLIGHT_FAMILIES)
    pool_stats = sorted(pool_dir.rglob("*.stats.json"))
    mix_raw_stats = sorted(mix_raw_dir.glob("task_*.jsonl.zst.stats.json"))
    merge_stats_path = mix_dir / "merge.stats.json"
    if not merge_stats_path.exists():
        raise ValueError(f"Missing merge stats at {merge_stats_path}")
    merge_stats = json.loads(merge_stats_path.read_text(encoding="utf-8"))
    data_files = sorted(pool_dir.rglob("*.jsonl.zst")) + sorted(
        mix_dir.rglob("bucket_*.jsonl.zst")
    )
    manifest_rows, family_counts = write_manifest(
        data_files, args.work_root, manifest_output
    )
    expected_rows = sum_field(pool_stats, "records_kept") + int(
        merge_stats.get("records_written", 0)
    )
    if manifest_rows != expected_rows:
        raise ValueError(
            f"Manifest rows {manifest_rows} did not match expected rows {expected_rows}"
        )
    missing_rows = [
        family
        for family in REQUIRED_PREFLIGHT_FAMILIES
        if family_counts.get(family, 0) == 0
    ]
    if missing_rows:
        raise ValueError(
            f"Approved families produced zero rows: {', '.join(missing_rows)}"
        )
    stats = {
        "manifest_rows": manifest_rows,
        "expected_issue_rows": 1_260_000_000,
        "count_warning": manifest_rows != 1_260_000_000,
        "pool_records_kept": sum_field(pool_stats, "records_kept"),
        "mix_raw_records_kept": sum_field(mix_raw_stats, "records_kept"),
        "mix_unique_records_kept": int(merge_stats.get("records_written", 0)),
        "cross_task_duplicates_removed": int(
            merge_stats.get("cross_task_duplicates_removed", 0)
        ),
        "source_family_counts": dict(sorted(family_counts.items())),
        "preflight_file": str(preflight_file),
    }
    write_json(stats_output, stats)
    readme_output.parent.mkdir(parents=True, exist_ok=True)
    readme_output.write_text(build_readme(stats), encoding="utf-8")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())