mldataforge 0.2.0__tar.gz → 0.2.1__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mldataforge
-Version: 0.2.0
+Version: 0.2.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.2.0 → mldataforge-0.2.1}/mldataforge/trafos.py

@@ -1,34 +1,95 @@
 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__ = ['IndexedDatasetView', 'Transformation', 'Transformations', 'flatten_json', 'unflatten_json']
+
+class IndexedDatasetView:
+    def __init__(self, dataset, indices):
+        self.dataset = dataset
+        self.indices = list(indices)  # ensure repeatable access
+
+    def __iter__(self):
+        for idx in self.indices:
+            yield self.dataset[idx]
+
+    def __len__(self):
+        return len(self.indices)
+
+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]
+        self.indices = indices  # Optional index iterable
+
+    def __call__(self, dataset):
+        # Wrap dataset with IndexedDatasetView if indices are provided
+        if self.indices is not None:
+            dataset = IndexedDatasetView(dataset, self.indices)
+
+        result = dataset
+        for transform in self.pipeline:
+            result = transform(result)
+        return result
+
+    def __len__(self):
+        # Return the input length to the pipeline
+        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 identity(obj):
+    return obj
 
 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 +110,12 @@ def flatten_json(obj, parent_key='', sep='.', escape_char='\\'):
         items.append((parent_key, obj))
     return dict(items)
 
-
 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 +128,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 +162,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.1}/mldataforge/utils.py

@@ -12,7 +12,7 @@ from tqdm import tqdm
 from .compression import determine_compression, open_compression, pigz_compress
 from .mds import MDSBulkReader, MDSWriter
 from .pigz import pigz_open
-from .trafos import Trafo
+from .trafos import Transformations
 
 __all__ = [
     "check_arguments",
@@ -115,9 +115,8 @@ def load_mds_directories(mds_directories, split='.', batch_size=2**16, bulk=True
 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 +134,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 +177,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.1}/pyproject.toml

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