fix: formatting, output type

app.py

@@ -27,7 +27,7 @@ def convert_outputs_to_pdb(outputs):
             chain_index=outputs["chain_index"][i] if "chain_index" in outputs else None,
         )
         pdbs.append(to_pdb(pred))
-    return pdbs
+    return pdbs[0]
 
 def fold_prot_locally(sequence):
     logger.info("Folding: " + sequence)
@@ -38,30 +38,41 @@ def fold_prot_locally(sequence):
     pdb = convert_outputs_to_pdb(output)
     return pdb
 
+def suggest(option):
+   if option == "Plastic degradation protein":
+     suggestion = "MGSSHHHHHHSSGLVPRGSHMRGPNPTAASLEASAGPFTVRSFTVSRPSGYGAGTVYYPTNAGGTVGAIAIVPGYTARQSSIKWWGPRLASHGFVVITIDTNSTLDQPSSRSSQQMAALRQVASLNGTSSSPIYGKVDTARMGVMGWSMGGGGSLISAANNPSLKAAAPQAPWDSSTNFSSVTVPTLIFACENDSIAPVNSSALPIYDSMSRNAKQFLEINGGSHSCANSGNSNQALIGKKGVAWMKRFMDNDTRYSTFACENPNSTRVSDFRTANCSLEDPAANKARKEAELAAATAEQ"
+   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
+
 sample_code = """
-## Sample usage
+from gradio_client import Client
+import json
+
+client = Client("https://wwydmanski-esmfold.hf.space/")
+
+def fold_huggingface(sequence, fname=None):
+    result = client.predict(
+                    sequence,	# str in 'sequence' Textbox component
+                    api_name="/predict")
 
-    from gradio_client import Client
-    
-    client = Client("https://wwydmanski-esmfold.hf.space/")
-    
-    def fold_huggingface(sequence, fname=None):
-        result = client.predict(
-                        sequence,	# str in 'sequence' Textbox component
-                        api_name="/predict")
-            result = eval(result)[0]
-            if fname is None:
-                with tempfile.NamedTemporaryFile("w", delete=False, suffix=".pdb", prefix="esmfold_") as fp:
-                    fp.write(result)
-                    fp.flush()
-                    return fp.name
-            else:
-                with open(fname, "w") as fp:
-                    fp.write(result)
-                    fp.flush()
-                return fname
-    pdb_fname = fold_huggingface("MALWMRLLPLLALLALWGPDPAAAFVNQHLCGSHLVEALYLVCGERGFFYTPKTRREAEDLQVGQVELGGGPGAGSLQPLALEGSLQKRGIVEQCCTSICSLYQLENYCN")
-    
+    if fname is None:
+        with tempfile.NamedTemporaryFile("w", delete=False, suffix=".pdb", prefix="esmfold_") as fp:
+            fp.write(result)
+            fp.flush()
+            return fp.name
+    else:
+        with open(fname, "w") as fp:
+            fp.write(result)
+            fp.flush()
+        return fname
+pdb_fname = fold_huggingface("MALWMRLLPLLALLALWGPDPAAAFVNQHLCGSHLVEALYLVCGERGFFYTPKTRREAEDLQVGQVELGGGPGAGSLQPLALEGSLQKRGIVEQCCTSICSLYQLENYCN")
 """
 
 tokenizer = AutoTokenizer.from_pretrained("facebook/esmfold_v1")
@@ -69,6 +80,25 @@ model = EsmForProteinFolding.from_pretrained("facebook/esmfold_v1", low_cpu_mem_
 model.esm = model.esm.half()
 torch.backends.cuda.matmul.allow_tf32 = True
 
-iface = gr.Interface(fn=fold_prot_locally, inputs="text", outputs="text", article=sample_code, title="ESMFold")
+with gr.Blocks() as demo:
+    gr.Markdown("# ESMFold")
+    with gr.Row():
+        with gr.Column():
+            inp = gr.Textbox(lines=1, label="Sequence")
+            name = gr.Dropdown(label="Choose a Sample Protein", value="Plastic degradation protein", choices=["Antifreeze protein", "Plastic degradation protein",  "AI Generated protein", "7-bladed propeller fold", "custom"])
+            btn = gr.Button("🔬 Predict Structure ")
+        
+    with gr.Row():
+        with gr.Column():
+            gr.Markdown("## Sample code")
+            gr.Code(sample_code, label="Sample usage", language="python", interactive=False)
+
+    with gr.Row():
+        with gr.Column():
+            gr.Markdown("## Output")
+            out = gr.Code(label="Output", interactive=False)
+
+    name.change(fn=suggest, inputs=name, outputs=inp)
+    btn.click(fold_prot_locally, inputs=[inp], outputs=[out])
 
-iface.launch()
+demo.launch()