Spaces:
Sleeping
Sleeping
| #!/usr/bin/env python3 | |
| """ | |
| Model loader that handles joblib compatibility issues | |
| """ | |
| import os | |
| import joblib | |
| import pickle | |
| import sys | |
| from typing import Optional, Any | |
| # Import our models module | |
| from models import DataPreprocessor | |
| def load_models_safely(models_dir: str = 'models') -> tuple[Optional[Any], Optional[DataPreprocessor]]: | |
| """ | |
| Safely load models with proper class resolution for joblib | |
| Returns: | |
| tuple: (model, preprocessor) or (None, None) if failed | |
| """ | |
| # Set up the module namespace for joblib | |
| current_module = sys.modules[__name__] | |
| current_module.DataPreprocessor = DataPreprocessor | |
| # Also set it in the main module | |
| import __main__ | |
| __main__.DataPreprocessor = DataPreprocessor | |
| # And in the models module | |
| import models | |
| models.DataPreprocessor = DataPreprocessor | |
| try: | |
| # Try to load preprocessor | |
| preprocessor_path = os.path.join(models_dir, 'preprocessor.pkl') | |
| if not os.path.exists(preprocessor_path): | |
| print(f"❌ Preprocessor file not found: {preprocessor_path}") | |
| return None, None | |
| print(f"📥 Loading preprocessor from {preprocessor_path}") | |
| preprocessor = joblib.load(preprocessor_path) | |
| print("✅ Preprocessor loaded successfully") | |
| # Try to load Random Forest model | |
| rf_path = os.path.join(models_dir, 'random_forest_model.pkl') | |
| if not os.path.exists(rf_path): | |
| print(f"❌ Random Forest file not found: {rf_path}") | |
| return None, None | |
| print(f"📥 Loading Random Forest model from {rf_path}") | |
| model = joblib.load(rf_path) | |
| print("✅ Random Forest model loaded successfully") | |
| return model, preprocessor | |
| except Exception as e: | |
| print(f"❌ Error loading models: {e}") | |
| # Try alternative loading method | |
| try: | |
| print("🔄 Trying alternative loading method...") | |
| # Load with explicit class mapping | |
| with open(preprocessor_path, 'rb') as f: | |
| # Create a custom unpickler that knows about our classes | |
| import pickle | |
| class CustomUnpickler(pickle.Unpickler): | |
| def find_class(self, module, name): | |
| if name == 'DataPreprocessor': | |
| return DataPreprocessor | |
| return super().find_class(module, name) | |
| preprocessor = CustomUnpickler(f).load() | |
| print("✅ Preprocessor loaded with custom unpickler") | |
| # Load model normally (should work since it's sklearn) | |
| model = joblib.load(rf_path) | |
| print("✅ Random Forest model loaded successfully") | |
| return model, preprocessor | |
| except Exception as e2: | |
| print(f"❌ Alternative loading also failed: {e2}") | |
| return None, None | |
| def test_model_loading(): | |
| """Test function to verify model loading works""" | |
| print("🧪 Testing model loading...") | |
| model, preprocessor = load_models_safely() | |
| if model is not None and preprocessor is not None: | |
| print("✅ Model loading test passed!") | |
| # Test basic functionality | |
| try: | |
| print("🔍 Testing preprocessor functionality...") | |
| print(f" - Label encoders available: {list(preprocessor.label_encoders.keys()) if hasattr(preprocessor, 'label_encoders') else 'None'}") | |
| print(f" - Scaler available: {preprocessor.scaler is not None if hasattr(preprocessor, 'scaler') else 'Unknown'}") | |
| print("🔍 Testing model functionality...") | |
| print(f" - Model type: {type(model).__name__}") | |
| return True | |
| except Exception as e: | |
| print(f"❌ Functionality test failed: {e}") | |
| return False | |
| else: | |
| print("❌ Model loading test failed!") | |
| return False | |
| if __name__ == "__main__": | |
| test_model_loading() | |