import json import gradio as gr import pandas as pd import plotly.express as px import pyarrow.parquet as pq import os import requests from io import BytesIO import math # Define pipeline tags (keeping the same ones from the provided code) 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', ] # Model size categories in GB MODEL_SIZE_RANGES = { "Small (<1GB)": (0, 1), "Medium (1-5GB)": (1, 5), "Large (5-20GB)": (5, 20), "X-Large (20-50GB)": (20, 50), "XX-Large (>50GB)": (50, float('inf')) } # Filter functions for tags - keeping the same from provided code def is_audio_speech(repo_dct): res = (repo_dct.get("pipeline_tag", None) and "audio" in repo_dct.get("pipeline_tag", "").lower()) or \ (repo_dct.get("pipeline_tag", None) and "speech" in repo_dct.get("pipeline_tag", "").lower()) or \ (repo_dct.get("tags", None) and any("audio" in tag.lower() for tag in repo_dct.get("tags", []))) or \ (repo_dct.get("tags", None) and any("speech" in tag.lower() for tag in repo_dct.get("tags", []))) return res def is_music(repo_dct): res = (repo_dct.get("tags", None) and any("music" in tag.lower() for tag in repo_dct.get("tags", []))) return res def is_robotics(repo_dct): res = (repo_dct.get("tags", None) and any("robot" in tag.lower() for tag in repo_dct.get("tags", []))) return res def is_biomed(repo_dct): res = (repo_dct.get("tags", None) and any("bio" in tag.lower() for tag in repo_dct.get("tags", []))) or \ (repo_dct.get("tags", None) and any("medic" in tag.lower() for tag in repo_dct.get("tags", []))) return res def is_timeseries(repo_dct): res = (repo_dct.get("tags", None) and any("series" in tag.lower() for tag in repo_dct.get("tags", []))) return res def is_science(repo_dct): res = (repo_dct.get("tags", None) and any("science" in tag.lower() and not "bigscience" in tag for tag in repo_dct.get("tags", []))) return res def is_video(repo_dct): res = (repo_dct.get("tags", None) and any("video" in tag.lower() for tag in repo_dct.get("tags", []))) return res def is_image(repo_dct): res = (repo_dct.get("tags", None) and any("image" in tag.lower() for tag in repo_dct.get("tags", []))) return res def is_text(repo_dct): res = (repo_dct.get("tags", None) and any("text" in tag.lower() for tag in repo_dct.get("tags", []))) return res # Add model size filter function def is_in_size_range(repo_dct, size_range): if size_range is None: return True min_size, max_size = MODEL_SIZE_RANGES[size_range] # Get model size in GB from safetensors total (if available) if repo_dct.get("safetensors") and repo_dct["safetensors"].get("total"): # Convert bytes to GB size_gb = repo_dct["safetensors"]["total"] / (1024 * 1024 * 1024) return min_size <= size_gb < max_size return False TAG_FILTER_FUNCS = { "Audio & Speech": is_audio_speech, "Time series": is_timeseries, "Robotics": is_robotics, "Music": is_music, "Video": is_video, "Images": is_image, "Text": is_text, "Biomedical": is_biomed, "Sciences": is_science, } def make_org_stats(count_by, org_stats, top_k=20, filter_func=None, size_range=None): assert count_by in ["likes", "downloads"] # Apply both filter_func and size_range if provided def combined_filter(dct): passes_tag_filter = filter_func(dct) if filter_func else True passes_size_filter = is_in_size_range(dct, size_range) if size_range else True return passes_tag_filter and passes_size_filter # Sort organizations by total count sorted_stats = sorted( [( org_id, sum(model[count_by] for model in models if combined_filter(model)) ) for org_id, models in org_stats.items()], key=lambda x: x[1], reverse=True, ) # Top organizations + Others category res = sorted_stats[:top_k] + [("Others...", sum(st for auth, st in sorted_stats[top_k:]))] total_st = sum(st for o, st in res) # Prepare data for treemap res_plot_df = [] for org, st in res: if org == "Others...": res_plot_df += [("Others...", "other", st * 100 / total_st if total_st > 0 else 0)] else: for model in org_stats[org]: if combined_filter(model): res_plot_df += [(org, model["id"], model[count_by] * 100 / total_st if total_st > 0 else 0)] return ([(o, 100 * st / total_st if total_st > 0 else 0) for o, st in res if st > 0], res_plot_df) def make_figure(count_by, org_stats, tag_filter=None, pipeline_filter=None, size_range=None): assert count_by in ["downloads", "likes"] # Determine which filter function to use filter_func = None if tag_filter: filter_func = TAG_FILTER_FUNCS[tag_filter] elif pipeline_filter: filter_func = lambda dct: dct.get("pipeline_tag", None) and dct.get("pipeline_tag", "") == pipeline_filter else: filter_func = lambda dct: True # Generate stats with filters _, res_plot_df = make_org_stats(count_by, org_stats, top_k=25, filter_func=filter_func, size_range=size_range) # Create DataFrame for Plotly df = pd.DataFrame( dict( organizations=[o for o, _, _ in res_plot_df], model=[r for _, r, _ in res_plot_df], stats=[s for _, _, s in res_plot_df], ) ) df["models"] = "models" # Root node # Create treemap fig = px.treemap(df, path=["models", 'organizations', 'model'], values='stats', title=f"HuggingFace Models - {count_by.capitalize()} by Organization") fig.update_layout( margin=dict(t=50, l=25, r=25, b=25) ) return fig def download_and_process_models(): """Download and process the models data from HuggingFace dataset""" try: # Create a cache directory if not os.path.exists('data'): os.makedirs('data') # Check if we have cached data if os.path.exists('data/processed_models.json'): print("Loading from cache...") with open('data/processed_models.json', 'r') as f: return json.load(f) # URL to the models.parquet file url = "https://huggingface.co/datasets/cfahlgren1/hub-stats/resolve/main/models.parquet" print(f"Downloading models data from {url}...") response = requests.get(url) if response.status_code != 200: raise Exception(f"Failed to download data: HTTP {response.status_code}") # Read the parquet file table = pq.read_table(BytesIO(response.content)) df = table.to_pandas() print(f"Downloaded {len(df)} models") # Process the dataframe into the organization structure we need org_stats = {} for _, row in df.iterrows(): model_id = row['id'] # Extract the organization part of the model ID if '/' in model_id: org_id = model_id.split('/')[0] else: org_id = "unaffiliated" # Create model entry with needed fields model_entry = { "id": model_id, "downloads": row.get('downloads', 0), "likes": row.get('likes', 0), "pipeline_tag": row.get('pipeline_tag'), "tags": row.get('tags', []), } # Add safetensors information if available if 'safetensors' in row and row['safetensors']: if isinstance(row['safetensors'], dict) and 'total' in row['safetensors']: model_entry["safetensors"] = {"total": row['safetensors']['total']} elif isinstance(row['safetensors'], str): # Try to parse JSON string try: safetensors = json.loads(row['safetensors']) if isinstance(safetensors, dict) and 'total' in safetensors: model_entry["safetensors"] = {"total": safetensors['total']} except: pass # Add to organization stats if org_id not in org_stats: org_stats[org_id] = [] org_stats[org_id].append(model_entry) # Cache the processed data with open('data/processed_models.json', 'w') as f: json.dump(org_stats, f) return org_stats except Exception as e: print(f"Error downloading or processing data: {e}") # Return sample data for testing if real data unavailable return create_sample_data() def create_sample_data(): """Create sample data for testing when real data is unavailable""" print("Creating sample data for testing...") sample_orgs = ['openai', 'meta', 'google', 'microsoft', 'anthropic', 'stability', 'huggingface'] org_stats = {} for org in sample_orgs: org_stats[org] = [] num_models = 5 # Each org has 5 sample models for i in range(num_models): model_id = f"{org}/model-{i+1}" # Random pipeline tag pipeline_idx = i % len(PIPELINE_TAGS) pipeline_tag = PIPELINE_TAGS[pipeline_idx] # Random tags tags = [pipeline_tag, "sample-data"] # Random downloads and likes downloads = int(1000 * (10 ** (org_stats.keys().index(org) % 3))) # Different magnitudes likes = int(downloads * 0.05) # 5% like rate # Random model size in bytes (from 100MB to 100GB) model_size = (10**8) * (10 ** (i % 3)) # Different magnitudes org_stats[org].append({ "id": model_id, "downloads": downloads, "likes": likes, "pipeline_tag": pipeline_tag, "tags": tags, "safetensors": {"total": model_size} }) return org_stats # Create Gradio interface with gr.Blocks() as demo: models_data = gr.State(value=None) # To store loaded data with gr.Row(): gr.Markdown(""" ## HuggingFace Models TreeMap This app shows how different organizations contribute to the HuggingFace ecosystem with their models. Use the filters to explore models by different metrics, tags, pipelines, and model sizes. """) with gr.Row(): with gr.Column(scale=1): count_by_dropdown = gr.Dropdown( label="Metric", choices=["downloads", "likes"], value="downloads" ) filter_choice_radio = gr.Radio( label="Filter by", choices=["None", "Tag Filter", "Pipeline Filter"], value="None" ) tag_filter_dropdown = gr.Dropdown( label="Select Tag", choices=list(TAG_FILTER_FUNCS.keys()), value=None, visible=False ) pipeline_filter_dropdown = gr.Dropdown( label="Select Pipeline Tag", choices=PIPELINE_TAGS, value=None, visible=False ) size_filter_dropdown = gr.Dropdown( label="Model Size Filter", choices=["None"] + list(MODEL_SIZE_RANGES.keys()), value="None" ) generate_plot_button = gr.Button("Generate Plot") with gr.Column(scale=3): plot_output = gr.Plot() def generate_plot_on_click(count_by, filter_choice, tag_filter, pipeline_filter, size_filter, data): print(f"Generating plot with: Metric={count_by}, Filter={filter_choice}, Tag={tag_filter}, Pipeline={pipeline_filter}, Size={size_filter}") if data is None: print("Error: Data not loaded yet.") return None selected_tag_filter = None selected_pipeline_filter = None selected_size_filter = None if filter_choice == "Tag Filter": selected_tag_filter = tag_filter elif filter_choice == "Pipeline Filter": selected_pipeline_filter = pipeline_filter if size_filter != "None": selected_size_filter = size_filter fig = make_figure( count_by=count_by, org_stats=data, tag_filter=selected_tag_filter, pipeline_filter=selected_pipeline_filter, size_range=selected_size_filter ) return fig def update_filter_visibility(filter_choice): if filter_choice == "Tag Filter": return gr.update(visible=True), gr.update(visible=False) elif filter_choice == "Pipeline Filter": return gr.update(visible=False), gr.update(visible=True) else: # "None" return gr.update(visible=False), gr.update(visible=False) filter_choice_radio.change( fn=update_filter_visibility, inputs=[filter_choice_radio], outputs=[tag_filter_dropdown, pipeline_filter_dropdown] ) # Load data once at startup demo.load( fn=download_and_process_models, inputs=[], outputs=[models_data] ) # Button click event to generate plot generate_plot_button.click( fn=generate_plot_on_click, inputs=[ count_by_dropdown, filter_choice_radio, tag_filter_dropdown, pipeline_filter_dropdown, size_filter_dropdown, models_data ], outputs=[plot_output] ) if __name__ == "__main__": demo.launch()