The ML Pipeline & Explainable AI

The Architectural Primer: Why MLflow and SHAP?

Machine learning in production is vastly different from a Jupyter Notebook. It requires strict version control and, in the financial sector, absolute transparency.

1. MLflow (The Registry): Hardcoding a model file (like a model.pkl) directly into an API is a massive anti-pattern. Models degrade over time and need continuous retraining. MLflow acts as a centralized version control system for your models. It tracks metrics, logs parameters, and provides a dynamic registry. When the API boots up, it doesn’t look for a static file; it asks MLflow for the “latest” production model

2. SHAP (Explainable AI): XGBoost is incredibly accurate but notoriously opaque. SHAP (SHapley Additive exPlanations) uses cooperative game theory to solve this. It calculates exactly how much each specific feature (or “player”) contributed to the final prediction (the “game score”). This turns a vague “99% Fraud Risk” into a precise statement: “Fraud because behavioral vector V14 was abnormally high and the Amount was an outlier.”

Automated Training and Registration

To ensure the system is reproducible, the model isn’t trained manually. Instead, a dedicated script (train_in_docker.py) runs as a Kubernetes Job. It trains the XGBoost pipeline and registers it directly to the MLflow tracking server

# src/ml/train_in_docker.py (Simplified)

# 1. Define the Machine Learning Pipeline
pipeline = Pipeline([
    ('scaler', StandardScaler()),
    ('classifier', XGBClassifier(
        n_estimators=100,
        learning_rate=0.1,
        max_depth=5,
        eval_metric='logloss',
        n_jobs=-1
    ))
])

# 2. Train & Register with MLflow
with mlflow.start_run():
    logger.info("Training Model...")
    pipeline.fit(X_train, y_train)
    
    logger.info("Registering Model to MLflow...")
    # This automatically versions the model and saves it to the backend store
    mlflow.sklearn.log_model(
        pipeline, 
        "model", 
        registered_model_name="FraudDetectionSOTA"
    )
    
    accuracy = pipeline.score(X_test, y_test)
    mlflow.log_metric("accuracy", accuracy)

The Explainability Endpoint

Once the FastAPI service pulls the latest pipeline from MLflow, it also initializes a shap.Explainer. When the Streamlit dashboard or the AI Agent needs to know why a transaction was flagged, it hits the /explain endpoint.

Instead of running a full retraining cycle, the explainer uses the pre-computed mathematical weights of the XGBoost model to instantly isolate the top driving factors of that specific transaction

# src/api/app.py - Inside explain_transaction()

@app.post("/explain")
def explain_transaction(transaction: Transaction):
    explainer = ml_resources["explainer"]
    pipeline = ml_resources["model_pipeline"]
    feature_names = ml_resources["feature_names"]
        
    # Format and scale the incoming raw transaction data
    df = pd.DataFrame([transaction.dict(exclude={'transaction_id'})])
    scaled_data = pipeline.named_steps['scaler'].transform(df)
    
    # Generate SHAP values (Game Theory feature attribution)
    shap_results = explainer(scaled_data)
    vals = shap_results.values[0, :, 1] if len(shap_results.values.shape) == 3 else shap_results.values[0]
    
    # Map the SHAP values back to their human-readable feature names
    importance_map = {k: float(v) for k, v in zip(feature_names, vals)}
    
    # Sort to find the top 5 biggest contributing factors
    sorted_factors = sorted(importance_map.items(), key=lambda item: abs(item[1]), reverse=True)
    
    return {
        "transaction_id": transaction.transaction_id,
        "top_contributing_features": dict(sorted_factors[:5])
    }

By isolating the explainer logic into its own endpoint, the high-speed /predict endpoint remains unburdened, while the UI and the AI Agent can still request deep mathematical audits on demand