| import numpy as np |
| import pandas as pd |
| from sklearn.ensemble import RandomForestClassifier |
| from xgboost import XGBClassifier |
|
|
| class PurchasePropensityModel: |
| def __init__(self, use_xgboost=True, random_state=42): |
| self.use_xgboost = use_xgboost |
| if use_xgboost: |
| self.model = XGBClassifier( |
| n_estimators=100, |
| max_depth=4, |
| learning_rate=0.08, |
| random_state=random_state, |
| eval_metric="logloss" |
| ) |
| else: |
| self.model = RandomForestClassifier( |
| n_estimators=100, |
| max_depth=6, |
| random_state=random_state |
| ) |
|
|
| def fit(self, X: pd.DataFrame, y: pd.Series): |
| self.model.fit(X, y) |
| return self |
|
|
| def predict_proba(self, X: pd.DataFrame) -> np.ndarray: |
| return self.model.predict_proba(X)[:, 1] |
|
|
| def predict(self, X: pd.DataFrame) -> np.ndarray: |
| return self.model.predict(X) |
|
|
