Spaces:

as-cle-bert
/

proteins-with-esm

Sleeping

App Files Files Community

as-cle-bert commited on Jun 28

Commit

884b6ae

•

1 Parent(s): 1468add

Create apps.py

Browse files

Files changed (1) hide show

apps.py +106 -0

apps.py ADDED Viewed

	@@ -0,0 +1,106 @@

+from huggingface_hub import login
+from esm.models.esm3 import ESM3
+from esm.sdk.api import ESM3InferenceClient, ESMProtein, GenerationConfig
+import spaces
+import os
+from gradio_molecule3d import Molecule3D
+# This will prompt you to get an API key from huggingface hub, make one with
+# "Read" or "Write" permission and copy it back here.
+TOKEN = os.getenv("HF_TOKEN")
+login(TOKEN)
+# This will download the model weights and instantiate the model on your machine.
+model: ESM3InferenceClient = ESM3.from_pretrained("esm3_sm_open_v1").to("cuda")
+def molecule(input_pdb):
+    mol = read_mol(input_pdb)
+    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>"""
+@spaces.GPU(duration=300)
+def prediction(prompt, temperature, do_structure, enable_roundtrip):
+    protein = ESMProtein(sequence=prompt)
+    protein = model.generate(protein, GenerationConfig(track="sequence", num_steps=8, temperature=temperature))
+    if do_structure == "Yes":
+        protein = model.generate(protein, GenerationConfig(track="structure", num_steps=8))
+        protein.to_pdb("./generation.pdb")
+        html = molecule("./generation.pdb")
+        if enable_roundtrip == "Yes":
+            seq = protein.sequence
+            protein.sequence = None
+            protein = model.generate(protein, GenerationConfig(track="sequence", num_steps=8))
+            protein.coordinates = None
+            protein = model.generate(protein, GenerationConfig(track="structure", num_steps=8))
+            protein.to_pdb("./round_tripped.pdb")
+            html1 = molecule("./round_tripped.pdb")
+            return seq, protein.sequemce, html, html1, "./generation.pdb", "./round_tripped.pdb"
+        else:
+            html1 = "<h3>Inverse folding and re-generation not enabled</h3>"
+            f = open("./round_tripped.pdb", "w")
+            f.write("\n")
+            f.close()
+            return protein.sequence, "Inverse folding and re-generation not enabled", html, html1, "./generation.pdb", "./round_tripped.pdb"
+    else:
+        f = open("./empty.pdb", "w")
+        f.write("\n")
+        f.close()
+        return protein.sequence, "Inverse folding and re-generation not enabled", "<h3>Structure reconstruction not enabled</h3>", "<h3>Inverse folding and re-generation not enabled</h3>", "./empty.pdb", "./empty.pdb"
+demo = gr.Interface(fn = prediction, inputs = [gr.Textbox(label="Masked protein sequence", info="Use '_' as masking character", value="___________________________________________________DQATSLRILNNGHAFNVEFDDSQDKAVLKGGPLDGTYRLIQFHFHWGSLDGQGSEHTVDKKKYAAELHLVHWNTKYGDFGKAVQQPDGLAVLGIFLKVGSAKPGLQKVVDVLDSIKTKGKSADFTNFDPRGLLPESLDYWTYPGSLTTPP___________________________________________________________"), gr.Slider(0,1,label="Temperature"), gr.Radio(["Yes", "No"], label="Reconstruct structure", info="Choose wheter to reconstruct structure or not"), gr.Radio(["Yes", "No"], label="Allow inverse-folding", info="Choose wether to allow double check of prediction with inverse folding (ONLY when 'Reconstruct structure' is set to 'Yes')")], outputs = [gr.Textbox(label="Originally predicted sequence"),gr.Textbox(label="Inverse folding predicted sequence"),gr.HTML(label="Predicted 3D structure"),gr.HTML(label="Inverse-folding predicted 3D structure"), Molecule3D("Predicted molecular structure"), Molecule3D(label="Inverse-folding predicted molecular structure")], title="""<h1 align='center'>Proteins with ESM</h1>
+    <h2 align='center'>Predict the whole sequence and 3D structure of masked protein sequences!</h2>
+    <h3 align='center'>Support this space with a ⭐ on <a href='https://github.com/AstraBert/proteins-w-esm'>GitHub</a></h3>
+    <h3 align='center'>Support Evolutionary Scale's ESM with a ⭐ on <a href='https://github.com/evolutionaryscale/esm'>GitHub</a></h3>""",
+examples = [["___________________________________________________DQATSLRILNNGHAFNVEFDDSQDKAVLKGGPLDGTYRLIQFHFHWGSLDGQGSEHTVDKKKYAAELHLVHWNTKYGDFGKAVQQPDGLAVLGIFLKVGSAKPGLQKVVDVLDSIKTKGKSADFTNFDPRGLLPESLDYWTYPGSLTTPP___________________________________________________________", 0.7, "No", "No"], ["__________________________________________________________AGQEEYSAMRDQYMRTGEGFLCVFAINNTKSFEDIHQYREQIKRVKDSDDVPMVLVGNKCDLAARTVESRQAQDLARSYGIPYIETSAKTRQGVEDAFYTLVRE___________________________", 0.2, "Yes", "No"], ["__________KTITLEVEPSDTIENVKAKIQDKEGIPPDQQRLIFAGKQLEDGRTLSDYNIQKESTLH________", 0.5, "Yes", "Yes"]], cache_examples=False)
+demo.launch()