From 65296e75cbf68171c6e8ed52db0def106c293204 Mon Sep 17 00:00:00 2001 From: Tavish Date: Sun, 21 Jun 2026 15:44:14 +0800 Subject: [PATCH] Add generic converter adapter hooks Co-authored-by: Codex --- generic_converter/README.md | 16 +++++++++--- generic_converter/adapter.py | 36 +++++++++++++++++++++++--- generic_converter/pipeline.py | 48 +++++++++++++++++++---------------- 3 files changed, 71 insertions(+), 29 deletions(-) diff --git a/generic_converter/README.md b/generic_converter/README.md index 6614df2..b34744c 100644 --- a/generic_converter/README.md +++ b/generic_converter/README.md @@ -43,7 +43,13 @@ Optional attributes: Required methods: - `load_tasks(self) -> list[ConversionTask]` -- `load_subset(self, task: ConversionTask) -> Iterable[Sequence[dict]]` +- `load_subset(self, task: ConversionTask) -> Iterable[Any]` + +Optional hooks: + +- `create_dataset(self, task: ConversionTask)` +- `save_episode(self, dataset, episode_data, task) -> bool` +- `get_episode_length(self, episode_data) -> int` `run_converter` reads `adapter.output_path` and calls `adapter.load_tasks()` without arguments. Store paths, task manifests, or other adapter options on the @@ -53,9 +59,11 @@ Use `adapter.temp_output_path` when building task-level temporary output paths. `load_subset` receives the full `ConversionTask`, not just an input path. Use `task.input_path` for raw data and `task.metadata` for dataset-specific values -such as language instructions. Each yielded episode must be a sequence of frame -dictionaries accepted by `LeRobotDataset.add_frame`; each frame should include -the LeRobot `task` field when language tasks are needed. +such as language instructions. By default, each yielded episode must be a +sequence of frame dictionaries accepted by `LeRobotDataset.add_frame`; each +frame should include the LeRobot `task` field when language tasks are needed. +Adapters that need custom dataset classes or extra per-episode arguments can +override the optional hooks. ## ConversionTask diff --git a/generic_converter/adapter.py b/generic_converter/adapter.py index be58609..4bdd7d6 100644 --- a/generic_converter/adapter.py +++ b/generic_converter/adapter.py @@ -1,6 +1,7 @@ from abc import ABC, abstractmethod from collections.abc import Iterable, Sequence from pathlib import Path +from typing import Any from .utils import ConversionTask, FeatureSpec @@ -26,7 +27,36 @@ class BaseAdapter(ABC): """Build conversion tasks from dataset-specific inputs.""" @abstractmethod - def load_subset( - self, task: ConversionTask - ) -> Iterable[Sequence[dict]]: + def load_subset(self, task: ConversionTask) -> Iterable[Any]: """Yield LeRobot episodes for one raw input path.""" + + def create_dataset(self, task: ConversionTask): + """Create the temporary LeRobot dataset for one conversion task.""" + from lerobot.datasets import LeRobotDataset + + return LeRobotDataset.create( + repo_id=task.local_repo_id, + root=task.output_path, + fps=self.fps, + robot_type=self.robot_type, + features=self.features, + ) + + def save_episode( + self, + dataset: Any, + episode_data: Sequence[dict], + task: ConversionTask, + ) -> bool: + """Save one episode to the temporary dataset. + + Adapters can override this when a dataset needs extra per-episode + arguments or a non-standard writer. + """ + for frame in episode_data: + dataset.add_frame(frame) + dataset.save_episode() + return True + + def get_episode_length(self, episode_data: Any) -> int: + return len(episode_data) diff --git a/generic_converter/pipeline.py b/generic_converter/pipeline.py index b4963bc..9154b27 100644 --- a/generic_converter/pipeline.py +++ b/generic_converter/pipeline.py @@ -33,13 +33,7 @@ class SaveLeRobotDataset(PipelineStep): if task.output_path.exists(): shutil.rmtree(task.output_path) - dataset = LeRobotDataset.create( - repo_id=task.local_repo_id, - root=task.output_path, - fps=self.adapter.fps, - robot_type=self.adapter.robot_type, - features=self.adapter.features, - ) + dataset = self.adapter.create_dataset(task) logger.info( f"start processing for {task.input_path}, saving to {task.output_path}" @@ -47,11 +41,15 @@ class SaveLeRobotDataset(PipelineStep): raw_dataset = self.adapter.load_subset(task) for episode_index, episode_data in enumerate(raw_dataset): with self.track_time("saving episode"): - for frame in episode_data: - dataset.add_frame(frame) - dataset.save_episode() + saved = self.adapter.save_episode( + dataset, + episode_data, + task, + ) + status = "skipped" if saved is False else "process done" logger.info( - f"process done for {dataset.repo_id}, episode {episode_index}, len {len(episode_data)}" + f"{status} for {dataset.repo_id}, episode {episode_index}, " + f"len {self.adapter.get_episode_length(episode_data)}" ) dataset.finalize() @@ -97,10 +95,13 @@ def run_converter( if workers == -1 else workers ) - executor_cls, executor_config = LocalPipelineExecutor, { - "tasks": len(tasks), - "workers": resolved_workers, - } + executor_cls, executor_config = ( + LocalPipelineExecutor, + { + "tasks": len(tasks), + "workers": resolved_workers, + }, + ) case "ray": import ray from datatrove.executor import RayPipelineExecutor @@ -108,12 +109,15 @@ def run_converter( runtime_env = RuntimeEnv(env_vars=_build_ray_env_vars()) ray.init(runtime_env=runtime_env) - executor_cls, executor_config = RayPipelineExecutor, { - "tasks": len(tasks), - "workers": workers, - "cpus_per_task": cpus_per_task, - "tasks_per_job": tasks_per_job, - } + executor_cls, executor_config = ( + RayPipelineExecutor, + { + "tasks": len(tasks), + "workers": workers, + "cpus_per_task": cpus_per_task, + "tasks_per_job": tasks_per_job, + }, + ) case _: raise ValueError(f"Executor {executor} not supported") @@ -121,7 +125,7 @@ def run_converter( logging_dir = str(resume_dir) else: logging_dir = str(Path.cwd() / "logs" / f"{get_timestamp()}_{get_random_str()}") - + executor_cls( pipeline=[SaveLeRobotDataset(tasks, adapter)], **executor_config,