add cached data and preprocessing code

.gitattributes

@@ -35,3 +35,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
 org_to_artifacts_2l_stats.json filter=lfs diff=lfs merge=lfs -text
 models.csv filter=lfs diff=lfs merge=lfs -text
+models_processed.parquet filter=lfs diff=lfs merge=lfs -text

app.py

@@ -1,550 +1,541 @@
+# --- START OF FILE app.py ---
+
 import json
 import gradio as gr
 import pandas as pd
 import plotly.express as px
 import os
-import numpy as np
-import io
+import numpy as np # Make sure NumPy is imported
+import duckdb
+from tqdm.auto import tqdm # Standard tqdm for console, gr.Progress will track it
+import time
+import ast # For safely evaluating string representations of lists/dicts
+
+# --- 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'))
+}
+PROCESSED_PARQUET_FILE_PATH = "models_processed.parquet"
+HF_PARQUET_URL = 'https://huggingface.co/datasets/cfahlgren1/hub-stats/resolve/main/models.parquet'
+
+TAG_FILTER_CHOICES = [
+    "Audio & Speech", "Time series", "Robotics", "Music", "Video", "Images",
+    "Text", "Biomedical", "Sciences"
+]
 
-# Define pipeline tags
 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',
+ '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'))
-}
+# --- Utility Functions ---
+def extract_model_size(safetensors_data): # Renamed for consistency if used, preprocessor uses extract_model_file_size_gb
+    try:
+        if pd.isna(safetensors_data): return 0.0
+        data_to_parse = safetensors_data
+        if isinstance(safetensors_data, str):
+            try:
+                if (safetensors_data.startswith('{') and safetensors_data.endswith('}')) or \
+                   (safetensors_data.startswith('[') and safetensors_data.endswith(']')):
+                    data_to_parse = ast.literal_eval(safetensors_data)
+                else: data_to_parse = json.loads(safetensors_data)
+            except: return 0.0
+        if isinstance(data_to_parse, dict) and 'total' in data_to_parse:
+            try:
+                total_bytes_val = data_to_parse['total']
+                size_bytes = float(total_bytes_val)
+                return size_bytes / (1024 * 1024 * 1024)
+            except (ValueError, TypeError): pass
+        return 0.0
+    except: return 0.0
+
+def extract_org_from_id(model_id):
+    if pd.isna(model_id): return "unaffiliated"
+    model_id_str = str(model_id)
+    return model_id_str.split("/")[0] if "/" in model_id_str else "unaffiliated"
 
-# Filter functions for tags
-def is_audio_speech(row):
-    tags = row.get("tags", [])
-    pipeline_tag = row.get("pipeline_tag", "")
+# --- THIS IS THE CORRECTED process_tags_for_series from preprocess.py ---
+def process_tags_for_series(series_of_tags_values, tqdm_cls=None): # Added tqdm_cls for Gradio progress
+    processed_tags_accumulator = []
     
-    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)
-
-def is_music(row):
-    tags = row.get("tags", [])
-    return any("music" in tag.lower() for tag in tags)
-
-def is_robotics(row):
-    tags = row.get("tags", [])
-    return any("robot" in tag.lower() for tag in tags)
-
-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)
-
-def is_timeseries(row):
-    tags = row.get("tags", [])
-    return any("series" in tag.lower() for tag in tags)
-
-def is_science(row):
-    tags = row.get("tags", [])
-    return any("science" in tag.lower() and "bigscience" not in tag for tag in tags)
-
-def is_video(row):
-    tags = row.get("tags", [])
-    return any("video" in tag.lower() for tag in tags)
-
-def is_image(row):
-    tags = row.get("tags", [])
-    return any("image" in tag.lower() for tag in tags)
-
-def is_text(row):
-    tags = row.get("tags", [])
-    return any("text" in tag.lower() for tag in tags)
-
-# Add model size filter function
-def is_in_size_range(row, size_range):
-    if size_range is None:
-        return True
+    # Determine the iterable (use tqdm if tqdm_cls is provided, else direct iteration)
+    iterable = series_of_tags_values
+    if tqdm_cls and tqdm_cls != tqdm : # Check if it's Gradio's progress tracker
+        iterable = tqdm_cls(series_of_tags_values, desc="Standardizing Tags (App)", unit="row")
+    elif tqdm_cls == tqdm: # For direct console tqdm if passed
+        iterable = tqdm(series_of_tags_values, desc="Standardizing Tags (App)", unit="row", leave=False)
+
+
+    for i, tags_value_from_series in enumerate(iterable):
+        temp_processed_list_for_row = []
+        current_value_for_error_msg = str(tags_value_from_series)[:200]
+
+        try:
+            if isinstance(tags_value_from_series, list):
+                current_tags_in_list = []
+                for tag_item in tags_value_from_series:
+                    try:
+                        if pd.isna(tag_item): continue
+                        str_tag = str(tag_item)
+                        stripped_tag = str_tag.strip()
+                        if stripped_tag:
+                            current_tags_in_list.append(stripped_tag)
+                    except Exception as e_inner_list_proc:
+                        print(f"APP ERROR processing item '{tag_item}' (type: {type(tag_item)}) within a list for row {i}. Error: {e_inner_list_proc}. Original: {current_value_for_error_msg}")
+                temp_processed_list_for_row = current_tags_in_list
+
+            elif isinstance(tags_value_from_series, np.ndarray):
+                current_tags_in_list = []
+                for tag_item in tags_value_from_series.tolist():
+                    try:
+                        if pd.isna(tag_item): continue
+                        str_tag = str(tag_item)
+                        stripped_tag = str_tag.strip()
+                        if stripped_tag:
+                            current_tags_in_list.append(stripped_tag)
+                    except Exception as e_inner_array_proc:
+                        print(f"APP ERROR processing item '{tag_item}' (type: {type(tag_item)}) within a NumPy array for row {i}. Error: {e_inner_array_proc}. Original: {current_value_for_error_msg}")
+                temp_processed_list_for_row = current_tags_in_list
+            
+            elif tags_value_from_series is None or pd.isna(tags_value_from_series):
+                temp_processed_list_for_row = []
+
+            elif isinstance(tags_value_from_series, str):
+                processed_str_tags = []
+                if (tags_value_from_series.startswith('[') and tags_value_from_series.endswith(']')) or \
+                   (tags_value_from_series.startswith('(') and tags_value_from_series.endswith(')')):
+                    try:
+                        evaluated_tags = ast.literal_eval(tags_value_from_series)
+                        if isinstance(evaluated_tags, (list, tuple)):
+                            current_eval_list = []
+                            for tag_item in evaluated_tags:
+                                if pd.isna(tag_item): continue
+                                str_tag = str(tag_item).strip()
+                                if str_tag: current_eval_list.append(str_tag)
+                            processed_str_tags = current_eval_list
+                    except (ValueError, SyntaxError):
+                        pass 
+
+                if not processed_str_tags:
+                    try:
+                        json_tags = json.loads(tags_value_from_series)
+                        if isinstance(json_tags, list):
+                            current_json_list = []
+                            for tag_item in json_tags:
+                                if pd.isna(tag_item): continue
+                                str_tag = str(tag_item).strip()
+                                if str_tag: current_json_list.append(str_tag)
+                            processed_str_tags = current_json_list
+                    except json.JSONDecodeError:
+                        processed_str_tags = [tag.strip() for tag in tags_value_from_series.split(',') if tag.strip()]
+                    except Exception as e_json_other:
+                        print(f"APP ERROR during JSON processing for string '{current_value_for_error_msg}' for row {i}. Error: {e_json_other}")
+                        processed_str_tags = [tag.strip() for tag in tags_value_from_series.split(',') if tag.strip()]
+
+                temp_processed_list_for_row = processed_str_tags
+            
+            else: 
+                if pd.isna(tags_value_from_series):
+                     temp_processed_list_for_row = []
+                else:
+                    str_val = str(tags_value_from_series).strip()
+                    temp_processed_list_for_row = [str_val] if str_val else []
+            
+            processed_tags_accumulator.append(temp_processed_list_for_row)
+
+        except Exception as e_outer_tag_proc:
+            print(f"APP CRITICAL UNHANDLED ERROR processing row {i}: value '{current_value_for_error_msg}' (type: {type(tags_value_from_series)}). Error: {e_outer_tag_proc}. Appending [].")
+            processed_tags_accumulator.append([])
+            
+    return processed_tags_accumulator
+# --- END OF CORRECTED process_tags_for_series ---
+
+
+def load_models_data(force_refresh=False, tqdm_cls=None): # tqdm_cls for Gradio progress
+    # ... (initial part of load_models_data for loading pre-processed parquet is the same) ...
+    if tqdm_cls is None: tqdm_cls = tqdm # Default to standard tqdm if None
+    overall_start_time = time.time()
+    print(f"Gradio load_models_data called with force_refresh={force_refresh}")
+
+    expected_cols_in_processed_parquet = [
+        'id', 'downloads', 'downloadsAllTime', 'likes', 'pipeline_tag', 'tags', 'params',
+        'size_category', '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'
+    ]
+
+    if not force_refresh and os.path.exists(PROCESSED_PARQUET_FILE_PATH):
+        print(f"Attempting to load pre-processed data from: {PROCESSED_PARQUET_FILE_PATH}")
+        try:
+            df = pd.read_parquet(PROCESSED_PARQUET_FILE_PATH)
+            elapsed = time.time() - overall_start_time
+            missing_cols = [col for col in expected_cols_in_processed_parquet if col not in df.columns]
+            if missing_cols:
+                raise ValueError(f"Pre-processed Parquet is missing columns: {missing_cols}. Please run preprocessor or refresh data in app.")
+            
+            if 'has_robot' in df.columns:
+                robot_count_parquet = df['has_robot'].sum()
+                print(f"DIAGNOSTIC (App - Parquet Load): 'has_robot' column found. Number of True values: {robot_count_parquet}")
+            else:
+                print("DIAGNOSTIC (App - Parquet Load): 'has_robot' column NOT FOUND.")
+
+            msg = f"Successfully loaded pre-processed data in {elapsed:.2f}s. Shape: {df.shape}"
+            print(msg)
+            return df, True, msg
+        except Exception as e:
+            print(f"Could not load pre-processed Parquet: {e}. ")
+            if force_refresh: print("Proceeding to fetch fresh data as force_refresh=True.")
+            else:
+                 err_msg = (f"Pre-processed data could not be loaded: {e}. "
+                           "Please use 'Refresh Data from Hugging Face' button.")
+                 return pd.DataFrame(), False, err_msg
+
+    df_raw = None
+    raw_data_source_msg = ""
+    if force_refresh:
+        print("force_refresh=True (Gradio). Fetching fresh data...")
+        fetch_start = time.time()
+        try:
+            query = f"SELECT * FROM read_parquet('{HF_PARQUET_URL}')"
+            df_raw = duckdb.sql(query).df()
+            if df_raw is None or df_raw.empty: raise ValueError("Fetched data is empty or None.")
+            raw_data_source_msg = f"Fetched by Gradio in {time.time() - fetch_start:.2f}s. Rows: {len(df_raw)}"
+            print(raw_data_source_msg)
+        except Exception as e_hf:
+            return pd.DataFrame(), False, f"Fatal error fetching from Hugging Face (Gradio): {e_hf}"
+    else: 
+        err_msg = (f"Pre-processed data '{PROCESSED_PARQUET_FILE_PATH}' not found/invalid. "
+                   "Run preprocessor or use 'Refresh Data' button.")
+        return pd.DataFrame(), False, err_msg
+
+    print(f"Initiating processing for data newly fetched by Gradio. {raw_data_source_msg}")
+    df = pd.DataFrame() # This will be our processed DataFrame
+    proc_start = time.time()
     
