src/streamlit_app.py

import streamlit as st
import pandas as pd
import altair as alt
from huggingface_hub import HfApi, hf_hub_download
import json
from pathlib import Path
import os
from typing import Dict, List, Optional
import numpy as np

# Page config
st.set_page_config(
    page_title="Grounding Benchmark Leaderboard",
    page_icon="🎯",
    layout="wide"
)

# Constants
REPO_ID = "mlfoundations-cua-dev/leaderboard"
GROUNDING_PATH = "grounding"

# Baselines for different datasets
BASELINES = {
    "screenspot-v2": {
        "Qwen2-VL-7B": {
            "desktop_text": 52.01,
            "desktop_icon": 44.98,
            "web_text": 33.04,
            "web_icon": 21.84,
            "overall": 37.96
        },
        "UI-TARS-2B": {
            "desktop_text": 90.7,
            "desktop_icon": 68.6,
            "web_text": 87.2,
            "web_icon": 84.7,
            "overall": 82.8
        },
        "UI-TARS-7B": {
            "desktop_text": 95.4,
            "desktop_icon": 87.8,
            "web_text": 93.8,
            "web_icon": 91.6,
            "overall": 92.2
        },
        "UI-TARS-72B": {
            "desktop_text": 91.2,
            "desktop_icon": 87.8,
            "web_text": 87.7,
            "web_icon": 86.3,
            "overall": 88.3
        }
    }
}

@st.cache_data(ttl=300)  # Cache for 5 minutes
def fetch_leaderboard_data():
    """Fetch all grounding results from HuggingFace leaderboard."""
    api = HfApi()
    
    try:
        # List all files in the grounding directory
        files = api.list_repo_files(repo_id=REPO_ID, repo_type="dataset")
        grounding_files = [f for f in files if f.startswith(f"{GROUNDING_PATH}/") and f.endswith(".json")]
        
        results = []
        for file_path in grounding_files:
            try:
                # Download and parse each JSON file
                local_path = hf_hub_download(
                    repo_id=REPO_ID,
                    filename=file_path,
                    repo_type="dataset"
                )
                
                with open(local_path, 'r') as f:
                    data = json.load(f)
                
                # Extract key information
                metadata = data.get("metadata", {})
                metrics = data.get("metrics", {})
                detailed_results = data.get("detailed_results", {})
                
                # Parse the file path to get dataset and model info
                path_parts = file_path.split('/')
                dataset_name = path_parts[1] if len(path_parts) > 1 else "unknown"
                
                # Get model name from metadata or path
                model_name = metadata.get("model_checkpoint", "").split('/')[-1]
                if not model_name and len(path_parts) > 2:
                    model_name = path_parts[2].replace("results_", "").replace(".json", "")
                
                # Extract UI type results if available
                ui_type_results = detailed_results.get("by_ui_type", {})
                dataset_type_results = detailed_results.get("by_dataset_type", {})
                
                results.append({
                    "dataset": dataset_name,
                    "model": model_name,
                    "model_path": metadata.get("model_checkpoint", ""),
                    "overall_accuracy": metrics.get("accuracy", 0) * 100,  # Convert to percentage
                    "total_samples": metrics.get("total", 0),
                    "timestamp": metadata.get("evaluation_timestamp", ""),
                    "checkpoint_steps": metadata.get("checkpoint_steps"),
                    "training_loss": metadata.get("training_loss"),
                    "ui_type_results": ui_type_results,
                    "dataset_type_results": dataset_type_results,
                    "raw_data": data
                })
                
            except Exception as e:
                st.warning(f"Error loading {file_path}: {str(e)}")
                continue
        
        return pd.DataFrame(results)
    
    except Exception as e:
        st.error(f"Error fetching leaderboard data: {str(e)}")
        return pd.DataFrame()

