app.py

import json
import warnings
import pandas as pd
import numpy as np
import plotly.express as px
import plotly.graph_objects as go
import pycountry
import gradio as gr

# Suppress pydantic field alias warning from Gradio
warnings.filterwarnings("ignore", message=".*alias.*session_hash.*")
warnings.filterwarnings("ignore", message=".*UnsupportedFieldAttributeWarning.*")

# Load data
with open('dl_stats_0102_0109_top10k.json', 'r') as f:
    dl_by_date = json.load(f)

# Helper function: Convert ISO-2 to ISO-3
def iso2_to_iso3(code):
    """Convert 2-letter ISO code to 3-letter ISO code"""
    if pd.isna(code) or code == 'null':
        return None
    try:
        return pycountry.countries.get(alpha_2=code).alpha_3
    except (AttributeError, KeyError):
        return None

# Choropleth function
def plot_downloads_choropleth(repo_id=None, user_id=None, norm_type='region-log', date=None):
    """Create a choropleth map showing downloads by country."""
    if date not in dl_by_date:
        raise ValueError(f"Date '{date}' not found")
    
    repo_downloads = dl_by_date[date]['repo_downloads']
    org_downloads = dl_by_date[date]['org_downloads']
    total_downloads = dl_by_date[date]['total_downloads']
    
    # Get the appropriate downloads data
    if repo_id and repo_id != "None":
        if repo_id not in repo_downloads:
            raise ValueError(f"Repo '{repo_id}' not found")
        downloads_data = repo_downloads[repo_id]['downloads_by_region']
        title_prefix = f"Repo: {repo_id}"
    elif user_id and user_id != "None":
        if user_id not in org_downloads:
            raise ValueError(f"User/Org '{user_id}' not found")
        downloads_data = org_downloads[user_id]['downloads_by_region']
        title_prefix = f"User/Org: {user_id}"
    else:
        downloads_data = total_downloads['downloads_by_region']
        title_prefix = "Global"
    
    # Prepare data for choropleth
    use_log = norm_type.endswith('-log')
    is_global_norm = norm_type.startswith('global')
    
    # Calculate ranks by country if showing specific org/repo
    ranks_by_country = {}
    if repo_id and repo_id != "None":
        # Calculate repo ranks by country
        for code_2 in downloads_data.keys():
            if code_2 not in total_downloads['downloads_by_region']:
                continue
            repos_in_country = sorted([
                (name, repo_dct['downloads_by_region'].get(code_2, 0))
                for name, repo_dct in repo_downloads.items()
            ], key=lambda x: x[1], reverse=True)
            # Find rank (1-indexed)
            for rank, (name, _) in enumerate(repos_in_country, 1):
                if name == repo_id:
                    ranks_by_country[code_2] = rank
                    break
    elif user_id and user_id != "None":
        # Calculate org ranks by country
        for code_2 in downloads_data.keys():
            if code_2 not in total_downloads['downloads_by_region']:
                continue
            orgs_in_country = sorted([
                (name, org_dct['downloads_by_region'].get(code_2, 0))
                for name, org_dct in org_downloads.items()
            ], key=lambda x: x[1], reverse=True)
            # Find rank (1-indexed)
            for rank, (name, _) in enumerate(orgs_in_country, 1):
                if name == user_id:
                    ranks_by_country[code_2] = rank
                    break
    
    choropleth_data = []
    for code_2, downloads in downloads_data.items():
        code_3 = iso2_to_iso3(code_2)
        if code_3 is None:
            continue
        
        # Calculate normalized value
        if is_global_norm:
            total = total_downloads['total_downloads']
        else:
            total = total_downloads['downloads_by_region'].get(code_2, 0)
        
        if total <= 0:
            continue
        
        value = downloads / total
        color_value = np.log10(value + 1e-10) if use_log else value
        
        # Get rank if available
        rank = ranks_by_country.get(code_2, None)
        
        choropleth_data.append({
            'country_code': code_3,
            'value': color_value,
            'proportion': value,
            'downloads': downloads,
            'rank': rank
        })
    
    if not choropleth_data:
        raise ValueError("No valid data to plot")
    
    df_map = pd.DataFrame(choropleth_data)
    
    # Create color label
    color_label = 'Proportion of Global Downloads' if is_global_norm else "Proportion of Country's Downloads"
    if use_log:
        color_label += ' (log10)'
    
    # Prepare hover data
    hover_data_dict = {
        'downloads': ':,.0f',
        'proportion': ':.4%',
        'value': False
    }
    # Add rank if available (only for specific org/repo)
    has_rank = 'rank' in df_map.columns and df_map['rank'].notna().any()
    if has_rank:
        # Format rank as string with # prefix, handling None values
        df_map['rank_display'] = df_map['rank'].apply(lambda x: f"#{int(x)}" if pd.notna(x) else "N/A")
        hover_data_dict['rank_display'] = True
    
    # Create choropleth map
    fig = px.choropleth(
        df_map,
        locations='country_code',
        color='value',
        hover_name='country_code',
        hover_data=hover_data_dict,
        color_continuous_scale='Viridis',
        title=f'{title_prefix} - {color_label} ({date})',
        labels={'value': color_label, 'country_code': 'Country'}
    )
    
    fig.update_layout(
        geo=dict(showframe=False, showcoastlines=True),
        height=600
    )
    
    return fig

# Top downloads function
def plot_top_downloads(date, data_type='orgs', country_code_iso3=None, top_n=30, use_log=True, show_percentage=True):
    """Create a plot showing top orgs or repos for a given date."""
    if date not in dl_by_date:
        raise ValueError(f"Date '{date}' not found")
    
    date_data = dl_by_date[date]
    downloads_dict = date_data['org_downloads'] if data_type == 'orgs' else date_data['repo_downloads']
    total_dl_dict = date_data['total_downloads']
    
    # Convert ISO-3 to ISO-2 if country code provided
    if country_code_iso3 and country_code_iso3 != "None":
        try:
            country = pycountry.countries.get(alpha_3=country_code_iso3)
            country_code_iso2 = country.alpha_2
            country_name = country.name
        except (AttributeError, KeyError):
            raise ValueError(f"Invalid ISO-3 country code: {country_code_iso3}")
    else:
        country_code_iso2 = None
        country_name = "Global"
    
    # Get total downloads for percentage calculation
    if country_code_iso2 is None:
        total_dl = total_dl_dict['total_downloads']
    else:
        total_dl = total_dl_dict['downloads_by_region'].get(country_code_iso2, 0)
    
    # Get top items
    if country_code_iso2 is None:
        data = sorted([
            (name, dl_dct['total_downloads'])
            for name, dl_dct in downloads_dict.items()
        ], key=lambda x: x[1], reverse=True)[:top_n]
    else:
        data = sorted([
            (name, dl_dct['downloads_by_region'].get(country_code_iso2, 0))
            for name, dl_dct in downloads_dict.items()
        ], key=lambda x: x[1], reverse=True)
        data = [(name, dl) for name, dl in data if dl > 0][:top_n]
    
    if not data:
        raise ValueError(f"No data found for {data_type} in {country_name}")
    
    # Prepare data for plotting
    names = [name for name, _ in data]
    downloads = [dl for _, dl in data]
    percentages = [dl / total_dl * 100 if total_dl > 0 else 0 for dl in downloads]
    
    # Prepare text labels
    if show_percentage:
        text_labels = [f'{dl:,.0f}<br>({pct:.2f}%)' for dl, pct in zip(downloads, percentages)]
    else:
        text_labels = [f'{dl:,.0f}' for dl in downloads]
    
    # Determine x-axis values
    x_values = [np.log10(dl + 1) if use_log else dl for dl in downloads]
    
    # Calculate x-axis range
    if x_values:
        x_min = min(x_values)
        x_max = max(x_values)
        x_range = x_max - x_min
        padding = x_range * 0.1
        x_range_min = max(0, x_min - padding) if not use_log else x_min - padding
        x_range_max = x_max + padding
    else:
        x_range_min = 0
        x_range_max = 1
    
    # Create figure
    fig = go.Figure()
    
    fig.add_trace(go.Bar(
        y=names,
        x=x_values,
        orientation='h',
        text=text_labels,
        textposition='outside',
        textfont=dict(size=11),
        marker_color='lightblue' if data_type == 'orgs' else 'lightcoral',
        hovertemplate='<b>%{y}</b><br>Downloads: %{customdata:,.0f}<extra></extra>',
        customdata=downloads
    ))
    
    # Update layout
    xaxis_title = 'Downloads (log10)' if use_log else 'Downloads'
    title_type = 'Users/Orgs' if data_type == 'orgs' else 'Repos'
    
    fig.update_layout(
        title_text=f'Top {len(data)} {title_type} - {country_name} ({date})',
        height=max(600, len(data) * 40 + 100),
        xaxis_title=xaxis_title,
        font=dict(size=14)
    )
    
    fig.update_yaxes(
        autorange="reversed",
        tickangle=-45,
        tickfont=dict(size=14)
    )
    
    fig.update_xaxes(range=[x_range_min, x_range_max])
    
    return fig

