Update app.py

app.py

@@ -52,20 +52,6 @@ def create_reranking_interface(task_data):
             # Return specific error message
             return f"Error: {str(e)}", f"Progress: {sum(completed_samples.values())}/{len(samples)}"
     
-    def move_document_up(idx):
-        nonlocal current_order
-        if idx > 0:
-            # Swap with the document above
-            current_order[idx], current_order[idx-1] = current_order[idx-1], current_order[idx]
-        return gr.update(value=f"Moved document up: {idx}")
-    
-    def move_document_down(idx):
-        nonlocal current_order
-        if idx < len(current_order) - 1:
-            # Swap with the document below
-            current_order[idx], current_order[idx+1] = current_order[idx+1], current_order[idx]
-        return gr.update(value=f"Moved document down: {idx}")
-    
     with gr.Blocks(theme=gr.themes.Soft()) as demo:
         gr.Markdown(f"# {task_data['task_name']} - Human Reranking Evaluation")
         
@@ -117,14 +103,24 @@ def create_reranking_interface(task_data):
                 # Check if this sample has already been annotated to restore ordering
                 existing_annotation = next((a for a in results["annotations"] if a["sample_id"] == sample_id), None)
                 if existing_annotation and "rankings" in existing_annotation:
-                    # Sort documents based on previous rankings
-                    # This is a simplified version that just keeps the original order for now
-                    pass
+                    # Create pairs of (doc_idx, rank)
+                    ranked_docs = []
+                    for doc_idx, rank in enumerate(existing_annotation["rankings"]):
+                        ranked_docs.append((doc_idx, rank))
+                    
+                    # Sort by rank (ascending)
+                    ranked_docs.sort(key=lambda x: x[1])
+                    
+                    # Extract document indices in rank order
+                    current_order = [doc[0] for doc in ranked_docs]
                 
                 # Update UI
-                for i, doc_idx in enumerate(current_order):
-                    if i < len(document_containers) and doc_idx < len(docs):
+                for i in range(len(document_containers)):
+                    if i < len(docs):
+                        doc_idx = current_order[i]
                         document_containers[i].value = f"Document {doc_idx+1} (Rank: {i+1}): {docs[doc_idx]}"
+                    else:
+                        document_containers[i].value = ""  # Clear unused containers
                 
                 # Status message
                 status = f"Viewing query {samples.index(sample) + 1} of {len(samples)}"
@@ -135,37 +131,75 @@ def create_reranking_interface(task_data):
             
             # Create document display containers with up/down buttons
             with gr.Column():
-                for i in range(10):  # Assuming max 10 documents per sample
+                for i in range(20):  # Now handling up to 20 documents per sample
                     with gr.Group():
                         with gr.Row():
                             doc_text = gr.Textbox(label=f"Document {i+1}", interactive=False)
                             document_containers.append(doc_text)
                         
                         with gr.Row():
-                            up_btn = gr.Button(f"⬆️ Move Up", size="sm")
-                            down_btn = gr.Button(f"⬇️ Move Down", size="sm")
+                            # Create unique buttons for each document to avoid closure issues
+                            up_btn = gr.Button(f"⬆️ Move Up #{i+1}", size="sm")
+                            down_btn = gr.Button(f"⬇️ Move Down #{i+1}", size="sm")
                         
-                        # Connect the up/down buttons
+                        # Create a closure that properly captures the current index
+                        def make_up_handler(idx):
+                            def up_handler():
+                                nonlocal current_order
+                                
+                                # Only move if index is valid
+                                if idx < len(current_order) and idx > 0:
+                                    # Swap with the document above
+                                    current_order[idx], current_order[idx-1] = current_order[idx-1], current_order[idx]
+                                    
+                                    # Update all document displays with new order
+                                    sample = next((s for s in samples if s["id"] == current_sample_id.value), None)
+                                    if sample:
+                                        docs = sample["candidates"]
+                                        updates = []
+                                        for j, doc_idx in enumerate(current_order):
+                                            if j < len(document_containers) and doc_idx < len(docs):
+                                                updates.append(gr.update(value=f"Document {doc_idx+1} (Rank: {j+1}): {docs[doc_idx]}"))
+                                            else:
+                                                updates.append(gr.update(value=""))
+                                        return updates
+                                return [gr.update() for _ in document_containers]
+                            return up_handler
+                        
+                        def make_down_handler(idx):
+                            def down_handler():
+                                nonlocal current_order
+                                
+                                # Only move if index is valid
+                                if idx < len(current_order) - 1:
+                                    # Swap with the document below
+                                    current_order[idx], current_order[idx+1] = current_order[idx+1], current_order[idx]
+                                    
+                                    # Update all document displays with new order
+                                    sample = next((s for s in samples if s["id"] == current_sample_id.value), None)
+                                    if sample:
+                                        docs = sample["candidates"]
+                                        updates = []
+                                        for j, doc_idx in enumerate(current_order):
+                                            if j < len(document_containers) and doc_idx < len(docs):
+                                                updates.append(gr.update(value=f"Document {doc_idx+1} (Rank: {j+1}): {docs[doc_idx]}"))
+                                            else:
+                                                updates.append(gr.update(value=""))
+                                        return updates
+                                return [gr.update() for _ in document_containers]
+                            return down_handler
+                        
+                        # Connect buttons with properly created handlers
                         up_btn.click(
-                            lambda idx=i: move_document_up(idx),
+                            make_up_handler(i),
                             inputs=[],
-                            outputs=[status_box]
-                        ).then(
-                            lambda: [gr.update(value=f"Document {current_order[j]+1} (Rank: {j+1}): {samples[0]['candidates'][current_order[j]]}") 
-                                    for j in range(len(current_order))],
-                            inputs=[],
-                            outputs=document_containers[:len(current_order)]
+                            outputs=document_containers
                         )
                         
                         down_btn.click(
-                            lambda idx=i: move_document_down(idx),
-                            inputs=[],
-                            outputs=[status_box]
-                        ).then(
-                            lambda: [gr.update(value=f"Document {current_order[j]+1} (Rank: {j+1}): {samples[0]['candidates'][current_order[j]]}") 
-                                    for j in range(len(current_order))],
+                            make_down_handler(i),
                             inputs=[],
-                            outputs=document_containers[:len(current_order)]
+                            outputs=document_containers
                         )
             
             with gr.Row():