Update app.py

app.py

@@ -390,7 +390,7 @@ def _format_component_table(
         )
 
     return (
-        "<h4 style='margin:12px 0 6px'></h4>"
+        "<h4 style='margin:12px 0 6px'>Highlighted Binding site</h4>"
         "<table style='border-collapse:collapse;margin:6px 0 16px;width:100%'>"
         "<thead><tr style='background:#f5f5f5'>"
         "<th style='border:1px solid #ddd;padding:6px'>Protein range</th>"
@@ -540,11 +540,15 @@ def visualize_attention_and_ranges(
         # --------------------------------------------------
         # Connected components (visual coherence)
         # --------------------------------------------------
+        p_tokens_orig = p_tokens.copy()
+        d_tokens_orig = d_tokens.copy()
+        
         components = _connected_components_2d(mask_top)
+        
         ranges_html = _format_component_table(
             components,
-            p_tokens,
-            d_tokens,
+            p_tokens_orig,
+            d_tokens_orig,
             mode=mode,
         )
 
@@ -857,17 +861,25 @@ def inference_cb(prot_seq, drug_seq, head_choice, topk_choice, mode_choice):
     # ------------------------------
     # Parse mode (pair vs residue)
     # ------------------------------
-    mode_choice = str(mode_choice).lower()
-
-    if "residue" in mode_choice:
+    mode_choice_clean = str(mode_choice).strip()
+    
+    if mode_choice_clean == "Top-K residues":
         mode = "residue"
         topk_pairs = 1
         topk_residues = min(20, topk)
+    
+    elif mode_choice_clean == "Top-K residues-atom pairs":
+        mode = "pair"
+        topk_pairs = topk
+        topk_residues = 1
+    
     else:
+        # fallback
         mode = "pair"
         topk_pairs = topk
         topk_residues = 1
 
+
     # ------------------------------
     # Visualisation
     # ------------------------------
@@ -1011,8 +1023,12 @@ h1{
   margin: 0 auto;
   justify-content: center;
 }
+
+
+}
 """
 with gr.Blocks() as demo:
+
     gr.Markdown("<h1>ExplainBind: Token-level Protein–Ligand Interaction Visualiser</h1>")
     gr.Markdown(
         '<p class="subtle">'
@@ -1021,8 +1037,6 @@ with gr.Blocks() as demo:
         '</p>'
     )
 
-
-    # ===== Loscalzo logo + group name =====
     gr.HTML(f"""
     <div class="loscalzo-block">
       <img src="data:image/png;base64,{LOSCAZLO_B64}"