Update app.py

app.py

@@ -81,7 +81,39 @@ class WeightedTrainer(Trainer):
         loss = compute_loss(model, inputs) 
         return (loss, outputs) if return_outputs else loss
 
-#         
+# Predict binding site with finetuned PEFT model
+def predict_bind(base_model_path,PEFT_model_path,input_seq):
+    # Load the model
+    base_model = AutoModelForTokenClassification.from_pretrained(base_model_path)
+    loaded_model = PeftModel.from_pretrained(base_model, PEFT_model_path)
+
+    # Ensure the model is in evaluation mode
+    loaded_model.eval()
+    
+    # Tokenization
+    tokenizer = AutoTokenizer.from_pretrained(base_model_path) 
+
+    # Tokenize the sequence
+    inputs = tokenizer(input_sequence, return_tensors="pt", truncation=True, max_length=1024, padding='max_length')
+
+    # Run the model
+    with torch.no_grad():
+        logits = loaded_model(**inputs).logits
+
+    # Get predictions
+    tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])  # Convert input ids back to tokens
+    predictions = torch.argmax(logits, dim=2)
+
+    binding_site=[]
+    # Print the predicted labels for each token
+    for n, token, prediction in enumerate(zip(tokens, predictions[0].numpy())):
+        if token not in ['<pad>', '<cls>', '<eos>']:
+            print((token, id2label[prediction]))
+        if prediction == 1:
+            print((n+1,token, id2label[prediction]))
+            binding_site.append(n+1,token, id2label[prediction])
+    return binding_site
+
 # fine-tuning function
 def train_function_no_sweeps(base_model_path):   #, train_dataset, test_dataset):
     
@@ -102,8 +134,10 @@ def train_function_no_sweeps(base_model_path):   #, train_dataset, test_dataset)
     #base_model_path = "facebook/esm2_t12_35M_UR50D"  
     
     # Define labels and model
-    id2label = {0: "No binding site", 1: "Binding site"}
-    label2id = {v: k for k, v in id2label.items()}
+    #id2label = {0: "No binding site", 1: "Binding site"}
+    #label2id = {v: k for k, v in id2label.items()}
+
+    
     base_model = AutoModelForTokenClassification.from_pretrained(base_model_path, num_labels=len(id2label), id2label=id2label, label2id=label2id)
 
     '''
@@ -289,12 +323,14 @@ with torch.no_grad():
 tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])  # Convert input ids back to tokens
 predictions = torch.argmax(logits, dim=2)
 
+'''
 # Define labels
 id2label = {
     0: "No binding site",
     1: "Binding site"
 }
 
+'''
 # Print the predicted labels for each token
 for token, prediction in zip(tokens, predictions[0].numpy()):
     if token not in ['<pad>', '<cls>', '<eos>']: