Model

churn_prediction.model.training

XGBoost Training Module.

This module provides functionality to train an XGBoost model for churn prediction, evaluate its performance, and generate a comprehensive model card in JSON format.

train_xgb_model(train_data_path, val_data_path, model_filepath, model_card_path, report_path)

Train, evaluate, and export the churn prediction XGBoost model.

Orchestrates the complete model training stage by loading processed datasets, constructing XGBoost matrices, calculating class-imbalance weighting, training the booster, evaluating validation performance, and exporting the model, model card, and training report.

Parameters:

Name	Type	Description	Default
train_data_path	str or Path	Path to the processed training Parquet dataset.	required
val_data_path	str or Path	Path to the processed validation Parquet dataset.	required
model_filepath	str or Path	Destination path for the trained XGBoost model.	required
model_card_path	str or Path	Destination path for the machine-readable model card.	required
report_path	str or Path	Destination path for the comprehensive training report.	required

Raises:

Type	Description
FileNotFoundError	If either processed dataset cannot be found.
RuntimeError	If model training, model serialization, model-card generation, or training-report generation fails.
ValueError	If the processed datasets are incompatible with the configured training schema or model.

Source code in src/churn_prediction/model/training.py

def train_xgb_model(
    train_data_path: str | Path,
    val_data_path: str | Path,
    model_filepath: str | Path,
    model_card_path: str | Path,
    report_path: str | Path,
) -> None:
    """Train, evaluate, and export the churn prediction XGBoost model.

    Orchestrates the complete model training stage by loading processed
    datasets, constructing XGBoost matrices, calculating class-imbalance
    weighting, training the booster, evaluating validation performance,
    and exporting the model, model card, and training report.

    Parameters
    ----------
    train_data_path : str or pathlib.Path
        Path to the processed training Parquet dataset.
    val_data_path : str or pathlib.Path
        Path to the processed validation Parquet dataset.
    model_filepath : str or pathlib.Path
        Destination path for the trained XGBoost model.
    model_card_path : str or pathlib.Path
        Destination path for the machine-readable model card.
    report_path : str or pathlib.Path
        Destination path for the comprehensive training report.

    Raises
    ------
    FileNotFoundError
        If either processed dataset cannot be found.
    RuntimeError
        If model training, model serialization, model-card generation, or
        training-report generation fails.
    ValueError
        If the processed datasets are incompatible with the configured
        training schema or model.
    """
    logger = get_logger()
    settings = get_settings()

    # Resolve paths to absolute paths
    train_data_path = Path(train_data_path).expanduser().resolve()
    val_data_path = Path(val_data_path).expanduser().resolve()
    model_filepath = Path(model_filepath).expanduser().resolve()
    model_card_path = Path(model_card_path).expanduser().resolve()
    report_path = Path(report_path).expanduser().resolve()

    logger.debug(
        "Loading configuration from settings. Target: %s",
        settings.PARAMS.schema_config.target_column,
    )

    # Extract parameters from the validated pipeline configuration
    target_col = settings.PARAMS.schema_config.target_column
    feature_cols = settings.PARAMS.schema_config.feature_columns
    verbose_eval = settings.PARAMS.train.verbose_eval
    num_boost_round = settings.PARAMS.train.num_boost_round
    early_stopping_rounds = settings.PARAMS.train.early_stopping_rounds
    classification_threshold = settings.PARAMS.evaluate.classification_threshold
    xgb_params = dict(settings.PARAMS.train.xgb_params)

    # 0. Load the datasets
    logger.debug("Reading Parquet files: train=%s, val=%s", train_data_path, val_data_path)
    train_df = pd.read_parquet(train_data_path)
    val_df = pd.read_parquet(val_data_path)
    train_rows, train_cols = train_df.shape
    val_rows, val_cols = val_df.shape

    # 1. Build DMatrix objects
    dtrain, dval, scale_pos_weight = _build_dmatrices(
        train_df,
        val_df,
        target_col,
        feature_cols,
    )

    # 2. Train the model
    booster, evals_result, params = _train_model(
        dtrain,
        dval,
        xgb_params=xgb_params,
        scale_pos_weight=scale_pos_weight,
        n_estimators=num_boost_round,
        early_stopping_rounds=early_stopping_rounds,
        verbose_eval=verbose_eval,
    )

    # 3. Evaluate the model
    metrics, clf_report, cm, feature_importances = _evaluate_model(
        booster,
        dval,
        scale_pos_weight=scale_pos_weight,
        classification_threshold=classification_threshold,
    )

    # Unpack confusion matrix for clarity
    (tn, fp), (fn, tp) = cm

    sorted_importances = dict(sorted(feature_importances.items(), key=lambda x: x[1], reverse=True))

    # 4. Save the trained model
    try:
        logger.debug("Saving booster model to %s", model_filepath)
        booster.save_model(str(model_filepath))
        logger.info("Trained model saved successfully to: %s", model_filepath)
    except Exception as err:
        logger.error("Failed to save the trained model: %s", err)
        raise RuntimeError(
            f"Failed to save the trained model to '{model_filepath}': {err}"
        ) from err

    if not model_filepath.exists() or model_filepath.stat().st_size == 0:
        logger.error("Model file is empty or does not exist: %s", model_filepath)
        raise RuntimeError(f"Model file is empty or does not exist: '{model_filepath}'")

    # 5. Generate and save model card
    _export_model_card(
        output_path=model_card_path,
        params=params,
        metrics=metrics,
        clf_report=clf_report,
        feature_importances=sorted_importances,
        feature_columns=feature_cols,
        target_column=target_col,
        train_shape=train_df.shape,
        val_shape=val_df.shape,
        train_data_path=train_data_path,
        val_data_path=val_data_path,
    )

    # 6. Export comprehensive training report
    try:
        logger.debug("Exporting training report to %s", report_path)

        best_iter = getattr(booster, "best_iteration", num_boost_round - 1)
        best_score_raw = getattr(booster, "best_score", float("nan"))
        best_score = float(best_score_raw) if not np.isnan(best_score_raw) else None

        report_data: dict[str, Any] = {
            "metadata": {
                "generated_at": datetime.now(UTC).isoformat(),
                "target_column": target_col,
                "feature_count": len(feature_cols),
                "feature_names": feature_cols,
                "dataset_dimensions": {
                    "train_rows": train_rows,
                    "train_cols": train_cols,
                    "val_rows": val_rows,
                    "val_cols": val_cols,
                },
                "class_imbalance": {
                    "scale_pos_weight": scale_pos_weight,
                },
            },
            "training_summary": {
                "num_boost_round": num_boost_round,
                "early_stopping_rounds": early_stopping_rounds,
                "best_iteration": best_iter,
                "best_score": best_score,
                "hyperparameters": params,
                "evals_result": evals_result,
            },
            "performance_evaluation": {
                **metrics,
                "confusion_matrix_optimal_2x2": cm,
                "confusion_matrix_unpacked": {
                    "true_negatives": tn,
                    "false_positives": fp,
                    "false_negatives": fn,
                    "true_positives": tp,
                },
                "classification_report": clf_report,
                "feature_importance_gain": sorted_importances,
            },
        }

        report_json_bytes = orjson.dumps(
            report_data,
            option=orjson.OPT_INDENT_2 | orjson.OPT_SERIALIZE_NUMPY | orjson.OPT_NON_STR_KEYS,
        )
        report_path.write_bytes(report_json_bytes)
        logger.info("Training report exported successfully to: %s", report_path)
    except Exception as err:
        logger.error("Failed to export training report: %s", err)
        raise RuntimeError(f"Failed to export training report to '{report_path}': {err}") from err

