Added Context preload and inference interruption (#2)

- Added Context preload and inference interruption (189d8573833ade641554ffcb21e5f5e30412ea75)


Co-authored-by: Kai <kai-aizip@users.noreply.huggingface.co>

app.py

@@ -2,17 +2,17 @@ import gradio as gr
 import random
 import pandas as pd
 import os
+import threading
+from threading import Event
 from utils.data_loader import get_random_example
 from utils.models import generate_summaries, model_names
 from utils.ui_helpers import toggle_context_display, update_feedback, get_context_html
 from utils.leaderboard import load_leaderboard_data, save_leaderboard_data
 
-# Read CSS from file
-css_path = os.path.join(os.getcwd(), 'static', 'styles.css')
-with open(css_path, 'r') as f:
-    css_content = f.read()
+# Global interrupt mechanism for model generation
+generation_interrupt = Event()
 
-# Feedback options
+# Feedback options for different voting outcomes
 feedback_options = {
     "left": ["Model A: More complete", "Model A: More accurate", "Model A: More relevant", "Model A: Better written", "Model A: Better refusal (if applicable)"],
     "right": ["Model B: More complete", "Model B: More accurate", "Model B: More relevant", "Model B: Better written", "Model B: Better refusal (if applicable)"],
@@ -20,15 +20,11 @@ feedback_options = {
     "neither": ["Both incomplete", "Both hallucinate", "Both irrelevant", "Both incorrectly refuse (if applicable)", "A is bad", "B is bad"]
 }
 
-def load_new_question_improved(agg_results=None, show_full=False):
-    """Loads a new random question, contexts, and model summaries."""
-    if agg_results is None:
-        agg_results = load_leaderboard_data()
-    
+def load_context():
+    """Load a new question and context (fast operation)"""
+    generation_interrupt.clear()
     example = get_random_example()
-    m_a_name, m_b_name = random.sample(model_names, 2)
-    s_a, s_b = generate_summaries(example, m_a_name, m_b_name)
-
+    
     context_desc = example.get('processed_context_desc', '')
     if context_desc:
         context_desc = f"<div class='context-topic'><span class='topic-label'>The question and context are about:</span> {context_desc}</div>"
@@ -37,36 +33,98 @@ def load_new_question_improved(agg_results=None, show_full=False):
     context_html = get_context_html(example, show_full=show_full)
     
     return [
-        example,  # current_example
-        m_a_name,  # model_a_name
-        m_b_name,  # model_b_name
-        s_a,  # summary_a_text
-        s_b,  # summary_b_text
-        None,  # selected_winner
-        [],  # feedback_list
-        False,  # show_results_state
-        agg_results,  # results_agg
-        show_full,  # show_full_context
-        gr.update(value=example['question']),  # query_display
-        gr.update(value=context_desc, visible=bool(context_desc)),  # context_description
-        gr.update(value=context_html),  # context_display
-        gr.update(value="Show Full Context", elem_classes=["context-toggle-button"]),  # context_toggle_btn
-        gr.update(value=s_a),  # summary_a_display
-        gr.update(value=s_b),  # summary_b_display
-        gr.update(interactive=True, elem_classes=["vote-button"]),  # vote_button_a
-        gr.update(interactive=True, elem_classes=["vote-button"]),  # vote_button_b
-        gr.update(interactive=True, elem_classes=["vote-button"]),  # vote_button_tie
-        gr.update(interactive=True, elem_classes=["vote-button", "vote-button-neither"]),  # vote_button_neither
-        gr.update(choices=[], value=[], interactive=False, visible=False),  # feedback_checkboxes
-        gr.update(visible=False),  # feedback_section
-        gr.update(interactive=False, visible=True),  # submit_button
-        gr.update(visible=False),  # results_reveal_area
-        gr.update(interactive=True),  # random_question_btn
-        gr.update(elem_classes=[])  # main_interface_area
+        example,
+        gr.update(value=example['question']),
+        gr.update(value=context_desc, visible=bool(context_desc)),
+        gr.update(value=context_html),
+        gr.update(value="Show Full Context", elem_classes=["context-toggle-button"]),
+        show_full
+    ]
+
+def generate_model_summaries_with_timeout(example, timeout=30):
+    """Run model inference in a separate thread with timeout for interruptibility"""
+    import threading
+    import time
+    
+    result = {
+        "model_a": "",
+        "model_b": "",
+        "summary_a": "",
+        "summary_b": "",
+        "completed": False
+    }
+    
+    if generation_interrupt.is_set():
+        return result
+
+    def run_generation():
+        try:
+            m_a_name, m_b_name = random.sample(model_names, 2)
+            s_a, s_b = generate_summaries(example, m_a_name, m_b_name)
+            
+            if not generation_interrupt.is_set():
+                result["model_a"] = m_a_name
+                result["model_b"] = m_b_name
+                result["summary_a"] = s_a
+                result["summary_b"] = s_b
+                result["completed"] = True
+        except Exception as e:
+            print(f"Error in generation thread: {e}")
+    
+    generation_thread = threading.Thread(target=run_generation)
+    generation_thread.daemon = True
+    generation_thread.start()
+    
+    start_time = time.time()
+    while time.time() - start_time < timeout:
+        if generation_interrupt.is_set() or not generation_thread.is_alive() or result["completed"]:
+            break
+        time.sleep(0.1)
+        
+    return result
+
+def process_generation_result(result):
+    """Process the results from the threaded generation function"""
+    if not result["completed"]:
+        # Generation was interrupted or failed
+        return [
+            "", "", "", "", None, [], False, load_leaderboard_data(),
+            gr.update(value="Generation was interrupted or timed out. Please try again."),
+            gr.update(value="Generation was interrupted or timed out. Please try again."),
+            gr.update(interactive=True, elem_classes=["vote-button"]),
+            gr.update(interactive=True, elem_classes=["vote-button"]),
+            gr.update(interactive=True, elem_classes=["vote-button"]),
+            gr.update(interactive=True, elem_classes=["vote-button", "vote-button-neither"]),
+            gr.update(choices=[], value=[], interactive=False, visible=False),
+            gr.update(visible=False),
+            gr.update(interactive=False, visible=True),
+            gr.update(visible=False),
+            gr.update(interactive=True),
+            gr.update(elem_classes=[])
+        ]
+    
+    # Generation completed successfully
+    agg_results = load_leaderboard_data()
+    return [
+        result["model_a"], result["model_b"], 
+        result["summary_a"], result["summary_b"],
+        None, [], False, agg_results,
+        gr.update(value=result["summary_a"]),
+        gr.update(value=result["summary_b"]),
+        gr.update(interactive=True, elem_classes=["vote-button"]),
+        gr.update(interactive=True, elem_classes=["vote-button"]),
+        gr.update(interactive=True, elem_classes=["vote-button"]),
+        gr.update(interactive=True, elem_classes=["vote-button", "vote-button-neither"]),
+        gr.update(choices=[], value=[], interactive=False, visible=False),
+        gr.update(visible=False),
+        gr.update(interactive=False, visible=True),
+        gr.update(visible=False),
+        gr.update(interactive=True),
+        gr.update(elem_classes=[])
     ]
 
 def select_vote_improved(winner_choice):
-    """Handles vote button selections."""
+    """Updates UI based on vote selection"""
     feedback_choices = feedback_options.get(winner_choice, [])
 
     btn_a_classes = ["vote-button"]
@@ -84,18 +142,18 @@ def select_vote_improved(winner_choice):
         btn_neither_classes.append("selected")
 
     return [
-        winner_choice,  # selected_winner
-        gr.update(choices=feedback_choices, value=[], interactive=True, visible=True),  # feedback_checkboxes
-        gr.update(visible=True),  # feedback_section
-        gr.update(interactive=True),  # submit_button
-        gr.update(elem_classes=btn_a_classes),  # vote_button_a
-        gr.update(elem_classes=btn_b_classes),  # vote_button_b
-        gr.update(elem_classes=btn_tie_classes),  # vote_button_tie
-        gr.update(elem_classes=btn_neither_classes)  # vote_button_neither
+        winner_choice,
+        gr.update(choices=feedback_choices, value=[], interactive=True, visible=True),
+        gr.update(visible=True),
+        gr.update(interactive=True),
+        gr.update(elem_classes=btn_a_classes),
+        gr.update(elem_classes=btn_b_classes),
+        gr.update(elem_classes=btn_tie_classes),
+        gr.update(elem_classes=btn_neither_classes)
     ]
 
 def submit_vote_fixed(m_a, m_b, winner, feedback, current_results):
-    """Processes vote submission and updates results."""
+    """Processes vote and updates leaderboard"""
     if winner is None:
         print("Warning: Submit called without a winner selected.")
         return {}
@@ -123,7 +181,9 @@ def submit_vote_fixed(m_a, m_b, winner, feedback, current_results):
 
     # Prepare Results Table
     results_list = []
-    all_models = list(set(list(updated_results["wins"].keys()) + list(updated_results["losses"].keys()) + list(updated_results["ties"].keys())))
+    all_models = list(set(list(updated_results["wins"].keys()) + 
+                          list(updated_results["losses"].keys()) + 
+                          list(updated_results["ties"].keys())))
 
     for model in sorted(all_models):
         wins = updated_results["wins"].get(model, 0)
@@ -146,34 +206,26 @@ def submit_vote_fixed(m_a, m_b, winner, feedback, current_results):
         results_df = results_df.sort_values(by='Win Rate Value', ascending=False).drop(columns=['Win Rate Value'])
 
     return [
-        True,  # show_results_state
-        updated_results,  # results_agg
-        gr.update(interactive=False),  # vote_button_a
-        gr.update(interactive=False),  # vote_button_b
-        gr.update(interactive=False),  # vote_button_tie
-        gr.update(interactive=False),  # vote_button_neither
-        gr.update(interactive=False),  # feedback_checkboxes
-        gr.update(visible=True),  # feedback_section
-        gr.update(visible=False),  # submit_button
-        gr.update(visible=True),  # results_reveal_area
-        gr.update(interactive=False),  # random_question_btn
-        gr.update(value=results_df, visible=True),  # results_table_display
-        gr.update(elem_classes=["results-revealed"]),  # main_interface_area
-        gr.update(interactive=True),  # context_toggle_btn
-        gr.update(value=m_a),  # model_a_reveal
-        gr.update(value=m_b)  # model_b_reveal
+        True, updated_results,
+        gr.update(interactive=False), gr.update(interactive=False),
+        gr.update(interactive=False), gr.update(interactive=False),
+        gr.update(interactive=False), gr.update(visible=True),
+        gr.update(visible=False), gr.update(visible=True),
+        gr.update(interactive=False), gr.update(value=results_df, visible=True),
+        gr.update(elem_classes=["results-revealed"]),
+        gr.update(interactive=True), gr.update(value=m_a), gr.update(value=m_b)
     ]
 
-# Create embedded CSS
-css_html = f"<style>{css_content}</style>"
-
 # Create Gradio interface
 with gr.Blocks(theme=gr.themes.Default(
     primary_hue=gr.themes.colors.orange,
     secondary_hue=gr.themes.colors.slate
 )) as demo:
-    # Embed CSS directly in HTML
-    gr.HTML(css_html)
+    # Load CSS
+    css_path = os.path.join(os.getcwd(), 'static', 'styles.css')
+    with open(css_path, 'r') as f:
+        css_content = f.read()
+    gr.HTML(f"<style>{css_content}</style>")
 
     # State Variables
     current_example = gr.State({})
@@ -191,7 +243,6 @@ with gr.Blocks(theme=gr.themes.Default(
     with gr.Tabs() as tabs:
         # Main Arena Tab
         with gr.TabItem("Arena", id="arena-tab"):
-            # Main title and description
             gr.Markdown("# RAG Summarizer Arena")
             gr.Markdown("Compare summaries generated by different models based on the provided context and query. Select the better summary, or choose 'Tie' or 'Neither'. Your feedback helps evaluate model performance.")
 
@@ -206,10 +257,9 @@ with gr.Blocks(theme=gr.themes.Default(
                         query_display = gr.Markdown(value="Loading question...", elem_classes="query-text")
                     random_question_btn = gr.Button("🔄 Get Random Question", elem_classes="query-button")
                 
-                # Context description
+                # Context description and display
                 context_description = gr.Markdown("", elem_classes="context-description")
-
-                # Context section
+                
                 with gr.Row(elem_id="context-header-row"):
                     gr.Markdown("### Context Provided", elem_classes="context-title")
                     context_toggle_btn = gr.Button("Show Full Context", elem_classes=["context-toggle-button"])
@@ -236,11 +286,9 @@ with gr.Blocks(theme=gr.themes.Default(
                     vote_button_b = gr.Button("➡️ Summary B is Better", elem_classes=["vote-button"])
                     vote_button_neither = gr.Button("❌ Neither is Adequate", elem_classes=["vote-button", "vote-button-neither"])
 
-                # Feedback section
+                # Feedback and Submit sections
                 with gr.Group(elem_classes=["feedback-section"], visible=False) as feedback_section:
                     feedback_checkboxes = gr.CheckboxGroup(label="Feedback (optional)", choices=[], interactive=False)
-
-                # Submit button
                 submit_button = gr.Button("Submit Vote", variant="primary", interactive=False, elem_id="submit-button")
 
                 # Results area
@@ -258,7 +306,7 @@ with gr.Blocks(theme=gr.themes.Default(
                             model_b_reveal = gr.Markdown("", elem_classes="model-reveal model-b-reveal")
                      
                     gr.HTML("<div style='height: 10px;'></div>")
-                     
+                    
                     # Try another button
                     with gr.Row(elem_classes=["control-buttons"]):
                         try_another_btn = gr.Button("🔄 Try Another Question", elem_id="try-another-btn")
@@ -269,99 +317,93 @@ with gr.Blocks(theme=gr.themes.Default(
             gr.Markdown("View aggregate performance statistics for all models. The table below shows win rates, wins, losses, and ties for each model based on all evaluations.")
             results_table_display = gr.DataFrame(label="Model Performance", interactive=False, wrap=True)
 
-    # Event Listeners
+    # Generic function to handle starting a new example
+    def handle_new_example_click():
+        generation_interrupt.set()  # Interrupt any ongoing generation
+        return load_context()[0]
+        
+    def update_ui_for_new_context(example):
+        return [
+            gr.update(value=example['question']),
+            gr.update(value=example.get('processed_context_desc', ''), visible=bool(example.get('processed_context_desc', ''))),
+            gr.update(value=get_context_html(example, False)),
+            gr.update(value="Show Full Context", elem_classes=["context-toggle-button"]),
+            False
+        ]
+
+    # Event handling
+    # Toggle context display
     context_toggle_btn.click(
         fn=toggle_context_display,
         inputs=[current_example, show_full_context],
         outputs=[show_full_context, context_display, context_toggle_btn]
     )
     
+    # Initial loading - context first, then summaries
     demo.load(
-        fn=load_new_question_improved,
+        fn=load_context,
         inputs=[],
-        outputs=[
-            current_example, model_a_name, model_b_name, summary_a_text, summary_b_text,
-            selected_winner, feedback_list, show_results_state, results_agg, show_full_context,
-            query_display, context_description, context_display, context_toggle_btn,
-            summary_a_display, summary_b_display,
-            vote_button_a, vote_button_b, vote_button_tie, vote_button_neither,
-            feedback_checkboxes, feedback_section, submit_button, results_reveal_area, random_question_btn,
-            main_interface_area
-        ]
-    )
-
-    random_question_btn.click(
-        fn=load_new_question_improved,
-        inputs=[],
-        outputs=[
-            current_example, model_a_name, model_b_name, summary_a_text, summary_b_text,
-            selected_winner, feedback_list, show_results_state, results_agg, show_full_context,
-            query_display, context_description, context_display, context_toggle_btn,
-            summary_a_display, summary_b_display,
-            vote_button_a, vote_button_b, vote_button_tie, vote_button_neither,
-            feedback_checkboxes, feedback_section, submit_button, results_reveal_area, random_question_btn,
-            main_interface_area
-        ]
-    )
-
-    vote_button_a.click(
-        fn=lambda: select_vote_improved('left'),
-        inputs=None,
-        outputs=[selected_winner, feedback_checkboxes, feedback_section, submit_button, vote_button_a, vote_button_b, vote_button_tie, vote_button_neither]
-    )
-    vote_button_b.click(
-        fn=lambda: select_vote_improved('right'),
-        inputs=None,
-        outputs=[selected_winner, feedback_checkboxes, feedback_section, submit_button, vote_button_a, vote_button_b, vote_button_tie, vote_button_neither]
-    )
-    vote_button_tie.click(
-        fn=lambda: select_vote_improved('tie'),
-        inputs=None,
-        outputs=[selected_winner, feedback_checkboxes, feedback_section, submit_button, vote_button_a, vote_button_b, vote_button_tie, vote_button_neither]
-    )
-    vote_button_neither.click(
-        fn=lambda: select_vote_improved('neither'),
-        inputs=None,
-        outputs=[selected_winner, feedback_checkboxes, feedback_section, submit_button, vote_button_a, vote_button_b, vote_button_tie, vote_button_neither]
+        outputs=[current_example, query_display, context_description, context_display, 
+                context_toggle_btn, show_full_context]
+    ).then(
+        fn=lambda example: process_generation_result(generate_model_summaries_with_timeout(example)),
+        inputs=[current_example],
+        outputs=[model_a_name, model_b_name, summary_a_text, summary_b_text,
+                selected_winner, feedback_list, show_results_state, results_agg,
+                summary_a_display, summary_b_display, vote_button_a, vote_button_b, 
+                vote_button_tie, vote_button_neither, feedback_checkboxes, feedback_section, 
+                submit_button, results_reveal_area, random_question_btn, main_interface_area]
     )
 
+    # Random Question and Try Another buttons with interruption
+    for btn in [random_question_btn, try_another_btn]:
+        btn.click(
+            fn=handle_new_example_click,
+            inputs=[],
+            outputs=[current_example]
+        ).then(
+            fn=update_ui_for_new_context,
+            inputs=[current_example],
+            outputs=[query_display, context_description, context_display, 
+                    context_toggle_btn, show_full_context]
+        ).then(
+            fn=lambda example: process_generation_result(generate_model_summaries_with_timeout(example)),
+            inputs=[current_example],
+            outputs=[model_a_name, model_b_name, summary_a_text, summary_b_text,
+                    selected_winner, feedback_list, show_results_state, results_agg,
+                    summary_a_display, summary_b_display, vote_button_a, vote_button_b, 
+                    vote_button_tie, vote_button_neither, feedback_checkboxes, feedback_section, 
+                    submit_button, results_reveal_area, random_question_btn, main_interface_area]
+        )
+
+    # Vote button handlers
+    for btn, choice in zip(
+        [vote_button_a, vote_button_b, vote_button_tie, vote_button_neither],
+        ['left', 'right', 'tie', 'neither']
+    ):
+        btn.click(
+            fn=lambda choice=choice: select_vote_improved(choice),
+            inputs=None,
+            outputs=[selected_winner, feedback_checkboxes, feedback_section, submit_button, 
+                    vote_button_a, vote_button_b, vote_button_tie, vote_button_neither]
+        )
+
+    # Update feedback when checkboxes change
     feedback_checkboxes.change(
         fn=update_feedback,
         inputs=[feedback_checkboxes],
         outputs=[feedback_list]
     )
 
+    # Process vote submission and reveal results
     submit_button.click(
         fn=submit_vote_fixed,
         inputs=[model_a_name, model_b_name, selected_winner, feedback_list, results_agg],
-        outputs=[
-            show_results_state, results_agg,
-            vote_button_a, vote_button_b, vote_button_tie, vote_button_neither,
-            feedback_checkboxes,
-            feedback_section,
-            submit_button,
-            results_reveal_area,
-            random_question_btn,
-            results_table_display,
-            main_interface_area,
-            context_toggle_btn,
-            model_a_reveal,
-            model_b_reveal
-        ]
-    )
-
-    try_another_btn.click(
-        fn=load_new_question_improved,
-        inputs=[],
-        outputs=[
-            current_example, model_a_name, model_b_name, summary_a_text, summary_b_text,
-            selected_winner, feedback_list, show_results_state, results_agg, show_full_context,
-            query_display, context_description, context_display, context_toggle_btn,
-            summary_a_display, summary_b_display,
-            vote_button_a, vote_button_b, vote_button_tie, vote_button_neither,
-            feedback_checkboxes, feedback_section, submit_button, results_reveal_area, random_question_btn,
-            main_interface_area
-        ]
+        outputs=[show_results_state, results_agg, vote_button_a, vote_button_b, 
+                vote_button_tie, vote_button_neither, feedback_checkboxes,
+                feedback_section, submit_button, results_reveal_area,
+                random_question_btn, results_table_display, main_interface_area,
+                context_toggle_btn, model_a_reveal, model_b_reveal]
     )
 
 if __name__ == "__main__":