fkat 0.1.2__py3-none-any.whl

fkat/data/datasets/map.py

@@ -0,0 +1,90 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+from typing import TypeVar
+from collections.abc import Callable, Iterator
+
+from fkat.data.datasets import SizedDataset
+from torch.utils.data import IterableDataset
+
+T_in = TypeVar("T_in", contravariant=True)
+T_from = TypeVar("T_from", covariant=True)
+T_to = TypeVar("T_to", covariant=True)
+
+
+class MapDataset(SizedDataset[T_in, T_to]):
+    """A :class:`Dataset` that transforms the samples from another :class:`Dataset` using a function."""
+
+    def __init__(
+        self,
+        dataset: SizedDataset[T_in, T_from],
+        fn: Callable[[T_from], T_to],
+    ) -> None:
+        """Create a :class:`Dataset` that maps samples of another :class:`Dataset` using a function.
+
+        Args:
+            dataset (SizedDataset): Source :class:`Dataset`.
+            fn (Callable[[T_from], T_to]): Sample transformation function.
+
+        Returns:
+            None
+        """
+        self.dataset = dataset
+        self.fn = fn
+
+    def __len__(self) -> int:
+        """Get :class:`Dataset` size.
+
+        Returns:
+            int: :class:`Dataset` size.
+        """
+        return len(self.dataset)
+
+    def __getitems__(self, idxs: list[T_in]) -> list[T_to]:
+        """Get a batch of samples at the specified indices.
+
+        Args:
+            idxs (List[T_in]): Samples' indices.
+
+        Returns:
+            List[T_to]: A batch of samples.
+        """
+        if getitems := getattr(self.dataset, "__getitems__", None):
+            batch = getitems(idxs)
+        else:
+            batch = [self.dataset[idx] for idx in idxs]
+        samples = [self.fn(sample) for sample in batch]
+        return samples
+
+    def __getitem__(self, idx: T_in) -> T_to:
+        """Get a sample at the specified index.
+
+        Args:
+            idx (T_in): Sample index.
+
+        Returns:
+            T_to: A sample.
+        """
+        sample = self.fn(self.dataset[idx])
+        return sample
+
+
+class IterableMapDataset(IterableDataset[T_to]):
+    """An :class:`IterableDataset` that transforms the samples from another
+    :class:`IterableDataset` using a function."""
+
+    def __init__(
+        self,
+        dataset: IterableDataset[T_from],
+        fn: Callable[[T_from], T_to],
+    ) -> None:
+        self.dataset = dataset
+        self.fn = fn
+
+    def __iter__(self) -> Iterator[T_to]:
+        """Get :class:`IterableDataset` iterator.
+
+        Yields:
+            T_to: A sample.
+        """
+        for sample in iter(self.dataset):
+            yield self.fn(sample)

fkat/data/datasets/parquet.py