unscale_probabilities(y_prob_scaled, scale_pos_weight)

Convert class-weighted probabilities to calibrated likelihoods.

Reverses the probability distortion introduced by scale_pos_weight during model training. When no effective class weighting is present, the original probabilities are returned unchanged.

Parameters:

Name	Type	Description	Default
y_prob_scaled	ndarray	Probabilities predicted by the class-weighted model.	required
scale_pos_weight	float	Positive-class weighting factor used during model training.	required

Returns:

Type	Description
ndarray	Probabilities adjusted to represent the estimated likelihood under the original class distribution.

Source code in src/churn_prediction/model/training.py

def unscale_probabilities(
    y_prob_scaled: np.ndarray,
    scale_pos_weight: float,
) -> np.ndarray:
    """Convert class-weighted probabilities to calibrated likelihoods.

    Reverses the probability distortion introduced by ``scale_pos_weight``
    during model training. When no effective class weighting is present, the
    original probabilities are returned unchanged.

    Parameters
    ----------
    y_prob_scaled : numpy.ndarray
        Probabilities predicted by the class-weighted model.
    scale_pos_weight : float
        Positive-class weighting factor used during model training.

    Returns
    -------
    numpy.ndarray
        Probabilities adjusted to represent the estimated likelihood under
        the original class distribution.
    """
    logger = get_logger()

    if scale_pos_weight == 1.0 or scale_pos_weight <= 0:
        logger.debug("scale_pos_weight is %.4f; skipping probability unscaling", scale_pos_weight)
        return y_prob_scaled

    logger.debug("Unscaling probabilities using scale_pos_weight: %.4f", scale_pos_weight)
    p: np.ndarray = np.clip(y_prob_scaled, EPSILON, 1.0 - EPSILON)
    calibrated_p = p / (p + scale_pos_weight * (1.0 - p))

    logger.debug(
        "Probability range - Scaled: [%.4f, %.4f], Unscaled: [%.4f, %.4f]",
        y_prob_scaled.min(),
        y_prob_scaled.max(),
        calibrated_p.min(),
        calibrated_p.max(),
    )

    return calibrated_p

