mldataforge 0.1.0__tar.gz → 0.1.1__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mldataforge
-Version: 0.1.0
+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.0 → mldataforge-0.1.1}/mldataforge/commands/convert/mds.py

@@ -17,10 +17,11 @@ def mds():
 @overwrite_option()
 @yes_option()
 @batch_size_option()
-def jsonl(output_file, mds_directories, compression, processes, overwrite, yes, batch_size):
+@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),
+        load_mds_directories(mds_directories, batch_size=batch_size, bulk=not no_bulk),
         output_file,
         compression=compression,
         processes=processes,
@@ -28,15 +29,16 @@ 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("parquet_files", type=click.Path(exists=True), required=True, nargs=-1)
+@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()
-def parquet(output_file, parquet_files, compression, overwrite, yes, batch_size):
-    check_arguments(output_file, overwrite, yes, parquet_files)
+@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(parquet_files, batch_size=batch_size),
+        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/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.0 → mldataforge-0.1.1}/mldataforge/options.py

@@ -29,6 +29,16 @@ def buf_size_option(default=2**24):
         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.

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

@@ -12,6 +12,7 @@ import shutil
 from streaming import MDSWriter, StreamingDataset
 from tqdm import tqdm
 
+from .mds import MDSBulkReader
 from .pigz import pigz_open
 
 __all__ = [
@@ -98,7 +99,9 @@ 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 mds_directories:
         ds = StreamingDataset(

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

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