Update app.py

app.py

@@ -23,9 +23,12 @@ def load_models_data():
         dataset_dict = load_dataset(HF_DATASET_ID)
         df = dataset_dict[list(dataset_dict.keys())[0]].to_pandas()
         if 'params' in df.columns:
-            df['params'] = pd.to_numeric(df['params'], errors='coerce').fillna(0)
+            # IMPORTANT CHANGE: Fill NaN/coerce errors with -1 to signify unknown size
+            # This aligns with the utility function's return of -1.0 for unknown sizes.
+            df['params'] = pd.to_numeric(df['params'], errors='coerce').fillna(-1)
         else:
-            df['params'] = 0
+            # If 'params' column doesn't exist, assume all are unknown
+            df['params'] = -1
         msg = f"Successfully loaded dataset in {time.time() - overall_start_time:.2f}s."
         print(msg)
         return df, True, msg
@@ -40,9 +43,15 @@ def get_param_range_values(param_range_labels):
     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, param_range=None, skip_orgs=None):
+def make_treemap_data(df, count_by, top_k=25, tag_filter=None, pipeline_filter=None, param_range=None, skip_orgs=None, include_unknown_param_size=True):
     if df is None or df.empty: return pd.DataFrame()
     filtered_df = df.copy()
+
+    # New: Filter based on unknown parameter size
+    # If include_unknown_param_size is False, exclude models where params is -1 (unknown)
+    if not include_unknown_param_size and 'params' in filtered_df.columns:
+        filtered_df = filtered_df[filtered_df['params'] != -1]
+
     col_map = { "Audio & Speech": "is_audio_speech", "Music": "has_music", "Robotics": "has_robot", "Biomedical": "is_biomed", "Time series": "has_series", "Sciences": "has_science", "Video": "has_video", "Images": "has_image", "Text": "has_text" }
     if tag_filter and tag_filter in col_map and col_map[tag_filter] in filtered_df.columns:
         filtered_df = filtered_df[filtered_df[col_map[tag_filter]]]
