app.py

# --- START OF FIXED FILE app.py ---

import gradio as gr
import pandas as pd
import plotly.express as px
import time
import json
from datasets import load_dataset

# --- Constants ---
PARAM_CHOICES = ['< 1B', '1B', '5B', '12B', '32B', '64B', '128B', '256B', '> 500B']
PARAM_CHOICES_DEFAULT_INDICES_JSON = json.dumps([0, len(PARAM_CHOICES) - 1])

TOP_K_CHOICES = list(range(5, 51, 5))
HF_DATASET_ID = "evijit/orgstats_daily_data"
TAG_FILTER_CHOICES = [ "Audio & Speech", "Time series", "Robotics", "Music", "Video", "Images", "Text", "Biomedical", "Sciences" ]
PIPELINE_TAGS = [ 'text-generation', 'text-to-image', 'text-classification', 'text2text-generation', 'audio-to-audio', 'feature-extraction', 'image-classification', 'translation', 'reinforcement-learning', 'fill-mask', 'text-to-speech', 'automatic-speech-recognition', 'image-text-to-text', 'token-classification', 'sentence-similarity', 'question-answering', 'image-feature-extraction', 'summarization', 'zero-shot-image-classification', 'object-detection', 'image-segmentation', 'image-to-image', 'image-to-text', 'audio-classification', 'visual-question-answering', 'text-to-video', 'zero-shot-classification', 'depth-estimation', 'text-ranking', 'image-to-video', 'multiple-choice', 'unconditional-image-generation', 'video-classification', 'text-to-audio', 'time-series-forecasting', 'any-to-any', 'video-text-to-text', 'table-question-answering' ]

def load_models_data():
    overall_start_time = time.time()
    print(f"Attempting to load dataset from Hugging Face Hub: {HF_DATASET_ID}")
    try:
        dataset_dict = load_dataset(HF_DATASET_ID)
        df = dataset_dict[list(dataset_dict.keys())[0]].to_pandas()
        if 'params' in df.columns:
            df['params'] = pd.to_numeric(df['params'], errors='coerce').fillna(0)
        else:
            df['params'] = 0
        msg = f"Successfully loaded dataset in {time.time() - overall_start_time:.2f}s."
        print(msg)
        return df, True, msg
    except Exception as e:
        err_msg = f"Failed to load dataset. Error: {e}"
        print(err_msg)
        return pd.DataFrame(), False, err_msg

def get_param_range_values(param_range_labels):
    min_label, max_label = param_range_labels
    min_val = 0.0 if '<' in min_label else float(min_label.replace('B', ''))
    max_val = float('inf') if '>' in max_label else float(max_label.replace('B', ''))
    return min_val, max_val

def make_treemap_data(df, count_by, top_k=25, tag_filter=None, pipeline_filter=None, param_range=None, skip_orgs=None):
    if df is None or df.empty: return pd.DataFrame()
    filtered_df = df.copy()
    col_map = { "Audio & Speech": "is_audio_speech", "Music": "has_music", "Robotics": "has_robot", "Biomedical": "is_biomed", "Time series": "has_series", "Sciences": "has_science", "Video": "has_video", "Images": "has_image", "Text": "has_text" }
    if tag_filter and tag_filter in col_map and col_map[tag_filter] in filtered_df.columns:
        filtered_df = filtered_df[filtered_df[col_map[tag_filter]]]
    if pipeline_filter and "pipeline_tag" in filtered_df.columns:
        filtered_df = filtered_df[filtered_df["pipeline_tag"].astype(str) == pipeline_filter]
    if param_range:
        min_params, max_params = get_param_range_values(param_range)
        is_default_range = (param_range[0] == PARAM_CHOICES[0] and param_range[1] == PARAM_CHOICES[-1])
        if not is_default_range and 'params' in filtered_df.columns:
            if min_params is not None: filtered_df = filtered_df[filtered_df['params'] >= min_params]
            if max_params is not None and max_params != float('inf'): filtered_df = filtered_df[filtered_df['params'] < max_params]
    if skip_orgs and len(skip_orgs) > 0 and "organization" in filtered_df.columns:
        filtered_df = filtered_df[~filtered_df["organization"].isin(skip_orgs)]
    if filtered_df.empty: return pd.DataFrame()
    if count_by not in filtered_df.columns: filtered_df[count_by] = 0.0
    filtered_df[count_by] = pd.to_numeric(filtered_df[count_by], errors='coerce').fillna(0.0)
    org_totals = filtered_df.groupby("organization")[count_by].sum().nlargest(top_k, keep='first')
    top_orgs_list = org_totals.index.tolist()
    treemap_data = filtered_df[filtered_df["organization"].isin(top_orgs_list)][["id", "organization", count_by]].copy()
    treemap_data["root"] = "models"
    return treemap_data

