Add 'load first example' button

app.py

@@ -154,6 +154,20 @@ def format_stats_display(judgment_times, num_judgments, num_skips):
     
     return stats
 
+def load_first_example():
+    """Load the first example and start the annotation interface"""
+    prompt, comp_1, comp_2, idx, instruction, completion_1, completion_2 = get_random_example()
+    start_time = time.time()
+    
+    return (
+        prompt, comp_1, comp_2, idx, instruction, completion_1, completion_2,
+        start_time,
+        gr.update(visible=False),  # Hide load button
+        gr.update(visible=True),   # Show prompt
+        gr.update(visible=True),   # Show completion row
+        gr.update(visible=True),   # Show action buttons
+    )
+
 def handle_left_better(prompt, completion_1_display, completion_2_display, current_idx, instruction, completion_1, completion_2, 
                       start_time, judgment_times, num_judgments, num_skips):
     """Handle when user selects left completion as better"""
@@ -173,15 +187,15 @@ def handle_left_better(prompt, completion_1_display, completion_2_display, curre
     # Get new random example
     new_prompt, new_comp_1, new_comp_2, new_idx, new_instruction, new_completion_1, new_completion_2 = get_random_example()
     
+    # Reset timer for new example
+    new_start_time = time.time()
+    
     # Update stats display
     stats_display = format_stats_display(judgment_times, num_judgments, num_skips)
 
     message = "✅ Annotation saved! Left completion selected as better." if success else "✅ Left completion selected (save failed - check console)"
     gr.Info(message)
     
-    # Reset timer for new example
-    new_start_time = time.time()
-    
     return (
         new_prompt, 
         new_comp_1, 
@@ -216,15 +230,15 @@ def handle_right_better(prompt, completion_1_display, completion_2_display, curr
     # Get new random example
     new_prompt, new_comp_1, new_comp_2, new_idx, new_instruction, new_completion_1, new_completion_2 = get_random_example()
     
+    # Reset timer for new example
+    new_start_time = time.time()
+    
     # Update stats display
     stats_display = format_stats_display(judgment_times, num_judgments, num_skips)
 
     message = "✅ Annotation saved! Right completion selected as better." if success else "✅ Right completion selected (save failed - check console)"
     gr.Info(message)
     
-    # Reset timer for new example
-    new_start_time = time.time()
-    
     return (
         new_prompt, 
         new_comp_1, 
@@ -280,11 +294,23 @@ if HF_TOKEN:
 else:
     print("Warning: No HF_TOKEN found. Annotations will not be saved.")
 
-# Initialize with first random example
-init_prompt, init_comp_1, init_comp_2, init_idx, init_instruction, init_completion_1, init_completion_2 = get_random_example()
+def load_first_example():
+    """Load the first example and start the annotation interface"""
+    prompt, comp_1, comp_2, idx, instruction, completion_1, completion_2 = get_random_example()
+    start_time = time.time()
+    
+    return (
+        prompt, comp_1, comp_2, idx, instruction, completion_1, completion_2,
+        start_time,
+        gr.update(visible=False),  # Hide load button
+        gr.update(visible=True),   # Show prompt
+        gr.update(visible=True),   # Show completion row
+        gr.update(visible=True),   # Show action buttons
+    )
 
 # Create Gradio interface
-with gr.Blocks(title="AI Alignment: Binary Preference Annotation") as demo:
+
+with gr.Blocks(title="AI Alignment: Binary Preference Annotation", css=".square-button { height: 80px !important; }") as demo:
     gr.Markdown(f"""
 # 🎯 AI Alignment: Binary Preference Annotation
 
@@ -297,29 +323,34 @@ This simulates the data annotation process used in RLHF (Reinforcement Learning
 """)
     
     # State to track current example and its components
-    current_idx = gr.State(init_idx)
-    current_instruction = gr.State(init_instruction)
-    current_completion_1 = gr.State(init_completion_1)
-    current_completion_2 = gr.State(init_completion_2)
+    current_idx = gr.State(0)
+    current_instruction = gr.State("")
+    current_completion_1 = gr.State("")
+    current_completion_2 = gr.State("")
     
     # State to track timing and statistics
-    start_time = gr.State(time.time())  # When current example was loaded
+    start_time = gr.State(0.0)  # When current example was loaded
     judgment_times = gr.State([])  # List of times taken for each judgment
     num_judgments = gr.State(0)  # Number of judgments made
     num_skips = gr.State(0)  # Number of examples skipped
     
-    # Display prompt
-    prompt_display = gr.Markdown(init_prompt, label="Prompt")
+    # Load first example button (shown initially)
+    load_first_btn = gr.Button("🚀 Load First Example", variant="primary", size="lg", elem_classes="square-button")
+    
+    # Display prompt (hidden initially)
+    prompt_display = gr.Markdown("", label="Prompt", visible=False)
     
-    # Display completions side by side
-    with gr.Row():
+    # Display completions side by side (hidden initially)
+    completion_row = gr.Row(visible=False)
+    with completion_row:
         with gr.Column():
-            completion_1_display = gr.Markdown(init_comp_1, label="Completion A (Left)")
+            completion_1_display = gr.Markdown("", label="Completion A (Left)")
         with gr.Column():
-            completion_2_display = gr.Markdown(init_comp_2, label="Completion B (Right)")
+            completion_2_display = gr.Markdown("", label="Completion B (Right)")
     
-    # Action buttons
-    with gr.Row():
+    # Action buttons (hidden initially)
+    action_buttons = gr.Row(visible=False)
+    with action_buttons:
         left_better_btn = gr.Button("👈 Left is Better", variant="primary", size="lg")
         skip_btn = gr.Button("⏭️ Skip This Example", variant="secondary", size="lg")
         right_better_btn = gr.Button("👉 Right is Better", variant="primary", size="lg")
@@ -333,7 +364,15 @@ This simulates the data annotation process used in RLHF (Reinforcement Learning
     # Statistics display
     stats_display = gr.Markdown("📊 **Session Statistics:** No judgments made yet.", label="Performance Stats")
     
-    # Wire up the buttons
+    # Wire up the load first example button
+    load_first_btn.click(
+        load_first_example,
+        inputs=[],
+        outputs=[prompt_display, completion_1_display, completion_2_display, current_idx, current_instruction, current_completion_1, current_completion_2,
+                 start_time, load_first_btn, prompt_display, completion_row, action_buttons]
+    )
+    
+    # Wire up the action buttons
     left_better_btn.click(
         handle_left_better,
         inputs=[prompt_display, completion_1_display, completion_2_display, current_idx, current_instruction, current_completion_1, current_completion_2,