Update app.py

app.py

@@ -54,162 +54,288 @@ def create_reranking_interface(task_data):
             gr.Markdown("## Query:")
             query_text = gr.Textbox(value=samples[0]["query"], label="", interactive=False, lines=3)
             
-            # Validation
-            with gr.Row():
-                validate_btn = gr.Button("Validate Rankings", variant="secondary")
-                validation_text = gr.Textbox(label="Validation", interactive=False)
-            
             # Document ranking section
             gr.Markdown("## Documents to Rank:")
             
-            # Container for document elements
-            doc_containers = []
-            rank_inputs = []
-            doc_texts = []
+            # Create simple data structure for documents
+            doc_state = gr.State(value=samples[0]["candidates"])
             
-            # Create a container for up to 10 documents
-            max_docs = 10
-            for i in range(max_docs):
-                with gr.Group(visible=(i < len(samples[0]["candidates"]))) as doc_container:
-                    doc_containers.append(doc_container)
+            # Create dynamic HTML for the ranking interface
+            def generate_ranking_html(docs, existing_ranks=None):
+                """Generate HTML for number-based ranking interface."""
+                if not docs:
+                    return ""
+                
+                # Use existing ranks if available
+                ranks = list(range(1, len(docs) + 1))
+                if existing_ranks and len(existing_ranks) == len(docs):
+                    ranks = existing_ranks
+                
+                html = """
+                <style>
+                .doc-container {
+                    margin-bottom: 15px;
+                    border: 1px solid #ddd;
+                    border-radius: 8px;
+                    padding: 15px;
+                    background-color: #f9f9f9;
+                }
+                .doc-header {
+                    display: flex;
+                    align-items: center;
+                    margin-bottom: 10px;
+                }
+                .doc-rank {
+                    display: flex;
+                    align-items: center;
+                    margin-right: 15px;
+                }
+                .rank-label {
+                    font-weight: bold;
+                    margin-right: 8px;
+                    min-width: 80px;
+                }
+                .rank-input {
+                    width: 60px;
+                    padding: 5px;
+                    border: 2px solid #007bff;
+                    border-radius: 4px;
+                    text-align: center;
+                    font-size: 16px;
+                }
+                .doc-content {
+                    padding: 10px;
+                    background-color: white;
+                    border-radius: 4px;
+                    border-left: 4px solid #007bff;
+                    white-space: pre-wrap;
+                    font-family: sans-serif;
+                    line-height: 1.5;
+                }
+                </style>
+                
+                <div id="ranking-form">
+                <input type="hidden" id="ranking-state" value="">
+                """
+                
+                # Add each document with a number input
+                for i, doc in enumerate(docs):
+                    import html as html_lib
+                    escaped_doc = html_lib.escape(doc)
+                    current_rank = ranks[i] if i < len(ranks) else i + 1
                     
