File size: 5,989 Bytes
93a77af 178b187 33e1d55 9329e39 a418c62 6e3dfae a418c62 9329e39 93a77af 9329e39 4d6e83e 33e1d55 9329e39 4d6e83e 9329e39 4d6e83e 9329e39 aad23a2 4d6e83e 93a77af 178b187 64c1665 178b187 9329e39 93a77af c998ccf 93a77af 8f7fecd 9329e39 8f7fecd f01f3e2 93a77af 178b187 93a77af 7f1194a 93a77af 8f7fecd 93a77af 8f7fecd 93a77af f01f3e2 9329e39 8f7fecd f01f3e2 93a77af 8f7fecd |
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 169 170 171 172 173 174 175 176 177 178 179 180 |
import gradio as gr
import sys
import random
import os
import pandas as pd
import torch
from torch.utils.data import DataLoader
from transformers import AutoTokenizer
sys.path.append("/home/user/app/")
from scripts.foldseek_util import get_struc_seq
from scripts.utils import seed_everything
from scripts.models import PLTNUM_PreTrainedModel
from scripts.datasets import PLTNUMDataset
class Config:
batch_size = 2
use_amp = False
num_workers = 1
max_length = 512
used_sequence = "left"
padding_side = "right"
task = "classification"
sequence_col = "sequence"
# Assuming 'predict_stability' is your function that predicts protein stability
def predict_stability(cfg, model_choice, organism_choice, pdb_file=None, sequence=None):
# Check if pdb_file is provided
if pdb_file:
pdb_path = pdb_file.name # Get the path of the uploaded PDB file
os.system("chmod 777 bin/foldseek")
sequences = get_foldseek_seq(pdb_path)
if not sequences:
return "Failed to extract sequence from the PDB file."
if model_choice == "SaProt":
sequence = sequences[2]
else:
sequence = sequences[0]
if organism_choice == "Human":
cell_line = "HeLa"
else:
cell_line = "NIH3T3"
# If sequence is provided directly
if sequence:
cfg.model = f"sagawa/PLTNUM-{model_choice}-{cell_line}"
cfg.architecture = model_choice
cfg.model_path = f"sagawa/PLTNUM-{model_choice}-{cell_line}"
output = predict(cfg, sequence)
return f"Predicted Stability using {model_choice} for {organism_choice}: Example Output with sequence {output}..."
else:
return "No valid input provided."
def get_foldseek_seq(pdb_path):
parsed_seqs = get_struc_seq(
"bin/foldseek",
pdb_path,
["A"],
process_id=random.randint(0, 10000000),
)["A"]
return parsed_seqs
def predict(cfg, sequence):
cfg.token_length = 2 if cfg.architecture == "SaProt" else 1
cfg.device = "cuda" if torch.cuda.is_available() else "cpu"
if cfg.used_sequence == "both":
cfg.max_length += 1
seed_everything(cfg.seed)
df = pd.DataFrame({cfg.sequence_col: [sequence]})
tokenizer = AutoTokenizer.from_pretrained(
cfg.model_path, padding_side=cfg.padding_side
)
cfg.tokenizer = tokenizer
dataset = PLTNUMDataset(cfg, df, train=False)
dataloader = DataLoader(
dataset,
batch_size=cfg.batch_size,
shuffle=False,
num_workers=cfg.num_workers,
pin_memory=True,
drop_last=False,
)
model = PLTNUM_PreTrainedModel.from_pretrained(cfg.model_path, cfg=cfg)
model.to(cfg.device)
# predictions = predict_fn(loader, model, cfg)
model.eval()
predictions = []
for inputs, _ in dataloader:
inputs = inputs.to(cfg.device)
with torch.no_grad():
with torch.amp.autocast(enabled=cfg.use_amp):
preds = (
torch.sigmoid(model(inputs))
if cfg.task == "classification"
else model(inputs)
)
predictions += preds.cpu().tolist()
outputs = {}
outputs["raw prediction values"] = predictions
outputs["binary prediction values"] = [1 if x > 0.5 else 0 for x in predictions]
return outputs
# Gradio Interface
with gr.Blocks() as demo:
gr.Markdown(
"""
# PLTNUM: Protein LifeTime Neural Model
**Predict the protein half-life from its sequence or PDB file.**
"""
)
gr.Image("https://github.com/sagawatatsuya/PLTNUM/blob/main/model-image.png?raw=true", label="Model Image")
# Model and Organism selection in the same row to avoid layout issues
with gr.Row():
model_choice = gr.Radio(
choices=["SaProt", "ESM2"],
label="Select PLTNUM's base model.",
value="SaProt"
)
organism_choice = gr.Radio(
choices=["Mouse", "Human"],
label="Select the target organism.",
value="Mouse"
)
with gr.Tabs():
with gr.TabItem("Upload PDB File"):
gr.Markdown("### Upload your PDB file:")
pdb_file = gr.File(label="Upload PDB File")
predict_button = gr.Button("Predict Stability")
prediction_output = gr.Textbox(label="Stability Prediction", interactive=False)
predict_button.click(fn=predict_stability, inputs=[model_choice, organism_choice, pdb_file], outputs=prediction_output)
with gr.TabItem("Enter Protein Sequence"):
gr.Markdown("### Enter the protein sequence:")
sequence = gr.Textbox(
label="Protein Sequence",
placeholder="Enter your protein sequence here...",
lines=8,
)
predict_button = gr.Button("Predict Stability")
prediction_output = gr.Textbox(label="Stability Prediction", interactive=False)
predict_button.click(fn=predict_stability, inputs=[model_choice, organism_choice, sequence], outputs=prediction_output)
gr.Markdown(
"""
### How to Use:
- **Select Model**: Choose between 'SaProt' or 'ESM2' for your prediction.
- **Select Organism**: Choose between 'Mouse' or 'Human'.
- **Upload PDB File**: Choose the 'Upload PDB File' tab and upload your file.
- **Enter Sequence**: Alternatively, switch to the 'Enter Protein Sequence' tab and input your sequence.
- **Predict**: Click 'Predict Stability' to receive the prediction.
"""
)
gr.Markdown(
"""
### About the Tool
This tool allows researchers and scientists to predict the stability of proteins using advanced algorithms. It supports both PDB file uploads and direct sequence input.
"""
)
demo.launch()
|