-    min_size, max_size = MODEL_SIZE_RANGES[size_range]
+    core_cols = {'id': str, 'downloads': float, 'downloadsAllTime': float, 'likes': float,
+                 'pipeline_tag': str, 'tags': object, 'safetensors': object}
+    for col, dtype in core_cols.items():
+        if col in df_raw.columns:
+            df[col] = df_raw[col] # Assign raw data first
+            if dtype == float: df[col] = pd.to_numeric(df[col], errors='coerce').fillna(0.0)
+            elif dtype == str: df[col] = df[col].astype(str).fillna('')
+            # For 'tags' and 'safetensors' (object type), no specific conversion here, done later
+        else: # If column is missing in raw data
+            if col in ['downloads', 'downloadsAllTime', 'likes']: df[col] = 0.0
+            elif col == 'pipeline_tag': df[col] = ''
+            elif col == 'tags': df[col] = pd.Series([[] for _ in range(len(df_raw))]) # Default to empty lists
+            elif col == 'safetensors': df[col] = None # Default to None
+            elif col == 'id': return pd.DataFrame(), False, "Critical: 'id' column missing."
     
-    # 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
+    output_filesize_col_name = 'params'
+    if output_filesize_col_name in df_raw.columns and pd.api.types.is_numeric_dtype(df_raw[output_filesize_col_name]):
+        df[output_filesize_col_name] = pd.to_numeric(df_raw[output_filesize_col_name], errors='coerce').fillna(0.0)
+    elif 'safetensors' in df.columns:
+        # Use tqdm_cls for progress tracking if available (Gradio's gr.Progress.tqdm)
+        safetensors_iter = df['safetensors']
+        if tqdm_cls and tqdm_cls != tqdm: # Check if it's Gradio's progress tracker
+             safetensors_iter = tqdm_cls(df['safetensors'], desc="Extracting model sizes (GB)", unit="row")
+        elif tqdm_cls == tqdm: # For direct console tqdm if passed
+            safetensors_iter = tqdm(df['safetensors'], desc="Extracting model sizes (GB)", unit="row", leave=False)
+            
+        df[output_filesize_col_name] = [extract_model_size(s) for s in safetensors_iter]
+        df[output_filesize_col_name] = pd.to_numeric(df[output_filesize_col_name], errors='coerce').fillna(0.0)
+    else:
+        df[output_filesize_col_name] = 0.0
+
+    def get_size_category_gradio(size_gb_val):
+        try: numeric_size_gb = float(size_gb_val)
+        except (ValueError, TypeError): numeric_size_gb = 0.0
+        if pd.isna(numeric_size_gb): numeric_size_gb = 0.0
+        if 0 <= numeric_size_gb < 1: return "Small (<1GB)"
+        elif 1 <= numeric_size_gb < 5: return "Medium (1-5GB)"
+        elif 5 <= numeric_size_gb < 20: return "Large (5-20GB)"
+        elif 20 <= numeric_size_gb < 50: return "X-Large (20-50GB)"
+        elif numeric_size_gb >= 50: return "XX-Large (>50GB)"
+        else: return "Small (<1GB)" # Default
+    df['size_category'] = df[output_filesize_col_name].apply(get_size_category_gradio)
+
+    # >>> USE THE CORRECTED process_tags_for_series HERE <<<
+    df['tags'] = process_tags_for_series(df['tags'], tqdm_cls=tqdm_cls) 
     
