def __init__(self,
booster: Booster,
tree_index: int,
x_data: (pd.DataFrame, np.ndarray),
y_data: (pd.Series, np.ndarray),
feature_names: List[str] = None,
target_name: str = None,
class_names: (List[str], Mapping[int, str]) = None):
utils.check_tree_index(tree_index, booster.num_trees())
self.booster = booster
self.tree_index = tree_index
self.tree_nodes, self.children_left, self.children_right = self._get_nodes_info()
self.thresholds = None # lazy evaluation
self.features = None # lazy evaluation
self.node_to_samples = None
super().__init__(booster, x_data, y_data, feature_names, target_name, class_names)
def _create_checkpoint(model: Booster, epoch: int, filename: str,
frequency: int):
if epoch % frequency > 0:
return
with tune.checkpoint_dir(step=epoch) as checkpoint_dir:
model.save_model(os.path.join(checkpoint_dir, filename))
def calculate_auc(model: Booster, X_eval, y_eval) -> float:
y_pred = model.predict(X_eval, num_iteration=model.best_iteration)
return roc_auc_score(y_eval, y_pred)