Update app.py

app.py

@@ -1,4 +1,4 @@
-# --- START OF FILE app.py ---
+# --- START OF MODIFIED FILE app.py ---
 
 import gradio as gr
 import pandas as pd
@@ -7,10 +7,10 @@ import time
 from datasets import load_dataset # Import the datasets library
 
 # --- Constants ---
-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'))
-}
+# REMOVED the old MODEL_SIZE_RANGES dictionary.
+# NEW: Define the discrete steps for the parameter range slider.
+PARAM_CHOICES = ['< 1B', '1B', '5B', '12B', '32B', '64B', '128B', '256B', '> 500B']
+PARAM_CHOICES_DEFAULT = [PARAM_CHOICES[0], PARAM_CHOICES[-1]]
 
 # The Hugging Face dataset ID to load.
 HF_DATASET_ID = "evijit/orgstats_daily_data"
@@ -42,44 +42,40 @@ def load_models_data():
     overall_start_time = time.time()
     print(f"Attempting to load dataset from Hugging Face Hub: {HF_DATASET_ID}")
 
-    # These are the columns expected to be in the pre-processed dataset.
     expected_cols = [
         'id', 'downloads', 'downloadsAllTime', 'likes', 'pipeline_tag', 'tags', 'params',
-        'size_category', 'organization', 'has_audio', 'has_speech', 'has_music',
+        'organization', 'has_audio', 'has_speech', 'has_music',
         'has_robot', 'has_bio', 'has_med', 'has_series', 'has_video', 'has_image',
         'has_text', 'has_science', 'is_audio_speech', 'is_biomed',
         'data_download_timestamp'
     ]
 
     try:
-        # Load the dataset using the datasets library
-        # It will be cached locally after the first run.
         dataset_dict = load_dataset(HF_DATASET_ID)
         
         if not dataset_dict:
              raise ValueError(f"Dataset '{HF_DATASET_ID}' loaded but appears empty.")
 
-        # Get the name of the first split (e.g., 'train')
         split_name = list(dataset_dict.keys())[0]
         print(f"Using dataset split: '{split_name}'. Converting to Pandas.")
-
-        # Convert the dataset object to a Pandas DataFrame
         df = dataset_dict[split_name].to_pandas()
-
         elapsed = time.time() - overall_start_time
         
-        # Validate that the loaded data has the columns we expect.
         missing_cols = [col for col in expected_cols if col not in df.columns]
         if missing_cols:
-            raise ValueError(f"Loaded dataset is missing expected columns: {missing_cols}.")
-
-        # --- Diagnostic for 'has_robot' after loading ---
-        if 'has_robot' in df.columns:
-            robot_count = df['has_robot'].sum()
-            print(f"DIAGNOSTIC (Dataset Load): 'has_robot' column found. Number of True values: {robot_count}")
+            # The 'params' column is crucial for the new slider.
+            if 'params' in missing_cols:
+                raise ValueError(f"FATAL: Loaded dataset is missing the crucial 'params' column.")
+            print(f"Warning: Loaded dataset is missing some expected columns: {missing_cols}.")
+
+        # Ensure 'params' column is numeric, coercing errors to NaN and then filling with 0.
+        # This is important for filtering. Assumes params are in billions.
+        if 'params' in df.columns:
+            df['params'] = pd.to_numeric(df['params'], errors='coerce').fillna(0)
         else:
-            print("DIAGNOSTIC (Dataset Load): 'has_robot' column NOT FOUND.")
-        # --- End Diagnostic ---
+            # If 'params' is missing after all, create a dummy column to prevent crashes.
+            df['params'] = 0
+            print("CRITICAL WARNING: 'params' column not found in data. Parameter filtering will not work.")
 
         msg = f"Successfully loaded dataset '{HF_DATASET_ID}' (split: {split_name}) from HF Hub in {elapsed:.2f}s. Shape: {df.shape}"
         print(msg)
@@ -90,8 +86,24 @@ def load_models_data():
         print(err_msg)
         return pd.DataFrame(), False, err_msg
 
