File size: 4,627 Bytes
a444e27 |
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 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 |
"""Folding Studio Demo App."""
import logging
import gradio as gr
from folding_studio_data_models import FoldingModel
from gradio_molecule3d import Molecule3D
from folding_studio_demo.predict import predict
logger = logging.getLogger(__name__)
MOLECULE_REPS = [
{
"model": 0,
"chain": "",
"resname": "",
"style": "cartoon",
"color": "alphafold",
# "residue_range": "",
"around": 0,
"byres": False,
# "visible": False,
# "opacity": 0.5
}
]
DEFAULT_PROTEIN_SEQ = ">protein description\nMALWMRLLPLLALLALWGPDPAAA"
MODEL_CHOICES = [
# ("AlphaFold2", FoldingModel.AF2),
# ("OpenFold", FoldingModel.OPENFOLD),
# ("SoloSeq", FoldingModel.SOLOSEQ),
("Boltz-1", FoldingModel.BOLTZ),
("Chai-1", FoldingModel.CHAI),
("Protenix", FoldingModel.PROTENIX),
]
def sequence_input() -> gr.Textbox:
"""Sequence input component.
Returns:
gr.Textbox: Sequence input component
"""
sequence = gr.Textbox(
label="Protein Sequence",
value=DEFAULT_PROTEIN_SEQ,
lines=2,
placeholder="Enter a protein sequence or upload a FASTA file",
)
file_input = gr.File(
label="Upload a FASTA file",
file_types=[".fasta", ".fa"],
)
def _process_file(file: gr.File | None) -> gr.Textbox:
if file is None:
return gr.Textbox()
try:
with open(file.name, "r") as f:
content = f.read().strip()
return gr.Textbox(value=content)
except Exception as e:
logger.error(f"Error reading file: {e}")
return gr.Textbox()
file_input.change(fn=_process_file, inputs=[file_input], outputs=[sequence])
return sequence
def simple_prediction(api_key: str) -> None:
"""Simple prediction tab.
Args:
api_key (str): Folding Studio API key
"""
gr.Markdown(
"""
### Predict a Protein Structure
It will be run in the background and the results will be displayed in the output section.
The output will contain the protein structure and the pLDDT plot.
Select a model to run the inference with and enter a protein sequence or upload a FASTA file.
"""
)
with gr.Row():
dropdown = gr.Dropdown(
label="Model",
choices=MODEL_CHOICES,
scale=0,
value=FoldingModel.BOLTZ,
)
with gr.Column():
sequence = sequence_input()
predict_btn = gr.Button("Predict")
with gr.Row():
mol_output = Molecule3D(label="Protein Structure", reps=MOLECULE_REPS)
metrics_plot = gr.Plot(label="pLDDT")
predict_btn.click(
fn=predict,
inputs=[sequence, api_key, dropdown],
outputs=[mol_output, metrics_plot],
)
def model_comparison(api_key: str) -> None:
"""Model comparison tab.
Args:
api_key (str): Folding Studio API key
"""
with gr.Row():
model = gr.Dropdown(
label="Model",
choices=MODEL_CHOICES,
multiselect=True,
scale=0,
min_width=300,
value=[FoldingModel.BOLTZ, FoldingModel.CHAI, FoldingModel.PROTENIX],
)
with gr.Column():
sequence = sequence_input()
predict_btn = gr.Button("Compare Models")
with gr.Row():
mol_output = Molecule3D(label="Protein Structure", reps=MOLECULE_REPS)
metrics_plot = gr.Plot(label="pLDDT")
predict_btn.click(
fn=predict,
inputs=[sequence, api_key, model],
outputs=[mol_output, metrics_plot],
)
def __main__():
with gr.Blocks(title="Folding Studio Demo") as demo:
gr.Markdown(
"""
# Folding Studio: Harness the Power of Protein Folding 🧬
Folding Studio is a platform for protein structure prediction.
It uses the latest AI-powered folding models to predict the structure of a protein.
Available models are : AlphaFold2, OpenFold, SoloSeq, Boltz-1, Chai and Protenix.
## API Key
To use the Folding Studio API, you need to provide an API key.
You can get your API key by asking to the Folding Studio team.
"""
)
api_key = gr.Textbox(label="Folding Studio API Key", type="password")
gr.Markdown("## Demo Usage")
with gr.Tab("🚀 Simple Prediction"):
simple_prediction(api_key)
with gr.Tab("📊 Model Comparison"):
model_comparison(api_key)
demo.launch()
|