mldataforge 0.1.6__tar.gz → 0.2.0__tar.gz

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

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

@@ -1,16 +1,26 @@
+from copy import deepcopy
 import json
 import numpy as np
 import os
 import shutil
 from streaming.base.compression import compress, decompress, get_compression_extension, is_compression
+from streaming.base.format import _readers
+from streaming.base.format.base.reader import FileInfo, JointReader
 from streaming.base.format.index import get_index_basename
-from streaming.base.format.mds.encodings import mds_decode, mds_encode, is_mds_encoding, get_mds_encodings, get_mds_encoded_size
+from streaming.base.format.mds.encodings import mds_decode, mds_encode, is_mds_encoding, is_mds_encoding_safe, get_mds_encodings, get_mds_encoded_size
 from streaming.base.hashing import get_hash, is_hash
 from streaming.base.util import bytes_to_int
 from typing import Any, Optional, Generator, Self, Union
 
 from .utils import open_compression
 
+__all__ = [
+    "MDSBulkReader",
+    "MDSBulkShardReader",
+    "MDSReader",
+    "MDSWriter",
+]
+
 class MDSBulkReader:
     def __init__(
         self,
@@ -37,11 +47,11 @@ class MDSBulkReader:
 
     def __iter__(self) -> Generator[dict[str, Any], None, None]:
         for shard in self.shards:
-            with MDSShardReader(**shard) as reader:
+            with MDSBulkShardReader(**shard) as reader:
                 for sample in reader:
                     yield sample
 
-class MDSShardReader:
+class MDSBulkShardReader:
     def __init__(
         self,
         filename: str,
@@ -94,7 +104,7 @@ class MDSShardReader:
         for i in range(self.samples):
             yield self.get_item(i)
 
-    def __enter__(self) -> "MDSShardReader":
+    def __enter__(self) -> "MDSBulkShardReader":
         return self
 
     def __exit__(self, exc_type, exc_value, traceback) -> None:
@@ -315,3 +325,124 @@ class MDSWriter:
 
     def __exit__(self, exc_type, exc, traceback):
         self.finish()
+
+class MDSReader(JointReader):
+
+    def __init__(
+        self,
+        dirname: str,
+        split: Optional[str],
+        column_encodings: list[str],
+        column_names: list[str],
+        column_sizes: list[Optional[int]],
+        compression: Optional[str],
+        hashes: list[str],
+        raw_data: FileInfo,
+        samples: int,
+        size_limit: Optional[Union[int, str]],
+        zip_data: Optional[FileInfo],
+    ) -> None:
+        self.sample_compression = None
+        if compression and compression.startswith("sample::"):
+            compression, self.sample_compression = None, compression.removeprefix("sample::")
+        super().__init__(dirname, split, compression, hashes, raw_data, samples, size_limit,
+                         zip_data)
+        self.column_encodings = column_encodings
+        self.column_names = column_names
+        self.column_sizes = column_sizes
+
+    @classmethod
+    def from_json(cls, dirname: str, split: Optional[str], obj: dict[str, Any]) -> Self:
+        """Initialize from JSON object.
+
+        Args:
+            dirname (str): Local directory containing shards.
+            split (str, optional): Which dataset split to use, if any.
+            obj (Dict[str, Any]): JSON object to load.
+
+        Returns:
+            Self: Loaded MDSReader.
+        """
+        args = deepcopy(obj)
+        args_version = args['version']
+        if args_version != 2:
+            raise ValueError(
+                f'Unsupported streaming data version: {args_version}. Expected version 2.')
+        del args['version']
+        args_format = args['format']
+        if args_format != 'mds':
+            raise ValueError(f'Unsupported data format: {args_format}. Expected to be `mds`.')
+        del args['format']
+        args['dirname'] = dirname
+        args['split'] = split
+        for key in ['raw_data', 'zip_data']:
+            arg = args[key]
+            args[key] = FileInfo(**arg) if arg else None
+        return cls(**args)
+
+    def validate(self, allow_unsafe_types: bool) -> None:
+        """Check whether this shard is acceptable to be part of some Stream.
+
+        Args:
+            allow_unsafe_types (bool): If a shard contains Pickle, which allows arbitrary code
+                execution during deserialization, whether to keep going if ``True`` or raise an
+                error if ``False``.
+        """
+        if not allow_unsafe_types:
+            for column_id, encoding in enumerate(self.column_encodings):
+                if not is_mds_encoding_safe(encoding):
+                    name = self.column_names[column_id]
+                    raise ValueError(f'Column {name} contains an unsafe type: {encoding}. To ' +
+                                     f'proceed anyway, set ``allow_unsafe_types=True``.')
+
+    def decode_sample(self, data: bytes) -> dict[str, Any]:
+        """Decode a sample dict from bytes.
+
+        Args:
+            data (bytes): The sample encoded as bytes.
+
+        Returns:
+            Dict[str, Any]: Sample dict.
+        """
+        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:
+        """Get the raw sample data at the index.
+
+        Args:
+            idx (int): Sample index.
+
+        Returns:
+            bytes: Sample data.
+        """
+        filename = os.path.join(self.dirname, self.split, self.raw_data.basename)
+        offset = (1 + idx) * 4
+        with open(filename, 'rb', 0) as fp:
+            fp.seek(offset)
+            pair = fp.read(8)
+            begin, end = np.frombuffer(pair, np.uint32)
+            fp.seek(begin)
+            data = fp.read(end - begin)
+        if not data:
+            raise IndexError(
+                f'Relative sample index {idx} is not present in the {self.raw_data.basename} file.'
+            )
+        if self.sample_compression:
+            data = decompress(self.sample_compression, data)
+        return data
+
+_readers["mds"] = MDSReader

{mldataforge-0.1.6 → 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.6 → 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",
@@ -23,6 +24,11 @@ __all__ = [
     "save_parquet",
 ]
 
+_NO_PROGESS = False
+def set_progress(value):
+    global _NO_PROGESS
+    _NO_PROGESS = value
+
 def _batch_iterable(iterable, batch_size):
     batch = []
     for item in iterable:
@@ -73,7 +79,7 @@ def _infer_mds_encoding(value):
     return 'pkl'
 
 def _streaming_jsonl(jsonl_files, compressions):
-    for jsonl_file, compression in tqdm(zip(jsonl_files, compressions), desc="Loading JSONL files", unit="file"):
+    for jsonl_file, compression in tqdm(zip(jsonl_files, compressions), desc="Loading JSONL files", unit="file", disable=_NO_PROGESS):
         for line in open_compression(jsonl_file, mode="rt", compression=compression):
             yield json.loads(line)
 
@@ -106,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
-    for item in tqdm(iterable, desc="Writing to JSONL", unit="sample"):
+    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)
@@ -122,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 = []
-    for sample in tqdm(it, desc="Writing to MDS", unit="sample"):
+    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)
@@ -151,7 +161,7 @@ def save_mds(it, output_dir, processes=64, compression=None, buf_size=2**24, pig
             name2info = {shard["raw_data"]["basename"]: shard for shard in index["shards"]}
             file_names = [file for file in os.listdir(output_dir) if file.endswith(".mds")]
             assert set(file_names) == set(name2info.keys())
-            for file_name in tqdm(file_names, desc="Compressing with pigz", unit="file"):
+            for file_name in tqdm(file_names, desc="Compressing with pigz", unit="file", disable=_NO_PROGESS):
                 compressed_file_name = file_name + ".gz"
                 file_path = os.path.join(output_dir, file_name)
                 compressed_file_path = os.path.join(output_dir, compressed_file_name)
@@ -165,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
-    it = tqdm(it, desc="Writing to Parquet", unit="sample")
+    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.6 → mldataforge-0.2.0}/pyproject.toml

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