+# --- NEW: Helper function to parse slider labels into numerical values ---
+def get_param_range_values(param_range_labels):
+    """Converts a list of two string labels from the slider into a numerical min/max tuple."""
+    if not param_range_labels or len(param_range_labels) != 2:
+        return None, None
+        
+    min_label, max_label = param_range_labels
+    
+    # Min value logic: '< 1B' becomes 0, otherwise parse the number.
+    min_val = 0.0 if '<' in min_label else float(min_label.replace('B', ''))
+    
+    # Max value logic: '> 500B' becomes infinity, otherwise parse the number.
+    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, size_filter=None, skip_orgs=None):
+# --- MODIFIED: Function signature and filtering logic updated for parameter range ---
+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",
@@ -107,16 +119,26 @@ def make_treemap_data(df, count_by, top_k=25, tag_filter=None, pipeline_filter=N
 
     if pipeline_filter:
         if "pipeline_tag" in filtered_df.columns:
-            # Ensure the comparison works even if pipeline_tag has NaNs or mixed types
             filtered_df = filtered_df[filtered_df["pipeline_tag"].astype(str) == pipeline_filter]
         else:
             print(f"Warning: 'pipeline_tag' column not found for filtering.")
 
-    if size_filter and size_filter != "None" and size_filter in MODEL_SIZE_RANGES.keys():
-        if 'size_category' in filtered_df.columns:
-            filtered_df = filtered_df[filtered_df['size_category'] == size_filter]
-        else:
-            print("Warning: 'size_category' column not found for filtering.")
+    # --- MODIFIED: Filtering logic now uses the numerical parameter range ---
+    if param_range:
+        min_params, max_params = get_param_range_values(param_range)
+        is_default_range = (param_range == PARAM_CHOICES_DEFAULT)
+
+        # Only filter if the range is not the default full range
+        if not is_default_range and 'params' in filtered_df.columns:
+            # The 'params' column is in billions, so the values match our slider
+            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'):
+                # The upper bound is exclusive, e.g., 5B to 64B is [5, 64)
+                filtered_df = filtered_df[filtered_df['params'] < max_params]
+        elif 'params' not in filtered_df.columns:
+             print("Warning: 'params' column not found for filtering.")
+
 
     if skip_orgs and len(skip_orgs) > 0:
         if "organization" in filtered_df.columns:
@@ -126,20 +148,16 @@ def make_treemap_data(df, count_by, top_k=25, tag_filter=None, pipeline_filter=N
 
     if filtered_df.empty: return pd.DataFrame()
 
-    # Ensure the metric column is numeric and handle potential missing values
     if count_by not in filtered_df.columns:
          print(f"Warning: Metric column '{count_by}' not found. Using 0.")
          filtered_df[count_by] = 0.0
     filtered_df[count_by] = pd.to_numeric(filtered_df[count_by], errors="coerce").fillna(0.0)
     
-    # Group and get top organizations
     org_totals = filtered_df.groupby("organization")[count_by].sum().nlargest(top_k, keep='first')
     top_orgs_list = org_totals.index.tolist()
     
-    # Prepare data for treemap
     treemap_data = filtered_df[filtered_df["organization"].isin(top_orgs_list)][["id", "organization", count_by]].copy()
     treemap_data["root"] = "models"
-    # Ensure numeric again for the final slice
     treemap_data[count_by] = pd.to_numeric(treemap_data[count_by], errors="coerce").fillna(0.0) 
     return treemap_data
 
@@ -168,7 +186,19 @@ with gr.Blocks(title="HuggingFace Model Explorer", fill_width=True) as demo:
             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)
-            size_filter_dropdown = gr.Dropdown(label="Model Size Filter", choices=["None"] + list(MODEL_SIZE_RANGES.keys()), value="None")
+            
+            # --- MODIFIED: Replaced Dropdown with RangeSlider and a Reset Button ---
+            with gr.Group():
+                with gr.Row():
+                    gr.Markdown("<div style='padding-top: 10px; font-weight: 500;'>Parameters</div>")
+                    reset_params_button = gr.Button("🔄 Reset", visible=False, size="sm", min_width=80)
+                param_range_slider = gr.RangeSlider(
+                    label=None, # Label is handled by Markdown above
+                    choices=PARAM_CHOICES,
+                    value=PARAM_CHOICES_DEFAULT,
+                )
+            # --- END OF MODIFICATION ---
+            
             top_k_slider = gr.Slider(label="Number of Top Organizations", minimum=5, maximum=50, value=25, step=5)
             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)
@@ -178,6 +208,20 @@ with gr.Blocks(title="HuggingFace Model Explorer", fill_width=True) as demo:
             status_message_md = gr.Markdown("Initializing...")
             data_info_md = gr.Markdown("")
 
+    # --- NEW: Event handlers for the new parameter slider and reset button ---
+    def _update_reset_button_visibility(current_range):
+        """Shows the reset button only if the slider is not at its default full range."""
+        is_default = (current_range == PARAM_CHOICES_DEFAULT)
+        return gr.update(visible=not is_default)
+
+    def _reset_param_slider_and_button():
+        """Resets the slider to its default value and hides the reset button."""
+        return gr.update(value=PARAM_CHOICES_DEFAULT), gr.update(visible=False)
+
+    param_range_slider.release(fn=_update_reset_button_visibility, inputs=param_range_slider, outputs=reset_params_button)
+    reset_params_button.click(fn=_reset_param_slider_and_button, outputs=[param_range_slider, reset_params_button])
+    # --- END OF NEW 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)
@@ -193,37 +237,24 @@ with gr.Blocks(title="HuggingFace Model Explorer", fill_width=True) as demo:
         data_info_text = ""
         current_df = pd.DataFrame()
         load_success_flag = False
-        data_as_of_date_display = "N/A"
         try:
-            # Call the load function that uses the datasets library.
             current_df, load_success_flag, status_msg_from_load = load_models_data()
             if load_success_flag:
                 progress(0.9, desc="Processing loaded data...")