-                    with gr.Row():
-                        # Rank selection
-                        with gr.Column(scale=1, min_width=100):
-                            rank_input = gr.Number(
-                                value=i+1, 
-                                label=f"Rank", 
-                                minimum=1,
-                                maximum=len(samples[0]["candidates"]),
-                                step=1,
-                                interactive=True
-                            )
-                            rank_inputs.append(rank_input)
+                    html += f"""
+                    <div class="doc-container" id="doc-{i}">
+                        <div class="doc-header">
+                            <div class="doc-rank">
+                                <span class="rank-label">Document {i+1} Rank:</span>
+                                <input type="number" class="rank-input" id="rank-{i}" value="{current_rank}" 
+                                       min="1" max="{len(docs)}" data-doc-id="{i}" 
+                                       onchange="updateRankings()">
+                            </div>
+                        </div>
+                        <div class="doc-content">{escaped_doc}</div>
+                    </div>
+                    """
+                
+                # Add validation and state tracking JS
+                html += """
+                <script>
+                function updateRankings() {
+                    // Collect all rank inputs
+                    const inputs = document.querySelectorAll('.rank-input');
+                    const rankings = [];
+                    
+                    // Get values and highlight duplicates
+                    const values = new Map();
+                    const duplicates = new Set();
+                    
+                    inputs.forEach(input => {
+                        const docId = parseInt(input.getAttribute('data-doc-id'));
+                        const rank = parseInt(input.value);
+                        
+                        // Store value
+                        rankings.push({
+                            docId: docId,
+                            rank: rank
+                        });
+                        
+                        // Check for duplicates
+                        if (values.has(rank)) {
+                            duplicates.add(rank);
+                        } else {
+                            values.set(rank, docId);
+                        }
                         
-                        # Document text
-                        with gr.Column(scale=4):
-                            doc_text = gr.Textbox(
-                                value=samples[0]["candidates"][i] if i < len(samples[0]["candidates"]) else "",
-                                label=f"Document {i+1}",
-                                lines=4,
-                                interactive=False
-                            )
-                            doc_texts.append(doc_text)
+                        // Reset styling
+                        input.style.borderColor = '#007bff';
+                    });
                     
-                    gr.Markdown("---")
+                    // Highlight duplicates
+                    inputs.forEach(input => {
+                        const rank = parseInt(input.value);
+                        if (duplicates.has(rank)) {
+                            input.style.borderColor = '#ff3860';
+                        }
+                    });
+                    
+                    // Store to hidden input
+                    const stateInput = document.getElementById('ranking-state');
+                    if (stateInput) {
+                        stateInput.value = JSON.stringify(rankings);
+                    }
+                    
+                    // Update gradio text area
+                    const textArea = document.querySelector('#rankings-state-input textarea');
+                    if (textArea) {
+                        textArea.value = JSON.stringify(rankings);
+                        const event = new Event('input', { bubbles: true });
+                        textArea.dispatchEvent(event);
+                    }
+                }
+                
+                // Initialize on page load
+                document.addEventListener('DOMContentLoaded', updateRankings);
+                // Also use a delay as a backup
+                setTimeout(updateRankings, 500);
+                </script>
+                </div>
+                """
+                
+                return html
+            
+            # Initial ranking HTML
+            ranking_html = gr.HTML(
+                generate_ranking_html(samples[0]["candidates"]),
+                elem_id="ranking-container"
+            )
+            
+            # Hidden input for state
+            rankings_state = gr.Textbox(
+                value="[]", 
+                visible=False, 
+                elem_id="rankings-state-input"
+            )
+            
+            # Validation message
+            validation_msg = gr.Textbox(
+                label="Validation", 
+                interactive=False
+            )
             
             # Navigation and submission buttons
             with gr.Row():
                 prev_btn = gr.Button("← Previous Query", size="sm")
+                validate_btn = gr.Button("Validate Rankings", variant="secondary")
                 submit_btn = gr.Button("Submit Rankings", size="lg", variant="primary")
-                next_btn = gr.Button("Next →", size="sm")
+                next_btn = gr.Button("Next Query →", size="sm")
             save_btn = gr.Button("💾 Save All Results", variant="secondary")
         
-        # Function to validate rankings
-        def validate_rankings(*ranks):
-            try:
-                # Filter out None values
-                valid_ranks = [int(r) for r in ranks if r is not None]
-                
-                # Check for duplicates
-                if len(set(valid_ranks)) != len(valid_ranks):
-                    # Find duplicate ranks
-                    dupes = {}
-                    for r in valid_ranks:
-                        dupes[r] = dupes.get(r, 0) + 1
-                    duplicates = [r for r, count in dupes.items() if count > 1]
-                    return f"⚠️ Duplicate ranks found: {', '.join(str(d) for d in sorted(duplicates))}. Each document must have a unique rank."
-                
-                # Check for complete ranking
-                max_rank = max(valid_ranks) if valid_ranks else 0
-                expected_ranks = set(range(1, max_rank + 1))
-                if set(valid_ranks) != expected_ranks:
-                    missing = sorted(expected_ranks - set(valid_ranks))
-                    if missing:
-                        return f"⚠️ Missing ranks: {', '.join(str(m) for m in missing)}. Ranks must be consecutive integers from 1 to {max_rank}."
-                
-                return "✓ Rankings are valid! Ready to submit."
-            except Exception as e:
-                return f"Error validating rankings: {str(e)}"
-        
         # Function to load a sample
         def load_sample(sample_id):
             try:
                 sample = next((s for s in samples if s["id"] == sample_id), None)
                 if not sample:
-                    return [gr.update()] * (3 + 2*max_docs)
-                
-                candidates = sample["candidates"]
-                num_docs = len(candidates)
+                    return sample_id, gr.update(), gr.update(), gr.update(), gr.update(), "[]", gr.update()
                 
                 # Get existing ranking if available
                 existing_ranking = next((anno["rankings"] for anno in results["annotations"] if anno["sample_id"] == sample_id), None)
                 
