Upload 2 files

utils/data_loader.py

@@ -7,7 +7,7 @@ import config
 
 class ProteinDataset(Dataset):
     def __init__(self, csv_file, tokenizer, model):
-        self.data = pd.read_csv(csv_file)
+        self.data = pd.read_csv(csv_file).head(4)
         self.tokenizer = tokenizer
         self.model = model
 
@@ -30,11 +30,11 @@ def collate_fn(batch):
     return latents_padded, attention_mask_padded
 
 def get_dataloaders(config):
-    tokenizer, model = load_esm2_model(config.MODEL_NAME)
+    tokenizer, masked_model, embedding_model = load_esm2_model(config.MODEL_NAME)
     
-    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_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)

utils/esm_utils.py

@@ -1,14 +1,15 @@
 import torch
-from transformers import AutoTokenizer, AutoModel
+import config
+from transformers import AutoTokenizer, AutoModel, AutoModelForMaskedLM
 
 def load_esm2_model(model_name):
     tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModel.from_pretrained(model_name)
-    return tokenizer, model
+    masked_model = AutoModelForMaskedLM.from_pretrained(model_name)
+    embedding_model = AutoModel.from_pretrained(model_name)
+    return tokenizer, masked_model, embedding_model
 
 def get_latents(model, tokenizer, sequence):
     inputs = tokenizer(sequence, return_tensors="pt")
     with torch.no_grad():
         outputs = model(**inputs)
-    return outputs.last_hidden_state.squeeze(0)
-
+    return outputs.last_hidden_state.squeeze(0)