app.py

"""
Gradio interface for Rabbinic Hebrew/Aramaic Embedding Evaluation.

A Hugging Face Space for evaluating embedding models on cross-lingual
retrieval between Hebrew/Aramaic source texts and English translations.
"""

import os
from datetime import datetime

import gradio as gr
import pandas as pd
import plotly.graph_objects as go

from data_loader import load_benchmark_dataset, get_benchmark_stats
from models import (
    CURATED_MODELS,
    API_MODELS,
    ALL_MODELS,
    get_curated_model_choices,
    get_api_model_choices,
    get_all_model_choices,
    load_model,
    validate_model_id,
    is_api_model,
    requires_api_key,
    api_key_optional,
    get_api_key_type,
    get_api_key_env_var,
)
from evaluation import (
    EvaluationResults,
    evaluate_model,
    evaluate_model_streaming,
    compute_similarity_matrix,
    get_rank_distribution,
)
from leaderboard import (
    load_leaderboard as load_leaderboard_from_hub,
    add_result as add_result_to_hub,
)

# HuggingFace Dataset ID for benchmark data
BENCHMARK_DATASET_ID = "Sefaria/Rabbinic-Hebrew-English-Pairs"

# Global state
_benchmark_data = None


def load_benchmark():
    """Load benchmark data from HuggingFace Hub, with fallback to sample data."""
    global _benchmark_data
    
    if _benchmark_data is not None:
        return _benchmark_data
    
    try:
        _benchmark_data = load_benchmark_dataset(BENCHMARK_DATASET_ID)
        print(f"Loaded {len(_benchmark_data)} benchmark pairs from {BENCHMARK_DATASET_ID}")
    except Exception as e:
        print(f"Failed to load benchmark: {e}")
        print("Using sample data for testing")
        # Create minimal sample data for testing
        _benchmark_data = [
            {
                "ref": "Sample.1",
                "he": "בראשית ברא אלהים את השמים ואת הארץ",
                "en": "In the beginning God created the heaven and the earth",
                "category": "Sample",
            },
            {
                "ref": "Sample.2",
                "he": "והארץ היתה תהו ובהו וחשך על פני תהום",
                "en": "And the earth was without form, and void; and darkness was upon the face of the deep",
                "category": "Sample",
            },
        ]
    
    return _benchmark_data


def load_leaderboard():
    """Load leaderboard from HuggingFace Hub."""
    return load_leaderboard_from_hub()


def add_to_leaderboard(results: EvaluationResults):
    """Add evaluation results to leaderboard on HuggingFace Hub."""
    entry = results.to_dict()
    entry["timestamp"] = datetime.now().isoformat()
    
    # Add to Hub (handles deduplication and sorting internally)
    success = add_result_to_hub(entry)
    
    if not success:
        print("Note: Results saved locally but not persisted to Hub (no HF_TOKEN)")


def format_leaderboard_df():
    """Format leaderboard as pandas DataFrame for display."""
    leaderboard = load_leaderboard()
    
    if not leaderboard:
        return pd.DataFrame(columns=[
            "#", "Model", "MRR", "R@1", "R@5", "R@10", 
            "Bitext", "TrueSim", "RandSim", "N"
        ])
    
    rows = []
    for i, entry in enumerate(leaderboard, 1):
        rows.append({
            "#": i,
            "Model": entry.get("model_name", entry["model_id"]),
            "MRR": f"{entry['mrr']:.3f}",
            "R@1": f"{entry['recall_at_1']:.1%}",
            "R@5": f"{entry['recall_at_5']:.1%}",
            "R@10": f"{entry['recall_at_10']:.1%}",
            "Bitext": f"{entry['bitext_accuracy']:.1%}",
            "TrueSim": f"{entry['avg_true_pair_similarity']:.3f}",
            "RandSim": f"{entry['avg_random_pair_similarity']:.3f}",
            "N": entry["num_pairs"],
        })
    
    return pd.DataFrame(rows)


