abersbail's picture
Deploy Customer Churn ML Predictor & Demo Video to Hugging Face
d0bac84 verified
Raw
History Blame Contribute Delete
1.68 kB
import os
import joblib
import numpy as np
from sklearn.ensemble import RandomForestClassifier
class ChurnClassifier:
"""
Churn Classifier wrapper class using RandomForestClassifier.
Includes model training, prediction, evaluation, and saving/loading functionality.
"""
def __init__(self, n_estimators: int = 100, max_depth: int = None, random_state: int = 42):
self.model = RandomForestClassifier(
n_estimators=n_estimators,
max_depth=max_depth,
random_state=random_state,
class_weight='balanced'
)
self.preprocessor = None
def fit(self, X: np.ndarray, y: np.ndarray, preprocessor=None):
self.model.fit(X, y)
self.preprocessor = preprocessor
return self
def predict(self, X: np.ndarray) -> np.ndarray:
return self.model.predict(X)
def predict_proba(self, X: np.ndarray) -> np.ndarray:
return self.model.predict_proba(X)
def save(self, filepath: str):
"""
Saves the complete model bundle (model classifier + preprocessor).
"""
bundle = {
'model': self.model,
'preprocessor': self.preprocessor
}
os.makedirs(os.path.dirname(os.path.abspath(filepath)), exist_ok=True)
joblib.dump(bundle, filepath)
@classmethod
def load(cls, filepath: str):
"""
Loads a saved model bundle and returns a fully restored ChurnClassifier instance.
"""
bundle = joblib.load(filepath)
instance = cls()
instance.model = bundle['model']
instance.preprocessor = bundle['preprocessor']
return instance