mldataforge 0.1.4__tar.gz → 0.1.6__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mldataforge
-Version: 0.1.4
+Version: 0.1.6
 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
@@ -10,10 +10,15 @@ Classifier: License :: OSI Approved :: MIT License
 Classifier: Operating System :: OS Independent
 Classifier: Programming Language :: Python :: 3
 Requires-Python: >=3.12
+Requires-Dist: brotlicffi
 Requires-Dist: click
 Requires-Dist: datasets
+Requires-Dist: isal
+Requires-Dist: lz4
 Requires-Dist: mltiming
 Requires-Dist: mosaicml-streaming
+Requires-Dist: python-snappy
+Requires-Dist: zstandard
 Provides-Extra: all
 Requires-Dist: build; extra == 'all'
 Requires-Dist: pytest; extra == 'all'

mldataforge-0.1.6/mldataforge/brotli.py

@@ -0,0 +1,82 @@
+import brotlicffi as brotli
+import io
+
+__all__ = ["brotli_open"]
+
+def brotli_open(filename, mode='rb', encoding='utf-8', compress_level=11):
+    return BrotliFile(filename, mode=mode, encoding=encoding, compress_level=11)
+
+import brotlicffi as brotli
+import io
+
+__all__ = ["brotli_open"]
+
+class BrotliFile:
+    def __init__(self, filename, mode='rb', encoding='utf-8', compress_level=11):
+        self.filename = filename
+        self.mode = mode
+        self.encoding = encoding
+        self.compress_level = compress_level
+
+        self.binary = 'b' in mode
+        file_mode = mode.replace('t', 'b')
+        self.file = open(filename, file_mode)
+
+        if 'r' in mode:
+            self._decompressor = brotli.Decompressor()
+            self._stream = self._wrap_reader()
+        elif 'w' in mode:
+            self._compressor = brotli.Compressor(quality=compress_level)
+            self._stream = self._wrap_writer()
+        else:
+            raise ValueError("Unsupported mode (use 'rb', 'wb', 'rt', or 'wt')")
+
+    def _wrap_reader(self):
+        buffer = io.BytesIO()
+        while True:
+            chunk = self.file.read(8192)
+            if not chunk:
+                break
+            buffer.write(self._decompressor.process(chunk))
+        buffer.seek(0)
+        return buffer if self.binary else io.TextIOWrapper(buffer, encoding=self.encoding)
+
+    def _wrap_writer(self):
+        return self if self.binary else io.TextIOWrapper(self, encoding=self.encoding)
+
+    def write(self, data):
+        if isinstance(data, str):
+            data = data.encode(self.encoding)
+        compressed = self._compressor.process(data)
+        self.file.write(compressed)
+        return len(data)
+
+    def flush(self):
+        if hasattr(self, '_compressor'):
+            self.file.write(self._compressor.finish())
+        self.file.flush()
+
+    def read(self, *args, **kwargs):
+        return self._stream.read(*args, **kwargs)
+
+    def readline(self, *args, **kwargs):
+        return self._stream.readline(*args, **kwargs)
+
+    def __iter__(self):
+        return iter(self._stream)
+
+    def close(self):
+        try:
+            if hasattr(self._stream, 'flush'):
+                self._stream.flush()
+        finally:
+            self.file.close()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.close()
+
+    def tell(self):
+        return self._stream.tell()

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

@@ -1,10 +1,11 @@
 import click
 from datasets import load_dataset
 
+from ...compression import *
 from ...options import *
 from ...utils import *
 
-__all__ = ["jsonl"]
+__all__ = ["jsonl_to_mds", "jsonl_to_parquet"]
 
 @click.group()
 def jsonl():
@@ -13,35 +14,40 @@ def jsonl():
 @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'])
+@compression_option(MDS_COMPRESSIONS)
 @overwrite_option()
 @yes_option()
 @processes_option()
 @buf_size_option()
 @shard_size_option()
-def mds(output_dir, jsonl_files, compression, processes, overwrite, yes, buf_size, shard_size):
+@no_pigz_option()
+def mds(**kwargs):
+    jsonl_to_mds(**kwargs)
+def jsonl_to_mds(output_dir, jsonl_files, compression, processes, overwrite, yes, buf_size, shard_size, no_pigz):
     check_arguments(output_dir, overwrite, yes, jsonl_files)
     save_mds(
-        load_dataset("json", data_files=jsonl_files, split="train"),
+        load_jsonl_files(jsonl_files),
         output_dir,
         processes=processes,
         compression=compression,
         buf_size=buf_size,
-        pigz=use_pigz(compression),
+        pigz=use_pigz(compression, no_pigz),
         shard_size=shard_size,
     )
 
 @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"])
