I’m using from databricks.feature_engineering import FeatureEngineering functions for the future store.

I’d like to perform a score batch inference.

After logging a simple RF classifier:

fe = FeatureEngineeringClient()
fe.log_model(
       model=model,
       artifact_path=artifact_path,
       flavor=flavor,
       training_set=training_set,
       registered_model_name=model_name
)

I want to do score batch inference, but I need the probabilities, not just the predicted label.

The score_batch function only retrieves the predicted label. If changed the name of the prediction columns in the classifier, still, the prediction label expected by score_batch has to be double.

prediction_df = fe.score_batch(
  model_uri=uc_modeling.get_lastest_model_uri(), 
  df=batch_input_df)

Thanks!