# Prepare dropdown options
# Separate week aggregations from daily dates
all_dates = list(dl_by_date.keys())
week_dates = [d for d in all_dates if 'week' in d]
daily_dates = [d for d in all_dates if 'week' not in d]

# Sort each group
week_dates = sorted(week_dates)
daily_dates = sorted(daily_dates)

# Combine: week aggregations first, then daily dates
dates = week_dates + daily_dates

# Set default date (prefer week aggregation if available)
default_date = "Jan-2026-week-02-to-09" if "Jan-2026-week-02-to-09" in dates else (dates[-1] if dates else None)

# Country dropdown options
country_options = ["None"]
for country in pycountry.countries:
    iso2 = country.alpha_2
    iso3 = country.alpha_3
    name = country.name
    country_options.append(f"{iso2} / {iso3} - {name}")

# Gradio interface functions
def create_map(map_type, org_textbox, repo_textbox, date, norm_type):
    """Create choropleth map."""
    repo_id = None if map_type != "by repo" else (None if not repo_textbox or repo_textbox.strip() == "" else repo_textbox.strip())
    user_id = None if map_type != "by org" else (None if not org_textbox or org_textbox.strip() == "" else org_textbox.strip())
    
    if date is None:
        return None
    
    try:
        fig = plot_downloads_choropleth(
            repo_id=repo_id,
            user_id=user_id,
            norm_type=norm_type,
            date=date
        )
        return fig
    except Exception as e:
        return f"Error: {str(e)}"

def create_top_downloads(region, repo_type, date, top_n, use_log, show_percentage):
    """Create top downloads plot."""
    if date is None:
        return None
    
    # Parse region
    country_code_iso3 = None
    if region and region != "None":
        # Extract ISO-3 code from "ISO2 / ISO3 - Name" format
        parts = region.split(" / ")
        if len(parts) >= 2:
            country_code_iso3 = parts[1].split(" - ")[0]
    
    try:
        fig = plot_top_downloads(
            date=date,
            data_type=repo_type,
            country_code_iso3=country_code_iso3,
            top_n=top_n,
            use_log=use_log,
            show_percentage=show_percentage
        )
        return fig
    except Exception as e:
        return f"Error: {str(e)}"

