aspect-data 0.0.1__py3-none-any.whl

aspect/__init__.py

@@ -0,0 +1,5 @@
+from importlib.metadata import version
+
+app_name = "aspect-data"
+__author__ = "Eachan Johnson"
+__version__ = version(app_name)

aspect/checkpoint_utils.py

@@ -0,0 +1,95 @@
+""""Utilities for loading and saving checkpoints."""
+
+from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Union
+from tempfile import TemporaryDirectory
+import os
+import json
+
+from carabiner import print_err
+
+if TYPE_CHECKING:
+    from datasets import Dataset, IterableDataset
+else:
+    Dataset, IterableDataset = Any, Any
+
+
+def _load_json(checkpoint: str, filename: str) -> Dict[str, Any]:
+    with open(os.path.join(checkpoint, filename), "r") as f:
+            obj = json.load(f)
+    return obj
+   
+
+def save_json(obj, filename: str) -> None:
+    with open(filename, "w") as f:
+        try:
+            json.dump(obj, f, sort_keys=True, indent=4)
+        except TypeError as e:
+            print_err(f"{obj=}")
+            raise e
+    return None
+
+
+def _load_hf_dataset(checkpoint, filename) -> Union[Dataset, IterableDataset]:
+    from datasets import load_from_disk
+    return load_from_disk(os.path.join(checkpoint, filename))
+
+
+FILE_LOADING_CALLBACKS = {
+    "json": _load_json,
+    "hf-dataset": _load_hf_dataset,
+}
+
+
+def load_checkpoint_file(
+    checkpoint: str,
+    filename: str,
+    *args, 
+    callback: Union[str, Callable] = "json",
+    none_on_error: bool = False,
+    cache_dir: Optional[str] = None,
+    **kwargs
+) -> Union[Any, None]:
+    from huggingface_hub import snapshot_download
+
+    obj = None
+    if isinstance(callback, str):
+        try:
+            callback = FILE_LOADING_CALLBACKS[callback.casefold()]
+        except KeyError:
+            raise ValueError(
+                """
+                File loading callback must be callable or name.
+                """
+            )
+    if os.path.exists(checkpoint):
+        obj = callback(checkpoint, filename)
+    elif checkpoint.startswith("hf://"):
+        checkpoint = checkpoint.split("hf://")[-1]
+        if filename.endswith(".hf"):
+            filename_pattern = [filename + '/*.arrow', filename + '/*.json']
+        else:
+            filename_pattern = filename
+        with TemporaryDirectory() as tmpdirname:
+            try:
+                print_err(f"Looking up: {checkpoint} :: {filename}")
+                snapshot_download(
+                    repo_id=checkpoint,
+                    allow_patterns=filename_pattern,
+                    local_dir=tmpdirname,
+                    cache_dir=cache_dir,
+                    *args, **kwargs
+                )
+            except Exception as e:
+                print_err(e)
+                if none_on_error:
+                    return None
+                else:
+                    raise e
+            else:
+                obj = callback(tmpdirname, filename)
+    if obj is not None:
+        return obj
+    else:
+        raise AttributeError(
+            f"Could not load anything from {checkpoint=}, {filename=} with {callback=}."
+        )

aspect/cli_module/cli.py

@@ -0,0 +1,129 @@
+"""Command-line interface for aspect."""
+
+from argparse import FileType
+import os
+import sys
+
+from carabiner.cliutils import CLIOption, CLICommand, CLIApp
+
+from .. import app_name, __version__
+from .featurize import _featurize, _serialize
+
+def main() -> None:
+
+    input_file = CLIOption(
+        'input_file',
+        type=FileType('r'),
+        default=sys.stdin,
+        nargs='?',
+        help='Input file. Default: STDIN',
+    )
+    input_filename = CLIOption(
+        'input_file',
+        type=str,
+        help='Input file.',
+    )
+    feature_cols = CLIOption(
+        '--features', '-x',
+        type=str,
+        nargs='*',
+        default=None,
+        help='Featurization spec: column_name[:transform[(kwargs)]:...][@output_name]...',
+    )
+    cache = CLIOption(
+        '--cache',
+        type=str,
+        default=None,
+        help='Where to cache data.',
+    )
+    _config = CLIOption(
+        '--config',
+        type=str,
+        default=None,
+        help='Load pipeline from this config or checkpoint. Default: do not use, process from scratch.',
+    )
+    _checkpoint = CLIOption(
+        '--checkpoint',
+        type=str,
+        default=None,
+        help='Save data at this checkpoint. Default: do not save checkpoint.',
+    )
+
+    output_name = CLIOption(
+        '--output', '-o', 
+        type=str,
+        required=True,
+        help='Output filename.',
+    )
+
+    # slice dataset
+    slice_start = CLIOption(
+        '--start', 
+        type=int,
+        default=0,
+        help='First row of dataset to process.',
+    )
+    slice_end = CLIOption(
+        '--end', 
+        type=int,
+        default=None,
+        help='Last row of dataset to process. Default: end of dataset.',
+    )
+    extra_cols = CLIOption(
+        '--extras',
+        type=str,
+        nargs="*",
+        default=None,
+        help='Extra columns to retain without transformation.',
+    )
+    random_seed = CLIOption(
+        '--seed', '-e', 
+        type=int,
+        default=None,
+        help='Random seed. Default: determininstic.',
+    )
+
+    serialize = CLICommand(
+        "serialize",
+        description="Checkpoint a feature spec.",
+        options=[
+            output_name,
+            feature_cols,
+            extra_cols,
+        ],
+        main=_serialize,
+    )
+
+    featurize = CLICommand(
+        "featurize",
+        description="Featurize a table.",
+        options=[
+            input_filename, 
+            output_name,
+            feature_cols,
+            extra_cols,
+            slice_start,
+            slice_end,
+            _config,
+            _checkpoint,
+            random_seed,
+            cache,
+        ],
+        main=_featurize,
+    )
+
+    app = CLIApp(
+        app_name, 
+        description="Serializable featurization pipelines for ML/AI on chemistry, taxonomy, and general tabular data.", 
+        version=__version__, 
+        commands=[
+            serialize,
+            featurize,
+        ],
+    )
+    app.run()
+    return None
+
+
+if __name__ == '__main__':
+    main()

