app.py

import gradio as gr
import pandas as pd
from plotly import graph_objects as go
import plotly.io as pio
import plotly.express as px

# Set the default theme to "plotly_dark"
pio.templates.default = "plotly_dark"


def process_dataset():
    """
    Process the dataset and perform the following operations:
    1. Read the file_counts_and_sizes, repo_by_size_df, unique_files_df, and file_extensions data from parquet files.
    2. Convert the total size to petabytes and format it to two decimal places.
    3. Capitalize the 'type' column in the file_counts_and_sizes dataframe.
    4. Rename the columns in the file_counts_and_sizes dataframe.
    5. Sort the file_counts_and_sizes dataframe by total size in descending order.
    6. Drop rows with missing values in the 'extension' column of the file_extensions dataframe.
    7. Return the repo_by_size_df, unique_files_df, file_counts_and_sizes, and file_extensions dataframes.
    """

    file_counts_and_sizes = pd.read_parquet(
        "hf://datasets/xet-team/lfs-analysis-data/transformed/file_counts_and_sizes.parquet"
    )
    repo_by_size_df = pd.read_parquet(
        "hf://datasets/xet-team/lfs-analysis-data/transformed/repo_by_size.parquet"
    )
    unique_files_df = pd.read_parquet(
        "hf://datasets/xet-team/lfs-analysis-data/transformed/repo_by_size_file_dedupe.parquet"
    )
    file_extensions = pd.read_parquet(
        "hf://datasets/xet-team/lfs-analysis-data/transformed/file_extensions.parquet"
    )

    # read the file_extensions_by_month.parquet file
    file_extensions_by_month = pd.read_parquet(
        "hf://datasets/xet-team/lfs-analysis-data/transformed/file_extensions_by_month.parquet"
    )
    # drop any nas
    file_extensions_by_month = file_extensions_by_month.dropna()

    # Convert the total size to petabytes and format to two decimal places
    file_counts_and_sizes = format_dataframe_size_column(
        file_counts_and_sizes, "total_size"
    )

    file_counts_and_sizes["type"] = file_counts_and_sizes["type"].str.capitalize()
    # update the column name to 'total size (PB)'
    file_counts_and_sizes = file_counts_and_sizes.rename(
        columns={
            "type": "Repository Type",
            "num_files": "Number of Files",
            "total_size": "Total Size (PBs)",
        }
    )
    # sort the dataframe by total size in descending order
    file_counts_and_sizes = file_counts_and_sizes.sort_values(
        by="Total Size (PBs)", ascending=False
    )

    # drop nas from the extension column
    file_extensions = file_extensions.dropna(subset=["extension"])

    return (
        repo_by_size_df,
        unique_files_df,
        file_counts_and_sizes,
        file_extensions,
        file_extensions_by_month,
    )


def format_dataframe_size_column(_df, column_name):
    """
    Format the size to petabytes and return the formatted size.
    """
    _df[column_name] = _df[column_name] / 1e15
    _df[column_name] = _df[column_name].map("{:.2f}".format)
    return _df


def cumulative_growth_plot_analysis(df, df_compressed):
    """
    Calculates the cumulative growth of models, spaces, and datasets over time and generates a plot and dataframe from the analysis.

    Args:
        df (DataFrame): The input dataframe containing the data.
        df_compressed (DataFrame): The input dataframe containing the compressed data.

    Returns:
        tuple: A tuple containing two elements:
            - fig (Figure): The Plotly figure showing the cumulative growth of models, spaces, and datasets over time.
            - last_10_months (DataFrame): The last 10 months of data showing the month-to-month growth in petabytes.

    Raises:
        None
    """
    # Convert year and month into a datetime column
    df["date"] = pd.to_datetime(df[["year", "month"]].assign(day=1))
    df_compressed["date"] = pd.to_datetime(
        df_compressed[["year", "month"]].assign(day=1)
    )

    # Sort by date to ensure correct cumulative sum
    df = df.sort_values(by="date")
    df_compressed = df_compressed.sort_values(by="date")

    # Pivot the dataframe to get the totalsize for each type
    pivot_df = df.pivot_table(
        index="date", columns="type", values="totalsize", aggfunc="sum"
    ).fillna(0)
    pivot_df_compressed = df_compressed.pivot_table(
        index="date", columns="type", values="totalsize", aggfunc="sum"
    ).fillna(0)

    # Calculate cumulative sum for each type
    cumulative_df = pivot_df.cumsum()
    cumulative_df_compressed = pivot_df_compressed.cumsum()

    last_10_months = cumulative_df.tail(10).copy()
    last_10_months["total"] = last_10_months.sum(axis=1)
    last_10_months["total_change"] = last_10_months["total"].diff()
    last_10_months = format_dataframe_size_column(last_10_months, "total_change")
    last_10_months["date"] = cumulative_df.tail(10).index
    # drop the dataset, model, and space
    last_10_months = last_10_months.drop(columns=["model", "space", "dataset"])
    # pretiffy the date column to not have 00:00:00
    last_10_months["date"] = last_10_months["date"].dt.strftime("%Y-%m")
    # drop the first row
    last_10_months = last_10_months.drop(last_10_months.index[0])
    # order the columns date, total, total_change
    last_10_months = last_10_months[["date", "total_change"]]
    # rename the columns
    last_10_months = last_10_months.rename(
        columns={"date": "Date", "total_change": "Month-to-Month Growth (PBs)"}
    )

    # Create a Plotly figure
    fig = go.Figure()

    # Define a color map for each type
    color_map = {
        "model": px.colors.qualitative.Alphabet[3],
        "space": px.colors.qualitative.Alphabet[2],
        "dataset": px.colors.qualitative.Alphabet[9],
    }

    # Add a scatter trace for each type
    for column in cumulative_df.columns:
        fig.add_trace(
            go.Scatter(
                x=cumulative_df.index,
                y=cumulative_df[column] / 1e15,  # Convert to petabytes
                mode="lines",
                name=column.capitalize(),
                line=dict(color=color_map.get(column, "black")),  # Use color map
            )
        )

    # Add a scatter trace for each type
    for column in cumulative_df_compressed.columns:
        fig.add_trace(
            go.Scatter(
                x=cumulative_df_compressed.index,
                y=cumulative_df_compressed[column] / 1e15,  # Convert to petabytes
                mode="lines",
                name=column.capitalize() + " (Compressed)",
                line=dict(color=color_map.get(column, "black"), dash="dash"),
            )
        )

    # Update layout
    fig.update_layout(
        title="Cumulative Growth of Models, Spaces, and Datasets Over Time<br><sup>Dotted lines represent growth with file-level deduplication</sup>",
        xaxis_title="Date",
        yaxis_title="Cumulative Size (PBs)",
        legend_title="Type",
        yaxis=dict(tickformat=".2f"),  # Format y-axis labels to 2 decimal places
    )
    return fig, last_10_months


def plot_total_sum(by_type_arr):
    # Sort the array by size in decreasing order
    by_type_arr = sorted(by_type_arr, key=lambda x: x[1], reverse=True)

    # Create a Plotly figure
    fig = go.Figure()

    # Add a bar trace for each type
    for type, size in by_type_arr:
        fig.add_trace(
            go.Bar(
                x=[type],
                y=[size / 1e15],  # Convert to petabytes
                name=type.capitalize(),
            )
        )

    # Update layout
    fig.update_layout(
        title="Top 20 File Extensions by Total Size",
        xaxis_title="File Extension",
        yaxis_title="Total Size (PBs)",
        yaxis=dict(tickformat=".2f"),  # Format y-axis labels to 2 decimal places
        colorway=px.colors.qualitative.Alphabet,  # Use Plotly color palette
    )
    return fig


def filter_by_extension_month(_df, _extension):
    """
    Filters the given DataFrame (_df) by the specified extension and creates a line plot using Plotly.

    Parameters:
    _df (DataFrame): The input DataFrame containing the data.
    extension (str): The extension to filter the DataFrame by. If set to "All", no filtering is applied.

    Returns:
    fig (Figure): The Plotly figure object representing the line plot.
    """
    # Filter the DataFrame by the specified extension or extensions
    if len(_extension) == 1 and "All" in _extension or len(_extension) == 0:
        pass
    else:
        _df = _df[_df["extension"].isin(_extension)].copy()

    # Convert year and month into a datetime column and sort by date
    _df["date"] = pd.to_datetime(_df[["year", "month"]].assign(day=1))
    _df = _df.sort_values(by="date")

    # Pivot the DataFrame to get the total size for each extension and make this plotable as a time series
    pivot_df = _df.pivot_table(
        index="date", columns="extension", values="total_size"
    ).fillna(0)

    # Plot!!
    fig = go.Figure()
    for i, column in enumerate(pivot_df.columns):
        if column != "":
            fig.add_trace(
                go.Scatter(
                    x=pivot_df.index,
                    y=pivot_df[column] / 1e12,  # Convert to petabytes
                    mode="lines",
                    name=column,
                    line=dict(color=px.colors.qualitative.Alphabet[i]),
                )
            )

    # Update layout
    fig.update_layout(
        title="Monthly Additions of LFS Files by Extension (in TBs)",
        xaxis_title="Date",
        yaxis_title="Size (TBs)",
        legend_title="Type",
        yaxis=dict(tickformat=".2f"),  # Format y-axis labels to 2 decimal places
    )

    return fig


