File size: 1,182 Bytes
fbdaace |
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 |
import gradio as gr
import torch
from transformers import EsmForProteinFolding, EsmTokenizer
import py3Dmol
# Load model
model = EsmForProteinFolding.from_pretrained("facebook/esmfold_v1")
tokenizer = EsmTokenizer.from_pretrained("facebook/esmfold_v1")
def predict_structure(sequence):
# Tokenize sequence
inputs = tokenizer(sequence, return_tensors="pt", add_special_tokens=False)
# Predict structure (ML happens here)
with torch.no_grad():
output = model(**inputs)
# Extract PDB string
pdb_str = output["pdb"]
# Visualize structure
viewer = py3Dmol.view(width=400, height=400)
viewer.addModel(pdb_str, "pdb")
viewer.setStyle({"cartoon": {"color": "spectrum"}})
viewer.zoomTo()
return viewer.show(), pdb_str
# Gradio UI
iface = gr.Interface(
fn=predict_structure,
inputs=gr.Textbox(label="Enter Amino Acid Sequence", placeholder="MKTAYIAKQRQISFVK..."),
outputs=[
gr.HTML(label="3D Structure"),
gr.Textbox(label="PDB Output")
],
title="Protein Structure Prediction",
description="Enter a protein sequence to predict its 3D structure using ESMFold."
)
iface.launch()
|