wehub-resource-sync
59 lines
1.7 KiB
Python
59 lines
1.7 KiB
Python
import os
|
|
import pickle
|
|
|
|
import torch
|
|
from torch.utils.data import Dataset
|
|
|
|
from utils.hparams import hparams
|
|
from utils.indexed_datasets import IndexedDataset
|
|
|
|
|
|
class BaseDataset(Dataset):
|
|
"""
|
|
Base class for datasets.
|
|
1. *sizes*:
|
|
clipped length if "max_frames" is set;
|
|
2. *num_frames*:
|
|
unclipped length.
|
|
|
|
Subclasses should define:
|
|
1. *collate*:
|
|
take the longest data, pad other data to the same length;
|
|
2. *__getitem__*:
|
|
the index function.
|
|
"""
|
|
|
|
def __init__(self, prefix, size_key='lengths', preload=False):
|
|
super().__init__()
|
|
self.prefix = prefix
|
|
self.data_dir = hparams['binary_data_dir']
|
|
with open(os.path.join(self.data_dir, f'{self.prefix}.meta'), 'rb') as f:
|
|
self.metadata = pickle.load(f)
|
|
self.sizes = self.metadata[size_key]
|
|
self._indexed_ds = IndexedDataset(self.data_dir, self.prefix)
|
|
if preload:
|
|
self.indexed_ds = [self._indexed_ds[i] for i in range(len(self._indexed_ds))]
|
|
del self._indexed_ds
|
|
else:
|
|
self.indexed_ds = self._indexed_ds
|
|
|
|
def __getitem__(self, index):
|
|
return {'_idx': index, **self.indexed_ds[index]}
|
|
|
|
def __len__(self):
|
|
return len(self.sizes)
|
|
|
|
def num_frames(self, index):
|
|
return self.sizes[index]
|
|
|
|
def size(self, index):
|
|
"""Return an example's size as a float or tuple. This value is used when
|
|
filtering a dataset with ``--max-positions``."""
|
|
return self.sizes[index]
|
|
|
|
def collater(self, samples):
|
|
return {
|
|
'size': len(samples),
|
|
'indices': torch.LongTensor([s['_idx'] for s in samples])
|
|
}
|