goathead777/Zero_Shot_Inference

1# modified from https://github.com/feng-yufei/shared_debugging_code/blob/main/data_module.py2from pytorch_lightning import LightningDataModule3from AR.data.bucket_sampler import DistributedBucketSampler4from AR.data.dataset import Text2SemanticDataset5from torch.utils.data import DataLoader6 7 8class Text2SemanticDataModule(LightningDataModule):9    def __init__(10        self,11        config,12        train_semantic_path,13        train_phoneme_path,14        dev_semantic_path=None,15        dev_phoneme_path=None,16    ):17        super().__init__()18        self.config = config19        self.train_semantic_path = train_semantic_path20        self.train_phoneme_path = train_phoneme_path21        self.dev_semantic_path = dev_semantic_path22        self.dev_phoneme_path = dev_phoneme_path23        self.num_workers = self.config["data"]["num_workers"]24 25    def prepare_data(self):26        pass27 28    def setup(self, stage=None, output_logs=False):29        self._train_dataset = Text2SemanticDataset(30            phoneme_path=self.train_phoneme_path,31            semantic_path=self.train_semantic_path,32            max_sec=self.config["data"]["max_sec"],33            pad_val=self.config["data"]["pad_val"],34        )35        self._dev_dataset = self._train_dataset36        # self._dev_dataset = Text2SemanticDataset(37        #     phoneme_path=self.dev_phoneme_path,38        #     semantic_path=self.dev_semantic_path,39        #     max_sample=self.config['data']['max_eval_sample'],40        #     max_sec=self.config['data']['max_sec'],41        #     pad_val=self.config['data']['pad_val'])42 43    def train_dataloader(self):44        batch_size = self.config["train"]["batch_size"]45        sampler = DistributedBucketSampler(self._train_dataset, batch_size=batch_size)46        return DataLoader(47            self._train_dataset,48            batch_size=batch_size,49            sampler=sampler,50            collate_fn=self._train_dataset.collate,51            num_workers=self.num_workers,52            persistent_workers=True,53            prefetch_factor=16,54        )55 56    def val_dataloader(self):57        return DataLoader(58            self._dev_dataset,59            batch_size=1,60            shuffle=False,61            collate_fn=self._train_dataset.collate,62            num_workers=max(self.num_workers, 12),63            persistent_workers=True,64            prefetch_factor=16,65        )66 67    # 这个会使用到嘛？68    def test_dataloader(self):69        return DataLoader(70            self._dev_dataset,71            batch_size=1,72            shuffle=False,73            collate_fn=self._train_dataset.collate,74        )75