@@ -0,0 +1,242 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+from typing import Any
+from collections.abc import Iterator
+import awswrangler as s3wr
+from pyarrow.fs import FileSystem, S3FileSystem  # type: ignore[possibly-unbound-import]
+import numpy as np
+import pyarrow as pa
+from torch.utils.data import IterableDataset
+
+from fkat.data.datasets import SizedDataset
+from fkat.utils.pyarrow import iter_rows as pa_iter_rows
+from fkat.utils.pandas import iter_rows as pd_iter_rows
+from fkat.utils.boto3 import session
+
+
+class IterableParquetDataset(IterableDataset[dict[str, Any]]):
+    """
+    An :class:`IterableDataset` backed by Parquet data.
+
+    .. note:: If you want to keep the original type from reading parquet,
+        you should set ``dtype_backend='pyarrow'``.
+
+        example config:
+
+        .. code-block:: yaml
+
+            _target_: fkat.data.datasets.parquet.IterableParquetDataset
+            uri: s3://path/to/fkat.parquet
+            dtype_backend: pyarrow
+
+    Args:
+        uri (str or list[str]): URI of Parquet data.
+        columns (List[str], optional): Columns to load.
+            Default to ``None``
+        use_threads (bool): Use multi-threaded processing.
+            Default to ``True``.
+        chunk_size (int): An iterable of DataFrames is returned with maximum rows equal to the received INTEGER.
+            Default to ``10000``
+        replace_nan (bool): Whether to replace np.nan as None.
+            Default to ``True``
+        s3wr_args (dict): config for s3wr.s3.read_parquet,
+            refer to https://aws-sdk-pandas.readthedocs.io/en/3.5.1/stubs/awswrangler.s3.read_parquet.html
+    """
+
+    def __init__(
+        self,
+        uri: str | list[str],
+        columns: list[str] | None = None,
+        use_threads: bool = True,
+        chunk_size: int = 10000,
+        replace_nan: bool = True,
+        **s3wr_args: Any,
+    ) -> None:
+        fs: FileSystem
+        path: str
+        if isinstance(uri, str):
+            fs, path = FileSystem.from_uri(uri)
+        else:
+            fs, path = FileSystem.from_uri(uri[0])
+        self.s3_file = isinstance(fs, S3FileSystem)
+        if self.s3_file:
+            self.uri = uri
+            self.use_threads = use_threads
+            self.columns = columns
+            self.chunk_size = chunk_size
+            self.replace_nan = replace_nan
+            self.s3wr_args = s3wr_args
+        else:
+            # otherwise, use pyarrow
+            path_list = []
+            if isinstance(uri, str):
+                _, path = FileSystem.from_uri(uri)
+                path_list.append(path)
+            else:
+                path_list = []
+                for each in uri:
+                    _, path = FileSystem.from_uri(each)
+                    path_list.append(path)
+            pds = pa.parquet.ParquetDataset(path_list, filesystem=fs)  # type: ignore
+            self.tbl = pds.read(columns, use_threads=use_threads, use_pandas_metadata=False)
+            self.chunk_size = chunk_size
+
+    def __iter__(self) -> Iterator[dict[str, Any]]:
+        """Creates dataset iterator.
+        Returns:
+            Iterator[dict[str, Any]]: dataset iterator
+        """
+        if self.s3_file:
+            return pd_iter_rows(
+                s3wr.s3.read_parquet(
+                    self.uri,
+                    use_threads=self.use_threads,
+                    columns=self.columns,
+                    chunked=self.chunk_size,
+                    boto3_session=session(clients=["s3"]),
+                    path_suffix="parquet",
+                    **self.s3wr_args,
+                ),
+                self.replace_nan,
+            )
+        else:
+            return pa_iter_rows(self.tbl, chunk_size=self.chunk_size)
+
+
+class ParquetDataset(SizedDataset[int, dict[str, Any]]):
+    """
+    A :class:`Dataset` backed by Parquet data.
+
+    Create a :class:`Dataset` from Parquet data at the specified URI.
+
+    .. note:: If you want to keep the original type from reading parquet,
+        you should set ``dtype_backend='pyarrow'``.
+
+        example config:
+
+        .. code-block:: yaml
+
+            _target_: fkat.data.datasets.parquet.ParquetDataset
+            uri: s3://path/to/fkat.parquet
+            dtype_backend: pyarrow
+
+        Difference for ``dtype_backend`` between ``pyarrow`` and ``numpy_nullable``
+
+        .. code-block:: python
+
+            from fkat.data.datasets.parquet import ParquetDataset
+            from fkat.utils.s3_utils import fs_save_prediction_output_parquet
+
+            uri = "s3://path/to/fkat.parquet"
+
+            saved_data = {
+                "purchased_items": [
+                    [
+                        {"product_id": "PROD001", "item_index": 12345, "quantity": "1"},
+                        {"product_id": "PROD002", "item_index": None, "quantity": "1"},
+                    ],
+                    [{"product_id": "PROD001", "item_index": 12345, "quantity": "1"}],
+                ],
+                "ground_truth": [[1, 2, 3], [1, 2]],
+                "embeddings": [np.random.randn(128), np.random.randn(128)],
+            }
+            fs_save_prediction_output_parquet()(saved_data, uri)
+
+            dataset = ParquetDataset(uri)  # dtype_backend: numpy_nullable
+            print(type(dataset[0]["embeddings"]))  # type: numpy.ndarray
+            print(type(dataset[0]["purchased_items"]))  # type: numpy.ndarray of object
+            print(type(dataset[0]["ground_truth"]))  # type: numpy.ndarray of object
+
+            pyarrow_dataset = ParquetDataset(uri, dtype_backend="pyarrow")  # dtype_backend: pyarrow
+            print(type(pyarrow_dataset[0]["embeddings"]))  # type: list
+            print(type(pyarrow_dataset[0]["purchased_items"]))  # type: list of dictionary
+            print(type(pyarrow_dataset[0]["ground_truth"]))  # type: list of int
+
+    Args:
+        uri (str | list[str]): URI of Parquet data.
+        columns (list[str] | None): Columns to load.
+        use_threads (bool): Use multi-threaded processing.
+            Defaults to ``True``.
+        replace_nan (bool): Whether to replace np.nan as None.
+            Default to ``True``
+        s3wr_args (dict): config for s3wr.s3.read_parquet,
+            refer to https://aws-sdk-pandas.readthedocs.io/en/3.5.1/stubs/awswrangler.s3.read_parquet.html
+    """
+
+    def __init__(
+        self,
+        uri: str | list[str],
+        columns: list[str] | None = None,
+        use_threads: bool = True,
+        replace_nan: bool = True,
+        **s3wr_args: Any,
+    ) -> None:
+        # check FileSystem
+        fs: FileSystem
+        path: str
+        if isinstance(uri, str):
+            fs, path = FileSystem.from_uri(uri)
+        else:
+            fs, path = FileSystem.from_uri(uri[0])
+        if isinstance(fs, S3FileSystem):
+            # if file is in S3, then use awswrangler
+            self.df = s3wr.s3.read_parquet(
+                uri,
+                use_threads=use_threads,
+                columns=columns,
+                boto3_session=session(clients=["s3"]),
+                path_suffix="parquet",
+                **s3wr_args,
+            )
+            if replace_nan:
+                self.df = self.df.replace({np.nan: None})
+        else:
+            # otherwise, use pyarrow
+            path_list = []
+            if isinstance(uri, str):
+                fs, path = FileSystem.from_uri(uri)
+                path_list.append(path)
+            elif isinstance(uri, list):
+                path_list = []
+                for each in uri:
+                    _, path = FileSystem.from_uri(each)
+                    path_list.append(path)
+            else:
+                raise Exception(f"ParquetDataset can't support uri as {type(uri)}")
+            pds = pa.parquet.ParquetDataset(path_list, filesystem=fs)  # type: ignore
+            tbl = pds.read(columns, use_threads=use_threads, use_pandas_metadata=False)
+            self.df = tbl.to_pandas()
+
+    def __len__(self) -> int:
+        """Get :class:`Dataset` size.
+
+        Returns:
+            int: :class:`Dataset` size.
+        """
+        return len(self.df)
+
+    def __getitems__(self, idxs: list[int]) -> list[dict[str, Any]]:
+        """Get a batch of samples at the specified indices.
+
+        Args:
+            idxs (list[int]): Samples' indices.
+
+        Returns:
+            list[dict[str, Any]]: A batch of samples.
+        """
+        series = self.df.iloc[idxs]
+        samples = [series.iloc[i].to_dict() for i in range(len(idxs))]
+        return samples
+
+    def __getitem__(self, idx: int) -> dict[str, Any]:
+        """Get a sample at the specified index.
+
+        Args:
+            idx (int): Sample index.
+
+        Returns:
+            dict[str, Any]: A sample.
+        """
+        series = self.df.iloc[idx]
+        sample = series.to_dict()
+        return sample

