diff --git a/community_v3_migration/migrate_molmoact.py b/community_v3_migration/migrate_molmoact.py new file mode 100644 index 000000000..2c2a4e50a --- /dev/null +++ b/community_v3_migration/migrate_molmoact.py @@ -0,0 +1,133 @@ +"""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() diff --git a/community_v3_migration/run_migration.py b/community_v3_migration/run_migration.py index 30884eeb4..298dc2a21 100644 --- a/community_v3_migration/run_migration.py +++ b/community_v3_migration/run_migration.py @@ -80,10 +80,11 @@ def already_done(api: HfApi, dst: str, sub: str, dst_files: set[str]) -> bool: return f"{sub}/meta/info.json" in dst_files # present in target => skip (resume) -def _write_dataset_card(local: Path, sub: str, result: dict) -> None: +def _write_dataset_card(local: Path, sub: str, result: dict, standalone: bool = False) -> None: """Regenerate the sub-dataset's card the way LeRobot does (create_lerobot_dataset_card from meta/info.json), then append a migration section documenting provenance and the - joint-encoding fix.""" + joint-encoding fix. When ``standalone`` is set, ``sub`` is itself the source dataset's HF + repo id (rather than a folder inside the ``SRC_REPO`` monorepo).""" enc = result.get("encoding") converted_degrees = bool(result.get("converted")) approx = enc == "normalized" and not so_arm_frame.CANON_IS_CALIBRATED @@ -112,7 +113,8 @@ def _write_dataset_card(local: Path, sub: str, result: dict) -> None: + (" with SO-100/101 joint state/action mapped to the post-#777 physical frame (in degrees)." if converted_degrees else "."), "", - f"- Source: [`{SRC_REPO}`](https://huggingface.co/datasets/{SRC_REPO}/tree/main/{sub}) (`{sub}`)", + (f"- Source: [`{sub}`](https://huggingface.co/datasets/{sub})" if standalone else + f"- Source: [`{SRC_REPO}`](https://huggingface.co/datasets/{SRC_REPO}/tree/main/{sub}) (`{sub}`)"), "- Codebase version: v2.1 -> v3.0", ] if result.get("is_so"): @@ -165,11 +167,17 @@ def _write_dataset_card(local: Path, sub: str, result: dict) -> None: readme.write_text(f"# {sub}\n\n" + section) -def migrate_one(api, dst_repo, sub, work_dir, no_upload) -> dict: +def migrate_one(api, dst_repo, sub, work_dir, no_upload, src_repo: str = SRC_REPO, + standalone: bool = False) -> dict: local = Path(work_dir) / sub if local.parent.exists(): shutil.rmtree(local.parent, ignore_errors=True) # clean any partial - download_subfolder(sub, work_dir) + if standalone: + # ``sub`` is a self-contained HF dataset repo (not a monorepo folder): pull it whole. + from huggingface_hub import snapshot_download + snapshot_download(repo_id=sub, repo_type="dataset", local_dir=str(local)) + else: + download_subfolder(sub, work_dir, repo=src_repo) info = load_info(local) if info.get("codebase_version") != "v2.1": @@ -195,7 +203,7 @@ def migrate_one(api, dst_repo, sub, work_dir, no_upload) -> dict: from lerobot.scripts.convert_dataset_v21_to_v30 import convert_dataset convert_dataset(repo_id=sub, root=str(local), push_to_hub=False) # v2.1 -> v3.0, in place - _write_dataset_card(local, sub, result) # document the conversion in the dataset card + _write_dataset_card(local, sub, result, standalone=standalone) # document conversion in the card base = {k: result.get(k) for k in ("robot_type", "is_so", "encoding", "action_dim", "maxabs", "ambiguous", "action")} diff --git a/community_v3_migration/slurm_migrate.py b/community_v3_migration/slurm_migrate.py index 7047e4ad8..4be264060 100644 --- a/community_v3_migration/slurm_migrate.py +++ b/community_v3_migration/slurm_migrate.py @@ -53,6 +53,7 @@ class MigrateShard(PipelineStep): migration_dir, no_push=False, only_classify=False, + standalone=False, ): super().__init__() self.subs = subs @@ -62,6 +63,7 @@ class MigrateShard(PipelineStep): self.migration_dir = migration_dir self.no_push = no_push self.only_classify = only_classify + self.standalone = standalone def run(self, data=None, rank: int = 0, world_size: int = 1): # Pickled onto the worker: keep self-contained. The migration package dir must be on @@ -116,13 +118,21 @@ class MigrateShard(PipelineStep): for i, sub in enumerate(my_subs): try: if self.only_classify: - download_subfolder(sub, work_dir, patterns=[f"{sub}/meta/*"]) + if self.standalone: + from huggingface_hub import snapshot_download + + snapshot_download(repo_id=sub, repo_type="dataset", + local_dir=str(Path(work_dir) / sub), + allow_patterns=["meta/*"]) + else: + download_subfolder(sub, work_dir, patterns=[f"{sub}/meta/*"]) row = {"root": sub, **classify(Path(work_dir) / sub)} shutil.rmtree(Path(work_dir) / sub.split("/")[0], ignore_errors=True) elif not self.no_push and already_done(api, self.dst_repo, sub, dst_files): row = {"root": sub, "action": "skipped: already present in destination repo"} else: - row = migrate_one(api, self.dst_repo, sub, work_dir, self.no_push) + row = migrate_one(api, self.dst_repo, sub, work_dir, self.no_push, + standalone=self.standalone) except Exception as e: row = {"root": sub, "action": f"ERROR: {e}"} traceback.print_exc() @@ -168,12 +178,20 @@ def main(): if MIGRATION_DIR not in sys.path: sys.path.insert(0, MIGRATION_DIR) from huggingface_hub import HfApi + from migrate_molmoact import REFERENCE_REPO, list_molmoact_datasets, pending_datasets from run_migration import SRC_REPO, list_datasets, resolve_folders p = argparse.ArgumentParser( - description="SLURM-distributed community_dataset_v3 -> v3.0 migration (map-only).", + description="SLURM-distributed migration to LeRobotDataset v3.0 (map-only). Source is " + "either the community_dataset_v3 monorepo or the standalone datasets listed " + "by allenai/MolmoAct2-SO100_101-Dataset.", formatter_class=argparse.RawDescriptionHelpFormatter, ) + p.add_argument("--source", choices=("monorepo", "molmoact"), default="monorepo", + help="'monorepo': HuggingFaceVLA/community_dataset_v3 subfolders. " + "'molmoact': standalone datasets listed by allenai/MolmoAct2-SO100_101-Dataset.") + p.add_argument("--reference-repo", default=REFERENCE_REPO, metavar="ORG/NAME", + help="(--source molmoact) Repo whose datasets are already migrated and skipped.") p.add_argument("--dst-repo", default=None, metavar="ORG/NAME", help="Destination HF dataset repo.") p.add_argument("--work-dir", default="./cdv3_work", help="Scratch root; each rank gets a subdir.") p.add_argument("--manifest-dir", default="./cdv3_manifests", help="Per-rank manifest CSVs land here.") @@ -198,7 +216,18 @@ def main(): p.error("--dst-repo is required unless --no-push or --only-classify is set.") api = HfApi() - if args.folder_name: + standalone = args.source == "molmoact" + if standalone: + 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] + else: + subs = pending_datasets(api, all_ids, args.dst_repo, args.reference_repo, + args.no_push, args.only_classify) + if args.limit: + subs = subs[: args.limit] + elif args.folder_name: subs = resolve_folders(api, SRC_REPO, args.folder_name) else: subs = list_datasets(api, SRC_REPO) @@ -222,6 +251,7 @@ def main(): MIGRATION_DIR, no_push=args.no_push, only_classify=args.only_classify, + standalone=standalone, ) ], logs_dir=args.logs_dir,