File size: 6,023 Bytes
884b6ae
 
 
 
8165629
884b6ae
 
 
 
 
 
 
 
 
 
6e9c18e
884b6ae
8165629
 
 
 
 
 
 
 
884b6ae
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
bbe157e
884b6ae
 
0cf7233
 
 
 
bbe157e
 
 
 
 
 
 
 
0cf7233
 
 
 
 
884b6ae
469bcfa
 
 
 
 
 
 
 
 
 
 
 
 
 
bbe157e
884b6ae
 
bbe157e
623a8df
0cf7233
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
from huggingface_hub import login
from esm.models.esm3 import ESM3
from esm.sdk.api import ESM3InferenceClient, ESMProtein, GenerationConfig
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("cpu")

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 = (
        """<!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>"""

def prediction(prompt, temperature, do_structure):
    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")
        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, "./round_tripped.pdb", "./generation.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"

reps =    [
    {
      "model": 0,
      "chain": "",
      "resname": "",
      "style": "stick",
      "color": "whiteCarbon",
      "residue_range": "",
      "around": 0,
      "byres": False,
      "visible": False
    }
]

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, allowing also inverse folding-powered double check")], outputs = [gr.Textbox(label="Originally predicted sequence", show_copy_button=True),gr.Textbox(label="Inverse folding predicted sequence", show_copy_button=True),gr.HTML(label="Predicted 3D structure"),gr.HTML(label="Inverse-folding predicted 3D structure"), Molecule3D(label="Inverse-folding predicted molecular structure", reps=reps), Molecule3D(label="Predicted molecular structure", reps=reps)], 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"], ["__________________________________________________________AGQEEYSAMRDQYMRTGEGFLCVFAINNTKSFEDIHQYREQIKRVKDSDDVPMVLVGNKCDLAARTVESRQAQDLARSYGIPYIETSAKTRQGVEDAFYTLVRE___________________________", 0.2, "Yes"], ["__________KTITLEVEPSDTIENVKAKIQDKEGIPPDQQRLIFAGKQLEDGRTLSDYNIQKESTLH________", 0.5, "Yes",]], cache_examples=False)

demo.launch()