Update app.py

app.py

@@ -30,23 +30,37 @@ class Config:
 
 
 def predict_stability_with_pdb(model_choice, organism_choice, pdb_files, cfg=Config()):
-    try:
-        results = []
-        for pdb_file in pdb_files:
+    results = []
+    for pdb_file in pdb_files:
+        try:
             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}.")
+                results.append({"file_name": pdb_path,
+                                "raw prediction value": None,
+                                "binary prediction value": None
+                                })
                 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)}"
 
+            sequence = sequences[2] if model_choice == "SaProt" else sequences[0]
+            output = predict_stability_core(model_choice, organism_choice, sequence, cfg)
+
+            results.append({"file_name": pdb_path,
+                            "raw prediction value": output["raw prediction values"][0],
+                            "binary prediction value": output["binary prediction values"][0]
+                            })
+        except Exception as e:
+            results.append({"file_name": pdb_file.name,
+                            "raw prediction value": None,
+                            "binary prediction value": None
+                            })
+            
+    df = pd.DataFrame(results)
+    output_csv = "/tmp/predictions.csv"
+    df.to_csv(output_csv, index=False)
+
+    return output_csv
 
 def predict_stability_with_sequence(model_choice, organism_choice, sequence, cfg=Config()):
     try:
@@ -167,8 +181,8 @@ with gr.Blocks() as demo:
             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"
+            prediction_output = gr.File(
+                label="Download Predictions"
             )
 
             predict_button.click(