def parse_ui_type_metrics(df: pd.DataFrame, dataset_filter: str) -> pd.DataFrame:
    """Parse UI type metrics from the results dataframe."""
    metrics_list = []
    
    for _, row in df.iterrows():
        if row['dataset'] != dataset_filter:
            continue
            
        model = row['model']
        ui_results = row['ui_type_results']
        
        # For ScreenSpot datasets, we have desktop/web and text/icon
        if 'screenspot' in dataset_filter.lower():
            # Calculate aggregated metrics
            desktop_text = ui_results.get('desktop_text', {}).get('correct', 0) / max(ui_results.get('desktop_text', {}).get('total', 1), 1) * 100
            desktop_icon = ui_results.get('desktop_icon', {}).get('correct', 0) / max(ui_results.get('desktop_icon', {}).get('total', 1), 1) * 100
            web_text = ui_results.get('web_text', {}).get('correct', 0) / max(ui_results.get('web_text', {}).get('total', 1), 1) * 100
            web_icon = ui_results.get('web_icon', {}).get('correct', 0) / max(ui_results.get('web_icon', {}).get('total', 1), 1) * 100
            
            # Calculate averages
            desktop_avg = (desktop_text + desktop_icon) / 2 if desktop_text or desktop_icon else 0
            web_avg = (web_text + web_icon) / 2 if web_text or web_icon else 0
            text_avg = (desktop_text + web_text) / 2 if desktop_text or web_text else 0
            icon_avg = (desktop_icon + web_icon) / 2 if desktop_icon or web_icon else 0
            
            metrics_list.append({
                'model': model,
                'desktop_text': desktop_text,
                'desktop_icon': desktop_icon,
                'web_text': web_text,
                'web_icon': web_icon,
                'desktop_avg': desktop_avg,
                'web_avg': web_avg,
                'text_avg': text_avg,
                'icon_avg': icon_avg,
                'overall': row['overall_accuracy']
            })
    
    return pd.DataFrame(metrics_list)

def create_bar_chart(data: pd.DataFrame, metric: str, title: str):
    """Create a bar chart for a specific metric."""
    # Prepare data for the chart
    chart_data = []
    
    # Add model results
    for _, row in data.iterrows():
        if metric in row and row[metric] > 0:
            chart_data.append({
                'Model': row['model'],
                'Score': row[metric],
                'Type': 'Evaluated'
            })
    
    # Add baselines if available
    dataset = st.session_state.get('selected_dataset', '')
    if dataset in BASELINES:
        for baseline_name, baseline_metrics in BASELINES[dataset].items():
            metric_key = metric.replace('_avg', '').replace('avg', 'overall')
            if metric_key in baseline_metrics:
                chart_data.append({
                    'Model': baseline_name,
                    'Score': baseline_metrics[metric_key],
                    'Type': 'Baseline'
                })
    
    if not chart_data:
        return None
    
    df_chart = pd.DataFrame(chart_data)
    
    # Create the bar chart
    chart = alt.Chart(df_chart).mark_bar().encode(
        x=alt.X('Model:N', 
                sort=alt.EncodingSortField(field='Score', order='descending'),
                axis=alt.Axis(labelAngle=-45)),
        y=alt.Y('Score:Q', 
                scale=alt.Scale(domain=[0, 100]),
                axis=alt.Axis(title='Score (%)')),
        color=alt.Color('Type:N', 
                       scale=alt.Scale(domain=['Evaluated', 'Baseline'],
                                     range=['#4ECDC4', '#FFA726'])),
        tooltip=['Model', 'Score', 'Type']
    ).properties(
        title=title,
        width=400,
        height=300
    )
    
    # Add value labels
    text = chart.mark_text(
        align='center',
        baseline='bottom',
        dy=-5
    ).encode(
        text=alt.Text('Score:Q', format='.1f')
    )
    
    return chart + text