aspect/cli_module/featurize.py

@@ -0,0 +1,187 @@
+from typing import Optional, Iterable, Union
+
+from argparse import Namespace
+import os
+
+from carabiner import print_err
+from carabiner.cliutils import clicommand
+
+from .io import _resolve_and_slice_data, _save_dataset
+
+
+def _parse_feature_spec(x: str, i: int):
+    if ":" in x:
+        input_col, _, remainder = x.partition(":") 
+        transforms, _, output_col = remainder.rpartition("@")
+    else:
+        input_col, _, output_col = x.partition("@")
+        transforms, remainder = "", ""
+    # print(x, " | ".join([input_col, remainder, transforms, output_col]))
+    if not output_col:
+        if not transforms:
+            output_col = input_col
+            extra = True
+        else:
+            output_col = f"col_{i:03d}"
+    else:
+        extra = False
+    transforms = transforms.split(":")
+    if transforms == [""]:
+        return extra, {output_col: (input_col, {"name": "identity"})}
+    parsed_transforms = []
+    for transform in transforms:
+        name, _, kwargs = transform.partition("(")
+        if not kwargs:
+            parsed_transforms.append({"name": name})
+            continue
+        kwargs = kwargs.removesuffix(")").split(",")
+        key_vals = [k.strip().split("=") for k in kwargs]
+        wrong_number = [k for k in key_vals if len(k) != 2]
+        if wrong_number:
+            raise ValueError(f"Badly formatted kwargs: {kwargs=}, parsed to {key_vals}")
+        kwargs = {
+            key.strip(): value.strip() 
+            for key, value in key_vals
+        }
+        parsed_kwargs = {}
+        for k, v in kwargs.items():
+            if v.isdigit():
+                parsed_v = int(v)
+            elif v.casefold() == "true":
+                parsed_v = True
+            elif v.casefold() == "false":
+                parsed_v = False
+            else:
+                try: 
+                    parsed_v = float(v)
+                except:
+                    parsed_v = v
+            parsed_kwargs[k] = parsed_v
+        parsed_transforms.append({
+            "name": name,
+        } | parsed_kwargs)
+    return extra, {output_col: (input_col, tuple(parsed_transforms))}
+
+
+def parse_feature_specs(x: Union[str, Iterable[str]]):
+    if isinstance(x, str):
+        x = [x]
+    column_transforms = []
+    extra_cols = []
+    for i, _x in enumerate(x):
+        extra, column_transform = _parse_feature_spec(_x, i)
+        if extra:
+            extra_cols.append(_x)
+        else:
+            column_transforms.append(column_transform)
+    return extra_cols, {k: v for d in column_transforms for k, v in d.items()}
+
+
+def _common_feature_spec_routine(x: Union[str, Iterable[str]], args_extras=None):
+    extras, column_transforms = parse_feature_specs(x)
+    extras += (args_extras or [])
+    extras = sorted(set(extras))
+
+    print_err(
+        f"""
+        Parsed the feature spec:
+            - Column transforms: {column_transforms}
+            - Other columns to retain: {extras}
+
+        """
+    )
+    return extras, column_transforms
+
+
+def _validate_checkpoints(chk1, path):
+    from ..data import DataPipeline
+    print_err(f"[INFO] Validating checkpoint at {path}...", end="")
+    pipeline2 = DataPipeline().load_checkpoint(path)
+    if chk1 != pipeline2:
+        attr1, attr2 = vars(chk1), vars(pipeline2)
+        wrong_attributes = {
+            k: (attr1[k], "!=", v) 
+            for k, v in attr2.items() 
+            if v != attr1[k]
+        }
+        raise IOError(
+            f"Checkpoint at {path} does not recreate the same object. "
+            f"These attributes were different: {wrong_attributes}"
+        )
+    print_err(" all good!")
+    return None
+
+
+@clicommand("Serializing featurization with the following parameters")
+def _serialize(args: Namespace) -> None:
+
+    from ..data import DataPipeline
+
+    output = args.output
+    out_dir = os.path.dirname(output)
+    base = os.path.basename(output)
+    if len(out_dir) > 0:
+        os.makedirs(out_dir, exist_ok=True)
+    
+    extras, column_transforms = _common_feature_spec_routine(
+        args.features, 
+        args.extras,
+    )
+    pipeline = DataPipeline(
+        column_transforms=column_transforms,
+        columns_to_keep=extras,
+    )
+    pipeline.save_checkpoint(output)
+    _validate_checkpoints(pipeline, output)
+    return None
+
+
+@clicommand("Featurizing data with the following parameters")
+def _featurize(args: Namespace) -> None:
+
+    from ..data import DataPipeline
+    
+    if args.features:
+        extras, column_transforms = _common_feature_spec_routine(args.features)
+        pipeline = DataPipeline(
+            column_transforms=column_transforms,
+            columns_to_keep=extras,
+            cache_dir=args.cache,
+        )
+    elif args.config:
+        extras = args.extras or []
+        pipeline = DataPipeline(
+            cache_dir=args.cache,
+        ).load_checkpoint(args.config)
+    else:
+        raise ValueError(f"One of --features or --config must be provided.")
+    
+    ds = _resolve_and_slice_data(
+        args.input_file,
+        start=args.start,
+        end=args.end,
+        cache_dir=args.cache,
+    )
+    ds = pipeline(
+        ds, 
+        keep_extra_columns=extras,
+        drop_unused_columns=True,
+    )
+    if args.checkpoint:
+        output = args.checkpoint
+        out_dir = os.path.dirname(output)
+        base = os.path.basename(output)
+        pipeline.save_checkpoint(output)
+        _validate_checkpoints(pipeline, output)
+
+    if args.output:
+        output = args.output
+        out_dir = os.path.dirname(output)
+        base = os.path.basename(output)
+        if len(out_dir) > 0:
+            os.makedirs(out_dir, exist_ok=True)
+        _save_dataset(
+            ds, 
+            output,
+        )
+    return None