-                # Set default ranks (from existing or sequential)
-                ranks = []
-                for i in range(num_docs):
-                    if existing_ranking and i < len(existing_ranking):
-                        ranks.append(existing_ranking[i])
-                    else:
-                        ranks.append(i + 1)
-                
-                # Set container visibility
-                container_visibility = [i < num_docs for i in range(max_docs)]
+                # Generate HTML with existing rankings if available
+                html = generate_ranking_html(sample["candidates"], existing_ranking)
                 
-                # Update maximum values for number inputs
-                for input_field in rank_inputs:
-                    input_field.maximum = num_docs
-                
-                # Fill in document contents
-                docs = [candidates[i] if i < num_docs else "" for i in range(max_docs)]
-                
-                # Update visuals based on completed status
+                # Update status
                 status = "Already ranked" if completed_samples.get(sample_id, False) else "Ready to rank"
                 progress = f"Progress: {sum(completed_samples.values())}/{len(samples)}"
                 
-                # Prepare all outputs
-                outputs = [sample["query"], progress, status]
-                outputs.extend(ranks)  # Rank values
-                outputs.extend(docs)   # Document texts
-                outputs.extend(container_visibility)  # Container visibilities
-                
-                return outputs
+                return sample_id, sample["query"], html, progress, status, "[]", ""
             except Exception as e:
                 import traceback
                 print(traceback.format_exc())
-                return [gr.update(value=f"Error loading sample: {str(e)}")] + [gr.update()] * (2 + 2*max_docs)
+                return sample_id, gr.update(), gr.update(), gr.update(), f"Error: {str(e)}", "[]", ""
         
-        # Function to save rankings
-        def save_rankings(sample_id, *ranks):
+        # Function to validate rankings from JSON state
+        def validate_ranking_state(state_json):
             try:
+                if not state_json or state_json == "[]":
+                    return "Please rank all documents before submitting."
+                
+                # Parse the state
+                state = json.loads(state_json)
+                if not state:
+                    return "No ranking data found."
+                
+                # Extract ranks
+                ranks = [item.get("rank") for item in state if "rank" in item]
+                if not ranks:
+                    return "No valid ranks found."
+                
+                # Check for duplicates
+                if len(set(ranks)) != len(ranks):
+                    # Find duplicates
+                    dupes = {}
+                    for r in ranks:
+                        dupes[r] = dupes.get(r, 0) + 1
+                    duplicates = [r for r, count in dupes.items() if count > 1]
+                    return f"⚠️ Duplicate ranks found: {', '.join(map(str, sorted(duplicates)))}. Each document must have a unique rank."
+                
+                # Check for complete sequence
+                max_rank = max(ranks)
+                expected = set(range(1, max_rank + 1))
+                if set(ranks) != expected:
+                    missing = sorted(expected - set(ranks))
+                    return f"⚠️ Missing ranks: {', '.join(map(str, missing))}. Ranks must be consecutive from 1 to {max_rank}."
+                
+                return "✅ Rankings are valid. Ready to submit."
+                
+            except json.JSONDecodeError:
+                return "Error parsing ranking data."
+            except Exception as e:
+                return f"Error validating rankings: {str(e)}"
+        
+        # Function to save rankings from JSON state
+        def save_ranking_state(sample_id, state_json):
+            try:
+                if not state_json or state_json == "[]":
+                    return "Please rank all documents before submitting.", progress_text.value
+                
                 # Get the sample
                 sample = next((s for s in samples if s["id"] == sample_id), None)
                 if not sample:
                     return "⚠️ Sample not found", progress_text.value
                 
+                # Parse the state
+                state = json.loads(state_json)
+                
+                # Create a rankings array in the correct order
                 num_candidates = len(sample["candidates"])
+                rankings = [0] * num_candidates
                 
-                # Get the rankings for just this sample
-                valid_ranks = [int(r) for r in ranks[:num_candidates] if r is not None]
+                # Fill in rankings from state
+                for item in state:
+                    doc_id = item.get("docId")
+                    rank = item.get("rank")
+                    if doc_id is not None and doc_id < num_candidates and rank is not None:
+                        rankings[doc_id] = rank
                 