@@ -51,9 +60,12 @@ def make_treemap_data(df, count_by, top_k=25, tag_filter=None, pipeline_filter=N
     if param_range:
         min_params, max_params = get_param_range_values(param_range)
         is_default_range = (param_range[0] == PARAM_CHOICES[0] and param_range[1] == PARAM_CHOICES[-1])
+        # Apply parameter range filter only if it's not the default (all range) AND params column exists
+        # This filter will naturally exclude -1 if the min_params is >= 0, as it should.
         if not is_default_range and 'params' in filtered_df.columns:
             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'): filtered_df = filtered_df[filtered_df['params'] < max_params]
+
     if skip_orgs and len(skip_orgs) > 0 and "organization" in filtered_df.columns:
         filtered_df = filtered_df[~filtered_df["organization"].isin(skip_orgs)]
     if filtered_df.empty: return pd.DataFrame()
@@ -82,7 +94,6 @@ custom_css = """
 #param-slider-wrapper div[data-testid="range-slider"] > span {
     display: none !important;
 }
-
 /* 
   THIS IS THE KEY FIX:
   We target all the individual component containers (divs with class .block)
@@ -129,6 +140,8 @@ with gr.Blocks(title="🤗 ModelVerse Explorer", fill_width=True, css=custom_css
                     elem_id="param-slider-wrapper"
                 )
                 param_range_display = gr.Markdown(f"Range: `{PARAM_CHOICES[0]}` to `{PARAM_CHOICES[-1]}`")
+                # New: Checkbox for including unknown parameter sizes
+                include_unknown_params_checkbox = gr.Checkbox(label="Include models with unknown parameter size", value=True)
             
             # This section remains un-grouped
             top_k_dropdown = gr.Dropdown(label="Number of Top Organizations", choices=TOP_K_CHOICES, value=25)
@@ -166,8 +179,11 @@ with gr.Blocks(title="🤗 ModelVerse Explorer", fill_width=True, css=custom_css
                 if 'data_download_timestamp' in current_df.columns and pd.notna(current_df['data_download_timestamp'].iloc[0]):
                     ts = pd.to_datetime(current_df['data_download_timestamp'].iloc[0], utc=True)
                     date_display = ts.strftime('%B %d, %Y, %H:%M:%S %Z')
-                param_count = (current_df['params'] > 0).sum() if 'params' in current_df.columns else 0
-                data_info_text = f"### Data Information\n- Source: `{HF_DATASET_ID}`\n- Status: {status_msg_from_load}\n- Total models loaded: {len(current_df):,}\n- Models with parameter counts: {param_count:,}\n- Data as of: {date_display}\n"
+                # Count models where params is not -1 (known size)
+                param_count = (current_df['params'] != -1).sum() if 'params' in current_df.columns else 0
+                unknown_param_count = (current_df['params'] == -1).sum() if 'params' in current_df.columns else 0
+
+                data_info_text = f"### Data Information\n- Source: `{HF_DATASET_ID}`\n- Status: {status_msg_from_load}\n- Total models loaded: {len(current_df):,}\n- Models with known parameter counts: {param_count:,}\n- Models with unknown parameter counts: {unknown_param_count:,}\n- Data as of: {date_display}\n"
             else:
                 data_info_text = f"### Data Load Failed\n- {status_msg_from_load}"
         except Exception as e:
@@ -178,7 +194,6 @@ with gr.Blocks(title="🤗 ModelVerse Explorer", fill_width=True, css=custom_css
             print(f"Critical error in load_and_generate_initial_plot: {e}")
 
         # --- Part 2: Generate Initial Plot ---
-        # We call the existing plot generation function with the default values from the UI
         progress(0.6, desc="Generating initial plot...")
         # Get default values directly from the UI component definitions
         default_metric = "downloads"
@@ -188,18 +203,20 @@ with gr.Blocks(title="🤗 ModelVerse Explorer", fill_width=True, css=custom_css
         default_param_indices = PARAM_CHOICES_DEFAULT_INDICES
         default_k = 25
         default_skip_orgs = "TheBloke,MaziyarPanahi,unsloth,modularai,Gensyn,bartowski"
+        # New default: include unknown params initially (matches checkbox default)
+        default_include_unknown_params = True 
 
         # Reuse the existing controller function for plotting
         initial_plot, initial_status = ui_generate_plot_controller(
             default_metric, default_filter_type, default_tag, default_pipeline, 
-            default_param_indices, default_k, default_skip_orgs, current_df, progress
+            default_param_indices, default_k, default_skip_orgs, default_include_unknown_params, current_df, progress
         )
         
         # Return all the necessary updates for the UI
         return current_df, load_success_flag, data_info_text, initial_status, initial_plot
 
     def ui_generate_plot_controller(metric_choice, filter_type, tag_choice, pipeline_choice, 
-                                   param_range_indices, k_orgs, skip_orgs_input, df_current_models, progress=gr.Progress()):
+                                   param_range_indices, k_orgs, skip_orgs_input, include_unknown_param_size_flag, df_current_models, progress=gr.Progress()):
         if df_current_models is None or df_current_models.empty:
             return create_treemap(pd.DataFrame(), metric_choice, "Error: Model Data Not Loaded"), "Model data is not loaded. Cannot generate plot."
         
@@ -212,7 +229,16 @@ with gr.Blocks(title="🤗 ModelVerse Explorer", fill_width=True, css=custom_css
         max_label = PARAM_CHOICES[int(param_range_indices[1])]
         param_labels_for_filtering = [min_label, max_label]
         
-        treemap_df = make_treemap_data(df_current_models, metric_choice, k_orgs, tag_to_use, pipeline_to_use, param_labels_for_filtering, orgs_to_skip)
+        treemap_df = make_treemap_data(
+            df_current_models, 
+            metric_choice, 
+            k_orgs, 
+            tag_to_use, 
+            pipeline_to_use, 
+            param_labels_for_filtering, 
+            orgs_to_skip, 
+            include_unknown_param_size_flag # Pass the new flag
+        )
         
         progress(0.7, desc="Generating plot...")
         title_labels = {"downloads": "Downloads (last 30 days)", "downloadsAllTime": "Downloads (All Time)", "likes": "Likes"}
@@ -237,7 +263,7 @@ with gr.Blocks(title="🤗 ModelVerse Explorer", fill_width=True, css=custom_css
     generate_plot_button.click(
         fn=ui_generate_plot_controller,
         inputs=[count_by_dropdown, filter_choice_radio, tag_filter_dropdown, pipeline_filter_dropdown,
-                param_range_slider, top_k_dropdown, skip_orgs_textbox, models_data_state],
+                param_range_slider, top_k_dropdown, skip_orgs_textbox, include_unknown_params_checkbox, models_data_state], # Add checkbox to inputs
         outputs=[plot_output, status_message_md]
     )