followup

app.py

@@ -120,7 +120,7 @@ def tensor_to_pil(tensor):
     arr = (arr * 255).astype(np.uint8)
     return Image.fromarray(arr, mode='L')
 
-def format_top_k(logits, top_k=5):
+def format_top_k(logits, top_k=20):
     """Return list of (token, probability) tuples."""
     probs = F.softmax(logits, dim=-1)
     top_probs, top_indices = torch.topk(probs, top_k, dim=-1)
@@ -133,7 +133,7 @@ def format_top_k(logits, top_k=5):
 
 def run_inference(sample_idx):
     if dataset is None:
-        return None, "Dataset not loaded", [], [], [], [], None
+        return None, "Dataset not loaded", None, {}
     
     # Load sample
     sample_idx = int(sample_idx) # ensure int
@@ -217,19 +217,19 @@ def run_inference(sample_idx):
     else:
         mmrm_res = [("Model not loaded (custom weight specific)", 0.0)]
     
-    # Format text outputs as string or dictionary for Label
-    def format_output(results):
-        out_dict = {k: v for k, v in results}
-        return out_dict
+    # Format raw results into a dictionary for State
+    raw_results = {
+        'zs': zs_res,
+        'text': text_res,
+        'visual': visual_res,
+        'mmrm': mmrm_res
+    }
 
     return (
         input_display_image,
         f"Context: {context_text}\nGround Truth: {ground_truth_text}",
-        format_output(zs_res),
-        format_output(text_res),
-        format_output(visual_res),
-        format_output(mmrm_res),
-        restored_pil
+        restored_pil,
+        raw_results
     )
 
 
@@ -240,12 +240,12 @@ with gr.Blocks(title="MMRM Demo", theme=gr.themes.Soft(spacing_size="sm", text_s
     gr.Markdown("Comparing MMRM with baselines on real-world damaged characters.")
     
     with gr.Row():
-        # --- Left Column: Inputs ---
+    # --- Left Column: Inputs ---
         with gr.Column(scale=1):
             gr.Markdown("### Input Selection")
             with gr.Row():
                 sample_dropdown = gr.Dropdown(
-                    choices=[x[1] for x in sample_options], # Use index as value
+                    choices=[x[1] for x in sample_options], 
                     type="value",
                     label="Select Sample",
                     container=False,
@@ -254,6 +254,9 @@ with gr.Blocks(title="MMRM Demo", theme=gr.themes.Soft(spacing_size="sm", text_s
                 sample_dropdown.choices = sample_options
                 run_btn = gr.Button("Run", variant="primary", scale=1, min_width=60)
             
+            with gr.Row():
+                top_k_slider = gr.Slider(minimum=1, maximum=20, value=5, step=1, label="Top K Predictions")
+
             with gr.Row():
                 input_image = gr.Image(label="Damaged Input", type="pil", height=250)
             
@@ -265,24 +268,59 @@ with gr.Blocks(title="MMRM Demo", theme=gr.themes.Soft(spacing_size="sm", text_s
             gr.Markdown("### Model Predictions")
             with gr.Row():
                 with gr.Column(min_width=80):
-                    zs_output = gr.Label(num_top_classes=3, label="Zero-shot")
+                    zs_output = gr.Label(num_top_classes=20, label="Zero-shot")
                 with gr.Column(min_width=80):
-                    text_output = gr.Label(num_top_classes=3, label="Textual")
+                    text_output = gr.Label(num_top_classes=20, label="Textual")
                 with gr.Column(min_width=80):
-                    visual_output = gr.Label(num_top_classes=3, label="Visual")
+                    visual_output = gr.Label(num_top_classes=20, label="Visual")
                 with gr.Column(min_width=80):
-                    mmrm_output = gr.Label(num_top_classes=3, label="MMRM")
+                    mmrm_output = gr.Label(num_top_classes=20, label="MMRM")
             
             with gr.Row():
                 with gr.Column():
                     gr.Markdown("### Visual Restoration")
                     restored_output = gr.Image(label="MMRM Output", type="pil", height=250)
+    
+    # State to hold raw top-20 results for all models
+    # Structure: {"zs": [...], "text": [...], "visual": [...], "mmrm": [...]}
+    raw_results_state = gr.State()
+
+    def update_views(raw_results, k):
+        if not raw_results:
+            return {}, {}, {}, {}
+        
+        k = int(k)
+        def slice_res(key):
+            # Take top k from list of tuples
+            full_list = raw_results.get(key, [])
+            return {term: score for term, score in full_list[:k]}
+            
+        return (
+            slice_res('zs'),
+            slice_res('text'),
+            slice_res('visual'),
+            slice_res('mmrm')
+        )
 
-    # Event
-    run_btn.click(
+    # Event Chain
+    # 1. Run inference -> updates State and Images/Text
+    run_event = run_btn.click(
         fn=run_inference,
         inputs=[sample_dropdown],
-        outputs=[input_image, input_text, zs_output, text_output, visual_output, mmrm_output, restored_output]
+        outputs=[input_image, input_text, restored_output, raw_results_state]
+    )
+    
+    # 2. Update Labels based on State and Slider (triggered by Run success OR Slider change)
+    run_event.success(
+        fn=update_views,
+        inputs=[raw_results_state, top_k_slider],
+        outputs=[zs_output, text_output, visual_output, mmrm_output]
+    )
+    
+    top_k_slider.change(
+        fn=update_views,
+        inputs=[raw_results_state, top_k_slider],
+        outputs=[zs_output, text_output, visual_output, mmrm_output]
     )