| import os |
| import pandas as pd |
| import numpy as np |
|
|
| from typing import Optional |
| from itertools import cycle, product |
|
|
| import torch |
| from torch.utils.data import Dataset, DataLoader |
|
|
| import pytorch_lightning as pl |
| from pytorch_lightning.utilities.types import EVAL_DATALOADERS |
|
|
| from sklearn.model_selection import KFold |
|
|
| from dataclasses import dataclass |
|
|
|
|
| class BiomakerDataset(Dataset): |
| def __init__(self, chem_obj, prot_obj): |
| self.pair_ids, self.pair_mask = self.make_pairset(chem_obj, prot_obj) |
|
|
| def make_pairset(self, chem_obj, prot_obj): |
| chem_ids = list(range(len(chem_obj["input_ids"]))) |
| marked_chem_obj = list(zip(chem_ids, chem_obj["input_ids"])) |
|
|
| pair_ids = list(product(*[marked_chem_obj, prot_obj["input_ids"]])) |
| pair_mask = list(product(*[chem_obj["attention_mask"], prot_obj["attention_mask"]])) |
|
|
| return pair_ids, pair_mask |
| |
| def __len__(self): |
| return len(self.pair_ids) |
|
|
| def __getitem__(self, idx): |
| ids = self.pair_ids[idx][0][0] |
| chem_obj = {"input_ids" : self.pair_ids[idx][0][1], "attention_mask": self.pair_mask[idx][0]} |
| prot_obj = {"input_ids" : self.pair_ids[idx][1], "attention_mask": self.pair_mask[idx][1]} |
|
|
| return ids, chem_obj, prot_obj |
|
|
|
|
| class BiomakerDataModule(pl.LightningDataModule): |
| def __init__(self, chem_obj:pd.DataFrame, prot_obj:pd.DataFrame, |
| config): |
| super().__init__() |
| self.num_seed = config.num_seed |
|
|
| self.chem_obj = chem_obj |
| self.prot_obj = prot_obj |
|
|
| self.batch_size = config.batch_size |
| self.num_workers = config.num_workers |
|
|
| self.df_prepared = False |
| |
| def setup(self, stage: str) -> None: |
| |
| self.predict_dataset = BiomakerDataset(self.chem_obj, self.prot_obj) |
|
|
| def predict_dataloader(self) -> EVAL_DATALOADERS: |
| return DataLoader(self.predict_dataset, batch_size=self.batch_size, num_workers=self.num_workers) |
| |
|
|
|
|
| @dataclass |
| class clearanceDataset(Dataset): |
| def __init__(self, data, labels = None, features = None): |
| self.data = data |
| self.labels = labels |
| self.features = features |
|
|
| def __len__(self): |
| return len(self.data["input_ids"]) |
|
|
| def __getitem__(self, idx): |
| data = {"input_ids": self.data["input_ids"][idx], |
| "attention_mask": self.data["attention_mask"][idx]} |
| |
| if self.features is not None and self.labels is not None: |
| labels = self.labels[idx] |
| features = self.features[idx] |
| return data, labels, features |
| elif self.features is not None: |
| features = self.features[idx] |
| return data, features |
| elif self.labels is not None: |
| labels = self.labels[idx] |
| return data, labels |
| else: |
| return data |
| |
|
|
| class clearanceDatamodule(pl.LightningDataModule): |
| def __init__(self, chem_tokenizer, config, |
| df_trainData=None, df_testData=None, df_predictData=None, |
| df_trainLabel=None, df_testLabel=None, df_predictLabel=None, |
| train_scaler=None, test_scaler=None, predict_scaler=None, |
| df_trainfeatures=None, df_testfeatures=None, df_predictfeatures=None): |
| super().__init__() |
| self.chem_tok = chem_tokenizer |
| |
| self.max_length = config.chem_max |
| self.batch_size = config.batch_size |
| self.num_workers = config.num_workers |
|
|
| self.df_trainData, self.df_testData, self.df_predictData = df_trainData, df_testData, df_predictData |
| self.df_trainLabel, self.df_testLabel, self.df_predictLabel = df_trainLabel, df_testLabel, df_predictLabel |
| self.train_scaler, self.test_scaler, self.predict_scaler = train_scaler, test_scaler, predict_scaler |
| |
| self.df_trainfeatures, self.df_testfeatures, self.df_predictfeatures = df_trainfeatures, df_testfeatures, df_predictfeatures |
| |
| self.is_load_file = False |
|
|
| def prepare_data(self): |
| if not self.is_load_file: |
| if self.df_trainData is not None: |
| train_set = np.array(self.df_trainData["SMILES"]).tolist() |
| self.train_pos = int(len(self.df_trainData) * 0.8) |
| self.train_obj = self.chem_tok(train_set[:self.train_pos], padding='max_length', |
| max_length=self.max_length, truncation=True, return_tensors="pt") |
| self.valid_obj = self.chem_tok(train_set[self.train_pos:], padding='max_length', |
| max_length=self.max_length, truncation=True, return_tensors="pt") |
| self.train_label = torch.from_numpy(np.array(self.df_trainLabel[:self.train_pos])) |
| self.valid_label = torch.from_numpy(np.array(self.df_trainLabel[self.train_pos:])) |
| self.train_features = torch.from_numpy(np.array(self.df_trainfeatures[:self.train_pos])) if self.df_trainfeatures is not None else None |
| self.valid_features = torch.from_numpy(np.array(self.df_trainfeatures[self.train_pos:])) if self.df_trainfeatures is not None else None |
| |
| if self.df_testData is not None: |
| test_set = np.array(self.df_testData["SMILES"]).tolist() |
| self.test_obj = self.chem_tok(test_set, padding='max_length', |
| max_length=self.max_length, truncation=True, return_tensors="pt") |
| self.test_label = torch.from_numpy(np.array(self.df_testLabel)) |
| self.test_features = torch.from_numpy(np.array(self.df_testfeatures)) if self.df_testfeatures is not None else None |
|
|
| if self.df_predictData is not None: |
| predict_set = np.array(self.df_predictData["SMILES"]).tolist() |
| self.predict_obj = self.chem_tok(predict_set, padding='max_length', |
| max_length=self.max_length, truncation=True, return_tensors="pt") |
| self.predict_label = torch.from_numpy(np.array(self.df_predictLabel)) if self.df_predictLabel is not None else None |
| self.predict_features = torch.from_numpy(np.array(self.df_predictfeatures)) if self.df_predictfeatures is not None else None |
|
|
| self.is_load_file = True |
|
|
| def setup(self, stage: Optional[str] = None) -> None: |
| if stage == 'fit' or stage is None: |
| if self.df_trainData is not None: |
| self.train_dataset = clearanceDataset(self.train_obj, self.train_label, self.train_features) |
| if self.df_trainData is not None: |
| self.valid_dataset = clearanceDataset(self.valid_obj, self.valid_label, self.valid_features) |
| |
| if self.df_testData is not None: |
| self.test_dataset = clearanceDataset(self.test_obj, self.test_label, self.test_features) |
|
|
| if stage == 'predict' or stage is None: |
| if self.df_predictData is not None: |
| self.pred_dataset = clearanceDataset(self.predict_obj, self.predict_label, self.predict_features) |
|
|
| def train_dataloader(self): |
| if self.df_trainData is not None: |
| return DataLoader(self.train_dataset, batch_size=self.batch_size, num_workers=self.num_workers, shuffle=True) |
| return None |
|
|
| def val_dataloader(self): |
| if self.df_trainData is not None: |
| return DataLoader(self.valid_dataset, batch_size=self.batch_size, num_workers=self.num_workers) |
| return None |
| |
| def test_dataloader(self): |
| if self.df_testData is not None: |
| return DataLoader(self.test_dataset, batch_size=self.batch_size, num_workers=self.num_workers) |
| return None |
|
|
| def predict_dataloader(self): |
| if self.df_predictData is not None: |
| return DataLoader(self.pred_dataset, batch_size=self.batch_size, num_workers=self.num_workers) |
| return None |
| |
| |
| class clearanceKfoldDatamodule(pl.LightningDataModule): |
| def __init__(self, chem_tokenizer, config, |
| df_trainData, df_testData, |
| df_trainLabel, df_testLabel, |
| train_scaler, test_scaler, |
| df_trainfeatures = None, df_testfeatures = None, kfold = 5): |
| super().__init__() |
| self.chem_tok = chem_tokenizer |
| |
| self.max_length = config.chem_max |
| self.batch_size = config.batch_size |
| self.num_workers = config.num_workers |
|
|
| self.df_trainData, self.df_testData = df_trainData, df_testData |
| self.df_trainLabel, self.df_testLabel = df_trainLabel, df_testLabel |
| self.train_scaler, self.test_scaler = train_scaler, test_scaler |
| |
| self.df_trainfeatures, self.df_testfeatures = df_trainfeatures, df_testfeatures |
| self.kfold = kfold |
|
|
| def setup(self, stage: Optional[str] = None) -> None: |
| |
| train_set, test_set = np.array(self.df_trainData["SMILES"]).tolist(), np.array(self.df_testData["SMILES"]).tolist() |
| |
| |
| train_encoding = self.chem_tok(train_set, padding='max_length', |
| max_length=self.max_length, truncation=True, return_tensors="pt") |
| test_encoding = self.chem_tok(test_set, padding='max_length', |
| max_length=self.max_length, truncation=True, return_tensors="pt") |
| |
| |
| train_label = torch.from_numpy(np.array(self.df_trainLabel)) |
| test_label = torch.from_numpy(np.array(self.df_testLabel)) |
| |
| |
| if self.df_trainfeatures is not None: |
| train_features = torch.from_numpy(np.array(self.df_trainfeatures)) |
| self.trainset = clearanceDataset(train_encoding, train_label, train_features) |
| else: |
| self.trainset = clearanceDataset(train_encoding, train_label) |
| |
| if self.df_testfeatures is not None: |
| test_features = torch.from_numpy(np.array(self.df_testfeatures)) |
| self.testset = clearanceDataset(test_encoding, test_label, test_features) |
| else: |
| self.testset = clearanceDataset(test_encoding, test_label) |
| |
| self.splits = list(KFold(n_splits=self.kfold, shuffle=True).split(self.trainset)) |
|
|
| |
| def train_dataloader(self, fold_index): |
| train_idx, val_idx = self.splits[fold_index] |
| train_dataset = torch.utils.data.Subset(self.trainset, train_idx) |
| val_dataset = torch.utils.data.Subset(self.trainset, val_idx) |
| self.train_loader = DataLoader(train_dataset, batch_size=self.batch_size, num_workers=self.num_workers, shuffle=True) |
| self.val_loader = DataLoader(val_dataset, batch_size=self.batch_size, num_workers=self.num_workers) |
| |
| return self.train_loader |
|
|
| def val_dataloader(self): |
| return self.val_loader |
| |
| def test_dataloader(self): |
| self.test_loader = DataLoader(self.testset, batch_size=self.batch_size, num_workers=self.num_workers) |
| return self.test_loader |