def main():
    st.title("🎯 Grounding Benchmark Leaderboard")
    st.markdown("Visualization of model performance on grounding benchmarks")
    
    # Fetch data
    with st.spinner("Loading leaderboard data..."):
        df = fetch_leaderboard_data()
    
    if df.empty:
        st.warning("No data available in the leaderboard.")
        return
    
    # Sidebar filters
    st.sidebar.header("Filters")
    
    # Dataset filter
    datasets = sorted(df['dataset'].unique())
    selected_dataset = st.sidebar.selectbox("Select Dataset", datasets)
    st.session_state['selected_dataset'] = selected_dataset
    
    # Filter data
    filtered_df = df[df['dataset'] == selected_dataset]
    
    # Model filter (optional)
    models = ['All'] + sorted(filtered_df['model'].unique())
    selected_model = st.sidebar.selectbox("Select Model", models)
    
    if selected_model != 'All':
        filtered_df = filtered_df[filtered_df['model'] == selected_model]
    
    # Main content
    st.header(f"Results for {selected_dataset}")
    
    # Overall metrics
    col1, col2, col3 = st.columns(3)
    with col1:
        st.metric("Models Evaluated", len(filtered_df))
    with col2:
        if not filtered_df.empty:
            best_acc = filtered_df['overall_accuracy'].max()
            best_model = filtered_df[filtered_df['overall_accuracy'] == best_acc]['model'].iloc[0]
            st.metric("Best Overall Accuracy", f"{best_acc:.1f}%", help=f"Model: {best_model}")
    with col3:
        total_samples = filtered_df['total_samples'].sum()
        st.metric("Total Samples Evaluated", f"{total_samples:,}")
    
    # Parse UI type metrics
    ui_metrics_df = parse_ui_type_metrics(filtered_df, selected_dataset)
    
    if not ui_metrics_df.empty and 'screenspot' in selected_dataset.lower():
        st.subheader("Performance by UI Type")
        
        # Create charts in a grid
        col1, col2 = st.columns(2)
        
        with col1:
            # Overall Average
            chart = create_bar_chart(ui_metrics_df, 'overall', 'Overall Average')
            if chart:
                st.altair_chart(chart, use_container_width=True)
            
            # Desktop Average
            chart = create_bar_chart(ui_metrics_df, 'desktop_avg', 'Desktop Average')
            if chart:
                st.altair_chart(chart, use_container_width=True)
            
            # Text Average
            chart = create_bar_chart(ui_metrics_df, 'text_avg', 'Text Average (UI-Type)')
            if chart:
                st.altair_chart(chart, use_container_width=True)
        
        with col2:
            # Web Average
            chart = create_bar_chart(ui_metrics_df, 'web_avg', 'Web Average')
            if chart:
                st.altair_chart(chart, use_container_width=True)
            
            # Icon Average
            chart = create_bar_chart(ui_metrics_df, 'icon_avg', 'Icon Average (UI-Type)')
            if chart:
                st.altair_chart(chart, use_container_width=True)
        
        # Detailed breakdown
        with st.expander("Detailed UI Type Breakdown"):
            # Create a heatmap-style table
            detailed_metrics = []
            for _, row in ui_metrics_df.iterrows():
                detailed_metrics.append({
                    'Model': row['model'],
                    'Desktop Text': f"{row['desktop_text']:.1f}%",
                    'Desktop Icon': f"{row['desktop_icon']:.1f}%",
                    'Web Text': f"{row['web_text']:.1f}%",
                    'Web Icon': f"{row['web_icon']:.1f}%",
                    'Overall': f"{row['overall']:.1f}%"
                })
            
            if detailed_metrics:
                st.dataframe(pd.DataFrame(detailed_metrics), use_container_width=True)
    
    else:
        # For non-ScreenSpot datasets, show a simple bar chart
        st.subheader("Model Performance")
        
        chart_data = filtered_df[['model', 'overall_accuracy']].copy()
        chart_data.columns = ['Model', 'Accuracy']
        
        chart = alt.Chart(chart_data).mark_bar().encode(
            x=alt.X('Model:N', sort='-y', axis=alt.Axis(labelAngle=-45)),
            y=alt.Y('Accuracy:Q', scale=alt.Scale(domain=[0, 100])),
            tooltip=['Model', 'Accuracy']
        ).properties(
            width=800,
            height=400
        )
        
        st.altair_chart(chart, use_container_width=True)
    
    # Model details table
    with st.expander("Model Details"):
        display_df = filtered_df[['model', 'overall_accuracy', 'total_samples', 'checkpoint_steps', 'training_loss', 'timestamp']].copy()
        display_df.columns = ['Model', 'Accuracy (%)', 'Samples', 'Checkpoint Steps', 'Training Loss', 'Timestamp']
        display_df['Accuracy (%)'] = display_df['Accuracy (%)'].apply(lambda x: f"{x:.2f}")
        display_df['Training Loss'] = display_df['Training Loss'].apply(lambda x: f"{x:.4f}" if pd.notna(x) else "N/A")
        st.dataframe(display_df, use_container_width=True)
    
    # Raw data viewer
    with st.expander("Raw Data"):
        if selected_model != 'All' and len(filtered_df) == 1:
            st.json(filtered_df.iloc[0]['raw_data'])
        else:
            st.info("Select a specific model to view raw data")

if __name__ == "__main__":
    main()