Spaces:
Sleeping
Sleeping
File size: 4,153 Bytes
bbd5f9c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 | #!/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()
|