from huggingface_hub import login
from esm.models.esm3 import ESM3
from esm.sdk.api import ESM3InferenceClient, ESMProtein, GenerationConfig
import spaces
import os
import gradio as gr
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 read_mol(molpath):
with open(molpath, "r") as fp:
lines = fp.readlines()
mol = ""
for l in lines:
mol += l
return mol
def molecule(input_pdb):
mol = read_mol(input_pdb)
x = (
"""
"""
)
return f""""""
@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.sequence, html, html1, "./generation.pdb", "./round_tripped.pdb"
else:
html1 = "Inverse folding and re-generation not enabled
"
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", "Structure reconstruction not enabled
", "Inverse folding and re-generation not enabled
", "./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="""Proteins with ESM
Predict the whole sequence and 3D structure of masked protein sequences!
Support this space with a ⭐ on GitHub
Support Evolutionary Scale's ESM with a ⭐ on GitHub
""", examples = [["___________________________________________________DQATSLRILNNGHAFNVEFDDSQDKAVLKGGPLDGTYRLIQFHFHWGSLDGQGSEHTVDKKKYAAELHLVHWNTKYGDFGKAVQQPDGLAVLGIFLKVGSAKPGLQKVVDVLDSIKTKGKSADFTNFDPRGLLPESLDYWTYPGSLTTPP___________________________________________________________", 0.7, "No", "No"], ["__________________________________________________________AGQEEYSAMRDQYMRTGEGFLCVFAINNTKSFEDIHQYREQIKRVKDSDDVPMVLVGNKCDLAARTVESRQAQDLARSYGIPYIETSAKTRQGVEDAFYTLVRE___________________________", 0.2, "Yes", "No"], ["__________KTITLEVEPSDTIENVKAKIQDKEGIPPDQQRLIFAGKQLEDGRTLSDYNIQKESTLH________", 0.5, "Yes", "Yes"]], cache_examples=False)
demo.launch()