mldataforge 0.0.1__tar.gz → 0.0.2__tar.gz

{mldataforge-0.0.1 → mldataforge-0.0.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mldataforge
-Version: 0.0.1
+Version: 0.0.2
 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
@@ -13,7 +13,7 @@ Requires-Python: >=3.12
 Requires-Dist: click
 Requires-Dist: datasets
 Requires-Dist: mltiming
-Requires-Dist: pygz
+Requires-Dist: mosaicml-streaming
 Description-Content-Type: text/markdown
 
 # mldatasets

{mldataforge-0.0.1 → mldataforge-0.0.2}/mldataforge/commands/convert/__init__.py

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

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

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

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

@@ -0,0 +1,50 @@
+#!/usr/bin/env python
+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:
+        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")

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

@@ -0,0 +1,26 @@
+#!/usr/bin/env python
+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.2/mldataforge/utils.py

@@ -0,0 +1,212 @@
+import atexit
+import bz2
+import click
+from datasets import concatenate_datasets, load_dataset
+import gzip
+import lzma
+from mltiming import timing
+import os
+import shutil
+import subprocess
+import tempfile
+from tqdm import tqdm
+
+__all__ = [
+    "check_overwrite",
+    "create_temp_file",
+    "determine_compression",
+    "infer_mds_encoding",
+    "infer_compression",
+    "load_parquet_files",
+    "open_jsonl",
+    "pigz_available",
+    "pigz_compress",
+    "use_pigz",
+]
+
+class PigzFile(object):
+    """A wrapper for pigz to handle gzip compression and decompression."""
+    def __init__(self, path, mode="rt", processes=4, encoding=None):
+        if mode not in ("rt", "wt", "rb", "wb"):
+            raise ValueError("Mode must be one of rt, wt, rb, or wb.")
+        self.path = path
+        self.mode = mode
+        self.processes = processes
+        self.encoding = "latin1" if mode[1] == "b" else ("utf-8" if encoding is None else encoding)
+        self._process = None
+        self._fw = None
+        if self.mode[0] == "r":
+            args = ["pigz", "-d", "-c", "-p", str(self.processes), self.path]
+            self._process = subprocess.Popen(
+                args, 
+                stdout=subprocess.PIPE,
+                encoding=encoding,
+            )
+        elif self.mode[0] == "w":
+            args = ["pigz", "-p", str(self.processes), "-c"]
+            self._fw = open(self.path, "w+")
+            self._process = subprocess.Popen(
+                args, 
+                stdout=self._fw,
+                stdin=subprocess.PIPE,
+                encoding=encoding,
+            )
+        
+    def __iter__(self):
+        assert self.mode[0] == "r"
+        for line in self._process.stdout:
+            yield line
+        self._process.wait()
+        assert self._process.returncode == 0
+        self._process.stdout.close()
+        self._process = None
+        
+    def write(self, line):
+        assert self.mode[0] == "w"
+        self._process.stdin.write(line if self.mode[1] == "t" else line.encode(self.encoding))
+    
+    def close(self): 
+        if self._process:
+            if self.mode[0] == "r":
+                self._process.kill()
+                self._process.stdout.close()
+                self._process = None
+            elif self.mode[1] == "w":
+                self._process.stdin.close()
+                self._process.wait()
+                self._process = None
+                self._fw.close()
+                self._fw = None
+    
+    def __enter__(self):
+        return self
+    
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.close()
+
+def check_overwrite(output_path, overwrite, yes):
+    if os.path.exists(output_path):
+        if os.path.isfile(output_path):
+            if not overwrite:
+                raise click.BadParameter(f"Output file '{output_path}' already exists. Use --overwrite to overwrite.")
+            if not yes:
+                confirm_overwrite(f"Output file '{output_path}' already exists. Do you want to delete it?")
+            with timing(message=f"Deleting existing file '{output_path}'"):
+                os.remove(output_path)
+        elif os.path.isdir(output_path):
+            if not overwrite:
+                raise click.BadParameter(f"Output directory '{output_path}' already exists. Use --overwrite to overwrite.")
+            if not yes:
+                confirm_overwrite(f"Output directory '{output_path}' already exists. Do you want to delete this directory and all its contents?")
+            with timing(message=f"Deleting existing directory '{output_path}'"):
+                shutil.rmtree(output_path)
+        else:
+            raise click.BadParameter(f"Output path '{output_path}' exists but is neither a file nor a directory.")
+
+def confirm_overwrite(message):
+    print(message)
+    response = input("Are you sure you want to proceed? (yes/no): ")
+    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"):
+    if compression == "infer":
+        compression = infer_compression(file_path)
+    if compression == "none":
+        compression = None
+    return compression
+
+def infer_mds_encoding(value):
+    """Determine the MDS encoding for a given value."""
+    if isinstance(value, str):
+        return 'str'
+    if isinstance(value, int):
+        return 'int'
+    if isinstance(value, float):
+        return 'float32'
+    if isinstance(value, bool):
+        return 'bool'
+    return 'pkl'
+
+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():
+            return 'pigz'
+        return 'gzip'
+    if extension.endswith('.bz2'):
+        return 'bz2'
+    if extension.endswith('.xz'):
+        return 'xz'
+    if extension.endswith('.zip'):
+        return 'zip'
+    if extension.endswith('.zst'):
+        return 'zstd'
+    return None
+
+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"):
+    """Open a JSONL file, handling gzip compression if necessary."""
+    compression = determine_compression(file_path, compression)
+    if compression in ("gzip", "pigz"):
+        return gzip.open(file_path, mode)
+    if compression == "bz2":
+        return bz2.open(file_path, mode)
+    if compression == "xz":
+        return lzma.open(file_path, mode)
+    if compression is None:
+        return open(file_path, mode)
+    raise ValueError(f"Unsupported compression type: {compression}")
+
+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):
+    """Compress a file using pigz."""
+    size = os.stat(input_file).st_size
+    num_blocks = (size+buf_size-1) // buf_size
+    with open(input_file, "rt", encoding="latin1") as f_in, PigzFile(output_file, "wb", processes=processes) as f_out:
+        for _ in tqdm(range(num_blocks), desc="Compressing with pigz", unit="block", disable=quiet):
+            buf = f_in.read(buf_size)
+            assert buf
+            f_out.write(buf)
+        buf = f_in.read()
+        assert not buf
+    if not keep:
+        os.remove(input_file)
+        if not quiet:
+            print(f"Removed {input_file}")
+
+def use_pigz(compression):
+    """Determine if pigz should be used based on the compression type."""
+    return compression == "pigz" or (compression == "gzip" and pigz_available())

