mldataforge 0.0.5__tar.gz → 0.1.1__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mldataforge
-Version: 0.0.5
+Version: 0.1.1
 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.1.1/mldataforge/commands/convert/jsonl.py

@@ -0,0 +1,46 @@
+import click
+from datasets import load_dataset
+
+from ...options import *
+from ...utils import *
+
+__all__ = ["jsonl"]
+
+@click.group()
+def jsonl():
+    pass
+
+@jsonl.command()
+@click.argument('output_dir', type=click.Path(exists=False))
+@click.argument('jsonl_files', nargs=-1, type=click.Path(exists=True))
+@compression_option(None, ['none', 'br', 'bz2', 'gzip', 'pigz', 'snappy', 'zstd'])
+@overwrite_option()
+@yes_option()
+@processes_option()
+@buf_size_option()
+def mds(output_dir, jsonl_files, compression, processes, overwrite, yes, buf_size):
+    check_arguments(output_dir, overwrite, yes, jsonl_files)
+    save_mds(
+        load_dataset("json", data_files=jsonl_files, split="train"),
+        output_dir,
+        processes=processes,
+        compression=compression,
+        buf_size=buf_size,
+        pigz=use_pigz(compression),
+    )
+
+@jsonl.command()
+@click.argument('output_file', type=click.Path(exists=False))
+@click.argument('jsonl_files', nargs=-1, type=click.Path(exists=True))
+@compression_option("snappy", ["snappy", "gzip", "zstd"])
+@overwrite_option()
+@yes_option()
+@batch_size_option()
+def parquet(output_file, jsonl_files, compression, overwrite, yes, batch_size):
+    check_arguments(output_file, overwrite, yes, jsonl_files)
+    save_parquet(
+        load_dataset("json", data_files=jsonl_files, split="train"),
+        output_file,
+        compression=compression,
+        batch_size=batch_size,
+    )

mldataforge-0.1.1/mldataforge/commands/convert/mds.py

@@ -0,0 +1,45 @@
+import click
+
+from ...options import *
+from ...utils import *
+
+__all__ = ["mds"]
+
+@click.group()
+def mds():
+    pass
+
+@mds.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)
+@compression_option("infer", ["none", "infer", "pigz", "gzip", "bz2", "xz"])
+@processes_option()
+@overwrite_option()
+@yes_option()
+@batch_size_option()
+@no_bulk_option()
+def jsonl(output_file, mds_directories, compression, processes, overwrite, yes, batch_size, no_bulk):
+    check_arguments(output_file, overwrite, yes, mds_directories)
+    save_jsonl(
+        load_mds_directories(mds_directories, batch_size=batch_size, bulk=not no_bulk),
+        output_file,
+        compression=compression,
+        processes=processes,
+    )
+
+@mds.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)
+@compression_option("snappy", ["snappy", "gzip", "zstd"])
+@overwrite_option()
+@yes_option()
+@batch_size_option()
+@no_bulk_option()
+def parquet(output_file, mds_directories, compression, overwrite, yes, batch_size, no_bulk):
+    check_arguments(output_file, overwrite, yes, mds_directories)
+    save_parquet(
+        load_mds_directories(mds_directories, batch_size=batch_size, bulk=not no_bulk),
+        output_file,
+        compression=compression,
+        batch_size=batch_size,
+    )

mldataforge-0.1.1/mldataforge/commands/convert/parquet.py

@@ -0,0 +1,46 @@
+import click
+from datasets import load_dataset
+
+from ...options import *
+from ...utils import *
+
+__all__ = ["parquet"]
+
+@click.group()
+def parquet():
+    pass
+
+@parquet.command()
+@click.argument("output_file", type=click.Path(exists=False), required=True)
+@click.argument("parquet_files", type=click.Path(exists=True), required=True, nargs=-1)
+@compression_option("infer", ["none", "infer", "pigz", "gzip", "bz2", "xz"])
+@processes_option()
+@overwrite_option()
+@yes_option()
+def jsonl(output_file, parquet_files, compression, processes, overwrite, yes):
+    check_arguments(output_file, overwrite, yes, parquet_files)
+    save_jsonl(
+        load_dataset("parquet", data_files=parquet_files, split="train"),
+        output_file,
+        compression=compression,
+        processes=processes,
+    )
+
+@parquet.command()
+@click.argument('output_dir', type=click.Path(exists=False))
+@click.argument('parquet_files', nargs=-1, type=click.Path(exists=True))
+@compression_option(None, ['none', 'br', 'bz2', 'gzip', 'pigz', 'snappy', 'zstd'])
+@processes_option()
+@overwrite_option()
+@yes_option()
+@buf_size_option()
+def mds(output_dir, parquet_files, compression, processes, overwrite, yes, buf_size):
+    check_arguments(output_dir, overwrite, yes, parquet_files)
+    save_mds(
+        load_dataset("parquet", data_files=parquet_files, split="train"),
+        output_dir,
+        processes=processes,
+        compression=compression,
+        buf_size=buf_size,
+        pigz=use_pigz(compression),
+    )

