Update app.py

app.py

@@ -38,23 +38,19 @@ def predict_stability_with_pdb(model_choice, organism_choice, pdb_files, cfg=Con
     sequences = []
 
     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["file_name"].append(pdb_file.name)
-                results["raw prediction value"].append(None)
-                results["binary prediction value"].append(None)
-                continue
-
-            sequence = sequences[2] if model_choice == "SaProt" else sequences[0]
-            file_names.append(pdb_file.name)
-            sequences.append(sequence)
-        except Exception as e:
-                results["file_name"].append(pdb_file.name)
-                results["raw prediction value"].append(None)
-                results["binary prediction value"].append(None)
+        pdb_path = pdb_file.name
+        os.system("chmod 777 bin/foldseek")
+        sequences = get_foldseek_seq(pdb_path)
+        if not sequences:
+            results["file_name"].append(pdb_file.name.split("/")[-1])
+            results["raw prediction value"].append(None)
+            results["binary prediction value"].append(None)
+            continue
+
+        sequence = sequences[2] if model_choice == "SaProt" else sequences[0]
+        file_names.append(pdb_file.name.split("/")[-1])
+        sequences.append(sequence)
+
     raw_prediction, binary_prediction = predict_stability_core(model_choice, organism_choice, sequences, cfg)
     results["file_name"] = results["file_name"] + file_names
     results["raw prediction value"] = results["raw prediction value"] + raw_prediction
@@ -70,8 +66,12 @@ def predict_stability_with_sequence(model_choice, organism_choice, sequence, cfg
     try:
         if not sequence:
             return "No valid sequence provided."
-        
-        return predict_stability_core(model_choice, organism_choice, sequence, cfg)
+        raw_prediction, binary_prediction = predict_stability_core(model_choice, organism_choice, sequence, cfg)
+        df = pd.DataFrame({"sequence": sequence, "raw prediction value": raw_prediction, "binary prediction value": binary_prediction})
+        output_csv = "/tmp/predictions.csv"
+        df.to_csv(output_csv, index=False)
+
+        return output_csv 
     except Exception as e:
         return f"An error occurred: {str(e)}"
 
@@ -192,8 +192,8 @@ with gr.Blocks() as demo:
                 lines=8,
             )
             predict_button = gr.Button("Predict Stability")
-            prediction_output = gr.HTML(
-                label="Stability Prediction"
+            prediction_output = gr.File(
+                label="Download Predictions"
             )
 
             predict_button.click(