churn_prediction.model.inference

XGBoost Inference Module.

This module provides the run_inference function, which takes a preprocessed DataFrame, a trained XGBoost booster model, and a path to a model card JSON file. It generates churn probabilities, calibrates them based on class imbalance, and applies a decision threshold to produce final predictions.

run_inference(booster, raw_filepath, input_filepath, output_filepath, model_card_path)

Generate churn predictions for a processed inference dataset.

Loads model metadata from the model card, generates raw XGBoost probabilities, reverses class-imbalance probability scaling, applies the configured decision threshold, and persists the resulting predictions alongside the original customer records.

Parameters:

Name	Type	Description	Default
booster	Booster	Trained XGBoost model used to generate predictions.	required
raw_filepath	str or Path	Path to the original raw input dataset.	required
input_filepath	str or Path	Path to the processed inference dataset supplied to the model.	required
output_filepath	str or Path	Destination path for the generated prediction dataset.	required
model_card_path	str or Path	Path to the JSON model card containing the class-imbalance weight and classification threshold.	required

Raises:

Type	Description
FileNotFoundError	If the model card or required input files do not exist.
ValueError	If required model metadata is missing or the inference data cannot be converted into an XGBoost input matrix.
RuntimeError	If prediction generation or output persistence fails.

Source code in src/churn_prediction/model/inference.py