mldataforge-0.1.1/mldataforge/mds.py

@@ -0,0 +1,97 @@
+import gzip
+import json
+from mltiming import timing
+import numpy as np
+import os
+from streaming.base.format.mds.encodings import mds_decode
+from typing import Any, Optional, Generator
+
+class MDSBulkReader:
+    def __init__(
+        self,
+        dirnames: list[str],
+        split: Optional[str],
+    ) -> None:
+        self.shards = []
+        self.samples = 0
+        for dirname in dirnames:
+            if split is not None:
+                dirname = os.path.join(dirname, split)
+            index = json.load(open(os.path.join(dirname, "index.json"), 'rt'))
+            for shard in index["shards"]:
+                basename = shard['raw_data']['basename'] if shard['zip_data'] is None else shard['zip_data']['basename']
+                filename = os.path.join(dirname, basename)
+                self.shards.append({
+                    "filename": filename,
+                    "compression": shard['compression'],
+                })
+                self.samples += shard['samples']
+
+    def __len__(self) -> int:
+        return self.samples
+
+    def __iter__(self) -> Generator[dict[str, Any], None, None]:
+        for shard in self.shards:
+            with MDSShardReader(**shard) as reader:
+                for sample in reader:
+                    yield sample
+
+class MDSShardReader:
+    def __init__(
+        self,
+        filename: str,
+        compression: Optional[str],
+    ) -> None:
+        if compression is None:
+            _open = open
+        elif compression == 'gz':
+            _open = gzip.open
+        else:
+            raise ValueError(f'Unsupported compression type: {compression}. Supported types: None, gzip.')
+        self.fp = _open(filename, "rb")
+        self.samples = np.frombuffer(self.fp.read(4), np.uint32)[0]
+        self.index = np.frombuffer(self.fp.read((1+self.samples)*4), np.uint32)
+        info = json.loads(self.fp.read(self.index[0]-self.fp.tell()))
+        self.column_encodings = info["column_encodings"]
+        self.column_names = info["column_names"]
+        self.column_sizes = info["column_sizes"]
+        assert self.fp.tell() == self.index[0]
+
+    def decode_sample(self, data: bytes) -> dict[str, Any]:
+        sizes = []
+        idx = 0
+        for key, size in zip(self.column_names, self.column_sizes):
+            if size:
+                sizes.append(size)
+            else:
+                size, = np.frombuffer(data[idx:idx + 4], np.uint32)
+                sizes.append(size)
+                idx += 4
+        sample = {}
+        for key, encoding, size in zip(self.column_names, self.column_encodings, sizes):
+            value = data[idx:idx + size]
+            sample[key] = mds_decode(encoding, value)
+            idx += size
+        return sample
+
+    def get_sample_data(self, idx: int) -> bytes:
+        begin, end = self.index[idx:idx+2]
+        assert self.fp.tell() == begin
+        data = self.fp.read(end - begin)
+        assert self.fp.tell() == end
+        assert data
+        return data
+
+    def get_item(self, idx: int) -> dict[str, Any]:
+        data = self.get_sample_data(idx)
+        return self.decode_sample(data)
+
+    def __iter__(self) -> Generator[dict[str, Any], None, None]:
+        for i in range(self.samples):
+            yield self.get_item(i)
+
+    def __enter__(self) -> "MDSShardReader":
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback) -> None:
+        self.fp.close()

mldataforge-0.1.1/mldataforge/options.py

