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 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')
return generated_sequence
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')
return generated_sequence
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 predict_protein(sequence):
model_type = "alphafold2_ptm"
download_alphafold_params(model_type, Path("."))
results = run(
queries=queries,
result_dir=result_dir,
use_templates=use_templates,
custom_template_path=custom_template_path,
num_relax=0,
msa_mode=msa_mode,
model_type=model_type,
num_models=1,
num_recycles=1,
recycle_early_stop_tolerance=recycle_early_stop_tolerance,
num_seeds=num_seeds,
use_dropout=use_dropout,
model_order=[1],
is_complex=False,
data_dir=Path("."),
keep_existing_results=False,
rank_by="auto",
pair_mode=pair_mode,
pairing_strategy=pairing_strategy,
stop_at_score=float(100),
prediction_callback=prediction_callback,
dpi=dpi,
zip_results=False,
save_all=save_all,
max_msa=max_msa,
use_cluster_profile=use_cluster_profile,
input_features_callback=input_features_callback,
save_recycles=save_recycles,
user_agent="colabfold/google-colab-main",
)
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",
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",
# 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()