python-wml 3.0.0__py3-none-any.whl

wml/wtorch/data/datapipes/iter/batch.py

@@ -0,0 +1,126 @@
+import warnings
+from torch.utils.data import IterDataPipe
+from typing import TypeVar, Optional, Iterator, List, Sized, Callable
+
+T_co = TypeVar('T_co', covariant=True)
+
+
+class BatchIterDataPipe(IterDataPipe[List[T_co]]):
+    r""" :class:`BatchIterDataPipe`.
+
+    Iterable DataPipe to create mini-batches of data. An outer dimension will be added as
+    `batch_size` if `drop_last` is set to `True`, or `length % batch_size` for the
+    last batch if `drop_last` is set to `False`.
+    args:
+        datapipe: Iterable DataPipe being batched
+        batch_size: The size of each batch
+        drop_last: Option to drop the last batch if it's not full
+    """
+    datapipe: IterDataPipe[T_co]
+    batch_size: int
+    drop_last: bool
+    length: Optional[int]
+
+    def __init__(self,
+                 datapipe: IterDataPipe[T_co],
+                 *,
+                 batch_size: int,
+                 drop_last: bool = False,
+                 ) -> None:
+        assert batch_size > 0, "Batch size is required to be larger than 0!"
+        super().__init__()
+        self.datapipe = datapipe
+        self.batch_size = batch_size
+        self.drop_last = drop_last
+        self.length = None
+
+    def __iter__(self) -> Iterator[List[T_co]]:
+        batch: List[T_co] = []
+        for x in self.datapipe:
+            batch.append(x)
+            if len(batch) == self.batch_size:
+                yield batch
+                batch.clear()
+        if len(batch) > 0:
+            if not self.drop_last:
+                yield batch
+            batch.clear()
+
+    def __len__(self) -> int:
+        if self.length is not None:
+            return self.length
+        if isinstance(self.datapipe, Sized) and len(self.datapipe) >= 0:
+            if self.drop_last:
+                self.length = len(self.datapipe) // self.batch_size
+            else:
+                self.length = (len(self.datapipe) + self.batch_size - 1) // self.batch_size
+            return self.length
+        raise NotImplementedError
+
+
+class BucketBatchIterDataPipe(IterDataPipe[List[T_co]]):
+    r""" :class:`BucketBatchIterDataPipe`.
+
+    Iterable DataPipe to create mini-batches of data from sorted bucket. An outer
+    dimension will be added as `batch_size` if `drop_last` is set to `True`,
+    or `length % batch_size` for the last batch if `drop_last` is set to `False`.
+        args:
+        datapipe: Iterable DataPipe being batched
+        batch_size: The size of each batch
+        drop_last: Option to drop the last batch if it's not full
+        bucket_size_mul: The multiplier to specify the size of bucket
+        sort_key: Callable to specify the comparison key for sorting within bucket
+    """
+    datapipe: IterDataPipe[T_co]
+    batch_size: int
+    drop_last: bool
+    bucket_size_mul: int
+    sort_key: Optional[Callable]
+    length: Optional[int]
+
+    def __init__(self,
+                 datapipe: IterDataPipe[T_co],
+                 *,
+                 batch_size: int,
+                 drop_last: bool = False,
+                 bucket_size_mul: int = 100,
+                 sort_key: Optional[Callable] = None,
+                 ) -> None:
+        assert batch_size > 0, "Batch size is required to be larger than 0!"
+        super().__init__()
+        self.datapipe = datapipe
+        self.batch_size = batch_size
+        self.drop_last = drop_last
+        self.bucket_size = batch_size * bucket_size_mul
+        self.sort_key = sort_key
+        if sort_key is not None and sort_key.__name__ == '<lambda>':
+            warnings.warn("Lambda function is not supported for pickle, "
+                          "please use regular python function instead.")
+        self.bucket_ds = BatchIterDataPipe(datapipe, batch_size=self.bucket_size, drop_last=False)
+        self.length = None
+
+    def __iter__(self) -> Iterator[List[T_co]]:
+        # Bucket without sorting remains same order, directly returns BatchDataset
+        if self.sort_key is None:
+            yield from BatchIterDataPipe(self.datapipe, batch_size=self.batch_size, drop_last=self.drop_last)
+        else:
+            bucket: List[T_co]
+            batch: List[T_co] = []
+            for bucket in self.bucket_ds:
+                # In-place sort within bucket
+                bucket.sort(key=self.sort_key)
+                for start in range(0, len(bucket), self.batch_size):
+                    batch = bucket[start: start + self.batch_size]
+                    if len(batch) == self.batch_size or not self.drop_last:
+                        yield batch
+
+    def __len__(self) -> int:
+        if self.length is not None:
+            return self.length
+        if isinstance(self.datapipe, Sized) and len(self.datapipe) >= 0:
+            if self.drop_last:
+                self.length = len(self.datapipe) // self.batch_size
+            else:
+                self.length = (len(self.datapipe) + self.batch_size - 1) // self.batch_size
+            return self.length
+        raise NotImplementedError