def create_treemap(treemap_data, count_by, title=None):
    if treemap_data.empty:
        fig = px.treemap(names=["No data matches filters"], parents=[""], values=[1])
        fig.update_layout(title="No data matches the selected filters", margin=dict(t=50, l=25, r=25, b=25))
        return fig
    fig = px.treemap(treemap_data, path=["root", "organization", "id"], values=count_by, title=title, color_discrete_sequence=px.colors.qualitative.Plotly)
    fig.update_layout(margin=dict(t=50, l=25, r=25, b=25))
    fig.update_traces(textinfo="label+value+percent root", hovertemplate="<b>%{label}</b><br>%{value:,} " + count_by + "<br>%{percentRoot:.2%} of total<extra></extra>")
    return fig

# Custom head with noUiSlider CSS and JS
custom_head = """
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/noUiSlider/15.7.1/nouislider.min.css">
<script src="https://cdnjs.cloudflare.com/ajax/libs/noUiSlider/15.7.1/nouislider.min.js"></script>
"""

# JavaScript for creating the slider - this will be injected properly
def create_slider_js():
    return f"""
function initializeSlider() {{
    const paramChoices = {json.dumps(PARAM_CHOICES)};
    const sliderContainer = document.getElementById('param-slider');
    const hiddenInput = document.querySelector('#param-range-hidden input');
    
    if (!sliderContainer || !hiddenInput) {{
        console.log('Slider elements not found, retrying...');
        setTimeout(initializeSlider, 100);
        return;
    }}
    
    // Clear any existing slider
    if (sliderContainer.noUiSlider) {{
        sliderContainer.noUiSlider.destroy();
    }}
    
    // Create the slider
    noUiSlider.create(sliderContainer, {{
        start: [0, paramChoices.length - 1],
        connect: true,
        step: 1,
        range: {{
            'min': 0,
            'max': paramChoices.length - 1
        }},
        pips: {{
            mode: 'values',
            values: Array.from({{length: paramChoices.length}}, (_, i) => i),
            density: 100 / (paramChoices.length - 1),
            format: {{
                to: function(value) {{
                    return paramChoices[Math.round(value)];
                }}
            }}
        }}
    }});
    
    // Update hidden input when slider changes
    sliderContainer.noUiSlider.on('update', function(values) {{
        const indices = values.map(v => Math.round(parseFloat(v)));
        hiddenInput.value = JSON.stringify(indices);
        hiddenInput.dispatchEvent(new Event('input', {{ bubbles: true }}));
        
        // Highlight selected range
        document.querySelectorAll('.noUi-value').forEach((pip, index) => {{
            const isSelected = index >= indices[0] && index <= indices[1];
            pip.style.fontWeight = isSelected ? 'bold' : 'normal';
            pip.style.color = isSelected ? '#2563eb' : '#6b7280';
        }});
    }});
    
    // Initial highlight
    document.querySelectorAll('.noUi-value').forEach((pip, index) => {{
        const isSelected = index >= 0 && index <= paramChoices.length - 1;
        pip.style.fontWeight = isSelected ? 'bold' : 'normal';
        pip.style.color = isSelected ? '#2563eb' : '#6b7280';
    }});
    
    console.log('Slider initialized successfully');
}}

// Initialize when DOM is ready
if (document.readyState === 'loading') {{
    document.addEventListener('DOMContentLoaded', initializeSlider);
}} else {{
    initializeSlider();
}}
"""

