Alberto_Umana
Databricks Employee
Databricks Employee

To address your concerns about logging LightGBM feature importance and modifying the AutoML-generated LightGBM model to use the mlflow.lightgbm flavor, you'll need to make some changes to the AutoML notebook. Here's an approach to achieve what you're looking for:

Logging Feature Importance

LightGBM's feature importance is not logged by default in MLflow's autologging. To log this information, you can manually add it to the MLflow run after the model is trained. Here's how you can do this:

 

import mlflow

import lightgbm as lgb

 

# Assuming 'model' is your trained LightGBM model

feature_importance = model.feature_importance(importance_type='gain')

feature_names = model.feature_name()

 

# Log feature importance

for feature, importance in zip(feature_names, feature_importance):

    mlflow.log_metric(f"feature_importance_{feature}", importance)

 

Changing Model Flavor to LightGBM

To log the model with the LightGBM flavor instead of scikit-learn, you need to modify the model logging process in the AutoML notebook. Here's how you can do it:

  1. Find the part of the notebook where the model is being logged to MLflow.
  2. Replace the existing logging code with mlflow.lightgbm.log_model(). Here's an example:

 

import mlflow.lightgbm

 

# Assuming 'model' is your trained LightGBM model

mlflow.lightgbm.log_model(model, "model", registered_model_name="your_model_name")

 

  1. If you're using a pipeline that includes preprocessing steps, you'll need to log the LightGBM model separately from the pipeline. You can do this by extracting the LightGBM model from the pipeline:

 

lightgbm_model = pipeline.named_steps['lightgbm']

mlflow.lightgbm.log_model(lightgbm_model, "lightgbm_model")

 

  1. You may also need to log the preprocessing steps separately if they're required for making predictions.

 

Loading the Model

After making these changes, you should be able to load the model using:

 

loaded_model = mlflow.lightgbm.load_model("runs:/your_run_id/lightgbm_model")

This loaded model will have access to LightGBM-specific attributes like feature_importances_.

View solution in original post