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

"""CLI for enriching Dolma-style datasets with format metadata."""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

from data_attribution.cli.config import configure_logging
from dolma.constants import DEFAULT_SPLIT, DOLMA_DATASET_ID
from dolma.hf_io import iter_dolma3_hf
from dolma.paths import expand_inputs, hf_output_name, output_path
from dolma.pipeline import iter_local_records, run_enrichment


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Dolma format enrichment CLI")
    subparsers = parser.add_subparsers(dest="command", required=True)

    enrich_parser = subparsers.add_parser(
        "enrich-format", help="Add format classification metadata"
    )
    enrich_parser.add_argument("--input", choices=("hf", "local"), required=True)
    enrich_parser.add_argument("--hf-dataset", default=DOLMA_DATASET_ID)
    enrich_parser.add_argument("--split", default=DEFAULT_SPLIT)
    enrich_parser.add_argument("--data-files", nargs="+")
    enrich_parser.add_argument("--input-files", nargs="+")
    enrich_parser.add_argument(
        "--output-dir", type=Path, default=Path("runs/dolma_enriched")
    )
    enrich_parser.add_argument("--cache-dir", type=Path, default=Path(".hf_cache"))
    enrich_parser.add_argument("--batch-size", type=int, default=16)
    enrich_parser.add_argument("--max-length", type=int, default=1024)
    enrich_parser.add_argument("--device", default="cuda")
    enrich_parser.add_argument("--dtype", default="bf16")
    enrich_parser.add_argument(
        "--format-model", default="WebOrganizer/FormatClassifier"
    )
    enrich_parser.add_argument(
        "--format-model-nourl", default="WebOrganizer/FormatClassifier-NoURL"
    )
    enrich_parser.add_argument("--topic-model", default="WebOrganizer/TopicClassifier")
    enrich_parser.add_argument(
        "--topic-model-nourl", default="WebOrganizer/TopicClassifier-NoURL"
    )
    enrich_parser.add_argument(
        "--prefer-url",
        action=argparse.BooleanOptionalAction,
        default=True,
        help="Prefer URL-based classifier; fallback to NoURL when missing",
    )
    enrich_parser.add_argument(
        "--fill-word-count",
        action=argparse.BooleanOptionalAction,
        default=True,
        help="Compute original_word_count when missing",
    )
    enrich_parser.add_argument(
        "--overwrite",
        action=argparse.BooleanOptionalAction,
        default=False,
        help="Replace existing enriched outputs (backs up then deletes backup on success)",
    )
    enrich_parser.add_argument(
        "--use-nourl-fallback",
        action=argparse.BooleanOptionalAction,
        default=True,
        help="Enable NoURL model when URL is missing",
    )
    enrich_parser.add_argument("--num-docs-limit", type=int)
    enrich_parser.add_argument("--target-tokens", type=int)
    enrich_parser.add_argument("--log-every", type=int, default=1000)
    enrich_parser.add_argument(
        "--resume", action=argparse.BooleanOptionalAction, default=False
    )
    enrich_parser.add_argument(
        "--output-compression", choices=("zst", "none"), default="zst"
    )
    enrich_parser.add_argument("--num-workers", type=int, default=2)
    enrich_parser.add_argument(
        "--compile-model", action=argparse.BooleanOptionalAction, default=False
    )
    enrich_parser.add_argument(
        "--use-memory-efficient-attention",
        action=argparse.BooleanOptionalAction,
        default=True,
    )
    enrich_parser.add_argument(
        "--unpad-inputs",
        action=argparse.BooleanOptionalAction,
        default=True,
    )
    enrich_parser.add_argument(
        "--mock-model", action=argparse.BooleanOptionalAction, default=False
    )
    enrich_parser.add_argument(
        "--verbose", action=argparse.BooleanOptionalAction, default=False
    )
    enrich_parser.add_argument(
        "--min-words",
        type=int,
        default=0,
        help=(
            "Drop documents with fewer whitespace-split words before classification. "
            "0 disables the filter (default). Recommended: 3."
        ),
    )
    enrich_parser.add_argument(
        "--fasttext-format",
        nargs="?",
        const="allenai/dolma3-fasttext-weborganizer-format-classifier",
        default=None,
        help=(
            "Use FastText for format classification instead of the GPU transformer. "
            "Optionally pass a custom HF repo ID."
        ),
    )
    enrich_parser.add_argument(
        "--fasttext-topic",
        nargs="?",
        const="allenai/dolma3-fasttext-weborganizer-topic-classifier",
        default=None,
        help=(
            "Use FastText for topic classification instead of the GPU transformer. "
            "Optionally pass a custom HF repo ID."
        ),
    )
    return parser


def main(argv: list[str] | None = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)
    configure_logging(args.verbose)

    if args.command == "enrich-format":
        enrich_format(args)
        return 0

    raise ValueError(f"Unknown command: {args.command}")


def enrich_format(args: argparse.Namespace) -> None:
    if args.input == "local":
        if not args.input_files:
            raise ValueError("--input-files is required for local input")
        input_files = expand_inputs(args.input_files)
        compress = (
            None if args.output_compression == "none" else args.output_compression
        )
        for input_path in input_files:
            target_path = output_path(args.output_dir, input_path, compress=compress)
            run_enrichment(
                args,
                iter_local_records(input_path),
                target_path,
                name_hint=str(input_path),
            )
        return

    if args.input == "hf":
        data_files = args.data_files
        output_name = hf_output_name(args.hf_dataset, data_files)
        compress = (
            None if args.output_compression == "none" else args.output_compression
        )
        target_path = output_path(args.output_dir, output_name, compress=compress)
        records = iter_dolma3_hf(
            args.hf_dataset,
            split=args.split,
            data_files=data_files,
            streaming=True,
            cache_dir=args.cache_dir,
        )
        run_enrichment(args, records, target_path, name_hint=output_name)
        return

    raise ValueError(f"Unsupported input mode: {args.input}")


__all__ = ["build_parser", "enrich_format", "main"]


if __name__ == "__main__":
    sys.exit(main())