train.py

import os
import warnings
import numpy as np
import pandas as pd
import joblib
import shap
import mlflow
import mlflow.sklearn
from google.cloud import storage
from sklearn.ensemble import RandomForestClassifier
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import accuracy_score, classification_report, confusion_matrix
from fairlearn.metrics import MetricFrame, demographic_parity_difference
from fairlearn.metrics import equalized_odds_difference

warnings.filterwarnings("ignore")

# ── Config ──────────────────────────────────────────────────────────────────
GCS_BUCKET   = "mlops_f12026"
GCS_DATA_PATH = "oppe2/data.csv"
MLFLOW_URI = "http://127.0.0.1:8101/"
EXPERIMENT   = "heart_disease"
LOCAL_DATA   = "/tmp/data.csv"

FEATURES = ["age","gender","cp","trestbps","chol","fbs","restecg",
            "thalach","exang","oldpeak","slope","ca","thal"]
TARGET = "target"

# ── 1. Download data from GCS ────────────────────────────────────────────────
def download_data():
    print("[INFO] Downloading data from GCS …")
    client = storage.Client()
    bucket = client.bucket(GCS_BUCKET)
    blob   = bucket.blob(GCS_DATA_PATH)
    blob.download_to_filename(LOCAL_DATA)
    print(f"[INFO] Saved to {LOCAL_DATA}")

# ── 2. Preprocess ────────────────────────────────────────────────────────────
def preprocess(path):
    df = pd.read_csv(path)
    df["gender"] = (df["gender"] == "male").astype(int)
    df["target"] = (df["target"] == "yes").astype(int)
    df = df.dropna(subset=FEATURES + [TARGET])
    X = df[FEATURES]
    y = df[TARGET]
    return X, y, df

# ── 3. Train ─────────────────────────────────────────────────────────────────
def train(X_train, y_train):
    model = LogisticRegression(max_iter=1000, random_state=42)
    model.fit(X_train, y_train)
    return model

# ── 4. SHAP Explainability on False Negatives ────────────────────────────────
def explain_false_negatives(model, scaler, X_test, y_test, feature_names):
    print("\n[EXPLAINABILITY] Analysing false negatives with SHAP …")
    X_scaled = scaler.transform(X_test)
    preds = model.predict(X_scaled)
    fn_mask = (preds == 0) & (y_test.values == 1)
    fn_X = X_scaled[fn_mask]
    if fn_X.shape[0] == 0:
        print("[INFO] No false negatives found in test set.")
        return

    explainer   = shap.LinearExplainer(model, X_scaled, feature_perturbation="interventional")
    shap_values = explainer.shap_values(fn_X)

    print(f"\n[FALSE NEGATIVES] {fn_X.shape[0]} samples misclassified as healthy:")
    for i in range(min(3, fn_X.shape[0])):
        sv  = shap_values[i] if isinstance(shap_values, np.ndarray) else shap_values[1][i]
        top = np.argsort(np.abs(sv))[::-1][:5]
        print(f"\n  Sample {i+1}:")
        for idx in top:
            direction = "↑ risk" if sv[idx] > 0 else "↓ risk"
            print(f"    {feature_names[idx]:12s}: SHAP={sv[idx]:+.3f}  [{direction}]")
    print("\n[PLAIN ENGLISH] These patients were missed (false negatives).")
    print("  The top contributing features pushed the model toward 'no disease'.")
    print("  E.g. a high thalach (max heart rate) or low cp (chest pain type)")
    print("  may have made the model confident they were healthy — but they weren't.")

# ── 5. Fairness Check ─────────────────────────────────────────────────────────
def fairness_check(model, scaler, X_test, y_test):
    print("\n[FAIRNESS] Running Fairlearn fairness check …")
    age_bins = pd.cut(X_test["age"], bins=[0, 40, 60, 80, 120],
                      labels=["<40", "40-60", "60-80", "80+"])
    X_scaled = scaler.transform(X_test)
    preds = model.predict(X_scaled)

    mf = MetricFrame(
        metrics={"accuracy": accuracy_score},
        y_true=y_test,
        y_pred=preds,
        sensitive_features=age_bins
    )
    print("\n  Accuracy by age group:")
    print(mf.by_group.to_string())

    dp_diff = demographic_parity_difference(y_test, preds, sensitive_features=age_bins)
    eo_diff = equalized_odds_difference(y_test, preds, sensitive_features=age_bins)
    print(f"\n  Demographic Parity Difference : {dp_diff:.4f}  (0=perfectly fair)")
    print(f"  Equalized Odds Difference     : {eo_diff:.4f}  (0=perfectly fair)")

    if abs(dp_diff) > 0.1:
        print("  ⚠  Demographic parity gap > 0.10 — potential age bias detected.")
    else:
        print("  ✓  Demographic parity within acceptable range.")
    return {"dp_diff": dp_diff, "eo_diff": eo_diff}

