app.py

#!/usr/bin/env python3
"""
VIBE Multi-Stage Evaluation Dashboard - FastHTML
Interactive dashboard for exploring checkpoint evaluation results across training stages.

Usage:
    python vibe_multi_stage_dashboard.py --results-dir data/multi_stage_results
"""

import json
import argparse
import os
from pathlib import Path
from typing import List, Dict, Any, Tuple
import plotly.graph_objects as go
import plotly.express as px
from fasthtml.common import *
import pandas as pd
from collections import defaultdict

# Global data store
DATA = {}  # {stage_id: {checkpoint_label: data}}
STAGE_METADATA = {}  # {stage_id: {name, attention_window, checkpoint_count}}

def load_multi_stage_data(results_base_dir: Path) -> Tuple[Dict, Dict]:
    """Load all checkpoint results organized by stage."""
    stage_data = defaultdict(dict)
    stage_metadata = {}

    # Find all stage subdirectories
    for stage_dir in sorted(results_base_dir.iterdir()):
        if not stage_dir.is_dir():
            continue

        stage_id = stage_dir.name.replace('_results', '')

        # Load all JSON files in this stage directory
        json_files = list(stage_dir.glob("*_evaluation_results.json"))

        if not json_files:
            continue

        print(f"Loading {stage_id}: {len(json_files)} checkpoints")

        for json_file in sorted(json_files):
            with open(json_file, 'r') as f:
                data = json.load(f)
                checkpoint_label = data.get('checkpoint_id', json_file.stem)
                checkpoint_step = data['checkpoint_step']

                # Store with unique key
                key = f"{checkpoint_label}_step_{checkpoint_step}"
                stage_data[stage_id][key] = data

        # Extract stage metadata from first checkpoint
        if stage_data[stage_id]:
            first_checkpoint = list(stage_data[stage_id].values())[0]

            # Infer stage info from stage_id
            stage_name = stage_id.replace('_', ' ').title()
            attention_window = None

            if 'pre_training' in stage_id:
                stage_name = 'Pre-training'
            elif 'stage1_8k' in stage_id:
                stage_name = 'Stage 1 - 8k Attention'
                attention_window = '8k'
            elif 'stage2_64k' in stage_id:
                stage_name = 'Stage 2 - 64k Attention'
                attention_window = '64k'
            elif 'stage3_128k' in stage_id:
                stage_name = 'Stage 3 - 128k Attention'
                attention_window = '128k'
            elif 'stage4_512k' in stage_id:
                stage_name = 'Stage 4 - 512k Attention'
                attention_window = '512k'

            stage_metadata[stage_id] = {
                'name': stage_name,
                'attention_window': attention_window,
                'checkpoint_count': len(stage_data[stage_id]),
                'checkpoints': sorted(stage_data[stage_id].keys(), key=lambda k: stage_data[stage_id][k]['checkpoint_step'])
            }

    print(f"\n✓ Loaded {len(stage_data)} stages with {sum(len(v) for v in stage_data.values())} total checkpoints")
    return dict(stage_data), stage_metadata

def get_all_categories() -> List[str]:
    """Get all unique categories across all stages and checkpoints."""
    categories = set()
    for stage_checkpoints in DATA.values():
        for checkpoint_data in stage_checkpoints.values():
            for eval in checkpoint_data['evaluations']:
                categories.add(eval['category'])
    return sorted(categories)

def get_stage_average_score(stage_id: str) -> float:
    """Calculate average score across all checkpoints in a stage."""
    all_scores = []
    for checkpoint_data in DATA[stage_id].values():
        all_scores.append(checkpoint_data['mean_score'])
    return sum(all_scores) / len(all_scores) if all_scores else 0.0

def get_category_scores_by_stage(category: str) -> Dict[str, List[float]]:
    """Get scores for a specific category across all stages."""
    stage_scores = {}

    for stage_id in DATA.keys():
        scores = []
        for checkpoint_data in DATA[stage_id].values():
            cat_scores = [e['overall_score'] for e in checkpoint_data['evaluations']
                         if e['category'] == category]
            if cat_scores:
                scores.extend(cat_scores)

        if scores:
            stage_scores[stage_id] = scores

    return stage_scores

