Add Megatron-Bridge pretrain launcher

dataset/pretrain/scripts/export_pretrain_parquet_text_jsonl.py

@@ -0,0 +1,75 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import argparse
+import glob
+import json
+from pathlib import Path
+
+import pyarrow.parquet as pq
+
+
+def iter_paths(patterns: list[str]) -> list[Path]:
+    paths: list[Path] = []
+    for pattern in patterns:
+        path = Path(pattern)
+        if path.is_dir():
+            paths.extend(path.glob("*.parquet"))
+            continue
+        matches = glob.glob(pattern)
+        if matches:
+            paths.extend(Path(match) for match in matches)
+            continue
+        if path.exists():
+            paths.append(path)
+            continue
+        raise FileNotFoundError(pattern)
+    return sorted(set(paths), key=str)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Export pretrain parquet text rows to Megatron JSONL.")
+    parser.add_argument("--input", action="append", required=True, help="Parquet directory/file/glob.")
+    parser.add_argument("--output", required=True, help="Output JSONL path.")
+    parser.add_argument("--text-field", default="text")
+    parser.add_argument("--batch-size", type=int, default=8192)
+    parser.add_argument("--max-docs", type=int, default=0, help="0 means no limit.")
+    parser.add_argument("--progress-every", type=int, default=25)
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    paths = iter_paths(args.input)
+    if not paths:
+        raise SystemExit("no parquet inputs resolved")
+
+    output = Path(args.output)
+    output.parent.mkdir(parents=True, exist_ok=True)
+    docs = 0
+    with output.open("w", encoding="utf-8") as handle:
+        for file_idx, path in enumerate(paths, start=1):
+            parquet_file = pq.ParquetFile(path)
+            if args.text_field not in parquet_file.schema_arrow.names:
+                continue
+            for batch in parquet_file.iter_batches(
+                batch_size=args.batch_size,
+                columns=[args.text_field],
+                use_threads=True,
+            ):
+                for row in batch.to_pylist():
+                    text = row.get(args.text_field)
+                    if not isinstance(text, str) or not text.strip():
+                        continue
+                    handle.write(json.dumps({"text": text}, ensure_ascii=False) + "\n")
+                    docs += 1
+                    if args.max_docs and docs >= args.max_docs:
+                        print(f"export_done docs={docs} files_seen={file_idx}/{len(paths)} output={output}")
+                        return
+            if file_idx % args.progress_every == 0:
+                print(f"export_progress files={file_idx}/{len(paths)} docs={docs}")
+    print(f"export_done docs={docs} files_seen={len(paths)} output={output}")
+
+
+if __name__ == "__main__":
+    main()