Update app.py

app.py

@@ -277,6 +277,7 @@ MODEL_OPTIONS = [
 
 PEFT_MODEL_OPTIONS = [
     "AmelieSchreiber/esm2_t12_35M_lora_binding_sites_v2_cp3",
+    "wangjin2000/esm2_t6_8M-lora-binding-sites_2024-07-02_09-26-54",
 ]  # finetuned models 
 
 
@@ -372,7 +373,7 @@ with demo:
     gr.Markdown("# DEMO FOR ESM2Bind")
     #gr.Textbox(dubug_result)
     
-    gr.Markdown("## Finetune Pre-trained Model")
+    #gr.Markdown("## Finetune Pre-trained Model")
     with gr.Column():
         gr.Markdown("## Select a base model")
         gr.Markdown(
@@ -392,7 +393,7 @@ with demo:
                     label="PEFT Model Name",
                     interactive = True,
                 )
-            with gr.Column(variant="compact"): 
+            with gr.Column(scale=0.5, variant="compact"): 
                     name = gr.Dropdown(
                         label="Choose a Sample Protein", 
                         value="Default protein", 
@@ -411,20 +412,27 @@ with demo:
                     placeholder="Paste your protein sequence here...",
                     interactive = True,
                 ) 
-            with gr.Column(variant="compact", scale = 2):  
-                finetune_button = gr.Button(
-                    value="Finetune Pre-trained Model", 
+        with gr.Column(variant="compact", scale = 2):
+                predict_btn = gr.Button(
+                    value="Predict binding site", 
                     interactive=True, 
                     variant="primary",
                 )
                 plot_struc_btn = gr.Button(value = "Plot ESMFold Predicted Structure ", variant="primary")
         with gr.Row():  
-            finetune_output_text = gr.Textbox(
-                lines=1,
-                max_lines=12,
-                label="Finetune Status",
-                placeholder="Finetune Status Shown Here",
-            )               
+            with gr.Column(variant="compact", scale = 5):
+                output_text = gr.Textbox(
+                    lines=1,
+                    max_lines=12,
+                    label="Output",
+                    placeholder="Output",
+                )   
+            with gr.Column(variant="compact", scale = 5):  
+                finetune_button = gr.Button(
+                    value="Finetune Pre-trained Model", 
+                    interactive=True, 
+                    variant="primary",
+                )
         with gr.Row(): 
             output_viewer = gr.HTML()
             output_file = gr.File(
@@ -437,14 +445,21 @@ with demo:
     # select protein sample
     name.change(fn=suggest, inputs=name, outputs=input_seq)   
 
+    # "Predict binding site" actions
+    predict_btn.click(
+        fn = predict_bind,
+        inputs=[base_model_name,PEFT_model_name,input_seq], 
+        outputs = [output_text],
+    )
     # "Finetune Pre-trained Model" actions
     finetune_button.click(
         fn = train_function_no_sweeps,
-        inputs=[base_model_name], #finetune_dataset_name],
-        outputs = [finetune_output_text],
+        inputs=[base_model_name],
+        outputs = [output_text],
     )
            
     # plot protein structure
     plot_struc_btn.click(fn=plot_struc, inputs=input_seq, outputs=[output_file, output_viewer])
+
     
 demo.launch()