import re
from pathlib import Path
import gradio as gr
from evodiff.pretrained import OA_DM_38M, D3PM_UNIFORM_38M, MSA_OA_DM_MAXSUB
from evodiff.generate import generate_oaardm, generate_d3pm
from evodiff.generate_msa import generate_query_oadm_msa_simple
import py3Dmol
from colabfold.download import download_alphafold_params, default_data_dir
from colabfold.utils import setup_logging
from colabfold.batch import get_queries, run, set_model_type
from colabfold.plot import plot_msa_v2
def a3m_file(file):
return "tmp.a3m"
def predict_protein(sequence):
download_alphafold_params("alphafold2_ptm", Path("."))
results = run(
queries=[('evodiff_protein',sequence, None)],
result_dir='evodiff_protein',
use_templates=False,
num_relax=0,
msa_mode="mmseqs2_uniref_env",
model_type="alphafold2_ptm",
num_models=1,
num_recycles=1,
model_order=[1],
is_complex=False,
data_dir=Path("."),
keep_existing_results=False,
rank_by="auto",
stop_at_score=float(100),
zip_results=False,
user_agent="colabfold/google-colab-main",
)
return f"evodiff_protein/evodiff_protein_unrelaxed_rank_001_alphafold2_ptm_model_1_seed_000.pdb"
def display_pdb(path_to_pdb):
'''
#function to display pdb in py3dmol
SOURCE: https://huggingface.co/spaces/merle/PROTEIN_GENERATOR/blob/main/app.py
'''
pdb = open(path_to_pdb, "r").read()
view = py3Dmol.view(width=500, height=500)
view.addModel(pdb, "pdb")
view.setStyle({'model': -1}, {"cartoon": {'colorscheme':{'prop':'b','gradient':'roygb','min':0,'max':1}}})#'linear', 'min': 0, 'max': 1, 'colors': ["#ff9ef0","#a903fc",]}}})
view.zoomTo()
output = view._make_html().replace("'", '"')
print(view._make_html())
x = f""" {output} """ # do not use ' in this input
return f""""""
'''
return f""""""
'''
def make_uncond_seq(seq_len, model_type):
if model_type == "EvoDiff-Seq-OADM 38M":
checkpoint = OA_DM_38M()
model, collater, tokenizer, scheme = checkpoint
tokeinzed_sample, generated_sequence = generate_oaardm(model, tokenizer, seq_len, batch_size=1, device='cpu')
if model_type == "EvoDiff-D3PM-Uniform 38M":
checkpoint = D3PM_UNIFORM_38M(return_all=True)
model, collater, tokenizer, scheme, timestep, Q_bar, Q = checkpoint
tokeinzed_sample, generated_sequence = generate_d3pm(model, tokenizer, Q, Q_bar, timestep, seq_len, batch_size=1, device='cpu')
path_to_pdb = predict_protein(generated_sequence)
molhtml = display_pdb(path_to_pdb)
return generated_sequence, molhtml
def make_cond_seq(seq_len, msa_file, model_type):
if model_type == "EvoDiff-MSA":
checkpoint = MSA_OA_DM_MAXSUB()
model, collater, tokenizer, scheme = checkpoint
tokeinzed_sample, generated_sequence = generate_query_oadm_msa_simple(msa_file.name, model, tokenizer, n_sequences=64, seq_length=seq_len, device='cpu', selection_type='random')
path_to_pdb = predict_protein(generated_sequence)
molhtml = display_pdb(path_to_pdb)
return generated_sequence, molhtml
usg_app = gr.Interface(
fn=make_uncond_seq,
inputs=[
gr.Slider(10, 100, label = "Sequence Length"),
gr.Dropdown(["EvoDiff-Seq-OADM 38M", "EvoDiff-D3PM-Uniform 38M"], type="value", label = "Model")
],
outputs=[
"text",
gr.HTML()
],
title = "Unconditional sequence generation",
description="Generate a sequence with `EvoDiff-Seq-OADM 38M` (smaller/faster) or `EvoDiff-D3PM-Uniform 38M` (larger/slower) models."
)
csg_app = gr.Interface(
fn=make_cond_seq,
inputs=[
gr.Slider(10, 100, label = "Sequence Length"),
gr.File(file_types=["a3m"], label = "MSA File"),
gr.Dropdown(["EvoDiff-MSA"], type="value", label = "Model")
],
outputs=[
"text",
gr.HTML()
],
# examples=[["https://github.com/microsoft/evodiff/raw/main/examples/example_files/bfd_uniclust_hits.a3m"]],
title = "Conditional sequence generation",
description="Evolutionary guided sequence generation with the `EvoDiff-MSA` model."
)
with gr.Blocks() as edapp:
with gr.Row():
gr.Markdown(
"""
# EvoDiff
## Generation of protein sequences and evolutionary alignments via discrete diffusion models
Created By: Microsoft Research [Sarah Alamdari, Nitya Thakkar, Rianne van den Berg, Alex X. Lu, Nicolo Fusi, ProfileAva P. Amini, and Kevin K. Yang]
Spaces App By: [Colby T. Ford](httos://github.com/colbyford)
"""
)
with gr.Row():
gr.TabbedInterface([usg_app, csg_app], ["Unconditional sequence generation", "Conditional generation"])
if __name__ == "__main__":
edapp.launch()