def create_overview_charts():
    """Create overview charts showing progression across stages."""
    if not DATA:
        # Return empty charts if no data
        empty_fig = go.Figure()
        empty_fig.update_layout(title='No data loaded')
        return empty_fig, empty_fig

    stage_ids = sorted(DATA.keys())

    # Build checkpoint list with stage info
    checkpoint_list = []
    stage_colors = {
        'pre_training': '#64748b',
        'stage1_8k': '#3b82f6',
        'stage2_64k': '#8b5cf6',
        'stage3_128k': '#ec4899',
        'stage4_512k': '#f59e0b'
    }

    for stage_id in stage_ids:
        stage_name = STAGE_METADATA[stage_id]['name']
        for checkpoint_key in sorted(DATA[stage_id].keys(), key=lambda k: DATA[stage_id][k]['checkpoint_step']):
            checkpoint_data = DATA[stage_id][checkpoint_key]
            checkpoint_list.append({
                'stage_id': stage_id,
                'stage_name': stage_name,
                'checkpoint_key': checkpoint_key,
                'checkpoint_step': checkpoint_data['checkpoint_step'],
                'mean_score': checkpoint_data['mean_score'],
                'color': stage_colors.get(stage_id, '#666'),
                'data': checkpoint_data
            })

    # 1. Stage progression chart - showing individual checkpoints
    checkpoint_labels = [f"Step {cp['checkpoint_step']}" for cp in checkpoint_list]
    checkpoint_scores = [cp['mean_score'] for cp in checkpoint_list]
    checkpoint_colors = [cp['color'] for cp in checkpoint_list]
    hover_text = [f"{cp['stage_name']}<br>Step: {cp['checkpoint_step']}<br>Score: {cp['mean_score']:.3f}"
                  for cp in checkpoint_list]

    fig_stage_progression = go.Figure()
    fig_stage_progression.add_trace(go.Bar(
        x=checkpoint_labels,
        y=checkpoint_scores,
        marker_color=checkpoint_colors,
        text=[f'{score:.3f}' for score in checkpoint_scores],
        textposition='outside',
        hovertext=hover_text,
        hoverinfo='text'
    ))
    fig_stage_progression.update_layout(
        title='Checkpoint Performance Across Training Stages',
        xaxis_title='Checkpoint (Training Step)',
        yaxis_title='Mean Score',
        yaxis=dict(range=[0, 1.1]),
        height=500,
        template='plotly_white',
        xaxis=dict(tickangle=-45)
    )

    # 2. Category performance across checkpoints (heatmap)
    categories = get_all_categories()
    heatmap_data = []

    for cat in categories:
        row = []
        for cp in checkpoint_list:
            cat_scores = [e['overall_score'] for e in cp['data']['evaluations']
                         if e['category'] == cat]
            avg_score = sum(cat_scores) / len(cat_scores) if cat_scores else 0
            row.append(avg_score)
        heatmap_data.append(row)

    fig_category_heatmap = go.Figure(data=go.Heatmap(
        z=heatmap_data,
        x=checkpoint_labels,
        y=categories,
        colorscale='RdYlGn',
        zmin=0,
        zmax=1,
        text=[[f'{val:.2f}' for val in row] for row in heatmap_data],
        texttemplate='%{text}',
        textfont={"size": 9},
        colorbar=dict(title="Score"),
        hovertemplate='Checkpoint: %{x}<br>Category: %{y}<br>Score: %{z:.3f}<extra></extra>'
    ))
    fig_category_heatmap.update_layout(
        title='Category Performance Across All Checkpoints',
        xaxis_title='Checkpoint (Training Step)',
        yaxis_title='Category',
        height=600,
        template='plotly_white',
        xaxis=dict(tickangle=-45)
    )

    return fig_stage_progression, fig_category_heatmap