-    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,
-}
+    df['temp_tags_joined'] = df['tags'].apply(
+        lambda tl: '~~~'.join(str(t).lower().strip() for t in tl if pd.notna(t) and str(t).strip()) if isinstance(tl, list) else ''
+    )
+    tag_map = {
+        'has_audio': ['audio'], 'has_speech': ['speech'], 'has_music': ['music'],
+        'has_robot': ['robot', 'robotics'], 
+        'has_bio': ['bio'], 'has_med': ['medic', 'medical'],
+        'has_series': ['series', 'time-series', 'timeseries'],
+        'has_video': ['video'], 'has_image': ['image', 'vision'],
+        'has_text': ['text', 'nlp', 'llm']
+    }
+    for col, kws in tag_map.items():
+        pattern = '|'.join(kws)
+        df[col] = df['temp_tags_joined'].str.contains(pattern, na=False, case=False, regex=True)
+    df['has_science'] = (
+        df['temp_tags_joined'].str.contains('science', na=False, case=False, regex=True) &
+        ~df['temp_tags_joined'].str.contains('bigscience', na=False, case=False, regex=True)
+    )
+    del df['temp_tags_joined']
+    df['is_audio_speech'] = (df['has_audio'] | df['has_speech'] |
+                            df['pipeline_tag'].str.contains('audio|speech', case=False, na=False, regex=True))
+    df['is_biomed'] = df['has_bio'] | df['has_med']
+    df['organization'] = df['id'].apply(extract_org_from_id)
 
-def extract_org_from_id(model_id):
-    """Extract organization name from model ID"""
-    if "/" in model_id:
-        return model_id.split("/")[0]
-    return "unaffiliated"
+    # Drop safetensors if params was calculated from it, and params didn't pre-exist as numeric
+    if 'safetensors' in df.columns and \
+       not (output_filesize_col_name in df_raw.columns and pd.api.types.is_numeric_dtype(df_raw[output_filesize_col_name])):
+        df = df.drop(columns=['safetensors'], errors='ignore')
+    
+    if force_refresh and 'has_robot' in df.columns:
+        robot_count_app_proc = df['has_robot'].sum()
+        print(f"DIAGNOSTIC (App - Force Refresh Processing): 'has_robot' column processed. Number of True values: {robot_count_app_proc}")
 
