app.py

import gradio as gr
import json
import os
from pathlib import Path

def create_reranking_interface(task_data):
    """Create a Gradio interface for reranking evaluation."""
    samples = task_data["samples"]
    results = {"task_name": task_data["task_name"], "task_type": "reranking", "annotations": []}
    completed_samples = {s["id"]: False for s in samples}
    
    # Try to load existing results
    output_path = f"{task_data['task_name']}_human_results.json"
    if os.path.exists(output_path):
        try:
            with open(output_path, "r") as f:
                existing_results = json.load(f)
                results = existing_results
                # Update completed samples based on existing annotations
                for anno in results.get("annotations", []):
                    if "sample_id" in anno:
                        completed_samples[anno["sample_id"]] = True
        except Exception as e:
            print(f"Error loading existing results: {str(e)}")
    
    # Create the main interface
    with gr.Blocks(theme=gr.themes.Soft()) as demo:
        gr.Markdown(f"# {task_data['task_name']} - Human Reranking Evaluation")
        with gr.Accordion("Instructions", open=True):
            gr.Markdown("""
            ## Task Instructions
            
            {instructions}
            
            ### How to use this interface:
            1. Read the query at the top
            2. For each document, select its rank (1 = most relevant)
            3. Make sure each document has a unique rank (1 to N)
            4. Click "Submit Rankings" when you're done with the current query
            5. Use "Previous" and "Next" to navigate between queries
            6. Click "Save All Results" periodically to ensure your work is saved
            """.format(instructions=task_data["instructions"]))
        
        # State variables
        current_sample_id = gr.State(value=samples[0]["id"])
        
        # Progress tracking
        with gr.Row():
            progress_text = gr.Textbox(label="Progress", value=f"Progress: {sum(completed_samples.values())}/{len(samples)}", interactive=False)
            status_box = gr.Textbox(label="Status", value="Ready to start evaluation", interactive=False)
        
        # Query display
        with gr.Group():
            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 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)
                    
                    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)
                        
                        # 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)
                    
                    gr.Markdown("---")
            
            # Navigation and submission buttons
            with gr.Row():
                prev_btn = gr.Button("← Previous Query", size="sm")
                submit_btn = gr.Button("Submit Rankings", size="lg", variant="primary")
                next_btn = gr.Button("Next →", 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)
                
                # 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)]
                
                # 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
                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
            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)
        
        # Function to save rankings
        def save_rankings(sample_id, *ranks):
            try:
                # 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
                
                num_candidates = len(sample["candidates"])
                
                # Get the rankings for just this sample
                valid_ranks = [int(r) for r in ranks[:num_candidates] if r is not None]
                
                # 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
                
                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
                
                # Create annotation
                annotation = {"sample_id": sample_id, "rankings": valid_ranks}
                
                # Update or add the annotation
                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
                completed_samples[sample_id] = True
                
                # Save to file
                with open(output_path, "w") as f:
                    json.dump(results, f, indent=2)
                
                # Update progress
                progress = f"Progress: {sum(completed_samples.values())}/{len(samples)}"
                
                return f"✅ Rankings saved successfully! ({sum(completed_samples.values())}/{len(samples)} completed)", progress
            except Exception as e:
                import traceback
                print(traceback.format_exc())
                return f"Error saving rankings: {str(e)}", progress_text.value
        
        # Function to navigate to next sample
        def next_sample_id(current_id):
            current_idx = next((i for i, s in enumerate(samples) if s["id"] == current_id), -1)
            if current_idx == -1:
                return current_id
            next_idx = min(current_idx + 1, len(samples) - 1)
            return samples[next_idx]["id"]
        
        # Function to navigate to previous sample
        def prev_sample_id(current_id):
            current_idx = next((i for i, s in enumerate(samples) if s["id"] == current_id), -1)
            if current_idx == -1:
                return current_id
            prev_idx = max(current_idx - 1, 0)
            return samples[prev_idx]["id"]
        
        # Function to save all results
        def save_results():
            try:
                with open(output_path, "w") as f:
                    json.dump(results, f, indent=2)
                return f"✅ Results saved to {output_path} ({len(results['annotations'])} annotations)"
            except Exception as e:
                return f"⚠️ Error saving results file: {str(e)}"
        
        # Connect validation button
        validate_btn.click(
            validate_rankings,
            inputs=rank_inputs,
            outputs=validation_text
        )
        
        # Connect submission button
        submit_btn.click(
            save_rankings,
            inputs=[current_sample_id] + rank_inputs,
            outputs=[status_box, progress_text]
        )
        
        # Connect navigation buttons
        next_btn.click(
            next_sample_id, 
            inputs=[current_sample_id], 
            outputs=[current_sample_id]
        ).then(
            load_sample,
            inputs=[current_sample_id],
            outputs=[query_text, progress_text, status_box] + 
                    rank_inputs + 
                    doc_texts +
                    doc_containers
        )
        
        prev_btn.click(
            prev_sample_id, 
            inputs=[current_sample_id], 
            outputs=[current_sample_id]
        ).then(
            load_sample,
            inputs=[current_sample_id],
            outputs=[query_text, progress_text, status_box] + 
                    rank_inputs + 
                    doc_texts +
                    doc_containers
        )
        
        # Connect save button
        save_btn.click(save_results, outputs=[status_box])
        
        # Initialize interface with first sample
        demo.load(
            lambda: load_sample(samples[0]['id']), 
            outputs=[query_text, progress_text, status_box] + 
                    rank_inputs + 
                    doc_texts +
                    doc_containers
        )
        
        # Add CSS styling
        demo.load(lambda: gr.Accordion.update(open=True), outputs=[])
        
    return demo

