Update app.py

app.py

@@ -23,14 +23,11 @@ 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:
-            # 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:
-            # If 'params' column doesn't exist, assume all are unknown
             df['params'] = -1
         
-        # NEW: Ensure createdAt is in datetime format, coercing errors
+        # Ensure createdAt is in datetime format, coercing errors
         if 'createdAt' in df.columns:
             df['createdAt'] = pd.to_datetime(df['createdAt'], errors='coerce')
 
@@ -48,12 +45,10 @@ 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, include_unknown_param_size=True, created_after_date=None): # ADDED created_after_date
+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, created_after_date=None):
     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]
 
@@ -65,17 +60,17 @@ 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]
 
-    # NEW: Filter by creation date
+    # Filter by creation date
     if created_after_date is not None and 'createdAt' in filtered_df.columns:
-        # The 'createdAt' column should already be datetime from the loading function
-        # We also convert the input date to ensure timezone consistency (or lack thereof)
-        filtered_df = filtered_df[filtered_df['createdAt'] > pd.to_datetime(created_after_date)]
+        # Filter out rows where createdAt is NaT (Not a Time) before comparison
+        filtered_df = filtered_df.dropna(subset=['createdAt'])
+        # The 'createdAt' column is already datetime from the loading function
+        filtered_df = filtered_df[filtered_df['createdAt'].dt.date > created_after_date]
+
 
     if skip_orgs and len(skip_orgs) > 0 and "organization" in filtered_df.columns:
         filtered_df = filtered_df[~filtered_df["organization"].isin(skip_orgs)]
@@ -98,27 +93,16 @@ def create_treemap(treemap_data, count_by, title=None):
     fig.update_traces(textinfo="label+value+percent root", hovertemplate="<b>%{label}</b><br>%{value:,} " + count_by + "<br>%{percentRoot:.2%} of total<extra></extra>")
     return fig
 
-# --- FINAL, CORRECTED CSS ---
 custom_css = """
-/* Hide the extra UI elements from the RangeSlider component */
-#param-slider-wrapper .head,
-#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)
-  that are *direct children* of our custom-classed group.
-  
-  This removes the "box-in-a-box" effect by making the inner component
-  containers transparent. The parent gr.Group now acts as the single card,
-  which is exactly what we want.
-*/
 .model-parameters-group > .block {
     background: none !important;
     border: none !important;
     box-shadow: none !important;
 }
+#param-slider-wrapper .head,
+#param-slider-wrapper div[data-testid="range-slider"] > span {
+    display: none !important;
+}
 """
 
 with gr.Blocks(title="🤗 ModelVerse Explorer", fill_width=True, css=custom_css) as demo:
@@ -131,32 +115,23 @@ with gr.Blocks(title="🤗 ModelVerse Explorer", fill_width=True, css=custom_css
     with gr.Row():
         with gr.Column(scale=1):
             
-            # This section remains un-grouped for a consistent flat look
             count_by_dropdown = gr.Dropdown(label="Metric", choices=[("Downloads (last 30 days)", "downloads"), ("Downloads (All Time)", "downloadsAllTime"), ("Likes", "likes")], value="downloads")
             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)
             
