providence 0.0.1__py3-none-any.whl

providence/__init__.py

@@ -0,0 +1,8 @@
+from importlib.metadata import version
+
+app_name = "providence"
+__author__ = "Eachan Johnson"
+__version__ = version(app_name)
+
+from .data import DataResolver
+from .io import DataSource

providence/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=}."
+        )

providence/cli_module/cli.py

@@ -0,0 +1,92 @@
+"""Command-line interface for providence."""
+
+from argparse import FileType
+import os
+import sys
+
+from carabiner.cliutils import CLIOption, CLICommand, CLIApp
+
+from .. import app_name, __version__
+from .fetch import _fetch
+
+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.',
+    )
+    force_write = CLIOption(
+        '--force-write-data', '-f',
+        action='store_true',
+        help='Save data to disk from remote repos, instead of simply saving a reference.',
+    )
+    cache = CLIOption(
+        '--cache',
+        type=str,
+        default=None,
+        help='Where to cache data.',
+    )
+    output_name = CLIOption(
+        '--output', '-o', 
+        type=str,
+        default=False,
+        help='Output filename for data table.',
+    )
+    checkpoint = CLIOption(
+        '--checkpoint', '-c', 
+        type=str,
+        default=False,
+        help='Output filename for Providence checkpoint.',
+    )
+
+    # slice dataset
+    slice_start = CLIOption(
+        '--start', 
+        type=int,
+        default=None,
+        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.',
+    )
+
+    get_command = CLICommand(
+        "get",
+        description="Load a local or remote dataset, and save as a reloadable and verifiable providence checkpoint.",
+        options=[
+            input_filename,
+            slice_start,
+            slice_end,
+            force_write,
+            checkpoint,
+            output_name,
+            cache,
+        ],
+        main=_fetch,
+    )
+    
+    app = CLIApp(
+        app_name, 
+        description="Data ingestion, validation, and provenance.", 
+        version=__version__, 
+        commands=[
+            get_command,
+        ],
+    )
+    app.run()
+    return None
+
+
+if __name__ == '__main__':
+    main()

providence/cli_module/fetch.py

@@ -0,0 +1,43 @@
+from argparse import Namespace
+import os
+
+from carabiner.cliutils import clicommand
+
+
+@clicommand("Fetching data with the following parameters")
+def _fetch(args: Namespace) -> None:
+
+    from .io import _save_dataset
+    from ..data import DataResolver
+    
+    resolved = DataResolver(
+        args.input_file,
+        cache_dir=args.cache,
+    )
+    if args.start is not None or args.end is not None:
+        resolved = resolved.select_rows(
+            args.start,
+            args.end,
+        )
+    
+    if not args.checkpoint and not args.output:
+        raise ValueError("At least one of --checkpoint and --output should be set.")
+    
+    if args.checkpoint:
+        output = args.checkpoint
+        out_dir = os.path.dirname(output)
+        base = os.path.basename(output)
+        
+        resolved.save_checkpoint(output, save_source_data=args.force_write_data)
+
+    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(
+            resolved.data, 
+            output,
+        )
+    return None

providence/cli_module/io.py

@@ -0,0 +1,51 @@
+
+from carabiner import print_err
+
+
+def _resolve_and_slice_data(
+    data: str,
+    start: int | None = None,
+    end: int | None = None,
+    batch_size: int = 1024,
+    cache_dir: str = None
+):
+    from ..io import AutoDataset
+
+    ds = AutoDataset.load(data, cache_dir)
+    nrows = ds._dataset.num_rows
+    skip = 0 if start is None else start
+    take = (nrows if end is None else end) - skip
+    if (take - skip) < nrows:
+        print_err(f"[INFO] Reading dataset from row {skip} to row {take + skip} / {nrows}.")
+    return ds._dataset.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

providence/data.py

