mldataforge 0.1.1__tar.gz → 0.1.3__tar.gz

mldataforge-0.1.3/PKG-INFO

@@ -0,0 +1,62 @@
+Metadata-Version: 2.4
+Name: mldataforge
+Version: 0.1.3
+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
+Author: Peter Schneider-Kamp
+License-File: LICENSE
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Requires-Python: >=3.12
+Requires-Dist: click
+Requires-Dist: datasets
+Requires-Dist: mltiming
+Requires-Dist: mosaicml-streaming
+Provides-Extra: all
+Requires-Dist: build; extra == 'all'
+Requires-Dist: pytest; extra == 'all'
+Requires-Dist: pytest-dependency; extra == 'all'
+Requires-Dist: twine; extra == 'all'
+Provides-Extra: dev
+Requires-Dist: build; extra == 'dev'
+Requires-Dist: twine; extra == 'dev'
+Provides-Extra: test
+Requires-Dist: pytest; extra == 'test'
+Requires-Dist: pytest-dependency; extra == 'test'
+Description-Content-Type: text/markdown
+
+# mldatasets
+swiss army knife of scripts for transforming and processing datasets for machine learning
+
+## 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
+python -m mldataforge --help
+```
+
+## usage example: converting MosaiclML Dataset (MDS) to Parquet format
+```
+Usage: python -m mldataforge convert mds parquet [OPTIONS] OUTPUT_FILE
+                                                 MDS_DIRECTORIES...
+
+Options:
+  --compression [snappy|gzip|zstd]
+                                  Compress the output file (default: snappy).
+  --overwrite                     Overwrite existing path.
+  --yes                           Assume yes to all prompts. Use with caution
+                                  as it will remove files or even entire
+                                  directories without confirmation.
+  --batch-size INTEGER            Batch size for loading data and writing
+                                  files (default: 65536).
+  --no-bulk                       Use a custom space and time-efficient bulk
+                                  reader (only gzip and no compression).
+  --help                          Show this message and exit.
+```

mldataforge-0.1.3/README.md

@@ -0,0 +1,33 @@
+# mldatasets
+swiss army knife of scripts for transforming and processing datasets for machine learning
+
+## 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
+python -m mldataforge --help
+```
+
+## usage example: converting MosaiclML Dataset (MDS) to Parquet format
+```
+Usage: python -m mldataforge convert mds parquet [OPTIONS] OUTPUT_FILE
+                                                 MDS_DIRECTORIES...
+
+Options:
+  --compression [snappy|gzip|zstd]
+                                  Compress the output file (default: snappy).
+  --overwrite                     Overwrite existing path.
+  --yes                           Assume yes to all prompts. Use with caution
+                                  as it will remove files or even entire
+                                  directories without confirmation.
+  --batch-size INTEGER            Batch size for loading data and writing
+                                  files (default: 65536).
+  --no-bulk                       Use a custom space and time-efficient bulk
+                                  reader (only gzip and no compression).
+  --help                          Show this message and exit.
+```

mldataforge-0.1.3/mldataforge/__main__.py

@@ -0,0 +1,4 @@
+from .commands import cli
+
+if __name__ == "__main__":
+    cli()

mldataforge-0.1.3/mldataforge/commands/__init__.py

@@ -0,0 +1,15 @@
+import click
+
+from .convert import convert
+from .join import join
+from .split import split
+
+__all__ = ["cli"]
+
+@click.group()
+def cli():
+    pass
+
+cli.add_command(convert)
+cli.add_command(join)
+cli.add_command(split)

{mldataforge-0.1.1 → mldataforge-0.1.3}/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.1 → mldataforge-0.1.3}/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.3/mldataforge/commands/join.py

