| 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 |
|
|