+@compression_option(PARQUET_COMPRESSIONS)
 @overwrite_option()
 @yes_option()
 @batch_size_option()
-def parquet(output_file, jsonl_files, compression, overwrite, yes, batch_size):
+def parquet(**kwargs):
+    jsonl_to_parquet(**kwargs)
+def jsonl_to_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"),
+        load_jsonl_files(jsonl_files),
         output_file,
         compression=compression,
         batch_size=batch_size,

{mldataforge-0.1.4 → mldataforge-0.1.6}/mldataforge/commands/convert/mds.py

@@ -1,9 +1,10 @@
 import click
 
+from ...compression import *
 from ...options import *
 from ...utils import *
 
-__all__ = ["mds"]
+__all__ = ["mds_to_jsonl", "mds_to_parquet"]
 
 @click.group()
 def mds():
@@ -12,13 +13,15 @@ def mds():
 @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"])
+@compression_option(JSONL_COMPRESSIONS)
 @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):
+def jsonl(**kwargs):
+    mds_to_jsonl(**kwargs)
+def mds_to_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),
@@ -30,12 +33,14 @@ def jsonl(output_file, mds_directories, compression, processes, overwrite, yes,
 @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"])
+@compression_option(PARQUET_COMPRESSIONS)
 @overwrite_option()
 @yes_option()
 @batch_size_option()
 @no_bulk_option()
-def parquet(output_file, mds_directories, compression, overwrite, yes, batch_size, no_bulk):
+def parquet(**kwargs):
+    mds_to_parquet(**kwargs)
+def mds_to_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),

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

@@ -1,10 +1,11 @@
 import click
 from datasets import load_dataset
 
+from ...compression import *
 from ...options import *
 from ...utils import *
 
-__all__ = ["parquet"]
+__all__ = ["parquet_to_jsonl", "parquet_to_mds"]
 
 @click.group()
 def parquet():
@@ -13,11 +14,13 @@ def parquet():
 @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"])
+@compression_option(JSONL_COMPRESSIONS)
 @processes_option()
 @overwrite_option()
 @yes_option()
-def jsonl(output_file, parquet_files, compression, processes, overwrite, yes):
+def jsonl(**kwargs):
+    parquet_to_jsonl(**kwargs)
+def parquet_to_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"),
@@ -29,13 +32,16 @@ def jsonl(output_file, parquet_files, compression, processes, overwrite, yes):
 @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'])
+@compression_option(MDS_COMPRESSIONS)
 @processes_option()
 @overwrite_option()
 @yes_option()
 @buf_size_option()
 @shard_size_option()
-def mds(output_dir, parquet_files, compression, processes, overwrite, yes, buf_size, shard_size):
+@no_pigz_option()
+def mds(**kwargs):
+    parquet_to_mds(**kwargs)
+def parquet_to_mds(output_dir, parquet_files, compression, processes, overwrite, yes, buf_size, shard_size, no_pigz):
     check_arguments(output_dir, overwrite, yes, parquet_files)
     save_mds(
         load_dataset("parquet", data_files=parquet_files, split="train"),
@@ -43,6 +49,6 @@ def mds(output_dir, parquet_files, compression, processes, overwrite, yes, buf_s
         processes=processes,
         compression=compression,
         buf_size=buf_size,
-        pigz=use_pigz(compression),
+        pigz=use_pigz(compression, no_pigz=no_pigz),
         shard_size=shard_size,
     )

{mldataforge-0.1.4 → mldataforge-0.1.6}/mldataforge/commands/join.py

@@ -1,10 +1,11 @@
 import click
 from datasets import load_dataset
 
+from ..compression import *
 from ..options import *
 from ..utils import *
 
-__all__ = ["join"]
+__all__ = ["join_jsonl", "join_mds", "join_parquet"]
 
 @click.group()
 def join():
@@ -13,14 +14,16 @@ def join():
 @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"])
+@compression_option(JSONL_COMPRESSIONS)
 @processes_option()
 @overwrite_option()
 @yes_option()
-def jsonl(output_file, jsonl_files, compression, processes, overwrite, yes):
+def jsonl(**kwargs):
+    join_jsonl(**kwargs)
+def join_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"),
+        load_jsonl_files(jsonl_files),
         output_file,
         compression=compression,
         processes=processes,
@@ -29,14 +32,19 @@ def jsonl(output_file, jsonl_files, compression, processes, overwrite, yes):
 @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'])
+@compression_option(MDS_COMPRESSIONS)
 @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):