wml/wtorch/data/datapipes/iter/callable.py

@@ -0,0 +1,92 @@
+import warnings
+from torch.utils.data import IterDataPipe, _utils
+from typing import TypeVar, Callable, Iterator, Sized
+
+T_co = TypeVar('T_co', covariant=True)
+
+
+# Default function to return each item directly
+# In order to keep datapipe picklable, eliminates the usage
+# of python lambda function
+def default_fn(data):
+    return data
+
+
+class CallableIterDataPipe(IterDataPipe[T_co]):
+    r""" :class:`CallableIterDataPipe`.
+
+    Iterable DataPipe to run a function over each item from the source DataPipe.
+    args:
+        datapipe: Source Iterable DataPipe
+        fn: Function called over each item
+    """
+    datapipe: IterDataPipe
+    fn: Callable
+
+    def __init__(self,
+                 datapipe: IterDataPipe,
+                 *args,
+                 fn: Callable = default_fn,
+                 **kwargs,
+                 ) -> None:
+        super().__init__()
+        self.datapipe = datapipe
+        if fn.__name__ == '<lambda>':
+            warnings.warn("Lambda function is not supported for pickle, "
+                          "please use regular python function instead.")
+        self.fn = fn  # type: ignore
+        self.args = args
+        self.kwargs = kwargs
+
+    def __iter__(self) -> Iterator[T_co]:
+        for data in self.datapipe:
+            yield self.fn(data, *self.args, **self.kwargs)
+
+    def __len__(self) -> int:
+        if isinstance(self.datapipe, Sized) and len(self.datapipe) >= 0:
+            return len(self.datapipe)
+        raise NotImplementedError
+
+
+class CollateIterDataPipe(CallableIterDataPipe):
+    r""" :class:`CollateIterDataPipe`.
+
+    Iterable DataPipe to collate samples from datapipe to Tensor(s) by `util_.collate.default_collate`,
+    or customized Data Structure by collate_fn.
+    args:
+        datapipe: Iterable DataPipe being collated
+        collate_fn: Customized collate function to collect and combine data or a batch of data.
+                    Default function collates to Tensor(s) based on data type.
+
+    Example: Convert integer data to float Tensor
+        >>> class MyIterDataPipe(torch.utils.data.IterDataPipe):
+        ...     def __init__(self, start, end):
+        ...         super(MyIterDataPipe).__init__()
+        ...         assert end > start, "this example code only works with end >= start"
+        ...         self.start = start
+        ...         self.end = end
+        ...
+        ...     def __iter__(self):
+        ...         return iter(range(self.start, self.end))
+        ...
+        ...     def __len__(self):
+        ...         return self.end - self.start
+        ...
+        >>> ds = MyIterDataPipe(start=3, end=7)
+        >>> print(list(ds))
+        [3, 4, 5, 6]
+
+        >>> def collate_fn(batch):
+        ...     return torch.tensor(batch, dtype=torch.float)
+        ...
+        >>> collated_ds = CollateIterDataPipe(ds, collate_fn=collate_fn)
+        >>> print(list(collated_ds))
+        [tensor(3.), tensor(4.), tensor(5.), tensor(6.)]
+    """
+    def __init__(self,
+                 datapipe: IterDataPipe,
+                 *args,
+                 collate_fn: Callable = _utils.collate.default_collate,
+                 **kwargs,
+                 ) -> None:
+        super().__init__(datapipe, *args, fn=collate_fn, **kwargs)

wml/wtorch/data/datapipes/iter/listdirfiles.py