def run_evaluation(
    model_choice: str,
    custom_model_id: str,
    api_key: str,
    max_pairs: int,
):
    """
    Run evaluation for the selected model (generator for streaming status updates).
    
    Args:
        model_choice: Selected curated model or "custom"
        custom_model_id: Custom model ID if selected
        api_key: API key for API-based models
        max_pairs: Maximum pairs to evaluate
        
    Yields:
        Tuples of (status, results, leaderboard)
    """
    # Helper to yield status updates
    def status_update(msg):
        return (msg, gr.update(), gr.update())
    
    # Determine which model to use
    if model_choice == "custom":
        model_id = custom_model_id.strip()
        is_valid, error = validate_model_id(model_id)
        if not is_valid:
            yield (
                f"❌ {error}",
                f"❌ Invalid model ID: {error}",
                format_leaderboard_df(),
            )
            return
    else:
        model_id = model_choice
    
    # Check if API key is required but not provided
    if requires_api_key(model_id):
        api_key = api_key.strip() if api_key else ""
        env_var = get_api_key_env_var(model_id)
        key_type = get_api_key_type(model_id)
        
        # Skip API key check for models that support Application Default Credentials
        if not api_key and not os.environ.get(env_var) and not api_key_optional(model_id):
            yield (
                "❌ API key required",
                f"❌ API key required for {model_id}. Please enter your {key_type.upper()} API key or set the {env_var} environment variable.",
                format_leaderboard_df(),
            )
            return
    
    yield status_update(f"⏳ Loading benchmark data...")
    benchmark = load_benchmark()
    
    if max_pairs and max_pairs < len(benchmark):
        benchmark = benchmark[:max_pairs]
    
    yield status_update(f"⏳ Loading model: {model_id}...")
    
    try:
        # Pass API key for API-based models
        model = load_model(model_id, api_key=api_key if api_key else None)
    except Exception as e:
        yield (
            "❌ Model load failed",
            f"❌ Failed to load model: {str(e)}",
            format_leaderboard_df(),
        )
        return
    
    # Stream progress updates during evaluation
    try:
        results = None
        for item in evaluate_model_streaming(model, benchmark, batch_size=32):
            if isinstance(item, str):
                # Progress update
                yield status_update(item)
            else:
                # Final results
                results = item
    except Exception as e:
        yield (
            "❌ Evaluation failed",
            f"❌ Evaluation failed: {str(e)}",
            format_leaderboard_df(),
        )
        return
    
    yield status_update("⏳ Saving results...")
    add_to_leaderboard(results)
    
    # Format results summary
    summary = f"""## Results for {results.model_name}

| Metric | Value |
|--------|-------|
| **MRR** | {results.mrr:.4f} |
| **Recall@1** | {results.recall_at_1:.1%} |
| **Recall@5** | {results.recall_at_5:.1%} |
| **Recall@10** | {results.recall_at_10:.1%} |
| **Bitext Accuracy** | {results.bitext_accuracy:.1%} |
| **Avg True Pair Sim** | {results.avg_true_pair_similarity:.4f} |
| **Avg Random Pair Sim** | {results.avg_random_pair_similarity:.4f} |
| **Pairs Evaluated** | {results.num_pairs:,} |
"""
    
    # Final yield with all results (clear status)
    yield (
        "✅ Complete!",
        summary,
        format_leaderboard_df(),
    )


def create_leaderboard_comparison():
    """Create comparison chart of all models on leaderboard."""
    leaderboard = load_leaderboard()
    
    if len(leaderboard) < 2:
        return None
    
    models = [e.get("model_name", e["model_id"]) for e in leaderboard]
    mrr = [e["mrr"] for e in leaderboard]
    r1 = [e["recall_at_1"] for e in leaderboard]
    r5 = [e["recall_at_5"] for e in leaderboard]
    r10 = [e["recall_at_10"] for e in leaderboard]
    bitext = [e["bitext_accuracy"] for e in leaderboard]
    
    fig = go.Figure()
    
    fig.add_trace(go.Bar(name="MRR", x=models, y=mrr, marker_color="#2E86AB"))
    fig.add_trace(go.Bar(name="R@1", x=models, y=r1, marker_color="#A23B72"))
    fig.add_trace(go.Bar(name="R@5", x=models, y=r5, marker_color="#F18F01"))
    fig.add_trace(go.Bar(name="R@10", x=models, y=r10, marker_color="#C73E1D"))
    fig.add_trace(go.Bar(name="Bitext Acc", x=models, y=bitext, marker_color="#6B5B95"))
    
    fig.update_layout(
        title="Model Comparison",
        yaxis_title="Score",
        yaxis_range=[0, 1],
        barmode="group",
        template="plotly_white",
        height=400,
    )
    
    return fig


def update_model_inputs_visibility(choice):
    """Show/hide custom model input and API key based on selection."""
    show_custom = (choice == "custom")
    show_api_key = requires_api_key(choice) if choice != "custom" else False
    
    # Update API key label based on model type
    if show_api_key:
        key_type = get_api_key_type(choice)
        env_var = get_api_key_env_var(choice)
        is_optional = api_key_optional(choice)
        
        if key_type == "voyage":
            label = "Voyage AI API Key"
            placeholder = f"Enter your Voyage AI API key (or set {env_var} env var)"
        elif key_type == "gemini":
            label = "Gemini API Key (optional if using gcloud)"
            placeholder = f"Leave blank if using gcloud ADC, or enter API key / set {env_var}"
        else:
            label = "OpenAI API Key"
            placeholder = f"Enter your OpenAI API key (or set {env_var} env var)"
        return (
            gr.update(visible=show_custom),
            gr.update(visible=show_api_key, label=label, placeholder=placeholder),
        )
    
    return (
        gr.update(visible=show_custom),
        gr.update(visible=show_api_key),
    )