def run_inference(
    booster: xgb.Booster,
    raw_filepath: str | Path,
    input_filepath: str | Path,
    output_filepath: str | Path,
    model_card_path: str | Path,
) -> None:
    """Generate churn predictions for a processed inference dataset.

    Loads model metadata from the model card, generates raw XGBoost
    probabilities, reverses class-imbalance probability scaling, applies
    the configured decision threshold, and persists the resulting
    predictions alongside the original customer records.

    Parameters
    ----------
    booster : xgboost.Booster
        Trained XGBoost model used to generate predictions.
    raw_filepath : str or pathlib.Path
        Path to the original raw input dataset.
    input_filepath : str or pathlib.Path
        Path to the processed inference dataset supplied to the model.
    output_filepath : str or pathlib.Path
        Destination path for the generated prediction dataset.
    model_card_path : str or pathlib.Path
        Path to the JSON model card containing the class-imbalance weight
        and classification threshold.

    Raises
    ------
    FileNotFoundError
        If the model card or required input files do not exist.
    ValueError
        If required model metadata is missing or the inference data cannot
        be converted into an XGBoost input matrix.
    RuntimeError
        If prediction generation or output persistence fails.
    """
    logger = get_logger()
    settings = get_settings()

    # Resolve paths to absolute paths
    raw_filepath = Path(raw_filepath).expanduser().resolve()
    input_filepath = Path(input_filepath).expanduser().resolve()
    output_filepath = Path(output_filepath).expanduser().resolve()
    model_card_path = Path(model_card_path).expanduser().resolve()

    if not model_card_path.is_file():
        logger.error("Model card file missing at: %s", model_card_path)
        raise FileNotFoundError(f"Model card not found at path: {model_card_path}")

    # 1. Load model card to retrieve threshold and scale_pos_weight
    logger.info("Loading model metadata from card: %s", model_card_path.name)
    with model_card_path.open("r", encoding="utf-8") as f:
        model_card: dict[str, Any] = json.load(f)

    try:
        scale_pos_weight = float(model_card["model_details"]["hyperparameters"]["scale_pos_weight"])
        decision_thresholds = model_card["metrics"]["decision_thresholds"]
        threshold = float(
            decision_thresholds.get(
                "optimal_threshold", decision_thresholds.get("default_threshold", 0.5)
            )
        )
        logger.debug(
            "Extracted model params - scale_pos_weight: %s, threshold: %s",
            scale_pos_weight,
            threshold,
        )
    except (KeyError, TypeError, ValueError) as err:
        logger.error("Model card schema mismatch or missing values: %s", err)
        raise ValueError(f"Failed to extract required parameters from model card: {err}") from err

    # 2. Load the dataset
    logger.info("Reading Parquet file: %s", input_filepath)
    df = pd.read_parquet(input_filepath)

    # 3. Create DMatrix for inference
    try:
        feature_cols = settings.PARAMS.schema_config.feature_columns
        logger.info("Creating DMatrix for %d features and %d rows", len(feature_cols), len(df))
        dtest = xgb.DMatrix(df[feature_cols], feature_names=feature_cols)
        logger.info("DMatrix created successfully. Shapes: dtest=%s", dtest.num_row())
    except Exception as err:
        logger.error("Failed to create DMatrix: %s", err)
        raise ValueError(f"Failed to create DMatrix: {err}") from err

    # 4. Run inference using the booster model
    logger.info("Executing model prediction via XGBoost booster")
    raw_probs: np.ndarray = booster.predict(dtest)
    logger.debug("Raw probabilities generated. Sample (first 3): %s", raw_probs[:3])

    # 5. Calibrate probabilities using scale_pos_weight
    logger.info("Calibrating raw probabilities using scale_pos_weight=%s", scale_pos_weight)
    calibrated_probs = unscale_probabilities(raw_probs, scale_pos_weight)
    logger.debug("Calibrated probabilities sample (first 3): %s", calibrated_probs[:3])

    # 6. Make predictions based on the calibrated probabilities and threshold
    logger.info("Applying decision threshold: %s", threshold)
    predictions = (calibrated_probs >= threshold).astype(int)

    # Log distribution of predictions for observability
    pos_count = np.sum(predictions)
    logger.info(
        "Inference complete. Predictions: %d churn, %d non-churn",
        pos_count,
        len(predictions) - pos_count,
    )

    # 7. Prepare results DataFrame with churn probabilities and predictions
    logger.debug("Merging predictions back into results DataFrame")
    results_df = pd.read_parquet(raw_filepath)
    results_df["churn_probability"] = np.round(calibrated_probs, 6)
    results_df["churn_prediction"] = predictions

    # 8. Save the results to the specified output file
    logger.info("Saving predictions to CSV file: %s", output_filepath)
    results_df.to_csv(output_filepath, index=False)

churn_prediction.model.explainability

Module for exporting SHAP explainability artifacts for dynamic dashboards.

This module computes SHAP values for XGBoost models and exports structured, high-performance data artifacts (Parquet, JSON) designed for direct consumption by interactive frontend dashboards (e.g., Dash, React, FastAPI backends).

SHAPArtifactExporter(batch_id, df, booster, output_dir, *, max_samples=10000, random_state=42)

Compute and export SHAP artifacts for churn model interpretation.

This class transforms model predictions and SHAP explanations into dashboard-ready Parquet and JSON artifacts. It supports deterministic sampling for large datasets and produces both global feature importance metrics and customer-level churn/retention drivers.