@@ -0,0 +1,37 @@
+from torch.utils.data import IterDataPipe
+from torch.utils.data.datapipes.utils.common import get_file_pathnames_from_root
+from typing import List, Union, Iterator
+
+class ListDirFilesIterDataPipe(IterDataPipe):
+    r""" :class:`ListDirFilesIterDataPipe`
+
+    Iterable DataPipe to load file pathname(s) (path + filename), yield pathname from given disk root dir.
+    args:
+        root : root dir
+        mask : a unix style filter string or string list for filtering file name(s)
+        abspath : whether to return relative pathname or absolute pathname
+        length : a nominal length of the datapipe
+    """
+
+    def __init__(
+            self,
+            root: str = '.',
+            masks: Union[str, List[str]] = '',
+            *,
+            recursive: bool = False,
+            abspath: bool = False,
+            length: int = -1):
+        super().__init__()
+        self.root : str = root
+        self.masks : Union[str, List[str]] = masks
+        self.recursive : bool = recursive
+        self.abspath : bool = abspath
+        self.length : int = length
+
+    def __iter__(self) -> Iterator[str] :
+        yield from get_file_pathnames_from_root(self.root, self.masks, self.recursive, self.abspath)
+
+    def __len__(self):
+        if self.length == -1:
+            raise NotImplementedError
+        return self.length

wml/wtorch/data/datapipes/iter/loadfilesfromdisk.py

@@ -0,0 +1,30 @@
+from torch.utils.data import IterDataPipe
+from torch.utils.data.datapipes.utils.common import get_file_binaries_from_pathnames
+from typing import Iterable, Iterator, Tuple
+from io import BufferedIOBase
+
+class LoadFilesFromDiskIterDataPipe(IterDataPipe):
+    r""" :class:`LoadFilesFromDiskIterDataPipe`.
+
+    Iterable Datapipe to load file binary streams from given pathnames,
+    yield pathname and binary stream in a tuple.
+    args:
+        datapipe: Iterable datapipe that provides pathnames
+        length: a nominal length of the datapipe
+    """
+
+    def __init__(
+            self,
+            datapipe : Iterable[str],
+            length : int = -1):
+        super().__init__()
+        self.datapipe : Iterable = datapipe
+        self.length : int = length
+
+    def __iter__(self) -> Iterator[Tuple[str, BufferedIOBase]] :
+        yield from get_file_binaries_from_pathnames(self.datapipe)
+
+    def __len__(self):
+        if self.length == -1:
+            raise NotImplementedError
+        return self.length

wml/wtorch/data/datapipes/iter/readfilesfromtar.py

@@ -0,0 +1,60 @@
+from torch.utils.data import IterDataPipe
+from torch.utils.data.datapipes.utils.common import validate_pathname_binary_tuple
+from typing import Iterable, Iterator, Tuple, Optional, IO, cast
+from io import BufferedIOBase
+
+import os
+import tarfile
+import warnings
+
+class ReadFilesFromTarIterDataPipe(IterDataPipe):
+    r""" :class:`ReadFilesFromTarIDP`.
+
+    Iterable datapipe to extract tar binary streams from input iterable which contains tuples of
+    pathname and tar binary stream, yields pathname and extracted binary stream in a tuple.
+    args:
+        datapipe: Iterable datapipe that provides pathname and tar binary stream in tuples
+        length: a nominal length of the datapipe
+    """
+    def __init__(
+            self,
+            datapipe : Iterable[Tuple[str, BufferedIOBase]],
+            length : int = -1):
+        super().__init__()
+        self.datapipe : Iterable[Tuple[str, BufferedIOBase]] = datapipe
+        self.length : int = length
+
+
+    def __iter__(self) -> Iterator[Tuple[str, BufferedIOBase]]:
+        if not isinstance(self.datapipe, Iterable):
+            raise TypeError("datapipe must be Iterable type but got {}".format(type(self.datapipe)))
+        for data in self.datapipe:
+            validate_pathname_binary_tuple(data)
+            pathname, data_stream = data
+            try:
+                # typing.cast is used here to silence mypy's type checker
+                tar = tarfile.open(fileobj=cast(Optional[IO[bytes]], data_stream), mode="r:*")
+                for tarinfo in tar:
+                    if not tarinfo.isfile():
+                        continue
+                    extracted_fobj = tar.extractfile(tarinfo)
+                    if extracted_fobj is None:
+                        warnings.warn("failed to extract file {} from source tarfile {}".format(tarinfo.name, pathname))
+                        raise tarfile.ExtractError
+                    inner_pathname = os.path.normpath(os.path.join(pathname, tarinfo.name))
+                    # Add a reference of the source tarfile into extracted_fobj, so the source
+                    # tarfile handle won't be released until all the extracted file objs are destroyed.
+                    # Add `# type: ignore` to silence mypy's type checker
+                    extracted_fobj.source_tarfile_ref = tar  # type: ignore
+                    # typing.cast is used here to silence mypy's type checker
+                    yield (inner_pathname, cast(BufferedIOBase, extracted_fobj))
+            except Exception as e:
+                warnings.warn(
+                    "Unable to extract files from corrupted tarfile stream {} due to: {}, abort!".format(pathname, e))
+                raise e
+
+
+    def __len__(self):
+        if self.length == -1:
+            raise NotImplementedError
+        return self.length