@@ -0,0 +1,64 @@
+import click
+from datasets import load_dataset
+
+from ..options import *
+from ..utils import *
+
+__all__ = ["join"]
+
+@click.group()
+def join():
+    pass
+
+@join.command()
+@click.argument("output_file", type=click.Path(exists=False), required=True)
+@click.argument("jsonl_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, jsonl_files, compression, processes, overwrite, yes):
+    check_arguments(output_file, overwrite, yes, jsonl_files)
+    save_jsonl(
+        load_dataset("json", data_files=jsonl_files, split="train"),
+        output_file,
+        compression=compression,
+        processes=processes,
+    )
+
+@join.command()
+@click.argument("output_dir", type=click.Path(exists=False), required=True)
+@click.argument("mds_directories", type=click.Path(exists=True), required=True, nargs=-1)
+@compression_option(None, ['none', 'br', 'bz2', 'gzip', 'pigz', 'snappy', 'zstd'])
+@processes_option()
+@overwrite_option()
+@yes_option()
+@batch_size_option()
+@buf_size_option()
+@no_bulk_option()
+def mds(output_dir, mds_directories, compression, processes, overwrite, yes, batch_size, buf_size, no_bulk):
+    check_arguments(output_dir, overwrite, yes, mds_directories)
+    save_mds(
+        load_mds_directories(mds_directories, batch_size=batch_size, bulk=not no_bulk),
+        output_dir,
+        processes=processes,
+        compression=compression,
+        buf_size=buf_size,
+        pigz=use_pigz(compression),
+    )
+
+@join.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("snappy", ["snappy", "gzip", "zstd"])
+@overwrite_option()
+@yes_option()
+@batch_size_option()
+def parquet(output_file, parquet_files, compression, overwrite, yes, batch_size):
+    check_arguments(output_file, overwrite, yes, parquet_files)
+    save_parquet(
+        load_dataset("parquet", data_files=parquet_files, split="train"),
+        output_file,
+        compression=compression,
+        batch_size=batch_size,
+    )

mldataforge-0.1.3/mldataforge/commands/split.py

@@ -0,0 +1,58 @@
+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,
+    )

{mldataforge-0.1.1 → mldataforge-0.1.3}/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.1 → mldataforge-0.1.3}/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.1 → mldataforge-0.1.3}/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,43 +176,67 @@ 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"))
-
-def save_mds(it, output_dir, processes=64, compression=None, buf_size=2**24, pigz=False):
+    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, 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)
+        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
     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")
+        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):
     writer = None

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

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "mldataforge"
-version = "0.1.1"
+version = "0.1.3"
 authors = [
   { name = "Peter Schneider-Kamp" }
 ]
@@ -25,6 +25,11 @@ dependencies = [
     'mosaicml-streaming'
 ]
 
+[project.optional-dependencies]
+test = ["pytest", "pytest-dependency"]
+dev = ["build", "twine"]
+all = ["build", "twine", "pytest", "pytest-dependency"]
+
 [project.urls]
 "Homepage" = "https://github.com/schneiderkamplab/mldataforge"
 "Bug Tracker" = "https://github.com/schneiderkamplab/mldataforge/issues"

mldataforge-0.1.1/PKG-INFO

@@ -1,20 +0,0 @@
-Metadata-Version: 2.4
-Name: mldataforge
-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
-Author: Peter Schneider-Kamp
-License-File: LICENSE
-Classifier: License :: OSI Approved :: MIT License
-Classifier: Operating System :: OS Independent
-Classifier: Programming Language :: Python :: 3
-Requires-Python: >=3.12
-Requires-Dist: click
-Requires-Dist: datasets
-Requires-Dist: mltiming
-Requires-Dist: mosaicml-streaming
-Description-Content-Type: text/markdown
-
-# mldatasets
-swiss army knife of scripts for transforming and processing datasets for machine learning

mldataforge-0.1.1/README.md

@@ -1,2 +0,0 @@
-# mldatasets
-swiss army knife of scripts for transforming and processing datasets for machine learning

mldataforge-0.1.1/mldataforge/__main__.py

@@ -1,12 +0,0 @@
-import click
-
-from .commands import convert
-
-@click.group()
-def cli():
-    pass
-
-cli.add_command(convert)
-
-if __name__ == "__main__":
-    cli()

mldataforge-0.1.1/mldataforge/commands/__init__.py

@@ -1,3 +0,0 @@
-from .convert import convert
-
-__all__ = ["convert"]