# Main app with file upload capability
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# MTEB Human Evaluation Demo")
    
    with gr.Tabs():
        with gr.TabItem("Demo"):
            gr.Markdown("""
            ## MTEB Human Evaluation Interface
            
            This interface allows you to evaluate the relevance of documents for reranking tasks.
            """)
            
            # Function to get the most recent task file
            def get_latest_task_file():
                # Check first in uploaded_tasks directory
                os.makedirs("uploaded_tasks", exist_ok=True)
                uploaded_tasks = [f for f in os.listdir("uploaded_tasks") if f.endswith(".json")]
                
                if uploaded_tasks:
                    # Sort by modification time, newest first
                    uploaded_tasks.sort(key=lambda x: os.path.getmtime(os.path.join("uploaded_tasks", x)), reverse=True)
                    return os.path.join("uploaded_tasks", uploaded_tasks[0])
                
                # Fall back to default example
                return "AskUbuntuDupQuestions_human_eval.json"
            
            # Load the task file
            task_file = get_latest_task_file()
            
            try:
                with open(task_file, "r") as f:
                    task_data = json.load(f)
                
                # Show which task is currently loaded
                gr.Markdown(f"**Current Task: {task_data['task_name']}** ({len(task_data['samples'])} samples)")
                
                # Display the interface
                reranking_demo = create_reranking_interface(task_data)
            except Exception as e:
                gr.Markdown(f"**Error loading task: {str(e)}**")
                gr.Markdown("Please upload a valid task file in the 'Upload & Evaluate' tab.")
        
        with gr.TabItem("Upload & Evaluate"):
            gr.Markdown("""
            ## Upload Your Own Task File
            
            If you have a prepared task file, you can upload it here to create an evaluation interface.
            """)
            
            with gr.Row():
                with gr.Column():
                    file_input = gr.File(label="Upload a task file (JSON)")
                    load_btn = gr.Button("Load Task")
                    message = gr.Textbox(label="Status", interactive=False)
                    
                    # Add task list for previously uploaded tasks
                    gr.Markdown("### Previous Uploads")
                    
                    # Function to list existing task files in the tasks directory
                    def list_task_files():
                        os.makedirs("uploaded_tasks", exist_ok=True)
                        tasks = [f for f in os.listdir("uploaded_tasks") if f.endswith(".json")]
                        if not tasks:
                            return "No task files uploaded yet."
                        return "\n".join([f"- [{t}](javascript:selectTask('{t}'))" for t in tasks])
                    
                    task_list = gr.Markdown(list_task_files())
                    refresh_btn = gr.Button("Refresh List")
                    
                    # Add results management section
                    gr.Markdown("### Results Management")
                    
                    # Function to list existing result files
                    def list_result_files():
                        results = [f for f in os.listdir(".") if f.endswith("_human_results.json")]
                        if not results:
                            return "No result files available yet."
                        
                        result_links = []
                        for r in results:
                            # Calculate completion stats
                            try:
                                with open(r, "r") as f:
                                    result_data = json.load(f)
                                annotation_count = len(result_data.get("annotations", []))
                                task_name = result_data.get("task_name", "Unknown")
                                result_links.append(f"- {r} ({annotation_count} annotations for {task_name})")
                            except:
                                result_links.append(f"- {r}")
                        
                        return "\n".join(result_links)
                    
                    results_list = gr.Markdown(list_result_files())
                    download_results_btn = gr.Button("Download Results")
                    
                # Right side - will contain the actual interface
                with gr.Column():
                    task_container = gr.HTML()
            
            # Handle file upload and storage
            def handle_upload(file):
                if not file:
                    return "Please upload a task file", task_list.value, task_container.value
                
                try:
                    # Create directory if it doesn't exist
                    os.makedirs("uploaded_tasks", exist_ok=True)
                    
                    # Read the uploaded file
                    with open(file.name, "r") as f:
                        task_data = json.load(f)
                    
                    # Validate task format
                    if "task_name" not in task_data or "samples" not in task_data:
                        return "Invalid task file format. Must contain 'task_name' and 'samples' fields.", task_list.value, task_container.value
                    
                    # Save to a consistent location
                    task_filename = f"uploaded_tasks/{task_data['task_name']}_task.json"
                    with open(task_filename, "w") as f:
                        json.dump(task_data, f, indent=2)
                    
                    # Instead of trying to create the interface here,
                    # we'll return a message with instructions
                    return f"Task '{task_data['task_name']}' uploaded successfully with {len(task_data['samples'])} samples. Please refresh the app and use the Demo tab to evaluate it.", list_task_files(), f"""
                    <div style="padding: 20px; background-color: #f0f0f0; border-radius: 10px;">
                        <h3>Task uploaded successfully!</h3>
                        <p>Task Name: {task_data['task_name']}</p>
                        <p>Samples: {len(task_data['samples'])}</p>
                        <p>To evaluate this task:</p>
                        <ol>
                            <li>Refresh the app</li>
                            <li>The Demo tab will now use your uploaded task</li>
                            <li>Complete your evaluations</li>
                            <li>Results will be saved as {task_data['task_name']}_human_results.json</li>
                        </ol>
                    </div>
                    """
                except Exception as e:
                    return f"Error processing task file: {str(e)}", task_list.value, task_container.value
            
            # Function to prepare results for download
            def prepare_results_for_download():
                results = [f for f in os.listdir(".") if f.endswith("_human_results.json")]
                if not results:
                    return None
                
                # Create a zip file with all results
                import zipfile
                zip_path = "mteb_human_eval_results.zip"
                with zipfile.ZipFile(zip_path, 'w') as zipf:
                    for r in results:
                        zipf.write(r)
                
                return zip_path
            
            # Connect events
            load_btn.click(handle_upload, inputs=[file_input], outputs=[message, task_list, task_container])
            refresh_btn.click(list_task_files, outputs=[task_list])
            download_results_btn.click(prepare_results_for_download, outputs=[gr.File(label="Download Results")])
        
        with gr.TabItem("Results Management"):
            gr.Markdown("""
            ## Manage Evaluation Results
            
            View, download, and analyze your evaluation results.
            """)
            
            # Function to load and display result stats
            def get_result_stats():
                results = [f for f in os.listdir(".") if f.endswith("_human_results.json")]
                if not results:
                    return "No result files available yet."
                
                stats = []
                for r in results:
                    try:
                        with open(r, "r") as f:
                            result_data = json.load(f)
                        
                        task_name = result_data.get("task_name", "Unknown")
                        annotations = result_data.get("annotations", [])
                        annotation_count = len(annotations)
                        
                        # Calculate completion percentage
                        sample_ids = set(a.get("sample_id") for a in annotations)
                        
                        # Try to get the total sample count from the corresponding task file
                        total_samples = 0
                        task_file = f"uploaded_tasks/{task_name}_task.json"
                        if os.path.exists(task_file):
                            with open(task_file, "r") as f:
                                task_data = json.load(f)
                            total_samples = len(task_data.get("samples", []))
                        
                        completion = f"{len(sample_ids)}/{total_samples}" if total_samples else f"{len(sample_ids)} samples"
                        
                        stats.append(f"### {task_name}\n- Annotations: {annotation_count}\n- Completion: {completion}\n- File: {r}")
                    except Exception as e:
                        stats.append(f"### {r}\n- Error loading results: {str(e)}")
                
                return "\n\n".join(stats)
            
            result_stats = gr.Markdown(get_result_stats())
            refresh_results_btn = gr.Button("Refresh Results")
            
            # Add download options
            with gr.Row():
                with gr.Column():
                    download_all_btn = gr.Button("Download All Results (ZIP)")
                with gr.Column():
                    result_select = gr.Dropdown(choices=[f for f in os.listdir(".") if f.endswith("_human_results.json")], label="Select Result to Download", value=None)
                    download_selected_btn = gr.Button("Download Selected")
            
            # Add results visualization placeholder
            gr.Markdown("### Results Visualization")
            gr.Markdown("*Visualization features will be added in a future update.*")
            
            # Connect events
            refresh_results_btn.click(get_result_stats, outputs=[result_stats])
            
            # Function to prepare all results for download as ZIP
            def prepare_all_results():
                import zipfile
                zip_path = "mteb_human_eval_results.zip"
                with zipfile.ZipFile(zip_path, 'w') as zipf:
                    for r in [f for f in os.listdir(".") if f.endswith("_human_results.json")]:
                        zipf.write(r)
                return zip_path
            
            # Function to return a single result file
            def get_selected_result(filename):
                if not filename:
                    return None
                if os.path.exists(filename):
                    return filename
                return None
            
            # Update dropdown when refreshing results
            def update_result_dropdown():
                return gr.Dropdown.update(choices=[f for f in os.listdir(".") if f.endswith("_human_results.json")])
            
            refresh_results_btn.click(update_result_dropdown, outputs=[result_select])
            download_all_btn.click(prepare_all_results, outputs=[gr.File(label="Download All Results")])
            download_selected_btn.click(get_selected_result, inputs=[result_select], outputs=[gr.File(label="Download Selected Result")])

if __name__ == "__main__":
    demo.launch()