Spaces:
Running
Running
| import re, os | |
| 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 | |
| from evodiff.conditional_generation import inpaint_simple, generate_scaffold | |
| 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, int(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, int(seq_len), batch_size=1, device='cpu') | |
| return generated_sequence | |
| def make_cond_seq(seq_len, msa_file, n_sequences, model_type): | |
| if model_type == "EvoDiff-MSA": | |
| checkpoint = MSA_OA_DM_MAXSUB() | |
| model, collater, tokenizer, scheme = checkpoint | |
| print(f"MSA File Path: {msa_file.name}") | |
| tokeinzed_sample, generated_sequence = generate_query_oadm_msa_simple(msa_file.name, model, tokenizer, int(n_sequences), seq_length=int(seq_len), device='cpu', selection_type='random') | |
| return generated_sequence | |
| def make_inpainted_idrs(sequence, start_idx, end_idx, model_type): | |
| if model_type == "EvoDiff-Seq": | |
| checkpoint = OA_DM_38M() | |
| model, collater, tokenizer, scheme = checkpoint | |
| sample, entire_sequence, generated_idr = inpaint_simple(model, sequence, int(start_idx), int(end_idx), tokenizer=tokenizer, device='cpu') | |
| generated_idr_output = { | |
| "original_sequence": sequence, | |
| "generated_sequence": entire_sequence, | |
| "original_region": sequence[start_idx:end_idx], | |
| "generated_region": generated_idr | |
| } | |
| return generated_idr_output | |
| def make_scaffold_motifs(pdb_code, start_idx, end_idx, scaffold_length, model_type): | |
| if model_type == "EvoDiff-Seq": | |
| checkpoint = OA_DM_38M() | |
| model, collater, tokenizer, scheme = checkpoint | |
| data_top_dir = '/home/user/.cache/huggingface/datasets/' | |
| os.makedirs(data_top_dir, exist_ok=True) | |
| # print("Folders in User Cache Directory:", os.listdir("/home/user/.cache")) | |
| start_idx = list(map(int, start_idx.strip('][').split(','))) | |
| end_idx = list(map(int, end_idx.strip('][').split(','))) | |
| generated_sequence, new_start_idx, new_end_idx = generate_scaffold(model, pdb_code, start_idx, end_idx, scaffold_length, data_top_dir, tokenizer, device='cpu') | |
| generated_scaffold_output = { | |
| "generated_sequence": generated_sequence, | |
| "new_start_index": new_start_idx, | |
| "new_end_index": new_end_idx | |
| } | |
| return generated_scaffold_output | |
| usg_app = gr.Interface( | |
| fn=make_uncond_seq, | |
| inputs=[ | |
| gr.Slider(10, 250, step=1, label = "Sequence Length"), | |
| gr.Dropdown(["EvoDiff-Seq-OADM 38M", "EvoDiff-D3PM-Uniform 38M"], value="EvoDiff-Seq-OADM 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, 250, label = "Sequence Length"), | |
| gr.File(file_types=["a3m"], label = "MSA File"), | |
| gr.Number(value=64, precision=0, label = "Number of Sequences to Sample"), | |
| gr.Dropdown(["EvoDiff-MSA"], value="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." | |
| ) | |
| idr_app = gr.Interface( | |
| fn=make_inpainted_idrs, | |
| inputs=[ | |
| gr.Textbox(value = "DQTERTVRSFEGRRTAPYLDSRNVLTIGYGHLLNRPGANKSWEGRLTSALPREFKQRLTELAASQLHETDVRLATARAQALYGSGAYFESVPVSLNDLWFDSVFNLGERKLLNWSGLRTKLESRDWGAAAKDLGRHTFGREPVSRRMAESMRMRRGIDLNHYNI", | |
| label = "Sequence"), | |
| gr.Number(value=20, precision=0, label = "Start Index"), | |
| gr.Number(value=50, precision=0, label = "End Index"), | |
| gr.Dropdown(["EvoDiff-Seq"], value="EvoDiff-Seq", type="value", label = "Model") | |
| ], | |
| outputs=["text"], | |
| title = "Inpainting IDRs", | |
| description="Inpaining a new region inside a given sequence using the `EvoDiff-Seq` model." | |
| ) | |
| scaffold_app = gr.Interface( | |
| fn=make_scaffold_motifs, | |
| inputs=[ | |
| gr.Textbox(value="1prw", label = "PDB Code"), | |
| gr.Textbox(value="[15, 51]", label = "Start Index (as list)"), | |
| gr.Textbox(value="[34, 70]", label = "End Index (as list)"), | |
| gr.Number(value=75, precision=0, label = "Scaffold Length"), | |
| gr.Dropdown(["EvoDiff-Seq", "EvoDiff-MSA"], value="EvoDiff-Seq", type="value", label = "Model") | |
| ], | |
| outputs=["text"], | |
| title = "Scaffolding functional motifs", | |
| description="Scaffolding a new functional motif inside a given PDB structure using the `EvoDiff-Seq` 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: Tuple, The Cloud Genomics Company [Colby T. Ford] | |
| """ | |
| ) | |
| with gr.Row(): | |
| gr.TabbedInterface([ | |
| usg_app, | |
| csg_app, | |
| idr_app#, | |
| # scaffold_app | |
| ], | |
| [ | |
| "Unconditional sequence generation", | |
| "Conditional generation", | |
| "Inpainting IDRs"#, | |
| # "Scaffolding functional motifs" | |
| ]) | |
| if __name__ == "__main__": | |
| edapp.launch() |