social-data-attribution-w2/scripts/query_data/upload_query_data.py

"""Upload OLMES query JSONL manifests to HF as raw files.

Generates a dataset card README.md so the HF viewer can render the data.

Usage:
    python scripts/upload_query_data.py --variant base
    python scripts/upload_query_data.py --variant instruct
    python scripts/upload_query_data.py --variant instruct_cot
    python scripts/upload_query_data.py --variant all
"""

import argparse
import io
import logging
from pathlib import Path

from huggingface_hub import HfApi

from data_attribution.hf_hub import get_hf_token, load_env_secret

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)s %(message)s",
)
log = logging.getLogger(__name__)

SUBSETS = (
    "gsm8k",
    "mmlu_social_science",
    "mmlu_stem",
    "socialiqa",
    "arc_easy",
    "arc_challenge",
    "bbh_snarks",
    "bbh_causal_judgement",
    "bbh_sports_understanding",
)

VARIANTS = {
    "base": {
        "repo_id": "HCAI-Lab/base-query-data",
        "file_prefix": "olmes_",
    },
    "instruct": {
        "repo_id": "HCAI-Lab/instruct-query-data",
        "file_prefix": "olmes_instruct_",
    },
    "instruct_cot": {
        "repo_id": "HCAI-Lab/instruct-cot-query-data",
        "file_prefix": "olmes_instruct_cot_",
    },
}


def _build_readme(file_prefix: str) -> str:
    lines = ["---", "configs:"]
    for subset in SUBSETS:
        filename = f"{file_prefix}{subset}.jsonl"
        lines.append(f"- config_name: {subset}")
        lines.append(f"  data_files: {filename}")
    lines.append("---")
    lines.append("")
    return "\n".join(lines)


def _upload_variant(
    api: HfApi,
    token: str,
    input_dir: Path,
    variant: str,
    private: bool,
) -> None:
    cfg = VARIANTS[variant]
    repo_id = cfg["repo_id"]
    file_prefix = cfg["file_prefix"]

    api.create_repo(
        repo_id=repo_id,
        repo_type="dataset",
        private=private,
        exist_ok=True,
        token=token,
    )
    log.info("Target repo: %s", repo_id)

    existing = set(api.list_repo_files(repo_id, repo_type="dataset", token=token))
    parquets = [f for f in existing if f.endswith(".parquet")]
    if parquets:
        from huggingface_hub import CommitOperationDelete

        ops = [CommitOperationDelete(path_in_repo=p) for p in parquets]
        api.create_commit(
            repo_id=repo_id,
            repo_type="dataset",
            operations=ops,
            commit_message="Remove parquet files",
            token=token,
        )
        log.info("Deleted %d parquet files", len(parquets))

    for subset in SUBSETS:
        jsonl_path = input_dir / f"{file_prefix}{subset}.jsonl"
        if not jsonl_path.exists():
            log.error("Missing %s", jsonl_path)
            raise SystemExit(1)

        log.info("Uploading %s -> %s", jsonl_path.name, repo_id)
        api.upload_file(
            path_or_fileobj=str(jsonl_path),
            path_in_repo=jsonl_path.name,
            repo_id=repo_id,
            repo_type="dataset",
            commit_message=f"Upload {jsonl_path.name}",
            token=token,
        )
        log.info("Uploaded %s", jsonl_path.name)

    readme = _build_readme(file_prefix)
    api.upload_file(
        path_or_fileobj=io.BytesIO(readme.encode("utf-8")),
        path_in_repo="README.md",
        repo_id=repo_id,
        repo_type="dataset",
        commit_message="Add dataset card",
        token=token,
    )
    log.info("Uploaded README.md to %s", repo_id)


def main() -> None:
    parser = argparse.ArgumentParser(description="Upload query manifests to HF")
    parser.add_argument("--input-dir", type=Path, default=Path("runs/manifests"))
    parser.add_argument(
        "--variant",
        choices=list(VARIANTS.keys()) + ["all"],
        required=True,
        help="Which variant to upload (base, instruct, instruct_cot, or all)",
    )
    parser.add_argument("--private", action="store_true", default=True)
    args = parser.parse_args()

    load_env_secret()
    token = get_hf_token()
    api = HfApi()

    variants = list(VARIANTS.keys()) if args.variant == "all" else [args.variant]

    for variant in variants:
        log.info("Uploading variant: %s", variant)
        _upload_variant(api, token, args.input_dir, variant, args.private)

    log.info("Done")


if __name__ == "__main__":
    main()