lerobot/community_v3_migration/migrate_molmoact.py

"""Migrate the SO-100/101 datasets referenced by ``allenai/MolmoAct2-SO100_101-Dataset``.

That repo does NOT store the datasets themselves; it lists them. Each
``language_annotations/{user}/{dataset}/...`` folder name is the HF repo id of a *standalone*
LeRobotDataset. This script derives those repo ids, drops any already present in
``lerobot/community_dataset_v3`` (already migrated) and in the destination (resume), then runs
the exact same per-dataset pipeline as ``run_migration.py`` on each remaining standalone repo
(download whole repo -> SO-arm joint fix -> v2.1->v3.0 convert -> card -> upload -> cleanup).

  python migrate_molmoact.py --dst-repo lerobot/community_dataset_v3 --work-dir ./molmo_work
Flags mirror run_migration.py: --only-classify, --no-push, --folder-name USER/DATASET [...],
--limit N, --reference-repo (the "already migrated" set to skip against).
"""
import argparse
import csv
import shutil
import sys
import traceback
from pathlib import Path

from huggingface_hub import HfApi

sys.path.insert(0, str(Path(__file__).resolve().parent))
from classify import classify  # noqa: E402
from run_migration import already_done, list_datasets, migrate_one  # noqa: E402

LIST_REPO = "allenai/MolmoAct2-SO100_101-Dataset"
REFERENCE_REPO = "lerobot/community_dataset_v3"  # the "already migrated" set to skip against
ANNOTATIONS_PREFIX = "language_annotations/"


def list_molmoact_datasets(api: HfApi, repo: str = LIST_REPO) -> list[str]:
    """Standalone dataset repo ids (``{user}/{dataset}``) derived from the folder names under
    ``language_annotations/`` in the MolmoAct listing repo."""
    files = api.list_repo_files(repo, repo_type="dataset")
    return sorted({"/".join(f.split("/")[1:3]) for f in files
                   if f.startswith(ANNOTATIONS_PREFIX) and len(f.split("/")) >= 3})


def pending_datasets(api: HfApi, subs: list[str], dst_repo: str | None,
                     reference_repo: str, no_upload: bool, only_classify: bool) -> list[str]:
    """Drop ids already in the reference repo (already migrated) and, unless classify/no-push,
    ids already in the destination repo (resume)."""
    skip = set(list_datasets(api, reference_repo))
    if not only_classify and not no_upload and dst_repo and dst_repo != reference_repo:
        skip |= {p[: -len("/meta/info.json")] for p in api.list_repo_files(dst_repo, repo_type="dataset")
                 if p.endswith("/meta/info.json")}
    return [s for s in subs if s not in skip]


def main():
    ap = argparse.ArgumentParser(
        description="Migrate the standalone SO-100/101 datasets listed by "
                    f"{LIST_REPO} to LeRobotDataset v3.0 (degrees), skipping any already present "
                    f"in --reference-repo. One dataset at a time (download -> fix -> convert -> "
                    "upload -> cleanup); resumable.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    ap.add_argument("--dst-repo", default=REFERENCE_REPO, metavar="ORG/NAME",
                    help="Destination HF dataset repo to push the converted v3.0 datasets to "
                         "(created if missing).")
    ap.add_argument("--reference-repo", default=REFERENCE_REPO, metavar="ORG/NAME",
                    help="Repo whose datasets are considered already migrated and skipped.")
    ap.add_argument("--work-dir", default="./molmo_work", metavar="DIR",
                    help="Local scratch directory (one dataset lives here at a time on a push run).")
    ap.add_argument("--manifest", default="manifest_molmoact.csv", metavar="CSV",
                    help="CSV log appended to as datasets are processed. Reused across resumed runs.")
    ap.add_argument("--limit", type=int, default=None, metavar="N",
                    help="Process only the first N pending datasets (alphabetical). Ignored with "
                         "--folder-name.")
    ap.add_argument("--folder-name", nargs="+", default=None, metavar="USER/DATASET",
                    help="One or more specific standalone repo ids to process (must appear in the "
                         f"{LIST_REPO} listing).")
    ap.add_argument("--only-classify", action="store_true",
                    help="Detect robot type + joint encoding and write the manifest only; no "
                         "download of data, convert, or push.")
    ap.add_argument("--no-push", action="store_true",
                    help="Fix + convert locally but do NOT upload; output kept under --work-dir.")
    args = ap.parse_args()
    no_upload = args.no_push

    api = HfApi()
    all_ids = list_molmoact_datasets(api)
    if args.folder_name:
        wanted = {n.strip("/") for n in args.folder_name}
        subs = [s for s in all_ids if s in wanted]
        missing = wanted - set(subs)
        if missing:
            print(f"warning: not in {LIST_REPO} listing: {', '.join(sorted(missing))}", file=sys.stderr)
    else:
        subs = pending_datasets(api, all_ids, args.dst_repo, args.reference_repo, no_upload, args.only_classify)
        if args.limit:
            subs = subs[: args.limit]
    print(f"{len(subs)} dataset(s) to process (of {len(all_ids)} listed)", file=sys.stderr)
    if not subs:
        return

    if not args.only_classify and not no_upload:
        api.create_repo(args.dst_repo, repo_type="dataset", exist_ok=True)
    dst_files = set() if (args.only_classify or no_upload) else set(
        api.list_repo_files(args.dst_repo, repo_type="dataset"))

    first = not Path(args.manifest).exists()
    with open(args.manifest, "a", newline="") as mf:
        w = None
        for i, sub in enumerate(subs):
            try:
                if args.only_classify:
                    from huggingface_hub import snapshot_download
                    local = Path(args.work_dir) / sub
                    snapshot_download(repo_id=sub, repo_type="dataset", local_dir=str(local),
                                      allow_patterns=["meta/*"])
                    row = {"root": sub, **classify(local)}
                    shutil.rmtree(Path(args.work_dir) / sub.split("/")[0], ignore_errors=True)
                elif not no_upload and already_done(api, args.dst_repo, sub, dst_files):
                    row = {"root": sub, "action": "skipped: already present in destination repo"}
                else:
                    row = migrate_one(api, args.dst_repo, sub, args.work_dir, no_upload, standalone=True)
            except Exception as e:
                row = {"root": sub, "action": f"ERROR: {e}"}
                traceback.print_exc()
            if w is None:
                w = csv.DictWriter(mf, fieldnames=sorted(
                    {"root", "robot_type", "is_so", "encoding", "action_dim",
                     "maxabs", "ambiguous", "action", "codebase_version", "note"}))
                if first:
                    w.writeheader()
            w.writerow({k: row.get(k) for k in w.fieldnames})
            mf.flush()
            print(f"[{i+1}/{len(subs)}] {sub}: {row.get('action')}", file=sys.stderr)


if __name__ == "__main__":
    main()