Update app.py

app.py

@@ -18,6 +18,7 @@ with gr.Blocks() as demo:
     model_choice = gr.Radio(choices=["SaProt", "ESM-2"], label="Select PLTNUM's base model.", value="SaProt")
     organism_choice = gr.Radio(choices=["Mouse", "Human"], label="Select the target organism.", value="Mouse")
     
+    
     with gr.Tabs():
         with gr.TabItem("Upload PDB File"):
             gr.Markdown("### Upload your PDB file:")
@@ -25,7 +26,7 @@ with gr.Blocks() as demo:
             predict_button = gr.Button("Predict Stability")
             prediction_output = gr.Textbox(label="Stability Prediction", interactive=False)
 
-            predict_button.click(fn=predict_stability, inputs=pdb_file, outputs=prediction_output)
+            predict_button.click(fn=predict_stability, inputs=[model_choice, pdb_file], outputs=prediction_output)
         
         with gr.TabItem("Enter Protein Sequence"):
             gr.Markdown("### Enter the protein sequence:")
@@ -37,17 +38,17 @@ with gr.Blocks() as demo:
             predict_button = gr.Button("Predict Stability")
             prediction_output = gr.Textbox(label="Stability Prediction", interactive=False)
 
-            predict_button.click(fn=predict_stability, inputs=sequence, outputs=prediction_output)
+            predict_button.click(fn=predict_stability, inputs=[model_choice, sequence], outputs=prediction_output)
     
-    with gr.Row():
-        gr.Markdown(
-            """
-            ### How to Use:
-            - **Upload PDB File**: Choose the 'Upload PDB File' tab and upload your file.
-            - **Enter Sequence**: Alternatively, switch to the 'Enter Protein Sequence' tab and input your sequence.
-            - **Predict**: Click 'Predict Stability' to receive the prediction.
-            """
-        )
+    gr.Markdown(
+        """
+        ### How to Use:
+        - **Select Model**: Choose between 'Model A' or 'Model B' for your prediction.
+        - **Upload PDB File**: Choose the 'Upload PDB File' tab and upload your file.
+        - **Enter Sequence**: Alternatively, switch to the 'Enter Protein Sequence' tab and input your sequence.
+        - **Predict**: Click 'Predict Stability' to receive the prediction.
+        """
+    )
     
     gr.Markdown(
         """
@@ -56,4 +57,4 @@ with gr.Blocks() as demo:
         """
     )
 
-demo.launch()
+demo.launch()