Preserve per-row text metadata in collapsed WebDataset segments

Author: VibhuJawa

nemo_curator/stages/multimodal/io/writers/multimodal.py

@@ -200,22 +200,17 @@ def _write_webdataset_to_fileobj(task: MultimodalBatch, fileobj: BinaryIO) -> No
 
         first_text_index: dict[str, int] = {}
         text_row_count: dict[str, int] = {}
-        merged_text_payload: dict[str, bytes] = {}
         text_segments: dict[str, list[dict[str, object]]] = {}
         has_text_segment_metadata: dict[str, bool] = {}
         for idx, (sample_id, modality) in enumerate(zip(sample_ids, modalities, strict=True)):
             sid = str(sample_id)
             if str(modality) == "text":
                 if sid not in first_text_index:
                     first_text_index[sid] = idx
-                    merged_text_payload[sid] = b""
                     text_row_count[sid] = 0
                     text_segments[sid] = []
                     has_text_segment_metadata[sid] = False
                 text_row_count[sid] += 1
-                current = merged_text_payload[sid]
-                text_bytes = str(text_contents[idx] or "").encode("utf-8")
-                merged_text_payload[sid] = text_bytes if current == b"" else current + b"\n" + text_bytes
                 segment: dict[str, object] = {"modality": "text", "text": str(text_contents[idx] or "")}
                 text_row_metadata = MultimodalWriterStage._parse_json_or_raw(element_metadata_jsons[idx])
                 if text_row_metadata is not None:
@@ -233,7 +228,6 @@ def _write_webdataset_to_fileobj(task: MultimodalBatch, fileobj: BinaryIO) -> No
                         suffix, payload = MultimodalWriterStage._text_suffix_and_payload(
                             sample_id=sid,
                             content_type=content_types[idx],
-                            merged_text_payload=merged_text_payload[sid],
                             text_segments=text_segments[sid],
                             include_segment_metadata=has_text_segment_metadata[sid],
                         )
@@ -257,7 +251,6 @@ def _text_suffix_and_payload(
         *,
         sample_id: str,
         content_type: object | None,
-        merged_text_payload: bytes,
         text_segments: list[dict[str, object]],
         include_segment_metadata: bool,
     ) -> tuple[str, bytes]:
@@ -273,7 +266,8 @@ def _text_suffix_and_payload(
             return "json", json.dumps(payload, ensure_ascii=True).encode("utf-8")
         ctype = str(content_type) if content_type is not None else "text/plain"
         suffix = "json" if ctype == "application/json" else "txt"
-        return suffix, merged_text_payload
+        text_payload = "\n".join(str(segment.get("text", "")) for segment in text_segments)
+        return suffix, text_payload.encode("utf-8")
 
     @staticmethod
     def _parse_json_or_raw(value: object | None) -> object | None:

tests/stages/multimodal/test_writer_output_formats.py

@@ -614,6 +614,62 @@ def test_webdataset_writer_collapsed_text_preserves_element_metadata_json(tmp_pa
     assert json.loads(str(rows[1]["element_metadata_json"]))["element_metadata_json"]["lang"] == "en"
 
 
+def test_webdataset_writer_collapsed_text_writes_full_segment_metadata_payload(tmp_path: Path) -> None:
+    out = tmp_path / "collapsed_text_full_metadata.tar"
+    table = pa.table(
+        {
+            "sample_id": ["doc", "doc", "doc", "doc"],
+            "position": [0, 1, 2, 3],
+            "modality": ["text", "text", "text", "image"],
+            "content_type": ["text/plain", "text/plain", "text/plain", "image/jpeg"],
+            "text_content": ["alpha", "beta", "gamma", None],
+            "binary_content": [None, None, None, b"img"],
+            "element_metadata_json": [
+                '{"quality": 0.91, "token_count": 1}',
+                '{"quality": 0.77, "lang": "en", "attrs": {"source": "ocr"}}',
+                '{"quality": 0.55, "tags": ["x", "y"]}',
+                None,
+            ],
+            "source_id": ["src", "src", "src", "src"],
+            "source_shard": ["shard", "shard", "shard", "shard"],
+            "content_path": [None, None, None, None],
+            "content_key": [None, None, None, "doc.jpg"],
+        },
+        schema=MULTIMODAL_SCHEMA,
+    )
+    task = MultimodalBatch(task_id="t-full-meta", dataset_name="ds", data=table)
+    result = MultimodalWriterStage(output_path=str(out), output_format="webdataset").process(task)
+    names, members = _read_tar_members(Path(result.data[0]))
+
+    assert names == ["doc.000000.json", "doc.000003.jpg"]
+    payload = json.loads(members["doc.000000.json"].decode("utf-8"))
+    assert payload["sample_id"] == "doc"
+    assert [segment["text"] for segment in payload["segments"]] == ["alpha", "beta", "gamma"]
+    assert payload["segments"][0]["element_metadata_json"] == {"quality": 0.91, "token_count": 1}
+    assert payload["segments"][1]["element_metadata_json"] == {
+        "quality": 0.77,
+        "lang": "en",
+        "attrs": {"source": "ocr"},
+    }
+    assert payload["segments"][2]["element_metadata_json"] == {"quality": 0.55, "tags": ["x", "y"]}
+
+    roundtrip = WebDatasetReaderStage(load_binary=False, sample_format="auto").process(
+        FileGroupTask(task_id="rt-full-meta", dataset_name="ds", data=[result.data[0]])
+    )
+    rows = sorted(
+        [row for row in roundtrip.data.to_pylist() if row["modality"] == "text"],
+        key=lambda row: int(row["position"]),
+    )
+    assert [row["text_content"] for row in rows] == ["alpha", "beta", "gamma"]
+    assert json.loads(str(rows[0]["element_metadata_json"]))["element_metadata_json"] == {"quality": 0.91, "token_count": 1}
+    assert json.loads(str(rows[1]["element_metadata_json"]))["element_metadata_json"] == {
+        "quality": 0.77,
+        "lang": "en",
+        "attrs": {"source": "ocr"},
+    }
+    assert json.loads(str(rows[2]["element_metadata_json"]))["element_metadata_json"] == {"quality": 0.55, "tags": ["x", "y"]}
+
+
 def test_webdataset_writer_allows_text_only_batch(tmp_path: Path) -> None:
     out = tmp_path / "text-only.tar"
     table = pa.table(