File size: 1,975 Bytes
ed920f9 42149de ed920f9 42149de ed920f9 d6c63a1 ed920f9 42149de ed920f9 42149de ed920f9 d6c63a1 ed920f9 d6c63a1 ed920f9 42149de ed920f9 42149de |
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 |
import pandas as pd
import torch
from torch.utils.data import Dataset, DataLoader
from torch.nn.utils.rnn import pad_sequence
from esm_utils import get_latents, load_esm2_model
import config
class ProteinDataset(Dataset):
def __init__(self, csv_file, tokenizer, model):
self.data = pd.read_csv(csv_file).head(4)
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)
attention_mask = torch.ones_like(latents)
attention_mask = torch.mean(attention_mask, dim=-1)
return latents, attention_mask
def collate_fn(batch):
latents, attention_mask = zip(*batch)
latents_padded = pad_sequence([torch.tensor(latent) for latent in latents], batch_first=True, padding_value=0)
attention_mask_padded = pad_sequence([torch.tensor(mask) for mask in attention_mask], batch_first=True, padding_value=0)
return latents_padded, attention_mask_padded
def get_dataloaders(config):
tokenizer, masked_model, embedding_model = load_esm2_model(config.MODEL_NAME)
train_dataset = ProteinDataset(config.Loader.DATA_PATH + "/train.csv", tokenizer, embedding_model)
val_dataset = ProteinDataset(config.Loader.DATA_PATH + "/val.csv", tokenizer, embedding_model)
test_dataset = ProteinDataset(config.Loader.DATA_PATH + "/test.csv", tokenizer, embedding_model)
train_loader = DataLoader(train_dataset, batch_size=config.Loader.BATCH_SIZE, num_workers=0, shuffle=True, collate_fn=collate_fn)
val_loader = DataLoader(val_dataset, batch_size=config.Loader.BATCH_SIZE, num_workers=0, shuffle=False, collate_fn=collate_fn)
test_loader = DataLoader(test_dataset, batch_size=config.Loader.BATCH_SIZE, num_workers=0, shuffle=False, collate_fn=collate_fn)
return train_loader, val_loader, test_loader |