Update app.py

app.py

@@ -28,27 +28,39 @@ class Config:
         self.seed = 42
 
 
-def predict_stability(model_choice, organism_choice, pdb_file=None, sequence=None, cfg=Config()):
-    return pdb_file, sequence
-    # if pdb_file:
-    #     pdb_path = pdb_file.name  # Get the path of the uploaded PDB file
-    #     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]
-
-    # if not sequence:
-    #     return "No valid input provided."
-
-    # cell_line = "HeLa" if organism_choice == "Human" else "NIH3T3"
-    # cfg.model = f"sagawa/PLTNUM-{model_choice}-{cell_line}"
-    # cfg.architecture = model_choice
-    # cfg.model_path = f"sagawa/PLTNUM-{model_choice}-{cell_line}"
-    
-    # output = predict(cfg, sequence)
-    # return output
 
+def predict_stability_with_pdb(model_choice, organism_choice, pdb_file, 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]
+
+        return predict_stability_core(model_choice, organism_choice, sequence, cfg)
+    except Exception as e:
+        return f"An error occurred: {str(e)}"
+
+
+def predict_stability_with_sequence(model_choice, organism_choice, sequence, cfg=Config()):
+    try:
+        if not sequence:
+            return "No valid sequence provided."
+        
+        return predict_stability_core(model_choice, organism_choice, sequence, cfg)
+    except Exception as e:
+        return f"An error occurred: {str(e)}"
+
+
+def predict_stability_core(model_choice, organism_choice, sequence, cfg=Config()):
+    cell_line = "HeLa" if organism_choice == "Human" else "NIH3T3"
+    cfg.model = f"sagawa/PLTNUM-{model_choice}-{cell_line}"
+    cfg.architecture = model_choice
+    cfg.model_path = f"sagawa/PLTNUM-{model_choice}-{cell_line}"
+
+    output = predict(cfg, sequence)
+    return output
 
 
 def get_foldseek_seq(pdb_path):
@@ -155,7 +167,7 @@ with gr.Blocks() as demo:
             )
 
             predict_button.click(
-                fn=predict_stability,
+                fn=predict_stability_with_pdb,
                 inputs=[model_choice, organism_choice, pdb_file],
                 outputs=prediction_output,
             )
@@ -173,8 +185,8 @@ with gr.Blocks() as demo:
             )
 
             predict_button.click(
-                fn=predict_stability,
-                inputs=[model_choice, organism_choice, None, sequence],
+                fn=predict_stability_with_sequence,
+                inputs=[model_choice, organism_choice, sequence],
                 outputs=prediction_output,
             )