Fix ruff formatting for interleaved module

benchmarking/scripts/multimodal_mint1t_benchmark.py

@@ -54,7 +54,9 @@ def create_pipeline(args: argparse.Namespace) -> Pipeline:
             per_text_fields=tuple(args.per_text_fields) if args.per_text_fields else (),
         )
     )
-    pipeline.add_stage(InterleavedAspectRatioFilterStage(drop_invalid_rows=True, min_aspect_ratio=1.0, max_aspect_ratio=2.0))
+    pipeline.add_stage(
+        InterleavedAspectRatioFilterStage(drop_invalid_rows=True, min_aspect_ratio=1.0, max_aspect_ratio=2.0)
+    )
     pipeline.add_stage(
         InterleavedParquetWriterStage(
             path=args.output_path,

nemo_curator/stages/interleaved/io/readers/webdataset.py

@@ -90,7 +90,11 @@ def __post_init__(self) -> None:
     # -- source_ref construction --
 
     def _build_source_ref(
-        self, ctx: _SampleContext, content_key: str | None, *, frame_index: int | None = None,
+        self,
+        ctx: _SampleContext,
+        content_key: str | None,
+        *,
+        frame_index: int | None = None,
     ) -> str:
         if content_key is None:
             return InterleavedBatch.build_source_ref(path=None, member=None)
@@ -101,8 +105,11 @@ def _build_source_ref(
             byte_offset = info.offset_data
             byte_size = info.size
         return InterleavedBatch.build_source_ref(
-            path=ctx.tar_path, member=content_key,
-            byte_offset=byte_offset, byte_size=byte_size, frame_index=frame_index,
+            path=ctx.tar_path,
+            member=content_key,
+            byte_offset=byte_offset,
+            byte_size=byte_size,
+            frame_index=frame_index,
         )
 
     # -- row builders (override in subclasses for custom formats) --
@@ -121,16 +128,24 @@ def _build_row(ctx: _SampleContext, row_fields: dict[str, Any]) -> dict[str, Any
         }
 
     def _metadata_row(self, ctx: _SampleContext) -> dict[str, Any]:
-        return {**self._build_row(ctx, {
-            "position": -1,
-            "modality": "metadata",
-            "content_type": "application/json",
-            "source_ref": self._build_source_ref(ctx, ctx.json_member_name),
-        }), **ctx.passthrough}
+        return {
+            **self._build_row(
+                ctx,
+                {
+                    "position": -1,
+                    "modality": "metadata",
+                    "content_type": "application/json",
+                    "source_ref": self._build_source_ref(ctx, ctx.json_member_name),
+                },
+            ),
+            **ctx.passthrough,
+        }
 
     @staticmethod
     def _apply_per_modality_fields(
-        row: dict[str, Any], passthrough: dict[str, list[Any]], index: int,
+        row: dict[str, Any],
+        passthrough: dict[str, list[Any]],
+        index: int,
     ) -> None:
         for field_name, values in passthrough.items():
             if index < len(values):
@@ -139,13 +154,21 @@ def _apply_per_modality_fields(
 
     @staticmethod
     def _warn_per_modality_length_mismatch(
-        sample_id: str, passthrough: dict[str, list[Any]], actual_count: int, modality: str,
+        sample_id: str,
+        passthrough: dict[str, list[Any]],
+        actual_count: int,
+        modality: str,
     ) -> None:
         for field_name, values in passthrough.items():
             if actual_count != len(values):
                 logger.warning(
                     "sample_id={}: per_{}_field '{}' has {} values but {} non-None {}s",
-                    sample_id, modality, field_name, len(values), actual_count, modality,
+                    sample_id,
+                    modality,
+                    field_name,
+                    len(values),
+                    actual_count,
+                    modality,
                 )
 
     def _text_rows(self, ctx: _SampleContext) -> list[dict[str, Any]]:
@@ -158,13 +181,16 @@ def _text_rows(self, ctx: _SampleContext) -> list[dict[str, Any]]:
         for idx, text_value in enumerate(texts):
             if text_value is None:
                 continue
-            row = self._build_row(ctx, {
-                "position": idx,
-                "modality": "text",
-                "content_type": "text/plain",
-                "text_content": str(text_value),
-                "source_ref": source_ref,
-            })
+            row = self._build_row(
+                ctx,
+                {
+                    "position": idx,
+                    "modality": "text",
+                    "content_type": "text/plain",
+                    "text_content": str(text_value),
+                    "source_ref": source_ref,
+                },
+            )
             self._apply_per_modality_fields(row, ctx.per_text_passthrough, non_none_counter)
             non_none_counter += 1
             rows.append(row)
@@ -176,7 +202,10 @@ def _image_rows(self, ctx: _SampleContext) -> list[dict[str, Any]]:
         if not isinstance(images, list):
             return []
         image_member_name = self._resolve_default_image_member_name(
-            ctx.sample_id, ctx.sample, images, ctx.member_names,
+            ctx.sample_id,
+            ctx.sample,
+            images,
+            ctx.member_names,
         )
         rows: list[dict[str, Any]] = []
         frame_counters: dict[str, int] = {}
@@ -191,12 +220,15 @@ def _image_rows(self, ctx: _SampleContext) -> list[dict[str, Any]]:
             if content_key is not None and is_multiframe_candidate:
                 frame_index = frame_counters.get(content_key, 0)
                 frame_counters[content_key] = frame_index + 1
-            row = self._build_row(ctx, {
-                "position": idx,
-                "modality": "image",
-                "content_type": content_type or ("application/octet-stream" if image_member_name else None),
-                "source_ref": self._build_source_ref(ctx, content_key, frame_index=frame_index),
-            })
+            row = self._build_row(
+                ctx,
+                {
+                    "position": idx,
+                    "modality": "image",
+                    "content_type": content_type or ("application/octet-stream" if image_member_name else None),
+                    "source_ref": self._build_source_ref(ctx, content_key, frame_index=frame_index),
+                },
+            )
             self._apply_per_modality_fields(row, ctx.per_image_passthrough, non_none_counter)
             non_none_counter += 1
             rows.append(row)
@@ -234,7 +266,8 @@ def _build_passthrough_row(self, sample: dict[str, Any]) -> dict[str, Any]:
 
     @staticmethod
     def _extract_per_modality_fields(
-        sample: dict[str, Any], field_names: tuple[str, ...],
+        sample: dict[str, Any],
+        field_names: tuple[str, ...],
     ) -> dict[str, list[Any]]:
         result: dict[str, list[Any]] = {}
         for field_name in field_names:
@@ -245,10 +278,7 @@ def _extract_per_modality_fields(
             if isinstance(value, list):
                 result[field_name] = value
             else:
-                msg = (
-                    f"per-modality field '{field_name}' must be a list, "
-                    f"got {type(value).__name__}"
-                )
+                msg = f"per-modality field '{field_name}' must be a list, got {type(value).__name__}"
                 raise TypeError(msg)
         return result
 

nemo_curator/tasks/interleaved.py

@@ -182,7 +182,10 @@ def build_source_ref(
     ) -> str:
         """Build a ``source_ref`` JSON locator string."""
         ref: dict[str, object] = {
-            "path": path, "member": member, "byte_offset": byte_offset, "byte_size": byte_size,
+            "path": path,
+            "member": member,
+            "byte_offset": byte_offset,
+            "byte_size": byte_size,
         }
         if frame_index is not None:
             ref["frame_index"] = frame_index

tests/stages/interleaved/test_interleaved_task.py

@@ -158,7 +158,11 @@ def test_parse_source_ref_with_frame_index() -> None:
 )
 def test_build_source_ref_frame_index(frame_index: int | None, key_present: bool) -> None:
     ref_str = InterleavedBatch.build_source_ref(
-        path="/a.tar", member="m.jpg", byte_offset=10, byte_size=20, frame_index=frame_index,
+        path="/a.tar",
+        member="m.jpg",
+        byte_offset=10,
+        byte_size=20,
+        frame_index=frame_index,
     )
     parsed = json.loads(ref_str)
     assert ("frame_index" in parsed) is key_present

tests/stages/interleaved/test_materialization.py

@@ -58,7 +58,10 @@ def _image_row(
         "text_content": None,
         "binary_content": None,
         "source_ref": InterleavedBatch.build_source_ref(
-            path=path, member=member, byte_offset=byte_offset, byte_size=byte_size,
+            path=path,
+            member=member,
+            byte_offset=byte_offset,
+            byte_size=byte_size,
         ),
         "materialize_error": None,
     }
@@ -87,23 +90,27 @@ def test_get_frame_index_returns_none_for_missing_values(val: object, expected:
     [pytest.param(float("nan"), id="nan_path"), pytest.param("", id="empty_path")],
 )
 def test_classify_rows_missing_path_variants(path_val: object) -> None:
-    df = pd.DataFrame({
-        "_src_path": [path_val],
-        "_src_member": [None],
-        "_src_byte_offset": [None],
-        "_src_byte_size": [None],
-    })
+    df = pd.DataFrame(
+        {
+            "_src_path": [path_val],
+            "_src_member": [None],
+            "_src_byte_offset": [None],
+            "_src_byte_size": [None],
+        }
+    )
     result = _classify_rows(df, pd.Series([True]))
     assert result.missing == [0]
 
 
 def test_classify_rows_range_with_zero_size() -> None:
-    df = pd.DataFrame({
-        "_src_path": ["/shard.tar"],
-        "_src_member": ["img.jpg"],
-        "_src_byte_offset": [100],
-        "_src_byte_size": [0],
-    })
+    df = pd.DataFrame(
+        {
+            "_src_path": ["/shard.tar"],
+            "_src_member": ["img.jpg"],
+            "_src_byte_offset": [100],
+            "_src_byte_size": [0],
+        }
+    )
     result = _classify_rows(df, pd.Series([True]))
     assert "/shard.tar" in result.tar_extract
     assert not result.range_read
@@ -322,18 +329,21 @@ def test_materialize_with_only_missing_binary_false(tmp_path: Path) -> None:
     img_path = tmp_path / "img.jpg"
     img_path.write_bytes(new_bytes)
 
-    rows = [{
-        "sample_id": "s1",
-        "position": 0,
-        "modality": "image",
-        "content_type": "image/jpeg",
-        "text_content": None,
-        "binary_content": b"old-bytes",
-        "source_ref": InterleavedBatch.build_source_ref(path=str(img_path), member=None),
-        "materialize_error": None,
-    }]
+    rows = [
+        {
+            "sample_id": "s1",
+            "position": 0,
+            "modality": "image",
+            "content_type": "image/jpeg",
+            "text_content": None,
+            "binary_content": b"old-bytes",
+            "source_ref": InterleavedBatch.build_source_ref(path=str(img_path), member=None),
+            "materialize_error": None,
+        }
+    ]
     task = InterleavedBatch(
-        task_id="re_mat", dataset_name="d",
+        task_id="re_mat",
+        dataset_name="d",
         data=pa.Table.from_pylist(rows, schema=INTERLEAVED_SCHEMA),
     )
     result = materialize_task_binary_content(task, only_missing_binary=False)