wisesa77/stratum-api

1# backend/pipeline.py2"""3Stratum ML Pipeline — UCI Wine Quality (proxy dataset for SPL Biscuit QA)4---------------------------------------------------------------------5Trains an XGBoost classifier with Stratified 5-Fold CV, saves:6  artifacts/model.joblib7  artifacts/metrics.json8  artifacts/shap_beeswarm.png9  artifacts/shap_waterfall_fail.png10  artifacts/shap_waterfall_pass.png11  artifacts/shap_summary.json12 13NOTE: This MVP uses the UCI Wine Quality dataset as a public proxy dataset14to validate Stratum's ML and explainability pipeline. The architecture is15designed to be retrained with actual SPL biscuit lab data when available.16"""17 18import os19import json20import warnings21import numpy as np22import pandas as pd23import matplotlib24matplotlib.use("Agg")25import matplotlib.pyplot as plt26import joblib27import shap28from xgboost import XGBClassifier29from sklearn.model_selection import StratifiedKFold, cross_validate30from sklearn.metrics import make_scorer, accuracy_score, precision_score, recall_score, f1_score31 32from utils.preprocessing import normalize_columns, FEATURE_COLUMNS33 34warnings.filterwarnings("ignore")35 36# Paths37BASE_DIR    = os.path.dirname(__file__)38DATA_PATH   = os.path.join(BASE_DIR, "data", "winequality-red.csv")39ARTIFACTS   = os.path.join(BASE_DIR, "artifacts")40MODEL_PATH  = os.path.join(ARTIFACTS, "model.joblib")41METRICS_PATH= os.path.join(ARTIFACTS, "metrics.json")42SHAP_BEESWARM   = os.path.join(ARTIFACTS, "shap_beeswarm.png")43SHAP_FAIL       = os.path.join(ARTIFACTS, "shap_waterfall_fail.png")44SHAP_PASS       = os.path.join(ARTIFACTS, "shap_waterfall_pass.png")45SHAP_SUMMARY    = os.path.join(ARTIFACTS, "shap_summary.json")46 47 48def load_data():49    df = pd.read_csv(DATA_PATH, sep=";")50    df = normalize_columns(df)51    df["status_binary"] = (df["quality"] >= 6).astype(int)52    X = df[FEATURE_COLUMNS]53    y = df["status_binary"]54    print(f"Loaded {len(df)} rows | PASS: {y.sum()} | FAIL: {(y == 0).sum()}")55    return X, y56 57 58def train_and_evaluate(X, y):59    model = XGBClassifier(60        n_estimators=300,61        max_depth=5,62        learning_rate=0.05,63        subsample=0.8,64        colsample_bytree=0.8,65        eval_metric="logloss",66        random_state=42,67        use_label_encoder=False,68    )69 70    cv = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)71    scoring = {72        "accuracy":  make_scorer(accuracy_score),73        "precision": make_scorer(precision_score, zero_division=0),74        "recall":    make_scorer(recall_score, zero_division=0),75        "f1":        make_scorer(f1_score, zero_division=0),76    }77    scores = cross_validate(model, X, y, cv=cv, scoring=scoring, return_train_score=False)78 79    metrics = {80        "model":      "XGBoost Classifier",81        "dataset":    "UCI Wine Quality - Red Wine",82        "validation": "Stratified 5-Fold Cross Validation",83        "accuracy":   round(float(scores["test_accuracy"].mean()), 4),84        "precision":  round(float(scores["test_precision"].mean()), 4),85        "recall":     round(float(scores["test_recall"].mean()), 4),86        "f1_score":   round(float(scores["test_f1"].mean()), 4),87    }88    print("CV Metrics:", metrics)89 90    # Train final model on full data91    model.fit(X, y)92    return model, metrics93 94 95def generate_shap(model, X, y):96    explainer = shap.TreeExplainer(model)97    shap_values = explainer.shap_values(X)  # shape (n, features)98 99    # --- Beeswarm plot ---100    plt.figure(figsize=(10, 6))101    shap.summary_plot(shap_values, X, show=False, plot_type="dot")102    plt.title("SHAP Feature Impact (Beeswarm)", fontsize=13, pad=12)103    plt.tight_layout()104    plt.savefig(SHAP_BEESWARM, dpi=150, bbox_inches="tight")105    plt.close()106    print(f"Saved {SHAP_BEESWARM}")107 108    # --- Waterfall — FAIL sample ---109    fail_idx = y[y == 0].index[0]110    _save_waterfall(explainer, X, fail_idx, SHAP_FAIL, "SHAP Waterfall — FAIL Sample")111 112    # --- Waterfall — PASS sample ---113    pass_idx = y[y == 1].index[0]114    _save_waterfall(explainer, X, pass_idx, SHAP_PASS, "SHAP Waterfall — PASS Sample")115 116    # --- SHAP summary JSON (mean |SHAP| per feature, for frontend bar chart) ---117    mean_abs = np.abs(shap_values).mean(axis=0)118    summary = [119        {"feature": feat, "importance": round(float(v), 5)}120        for feat, v in sorted(121            zip(FEATURE_COLUMNS, mean_abs),122            key=lambda x: x[1], reverse=True123        )124    ]125    with open(SHAP_SUMMARY, "w") as f:126        json.dump(summary, f, indent=2)127    print(f"Saved {SHAP_SUMMARY}")128 129    return shap_values130 131 132def _save_waterfall(explainer, X, idx, path, title):133    row = X.iloc[[idx]]134    sv  = explainer(row)135    plt.figure(figsize=(8, 5))136    shap.plots.waterfall(sv[0], show=False)137    plt.title(title, fontsize=11, pad=10)138    plt.tight_layout()139    plt.savefig(path, dpi=150, bbox_inches="tight")140    plt.close()141    print(f"Saved {path}")142 143 144def run():145    os.makedirs(ARTIFACTS, exist_ok=True)146    print("=== Stratum ML Pipeline ===")147    X, y = load_data()148    model, metrics = train_and_evaluate(X, y)149 150    # Save model151    joblib.dump(model, MODEL_PATH)152    print(f"Saved model → {MODEL_PATH}")153 154    # Save metrics155    with open(METRICS_PATH, "w") as f:156        json.dump(metrics, f, indent=2)157    print(f"Saved metrics → {METRICS_PATH}")158 159    # Generate SHAP plots160    print("Generating SHAP visualizations…")161    generate_shap(model, X, y)162 163    print("\n✓ Pipeline complete. All artifacts saved to backend/artifacts/")164    return metrics165 166 167if __name__ == "__main__":168    run()169