social-data-attribution-w2/src/data_attribution/cli/parser.py

"""Parser builders for the data attribution CLI."""

from __future__ import annotations

import argparse
from pathlib import Path
from typing import Iterable

from data_attribution.cli.handlers import _attribution_handler, _log_only_handler
from data_attribution.cli.ingest import register_ingest_subcommand


def _register_run_subcommand(
    subparsers: argparse._SubParsersAction[argparse.ArgumentParser],
) -> None:
    parser = subparsers.add_parser("run", help="Run attribution from saved gradients")
    parser.add_argument("--query-gradients", type=Path, required=True)
    parser.add_argument("--index-path", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path)
    parser.add_argument("--metadata-path", type=Path)
    parser.add_argument("--query-id-field", default="query_id")
    parser.add_argument("--doc-id-field", default="doc_id")
    parser.add_argument("--metadata-join-keys", nargs="+", action="extend")
    parser.add_argument(
        "--output-format", choices=("jsonl", "parquet"), default="parquet"
    )
    parser.add_argument("--top-k", type=int, default=10)
    parser.add_argument("--device")
    parser.add_argument(
        "--unit-norm", action=argparse.BooleanOptionalAction, default=True
    )
    parser.set_defaults(func=_attribution_handler)


def _register_subcommands(
    subparsers: argparse._SubParsersAction[argparse.ArgumentParser],
) -> None:
    register_ingest_subcommand(subparsers)
    _register_run_subcommand(subparsers)
    commands: Iterable[str] = (
        "index",
        "queries",
        "attribution",
        "analysis",
    )

    for name in commands:
        parser = subparsers.add_parser(name, help=f"Run the {name} workflow")
        parser.set_defaults(func=_log_only_handler)


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Data attribution CLI")
    parser.add_argument(
        "--config",
        type=Path,
        help="Path to YAML or TOML configuration file",
    )
    parser.add_argument(
        "--run-id",
        dest="run_id",
        help="Override autogenerated run identifier",
    )
    parser.add_argument(
        "--verbose",
        action="store_true",
        help="Enable debug logging",
    )

    subparsers = parser.add_subparsers(dest="command", required=True)
    _register_subcommands(subparsers)
    return parser


__all__ = ["build_parser"]