feat: use pipefunc.helpers.collect_kwargs (#30)

yaak-ai · Dec 31, 2024 · 2f8e584 · 2f8e584
1 parent a61a037
commit 2f8e584
Show file tree

Hide file tree

Showing 7 changed files with 19 additions and 69 deletions.
diff --git a/config/_templates/dataset/yaak.yaml b/config/_templates/dataset/yaak.yaml
@@ -35,18 +35,17 @@ inputs:
           cache_dir: /tmp/rbyte-cache
         functions:
           - _target_: pipefunc.PipeFunc
-            scope: meta
-            output_name: data
+            scope: yaak_metadata
+            output_name: output
             cache: true
             func:
               _target_: hydra.utils.get_method
               path: rbyte.io.build_yaak_metadata_dataframe
 
           - _target_: pipefunc.PipeFunc
-            scope: mcap
             bound:
               path: ${data_dir}/(@=input_id@)/ai.mcap
-            output_name: data
+            output_name: mcap
             func:
               _target_: rbyte.io.McapDataFrameBuilder
               decoder_factories: [rbyte.utils._mcap.ProtobufDecoderFactory]
@@ -61,14 +60,10 @@ inputs:
 
           - _target_: pipefunc.PipeFunc
             func:
-              _target_: hydra.utils.get_method
-              path: rbyte.utils.make_dict
-            bound:
-              k0: meta
-              k1: mcap
+              _target_: pipefunc.helpers.collect_kwargs
+              parameters: [meta, mcap]
             renames:
-              v0: meta.data
-              v1: mcap.data
+              meta: yaak_metadata.output
             output_name: data
 
           - _target_: pipefunc.PipeFunc
@@ -137,7 +132,7 @@ inputs:
               length: 6
 
       kwargs:
-        meta:
+        yaak_metadata:
           path: ${data_dir}/(@=input_id@)/metadata.log
           fields:
             rbyte.io.yaak.proto.sensor_pb2.ImageMetadata:

diff --git a/config/_templates/dataset/zod.yaml b/config/_templates/dataset/zod.yaml
@@ -32,7 +32,7 @@ inputs:
           - _target_: pipefunc.PipeFunc
             bound:
               path: "${data_dir}/sequences/000002_short/camera_front_blur/000002_romeo_{timestamp:%Y-%m-%dT%H:%M:%S.%f}Z.jpg"
-            output_name: camera_front_blur_data
+            output_name: camera_front_blur
             func:
               _target_: rbyte.io.PathDataFrameBuilder
               fields:
@@ -43,7 +43,7 @@ inputs:
           - _target_: pipefunc.PipeFunc
             bound:
               path: "${data_dir}/sequences/000002_short/lidar_velodyne/000002_romeo_{timestamp:%Y-%m-%dT%H:%M:%S.%f}Z.npy"
-            output_name: lidar_velodyne_data
+            output_name: lidar_velodyne
             func:
               _target_: rbyte.io.PathDataFrameBuilder
               fields:
@@ -74,18 +74,10 @@ inputs:
                   speed/meters_per_second/value:
 
           - _target_: pipefunc.PipeFunc
-            bound:
-              k0: camera_front_blur
-              k1: lidar_velodyne
-              k2: vehicle_data
-            renames:
-              v0: camera_front_blur_data
-              v1: lidar_velodyne_data
-              v2: vehicle_data
-            output_name: data
             func:
-              _target_: hydra.utils.get_method
-              path: rbyte.utils.make_dict
+              _target_: pipefunc.helpers.collect_kwargs
+              parameters: [camera_front_blur, lidar_velodyne, vehicle_data]
+            output_name: data
 
           - _target_: pipefunc.PipeFunc
             renames:

diff --git a/pyproject.toml b/pyproject.toml
@@ -1,14 +1,14 @@
 [project]
 name = "rbyte"
-version = "0.10.1"
+version = "0.10.2"
 description = "Multimodal PyTorch dataset library"
 authors = [{ name = "Evgenii Gorchakov", email = "evgenii@yaak.ai" }]
 maintainers = [{ name = "Evgenii Gorchakov", email = "evgenii@yaak.ai" }]
 dependencies = [
   "tensordict>=0.6.2",
   "torch",
   "numpy",
-  "polars>=1.16.0",
+  "polars>=1.18.0",
   "pydantic>=2.10.2",
   "more-itertools>=10.5.0",
   "hydra-core>=1.3.2",
@@ -17,9 +17,8 @@ dependencies = [
   "diskcache>=5.6.3",
   "parse>=1.20.2",
   "structlog>=24.4.0",
-  "xxhash>=3.5.0",
   "tqdm>=4.66.5",
-  "pipefunc>=0.41.0",
+  "pipefunc>=0.46.0",
 ]
 readme = "README.md"
 requires-python = ">=3.12,<3.13"
@@ -39,7 +38,7 @@ repo = "https://github.com/yaak-ai/rbyte"
 
 [project.optional-dependencies]
 build = ["hatchling>=1.25.0", "grpcio-tools>=1.62.0", "protoletariat==3.2.19"]
-visualize = ["rerun-sdk[notebook]>=0.20.2"]
+visualize = ["rerun-sdk[notebook]>=0.21.0"]
 mcap = [
   "mcap>=1.2.1",
   "mcap-ros2-support>=0.5.5",
@@ -53,7 +52,7 @@ video = [
   "video-reader-rs>=0.2.1",
 ]
 hdf5 = ["h5py>=3.12.1"]
-rrd = ["rerun-sdk>=0.20.2", "pyarrow-stubs"]
+rrd = ["rerun-sdk>=0.21.0", "pyarrow-stubs"]
 
 [project.scripts]
 rbyte-visualize = 'rbyte.scripts.visualize:main'

diff --git a/src/rbyte/io/yaak/idl-repo b/src/rbyte/io/yaak/idl-repo
diff --git a/src/rbyte/utils/__init__.py b/src/rbyte/utils/__init__.py
@@ -1,3 +0,0 @@
-from ._pipefunc import make_dict
-
-__all__ = ["make_dict"]

diff --git a/src/rbyte/utils/_pipefunc.py b/src/rbyte/utils/_pipefunc.py
diff --git a/src/rbyte/viz/loggers/rerun_logger.py b/src/rbyte/viz/loggers/rerun_logger.py
@@ -209,6 +209,6 @@ def log(self, batch_idx: int, batch: Batch) -> None:
                         rr.send_columns(
                             entity_path=entity_path,
                             times=times,
-                            components=components,  # pyright: ignore[reportArgumentType]
+                            components=components,
                             strict=True,
                         )
Original file line number	Diff line number	Diff line change
		@@ -1,3 +0,0 @@
		from ._pipefunc import make_dict

		__all__ = ["make_dict"]