# ── 6. Data Drift Detection ───────────────────────────────────────────────────
def data_drift_check(X_train, X_test, feature_names):
    print("\n[DATA DRIFT] Comparing training vs synthetic generated data …")
    np.random.seed(42)
    n = 100
    generated = pd.DataFrame({
        "age"      : np.random.randint(45, 75, n),
        "gender"   : np.random.randint(0, 2, n),
        "cp"       : np.random.randint(0, 4, n),
        "trestbps" : np.random.normal(140, 20, n),
        "chol"     : np.random.normal(260, 50, n),
        "fbs"      : np.random.randint(0, 2, n),
        "restecg"  : np.random.randint(0, 3, n),
        "thalach"  : np.random.normal(140, 25, n),
        "exang"    : np.random.randint(0, 2, n),
        "oldpeak"  : np.random.uniform(0, 4, n),
        "slope"    : np.random.randint(0, 3, n),
        "ca"       : np.random.randint(0, 5, n),
        "thal"     : np.random.randint(0, 4, n),
    })

    from scipy.stats import ks_2samp
    print(f"\n  {'Feature':<12} {'Train Mean':>12} {'Gen Mean':>12} {'KS Stat':>10} {'p-value':>10} {'Drift?':>8}")
    print("  " + "-"*66)
    drifted = []
    for feat in feature_names:
        stat, pval = ks_2samp(X_train[feat], generated[feat])
        drift = "YES ⚠" if pval < 0.05 else "no"
        if pval < 0.05:
            drifted.append(feat)
        print(f"  {feat:<12} {X_train[feat].mean():>12.3f} {generated[feat].mean():>12.3f} "
              f"{stat:>10.4f} {pval:>10.4f} {drift:>8}")
    print(f"\n  Drifted features: {drifted if drifted else 'None'}")
    return drifted

# ── 7. Inference on 100 samples ───────────────────────────────────────────────
def run_inference(model, scaler, X_test, y_test):
    print("\n[INFERENCE] Running on 100 random samples …")
    import time
    sample = X_test.sample(min(100, len(X_test)), random_state=42)
    latencies = []
    for _, row in sample.iterrows():
        t0 = time.time()
        model.predict(scaler.transform([row.values]))
        latencies.append((time.time() - t0) * 1000)
    print(f"  Avg latency : {np.mean(latencies):.3f} ms")
    print(f"  P95 latency : {np.percentile(latencies, 95):.3f} ms")
    print(f"  Max latency : {np.max(latencies):.3f} ms")

# ── Main ──────────────────────────────────────────────────────────────────────
if __name__ == "__main__":
    download_data()
    X, y, df = preprocess(LOCAL_DATA)

    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.2, random_state=42, stratify=y)

    scaler = StandardScaler()
    X_train_s = scaler.fit_transform(X_train)

    mlflow.set_tracking_uri(MLFLOW_URI)
    mlflow.set_experiment(EXPERIMENT)

    with mlflow.start_run(run_name="logistic_baseline"):
        model = train(X_train_s, y_train)
        X_test_s = scaler.transform(X_test)
        preds = model.predict(X_test_s)
        acc   = accuracy_score(y_test, preds)
        print(f"\n[METRICS] Accuracy: {acc:.4f}")
        print(classification_report(y_test, preds))

        mlflow.log_param("model_type", "LogisticRegression")
        mlflow.log_metric("accuracy", acc)
        mlflow.sklearn.log_model(model, "model")

        joblib.dump(model,  "model.pkl")
        joblib.dump(scaler, "scaler.pkl")
        print("[INFO] model.pkl and scaler.pkl saved.")

        # Upload to GCS
        client = storage.Client()
        bucket = client.bucket(GCS_BUCKET)
        for fname in ["model.pkl", "scaler.pkl"]:
            bucket.blob(f"oppe2final/{fname}").upload_from_filename(fname)
        print("[INFO] Artifacts uploaded to GCS.")

        # Deliverables
        explain_false_negatives(model, scaler, X_test, y_test, FEATURES)
        fair_metrics = fairness_check(model, scaler, X_test, y_test)
        mlflow.log_metric("dp_diff", fair_metrics["dp_diff"])
        mlflow.log_metric("eo_diff", fair_metrics["eo_diff"])

        data_drift_check(X_train, X_test, FEATURES)
        run_inference(model, scaler, X_test, y_test)

    print("\n[DONE] Training pipeline complete.")