| import pytorch_lightning as pl | |
| from torch.utils.data import DataLoader | |
| import torch | |
| from .attribute_selector import AttributeSelector | |
| from .similarity_vector_dataset import SimilarityVectorDataset | |
| from typing import List | |
| class MARCDataModule(pl.LightningDataModule): | |
| def __init__( | |
| self, | |
| train_processed_path: str, | |
| val_processed_path: str, | |
| test_processed_path: str, | |
| attrs: List[str], | |
| batch_size: int, | |
| ): | |
| super().__init__() | |
| self.train_processed_path = train_processed_path | |
| self.val_processed_path = val_processed_path | |
| self.test_processed_path = test_processed_path | |
| self.batch_size = batch_size | |
| self.transform = torch.nn.Sequential(AttributeSelector(attrs)) | |
| self.train_set = None | |
| self.val_set = None | |
| self.test_set = None | |
| def setup(self, stage=None): | |
| self.train_set = SimilarityVectorDataset( | |
| self.train_processed_path, transform=self.transform | |
| ) | |
| self.val_set = SimilarityVectorDataset( | |
| self.val_processed_path, transform=self.transform | |
| ) | |
| self.test_set = SimilarityVectorDataset( | |
| self.test_processed_path, transform=self.transform | |
| ) | |
| def train_dataloader(self): | |
| return DataLoader( | |
| self.train_set, batch_size=self.batch_size, num_workers=0, shuffle=True | |
| ) | |
| def val_dataloader(self): | |
| return DataLoader(self.val_set, batch_size=self.batch_size, num_workers=0) | |
| def test_dataloader(self): | |
| return DataLoader(self.test_set, batch_size=self.batch_size, num_workers=0) | |