# Build the Gradio interface
def create_app():
    """Create and return the Gradio app."""
    
    # Get all model choices - local models first, then API models
    model_choices = []
    
    # Local models
    for model_id, info in CURATED_MODELS.items():
        model_choices.append((f"🖥️ {info['name']}", model_id))
    
    # API models
    for model_id, info in API_MODELS.items():
        model_choices.append((f"🌐 {info['name']}", model_id))
    
    # Custom option
    model_choices.append(("⚙️ Custom Model (enter ID below)", "custom"))
    
    # Load initial data
    load_benchmark()
    load_leaderboard()
    benchmark_stats = get_benchmark_stats(_benchmark_data) if _benchmark_data else {}
    
    with gr.Blocks(
        title="Rabbinic Embedding Benchmark",
        theme=gr.themes.Soft(
            primary_hue="blue",
            secondary_hue="orange",
            font=gr.themes.GoogleFont("Source Sans Pro"),
        ),
        css="""
        .main-header {
            text-align: center;
            margin-bottom: 1rem;
        }
        .stats-box {
            background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
            color: white;
            padding: 1rem;
            border-radius: 8px;
            margin: 0.5rem 0;
        }
        """,
    ) as app:
        
        gr.Markdown(
            """
            # 📚 Rabbinic Hebrew/Aramaic Embedding Benchmark
            
            Evaluate embedding models on cross-lingual retrieval between Hebrew/Aramaic 
            source texts and their English translations from Sefaria.
            
            **How it works:** Given a Hebrew/Aramaic text, can the model find its correct 
            English translation from a pool of candidates? Models that excel at this task 
            produce high-quality embeddings for Rabbinic literature.
            """,
            elem_classes=["main-header"],
        )
        
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown(f"""
                ### 📊 Benchmark Stats
                - **Total Pairs:** {benchmark_stats.get('total_pairs', 'N/A'):,}
                - **Categories:** {len(benchmark_stats.get('categories', {}))}
                - **Avg Hebrew Length:** {benchmark_stats.get('avg_he_length', 0):.0f} chars
                """)
            
            with gr.Column(scale=1):
                gr.Markdown("""
                ### 📏 Metrics
                - **MRR:** Mean Reciprocal Rank
                - **R@k:** Recall at k (correct in top k)
                - **Bitext Acc:** True vs random pair classification
                """)
        
        gr.Markdown("---")
        
        with gr.Tabs():
            with gr.TabItem("🔬 Evaluate Model"):
                with gr.Row():
                    with gr.Column(scale=2):
                        model_dropdown = gr.Dropdown(
                            choices=model_choices,
                            value=model_choices[0][1],
                            label="Select Model",
                            info="Choose a curated model or enter a custom Hugging Face model ID",
                        )
                        
                        custom_model_input = gr.Textbox(
                            label="Custom Model ID",
                            placeholder="e.g., organization/model-name",
                            visible=False,
                        )
                        
                        api_key_input = gr.Textbox(
                            label="API Key",
                            placeholder="Enter your API key (or set appropriate env var)",
                            type="password",
                            visible=False,
                            info="Required for API-based models (OpenAI, Voyage AI). Your key is not stored.",
                        )
                        
                        total_pairs = benchmark_stats.get('total_pairs', 1000)
                        max_pairs_slider = gr.Slider(
                            minimum=100,
                            maximum=total_pairs,
                            value=total_pairs,
                            step=100,
                            label="Max Pairs to Evaluate",
                            info="Use fewer pairs for faster evaluation",
                        )
                    
                    with gr.Column(scale=3):
                        evaluate_btn = gr.Button(
                            "🚀 Run Evaluation",
                            variant="primary",
                            size="lg",
                        )
                        
                        status_text = gr.Markdown("")
                        
                        results_markdown = gr.Markdown("")
            
            with gr.TabItem("🏆 Leaderboard"):
                leaderboard_table = gr.Dataframe(
                    value=format_leaderboard_df(),
                    label="Model Rankings",
                    interactive=False,
                )
                
                refresh_btn = gr.Button("🔄 Refresh Leaderboard")
                
                comparison_plot = gr.Plot(label="Model Comparison")
        
        gr.Markdown("""
        ---
        ### About
        
        This benchmark evaluates embedding models for Rabbinic Hebrew and Aramaic texts using 
        cross-lingual retrieval. 

        All texts and translations sourced from [Sefaria](https://www.sefaria.org).
        """)
        
        # Event handlers
        model_dropdown.change(
            fn=update_model_inputs_visibility,
            inputs=[model_dropdown],
            outputs=[custom_model_input, api_key_input],
        )
        
        evaluate_btn.click(
            fn=run_evaluation,
            inputs=[model_dropdown, custom_model_input, api_key_input, max_pairs_slider],
            outputs=[status_text, results_markdown, leaderboard_table],
            show_progress="hidden",
        )
        
        refresh_btn.click(
            fn=lambda: (format_leaderboard_df(), create_leaderboard_comparison()),
            outputs=[leaderboard_table, comparison_plot],
        )
    
    return app


# Main entry point
if __name__ == "__main__":
    app = create_app()
    app.launch()