diff --git a/community_v3_migration/extract_dataset.py b/community_v3_migration/extract_dataset.py new file mode 100644 index 000000000..6e17ae6e3 --- /dev/null +++ b/community_v3_migration/extract_dataset.py @@ -0,0 +1,55 @@ +"""Extract one sub-dataset from a LeRobotDataset monorepo into a standalone HF dataset repo. + +Downloads only ``{src-repo}/{folder}/...`` (scoped listing, no whole-repo walk) and re-uploads +its contents at the ROOT of a NEW dataset repo, so the result is a self-contained LeRobotDataset. + + python extract_dataset.py --folder wannrrr/etnai --dst-repo CarolinePascal/etnai +""" +import argparse +import shutil +import sys +from pathlib import Path + +from huggingface_hub import HfApi + +sys.path.insert(0, str(Path(__file__).resolve().parent)) +from run_migration import download_subfolder # noqa: E402 + + +def main(): + ap = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + ap.add_argument("--folder", required=True, metavar="USER/DATASET", + help="Sub-dataset path within the source monorepo, e.g. 'wannrrr/etnai'.") + ap.add_argument("--dst-repo", required=True, metavar="ORG/NAME", + help="New standalone destination dataset repo (must differ from the source).") + ap.add_argument("--src-repo", default="lerobot/community_dataset_v3", metavar="ORG/NAME", + help="Source monorepo to pull the sub-dataset from.") + ap.add_argument("--work-dir", default="./extract_work", help="Local scratch directory.") + ap.add_argument("--private", action="store_true", help="Create the destination repo as private.") + args = ap.parse_args() + + if args.dst_repo in (args.src_repo, args.folder): + ap.error("--dst-repo must be a new repo name, distinct from the source repo/folder.") + + local = Path(args.work_dir) / args.folder + if local.parent.exists(): + shutil.rmtree(local.parent, ignore_errors=True) + + print(f"downloading {args.src_repo}/{args.folder} ...", file=sys.stderr) + download_subfolder(args.folder, args.work_dir, repo=args.src_repo) + if not (local / "meta" / "info.json").exists(): + ap.error(f"'{args.folder}' is not a LeRobotDataset (no meta/info.json) in {args.src_repo}") + + api = HfApi() + api.create_repo(args.dst_repo, repo_type="dataset", private=args.private, exist_ok=True) + print(f"uploading -> {args.dst_repo} ...", file=sys.stderr) + api.upload_folder(repo_id=args.dst_repo, repo_type="dataset", folder_path=str(local), + commit_message=f"Standalone copy of {args.folder} from {args.src_repo}") + shutil.rmtree(Path(args.work_dir) / args.folder.split("/")[0], ignore_errors=True) + print(f"done: https://huggingface.co/datasets/{args.dst_repo}", file=sys.stderr) + + +if __name__ == "__main__": + main() diff --git a/community_v3_migration/run_migration.py b/community_v3_migration/run_migration.py index 370e48a69..754cc3757 100644 --- a/community_v3_migration/run_migration.py +++ b/community_v3_migration/run_migration.py @@ -22,8 +22,8 @@ from fix_dataset import fix_dataset_in_place SRC_REPO = "HuggingFaceVLA/community_dataset_v3" -def download_subfolder(sub: str, work_dir: str, patterns: list[str] | None = None) -> None: - """Download only ``SRC_REPO/{sub}/...`` into ``work_dir``. +def download_subfolder(sub: str, work_dir: str, patterns: list[str] | None = None, repo: str = SRC_REPO) -> None: + """Download only ``{repo}/{sub}/...`` into ``work_dir``. ``snapshot_download`` walks the entire repo tree (``list_repo_tree(recursive=True)`` with no path scope) before applying ``allow_patterns``. On this 791-dataset monorepo @@ -36,12 +36,12 @@ def download_subfolder(sub: str, work_dir: str, patterns: list[str] | None = Non from huggingface_hub.hf_api import RepoFile api = HfApi() - for entry in api.list_repo_tree(SRC_REPO, path_in_repo=sub, repo_type="dataset", recursive=True): + for entry in api.list_repo_tree(repo, path_in_repo=sub, repo_type="dataset", recursive=True): if not isinstance(entry, RepoFile): continue if patterns and not any(fnmatch(entry.path, pat) for pat in patterns): continue - hf_hub_download(SRC_REPO, filename=entry.path, repo_type="dataset", local_dir=work_dir) + hf_hub_download(repo, filename=entry.path, repo_type="dataset", local_dir=work_dir) def list_datasets(api: HfApi, repo: str) -> list[str]: