mldataforge 0.1.7__tar.gz → 0.2.0__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mldataforge
-Version: 0.1.7
+Version: 0.2.0
 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.7 → mldataforge-0.2.0}/mldataforge/commands/convert/jsonl.py

@@ -21,9 +21,10 @@ def jsonl():
 @buf_size_option()
 @shard_size_option()
 @no_pigz_option()
+@trafo_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):
+def jsonl_to_mds(output_dir, jsonl_files, compression, processes, overwrite, yes, buf_size, shard_size, no_pigz, trafo):
     check_arguments(output_dir, overwrite, yes, jsonl_files)
     save_mds(
         load_jsonl_files(jsonl_files),
@@ -33,6 +34,7 @@ def jsonl_to_mds(output_dir, jsonl_files, compression, processes, overwrite, yes
         buf_size=buf_size,
         pigz=use_pigz(compression, no_pigz),
         shard_size=shard_size,
+        trafo=trafo,
     )
 
 @jsonl.command()
@@ -42,13 +44,15 @@ def jsonl_to_mds(output_dir, jsonl_files, compression, processes, overwrite, yes
 @overwrite_option()
 @yes_option()
 @batch_size_option()
+@trafo_option()
 def parquet(**kwargs):
     jsonl_to_parquet(**kwargs)
-def jsonl_to_parquet(output_file, jsonl_files, compression, overwrite, yes, batch_size):
+def jsonl_to_parquet(output_file, jsonl_files, compression, overwrite, yes, batch_size, trafo):
     check_arguments(output_file, overwrite, yes, jsonl_files)
     save_parquet(
         load_jsonl_files(jsonl_files),
         output_file,
         compression=compression,
         batch_size=batch_size,
+        trafo=trafo,
     )

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

@@ -19,15 +19,17 @@ def mds():
 @yes_option()
 @batch_size_option()
 @no_bulk_option()
+@trafo_option()
 def jsonl(**kwargs):
     mds_to_jsonl(**kwargs)
-def mds_to_jsonl(output_file, mds_directories, compression, processes, overwrite, yes, batch_size, no_bulk):
+def mds_to_jsonl(output_file, mds_directories, compression, processes, overwrite, yes, batch_size, no_bulk, trafo):
     check_arguments(output_file, overwrite, yes, mds_directories)
     save_jsonl(
         load_mds_directories(mds_directories, batch_size=batch_size, bulk=not no_bulk),
         output_file,
         compression=compression,
         processes=processes,
+        trafo=trafo,
     )
 
 @mds.command()
@@ -38,13 +40,15 @@ def mds_to_jsonl(output_file, mds_directories, compression, processes, overwrite
 @yes_option()
 @batch_size_option()
 @no_bulk_option()
+@trafo_option()
 def parquet(**kwargs):
     mds_to_parquet(**kwargs)
-def mds_to_parquet(output_file, mds_directories, compression, overwrite, yes, batch_size, no_bulk):
+def mds_to_parquet(output_file, mds_directories, compression, overwrite, yes, batch_size, no_bulk, trafo):
     check_arguments(output_file, overwrite, yes, mds_directories)
     save_parquet(
         load_mds_directories(mds_directories, batch_size=batch_size, bulk=not no_bulk),
         output_file,
         compression=compression,
         batch_size=batch_size,
+        trafo=trafo,
     )

{mldataforge-0.1.7 → mldataforge-0.2.0}/mldataforge/commands/convert/parquet.py

@@ -18,15 +18,17 @@ def parquet():
 @processes_option()
 @overwrite_option()
 @yes_option()
+@trafo_option()
 def jsonl(**kwargs):
     parquet_to_jsonl(**kwargs)
-def parquet_to_jsonl(output_file, parquet_files, compression, processes, overwrite, yes):
+def parquet_to_jsonl(output_file, parquet_files, compression, processes, overwrite, yes, trafo):
     check_arguments(output_file, overwrite, yes, parquet_files)
     save_jsonl(
         load_dataset("parquet", data_files=parquet_files, split="train"),
         output_file,
         compression=compression,
         processes=processes,
+        trafo=trafo,
     )
 
 @parquet.command()
@@ -39,9 +41,10 @@ def parquet_to_jsonl(output_file, parquet_files, compression, processes, overwri
 @buf_size_option()
 @shard_size_option()
 @no_pigz_option()
+@trafo_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):
+def parquet_to_mds(output_dir, parquet_files, compression, processes, overwrite, yes, buf_size, shard_size, no_pigz, trafo):
     check_arguments(output_dir, overwrite, yes, parquet_files)
     save_mds(
         load_dataset("parquet", data_files=parquet_files, split="train"),
@@ -51,4 +54,5 @@ def parquet_to_mds(output_dir, parquet_files, compression, processes, overwrite,
         buf_size=buf_size,
         pigz=use_pigz(compression, no_pigz=no_pigz),
         shard_size=shard_size,
+        trafo=trafo,
     )

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

@@ -18,9 +18,10 @@ def join():
 @processes_option()
 @overwrite_option()
 @yes_option()
+@trafo_option()
 def jsonl(**kwargs):
     join_jsonl(**kwargs)
-def join_jsonl(output_file, jsonl_files, compression, processes, overwrite, yes):
+def join_jsonl(output_file, jsonl_files, compression, processes, overwrite, yes, trafo):
     check_arguments(output_file, overwrite, yes, jsonl_files)
     save_jsonl(
         load_jsonl_files(jsonl_files),
@@ -41,10 +42,11 @@ def join_jsonl(output_file, jsonl_files, compression, processes, overwrite, yes)
 @no_bulk_option()
 @shard_size_option()
 @no_pigz_option()
+@trafo_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):
+def join_mds(output_dir, mds_directories, compression, processes, overwrite, yes, batch_size, buf_size, no_bulk, shard_size, no_pigz, trafo):
     check_arguments(output_dir, overwrite, yes, mds_directories)
     save_mds(
         load_mds_directories(mds_directories, batch_size=batch_size, bulk=not no_bulk),
@@ -54,6 +56,7 @@ def join_mds(output_dir, mds_directories, compression, processes, overwrite, yes
         buf_size=buf_size,
         shard_size=shard_size,
         pigz=use_pigz(compression, no_pigz),
+        trafo=trafo,
     )
 
 @join.command()
@@ -63,13 +66,15 @@ def join_mds(output_dir, mds_directories, compression, processes, overwrite, yes
 @overwrite_option()
 @yes_option()
 @batch_size_option()
+@trafo_option()
 def parquet(**kwargs):
     join_parquet(**kwargs)
-def join_parquet(output_file, parquet_files, compression, overwrite, yes, batch_size):
+def join_parquet(output_file, parquet_files, compression, overwrite, yes, batch_size, trafo):
     check_arguments(output_file, overwrite, yes, parquet_files)
     save_parquet(
         load_dataset("parquet", data_files=parquet_files, split="train"),
         output_file,
         compression=compression,
         batch_size=batch_size,
+        trafo=trafo,
     )

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

@@ -20,7 +20,10 @@ def split():
 @processes_option()
 @overwrite_option()
 @yes_option()
-def jsonl(jsonl_files, prefix, output_dir, size_hint, compression, processes, overwrite, yes):
+@trafo_option()
+def jsonl(*args, **kwargs):
+    split_jsonl(*args, **kwargs)
+def split_jsonl(jsonl_files, prefix, output_dir, size_hint, compression, processes, overwrite, yes, trafo):
     save_jsonl(
         load_jsonl_files(jsonl_files),
         output_file=f"{output_dir}/{prefix}{{part:04d}}.jsonl{extension_compression(compression, jsonl_files[0])}",
@@ -29,6 +32,7 @@ def jsonl(jsonl_files, prefix, output_dir, size_hint, compression, processes, ov
         size_hint=size_hint,
         overwrite=overwrite,
         yes=yes,
+        trafo=trafo,
     )
 
 @split.command()
@@ -45,7 +49,10 @@ def jsonl(jsonl_files, prefix, output_dir, size_hint, compression, processes, ov
 @no_bulk_option()
 @shard_size_option()
 @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):
+@trafo_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):
     save_mds(
         load_mds_directories(mds_directories, batch_size=batch_size, bulk=not no_bulk),
         output_dir=f"{output_dir}/{prefix}{{part:04d}}",
@@ -57,6 +64,7 @@ def mds(mds_directories, prefix, output_dir, size_hint, compression, processes,
         size_hint=size_hint,
         overwrite=overwrite,
         yes=yes,
+        trafo=trafo,
     )
 
 @split.command()
@@ -68,7 +76,10 @@ def mds(mds_directories, prefix, output_dir, size_hint, compression, processes,
 @overwrite_option()
 @yes_option()
 @batch_size_option()
-def parquet(parquet_files, prefix, output_dir, size_hint, compression, overwrite, yes, batch_size):
+@trafo_option()
+def parquet(*args, **kwargs):
+    split_parquet(*args, **kwargs)
+def split_parquet(parquet_files, prefix, output_dir, size_hint, compression, overwrite, yes, batch_size, trafo):
     save_parquet(
         load_dataset("parquet", data_files=parquet_files, split="train"),
         output_file=f"{output_dir}/{prefix}{{part:04d}}.parquet",
@@ -77,4 +88,5 @@ def parquet(parquet_files, prefix, output_dir, size_hint, compression, overwrite
         size_hint=size_hint,
         overwrite=overwrite,
         yes=yes,
+        trafo=trafo,
     )

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

@@ -14,6 +14,7 @@ __all__ = [
     "prefix_option",
     "shard_size_option",
     "size_hint_option",
+    "trafo_option",
     "yes_option",
 ]
 
@@ -129,6 +130,17 @@ def size_hint_option(default=2**26):
         help=f"Size hint for the dataset (default: {default}).",
     )
 
+def trafo_option():
+    """
+    Option for specifying the transformation function.
+    """
+    return click.option(
+        "--trafo",
+        default=None,
+        type=str,
+        help="Transformation function to apply to the dataset.",
+    )
+
 def yes_option():
     """
     Option for specifying whether to assume yes to all prompts.

mldataforge-0.2.0/mldataforge/trafos.py

@@ -0,0 +1,111 @@
+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})"
+
+
+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 + ']')
+
+    if isinstance(obj, dict):
+        if not obj:
+            # explicitly handle empty dict
+            items.append((parent_key, {}))
+        else:
+            for k, v in obj.items():
+                new_key = f"{parent_key}{sep}{escape(k)}" if parent_key else escape(k)
+                items.extend(flatten_json(v, new_key, sep, escape_char).items())
+    elif isinstance(obj, list):
+        if not obj:
+            # explicitly handle empty list
+            items.append((parent_key, []))
+        else:
+            for idx, v in enumerate(obj):
+                new_key = f"{parent_key}[{idx}]"
+                items.extend(flatten_json(v, new_key, sep, escape_char).items())
+    else:
+        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 = []
+        for token in tokens:
+            if token.startswith('['):
+                parsed.append(int(token[1:-1]))
+            else:
+                parsed.append(token.replace(escape_char + sep, sep)
+                                  .replace(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
+        for idx, key in enumerate(keys):
+            if idx == len(keys) - 1:
+                if isinstance(key, int):
+                    if not isinstance(current, list):
+                        current_parent[last_key] = []
+                        current = current_parent[last_key]
+                    while len(current) <= key:
+                        current.append(None)
+                    current[key] = value
+                else:
+                    current[key] = value
+            else:
+                next_key = keys[idx + 1]
+                if isinstance(key, int):
+                    if not isinstance(current, list):
+                        current_parent[last_key] = []
+                        current = current_parent[last_key]
+                    while len(current) <= key:
+                        current.append(None)
+                    if current[key] is None:
+                        current[key] = [] if isinstance(next_key, int) else {}
+                    current_parent = current
+                    current = current[key]
+                else:
+                    if key not in current:
+                        current[key] = [] if isinstance(next_key, int) else {}
+                    current_parent = current
+                    current = current[key]
+            last_key = key
+
+    return result

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

@@ -12,6 +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
 
 __all__ = [
     "check_arguments",
@@ -111,10 +112,12 @@ def load_mds_directories(mds_directories, split='.', batch_size=2**16, bulk=True
             ds = concatenate_datasets(dsets=dss)
     return ds
 
-def save_jsonl(iterable, output_file, compression=None, processes=64, size_hint=None, overwrite=True, yes=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)
         if f is None:
             part_file = output_file.format(part=part)
             check_arguments(part_file, overwrite, yes)
@@ -127,12 +130,14 @@ def save_jsonl(iterable, output_file, compression=None, processes=64, size_hint=
     if f is not None:
         f.close()
 
-def save_mds(it, output_dir, processes=64, compression=None, buf_size=2**24, pigz=True, shard_size=None, size_hint=None, overwrite=True, yes=True):
+def save_mds(it, output_dir, processes=64, compression=None, buf_size=2**24, pigz=True, shard_size=None, size_hint=None, overwrite=True, yes=True, trafo=None):
     compression = determine_compression("mds", output_dir, compression, no_pigz=not pigz)
     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)
         if writer is None:
             part_dir = output_dir.format(part=part)
             check_arguments(part_dir, overwrite, yes)
@@ -170,12 +175,14 @@ def save_mds(it, output_dir, processes=64, compression=None, buf_size=2**24, pig
             json.dump(index, open(index_path, "wt"))
             print(f"Compressed {output_dir} with pigz")
 
-def save_parquet(it, output_file, compression=None, batch_size=2**16, size_hint=None, overwrite=True, yes=True):
+def save_parquet(it, output_file, compression=None, batch_size=2**16, size_hint=None, overwrite=True, yes=True, trafo=None):
     compression = determine_compression("parquet", output_file, compression)
     writer = None
     part = 0
+    trafo = Trafo(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]
         table = pa.Table.from_pylist(batch)
         if writer is None:
             part_file = output_file.format(part=part)

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

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