def create_category_progression_chart(category: str):
    """Create line chart showing category progression across stages."""
    stage_ids = sorted(DATA.keys())
    stage_names = [STAGE_METADATA[s]['name'] for s in stage_ids]

    avg_scores = []
    for stage_id in stage_ids:
        scores = []
        for checkpoint_data in DATA[stage_id].values():
            cat_scores = [e['overall_score'] for e in checkpoint_data['evaluations']
                         if e['category'] == category]
            scores.extend(cat_scores)
        avg_scores.append(sum(scores) / len(scores) if scores else 0)

    fig = go.Figure()
    fig.add_trace(go.Scatter(
        x=stage_names,
        y=avg_scores,
        mode='lines+markers',
        name=category,
        line=dict(width=3),
        marker=dict(size=10)
    ))
    fig.update_layout(
        title=f'{category} Performance Across Stages',
        xaxis_title='Training Stage',
        yaxis_title='Average Score',
        yaxis=dict(range=[0, 1]),
        height=400,
        template='plotly_white'
    )

    return fig.to_html(include_plotlyjs=False, div_id=f'category_{category}')

# FastHTML App
css_link = Link(rel='stylesheet', href='https://cdn.jsdelivr.net/npm/@picocss/pico@2/css/pico.min.css')
plotly_script = Script(src='https://cdn.plot.ly/plotly-2.27.0.min.js')
custom_style = Style("""
    .score-badge {
        display: inline-block;
        padding: 4px 12px;
        border-radius: 12px;
        font-weight: bold;
        margin: 2px;
    }
    .score-high { background: #10b981; color: white; }
    .score-medium { background: #f59e0b; color: white; }
    .score-low { background: #ef4444; color: white; }
    .stage-badge {
        display: inline-block;
        padding: 6px 16px;
        border-radius: 16px;
        font-weight: bold;
        margin: 4px;
        font-size: 0.9rem;
    }
    .stage-pre { background: #64748b; color: white; }
    .stage-1 { background: #3b82f6; color: white; }
    .stage-2 { background: #8b5cf6; color: white; }
    .stage-3 { background: #ec4899; color: white; }
    .stage-4 { background: #f59e0b; color: white; }
    nav { background: #1e293b; padding: 1rem; margin-bottom: 2rem; }
    nav a { color: white; margin: 0 1rem; text-decoration: none; }
    nav a:hover { text-decoration: underline; }
    .stats-grid {
        display: grid;
        grid-template-columns: repeat(auto-fit, minmax(200px, 1fr));
        gap: 1rem;
        margin: 2rem 0;
    }
    .stat-card {
        background: #f8fafc;
        padding: 1rem;
        border-radius: 8px;
        border-left: 4px solid #3b82f6;
    }
    .stat-value {
        font-size: 2rem;
        font-weight: bold;
        color: #1e293b;
    }
    .stat-label {
        font-size: 0.875rem;
        color: #64748b;
        text-transform: uppercase;
    }
""")

app = FastHTML(hdrs=(css_link, plotly_script, custom_style))

def ensure_data_loaded():
    """Ensure DATA is loaded, reloading from env var if needed."""
    global DATA, STAGE_METADATA
    if not DATA and 'VIBE_RESULTS_DIR' in os.environ:
        results_dir = Path(os.environ['VIBE_RESULTS_DIR'])
        DATA, STAGE_METADATA = load_multi_stage_data(results_dir)