-            # This group's styling will be modified by the custom CSS
             with gr.Group(elem_classes="model-parameters-group"):
                 gr.Markdown("<div style='font-weight: 500;'>Model Parameters</div>")
                 param_range_slider = RangeSlider(
-                    minimum=0,
-                    maximum=len(PARAM_CHOICES) - 1,
-                    value=PARAM_CHOICES_DEFAULT_INDICES,
-                    step=1,
-                    label=None,
-                    show_label=False,
-                    elem_id="param-slider-wrapper"
+                    minimum=0, maximum=len(PARAM_CHOICES) - 1, value=PARAM_CHOICES_DEFAULT_INDICES,
+                    step=1, label=None, show_label=False, 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
-            # NEW: Date picker for filtering by creation date
-            created_after_datepicker = gr.DatePicker(label="Created After", value=None)
+            # --- CORRECTED LINE ---
+            created_after_datepicker = gr.Date(label="Created After", value=None)
             
             top_k_dropdown = gr.Dropdown(label="Number of Top Organizations", choices=TOP_K_CHOICES, value=25)
             skip_orgs_textbox = gr.Textbox(label="Organizations to Skip (comma-separated)", value="TheBloke,MaziyarPanahi,unsloth,modularai,Gensyn,bartowski")
@@ -172,141 +147,80 @@ with gr.Blocks(title="🤗 ModelVerse Explorer", fill_width=True, css=custom_css
         min_idx, max_idx = int(value[0]), int(value[1])
         return f"Range: `{PARAM_CHOICES[min_idx]}` to `{PARAM_CHOICES[max_idx]}`"
     
-    # New function to toggle the unknown params checkbox interactivity
     def _toggle_unknown_params_checkbox(param_range_indices):
         min_idx, max_idx = int(param_range_indices[0]), int(param_range_indices[1])
-        is_default_range = (min_idx == PARAM_CHOICES_DEFAULT_INDICES[0] and 
-                            max_idx == PARAM_CHOICES_DEFAULT_INDICES[1])
-        
-        # If a specific range is selected (not the default all-inclusive range), disable the checkbox
-        # and uncheck it to ensure consistency.
-        if not is_default_range:
-            return gr.update(interactive=False, value=False) # Disable and uncheck
-        else:
-            return gr.update(interactive=True) # Enable
+        is_default_range = (min_idx == PARAM_CHOICES_DEFAULT_INDICES[0] and max_idx == PARAM_CHOICES_DEFAULT_INDICES[1])
+        return gr.update(interactive=is_default_range, value=True if is_default_range else False)
             
     param_range_slider.change(update_param_display, param_range_slider, param_range_display)
-    # Connect the new toggle function to the param_range_slider's change event
-    param_range_slider.change(
-        fn=_toggle_unknown_params_checkbox, 
-        inputs=[param_range_slider], 
-        outputs=[include_unknown_params_checkbox]
-    )
+    param_range_slider.change(_toggle_unknown_params_checkbox, param_range_slider, include_unknown_params_checkbox)
 
-    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)
+    loading_complete_state.change(lambda is_loaded: gr.update(interactive=is_loaded), loading_complete_state, generate_plot_button)
 
-    def _toggle_filters_visibility(choice): 
-        return gr.update(visible=choice == "Tag Filter"), gr.update(visible=choice == "Pipeline Filter")
-    filter_choice_radio.change(fn=_toggle_filters_visibility, inputs=filter_choice_radio, outputs=[tag_filter_dropdown, pipeline_filter_dropdown])
+    filter_choice_radio.change(lambda choice: (gr.update(visible=choice == "Tag Filter"), gr.update(visible=choice == "Pipeline Filter")),
+                               filter_choice_radio, [tag_filter_dropdown, pipeline_filter_dropdown])
 
-    ## CHANGE: Renamed and modified ui_load_data_controller to also generate the initial plot
     def load_and_generate_initial_plot(progress=gr.Progress()):
         progress(0, desc=f"Loading dataset '{HF_DATASET_ID}'...")
-        # --- Part 1: Data Loading ---
+        current_df, load_success_flag, status_msg_from_load = pd.DataFrame(), False, ""
         try:
             current_df, load_success_flag, status_msg_from_load = load_models_data()
             if load_success_flag:
                 progress(0.5, desc="Processing data...")
-                date_display = "Pre-processed (date unavailable)"
-                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')
-                # 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"
+                ts = pd.to_datetime(current_df['data_download_timestamp'].iloc[0], utc=True) if 'data_download_timestamp' in current_df.columns and pd.notna(current_df['data_download_timestamp'].iloc[0]) else None
+                date_display = ts.strftime('%B %d, %Y, %H:%M:%S %Z') if ts else "Pre-processed (date unavailable)"
+                
+                param_count = (current_df['params'] != -1).sum()
+                data_info_text = (f"### Data Information\n- Source: `{HF_DATASET_ID}`\n- Status: {status_msg_from_load}\n"
+                                  f"- Total models loaded: {len(current_df):,}\n- Models with known parameter counts: {param_count:,}\n"
+                                  f"- Models with unknown parameter counts: {len(current_df) - 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:
             status_msg_from_load = f"An unexpected error occurred: {str(e)}"
             data_info_text = f"### Critical Error\n- {status_msg_from_load}"
-            load_success_flag = False
-            current_df = pd.DataFrame()
             print(f"Critical error in load_and_generate_initial_plot: {e}")
 
-        # --- Part 2: Generate Initial Plot ---
         progress(0.6, desc="Generating initial plot...")
-        # Get default values directly from the UI component definitions
-        default_metric = "downloads"
-        default_filter_type = "None"
-        default_tag = None
-        default_pipeline = None
-        default_param_indices = PARAM_CHOICES_DEFAULT_INDICES
-        default_k = 25
-        default_skip_orgs = "TheBloke,MaziyarPanahi,unsloth,modularai,Gensyn,bartowski"
-        default_include_unknown_params = True
-        default_created_after_date = None # NEW: Default for date picker is no filter
-
-        # 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, default_include_unknown_params,
-            default_created_after_date, # Pass new default date
-            current_df, progress
+            "downloads", "None", None, None, PARAM_CHOICES_DEFAULT_INDICES, 25, 
+            "TheBloke,MaziyarPanahi,unsloth,modularai,Gensyn,bartowski", True, None, 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, include_unknown_param_size_flag,
-                                   created_after_date, # ADDED created_after_date
-                                   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."
+                                   created_after_date, df_current_models, progress=gr.Progress()):
+        if df_current_models.empty:
+            return create_treemap(pd.DataFrame(), metric_choice, "Error: Model Data Not Loaded"), "Model data is not loaded."
         
         progress(0.1, desc="Preparing data...")
-        tag_to_use = tag_choice if filter_type == "Tag Filter" else None
-        pipeline_to_use = pipeline_choice if filter_type == "Pipeline Filter" else None
-        orgs_to_skip = [org.strip() for org in skip_orgs_input.split(',') if org.strip()]
-        
-        min_label = PARAM_CHOICES[int(param_range_indices[0])]
-        max_label = PARAM_CHOICES[int(param_range_indices[1])]
-        param_labels_for_filtering = [min_label, max_label]
+        param_labels = [PARAM_CHOICES[int(param_range_indices[0])], PARAM_CHOICES[int(param_range_indices[1])]]
         
         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,
-            created_after_date # Pass the new date filter value
+            df_current_models, metric_choice, k_orgs, 
+            tag_choice if filter_type == "Tag Filter" else None, 
+            pipeline_choice if filter_type == "Pipeline Filter" else None,
+            param_labels, [org.strip() for org in skip_orgs_input.split(',') if org.strip()], 
+            include_unknown_param_size_flag, created_after_date
         )
         
         progress(0.7, desc="Generating plot...")
         title_labels = {"downloads": "Downloads (last 30 days)", "downloadsAllTime": "Downloads (All Time)", "likes": "Likes"}
