update

app.py

@@ -3,16 +3,11 @@ import time
 
 import torch
 import gradio as gr
-import numpy as np
+import pandas as pd
 from loguru import logger
 
 from stoic.model import Stoic
 
-CHAIN_COLORS = [
-    "#636EFA", "#EF553B", "#00CC96", "#AB63FA", "#FFA15A",
-    "#19D3F3", "#FF6692", "#B6E880", "#FF97FF", "#FECB52",
-]
-
 
 @functools.lru_cache(maxsize=1)
 def get_model():
@@ -46,13 +41,15 @@ def predict(sequences_text: str, top_n: int, return_weights: bool):
 
     chain_labels = [chr(ord("A") + i) for i in range(len(sequences))]
 
-    header = "| Rank | " + " | ".join(f"Chain {l}" for l in chain_labels) + " | Stoichiometry |"
-    separator = "|------|" + "|".join("-----" for _ in chain_labels) + "|---------------|"
+    header = "| Rank | " + " | ".join(f"Chain {l}" for l in chain_labels) + " | Stoichiometry | Score | Probability |"
+    separator = "|------|" + "|".join("-----" for _ in chain_labels) + "|---------------|-------|-------------|"
     rows = []
     for rank, candidate in enumerate(results, 1):
         copies = [candidate.get(seq, 0) for seq in sequences]
         stoich = "".join(f"{l}<sub>{c}</sub>" for l, c in zip(chain_labels, copies))
-        row = f"| {rank} | " + " | ".join(str(c) for c in copies) + f" | {stoich} |"
+        score = candidate.get("score", 0)
+        prob = candidate.get("probability", 0)
+        row = f"| {rank} | " + " | ".join(str(c) for c in copies) + f" | {stoich} | {score:.2f} | {prob:.2e} |"
         rows.append(row)
 
     table = "\n".join([header, separator] + rows)
@@ -65,54 +62,28 @@ def predict(sequences_text: str, top_n: int, return_weights: bool):
     stoich_md = table + "\n".join(legend_lines)
 
     if return_weights:
-        fig = build_weights_plot(residue_predictions, chain_labels)
-        return stoich_md, gr.update(value=fig, visible=True), f"{elapsed:.2f}s"
+        df = build_weights_df(residue_predictions, chain_labels)
+        return stoich_md, gr.update(value=df, visible=True), f"{elapsed:.2f}s"
 
     return stoich_md, gr.update(value=None, visible=False), f"{elapsed:.2f}s"
 
 
-def build_weights_plot(residue_predictions, chain_labels):
-    import plotly.graph_objects as go
-    from plotly.subplots import make_subplots
-
+def build_weights_df(residue_predictions, chain_labels):
     pred_residues = residue_predictions["pred_residues"]
     attention_mask = residue_predictions["attention_mask"]
     seqs = residue_predictions["sequences"]
-    n_chains = len(seqs)
-
-    fig = make_subplots(
-        rows=n_chains, cols=1,
-        subplot_titles=[f"Chain {chain_labels[i]}" for i in range(n_chains)],
-        vertical_spacing=0.12 / max(n_chains, 1),
-    )
 
+    records = []
     for i, seq in enumerate(seqs):
         mask = attention_mask[i].astype(bool)
         weights = pred_residues[i][mask]
-        residues = list(seq[:len(weights)])
-        positions = list(range(1, len(weights) + 1))
-        color = CHAIN_COLORS[i % len(CHAIN_COLORS)]
-
-        fig.add_trace(
-            go.Bar(
-                x=positions,
-                y=weights,
-                marker_color=color,
-                name=f"Chain {chain_labels[i]}",
-                hovertemplate="Pos %{x}: %{customdata}<br>Weight: %{y:.4f}<extra></extra>",
-                customdata=residues,
-            ),
-            row=i + 1, col=1,
-        )
-        fig.update_xaxes(title_text="Residue position", row=i + 1, col=1)
-        fig.update_yaxes(title_text="Weight", row=i + 1, col=1)
-
-    fig.update_layout(
-        title="Residue-level Interface Prediction Weights",
-        height=max(300 * n_chains, 400),
-        showlegend=False,
-    )
-    return fig
+        for pos, w in enumerate(weights, 1):
+            records.append({
+                "Position": pos,
+                "Weight": float(w),
+                "Chain": f"Chain {chain_labels[i]}",
+            })
+    return pd.DataFrame(records)
 
 
 with gr.Blocks(title="Stoic - Protein Stoichiometry Prediction") as app:
@@ -144,7 +115,15 @@ with gr.Blocks(title="Stoic - Protein Stoichiometry Prediction") as app:
             results_output = gr.Markdown(value="Results will appear here.")
             run_time = gr.Textbox(label="Runtime")
 
-    weights_plot = gr.Plot(label="Residue-level Interface Weights", visible=False)
+    weights_plot = gr.LinePlot(
+        x="Position",
+        y="Weight",
+        color="Chain",
+        x_title="Residue Position",
+        y_title="Weight",
+        label="Residue-level Interface Weights",
+        visible=False,
+    )
 
     btn.click(
         predict,