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