-                # Validate rankings
-                if len(valid_ranks) != num_candidates:
-                    return f"⚠️ Not all documents have ranks. Expected {num_candidates}, got {len(valid_ranks)}.", progress_text.value
+                # Validate rankings 
+                if any(r == 0 for r in rankings):
+                    return "⚠️ Not all documents have rankings", progress_text.value
                 
-                if sorted(valid_ranks) != list(range(1, num_candidates + 1)):
-                    return "⚠️ Rankings must include all integers from 1 to " + str(num_candidates), progress_text.value
+                if sorted(rankings) != list(range(1, num_candidates + 1)):
+                    return f"⚠️ Invalid ranking sequence. Please use each number from 1 to {num_candidates} exactly once.", progress_text.value
                 
                 # Create annotation
-                annotation = {"sample_id": sample_id, "rankings": valid_ranks}
+                annotation = {"sample_id": sample_id, "rankings": rankings}
                 
-                # Update or add the annotation
+                # Update or add to results
                 existing_idx = next((i for i, a in enumerate(results["annotations"]) if a["sample_id"] == sample_id), None)
                 if existing_idx is not None:
                     results["annotations"][existing_idx] = annotation
                 else:
                     results["annotations"].append(annotation)
                 
-                # Mark sample as completed
+                # Mark as completed
                 completed_samples[sample_id] = True
                 
                 # Save to file
@@ -223,7 +349,7 @@ def create_reranking_interface(task_data):
             except Exception as e:
                 import traceback
                 print(traceback.format_exc())
-                return f"Error saving rankings: {str(e)}", progress_text.value
+                return f"⚠️ Error saving rankings: {str(e)}", progress_text.value
         
         # Function to navigate to next sample
         def next_sample_id(current_id):
@@ -250,21 +376,19 @@ def create_reranking_interface(task_data):
             except Exception as e:
                 return f"⚠️ Error saving results file: {str(e)}"
         
-        # Connect validation button
+        # Connect buttons
         validate_btn.click(
-            validate_rankings,
-            inputs=rank_inputs,
-            outputs=validation_text
+            validate_ranking_state,
+            inputs=[rankings_state],
+            outputs=[validation_msg]
         )
         
-        # Connect submission button
         submit_btn.click(
-            save_rankings,
-            inputs=[current_sample_id] + rank_inputs,
+            save_ranking_state,
+            inputs=[current_sample_id, rankings_state],
             outputs=[status_box, progress_text]
         )
         
-        # Connect navigation buttons
         next_btn.click(
             next_sample_id, 
             inputs=[current_sample_id], 
@@ -272,10 +396,15 @@ def create_reranking_interface(task_data):
         ).then(
             load_sample,
             inputs=[current_sample_id],
-            outputs=[query_text, progress_text, status_box] + 
-                    rank_inputs + 
-                    doc_texts +
-                    doc_containers
+            outputs=[
+                current_sample_id,
+                query_text, 
+                ranking_html,
+                progress_text,
+                status_box,
+                rankings_state,
+                validation_msg
+            ]
         )
         
         prev_btn.click(
@@ -285,27 +414,33 @@ def create_reranking_interface(task_data):
         ).then(
             load_sample,
             inputs=[current_sample_id],
-            outputs=[query_text, progress_text, status_box] + 
-                    rank_inputs + 
-                    doc_texts +
-                    doc_containers
+            outputs=[
+                current_sample_id,
+                query_text, 
+                ranking_html,
+                progress_text,
+                status_box,
+                rankings_state,
+                validation_msg
+            ]
         )
         
-        # Connect save button
         save_btn.click(save_results, outputs=[status_box])
         
-        # Initialize interface with first sample
+        # Initialize with first sample
         demo.load(
-            lambda: load_sample(samples[0]['id']), 
-            outputs=[query_text, progress_text, status_box] + 
-                    rank_inputs + 
-                    doc_texts +
-                    doc_containers
+            lambda: load_sample(samples[0]['id']),
+            outputs=[
+                current_sample_id,
+                query_text, 
+                ranking_html,
+                progress_text,
+                status_box,
+                rankings_state,
+                validation_msg
+            ]
         )
         
-        # Add CSS styling
-        demo.load(lambda: gr.Accordion.update(open=True), outputs=[])
-        
     return demo
 
 # Main app with file upload capability