mldataforge 0.0.4__tar.gz → 0.0.5__tar.gz

{mldataforge-0.0.4 → mldataforge-0.0.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mldataforge
-Version: 0.0.4
+Version: 0.0.5
 Summary: swiss army knife of scripts for transforming and processing datasets for machine learning.
 Project-URL: Homepage, https://github.com/schneiderkamplab/mldataforge
 Project-URL: Bug Tracker, https://github.com/schneiderkamplab/mldataforge/issues

{mldataforge-0.0.4 → mldataforge-0.0.5}/mldataforge/commands/convert/jsonl/__init__.py

@@ -1,6 +1,7 @@
 import click
 
 from .mds import mds
+from .parquet import parquet
 
 __all__ = ["jsonl"]
 
@@ -9,3 +10,4 @@ def jsonl():
     pass
 
 jsonl.add_command(mds)
+jsonl.add_command(parquet)

mldataforge-0.0.5/mldataforge/commands/convert/jsonl/parquet.py

@@ -0,0 +1,39 @@
+import click
+import json
+import pyarrow as pa
+import pyarrow.parquet as pq
+from tqdm import tqdm
+
+from ....utils import batch_iterable, check_overwrite, open_jsonl
+
+def _iterate(jsonl_files):
+    lines = 0
+    for jsonl_file in tqdm(jsonl_files, desc="Processing JSONL files", unit="file"):
+        with open_jsonl(jsonl_file, compression="infer") as f:
+            for line_num, line in enumerate(f, start=1):
+                try:
+                    item = json.loads(line)
+                    yield item
+                except json.JSONDecodeError as e:
+                    print(f"Skipping line {line_num} in {jsonl_file} due to JSON error: {e}")
+                lines += 1
+    print(f"Wrote {lines} lines from {len(jsonl_files)} files")
+
+@click.command()
+@click.argument('output_file', type=click.Path(exists=False))
+@click.argument('jsonl_files', nargs=-1, type=click.Path(exists=True))
+@click.option("--compression", default="snappy", type=click.Choice(["snappy", "gzip", "zstd"]), help="Compress the Parquet file (default: snappy).")
+@click.option("--overwrite", is_flag=True, help="Overwrite existing MDS directory.")
+@click.option("--yes", is_flag=True, help="Assume yes to all prompts. Use with caution as it will remove entire directory trees without confirmation.")
+@click.option("--batch-size", default=2**16, help="Batch size for loading MDS directories and writing Parquet files (default: 65536).")
+def parquet(output_file, jsonl_files, compression, overwrite, yes, batch_size):
+    check_overwrite(output_file, overwrite, yes)
+    if not jsonl_files:
+        raise click.BadArgumentUsage("No JSONL files provided.")
+    writer = None
+    for batch in batch_iterable(_iterate(jsonl_files), batch_size):
+        table = pa.Table.from_pylist(batch)
+        if writer is None:
+            writer = pq.ParquetWriter(output_file, table.schema, compression=compression)
+        writer.write_table(table)
+    writer.close()

{mldataforge-0.0.4 → mldataforge-0.0.5}/mldataforge/commands/convert/mds/__init__.py

@@ -1,6 +1,7 @@
 import click
 
 from .jsonl import jsonl
+from .parquet import parquet
 
 __all__ = ["mds"]
 
@@ -9,3 +10,4 @@ def mds():
     pass
 
 mds.add_command(jsonl)
+mds.add_command(parquet)

mldataforge-0.0.5/mldataforge/commands/convert/mds/parquet.py

@@ -0,0 +1,26 @@
+import click
+import pyarrow as pa
+import pyarrow.parquet as pq
+from tqdm import tqdm
+
+from ....utils import batch_iterable, check_overwrite, load_mds_directories
+
+@click.command()
+@click.argument("output_file", type=click.Path(exists=False), required=True)
+@click.argument("mds_directories", type=click.Path(exists=True), required=True, nargs=-1)
+@click.option("--compression", default="snappy", type=click.Choice(["snappy", "gzip", "zstd"]), help="Compress the Parquet file (default: snappy).")
+@click.option("--overwrite", is_flag=True, help="Overwrite existing Parquet files.")
+@click.option("--yes", is_flag=True, help="Assume yes to all prompts. Use with caution as it will remove files without confirmation.")
+@click.option("--batch-size", default=2**16, help="Batch size for loading MDS directories and writing Parquet files (default: 65536).")
+def parquet(output_file, mds_directories, compression, overwrite, yes, batch_size):
+    check_overwrite(output_file, overwrite, yes)
+    if not mds_directories:
+        raise click.BadArgumentUsage("No MDS files provided.")
+    ds = load_mds_directories(mds_directories, batch_size=batch_size)
+    writer = None
+    for batch in tqdm(batch_iterable(ds, batch_size), desc="Writing to Parquet", unit="batch", total=(len(ds)+batch_size-1) // batch_size):
+        table = pa.Table.from_pylist(batch)
+        if writer is None:
+            writer = pq.ParquetWriter(output_file, table.schema, compression=compression)
+        writer.write_table(table)
+    writer.close()

{mldataforge-0.0.4 → mldataforge-0.0.5}/mldataforge/utils.py

@@ -27,6 +27,16 @@ __all__ = [
     "use_pigz",
 ]
 
+def batch_iterable(iterable, batch_size):
+    batch = []
+    for item in iterable:
+        batch.append(item)
+        if len(batch) == batch_size:
+            yield batch
+            batch.clear()
+    if batch:
+        yield batch
+
 def check_overwrite(output_path, overwrite, yes):
     if os.path.exists(output_path):
         if os.path.isfile(output_path):

{mldataforge-0.0.4 → mldataforge-0.0.5}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "mldataforge"
-version = "0.0.4"
+version = "0.0.5"
 authors = [
   { name = "Peter Schneider-Kamp" }
 ]