Upload 2 files

utils/data_loader.py

@@ -1,6 +1,9 @@
 import pandas as pd
+import torch
 from torch.utils.data import Dataset, DataLoader
-from utils.esm_utils import get_latents, load_esm2_model
+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):
@@ -12,19 +15,29 @@ class ProteinDataset(Dataset):
         return len(self.data)
 
     def __getitem__(self, idx):
-        sequence = self.data.iloc[idx]['sequence']
+        sequence = self.data.iloc[idx]['Sequence']
         latents = get_latents(self.model, self.tokenizer, sequence)
-        return latents
+
+        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, model = load_esm2_model(config.model_name)
+    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_dataset = ProteinDataset(config.Loader.DATA_PATH + "/train.csv", tokenizer, model)
+    val_dataset = ProteinDataset(config.Loader.DATA_PATH + "/val.csv", tokenizer, model)
+    test_dataset = ProteinDataset(config.Loader.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)
+    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
+    return train_loader, val_loader, test_loader

utils/esm_utils.py

@@ -11,3 +11,4 @@ def get_latents(model, tokenizer, sequence):
     with torch.no_grad():
         outputs = model(**inputs)
     return outputs.last_hidden_state.squeeze(0)
+