+@shard_size_option()
+@no_pigz_option()
+def mds(**kwargs):
+    print(kwargs)
+    join_mds(**kwargs)
+def join_mds(output_dir, mds_directories, compression, processes, overwrite, yes, batch_size, buf_size, no_bulk, shard_size, no_pigz):
     check_arguments(output_dir, overwrite, yes, mds_directories)
     save_mds(
         load_mds_directories(mds_directories, batch_size=batch_size, bulk=not no_bulk),
@@ -44,17 +52,20 @@ def mds(output_dir, mds_directories, compression, processes, overwrite, yes, bat
         processes=processes,
         compression=compression,
         buf_size=buf_size,
-        pigz=use_pigz(compression),
+        shard_size=shard_size,
+        pigz=use_pigz(compression, no_pigz),
     )
 
 @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"])
+@compression_option(PARQUET_COMPRESSIONS)
 @overwrite_option()
 @yes_option()
 @batch_size_option()
-def parquet(output_file, parquet_files, compression, overwrite, yes, batch_size):
+def parquet(**kwargs):
+    join_parquet(**kwargs)
+def join_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"),

{mldataforge-0.1.4 → mldataforge-0.1.6}/mldataforge/commands/split.py

@@ -1,10 +1,11 @@
 import click
 from datasets import load_dataset
 
+from ..compression import *
 from ..options import *
 from ..utils import *
 
-__all__ = ["split"]
+__all__ = ["split_jsonl", "split_mds", "split_parquet"]
 
 @click.group()
 def split():
@@ -15,14 +16,14 @@ def split():
 @prefix_option()
 @output_dir_option()
 @size_hint_option()
-@compression_option("infer", ["none", "infer", "pigz", "gzip", "bz2", "xz"])
+@compression_option(JSONL_COMPRESSIONS)
 @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])}",
+        load_jsonl_files(jsonl_files),
+        output_file=f"{output_dir}/{prefix}{{part:04d}}.jsonl{extension_compression(compression, jsonl_files[0])}",
         compression=compression,
         processes=processes,
         size_hint=size_hint,
@@ -35,7 +36,7 @@ def jsonl(jsonl_files, prefix, output_dir, size_hint, compression, processes, ov
 @prefix_option()
 @output_dir_option()
 @size_hint_option()
-@compression_option(None, ['none', 'br', 'bz2', 'gzip', 'pigz', 'snappy', 'zstd'])
+@compression_option(MDS_COMPRESSIONS)
 @processes_option()
 @overwrite_option()
 @yes_option()
@@ -43,14 +44,15 @@ def jsonl(jsonl_files, prefix, output_dir, size_hint, compression, processes, ov
 @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):
+@no_pigz_option()
+def mds(mds_directories, prefix, output_dir, size_hint, compression, processes, overwrite, yes, buf_size, batch_size, no_bulk, shard_size, no_pigz):
     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),
+        pigz=use_pigz(compression, no_pigz),
         shard_size=shard_size,
         size_hint=size_hint,
         overwrite=overwrite,
