# --- 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())")