Update app.py

app.py

@@ -1,86 +1,337 @@
-"""
-SERAPH Playground — Gradio app entry point.
-
-Run with: python 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 config import PORTFOLIO_URL
-from dataset import PROTEIN_DATASET
-from html_builder import build_result_html
-from inference import predict_structure, on_preset_change
-from styles import CUSTOM_CSS
-
-with gr.Blocks(css=CUSTOM_CSS, title="SERAPH — Protein Playground") as demo:
-
-    # Top Branding Navigation Bar (back button + title + status)
-    gr.HTML(f"""
-    <div style='display: flex; justify-content: space-between; align-items: center; padding: 12px 0; margin-bottom: 20px; border-bottom: 1px solid rgba(255,255,255,0.08);'>
-        <a href="{PORTFOLIO_URL}" target="_top" class="back-to-portfolio">
-            &larr; Back to Portfolio
-        </a>
-        <div style='font-family: "Space Grotesk", sans-serif; font-weight: 700; font-size: 1.1rem; color: #f5f5f7;'>
-            SERAPH <span style='font-size: 0.75rem; font-family: "JetBrains Mono"; color: #8e8e93; font-weight: 400; margin-left: 8px;'>v1.0 • ESM-2 + BiLSTM</span>
-        </div>
-        <div style='font-family: "JetBrains Mono", monospace; font-size: 0.72rem; padding: 4px 12px; background: rgba(255, 255, 255, 0.04); border: 1px solid rgba(255, 255, 255, 0.08); border-radius: 99px; color: #8e8e93;'>
-            🟢 Model Live
+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>
-    """)
+    """
 
-    gr.HTML("""
-    <h1 class="hero-title">SERAPH Playground</h1>
-    <p class="hero-subtitle">Predict 3-state protein secondary structure (Alpha Helix, Beta Sheet, Coil) using fine-tuned ESM-2 protein language embeddings.</p>
-    """)
 
-    with gr.Row():
-        preset_dropdown = gr.Dropdown(
-            choices=["-- Select a Preloaded Protein Preset (50 Available) --"] + [p["name"] for p in PROTEIN_DATASET],
-            value="01. Human Myoglobin (Oxygen Storage)",
-            label="SELECT FROM 50 PRELOADED PROTEIN STRUCTURES",
-            interactive=True
-        )
+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
 
-    with gr.Row():
-        sequence_input = gr.Textbox(
-            value=PROTEIN_DATASET[0]["sequence"],
-            label="AMINO ACID SEQUENCE (FASTA / IUPAC)",
-            lines=3,
-            placeholder="Type or paste amino acid sequence..."
-        )
 
-    with gr.Row():
-        predict_btn = gr.Button("Run SERAPH Prediction ⚡", elem_classes=["btn-magnetic"])
-
-    output_html = gr.HTML(
-        value=build_result_html(
-            PROTEIN_DATASET[0]["sequence"],
-            "C" * len(PROTEIN_DATASET[0]["sequence"]),  # placeholder before click
-            PROTEIN_DATASET[0]["true_ss"],
-            None,
-            PROTEIN_DATASET[0]["description"],
-            PROTEIN_DATASET[0]["fun_fact"]
-        )
-    )
+# ──────────────────────────────────────────────────────────────────────────
+# 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)",
+)
 
-    # Event Handlers
-    preset_dropdown.change(
-        fn=on_preset_change,
-        inputs=[preset_dropdown],
-        outputs=[sequence_input]
-    )
+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; }
 
-    predict_btn.click(
-        fn=predict_structure,
-        inputs=[sequence_input, preset_dropdown],
-        outputs=[output_html, gr.State(), gr.State()]
+.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="_self">'
+        '← 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>'
     )
 
-    gr.HTML("""
-    <div style='border-top: 1px solid rgba(255,255,255,0.08); margin-top: 40px; padding-top: 20px; text-align: center; font-size: 0.8rem; color: #55555a; font-family: "Plus Jakarta Sans";'>
-        Built with PyTorch &amp; Hugging Face Spaces
-    </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)
 
-if __name__ == "__main__":
-    demo.launch()
+demo.launch()