wehub-resource-sync
841 lines
33 KiB
Python
841 lines
33 KiB
Python
# -*- coding:utf-8 -*-
|
|
# Author: hankcs
|
|
# Date: 2020-05-09 20:27
|
|
import math
|
|
import os
|
|
import random
|
|
import tempfile
|
|
import warnings
|
|
from abc import ABC, abstractmethod
|
|
from copy import copy
|
|
from logging import Logger
|
|
from typing import Union, List, Callable, Iterable, Dict, Any
|
|
|
|
import torch
|
|
import torch.multiprocessing as mp
|
|
from hanlp.common.transform import TransformList, VocabDict, EmbeddingNamedTransform
|
|
from hanlp.common.vocab import Vocab
|
|
from hanlp.components.parsers.alg import kmeans
|
|
from hanlp.utils.io_util import read_cells, get_resource
|
|
from hanlp.utils.time_util import CountdownTimer
|
|
from hanlp.utils.torch_util import dtype_of
|
|
from hanlp_common.configurable import AutoConfigurable
|
|
from hanlp_common.constant import IDX, HANLP_VERBOSE
|
|
from hanlp_common.util import isdebugging, merge_list_of_dict, k_fold
|
|
from torch.nn.utils.rnn import pad_sequence
|
|
from torch.utils.data import Dataset, DataLoader, Sampler
|
|
from torch.utils.data.dataset import IterableDataset
|
|
|
|
|
|
class Transformable(ABC):
|
|
def __init__(self, transform: Union[Callable, List] = None) -> None:
|
|
"""An object which can be transformed with a list of functions. It is the final result of an object being passed
|
|
through a list of functions, while these functions are kept in a list.
|
|
|
|
Args:
|
|
transform: A transform function or a list of functions.
|
|
"""
|
|
super().__init__()
|
|
if isinstance(transform, list) and not isinstance(transform, TransformList):
|
|
transform = TransformList(*transform)
|
|
self.transform: Union[Callable, TransformList] = transform
|
|
|
|
def append_transform(self, transform: Callable):
|
|
"""Append a transform to its list of transforms.
|
|
|
|
Args:
|
|
transform: A new transform to be appended.
|
|
|
|
Returns:
|
|
Itself.
|
|
|
|
"""
|
|
assert transform is not None, 'None transform not allowed'
|
|
if not self.transform:
|
|
self.transform = TransformList(transform)
|
|
elif not isinstance(self.transform, TransformList):
|
|
if self.transform != transform:
|
|
self.transform = TransformList(self.transform, transform)
|
|
else:
|
|
if transform not in self.transform:
|
|
self.transform.append(transform)
|
|
return self
|
|
|
|
def insert_transform(self, index: int, transform: Callable):
|
|
"""Insert a transform to a certain position.
|
|
|
|
Args:
|
|
index: A certain position.
|
|
transform: A new transform.
|
|
|
|
Returns:
|
|
Itself.
|
|
|
|
"""
|
|
assert transform is not None, 'None transform not allowed'
|
|
if not self.transform:
|
|
self.transform = TransformList(transform)
|
|
elif not isinstance(self.transform, TransformList):
|
|
if self.transform != transform:
|
|
self.transform = TransformList(self.transform)
|
|
self.transform.insert(index, transform)
|
|
else:
|
|
if transform not in self.transform:
|
|
self.transform.insert(index, transform)
|
|
return self
|
|
|
|
def transform_sample(self, sample: dict, inplace=False) -> dict:
|
|
"""Apply transforms to a sample.
|
|
|
|
Args:
|
|
sample: A sample, which is a ``dict`` holding features.
|
|
inplace: ``True`` to apply transforms inplace.
|
|
|
|
.. Attention::
|
|
If any transform modifies existing features, it will modify again and again when ``inplace=True``.
|
|
For example, if a transform insert a ``BOS`` token to a list inplace, and it is called twice,
|
|
then 2 ``BOS`` will be inserted which might not be an intended result.
|
|
|
|
Returns:
|
|
Transformed sample.
|
|
"""
|
|
if not inplace:
|
|
sample = copy(sample)
|
|
if self.transform:
|
|
sample = self.transform(sample)
|
|
return sample
|
|
|
|
|
|
class TransformableDataset(Transformable, Dataset, ABC):
|
|
|
|
def __init__(self,
|
|
data: Union[str, List],
|
|
transform: Union[Callable, List] = None,
|
|
cache=None,
|
|
generate_idx=None) -> None:
|
|
"""A :class:`~torch.utils.data.Dataset` which can be applied with a list of transform functions.
|
|
|
|
Args:
|
|
data: The local or remote path to a dataset, or a list of samples where each sample is a dict.
|
|
transform: Predefined transform(s).
|
|
cache: ``True`` to enable caching, so that transforms won't be called twice.
|
|
generate_idx: Create a :const:`~hanlp_common.constants.IDX` field for each sample to store its order in dataset. Useful for prediction when
|
|
samples are re-ordered by a sampler.
|
|
"""
|
|
super().__init__(transform)
|
|
if generate_idx is None:
|
|
generate_idx = isinstance(data, list)
|
|
data_ = self.load_data(data, generate_idx)
|
|
# assert data_, f'No samples loaded from {data}'
|
|
if data_:
|
|
assert isinstance(data_[0], dict
|
|
), f'TransformDataset expects each sample to be a dict but got {type(data_[0])} instead.'
|
|
self.data = data_
|
|
if cache:
|
|
self.cache = [None] * len(data_)
|
|
else:
|
|
self.cache = None
|
|
|
|
def load_data(self, data, generate_idx=False):
|
|
"""A intermediate step between constructor and calling the actual file loading method.
|
|
|
|
Args:
|
|
data: If data is a file, this method calls :meth:`~hanlp.common.dataset.TransformableDataset.load_file`
|
|
to load it.
|
|
generate_idx: Create a :const:`~hanlp_common.constants.IDX` field for each sample to store its order in dataset. Useful for prediction when
|
|
samples are re-ordered by a sampler.
|
|
|
|
Returns: Loaded samples.
|
|
|
|
"""
|
|
if self.should_load_file(data):
|
|
if isinstance(data, str):
|
|
data = get_resource(data)
|
|
data = list(self.load_file(data))
|
|
if generate_idx:
|
|
for i, each in enumerate(data):
|
|
each[IDX] = i
|
|
# elif isinstance(data, list):
|
|
# data = self.load_list(data)
|
|
return data
|
|
|
|
# noinspection PyMethodMayBeStatic
|
|
# def load_list(self, data: list) -> List[Dict[str, Any]]:
|
|
# return data
|
|
|
|
def should_load_file(self, data) -> bool:
|
|
"""Determines whether data is a filepath.
|
|
|
|
Args:
|
|
data: Data to check.
|
|
|
|
Returns: ``True`` to indicate it's a filepath.
|
|
|
|
"""
|
|
return isinstance(data, str)
|
|
|
|
@abstractmethod
|
|
def load_file(self, filepath: str):
|
|
"""The actual file loading logic.
|
|
|
|
Args:
|
|
filepath: The path to a dataset.
|
|
"""
|
|
pass
|
|
|
|
def __getitem__(self, index: Union[int, slice]) -> Union[dict, List[dict]]:
|
|
""" Get the index-th sample in this dataset.
|
|
|
|
Args:
|
|
index: Either a integer index of a list of indices.
|
|
|
|
Returns: Either a sample or or list of samples depending on how many indices are passed in.
|
|
|
|
"""
|
|
# if isinstance(index, (list, tuple)):
|
|
# assert len(index) == 1
|
|
# index = index[0]
|
|
if isinstance(index, slice):
|
|
indices = range(*index.indices(len(self)))
|
|
return [self[i] for i in indices]
|
|
|
|
if self.cache:
|
|
cache = self.cache[index]
|
|
if cache:
|
|
return cache
|
|
sample = self.data[index]
|
|
sample = self.transform_sample(sample)
|
|
if self.cache:
|
|
self.cache[index] = sample
|
|
return sample
|
|
|
|
def __len__(self) -> int:
|
|
return len(self.data)
|
|
|
|
def __repr__(self) -> str:
|
|
return f'{len(self)} samples: {self[0] if len(self) else ""} ...'
|
|
|
|
def purge_cache(self):
|
|
"""Purges all cache. If cache is not enabled, this method enables it.
|
|
"""
|
|
self.cache = [None] * len(self.data)
|
|
|
|
def split(self, *ratios):
|
|
"""Split dataset into subsets.
|
|
|
|
Args:
|
|
*ratios: The ratios for each subset. They can be any type of numbers which will be normalized. For example,
|
|
``8, 1, 1`` are equivalent to ``0.8, 0.1, 0.1``.
|
|
|
|
Returns:
|
|
list[TransformableDataset]: A list of subsets.
|
|
"""
|
|
ratios = [x / sum(ratios) for x in ratios]
|
|
chunks = []
|
|
prev = 0
|
|
for r in ratios:
|
|
cur = prev + math.ceil(len(self) * r)
|
|
chunks.append([prev, cur])
|
|
prev = cur
|
|
chunks[-1][1] = len(self)
|
|
outputs = []
|
|
for b, e in chunks:
|
|
dataset = copy(self)
|
|
dataset.data = dataset.data[b:e]
|
|
if dataset.cache:
|
|
dataset.cache = dataset.cache[b:e]
|
|
outputs.append(dataset)
|
|
return outputs
|
|
|
|
def k_fold(self, k, i):
|
|
"""Perform k-fold sampling.
|
|
|
|
Args:
|
|
k (int): Number of folds.
|
|
i (int): The i-th fold.
|
|
|
|
Returns:
|
|
TransformableDataset: The i-th fold subset of this dataset.
|
|
|
|
"""
|
|
assert 0 <= i <= k, f'Invalid split {i}'
|
|
train_indices, test_indices = k_fold(k, len(self), i)
|
|
return self.subset(train_indices), self.subset(test_indices)
|
|
|
|
def subset(self, indices):
|
|
"""Create a subset given indices of samples.
|
|
|
|
Args:
|
|
indices: Indices of samples.
|
|
|
|
Returns:
|
|
TransformableDataset: The a subset of this dataset.
|
|
"""
|
|
dataset = copy(self)
|
|
dataset.data = [dataset.data[i] for i in indices]
|
|
if dataset.cache:
|
|
dataset.cache = [dataset.cache[i] for i in indices]
|
|
return dataset
|
|
|
|
def shuffle(self):
|
|
"""Shuffle this dataset inplace.
|
|
"""
|
|
if not self.cache:
|
|
random.shuffle(self.data)
|
|
else:
|
|
z = list(zip(self.data, self.cache))
|
|
random.shuffle(z)
|
|
self.data, self.cache = zip(*z)
|
|
|
|
def prune(self, criterion: Callable, logger: Logger = None):
|
|
"""Prune (to discard) samples according to a criterion.
|
|
|
|
Args:
|
|
criterion: A functions takes a sample as input and output ``True`` if the sample needs to be pruned.
|
|
logger: If any, log statistical messages using it.
|
|
|
|
Returns:
|
|
int: Size before pruning.
|
|
"""
|
|
# noinspection PyTypeChecker
|
|
size_before = len(self)
|
|
good_ones = [i for i, s in enumerate(self) if not criterion(s)]
|
|
self.data = [self.data[i] for i in good_ones]
|
|
if self.cache:
|
|
self.cache = [self.cache[i] for i in good_ones]
|
|
if logger:
|
|
size_after = len(self)
|
|
num_pruned = size_before - size_after
|
|
logger.info(f'Pruned [yellow]{num_pruned} ({num_pruned / size_before:.1%})[/yellow] '
|
|
f'samples out of {size_before}.')
|
|
return size_before
|
|
|
|
|
|
class TransformSequentialDataset(Transformable, IterableDataset, ABC):
|
|
pass
|
|
|
|
|
|
class DeviceDataLoader(DataLoader):
|
|
def __init__(self, dataset, batch_size=32, shuffle=False, sampler=None,
|
|
batch_sampler=None, num_workers=None, collate_fn=None,
|
|
pin_memory=False, drop_last=False, timeout=0,
|
|
worker_init_fn=None, multiprocessing_context=None,
|
|
device=None, **kwargs):
|
|
if batch_sampler is not None:
|
|
batch_size = 1
|
|
if num_workers is None:
|
|
if isdebugging():
|
|
num_workers = 0
|
|
else:
|
|
num_workers = 2
|
|
# noinspection PyArgumentList
|
|
super(DeviceDataLoader, self).__init__(dataset=dataset, batch_size=batch_size, shuffle=shuffle,
|
|
sampler=sampler,
|
|
batch_sampler=batch_sampler, num_workers=num_workers,
|
|
collate_fn=collate_fn,
|
|
pin_memory=pin_memory, drop_last=drop_last, timeout=timeout,
|
|
worker_init_fn=worker_init_fn,
|
|
multiprocessing_context=multiprocessing_context, **kwargs)
|
|
self.device = device
|
|
|
|
def __iter__(self):
|
|
for raw_batch in super(DeviceDataLoader, self).__iter__():
|
|
if self.device is not None:
|
|
for field, data in raw_batch.items():
|
|
if isinstance(data, torch.Tensor):
|
|
data = data.to(self.device)
|
|
raw_batch[field] = data
|
|
yield raw_batch
|
|
|
|
def collate_fn(self, samples):
|
|
return merge_list_of_dict(samples)
|
|
|
|
|
|
class PadSequenceDataLoader(DataLoader):
|
|
|
|
def __init__(self, dataset, batch_size=32, shuffle=False, sampler=None,
|
|
batch_sampler=None, num_workers=0, collate_fn=None,
|
|
pin_memory=False, drop_last=False, timeout=0,
|
|
worker_init_fn=None, multiprocessing_context=None,
|
|
pad: dict = None, vocabs: VocabDict = None, device=None, **kwargs):
|
|
""" A dataloader commonly used for NLP tasks. It offers the following convenience.
|
|
|
|
- Bachify each field of samples into a :class:`~torch.Tensor` if the field name satisfies the following criterion.
|
|
- Name ends with _id, _ids, _count, _offset, _span, mask
|
|
- Name is in `pad` dict.
|
|
|
|
- Pad each field according to field name, the vocabs and pad dict.
|
|
- Move :class:`~torch.Tensor` onto device.
|
|
|
|
Args:
|
|
dataset: A :class:`~torch.utils.data.Dataset` to be bachified.
|
|
batch_size: Max size of each batch.
|
|
shuffle: ``True`` to shuffle batches.
|
|
sampler: A :class:`~torch.utils.data.Sampler` to sample samples from data.
|
|
batch_sampler: A :class:`~torch.utils.data.Sampler` to sample batches form all batches.
|
|
num_workers: Number of workers for multi-thread loading. Note that multi-thread loading aren't always
|
|
faster.
|
|
collate_fn: A function to perform batchifying. It must be set to ``None`` in order to make use of the
|
|
features this class offers.
|
|
pin_memory: If samples are loaded in the Dataset on CPU and would like to be pushed to
|
|
the GPU, enabling pin_memory can speed up the transfer. It's not useful since most data field are
|
|
not in Tensor type.
|
|
drop_last: Drop the last batch since it could be half-empty.
|
|
timeout: For multi-worker loading, set a timeout to wait for a worker.
|
|
worker_init_fn: Init function for multi-worker.
|
|
multiprocessing_context: Context for multiprocessing.
|
|
pad: A dict holding field names and their padding values.
|
|
vocabs: A dict of vocabs so padding value can be fetched from it.
|
|
device: The device tensors will be moved onto.
|
|
**kwargs: Other arguments will be passed to :meth:`torch.utils.data.Dataset.__init__`
|
|
"""
|
|
if device == -1:
|
|
device = None
|
|
if collate_fn is None:
|
|
collate_fn = self.collate_fn
|
|
if num_workers is None:
|
|
if isdebugging():
|
|
num_workers = 0
|
|
else:
|
|
num_workers = 2
|
|
if batch_sampler is None:
|
|
assert batch_size, 'batch_size has to be specified when batch_sampler is None'
|
|
else:
|
|
batch_size = 1
|
|
shuffle = None
|
|
drop_last = None
|
|
# noinspection PyArgumentList
|
|
super(PadSequenceDataLoader, self).__init__(dataset=dataset, batch_size=batch_size, shuffle=shuffle,
|
|
sampler=sampler,
|
|
batch_sampler=batch_sampler, num_workers=num_workers,
|
|
collate_fn=collate_fn,
|
|
pin_memory=pin_memory, drop_last=drop_last, timeout=timeout,
|
|
worker_init_fn=worker_init_fn,
|
|
multiprocessing_context=multiprocessing_context, **kwargs)
|
|
self.vocabs = vocabs
|
|
if isinstance(dataset, TransformableDataset) and dataset.transform:
|
|
transform = dataset.transform
|
|
if not isinstance(transform, TransformList):
|
|
transform = []
|
|
for each in transform:
|
|
if isinstance(each, EmbeddingNamedTransform):
|
|
if pad is None:
|
|
pad = {}
|
|
if each.dst not in pad:
|
|
pad[each.dst] = 0
|
|
self.pad = pad
|
|
self.device = device
|
|
|
|
def __iter__(self):
|
|
for raw_batch in super(PadSequenceDataLoader, self).__iter__():
|
|
yield self.tensorize(raw_batch, vocabs=self.vocabs, pad_dict=self.pad, device=self.device)
|
|
|
|
@staticmethod
|
|
def tensorize(raw_batch: Dict[str, Any], vocabs: VocabDict, pad_dict: Dict[str, int] = None, device=None):
|
|
for field, data in raw_batch.items():
|
|
if isinstance(data, torch.Tensor):
|
|
continue
|
|
vocab_key = field[:-len('_id')] if field.endswith('_id') else None
|
|
vocab: Vocab = vocabs.get(vocab_key, None) if vocabs and vocab_key else None
|
|
if vocab:
|
|
pad = vocab.safe_pad_token_idx
|
|
dtype = torch.long
|
|
elif pad_dict is not None and pad_dict.get(field, None) is not None:
|
|
pad = pad_dict[field]
|
|
dtype = dtype_of(pad)
|
|
elif field.endswith('_offset') or field.endswith('_id') or field.endswith(
|
|
'_count') or field.endswith('_ids') or field.endswith('_score') or field.endswith(
|
|
'_length') or field.endswith('_span'):
|
|
# guess some common fields to pad
|
|
pad = 0
|
|
dtype = torch.long
|
|
elif field.endswith('_mask'):
|
|
pad = False
|
|
dtype = torch.bool
|
|
else:
|
|
# no need to pad
|
|
continue
|
|
data = PadSequenceDataLoader.pad_data(data, pad, dtype)
|
|
raw_batch[field] = data
|
|
if device is not None:
|
|
for field, data in raw_batch.items():
|
|
if isinstance(data, torch.Tensor):
|
|
data = data.to(device)
|
|
raw_batch[field] = data
|
|
return raw_batch
|
|
|
|
@staticmethod
|
|
def pad_data(data: Union[torch.Tensor, Iterable], pad, dtype=None, device=None):
|
|
"""Perform the actual padding for a given data.
|
|
|
|
Args:
|
|
data: Data to be padded.
|
|
pad: Padding value.
|
|
dtype: Data type.
|
|
device: Device to be moved onto.
|
|
|
|
Returns:
|
|
torch.Tensor: A ``torch.Tensor``.
|
|
"""
|
|
if isinstance(data[0], torch.Tensor):
|
|
data = pad_sequence(data, True, pad)
|
|
elif isinstance(data[0], Iterable):
|
|
inner_is_iterable = False
|
|
for each in data:
|
|
if len(each):
|
|
if isinstance(each[0], Iterable):
|
|
inner_is_iterable = True
|
|
if len(each[0]):
|
|
if not dtype:
|
|
dtype = dtype_of(each[0][0])
|
|
else:
|
|
inner_is_iterable = False
|
|
if not dtype:
|
|
dtype = dtype_of(each[0])
|
|
break
|
|
if inner_is_iterable:
|
|
max_seq_len = len(max(data, key=len))
|
|
max_word_len = len(max([chars for words in data for chars in words], key=len))
|
|
ids = torch.zeros(len(data), max_seq_len, max_word_len, dtype=dtype, device=device)
|
|
for i, words in enumerate(data):
|
|
for j, chars in enumerate(words):
|
|
ids[i][j][:len(chars)] = torch.tensor(chars, dtype=dtype, device=device)
|
|
data = ids
|
|
else:
|
|
data = pad_sequence([torch.tensor(x, dtype=dtype, device=device) for x in data], True, pad)
|
|
elif isinstance(data, list):
|
|
data = torch.tensor(data, dtype=dtype, device=device)
|
|
return data
|
|
|
|
def collate_fn(self, samples):
|
|
return merge_list_of_dict(samples)
|
|
|
|
|
|
class CachedDataLoader(object):
|
|
def __init__(self, dataloader: torch.utils.data.DataLoader, filename=None):
|
|
if not filename:
|
|
filename = tempfile.NamedTemporaryFile(prefix='hanlp-cache-', delete=False).name
|
|
self.filename = filename
|
|
self.size = len(dataloader)
|
|
self._build_cache(dataloader)
|
|
|
|
def _build_cache(self, dataset, verbose=HANLP_VERBOSE):
|
|
timer = CountdownTimer(self.size)
|
|
with open(self.filename, "wb") as f:
|
|
for i, batch in enumerate(dataset):
|
|
torch.save(batch, f, _use_new_zipfile_serialization=False)
|
|
if verbose:
|
|
timer.log(f'Caching {self.filename} [blink][yellow]...[/yellow][/blink]')
|
|
|
|
def close(self):
|
|
if os.path.isfile(self.filename):
|
|
os.remove(self.filename)
|
|
|
|
def __iter__(self):
|
|
with open(self.filename, "rb") as f:
|
|
for i in range(self.size):
|
|
batch = torch.load(f)
|
|
yield batch
|
|
|
|
def __len__(self):
|
|
return self.size
|
|
|
|
|
|
def _prefetch_generator(dataloader, queue, batchify=None):
|
|
while True:
|
|
for batch in dataloader:
|
|
if batchify:
|
|
batch = batchify(batch)
|
|
queue.put(batch)
|
|
|
|
|
|
class PrefetchDataLoader(DataLoader):
|
|
def __init__(self, dataloader: torch.utils.data.DataLoader, prefetch: int = 10, batchify: Callable = None) -> None:
|
|
""" A dataloader wrapper which speeds up bachifying using multi-processing. It works best for dataloaders
|
|
of which the bachify takes very long time. But it introduces extra GPU memory consumption since prefetched
|
|
batches are stored in a ``Queue`` on GPU.
|
|
|
|
.. Caution::
|
|
|
|
PrefetchDataLoader only works in spawn mode with the following initialization code:
|
|
|
|
Examples::
|
|
|
|
if __name__ == '__main__':
|
|
import torch
|
|
|
|
torch.multiprocessing.set_start_method('spawn')
|
|
|
|
And these 2 lines **MUST** be put into ``if __name__ == '__main__':`` block.
|
|
|
|
Args:
|
|
dataloader: A :class:`~torch.utils.data.DatasetLoader` to be prefetched.
|
|
prefetch: Number of batches to prefetch.
|
|
batchify: A bachify function called on each batch of samples. In which case, the inner dataloader shall
|
|
return samples without really bachify them.
|
|
"""
|
|
super().__init__(dataset=dataloader)
|
|
self._batchify = batchify
|
|
self.prefetch = None if isdebugging() else prefetch
|
|
if self.prefetch:
|
|
self._fire_process(dataloader, prefetch)
|
|
|
|
def _fire_process(self, dataloader, prefetch):
|
|
self.queue = mp.Queue(prefetch)
|
|
self.process = mp.Process(target=_prefetch_generator, args=(dataloader, self.queue, self._batchify))
|
|
self.process.start()
|
|
|
|
def __iter__(self):
|
|
if not self.prefetch:
|
|
for batch in self.dataset:
|
|
if self._batchify:
|
|
batch = self._batchify(batch)
|
|
yield batch
|
|
else:
|
|
size = len(self)
|
|
while size:
|
|
batch = self.queue.get()
|
|
yield batch
|
|
size -= 1
|
|
|
|
def close(self):
|
|
"""Close this dataloader and terminates internal processes and queue. It's recommended to call this method to
|
|
ensure a program can gracefully shutdown.
|
|
"""
|
|
if self.prefetch:
|
|
self.queue.close()
|
|
self.process.terminate()
|
|
|
|
@property
|
|
def batchify(self):
|
|
return self._batchify
|
|
|
|
@batchify.setter
|
|
def batchify(self, batchify):
|
|
self._batchify = batchify
|
|
if not self.prefetch:
|
|
prefetch = vars(self.queue).get('maxsize', 10)
|
|
self.close()
|
|
self._fire_process(self.dataset, prefetch)
|
|
|
|
|
|
class BucketSampler(Sampler):
|
|
# noinspection PyMissingConstructor
|
|
def __init__(self, buckets: Dict[float, List[int]], batch_max_tokens, batch_size=None, shuffle=False):
|
|
"""A bucketing based sampler which groups samples into buckets then creates batches from each bucket.
|
|
|
|
Args:
|
|
buckets: A dict of which keys are some statistical numbers of each bucket, and values are the indices of
|
|
samples in each bucket.
|
|
batch_max_tokens: Maximum tokens per batch.
|
|
batch_size: Maximum samples per batch.
|
|
shuffle: ``True`` to shuffle batches and samples in a batch.
|
|
"""
|
|
self.shuffle = shuffle
|
|
self.sizes, self.buckets = zip(*[
|
|
(size, bucket) for size, bucket in buckets.items()
|
|
])
|
|
# the number of chunks in each bucket, which is clipped by
|
|
# range [1, len(bucket)]
|
|
if batch_size:
|
|
self.chunks = [
|
|
max(batch_size, min(len(bucket), max(round(size * len(bucket) / batch_max_tokens), 1)))
|
|
for size, bucket in zip(self.sizes, self.buckets)
|
|
]
|
|
else:
|
|
self.chunks = [
|
|
min(len(bucket), max(round(size * len(bucket) / batch_max_tokens), 1))
|
|
for size, bucket in zip(self.sizes, self.buckets)
|
|
]
|
|
|
|
def __iter__(self):
|
|
# if shuffle, shuffle both the buckets and samples in each bucket
|
|
range_fn = torch.randperm if self.shuffle else torch.arange
|
|
for i in range_fn(len(self.buckets)).tolist():
|
|
split_sizes = [(len(self.buckets[i]) - j - 1) // self.chunks[i] + 1 for j in range(self.chunks[i])]
|
|
# DON'T use `torch.chunk` which may return wrong number of chunks
|
|
for batch in range_fn(len(self.buckets[i])).split(split_sizes):
|
|
yield [self.buckets[i][j] for j in batch.tolist()]
|
|
|
|
def __len__(self):
|
|
return sum(self.chunks)
|
|
|
|
|
|
class KMeansSampler(BucketSampler):
|
|
def __init__(self, lengths, batch_max_tokens, batch_size=None, shuffle=False, n_buckets=1):
|
|
"""A bucket sampler which groups samples using KMeans on their lengths.
|
|
|
|
Args:
|
|
lengths: Lengths of each sample, usually measured by number of tokens.
|
|
batch_max_tokens: Maximum tokens per batch.
|
|
batch_size: Maximum samples per batch.
|
|
shuffle: ``True`` to shuffle batches. Samples in the same batch won't be shuffled since the ordered sequence
|
|
is helpful to speed up RNNs.
|
|
n_buckets: Number of buckets. Clusters in terms of KMeans.
|
|
"""
|
|
if n_buckets > len(lengths):
|
|
n_buckets = 1
|
|
self.n_buckets = n_buckets
|
|
self.lengths = lengths
|
|
buckets = dict(zip(*kmeans(self.lengths, n_buckets)))
|
|
super().__init__(buckets, batch_max_tokens, batch_size, shuffle)
|
|
|
|
|
|
class SortingSampler(Sampler):
|
|
# noinspection PyMissingConstructor
|
|
def __init__(self, lengths: List[int], batch_size=None, batch_max_tokens=None, use_effective_tokens=False,
|
|
shuffle=False) -> None:
|
|
"""A sampler which sorts samples according to their lengths. It takes a continuous chunk of sorted samples to
|
|
make a batch. The effective batch size is determined by ``batch_size``, ``batch_max_tokens`` and
|
|
``use_effective_tokens``.
|
|
|
|
Args:
|
|
lengths: Lengths of each sample, usually measured by number of tokens.
|
|
batch_max_tokens: Maximum tokens per batch.
|
|
use_effective_tokens: Whether to calculate the effective number of tokens after padding when applying the
|
|
``batch_max_tokens``.
|
|
batch_size: Maximum samples per batch.
|
|
shuffle: ``True`` to shuffle batches and samples in a batch.
|
|
"""
|
|
# assert any([batch_size, batch_max_tokens]), 'At least one of batch_size and batch_max_tokens is required'
|
|
self.shuffle = shuffle
|
|
self.batch_size = batch_size
|
|
# self.batch_max_tokens = batch_max_tokens
|
|
self.batch_indices = []
|
|
num_tokens = 0
|
|
mini_batch = []
|
|
for i in torch.argsort(torch.tensor(lengths), descending=True).tolist():
|
|
# if batch_max_tokens:
|
|
effective_tokens = lengths[i] if (not mini_batch or not use_effective_tokens) else lengths[mini_batch[0]]
|
|
if (batch_max_tokens is None or num_tokens + effective_tokens <= batch_max_tokens) and (
|
|
batch_size is None or len(mini_batch) < batch_size):
|
|
mini_batch.append(i)
|
|
num_tokens += effective_tokens
|
|
else:
|
|
if not mini_batch: # this sequence is longer than batch_max_tokens
|
|
mini_batch.append(i)
|
|
self.batch_indices.append(mini_batch)
|
|
mini_batch = []
|
|
num_tokens = 0
|
|
else:
|
|
self.batch_indices.append(mini_batch)
|
|
mini_batch = [i]
|
|
num_tokens = effective_tokens
|
|
if mini_batch:
|
|
self.batch_indices.append(mini_batch)
|
|
# print(len(max(self.batch_indices, key=len)))
|
|
|
|
def __iter__(self):
|
|
if self.shuffle:
|
|
random.shuffle(self.batch_indices)
|
|
for batch in self.batch_indices:
|
|
yield batch
|
|
|
|
def __len__(self) -> int:
|
|
return len(self.batch_indices)
|
|
|
|
|
|
class SamplerBuilder(AutoConfigurable, ABC):
|
|
@abstractmethod
|
|
def build(self, lengths: List[int], shuffle=False, gradient_accumulation=1, **kwargs) -> Sampler:
|
|
"""Build a ``Sampler`` given statistics of samples and other arguments.
|
|
|
|
Args:
|
|
lengths: The lengths of samples.
|
|
shuffle: ``True`` to shuffle batches. Note samples in each mini-batch are not necessarily shuffled.
|
|
gradient_accumulation: Number of mini-batches per update step.
|
|
**kwargs: Other arguments to be passed to the constructor of the sampler.
|
|
"""
|
|
pass
|
|
|
|
def __call__(self, lengths: List[int], shuffle=False, **kwargs) -> Sampler:
|
|
return self.build(lengths, shuffle, **kwargs)
|
|
|
|
def scale(self, gradient_accumulation):
|
|
r"""Scale down the ``batch_size`` and ``batch_max_tokens`` to :math:`\frac{1}{\text{gradient_accumulation}}`
|
|
of them respectively.
|
|
|
|
Args:
|
|
gradient_accumulation: Number of mini-batches per update step.
|
|
|
|
Returns:
|
|
tuple(int,int): batch_size, batch_max_tokens
|
|
"""
|
|
batch_size = self.batch_size
|
|
batch_max_tokens = self.batch_max_tokens
|
|
if gradient_accumulation:
|
|
if batch_size:
|
|
batch_size //= gradient_accumulation
|
|
if batch_max_tokens:
|
|
batch_max_tokens //= gradient_accumulation
|
|
return batch_size, batch_max_tokens
|
|
|
|
|
|
class SortingSamplerBuilder(SortingSampler, SamplerBuilder):
|
|
# noinspection PyMissingConstructor
|
|
def __init__(self, batch_size=None, batch_max_tokens=None, use_effective_tokens=False) -> None:
|
|
"""Builds a :class:`~hanlp.common.dataset.SortingSampler`.
|
|
|
|
Args:
|
|
batch_max_tokens: Maximum tokens per batch.
|
|
use_effective_tokens: Whether to calculate effective number of tokens when applying the `batch_max_tokens`.
|
|
batch_size: Maximum samples per batch.
|
|
"""
|
|
self.use_effective_tokens = use_effective_tokens
|
|
self.batch_max_tokens = batch_max_tokens
|
|
self.batch_size = batch_size
|
|
|
|
def build(self, lengths: List[int], shuffle=False, gradient_accumulation=1, **kwargs) -> Sampler:
|
|
batch_size, batch_max_tokens = self.scale(gradient_accumulation)
|
|
return SortingSampler(lengths, batch_size, batch_max_tokens, shuffle)
|
|
|
|
def __len__(self) -> int:
|
|
return 1
|
|
|
|
|
|
class KMeansSamplerBuilder(KMeansSampler, SamplerBuilder):
|
|
# noinspection PyMissingConstructor
|
|
def __init__(self, batch_max_tokens, batch_size=None, n_buckets=1):
|
|
"""Builds a :class:`~hanlp.common.dataset.KMeansSampler`.
|
|
|
|
Args:
|
|
batch_max_tokens: Maximum tokens per batch.
|
|
batch_size: Maximum samples per batch.
|
|
n_buckets: Number of buckets. Clusters in terms of KMeans.
|
|
"""
|
|
self.n_buckets = n_buckets
|
|
self.batch_size = batch_size
|
|
self.batch_max_tokens = batch_max_tokens
|
|
|
|
def build(self, lengths: List[int], shuffle=False, gradient_accumulation=1, **kwargs) -> Sampler:
|
|
batch_size, batch_max_tokens = self.scale(gradient_accumulation)
|
|
return KMeansSampler(lengths, batch_max_tokens, batch_size, shuffle, self.n_buckets)
|
|
|
|
def __len__(self) -> int:
|
|
return 1
|
|
|
|
|
|
class TableDataset(TransformableDataset):
|
|
def __init__(self,
|
|
data: Union[str, List],
|
|
transform: Union[Callable, List] = None,
|
|
cache=None,
|
|
delimiter='auto',
|
|
strip=True,
|
|
headers=None) -> None:
|
|
self.headers = headers
|
|
self.strip = strip
|
|
self.delimiter = delimiter
|
|
super().__init__(data, transform, cache)
|
|
|
|
def load_file(self, filepath: str):
|
|
for idx, cells in enumerate(read_cells(filepath, strip=self.strip, delimiter=self.delimiter)):
|
|
if not idx and not self.headers:
|
|
self.headers = cells
|
|
if any(len(h) > 32 for h in self.headers):
|
|
warnings.warn('As you did not pass in `headers` to `TableDataset`, the first line is regarded as '
|
|
'headers. However, the length for some headers are too long (>32), which might be '
|
|
'wrong. To make sure, pass `headers=...` explicitly.')
|
|
else:
|
|
yield dict(zip(self.headers, cells))
|