@app.get("/")
def home():
    """Overview page showing all stages."""
    ensure_data_loaded()

    if not DATA:
        return Title("Error"), Main(
            H1("No Data Loaded"),
            P("Could not load evaluation results. Please check the results directory."),
            cls="container"
        )

    fig_stage_prog, fig_cat_heat = create_overview_charts()

    total_checkpoints = sum(meta['checkpoint_count'] for meta in STAGE_METADATA.values())
    total_stages = len(STAGE_METADATA)
    total_categories = len(get_all_categories())

    # Overall average score
    all_scores = []
    for stage_checkpoints in DATA.values():
        for checkpoint_data in stage_checkpoints.values():
            all_scores.append(checkpoint_data['mean_score'])
    overall_avg = sum(all_scores) / len(all_scores) if all_scores else 0

    return Title("VIBE Multi-Stage Dashboard"), Main(
        Nav(
            A("Overview", href="/"),
            A("Stage Details", href="/stages"),
            A("Category Analysis", href="/categories"),
            A("Checkpoint Explorer", href="/explorer"),
            style="background: #1e293b; padding: 1rem;"
        ),

        H1("🚀 VIBE Multi-Stage Evaluation Dashboard"),
        P("Analyzing model performance across training stages with progressive attention scaling"),

        # Summary stats
        Div(
            Div(
                Div(f"{total_stages}", cls="stat-value"),
                Div("Training Stages", cls="stat-label"),
                cls="stat-card"
            ),
            Div(
                Div(f"{total_checkpoints}", cls="stat-value"),
                Div("Total Checkpoints", cls="stat-label"),
                cls="stat-card"
            ),
            Div(
                Div(f"{total_categories}", cls="stat-value"),
                Div("Categories", cls="stat-label"),
                cls="stat-card"
            ),
            Div(
                Div(f"{overall_avg:.3f}", cls="stat-value"),
                Div("Overall Avg Score", cls="stat-label"),
                cls="stat-card"
            ),
            cls="stats-grid"
        ),

        H2("📊 Stage Progression"),
        Div(NotStr(fig_stage_prog.to_html(include_plotlyjs=False, div_id='chart1'))),

        H2("🎯 Category Performance Across Stages"),
        Div(NotStr(fig_cat_heat.to_html(include_plotlyjs=False, div_id='chart2'))),

        cls="container"
    )

@app.get("/stages")
def stages(stage: str = None):
    """Detailed view of individual stages."""
    ensure_data_loaded()
    stage_ids = sorted(DATA.keys())

    if not stage or stage == '':
        stage = stage_ids[0]

    stage_meta = STAGE_METADATA[stage]
    checkpoints = DATA[stage]

    # Calculate stage statistics
    all_scores = []
    for checkpoint_data in checkpoints.values():
        all_scores.extend([e['overall_score'] for e in checkpoint_data['evaluations']])

    avg_score = sum(all_scores) / len(all_scores) if all_scores else 0
    min_score = min(all_scores) if all_scores else 0
    max_score = max(all_scores) if all_scores else 0

    # Per-category stats for this stage
    category_stats = {}
    for checkpoint_data in checkpoints.values():
        for eval in checkpoint_data['evaluations']:
            cat = eval['category']
            if cat not in category_stats:
                category_stats[cat] = []
            category_stats[cat].append(eval['overall_score'])

    return Title(f"Stage Details - {stage_meta['name']}"), Main(
        Nav(
            A("Overview", href="/"),
            A("Stage Details", href="/stages"),
            A("Category Analysis", href="/categories"),
            A("Checkpoint Explorer", href="/explorer"),
            style="background: #1e293b; padding: 1rem;"
        ),

        H1(f"Stage: {stage_meta['name']}"),

        Form(
            Label("Select Stage:", _for="stage"),
            Select(
                *[Option(STAGE_METADATA[s]['name'], value=s, selected=(s == stage))
                  for s in stage_ids],
                name="stage",
                id="stage",
                onchange="this.form.submit()"
            ),
            method="get",
            style="margin: 2rem 0;"
        ),

        # Stage metadata
        Div(
            Div(
                Div(f"{stage_meta['checkpoint_count']}", cls="stat-value"),
                Div("Checkpoints", cls="stat-label"),
                cls="stat-card"
            ),
            Div(
                Div(f"{stage_meta['attention_window'] or 'N/A'}", cls="stat-value"),
                Div("Attention Window", cls="stat-label"),
                cls="stat-card"
            ),
            Div(
                Div(f"{avg_score:.3f}", cls="stat-value"),
                Div("Average Score", cls="stat-label"),
                cls="stat-card"
            ),
            Div(
                Div(f"{min_score:.3f} - {max_score:.3f}", cls="stat-value"),
                Div("Score Range", cls="stat-label"),
                cls="stat-card"
            ),
            cls="stats-grid"
        ),

        H2("Category Performance in This Stage"),
        Table(
            Thead(
                Tr(
                    Th("Category"),
                    Th("Avg Score"),
                    Th("Min Score"),
                    Th("Max Score"),
                    Th("# Evaluations")
                )
            ),
            Tbody(
                *[Tr(
                    Td(cat),
                    Td(Span(f"{sum(scores)/len(scores):.3f}",
                           cls=f"score-badge {'score-high' if sum(scores)/len(scores) >= 0.7 else 'score-medium' if sum(scores)/len(scores) >= 0.4 else 'score-low'}")),
                    Td(f"{min(scores):.3f}"),
                    Td(f"{max(scores):.3f}"),
                    Td(str(len(scores)))
                ) for cat, scores in sorted(category_stats.items())]
            )
        ),

        H2("Checkpoints in This Stage"),
        Ul(
            *[Li(f"{checkpoint_key}: Mean Score = {checkpoints[checkpoint_key]['mean_score']:.3f}")
              for checkpoint_key in sorted(checkpoints.keys(), key=lambda k: checkpoints[k]['checkpoint_step'])]
        ),

        cls="container"
    )