# Create Gradio interface
with gr.Blocks(title="Hugging Face Downloads Analysis") as app:
    gr.Markdown("""
    # Hugging Face Repository Downloads Analysis
    
    This dashboard visualizes download statistics for Hugging Face repositories, datasets, and models.
    Explore download patterns by organization, repository, country, and date.
    """)
    
    with gr.Row():
        # Main content area
        with gr.Column(scale=4):
            # Map section
            with gr.Accordion("Choropleth Map - Downloads by Country", open=True):
                with gr.Row():
                    map_type = gr.Radio(
                        choices=["everything", "by org", "by repo"],
                        value="everything",
                        label="Map Type"
                    )
                    map_date = gr.Dropdown(
                        choices=dates,
                        value=default_date,
                        label="Date"
                    )
                
                with gr.Row():
                    map_org = gr.Textbox(
                        value="",
                        label="Organization (if 'by org' selected)",
                        placeholder="e.g., Qwen, sentence-transformers",
                        visible=False,
                        info="Enter organization/user name and press Enter (e.g., 'Qwen' or 'sentence-transformers')"
                    )
                    map_repo = gr.Textbox(
                        value="",
                        label="Repository (if 'by repo' selected)",
                        placeholder="e.g., sentence-transformers/all-MiniLM-L6-v2",
                        visible=False,
                        info="Enter full repo name and press Enter (e.g., 'sentence-transformers/all-MiniLM-L6-v2')"
                    )
                
                map_output = gr.Plot(label="Choropleth Map")
            
            # Top downloads section
            with gr.Accordion("Top Downloads - By Organization or Repository", open=True):
                with gr.Row():
                    top_region = gr.Dropdown(
                        choices=country_options,
                        value="None",
                        label="Region (ISO-2 / ISO-3 - Country Name)"
                    )
                    top_repo_type = gr.Radio(
                        choices=["orgs", "repos"],
                        value="orgs",
                        label="Show"
                    )
                    top_date = gr.Dropdown(
                        choices=dates,
                        value=default_date,
                        label="Date"
                    )
                
                top_output = gr.Plot(label="Top Downloads")
        
        # Sidebar with config options
        with gr.Column(scale=1):
            with gr.Group():
                gr.Markdown("### Map Configuration")
                map_norm_type = gr.Radio(
                    choices=["global", "global-log", "region", "region-log"],
                    value="global-log",
                    label="Normalization Type"
                )
            
            with gr.Group():
                gr.Markdown("### Top Downloads Configuration")
                top_n = gr.Slider(
                    minimum=5,
                    maximum=100,
                    value=30,
                    step=5,
                    label="Number of Items"
                )
                top_use_log = gr.Checkbox(
                    value=True,
                    label="Use Log Scale"
                )
                top_show_pct = gr.Checkbox(
                    value=True,
                    label="Show Percentage"
                )
    
    # Helper function to update textbox visibility
    def update_map_textboxes(map_type):
        """Update visibility of org and repo textboxes."""
        return (
            gr.Textbox(visible=(map_type == "by org")),
            gr.Textbox(visible=(map_type == "by repo"))
        )
    
    # Update map org/repo textboxes visibility based on map type
    map_type.change(
        fn=update_map_textboxes,
        inputs=[map_type],
        outputs=[map_org, map_repo]
    )
    
    # Create map
    map_type.change(fn=create_map, inputs=[map_type, map_org, map_repo, map_date, map_norm_type], outputs=[map_output])
    map_org.submit(fn=create_map, inputs=[map_type, map_org, map_repo, map_date, map_norm_type], outputs=[map_output])
    map_repo.submit(fn=create_map, inputs=[map_type, map_org, map_repo, map_date, map_norm_type], outputs=[map_output])
    map_date.change(fn=create_map, inputs=[map_type, map_org, map_repo, map_date, map_norm_type], outputs=[map_output])
    map_norm_type.change(fn=create_map, inputs=[map_type, map_org, map_repo, map_date, map_norm_type], outputs=[map_output])
    
    # Create top downloads
    top_region.change(fn=create_top_downloads, inputs=[top_region, top_repo_type, top_date, top_n, top_use_log, top_show_pct], outputs=[top_output])
    top_repo_type.change(fn=create_top_downloads, inputs=[top_region, top_repo_type, top_date, top_n, top_use_log, top_show_pct], outputs=[top_output])
    top_date.change(fn=create_top_downloads, inputs=[top_region, top_repo_type, top_date, top_n, top_use_log, top_show_pct], outputs=[top_output])
    top_n.change(fn=create_top_downloads, inputs=[top_region, top_repo_type, top_date, top_n, top_use_log, top_show_pct], outputs=[top_output])
    top_use_log.change(fn=create_top_downloads, inputs=[top_region, top_repo_type, top_date, top_n, top_use_log, top_show_pct], outputs=[top_output])
    top_show_pct.change(fn=create_top_downloads, inputs=[top_region, top_repo_type, top_date, top_n, top_use_log, top_show_pct], outputs=[top_output])
    
    # Generate initial plots
    def init_app():
        if default_date:
            map_fig = create_map("everything", "", "", default_date, "region-log")
            top_fig = create_top_downloads("None", "orgs", default_date, 30, True, True)
            return map_fig, top_fig
        return None, None
    
    app.load(fn=init_app, outputs=[map_output, top_output])

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