FunASR/funasr/datasets/large_datasets/datapipes/map.py
2022-11-26 21:56:51 +08:00

23 lines
453 B
Python

from torch.utils.data import IterableDataset
def default_fn(data):
return data
class MapperIterDataPipe(IterableDataset):
def __init__(self,
datapipe,
fn=default_fn):
self.datapipe = datapipe
self.fn = fn
def set_epoch(self, epoch):
self.epoch = epoch
def __iter__(self):
assert callable(self.fn)
for data in self.datapipe:
yield self.fn(data)