@@ -0,0 +1,81 @@
+import click
+
+__alll__ = [
+    "batch_size_option",
+    "buf_size_option",
+    "compression_option",
+    "overwrite_option",
+    "processes_option",
+    "yes_option",
+]
+
+def batch_size_option(default=2**16):
+    """
+    Option for specifying the batch size.
+    """
+    return click.option(
+        "--batch-size",
+        default=default,
+        help=f"Batch size for loading data and writing files (default: {default}).",
+    )
+
+def buf_size_option(default=2**24):
+    """
+    Option for specifying the buffer size.
+    """
+    return click.option(
+        "--buf-size",
+        default=default,
+        help=f"Buffer size for pigz compression (default: {default}).",
+    )
+
+def no_bulk_option():
+    """
+    Option for specifying whether to use a custom space and time-efficient bulk reader (only gzip and no compression).
+    """
+    return click.option(
+        "--no-bulk",
+        is_flag=True,
+        help="Use a custom space and time-efficient bulk reader (only gzip and no compression).",
+    )
+
+def compression_option(default, choices):
+    """
+    Option for specifying the compression type.
+    """
+    return click.option(
+        "--compression",
+        default=default,
+        type=click.Choice(choices, case_sensitive=False),
+        help=f"Compress the output file (default: {default}).",
+    )
+
+def overwrite_option():
+    """
+    Option for specifying whether to overwrite existing files.
+    """
+    return click.option(
+        "--overwrite",
+        is_flag=True,
+        help="Overwrite existing path.",
+    )
+
+def processes_option(default=64):
+    """
+    Option for specifying the number of processes to use.
+    """
+    return click.option(
+        "--processes",
+        default=default,
+        help=f"Number of processes to use (default: {default}).",
+    )
+
+def yes_option():
+    """
+    Option for specifying whether to assume yes to all prompts.
+    """
+    return click.option(
+        "--yes",
+        is_flag=True,
+        help="Assume yes to all prompts. Use with caution as it will remove files or even entire directories without confirmation.",
+    )

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

@@ -1,29 +1,28 @@
-import atexit
 import bz2
 import click
-from datasets import concatenate_datasets, load_dataset
+from datasets import concatenate_datasets
 import gzip
+import json
 import lzma
 from mltiming import timing
+import pyarrow as pa
+import pyarrow.parquet as pq
 import os
 import shutil
-from streaming import StreamingDataset
-import tempfile
+from streaming import MDSWriter, StreamingDataset
 from tqdm import tqdm
 
+from .mds import MDSBulkReader
 from .pigz import pigz_open
 
 __all__ = [
-    "check_overwrite",
-    "create_temp_file",
-    "determine_compression",
-    "infer_mds_encoding",
-    "infer_compression",
-    "load_parquet_files",
+    "batch_iterable",
+    "check_arguments",
+    "confirm_overwrite",
     "load_mds_directories",
-    "open_jsonl",
-    "pigz_available",
-    "pigz_compress",
+    "save_jsonl",
+    "save_mds",
+    "save_parquet",
     "use_pigz",
 ]
 
@@ -37,7 +36,9 @@ def batch_iterable(iterable, batch_size):
     if batch:
         yield batch
 
-def check_overwrite(output_path, overwrite, yes):
+def check_arguments(output_path, overwrite, yes, input_paths):
+    if not input_paths:
+        raise click.BadArgumentUsage("No input paths provided.")
     if os.path.exists(output_path):
         if os.path.isfile(output_path):
             if not overwrite:
@@ -62,31 +63,14 @@ def confirm_overwrite(message):
     if response.lower() != 'yes':
         raise click.Abort()
 
-def create_temp_file():
-    def _cleanup_file(file_path):
-        try:
-            os.remove(file_path)
-        except OSError:
-            pass
-    # Create a named temp file, don't delete right away
-    temp = tempfile.NamedTemporaryFile(delete=False)
-    temp_name = temp.name
-    # Close so others can open it again without conflicts (especially on Windows)
-    temp.close()
-
-    # Schedule its deletion at exit
-    atexit.register(_cleanup_file, temp_name)
-
-    return temp_name
-
-def determine_compression(file_path, compression="infer"):
+def _determine_compression(file_path, compression="infer"):
     if compression == "infer":
-        compression = infer_compression(file_path)
+        compression = _infer_compression(file_path)
     if compression == "none":
         compression = None
     return compression
 
-def infer_mds_encoding(value):
+def _infer_mds_encoding(value):
     """Determine the MDS encoding for a given value."""
     if isinstance(value, str):
         return 'str'
@@ -98,11 +82,11 @@ def infer_mds_encoding(value):
         return 'bool'
     return 'pkl'
 
-def infer_compression(file_path):
+def _infer_compression(file_path):
     """Infer the compression type from the file extension."""
     extension = os.path.splitext(file_path)[1]
     if extension.endswith('.gz'):
-        if pigz_available():
+        if _pigz_available():
             return 'pigz'
         return 'gzip'
     if extension.endswith('.bz2'):
