Update app.py

app.py

@@ -5,6 +5,7 @@ import plotly.express as px
 import os
 import numpy as np
 import io
+import duckdb
 
 # Define pipeline tags
 PIPELINE_TAGS = [
@@ -57,65 +58,120 @@ MODEL_SIZE_RANGES = {
     "XX-Large (>50GB)": (50, float('inf'))
 }
 
-# Filter functions for tags
+# Filter functions for tags - UPDATED to use cached columns
 def is_audio_speech(row):
-    tags = row.get("tags", [])
-    pipeline_tag = row.get("pipeline_tag", "")
-    
-    return (pipeline_tag and ("audio" in pipeline_tag.lower() or "speech" in pipeline_tag.lower())) or \
-           any("audio" in tag.lower() for tag in tags) or \
-           any("speech" in tag.lower() for tag in tags)
+    # Use cached column instead of recalculating
+    return row['is_audio_speech']
 
 def is_music(row):
-    tags = row.get("tags", [])
-    return any("music" in tag.lower() for tag in tags)
+    # Use cached column instead of recalculating
+    return row['has_music']
 
 def is_robotics(row):
-    tags = row.get("tags", [])
-    return any("robot" in tag.lower() for tag in tags)
+    # Use cached column instead of recalculating
+    return row['has_robot']
 
 def is_biomed(row):
-    tags = row.get("tags", [])
-    return any("bio" in tag.lower() for tag in tags) or \
-           any("medic" in tag.lower() for tag in tags)
+    # Use cached column instead of recalculating
+    return row['is_biomed']
 
 def is_timeseries(row):
-    tags = row.get("tags", [])
-    return any("series" in tag.lower() for tag in tags)
+    # Use cached column instead of recalculating
+    return row['has_series']
 
 def is_science(row):
-    tags = row.get("tags", [])
-    return any("science" in tag.lower() and "bigscience" not in tag for tag in tags)
+    # Use cached column instead of recalculating
+    return row['has_science']
 
 def is_video(row):
-    tags = row.get("tags", [])
-    return any("video" in tag.lower() for tag in tags)
+    # Use cached column instead of recalculating
+    return row['has_video']
+
+def is_image(row):
+    # Use cached column instead of recalculating
+    return row['has_image']
+
+def is_text(row):
+    # Use cached column instead of recalculating
+    return row['has_text']
 
 def is_image(row):
     tags = row.get("tags", [])
-    return any("image" in tag.lower() for tag in tags)
+    
+    # Check if tags exists and is not empty
+    if tags is not None:
+        # For numpy arrays
+        if hasattr(tags, 'dtype') and hasattr(tags, 'tolist'):
+            # Convert numpy array to list
+            tags_list = tags.tolist()
+            return any("image" in str(tag).lower() for tag in tags_list)
+        # For regular lists
+        elif isinstance(tags, list):
+            return any("image" in str(tag).lower() for tag in tags)
+        # For string tags
+        elif isinstance(tags, str):
+            return "image" in tags.lower()
+    return False
 
 def is_text(row):
     tags = row.get("tags", [])
-    return any("text" in tag.lower() for tag in tags)
+    
+    # Check if tags exists and is not empty
+    if tags is not None:
+        # For numpy arrays
+        if hasattr(tags, 'dtype') and hasattr(tags, 'tolist'):
+            # Convert numpy array to list
+            tags_list = tags.tolist()
+            return any("text" in str(tag).lower() for tag in tags_list)
+        # For regular lists
+        elif isinstance(tags, list):
+            return any("text" in str(tag).lower() for tag in tags)
+        # For string tags
+        elif isinstance(tags, str):
+            return "text" in tags.lower()
+    return False
+
+def extract_model_size(safetensors_data):
+    """Extract model size in GB from safetensors data"""
+    try:
+        if pd.isna(safetensors_data):
+            return 0
+        
+        # If it's already a dictionary, use it directly
+        if isinstance(safetensors_data, dict):
+            if 'total' in safetensors_data:
+                try:
+                    size_bytes = float(safetensors_data['total'])
+                    return size_bytes / (1024 * 1024 * 1024)  # Convert to GB
+                except (ValueError, TypeError):
+                    pass
+        
+        # If it's a string, try to parse it as JSON
+        elif isinstance(safetensors_data, str):
+            try:
+                data_dict = json.loads(safetensors_data)
+                if 'total' in data_dict:
+                    try:
+                        size_bytes = float(data_dict['total'])
+                        return size_bytes / (1024 * 1024 * 1024)  # Convert to GB
+                    except (ValueError, TypeError):
+                        pass
+            except:
+                pass
+        
+        return 0
+    except Exception as e:
+        print(f"Error extracting model size: {e}")
+        return 0
 
-# Add model size filter function
+# Add model size filter function - UPDATED to use cached size_category column
 def is_in_size_range(row, size_range):
-    if size_range is None:
+    """Check if a model is in the specified size range using pre-calculated size category"""
+    if size_range is None or size_range == "None":
         return True
     
-    min_size, max_size = MODEL_SIZE_RANGES[size_range]
-    
-    # Get model size in GB from params column
-    if "params" in row and pd.notna(row["params"]):
-        try:
-            # Convert to GB (assuming params are in bytes or scientific notation)
-            size_gb = float(row["params"]) / (1024 * 1024 * 1024)
-            return min_size <= size_gb < max_size
-        except (ValueError, TypeError):
-            return False
-    
-    return False
+    # Simply compare with cached size_category
+    return row['size_category'] == size_range
 
 TAG_FILTER_FUNCS = {
     "Audio & Speech": is_audio_speech,
@@ -136,24 +192,64 @@ def extract_org_from_id(model_id):
     return "unaffiliated"
 
 def make_treemap_data(df, count_by, top_k=25, tag_filter=None, pipeline_filter=None, size_filter=None, skip_orgs=None):
-    """Process DataFrame into treemap format with filters applied"""
+    """Process DataFrame into treemap format with filters applied - OPTIMIZED with cached columns"""
     # Create a copy to avoid modifying the original
     filtered_df = df.copy()
     
     # Apply filters
+    filter_stats = {"initial": len(filtered_df)}
+    start_time = pd.Timestamp.now()
+    
+    # Apply tag filter - OPTIMIZED to use cached columns
     if tag_filter and tag_filter in TAG_FILTER_FUNCS:
-        filter_func = TAG_FILTER_FUNCS[tag_filter]
-        filtered_df = filtered_df[filtered_df.apply(filter_func, axis=1)]
+        print(f"Applying tag filter: {tag_filter}")
+        
+        # Use direct column filtering instead of applying a function to each row
+        if tag_filter == "Audio & Speech":
+            filtered_df = filtered_df[filtered_df['is_audio_speech']]
+        elif tag_filter == "Music":
+            filtered_df = filtered_df[filtered_df['has_music']]
+        elif tag_filter == "Robotics":
+            filtered_df = filtered_df[filtered_df['has_robot']]
+        elif tag_filter == "Biomedical":
+            filtered_df = filtered_df[filtered_df['is_biomed']]
+        elif tag_filter == "Time series":
+            filtered_df = filtered_df[filtered_df['has_series']]
+        elif tag_filter == "Sciences":
+            filtered_df = filtered_df[filtered_df['has_science']]
+        elif tag_filter == "Video":
+            filtered_df = filtered_df[filtered_df['has_video']]
+        elif tag_filter == "Images":
+            filtered_df = filtered_df[filtered_df['has_image']]
+        elif tag_filter == "Text":
+            filtered_df = filtered_df[filtered_df['has_text']]
+        
+        filter_stats["after_tag_filter"] = len(filtered_df)
+        print(f"Tag filter applied in {(pd.Timestamp.now() - start_time).total_seconds():.3f} seconds")
+        start_time = pd.Timestamp.now()
     
+    # Apply pipeline filter
     if pipeline_filter:
+        print(f"Applying pipeline filter: {pipeline_filter}")
         filtered_df = filtered_df[filtered_df["pipeline_tag"] == pipeline_filter]
+        filter_stats["after_pipeline_filter"] = len(filtered_df)
+        print(f"Pipeline filter applied in {(pd.Timestamp.now() - start_time).total_seconds():.3f} seconds")
+        start_time = pd.Timestamp.now()
     
+    # Apply size filter - OPTIMIZED to use cached size_category column
     if size_filter and size_filter in MODEL_SIZE_RANGES:
-        # Create a function to check if a model is in the size range
-        def check_size(row):
-            return is_in_size_range(row, size_filter)
+        print(f"Applying size filter: {size_filter}")
+        
+        # Use the cached size_category column directly
+        filtered_df = filtered_df[filtered_df['size_category'] == size_filter]
+        
+        # Debug info
+        print(f"Size filter '{size_filter}' applied.")
+        print(f"Models after size filter: {len(filtered_df)}")
         
-        filtered_df = filtered_df[filtered_df.apply(check_size, axis=1)]
+        filter_stats["after_size_filter"] = len(filtered_df)
+        print(f"Size filter applied in {(pd.Timestamp.now() - start_time).total_seconds():.3f} seconds")
+        start_time = pd.Timestamp.now()
     
     # Add organization column
     filtered_df["organization"] = filtered_df["id"].apply(extract_org_from_id)
@@ -161,6 +257,17 @@ def make_treemap_data(df, count_by, top_k=25, tag_filter=None, pipeline_filter=N
     # Skip organizations if specified
     if skip_orgs and len(skip_orgs) > 0:
         filtered_df = filtered_df[~filtered_df["organization"].isin(skip_orgs)]
+        filter_stats["after_skip_orgs"] = len(filtered_df)
+    
+    # Print filter stats
+    print("Filter statistics:")
+    for stage, count in filter_stats.items():
+        print(f"  {stage}: {count} models")
+    
+    # Check if we have any data left
+    if filtered_df.empty:
+        print("Warning: No data left after applying filters!")
+        return pd.DataFrame()  # Return empty DataFrame
     
     # Aggregate by organization
     org_totals = filtered_df.groupby("organization")[count_by].sum().reset_index()
@@ -181,6 +288,7 @@ def make_treemap_data(df, count_by, top_k=25, tag_filter=None, pipeline_filter=N
     # Ensure numeric values
     treemap_data[count_by] = pd.to_numeric(treemap_data[count_by], errors="coerce").fillna(0)
     
+    print(f"Treemap data prepared in {(pd.Timestamp.now() - start_time).total_seconds():.3f} seconds")
     return treemap_data
 
 def create_treemap(treemap_data, count_by, title=None):
@@ -219,23 +327,271 @@ def create_treemap(treemap_data, count_by, title=None):
     
     return fig
 
-def load_models_csv():    
-    # Read the CSV file
-    df = pd.read_csv('models.csv')
-    
-    # Process the tags column 
-    def process_tags(tags_str):
-        if pd.isna(tags_str):
-            return []
+def load_models_data():
+    """Load models data from Hugging Face using DuckDB with caching for improved performance"""
+    try:
+        # The URL to the parquet file
+        parquet_url = "https://huggingface.co/datasets/cfahlgren1/hub-stats/resolve/main/models.parquet"
         
-        # Clean the string and convert to a list
-        tags_str = tags_str.strip("[]").replace("'", "")
-        tags = [tag.strip() for tag in tags_str.split() if tag.strip()]
-        return tags
-    
-    df['tags'] = df['tags'].apply(process_tags)
-    
-    return df
+        print("Fetching data from Hugging Face models.parquet...")
+        
+        # Based on the column names provided, we can directly select the columns we need
+        # Note: We need to select safetensors to get the model size information
+        try:
+            query = """
+            SELECT 
+                id, 
+                downloads, 
+                downloadsAllTime, 
+                likes, 
+                pipeline_tag, 
+                tags,
+                safetensors
+            FROM read_parquet('https://huggingface.co/datasets/cfahlgren1/hub-stats/resolve/main/models.parquet')
+            """
+            df = duckdb.sql(query).df()
+        except Exception as sql_error:
+            print(f"Error with specific column selection: {sql_error}")
+            # Fallback to just selecting everything and then filtering
+            print("Falling back to select * query...")
+            query = "SELECT * FROM read_parquet('https://huggingface.co/datasets/cfahlgren1/hub-stats/resolve/main/models.parquet')"
+            raw_df = duckdb.sql(query).df()
+            
+            # Now extract only the columns we need
+            needed_columns = ['id', 'downloads', 'downloadsAllTime', 'likes', 'pipeline_tag', 'tags', 'safetensors']
+            available_columns = set(raw_df.columns)
+            df = pd.DataFrame()
+            
+            # Copy over columns that exist
+            for col in needed_columns:
+                if col in available_columns:
+                    df[col] = raw_df[col]
+                else:
+                    # Create empty columns for missing data
+                    if col in ['downloads', 'downloadsAllTime', 'likes']:
+                        df[col] = 0
+                    elif col == 'pipeline_tag':
+                        df[col] = ''
+                    elif col == 'tags':
+                        df[col] = [[] for _ in range(len(raw_df))]
+                    elif col == 'safetensors':
+                        df[col] = None
+                    elif col == 'id':
+                        # Create IDs based on index if missing
+                        df[col] = [f"model_{i}" for i in range(len(raw_df))]
+        
+        print(f"Data fetched successfully. Shape: {df.shape}")
+        
+        # Check if safetensors column exists before trying to process it
+        if 'safetensors' in df.columns:
+            # Add params column derived from safetensors.total (model size in GB)
+            df['params'] = df['safetensors'].apply(extract_model_size)
+            
+            # Debug model sizes
+            size_ranges = {
+                "Small (<1GB)": 0,
+                "Medium (1-5GB)": 0,
+                "Large (5-20GB)": 0,
+                "X-Large (20-50GB)": 0,
+                "XX-Large (>50GB)": 0
+            }
+            
+            # Count models in each size range
+            for idx, row in df.iterrows():
+                size_gb = row['params']
+                if 0 <= size_gb < 1:
+                    size_ranges["Small (<1GB)"] += 1
+                elif 1 <= size_gb < 5:
+                    size_ranges["Medium (1-5GB)"] += 1
+                elif 5 <= size_gb < 20:
+                    size_ranges["Large (5-20GB)"] += 1
+                elif 20 <= size_gb < 50:
+                    size_ranges["X-Large (20-50GB)"] += 1
+                elif size_gb >= 50:
+                    size_ranges["XX-Large (>50GB)"] += 1
+            
+            print("Model size distribution:")
+            for size_range, count in size_ranges.items():
+                print(f"  {size_range}: {count} models")
+            
+            # CACHE SIZE CATEGORY: Add a size_category column for faster filtering
+            def get_size_category(size_gb):
+                if 0 <= size_gb < 1:
+                    return "Small (<1GB)"
+                elif 1 <= size_gb < 5:
+                    return "Medium (1-5GB)"
+                elif 5 <= size_gb < 20:
+                    return "Large (5-20GB)"
+                elif 20 <= size_gb < 50:
+                    return "X-Large (20-50GB)"
+                elif size_gb >= 50:
+                    return "XX-Large (>50GB)"
+                return None
+                
+            # Add cached size category column
+            df['size_category'] = df['params'].apply(get_size_category)
+            
+            # Remove the safetensors column as we don't need it anymore
+            df = df.drop(columns=['safetensors'])
+        else:
+            # If no safetensors column, add empty params column
+            df['params'] = 0
+            df['size_category'] = None
+        
+        # Process tags to ensure it's in the right format - FIXED
+        def process_tags(tags_value):
+            try:
+                if pd.isna(tags_value) or tags_value is None:
+                    return []
+                
+                # If it's a numpy array, convert to a list of strings
+                if hasattr(tags_value, 'dtype') and hasattr(tags_value, 'tolist'):
+                    # Note: This is the fix for the error
+                    return [str(tag) for tag in tags_value.tolist()]
+                
+                # If already a list, ensure all elements are strings
+                if isinstance(tags_value, list):
+                    return [str(tag) for tag in tags_value]
+                
+                # If string, try to parse as JSON or split by comma
+                if isinstance(tags_value, str):
+                    try:
+                        tags_list = json.loads(tags_value)
+                        if isinstance(tags_list, list):
+                            return [str(tag) for tag in tags_list]
+                    except:
+                        # Split by comma if JSON parsing fails
+                        return [tag.strip() for tag in tags_value.split(',') if tag.strip()]
+                
+                # Last resort, convert to string and return as a single tag
+                return [str(tags_value)]
+                
+            except Exception as e:
+                print(f"Error processing tags: {e}")
+                return []
+        
+        # Check if tags column exists before trying to process it
+        if 'tags' in df.columns:
+            # Process tags column
+            df['tags'] = df['tags'].apply(process_tags)
+            
+            # CACHE TAG CATEGORIES: Pre-calculate tag categories for faster filtering
+            print("Pre-calculating cached tag categories...")
+            
+            # Helper functions to check for specific tags (simplified for caching)
+            def has_audio_tag(tags):
+                if tags and isinstance(tags, list):
+                    return any("audio" in str(tag).lower() for tag in tags)
+                return False
+                
+            def has_speech_tag(tags):
+                if tags and isinstance(tags, list):
+                    return any("speech" in str(tag).lower() for tag in tags)
+                return False
+                
+            def has_music_tag(tags):
+                if tags and isinstance(tags, list):
+                    return any("music" in str(tag).lower() for tag in tags)
+                return False
+                
+            def has_robot_tag(tags):
+                if tags and isinstance(tags, list):
+                    return any("robot" in str(tag).lower() for tag in tags)
+                return False
+                
+            def has_bio_tag(tags):
+                if tags and isinstance(tags, list):
+                    return any("bio" in str(tag).lower() for tag in tags)
+                return False
+                
+            def has_med_tag(tags):
+                if tags and isinstance(tags, list):
+                    return any("medic" in str(tag).lower() for tag in tags)
+                return False
+                
+            def has_series_tag(tags):
+                if tags and isinstance(tags, list):
+                    return any("series" in str(tag).lower() for tag in tags)
+                return False
+                
+            def has_science_tag(tags):
+                if tags and isinstance(tags, list):
+                    return any("science" in str(tag).lower() and "bigscience" not in str(tag).lower() for tag in tags)
+                return False
+                
+            def has_video_tag(tags):
+                if tags and isinstance(tags, list):
+                    return any("video" in str(tag).lower() for tag in tags)
+                return False
+                
+            def has_image_tag(tags):
+                if tags and isinstance(tags, list):
+                    return any("image" in str(tag).lower() for tag in tags)
+                return False
+                
+            def has_text_tag(tags):
+                if tags and isinstance(tags, list):
+                    return any("text" in str(tag).lower() for tag in tags)
+                return False
+            
+            # Add cached columns for tag categories
+            print("Creating cached tag columns...")
+            df['has_audio'] = df['tags'].apply(has_audio_tag)
+            df['has_speech'] = df['tags'].apply(has_speech_tag)
+            df['has_music'] = df['tags'].apply(has_music_tag)
+            df['has_robot'] = df['tags'].apply(has_robot_tag)
+            df['has_bio'] = df['tags'].apply(has_bio_tag)
+            df['has_med'] = df['tags'].apply(has_med_tag)
+            df['has_series'] = df['tags'].apply(has_series_tag)
+            df['has_science'] = df['tags'].apply(has_science_tag)
+            df['has_video'] = df['tags'].apply(has_video_tag)
+            df['has_image'] = df['tags'].apply(has_image_tag)
+            df['has_text'] = df['tags'].apply(has_text_tag)
+            
+            # Create combined category flags for faster filtering
+            df['is_audio_speech'] = (df['has_audio'] | df['has_speech'] | 
+                                    df['pipeline_tag'].str.contains('audio', case=False, na=False) | 
+                                    df['pipeline_tag'].str.contains('speech', case=False, na=False))
+            df['is_biomed'] = df['has_bio'] | df['has_med']
+            
+            print("Cached tag columns created successfully!")
+        else:
+            # If no tags column, add empty tags and set all category flags to False
+            df['tags'] = [[] for _ in range(len(df))]
+            for col in ['has_audio', 'has_speech', 'has_music', 'has_robot',
+                        'has_bio', 'has_med', 'has_series', 'has_science',
+                        'has_video', 'has_image', 'has_text',
+                        'is_audio_speech', 'is_biomed']:
+                df[col] = False
+        
+        # Fill NaN values
+        df.fillna({'downloads': 0, 'downloadsAllTime': 0, 'likes': 0, 'params': 0}, inplace=True)
+        
+        # Ensure pipeline_tag is a string
+        if 'pipeline_tag' in df.columns:
+            df['pipeline_tag'] = df['pipeline_tag'].fillna('')
+        else:
+            df['pipeline_tag'] = ''
+        
+        # Make sure all required columns exist
+        for col in ['id', 'downloads', 'downloadsAllTime', 'likes', 'pipeline_tag', 'tags', 'params']:
+            if col not in df.columns:
+                if col in ['downloads', 'downloadsAllTime', 'likes', 'params']:
+                    df[col] = 0
+                elif col == 'pipeline_tag':
+                    df[col] = ''
+                elif col == 'tags':
+                    df[col] = [[] for _ in range(len(df))]
+                elif col == 'id':
+                    df[col] = [f"model_{i}" for i in range(len(df))]
+        
+        print(f"Successfully processed {len(df)} models with cached tag and size information")
+        return df, True
+        
+    except Exception as e:
+        print(f"Error loading data: {e}")
+        # Return an empty DataFrame and False to indicate loading failure
+        return pd.DataFrame(), False
 
 # Create Gradio interface
 with gr.Blocks() as demo:
@@ -250,6 +606,7 @@ with gr.Blocks() as demo:
             Use the filters to explore models by different metrics, tags, pipelines, and model sizes.
 
             The treemap visualizes models grouped by organization, with the size of each box representing the selected metric.
+            
         """)
 
     with gr.Row():
@@ -311,24 +668,12 @@ with gr.Blocks() as demo:
             )
 
             generate_plot_button = gr.Button("Generate Plot", variant="primary", interactive=False)
+            refresh_data_button = gr.Button("Refresh Data from Hugging Face", variant="secondary")
 
         with gr.Column(scale=3):
             plot_output = gr.Plot()
-            stats_output = gr.Markdown("*Generate a plot to see statistics*")
-
-    # Updated load function returning both the data and loading flag
-    def load_models_csv():    
-        df = pd.read_csv('models.csv')
-
-        def process_tags(tags_str):
-            if pd.isna(tags_str):
-                return []
-            tags_str = tags_str.strip("[]").replace("'", "")
-            tags = [tag.strip() for tag in tags_str.split() if tag.strip()]
-            return tags
-
-        df['tags'] = df['tags'].apply(process_tags)
-        return df, True
+            stats_output = gr.Markdown("*Loading data from Hugging Face...*")
+            data_info = gr.Markdown("")
 
     # Button enablement after data load
     def enable_plot_button(loaded):
@@ -355,6 +700,25 @@ with gr.Blocks() as demo:
         outputs=[tag_filter_dropdown, pipeline_filter_dropdown]
     )
 
+    # Function to handle data load and provide data info
+    def load_and_provide_info():
+        df, success = load_models_data()
+        
+        if success:
+            # Generate information about the loaded data
+            info_text = f"""
+### Data Information
+- **Total models loaded**: {len(df):,}
+- **Last update**: {pd.Timestamp.now().strftime('%Y-%m-%d %H:%M:%S')}
+- **Data source**: [Hugging Face Hub Stats](https://huggingface.co/datasets/cfahlgren1/hub-stats) (models.parquet)
+            """
+            
+            # Return the data, loading status, and info text
+            return df, True, info_text, "*Data loaded successfully. Use the controls to generate a plot.*"
+        else:
+            # Return empty data, failed loading status, and error message
+            return pd.DataFrame(), False, "*Error loading data from Hugging Face.*", "*Failed to load data. Please try again.*"
+
     # Main generate function
     def generate_plot_on_click(count_by, filter_choice, tag_filter, pipeline_filter, size_filter, top_k, skip_orgs_text, data_df):
         if data_df is None or not isinstance(data_df, pd.DataFrame) or data_df.empty:
@@ -451,9 +815,16 @@ with gr.Blocks() as demo:
 
     # Load data at startup
     demo.load(
-        fn=load_models_csv,
+        fn=load_and_provide_info,
         inputs=[], 
-        outputs=[models_data, loading_complete]
+        outputs=[models_data, loading_complete, data_info, stats_output]
+    )
+
+    # Refresh data when button is clicked
+    refresh_data_button.click(
+        fn=load_and_provide_info,
+        inputs=[],
+        outputs=[models_data, loading_complete, data_info, stats_output]
     )
 
     generate_plot_button.click(
@@ -471,7 +842,5 @@ with gr.Blocks() as demo:
         outputs=[plot_output, stats_output]
     )
 
-
-
 if __name__ == "__main__":
     demo.launch()