MLflow Integration (dependent on Surrogate Modelling)

[tim-huntley]

Dependency

#83

What we want to achieve

We would like to integrate xplainable into the MLflow workflow. This involves the following workflows:

• If the model is an xplainable model, the model is logged to MLflow and Xplainable Cloud (with the same experiment_id)
• If the model is NOT an xplainable model, a surrogate model is created and logged to MLflow and Xplainable Cloud

Key considerations

This should be as uninvasive as possible. We don't want people to have to change their workflow drastically, rather just import xplainable and automate the explainer creation and logging.

Flexibility in approach

Please run with your own creative freedom here. We would love to discuss different approaches, as long as the disruption to user workflow is minimal.

Vision 1

This idea involves implementing a manual logging step. Behind the scenes it will train an xplainable surrogate model and log to both MLflow and Xplainable Cloud.

import xplainable as xp
import mlflow

# <- Build any model here

with mlflow.start_run():
    signature = infer_signature(X_train, model.predict(X_train))

    # <- other logging here

    mlflow.log_model(
        model=model,
        signature=signature,
        input_example=X_train,
        registered_model_name="model-with-xplainable"
    )

    # This should log to mlflow and to xplainable cloud (if an api key is active)
    xp.mlflow.log_explanation(model.predict, X)

Vision 2

This is similar to vision 1, but without the need to manually log explainers. It's not clear how this would be achieved, but it is an idea worth fleshing out.

import xplainable as xp
import mlflow

# <- Build any model here

xp.mlflow.auto_logging = True

with mlflow.start_run():
    signature = infer_signature(X_train, model.predict(X_train))

    # <- other logging here

    mlflow.log_model(
        model=model,
        signature=signature,
        input_example=X_train,
        registered_model_name="model-with-xplainable"
    ) # < -- This should auto-log to mlflow and to xplainable cloud (if an api key is active)