File size: 1,979 Bytes
002bd9b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 |
# Deprecated, Not compatible with HuggingFace Datasets
from torch.utils.data import Dataset, IterableDataset
from datasets import HFDataset, HFIterableDataset
from math import ceil
def prepare_regional_chunkified_dataset(dataset: Dataset, regional_chunk_size: int) -> Dataset:
if not isinstance(dataset, (HFDataset, Dataset)):
raise ValueError(f"Currently dataset must be a map-style Dataset, got {type(dataset)}")
class RegionalChunkifiedDataset(Dataset):
def __init__(self, dataset: Dataset, regional_chunk_size: int):
self.dataset = dataset
self.regional_chunk_size = regional_chunk_size
self.get_num_regions_and_num_samples_after_regional_chunkify()
def get_num_regions_and_num_samples_after_regional_chunkify(self):
self.num_regions = 0
self.num_samples_after_regional_chunkify = 0
self.chunkified_sample_ids_to_sample_idx = {}
for sample_idx, example in enumerate(self.dataset):
for chunkified_sample_idx in range(ceil(len(example["regions"]) / self.regional_chunk_size)):
self.chunkified_sample_ids_to_sample_idx[
self.num_samples_after_regional_chunkify + chunkified_sample_idx
] = (sample_idx, chunkified_sample_idx)
self.num_regions += len(example["regions"])
self.num_samples_after_regional_chunkify += ceil(len(example["regions"]) / self.regional_chunk_size)
def __len__(self):
return self.num_samples_after_regional_chunkify
def __getitem__(self, idx):
sample_idx, chunkified_sample_idx = self.chunkified_sample_ids_to_sample_idx[idx]
example = self.dataset[sample_idx]
regions = example.pop("regions")
return {
**example,
"regions": regions[
chunkified_sample_idx
* self.regional_chunk_size : (chunkified_sample_idx + 1)
* self.regional_chunk_size
],
}
|