-                # Get the data timestamp from the loaded file
                 if 'data_download_timestamp' in current_df.columns and not current_df.empty and pd.notna(current_df['data_download_timestamp'].iloc[0]):
-                    timestamp_from_parquet = pd.to_datetime(current_df['data_download_timestamp'].iloc[0])
-                    # Ensure the timestamp is timezone-aware for consistent formatting
-                    if timestamp_from_parquet.tzinfo is None:
-                        timestamp_from_parquet = timestamp_from_parquet.tz_localize('UTC')
+                    timestamp_from_parquet = pd.to_datetime(current_df['data_download_timestamp'].iloc[0]).tz_localize('UTC')
                     data_as_of_date_display = timestamp_from_parquet.strftime('%B %d, %Y, %H:%M:%S %Z')
                 else:
                     data_as_of_date_display = "Pre-processed (date unavailable)"
                 
-                # Create summary text for the UI
-                size_dist_lines = []
-                if 'size_category' in current_df.columns:
-                    for cat in MODEL_SIZE_RANGES.keys():
-                        count = (current_df['size_category'] == cat).sum()
-                        size_dist_lines.append(f"  - {cat}: {count:,} models")
-                else: size_dist_lines.append("  - Size category information not available.")
-                size_dist = "\n".join(size_dist_lines)
-                
+                # --- MODIFIED: Removed the old size category distribution text ---
+                param_count = (current_df['params'] > 0).sum() if 'params' in current_df.columns else 0
                 data_info_text = (f"### Data Information\n"
                                   f"- Source: `{HF_DATASET_ID}`\n"
                                   f"- Overall Status: {status_msg_from_load}\n" 
                                   f"- Total models loaded: {len(current_df):,}\n"
-                                  f"- Data as of: {data_as_of_date_display}\n"
-                                  f"- Size categories:\n{size_dist}")
+                                  f"- Models with parameter counts: {param_count:,}\n"
+                                  f"- Data as of: {data_as_of_date_display}\n")
                 
                 status_msg_ui = "Data loaded successfully. Ready to generate plot."
             else: 
@@ -236,8 +267,9 @@ with gr.Blocks(title="HuggingFace Model Explorer", fill_width=True) as demo:
             load_success_flag = False 
         return current_df, load_success_flag, data_info_text, status_msg_ui
 
+    # --- MODIFIED: Updated controller signature and logic to handle new slider ---
     def ui_generate_plot_controller(metric_choice, filter_type, tag_choice, pipeline_choice, 
-                                   size_choice, k_orgs, skip_orgs_input, df_current_models, progress=gr.Progress()):
+                                   param_range_choice, k_orgs, skip_orgs_input, df_current_models, progress=gr.Progress()):
         if df_current_models is None or df_current_models.empty:
             empty_fig = create_treemap(pd.DataFrame(), metric_choice, "Error: Model Data Not Loaded")
             error_msg = "Model data is not loaded or is empty. Please wait for data to load."
@@ -248,11 +280,10 @@ with gr.Blocks(title="HuggingFace Model Explorer", fill_width=True) as demo:
         
         tag_to_use = tag_choice if filter_type == "Tag Filter" else None
         pipeline_to_use = pipeline_choice if filter_type == "Pipeline Filter" else None
-        size_to_use = size_choice if size_choice != "None" else None
         orgs_to_skip = [org.strip() for org in skip_orgs_input.split(',') if org.strip()] if skip_orgs_input else []
         
-
-        treemap_df = make_treemap_data(df_current_models, metric_choice, k_orgs, tag_to_use, pipeline_to_use, size_to_use, orgs_to_skip)
+        # Pass the param_range_choice directly to make_treemap_data
+        treemap_df = make_treemap_data(df_current_models, metric_choice, k_orgs, tag_to_use, pipeline_to_use, param_range_choice, orgs_to_skip)
         
         progress(0.7, desc="Generating Plotly visualization...")
 
@@ -269,23 +300,22 @@ with gr.Blocks(title="HuggingFace Model Explorer", fill_width=True) as demo:
         
         return plotly_fig, plot_stats_md
 
-    # On app load, call the controller to fetch data using the datasets library.
     demo.load(
         fn=ui_load_data_controller,
         inputs=[],
         outputs=[models_data_state, loading_complete_state, data_info_md, status_message_md]
     )
 
+    # --- MODIFIED: Updated the inputs list for the click event ---
     generate_plot_button.click(
         fn=ui_generate_plot_controller,
         inputs=[count_by_dropdown, filter_choice_radio, tag_filter_dropdown, pipeline_filter_dropdown,
-                size_filter_dropdown, top_k_slider, skip_orgs_textbox, models_data_state],
+                param_range_slider, top_k_slider, skip_orgs_textbox, models_data_state],
         outputs=[plot_output, status_message_md]
     )
 
 if __name__ == "__main__":
     print(f"Application starting. Data will be loaded from Hugging Face dataset: {HF_DATASET_ID}")
-    # Increase the queue size for potentially busy traffic if hosted
     demo.queue().launch()
 
-# --- END OF FILE app.py ---
+# --- END OF MODIFIED FILE app.py ---