Update app.py

app.py

@@ -33,7 +33,24 @@ from accelerate import Accelerator
 # Imports specific to the custom peft lora model
 from peft import get_peft_config, PeftModel, PeftConfig, get_peft_model, LoraConfig, TaskType
 
-
+from plot_pdb import plot_struc
+
+def suggest(option):
+    if option == "Plastic degradation protein":
+        suggestion = "MGSSHHHHHHSSGLVPRGSHMRGPNPTAASLEASAGPFTVRSFTVSRPSGYGAGTVYYPTNAGGTVGAIAIVPGYTARQSSIKWWGPRLASHGFVVITIDTNSTLDQPSSRSSQQMAALRQVASLNGTSSSPIYGKVDTARMGVMGWSMGGGGSLISAANNPSLKAAAPQAPWDSSTNFSSVTVPTLIFACENDSIAPVNSSALPIYDSMSRNAKQFLEINGGSHSCANSGNSNQALIGKKGVAWMKRFMDNDTRYSTFACENPNSTRVSDFRTANCSLEDPAANKARKEAELAAATAEQ"
+    elif option == "Default protein":
+        #suggestion = "MAPLRKTYVLKLYVAGNTPNSVRALKTLNNILEKEFKGVYALKVIDVLKNPQLAEEDKILATPTLAKVLPPPVRRIIGDLSNREKVLIGLDLLYEEIGDQAEDDLGLE"
+        suggestion = "MAVPETRPNHTIYINNLNEKIKKDELKKSLHAIFSRFGQILDILVSRSLKMRGQAFVIFKEVSSATNALRSMQGFPFYDKPMRIQYAKTDSDIIAKMKGT" 
+    elif option == "Antifreeze protein":
+        suggestion = "QCTGGADCTSCTGACTGCGNCPNAVTCTNSQHCVKANTCTGSTDCNTAQTCTNSKDCFEANTCTDSTNCYKATACTNSSGCPGH"
+    elif option == "AI Generated protein":
+        suggestion = "MSGMKKLYEYTVTTLDEFLEKLKEFILNTSKDKIYKLTITNPKLIKDIGKAIAKAAEIADVDPKEIEEMIKAVEENELTKLVITIEQTDDKYVIKVELENEDGLVHSFEIYFKNKEEMEKFLELLEKLISKLSGS"
+    elif option == "7-bladed propeller fold":
+        suggestion = "VKLAGNSSLCPINGWAVYSKDNSIRIGSKGDVFVIREPFISCSHLECRTFFLTQGALLNDKHSNGTVKDRSPHRTLMSCPVGEAPSPYNSRFESVAWSASACHDGTSWLTIGISGPDNGAVAVLKYNGIITDTIKSWRNNILRTQESECACVNGSCFTVMTDGPSNGQASYKIFKMEKGKVVKSVELDAPNYHYEECSCYPNAGEITCVCRDNWHGSNRPWVSFNQNLEYQIGYICSGVFGDNPRPNDGTGSCGPVSSNGAYGVKGFSFKYGNGVWIGRTKSTNSRSGFEMIWDPNGWTETDSSFSVKQDIVAITDWSGYSGSFVQHPELTGLDCIRPCFWVELIRGRPKESTIWTSGSSISFCGVNSDTVGWSWPDGAELPFTIDK"
+    else:
+        suggestion = ""
+    return suggestion
+    
 # Helper Functions and Data Preparation
 def truncate_labels(labels, max_length):
     """Truncate labels to the specified max_length."""
@@ -374,13 +391,33 @@ with demo:
                     value=PEFT_MODEL_OPTIONS[0],
                     label="PEFT Model Name",
                     interactive = True,
-                )  
-            with gr.Column(scale=0.5, variant="compact"):    
+                )
+            with gr.Column(variant="compact"): 
+                    name = gr.Dropdown(
+                        label="Choose a Sample Protein", 
+                        value="Default protein", 
+                        choices=["Default protein", "Antifreeze protein", "Plastic degradation protein",  "AI Generated protein", "7-bladed propeller fold", "custom"]
+                    )
+        gr.Markdown(
+                """Predict binding site and Plot structure for selected protein sequence:"""
+                )
+        with gr.Row():
+            with gr.Column(variant="compact", scale = 8):  
+                input_seq = gr.Textbox(
+                    lines=1,
+                    max_lines=12,
+                    label="Protein sequency to be predicted:",
+                    value="MAVPETRPNHTIYINNLNEKIKKDELKKSLHAIFSRFGQILDILVSRSLKMRGQAFVIFKEVSSATNALRSMQGFPFYDKPMRIQYAKTDSDIIAKMKGT",
+                    placeholder="Paste your protein sequence here...",
+                    interactive = True,
+                ) 
+            with gr.Column(variant="compact", scale = 2):  
                 finetune_button = gr.Button(
                     value="Finetune Pre-trained Model", 
                     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,
@@ -388,7 +425,18 @@ with demo:
                 label="Finetune Status",
                 placeholder="Finetune Status Shown Here",
             )               
-                    
+        with gr.Row(): 
+            output_viewer = gr.HTML()
+            output_file = gr.File(
+                label="Download as Text File",
+                file_count="single",
+                type="filepath",  
+                interactive=False,
+            )  
+            
+    # select protein sample
+    name.change(fn=suggest, inputs=name, outputs=input_seq)   
+
     # "Finetune Pre-trained Model" actions
     finetune_button.click(
         fn = train_function_no_sweeps,
@@ -396,5 +444,7 @@ with demo:
         outputs = [finetune_output_text],
     )
            
-
+    # plot protein structure
+    plot_struc_btn.click(fn=plot_struc, inputs=input_seq, outputs=[output_file, output_viewer])
+    
 demo.launch()