with gr.Blocks(title="🤗 ModelVerse Explorer", fill_width=True, head=custom_head) as demo:
    models_data_state = gr.State(pd.DataFrame())
    loading_complete_state = gr.State(False)
    
    with gr.Row():
        with gr.Column(scale=1):
            count_by_dropdown = gr.Dropdown(
                label="Metric", 
                choices=[("Downloads (last 30 days)", "downloads"), ("Downloads (All Time)", "downloadsAllTime"), ("Likes", "likes")], 
                value="downloads"
            )
            
            filter_choice_radio = gr.Radio(
                label="Filter Type", 
                choices=["None", "Tag Filter", "Pipeline Filter"], 
                value="None"
            )
            
            tag_filter_dropdown = gr.Dropdown(
                label="Select Tag", 
                choices=TAG_FILTER_CHOICES, 
                value=None, 
                visible=False
            )
            
            pipeline_filter_dropdown = gr.Dropdown(
                label="Select Pipeline Tag", 
                choices=PIPELINE_TAGS, 
                value=None, 
                visible=False
            )
            
            # Parameter range slider section
            with gr.Group():
                gr.Markdown("### Parameters")
                
                # Custom HTML for the slider
                gr.HTML(f"""
                    <div id="param-slider" style="margin: 20px 10px 60px 10px; height: 20px;"></div>
                    <style>
                        #param-slider {{
                            height: 20px;
                        }}
                        .noUi-target {{
                            background: #f1f5f9;
                            border-radius: 10px;
                            border: 1px solid #e2e8f0;
                            box-shadow: none;
                        }}
                        .noUi-connect {{
                            background: #3b82f6;
                            border-radius: 10px;
                        }}
                        .noUi-handle {{
                            width: 20px;
                            height: 20px;
                            right: -10px;
                            top: -5px;
                            background: white;
                            border: 2px solid #3b82f6;
                            border-radius: 50%;
                            box-shadow: 0 2px 4px rgba(0,0,0,0.1);
                            cursor: pointer;
                        }}
                        .noUi-handle:before,
                        .noUi-handle:after {{
                            display: none;
                        }}
                        .noUi-handle:focus {{
                            outline: none;
                        }}
                        .noUi-pips {{
                            color: #6b7280;
                            font-size: 12px;
                        }}
                        .noUi-pips-horizontal {{
                            padding: 10px 0;
                            height: 60px;
                        }}
                        .noUi-value {{
                            font-size: 11px;
                            padding-top: 5px;
                            cursor: pointer;
                        }}
                        .noUi-marker-horizontal.noUi-marker {{
                            background: #e2e8f0;
                            height: 5px;
                            width: 1px;
                        }}
                    </style>
                """)
                
                # Hidden input to store slider values
                param_range_hidden = gr.Textbox(
                    value=PARAM_CHOICES_DEFAULT_INDICES_JSON, 
                    visible=False, 
                    elem_id="param-range-hidden"
                )
            
            top_k_dropdown = gr.Dropdown(
                label="Number of Top Organizations", 
                choices=TOP_K_CHOICES, 
                value=25
            )
            
            skip_orgs_textbox = gr.Textbox(
                label="Organizations to Skip (comma-separated)", 
                value="TheBloke,MaziyarPanahi,unsloth,modularai,Gensyn,bartowski"
            )
            
            generate_plot_button = gr.Button(
                value="Generate Plot", 
                variant="primary", 
                interactive=False
            )

        with gr.Column(scale=3):
            plot_output = gr.Plot()
            status_message_md = gr.Markdown("Initializing...")
            data_info_md = gr.Markdown("")

    # Event handlers
    def _update_button_interactivity(is_loaded_flag): 
        return gr.update(interactive=is_loaded_flag)
    
    loading_complete_state.change(
        fn=_update_button_interactivity, 
        inputs=loading_complete_state, 
        outputs=generate_plot_button
    )

    def _toggle_filters_visibility(choice): 
        return (
            gr.update(visible=choice == "Tag Filter"), 
            gr.update(visible=choice == "Pipeline Filter")
        )
    
    filter_choice_radio.change(
        fn=_toggle_filters_visibility, 
        inputs=filter_choice_radio, 
        outputs=[tag_filter_dropdown, pipeline_filter_dropdown]
    )

    def ui_load_data_controller(progress=gr.Progress()):
        progress(0, desc=f"Loading dataset '{HF_DATASET_ID}'...")
        try:
            current_df, load_success_flag, status_msg_from_load = load_models_data()
            if load_success_flag:
                progress(0.9, desc="Processing data...")
                date_display = "Pre-processed (date unavailable)"
                if 'data_download_timestamp' in current_df.columns and pd.notna(current_df['data_download_timestamp'].iloc[0]):
                    ts = pd.to_datetime(current_df['data_download_timestamp'].iloc[0], utc=True)
                    date_display = ts.strftime('%B %d, %Y, %H:%M:%S %Z')
                param_count = (current_df['params'] > 0).sum() if 'params' in current_df.columns else 0
                data_info_text = f"### Data Information\n- Source: `{HF_DATASET_ID}`\n- Status: {status_msg_from_load}\n- Total models loaded: {len(current_df):,}\n- Models with parameter counts: {param_count:,}\n- Data as of: {date_display}\n"
                status_msg_ui = "Data loaded. Ready to generate plot."
            else:
                data_info_text = f"### Data Load Failed\n- {status_msg_from_load}"
                status_msg_ui = status_msg_from_load
        except Exception as e:
            status_msg_ui = f"An unexpected error occurred: {str(e)}"
            data_info_text = f"### Critical Error\n- {status_msg_ui}"
            load_success_flag = False
            print(f"Critical error in ui_load_data_controller: {e}")
        return current_df, load_success_flag, data_info_text, status_msg_ui

    def ui_generate_plot_controller(metric_choice, filter_type, tag_choice, pipeline_choice, 
                                   param_range_json, k_orgs, skip_orgs_input, df_current_models, progress=gr.Progress()):
        if df_current_models is None or df_current_models.empty:
            return create_treemap(pd.DataFrame(), metric_choice, "Error: Model Data Not Loaded"), "Model data is not loaded."
        
        progress(0.1, desc="Preparing data...")
        tag_to_use = tag_choice if filter_type == "Tag Filter" else None
        pipeline_to_use = pipeline_choice if filter_type == "Pipeline Filter" else None
        orgs_to_skip = [org.strip() for org in skip_orgs_input.split(',') if org.strip()]
        
        try:
            param_range_indices = json.loads(param_range_json)
        except:
            param_range_indices = [0, len(PARAM_CHOICES) - 1]
            
        min_label = PARAM_CHOICES[int(param_range_indices[0])]
        max_label = PARAM_CHOICES[int(param_range_indices[1])]
        param_labels_for_filtering = [min_label, max_label]
        
        treemap_df = make_treemap_data(df_current_models, metric_choice, k_orgs, tag_to_use, pipeline_to_use, param_labels_for_filtering, orgs_to_skip)
        
        progress(0.7, desc="Generating plot...")
        title_labels = {"downloads": "Downloads (last 30 days)", "downloadsAllTime": "Downloads (All Time)", "likes": "Likes"}
        chart_title = f"HuggingFace Models - {title_labels.get(metric_choice, metric_choice)} by Organization"
        plotly_fig = create_treemap(treemap_df, metric_choice, chart_title)
        
        if treemap_df.empty:
            plot_stats_md = "No data matches the selected filters. Please try different options."
        else:
            total_items_in_plot = len(treemap_df['id'].unique())
            total_value_in_plot = treemap_df[metric_choice].sum()
            plot_stats_md = f"## Plot Statistics\n- **Models shown**: {total_items_in_plot:,}\n- **Total {metric_choice}**: {int(total_value_in_plot):,}"
        return plotly_fig, plot_stats_md

    # Load data on startup and initialize slider
    demo.load(
        fn=ui_load_data_controller, 
        inputs=[], 
        outputs=[models_data_state, loading_complete_state, data_info_md, status_message_md]
    )
    
    # Initialize slider after page loads
    demo.load(
        fn=lambda: None,
        inputs=[],
        outputs=[],
        js=create_slider_js()
    )

    # Generate plot button click handler
    generate_plot_button.click(
        fn=ui_generate_plot_controller,
        inputs=[count_by_dropdown, filter_choice_radio, tag_filter_dropdown, pipeline_filter_dropdown,
                param_range_hidden, top_k_dropdown, skip_orgs_textbox, models_data_state],
        outputs=[plot_output, status_message_md]
    )

if __name__ == "__main__":
    print(f"Application starting...")
    demo.queue().launch()

# --- END OF FIXED FILE