app.py

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
# ──────────────────────────────────────────────────────────────────────────
@spaces.GPU
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()