@@ -62,7 +64,7 @@ def mds(mds_directories, prefix, output_dir, size_hint, compression, processes,
 @prefix_option()
 @output_dir_option()
 @size_hint_option()
-@compression_option("snappy", ["snappy", "gzip", "zstd"])
+@compression_option(PARQUET_COMPRESSIONS)
 @overwrite_option()
 @yes_option()
 @batch_size_option()

mldataforge-0.1.6/mldataforge/compression.py

@@ -0,0 +1,162 @@
+import bz2
+from isal import igzip as gzip
+import lz4
+import lzma
+import os
+import shutil
+from tqdm import tqdm
+import zstandard
+
+from .brotli import brotli_open
+from .pigz import pigz_open
+from .snappy import snappy_open
+
+__all__ = [
+    "JSONL_COMPRESSIONS",
+    "MDS_COMPRESSIONS",
+    "PARQUET_COMPRESSIONS",
+    "determine_compression",
+    "extension_compression",
+    "infer_compression",
+    "open_compression",
+    "pigz_available",
+    "pigz_compress",
+    "use_pigz",
+]
+
+JSONL_COMPRESSIONS = dict(
+    default="infer",
+    choices=["infer", "none", "bz2", "gzip", "lz4", "lzma", "pigz", "snappy", "xz", "zstd"],
+)
+MDS_COMPRESSIONS = dict(
+    default=None,
+    choices=["none", "brotli", "bz2", "gzip", "pigz", "snappy", "zstd", "sample::brotli", "sample::bz2", "sample::gzip", "sample::snappy", "sample::zstd"],
+)
+PARQUET_COMPRESSIONS = dict(
+    default="snappy",
+    choices=["snappy", "brotli", "gzip", "lz4", "zstd"],
+)
+
+def determine_compression(fmt, file_path, compression="infer", no_pigz=False):
+    if compression == "none":
+        return None
+    if fmt == "jsonl":
+        if compression == "infer":
+            compression = infer_compression(file_path)
+        if compression == "brotli":
+            return "br"
+        return compression
+    if fmt == "mds":
+        if compression == "infer":
+            raise ValueError()
+        if compression == "pigz" or (not no_pigz and compression == "gzip" and pigz_available()):
+            return None
+        if compression == "gzip":
+            return "gz"
+        if compression == "brotli":
+            return "br"
+        if compression == "sample::gzip":
+            return "gz"
+        if compression == "sample::brotli":
+            return "br"
+        return compression
+    if fmt == "parquet":
+        return compression
+    raise ValueError(f"Unsupported format: {format}")
+
+def extension_compression(compression, file_path):
+    """Get the file extension for the given compression type."""
+    if compression == "infer":
+        compression = infer_compression(file_path)
+    if compression == "brotli":
+        return ".br"
+    if compression == "bz2":
+        return ".bz2"
+    if compression in ("gzip", "pigz"):
+        return ".gz"
+    if compression == "lz4":
+        return ".lz4"
+    if compression == "lzma":
+        return ".lzma"
+    if compression == "snappy":
+        return ".snappy"
+    if compression == "xz":
+        return ".xz"
+    if compression == "zstd":
+        return ".zst"
+    if compression is None or compression == "none":
+        return ""
+    raise ValueError(f"Unsupported compression type: {compression}")
+
+def infer_compression(file_path, pigz=True):
+    """Infer the compression type from the file extension."""
+    extension = os.path.splitext(file_path)[1]
+    if extension.endswith('.br'):
+        return 'brotli'
+    if extension.endswith('.bz2'):
+        return 'bz2'
+    if extension.endswith('.gz'):
+        if pigz and pigz_available():
+            return 'pigz'
+        return 'gzip'
+    if extension.endswith('.lz4'):
+        return 'lz4'
+    if extension.endswith('.lzma'):
+        return 'lzma'
+    if extension.endswith('.snappy'):
+        return 'snappy'
+    if extension.endswith('.xz'):
+        return 'xz'
+    if extension.endswith('.zip'):
+        return 'zip'
+    if extension.endswith('.zst'):
+        return 'zstd'
+    return None
+
+def open_compression(file_path, mode="rt", compression="infer", processes=64):
+    """Open a file, handling compression if necessary."""
+    if compression == "infer":
+        compression = infer_compression(file_path)
+    if compression in ("brotli", "br"):
+        return brotli_open(file_path, mode)
+    if compression in ("gzip", "gz"):
+        return gzip.open(file_path, mode)
+    if compression == "pigz":
+        return pigz_open(file_path, mode, processes=processes) if mode[0] == "w" else gzip.open(file_path, mode)
+    if compression == "bz2":
+        return bz2.open(file_path, mode)
+    if compression == "lz4":
+        return lz4.frame.open(file_path, mode)
+    if compression in ("lzma", "xz"):
+        return lzma.open(file_path, mode)
+    if compression == "snappy":
+        return snappy_open(file_path, mode)
+    if compression == "zstd":
+        return zstandard.open(file_path, mode)
+    if compression is None or compression == "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, "rb") as f_in, pigz_open(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, no_pigz=False):
+    """Determine if pigz should be used based on the compression type."""
+    return compression == "pigz" or (not no_pigz and compression == "gzip" and pigz_available())