mldataforge 0.1.2__tar.gz → 0.1.4__tar.gz

{mldataforge-0.1.2 → mldataforge-0.1.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mldataforge
-Version: 0.1.2
+Version: 0.1.4
 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
@@ -30,9 +30,12 @@ Description-Content-Type: text/markdown
 # mldatasets
 swiss army knife of scripts for transforming and processing datasets for machine learning
 
-## scope
+## conversion
 Currently, mldataforge provides space- and time-efficient conversions between JSONL (with or without compression), MosaiclML Dataset (MDS format), and Parquet. The implementations handle conversions by individual samples or small batches of samples and make efficient use of multi-core architectures where possible. Consequently, mldataforge is an excellent choice when transforming TB-scale datasets on data processing nodes with many cores.
 
+## splitting
+Currently, mldataforge provides space- and time-efficient splitting of JSONL (with or without compression). The implementations handle conversions by individual samples or small batches of samples and make efficient use of multi-core architectures where possible. The splitting function can take an already splitted dataset and re-split it with a different granularity.
+
 ## installation and general usage
 ```
 pip install mldataforge

{mldataforge-0.1.2 → mldataforge-0.1.4}/README.md

@@ -1,9 +1,12 @@
 # mldatasets
 swiss army knife of scripts for transforming and processing datasets for machine learning
 
-## scope
+## conversion
 Currently, mldataforge provides space- and time-efficient conversions between JSONL (with or without compression), MosaiclML Dataset (MDS format), and Parquet. The implementations handle conversions by individual samples or small batches of samples and make efficient use of multi-core architectures where possible. Consequently, mldataforge is an excellent choice when transforming TB-scale datasets on data processing nodes with many cores.
 
+## splitting
+Currently, mldataforge provides space- and time-efficient splitting of JSONL (with or without compression). The implementations handle conversions by individual samples or small batches of samples and make efficient use of multi-core architectures where possible. The splitting function can take an already splitted dataset and re-split it with a different granularity.
+
 ## installation and general usage
 ```
 pip install mldataforge

{mldataforge-0.1.2 → mldataforge-0.1.4}/mldataforge/commands/__init__.py

@@ -2,6 +2,7 @@ import click
 
 from .convert import convert
 from .join import join
+from .split import split
 
 __all__ = ["cli"]
 
@@ -11,3 +12,4 @@ def cli():
 
 cli.add_command(convert)
 cli.add_command(join)
+cli.add_command(split)

{mldataforge-0.1.2 → mldataforge-0.1.4}/mldataforge/commands/convert/jsonl.py

@@ -18,7 +18,8 @@ def jsonl():
 @yes_option()
 @processes_option()
 @buf_size_option()
-def mds(output_dir, jsonl_files, compression, processes, overwrite, yes, buf_size):
+@shard_size_option()
+def mds(output_dir, jsonl_files, compression, processes, overwrite, yes, buf_size, shard_size):
     check_arguments(output_dir, overwrite, yes, jsonl_files)
     save_mds(
         load_dataset("json", data_files=jsonl_files, split="train"),
@@ -27,6 +28,7 @@ def mds(output_dir, jsonl_files, compression, processes, overwrite, yes, buf_siz
         compression=compression,
         buf_size=buf_size,
         pigz=use_pigz(compression),
+        shard_size=shard_size,
     )
 
 @jsonl.command()

{mldataforge-0.1.2 → mldataforge-0.1.4}/mldataforge/commands/convert/parquet.py

@@ -34,7 +34,8 @@ def jsonl(output_file, parquet_files, compression, processes, overwrite, yes):
 @overwrite_option()
 @yes_option()
 @buf_size_option()
-def mds(output_dir, parquet_files, compression, processes, overwrite, yes, buf_size):
+@shard_size_option()
+def mds(output_dir, parquet_files, compression, processes, overwrite, yes, buf_size, shard_size):
     check_arguments(output_dir, overwrite, yes, parquet_files)
     save_mds(
         load_dataset("parquet", data_files=parquet_files, split="train"),
@@ -43,4 +44,5 @@ def mds(output_dir, parquet_files, compression, processes, overwrite, yes, buf_s
         compression=compression,
         buf_size=buf_size,
         pigz=use_pigz(compression),
+        shard_size=shard_size,
     )

mldataforge-0.1.4/mldataforge/commands/split.py

@@ -0,0 +1,78 @@
+import click
+from datasets import load_dataset
+
+from ..options import *
+from ..utils import *
+
+__all__ = ["split"]
+
+@click.group()
+def split():
+    pass
+
+@split.command()
+@click.argument("jsonl_files", type=click.Path(exists=True), required=True, nargs=-1)
+@prefix_option()
+@output_dir_option()
+@size_hint_option()
+@compression_option("infer", ["none", "infer", "pigz", "gzip", "bz2", "xz"])
+@processes_option()
+@overwrite_option()
+@yes_option()
+def jsonl(jsonl_files, prefix, output_dir, size_hint, compression, processes, overwrite, yes):
+    save_jsonl(
+        load_dataset("json", data_files=jsonl_files, split="train"),
+        output_file=f"{output_dir}/{prefix}{{part:04d}}.jsonl{extension(compression, jsonl_files[0])}",
+        compression=compression,
+        processes=processes,
+        size_hint=size_hint,
+        overwrite=overwrite,
+        yes=yes,
+    )
+
+@split.command()
+@click.argument("mds_directories", type=click.Path(exists=True), required=True, nargs=-1)
+@prefix_option()
+@output_dir_option()
+@size_hint_option()
+@compression_option(None, ['none', 'br', 'bz2', 'gzip', 'pigz', 'snappy', 'zstd'])
+@processes_option()
+@overwrite_option()
+@yes_option()
+@buf_size_option()
+@batch_size_option()
+@no_bulk_option()
+@shard_size_option()
+def mds(mds_directories, prefix, output_dir, size_hint, compression, processes, overwrite, yes, buf_size, batch_size, no_bulk, shard_size):
+    save_mds(
+        load_mds_directories(mds_directories, batch_size=batch_size, bulk=not no_bulk),
+        output_dir=f"{output_dir}/{prefix}{{part:04d}}",
+        processes=processes,
+        compression=compression,
+        buf_size=buf_size,
+        pigz=use_pigz(compression),
+        shard_size=shard_size,
+        size_hint=size_hint,
+        overwrite=overwrite,
+        yes=yes,
+    )
+
+@split.command()
+@click.argument("parquet_files", type=click.Path(exists=True), required=True, nargs=-1)
+@prefix_option()
+@output_dir_option()
+@size_hint_option()
+@compression_option("snappy", ["snappy", "gzip", "zstd"])
+@overwrite_option()
+@yes_option()
+@batch_size_option()
+def parquet(parquet_files, prefix, output_dir, size_hint, compression, overwrite, yes, batch_size):
+    save_parquet(
+        load_dataset("parquet", data_files=parquet_files, split="train"),
+        output_file=f"{output_dir}/{prefix}{{part:04d}}.parquet",
+        compression=compression,
+        batch_size=batch_size,
+        size_hint=size_hint,
+        overwrite=overwrite,
+        yes=yes,
+    )

{mldataforge-0.1.2 → mldataforge-0.1.4}/mldataforge/options.py

@@ -50,6 +50,17 @@ def compression_option(default, choices):
         help=f"Compress the output file (default: {default}).",
     )
 
+def output_dir_option(default="."):
+    """
+    Option for specifying the output directory.
+    """
+    return click.option(
+        "--output-dir",
+        default=default,
+        type=click.Path(exists=False),
+        help="Output directory.",
+    )
+
 def overwrite_option():
     """
     Option for specifying whether to overwrite existing files.
@@ -60,6 +71,16 @@ def overwrite_option():
         help="Overwrite existing path.",
     )
 
+def prefix_option(default="part-"):
+    """
+    Option for specifying the prefix for output files.
+    """
+    return click.option(
+        "--prefix",
+        default=default,
+        help=f"Prefix for output files (default: {default}).",
+    )
+
 def processes_option(default=64):
     """
     Option for specifying the number of processes to use.
@@ -70,6 +91,26 @@ def processes_option(default=64):
         help=f"Number of processes to use (default: {default}).",
     )
 
+def shard_size_option(default=2**26):
+    """
+    Option for specifying the shard size.
+    """
+    return click.option(
+        "--shard-size",
+        default=default,
+        help=f"Shard size for the dataset (default: {default}).",
+    )
+
+def size_hint_option(default=2**26):
+    """
+    Option for specifying the size hint.
+    """
+    return click.option(
+        "--size-hint",
+        default=default,
+        help=f"Size hint for the dataset (default: {default}).",
+    )
+
 def yes_option():
     """
     Option for specifying whether to assume yes to all prompts.

{mldataforge-0.1.2 → mldataforge-0.1.4}/mldataforge/pigz.py

@@ -16,6 +16,7 @@ class PigzFile(object):
         self.encoding = encoding if self.is_text else None
         self._process = None
         self._fw = None
+        self.offset = 0
         args = ["pigz", "-p", str(self.processes), "-c"]
         if self.is_read:
             args.extend(("-d", self.path))
@@ -28,6 +29,7 @@ class PigzFile(object):
         assert self.is_read
         for line in self._process.stdout:
             assert isinstance(line, str) if self.is_text else isinstance(line, bytes)
+            self.offset += len(line)
             yield line
         self._process.wait()
         assert self._process.returncode == 0
@@ -39,6 +41,7 @@ class PigzFile(object):
         assert self._fw is not None
         assert isinstance(line, str) if self.is_text else isinstance(line, bytes)
         self._process.stdin.write(line)
+        self.offset += len(line)
     
     def close(self): 
         if self._process:
@@ -52,7 +55,10 @@ class PigzFile(object):
                 self._process = None
                 self._fw.close()
                 self._fw = None
-    
+
+    def tell(self):
+        return self.offset
+
     def __enter__(self):
         return self
     

{mldataforge-0.1.2 → mldataforge-0.1.4}/mldataforge/utils.py

@@ -19,6 +19,7 @@ __all__ = [
     "batch_iterable",
     "check_arguments",
     "confirm_overwrite",
+    "extension",
     "load_mds_directories",
     "save_jsonl",
     "save_mds",
@@ -36,8 +37,8 @@ def batch_iterable(iterable, batch_size):
     if batch:
         yield batch
 
-def check_arguments(output_path, overwrite, yes, input_paths):
-    if not input_paths:
+def check_arguments(output_path, overwrite, yes, input_paths=None):
+    if input_paths is not None and not input_paths:
         raise click.BadArgumentUsage("No input paths provided.")
     if os.path.exists(output_path):
         if os.path.isfile(output_path):
@@ -70,6 +71,24 @@ def _determine_compression(file_path, compression="infer"):
         compression = None
     return compression
 
+def extension(compression, file_path):
+    """Get the file extension for the given compression type."""
+    if compression == "infer":
+        compression = _infer_compression(file_path)
+    if compression in ("gzip", "pigz"):
+        return ".gz"
+    if compression == "bz2":
+        return ".bz2"
+    if compression == "xz":
+        return ".xz"
+    if compression == "zip":
+        return ".zip"
+    if compression == "zstd":
+        return ".zst"
+    if compression is None:
+        return ""
+    raise ValueError(f"Unsupported compression type: {compression}")
+
 def _infer_mds_encoding(value):
     """Determine the MDS encoding for a given value."""
     if isinstance(value, str):
@@ -157,53 +176,88 @@ def _pigz_compress(input_file, output_file, processes=64, buf_size=2**24, keep=F
         if not quiet:
             print(f"Removed {input_file}")
 
-def save_jsonl(iterable, output_file, compression=None, processes=64):
+def save_jsonl(iterable, output_file, compression=None, processes=64, size_hint=None, overwrite=True, yes=True):
     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"))
+    f = None
+    part = 0
+    for item in tqdm(iterable, desc="Writing to JSONL", unit="sample"):
+        if f is None:
+            part_file = output_file.format(part=part)
+            check_arguments(part_file, overwrite, yes)
+            f = _open_jsonl(part_file, mode="wb", compression=compression, processes=processes)
+        f.write(f"{json.dumps(item)}\n".encode("utf-8"))
+        if size_hint is not None and f.tell() >= size_hint:
+            f.close()
+            part += 1
+            f = None
+    if f is not None:
+        f.close()
 
-def save_mds(it, output_dir, processes=64, compression=None, buf_size=2**24, pigz=False):
+def save_mds(it, output_dir, processes=64, compression=None, buf_size=2**24, pigz=False, shard_size=None, size_hint=None, overwrite=True, yes=True):
     if compression == "none" or pigz:
         compression = None
     if compression == "gzip":
         compression = "gz"
     writer = None
+    part = 0
+    files = []
     for sample in tqdm(it, desc="Writing to MDS", unit="sample"):
         if writer is None:
+            part_dir = output_dir.format(part=part)
+            check_arguments(part_dir, overwrite, yes)
+            files.append(part_dir)
             columns = {key: _infer_mds_encoding(value) for key, value in sample.items()}
-            writer = MDSWriter(out=output_dir, columns=columns, compression=compression)
+            writer = MDSWriter(out=part_dir, columns=columns, compression=compression, size_limit=shard_size)
+            offset = 0
+        prev = writer.new_shard_size
         writer.write(sample)
-    writer.finish()
+        offset += (writer.new_shard_size - prev) if prev < writer.new_shard_size else writer.new_shard_size
+        if size_hint is not None and offset >= size_hint:
+            writer.finish()
+            part += 1
+            writer = None
+    if writer is not None:
+        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):
+        for output_dir in files:
+            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, size_hint=None, overwrite=True, yes=True):
     writer = None
+    part = 0
     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)
+            part_file = output_file.format(part=part)
+            check_arguments(part_file, overwrite, yes)
+            writer = pq.ParquetWriter(part_file, table.schema, compression=compression)
+            offset = 0
         writer.write_table(table)
-    writer.close()
+        offset += table.nbytes
+        if size_hint is not None and offset >= size_hint:
+            writer.close()
+            part += 1
+            writer = None
+    if writer is not None:
+        writer.close()
 
 def use_pigz(compression):
     """Determine if pigz should be used based on the compression type."""

{mldataforge-0.1.2 → mldataforge-0.1.4}/pyproject.toml

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