add: num workers to dataset tools while you're at it

This commit is contained in:
Francesco Capuano
2025-11-06 00:46:11 +00:00
parent bb5676ee5a
commit c87fd37736
+3 -1
View File
@@ -234,6 +234,7 @@ def merge_datasets(
datasets: list[LeRobotDataset],
output_repo_id: str,
output_dir: str | Path | None = None,
num_workers: int | None = None,
) -> LeRobotDataset:
"""Merge multiple LeRobotDatasets into a single dataset.
@@ -257,6 +258,7 @@ def merge_datasets(
aggr_repo_id=output_repo_id,
roots=roots,
aggr_root=output_dir,
num_workers=num_workers,
)
merged_dataset = LeRobotDataset(
@@ -329,7 +331,7 @@ def modify_features(
if repo_id is None:
repo_id = f"{dataset.repo_id}_modified"
output_dir = Path(output_dir) if output_dir is not None else HF_LEROBOT_HOME / repo_id
output_dir = Path(output_dir, exists_ok=True) if output_dir is not None else HF_LEROBOT_HOME / repo_id
new_features = dataset.meta.features.copy()