+    print(f"Data processing by Gradio completed in {time.time() - proc_start:.2f}s.")
+    
+    total_elapsed = time.time() - overall_start_time
+    final_msg = f"{raw_data_source_msg}. Processing by Gradio took {time.time() - proc_start:.2f}s. Total: {total_elapsed:.2f}s. Shape: {df.shape}"
+    print(final_msg)
+    return df, True, final_msg
+
+
+# ... (make_treemap_data, create_treemap functions remain unchanged) ...
 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"""
-    # Create a copy to avoid modifying the original
+    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",
+                "Biomedical": "is_biomed", "Time series": "has_series", "Sciences": "has_science",
+                "Video": "has_video", "Images": "has_image", "Text": "has_text"}
     
-    # Apply filters
-    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)]
-    
+    if 'has_robot' in filtered_df.columns:
+        initial_robot_count = filtered_df['has_robot'].sum()
+        # print(f"DIAGNOSTIC (make_treemap_data entry): Input df has {initial_robot_count} 'has_robot' models.") # Can be noisy
+    # else:
+        # print("DIAGNOSTIC (make_treemap_data entry): 'has_robot' column NOT in input df.")
+
+    if tag_filter and tag_filter in col_map:
+        target_col = col_map[tag_filter]
+        if target_col in filtered_df.columns:
+            # if tag_filter == "Robotics":
+                # count_before_robot_filter = filtered_df[target_col].sum()
+                # print(f"DIAGNOSTIC (make_treemap_data): Applying 'Robotics' filter. Models with '{target_col}'=True: {count_before_robot_filter}")
+            filtered_df = filtered_df[filtered_df[target_col]]
+            # if tag_filter == "Robotics":
+                 # print(f"DIAGNOSTIC (make_treemap_data): After 'Robotics' filter ({target_col}), df rows: {len(filtered_df)}")
+        else:
+            print(f"Warning: Tag filter column '{col_map[tag_filter]}' not found in DataFrame.")
     if pipeline_filter:
-        filtered_df = filtered_df[filtered_df["pipeline_tag"] == pipeline_filter]
-    
-    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)
-        
-        filtered_df = filtered_df[filtered_df.apply(check_size, axis=1)]
-    
-    # Add organization column
-    filtered_df["organization"] = filtered_df["id"].apply(extract_org_from_id)
-    
-    # Skip organizations if specified
+        if "pipeline_tag" in filtered_df.columns:
+            filtered_df = filtered_df[filtered_df["pipeline_tag"] == 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.")
     if skip_orgs and len(skip_orgs) > 0:
-        filtered_df = filtered_df[~filtered_df["organization"].isin(skip_orgs)]
-    
-    # Aggregate by organization
-    org_totals = filtered_df.groupby("organization")[count_by].sum().reset_index()
-    org_totals = org_totals.sort_values(by=count_by, ascending=False)
-    
-    # Get top organizations
-    top_orgs = org_totals.head(top_k)["organization"].tolist()
-    
-    # Filter to only include models from top organizations
-    filtered_df = filtered_df[filtered_df["organization"].isin(top_orgs)]
-    
-    # Prepare data for treemap
-    treemap_data = filtered_df[["id", "organization", count_by]].copy()
-    
-    # Add a root node
-    treemap_data["root"] = "models"
-    
-    # Ensure numeric values
-    treemap_data[count_by] = pd.to_numeric(treemap_data[count_by], errors="coerce").fillna(0)
-    
+        if "organization" in filtered_df.columns:
+            filtered_df = filtered_df[~filtered_df["organization"].isin(skip_orgs)]
+        else:
+            print("Warning: 'organization' column not found for filtering.")
+    if filtered_df.empty: return pd.DataFrame()
+    # Ensure count_by column is numeric, coercing if necessary
+    if count_by not in filtered_df.columns or not pd.api.types.is_numeric_dtype(filtered_df[count_by]):
+        # print(f"Warning: Column '{count_by}' for treemap values is not numeric or missing. Coercing to numeric, filling NaNs with 0.")
+        filtered_df[count_by] = pd.to_numeric(filtered_df.get(count_by), errors="coerce").fillna(0.0)
+        
+    org_totals = filtered_df.groupby("organization")[count_by].sum().nlargest(top_k, keep='first') # Default keep='first'
+    top_orgs_list = org_totals.index.tolist()
+    treemap_data = filtered_df[filtered_df["organization"].isin(top_orgs_list)][["id", "organization", count_by]].copy()
+    treemap_data["root"] = "models" # For treemap structure
+    treemap_data[count_by] = pd.to_numeric(treemap_data[count_by], errors="coerce").fillna(0.0) # Ensure numeric again after subsetting
     return treemap_data
 
 def create_treemap(treemap_data, count_by, title=None):
-    """Create a Plotly treemap from the prepared data"""
     if treemap_data.empty:
-        # Create an empty figure with a message
-        fig = px.treemap(
-            names=["No data matches the selected filters"],
-            values=[1]
-        )
-        fig.update_layout(
-            title="No data matches the selected filters",
-            margin=dict(t=50, l=25, r=25, b=25)
-        )
+        fig = px.treemap(names=["No data matches filters"], parents=[""], values=[1]) # Placeholder for empty data
+        fig.update_layout(title="No data matches the selected filters", margin=dict(t=50, l=25, r=25, b=25))
         return fig
-    
-    # Create the treemap
     fig = px.treemap(
-        treemap_data,
-        path=["root", "organization", "id"],
-        values=count_by,
+        treemap_data, path=["root", "organization", "id"], values=count_by,
         title=title or f"HuggingFace Models - {count_by.capitalize()} by Organization",
-        color_discrete_sequence=px.colors.qualitative.Plotly
-    )
-    
-    # Update layout
-    fig.update_layout(
-        margin=dict(t=50, l=25, r=25, b=25)
-    )
-    
-    # Update traces for better readability
-    fig.update_traces(
-        textinfo="label+value+percent root",
-        hovertemplate="<b>%{label}</b><br>%{value:,} " + count_by + "<br>%{percentRoot:.2%} of total<extra></extra>"
+        color_discrete_sequence=px.colors.qualitative.Plotly # Example color sequence
     )
-    
+    fig.update_layout(margin=dict(t=50, l=25, r=25, b=25))
+    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
 
-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 []
-        
-        # 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)
-    
-    # Add more sample data for better visualization
-    add_sample_data(df)
-    
-    return df
-
-def add_sample_data(df):
-    """Add more sample data to make the visualization more interesting"""
-    # Top organizations to include
-    orgs = ['openai', 'meta', 'google', 'microsoft', 'anthropic', 'nvidia', 'huggingface', 
-            'deepseek-ai', 'stability-ai', 'mistralai', 'cerebras', 'databricks', 'together',
-            'facebook', 'amazon', 'deepmind', 'cohere', 'bigscience', 'eleutherai']
-    
-    # Common model name formats
-    model_name_patterns = [
-        "model-{size}-{version}",
-        "{prefix}-{size}b",
-        "{prefix}-{size}b-{variant}",
-        "llama-{size}b-{variant}",
-        "gpt-{variant}-{size}b",
-        "{prefix}-instruct-{size}b",
-        "{prefix}-chat-{size}b",
-        "{prefix}-coder-{size}b",
-        "stable-diffusion-{version}",
-        "whisper-{size}",
-        "bert-{size}-{variant}",
-        "roberta-{size}",
-        "t5-{size}",
-        "{prefix}-vision-{size}b"
-    ]
-    
-    # Common name parts
-    prefixes = ["falcon", "llama", "mistral", "gpt", "phi", "gemma", "qwen", "yi", "mpt", "bloom"]
-    sizes = ["7", "13", "34", "70", "1", "3", "7b", "13b", "70b", "8b", "2b", "1b", "0.5b", "small", "base", "large", "huge"]
-    variants = ["chat", "instruct", "base", "v1.0", "v2", "beta", "turbo", "fast", "xl", "xxl"]
-    
-    # Generate sample data
-    sample_data = []
-    for org_idx, org in enumerate(orgs):
-        # Create 5-10 models per organization
-        num_models = np.random.randint(5, 11)
-        
-        for i in range(num_models):
-            # Create realistic model name
-            pattern = np.random.choice(model_name_patterns)
-            prefix = np.random.choice(prefixes)
-            size = np.random.choice(sizes)
-            version = f"v{np.random.randint(1, 4)}"
-            variant = np.random.choice(variants)
-            
-            model_name = pattern.format(
-                prefix=prefix,
-                size=size,
-                version=version,
-                variant=variant
-            )
-            
-            model_id = f"{org}/{model_name}"
-            
-            # Select a realistic pipeline tag based on name
-            if "diffusion" in model_name or "image" in model_name:
-                pipeline_tag = np.random.choice(["text-to-image", "image-to-image", "image-segmentation"])
-            elif "whisper" in model_name or "speech" in model_name:
-                pipeline_tag = np.random.choice(["automatic-speech-recognition", "text-to-speech"])
-            elif "coder" in model_name or "code" in model_name:
-                pipeline_tag = "text-generation"
-            elif "bert" in model_name or "roberta" in model_name:
-                pipeline_tag = np.random.choice(["fill-mask", "text-classification", "token-classification"])
-            elif "vision" in model_name:
-                pipeline_tag = np.random.choice(["image-classification", "image-to-text", "visual-question-answering"])
-            else:
-                pipeline_tag = "text-generation"  # Most common
-            
-            # Generate realistic tags
-            tags = [pipeline_tag]
-            
-            if "text-generation" in pipeline_tag:
-                tags.extend(["language-model", "text", "gpt", "llm"])
-                if "instruct" in model_name:
-                    tags.append("instruction-following")
-                if "chat" in model_name:
-                    tags.append("chat")
-            elif "speech" in pipeline_tag:
-                tags.extend(["audio", "speech", "voice"])
-            elif "image" in pipeline_tag:
-                tags.extend(["vision", "image", "diffusion"])
-            
-            # Add language tags
-            if np.random.random() < 0.8:  # 80% chance for English
-                tags.append("en")
-            if np.random.random() < 0.3:  # 30% chance for multilingual
-                tags.append("multilingual")
-            
-            # Generate downloads and likes (weighted by org position for variety)
-            # Earlier orgs get more downloads to make the visualization interesting
-            popularity_factor = (len(orgs) - org_idx) / len(orgs)  # 1.0 to 0.0
-            base_downloads = 10000 * (10 ** (2 * popularity_factor))
-            downloads = int(base_downloads * np.random.uniform(0.3, 3.0))
-            likes = int(downloads * np.random.uniform(0.01, 0.1))  # 1-10% like ratio
-            
-            # Generate model size (in bytes for params)
-            # Model size should correlate somewhat with the size in the name
-            size_indicator = 1
-            for s in ["70b", "13b", "7b", "3b", "2b", "1b", "large", "huge", "xl", "xxl"]:
-                if s in model_name.lower():
-                    size_indicator = float(s.replace("b", "")) if s[0].isdigit() else 3
-                    break
-            
-            # Size in bytes
-            params = int(np.random.uniform(0.5, 2.0) * size_indicator * 1e9)
-            
-            # Create model entry
-            model = {
-                "id": model_id,
-                "author": org,
-                "downloads": downloads,
-                "likes": likes,
-                "pipeline_tag": pipeline_tag,
-                "tags": tags,
-                "params": params
-            }
-            
-            sample_data.append(model)
-    
-    # Convert sample data to DataFrame and append to original
-    sample_df = pd.DataFrame(sample_data)
-    return pd.concat([df, sample_df], ignore_index=True)
+# --- Gradio UI and Controllers ---
+with gr.Blocks(title="HuggingFace Model Explorer") as demo:
+    models_data_state = gr.State(pd.DataFrame())
+    loading_complete_state = gr.State(False) # To control button interactivity
 
-# Create Gradio interface
-with gr.Blocks() as demo:
-    models_data = gr.State()  # To store loaded data
-    
     with gr.Row():
-        gr.Markdown("""
-            # HuggingFace Models TreeMap Visualization
-            
-            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.
-            
-            The treemap visualizes models grouped by organization, with the size of each box representing the selected metric (downloads or likes).
-        """)
-    
+        gr.Markdown("# HuggingFace Models TreeMap Visualization")
     with gr.Row():
-        with gr.Column(scale=1):
-            count_by_dropdown = gr.Dropdown(
-                label="Metric",
-                choices=["downloads", "likes"],
-                value="downloads",
-                info="Select the metric to determine box sizes"
-            )
-            
-            filter_choice_radio = gr.Radio(
-                label="Filter Type",
-                choices=["None", "Tag Filter", "Pipeline Filter"],
-                value="None",
-                info="Choose how to filter the models"
-            )
-            
-            tag_filter_dropdown = gr.Dropdown(
-                label="Select Tag",
-                choices=list(TAG_FILTER_FUNCS.keys()),
-                value=None,
-                visible=False,
-                info="Filter models by domain/category"
-            )
-            
-            pipeline_filter_dropdown = gr.Dropdown(
-                label="Select Pipeline Tag",
-                choices=PIPELINE_TAGS,
-                value=None,
-                visible=False,
-                info="Filter models by specific pipeline"
-            )
+        with gr.Column(scale=1): # Controls column
+            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)
+            size_filter_dropdown = gr.Dropdown(label="Model Size Filter", choices=["None"] + list(MODEL_SIZE_RANGES.keys()), value="None")
+            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") # Common large orgs
             
-            size_filter_dropdown = gr.Dropdown(
-                label="Model Size Filter",
-                choices=["None"] + list(MODEL_SIZE_RANGES.keys()),
-                value="None",
-                info="Filter models by their size (using params column)"
-            )
+            generate_plot_button = gr.Button(value="Generate Plot", variant="primary", interactive=False) # Starts disabled
+            refresh_data_button = gr.Button(value="Refresh Data from Hugging Face", variant="secondary")
 
-            top_k_slider = gr.Slider(
-                label="Number of Top Organizations",
-                minimum=5,
-                maximum=50,
-                value=25,
-                step=5,
-                info="Number of top organizations to include"
-            )
-            
-            skip_orgs_textbox = gr.Textbox(
-                label="Organizations to Skip (comma-separated)",
-                placeholder="e.g., openai, meta, huggingface",
-                info="Enter names of organizations to exclude from the visualization"
+        with gr.Column(scale=3): # Plot and info column
+            plot_output = gr.Plot()
+            status_message_md = gr.Markdown("Initializing...") # For general status updates
+            data_info_md = gr.Markdown("") # For detailed data stats
+
+    # Enable generate button only after data is loaded
+    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)
+
+    # Show/hide tag/pipeline filters based on radio choice
+    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])
+
+
+    def ui_load_data_controller(force_refresh_ui_trigger=False, progress=gr.Progress(track_tqdm=True)): # Gradio progress tracker
+        print(f"ui_load_data_controller called with force_refresh_ui_trigger={force_refresh_ui_trigger}")
+        status_msg_ui = "Loading data..."
+        data_info_text = ""
+        current_df = pd.DataFrame()
+        load_success_flag = False
+        data_as_of_date_display = "N/A"
+
+        try:
+            # Pass gr.Progress.tqdm to load_models_data if it's a Gradio call
+            current_df, load_success_flag, status_msg_from_load = load_models_data(
+                force_refresh=force_refresh_ui_trigger, tqdm_cls=progress.tqdm if progress else tqdm
             )
 
-            generate_plot_button = gr.Button("Generate Plot", variant="primary")
+            if load_success_flag:
+                if force_refresh_ui_trigger: # Data was just fetched by Gradio
+                    data_as_of_date_display = pd.Timestamp.now(tz='UTC').strftime('%B %d, %Y, %H:%M:%S %Z')
+                # If loaded from pre-processed parquet, check for its timestamp column
+                elif '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])
+                    if timestamp_from_parquet.tzinfo is None: # If no timezone, assume UTC from preprocessor
+                        timestamp_from_parquet = timestamp_from_parquet.tz_localize('UTC')
+                    data_as_of_date_display = timestamp_from_parquet.strftime('%B %d, %Y, %H:%M:%S %Z')
+                else: # Pre-processed data but no timestamp column or it's NaT
+                    data_as_of_date_display = "Pre-processed (date unavailable)"
 
-        with gr.Column(scale=3):
-            plot_output = gr.Plot()
-            stats_output = gr.Markdown("*Generate a plot to see statistics*")
+                # Build data info string
+                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)
 
-    def generate_plot_on_click(count_by, filter_choice, tag_filter, pipeline_filter, size_filter, top_k, skip_orgs_text, data_df):
-        print(f"Generating plot with: Metric={count_by}, Filter={filter_choice}, Tag={tag_filter}, Pipeline={pipeline_filter}, Size={size_filter}, Top K={top_k}")
-        
-        if data_df is None or len(data_df) == 0:
-            return None, "Error: No data available. Please try again."
-        
-        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
-        
-        # Process skip organizations list
-        skip_orgs = []
-        if skip_orgs_text and skip_orgs_text.strip():
-            skip_orgs = [org.strip() for org in skip_orgs_text.split(',') if org.strip()]
-            print(f"Skipping organizations: {skip_orgs}")
-        
-        # Process data for treemap
-        treemap_data = make_treemap_data(
-            df=data_df,
-            count_by=count_by,
-            top_k=top_k,
-            tag_filter=selected_tag_filter,
-            pipeline_filter=selected_pipeline_filter,
-            size_filter=selected_size_filter,
-            skip_orgs=skip_orgs
-        )
-        
-        # Create plot
-        fig = create_treemap(
-            treemap_data=treemap_data,
-            count_by=count_by,
-            title=f"HuggingFace Models - {count_by.capitalize()} by Organization"
-        )
+                data_info_text = (f"### Data Information\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}")
+                
+                if not current_df.empty and 'has_robot' in current_df.columns:
+                    robot_true_count = current_df['has_robot'].sum()
+                    data_info_text += f"\n- **Models flagged 'has_robot'**: {robot_true_count}"
+                    if 0 < robot_true_count <= 10:
+                        sample_robot_ids = current_df[current_df['has_robot']]['id'].head(5).tolist()
+                        data_info_text += f"\n  - Sample 'has_robot' model IDs: `{', '.join(sample_robot_ids)}`"
+                elif not current_df.empty:
+                    data_info_text += "\n- **Models flagged 'has_robot'**: 'has_robot' column not found."
+
+                status_msg_ui = "Data loaded successfully. Ready to generate plot."
+            else: # load_success_flag is False
+                data_info_text = f"### Data Load Failed\n- {status_msg_from_load}"
+                status_msg_ui = status_msg_from_load # Pass error message from load_models_data
         
-        # Generate statistics
-        if treemap_data.empty:
-            stats_md = "No data matches the selected filters."
-        else:
-            total_models = len(treemap_data)
-            total_value = treemap_data[count_by].sum()
-            top_5_orgs = treemap_data.groupby("organization")[count_by].sum().sort_values(ascending=False).head(5)
-            
-            # Format the statistics using clean markdown
-            stats_md = f"""
-## Statistics
-- **Total models shown**: {total_models:,}
-- **Total {count_by}**: {int(total_value):,}
+        except Exception as e:
+            status_msg_ui = f"An unexpected error occurred in ui_load_data_controller: {str(e)}"
+            data_info_text = f"### Critical Error\n- {status_msg_ui}"
+            print(f"Critical error in ui_load_data_controller: {e}")
+            load_success_flag = False # Ensure this is false on error
 