# Create a gradio blocks interface and launch a demo
with gr.Blocks() as demo:
    df, file_df, by_type, by_extension, by_extension_month = process_dataset()

    # Add a heading
    gr.Markdown("# Git LFS Analysis Across the Hub")
    gr.Markdown(
        "The Hugging Face Hub has just crossed 1,000,000 models - but where is all that data stored? The short answer is Git LFS. This analysis dives into the LFS storage on the Hub, breaking down the data by repository type, file extension, and growth over time."
    )

    gr.Markdown(
        "Now, you might ask yourself, 'Why are you doing this?' Well, the [Xet Team](https://huggingface.co/xet-team) is a [new addition to Hugging Face](https://huggingface.co/blog/xethub-joins-hf), bringing a new way to store massive datasets and models to enable ML teams to operate like software teams: Quickly and without friction. Because this story all starts with storage, that's where we've begun with our own deep dives into what the Hub holds. As part of this, we've included a look at what happens with just one simple deduplication strategy - deduplicating at the file level. Read on to see more!"
    )
    with gr.Row():
        # scale so that
        # group the data by month and year and compute a cumulative sum of the total_size column
        fig, last_10_months = cumulative_growth_plot_analysis(df, file_df)
        with gr.Column(scale=1):
            gr.Markdown("# Repository Growth")
            gr.Markdown(
                "The cumulative growth of models, spaces, and datasets over time can be seen in the adjacent chart. Beside that is a view of the total change, from the previous month to the current one, of LFS files stored on the hub over 2024. We're averaging nearly **2.3 PBs uploaded to LFS per month!**"
            )
            gr.Dataframe(last_10_months, height=250)
        with gr.Column(scale=3):
            gr.Plot(fig)
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown(
                "This table shows the total number of files and cumulative size of those files across all repositories on the Hub. These numbers might be hard to grok, so let's try to put them in context. The last [Common Crawl](https://commoncrawl.org/) download was [451 TBs](https://github.com/commoncrawl/cc-crawl-statistics/blob/master/stats/crawler/CC-MAIN-2024-38.json#L31). The Spaces repositories alone outpaces that. Meanwhile, between Datasets and Model repos, the Hub stores **64 Common Crawls** 🤯."
            )
        with gr.Column(scale=3):
            gr.Dataframe(by_type)

    # Add a heading
    gr.Markdown("## File Extension Analysis")
    gr.Markdown(
        "Breaking this down by file extension, some interesting trends emerge. [Safetensors](https://huggingface.co/docs/safetensors/en/index) are quickly becoming the defacto standard on the hub, accounting for over 7PBs (25%) of LFS storage. The top 20 file extensions seen here and in the table below account for 82% of all LFS storage on the hub."
    )
    # Get the top 10 file extnesions by size
    by_extension_size = by_extension.sort_values(by="size", ascending=False).head(22)
    # get the top 10 file extensions by count
    # by_extension_count = by_extension.sort_values(by="count", ascending=False).head(20)

    # make a pie chart of the by_extension_size dataframe
    gr.Plot(plot_total_sum(by_extension_size[["extension", "size"]].values))
    # drop the unnamed: 0 column
    by_extension_size = by_extension_size.drop(columns=["Unnamed: 0"])
    # average size
    by_extension_size["Average File Size (MBs)"] = (
        by_extension_size["size"].astype(float) / by_extension_size["count"]
    )
    by_extension_size["Average File Size (MBs)"] = (
        by_extension_size["Average File Size (MBs)"] / 1e6
    )
    by_extension_size["Average File Size (MBs)"] = by_extension_size[
        "Average File Size (MBs)"
    ].map("{:.2f}".format)
    # format the size column
    by_extension_size = format_dataframe_size_column(by_extension_size, "size")
    # Rename the other columns
    by_extension_size = by_extension_size.rename(
        columns={
            "extension": "File Extension",
            "count": "Number of Files",
            "size": "Total Size (PBs)",
        }
    )
    gr.Dataframe(by_extension_size)

    gr.Markdown("## File Extension Growth Over Time")
    gr.Markdown(
        "Want to dig a little deeper? Select a file extension to see how many bytes of that type were uploaded to the Hub each month."
    )

    # build a dropdown using the unique values in the extension column
    extension = gr.Dropdown(
        choices=by_extension["extension"].unique().tolist(),
        value="All",
        allow_custom_value=True,
        multiselect=True,
    )
    _by_extension_month = gr.State(by_extension_month)
    gr.Plot(filter_by_extension_month, inputs=[_by_extension_month, extension])

# launch the dang thing
demo.launch()