social-data-attribution-w2/src/dolma/eda/core.py

"""Dolma EDA aggregation and reporting."""

from __future__ import annotations

import argparse
import json
import logging
import shlex
from pathlib import Path
from typing import Iterable, Sequence

from data_attribution.cli.config import configure_logging
from dolma.eda.outputs import write_outputs
from dolma.eda.report import build_report
from dolma.eda.state import EdaState
from dolma.local_io import iter_jsonl, iter_jsonlzst

logger = logging.getLogger(__name__)


def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Dolma enriched dataset EDA")
    parser.add_argument("--mode", choices=("run", "worker", "merge"), default="run")
    parser.add_argument("--input-dir", type=Path, default=Path("runs/dolma_enriched"))
    parser.add_argument("--manifest", type=Path)
    parser.add_argument("--manifest-index", type=int)
    parser.add_argument(
        "--output-dir", type=Path, default=Path("runs/dolma_enriched/eda")
    )
    parser.add_argument("--worker-output-dir", type=Path)
    parser.add_argument("--workers-dir", type=Path)
    parser.add_argument("--resume", action=argparse.BooleanOptionalAction, default=True)
    parser.add_argument(
        "--value-fields",
        nargs="+",
        default=["lang", "warc_content_type", "cc_dump"],
    )
    parser.add_argument("--top-k", type=int, default=50)
    parser.add_argument("--joint-top-k", type=int, default=200)
    parser.add_argument("--heatmap-k", type=int, default=25)
    parser.add_argument("--report", action=argparse.BooleanOptionalAction, default=True)
    parser.add_argument(
        "--write-png", action=argparse.BooleanOptionalAction, default=True
    )
    parser.add_argument(
        "--verbose", action=argparse.BooleanOptionalAction, default=False
    )
    return parser


def _parse_manifest_groups(manifest: Path) -> list[list[Path]]:
    groups: list[list[Path]] = []
    for line in manifest.read_text().splitlines():
        if not line.strip():
            continue
        groups.append([Path(entry) for entry in shlex.split(line)])
    return groups


def _collect_inputs(
    input_dir: Path, manifest: Path | None, manifest_index: int | None
) -> list[Path]:
    if manifest:
        groups = _parse_manifest_groups(manifest)
        if manifest_index is None:
            return [path for group in groups for path in group]
        if manifest_index < 0 or manifest_index >= len(groups):
            raise ValueError("manifest-index is out of range")
        return groups[manifest_index]
    files = list(input_dir.rglob("*.enriched.jsonl.zst"))
    files.extend(input_dir.rglob("*.enriched.jsonl"))
    return sorted(files)


def _iter_records(path: Path) -> Iterable[dict[str, object]]:
    if path.suffix == ".zst":
        return iter_jsonlzst(path)
    return iter_jsonl(path)


def _load_state(output_dir: Path, resume: bool) -> EdaState:
    state_path = output_dir / "state.json"
    if resume and state_path.exists():
        return EdaState.from_dict(json.loads(state_path.read_text()))
    return EdaState()


def _load_done(output_dir: Path, resume: bool) -> set[str]:
    done_path = output_dir / "manifest_done.jsonl"
    if not (resume and done_path.exists()):
        return set()
    return {
        json.loads(line)["path"]
        for line in done_path.read_text().splitlines()
        if line.strip()
    }


def _append_done(output_dir: Path, entry: dict[str, object]) -> None:
    done_path = output_dir / "manifest_done.jsonl"
    with done_path.open("a", encoding="utf-8") as stream:
        stream.write(json.dumps(entry))
        stream.write("\n")


def _save_state(output_dir: Path, state: EdaState) -> None:
    state_path = output_dir / "state.json"
    state_path.write_text(json.dumps(state.to_dict(), indent=2), encoding="utf-8")


def _process_path(
    path: Path, state: EdaState, output_dir: Path, value_fields: Sequence[str]
) -> int:
    record_count = 0
    for record in _iter_records(path):
        if isinstance(record, dict):
            state.update_record(record, value_fields)
            record_count += 1
    _append_done(output_dir, {"path": str(path), "records": record_count})
    _save_state(output_dir, state)
    return record_count


def run_analysis(
    input_paths: Sequence[Path],
    output_dir: Path,
    *,
    value_fields: Sequence[str],
    resume: bool,
    report: bool,
    top_k: int,
    joint_top_k: int,
    heatmap_k: int,
    write_png: bool,
) -> EdaState:
    output_dir.mkdir(parents=True, exist_ok=True)
    state = _load_state(output_dir, resume)
    done = _load_done(output_dir, resume)
    for path in input_paths:
        if str(path) in done:
            continue
        record_count = _process_path(path, state, output_dir, value_fields)
        logger.info("Processed %s (%s records)", path, record_count)
    write_outputs(output_dir, state, top_k=top_k, joint_top_k=joint_top_k)
    if report:
        build_report(output_dir, top_k=top_k, heatmap_k=heatmap_k, write_png=write_png)
    return state


def run_worker(
    input_paths: Sequence[Path],
    output_dir: Path,
    *,
    value_fields: Sequence[str],
    resume: bool,
) -> EdaState:
    output_dir.mkdir(parents=True, exist_ok=True)
    state = _load_state(output_dir, resume)
    done = _load_done(output_dir, resume)
    for path in input_paths:
        if str(path) in done:
            continue
        record_count = _process_path(path, state, output_dir, value_fields)
        logger.info("Processed %s (%s records)", path, record_count)
    return state


def merge_states(workers_dir: Path) -> EdaState:
    state = EdaState()
    for state_path in workers_dir.rglob("state.json"):
        state.merge(EdaState.from_dict(json.loads(state_path.read_text())))
    return state


def main(argv: Sequence[str] | None = None) -> int:
    args = _build_parser().parse_args(argv)
    configure_logging(args.verbose)
    if args.mode == "merge":
        workers_dir = args.workers_dir or (args.output_dir / "workers")
        if not workers_dir.exists():
            raise ValueError(f"Workers directory not found: {workers_dir}")
        state = merge_states(workers_dir)
        write_outputs(
            args.output_dir, state, top_k=args.top_k, joint_top_k=args.joint_top_k
        )
        if args.report:
            build_report(
                args.output_dir,
                top_k=args.top_k,
                heatmap_k=args.heatmap_k,
                write_png=args.write_png,
            )
        return 0

    input_paths = _collect_inputs(args.input_dir, args.manifest, args.manifest_index)
    if not input_paths:
        raise ValueError("No enriched shards found to analyze.")
    if args.mode == "worker":
        output_dir = args.worker_output_dir or args.output_dir
        run_worker(
            input_paths,
            output_dir,
            value_fields=args.value_fields,
            resume=args.resume,
        )
        return 0

    run_analysis(
        input_paths,
        args.output_dir,
        value_fields=args.value_fields,
        resume=args.resume,
        report=args.report,
        top_k=args.top_k,
        joint_top_k=args.joint_top_k,
        heatmap_k=args.heatmap_k,
        write_png=args.write_png,
    )
    return 0


__all__ = ["main", "merge_states", "run_analysis", "run_worker"]


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