Update app.py

app.py

@@ -1,24 +1,13 @@
-# Improved version of your LocPred-Prok app
-# (Tailwind layout, responsive SVG, download PNG, dark/light auto theme)
-
-# NOTE: This is a placeholder structure.
-# I will generate the full, ready-to-run file in the next update.
-
-# --- START OF FILE ---
-
-# Improved, complete LocPred-Prok Gradio app
-# Features added/fixed:
-# - Responsive, centered SVG that supports horizontal/circular layouts and high-res rendering
-# - Dark / light automatic color adaptation (via prefers-color-scheme)
-# - Client-side SVG -> PNG download buttons (no server libs required)
-# - Tailwind CDN used for layout utilities (works inside Gradio HTML panel)
-# - Tailwind-like alignment applied to main layout
-# - Keeps original model loading and prediction logic
-
+# app_locpred_prok.py
 import os
-import json
 import re
+import json
 import uuid
+import io
+import shutil
+import base64
+from typing import Tuple
+
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -27,26 +16,31 @@ import matplotlib.pyplot as plt
 import numpy as np
 from transformers import AutoTokenizer, AutoModel
 
-# ---------- Environment & cache (same as original) ----------
+# Optional server-side PDF export dependency
+try:
+    import cairosvg
+    CAIROSVG_AVAILABLE = True
+except Exception:
+    CAIROSVG_AVAILABLE = False
+
+# ========== Environment (same cache handling as before) ==========
 plt.switch_backend('Agg')
 os.environ["HF_HOME"] = "/tmp/hf_cache"
 os.environ["TRANSFORMERS_CACHE"] = "/tmp/hf_cache"
 os.environ["HF_HUB_DISABLE_SYMLINKS_WARNING"] = "1"
 
-import shutil
 for path in ["/tmp/hf_cache", os.path.expanduser("~/.cache/huggingface")]:
     shutil.rmtree(path, ignore_errors=True)
     os.makedirs(path, exist_ok=True)
 
-# ---------- Model architecture (unchanged core, with attention output) ----------
+# ========== Model architecture (same as you had) ==========
 class AttentionPooling(nn.Module):
     def __init__(self, d_model):
         super().__init__()
         self.attention_net = nn.Linear(d_model, 1)
 
     def forward(self, x, mask):
-        # x: [B, L, D], mask: [B, L]
-        attn_logits = self.attention_net(x).squeeze(2)  # [B, L]
+        attn_logits = self.attention_net(x).squeeze(2)
         attn_logits = attn_logits.masked_fill(mask == 0, -1e9)
         attn_weights = F.softmax(attn_logits, dim=1)
         return torch.bmm(attn_weights.unsqueeze(1), x).squeeze(1), attn_weights
@@ -60,7 +54,11 @@ class ProtDualBranchEnhancedClassifier(nn.Module):
         self.tok_projector = nn.Linear(d_model, projection_dim)
         fused_dim = projection_dim * 2
         self.gate = nn.Sequential(nn.Linear(fused_dim, fused_dim), nn.Sigmoid())
-        self.classifier_head = nn.Sequential(nn.LayerNorm(fused_dim), nn.Linear(fused_dim, fused_dim * 2), nn.ReLU(), nn.Dropout(dropout), nn.Linear(fused_dim * 2, num_classes))
+        self.classifier_head = nn.Sequential(nn.LayerNorm(fused_dim),
+                                             nn.Linear(fused_dim, fused_dim * 2),
+                                             nn.ReLU(),
+                                             nn.Dropout(dropout),
+                                             nn.Linear(fused_dim * 2, num_classes))
 
     def forward(self, cls_embedding, token_embeddings, mask):
         z_cls = self.cls_projector(cls_embedding)
@@ -73,236 +71,384 @@ class ProtDualBranchEnhancedClassifier(nn.Module):
         z_fused_gated = z_fused_concat * gate_values
         return self.classifier_head(z_fused_gated), pooling_weights
 
-# ---------- Load PLM + classifier ----------
+# ========== Load models (keep same variable names so your logic works) ==========
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 PLM_MODEL_NAME = "facebook/esm2_t30_150M_UR50D"
 CLASSIFIER_PATH = "best_model_esm2_t30_150M_UR50D.pth"
 LABEL_MAP_PATH = "label_map.json"
 