{mldataforge-0.0.1 → mldataforge-0.0.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "mldataforge"
-version = "0.0.1"
+version = "0.0.2"
 authors = [
   { name = "Peter Schneider-Kamp" }
 ]
@@ -22,7 +22,7 @@ dependencies = [
     'click',
     'datasets',
     'mltiming',
-    'pygz'
+    'mosaicml-streaming'
 ]
 
 [project.urls]

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

@@ -1,55 +0,0 @@
-#!/usr/bin/env python
-import atexit
-import click
-from datasets import load_dataset
-from mltiming import timing
-import os
-from pygz import PigzFile
-from shutil import which
-import tempfile
-
-def create_temp_file():
-    # 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 _cleanup_file(file_path):
-    try:
-        os.remove(file_path)
-    except OSError:
-        pass
-
-@click.command()
-@click.argument("parquet_file", type=click.Path(exists=True))
-@click.argument("output_file", type=click.Path(exists=False), required=False)
-@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("--threads", default=64, help="Number of processes to use for pigz compression (default: 64).")
-@click.option("--overwrite", is_flag=True, help="Overwrite existing JSONL files.")
-def jsonl(parquet_file, output_file, compression, threads, overwrite):
-    if os.path.exists(output_file) and not overwrite:
-        raise click.ClickException(f"Output file {output_file} already exists. Use --overwrite to overwrite.")
-    with timing(message=f"Loading from {parquet_file}"):
-        ds = load_dataset("parquet", data_files=parquet_file)
-    orig_output_file = None
-    if compression == "none":
-        compression = None
-    elif compression == "infer":
-        if output_file.endswith(".gz") and which("pigz") is not None:
-            compression = "pigz"
-    if compression == "pigz":
-        compression = None
-        orig_output_file = output_file
-        output_file = create_temp_file()
-    with timing(message=f"Saving to {output_file} with compression {compression}"):
-        ds["train"].to_json(output_file, orient="records", lines=True, compression=compression)
-    if orig_output_file is not None:
-        with timing(message=f"Compressing {output_file} to {orig_output_file} with pigz using {threads} threads"):
-            with open(output_file, "rt") as f_in, PigzFile(f"{orig_output_file}.gz", "wt", threads=threads) as f_out:
-                f_out.write(f_in.read())