@app.get("/categories")
def categories(category: str = None):
    """Category-focused analysis across stages."""
    ensure_data_loaded()
    all_categories = get_all_categories()

    if not category or category == '':
        category = all_categories[0]

    stage_ids = sorted(DATA.keys())

    return Title(f"Category Analysis - {category}"), Main(
        Nav(
            A("Overview", href="/"),
            A("Stage Details", href="/stages"),
            A("Category Analysis", href="/categories"),
            A("Checkpoint Explorer", href="/explorer"),
            style="background: #1e293b; padding: 1rem;"
        ),

        H1(f"Category Analysis: {category}"),

        Form(
            Label("Select Category:", _for="category"),
            Select(
                *[Option(c, value=c, selected=(c == category)) for c in all_categories],
                name="category",
                id="category",
                onchange="this.form.submit()"
            ),
            method="get",
            style="margin: 2rem 0;"
        ),

        H2(f"{category} Performance Across Stages"),
        Div(NotStr(create_category_progression_chart(category))),

        H2("Stage-by-Stage Breakdown"),
        Table(
            Thead(
                Tr(
                    Th("Stage"),
                    Th("Avg Score"),
                    Th("# Evaluations"),
                    Th("Min"),
                    Th("Max")
                )
            ),
            Tbody(
                *[Tr(
                    Td(STAGE_METADATA[stage_id]['name']),
                    Td(Span(
                        f"{sum(scores)/len(scores):.3f}" if scores else "N/A",
                        cls=f"score-badge {'score-high' if scores and sum(scores)/len(scores) >= 0.7 else 'score-medium' if scores and sum(scores)/len(scores) >= 0.4 else 'score-low'}"
                    )),
                    Td(str(len(scores))),
                    Td(f"{min(scores):.3f}" if scores else "N/A"),
                    Td(f"{max(scores):.3f}" if scores else "N/A")
                ) for stage_id in stage_ids
                  for scores in [[e['overall_score']
                                 for checkpoint_data in DATA[stage_id].values()
                                 for e in checkpoint_data['evaluations']
                                 if e['category'] == category]]]
            )
        ),

        cls="container"
    )

