Hello!
Im trying to save my model with mlflow in databricks, it is a xgboost model, when I save it using code it saves with a sklearn flavor and not saves other parameters, also I'm using kedro with kedro-mlflow plugin.
def log_metrics_and_model(model, metrics: pd.DataFrame, params: Dict, X_train: pd.DataFrame):
model_name = params.get("model_name", "model")
mlflow.log_params(params)
model_params = model.get_params()
mlflow.log_params(model_params)
for metric_name, metric_values in metrics.items():
mlflow.log_metric(f"{model_name}_{metric_name}_train", metric_values.iloc[0])
mlflow.log_metric(f"{model_name}_{metric_name}_test", metric_values.iloc[1])
input_example = X_train.iloc[:5]
signature = infer_signature(X_train, model.predict(X_train.iloc[:5]))
mlflow.xgboost.log_model(
model,
model_name,
signature=signature,
input_example=input_example
)
In databricks it saves a simple sklearn model, but locally saves xgboost with all properties correctly, any idea?
