Update app.py

app.py

@@ -41,10 +41,10 @@ class ProteinDataset(Dataset):
     def __init__(self, file, tokenizer, peptide_length):
         data = pd.read_csv(file)
         self.tokenizer = tokenizer
-        self.proteins = data["Receptor Sequence"].tolist()
-        self.peptides = data["Binder"].tolist()
-        #self.proteins = data["P_Sequence"].tolist()  #header defined by Lin Qiao
-        #self.peptides = data["p_Sequence"].tolist()
+        #self.proteins = data["Receptor Sequence"].tolist()
+        #self.peptides = data["Binder"].tolist()
+        self.proteins = data["P_Sequence"].tolist()  #header defined by Lin Qiao
+        self.peptides = data["p_Sequence"].tolist()
         self.max_length_pm = 500 + 2 + peptide_length  #assume the maz length of protein is 500
         
     def __len__(self):
@@ -80,9 +80,11 @@ def finetune(base_model_path, peptide_length):   #, train_dataset, test_dataset)
     # Tokenization
     tokenizer = AutoTokenizer.from_pretrained(base_model_path) #("facebook/esm2_t12_35M_UR50D")
 
-    train_dataset = ProteinDataset("./datasets/pepnn_train.csv", tokenizer, peptide_length)
-    test_dataset = ProteinDataset("./datasets/pepnn_test.csv", tokenizer, peptide_length)
-
+    #train_dataset = ProteinDataset("./datasets/pepnn_train.csv", tokenizer, peptide_length)
+    #test_dataset = ProteinDataset("./datasets/pepnn_test.csv", tokenizer, peptide_length)
+    train_dataset = ProteinDataset("./datasets/peptide-protein-train.csv", tokenizer, peptide_length)
+    test_dataset = ProteinDataset("./datasets/peptide-protein-test.csv", tokenizer, peptide_length)
+    
     model_name_base = base_model_path.split("/")[1]
     timestamp = datetime.now().strftime('%Y-%m-%d_%H')
     lr = 0.0007984276816171436
@@ -269,10 +271,10 @@ def predict_peptide_from_file(base_model_path, finetuned_model_path, file_obj, m
     results = []
     
     for i, row in input.iterrows():
-        protein_seq = row['Receptor Sequence']
-        peptide_seq = row['Peptide Sequence']
-        #protein_seq = row['P_Sequence']
-        #peptide_seq = row['p_Sequence']
+        #protein_seq = row['Receptor Sequence']
+        #peptide_seq = row['Peptide Sequence']
+        protein_seq = row['P_Sequence']
+        peptide_seq = row['p_Sequence']
         peptide_length = min([len(peptide_seq), max_peptide_length])  # use the same length of ground truth peptide length for prediction limited to max_peptide_length
 
         #get metrics for ground truth peptide