|
from typing import Optional |
|
from pytorch_lightning import LightningDataModule |
|
from torch.utils.data import DataLoader |
|
from fengshen.data.mmap_index_dataset import MMapIndexDataset |
|
|
|
|
|
class MMapDataModule(LightningDataModule): |
|
@ staticmethod |
|
def add_data_specific_args(parent_args): |
|
parser = parent_args.add_argument_group('MMAP DataModule') |
|
parser.add_argument('--num_workers', default=8, type=int) |
|
parser.add_argument('--train_batchsize', default=32, type=int) |
|
parser.add_argument('--eval_batchsize', default=32, type=int) |
|
parser.add_argument('--test_batchsize', default=32, type=int) |
|
parser.add_argument('--train_datas', default=[ |
|
'./train_datas' |
|
], type=str, nargs='+') |
|
parser.add_argument('--valid_datas', default=[ |
|
'./valid_datas' |
|
], type=str, nargs='+') |
|
parser.add_argument('--test_datas', default=[ |
|
'./test_datas'], |
|
type=str, nargs='+') |
|
parser.add_argument('--input_tensor_name', default=['input_ids'], type=str, nargs='+') |
|
return parent_args |
|
|
|
def __init__( |
|
self, |
|
collate_fn, |
|
args, |
|
**kwargs, |
|
): |
|
super().__init__() |
|
self.collate_fn = collate_fn |
|
self.train_dataset = MMapIndexDataset(args.train_datas, args.input_tensor_name) |
|
self.valid_dataset = MMapIndexDataset(args.valid_datas, args.input_tensor_name) |
|
self.test_dataset = MMapIndexDataset(args.test_datas, args.input_tensor_name) |
|
self.save_hyperparameters(args) |
|
|
|
def setup(self, stage: Optional[str] = None) -> None: |
|
return super().setup(stage) |
|
|
|
def train_dataloader(self): |
|
return DataLoader( |
|
self.train_dataset, |
|
batch_size=self.hparams.train_batchsize, |
|
shuffle=True, |
|
num_workers=self.hparams.num_workers, |
|
collate_fn=self.collate_fn, |
|
) |
|
|
|
def val_dataloader(self): |
|
return DataLoader( |
|
self.valid_dataset, |
|
batch_size=self.hparams.eval_batchsize, |
|
shuffle=True, |
|
num_workers=self.hparams.num_workers, |
|
collate_fn=self.collate_fn, |
|
) |
|
|
|
def test_dataloader(self): |
|
return DataLoader( |
|
self.test_dataset, |
|
batch_size=self.hparams.test_batchsize, |
|
shuffle=True, |
|
num_workers=self.hparams.num_workers, |
|
collate_fn=self.collate_fn, |
|
) |
|
|
