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

#!/usr/bin/env python3
"""Materialize pool-backed SOC-127 documents."""

from __future__ import annotations

import argparse
import json
import logging
from pathlib import Path

from dolma.constants import DOLMA_POOL_DATASET_ID
from dolma.dedup.materialize import (
    done_path_for,
    download_shard,
    iter_shard_records,
    maybe_sleep_for_startup,
    open_zstd_writer,
    read_manifest,
    resolve_record_doc_id,
    resolve_task,
    shard_slice,
    stats_path_for,
    write_done_marker,
    write_json,
)
from dolma.provenance import BloomIndex, shard_folder_name, source_family

logger = logging.getLogger("materialize_pool_docs")

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


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Materialize pool documents")
    parser.add_argument("--dataset", default=DOLMA_POOL_DATASET_ID)
    parser.add_argument("--manifest", type=Path, required=True)
    parser.add_argument("--bloom-file", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, default=DEFAULT_OUTPUT)
    parser.add_argument("--download-dir", type=Path, default=None)
    parser.add_argument("--local-root", type=Path, default=None)
    parser.add_argument("--task-id", type=int, default=None)
    parser.add_argument("--task-count", type=int, default=None)
    parser.add_argument("--shards-per-task", type=int, default=5)
    parser.add_argument("--startup-delay", type=int, default=60)
    parser.add_argument(
        "--skip-existing", action=argparse.BooleanOptionalAction, default=True
    )
    parser.add_argument(
        "--drop-removed-text", action=argparse.BooleanOptionalAction, default=False
    )
    parser.add_argument("--verbose", action="store_true", default=False)
    return parser.parse_args(argv)


def materialize_shard(
    args: argparse.Namespace, bloom: BloomIndex, shard_path: str
) -> None:
    output_path = args.output_dir / shard_path
    done_path = done_path_for(output_path)
    if args.skip_existing and done_path.exists():
        logger.info("Skipping existing shard %s", shard_path)
        return
    local_path = download_shard(
        args.dataset, shard_path, args.download_dir, args.local_root
    )
    stats = {
        "dataset": args.dataset,
        "input_shard": shard_path,
        "source_family": source_family(shard_path),
        "source_folder": shard_folder_name(shard_path),
        "records_seen": 0,
        "records_kept": 0,
        "records_invalid_json": 0,
        "records_removed_text": 0,
        "records_missing_doc_id": 0,
    }
    output_path.parent.mkdir(parents=True, exist_ok=True)
    try:
        with open_zstd_writer(output_path) as writer:
            for _, record in iter_shard_records(local_path):
                stats["records_seen"] += 1
                if record is None or not isinstance(record, dict):
                    stats["records_invalid_json"] += 1
                    continue
                doc_id, doc_id_field = resolve_record_doc_id(record, shard_path)
                if doc_id is None:
                    stats["records_missing_doc_id"] += 1
                    continue
                if (
                    args.drop_removed_text
                    and str(record.get("text", "")).strip() == "[REMOVED]"
                ):
                    stats["records_removed_text"] += 1
                    continue
                if doc_id not in bloom:
                    continue
                enriched = dict(record)
                enriched["_soc_127"] = {
                    "doc_id": doc_id,
                    "doc_id_field": doc_id_field,
                    "input_shard": shard_path,
                    "phase": "pool",
                    "source_family": stats["source_family"],
                    "source_folder": stats["source_folder"],
                }
                writer.write(json.dumps(enriched, sort_keys=True) + "\n")
                stats["records_kept"] += 1
        write_json(stats_path_for(output_path), stats)
        write_done_marker(done_path)
    except Exception:
        if output_path.exists():
            output_path.unlink()
        raise


def main(argv: list[str] | None = None) -> int:
    args = parse_args(argv)
    logging.basicConfig(
        level=logging.DEBUG if args.verbose else logging.INFO,
        format="%(asctime)s %(levelname)s %(message)s",
    )
    task_id, _ = resolve_task(args.task_id, args.task_count)
    maybe_sleep_for_startup(task_id, args.startup_delay)
    shards = read_manifest(args.manifest)
    selected = shard_slice(shards, task_id, args.shards_per_task)
    if not selected:
        logger.info("No pool shards assigned for task %d", task_id)
        return 0
    bloom = BloomIndex.load(args.bloom_file)
    for shard_path in selected:
        materialize_shard(args, bloom, shard_path)
    return 0


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