| 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 |
