Update app.py

app.py

@@ -29,16 +29,74 @@ class Config:
 
 
 
-def predict_stability_with_pdb(model_choice, organism_choice, pdb_file, cfg=Config()):
+def predict_stability_with_pdb(model_choice, organism_choice, pdb_files, cfg=Config()):
     try:
-        pdb_path = pdb_file.name
-        os.system("chmod 777 bin/foldseek")
-        sequences = get_foldseek_seq(pdb_path)
-        if not sequences:
-            return "Failed to extract sequence from the PDB file."
-        sequence = sequences[2] if model_choice == "SaProt" else sequences[0]
+        results = []
+        for pdb_file in pdb_files:
+            pdb_path = pdb_file.name
+            os.system("chmod 777 bin/foldseek")
+            sequences = get_foldseek_seq(pdb_path)
+            if not sequences:
+                return "Failed to extract sequence from the PDB file."
+            sequence = sequences[2] if model_choice == "SaProt" else sequences[0]
+
+            prediction = predict_stability_core(model_choice, organism_choice, sequence, cfg)
+
+To modify the existing Gradio interface to accept multiple PDB files and process them, you need to adjust the file input component and modify the function that processes the PDB files. Here’s how you can do it:
+
+Adjusting the Gradio Interface
+Allow Multiple File Uploads: Modify the gr.File component to accept multiple files.
+Process Each File Separately: Adjust the prediction function to loop over all uploaded files and process each one.
+Here's the updated code:
+
+python
+Copy code
+import gradio as gr
+import sys
+import random
+import os
+import pandas as pd
+import torch
+import itertools
+from torch.utils.data import DataLoader
+from transformers import AutoTokenizer
 
-        return predict_stability_core(model_choice, organism_choice, sequence, cfg)
+sys.path.append("scripts/")
+from foldseek_util import get_struc_seq
+from utils import seed_everything
+from models import PLTNUM_PreTrainedModel
+from datasets_ import PLTNUMDataset
+
+
+class Config:
+    def __init__(self):
+        self.batch_size = 2
+        self.use_amp = False
+        self.num_workers = 1
+        self.max_length = 512
+        self.used_sequence = "left"
+        self.padding_side = "right"
+        self.task = "classification"
+        self.sequence_col = "sequence"
+        self.seed = 42
+
+
+def predict_stability_with_pdb(model_choice, organism_choice, pdb_files, cfg=Config()):
+    try:
+        results = []
+        for pdb_file in pdb_files:
+            pdb_path = pdb_file.name
+            os.system("chmod 777 bin/foldseek")
+            sequences = get_foldseek_seq(pdb_path)
+            if not sequences:
+                results.append(f"Failed to extract sequence from {pdb_file.name}.")
+                continue
+            sequence = sequences[2] if model_choice == "SaProt" else sequences[0]
+
+            prediction = predict_stability_core(model_choice, organism_choice, sequence, cfg)
+            results.append(f"Prediction for {pdb_file.name}: {prediction}")
+            
+            return "<br>".join(results)
     except Exception as e:
         return f"An error occurred: {str(e)}"
 
@@ -159,8 +217,8 @@ with gr.Blocks() as demo:
 
     with gr.Tabs():
         with gr.TabItem("Upload PDB File"):
-            gr.Markdown("### Upload your PDB file:")
-            pdb_file = gr.File(label="Upload PDB File")
+            gr.Markdown("### Upload your PDB files:")
+            pdb_files = gr.File(label="Upload PDB Files", file_count="multiple")
             predict_button = gr.Button("Predict Stability")
             prediction_output = gr.HTML(
                 label="Stability Prediction"
@@ -168,7 +226,7 @@ with gr.Blocks() as demo:
 
             predict_button.click(
                 fn=predict_stability_with_pdb,
-                inputs=[model_choice, organism_choice, pdb_file],
+                inputs=[model_choice, organism_choice, pdb_files],
                 outputs=prediction_output,
             )