wml/wtorch/data/datapipes/iter/readfilesfromzip.py

@@ -0,0 +1,63 @@
+from torch.utils.data import IterDataPipe
+from torch.utils.data.datapipes.utils.common import validate_pathname_binary_tuple
+from typing import Iterable, Iterator, Tuple, IO, cast
+from io import BufferedIOBase
+
+import os
+import sys
+import zipfile
+import warnings
+
+class ReadFilesFromZipIterDataPipe(IterDataPipe):
+    r""" :class:`ReadFilesFromZipIterDataPipe`.
+
+    Iterable data pipe to extract zip binary streams from input iterable which contains tuples of
+    pathname and zip binary stream, yields pathname and extracted binary stream in a tuple.
+    args:
+        datapipe: Iterable datapipe that provides pathname and zip binary stream in tuples
+        length: a nominal length of the datapipe
+    """
+    def __init__(
+            self,
+            datapipe : Iterable[Tuple[str, BufferedIOBase]],
+            length : int = -1):
+        super().__init__()
+        self.datapipe : Iterable[Tuple[str, BufferedIOBase]] = datapipe
+        self.length : int = length
+
+
+    def __iter__(self) -> Iterator[Tuple[str, BufferedIOBase]]:
+        if not isinstance(self.datapipe, Iterable):
+            raise TypeError("datapipe must be Iterable type but got {}".format(type(self.datapipe)))
+        for data in self.datapipe:
+            validate_pathname_binary_tuple(data)
+            pathname, data_stream = data
+            try:
+                # typing.cast is used here to silence mypy's type checker
+                zips = zipfile.ZipFile(cast(IO[bytes], data_stream))
+                for zipinfo in zips.infolist():
+                    # major version should always be 3 here.
+                    if sys.version_info[1] >= 6:
+                        if zipinfo.is_dir():
+                            continue
+                    elif zipinfo.filename.endswith('/'):
+                        continue
+
+                    extracted_fobj = zips.open(zipinfo)
+                    inner_pathname = os.path.normpath(os.path.join(pathname, zipinfo.filename))
+                    # Add a reference of the source zipfile into extracted_fobj, so the source
+                    # zipfile handle won't be released until all the extracted file objs are destroyed.
+                    # Add `# type: ignore` to silence mypy's type checker
+                    extracted_fobj.source_zipfile_ref = zips  # type: ignore
+                    # typing.cast is used here to silence mypy's type checker
+                    yield (inner_pathname, cast(BufferedIOBase, extracted_fobj))
+            except Exception as e:
+                warnings.warn(
+                    "Unable to extract files from corrupted zipfile stream {} due to: {}, abort!".format(pathname, e))
+                raise e
+
+
+    def __len__(self):
+        if self.length == -1:
+            raise NotImplementedError
+        return self.length

wml/wtorch/data/datapipes/iter/sampler.py

@@ -0,0 +1,94 @@
+from torch.utils.data import IterDataPipe, Sampler, SequentialSampler
+from typing import TypeVar, Type, Iterator, Sized,Optional
+import itertools
+import torch
+
+T_co = TypeVar('T_co', covariant=True)
+
+
+class SamplerIterDataPipe(IterDataPipe[T_co]):
+    r""" :class:`SamplerIterDataPipe`.
+
+    Iterable DataPipe to generate sample elements.
+    args:
+        datapipe: IterDataPipe sampled from
+        sampler: Sampler class to genereate sample elements from input DataPipe.
+                    Default is :class:`SequentialSampler` for IterDataPipe
+    """
+    datapipe: IterDataPipe
+    sampler: Sampler
+
+    def __init__(self,
+                 datapipe: IterDataPipe,
+                 *,
+                 sampler: Type[Sampler] = SequentialSampler,
+                 **kwargs
+                 ) -> None:
+        assert isinstance(datapipe, Sized), \
+            "Sampler class requires input datapipe implemented `__len__`"
+        super().__init__()
+        self.datapipe = datapipe
+        # https://github.com/python/mypy/pull/9629 will solve
+        self.sampler = sampler(data_source=self.datapipe, **kwargs)  # type: ignore
+
+    def __iter__(self) -> Iterator[T_co]:
+        return iter(self.sampler)
+
+    def __len__(self) -> int:
+        # Dataset has been tested as `Sized`
+        if isinstance(self.sampler, Sized) and len(self.sampler) >= 0:
+            return len(self.sampler)
+        raise NotImplementedError
+
+class InfiniteSampler(Sampler):
+    """
+    In training, we only care about the "infinite stream" of training data.
+    So this sampler produces an infinite stream of indices and
+    all workers cooperate to correctly shuffle the indices and sample different indices.
+    The samplers in each worker effectively produces `indices[worker_id::num_workers]`
+    where `indices` is an infinite stream of indices consisting of
+    `shuffle(range(size)) + shuffle(range(size)) + ...` (if shuffle is True)
+    or `range(size) + range(size) + ...` (if shuffle is False)
+    """
+
+    def __init__(
+        self,
+        size: int,
+        shuffle: bool = True,
+        seed: Optional[int] = 0,
+        rank=0,
+        world_size=1,
+    ):
+        """
+        Args:
+            size (int): the total number of data of the underlying dataset to sample from
+            shuffle (bool): whether to shuffle the indices or not
+            seed (int): the initial seed of the shuffle. Must be the same
+                across all workers. If None, will use a random seed shared
+                among workers (require synchronization among all workers).
+        """
+        self._size = size
+        assert size > 0
+        self._shuffle = shuffle
+        self._seed = int(seed)
+
+        self._rank = rank
+        self._world_size = world_size
+
+    def __iter__(self):
+        start = self._rank
+        yield from itertools.islice(
+            self._infinite_indices(), start, None, self._world_size
+        )
+
+    def _infinite_indices(self):
+        g = torch.Generator()
+        g.manual_seed(self._seed)
+        while True:
+            if self._shuffle:
+                yield from torch.randperm(self._size, generator=g)
+            else:
+                yield from torch.arange(self._size)
+
+    def __len__(self):
+        return self._size // self._world_size

wml/wtorch/data/datapipes/utils/common.py

@@ -0,0 +1,65 @@
+import os
+import fnmatch
+import warnings
+from typing import List, Union, Iterable
+from io import BufferedIOBase
+
+
+def match_masks(name : str, masks : Union[str, List[str]]) -> bool:
+    # empty mask matches any input name
+    if not masks:
+        return True
+
+    if isinstance(masks, str):
+        return fnmatch.fnmatch(name, masks)
+
+    for mask in masks:
+        if fnmatch.fnmatch(name, mask):
+            return True
+    return False
+
+
+def get_file_pathnames_from_root(
+        root: str,
+        masks: Union[str, List[str]],
+        recursive: bool = False,
+        abspath: bool = False) -> Iterable[str]:
+
+    # print out an error message and raise the error out
+    def onerror(err : OSError):
+        warnings.warn(err.filename + " : " + err.strerror)
+        raise err
+
+    for path, dirs, files in os.walk(root, onerror=onerror):
+        if abspath:
+            path = os.path.abspath(path)
+        for f in files:
+            if match_masks(f, masks):
+                yield os.path.join(path, f)
+        if not recursive:
+            break
+
+
+def get_file_binaries_from_pathnames(pathnames : Iterable):
+
+    if not isinstance(pathnames, Iterable):
+        warnings.warn("get_file_binaries_from_pathnames needs the input be an Iterable")
+        raise TypeError
+
+    for pathname in pathnames:
+        if not isinstance(pathname, str):
+            warnings.warn("file pathname must be string type, but got {}".format(type(pathname)))
+            raise TypeError
+
+        yield (pathname, open(pathname, 'rb'))
+
+
+def validate_pathname_binary_tuple(data):
+    if not isinstance(data, tuple):
+        raise TypeError("pathname binary data should be tuple type, but got {}".format(type(data)))
+    if len(data) != 2:
+        raise TypeError("pathname binary tuple length should be 2, but got {}".format(str(len(data))))
+    if not isinstance(data[0], str):
+        raise TypeError("pathname binary tuple should have string type pathname, but got {}".format(type(data[0])))
+    if not isinstance(data[1], BufferedIOBase):
+        raise TypeError("pathname binary tuple should have BufferedIOBase based binary type, but got {}".format(type(data[1])))