Initialize the SHAP artifact exporter.

Parameters:

Name	Type	Description	Default
batch_id	str	Identifier for the current batch of explanations.	required
df	DataFrame	Feature matrix containing the customer records to explain. The DataFrame index is treated as the customer identifier.	required
booster	Booster	Trained XGBoost booster for which SHAP explanations will be generated.	required
output_dir	Path	Directory in which explainability artifacts will be written.	required
max_samples	int	Maximum number of customer records to evaluate with SHAP. Larger datasets are sampled deterministically.	10000
random_state	int	Random seed used when sampling the evaluation dataset.	42

Raises:

Type	Description
Exception	If the SHAP tree explainer cannot be initialized for the supplied XGBoost model.

Source code in src/churn_prediction/model/explainability.py

def __init__(
    self,
    batch_id: str,
    df: pd.DataFrame,
    booster: xgb.Booster,
    output_dir: Path,
    *,
    max_samples: int = 10000,
    random_state: int = 42,
) -> None:
    """Initialize the SHAP artifact exporter.

    Parameters
    ----------
    batch_id : str
        Identifier for the current batch of explanations.
    df : pandas.DataFrame
        Feature matrix containing the customer records to explain. The
        DataFrame index is treated as the customer identifier.
    booster : xgboost.Booster
        Trained XGBoost booster for which SHAP explanations will be generated.
    output_dir : pathlib.Path
        Directory in which explainability artifacts will be written.
    max_samples : int, default=10000
        Maximum number of customer records to evaluate with SHAP. Larger
        datasets are sampled deterministically.
    random_state : int, default=42
        Random seed used when sampling the evaluation dataset.

    Raises
    ------
    Exception
        If the SHAP tree explainer cannot be initialized for the supplied
        XGBoost model.
    """
    # Initialize logging and settings
    self.logger = get_logger()
    self.settings = get_settings()

    self.logger.debug(
        "Initializing SHAPArtifactExporter: rows=%d, features=%d, output_dir=%s",
        len(df),
        df.shape[1],
        output_dir,
    )
    try:
        self.batch_id = batch_id
        self.output_dir = output_dir
        self.max_samples = max_samples
        self.random_state = random_state
        self.features_df = df.copy()
        self.customer_ids = df.index.values
        self.booster = booster
        self.explainer = shap.TreeExplainer(self.booster)
    except Exception:
        self.logger.exception("Failed to initialize SHAPArtifactExporter")
        raise

    self.logger.info(
        "SHAPArtifactExporter initialized: batch_id=%s",
        self.batch_id,
    )

export_artifacts(top_k_drivers=5)

Generate and persist dashboard-ready SHAP artifacts.

The export contains the raw SHAP matrix, the aligned input feature matrix, customer-level risk profiles, and global feature-importance metadata.

Parameters:

Name	Type	Description	Default
top_k_drivers	int	Maximum number of positive churn drivers and negative retention factors retained for each customer.	5

Returns:

Type	Description
Path	Directory containing all generated explainability artifacts.

Raises:

Type	Description
OSError	If an artifact cannot be written to disk.
Exception	If SHAP computation or artifact generation fails.

Source code in src/churn_prediction/model/explainability.py

