mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-23 17:56:07 +00:00
migration: add MolmoAct standalone-dataset migration path
Add migrate_molmoact.py to migrate the SO-100/101 datasets listed by allenai/MolmoAct2-SO100_101-Dataset (repo ids derived from language_annotations folder names), skipping any already in lerobot/community_dataset_v3. Generalize run_migration.migrate_one/_write_dataset_card with a backward-compatible standalone mode (download a whole standalone repo, provenance -> source repo) and extend slurm_migrate.py with --source molmoact / --reference-repo.
This commit is contained in:
@@ -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()
|
||||||
@@ -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)
|
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
|
"""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
|
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")
|
enc = result.get("encoding")
|
||||||
converted_degrees = bool(result.get("converted"))
|
converted_degrees = bool(result.get("converted"))
|
||||||
approx = enc == "normalized" and not so_arm_frame.CANON_IS_CALIBRATED
|
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)."
|
+ (" with SO-100/101 joint state/action mapped to the post-#777 physical frame (in degrees)."
|
||||||
if converted_degrees else "."),
|
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",
|
"- Codebase version: v2.1 -> v3.0",
|
||||||
]
|
]
|
||||||
if result.get("is_so"):
|
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)
|
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
|
local = Path(work_dir) / sub
|
||||||
if local.parent.exists():
|
if local.parent.exists():
|
||||||
shutil.rmtree(local.parent, ignore_errors=True) # clean any partial
|
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)
|
info = load_info(local)
|
||||||
if info.get("codebase_version") != "v2.1":
|
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
|
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
|
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
|
base = {k: result.get(k) for k in
|
||||||
("robot_type", "is_so", "encoding", "action_dim", "maxabs", "ambiguous", "action")}
|
("robot_type", "is_so", "encoding", "action_dim", "maxabs", "ambiguous", "action")}
|
||||||
|
|||||||
@@ -53,6 +53,7 @@ class MigrateShard(PipelineStep):
|
|||||||
migration_dir,
|
migration_dir,
|
||||||
no_push=False,
|
no_push=False,
|
||||||
only_classify=False,
|
only_classify=False,
|
||||||
|
standalone=False,
|
||||||
):
|
):
|
||||||
super().__init__()
|
super().__init__()
|
||||||
self.subs = subs
|
self.subs = subs
|
||||||
@@ -62,6 +63,7 @@ class MigrateShard(PipelineStep):
|
|||||||
self.migration_dir = migration_dir
|
self.migration_dir = migration_dir
|
||||||
self.no_push = no_push
|
self.no_push = no_push
|
||||||
self.only_classify = only_classify
|
self.only_classify = only_classify
|
||||||
|
self.standalone = standalone
|
||||||
|
|
||||||
def run(self, data=None, rank: int = 0, world_size: int = 1):
|
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
|
# 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):
|
for i, sub in enumerate(my_subs):
|
||||||
try:
|
try:
|
||||||
if self.only_classify:
|
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)}
|
row = {"root": sub, **classify(Path(work_dir) / sub)}
|
||||||
shutil.rmtree(Path(work_dir) / sub.split("/")[0], ignore_errors=True)
|
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):
|
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"}
|
row = {"root": sub, "action": "skipped: already present in destination repo"}
|
||||||
else:
|
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:
|
except Exception as e:
|
||||||
row = {"root": sub, "action": f"ERROR: {e}"}
|
row = {"root": sub, "action": f"ERROR: {e}"}
|
||||||
traceback.print_exc()
|
traceback.print_exc()
|
||||||
@@ -168,12 +178,20 @@ def main():
|
|||||||
if MIGRATION_DIR not in sys.path:
|
if MIGRATION_DIR not in sys.path:
|
||||||
sys.path.insert(0, MIGRATION_DIR)
|
sys.path.insert(0, MIGRATION_DIR)
|
||||||
from huggingface_hub import HfApi
|
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
|
from run_migration import SRC_REPO, list_datasets, resolve_folders
|
||||||
|
|
||||||
p = argparse.ArgumentParser(
|
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,
|
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("--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("--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.")
|
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.")
|
p.error("--dst-repo is required unless --no-push or --only-classify is set.")
|
||||||
|
|
||||||
api = HfApi()
|
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)
|
subs = resolve_folders(api, SRC_REPO, args.folder_name)
|
||||||
else:
|
else:
|
||||||
subs = list_datasets(api, SRC_REPO)
|
subs = list_datasets(api, SRC_REPO)
|
||||||
@@ -222,6 +251,7 @@ def main():
|
|||||||
MIGRATION_DIR,
|
MIGRATION_DIR,
|
||||||
no_push=args.no_push,
|
no_push=args.no_push,
|
||||||
only_classify=args.only_classify,
|
only_classify=args.only_classify,
|
||||||
|
standalone=standalone,
|
||||||
)
|
)
|
||||||
],
|
],
|
||||||
logs_dir=args.logs_dir,
|
logs_dir=args.logs_dir,
|
||||||
|
|||||||
Reference in New Issue
Block a user