social-data-attribution-w2/src/dolma/bin_analysis/cli.py

"""CLI entry point for bin population analysis."""

from __future__ import annotations

import argparse
import logging
from pathlib import Path

from dolma.bin_analysis.stats import (
    DEFAULT_DOC_THRESHOLD,
    DEFAULT_TOKEN_CAP,
    DEFAULT_TOKEN_FLOOR,
)
from dolma.constants import TARGET_DOCS_PER_BIN


def main(argv: list[str] | None = None) -> None:
    parser = argparse.ArgumentParser(
        description="Analyze bin populations in the 24x24 WebOrganizer grid."
    )
    parser.add_argument(
        "--manifest", type=Path, required=True, help="Path to manifest parquet"
    )
    parser.add_argument(
        "--output-dir",
        type=Path,
        default=Path("runs/bin_analysis"),
        help="Output directory (default: runs/bin_analysis)",
    )
    parser.add_argument(
        "--target-docs-per-bin",
        type=int,
        default=TARGET_DOCS_PER_BIN,
        help="Target docs per bin for classification",
    )
    parser.add_argument(
        "--doc-threshold",
        type=int,
        default=DEFAULT_DOC_THRESHOLD,
        help="Min docs for length mismatch detection",
    )
    parser.add_argument(
        "--token-floor",
        type=int,
        default=DEFAULT_TOKEN_FLOOR,
        help="Token floor for short_docs flag",
    )
    parser.add_argument(
        "--token-cap",
        type=int,
        default=DEFAULT_TOKEN_CAP,
        help="Token cap for long_docs flag",
    )
    parser.add_argument(
        "--report",
        action="store_true",
        default=True,
        dest="report",
        help="Generate HTML report (default)",
    )
    parser.add_argument("--no-report", action="store_false", dest="report")
    parser.add_argument("--verbose", action="store_true", default=False)
    args = parser.parse_args(argv)

    logging.basicConfig(
        level=logging.DEBUG if args.verbose else logging.INFO,
        format="%(asctime)s %(levelname)s [%(name)s] %(message)s",
    )

    from dolma.bin_analysis.core import run_bin_analysis

    result = run_bin_analysis(
        manifest_path=args.manifest,
        output_dir=args.output_dir,
        target_docs_per_bin=args.target_docs_per_bin,
        doc_threshold=args.doc_threshold,
        token_floor=args.token_floor,
        token_cap=args.token_cap,
    )

    if args.report:
        from dolma.bin_analysis.report import build_bin_report

        build_bin_report(
            result.bin_stats,
            result.mismatches,
            result.recommendation,
            result.output_dir,
        )


if __name__ == "__main__":
    main()