@@ -115,9 +99,11 @@ def infer_compression(file_path):
         return 'zstd'
     return None
 
-def load_mds_directories(mds_directories, split='.', batch_size=2**16):
+def load_mds_directories(mds_directories, split='.', batch_size=2**16, bulk=True):
+    if bulk:
+        return MDSBulkReader(mds_directories, split=split)
     dss = []
-    for mds_directory in tqdm(mds_directories, desc="Loading MDS directories", unit="directory"):
+    for mds_directory in mds_directories:
         ds = StreamingDataset(
             local=mds_directory,
             remote=None,
@@ -136,21 +122,9 @@ def load_mds_directories(mds_directories, split='.', batch_size=2**16):
             ds = concatenate_datasets(dsets=dss)
     return ds
 
-def load_parquet_files(parquet_files):
-    dss = []
-    for parquet_file in tqdm(parquet_files, desc="Loading parquet files", unit="file"):
-        ds = load_dataset("parquet", data_files=parquet_file, split="train")
-        dss.append(ds)
-    if len(dss) == 1:
-        ds = dss[0]
-    else:
-        with timing(message=f"Concatenating {len(dss)} datasets"):
-            ds = concatenate_datasets(dsets=dss)
-    return ds
-
-def open_jsonl(file_path, mode="rt", compression="infer", processes=64):
+def _open_jsonl(file_path, mode="rt", compression="infer", processes=64):
     """Open a JSONL file, handling gzip compression if necessary."""
-    compression = determine_compression(file_path, compression)
+    compression = _determine_compression(file_path, compression)
     if compression == "gzip":
         return gzip.open(file_path, mode)
     if compression == "pigz":
@@ -163,11 +137,11 @@ def open_jsonl(file_path, mode="rt", compression="infer", processes=64):
         return open(file_path, mode)
     raise ValueError(f"Unsupported compression type: {compression}")
 
-def pigz_available():
+def _pigz_available():
     """Check if pigz is available on the system."""
     return shutil.which("pigz") is not None
 
-def pigz_compress(input_file, output_file, processes=64, buf_size=2**24, keep=False, quiet=False):
+def _pigz_compress(input_file, output_file, processes=64, buf_size=2**24, keep=False, quiet=False):
     """Compress a file using pigz."""
     size = os.stat(input_file).st_size
     num_blocks = (size+buf_size-1) // buf_size
@@ -183,6 +157,54 @@ def pigz_compress(input_file, output_file, processes=64, buf_size=2**24, keep=Fa
         if not quiet:
             print(f"Removed {input_file}")
 
+def save_jsonl(iterable, output_file, compression=None, processes=64):
+    compression = _determine_compression(output_file, compression)
+    with _open_jsonl(output_file, mode="wb", compression=compression, processes=processes) as f:
+        for item in tqdm(iterable, desc="Writing to JSONL", unit="sample"):
+            f.write(f"{json.dumps(item)}\n".encode("utf-8"))
+
+def save_mds(it, output_dir, processes=64, compression=None, buf_size=2**24, pigz=False):
+    if compression == "none" or pigz:
+        compression = None
+    if compression == "gzip":
+        compression = "gz"
+    writer = None
+    for sample in tqdm(it, desc="Writing to MDS", unit="sample"):
+        if writer is None:
+            columns = {key: _infer_mds_encoding(value) for key, value in sample.items()}
+            writer = MDSWriter(out=output_dir, columns=columns, compression=compression)
+        writer.write(sample)
+    writer.finish()
+    if pigz:
+        index_path = os.path.join(output_dir, "index.json")
+        index = json.load(open(index_path, "rt"))
+        name2info = {shard["raw_data"]["basename"]: shard for shard in index["shards"]}
+        file_names = [file for file in os.listdir(output_dir) if file.endswith(".mds")]
+        assert set(file_names) == set(name2info.keys())
+        for file_name in tqdm(file_names, desc="Compressing with pigz", unit="file"):
+            compressed_file_name = file_name + ".gz"
+            file_path = os.path.join(output_dir, file_name)
+            compressed_file_path = os.path.join(output_dir, compressed_file_name)
+            _pigz_compress(file_path, compressed_file_path, processes, buf_size=buf_size, keep=False, quiet=True)
+            name2info[file_name]["compression"] = "gz"
+            name2info[file_name]["zip_data"] = {
+                "basename": compressed_file_name,
+                "bytes": os.stat(compressed_file_path).st_size,
+                "hashes": {},
+            }
+        json.dump(index, open(index_path, "wt"))
+        print(f"Compressed {output_dir} with pigz")
+
+def save_parquet(it, output_file, compression=None, batch_size=2**16):
+    writer = None
+    it = tqdm(it, desc="Writing to Parquet", unit="sample")
+    for batch in batch_iterable(it, 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()
+
 def use_pigz(compression):
     """Determine if pigz should be used based on the compression type."""
-    return compression == "pigz" or (compression == "gzip" and pigz_available())
+    return compression == "pigz" or (compression == "gzip" and _pigz_available())

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

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

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

@@ -1,13 +0,0 @@
-import click
-
-from .mds import mds
-from .parquet import parquet
-
-__all__ = ["jsonl"]
-
-@click.group()
-def jsonl():
-    pass
-
-jsonl.add_command(mds)
-jsonl.add_command(parquet)

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

@@ -1,59 +0,0 @@
-import click
-import json
-import os
-from streaming import MDSWriter
-from tqdm import tqdm
-
-from ....utils import check_overwrite, infer_mds_encoding, open_jsonl, pigz_compress, use_pigz
-
-@click.command()
-@click.argument('output_dir', type=click.Path(exists=False))
-@click.argument('jsonl_files', nargs=-1, type=click.Path(exists=True))
-@click.option('--compression', type=click.Choice(['none', 'br', 'bz2', 'gzip', 'pigz', 'snappy', 'zstd'], case_sensitive=False), default=None, help='Compression type for the output dataset (default: None).')
-@click.option("--processes", default=64, help="Number of processes to use for pigz compression (default: 64).")
-@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("--buf-size", default=2**24, help=f"Buffer size for pigz compression (default: {2**24}).")
-def mds(output_dir, jsonl_files, processes, compression, overwrite, yes, buf_size):
-    check_overwrite(output_dir, overwrite, yes)
-    if not jsonl_files:
-        raise click.BadArgumentUsage("No JSONL files provided.")
-    with open_jsonl(jsonl_files[0]) as f:
-        sample = json.loads(f.readline())
-    pigz = use_pigz(compression)
-    if compression == "none" or pigz:
-        compression = None
-    if compression == "gzip":
-        compression = "gz"
-    columns = {key: infer_mds_encoding(value) for key, value in sample.items()}
-    lines = 0
-    with MDSWriter(out=output_dir, columns=columns, compression=compression) as writer:
-        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)
-                        writer.write(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 to MDS files in {output_dir}")
-    if pigz:
-        index_path = os.path.join(output_dir, "index.json")
-        index = json.load(open(index_path, "rt"))
-        name2info = {shard["raw_data"]["basename"]: shard for shard in index["shards"]}
-        file_names = [file for file in os.listdir(output_dir) if file.endswith(".mds")]
-        assert set(file_names) == set(name2info.keys())
-        for file_name in tqdm(file_names, desc="Compressing with pigz", unit="file"):
-            compressed_file_name = file_name + ".gz"
-            file_path = os.path.join(output_dir, file_name)
-            compressed_file_path = os.path.join(output_dir, compressed_file_name)
-            pigz_compress(file_path, compressed_file_path, processes, buf_size=buf_size, keep=False, quiet=True)
-            name2info[file_name]["compression"] = "gz"
-            name2info[file_name]["zip_data"] = {
-                "basename": compressed_file_name,
-                "bytes": os.stat(compressed_file_path).st_size,
-                "hashes": {},
-            }
-        json.dump(index, open(index_path, "wt"))
-        print(f"Compressed {output_dir} with pigz")

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

@@ -1,39 +0,0 @@
-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.5/mldataforge/commands/convert/mds/__init__.py

@@ -1,13 +0,0 @@
-import click
-
-from .jsonl import jsonl
-from .parquet import parquet
-
-__all__ = ["mds"]
-
-@click.group()
-def mds():
-    pass
-
-mds.add_command(jsonl)
-mds.add_command(parquet)

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

@@ -1,23 +0,0 @@
-import click
-import json
-from tqdm import tqdm
-
-from ....utils import check_overwrite, create_temp_file, determine_compression, load_mds_directories, open_jsonl
-
-@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="infer", type=click.Choice(["none", "infer", "pigz", "gzip", "bz2", "xz"]), help="Compress the output JSONL file (default: infer; pigz for parallel gzip).")
-@click.option("--processes", default=64, help="Number of processes to use for pigz compression (default: 64).")
-@click.option("--overwrite", is_flag=True, help="Overwrite existing JSONL 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 (default: 65536).")
-def jsonl(output_file, mds_directories, compression, processes, 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)
-    compression = determine_compression(output_file, compression)
-    with open_jsonl(output_file, mode="wb", compression=compression, processes=processes) as f:
-        for item in tqdm(ds, desc="Writing to JSONL", unit="line"):
-            f.write(f"{json.dumps(item)}\n".encode("utf-8"))

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

@@ -1,26 +0,0 @@
-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.5/mldataforge/commands/convert/parquet/__init__.py

@@ -1,13 +0,0 @@
-import click
-
-from .jsonl import jsonl
-from .mds import mds
-
-__all__ = ["parquet"]
-
-@click.group()
-def parquet():
-    pass
-
-parquet.add_command(jsonl)
-parquet.add_command(mds)

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

@@ -1,25 +0,0 @@
-import click
-from mltiming import timing
-
-from ....utils import check_overwrite, create_temp_file, determine_compression, load_parquet_files, pigz_compress
-
-@click.command()
-@click.argument("output_file", type=click.Path(exists=False), required=True)
-@click.argument("parquet_files", type=click.Path(exists=True), required=True, nargs=-1)
-@click.option("--compression", default="infer", type=click.Choice(["none", "infer", "pigz", "gzip", "bz2", "xz"]), help="Compress the output JSONL file (default: infer; pigz for parallel gzip).")
-@click.option("--processes", default=64, help="Number of processes to use for pigz compression (default: 64).")
-@click.option("--overwrite", is_flag=True, help="Overwrite existing JSONL 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("--buf-size", default=2**24, help=f"Buffer size for pigz compression (default: {2**24}).")
-def jsonl(output_file, parquet_files, compression, processes, overwrite, yes, buf_size):
-    check_overwrite(output_file, overwrite, yes)
-    if not parquet_files:
-        raise click.BadArgumentUsage("No parquet files provided.")
-    ds = load_parquet_files(parquet_files)
-    compression = determine_compression(output_file, compression)
-    compressed_file = None
-    if compression == "pigz":
-        compression, compressed_file, output_file = None, output_file, create_temp_file()
-    ds.to_json(output_file, num_proc=processes, orient="records", lines=True, compression=compression)
-    if compressed_file is not None:
-        pigz_compress(output_file, compressed_file, processes, buf_size, keep=False)

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

@@ -1,43 +0,0 @@
-import click
-import json
-import os
-from streaming import MDSWriter
-from tqdm import tqdm
-
-from ....utils import check_overwrite, infer_mds_encoding, load_parquet_files, pigz_compress, use_pigz
-
-@click.command()
-@click.argument('output_dir', type=click.Path(exists=False))
-@click.argument('parquet_files', nargs=-1, type=click.Path(exists=True))
-@click.option('--compression', type=click.Choice(['none', 'br', 'bz2', 'gzip', 'pigz', 'snappy', 'zstd'], case_sensitive=False), default=None, help='Compression type for the output dataset (default: None).')
-@click.option("--processes", default=64, help="Number of processes to use for pigz compression (default: 64).")
-@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("--buf-size", default=2**24, help=f"Buffer size for pigz compression (default: {2**24}).")
-def mds(output_dir, parquet_files, processes, compression, overwrite, yes, buf_size):
-    check_overwrite(output_dir, overwrite, yes)
-    if not parquet_files:
-        raise click.BadArgumentUsage("No parquet files provided.")
-    ds = load_parquet_files(parquet_files)
-    pigz = use_pigz(compression)
-    sample = ds[0]
-    if compression == "none" or pigz:
-        compression = None
-    if compression == "gzip":
-        compression = "gz"
-    columns = {key: infer_mds_encoding(value) for key, value in sample.items()}
-    lines = 0
-    with MDSWriter(out=output_dir, columns=columns, compression=compression) as writer:
-        for item in tqdm(ds, desc="Processing samples", unit="sample"):
-            writer.write(item)
-            lines += 1
-    print(f"Wrote {lines} lines from {len(parquet_files)} files to MDS files in {output_dir}")
-    if pigz:
-        file_paths = []
-        for file in os.listdir(output_dir):
-            if file.endswith(".mds"):
-                file_paths.append(os.path.join(output_dir, file))
-        for file_path in tqdm(file_paths, desc="Compressing with pigz", unit="file"):
-            pigz_compress(file_path, file_path + ".gz", processes, buf_size=buf_size, keep=False, quiet=True)
-        output_dir
-        print(f"Compressed {output_dir} with pigz")