fkat/data/datasets/sized.py

@@ -0,0 +1,31 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+from typing import TypeVar
+
+from typing_extensions import Protocol
+
+T_in = TypeVar("T_in", contravariant=True)
+T_out = TypeVar("T_out", covariant=True)
+
+
+class SizedDataset(Protocol[T_in, T_out]):
+    """A :class:`Dataset`  with a known size."""
+
+    def __len__(self) -> int:
+        """Get :class:`Dataset` size.
+
+        Returns:
+            int: :class:`Dataset` size.
+        """
+        ...
+
+    def __getitem__(self, idx: T_in) -> T_out:
+        """Get a sample at the specified index.
+
+        Args:
+            idx (T_in): Sample index.
+
+        Returns:
+            T_out: A sample.
+        """
+        ...

fkat/data/dict.py

@@ -0,0 +1,42 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+from typing import Any
+from collections.abc import Iterable, Iterator
+
+from typing_extensions import override
+
+from fkat.data.samplers.strategies import SamplerStrategy
+from fkat.utils.config import to_primitive_container
+
+
+def wrap(key: str, name: str, batch: dict[str, Any]) -> dict[str, Any]:
+    if not isinstance(batch, dict) or key in batch:
+        raise RuntimeError(f"DataLoaders must return a batch dict without {key} key")
+    batch[key] = name
+    return batch
+
+
+class DictDataLoader(Iterable[dict[str, Any]]):
+    """A :class:`LightningDataModule` that manages multiple :class:`DataLoader`\\s for different stages."""
+
+    def __init__(
+        self,
+        dataloaders: dict[str, Iterable[dict[str, Any]]],
+        strategy: SamplerStrategy,
+        key: str = "dataset",
+    ) -> None:
+        self.dataloaders = to_primitive_container(dataloaders)
+        self.strategy = strategy
+        self.key = key
+
+    @override
+    def __iter__(self) -> Iterator[dict[str, Any]]:
+        iters = {k: iter(self.dataloaders[k]) for k in self.dataloaders}
+        for name in self.strategy:
+            if not iters:
+                return
+            if it := iters.get(name):
+                try:
+                    yield wrap(self.key, name, next(it))
+                except StopIteration:
+                    del iters[name]