-        chart_title = f"HuggingFace Models - {title_labels.get(metric_choice, metric_choice)} by Organization"
-        plotly_fig = create_treemap(treemap_df, metric_choice, chart_title)
+        plotly_fig = create_treemap(treemap_df, metric_choice, f"HuggingFace Models - {title_labels.get(metric_choice, metric_choice)} by Organization")
         
-        if treemap_df.empty:
-            plot_stats_md = "No data matches the selected filters. Please try different options."
-        else:
-            total_items_in_plot = len(treemap_df['id'].unique())
-            total_value_in_plot = treemap_df[metric_choice].sum()
-            plot_stats_md = f"## Plot Statistics\n- **Models shown**: {total_items_in_plot:,}\n- **Total {metric_choice}**: {int(total_value_in_plot):,}"
+        plot_stats_md = (f"## Plot Statistics\n- **Models shown**: {len(treemap_df['id'].unique()):,}\n"
+                         f"- **Total {metric_choice}**: {int(treemap_df[metric_choice].sum()):,}") if not treemap_df.empty else "No data matches the selected filters."
         return plotly_fig, plot_stats_md
 
-    ## CHANGE: Updated demo.load to call the new function and to add plot_output to the outputs list
-    demo.load(
-        fn=load_and_generate_initial_plot, 
-        inputs=[], 
-        outputs=[models_data_state, loading_complete_state, data_info_md, status_message_md, plot_output]
-    )
+    demo.load(load_and_generate_initial_plot, None, [models_data_state, loading_complete_state, data_info_md, status_message_md, plot_output])
 
     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, include_unknown_params_checkbox,
-            created_after_datepicker, # ADDED date picker to inputs
-            models_data_state
-        ],
-        outputs=[plot_output, status_message_md]
+        ui_generate_plot_controller,
+        [count_by_dropdown, filter_choice_radio, tag_filter_dropdown, pipeline_filter_dropdown,
+         param_range_slider, top_k_dropdown, skip_orgs_textbox, include_unknown_params_checkbox,
+         created_after_datepicker, models_data_state],
+        [plot_output, status_message_md]
     )
 
 if __name__ == "__main__":