@app.get("/explorer")
def explorer(stage: str = None, checkpoint: str = None, category: str = None):
    """Detailed checkpoint explorer."""
    ensure_data_loaded()
    stage_ids = sorted(DATA.keys())

    if not stage or stage == '':
        stage = stage_ids[0]

    checkpoints = sorted(DATA[stage].keys(), key=lambda k: DATA[stage][k]['checkpoint_step'])

    # If checkpoint is not specified or doesn't exist in current stage, use first checkpoint
    if not checkpoint or checkpoint == '' or checkpoint not in checkpoints:
        checkpoint = checkpoints[0]

    checkpoint_data = DATA[stage][checkpoint]
    categories = get_all_categories()

    # Filter evaluations
    evals = checkpoint_data['evaluations']
    if category and category != '':
        evals = [e for e in evals if e['category'] == category]

    return Title("Checkpoint Explorer"), Main(
        Nav(
            A("Overview", href="/"),
            A("Stage Details", href="/stages"),
            A("Category Analysis", href="/categories"),
            A("Checkpoint Explorer", href="/explorer"),
            style="background: #1e293b; padding: 1rem;"
        ),

        H1("🔍 Checkpoint Explorer"),

        Form(
            Div(
                Label("Stage:", _for="stage"),
                Select(
                    *[Option(STAGE_METADATA[s]['name'], value=s, selected=(s == stage))
                      for s in stage_ids],
                    name="stage",
                    id="stage",
                    onchange="this.form.submit()"
                ),
                style="display: inline-block; margin-right: 1rem;"
            ),
            Div(
                Label("Checkpoint:", _for="checkpoint"),
                Select(
                    *[Option(c, value=c, selected=(c == checkpoint))
                      for c in checkpoints],
                    name="checkpoint",
                    id="checkpoint",
                    onchange="this.form.submit()"
                ),
                style="display: inline-block; margin-right: 1rem;"
            ),
            Div(
                Label("Category:", _for="category"),
                Select(
                    Option("All", value="", selected=(not category)),
                    *[Option(c, value=c, selected=(c == category)) for c in categories],
                    name="category",
                    id="category",
                    onchange="this.form.submit()"
                ),
                style="display: inline-block;"
            ),
            method="get",
            style="margin: 2rem 0;"
        ),

        P(f"Showing {len(evals)} evaluation(s) | Mean Score: {checkpoint_data['mean_score']:.3f}"),

        # Display evaluations
        *[Article(
            Header(
                H3(f"Prompt: {eval['prompt_id']}", style="margin-bottom: 0.5rem;"),
                Div(
                    Span(f"Category: {eval['category']}", cls="score-badge score-medium"),
                    Span(
                        f"Score: {eval['overall_score']:.2f}",
                        cls=f"score-badge {'score-high' if eval['overall_score'] >= 0.7 else 'score-medium' if eval['overall_score'] >= 0.4 else 'score-low'}"
                    )
                )
            ),
            Details(
                Summary("View Details"),
                H4("Prompt:"),
                P(eval['prompt_text'], style="background: #f1f5f9; padding: 1rem; border-radius: 4px; color: #000;"),
                H4("Model Response:"),
                P(eval['model_response'], style="background: #f1f5f9; padding: 1rem; border-radius: 4px; color: #000;"),
                H4("Assessment:"),
                P(eval['overall_assessment'], style="font-style: italic; background: #fef3c7; padding: 1rem; border-radius: 4px;")
            ),
            style="margin-bottom: 2rem;"
        ) for eval in evals],

        cls="container"
    )

def main():
    parser = argparse.ArgumentParser(description='VIBE Multi-Stage Evaluation Dashboard')
    parser.add_argument('--results-dir', type=str, required=True,
                        help='Base directory containing stage result subdirectories')
    parser.add_argument('--port', type=int, default=5001,
                        help='Port to run dashboard on (default: 5001)')
    parser.add_argument('--host', type=str, default='127.0.0.1',
                        help='Host to run dashboard on (default: 127.0.0.1)')

    args = parser.parse_args()

    results_dir = Path(args.results_dir)
    if not results_dir.exists():
        print(f"❌ Error: Directory not found: {results_dir}")
        return

    print(f"Loading multi-stage data from: {results_dir}\n")

    # Set environment variable for worker processes
    os.environ['VIBE_RESULTS_DIR'] = str(results_dir.absolute())

    global DATA, STAGE_METADATA
    DATA, STAGE_METADATA = load_multi_stage_data(results_dir)

    if not DATA:
        print(f"❌ Error: No evaluation result files found in {results_dir}")
        return

    print(f"\n✓ Dashboard ready!")
    print(f"  Stages: {len(STAGE_METADATA)}")
    print(f"  Total Checkpoints: {sum(meta['checkpoint_count'] for meta in STAGE_METADATA.values())}")
    print(f"  Categories: {len(get_all_categories())}")
    print(f"\n🚀 Starting server at http://{args.host}:{args.port}")
    print(f"   Press Ctrl+C to stop\n")

    serve(port=args.port, host=args.host)

if __name__ == "__main__":
    main()