-if not os.path.exists(LABEL_MAP_PATH):
-    raise FileNotFoundError(f"Missing {LABEL_MAP_PATH}")
-if not os.path.exists(CLASSIFIER_PATH):
-    raise FileNotFoundError(f"Missing {CLASSIFIER_PATH}")
-
-with open(LABEL_MAP_PATH, 'r') as f:
-    label_to_idx = json.load(f)
-    idx_to_label = {v: k for k, v in label_to_idx.items()}
-NUM_CLASSES = len(idx_to_label)
-D_MODEL = 640
-
-print("🔹 Loading models...")
-tokenizer = AutoTokenizer.from_pretrained(PLM_MODEL_NAME)
-plm_model = AutoModel.from_pretrained(PLM_MODEL_NAME).to(DEVICE).eval()
-classifier = ProtDualBranchEnhancedClassifier(D_MODEL, 32, NUM_CLASSES, 0.3, 3).to(DEVICE)
-classifier.load_state_dict(torch.load(CLASSIFIER_PATH, map_location=DEVICE))
-classifier.eval()
-print("✅ Models loaded and ready.")
-
-# ---------- SVG Generator with layout options and responsive wrapper ----------
-def generate_bacterial_svg(target_class, layout='circular', high_res=False):
-    # Normalize target
-    target = target_class.lower() if target_class else ""
-
-    # Determine active compartments
-    is_sec  = "extracellular" in target or "secreted" in target
-    is_om   = "outer membrane" in target
-    is_peri = "periplasm" in target
-    is_cw   = "cell wall" in target
-    is_im   = "plasma membrane" in target or "inner membrane" in target
-    is_cyto = "cytoplasm" in target or "cytosol" in target
-
-    # Color tokens using CSS variables (support dark/light)
-    # Colors are referenced as var(--color-...)
-    c = {
-        "hl_stroke": "var(--hl-stroke)", "hl_fill": "var(--hl-fill)", "hl_text": "var(--hl-text)", "hl_dot": "var(--hl-dot)",
-        "bg_stroke": "var(--bg-stroke)", "bg_fill_om": "var(--bg-fill-om)", "bg_fill_im": "var(--bg-fill-im)",
-        "bg_text": "var(--bg-text)", "bg_line": "var(--bg-line)", "bg_dot": "var(--bg-dot)"
+# If you want to test UI without model files, set MOCK_MODE = True
+MOCK_MODE = False
+
+if not MOCK_MODE:
+    if not os.path.exists(LABEL_MAP_PATH):
+        raise FileNotFoundError(f"Missing {LABEL_MAP_PATH}")
+    if not os.path.exists(CLASSIFIER_PATH):
+        raise FileNotFoundError(f"Missing {CLASSIFIER_PATH}")
+
+    with open(LABEL_MAP_PATH, 'r') as f:
+        label_to_idx = json.load(f)
+        idx_to_label = {v: k for k, v in label_to_idx.items()}
+    NUM_CLASSES = len(idx_to_label)
+    D_MODEL = 640
+
+    print("🔹 Loading models...")
+    tokenizer = AutoTokenizer.from_pretrained(PLM_MODEL_NAME)
+    plm_model = AutoModel.from_pretrained(PLM_MODEL_NAME).to(DEVICE).eval()
+    classifier = ProtDualBranchEnhancedClassifier(D_MODEL, 32, NUM_CLASSES, 0.3, 3).to(DEVICE)
+    classifier.load_state_dict(torch.load(CLASSIFIER_PATH, map_location=DEVICE))
+    classifier.eval()
+    print("✅ Models loaded.")
+else:
+    # Mock objects for UI development
+    idx_to_label = {0: "Cytoplasm", 1: "Inner Membrane", 2: "Periplasm", 3: "Outer Membrane", 4: "Cell Wall", 5: "Extracellular"}
+    NUM_CLASSES = len(idx_to_label)
+    tokenizer = None
+    plm_model = None
+    classifier = None
+
+# ========== SVG generator (UniProt-like) ==========
+def generate_uniprot_style_svg(pred_label: str,
+                               gram: str = "negative",
+                               theme: str = "uniprot-blue",
+                               layout: str = "circular",
+                               high_res: bool = False,
+                               uid: str = None) -> str:
+    """
+    Create a UniProt-like bacterial localization diagram:
+      - pred_label: predicted top class name (used to highlight)
+      - gram: 'negative' or 'positive'
+      - theme: 'uniprot-blue', 'red-highlight', 'auto' (auto uses prefers-color-scheme)
+      - layout: 'circular' (capsule) or 'horizontal' (cell left, labels right)
+      - high_res: True -> larger viewBox for higher-quality PNG/PDF
+    Returns: HTML string containing responsive SVG and download buttons.
+    """
+
+    target = pred_label.lower() if pred_label else ""
+    is_active = {
+        "sec":  ("extracellular" in target) or ("secreted" in target),
+        "om":   ("outer membrane" in target),
+        "peri": ("periplasm" in target),
+        "cw":   ("cell wall" in target),
+        "im":   ("inner membrane" in target) or ("plasma membrane" in target),
+        "cyto": ("cytoplasm" in target) or ("cytosol" in target)
     }
 
-    def style(active, base_fill, base_stroke, w_act="4", w_norm="2"):
-        if active:
-            return c['hl_fill'], c['hl_stroke'], w_act
-        return base_fill, base_stroke, w_norm
-
-    om_f, om_s, om_w = style(is_peri, c['bg_fill_om'], c['hl_stroke'] if is_om else c['bg_stroke'])
-    cw_s = c['hl_stroke'] if is_cw else "var(--muted)"
-    cw_w, cw_d = ("3", "0") if is_cw else ("1.5", "6,4")
-    im_f, im_s, im_w = style(is_cyto, c['bg_fill_im'], c['hl_stroke'] if is_im else c['bg_stroke'])
-
-    def label_style(active):
-        if active: return c['hl_text'], 'bold', c['hl_stroke'], '2.5', c['hl_dot'], '5'
-        return c['bg_text'], 'normal', c['bg_line'], '1.5', c['bg_dot'], '3'
-
-    l_sec  = label_style(is_sec)
-    l_om   = label_style(is_om)
-    l_peri = label_style(is_peri)
-    l_cw   = label_style(is_cw)
-    l_im   = label_style(is_im)
-    l_cyto = label_style(is_cyto)
-
-    # Size and viewBox (increase resolution if high_res)
-    base_w, base_h = (1200, 600) if high_res else (800, 420)
-    viewbox = f"0 0 {base_w} {base_h}"
-
-    # Choose layout: circular or horizontal
-    if layout == 'horizontal':
-        # Place cell on left, labels on right in a row
-        bx, by = int(base_w * 0.35), int(base_h * 0.5)
-        tx = int(base_w * 0.75)
+    # If gram-positive, there is no outer membrane and cell wall is thicker
+    have_outer_membrane = (gram == "negative")
+
+    # color themes using CSS variables (supports prefers-color-scheme)
+    css_vars = {
+        "uniprot-blue": {
+            "--om-fill": "#F5F7FA", "--im-fill": "#FFFFFF", "--stroke": "#607D8B",
+            "--muted": "#B0BEC5", "--text": "#263238", "--highlight": "#0288D1"
+        },
+        "red-highlight": {
+            "--om-fill": "#FFEBEE", "--im-fill": "#FFFFFF", "--stroke": "#607D8B",
+            "--muted": "#B0BEC5", "--text": "#263238", "--highlight": "#D32F2F"
+        }
+    }
+
+    selected = css_vars.get(theme if theme in css_vars else "uniprot-blue")
+
+    # Unique IDs for DOM elements to allow multiple diagrams on page
+    uid = uid or str(uuid.uuid4()).replace("-", "")[:10]
+    svg_id = f"loc_svg_{uid}"
+
+    # Sizes
+    if high_res:
+        W, H = 1600, 800
     else:
-        # circular: center cell and labels on right
-        bx, by = int(base_w * 0.35), int(base_h * 0.5)
-        tx = int(base_w * 0.75)
-
-    # Anchor points relative to center
-    targets = {
-        'sec':  (bx,       by - 180),
-        'om':   (bx + 140, by - 120),
-        'peri': (bx + 120, by - 90),
-        'cw':   (bx + 100, by - 70),
-        'im':   (bx + 70,  by - 50),
-        'cyto': (bx,       by)
+        W, H = 800, 420
+
+    # anchor coords and label positions (tuned for viewBox)
+    center_x = int(W * 0.38)
+    center_y = int(H * 0.5)
+    label_x = int(W * 0.76)
+    label_x_left = int(W * 0.58)
+
+    label_y_map = {
+        "sec": int(H * 0.12),
+        "om": int(H * 0.22),
+        "peri": int(H * 0.32),
+        "cw": int(H * 0.42),
+        "im": int(H * 0.62),
+        "cyto": int(H * 0.78)
     }
 
-    # label Y positions
-    text_y = {
-        'sec':  int(base_h*0.12), 'om': int(base_h*0.22), 'peri': int(base_h*0.32),
-        'cw': int(base_h*0.42), 'im': int(base_h*0.62), 'cyto': int(base_h*0.78)
+    anchors = {
+        "sec": (center_x + 160, center_y - 160),
+        "om":  (center_x + 160, center_y - 100),
+        "peri":(center_x + 140, center_y - 40),
+        "cw":  (center_x + 120, center_y + 10),
+        "im":  (center_x + 80,  center_y + 60),
+        "cyto":(center_x + 20,  center_y + 70)
     }
 
-    def draw_connector(key, style_tuple, label_text):
-        txt_col, weight, line_col, width, dot_col, r = style_tuple
-        tx_pos, ty_pos = tx, text_y[key]
-        ex, ey = targets[key]
-        c1x, c1y = tx_pos - int(base_w*0.08), ty_pos
-        c2x, c2y = ex + int(base_w*0.06), ey
-        path = f"M {tx_pos - 10} {ty_pos - 5} C {c1x} {c1y}, {c2x} {c2y}, {ex} {ey}"
-        return f'''
-        <g>
-            <text x="{tx_pos}" y="{ty_pos}" fill="{txt_col}" font-weight="{weight}" font-size="14" font-family="Inter, Arial">{label_text}</text>
-            <path d="{path}" fill="none" stroke="{line_col}" stroke-width="{width}" stroke-linecap="round" stroke-linejoin="round" />
-            <circle cx="{ex}" cy="{ey}" r="{r}" fill="{dot_col}" stroke="white" stroke-width="1" />
+    # Build helper for connector group
+    def connector_svg(key, text):
+        ex, ey = anchors[key]
+        tx, ty = label_x, label_y_map[key]
+        # styling depends on active
+        active = is_active.get(key, False)
+        stroke_color = selected["--highlight"] if active else selected["--muted"]
+        stroke_w = "2.6" if active else "1.4"
+        fontw = "700" if active else "400"
+        dot_r = "6" if active else "4"
+        path = f"M {tx-20} {ty-6} C {tx-100} {ty-6}, {ex+60} {ey+10}, {ex} {ey}"
+        return f"""
+        <g class="connector connector-{key}">
+            <text x="{tx}" y="{ty}" fill="{selected['--text']}" font-weight="{fontw}" font-size="{14 if not high_res else 22}" font-family="Inter, Arial">{text}</text>
+            <path d="{path}" fill="none" stroke="{stroke_color}" stroke-width="{stroke_w}" stroke-linecap="round" stroke-linejoin="round" />
+            <circle cx="{ex}" cy="{ey}" r="{dot_r}" fill="{stroke_color}" stroke="white" stroke-width="{1 if not high_res else 1.6}" />
+        </g>
+        """
+
+    # Build cell shapes: capsule-like curves - simpler parametric shapes
+    # Outer membrane (or single outer layer for gram-positive), cell wall, inner membrane
+    # Different shapes when gram-positive: thicker cell wall, no outer membrane ring.
+    # We'll draw with bezier-like path strings tuned to look UniProt-ish.
+    if have_outer_membrane:
+        om_fill = "var(--om-fill)"
+        om_stroke = "var(--stroke)"
+        cw_fill = "none"
+        cw_stroke = "var(--muted)"
+        im_fill = "var(--im-fill)"
+        im_stroke = "var(--stroke)"
+    else:
+        # gram-positive: cell wall thicker and outer membrane absent
+        om_fill = "none"
+        om_stroke = "none"
+        cw_fill = "var(--om-fill)"
+        cw_stroke = "var(--stroke)"
+        im_fill = "var(--im-fill)"
+        im_stroke = "var(--stroke)"
+
+    # layer highlight override if active
+    def stroke_override(key, base):
+        if is_active.get(key, False):
+            return selected["--highlight"]
+        return base
+
+    # inline CSS for animations and hover effects
+    svg_style = f"""
+    <style>
+      /* theme vars */
+      :root {{
+        --om-fill: {selected['--om-fill']};
+        --im-fill: {selected['--im-fill']};
+        --stroke: {selected['--stroke']};
+        --muted: {selected['--muted']};
+        --text: {selected['--text']};
+        --highlight: {selected['--highlight']};
+      }}
+      @media (prefers-color-scheme: dark) {{
+        :root {{
+          --om-fill: #28343a;
+          --im-fill: #1f2b30;
+          --stroke: #90a4ae;
+          --muted: #546e7a;
+          --text: #e0f2f1;
+        }}
+      }}
+
+      /* connector hover: slightly thicken the path and enlarge dot */
+      .connector path {{ transition: stroke-width 180ms ease, stroke 180ms ease; opacity:0.95; }}
+      .connector circle {{ transition: r 160ms ease, transform 160ms ease; transform-origin: center; }}
+      .connector text {{ transition: fill 160ms ease; }}
+
+      /* on hover of group, emphasize */
+      .connector:hover path {{ stroke-width: calc(var(--hover-w, 3)); opacity:1; filter: drop-shadow(0 2px 2px rgba(0,0,0,0.06)); }}
+      .connector:hover circle {{ transform: scale(1.25); }}
+      /* subtle floating animation for lines */
+      .connector path {{}}
+      @keyframes floatx {{ 0% {{ transform: translateX(0px); }} 50% {{ transform: translateX(1px); }} 100% {{ transform: translateX(0px); }} }}
+      .connector path {{ animation: floatx 4s ease-in-out infinite; animation-delay: calc(var(--i, 0) * 0.12s); opacity:0.95; }}
+
+      /* make the whole svg responsive */
+      svg {{ max-width: 100%; height: auto; display:block; }}
+
+      /* layer highlight when active: add glow */
+      .layer-active {{ filter: drop-shadow(0 4px 8px rgba(0,0,0,0.08)); }}
+    </style>
+    """
+
+    # Compose SVG core shapes (simplified, but tuned coordinates)
+    # We use path shapes with translated center for convenience.
+    cell_shapes = ""
+    # Outer membrane / envelope
+    if have_outer_membrane:
+        cell_shapes += f'''
+        <g id="outer_membrane" class="layer {'layer-active' if is_active['om'] else ''}">
+          <ellipse cx="{center_x}" cy="{center_y}" rx="{220 if not high_res else 440}" ry="{170 if not high_res else 340}"
+                   fill="var(--om-fill)" stroke="{stroke_override('om', 'var(--stroke)')}" stroke-width="{3 if is_active['om'] else 2}"/>
         </g>
         '''
-
-    # Draw cell shapes: use rounded rects for membranes (keeping original geometry scaled)
-    svg_shapes = f'''
-    <g transform="translate({bx}, {by})">
-        <rect x="{-150}" y="{-150}" width="300" height="300" rx="150" ry="150" fill="{om_f}" stroke="{om_s}" stroke-width="{om_w}" />
-        <rect x="{-110}" y="{-110}" width="220" height="220" rx="110" ry="110" fill="none" stroke="{cw_s}" stroke-width="{cw_w}" stroke-dasharray="{cw_d}" />
-        <rect x="{-70}" y="{-70}" width="140" height="140" rx="70" ry="70" fill="{im_f}" stroke="{im_s}" stroke-width="{im_w}" />
-        <g opacity="0.45">
-            <path d="M -30 -20 Q 0 -60 30 -20 T 60 -10" fill="none" stroke="var(--muted)" stroke-width="3" />
-            <circle cx="-40" cy="30" r="3" fill="var(--muted)" /> <circle cx="20" cy="40" r="3" fill="var(--muted)" />
+        # Cell wall (dashed)
+        cell_shapes += f'''
+        <g id="cell_wall">
+          <ellipse cx="{center_x}" cy="{center_y}" rx="{190 if not high_res else 380}" ry="{150 if not high_res else 300}"
+                   fill="none" stroke="{stroke_override('cw','var(--muted)')}" stroke-width="{4 if is_active['cw'] else 2}" stroke-dasharray="10 6"/>
+        </g>
+        '''
+        # inner membrane
+        cell_shapes += f'''
+        <g id="inner_membrane" class="layer {'layer-active' if is_active['im'] else ''}">
+          <ellipse cx="{center_x}" cy="{center_y}" rx="{140 if not high_res else 280}" ry="{100 if not high_res else 200}"
+                   fill="var(--im-fill)" stroke="{stroke_override('im','var(--stroke)')}" stroke-width="{3 if is_active['im'] else 1.8}"/>
+        </g>
+        '''
+    else:
+        # Gram positive: thick cell wall as filled ellipse + inner membrane
+        cell_shapes += f'''
+        <g id="cell_wall_gp" class="layer {'layer-active' if is_active['cw'] else ''}">
+          <ellipse cx="{center_x}" cy="{center_y}" rx="{230 if not high_res else 460}" ry="{180 if not high_res else 360}"
+                   fill="{selected['--om-fill']}" stroke="{stroke_override('cw','var(--stroke)')}" stroke-width="{3 if is_active['cw'] else 2}"/>
+        </g>
+        <g id="inner_membrane" class="layer {'layer-active' if is_active['im'] else ''}">
+          <ellipse cx="{center_x}" cy="{center_y}" rx="{150 if not high_res else 300}" ry="{110 if not high_res else 220}"
+                   fill="var(--im-fill)" stroke="{stroke_override('im','var(--stroke)')}" stroke-width="{2 if is_active['im'] else 1.4}"/>
         </g>
+        '''
+
+    # cytoplasm ornament
+    cell_shapes += f'''
+    <g id="cytoplasm_wiggles" opacity="0.65">
+      <path d="M {center_x-60} {center_y+10} q 30 -50 70 0 q 30 50 70 0" stroke="var(--muted)" stroke-width="6" fill="none" stroke-linecap="round"/>
+      <circle cx="{center_x-40}" cy="{center_y+40}" r="{3 if not high_res else 6}" fill="var(--muted)"/>
+      <circle cx="{center_x+20}" cy="{center_y+50}" r="{3 if not high_res else 6}" fill="var(--muted)"/>
     </g>
     '''
 
-    # Compose connectors
-    connectors = "".join([
-        draw_connector('sec', l_sec, 'Extracellular / Secreted'),
-        draw_connector('om', l_om, 'Outer Membrane'),
-        draw_connector('peri', l_peri, 'Periplasm'),
-        draw_connector('cw', l_cw, 'Cell Wall'),
-        draw_connector('im', l_im, 'Inner Membrane'),
-        draw_connector('cyto', l_cyto, 'Cytoplasm')
-    ])
-
-    svg_core = f'''<svg id="svg_core" width="100%" height="auto" viewBox="{viewbox}" xmlns="http://www.w3.org/2000/svg" role="img" aria-label="Bacterial localization diagram">
+    # connectors
+    connectors = ""
+    connectors += connector_svg("sec", "Extracellular / Secreted")
+    # outer membrane only if present
+    if have_outer_membrane:
+        connectors += connector_svg("om", "Outer Membrane")
+    connectors += connector_svg("peri", "Periplasm")
+    connectors += connector_svg("cw", "Cell Wall")
+    connectors += connector_svg("im", "Inner Membrane")
+    connectors += connector_svg("cyto", "Cytoplasm")
+
+    # Build download buttons and client-side JS to download SVG and PNG
+    # PDF export will call a Gradio server endpoint (provided below)
+    html = f"""
+    <div style="width:100%; text-align:center;">
+    {svg_style}
+    <svg id="{svg_id}" viewBox="0 0 {W} {H}" xmlns="http://www.w3.org/2000/svg" role="img" aria-label="Bacterial localization diagram">
         <defs>
             <style><![CDATA[
                 text {{ font-family: Inter, Arial, sans-serif; }}
             ]]></style>
         </defs>
-        {svg_shapes}
+        {cell_shapes}
         {connectors}
-    </svg>'''
-
-    # Create unique wrapper id so multiple calls don't clash
-    uid = str(uuid.uuid4()).replace('-', '')
-    wrapper = f"loc_svg_{uid}"
-
-    # Build responsive HTML with inline JS to enable SVG/PNG download
-    html = f'''
-    <div id="{wrapper}" style="width:100%; text-align:center;">
-      <div style="display:inline-block; max-width:100%; width:900px;">
-        {svg_core}
-        <div style="margin-top:8px; display:flex; gap:8px; justify-content:center; align-items:center;">
-          <button id="btn_svg_{uid}" class="download-btn">Download SVG</button>
-          <button id="btn_png_{uid}" class="download-btn">Download PNG</button>
-          <div style="font-size:12px; color:var(--bg-text); align-self:center;">Layout: {layout.title()} {'· High-res' if high_res else ''}</div>
-        </div>
-      </div>
+    </svg>
+
+    <div style="margin-top:8px; display:flex; gap:8px; justify-content:center; align-items:center;">
+      <button id="download_svg_{uid}" class="download-btn">Download SVG</button>
+      <button id="download_png_{uid}" class="download-btn">Download PNG</button>
+      <button id="download_pdf_{uid}" class="download-btn">Download PDF</button>
+      <div style="font-size:12px; color:var(--text); align-self:center;">{gram.title()} · {layout.title()} {'· High-res' if high_res else ''}</div>
     </div>
+    </div>
+
     <script>
     (function(){{
-        const wrapper = document.getElementById('{wrapper}');
-        const svgEl = wrapper.querySelector('svg');
-        const btnSvg = document.getElementById('btn_svg_{uid}');
-        const btnPng = document.getElementById('btn_png_{uid}');
-
-        // Helper: download file
-        function download(filename, blob){{
-            const url = URL.createObjectURL(blob);
-            const a = document.createElement('a');
-            a.href = url; a.download = filename; document.body.appendChild(a); a.click();
-            setTimeout(()=>{{ URL.revokeObjectURL(url); a.remove(); }}, 100);
+      const svgEl = document.getElementById("{svg_id}");
+      const btnSvg = document.getElementById("download_svg_{uid}");
+      const btnPng = document.getElementById("download_png_{uid}");
+      const btnPdf = document.getElementById("download_pdf_{uid}");
+
+      function downloadFile(filename, blob) {{
+        const url = URL.createObjectURL(blob);
+        const a = document.createElement('a');
+        a.href = url; a.download = filename; document.body.appendChild(a); a.click();
+        setTimeout(()=>{{ URL.revokeObjectURL(url); a.remove(); }}, 200);
+      }}
+
+      btnSvg.addEventListener('click', ()=>{{
+        const serializer = new XMLSerializer();
+        let source = serializer.serializeToString(svgEl);
+        if(!source.match(/^<svg[^>]+xmlns="http:\\/\\/www.w3.org\\/2000\\/svg"/)) {{
+            source = source.replace(/^<svg/, '<svg xmlns="http://www.w3.org/2000/svg"');
         }}
-
-        btnSvg.addEventListener('click', ()=>{{
-            const serializer = new XMLSerializer();
-            let source = serializer.serializeToString(svgEl);
-            // Add name spaces.
-            if(!source.match(/^<svg[^>]+xmlns="http\:\/\/www\.w3\.org\/2000\/svg"/)){{
-                source = source.replace(/^<svg/, '<svg xmlns="http://www.w3.org/2000/svg"');
-            }}
-            if(!source.match(/^<svg[^>]+xmlns:xlink="http\:\/\/www\.w3\.org\/1999\/xlink"/)){{
-                source = source.replace(/^<svg/, '<svg xmlns:xlink="http://www.w3.org/1999/xlink"');
-            }}
-            const blob = new Blob([source], {{type: 'image/svg+xml;charset=utf-8'}});
-            download('locpred_diagram.svg', blob);
-        }});
-
-        btnPng.addEventListener('click', ()=>{{
-            const serializer = new XMLSerializer();
-            let source = serializer.serializeToString(svgEl);
-            if(!source.match(/^<svg[^>]+xmlns="http\:\/\/www\.w3\.org\/2000\/svg"/)){{
-                source = source.replace(/^<svg/, '<svg xmlns="http://www.w3.org/2000/svg"');
-            }}
-            const image = new Image();
-            const svgBlob = new Blob([source], {{type: 'image/svg+xml;charset=utf-8'}});
-            const url = URL.createObjectURL(svgBlob);
-            image.onload = function(){{
-                const canvas = document.createElement('canvas');
-                // scale canvas to natural image size (use 2x for better quality)
-                canvas.width = image.width * 2;
-                canvas.height = image.height * 2;
-                const ctx = canvas.getContext('2d');
-                // set background transparent-friendly
-                ctx.fillStyle = 'white';
-                ctx.fillRect(0,0,canvas.width, canvas.height);
-                ctx.drawImage(image, 0, 0, canvas.width, canvas.height);
-                canvas.toBlob(function(blob){{
-                    download('locpred_diagram.png', blob);
-                }}, 'image/png');
-                URL.revokeObjectURL(url);
-            }};
-            // In some environments the SVG does not have width/height set; set reasonable defaults
-            image.src = url;
-        }});
+        const blob = new Blob([source], {{type: 'image/svg+xml;charset=utf-8'}});
+        downloadFile('locpred_diagram.svg', blob);
+      }});
+
+      btnPng.addEventListener('click', ()=>{{
+        const serializer = new XMLSerializer();
+        let source = serializer.serializeToString(svgEl);
+        if(!source.match(/^<svg[^>]+xmlns="http:\\/\\/www.w3.org\\/2000\\/svg"/)) {{
+            source = source.replace(/^<svg/, '<svg xmlns="http://www.w3.org/2000/svg"');
+        }}
+        const svgBlob = new Blob([source], {{type: 'image/svg+xml;charset=utf-8'}});
+        const url = URL.createObjectURL(svgBlob);
+        const img = new Image();
+        img.onload = function() {{
+          const canvas = document.createElement('canvas');
+          // scale 2x for higher quality
+          const scale = 2;
+          canvas.width = img.width * scale;
+          canvas.height = img.height * scale;
+          const ctx = canvas.getContext('2d');
+          // optional white background
+          ctx.fillStyle = "white";
+          ctx.fillRect(0,0,canvas.width,canvas.height);
+          ctx.drawImage(img, 0, 0, canvas.width, canvas.height);
+          canvas.toBlob(function(blob) {{
+            downloadFile('locpred_diagram.png', blob);
+          }}, 'image/png');
+          URL.revokeObjectURL(url);
+        }};
+        img.onerror = function(e) {{
+          alert('Failed to render PNG in your browser.');
+          URL.revokeObjectURL(url);
+        }};
+        img.src = url;
+      }});
+
+      btnPdf.addEventListener('click', async ()=>{{
+        // send the SVG string to the server /gradio route for PDF conversion
+        const serializer = new XMLSerializer();
+        let source = serializer.serializeToString(svgEl);
+        if(!source.match(/^<svg[^>]+xmlns="http:\\/\\/www.w3.org\\/2000\\/svg"/)) {{
+            source = source.replace(/^<svg/, '<svg xmlns="http://www.w3.org/2000/svg"');
+        }}
+        // call Gradio server function via fetch to /convert_svg_to_pdf (provided below)
+        try {{
+          const resp = await fetch('/convert_svg_to_pdf', {{
+            method: 'POST',
+            headers: {{ 'Content-Type': 'application/json' }},
+            body: JSON.stringify({{ svg: source }})
+          }});
+          if(!resp.ok) {{
+            const txt = await resp.text();
+            alert('PDF conversion failed: ' + txt);
+            return;
+          }}
+          const blob = await resp.blob();
+          downloadFile('locpred_diagram.pdf', blob);
+        }} catch (err) {{
+          alert('PDF conversion failed: ' + err);
+        }}
+      }});
     }})();
     </script>
-    '''
+    """
+
     return html
 
-# ---------- Attention heatmap (unchanged) ----------
+# ========== Heatmap (same as before) ==========
 def draw_attention_heatmap_strip(weights, sequence):
     if weights.max() > 0:
         weights = (weights - weights.min()) / (weights.max() - weights.min())
     data = weights.reshape(1, -1)
     fig, ax = plt.subplots(figsize=(8, 1.5), dpi=150)
     im = ax.imshow(data, cmap='Reds', aspect='auto', vmin=0, vmax=1)
-    ax.set_title('Sequence Attention Heatmap (High Color = Key Feature)', fontsize=10, fontweight='bold', color='#37474F', pad=10)
+    ax.set_title('Sequence Attention Heatmap (High Color = Key Feature)', fontsize=10, fontweight='bold', color='#37474F', pad=6)
     ax.set_xlabel('Residue Position', fontsize=9)
     ax.set_yticks([])
     cbar = plt.colorbar(im, ax=ax, orientation='vertical', fraction=0.02, pad=0.02)
@@ -313,105 +459,139 @@ def draw_attention_heatmap_strip(weights, sequence):
     plt.tight_layout()
     return fig
 
-# ---------- Prediction logic (exposes layout + high_res options) ----------
-def predict(sequence_input, layout_choice, high_res_flag):
+# ========== Prediction function ==========
+def predict(sequence_input: str, gram_choice: str, theme_choice: str, layout_choice: str, high_res_flag: bool):
+    """
+    Returns:
+      - confidences dict (for Label)
+      - svg html string (for HTML output)
+      - attention heatmap figure (for Plot)
+    """
     if not sequence_input or sequence_input.isspace():
-        raise gr.Error('Empty Input')
+        raise gr.Error("Empty Input")
+
     seq = "".join(sequence_input.split('\n')[1:]) if sequence_input.startswith('>') else sequence_input
     seq = re.sub(r'[^A-Z]', '', seq.upper())[:1024]
     if not seq:
-        raise gr.Error('Invalid Sequence')
-
-    with torch.no_grad():
-        inputs = tokenizer(seq, return_tensors='pt', truncation=True, max_length=1024).to(DEVICE)
-        outputs = plm_model(**inputs)
-        hidden_states = outputs.last_hidden_state
-        cls_embedding = hidden_states[:, 0, :]
-        token_embeddings = hidden_states[:, 1:-1, :]
-        token_mask = inputs['attention_mask'][:, 1:-1]
-        logits, pooling_weights = classifier(cls_embedding, token_embeddings, token_mask)
-        probs = F.softmax(logits, dim=1)[0]
-
-    top_label = idx_to_label[torch.max(probs, dim=0)[1].item()]
-    confidences = {idx_to_label[i]: float(p) for i, p in enumerate(probs)}
-
-    svg = generate_bacterial_svg(top_label, layout=layout_choice, high_res=high_res_flag)
-    w_np = pooling_weights[0].cpu().numpy()
-    heatmap_plot = draw_attention_heatmap_strip(w_np, seq)
-
-    return confidences, svg, heatmap_plot
-
-# ---------- UI (Tailwind CDN + Auto dark mode via CSS variables) ----------
+        raise gr.Error("Invalid Sequence")
+
+    if MOCK_MODE:
+        # mock probabilities
+        probs = torch.softmax(torch.randn(NUM_CLASSES), dim=0)
+        logits = probs
+        pooling_weights = np.abs(np.random.randn(len(seq))).astype(float)
+        pooling_weights = pooling_weights / pooling_weights.sum()
+        top_label = idx_to_label[int(torch.argmax(probs).item())]
+    else:
+        with torch.no_grad():
+            inputs = tokenizer(seq, return_tensors="pt", truncation=True, max_length=1024).to(DEVICE)
+            outputs = plm_model(**inputs)
+            hidden_states = outputs.last_hidden_state
+            cls_embedding = hidden_states[:, 0, :]
+            token_embeddings = hidden_states[:, 1:-1, :]
+            token_mask = inputs['attention_mask'][:, 1:-1]
+            logits, pooling_weights = classifier(cls_embedding, token_embeddings, token_mask)
+            probs = F.softmax(logits, dim=1)[0]
+            top_label = idx_to_label[torch.argmax(probs).item()]
+            pooling_weights = pooling_weights[0].cpu().numpy()
+
+    confidences = { idx_to_label[i]: float(p) for i,p in enumerate(probs) } if not MOCK_MODE else { idx_to_label[i]: float(p) for i,p in enumerate(probs)}
+    svg_html = generate_uniprot_style_svg(top_label, gram=gram_choice, theme=theme_choice, layout=layout_choice, high_res=high_res_flag)
+    heatmap_fig = draw_attention_heatmap_strip(np.array(pooling_weights), seq)
+
+    return confidences, svg_html, heatmap_fig
+
+# ========== Server-side PDF conversion endpoint for Gradio ==========
+# This function will be exposed at /convert_svg_to_pdf when app launches.
+def convert_svg_to_pdf_endpoint(svg_str: str):
+    """
+    Convert an SVG string to PDF bytes using cairosvg (if available).
+    Return bytes-like object (PDF) or raise error.
+    """
+    if not CAIROSVG_AVAILABLE:
+        raise RuntimeError("Server-side PDF conversion requires 'cairosvg' package. Install with: pip install cairosvg")
+
+    # cairosvg.svg2pdf can take bytes or string
+    pdf_bytes = cairosvg.svg2pdf(bytestring=svg_str.encode('utf-8'))
+    return ("locpred_diagram.pdf", pdf_bytes)
+
+# ========== UI layout (Gradio) ==========
 layout_css = """
-/* Minimal overrides and CSS variables for dark/light */
-:root{
-  --bg-fill-om: #F5F5F5; --bg-fill-im: #FAFAFA; --bg-stroke: #90A4AE; --muted: #B0BEC5;
-  --hl-stroke: #D32F2F; --hl-fill: #FFEBEE; --hl-text: #B71C1C; --hl-dot: #D32F2F;
-  --bg-text: #37474F; --bg-line: #CFD8DC; --bg-dot: #B0BEC5;
-}
-@media (prefers-color-scheme: dark) {
-  :root{
-    --bg-fill-om: #263238; --bg-fill-im: #1E2930; --bg-stroke: #455A64; --muted: #37474F;
-    --hl-stroke: #FF8A80; --hl-fill: #3E2723; --hl-text: #FFCDD2; --hl-dot: #FF8A80;
-    --bg-text: #ECEFF1; --bg-line: #37474F; --bg-dot: #546E7A;
-  }
-}
-
-.download-btn{ padding:8px 12px; border-radius:6px; border:1px solid var(--bg-line); background:transparent; cursor:pointer; }
-.download-btn:hover{ box-shadow:0 2px 8px rgba(0,0,0,0.08); }
-
-/* Keep Gradio panels tidy */
-.gradio-container{ max-width:1100px; margin:0 auto; }
+/* small customizations and CSS variables fallback */
+:root { --panel-bg: white; }
+.gradio-container { max-width: 1200px; margin: 0 auto; }
+.download-btn { padding:8px 12px; border-radius:6px; border:1px solid #D1D5DB; background:transparent; cursor:pointer; }
+.download-btn:hover { box-shadow: 0 4px 14px rgba(16,24,40,0.08); }
+.panel-card { border:1px solid #e6eef5; border-radius:8px; padding:12px; background:var(--panel-bg); }
+.panel-header { font-weight:700; color:#475569; border-bottom:2px solid #f1f5f9; padding-bottom:8px; margin-bottom:10px; }
 """
 
-# Use Gradio theme but also inject Tailwind CDN for utility classes in HTML
 theme = gr.themes.Soft(primary_hue="sky").set(body_background_fill="white", block_background_fill="white", block_border_width="0px")
 
-gr_tailwind = """
-<link href="https://cdn.jsdelivr.net/npm/tailwindcss@2.2.19/dist/tailwind.min.css" rel="stylesheet">
-"""
-
-with gr.Blocks(theme=theme, css=layout_css, title="LocPred-Prok") as app:
-    # Inject Tailwind (works in Gradio HTML scope)
-    gr.HTML(gr_tailwind)
-
-    gr.HTML("""
-    <div class="w-full p-4 rounded-lg" style="background:linear-gradient(to right,#E0F7FA,#E1F5FE); border:1px solid #B3E5FC; text-align:center;">
-      <h1 style="font-family:Inter, Arial; font-size:28px; margin:0; color:#0288D1;">LocPred-Prok</h1>
-      <div style="color:#0277BD; margin-top:6px;">Deep Learning Framework for Prokaryotic Subcellular Localization</div>
-    </div>
-    """)
-
+with gr.Blocks(theme=theme, css=layout_css, title="LocPred-Prok (UniProt-style)") as app:
+    gr.Markdown("<div style='font-size:22px; font-weight:800; color:#0288D1;'>LocPred-Prok — UniProt-style visualization</div>")
     with gr.Row():
-        with gr.Column(elem_classes="panel-card", scale=6):
-            gr.Markdown("<div class='panel-header'><span class='panel-label'>A</span>Sequence Input</div>")
-            sequence_input = gr.Textbox(lines=8, show_label=False, placeholder=">Sequence... (single-letter AA)")
+        with gr.Column(scale=6):
+            gr.Markdown("<div class='panel-header'><span style='background:#E0F7FA;color:#0277BD;padding:3px 6px;border-radius:4px;font-weight:800;margin-right:8px;'>A</span>Sequence Input</div>")
+            sequence_input = gr.Textbox(lines=8, placeholder=">Sequence (single-letter amino acids) or paste raw sequence", show_label=False)
             with gr.Row():
                 clear_btn = gr.ClearButton(sequence_input, value="Clear")
                 submit_btn = gr.Button("Predict Analysis", variant="primary")
             with gr.Row():
-                layout_choice = gr.Radio(['circular', 'horizontal'], value='circular', label='Diagram Layout', info='Choose circular (default) or horizontal layout for the cell diagram')
-                high_res_flag = gr.Checkbox(label='High resolution render (larger SVG)', value=False)
-            gr.Examples([[">Outer Membrane\nAPKNTWYTGAKLGWSQYHDTGFINNNGPTHENQLGAGAF..." ]], inputs=sequence_input, label=None)
-
-        with gr.Column(elem_classes="panel-card", scale=6):
-            gr.Markdown("<div class='panel-header'><span class='panel-label'>B</span>Localization Visualization</div>")
+                gram_choice = gr.Radio(choices=["negative", "positive"], value="negative", label="Gram type")
+                theme_choice = gr.Radio(choices=["uniprot-blue", "red-highlight"], value="uniprot-blue", label="Color Theme")
+                layout_choice = gr.Radio(choices=["circular", "horizontal"], value="circular", label="Diagram Layout")
+                high_res_flag = gr.Checkbox(value=False, label="High resolution (bigger SVG/PDF)")
+            gr.Examples([[">Outer Membrane\nAPKNTWYTGAKLGWSQYHDTGFINNNGPTHENQLGAGAF..."]], inputs=sequence_input)
+
+        with gr.Column(scale=6):
+            gr.Markdown("<div class='panel-header'><span style='background:#E0F7FA;color:#0277BD;padding:3px 6px;border-radius:4px;font-weight:800;margin-right:8px;'>B</span>Localization Visualization</div>")
             output_svg = gr.HTML(label="Visual", show_label=False)
 
     with gr.Row():
-        with gr.Column(elem_classes="panel-card", scale=6):
-            gr.Markdown("<div class='panel-header'><span class='panel-label'>C</span>Prediction Confidence</div>")
+        with gr.Column(scale=6):
+            gr.Markdown("<div class='panel-header'><span style='background:#E0F7FA;color:#0277BD;padding:3px 6px;border-radius:4px;font-weight:800;margin-right:8px;'>C</span>Prediction Confidence</div>")
             output_label = gr.Label(num_top_classes=NUM_CLASSES, show_label=False)
-        with gr.Column(elem_classes="panel-card", scale=6):
-            gr.Markdown("<div class='panel-header'><span class='panel-label'>D</span>Learned Attention Heatmap</div>")
-            output_plot = gr.Plot(label="Attention", show_label=False)
+        with gr.Column(scale=6):
+            gr.Markdown("<div class='panel-header'><span style='background:#E0F7FA;color:#0277BD;padding:3px 6px;border-radius:4px;font-weight:800;margin-right:8px;'>D</span>Attention Heatmap</div>")
+            output_plot = gr.Plot(show_label=False)
 
-    submit_btn.click(fn=predict, inputs=[sequence_input, layout_choice, high_res_flag], outputs=[output_label, output_svg, output_plot])
+    submit_btn.click(fn=predict, inputs=[sequence_input, gram_choice, theme_choice, layout_choice, high_res_flag], outputs=[output_label, output_svg, output_plot])
     clear_btn.click(lambda: [None, None, None], outputs=[output_label, output_svg, output_plot])
 
-# ---------- Launch ----------
-if __name__ == '__main__':
-    app.launch()
-
-
-# --- END OF FILE ---
+    # Expose PDF conversion endpoint: Gradio allows adding a separate route handler via app.launch later.
+    # We'll attach the endpoint to the FastAPI app used by Gradio when launching.
+
+# ========== Run server with custom route for PDF conversion ==========
+if __name__ == "__main__":
+    from fastapi import FastAPI, Request, Response
+    import uvicorn
+
+    # Build Gradio app to get underlying FastAPI instance
+    demo = app
+    # Get the underlying FastAPI app (gradio >= 3.0)
+    # When launching, we'll mount a custom route /convert_svg_to_pdf handled by convert_svg_to_pdf_endpoint
+    # Gradio's launch will create a FastAPI object; to avoid internal changes, we use the `server_name` arg.
+
+    # Create a lightweight FastAPI for the PDF endpoint and mount the Gradio interface into it
+    fast_app = FastAPI()
+
+    @fast_app.post("/convert_svg_to_pdf")
+    async def convert_svg_to_pdf_api(request: Request):
+        payload = await request.json()
+        svg = payload.get("svg", None)
+        if not svg:
+            return Response(content="No svg provided", status_code=400)
+        if not CAIROSVG_AVAILABLE:
+            return Response(content="Server-side PDF conversion unavailable: install 'cairosvg' in the server environment.", status_code=501)
+        try:
+            pdf_bytes = cairosvg.svg2pdf(bytestring=svg.encode('utf-8'))
+            return Response(content=pdf_bytes, media_type="application/pdf")
+        except Exception as e:
+            return Response(content=f"PDF conversion error: {e}", status_code=500)
+
+    # Mount the Gradio interface at root
+    gr.mount_gradio_app(fast_app, demo, path="/")
+
+    # Launch uvicorn with the FastAPI app
+    uvicorn.run(fast_app, host="0.0.0.0", port=7860, log_level="info")