aspect/cli_module/io.py

@@ -0,0 +1,52 @@
+from typing import Mapping, Optional, Union
+
+from carabiner import print_err
+
+
+def _resolve_and_slice_data(
+    data: str,
+    start: Optional[int] = None,
+    end: Optional[int] = None,
+    batch_size: int = 1024,
+    cache_dir: str = None
+):
+    from ..io import AutoDataset
+
+    candidates_ds = AutoDataset.load(data, cache_dir)._dataset
+    nrows = candidates_ds.num_rows
+    skip = start or 0
+    take = (end or nrows) - skip
+    if (take - skip) < nrows:
+        print_err(f"[INFO] Reading dataset from row {skip} to row {take + skip} / {nrows}.")
+    return candidates_ds.skip(skip).take(take)
+
+
+def _save_dataset(
+    dataset,
+    output: str
+) -> None:
+    print_err("[INFO] Saving dataset:\n" + str(dataset) + "\n" + f"at {output} as", end=" ")
+    if output.removesuffix(".gz").endswith((".csv", ".tsv", ".txt")):
+        print_err("CSV.")
+        dataset.to_csv(
+            output, 
+            sep="," if output.removesuffix(".gz").endswith(".csv") else "\t",
+            compression='gzip' if output.endswith(".gz") else None,
+        )
+    elif output.endswith(".json"):
+        print_err("JSON.")
+        dataset.to_json(output)
+    elif output.endswith(".parquet"):
+        print_err("Parquet.")
+        dataset.to_parquet(output)
+    elif output.endswith(".sql"):
+        print_err("SQL.")
+        dataset.to_sql(output)
+    elif output.endswith(".hf"):
+        print_err("Hugging Face dataset.")
+        dataset.save_to_disk(output)
+    else:
+        print_err("Hugging Face dataset.")
+        dataset.save_to_disk(output + ".hf")
+        print_err(f"WARNING: Unsure what format to save as for filename {output}. Defaulted to Hugging Face dataset.")
+    return None