File size: 3,636 Bytes
9627ce0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 | from torch.utils.data import DataLoader
import torch
from torch.utils.data import DataLoader, DistributedSampler
from src.interface.data_interface import DInterface_base
from src.data.protein_dataset import ProteinDataset
from src.model.pretrain_model_interface import PretrainModelInterface
class DInterface(DInterface_base):
def __init__(self, **kwargs):
super().__init__(**kwargs)
self.save_hyperparameters()
def setup(self, stage=None):
pass
def data_setup(self, target="all"):
pretrain_model_interface = None
if self.finetune_type == "adapter":
pretrain_model_interface = PretrainModelInterface(self.hparams.pretrain_model_name, batch_size=self.hparams.pretrain_batch_size, max_length=self.hparams.seq_len, sequence_only=self.hparams.sequence_only, task_type=self.hparams.task_type)
if target == "all":
self.train_set = ProteinDataset(self.hparams.train_data_path, self.hparams.pretrain_model_name, self.hparams.seq_len, pretrain_model_interface=pretrain_model_interface, task_name=self.task_name, task_type=self.hparams.task_type, num_classes=self.hparams.num_classes)
self.val_set = ProteinDataset(self.hparams.val_data_path, self.hparams.pretrain_model_name, self.hparams.seq_len, pretrain_model_interface=pretrain_model_interface, task_name=self.task_name, task_type=self.hparams.task_type, num_classes=self.hparams.num_classes)
self.test_set = ProteinDataset(self.hparams.test_data_path, self.hparams.pretrain_model_name, self.hparams.seq_len, pretrain_model_interface=pretrain_model_interface, task_name=self.task_name, task_type=self.hparams.task_type, num_classes=self.hparams.num_classes)
elif target == "test":
self.test_set = ProteinDataset(self.hparams.test_data_path, self.hparams.pretrain_model_name, self.hparams.seq_len, pretrain_model_interface=pretrain_model_interface, task_name=self.task_name, task_type=self.hparams.task_type, num_classes=self.hparams.num_classes)
def train_dataloader(self):
return DataLoader(self.train_set, batch_size=self.hparams.batch_size, shuffle=True, num_workers=self.hparams.num_workers, pin_memory=True, collate_fn=self.data_process_fn)
def val_dataloader(self):
return DataLoader(self.val_set, batch_size=self.hparams.batch_size, shuffle=False, num_workers=self.hparams.num_workers, pin_memory=True, collate_fn=self.data_process_fn)
def test_dataloader(self):
return DataLoader(self.test_set, batch_size=self.hparams.batch_size, shuffle=False, num_workers=self.hparams.num_workers, pin_memory=True, collate_fn=self.data_process_fn)
def data_process_fn(self, data_list):
if self.hparams.finetune_type == 'adapter':
name_list = []
mask_list = []
label_list = []
embedding_list = []
smiles = []
for data in data_list:
name_list.append(data['name'])
mask_list.append(data['attention_mask'])
label_list.append(data['label'])
embedding_list.append(data['embedding'])
if data.get('smiles') is not None:
smiles.append(data['smiles'])
return {'name': name_list,
'attention_mask': torch.stack(mask_list, dim=0),
'label': torch.stack(label_list, dim=0),
'embedding': torch.stack(embedding_list, dim=0),
'smiles': torch.stack(smiles, dim=0) if len(smiles) > 0 else None,
}
else:
return data_list
|