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