Spaces:

colbyford
/

evodiff

Sleeping

App Files Files Community

colbyford commited on Oct 6, 2023

Commit

6bc4f0e

1 Parent(s): d7676d2

Update app.py

Browse files

Files changed (1) hide show

app.py +20 -112

app.py CHANGED Viewed

@@ -7,66 +7,8 @@ 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
-import py3Dmol
-from colabfold.download import download_alphafold_params
-from colabfold.batch import run
-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"""<!DOCTYPE html><html></center> {output} </center></html>"""  # do not use ' in this input
-    return f"""<iframe height="500px" width="100%"  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>"""
-'''
-    return f"""<iframe  style="width: 100%; height:700px" 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 make_uncond_seq(seq_len, model_type, pred_structure):
     if model_type == "EvoDiff-Seq-OADM 38M":
         checkpoint = OA_DM_38M()
         model, collater, tokenizer, scheme = checkpoint
@@ -77,29 +19,17 @@ def make_uncond_seq(seq_len, model_type, pred_structure):
         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')
-    if pred_structure:
-        path_to_pdb = predict_protein(generated_sequence)
-        molhtml = display_pdb(path_to_pdb)
-        return generated_sequence, molhtml
-    else:
-        return generated_sequence, None
-def make_cond_seq(seq_len, msa_file, n_sequences, model_type, pred_structure):
     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, int(n_sequences), seq_length=int(seq_len), device='cpu', selection_type='random')
-    if pred_structure:
-        path_to_pdb = predict_protein(generated_sequence)
-        molhtml = display_pdb(path_to_pdb)
-        return generated_sequence, molhtml
-    else:
-        return generated_sequence, None
-def make_inpainted_idrs(sequence, start_idx, end_idx, model_type, pred_structure):
     if model_type == "EvoDiff-Seq":
         checkpoint = OA_DM_38M()
         model, collater, tokenizer, scheme = checkpoint
@@ -112,15 +42,9 @@ def make_inpainted_idrs(sequence, start_idx, end_idx, model_type, pred_structure
             "generated_region": generated_idr
         }
-    if pred_structure:
-        path_to_pdb = predict_protein(entire_sequence)
-        molhtml = display_pdb(path_to_pdb)
-        return generated_idr_output, molhtml
-    else:
-        return generated_idr_output, None
-def make_scaffold_motifs(pdb_code, start_idx, end_idx, scaffold_length, model_type, pred_structure):
     if model_type == "EvoDiff-Seq":
         checkpoint = OA_DM_38M()
         model, collater, tokenizer, scheme = checkpoint
@@ -135,26 +59,16 @@ def make_scaffold_motifs(pdb_code, start_idx, end_idx, scaffold_length, model_ty
             "new_end_index": new_end_idx
         }
-    return generated_scaffold_output, None
-    # if pred_structure:
-    #     # path_to_pdb = predict_protein(generated_sequence)
-    #     path_to_pdb = f"scaffolding-pdbs/{pdb_code}.pdb"
-    #     molhtml = display_pdb(path_to_pdb)
-    #     return generated_scaffold_output, molhtml
-    # else:
-    #     return generated_scaffold_output, None
 usg_app = gr.Interface(
             fn=make_uncond_seq,
             inputs=[
-                gr.Slider(10, 100, step=1, label = "Sequence Length"),
-                gr.Dropdown(["EvoDiff-Seq-OADM 38M", "EvoDiff-D3PM-Uniform 38M"], value="EvoDiff-Seq-OADM 38M", type="value", label = "Model"),
-                gr.Checkbox(value=False, label = "Predict Structure?", visible=False)
                 ],
             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."
@@ -163,15 +77,13 @@ usg_app = gr.Interface(
 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.Number(value=1, placeholder=1, precision=0, label = "Number of Sequences"),
-                gr.Dropdown(["EvoDiff-MSA"], value="EvoDiff-MSA", type="value", label = "Model"),
-                gr.Checkbox(value=False, label = "Predict Structure?", visible=False)
                 ],
             outputs=[
-                "text",
-                gr.HTML()
             ],
             # examples=[["https://github.com/microsoft/evodiff/raw/main/examples/example_files/bfd_uniclust_hits.a3m"]],
             title = "Conditional sequence generation",
@@ -184,12 +96,10 @@ idr_app = gr.Interface(
                 gr.Textbox(placeholder="DQTERTVRSFEGRRTAPYLDSRNVLTIGYGHLLNRPGANKSWEGRLTSALPREFKQRLTELAASQLHETDVRLATARAQALYGSGAYFESVPVSLNDLWFDSVFNLGERKLLNWSGLRTKLESRDWGAAAKDLGRHTFGREPVSRRMAESMRMRRGIDLNHYNI", label = "Sequence"),
                 gr.Number(value=20, placeholder=20, precision=0, label = "Start Index"),
                 gr.Number(value=50, placeholder=50, precision=0, label = "End Index"),
-                gr.Dropdown(["EvoDiff-Seq"], value="EvoDiff-Seq", type="value", label = "Model"),
-                gr.Checkbox(value=False, label = "Predict Structure?", visible=False)
                 ],
             outputs=[
-                "text",
-                gr.HTML()
             ],
             title = "Inpainting IDRs",
             description="Inpaining a new region inside a given sequence using the `EvoDiff-Seq` model."
@@ -202,12 +112,10 @@ scaffold_app = gr.Interface(
                 gr.Textbox(value="[15, 51]", placeholder="[15, 51]", label = "Start Index (as list)"),
                 gr.Textbox(value="[34, 70]", placeholder="[34, 70]", label = "End Index (as list)"),
                 gr.Number(value=75, placeholder=75, precision=0, label = "Scaffold Length"),
-                gr.Dropdown(["EvoDiff-Seq", "EvoDiff-MSA"], value="EvoDiff-Seq", type="value", label = "Model"),
-                gr.Checkbox(value=False, label = "Predict Structure?", visible=False)
                 ],
             outputs=[
-                "text",
-                gr.HTML()
             ],
             title = "Scaffolding functional motifs",
             description="Scaffolding a new functional motif inside a given PDB structure using the `EvoDiff-Seq` model."

 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
         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
         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
             "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
             "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, placeholder=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",
                 gr.Textbox(placeholder="DQTERTVRSFEGRRTAPYLDSRNVLTIGYGHLLNRPGANKSWEGRLTSALPREFKQRLTELAASQLHETDVRLATARAQALYGSGAYFESVPVSLNDLWFDSVFNLGERKLLNWSGLRTKLESRDWGAAAKDLGRHTFGREPVSRRMAESMRMRRGIDLNHYNI", label = "Sequence"),
                 gr.Number(value=20, placeholder=20, precision=0, label = "Start Index"),
                 gr.Number(value=50, placeholder=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."
                 gr.Textbox(value="[15, 51]", placeholder="[15, 51]", label = "Start Index (as list)"),
                 gr.Textbox(value="[34, 70]", placeholder="[34, 70]", label = "End Index (as list)"),
                 gr.Number(value=75, placeholder=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."