File size: 1,243 Bytes
ed920f9 |
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 |
import pandas as pd
from torch.utils.data import Dataset, DataLoader
from utils.esm_utils import get_latents, load_esm2_model
class ProteinDataset(Dataset):
def __init__(self, csv_file, tokenizer, model):
self.data = pd.read_csv(csv_file)
self.tokenizer = tokenizer
self.model = model
def __len__(self):
return len(self.data)
def __getitem__(self, idx):
sequence = self.data.iloc[idx]['sequence']
latents = get_latents(self.model, self.tokenizer, sequence)
return latents
def get_dataloaders(config):
tokenizer, model = load_esm2_model(config.model_name)
train_dataset = ProteinDataset(config.data_path + "train.csv", tokenizer, model)
val_dataset = ProteinDataset(config.data_path + "val.csv", tokenizer, model)
test_dataset = ProteinDataset(config.data_path + "test.csv", tokenizer, model)
train_loader = DataLoader(train_dataset, batch_size=config.training["batch_size"], shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=config.training["batch_size"], shuffle=False)
test_loader = DataLoader(test_dataset, batch_size=config.training["batch_size"], shuffle=False)
return train_loader, val_loader, test_loader
|