fkat/data/samplers/__init__.py

@@ -0,0 +1,9 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+from .dict import DictBatchSampler
+from .sized import SizedSampler
+
+__all__ = [
+    "SizedSampler",
+    "DictBatchSampler",
+]

fkat/data/samplers/dict.py

@@ -0,0 +1,38 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+from collections.abc import Iterator
+
+from typing_extensions import override
+
+from fkat.data.samplers.sized import SizedSampler
+from fkat.data.samplers.strategies import SamplerStrategy
+
+
+class DictBatchSampler(SizedSampler[tuple[str, list[int]]]):
+    def __init__(self, strategy: SamplerStrategy, samplers: dict[str, SizedSampler[list[int]]]) -> None:
+        self.strategy = strategy
+        self.samplers = samplers
+        self.len = sum(len(sampler) for sampler in samplers.values())
+
+    @override
+    def __len__(self) -> int:
+        return self.len
+
+    @override
+    def __iter__(self) -> Iterator[tuple[str, list[int]]]:
+        rem_samplers = {name: iter(sampler) for name, sampler in self.samplers.items()}
+        for key in self.strategy:
+            if not rem_samplers:
+                # no more iterators left, stopping iteration
+                return
+            if key not in rem_samplers:
+                # this sampler is exhausted, skipping to sample the remaining ones
+                # this for example will adjust the effective weights for the remaining ones when sampling
+                # or sample next ones for sequential or frequency-based samplers
+                continue
+            try:
+                batch = next(rem_samplers[key])
+                yield (key, batch)
+            except StopIteration:
+                # this sampler is exhausted, removing from consideration
+                del rem_samplers[key]

fkat/data/samplers/sized.py

@@ -0,0 +1,16 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+from typing import TypeVar
+from collections.abc import Iterator
+
+from typing_extensions import Protocol
+
+T_co = TypeVar("T_co", covariant=True)
+
+
+class SizedSampler(Protocol[T_co]):
+    """A Sampler with a known size."""
+
+    def __len__(self) -> int: ...
+
+    def __iter__(self) -> Iterator[T_co]: ...

fkat/data/samplers/strategies.py

@@ -0,0 +1,68 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+from typing import Any
+from collections.abc import Iterator
+
+import numpy as np
+from typing_extensions import Protocol, override
+
+
+class SamplerStrategy(Protocol):
+    """
+    This strategy decides which Sampler to use next and returns its label
+    """
+
+    def __iter__(self) -> Iterator[str]: ...
+
+
+class Weighted(SamplerStrategy):
+    """
+    Sample the label to generate next microbatch from using the provided weight distribution.
+    For uniform distribution use the same weight for all labels.
+    """
+
+    def __init__(self, weights: dict[str, float]) -> None:
+        self.names = []
+        self.weights = []
+        for name, weight in weights.items():
+            self.names.append(name)
+            self.weights.append(weight)
+
+    @override
+    def __iter__(self) -> Iterator[str]:
+        while True:
+            yield np.random.choice(self.names, p=self.weights)
+
+
+class RoundRobin(SamplerStrategy):
+    """
+    Specifies the order of labels to generate microbatches from.
+    """
+
+    def __init__(self, names: list[str]) -> None:
+        self.names = names
+
+    @override
+    def __iter__(self) -> Iterator[str]:
+        i = -1
+        while True:
+            i = (i + 1) % len(self.names)
+            yield self.names[i]
+
+
+class Frequency(SamplerStrategy):
+    """
+    Specifies the order and number of microbatches to generate for specific labels.
+    E.g. [["first", 2], ["second", 1], ["first", 3], ["third", 1]]
+    """
+
+    def __init__(self, freq: list[list[Any]]) -> None:
+        assert all(isinstance(e[0], str) and isinstance(e[1], int) for e in freq)
+        self.freq = freq
+
+    @override
+    def __iter__(self) -> Iterator[str]:
+        while True:
+            for name, count in self.freq:
+                for _ in range(count):
+                    yield name