mirror of
https://github.com/huggingface/lerobot.git
synced 2026-06-17 16:27:04 +00:00
Mahbod
30790de178
* feat(edit-dataset): add `concatenate_videos` opt-out to merge
When merging datasets, source mp4s are concatenated into shards capped at
`video_files_size_in_mb` (default 200 MB). This is great for dataloader
throughput but destroys per-episode (or per-source) video boundaries,
which is undesirable when you want to inspect, ship, or reuse the
individual mp4s.
Add a `concatenate_videos: bool = True` knob plumbed through
`MergeConfig` → `merge_datasets` → `aggregate_datasets` → `aggregate_videos`.
When False, each source mp4 is copied 1:1 to its own destination mp4 with
no re-muxing, so the merge preserves source video boundaries.
Usage:
lerobot-edit-dataset \
--new_repo_id user/merged \
--operation.type=merge \
--operation.repo_ids "['user/a', 'user/b']" \
--operation.concatenate_videos=false
Defaults are unchanged; the dataloader path is unaffected because the
`episodes.parquet` `from_timestamp`/`to_timestamp` index keeps working
regardless of whether each mp4 holds one or many episodes.
* feat(edit-dataset): extend concatenate opt-out to data files
Following review, add a concatenate_data flag mirroring concatenate_videos,
threaded through MergeConfig, merge_datasets, aggregate_datasets, aggregate_data
and append_or_create_parquet_file. Metadata index files still always concatenate.
Also trim the verbose docstrings and comments since the names are
self-explanatory, and extend the existing merge test to cover data files.
106 lines
4.0 KiB
Python
106 lines
4.0 KiB
Python
#!/usr/bin/env python
|
|
|
|
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import draccus
|
|
import pytest
|
|
|
|
pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
|
|
|
|
from lerobot.scripts.lerobot_edit_dataset import (
|
|
ConvertImageToVideoConfig,
|
|
DeleteEpisodesConfig,
|
|
EditDatasetConfig,
|
|
InfoConfig,
|
|
MergeConfig,
|
|
ModifyTasksConfig,
|
|
OperationConfig,
|
|
RemoveFeatureConfig,
|
|
SplitConfig,
|
|
_validate_config,
|
|
)
|
|
|
|
|
|
def parse_cfg(cli_args: list[str]) -> EditDatasetConfig:
|
|
"""Helper to parse CLI args into an EditDatasetConfig via draccus."""
|
|
return draccus.parse(EditDatasetConfig, args=cli_args)
|
|
|
|
|
|
class TestOperationTypeParsing:
|
|
"""Test that --operation.type correctly selects the right config subclass."""
|
|
|
|
@pytest.mark.parametrize(
|
|
"type_name, expected_cls",
|
|
[
|
|
("delete_episodes", DeleteEpisodesConfig),
|
|
("split", SplitConfig),
|
|
("merge", MergeConfig),
|
|
("remove_feature", RemoveFeatureConfig),
|
|
("modify_tasks", ModifyTasksConfig),
|
|
("convert_image_to_video", ConvertImageToVideoConfig),
|
|
("info", InfoConfig),
|
|
],
|
|
)
|
|
def test_operation_type_resolves_correct_class(self, type_name, expected_cls):
|
|
cfg = parse_cfg(
|
|
["--repo_id", "test/repo", "--new_repo_id", "test/merged", "--operation.type", type_name]
|
|
)
|
|
assert isinstance(cfg.operation, expected_cls), (
|
|
f"Expected {expected_cls.__name__}, got {type(cfg.operation).__name__}"
|
|
)
|
|
|
|
def test_merge_requires_new_repo_id(self):
|
|
cfg = parse_cfg(["--operation.type", "merge"])
|
|
with pytest.raises(ValueError, match="--new_repo_id is required for merge"):
|
|
_validate_config(cfg)
|
|
|
|
@pytest.mark.parametrize("flag", ["concatenate_videos", "concatenate_data"])
|
|
def test_merge_concatenate_flag_defaults_true(self, flag):
|
|
cfg = parse_cfg(["--new_repo_id", "test/merged", "--operation.type", "merge"])
|
|
assert isinstance(cfg.operation, MergeConfig)
|
|
assert getattr(cfg.operation, flag) is True
|
|
|
|
@pytest.mark.parametrize("flag", ["concatenate_videos", "concatenate_data"])
|
|
def test_merge_concatenate_flag_can_be_disabled(self, flag):
|
|
cfg = parse_cfg(
|
|
["--new_repo_id", "test/merged", "--operation.type", "merge", f"--operation.{flag}", "false"]
|
|
)
|
|
assert isinstance(cfg.operation, MergeConfig)
|
|
assert getattr(cfg.operation, flag) is False
|
|
|
|
def test_non_merge_requires_repo_id(self):
|
|
cfg = parse_cfg(["--operation.type", "delete_episodes"])
|
|
with pytest.raises(ValueError, match="--repo_id is required for delete_episodes"):
|
|
_validate_config(cfg)
|
|
|
|
@pytest.mark.parametrize(
|
|
"type_name, expected_cls",
|
|
[
|
|
("delete_episodes", DeleteEpisodesConfig),
|
|
("split", SplitConfig),
|
|
("merge", MergeConfig),
|
|
("remove_feature", RemoveFeatureConfig),
|
|
("modify_tasks", ModifyTasksConfig),
|
|
("convert_image_to_video", ConvertImageToVideoConfig),
|
|
("info", InfoConfig),
|
|
],
|
|
)
|
|
def test_get_choice_name_roundtrips(self, type_name, expected_cls):
|
|
cfg = parse_cfg(
|
|
["--repo_id", "test/repo", "--new_repo_id", "test/merged", "--operation.type", type_name]
|
|
)
|
|
resolved_name = OperationConfig.get_choice_name(type(cfg.operation))
|
|
assert resolved_name == type_name
|