mldataforge 0.2.0__tar.gz → 0.2.2__tar.gz

{mldataforge-0.2.0 → mldataforge-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mldataforge
-Version: 0.2.0
+Version: 0.2.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

{mldataforge-0.2.0 → mldataforge-0.2.2}/mldataforge/commands/convert/mds.py

@@ -20,12 +20,13 @@ def mds():
 @batch_size_option()
 @no_bulk_option()
 @trafo_option()
+@shuffle_option()
 def jsonl(**kwargs):
     mds_to_jsonl(**kwargs)
-def mds_to_jsonl(output_file, mds_directories, compression, processes, overwrite, yes, batch_size, no_bulk, trafo):
+def mds_to_jsonl(output_file, mds_directories, compression, processes, overwrite, yes, batch_size, no_bulk, trafo, shuffle):
     check_arguments(output_file, overwrite, yes, mds_directories)
     save_jsonl(
-        load_mds_directories(mds_directories, batch_size=batch_size, bulk=not no_bulk),
+        load_mds_directories(mds_directories, batch_size=batch_size, bulk=not no_bulk, shuffle=shuffle),
         output_file,
         compression=compression,
         processes=processes,
@@ -41,12 +42,13 @@ def mds_to_jsonl(output_file, mds_directories, compression, processes, overwrite
 @batch_size_option()
 @no_bulk_option()
 @trafo_option()
+@shuffle_option()
 def parquet(**kwargs):
     mds_to_parquet(**kwargs)
-def mds_to_parquet(output_file, mds_directories, compression, overwrite, yes, batch_size, no_bulk, trafo):
+def mds_to_parquet(output_file, mds_directories, compression, overwrite, yes, batch_size, no_bulk, trafo, shuffle):
     check_arguments(output_file, overwrite, yes, mds_directories)
     save_parquet(
-        load_mds_directories(mds_directories, batch_size=batch_size, bulk=not no_bulk),
+        load_mds_directories(mds_directories, batch_size=batch_size, bulk=not no_bulk, shuffle=shuffle),
         output_file,
         compression=compression,
         batch_size=batch_size,

{mldataforge-0.2.0 → mldataforge-0.2.2}/mldataforge/commands/join.py

@@ -43,13 +43,14 @@ def join_jsonl(output_file, jsonl_files, compression, processes, overwrite, yes,
 @shard_size_option()
 @no_pigz_option()
 @trafo_option()
+@shuffle_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, trafo):
+def join_mds(output_dir, mds_directories, compression, processes, overwrite, yes, batch_size, buf_size, no_bulk, shard_size, no_pigz, trafo, shuffle):
     check_arguments(output_dir, overwrite, yes, mds_directories)
     save_mds(
-        load_mds_directories(mds_directories, batch_size=batch_size, bulk=not no_bulk),
+        load_mds_directories(mds_directories, batch_size=batch_size, bulk=not no_bulk, shuffle=shuffle),
         output_dir,
         processes=processes,
         compression=compression,

{mldataforge-0.2.0 → mldataforge-0.2.2}/mldataforge/commands/split.py

@@ -50,11 +50,12 @@ def split_jsonl(jsonl_files, prefix, output_dir, size_hint, compression, process
 @shard_size_option()
 @no_pigz_option()
 @trafo_option()
+@shuffle_option()
 def mds(*args, **kwargs):
     split_mds(*args, **kwargs)
-def split_mds(mds_directories, prefix, output_dir, size_hint, compression, processes, overwrite, yes, buf_size, batch_size, no_bulk, shard_size, no_pigz, trafo):
+def split_mds(mds_directories, prefix, output_dir, size_hint, compression, processes, overwrite, yes, buf_size, batch_size, no_bulk, shard_size, no_pigz, trafo, shuffle):
     save_mds(
-        load_mds_directories(mds_directories, batch_size=batch_size, bulk=not no_bulk),
+        load_mds_directories(mds_directories, batch_size=batch_size, bulk=not no_bulk, shuffle=shuffle),
         output_dir=f"{output_dir}/{prefix}{{part:04d}}",
         processes=processes,
         compression=compression,

mldataforge-0.2.2/mldataforge/indexing.py

@@ -0,0 +1,25 @@
+import numpy as np
+
+__all__ = ['IndexedDatasetView', 'shuffle_permutation']
+
+class IndexedDatasetView:
+    def __init__(self, dataset, indices):
+        self.dataset = dataset
+        self.indices = list(indices)  # ensure repeatable accessx
+
+    def __iter__(self):
+        for idx in self.indices:
+            yield self.dataset[idx]
+
+    def __len__(self):
+        return len(self.indices)
+
+def shuffle_permutation(n, seed=int):
+    rng = np.random.default_rng(seed)
+    return rng.permutation(n)
+
+def reverse_permutation(indices):
+    n = len(indices)
+    reverse_indices = np.empty(n, dtype=int)
+    reverse_indices[indices] = np.arange(n)
+    return reverse_indices

{mldataforge-0.2.0 → mldataforge-0.2.2}/mldataforge/options.py

@@ -13,6 +13,7 @@ __all__ = [
     "processes_option",
     "prefix_option",
     "shard_size_option",
+    "shuffle_option",
     "size_hint_option",
     "trafo_option",
     "yes_option",
@@ -120,6 +121,17 @@ def shard_size_option(default=2**26):
         help=f"Shard size for the dataset (default: {default}).",
     )
 
+def shuffle_option():
+    """
+    Option for specifying whether to shuffle the dataset by providing a random seed.
+    """
+    return click.option(
+        "--shuffle",
+        default=None,
+        type=int,
+        help="Shuffle the dataset by providing a random seed.",
+    )
+
 def size_hint_option(default=2**26):
     """
     Option for specifying the size hint.

{mldataforge-0.2.0 → mldataforge-0.2.2}/mldataforge/trafos.py

@@ -1,34 +1,73 @@
 import re
-from typing import Callable
-
-__all__ = ['Trafo', 'flatten_json', 'unflatten_json']
-
-class Trafo:
-    """
-    Base class for transformations.
-    """
-
-    def __init__(self, trafo: Callable | str | None):
-        self.trafo = trafo
-        if isinstance(trafo, str):
-            self.trafo = eval(trafo)
-
-    def __call__(self, obj):
-        return self.trafo(obj) if self.trafo else obj
-
-    def __repr__(self):
-        return f"{self.__class__.__name__}({self.trafo})"
 
+__all__ = ['Transformation', 'Transformations', 'flatten_json', 'identity', 'unflatten_json']
+
+class Transformation:
+    def __init__(self, code: str):
+        self.code = code
+        self._init_context()
+
+    def _init_context(self):
+        self.global_context = {}
+        exec(self.code, self.global_context)
+        if 'process' not in self.global_context or not callable(self.global_context['process']):
+            raise ValueError("code must define a callable named 'process'")
+        self.process = self.global_context['process']
+        self._flushable = hasattr(self.process, 'flushable') and self.process.flushable
+
+    def _normalize_outputs(self, result):
+        if result is None:
+            return []
+        if isinstance(result, (list, tuple, set)):
+            return list(result)
+        return [result]
+
+    def _flush(self):
+        if self._flushable:
+            while True:
+                flushed = self._normalize_outputs(self.process(None))
+                if not flushed:
+                    return
+                yield from flushed
+
+    def __call__(self, iterable):
+        for sample in iterable:
+            results = self._normalize_outputs(self.process(sample))
+            yield from results
+            if not results:
+                yield from self._flush()
+        if self._flushable:
+            yield from self._flush()
+
+    def __len__(self):
+        if self._last_input_len is not None:
+            return self._last_input_len
+        raise TypeError("Length is not available for this transformation.")
+
+class Transformations:
+    def __init__(self, codes: list[str], indices=None):
+        self.pipeline = [Transformation(code) for code in codes]
+
+    def __call__(self, dataset):
+        result = dataset
+        for transform in self.pipeline:
+            result = transform(result)
+        return result
+
+    def __len__(self):
+        if self.indices is not None:
+            return len(self.indices)
+        elif hasattr(self.pipeline[0], '_last_input_len') and self.pipeline[0]._last_input_len is not None:
+            return self.pipeline[0]._last_input_len
+        raise TypeError("Transformations length is not available until __call__ is used on a sized input.")
 
 def flatten_json(obj, parent_key='', sep='.', escape_char='\\'):
-    items = []
-
     def escape(key):
         return key.replace(escape_char, escape_char * 2)\
                   .replace(sep, escape_char + sep)\
                   .replace('[', escape_char + '[')\
                   .replace(']', escape_char + ']')
-
+    items = []
     if isinstance(obj, dict):
         if not obj:
             # explicitly handle empty dict
@@ -49,15 +88,15 @@ def flatten_json(obj, parent_key='', sep='.', escape_char='\\'):
         items.append((parent_key, obj))
     return dict(items)
 
+def identity(obj):
+    return obj
 
 def unflatten_json(flat_dict, sep='.', escape_char='\\'):
-
     def check_flat_json(obj):
         assert isinstance(obj, dict), "Input must be a dictionary"
         for k, v in obj.items():
             assert isinstance(k, str), f"Key {k} is not a string"
             assert isinstance(v, (str, int, float, bool)), f"Value {v} is not a valid JSON type"
-
     def parse_key(key):
         tokens = re.findall(r'(?:[^.\[\]\\]|\\.)+|\[\d+\]', key)
         parsed = []
@@ -70,11 +109,8 @@ def unflatten_json(flat_dict, sep='.', escape_char='\\'):
                                   .replace(escape_char + ']', ']')
                                   .replace(escape_char*2, escape_char))
         return parsed
-
     check_flat_json(flat_dict)
-
     result = {}
-
     for compound_key, value in flat_dict.items():
         keys = parse_key(compound_key)
         current = result
@@ -107,5 +143,4 @@ def unflatten_json(flat_dict, sep='.', escape_char='\\'):
                     current_parent = current
                     current = current[key]
             last_key = key
-
     return result

{mldataforge-0.2.0 → mldataforge-0.2.2}/mldataforge/utils.py

@@ -10,9 +10,10 @@ from streaming import StreamingDataset
 from tqdm import tqdm
 
 from .compression import determine_compression, open_compression, pigz_compress
+from .indexing import IndexedDatasetView, reverse_permutation, shuffle_permutation
 from .mds import MDSBulkReader, MDSWriter
 from .pigz import pigz_open
-from .trafos import Trafo
+from .trafos import Transformations
 
 __all__ = [
     "check_arguments",
@@ -89,7 +90,9 @@ def load_jsonl_files(jsonl_files):
         return _streaming_jsonl(jsonl_files, compressions)
     return load_dataset("json", data_files=jsonl_files, split="train")
 
-def load_mds_directories(mds_directories, split='.', batch_size=2**16, bulk=True):
+def load_mds_directories(mds_directories, split='.', batch_size=2**16, bulk=True, shuffle=None):
+    if bulk and shuffle is not None:
+        raise ValueError("Bulk reader does not support shuffling by design.")
     if bulk:
         return MDSBulkReader(mds_directories, split=split)
     dss = []
@@ -110,14 +113,19 @@ def load_mds_directories(mds_directories, split='.', batch_size=2**16, bulk=True
     else:
         with timing(message=f"Concatenating {len(dss)} datasets"):
             ds = concatenate_datasets(dsets=dss)
+    if shuffle is not None:
+        with timing(message="Creating shuffle indices"):
+            indices = shuffle_permutation(len(ds), seed=abs(shuffle))
+            if shuffle < 0:
+                indices = reverse_permutation(indices)
+        ds = IndexedDatasetView(ds, indices)
     return ds
 
 def save_jsonl(iterable, output_file, compression=None, processes=64, size_hint=None, overwrite=True, yes=True, trafo=None):
     f = None
     part = 0
-    trafo = Trafo(trafo)
-    for item in tqdm(iterable, desc="Writing to JSONL", unit="sample", disable=_NO_PROGESS):
-        item = trafo(item)
+    trafo = Transformations([] if trafo is None else [trafo])
+    for item in tqdm(trafo(iterable), desc="Writing to JSONL", unit="sample", disable=_NO_PROGESS):
         if f is None:
             part_file = output_file.format(part=part)
             check_arguments(part_file, overwrite, yes)
@@ -135,9 +143,8 @@ def save_mds(it, output_dir, processes=64, compression=None, buf_size=2**24, pig
     writer = None
     part = 0
     files = []
-    trafo = Trafo(trafo)
-    for sample in tqdm(it, desc="Writing to MDS", unit="sample", disable=_NO_PROGESS):
-        sample = trafo(sample)
+    trafo = Transformations([] if trafo is None else [trafo])
+    for sample in tqdm(trafo(it), desc="Writing to MDS", unit="sample", disable=_NO_PROGESS):
         if writer is None:
             part_dir = output_dir.format(part=part)
             check_arguments(part_dir, overwrite, yes)
@@ -179,10 +186,9 @@ def save_parquet(it, output_file, compression=None, batch_size=2**16, size_hint=
     compression = determine_compression("parquet", output_file, compression)
     writer = None
     part = 0
-    trafo = Trafo(trafo)
+    trafo = Transformations([] if trafo is None else [trafo])
     it = tqdm(it, desc="Writing to Parquet", unit="sample", disable=_NO_PROGESS)
-    for batch in _batch_iterable(it, batch_size):
-        batch = [trafo(sample) for sample in batch]
+    for batch in _batch_iterable(trafo(it), batch_size):
         table = pa.Table.from_pylist(batch)
         if writer is None:
             part_file = output_file.format(part=part)

{mldataforge-0.2.0 → mldataforge-0.2.2}/pyproject.toml

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