-## Top Organizations by {count_by.capitalize()}
+        return current_df, load_success_flag, data_info_text, status_msg_ui
 
-| Organization | {count_by.capitalize()} | % of Total |
-|--------------|--------:|--------:|"""
-            
-            # Add each organization as a row in the table
-            for org, value in top_5_orgs.items():
-                percentage = (value / total_value) * 100
-                stats_md += f"\n| {org} | {int(value):,} | {percentage:.2f}% |"
-            
-            # Add note about skipped organizations if any
-            if skip_orgs:
-                stats_md += f"\n\n*Note: {len(skip_orgs)} organization(s) excluded: {', '.join(skip_orgs)}*"
+    def ui_generate_plot_controller(metric_choice, filter_type, tag_choice, pipeline_choice, 
+                                   size_choice, k_orgs, skip_orgs_input, df_current_models):
+        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 load or refresh data first."
+            gr.Warning(error_msg) # Display a Gradio warning
+            return empty_fig, error_msg
+
+        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 # Handle "None" string
+        orgs_to_skip = [org.strip() for org in skip_orgs_input.split(',') if org.strip()] if skip_orgs_input else []
+
+        # if 'has_robot' in df_current_models.columns:
+        #     robot_count_before_treemap = df_current_models['has_robot'].sum()
+            # print(f"DIAGNOSTIC (ui_generate_plot_controller): df_current_models entering make_treemap_data has {robot_count_before_treemap} 'has_robot' models.")
+
+        treemap_df = make_treemap_data(df_current_models, metric_choice, k_orgs, tag_to_use, pipeline_to_use, size_to_use, orgs_to_skip)
         
-        return fig, stats_md
-
-    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
+        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)
+
+        if treemap_df.empty:
+            plot_stats_md = "No data matches the selected filters. Try adjusting your filters."
+        else:
+            total_items_in_plot = len(treemap_df['id'].unique()) # Count unique models in plot
+            total_value_in_plot = treemap_df[metric_choice].sum() # Sum of metric in plot
+            plot_stats_md = (f"## Plot Statistics\n- **Models shown**: {total_items_in_plot:,}\n- **Total {metric_choice}**: {int(total_value_in_plot):,}")
+        
+        return plotly_fig, plot_stats_md
+
+    # --- Event Handlers ---
+    # Initial data load on app start
     demo.load(
-        fn=load_models_csv,
+        fn=lambda progress=gr.Progress(track_tqdm=True): ui_load_data_controller(force_refresh_ui_trigger=False, progress=progress),
+        inputs=[], # No inputs for initial load
+        outputs=[models_data_state, loading_complete_state, data_info_md, status_message_md]
+    )
+
+    # Refresh data button
+    refresh_data_button.click(
+        fn=lambda progress=gr.Progress(track_tqdm=True): ui_load_data_controller(force_refresh_ui_trigger=True, progress=progress),
         inputs=[], 
-        outputs=[models_data]
+        outputs=[models_data_state, loading_complete_state, data_info_md, status_message_md]
     )
 
-    # Button click event to generate plot
+    # Generate plot button
     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,
-            top_k_slider,
-            skip_orgs_textbox,
-            models_data
-        ],
-        outputs=[plot_output, stats_output]
+        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],
+        outputs=[plot_output, status_message_md] # Update plot and status message
     )
 
-
 if __name__ == "__main__":
-    demo.launch()
+    if not os.path.exists(PROCESSED_PARQUET_FILE_PATH):
+        print(f"WARNING: Pre-processed data file '{PROCESSED_PARQUET_FILE_PATH}' not found.")
+        print("It is highly recommended to run the preprocessing script (preprocess.py) first.")
+    else:
+        print(f"Found pre-processed data file: '{PROCESSED_PARQUET_FILE_PATH}'.")
+    demo.launch()
+# --- END OF FILE app.py ---

models.csv → models_processed.parquet

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:72304e5cf712934720d0fb6e422ccee0d6d89331e9fa9c8a899dc8589d408654
-size 318247426
+oid sha256:998afad6c0c4c64f9e98efd8609d1cbab1dd2ac281b9c2e023878ad436c2fbde
+size 96033487

preprocess.py

@@ -0,0 +1,371 @@
+# --- START OF FILE preprocess.py ---
+
+import pandas as pd
+import numpy as np
+import json
+import ast
+from tqdm.auto import tqdm
+import time
+import os
+import duckdb
+import re # Import re for the manual regex check in debug
+
+# --- Constants ---
+PROCESSED_PARQUET_FILE_PATH = "models_processed.parquet"
+HF_PARQUET_URL = 'https://huggingface.co/datasets/cfahlgren1/hub-stats/resolve/main/models.parquet'
+
+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'))
+}
+
+# --- Debugging Constant ---
+# <<<<<<< SET THE MODEL ID YOU WANT TO DEBUG HERE >>>>>>>
+MODEL_ID_TO_DEBUG = "openvla/openvla-7b" 
+# Example: MODEL_ID_TO_DEBUG = "openai-community/gpt2" 
+# If you don't have a specific ID, the debug block will just report it's not found.
+
+# --- Utility Functions (extract_model_file_size_gb, extract_org_from_id, process_tags_for_series, get_file_size_category - unchanged from previous correct version) ---
+def extract_model_file_size_gb(safetensors_data):
+    try:
+        if pd.isna(safetensors_data): return 0.0
+        data_to_parse = safetensors_data
+        if isinstance(safetensors_data, str):
+            try:
+                if (safetensors_data.startswith('{') and safetensors_data.endswith('}')) or \
+                   (safetensors_data.startswith('[') and safetensors_data.endswith(']')):
+                    data_to_parse = ast.literal_eval(safetensors_data)
+                else: data_to_parse = json.loads(safetensors_data)
+            except Exception: return 0.0
+        if isinstance(data_to_parse, dict) and 'total' in data_to_parse:
+            total_bytes_val = data_to_parse['total']
+            try:
+                size_bytes = float(total_bytes_val)
+                return size_bytes / (1024 * 1024 * 1024) 
+            except (ValueError, TypeError): return 0.0
+        return 0.0
+    except Exception: return 0.0
+
+def extract_org_from_id(model_id):
+    if pd.isna(model_id): return "unaffiliated"
+    model_id_str = str(model_id)
+    return model_id_str.split("/")[0] if "/" in model_id_str else "unaffiliated"
+
+def process_tags_for_series(series_of_tags_values):
+    processed_tags_accumulator = []
+
+    for i, tags_value_from_series in enumerate(tqdm(series_of_tags_values, desc="Standardizing Tags", leave=False, unit="row")):
+        temp_processed_list_for_row = []
+        current_value_for_error_msg = str(tags_value_from_series)[:200] # Truncate for long error messages
+
+        try:
+            # Order of checks is important!
+            # 1. Handle explicit Python lists first
+            if isinstance(tags_value_from_series, list):
+                current_tags_in_list = []
+                for idx_tag, tag_item in enumerate(tags_value_from_series):
+                    try:
+                        # Ensure item is not NaN before string conversion if it might be a float NaN in a list
+                        if pd.isna(tag_item): continue 
+                        str_tag = str(tag_item)
+                        stripped_tag = str_tag.strip()
+                        if stripped_tag:
+                            current_tags_in_list.append(stripped_tag)
+                    except Exception as e_inner_list_proc:
+                        print(f"ERROR processing item '{tag_item}' (type: {type(tag_item)}) within a list for row {i}. Error: {e_inner_list_proc}. Original list: {current_value_for_error_msg}")
+                temp_processed_list_for_row = current_tags_in_list
+
+            # 2. Handle NumPy arrays
+            elif isinstance(tags_value_from_series, np.ndarray):
+                # Convert to list, then process elements, handling potential NaNs within the array
+                current_tags_in_list = []
+                for idx_tag, tag_item in enumerate(tags_value_from_series.tolist()): # .tolist() is crucial
+                    try:
+                        if pd.isna(tag_item): continue # Check for NaN after converting to Python type
+                        str_tag = str(tag_item)
+                        stripped_tag = str_tag.strip()
+                        if stripped_tag:
+                            current_tags_in_list.append(stripped_tag)
+                    except Exception as e_inner_array_proc:
+                        print(f"ERROR processing item '{tag_item}' (type: {type(tag_item)}) within a NumPy array for row {i}. Error: {e_inner_array_proc}. Original array: {current_value_for_error_msg}")
+                temp_processed_list_for_row = current_tags_in_list
+            
+            # 3. Handle simple None or pd.NA after lists and arrays (which might contain pd.NA elements handled above)
+            elif tags_value_from_series is None or pd.isna(tags_value_from_series): # Now pd.isna is safe for scalars
+                temp_processed_list_for_row = []
+
+            # 4. Handle strings (could be JSON-like, list-like, or comma-separated)
+            elif isinstance(tags_value_from_series, str):
+                processed_str_tags = []
+                # Attempt ast.literal_eval for strings that look like lists/tuples
+                if (tags_value_from_series.startswith('[') and tags_value_from_series.endswith(']')) or \
+                   (tags_value_from_series.startswith('(') and tags_value_from_series.endswith(')')):
+                    try:
+                        evaluated_tags = ast.literal_eval(tags_value_from_series)
+                        if isinstance(evaluated_tags, (list, tuple)): # Check if eval result is a list/tuple
+                            # Recursively process this evaluated list/tuple, as its elements could be complex
+                            # For simplicity here, assume elements are simple strings after eval
+                            current_eval_list = []
+                            for tag_item in evaluated_tags:
+                                if pd.isna(tag_item): continue
+                                str_tag = str(tag_item).strip()
+                                if str_tag: current_eval_list.append(str_tag)
+                            processed_str_tags = current_eval_list
+                    except (ValueError, SyntaxError):
+                        pass # If ast.literal_eval fails, let it fall to JSON or comma split
+
+                # If ast.literal_eval didn't populate, try JSON
+                if not processed_str_tags:
+                    try:
+                        json_tags = json.loads(tags_value_from_series)
+                        if isinstance(json_tags, list):
+                            # Similar to above, assume elements are simple strings after JSON parsing
+                            current_json_list = []
+                            for tag_item in json_tags:
+                                if pd.isna(tag_item): continue
+                                str_tag = str(tag_item).strip()
+                                if str_tag: current_json_list.append(str_tag)
+                            processed_str_tags = current_json_list
+                    except json.JSONDecodeError:
+                        # If not a valid JSON list, fall back to comma splitting as the final string strategy
+                        processed_str_tags = [tag.strip() for tag in tags_value_from_series.split(',') if tag.strip()]
+                    except Exception as e_json_other:
+                        print(f"ERROR during JSON processing for string '{current_value_for_error_msg}' for row {i}. Error: {e_json_other}")
+                        processed_str_tags = [tag.strip() for tag in tags_value_from_series.split(',') if tag.strip()] # Fallback
+
+                temp_processed_list_for_row = processed_str_tags
+            
+            # 5. Fallback for other scalar types (e.g., int, float that are not NaN)
+            else:
+                # This path is for non-list, non-ndarray, non-None/NaN, non-string types.
+                # Or for NaNs that slipped through if they are not None or pd.NA (e.g. float('nan'))
+                if pd.isna(tags_value_from_series): # Catch any remaining NaNs like float('nan')
+                     temp_processed_list_for_row = []
+                else:
+                    str_val = str(tags_value_from_series).strip()
+                    temp_processed_list_for_row = [str_val] if str_val else []
+            
+            processed_tags_accumulator.append(temp_processed_list_for_row)
+
+        except Exception as e_outer_tag_proc:
+            print(f"CRITICAL UNHANDLED ERROR processing row {i}: value '{current_value_for_error_msg}' (type: {type(tags_value_from_series)}). Error: {e_outer_tag_proc}. Appending [].")
+            processed_tags_accumulator.append([])
+            
+    return processed_tags_accumulator
+
+def get_file_size_category(file_size_gb_val):
+    try:
+        numeric_file_size_gb = float(file_size_gb_val)
+        if pd.isna(numeric_file_size_gb): numeric_file_size_gb = 0.0
+    except (ValueError, TypeError): numeric_file_size_gb = 0.0
+    if 0 <= numeric_file_size_gb < 1: return "Small (<1GB)"
+    elif 1 <= numeric_file_size_gb < 5: return "Medium (1-5GB)"
+    elif 5 <= numeric_file_size_gb < 20: return "Large (5-20GB)"
+    elif 20 <= numeric_file_size_gb < 50: return "X-Large (20-50GB)"
+    elif numeric_file_size_gb >= 50: return "XX-Large (>50GB)"
+    else: return "Small (<1GB)"
+
+
+def main_preprocessor():
+    print(f"Starting pre-processing script. Output: '{PROCESSED_PARQUET_FILE_PATH}'.")
+    overall_start_time = time.time()
+
+    print(f"Fetching fresh data from Hugging Face: {HF_PARQUET_URL}")
+    try:
+        fetch_start_time = time.time()
+        query = f"SELECT * FROM read_parquet('{HF_PARQUET_URL}')"
+        df_raw = duckdb.sql(query).df()
+        data_download_timestamp = pd.Timestamp.now(tz='UTC')
+        
+        if df_raw is None or df_raw.empty: raise ValueError("Fetched data is empty or None.")
+        if 'id' not in df_raw.columns: raise ValueError("Fetched data must contain 'id' column.")
+        
+        print(f"Fetched data in {time.time() - fetch_start_time:.2f}s. Rows: {len(df_raw)}. Downloaded at: {data_download_timestamp.strftime('%Y-%m-%d %H:%M:%S %Z')}")
+    except Exception as e_fetch:
+        print(f"ERROR: Could not fetch data from Hugging Face: {e_fetch}.")
+        return
+
+    df = pd.DataFrame()
+    print("Processing raw data...")
+    proc_start = time.time()
+
+    expected_cols_setup = {
+        'id': str, 'downloads': float, 'downloadsAllTime': float, 'likes': float,
+        'pipeline_tag': str, 'tags': object, 'safetensors': object
+    }
+    for col_name, target_dtype in expected_cols_setup.items():
+        if col_name in df_raw.columns:
+            df[col_name] = df_raw[col_name]
+            if target_dtype == float: df[col_name] = pd.to_numeric(df[col_name], errors='coerce').fillna(0.0)
+            elif target_dtype == str: df[col_name] = df[col_name].astype(str).fillna('')
+        else: 
+            if col_name in ['downloads', 'downloadsAllTime', 'likes']: df[col_name] = 0.0
+            elif col_name == 'pipeline_tag': df[col_name] = ''
+            elif col_name == 'tags': df[col_name] = pd.Series([[] for _ in range(len(df_raw))]) # Initialize with empty lists
+            elif col_name == 'safetensors': df[col_name] = None # Initialize with None
+            elif col_name == 'id': print("CRITICAL ERROR: 'id' column missing."); return
+    
+    output_filesize_col_name = 'params' 
+    if output_filesize_col_name in df_raw.columns and pd.api.types.is_numeric_dtype(df_raw[output_filesize_col_name]):
+        print(f"Using pre-existing '{output_filesize_col_name}' column as file size in GB.")
+        df[output_filesize_col_name] = pd.to_numeric(df_raw[output_filesize_col_name], errors='coerce').fillna(0.0)
+    elif 'safetensors' in df.columns:
+        print(f"Calculating '{output_filesize_col_name}' (file size in GB) from 'safetensors' data...")
+        df[output_filesize_col_name] = df['safetensors'].apply(extract_model_file_size_gb)
+        df[output_filesize_col_name] = pd.to_numeric(df[output_filesize_col_name], errors='coerce').fillna(0.0)
+    else:
+        print(f"Cannot determine file size. Setting '{output_filesize_col_name}' to 0.0.")
+        df[output_filesize_col_name] = 0.0
+    
+    df['data_download_timestamp'] = data_download_timestamp
+    print(f"Added 'data_download_timestamp' column.")
+
+    print("Categorizing models by file size...")
+    df['size_category'] = df[output_filesize_col_name].apply(get_file_size_category)
+
+    print("Standardizing 'tags' column...")
+    df['tags'] = process_tags_for_series(df['tags']) # This now uses tqdm internally
+
+    # --- START DEBUGGING BLOCK ---
+    # This block will execute before the main tag processing loop
+    if MODEL_ID_TO_DEBUG and MODEL_ID_TO_DEBUG in df['id'].values: # Check if ID exists
+        print(f"\n--- Pre-Loop Debugging for Model ID: {MODEL_ID_TO_DEBUG} ---")
+        
+        # 1. Check the 'tags' column content after process_tags_for_series
+        model_specific_tags_list = df.loc[df['id'] == MODEL_ID_TO_DEBUG, 'tags'].iloc[0]
+        print(f"1. Tags from df['tags'] (after process_tags_for_series): {model_specific_tags_list}")
+        print(f"   Type of tags: {type(model_specific_tags_list)}")
+        if isinstance(model_specific_tags_list, list):
+            for i, tag_item in enumerate(model_specific_tags_list):
+                print(f"   Tag item {i}: '{tag_item}' (type: {type(tag_item)}, len: {len(str(tag_item))})")
+                # Detailed check for 'robotics' specifically
+                if 'robotics' in str(tag_item).lower():
+                    print(f"     DEBUG: Found 'robotics' substring in '{tag_item}'")
+                    print(f"       - str(tag_item).lower().strip(): '{str(tag_item).lower().strip()}'")
+                    print(f"       - Is it exactly 'robotics'?: {str(tag_item).lower().strip() == 'robotics'}")
+                    print(f"       - Ordinals: {[ord(c) for c in str(tag_item)]}")
+
+        # 2. Simulate temp_tags_joined for this specific model
+        if isinstance(model_specific_tags_list, list):
+            simulated_temp_tags_joined = '~~~'.join(str(t).lower().strip() for t in model_specific_tags_list if pd.notna(t) and str(t).strip())
+        else:
+            simulated_temp_tags_joined = ''
+        print(f"2. Simulated 'temp_tags_joined' for this model: '{simulated_temp_tags_joined}'")
+
+        # 3. Simulate 'has_robot' check for this model
+        robot_keywords = ['robot', 'robotics']
+        robot_pattern = '|'.join(robot_keywords)
+        manual_robot_check = bool(re.search(robot_pattern, simulated_temp_tags_joined, flags=re.IGNORECASE))
+        print(f"3. Manual regex check for 'has_robot' ('{robot_pattern}' in '{simulated_temp_tags_joined}'): {manual_robot_check}")
+        print(f"--- End Pre-Loop Debugging for Model ID: {MODEL_ID_TO_DEBUG} ---\n")
+    elif MODEL_ID_TO_DEBUG:
+        print(f"DEBUG: Model ID '{MODEL_ID_TO_DEBUG}' not found in DataFrame for pre-loop debugging.")
+    # --- END DEBUGGING BLOCK ---
+
+
+    print("Vectorized creation of cached tag columns...")
+    tag_time = time.time()
+    # This is the original temp_tags_joined creation:
+    df['temp_tags_joined'] = df['tags'].apply(
+        lambda tl: '~~~'.join(str(t).lower().strip() for t in tl if pd.notna(t) and str(t).strip()) if isinstance(tl, list) else ''
+    )
+    
+    tag_map = {
+        'has_audio': ['audio'], 'has_speech': ['speech'], 'has_music': ['music'],
+        'has_robot': ['robot', 'robotics','openvla','vla'], 
+        'has_bio': ['bio'], 'has_med': ['medic', 'medical'], 
+        'has_series': ['series', 'time-series', 'timeseries'], 
+        'has_video': ['video'], 'has_image': ['image', 'vision'], 
+        'has_text': ['text', 'nlp', 'llm'] 
+    }
+    for col, kws in tag_map.items():
+        pattern = '|'.join(kws) 
+        df[col] = df['temp_tags_joined'].str.contains(pattern, na=False, case=False, regex=True)
+        
+    df['has_science'] = (
+        df['temp_tags_joined'].str.contains('science', na=False, case=False, regex=True) &
+        ~df['temp_tags_joined'].str.contains('bigscience', na=False, case=False, regex=True) 
+    )
+    del df['temp_tags_joined'] # Clean up temporary column
+    df['is_audio_speech'] = (df['has_audio'] | df['has_speech'] |
+                            df['pipeline_tag'].str.contains('audio|speech', case=False, na=False, regex=True))
+    df['is_biomed'] = df['has_bio'] | df['has_med']
+    print(f"Vectorized tag columns created in {time.time() - tag_time:.2f}s.")
+
+    # --- POST-LOOP DIAGNOSTIC for has_robot & a specific model ---
+    if 'has_robot' in df.columns:
+        print("\n--- 'has_robot' Diagnostics (Preprocessor - Post-Loop) ---")
+        print(df['has_robot'].value_counts(dropna=False))
+        
+        if MODEL_ID_TO_DEBUG and MODEL_ID_TO_DEBUG in df['id'].values:
+            model_has_robot_val = df.loc[df['id'] == MODEL_ID_TO_DEBUG, 'has_robot'].iloc[0]
+            print(f"Value of 'has_robot' for model '{MODEL_ID_TO_DEBUG}': {model_has_robot_val}")
+            if model_has_robot_val:
+                 print(f"  Original tags for '{MODEL_ID_TO_DEBUG}': {df.loc[df['id'] == MODEL_ID_TO_DEBUG, 'tags'].iloc[0]}")
+
+        if df['has_robot'].any():
+            print("Sample models flagged as 'has_robot':")
+            print(df[df['has_robot']][['id', 'tags', 'has_robot']].head(5))
+        else:
+            print("No models were flagged as 'has_robot' after processing.")
+        print("--------------------------------------------------------\n")
+    # --- END POST-LOOP DIAGNOSTIC ---
+
+
+    print("Adding organization column...")
+    df['organization'] = df['id'].apply(extract_org_from_id)
+
+    # Drop safetensors if params was calculated from it, and params didn't pre-exist as numeric
+    if 'safetensors' in df.columns and \
+       not (output_filesize_col_name in df_raw.columns and pd.api.types.is_numeric_dtype(df_raw[output_filesize_col_name])):
+        df = df.drop(columns=['safetensors'], errors='ignore')
+    
+    final_expected_cols = [
+        'id', 'downloads', 'downloadsAllTime', 'likes', 'pipeline_tag', 'tags',
+        'params', 'size_category', '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'
+    ]
+    # Ensure all final columns exist, adding defaults if necessary
+    for col in final_expected_cols:
+        if col not in df.columns:
+            print(f"Warning: Final expected column '{col}' is missing! Defaulting appropriately.")
+            if col == 'params': df[col] = 0.0
+            elif col == 'size_category': df[col] = "Small (<1GB)" # Default size category
+            elif 'has_' in col or 'is_' in col : df[col] = False # Default boolean flags to False
+            elif col == 'data_download_timestamp': df[col] = pd.NaT # Default timestamp to NaT
+
+    print(f"Data processing completed in {time.time() - proc_start:.2f}s.")
+    try:
+        print(f"Saving processed data to: {PROCESSED_PARQUET_FILE_PATH}")
+        df_to_save = df[final_expected_cols].copy() # Ensure only expected columns are saved
+        df_to_save.to_parquet(PROCESSED_PARQUET_FILE_PATH, index=False, engine='pyarrow')
+        print(f"Successfully saved processed data.")
+    except Exception as e_save:
+        print(f"ERROR: Could not save processed data: {e_save}")
+        return
+
+    total_elapsed_script = time.time() - overall_start_time
+    print(f"Pre-processing finished. Total time: {total_elapsed_script:.2f}s. Final Parquet shape: {df_to_save.shape}")
+
+if __name__ == "__main__":
+    if os.path.exists(PROCESSED_PARQUET_FILE_PATH):
+        print(f"Deleting existing '{PROCESSED_PARQUET_FILE_PATH}' to ensure fresh processing...")
+        try: os.remove(PROCESSED_PARQUET_FILE_PATH)
+        except OSError as e: print(f"Error deleting file: {e}. Please delete manually and rerun."); exit()
+    
+    main_preprocessor()
+    
+    if os.path.exists(PROCESSED_PARQUET_FILE_PATH):
+        print(f"\nTo verify, load parquet and check 'has_robot' and its 'tags':")
+        print(f"import pandas as pd; df_chk = pd.read_parquet('{PROCESSED_PARQUET_FILE_PATH}')")
+        print(f"print(df_chk['has_robot'].value_counts())")
+        if MODEL_ID_TO_DEBUG:
+            print(f"print(df_chk[df_chk['id'] == '{MODEL_ID_TO_DEBUG}'][['id', 'tags', 'has_robot']])")
+        else:
+            print(f"print(df_chk[df_chk['has_robot']][['id', 'tags', 'has_robot']].head())")