Spaces:
Running on Zero
Running on Zero
| import gradio as gr | |
| import spaces | |
| import torch | |
| import torch.nn as nn | |
| from huggingface_hub import hf_hub_download | |
| from transformers import EsmModel, EsmTokenizer | |
| from dataset import PROTEIN_DATASET, PROTEIN_MAP | |
| # ────────────────────────────────────────────────────────────────────────── | |
| # Model config | |
| # ────────────────────────────────────────────────────────────────────────── | |
| HF_REPO_ID = "PypCoder/SERAPH" | |
| WEIGHTS_FILE = "SERAPH.pth" | |
| ESM_MODEL_ID = "facebook/esm2_t6_8M_UR50D" | |
| IDX_TO_LABEL = {0: 'H', 1: 'E', 2: 'C'} | |
| LABEL_NAME = {'H': 'Alpha Helix', 'E': 'Beta Sheet', 'C': 'Coil / Loop'} | |
| CUSTOM_LABEL = "✎ Custom Sequence" | |
| PROTEIN_CHOICES = [CUSTOM_LABEL] + [p["name"] for p in PROTEIN_DATASET] | |
| class SERAPH(nn.Module): | |
| def __init__(self, esm_model, conv_channels=256, kernel_size=7, | |
| lstm_hidden=256, num_classes=3, dropout=0.3, freeze_esm=True): | |
| super().__init__() | |
| self.esm = esm_model | |
| if freeze_esm: | |
| for param in self.esm.encoder.layer[:-2].parameters(): | |
| param.requires_grad = False | |
| esm_embed_dim = self.esm.config.hidden_size | |
| self.conv = nn.Conv1d(esm_embed_dim, conv_channels, kernel_size=kernel_size, padding=kernel_size // 2) | |
| self.bn = nn.BatchNorm1d(conv_channels) | |
| self.dropout = nn.Dropout(dropout) | |
| self.bilstm = nn.LSTM(conv_channels, lstm_hidden, num_layers=2, batch_first=True, bidirectional=True) | |
| self.fc = nn.Linear(lstm_hidden * 2, num_classes) | |
| def forward(self, input_ids, attention_mask=None): | |
| x = self.esm(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state | |
| x = x.transpose(1, 2) | |
| x = torch.relu(self.bn(self.conv(x))) | |
| x = self.dropout(x) | |
| x = x.transpose(1, 2) | |
| x, _ = self.bilstm(x) | |
| x = self.dropout(x) | |
| return self.fc(x) | |
| print("Loading ESM2 backbone...") | |
| esm = EsmModel.from_pretrained(ESM_MODEL_ID) | |
| tokenizer = EsmTokenizer.from_pretrained(ESM_MODEL_ID) | |
| print("Downloading SERAPH weights...") | |
| weights_path = hf_hub_download(repo_id=HF_REPO_ID, filename=WEIGHTS_FILE) | |
| checkpoint = torch.load(weights_path, map_location="cpu") | |
| model = SERAPH(esm_model=esm) | |
| model.load_state_dict(checkpoint["model_state_dict"]) | |
| model.eval() | |
| print("SERAPH ready.") | |
| # ────────────────────────────────────────────────────────────────────────── | |
| # Inference | |
| # ────────────────────────────────────────────────────────────────────────── | |
| def run_model(sequence: str) -> str: | |
| tokens = tokenizer(sequence, return_tensors="pt", truncation=True, max_length=512) | |
| with torch.no_grad(): | |
| output = model(input_ids=tokens["input_ids"], attention_mask=tokens["attention_mask"]) | |
| preds = output.argmax(dim=-1)[0] | |
| labels = [IDX_TO_LABEL[p.item()] for p in preds[1:-1]] | |
| return "".join(labels) | |
| def render_alignment(sequence: str, prediction: str, true_ss: str = None) -> str: | |
| """Builds the responsive residue-by-residue alignment strip.""" | |
| cols = [] | |
| for i, aa in enumerate(sequence): | |
| pred_cls = f"ss-{prediction[i].lower()}" | |
| mismatch = " mismatch" if true_ss and i < len(true_ss) and prediction[i] != true_ss[i] else "" | |
| true_block = f'<div class="ss-block ss-{true_ss[i].lower()} true-row"></div>' if true_ss else "" | |
| cols.append( | |
| f'<div class="residue-col">' | |
| f'<span class="aa">{aa}</span>' | |
| f'<div class="ss-block {pred_cls}{mismatch}"></div>' | |
| f'{true_block}' | |
| f'</div>' | |
| ) | |
| return f'<div class="alignment-strip">{"".join(cols)}</div>' | |
| def render_legend(show_true: bool) -> str: | |
| rows = f""" | |
| <div class="legend"> | |
| <span class="legend-item"><span class="swatch ss-h"></span>Helix (H)</span> | |
| <span class="legend-item"><span class="swatch ss-e"></span>Sheet (E)</span> | |
| <span class="legend-item"><span class="swatch ss-c"></span>Coil (C)</span> | |
| </div> | |
| """ | |
| if show_true: | |
| rows += '<div class="legend-note">Top block = predicted · bottom block = ground truth · red outline = mismatch</div>' | |
| return rows | |
| def render_stats(prediction: str, true_ss: str = None) -> str: | |
| n = len(prediction) | |
| h, e, c = prediction.count("H"), prediction.count("E"), prediction.count("C") | |
| bars = f""" | |
| <div class="stat-bars"> | |
| <div class="stat-row"><span class="stat-label">Helix (H)</span><div class="bar-track"><div class="bar-fill ss-h" style="width:{h/n*100:.1f}%"></div></div><span class="stat-pct">{h/n*100:.1f}%</span></div> | |
| <div class="stat-row"><span class="stat-label">Sheet (E)</span><div class="bar-track"><div class="bar-fill ss-e" style="width:{e/n*100:.1f}%"></div></div><span class="stat-pct">{e/n*100:.1f}%</span></div> | |
| <div class="stat-row"><span class="stat-label">Coil (C)</span><div class="bar-track"><div class="bar-fill ss-c" style="width:{c/n*100:.1f}%"></div></div><span class="stat-pct">{c/n*100:.1f}%</span></div> | |
| </div> | |
| """ | |
| accuracy_html = "" | |
| if true_ss and len(true_ss) == n: | |
| matches = sum(1 for a, b in zip(prediction, true_ss) if a == b) | |
| acc = matches / n * 100 | |
| accuracy_html = f""" | |
| <div class="accuracy-badge"> | |
| <span class="accuracy-label">Q3 Accuracy vs. known structure</span> | |
| <span class="accuracy-value">{acc:.1f}%</span> | |
| </div> | |
| """ | |
| return f'<div class="stats-panel">{bars}{accuracy_html}</div>' | |
| def predict(sequence: str, selected_name: str): | |
| sequence = (sequence or "").upper().strip() | |
| if not sequence: | |
| return '<div class="placeholder-msg">Enter or select a sequence, then hit Predict.</div>' | |
| if len(sequence) > 512: | |
| sequence = sequence[:512] | |
| prediction = run_model(sequence) | |
| true_ss = None | |
| if selected_name and selected_name != CUSTOM_LABEL: | |
| entry = PROTEIN_MAP.get(selected_name) | |
| if entry and len(entry["true_ss"]) == len(prediction): | |
| true_ss = entry["true_ss"] | |
| alignment = render_alignment(sequence, prediction, true_ss) | |
| legend = render_legend(true_ss is not None) | |
| stats = render_stats(prediction, true_ss) | |
| return f""" | |
| <div class="result-card"> | |
| <div class="result-header"> | |
| <span>Predicted Structure</span> | |
| <span class="result-length">{len(sequence)} residues</span> | |
| </div> | |
| {alignment} | |
| {legend} | |
| {stats} | |
| </div> | |
| """ | |
| def load_preset(selected_name: str): | |
| if not selected_name or selected_name == CUSTOM_LABEL: | |
| return "", '<div class="info-card empty">Pick a preset above to see its background, or paste your own sequence.</div>' | |
| p = PROTEIN_MAP[selected_name] | |
| info_html = f""" | |
| <div class="info-card"> | |
| <div class="info-title">{p['name']}</div> | |
| <p class="info-desc">{p['description']}</p> | |
| <div class="fun-fact">💡 {p['fun_fact']}</div> | |
| </div> | |
| """ | |
| return p["sequence"], info_html | |
| # ────────────────────────────────────────────────────────────────────────── | |
| # Theme + CSS (matches playground.html design language) | |
| # ────────────────────────────────────────────────────────────────────────── | |
| THEME = gr.themes.Base( | |
| font=[gr.themes.GoogleFont("Plus Jakarta Sans"), "sans-serif"], | |
| font_mono=[gr.themes.GoogleFont("JetBrains Mono"), "monospace"], | |
| ).set( | |
| body_background_fill="#060608", | |
| body_background_fill_dark="#060608", | |
| body_text_color="#f5f5f7", | |
| body_text_color_dark="#f5f5f7", | |
| background_fill_primary="rgba(255,255,255,0.025)", | |
| background_fill_primary_dark="rgba(255,255,255,0.025)", | |
| background_fill_secondary="rgba(255,255,255,0.025)", | |
| block_background_fill="rgba(255,255,255,0.025)", | |
| block_background_fill_dark="rgba(255,255,255,0.025)", | |
| block_border_color="rgba(255,255,255,0.08)", | |
| block_border_color_dark="rgba(255,255,255,0.08)", | |
| block_label_text_color="#8e8e93", | |
| block_label_text_color_dark="#8e8e93", | |
| block_title_text_color="#f5f5f7", | |
| body_text_color_subdued="#8e8e93", | |
| input_background_fill="rgba(255,255,255,0.03)", | |
| input_background_fill_dark="rgba(255,255,255,0.03)", | |
| input_border_color="rgba(255,255,255,0.08)", | |
| input_border_color_dark="rgba(255,255,255,0.08)", | |
| button_primary_background_fill="#f5f5f7", | |
| button_primary_background_fill_hover="#ffffff", | |
| button_primary_text_color="#000000", | |
| button_secondary_background_fill="rgba(255,255,255,0.05)", | |
| button_secondary_background_fill_hover="rgba(255,255,255,0.09)", | |
| button_secondary_text_color="#f5f5f7", | |
| button_secondary_border_color="rgba(255,255,255,0.08)", | |
| border_color_primary="rgba(255,255,255,0.08)", | |
| color_accent_soft="rgba(255,255,255,0.05)", | |
| ) | |
| CSS = """ | |
| @import url('https://fonts.googleapis.com/css2?family=Space+Grotesk:wght@600;700&family=JetBrains+Mono:wght@400;500&display=swap'); | |
| :root{ | |
| --struct-helix:#7c9eff; | |
| --struct-sheet:#f59e0b; | |
| --struct-coil:rgba(255,255,255,0.18); | |
| --text-muted:#8e8e93; | |
| --text-dim:#55555a; | |
| --border-subtle:rgba(255,255,255,0.08); | |
| } | |
| .gradio-container{ max-width: 1020px !important; margin: 0 auto !important; } | |
| #header-row{ display:flex; justify-content:space-between; align-items:center; flex-wrap:wrap; gap:12px; margin-bottom: 6px; } | |
| #back-btn{ | |
| display:inline-flex; align-items:center; gap:8px; text-decoration:none; | |
| background: rgba(255,255,255,0.05); color:#f5f5f7; border:1px solid var(--border-subtle); | |
| padding:8px 16px; border-radius:99px; font-size:0.82rem; font-weight:500; | |
| transition: all 0.2s ease; white-space:nowrap; | |
| } | |
| #back-btn:hover{ background: rgba(255,255,255,0.09); border-color: rgba(255,255,255,0.22); } | |
| .hero-title{ | |
| font-family:'Space Grotesk', sans-serif; font-weight:700; | |
| font-size: clamp(1.9rem, 4.5vw, 2.8rem); letter-spacing:-0.02em; line-height:1.1; | |
| background: linear-gradient(180deg,#ffffff 0%, rgba(255,255,255,0.7) 100%); | |
| -webkit-background-clip:text; -webkit-text-fill-color:transparent; margin: 4px 0 2px 0; | |
| } | |
| .hero-subtitle{ color: var(--text-muted); font-size:0.95rem; max-width:640px; margin-bottom: 8px; } | |
| .info-card{ | |
| background: rgba(255,255,255,0.025); border:1px solid var(--border-subtle); border-radius:14px; | |
| padding:16px 18px; height:100%; | |
| } | |
| .info-card.empty{ display:flex; align-items:center; color: var(--text-dim); font-size:0.85rem; } | |
| .info-title{ font-family:'Space Grotesk', sans-serif; font-weight:700; font-size:1.02rem; margin-bottom:6px; } | |
| .info-desc{ color: var(--text-muted); font-size:0.85rem; line-height:1.5; margin-bottom:10px; } | |
| .fun-fact{ | |
| font-size:0.82rem; color:#f5f5f7; background: rgba(255,255,255,0.04); | |
| border-left:2px solid var(--struct-sheet); padding:8px 10px; border-radius:6px; line-height:1.5; | |
| } | |
| .placeholder-msg{ color: var(--text-dim); font-size:0.85rem; padding: 24px 8px; text-align:center; } | |
| .result-card{ border:1px solid var(--border-subtle); border-radius:14px; padding:18px; background: rgba(255,255,255,0.02); } | |
| .result-header{ | |
| display:flex; justify-content:space-between; align-items:baseline; font-family:'Space Grotesk', sans-serif; | |
| font-weight:700; font-size:1rem; margin-bottom:14px; | |
| } | |
| .result-length{ font-family:'JetBrains Mono', monospace; font-weight:400; font-size:0.75rem; color: var(--text-muted); } | |
| .alignment-strip{ display:flex; flex-wrap:wrap; gap:2px; margin-bottom:14px; max-height: 320px; overflow-y:auto; padding-right:4px; } | |
| .residue-col{ display:inline-flex; flex-direction:column; align-items:center; width:15px; font-family:'JetBrains Mono', monospace; } | |
| .residue-col .aa{ font-size:10px; color: var(--text-muted); line-height:1.4; } | |
| .ss-block{ width:100%; height:12px; border-radius:2px; margin-top:2px; } | |
| .ss-block.true-row{ margin-top:1px; opacity:0.55; } | |
| .ss-block.mismatch{ outline:1.5px solid #ef4444; outline-offset:-1px; } | |
| .ss-h{ background: var(--struct-helix); } | |
| .ss-e{ background: var(--struct-sheet); } | |
| .ss-c{ background: var(--struct-coil); } | |
| .legend{ display:flex; gap:16px; flex-wrap:wrap; margin-bottom:4px; } | |
| .legend-item{ display:flex; align-items:center; gap:6px; font-size:0.75rem; color: var(--text-muted); } | |
| .swatch{ width:10px; height:10px; border-radius:2px; display:inline-block; } | |
| .legend-note{ font-size:0.72rem; color: var(--text-dim); margin-bottom:14px; } | |
| .stats-panel{ margin-top:16px; padding-top:14px; border-top:1px solid var(--border-subtle); } | |
| .stat-row{ display:flex; align-items:center; gap:10px; margin-bottom:8px; } | |
| .stat-label{ width:70px; font-size:0.75rem; color: var(--text-muted); font-family:'JetBrains Mono', monospace; } | |
| .bar-track{ flex:1; background: rgba(255,255,255,0.06); border-radius:99px; height:7px; overflow:hidden; } | |
| .bar-fill{ height:100%; border-radius:99px; } | |
| .stat-pct{ width:44px; text-align:right; font-size:0.75rem; font-family:'JetBrains Mono', monospace; color: var(--text-muted); } | |
| .accuracy-badge{ | |
| display:flex; justify-content:space-between; align-items:center; margin-top:14px; | |
| background: rgba(124,158,255,0.08); border:1px solid rgba(124,158,255,0.25); | |
| border-radius:10px; padding:10px 14px; | |
| } | |
| .accuracy-label{ font-size:0.8rem; color: var(--text-muted); } | |
| .accuracy-value{ font-family:'Space Grotesk', sans-serif; font-weight:700; font-size:1.1rem; color:#7c9eff; } | |
| .footer-note{ text-align:center; color: var(--text-dim); font-size:0.75rem; margin-top: 24px; padding-top:18px; border-top:1px solid var(--border-subtle); } | |
| @media (max-width: 640px){ | |
| .residue-col{ width:13px; } | |
| .hero-title{ font-size: 1.7rem; } | |
| } | |
| """ | |
| # ────────────────────────────────────────────────────────────────────────── | |
| # UI | |
| # ────────────────────────────────────────────────────────────────────────── | |
| with gr.Blocks(theme=THEME, css=CSS, title="SERAPH — Protein Structure Prediction") as demo: | |
| gr.HTML( | |
| '<div id="header-row">' | |
| '<a id="back-btn" href="https://muhammad-asad-ullah.vercel.app/" target="_blank">' | |
| '← Back to Portfolio</a>' | |
| '</div>' | |
| '<div class="hero-title">SERAPH — Protein Secondary Structure Prediction</div>' | |
| '<div class="hero-subtitle">An ESM2 + Conv-BiLSTM model that predicts helix, sheet, and coil ' | |
| 'structure directly from an amino acid sequence. Pick one of 50 preloaded proteins or paste your own.</div>' | |
| ) | |
| with gr.Row(): | |
| with gr.Column(scale=3): | |
| protein_dropdown = gr.Dropdown( | |
| choices=PROTEIN_CHOICES, value=CUSTOM_LABEL, | |
| label="Preset Protein", elem_id="protein-select", | |
| ) | |
| sequence_box = gr.Textbox( | |
| label="Amino Acid Sequence", placeholder="e.g. GLSDGEWQLVLNVWGKV...", | |
| lines=4, elem_id="sequence-input", | |
| ) | |
| predict_btn = gr.Button("Predict Structure", variant="primary") | |
| with gr.Column(scale=2): | |
| info_box = gr.HTML('<div class="info-card empty">Pick a preset above to see its background, or paste your own sequence.</div>') | |
| result_box = gr.HTML('<div class="placeholder-msg">Enter or select a sequence, then hit Predict.</div>') | |
| gr.HTML('<div class="footer-note">SERAPH · Built by Muhammad Asad Ullah (PypCoder)</div>') | |
| protein_dropdown.change(fn=load_preset, inputs=protein_dropdown, outputs=[sequence_box, info_box]) | |
| predict_btn.click(fn=predict, inputs=[sequence_box, protein_dropdown], outputs=result_box) | |
| demo.launch() |