Spaces:

GaganaMD
/

Protein-Structure-Prediction

Runtime error

App Files Files Community

GaganaMD commited on Jul 17, 2024

Commit

50fcaf6

verified ·

1 Parent(s): 576036c

Upload 2 files

Browse files

Files changed (2) hide show

app.py +189 -0
client.py +14 -0

app.py ADDED Viewed

	@@ -0,0 +1,189 @@

+import gradio as gr
+from transformers import AutoTokenizer, EsmForProteinFolding
+from transformers.models.esm.openfold_utils.protein import to_pdb, Protein as OFProtein
+from transformers.models.esm.openfold_utils.feats import atom14_to_atom37
+import torch
+from logging import getLogger
+logger = getLogger(__name__)
+def convert_outputs_to_pdb(outputs):
+    final_atom_positions = atom14_to_atom37(outputs["positions"][-1], outputs)
+    outputs = {k: v.to("cpu").numpy() for k, v in outputs.items()}
+    final_atom_positions = final_atom_positions.cpu().numpy()
+    final_atom_mask = outputs["atom37_atom_exists"]
+    pdbs = []
+    for i in range(outputs["aatype"].shape[0]):
+        aa = outputs["aatype"][i]
+        pred_pos = final_atom_positions[i]
+        mask = final_atom_mask[i]
+        resid = outputs["residue_index"][i] + 1
+        pred = OFProtein(
+            aatype=aa,
+            atom_positions=pred_pos,
+            atom_mask=mask,
+            residue_index=resid,
+            b_factors=outputs["plddt"][i],
+            chain_index=outputs["chain_index"][i] if "chain_index" in outputs else None,
+        )
+        pdbs.append(to_pdb(pred))
+    return pdbs[0]
+def fold_prot_locally(sequence):
+    logger.info("Folding: " + sequence)
+    tokenized_input = tokenizer([sequence], return_tensors="pt", add_special_tokens=False)['input_ids'].cuda()
+    with torch.no_grad():
+        output = model(tokenized_input)
+    pdb = convert_outputs_to_pdb(output)
+    return pdb
+def get_esm2_embeddings(sequence):
+    logger.info("Getting embeddings for: " + sequence)
+    tokenized_input = tokenizer([sequence], return_tensors="pt", add_special_tokens=False)['input_ids'].cuda()
+    with torch.no_grad():
+        aa = tokenized_input
+        L = aa.shape[1]
+        device = tokenized_input.device
+        attention_mask = torch.ones_like(aa, device=device)
+        # === ESM ===
+        esmaa = model.af2_idx_to_esm_idx(aa, attention_mask)
+        esm_s = model.compute_language_model_representations(esmaa)
+    return {"res": esm_s.cpu().tolist()}
+def get_esmfold_embeddings(sequence):
+    logger.info("Getting embeddings for: " + sequence)
+    tokenized_input = tokenizer([sequence], return_tensors="pt", add_special_tokens=False)['input_ids'].cuda()
+    with torch.no_grad():
+        output = model(tokenized_input)
+    return {"res": output["s_s"].cpu().tolist()}
+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
+def molecule(mol):
+    x = (
+        """<!DOCTYPE html>
+        <html>
+        <head>
+    <meta http-equiv="content-type" content="text/html; charset=UTF-8" />
+    <style>
+    body{
+        font-family:sans-serif
+    }
+    .mol-container {
+    width: 100%;
+    height: 600px;
+    position: relative;
+    }
+    .mol-container select{
+        background-image:None;
+    }
+    </style>
+     <script src="https://cdnjs.cloudflare.com/ajax/libs/jquery/3.6.3/jquery.min.js" integrity="sha512-STof4xm1wgkfm7heWqFJVn58Hm3EtS31XFaagaa8VMReCXAkQnJZ+jEy8PCC/iT18dFy95WcExNHFTqLyp72eQ==" crossorigin="anonymous" referrerpolicy="no-referrer"></script>
+    <script src="https://3Dmol.csb.pitt.edu/build/3Dmol-min.js"></script>
+    </head>
+    <body>
+    <div id="container" class="mol-container"></div>
+            <script>
+               let pdb = `"""
+        + mol
+        + """`
+             $(document).ready(function () {
+                let element = $("#container");
+                let config = { backgroundColor: "white" };
+                let viewer = $3Dmol.createViewer(element, config);
+                viewer.addModel(pdb, "pdb");
+                viewer.getModel(0).setStyle({}, { cartoon: { colorscheme:"whiteCarbon" } });
+                viewer.zoomTo();
+                viewer.render();
+                viewer.zoom(0.8, 2000);
+              })
+        </script>
+        </body></html>"""
+    )
+    return f"""<iframe style="width: 100%; height: 600px" name="result" allow="midi; geolocation; microphone; camera;
+    display-capture; encrypted-media;" sandbox="allow-modals allow-forms
+    allow-scripts allow-same-origin allow-popups
+    allow-top-navigation-by-user-activation allow-downloads" allowfullscreen=""
+    allowpaymentrequest="" frameborder="0" srcdoc='{x}'></iframe>"""
+sample_code = """
+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="/pdb")
+    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")
+model = EsmForProteinFolding.from_pretrained("facebook/esmfold_v1", low_cpu_mem_usage=True).cuda()
+model.esm = model.esm.half()
+torch.backends.cuda.matmul.allow_tf32 = True
+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():
+        gr.Markdown("## Output")
+    with gr.Row():
+        with gr.Column():
+            out = gr.Code(label="Output", interactive=False)
+        with gr.Column():
+            out_mol = gr.HTML(label="3D Structure")
+    with gr.Row(visible=False):
+       with gr.Column():
+           gr.Markdown("## Embeddings")
+           embs = gr.JSON(label="Embeddings")
+    name.change(fn=suggest, inputs=name, outputs=inp)
+    btn.click(fold_prot_locally, inputs=[inp], outputs=[out], api_name="pdb")
+    btn.click(get_esmfold_embeddings, inputs=[inp], outputs=[embs], api_name="embeddings")
+    btn.click(get_esm2_embeddings, inputs=[inp], outputs=[embs], api_name="esm2_embeddings")
+    out.change(fn=molecule, inputs=[out], outputs=[out_mol], api_name="3d_fold")
+demo.launch()

client.py ADDED Viewed

	@@ -0,0 +1,14 @@

+#%%
+from gradio_client import Client
+#%%
+# client = Client("https://huggingface.co/spaces/GaganaMD/Protein-Structure-Prediction")
+client = Client("http://localhost:7860")
+# %%
+result = client.predict(
+    "MALWMRLLPLLALLALWGPDPAAAFVNQHLCGSHLVEALYLVCGERGFFYTPKTRREAEDLQVGQVELGGGPGAGSLQPLALEGSLQKRGIVEQCCTSICSLYQLENYCN",	# str in 'sequence' Textbox component
+    api_name="/esm2_embeddings")
+# %%
+result