@@ -0,0 +1,361 @@
+"""Data resolver class."""
+
+from typing import TYPE_CHECKING, Any
+from collections.abc import Iterable, Mapping
+import os
+from pathlib import Path
+
+from carabiner import print_err
+
+if TYPE_CHECKING:
+    from datasets import Dataset, DatasetDict, IterableDataset
+    from .collate import ColumnCollator
+else:
+    Dataset, DatasetDict, IterableDataset = Any, Any, Any
+    ColumnCollator = Any
+
+from . import app_name, __version__
+from .io import AutoDataset, DataSource, autoload, file_checksum, save_json, load_json
+from .package_data import CACHE_DIR
+from .typing import DataLike
+
+
+DEFAULT_FORMAT: str = "numpy"
+CONFIG_FILENAME: str = "config.json"
+DATA_FILENAME: str = "data.parquet"
+EXAMPLE_FILENAME: str = "example.parquet"
+SCHEMA_VERSION: str = "1"
+
+
+class DataResolver:
+    """Resolve datasets and track their provenance.
+
+    Examples
+    --------
+    Select rows without modifying the original dataset:
+
+    >>> data = DataResolver({"x": list(range(10))})
+    >>> subset = data.select_rows(2, 5)
+    >>> list(subset.data["x"])
+    [2, 3, 4]
+    >>> len(data.data)
+    10
+
+    Chain selections and restore a reference-only checkpoint:
+
+    >>> import tempfile
+    >>> from pathlib import Path
+    >>> import pandas as pd
+    >>> with tempfile.TemporaryDirectory() as tmp:
+    ...     source = Path(tmp) / "data.parquet"
+    ...     pd.DataFrame({"x": list(range(10))}).to_parquet(source)
+    ...     selected = DataResolver(source).select_rows(2, 8).select_rows(1, 3)
+    ...     selected.save(Path(tmp) / "checkpoint", save_source_data=False)
+    ...     restored = DataResolver.load(Path(tmp) / "checkpoint")
+    ...     print(list(restored.data["x"]))
+    [3, 4]
+
+    """
+    def __init__(
+        self,
+        data: str | Mapping[str, Any],
+        cache_dir: str | None = None,
+        _file_checksum: str | None = None,
+        _version: str = __version__,
+        _app: str = app_name
+    ):
+        self._version = _version
+        self._app = _app
+        self._data_requested = data
+        self.data = None
+        self._file_checksum = _file_checksum
+        self.selected_rows = None
+        self.data_source = None
+        self.data_example = None
+        self.cache_dir = cache_dir or CACHE_DIR
+        self._data_loaded = False
+        self._init()
+        
+
+    def _init(self) -> None:
+        self.data = self._resolve_data(
+            self._data_requested,
+        )
+        if self.data is None:
+            raise AttributeError(
+                f"Data could not be loaded from {self._data_requested}."
+            )
+        self._data_loaded = self.data is not None
+
+        self.data_example, self.data_shape = self._inspect_data(
+            self.data,
+            self.data_example,
+        )
+        return None
+
+    def _resolve_data(
+        self,
+        data: DataLike,
+        cache_dir: str | None = None
+    ) -> Dataset | IterableDataset:
+        resolved = AutoDataset.load(
+            data, 
+            cache_dir=cache_dir or self.cache_dir,
+        )
+        self.data_source = resolved.source
+        return resolved._dataset
+
+    def _apply_selection(
+        self,
+        start: int | None = None,
+        end: int | None = None
+    ) -> "DataResolver":
+        if self.selected_rows is None:
+            selected_rows = []
+        else:
+            selected_rows = self.selected_rows.copy()
+        
+        nrows = self.data.num_rows
+        indices = tuple(slice(start, end).indices(nrows))
+        skip = indices[0]
+        take = indices[1] - skip
+
+        selected = type(self)(
+            self.data.select(range(*indices)),
+            cache_dir=self.cache_dir,
+        )
+
+        selected.data_source = self.data_source
+        selected_rows.append((start, end))
+        selected.selected_rows = selected_rows
+        return selected
+
+    @staticmethod
+    def _normalize_selected_rows(
+        start: Iterable[Iterable[int]] | int | None = None,
+        end: int | None = None
+    ):
+        if start is None and end is None:
+            return None
+        if isinstance(start, int) or (start is None and isinstance(end, int)):
+            return [(start, end)]
+        elif isinstance(start, Iterable) and not isinstance(start, str):
+            start = list(start)
+            if len(start) == 0:
+                raise ValueError(f"start was iterable but length-0: {start}")
+            elif len(start) == 2 and all(isinstance(item, int) for item in start):
+                start = [start]
+            for i, item in enumerate(start):
+                if isinstance(item, Iterable) and not isinstance(item, str):
+                    if len(item) != 2:
+                        raise ValueError(f"start was iterable but item {i} was not a 2-tuple: {item}")
+                else:
+                    raise ValueError(f"start was iterable but item {i} was not: {item}\n{start=}")
+            return [tuple(item) for item in start]
+        else: 
+            raise ValueError(f"Cannot normalize inputs: {start=}, {end=}")
+        
+    def select_rows(
+        self,
+        start: Iterable[Iterable[int]] | int | None = None,
+        end: int | None = None
+    ) -> "DataResolver":
+        """Return a new resolver containing rows [start:end]."""
+
+        if start is None and end is None:
+            return self  # pass through
+        selections = self._normalize_selected_rows(start, end)
+        data = self
+        for item in selections:
+            data = data._apply_selection(*item)
+        return data
+
+    @staticmethod
+    def _inspect_data(data: Dataset, example=None, format: str = DEFAULT_FORMAT) -> None:
+        if example is None:
+            if data is not None:
+                data = data.with_format(format)
+                example = data.take(1)
+            else:
+                raise AttributeError("Cannot inspect data without loading data or having an example")
+        
+        first_item = example.with_format(format)[:1]
+        shape = {
+            col: first_item[col].shape[1:]
+            if not isinstance(first_item[col], dict)
+            else {
+                k: v.shape[1:] if v is not None else None 
+                for k, v in first_item[col].items()
+            }
+            for col in example.column_names
+        }
+        return example, shape
+
+    def save(
+        self,
+        path: str | os.PathLike,
+        *,
+        save_source_data: bool | None = None,
+        discard_example_data: bool = False
+    ) -> None:
+        """Save the pipeline and required training-data artefacts.
+
+        Parameters
+        ==========
+        path
+            Checkpoint directory.
+        save_source_data
+            Whether to embed the input dataset.
+
+            ``None`` chooses automatically: immutable remote sources are
+            referenced, while local or ephemeral sources are packaged.
+
+            ``True`` always packages the resolved input data.
+
+            ``False`` never packages it.
+        discard_example_data
+            Normally, save one processed example when available. Set 
+            discard_example_data=True to turn off.
+
+        """
+
+        if self.data_source is None:
+            source = DataSource()
+        else:
+            source = self.data_source
+
+
+        path = Path(path)
+        path.mkdir(parents=True, exist_ok=True)
+
+        if save_source_data is None:
+            save_source_data = not source.is_remote
+        if not save_source_data and source.uri is None:
+            raise ValueError(
+                "Reference-only checkpoint requires a source URI"
+            )
+        if save_source_data:
+            data_filename = str(path / DATA_FILENAME)
+            (
+                self.data
+                .with_format(None)
+                .to_parquet(data_filename)
+            )
+            checksum = file_checksum(data_filename)
+        else:
+            checksum = None
+
+        checkpoint_config = {
+            "source": source.to_config(),
+            "selected_rows": self.selected_rows,
+            "artefacts": {
+                "source_data": DATA_FILENAME if save_source_data else None,
+                "data_checksum": checksum,
+                "example": (
+                    EXAMPLE_FILENAME
+                    if (
+                        not discard_example_data
+                        and self.data_example is not None
+                    )
+                    else None
+                ),
+            },
+            "schema_version": SCHEMA_VERSION,
+        }
+
+        save_json(
+            checkpoint_config, 
+            str(path / CONFIG_FILENAME),
+        )
+
+        if (
+            not discard_example_data
+            and self.data_example is not None
+        ):
+            (
+                self.data_example
+                .with_format(None)
+                .to_parquet(str(path / EXAMPLE_FILENAME))
+            )
+
+    @classmethod
+    def load(
+        cls,
+        path: str | os.PathLike,
+        *,
+        cache_dir: str | None = None,
+    ) -> "DataResolver":
+        """Restore a saved DataResolver and its available data artefacts."""
+        path = Path(path)
+        config = load_json(path / CONFIG_FILENAME)
+
+        data_source = DataSource.from_config(
+            config.get("source")
+        )
+        artefacts = config.get("artefacts", {})
+        data_filename = artefacts.get("source_data")
+        data_checksum = artefacts.get("data_checksum") 
+        example_filename = artefacts.get("example")
+        selected_rows = config.get("selected_rows")
+        schema_version = config.get("schema_version")
+
+        if schema_version != SCHEMA_VERSION:
+            raise AttributeError(
+                f"Checkpoint schema version is {schema_version}, incompatible with package version {SCHEMA_VERSION}."
+            )
+
+        if data_filename is not None:
+            data_uri = str(path / data_filename)
+            this_checksum = file_checksum(data_uri)
+            if not this_checksum == data_checksum:
+                raise IOError(f"Checksum for {data_uri} should be {data_checksum}, but was {this_checksum}.")
+            has_embedded_data = True
+        elif data_source is not None:
+            if data_source.uri is None:
+                raise ValueError(
+                    f"Reference-only checkpoint has no source URI:\n{data_source}"
+                )
+            elif (
+                not data_source.is_remote
+                and not os.path.exists(data_source.uri)
+            ):
+                raise FileNotFoundError(
+                    "Checkpoint references non-remote data "
+                    f"at {data_source.uri}, but that source does not exist."
+                )
+            if not data_source.is_remote:
+                data_source.assert_verified()
+            data_uri = data_source.uri
+            has_embedded_data = False
+        else:
+            raise AttributeError(f"Checkpoint is missing a data definition!\n{config=}")
+
+        data = cls(
+            data=data_uri,
+            cache_dir=cache_dir,
+        )
+        if not has_embedded_data and selected_rows is not None:
+            data = data.select_rows(selected_rows)
+        data.selected_rows = cls._normalize_selected_rows(selected_rows)
+        if data_source is not None:
+            data.data_source = data_source
+
+        if example_filename is not None:
+            data.data_example = autoload(
+                path / example_filename,
+                cache_dir=cache_dir,
+            )
+        if data.data_example is not None:
+            data.data_example, data.data_shape = data._inspect_data(
+                data.data, 
+                data.data_example,
+            )
+        return data
+
+    # for compatibility
+    def save_checkpoint(self, *args, **kwargs):
+        return self.save(*args, **kwargs)
+
+    @classmethod
+    def load_checkpoint(cls, *args, **kwargs):
+        return cls.load(*args, **kwargs)