social-data-attribution-w2/scripts/sampling/upload_stratified_docs.py

"""Upload combined merged docs to HF, skipping already-uploaded files.

Deletes any non-merged files (per-worker artifacts) from the repo,
then uploads the combined job_N.jsonl.zst files that are not yet present.

Usage:
    python scripts/sampling/upload_stratified_docs.py \
        --input-dir stratified_data/merged_docs \
        --repo-id HCAI-Lab/archive-dolma3-pool-stratified
"""

import argparse
import logging
import re
from pathlib import Path

from huggingface_hub import HfApi

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

MERGED_PATTERN = re.compile(r"^job_\d\.jsonl\.zst$")


def main() -> None:
    parser = argparse.ArgumentParser(description="Upload merged docs to HF dataset")
    parser.add_argument("--input-dir", type=Path, required=True)
    parser.add_argument(
        "--repo-id", type=str, default="HCAI-Lab/archive-dolma3-pool-stratified"
    )
    args = parser.parse_args()

    api = HfApi()
    files = sorted(args.input_dir.glob("job_[0-9].jsonl.zst"))

    missing = [f for f in files if not f.exists()]
    if missing:
        log.error("Missing files: %s", missing)
        raise SystemExit(1)

    log.info("Target repo: %s", args.repo_id)
    log.info("Local files: %s", [f.name for f in files])

    existing = set(api.list_repo_files(args.repo_id, repo_type="dataset"))
    deletions = [
        f for f in existing if not f.startswith(".") and not MERGED_PATTERN.match(f)
    ]
    if deletions:
        from huggingface_hub import CommitOperationDelete

        ops = [CommitOperationDelete(path_in_repo=p) for p in deletions]
        api.create_commit(
            repo_id=args.repo_id,
            repo_type="dataset",
            operations=ops,
            commit_message="Remove per-worker files",
        )
        log.info("Deleted %d non-merged files", len(deletions))

    to_upload = [f for f in files if f.name not in existing]
    if not to_upload:
        log.info("All files already uploaded")
    else:
        log.info("Uploading %d files: %s", len(to_upload), [f.name for f in to_upload])

    for f in to_upload:
        log.info("Uploading %s...", f.name)
        api.upload_file(
            path_or_fileobj=str(f),
            path_in_repo=f.name,
            repo_id=args.repo_id,
            repo_type="dataset",
            commit_message=f"Upload {f.name}",
        )
        log.info("Uploaded %s", f.name)

    log.info("Done")


if __name__ == "__main__":
    main()