def export_artifacts(self, top_k_drivers: int = 5) -> Path:
    """Generate and persist dashboard-ready SHAP artifacts.

    The export contains the raw SHAP matrix, the aligned input feature
    matrix, customer-level risk profiles, and global feature-importance
    metadata.

    Parameters
    ----------
    top_k_drivers : int, default=5
        Maximum number of positive churn drivers and negative retention
        factors retained for each customer.

    Returns
    -------
    pathlib.Path
        Directory containing all generated explainability artifacts.

    Raises
    ------
    OSError
        If an artifact cannot be written to disk.
    Exception
        If SHAP computation or artifact generation fails.
    """
    batch_output_dir = self.output_dir
    batch_output_dir.mkdir(parents=True, exist_ok=True)

    self.logger.info("Starting SHAP artifact generation run: destination=%s", batch_output_dir)

    try:
        eval_df = self._sample_dataset()
        explanation = self._compute_shap_explanation(eval_df)

        # 1. Export SHAP Value Matrix
        shap_cols: list[str] = [f"shap_{col}" for col in explanation.feature_names]
        shap_values = np.asarray(explanation.values)
        shap_df = pd.DataFrame(
            data=shap_values,
            columns=pd.Index(shap_cols),
            index=pd.Index(self.sampled_customer_ids, name="customer_id"),
        )
        shap_parquet_path = batch_output_dir / "raw_shap_values.parquet"
        shap_df.to_parquet(
            shap_parquet_path,
            engine="pyarrow",
            compression="snappy",
        )
        self.logger.debug("Exported SHAP matrix: %s", shap_parquet_path)

        # 2. Export Feature Matrix (Aligned with SHAP Values)
        eval_df.index = self.sampled_customer_ids
        eval_df.index.name = "customer_id"
        features_parquet_path = batch_output_dir / "input_features.parquet"
        eval_df.to_parquet(features_parquet_path, engine="pyarrow", compression="snappy")
        self.logger.debug("Exported feature matrix: %s", features_parquet_path)

        # 3. Export Customer Risk Profiles (Top Drivers)
        risk_profiles_df = self._build_customer_risk_profiles(
            explanation=explanation,
            customer_ids=self.sampled_customer_ids,
            top_k=top_k_drivers,
        )
        risk_parquet_path = batch_output_dir / "customer_risk_profiles.parquet"
        risk_profiles_df.to_parquet(risk_parquet_path, engine="pyarrow", compression="snappy")
        self.logger.debug("Exported customer risk profiles: %s", risk_parquet_path)

        # 4. Export Metadata (Global Metrics)
        self._export_metadata(
            explanation=explanation,
            output_path=batch_output_dir / "metadata.json",
            total_samples=len(eval_df),
        )

    except Exception:
        self.logger.exception("Export run failed for output directory: %s", batch_output_dir)
        raise

    self.logger.info("Successfully exported all SHAP artifacts to: %s", batch_output_dir)
    return batch_output_dir

generate_shap_artifacts(batch_id, df, booster, shap_report_dir)

Generate SHAP explainability artifacts for a trained churn model.

This is the public entry point for the explainability workflow. It initializes :class:SHAPArtifactExporter, computes the requested SHAP outputs, and persists the resulting artifacts to the configured report directory.

Parameters:

Name	Type	Description	Default
batch_id	str	Identifier for the current batch of explanations.	required
df	DataFrame	Model feature matrix to explain.	required
booster	Booster	Trained XGBoost model.	required
shap_report_dir	Path	Directory in which SHAP artifacts will be generated.	required

Returns:

Type	Description
Path	Directory containing the generated SHAP artifacts.

Raises:

Type	Description
Exception	If SHAP initialization, computation, or artifact export fails.

Source code in src/churn_prediction/model/explainability.py

def generate_shap_artifacts(
    batch_id: str,
    df: pd.DataFrame,
    booster: xgb.Booster,
    shap_report_dir: Path,
) -> Path:
    """Generate SHAP explainability artifacts for a trained churn model.

    This is the public entry point for the explainability workflow. It
    initializes :class:`SHAPArtifactExporter`, computes the requested SHAP
    outputs, and persists the resulting artifacts to the configured report
    directory.

    Parameters
    ----------
    batch_id : str
        Identifier for the current batch of explanations.
    df : pandas.DataFrame
        Model feature matrix to explain.
    booster : xgboost.Booster
        Trained XGBoost model.
    shap_report_dir : pathlib.Path
        Directory in which SHAP artifacts will be generated.

    Returns
    -------
    pathlib.Path
        Directory containing the generated SHAP artifacts.

    Raises
    ------
    Exception
        If SHAP initialization, computation, or artifact export fails.
    """
    logger = get_logger()
    try:
        exporter = SHAPArtifactExporter(
            batch_id=batch_id,
            df=df,
            booster=booster,
            output_dir=shap_report_dir,
        )
        return exporter.export_artifacts()
    except Exception:
        logger.exception("Failed to complete SHAP artifact generation workflow.")
        raise