💥 Add generic converter adapter hooks (#107)

* Add generic converter adapter hooks

Co-authored-by: Codex <codex@openai.com>

* Require conversion task repo ids

Co-authored-by: Codex <codex@openai.com>

* Remove conversion task runtime repo id check

Co-authored-by: Codex <codex@openai.com>

* Apply suggestion from @gemini-code-assist[bot]

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

---------

Co-authored-by: Codex <codex@openai.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

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
+such as language instructions. By default, each yielded episode must be a
-dictionaries accepted by `LeRobotDataset.add_frame`; each frame should include
+sequence of frame dictionaries accepted by `LeRobotDataset.add_frame`; each
-the LeRobot `task` field when language tasks are needed.
+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
 
@@ -63,7 +71,7 @@ the LeRobot `task` field when language tasks are needed.
 
 - `input_path`: source file or directory
 - `output_path`: temporary LeRobot dataset directory for this task
-- `local_repo_id`: repo id used while writing the temporary dataset
+- `local_repo_id`: required repo id used while writing the temporary dataset
 - `metadata`: adapter-owned metadata
 
 Keep dataset-specific values in `metadata`; the generic pipeline does not know

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(
+    def load_subset(self, task: ConversionTask) -> Iterable[Any]:
-        self, task: ConversionTask
-    ) -> Iterable[Sequence[dict]]:
         """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: Any,
+        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)

generic_converter/pipeline.py

@@ -33,13 +33,7 @@ class SaveLeRobotDataset(PipelineStep):
         if task.output_path.exists():
             shutil.rmtree(task.output_path)
 
-        dataset = LeRobotDataset.create(
+        dataset = self.adapter.create_dataset(task)
-            repo_id=task.local_repo_id,
-            root=task.output_path,
-            fps=self.adapter.fps,
-            robot_type=self.adapter.robot_type,
-            features=self.adapter.features,
-        )
 
         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:
+                saved = self.adapter.save_episode(
-                    dataset.add_frame(frame)
+                    dataset,
-                dataset.save_episode()
+                    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, {
+            executor_cls, executor_config = (
-                "tasks": len(tasks),
+                LocalPipelineExecutor,
-                "workers": resolved_workers,
+                {
-            }
+                    "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, {
+            executor_cls, executor_config = (
-                "tasks": len(tasks),
+                RayPipelineExecutor,
-                "workers": workers,
+                {
-                "cpus_per_task": cpus_per_task,
+                    "tasks": len(tasks),
-                "tasks_per_job": tasks_per_job,
+                    "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,

generic_converter/utils.py

@@ -1,4 +1,3 @@
-import shutil
 from collections.abc import Mapping, Sequence
 from dataclasses import dataclass, field
 from pathlib import Path
@@ -14,7 +13,7 @@ class ConversionTask:
 
     input_path: Path
     output_path: Path
-    local_repo_id: str | None = None
+    local_repo_id: str
     metadata: TaskMetadata = field(default_factory=dict)